draw-things-chat

/**
 * Model Overlays Registry
 * Maps model identifiers to their parameter overlays for txt2img and img2img modes.
 *
 * - "auto": Use defaults (no overlay)
 * - "z-image": Use z-image-turbo optimized settings
 * - "qwen-image": Use qwen-image optimized settings
 * - "flux": Use flux optimized settings
 */

import {
  ImageGenerationParams,
  zImageOverlays,
  qwenImageOverlays,
  fluxOverlays,
} from "../core-bundle.mjs";
import { ltxOverlays } from "./ltxParamsDrawThings.js";

export type ModelId = "auto" | "z-image" | "qwen-image" | "qwen-edit" | "flux" | "ltx" | "custom";

/**
 * All model IDs including sub-categories (e.g. qwen-edit).
 * Used internally for overlay lookups.
 */
export const MODEL_IDS: readonly ModelId[] = [
  "auto",
  "z-image",
  "qwen-image",
  "qwen-edit",
  "flux",
  "ltx",
  "custom",
] as const;

/**
 * User-facing model families (no sub-categories).
 * Used for reports, snapshots, and UI model selection.
 * Sub-categories like "qwen-edit" are resolved internally but not exposed as separate families.
 */
export type ModelFamily = "auto" | "z-image" | "qwen-image" | "flux" | "ltx" | "custom";
export const MODEL_FAMILIES: readonly ModelFamily[] = [
  "auto",
  "z-image",
  "qwen-image",
  "flux",
  "ltx",
  "custom",
] as const;

export interface ModelOverlay {
  txt2img: Partial<ImageGenerationParams> | null;
  img2img: Partial<ImageGenerationParams> | null;
  edit: Partial<ImageGenerationParams> | null;
  txt2vid: Partial<ImageGenerationParams> | null;
  img2vid: Partial<ImageGenerationParams> | null;
}

/**
 * Registry of model overlays.
 * - `null` means "use defaults, no overlay"
 */
export const modelOverlays: Record<ModelId, ModelOverlay> = {
  auto: {
    txt2img: null,
    img2img: null,
    edit: null,
    txt2vid: null,
    img2vid: null,
  },
  "z-image": {
    txt2img: zImageOverlays.txt2img,
    img2img: zImageOverlays.img2img,
    edit: null, // z-image does not support edit mode
    txt2vid: null,
    img2vid: null,
  },
  "qwen-image": {
    txt2img: qwenImageOverlays.txt2img,
    img2img: qwenImageOverlays.img2img,
    edit: qwenImageOverlays.edit,
    txt2vid: null,
    img2vid: null,
  },
  // qwen-edit: dedicated edit/i2i preset using qwen_image_edit models
  // Uses same overlays as qwen-image for img2img/edit (no txt2img support)
  "qwen-edit": {
    txt2img: null, // qwen-edit does not support text2image
    img2img: qwenImageOverlays.img2img,
    edit: qwenImageOverlays.edit,
    txt2vid: null,
    img2vid: null,
  },
  flux: {
    txt2img: fluxOverlays.txt2img,
    img2img: fluxOverlays.img2img,
    edit: fluxOverlays.edit,
    txt2vid: null,
    img2vid: null,
  },
  ltx: {
    txt2img: ltxOverlays.txt2vid,    // LTX can also produce single images
    img2img: ltxOverlays.img2vid,    // LTX can also do image2image
    edit: null,                       // LTX does not support edit mode
    txt2vid: ltxOverlays.txt2vid,
    img2vid: ltxOverlays.img2vid,
  },
  custom: {
    txt2img: null, // custom configs only, no built-in overlay
    img2img: null, // custom configs only, no built-in overlay
    edit: null, // custom configs only, no built-in overlay
    txt2vid: null,
    img2vid: null,
  },
};

/**
 * Get overlay params for a given model and mode.
 * Returns null if no overlay should be applied.
 */
export function getModelOverlay(
  modelId: ModelId | string | undefined,
  mode: "txt2img" | "img2img" | "edit" | "txt2vid" | "img2vid"
): Partial<ImageGenerationParams> | null {
  if (!modelId || modelId === "auto") {
    return null;
  }
  const entry = modelOverlays[modelId as ModelId];
  if (!entry) {
    console.error(`Unknown model '${modelId}', using defaults`);
    return null;
  }
  return entry[mode];
}

/**
 * Check if a string is a valid ModelId
 */
export function isValidModelId(value: unknown): value is ModelId {
  return typeof value === "string" && MODEL_IDS.includes(value as ModelId);
}

/**
 * Get the actual .ckpt filename for a model preset.
 * Returns the filename that should exist on the Draw Things backend.
 * Returns null for "auto" or unknown models.
 */
export function getModelFilename(
  modelId: ModelId | string | undefined,
  mode: "txt2img" | "img2img" | "edit" | "txt2vid" | "img2vid"
): string | null {
  if (!modelId || modelId === "auto") {
    return null;
  }
  const entry = modelOverlays[modelId as ModelId];
  if (!entry) {
    return null;
  }
  const overlay = entry[mode];
  if (!overlay) {
    return null;
  }
  return typeof overlay.model === "string" ? overlay.model : null;
}

/**
 * Get all required .ckpt filenames for a model preset (model + LoRAs).
 * Returns an array of filenames that should exist on the Draw Things backend.
 */
export function getModelRequiredFiles(
  modelId: ModelId | string | undefined,
  mode: "txt2img" | "img2img" | "edit" | "txt2vid" | "img2vid"
): string[] {
  const files: string[] = [];
  if (!modelId || modelId === "auto") {
    return files;
  }
  const entry = modelOverlays[modelId as ModelId];
  if (!entry) {
    return files;
  }
  const overlay = entry[mode];
  if (!overlay) {
    return files;
  }
  if (typeof overlay.model === "string") {
    files.push(overlay.model);
  }
  if (Array.isArray(overlay.loras)) {
    for (const lora of overlay.loras) {
      if (
        typeof lora === "object" &&
        lora !== null &&
        typeof lora.file === "string"
      ) {
        files.push(lora.file);
      }
    }
  }
  return files;
}