build: 1.0.7

cssxsh · Oct 15, 2022 · 7bdd666 · 7bdd666
1 parent 618cb90
commit 7bdd666
Show file tree

Hide file tree

Showing 11 changed files with 152 additions and 42 deletions.
diff --git a/README.md b/README.md
@@ -19,15 +19,33 @@
 * `/nai <word>` 生成一张图片   
   例如 `/nai 连裤袜 双马尾` (只有部分词条会自动翻译)  
   例如 `/nai swimsuit #seed=12346` (设置种子)  
+  例如 `/nai swimsuit #steps=3` (AI迭代次数)  
   例如 `/nai "swimsuit, ahegao"` (如果需要以 `,` 分割词条, 请用 `"` 包裹)
+  可用的配置项有
+  `seed` 种子
+  `steps` 迭代次数
+  `width` 宽度
+  `height` 高度 
+  `samples` 出图数量
+  `scale` 比例
+  `sampler` 采样器 可选值 `k_euler_ancestral`, `k_euler`, `k_lms`, `plms`, `ddim`
+  `strength` 以图出图中对原图的更改程度 可选值 [0.00, 0.99]
+  `noise` 以图出图中的噪声 可选值 [0.00, 0.99]
 * `/nai-login <mail> <password>` 登录账号  
   例如 `/nai-login [email protected] 1919810`
+* `/naifu <word>` 生成一张图片
+  对接 `naifu`, `naifu` 是基于 novelai 官方 web 端的修改版，所以指令用法 和 `nai` 一致
 
 ## 配置
 
-* `config.yml` 配置文件 包括 `proxy`, `doh`, `ipv6` 等配置
+* `config.yml` 配置文件 包括 `proxy`, `doh`, `ipv6`, `naifu_api` 等配置
 * `ban.txt` 屏蔽的词条，可热编辑，保存后一段时间会自动启用
 
+## NaiFu
+
+`naifu` 是基于 novelai 官方 web 端的修改版  
+相关信息可以看这 <https://www.bilibili.com/video/BV14e4y1E74X>
+
 ## TODO
 
 * [ ] 更好的翻译

diff --git a/build.gradle.kts b/build.gradle.kts
@@ -6,7 +6,7 @@ plugins {
 }
 
 group = "xyz.cssxsh.mirai.novelai"
