diff --git a/app/comfy.test.ts b/app/comfy.test.ts index 70ac3fe..fbea511 100644 --- a/app/comfy.test.ts +++ b/app/comfy.test.ts @@ -1,5 +1,5 @@ import { describe, expect, test } from "bun:test"; -import { buildSdxlWorkflow, selectGeneratedOutputImage } from "./comfy"; +import { buildAnimaWorkflow, buildSdxlWorkflow, buildZImageTurboWorkflow, buildZImageWorkflow, selectGeneratedOutputImage } from "./comfy"; describe("Comfy adapter", () => { test("selects SaveImage output instead of uploaded input or mask images", () => { @@ -43,6 +43,48 @@ describe("Comfy adapter", () => { expect(workflow["13"]?.inputs).toMatchObject({ samples: ["5", 0], mask: ["12", 0] }); expect(workflow["6"]?.inputs.latent_image).toEqual(["13", 0]); }); + + test("builds Z-Image text-to-image with separated model loaders", () => { + const workflow = buildZImageWorkflow(textRequest({ architecture: "z-image", model: "z_image_bf16.safetensors", steps: 30, cfg: 4 })); + + expect(workflow["1"]).toMatchObject({ class_type: "UNETLoader", inputs: { unet_name: "z_image_bf16.safetensors", weight_dtype: "default" } }); + expect(workflow["2"]).toMatchObject({ class_type: "CLIPLoader", inputs: { clip_name: "qwen_3_4b.safetensors", type: "lumina2" } }); + expect(workflow["6"]?.class_type).toBe("EmptySD3LatentImage"); + expect(workflow["7"]).toMatchObject({ class_type: "ModelSamplingAuraFlow", inputs: { shift: 3 } }); + expect(workflow["8"]?.inputs).toMatchObject({ steps: 30, cfg: 4, sampler_name: "res_multistep", scheduler: "simple", model: ["7", 0] }); + expect(workflow["10"]?.class_type).toBe("SaveImage"); + }); + + test("builds Z-Image Turbo with zeroed negative conditioning", () => { + const workflow = buildZImageTurboWorkflow(textRequest({ architecture: "z-image-turbo", model: "z_image_turbo_bf16.safetensors", negativePrompt: "ignored" })); + + expect(workflow["1"]?.inputs.unet_name).toBe("z_image_turbo_bf16.safetensors"); + expect(workflow["5"]).toMatchObject({ class_type: "ConditioningZeroOut", inputs: { conditioning: ["4", 0] } }); + expect(workflow["8"]?.inputs).toMatchObject({ steps: 8, cfg: 1, sampler_name: "res_multistep", scheduler: "simple" }); + }); + + test("builds Anima text-to-image workflow", () => { + const workflow = buildAnimaWorkflow(textRequest({ architecture: "anima", model: "anima-base-v1.0.safetensors" })); + + expect(workflow["1"]).toMatchObject({ class_type: "UNETLoader", inputs: { unet_name: "anima-base-v1.0.safetensors" } }); + expect(workflow["2"]).toMatchObject({ class_type: "CLIPLoader", inputs: { clip_name: "qwen_3_06b_base.safetensors", type: "stable_diffusion" } }); + expect(workflow["3"]).toMatchObject({ class_type: "VAELoader", inputs: { vae_name: "qwen_image_vae.safetensors" } }); + expect(workflow["6"]?.class_type).toBe("EmptyLatentImage"); + expect(workflow["8"]?.inputs).toMatchObject({ steps: 30, cfg: 4, sampler_name: "er_sde", scheduler: "simple", model: ["1", 0] }); + }); + + test("builds Anima with model-specific text encoder and VAE", () => { + const workflow = buildAnimaWorkflow(textRequest({ + architecture: "anima", + model: "miaomiaoHarem_anima13.safetensors", + textEncoder: "miaomiaoHarem_anima13_txt.safetensors", + vae: "qwen_image_vae.safetensors", + })); + + expect(workflow["1"]?.inputs.unet_name).toBe("miaomiaoHarem_anima13.safetensors"); + expect(workflow["2"]?.inputs.clip_name).toBe("miaomiaoHarem_anima13_txt.safetensors"); + expect(workflow["3"]?.inputs.vae_name).toBe("qwen_image_vae.safetensors"); + }); }); function inpaintRequest(inpaint: { maskedContent: "neutral" | "original"; growMaskBy?: number }) { @@ -57,3 +99,17 @@ function inpaintRequest(inpaint: { maskedContent: "neutral" | "original"; growMa inpaint, }; } + +function textRequest(overrides: Partial[0]> = {}) { + return { + architecture: "sdxl" as const, + mode: "text-to-image" as const, + model: "model.safetensors", + prompt: "a studio portrait", + negativePrompt: "low quality", + width: 1024, + height: 1024, + seed: 123, + ...overrides, + }; +} diff --git a/app/comfy.ts b/app/comfy.ts index 3055927..4abe3ea 100644 --- a/app/comfy.ts +++ b/app/comfy.ts @@ -1,8 +1,13 @@ +type GenerateArchitecture = "sdxl" | "z-image" | "z-image-turbo" | "anima"; type GenerateMode = "text-to-image" | "image-to-image" | "inpaint" | "outpaint"; +type Workflow = Record }>; type ComfyGenerateRequest = { + architecture?: GenerateArchitecture; mode: GenerateMode; model: string; + textEncoder?: string; + vae?: string; prompt: string; negativePrompt?: string; strength?: number; @@ -35,7 +40,21 @@ type ComfyGenerateRequest = { maskImage?: string; }; +type ComfyObjectInfo = { + CheckpointLoaderSimple?: { input?: { required?: { ckpt_name?: [string[]] } } }; + KSampler?: { input?: { required?: { sampler_name?: [string[]]; scheduler?: [string[]] } } }; + UNETLoader?: { input?: { required?: { unet_name?: [string[]] } } }; + CLIPLoader?: { input?: { required?: { clip_name?: [string[]] } } }; + VAELoader?: { input?: { required?: { vae_name?: [string[]] } } }; +}; + const comfyBaseUrl = process.env.COMFYUI_URL ?? "http://127.0.0.1:8188"; +const defaultModels: Record = { + sdxl: "sd_xl_base_1.0.safetensors", + "z-image": "z_image_bf16.safetensors", + "z-image-turbo": "z_image_turbo_bf16.safetensors", + anima: "anima-base-v1.0.safetensors", +}; export async function handleComfyApi(request: Request) { try { @@ -51,14 +70,49 @@ export async function handleComfyApi(request: Request) { async function listGenerationOptions() { const response = await fetch(`${comfyBaseUrl}/object_info`); if (!response.ok) throw new Error(`ComfyUI option lookup failed: ${response.status}`); - const info = await response.json() as { - CheckpointLoaderSimple?: { input?: { required?: { ckpt_name?: [string[]] } } }; - KSampler?: { input?: { required?: { sampler_name?: [string[]]; scheduler?: [string[]] } } }; - }; + const info = await response.json() as ComfyObjectInfo; + const checkpointModels = info.CheckpointLoaderSimple?.input?.required?.ckpt_name?.[0] ?? []; + const diffusionModels = info.UNETLoader?.input?.required?.unet_name?.[0] ?? []; + const textEncoders = info.CLIPLoader?.input?.required?.clip_name?.[0] ?? []; + const vaes = info.VAELoader?.input?.required?.vae_name?.[0] ?? []; + return { - models: info.CheckpointLoaderSimple?.input?.required?.ckpt_name?.[0] ?? [], + models: checkpointModels, samplers: info.KSampler?.input?.required?.sampler_name?.[0] ?? [], schedulers: info.KSampler?.input?.required?.scheduler?.[0] ?? [], + diffusionModels, + textEncoders, + vaes, + architectures: [ + { + value: "sdxl", + label: "SDXL", + defaultModel: checkpointModels[0] ?? defaultModels.sdxl, + models: checkpointModels, + supportedModes: ["text-to-image", "image-to-image", "inpaint", "outpaint"], + }, + { + value: "z-image", + label: "Z-Image", + defaultModel: defaultModels["z-image"], + models: modelsForArchitecture(diffusionModels, defaultModels["z-image"], [/z[_-]?image(?!.*turbo)/i]), + supportedModes: ["text-to-image"], + }, + { + value: "z-image-turbo", + label: "Z-Image Turbo", + defaultModel: defaultModels["z-image-turbo"], + models: modelsForArchitecture(diffusionModels, defaultModels["z-image-turbo"], [/z[_-]?image.*turbo/i]), + supportedModes: ["text-to-image"], + }, + { + value: "anima", + label: "Anima", + defaultModel: defaultModels.anima, + models: modelsForArchitecture(diffusionModels, defaultModels.anima, [/anima/i]), + supportedModes: ["text-to-image"], + }, + ], }; } @@ -68,16 +122,15 @@ async function listCheckpointModels() { async function generate(request: ComfyGenerateRequest) { if (!request.prompt?.trim()) throw new Error("Prompt is required"); - if (!request.model || request.model === "auto") { - const models = await listCheckpointModels(); - request.model = models[0] ?? "sd_xl_base_1.0.safetensors"; - } + const architecture = normalizeArchitecture(request.architecture); + if (request.mode !== "text-to-image" && architecture !== "sdxl") throw new Error(`${architectureLabel(architecture)} currently supports text-to-image only`); + if (!request.model || request.model === "auto") request.model = await defaultModelForArchitecture(architecture); const clientId = crypto.randomUUID(); const uploaded = request.inputImage ? await uploadDataUrl(request.inputImage, `image-studio-${crypto.randomUUID()}.png`) : undefined; const mask = request.maskImage ? await uploadDataUrl(request.maskImage, `image-studio-mask-${crypto.randomUUID()}.png`) : undefined; if (request.mode === "inpaint" && (!uploaded || !mask)) throw new Error("Inpaint requires normalized input and mask images"); - const prompt = buildSdxlWorkflow({ ...request, inputImage: uploaded, maskImage: mask }); + const prompt = buildComfyWorkflow({ ...request, architecture, inputImage: uploaded, maskImage: mask }); const queued = await fetch(`${comfyBaseUrl}/prompt`, { method: "POST", @@ -143,7 +196,21 @@ function isGeneratedImage(image: { filename: string; subfolder?: string; type?: return image.type === undefined || image.type === "output"; } -export function buildSdxlWorkflow(request: ComfyGenerateRequest) { +export function buildComfyWorkflow(request: ComfyGenerateRequest): Workflow { + switch (normalizeArchitecture(request.architecture)) { + case "z-image": + return buildZImageWorkflow(request); + case "z-image-turbo": + return buildZImageTurboWorkflow(request); + case "anima": + return buildAnimaWorkflow(request); + case "sdxl": + default: + return buildSdxlWorkflow(request); + } +} + +export function buildSdxlWorkflow(request: ComfyGenerateRequest): Workflow { if (!request.width || !request.height) throw new Error("Generation width and height are required"); const width = Math.max(64, Math.round(request.width)); const height = Math.max(64, Math.round(request.height)); @@ -155,7 +222,7 @@ export function buildSdxlWorkflow(request: ComfyGenerateRequest) { const scheduler = request.scheduler ?? "normal"; const positive = request.prompt; const negative = request.negativePrompt ?? ""; - const workflow: Record }> = { + const workflow: Workflow = { "1": { class_type: "CheckpointLoaderSimple", inputs: { ckpt_name: request.model } }, "2": { class_type: "CLIPTextEncode", inputs: { text: positive, clip: ["1", 1] } }, "3": { class_type: "CLIPTextEncode", inputs: { text: negative, clip: ["1", 1] } }, @@ -195,6 +262,114 @@ export function buildSdxlWorkflow(request: ComfyGenerateRequest) { return workflow; } +export function buildZImageWorkflow(request: ComfyGenerateRequest): Workflow { + return buildSeparatedTextToImageWorkflow(request, { + architecture: "z-image", + filenamePrefix: "image-studio-z-image", + model: defaultModels["z-image"], + textEncoder: "qwen_3_4b.safetensors", + vae: "ae.safetensors", + clipType: "lumina2", + latentNode: "EmptySD3LatentImage", + modelSamplingAuraFlow: true, + negativeMode: "prompt", + steps: 30, + cfg: 4, + sampler: "res_multistep", + scheduler: "simple", + }); +} + +export function buildZImageTurboWorkflow(request: ComfyGenerateRequest): Workflow { + return buildSeparatedTextToImageWorkflow(request, { + architecture: "z-image-turbo", + filenamePrefix: "image-studio-z-image-turbo", + model: defaultModels["z-image-turbo"], + textEncoder: "qwen_3_4b.safetensors", + vae: "ae.safetensors", + clipType: "lumina2", + latentNode: "EmptySD3LatentImage", + modelSamplingAuraFlow: true, + negativeMode: "zero", + steps: 8, + cfg: 1, + sampler: "res_multistep", + scheduler: "simple", + }); +} + +export function buildAnimaWorkflow(request: ComfyGenerateRequest): Workflow { + return buildSeparatedTextToImageWorkflow(request, { + architecture: "anima", + filenamePrefix: "image-studio-anima", + model: defaultModels.anima, + textEncoder: "qwen_3_06b_base.safetensors", + vae: "qwen_image_vae.safetensors", + clipType: "stable_diffusion", + latentNode: "EmptyLatentImage", + modelSamplingAuraFlow: false, + negativeMode: "prompt", + steps: 30, + cfg: 4, + sampler: "er_sde", + scheduler: "simple", + }); +} + +function buildSeparatedTextToImageWorkflow(request: ComfyGenerateRequest, config: { + architecture: GenerateArchitecture; + filenamePrefix: string; + model: string; + textEncoder: string; + vae: string; + clipType: string; + latentNode: "EmptyLatentImage" | "EmptySD3LatentImage"; + modelSamplingAuraFlow: boolean; + negativeMode: "prompt" | "zero"; + steps: number; + cfg: number; + sampler: string; + scheduler: string; +}): Workflow { + if (request.mode !== "text-to-image") throw new Error(`${architectureLabel(config.architecture)} currently supports text-to-image only`); + const options = resolveSamplerOptions(request, config); + const modelOutput: [string, number] = config.modelSamplingAuraFlow ? ["7", 0] : ["1", 0]; + const workflow: Workflow = { + "1": { class_type: "UNETLoader", inputs: { unet_name: request.model && request.model !== "auto" ? request.model : config.model, weight_dtype: "default" } }, + "2": { class_type: "CLIPLoader", inputs: { clip_name: resolveSupportModelName(request.textEncoder, config.textEncoder), type: config.clipType, device: "default" } }, + "3": { class_type: "VAELoader", inputs: { vae_name: resolveSupportModelName(request.vae, config.vae) } }, + "4": { class_type: "CLIPTextEncode", inputs: { text: request.prompt, clip: ["2", 0] } }, + "6": { class_type: config.latentNode, inputs: { width: options.width, height: options.height, batch_size: 1 } }, + "8": { class_type: "KSampler", inputs: { seed: options.seed, steps: options.steps, cfg: options.cfg, sampler_name: options.sampler, scheduler: options.scheduler, denoise: 1, model: modelOutput, positive: ["4", 0], negative: ["5", 0], latent_image: ["6", 0] } }, + "9": { class_type: "VAEDecode", inputs: { samples: ["8", 0], vae: ["3", 0] } }, + "10": { class_type: "SaveImage", inputs: { filename_prefix: config.filenamePrefix, images: ["9", 0] } }, + }; + + if (config.modelSamplingAuraFlow) workflow["7"] = { class_type: "ModelSamplingAuraFlow", inputs: { model: ["1", 0], shift: 3 } }; + workflow["5"] = config.negativeMode === "zero" + ? { class_type: "ConditioningZeroOut", inputs: { conditioning: ["4", 0] } } + : { class_type: "CLIPTextEncode", inputs: { text: request.negativePrompt ?? "", clip: ["2", 0] } }; + + return workflow; +} + +function resolveSamplerOptions(request: ComfyGenerateRequest, defaults: { steps: number; cfg: number; sampler: string; scheduler: string }) { + if (!request.width || !request.height) throw new Error("Generation width and height are required"); + return { + width: Math.max(64, Math.round(request.width)), + height: Math.max(64, Math.round(request.height)), + seed: request.seed === undefined || request.seed < 0 ? Math.floor(Math.random() * 2 ** 32) : Math.round(request.seed), + steps: Math.max(1, Math.round(request.steps ?? defaults.steps)), + cfg: Math.max(0, request.cfg ?? defaults.cfg), + sampler: request.sampler ?? defaults.sampler, + scheduler: request.scheduler ?? defaults.scheduler, + }; +} + +function resolveSupportModelName(value: string | undefined, fallback: string) { + return value && value !== "auto" ? value : fallback; +} + function resolveGrowMaskBy(request: ComfyGenerateRequest): number { const value = request.inpaint?.growMaskBy ?? 6; if (!Number.isFinite(value)) return 6; @@ -205,6 +380,39 @@ function usesOriginalLatentContent(request: ComfyGenerateRequest): boolean { return request.inpaint?.maskedContent === "original" || request.inpaint?.maskedContent === "originalColor" || request.inpaint?.maskedContent === "edges"; } +async function defaultModelForArchitecture(architecture: GenerateArchitecture) { + if (architecture !== "sdxl") return defaultModels[architecture]; + const models = await listCheckpointModels(); + return models[0] ?? defaultModels.sdxl; +} + +function normalizeArchitecture(architecture: ComfyGenerateRequest["architecture"]): GenerateArchitecture { + if (architecture === "z-image" || architecture === "z-image-turbo" || architecture === "anima") return architecture; + return "sdxl"; +} + +function architectureLabel(architecture: GenerateArchitecture): string { + switch (architecture) { + case "z-image": + return "Z-Image"; + case "z-image-turbo": + return "Z-Image Turbo"; + case "anima": + return "Anima"; + case "sdxl": + default: + return "SDXL"; + } +} + +function modelsForArchitecture(models: string[], defaultModel: string, matchers: RegExp[]) { + return unique([defaultModel, ...models.filter((model) => model === defaultModel || matchers.some((matcher) => matcher.test(model)))]); +} + +function unique(values: T[]) { + return Array.from(new Set(values)); +} + function nodeErrorsMessage(nodeErrors: unknown): string | undefined { if (!nodeErrors) return undefined; if (Array.isArray(nodeErrors) && nodeErrors.length === 0) return undefined; diff --git a/commands/tool.test.ts b/commands/tool.test.ts index 7406fbf..590e7e0 100644 --- a/commands/tool.test.ts +++ b/commands/tool.test.ts @@ -5,7 +5,7 @@ import { toolEnterMaskEditCommand, toolEnterTemporaryPanCommand, toolExitMaskEdi const defaultBrush = { color: "#111827", size: 8, hardness: 100 }; const defaultChromaKey = { color: "#00ff00", tolerance: 32, softness: 24, feather: 0, choke: 0, despeckle: 0, spill: 50 }; const defaultMagicWand = { tolerance: 32, feather: 0, choke: 0, despeckle: 0, contiguous: true, mode: "replace" as const }; -const defaultGenerate = { mode: "text-to-image" as const, model: "auto" as const, prompt: "", negativePrompt: "", strength: 75, steps: 30, cfg: 7, seed: -1, sampler: "euler", scheduler: "normal", width: 1024, height: 1024, outpaint: { left: 128, top: 128, right: 128, bottom: 128, feathering: 32 }, inpaint: { maskedAreaOnly: true, cropPadding: 96, maskPolarity: "hidden" as const, maskedContent: "neutral" as const, growMaskBy: 6, maskExpand: 0, maskFeather: 0, maskBlur: 0, maskDespeckle: 0 } }; +const defaultGenerate = { architecture: "sdxl" as const, mode: "text-to-image" as const, model: "auto" as const, textEncoder: "auto", vae: "auto", prompt: "", negativePrompt: "", strength: 75, steps: 30, cfg: 7, seed: -1, sampler: "euler", scheduler: "normal", width: 1024, height: 1024, outpaint: { left: 128, top: 128, right: 128, bottom: 128, feathering: 32 }, inpaint: { maskedAreaOnly: true, cropPadding: 96, maskPolarity: "hidden" as const, maskedContent: "neutral" as const, growMaskBy: 6, maskExpand: 0, maskFeather: 0, maskBlur: 0, maskDespeckle: 0 } }; describe("tool commands", () => { test("sets active tool", () => { @@ -35,6 +35,23 @@ describe("tool commands", () => { expect(next.editor.tools.generate.inpaint).toEqual({ ...defaultGenerate.inpaint, cropPadding: 2048, growMaskBy: 0, maskExpand: -256, maskBlur: 256, maskPolarity: "revealed", maskedContent: "original" }); }); + test("applies architecture defaults and filters unsupported modes", () => { + const initial = toolSetGenerateSettingsCommand.execute({ state: createInitialAppState("Test") }, { mode: "inpaint" }); + const next = toolSetGenerateSettingsCommand.execute({ state: initial }, { architecture: "z-image-turbo" }); + + expect(next.editor.tools.generate).toMatchObject({ + architecture: "z-image-turbo", + mode: "text-to-image", + model: "auto", + textEncoder: "qwen_3_4b.safetensors", + vae: "ae.safetensors", + steps: 8, + cfg: 1, + sampler: "res_multistep", + scheduler: "simple", + }); + }); + test("sets and clears brush preview", () => { const showing = toolSetBrushPreviewCommand.execute({ state: createInitialAppState("Test") }, { position: { x: 10, y: 20 } }); const cleared = toolSetBrushPreviewCommand.execute({ state: showing }, undefined); diff --git a/commands/tool.ts b/commands/tool.ts index 7a406e8..968fa40 100644 --- a/commands/tool.ts +++ b/commands/tool.ts @@ -3,6 +3,7 @@ import type { Vec2D } from "@core/geometry"; import type { LayerId, ArtboardId, AssetId } from "@core/id"; import type { Layer } from "@core/layer"; import type { MaskViewMode } from "@editor/state"; +import { generateArchitectureDefaults } from "@editor/tools"; import type { BrushSettings, ChromaKeySettings, GenerateSettings, MagicWandSettings, ToolId } from "@editor/tools"; import type { Command } from "./command"; import { commandIds } from "./ids"; @@ -62,7 +63,12 @@ export const toolSetGenerateSettingsCommand: Command = { + sdxl: { model: "auto", textEncoder: "auto", vae: "auto", steps: 30, cfg: 7, sampler: "euler", scheduler: "normal", supportedModes: ["text-to-image", "image-to-image", "inpaint", "outpaint"] }, + "z-image": { model: "auto", textEncoder: "qwen_3_4b.safetensors", vae: "ae.safetensors", steps: 30, cfg: 4, sampler: "res_multistep", scheduler: "simple", supportedModes: ["text-to-image"] }, + "z-image-turbo": { model: "auto", textEncoder: "qwen_3_4b.safetensors", vae: "ae.safetensors", steps: 8, cfg: 1, sampler: "res_multistep", scheduler: "simple", supportedModes: ["text-to-image"] }, + anima: { model: "auto", textEncoder: "qwen_3_06b_base.safetensors", vae: "qwen_image_vae.safetensors", steps: 30, cfg: 4, sampler: "er_sde", scheduler: "simple", supportedModes: ["text-to-image"] }, +}; + export type MagicWandSettings = { tolerance: number; feather: number; @@ -87,8 +109,11 @@ export const initialToolState: ToolState = { chromaKey: { color: "#00ff00", tolerance: 32, softness: 24, feather: 0, choke: 0, despeckle: 0, spill: 50 }, magicWand: { tolerance: 32, feather: 0, choke: 0, despeckle: 0, contiguous: true, mode: "replace" }, generate: { + architecture: "sdxl", mode: "text-to-image", model: "auto", + textEncoder: "auto", + vae: "auto", prompt: "", negativePrompt: "", strength: 75, diff --git a/view/bottom-controls/GenerateControls.tsx b/view/bottom-controls/GenerateControls.tsx index bd02b08..c694511 100644 --- a/view/bottom-controls/GenerateControls.tsx +++ b/view/bottom-controls/GenerateControls.tsx @@ -2,10 +2,18 @@ import { useEffect, useRef, useState, type RefObject } from "react"; import { CaretDown, CaretUp } from "@phosphor-icons/react"; import { commandIds } from "@commands/ids"; import type { AppStore } from "@editor/store"; -import type { GenerateMode, GenerateModel, GenerateSettings } from "@editor/tools"; +import { generateArchitectureDefaults } from "@editor/tools"; +import type { GenerateArchitecture, GenerateMode, GenerateModel, GenerateSettings } from "@editor/tools"; import { BottomControlSelectMenu, type BottomControlSelectOption } from "./SelectMenu"; import { BottomControlSlider } from "./Slider"; +const architectures = [ + { value: "sdxl", label: "SDXL" }, + { value: "z-image", label: "Z-Image" }, + { value: "z-image-turbo", label: "Z-Image Turbo" }, + { value: "anima", label: "Anima" }, +] satisfies readonly BottomControlSelectOption[]; + const modes = [ { value: "text-to-image", label: "Text → image" }, { value: "image-to-image", label: "Image → image" }, @@ -38,26 +46,44 @@ export type GenerateControlsProps = { dispatch: AppStore["dispatch"]; }; +type ComfyArchitectureOption = { + value: GenerateArchitecture; + label: string; + defaultModel: string; + models: string[]; + supportedModes: GenerateMode[]; +}; + +type ComfyOptionsResponse = { + architectures?: ComfyArchitectureOption[]; + models?: string[]; + textEncoders?: string[]; + vaes?: string[]; + samplers?: string[]; + schedulers?: string[]; +}; + export function GenerateControls({ settings, dispatch }: GenerateControlsProps) { - const [models, setModels] = useState[]>([{ value: "auto", label: "Auto" }]); - const [samplers, setSamplers] = useState[]>([{ value: settings.sampler, label: settings.sampler }]); - const [schedulers, setSchedulers] = useState[]>([{ value: settings.scheduler, label: settings.scheduler }]); + const [comfyOptions, setComfyOptions] = useState(); const [advancedOpen, setAdvancedOpen] = useState(false); const [outpaintOpen, setOutpaintOpen] = useState(false); const [inpaintOpen, setInpaintOpen] = useState(false); const [sizeOpen, setSizeOpen] = useState(false); const sizeRef = useRef(null); const [error, setError] = useState(); + const modelOptions = resolveModelOptions(settings, comfyOptions); + const supportOptions = resolveSupportOptions(settings, comfyOptions); + const samplerOptions = resolveStringOptions(comfyOptions?.samplers, settings.sampler); + const schedulerOptions = resolveStringOptions(comfyOptions?.schedulers, settings.scheduler); + const modeOptions = resolveModeOptions(settings, comfyOptions); useEffect(() => { let cancelled = false; void fetch("/api/comfy/models") .then((response) => response.ok ? response.json() : Promise.reject(new Error("Unable to load ComfyUI models"))) - .then((body: { models?: string[]; samplers?: string[]; schedulers?: string[] }) => { + .then((body: ComfyOptionsResponse) => { if (cancelled) return; - setModels([{ value: "auto", label: "Auto" }, ...(body.models ?? []).map((model) => ({ value: model, label: model }))]); - if (body.samplers?.length) setSamplers(body.samplers.map((sampler) => ({ value: sampler, label: sampler }))); - if (body.schedulers?.length) setSchedulers(body.schedulers.map((scheduler) => ({ value: scheduler, label: scheduler }))); + setComfyOptions(body); }) .catch((reason: unknown) => { if (!cancelled) setError(reason instanceof Error ? reason.message : "Unable to load ComfyUI models"); @@ -107,7 +133,14 @@ export function GenerateControls({ settings, dispatch }: GenerateControlsProps)
- dispatch(commandIds.toolSetGenerateSettings, { model })} /> + dispatch(commandIds.toolSetGenerateSettings, { architecture })} /> + dispatch(commandIds.toolSetGenerateSettings, { model })} /> + {settings.architecture !== "sdxl" ? ( + <> + dispatch(commandIds.toolSetGenerateSettings, { textEncoder })} /> + dispatch(commandIds.toolSetGenerateSettings, { vae })} /> + + ) : null} dispatch(commandIds.toolSetGenerateSettings, { seed })} />
@@ -121,9 +154,9 @@ export function GenerateControls({ settings, dispatch }: GenerateControlsProps) {advancedOpen ? : }
- dispatch(commandIds.toolSetGenerateSettings, { mode })} /> - dispatch(commandIds.toolSetGenerateSettings, { sampler })} /> - dispatch(commandIds.toolSetGenerateSettings, { scheduler })} /> + dispatch(commandIds.toolSetGenerateSettings, { mode })} /> + dispatch(commandIds.toolSetGenerateSettings, { sampler })} /> + dispatch(commandIds.toolSetGenerateSettings, { scheduler })} />
dispatch(commandIds.toolSetGenerateSettings, { steps })} /> dispatch(commandIds.toolSetGenerateSettings, { cfg })} /> @@ -196,6 +229,38 @@ export function GenerateControls({ settings, dispatch }: GenerateControlsProps) ); } +function resolveModelOptions(settings: GenerateSettings, comfyOptions: ComfyOptionsResponse | undefined): readonly BottomControlSelectOption[] { + const architecture = comfyOptions?.architectures?.find((option) => option.value === settings.architecture); + const models = architecture?.models ?? (settings.architecture === "sdxl" ? comfyOptions?.models : undefined) ?? []; + const fallbackModel = architecture?.defaultModel ?? generateArchitectureDefaults[settings.architecture].model; + const values = unique(["auto", ...models, ...(models.length === 0 && fallbackModel !== "auto" ? [fallbackModel] : []), settings.model]); + return values.map((model) => ({ value: model, label: model === "auto" ? "Auto" : model })); +} + +function resolveSupportOptions(settings: GenerateSettings, comfyOptions: ComfyOptionsResponse | undefined) { + const defaults = generateArchitectureDefaults[settings.architecture]; + return { + textEncoders: resolveStringOptions([...(comfyOptions?.textEncoders ?? []), defaults.textEncoder].filter((value) => value !== "auto"), settings.textEncoder), + vaes: resolveStringOptions([...(comfyOptions?.vaes ?? []), defaults.vae].filter((value) => value !== "auto"), settings.vae), + }; +} + +function resolveStringOptions(values: string[] | undefined, current: string): readonly BottomControlSelectOption[] { + return unique([...(values ?? []), current]).map((value) => ({ value, label: value })); +} + +function resolveModeOptions(settings: GenerateSettings, comfyOptions: ComfyOptionsResponse | undefined): readonly BottomControlSelectOption[] { + const architecture = comfyOptions?.architectures?.find((option) => option.value === settings.architecture); + const supportedModes = architecture?.supportedModes?.length ? architecture.supportedModes : generateArchitectureDefaults[settings.architecture].supportedModes; + const availableModes = modes.filter((mode) => supportedModes.includes(mode.value)); + if (availableModes.some((mode) => mode.value === settings.mode)) return availableModes; + return [modes.find((mode) => mode.value === settings.mode), ...availableModes].filter((mode): mode is BottomControlSelectOption => Boolean(mode)); +} + +function unique(values: T[]): T[] { + return Array.from(new Set(values)); +} + function SectionTitle({ title }: { title: string }) { return
{title}
; } diff --git a/view/generate/runGenerate.ts b/view/generate/runGenerate.ts index 09aeee1..20a2f41 100644 --- a/view/generate/runGenerate.ts +++ b/view/generate/runGenerate.ts @@ -145,8 +145,11 @@ async function requestGenerate(options: { method: "POST", headers: { "content-type": "application/json" }, body: JSON.stringify({ + architecture: options.settings.architecture, mode: options.settings.mode, model: options.settings.model, + textEncoder: options.settings.textEncoder, + vae: options.settings.vae, prompt: options.settings.prompt, negativePrompt: options.settings.negativePrompt, strength: options.settings.strength,