Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

feat: support sd / stablehorde scheduler option #246

Merged
merged 2 commits into from
Aug 14, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
26 changes: 8 additions & 18 deletions data/sd-samplers.json
Original file line number Diff line number Diff line change
@@ -1,31 +1,21 @@
{
"k_dpmpp_2m_ka": "DPM++ 2M Karras",
"k_dpmpp_sde_ka": "DPM++ SDE Karras",
"k_dpmpp_2m_sde_exp": "DPM++ 2M SDE Exponential",
"k_dpmpp_2m_sde_ka": "DPM++ 2M SDE",
"k_dpmpp_2m": "DPM++ 2M",
"k_dpmpp_sde": "DPM++ SDE",
"k_dpmpp_2m_sde": "DPM++ 2M SDE",
"k_dpmpp_2m_sde_heun": "DPM++ 2M SDE Heun",
"k_dpmpp_2s_a": "DPM++ 2S a",
"k_dpmpp_3m_sde": "DPM++ 3M SDE",
"k_euler_a": "Euler a",
"k_euler": "Euler",
"k_lms": "LMS",
"k_heun": "Heun",
"k_dpm_2": "DPM2",
"k_dpm_2_a": "DPM2 a",
"k_dpmpp_2s_a": "DPM++ 2S a",
"k_dpmpp_2m": "DPM++ 2M",
"k_dpmpp_sde": "DPM++ SDE",
"k_dpmpp_2m_sde_heun": "DPM++ 2M SDE Heun",
"k_dpmpp_2m_sde_heun_ka": "DPM++ 2M SDE Heun Karras",
"k_dpmpp_2m_sde_heun_exp": "DPM++ 2M SDE Heun Exponential",
"k_dpmpp_3m_sde": "DPM++ 3M SDE",
"k_dpmpp_3m_sde_ka": "DPM++ 3M SDE Karras",
"k_dpmpp_3m_sde_exp": "DPM++ 3M SDE Exponential",
"k_dpm_fast": "DPM fast",
"k_dpm_ad": "DPM adaptive",
"k_lms_ka": "LMS Karras",
"k_dpm_2_ka": "DPM2 Karras",
"k_dpm_2_a_ka": "DPM2 a Karras",
"k_dpmpp_2s_a_ka": "DPM++ 2S a Karras",
"restart": "Restart",
"ddim": "DDIM",
"plms": "PLMS",
"unipc": "UniPC"
"unipc": "UniPC",
"k_lcm": "LCM"
}
24 changes: 8 additions & 16 deletions src/config.ts
Original file line number Diff line number Diff line change
Expand Up @@ -32,7 +32,11 @@ 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 scheduler {
export const nai = ['native', 'karras', 'exponential', 'polyexponential'] as const
export const sd = ['Automatic', 'Uniform', 'Karras', 'Exponential', 'Polyexponential', 'SGM Uniform'] as const
export const horde = ['karras'] as const
}
export const schedulerComfyUI = ['normal', 'karras', 'exponential', 'sgm_uniform', 'simple', 'ddim_uniform'] as const

export namespace sampler {
Expand Down Expand Up @@ -72,20 +76,6 @@ export namespace sampler {
dpmsolver: 'DPM solver',
lcm: 'LCM',
DDIM: 'DDIM',
k_lms_ka: 'LMS Karras',
k_heun_ka: 'Heun Karras',
k_euler_ka: 'Euler Karras',
k_euler_a_ka: 'Euler a Karras',
k_dpm_2_ka: 'DPM2 Karras',
k_dpm_2_a_ka: 'DPM2 a Karras',
k_dpm_fast_ka: 'DPM fast Karras',
k_dpm_adaptive_ka: 'DPM adaptive Karras',
k_dpmpp_2m_ka: 'DPM++ 2M Karras',
k_dpmpp_2s_a_ka: 'DPM++ 2S a Karras',
k_dpmpp_sde_ka: 'DPM++ SDE Karras',
dpmsolver_ka: 'DPM++ solver Karras',
lcm_ka: 'LCM Karras',
DDIM_ka: 'DDIM Karras',
}

export const comfyui = {
Expand Down Expand Up @@ -347,11 +337,13 @@ export const Config = Schema.intersect([
upscaler: Schema.union(upscalers).description('默认的放大算法。').default('Lanczos'),
restoreFaces: Schema.boolean().description('是否启用人脸修复。').default(false),
hiresFix: Schema.boolean().description('是否启用高分辨率修复。').default(false),
scheduler: Schema.union(scheduler.sd).description('默认的调度器。').default('Automatic'),
}),
Schema.object({
type: Schema.const('stable-horde').required(),
sampler: sampler.createSchema(sampler.horde),
model: Schema.union(hordeModels).loose().description('默认的生成模型。'),
scheduler: Schema.union(scheduler.horde).description('默认的调度器。').default('karras'),
}),
Schema.object({
type: Schema.const('naifu').required(),
Expand Down Expand Up @@ -381,7 +373,7 @@ export const Config = Schema.intersect([
sampler: sampler.createSchema(sampler.nai3),
smea: Schema.boolean().description('默认启用 SMEA。'),
smeaDyn: Schema.boolean().description('默认启用 SMEA 采样器的 DYN 变体。'),
scheduler: Schema.union(scheduler).description('默认的调度器。').default('native'),
scheduler: Schema.union(scheduler.nai).description('默认的调度器。').default('native'),
}),
Schema.object({ sampler: sampler.createSchema(sampler.nai) }),
]),
Expand Down
16 changes: 13 additions & 3 deletions src/index.ts
Original file line number Diff line number Diff line change
Expand Up @@ -112,7 +112,16 @@ export function apply(ctx: Context, config: Config) {
.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('scheduler', '-C <scheduler:string>', {
hidden: () => config.type === 'naifu',
type: ['token', 'login'].includes(config.type)
? scheduler.nai
: config.type === 'sd-webui'
? scheduler.sd
: config.type === 'stable-horde'
? scheduler.horde
: [],
})
.option('decrisper', '-D', { hidden: thirdParty })
.option('undesired', '-u <undesired>')
.option('noTranslator', '-T', { hidden: () => !ctx.translator || !config.translator })
Expand Down Expand Up @@ -352,6 +361,7 @@ export function apply(ctx: Context, config: Config) {
case 'sd-webui': {
return {
sampler_index: sampler.sd[options.sampler],
scheduler: options.scheduler,
init_images: image && [image.dataUrl], // sd-webui accepts data URLs with base64 encoded image
restore_faces: config.restoreFaces ?? false,
enable_hr: options.hiresFix ?? config.hiresFix ?? false,
Expand All @@ -373,14 +383,14 @@ export function apply(ctx: Context, config: Config) {
return {
prompt: parameters.prompt,
params: {
sampler_name: options.sampler.replace('_ka', ''),
sampler_name: options.sampler,
cfg_scale: parameters.scale,
denoising_strength: parameters.strength,
seed: parameters.seed.toString(),
height: parameters.height,
width: parameters.width,
post_processing: [],
karras: options.sampler.includes('_ka'),
karras: options.scheduler?.toLowerCase() === 'karras',
hires_fix: options.hiresFix ?? config.hiresFix ?? false,
steps: parameters.steps,
n: parameters.n_samples,
Expand Down
Loading