-version = "1.0.6"
+version = "1.0.7"
 
 repositories {
     mavenLocal()

diff --git a/src/main/kotlin/xyz/cssxsh/mirai/novelai/NovelAiHelper.kt b/src/main/kotlin/xyz/cssxsh/mirai/novelai/NovelAiHelper.kt
@@ -11,15 +11,30 @@ public object NovelAiHelper : KotlinPlugin(
     JvmPluginDescription(
         id = "xyz.cssxsh.mirai.plugin.novelai-helper",
         name = "novelai-helper",
-        version = "1.0.6",
+        version = "1.0.7",
     ) {
         author("cssxsh")
     }
 ) {
-    public val client: NovelAiClient = NovelAiClient(config = NovelAiHelperConfig)
+    public val client: NovelAiClient = NovelAiClient(config = NovelAiHelperConfig, default = true)
+
+    public fun translate(word: String): String? {
+        for (translation in NovelAiHelperConfig.database.data) {
+            for ((name, tag) in translation.data) {
+                if (tag.name == word) return name
+            }
+        }
+        for (translation in NovelAiHelperConfig.database.data) {
+            for ((name, tag) in translation.data) {
+                if (tag.name.startsWith(word)) return name
+            }
+        }
+        return null
+    }
 
     override fun onEnable() {
         NovelAiHelperConfig.reload()
+        NovelAiHelperConfig.save()
         NovelAiCommand.register()
         NovelAiLoginCommand.register()
     }

diff --git a/src/main/kotlin/xyz/cssxsh/mirai/novelai/command/NovelAiCommand.kt b/src/main/kotlin/xyz/cssxsh/mirai/novelai/command/NovelAiCommand.kt
@@ -11,7 +11,6 @@ import net.mamoe.mirai.message.data.Image.Key.queryUrl
 import net.mamoe.mirai.message.data.MessageSource.Key.quote
 import net.mamoe.mirai.utils.ExternalResource.Companion.toExternalResource
 import xyz.cssxsh.mirai.novelai.*
-import xyz.cssxsh.mirai.novelai.data.*
 import xyz.cssxsh.novelai.*
 import kotlin.random.*
 
@@ -21,20 +20,6 @@ public object NovelAiCommand : SimpleCommand(
     description = "生成图片"
 ) {
 
-    private fun translate(word: String): String? {
-        for (translation in NovelAiHelperConfig.database.data) {
-            for ((name, tag) in translation.data) {
-                if (tag.name == word) return name
-            }
-        }
-        for (translation in NovelAiHelperConfig.database.data) {
-            for ((name, tag) in translation.data) {
-                if (tag.name.startsWith(word)) return name
-            }
-        }
-        return null
-    }
-
     private val random = Random(seed = System.currentTimeMillis())
 
     @Handler
@@ -54,7 +39,7 @@ public object NovelAiCommand : SimpleCommand(
                 }
                 tag.startsWith("[") -> continue
                 else -> {
-                    input.add(translate(word = tag) ?: tag)
+                    input.add(NovelAiHelper.translate(word = tag) ?: tag)
                 }
             }
         }
@@ -88,7 +73,7 @@ public object NovelAiCommand : SimpleCommand(
             sendMessage(cause.error.message)
             return
         }
-        val image = subject.uploadImage(generate.data.toExternalResource().toAutoCloseable())
+        val image = generate.data.toExternalResource().use { subject.uploadImage(it) }
         sendMessage(fromEvent.message.quote() + image + "\n$seed")
     }
 }
diff --git a/src/main/kotlin/xyz/cssxsh/mirai/novelai/command/NovelAiFuCommand.kt b/src/main/kotlin/xyz/cssxsh/mirai/novelai/command/NovelAiFuCommand.kt
@@ -0,0 +1,81 @@
+package xyz.cssxsh.mirai.novelai.command
+
+import io.ktor.client.call.*
+import io.ktor.client.request.*
+import io.ktor.util.*
+import io.ktor.utils.io.core.*
+import kotlinx.serialization.json.*
+import net.mamoe.mirai.console.command.*
+import net.mamoe.mirai.message.data.*
+import net.mamoe.mirai.message.data.Image.Key.queryUrl
+import net.mamoe.mirai.message.data.MessageSource.Key.quote
+import net.mamoe.mirai.utils.ExternalResource.Companion.toExternalResource
+import xyz.cssxsh.mirai.novelai.*
+import xyz.cssxsh.mirai.novelai.data.*
+import xyz.cssxsh.novelai.*
+import kotlin.random.*
+
+public object NovelAiFuCommand : SimpleCommand(
+    owner = NovelAiHelper,
+    "nai-fu", "naifu",
+    description = "生成图片"
+) {
+    public val client: NovelAiClient = NovelAiClient(config = NovelAiHelperConfig, default = false)
+
+    private val random = Random(seed = System.currentTimeMillis())
+
+    @Handler
+    public suspend fun CommandSenderOnMessage<*>.handle(vararg tags: String) {
+        this as UserCommandSender
+        val input: MutableSet<String> = HashSet()
+        val params: MutableMap<String, String> = HashMap()
+        for (tag in tags) {
+            when {
+                tag.startsWith("#") -> {
+                    try {
+                        val (_, key, value) = tag.split('#', '=')
+                        params[key] = value
+                    } catch (_: Exception) {
+                        //
+                    }
+                }
+                tag.startsWith("[") -> continue
+                else -> {
+                    input.add(NovelAiHelper.translate(word = tag) ?: tag)
+                }
+            }
+        }
+        fromEvent.message.findIsInstance<Image>()?.let { source ->
+            try {
+                val url = source.queryUrl()
+                val response = client.http.get(url)
+                val packet = response.body<ByteReadPacket>()
+                params["image"] = packet.encodeBase64()
+            } catch (cause: Exception) {
+                NovelAiHelper.logger.warning("download image fail", cause)
+                null
+            }
+        }
+
+        val seed = random.nextLong(0, 2 shl 32 - 1)
+        NovelAiHelper.logger.info(input.joinToString(", ", "generate image seed: $seed, tags: "))
+        val generate = try {
+            client.ai.generateImage(input = input.joinToString(",")) {
+                put("seed", seed)
+                params.forEach { (key, value) ->
+                    when {
+                        key == "image" -> put(key, value)
+                        value.toDoubleOrNull() != null -> put(key, value.toDouble())
+                        value.toBooleanStrictOrNull() != null -> put(key, value.toBoolean())
+                        else -> put(key, value)
+                    }
+                }
+            }
+        } catch (cause: NovelAiApiException) {
+            sendMessage(cause.error.message)
+            return
+        }
+        val image = subject.uploadImage(generate.data.toExternalResource().toAutoCloseable())
+        sendMessage(fromEvent.message.quote() + image + "\n$seed")
+    }
+}
diff --git a/src/main/kotlin/xyz/cssxsh/mirai/novelai/data/NovelAiHelperConfig.kt b/src/main/kotlin/xyz/cssxsh/mirai/novelai/data/NovelAiHelperConfig.kt
@@ -29,6 +29,9 @@ public object NovelAiHelperConfig : ReadOnlyPluginConfig("config"), NovelAiClien
     override val timeout: Long by value(30_000L)
     override val image: ImageModel by value(ImageModel.SAFE_DIFFUSION)
 
