feat: enhance Comfy API handling and add tests for branded non-Z diffusion models

app/comfy.test.ts

@@ -1,5 +1,5 @@
 import { describe, expect, test } from "bun:test";
-import { buildAnimaWorkflow, buildSdxlWorkflow, buildZImageTurboWorkflow, buildZImageWorkflow, selectGeneratedOutputImage } from "./comfy";
+import { buildAnimaWorkflow, buildSdxlWorkflow, buildZImageTurboWorkflow, buildZImageWorkflow, handleComfyApi, selectGeneratedOutputImage } from "./comfy";
 
 describe("Comfy adapter", () => {
   test("selects SaveImage output instead of uploaded input or mask images", () => {
@@ -85,6 +85,44 @@ describe("Comfy adapter", () => {
     expect(workflow["2"]?.inputs.clip_name).toBe("miaomiaoHarem_anima13_txt.safetensors");
     expect(workflow["3"]?.inputs.vae_name).toBe("qwen_image_vae.safetensors");
   });
+
+  test("lists branded non-Z diffusion models under Anima", async () => {
+    const originalFetch = globalThis.fetch;
+    const mockFetch: typeof fetch = Object.assign(async () => new Response(JSON.stringify({
+      CheckpointLoaderSimple: { input: { required: { ckpt_name: [["sd_xl_base_1.0.safetensors"]] } } },
+      KSampler: { input: { required: { sampler_name: [["euler"]], scheduler: [["normal"]] } } },
+      UNETLoader: {
+        input: {
+          required: {
+            unet_name: [[
+              "anima-base-v1.0.safetensors",
+              "miaomiaoHarem_anima13.safetensors",
+              "novaAnimeAM_v30.safetensors",
+              "z_image_bf16.safetensors",
+              "z_image_turbo_bf16.safetensors",
+            ]],
+          },
+        },
+      },
+      CLIPLoader: { input: { required: { clip_name: [["qwen_3_06b_base.safetensors"]] } } },
+      VAELoader: { input: { required: { vae_name: [["qwen_image_vae.safetensors"]] } } },
+    }), { headers: { "content-type": "application/json" } }), { preconnect: originalFetch.preconnect });
+    globalThis.fetch = mockFetch;
+
+    try {
+      const response = await handleComfyApi(new Request("http://image-studio.test/api/comfy/models"));
+      const body = await response.json() as { architectures: { value: string; models: string[] }[] };
+      const anima = body.architectures.find((architecture) => architecture.value === "anima");
+
+      expect(response.status).toBe(200);
+      expect(anima?.models).toContain("novaAnimeAM_v30.safetensors");
+      expect(anima?.models).toContain("miaomiaoHarem_anima13.safetensors");
+      expect(anima?.models).not.toContain("z_image_bf16.safetensors");
+      expect(anima?.models).not.toContain("z_image_turbo_bf16.safetensors");
+    } finally {
+      globalThis.fetch = originalFetch;
+    }
+  });
 });
 
 function inpaintRequest(inpaint: { maskedContent: "neutral" | "original"; growMaskBy?: number }) {

app/comfy.ts

@@ -111,7 +111,7 @@ async function listGenerationOptions() {
         value: "anima",
         label: "Anima",
         defaultModel: defaultModels.anima,
-        models: modelsForArchitecture(diffusionModels, defaultModels.anima, [/anima/i]),
+        models: modelsForArchitecture(diffusionModels, defaultModels.anima, [], { exclude: [/z[_-]?image/i] }),
         supportedModes: ["text-to-image"],
       },
     ],
@@ -227,11 +227,12 @@ export function buildSdxlWorkflow(request: ComfyGenerateRequest): Workflow {
   const scheduler = request.scheduler ?? "normal";
   const positive = request.prompt;
   const negative = request.negativePrompt ?? "";
+  const samplerInputs: Record<string, unknown> = { seed, steps, cfg, sampler_name: sampler, scheduler, denoise, model: ["1", 0], positive: ["2", 0], negative: ["3", 0], latent_image: ["5", 0] };
   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] } },
-    "6": { class_type: "KSampler", inputs: { seed, steps, cfg, sampler_name: sampler, scheduler, denoise, model: ["1", 0], positive: ["2", 0], negative: ["3", 0], latent_image: ["5", 0] } },
+    "6": { class_type: "KSampler", inputs: samplerInputs },
     "7": { class_type: "VAEDecode", inputs: { samples: ["6", 0], vae: ["1", 2] } },
     "8": { class_type: "SaveImage", inputs: { filename_prefix: `image-studio-${request.mode}`, images: ["7", 0] } },
   };
@@ -250,7 +251,7 @@ export function buildSdxlWorkflow(request: ComfyGenerateRequest): Workflow {
       workflow["5"] = { class_type: "VAEEncode", inputs: { pixels: ["4", 0], vae: ["1", 2] } };
       workflow["12"] = { class_type: "GrowMask", inputs: { mask: ["11", 0], expand: resolveGrowMaskBy(request), tapered_corners: true } };
       workflow["13"] = { class_type: "SetLatentNoiseMask", inputs: { samples: ["5", 0], mask: ["12", 0] } };
-      workflow["6"].inputs.latent_image = ["13", 0];
+      samplerInputs.latent_image = ["13", 0];
     } else {
       workflow["5"] = { class_type: "VAEEncodeForInpaint", inputs: { pixels: ["4", 0], vae: ["1", 2], mask: ["11", 0], grow_mask_by: resolveGrowMaskBy(request) } };
     }
@@ -410,8 +411,13 @@ function architectureLabel(architecture: GenerateArchitecture): string {
   }
 }
 
-function modelsForArchitecture(models: string[], defaultModel: string, matchers: RegExp[]) {
-  return unique([defaultModel, ...models.filter((model) => model === defaultModel || matchers.some((matcher) => matcher.test(model)))]);
+function modelsForArchitecture(models: string[], defaultModel: string, matchers: RegExp[], options: { exclude?: RegExp[] } = {}) {
+  return unique([defaultModel, ...models.filter((model) => {
+    if (model === defaultModel) return true;
+    if (options.exclude?.some((matcher) => matcher.test(model))) return false;
+    if (matchers.length === 0) return true;
+    return matchers.some((matcher) => matcher.test(model));
+  })]);
 }
 
 function unique<T>(values: T[]) {