feat: support nai4 full (#274)

* chore: make linter happy * chore: update deps * feat: support nai4 full
koishijs · Mar 3, 2025 · 85cd7b9 · 85cd7b9
1 parent 4d09b73
commit 85cd7b9
Show file tree

Hide file tree

Showing 3 changed files with 42 additions and 32 deletions.
diff --git a/package.json b/package.json
@@ -63,24 +63,24 @@
     "generate"
   ],
   "peerDependencies": {
-    "koishi": "^4.17.8"
+    "koishi": "^4.18.7"
   },
   "devDependencies": {
-    "@cordisjs/vitepress": "^3.2.7",
-    "@koishijs/plugin-help": "^2.4.3",
+    "@cordisjs/vitepress": "^3.3.2",
+    "@koishijs/plugin-help": "^2.4.5",
     "@koishijs/translator": "^1.1.1",
-    "@types/adm-zip": "^0.5.5",
+    "@types/adm-zip": "^0.5.7",
     "@types/libsodium-wrappers-sumo": "^0.7.8",
-    "@types/node": "^20.14.1",
+    "@types/node": "^22.13.8",
     "atsc": "^1.2.2",
-    "koishi": "^4.17.8",
-    "sass": "^1.77.4",
-    "typescript": "^5.4.5",
+    "koishi": "^4.18.7",
+    "sass": "^1.85.1",
+    "typescript": "^5.8.2",
     "vitepress": "1.0.0-rc.40"
   },
   "dependencies": {
-    "adm-zip": "^0.5.13",
-    "image-size": "^1.1.1",
-    "libsodium-wrappers-sumo": "^0.7.13"
+    "adm-zip": "^0.5.16",
+    "image-size": "^1.2.0",
+    "libsodium-wrappers-sumo": "^0.7.15"
   }
 }