+    @ValueName("naifu_api")
+    override val baseUrl: String by value("http://127.0.0.1:6969/")
+
     override var token: String = ""
         get() {
             return field.ifEmpty {

diff --git a/src/main/kotlin/xyz/cssxsh/novelai/NovelAiClient.kt b/src/main/kotlin/xyz/cssxsh/novelai/NovelAiClient.kt
@@ -19,7 +19,7 @@ import xyz.cssxsh.novelai.subscription.*
 import xyz.cssxsh.novelai.user.*
 import java.net.*
 
-public open class NovelAiClient(internal val config: NovelAiClientConfig) {
+public open class NovelAiClient(internal val config: NovelAiClientConfig, internal val default: Boolean = false) {
     public open val http: HttpClient = HttpClient(OkHttp) {
         install(ContentNegotiation) {
             json(json = Json)
@@ -62,6 +62,9 @@ public open class NovelAiClient(internal val config: NovelAiClientConfig) {
         }
         BrowserUserAgent()
         ContentEncoding()
+        defaultRequest {
+            if (!default) url(config.baseUrl)
+        }
         engine {
             config {
                 dns(

diff --git a/src/main/kotlin/xyz/cssxsh/novelai/NovelAiClientConfig.kt b/src/main/kotlin/xyz/cssxsh/novelai/NovelAiClientConfig.kt
@@ -9,5 +9,6 @@ public interface NovelAiClientConfig {
     public val timeout: Long
     public val image: ImageModel
     public var token: String
+    public val baseUrl: String get() = "https://api.novelai.net"
     public val ban: String get() = "lowres, bad anatomy, bad hands, text, error, missing fingers, extra digit, fewer digits, cropped, worst quality, low quality, normal quality, jpeg artifacts, signature, watermark, username, blurry"
 }
diff --git a/src/main/kotlin/xyz/cssxsh/novelai/ai/AIController.kt b/src/main/kotlin/xyz/cssxsh/novelai/ai/AIController.kt
@@ -19,7 +19,7 @@ public class AIController(private val client: NovelAiClient) {
             model = model,
             parameters = buildJsonObject(block)
         )
-        val response = client.http.post("https://api.novelai.net/ai/generate") {
+        val response = client.http.post("/ai/generate") {
             setBody(body)
             contentType(ContentType.Application.Json)
         }
@@ -32,7 +32,7 @@ public class AIController(private val client: NovelAiClient) {
             model = model,
             parameters = buildJsonObject(block)
         )
-        val response = client.http.post("https://api.novelai.net/ai/generate-stream") {
+        val response = client.http.post("/ai/generate-stream") {
             setBody(body)
             contentType(ContentType.Application.Json)
         }
@@ -67,7 +67,7 @@ public class AIController(private val client: NovelAiClient) {
                 block.invoke(this)
             }
         )
-        val statement = client.http.preparePost("https://api.novelai.net/ai/generate-image") {
+        val statement = client.http.preparePost("/ai/generate-image") {
             setBody(body)
             contentType(ContentType.Application.Json)
         }
@@ -89,36 +89,40 @@ public class AIController(private val client: NovelAiClient) {
         var event = ""
         var id = 0L
         var data = ByteArray(0)
+        var key = ""
         while (packet.canRead()) {
-            val key = packet.readUTF8UntilDelimiter(":")
+            key = packet.readUTF8UntilDelimiter(":")
             packet.readByte()
             when (key.trim()) {
                 "event" -> event = packet.readUTF8UntilDelimiter("\n").trim()
                 "id" -> id = packet.readUTF8UntilDelimiter("\n").trim().toLong()
                 "data" -> data = packet.decodeBase64Bytes().readBytes()
             }
         }
+        if (event.isEmpty()) {
+            throw NovelAiApiException(error = Json.decodeFromString(key))
+        }
 
         return AiGenerateImage(event, id, data)
     }
 
     public suspend fun classify(): AiSequenceClassification {
-        val response = client.http.post("https://api.novelai.net/ai/classify") {
+        val response = client.http.post("/ai/classify") {
             contentType(ContentType.Application.Json)
         }
         return response.body()
     }
 
     public suspend fun generateImageTags(model: String, prompt: String): AiRequestImageGenerationTags {
-        val response = client.http.get("https://api.novelai.net/ai/generate-image/suggest-tags") {
+        val response = client.http.get("/ai/generate-image/suggest-tags") {
             parameter("model", model)
             parameter("prompt", prompt)
         }
         return response.body()
     }
 
     public suspend fun generateVoice(text: String, seed: String): AiSequenceClassification {
-        val response = client.http.get("https://api.novelai.net/ai/generate-voice") {
+        val response = client.http.get("/ai/generate-voice") {
             parameter("text", text)
             parameter("seed", seed)
             parameter("voice", "")

diff --git a/src/main/kotlin/xyz/cssxsh/novelai/subscription/SubscriptionController.kt b/src/main/kotlin/xyz/cssxsh/novelai/subscription/SubscriptionController.kt
@@ -8,7 +8,7 @@ import xyz.cssxsh.novelai.*
 public class SubscriptionController(private val client: NovelAiClient) {
     public suspend fun bind(id: String, processor: PaymentProcessor) {
         val body = BindSubscriptionRequest(paymentProcessor = processor, subscriptionId = id)
-        val response = client.http.post("https://api.novelai.net/user/subscription/bind") {
+        val response = client.http.post("/user/subscription/bind") {
             setBody(body)
             contentType(ContentType.Application.Json)
         }
@@ -17,7 +17,7 @@ public class SubscriptionController(private val client: NovelAiClient) {
 
     public suspend fun change(plan: SubscriptionPlan) {
         val body = ChangeSubscriptionPlanRequest(newSubscriptionPlan = plan)
-        val response = client.http.post("https://api.novelai.net/user/subscription/change") {
+        val response = client.http.post("/user/subscription/change") {
             setBody(body)
             contentType(ContentType.Application.Json)
         }