image-studio/server/comfy.test.ts

import { describe, expect, test } from "bun:test";
import { buildAnimaWorkflow, buildSdxlWorkflow, buildSemanticSelectionWorkflow, buildZImageTurboWorkflow, buildZImageWorkflow, selectGeneratedOutputImage, selectGeneratedOutputImages } from "./comfy";
import { handleComfyApi } from "./comfy-routes";

describe("Comfy adapter", () => {
  test("selects SaveImage output instead of uploaded input or mask images", () => {
    const image = selectGeneratedOutputImage({
      outputs: {
        "4": { images: [{ filename: "image-studio-input.png", type: "input" }] },
        "9": { images: [{ filename: "image-studio-mask.png", type: "input" }] },
        "8": { images: [{ filename: "image-studio-inpaint_00001_.png", subfolder: "", type: "output" }] },
      },
    });

    expect(image).toEqual({ filename: "image-studio-inpaint_00001_.png", subfolder: "", type: "output" });
  });

  test("falls back to generated filename prefixes when node ids differ", () => {
    const image = selectGeneratedOutputImage({
      outputs: {
        "12": { images: [{ filename: "image-studio-inpaint_00002_.png", type: "output" }] },
        "4": { images: [{ filename: "image-studio-input.png", type: "input" }] },
      },
    });

    expect(image?.filename).toBe("image-studio-inpaint_00002_.png");
  });

  test("retains every generated image from a batch", () => {
    expect(selectGeneratedOutputImages({ outputs: { "8": { images: [
      { filename: "image-studio-inpaint_1.png", type: "output" },
      { filename: "image-studio-inpaint_2.png", type: "output" },
    ] } } })).toHaveLength(2);
  });

  test("builds neutral inpaint with VAEEncodeForInpaint", () => {
    const workflow = buildSdxlWorkflow(inpaintRequest({ maskedContent: "neutral" }));

    expect(workflow["5"]?.class_type).toBe("VAEEncodeForInpaint");
    expect(workflow["5"]?.inputs).toMatchObject({ grow_mask_by: 6, mask: ["11", 0] });
    expect(workflow["6"]?.inputs.latent_image).toEqual(["5", 0]);
  });

  test("builds original-content inpaint with a latent noise mask", () => {
    const workflow = buildSdxlWorkflow(inpaintRequest({ maskedContent: "original", growMaskBy: 12 }));

    expect(workflow["5"]?.class_type).toBe("VAEEncode");
    expect(workflow["12"]?.class_type).toBe("GrowMask");
    expect(workflow["12"]?.inputs).toMatchObject({ mask: ["11", 0], expand: 12 });
    expect(workflow["13"]?.class_type).toBe("SetLatentNoiseMask");
    expect(workflow["13"]?.inputs).toMatchObject({ samples: ["5", 0], mask: ["12", 0] });
    expect(workflow["6"]?.inputs.latent_image).toEqual(["13", 0]);
  });

  test("repeats inpaint latents for batches and applies Canny ControlNet", () => {
    const workflow = buildSdxlWorkflow({
      ...inpaintRequest({ maskedContent: "neutral" }),
      batchSize: 4,
      inpaint: { maskedContent: "neutral", structureControl: "canny", controlModel: "controlnet-canny.safetensors", controlStrength: 0.6 },
    });
    expect(workflow["19"]).toMatchObject({ class_type: "RepeatLatentBatch", inputs: { amount: 4 } });
    expect(workflow["20"]).toMatchObject({ class_type: "ControlNetLoader", inputs: { control_net_name: "controlnet-canny.safetensors" } });
    expect(workflow["21"]?.class_type).toBe("Canny");
    expect(workflow["22"]).toMatchObject({ class_type: "ControlNetApplyAdvanced", inputs: { strength: 0.6 } });
    expect(workflow["6"]?.inputs).toMatchObject({ latent_image: ["19", 0], positive: ["22", 0], negative: ["22", 1] });
  });

  test("adds an optional low-denoise detail pass", () => {
    const workflow = buildSdxlWorkflow({ ...inpaintRequest({ maskedContent: "neutral" }), refinePass: true, refineStrength: 18, seed: 40 });
    expect(workflow["30"]).toMatchObject({ class_type: "KSampler", inputs: { seed: 41, denoise: 0.18, latent_image: ["6", 0] } });
    expect(workflow["7"]?.inputs.samples).toEqual(["30", 0]);
  });

  test("builds native SAM3 point selection as a mask output", () => {
    const workflow = buildSemanticSelectionWorkflow({ inputImage: "input.png", model: "sam3.safetensors", x: 24.4, y: 18.6 });
    expect(workflow["1"]).toMatchObject({ class_type: "UNETLoader", inputs: { unet_name: "sam3.safetensors" } });
    expect(workflow["3"]).toMatchObject({ class_type: "SAM3_Detect", inputs: { positive_coords: '[{"x":24,"y":19}]', refine_iterations: 2 } });
    expect(workflow["4"]).toMatchObject({ class_type: "MaskToImage", inputs: { mask: ["3", 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");
  });

  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 }) {
  return {
    mode: "inpaint" as const,
    model: "model.safetensors",
    prompt: "replace garment",
    width: 128,
    height: 128,
    inputImage: "input.png",
    maskImage: "mask.png",
    inpaint,
  };
}

function textRequest(overrides: Partial<Parameters<typeof buildSdxlWorkflow>[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,
  };
}