diff --git a/src/config.ts b/src/config.ts
@@ -11,6 +11,7 @@ export const modelMap = {
   furry: 'nai-diffusion-furry',
   'nai-v3': 'nai-diffusion-3',
   'nai-v4-curated-preview': 'nai-diffusion-4-curated-preview',
+  'nai-v4-full': 'nai-diffusion-4-full',
 } as const
 
 export const orientMap = {
@@ -266,6 +267,12 @@ export interface Config extends PromptConfig, ParamConfig {
   workflowImage2Image?: string
 }
 
+const NAI4ParamConfig = Schema.object({
+  sampler: sampler.createSchema(sampler.nai4).default('k_euler_a'),
+  scheduler: Schema.union(scheduler.nai4).description('默认的调度器。').default('karras'),
+  rescale: Schema.computed(Schema.number(), options).min(0).max(1).description('输入服从度调整规模。').default(0),
+})
+
 export const Config = Schema.intersect([
   Schema.object({
     type: Schema.union([
@@ -404,9 +411,11 @@ export const Config = Schema.intersect([
         }),
         Schema.object({
           model: Schema.const('nai-v4-curated-preview'),
-          sampler: sampler.createSchema(sampler.nai4).default('k_euler_a'),
-          scheduler: Schema.union(scheduler.nai4).description('默认的调度器。').default('karras'),
-          rescale: Schema.computed(Schema.number(), options).min(0).max(1).description('输入服从度调整规模。').default(0),
+          ...NAI4ParamConfig.dict,
+        }),
+        Schema.object({
+          model: Schema.const('nai-v4-full'),
+          ...NAI4ParamConfig.dict,
         }),
         Schema.object({ sampler: sampler.createSchema(sampler.nai) }),
       ]),

diff --git a/src/index.ts b/src/index.ts
@@ -1,4 +1,4 @@
-import { Computed, Context, Dict, h, Logger, omit, Quester, Session, SessionError, trimSlash } from 'koishi'
+import { Computed, Context, Dict, h, omit, Quester, Session, SessionError, trimSlash } from 'koishi'
 import { Config, modelMap, models, orientMap, parseInput, sampler, upscalers, scheduler } from './config'
 import { ImageData, NovelAI, StableDiffusionWebUI } from './types'
 import { closestMultiple, download, forceDataPrefix, getImageSize, login, NetworkError, project, resizeInput, Size } from './utils'
@@ -333,7 +333,9 @@ export function apply(ctx: Context, config: Config) {
               delete parameters.uc
             }
             parameters.dynamic_thresholding = options.decrisper ?? config.decrisper
-            if (model === 'nai-diffusion-3' || model === 'nai-diffusion-4-curated-preview') {
+            const isNAI3 = model === 'nai-diffusion-3'
+            const isNAI4 = model === 'nai-diffusion-4-curated-preview' || model === 'nai-diffusion-4-full'
+            if (isNAI3 || isNAI4) {
               parameters.params_version = 3
               parameters.legacy = false
               parameters.legacy_v3_extend = false
@@ -344,7 +346,7 @@ export function apply(ctx: Context, config: Config) {
               if (parameters.scale > 10) {
                 parameters.scale = parameters.scale / 2
               }
-              if (model === 'nai-diffusion-3') {
+              if (isNAI3) {
                 parameters.sm_dyn = options.smeaDyn ?? config.smeaDyn
                 parameters.sm = (options.smea ?? config.smea) || parameters.sm_dyn
                 if (['k_euler_ancestral', 'k_dpmpp_2s_ancestral'].includes(parameters.sampler)
@@ -356,19 +358,18 @@ export function apply(ctx: Context, config: Config) {
                   parameters.sm_dyn = false
                   delete parameters.noise_schedule
                 }
-              }
-              if (model === 'nai-diffusion-4-curated-preview') {
-                parameters.add_original_image = true  // unknown
+              } else if (isNAI4) {
+                parameters.add_original_image = true // unknown
                 parameters.cfg_rescale = session.resolve(config.rescale)
                 parameters.characterPrompts = [] satisfies NovelAI.V4CharacterPrompt[]
-                parameters.controlnet_strength = 1  // unknown
-                parameters.deliberate_euler_ancestral_bug = false  // unknown
-                parameters.prefer_brownian = true  // unknown
-                parameters.reference_image_multiple = []  // unknown
-                parameters.reference_information_extracted_multiple = []  // unknown
-                parameters.reference_strength_multiple = []  // unknown
-                parameters.skip_cfg_above_sigma = null  // unknown
-                parameters.use_coords = false  // unknown
+                parameters.controlnet_strength = 1 // unknown
+                parameters.deliberate_euler_ancestral_bug = false // unknown
+                parameters.prefer_brownian = true // unknown
+                parameters.reference_image_multiple = [] // unknown
+                parameters.reference_information_extracted_multiple = [] // unknown
+                parameters.reference_strength_multiple = [] // unknown
+                parameters.skip_cfg_above_sigma = null // unknown
+                parameters.use_coords = false // unknown
                 parameters.v4_prompt = {
                   caption: {
                     base_caption: prompt,
@@ -452,13 +453,13 @@ export function apply(ctx: Context, config: Config) {
             if (image) {
               const body = new FormData()
               const capture = /^data:([\w/.+-]+);base64,(.*)$/.exec(image.dataUrl)
-              const [, mime,] = capture
+              const [, mime] = capture
 
               let name = Date.now().toString()
               const ext = mime === 'image/jpeg' ? 'jpg' : mime === 'image/png' ? 'png' : ''
               if (ext) name += `.${ext}`
               const imageFile = new Blob([image.buffer], { type: mime })
-              body.append("image", imageFile, name)
+              body.append('image', imageFile, name)
               const res = await ctx.http(trimSlash(config.endpoint) + '/upload/image', {
                 method: 'POST',
                 headers: {
@@ -583,7 +584,7 @@ export function apply(ctx: Context, config: Config) {
               await sleep(config.pollInterval)
             }
             // get images by filename
-            const imagesOutput: { data: ArrayBuffer, mime: string }[] = [];
+            const imagesOutput: { data: ArrayBuffer; mime: string }[] = []
             for (const nodeId in outputs) {
               const nodeOutput = outputs[nodeId]
               if ('images' in nodeOutput) {
@@ -603,7 +604,7 @@ export function apply(ctx: Context, config: Config) {
           // data:
           //                                                                        ↓ nai-v3
           if (res.headers.get('content-type') === 'application/x-zip-compressed' || res.headers.get('content-disposition')?.includes('.zip')) {
-            const buffer = Buffer.from(res.data, 'binary')  // Ensure 'binary' encoding
+            const buffer = Buffer.from(res.data, 'binary') // Ensure 'binary' encoding
             const zip = new AdmZip(buffer)
 
             // Gets all files in the ZIP file