feat: support nai-diffusion-3 model (#221)

Author: MirrorCY

src/config.ts

@@ -32,6 +32,7 @@ type Orient = keyof typeof orientMap
 
 export const models = Object.keys(modelMap) as Model[]
 export const orients = Object.keys(orientMap) as Orient[]
+export const scheduler = ['native', 'karras', 'exponential', 'polyexponential'] as const
 
 export namespace sampler {
   export const nai = {
@@ -42,6 +43,15 @@ export namespace sampler {
     'plms': 'PLMS',
   }
 
+  export const nai3 = {
+    'k_euler': 'Euler',
+    'k_euler_a': 'Euler ancestral',
+    'k_dpmpp_2s_ancestral': 'DPM++ 2S ancestral',
+    'k_dpmpp_2m': 'DPM++ 2M',
+    'k_dpmpp_sde': 'DPM++ SDE',
+    'ddim_v3': 'DDIM V3',
+  }
+
   // samplers in stable-diffusion-webui
   // https://github.com/AUTOMATIC1111/stable-diffusion-webui/blob/master/modules/sd_samplers_compvis.py#L12
   // https://github.com/AUTOMATIC1111/stable-diffusion-webui/blob/master/modules/sd_samplers_kdiffusion.py#L12
@@ -103,9 +113,10 @@ export namespace sampler {
     })).loose().description('默认的采样器。').default('k_euler_a')
   }
 
-  export function sd2nai(sampler: string): string {
+  export function sd2nai(sampler: string, model: string): string {
     if (sampler === 'k_euler_a') return 'k_euler_ancestral'
-    if (sampler in nai) return sampler
+    if (model === 'nai-v3' && sampler in nai3) return sampler
+    else if (sampler in nai) return sampler
     return 'k_euler_ancestral'
   }
 }
@@ -193,6 +204,10 @@ const features = Schema.object({
 interface ParamConfig {
   model?: Model
   sampler?: string
+  smea?: boolean
+  smeaDyn?: boolean
+  scheduler?: string
+  decrisper?: boolean
   upscaler?: string
   restoreFaces?: boolean
   hiresFix?: boolean
@@ -325,10 +340,22 @@ export const Config = Schema.intersect([
       type: Schema.const('naifu').required(),
       sampler: sampler.createSchema(sampler.nai),
     }),
-    Schema.object({
-      sampler: sampler.createSchema(sampler.nai),
-      model: Schema.union(models).loose().description('默认的生成模型。').default('nai'),
-    }),
+    Schema.intersect([
+      Schema.object({
+        model: Schema.union(models).loose().description('默认的生成模型。').default('nai'),
+      }),
+      Schema.union([
+        Schema.object({
+          model: Schema.const('nai-v3').required(),
+          sampler: sampler.createSchema(sampler.nai3),
+          smea: Schema.boolean().description('默认启用 SMEA。'),
+          smeaDyn: Schema.boolean().description('默认启用 SMEA 采样器的 DYN 变体。'),
+          scheduler: Schema.union(scheduler).description('默认的调度器。').default('native'),
+        }),
+        Schema.object({ sampler: sampler.createSchema(sampler.nai) }),
+      ]),
+      Schema.object({ decrisper: Schema.boolean().description('默认启用 decrisper') }),
+    ]),
   ] as const),
 
   Schema.object({

src/index.ts

@@ -1,9 +1,9 @@
 import { Computed, Context, Dict, h, Logger, omit, Quester, Session, SessionError, trimSlash } from 'koishi'
-import { Config, modelMap, models, orientMap, parseInput, sampler, upscalers } from './config'
+import { Config, modelMap, models, orientMap, parseInput, sampler, upscalers, scheduler } from './config'
 import { ImageData, StableDiffusionWebUI } from './types'
 import { closestMultiple, download, forceDataPrefix, getImageSize, login, NetworkError, project, resizeInput, Size } from './utils'
-import {} from '@koishijs/translator'
-import {} from '@koishijs/plugin-help'
+import { } from '@koishijs/translator'
+import { } from '@koishijs/plugin-help'
 import AdmZip from 'adm-zip'
 
 export * from './config'
@@ -107,6 +107,10 @@ export function apply(ctx: Context, config: Config) {
     .option('noise', '-n <noise:number>', { hidden: some(restricted, thirdParty) })
     .option('strength', '-N <strength:number>', { hidden: restricted })
     .option('hiresFix', '-H', { hidden: () => config.type !== 'sd-webui' })
+    .option('smea', '-S', { hidden: () => config.model !== 'nai-v3' })
+    .option('smeaDyn', '-d', { hidden: () => config.model !== 'nai-v3' })
+    .option('scheduler', '-C <scheduler> ', { hidden: () => config.model !== 'nai-v3', type: scheduler })
+    .option('decrisper', '-D', { hidden: thirdParty })
     .option('undesired', '-u <undesired>')
     .option('noTranslator', '-T', { hidden: () => !ctx.translator || !config.translator })
     .option('iterations', '-i <iterations:posint>', { fallback: 1, hidden: () => config.maxIterations <= 1 })
@@ -306,12 +310,29 @@ export function apply(ctx: Context, config: Config) {
           case 'login':
           case 'token':
           case 'naifu': {
-            parameters.sampler = sampler.sd2nai(options.sampler)
+            parameters.sampler = sampler.sd2nai(options.sampler, model)
             parameters.image = image?.base64 // NovelAI / NAIFU accepts bare base64 encoded image
             if (config.type === 'naifu') return parameters
             // The latest interface changes uc to negative_prompt, so that needs to be changed here as well
-            parameters.negative_prompt = parameters.uc
-            delete parameters.uc
+            if (parameters.uc) {
+              parameters.negative_prompt = parameters.uc
+              delete parameters.uc
+            }
+            parameters.dynamic_thresholding = options.decrisper ?? config.decrisper
+            if (model === 'nai-diffusion-3') {
+              parameters.sm_dyn = options.smeaDyn ?? config.smeaDyn
+              parameters.sm = (options.smea ?? config.smea) || parameters.sm_dyn
+              parameters.noise_schedule = options.scheduler ?? config.scheduler
+              if (['k_euler_ancestral', 'k_dpmpp_2s_ancestral'].includes(parameters.sampler)
+                && parameters.noise_schedule === 'karras') {
+                parameters.noise_schedule = 'native'
+              }
+              if (parameters.sampler === 'ddim_v3') {
+                parameters.sm = false
+                parameters.sm_dyn = false
+                delete parameters.noise_schedule
+              }
+            }
             return { model, input: prompt, parameters: omit(parameters, ['prompt']) }
           }
           case 'sd-webui': {
@@ -436,7 +457,7 @@ export function apply(ctx: Context, config: Config) {
             const b64 = Buffer.from(firstImageBuffer).toString('base64')
             return forceDataPrefix(b64, 'image/png')
           }
-          
+
           return forceDataPrefix(res.data?.slice(27))
         }
 
@@ -522,7 +543,7 @@ export function apply(ctx: Context, config: Config) {
         case 'stable-horde':
           return sampler.horde
         default:
-          return sampler.nai
+          return { ...sampler.nai, ...sampler.nai3 }
       }
     }
 

src/locales/zh-CN.yml

@@ -26,6 +26,10 @@ commands:
       noTranslator: 禁用自动翻译
       iterations: 设置绘制次数
       batch: 设置绘制批次大小
+      smea: 启用 SMEA
+      smeaDyn: 启用 DYN
+      scheduler: 设置调度器
+      decrisper: 启用动态阈值
 
     messages:
       exceed-max-iteration: 超过最大绘制次数。