feat: enhance generate settings to support multiple architectures and their defaults

2026-07-05 11:21:56 +02:00
parent dd0c1df730
commit 0b6b064085
7 changed files with 435 additions and 52 deletions
--- 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<GenerateArchitecture>[];
+
 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<readonly BottomControlSelectOption<GenerateModel>[]>([{ value: "auto", label: "Auto" }]);
-  const [samplers, setSamplers] = useState<readonly BottomControlSelectOption<string>[]>([{ value: settings.sampler, label: settings.sampler }]);
-  const [schedulers, setSchedulers] = useState<readonly BottomControlSelectOption<string>[]>([{ value: settings.scheduler, label: settings.scheduler }]);
+  const [comfyOptions, setComfyOptions] = useState<ComfyOptionsResponse>();
  const [advancedOpen, setAdvancedOpen] = useState(false);
  const [outpaintOpen, setOutpaintOpen] = useState(false);
  const [inpaintOpen, setInpaintOpen] = useState(false);
  const [sizeOpen, setSizeOpen] = useState(false);
  const sizeRef = useRef<HTMLDivElement>(null);
  const [error, setError] = useState<string>();
+  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)

      <section className={panelSectionClass()}>
        <SectionTitle title="Essentials" />
-        <PanelSelect label="Model" value={settings.model} options={models} ariaLabel="Generate model" onValueChange={(model) => dispatch(commandIds.toolSetGenerateSettings, { model })} />
+        <PanelSelect label="Backend" value={settings.architecture} options={architectures} ariaLabel="Generate backend" onValueChange={(architecture) => dispatch(commandIds.toolSetGenerateSettings, { architecture })} />
+        <PanelSelect label="Model" value={settings.model} options={modelOptions} ariaLabel="Generate model" onValueChange={(model) => dispatch(commandIds.toolSetGenerateSettings, { model })} />
+        {settings.architecture !== "sdxl" ? (
+          <>
+            <PanelSelect label="Text enc." value={settings.textEncoder} options={supportOptions.textEncoders} ariaLabel="Generate text encoder" onValueChange={(textEncoder) => dispatch(commandIds.toolSetGenerateSettings, { textEncoder })} />
+            <PanelSelect label="VAE" value={settings.vae} options={supportOptions.vaes} ariaLabel="Generate VAE" onValueChange={(vae) => dispatch(commandIds.toolSetGenerateSettings, { vae })} />
+          </>
+        ) : null}
        <SizeControl refRoot={sizeRef} open={sizeOpen} setOpen={setSizeOpen} settings={settings} dispatch={dispatch} />
        <PanelNumber label="Seed" aria-label="Generate seed" min={-1} max={Number.MAX_SAFE_INTEGER} value={settings.seed} onValueChange={(seed) => dispatch(commandIds.toolSetGenerateSettings, { seed })} />
      </section>
@@ -121,9 +154,9 @@ export function GenerateControls({ settings, dispatch }: GenerateControlsProps)
          {advancedOpen ? <CaretUp size={18} weight="bold" /> : <CaretDown size={18} weight="bold" />}
        </button>
        <div id="generate-advanced-controls" className={`grid gap-2 overflow-hidden transition-all duration-200 ${advancedOpen ? "max-h-[32rem] pt-2 opacity-100" : "max-h-0 opacity-0"}`}>
-          <PanelSelect label="Mode" value={settings.mode} options={modes} ariaLabel="Generate mode" onValueChange={(mode) => dispatch(commandIds.toolSetGenerateSettings, { mode })} />
-          <PanelSelect label="Sampler" value={settings.sampler} options={samplers} ariaLabel="Generate sampler" onValueChange={(sampler) => dispatch(commandIds.toolSetGenerateSettings, { sampler })} />
-          <PanelSelect label="Scheduler" value={settings.scheduler} options={schedulers} ariaLabel="Generate scheduler" onValueChange={(scheduler) => dispatch(commandIds.toolSetGenerateSettings, { scheduler })} />
+          <PanelSelect label="Mode" value={settings.mode} options={modeOptions} ariaLabel="Generate mode" onValueChange={(mode) => dispatch(commandIds.toolSetGenerateSettings, { mode })} />
+          <PanelSelect label="Sampler" value={settings.sampler} options={samplerOptions} ariaLabel="Generate sampler" onValueChange={(sampler) => dispatch(commandIds.toolSetGenerateSettings, { sampler })} />
+          <PanelSelect label="Scheduler" value={settings.scheduler} options={schedulerOptions} ariaLabel="Generate scheduler" onValueChange={(scheduler) => dispatch(commandIds.toolSetGenerateSettings, { scheduler })} />
          <div className="grid grid-cols-2 gap-2">
            <PanelNumber label="Steps" aria-label="Generate steps" value={settings.steps} onValueChange={(steps) => dispatch(commandIds.toolSetGenerateSettings, { steps })} />
            <PanelNumber label="CFG" aria-label="Generate CFG" value={settings.cfg} onValueChange={(cfg) => dispatch(commandIds.toolSetGenerateSettings, { cfg })} />
@@ -196,6 +229,38 @@ export function GenerateControls({ settings, dispatch }: GenerateControlsProps)
  );
 }

+function resolveModelOptions(settings: GenerateSettings, comfyOptions: ComfyOptionsResponse | undefined): readonly BottomControlSelectOption<GenerateModel>[] {
+  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<string>[] {
+  return unique([...(values ?? []), current]).map((value) => ({ value, label: value }));
+}
+
+function resolveModeOptions(settings: GenerateSettings, comfyOptions: ComfyOptionsResponse | undefined): readonly BottomControlSelectOption<GenerateMode>[] {
+  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<GenerateMode> => Boolean(mode));
+}
+
+function unique<T>(values: T[]): T[] {
+  return Array.from(new Set(values));
+}
+
 function SectionTitle({ title }: { title: string }) {
  return <div className="px-1 text-xs font-semibold uppercase tracking-[0.18em] text-white/35">{title}</div>;
 }