draw-things-chat

import type { ImageGenerationParams } from "../core-bundle.mjs";
import * as flatbuffers from "flatbuffers";
import {
  GenerationConfiguration,
  CompressionMethod,
  SamplerType,
  SeedMode,
  LoRA,
  LoRAMode,
} from "../core-bundle.mjs";

export type DtConfigBuild = {
  view: Record<string, unknown>;
  bytes: Buffer;
};

export function buildDtGenerationConfiguration(
  p: Partial<ImageGenerationParams>
): DtConfigBuild {
  const w = typeof p.width === "number" ? p.width : 576;
  const h = typeof p.height === "number" ? p.height : 768;
  const start_width = Math.max(1, Math.round(w / 64));
  const start_height = Math.max(1, Math.round(h / 64));

  // Sampler: Accept number (SamplerType enum) directly, or string for legacy
  let samplerEnum = SamplerType.UniPCTrailing; // Default
  let samplerName = "UniPCTrailing";

  if (typeof p.sampler === "number") {
    // Direct number from custom_configs.json or explicit API call
    samplerEnum = p.sampler;
    // Reverse lookup for logging
    samplerName =
      Object.keys(SamplerType).find(
        (k) => (SamplerType as any)[k] === p.sampler
      ) || "Unknown";
  } else if (typeof p.sampler === "string") {
    // Legacy string format (e.g., "UniPC Trailing")
    samplerName = p.sampler.replace(/\s+/g, "");
    samplerEnum =
      (SamplerType as any)[samplerName] ?? SamplerType.UniPCTrailing;
  }
  const seedModeEnum =
    (SeedMode as any)[
      String(p.seed_mode || "Scale Alike").replace(/\s+/g, "")
    ] ?? SeedMode.ScaleAlike;

  const view: Record<string, unknown> = {
    id: 1,
    start_width,
    start_height,
    seed: typeof p.seed === "number" ? p.seed : undefined,
    steps: typeof p.steps === "number" ? p.steps : undefined,
    sampler: samplerName,
    batch_count:
      typeof (p as any).batch_count === "number"
        ? (p as any).batch_count
        : undefined,
    batch_size:
      typeof (p as any).batch_size === "number"
        ? Math.max(1, Math.min(4, Math.round((p as any).batch_size)))
        : 1,
    model: p.model || undefined,
    seed_mode: String(p.seed_mode || "Scale Alike").replace(/\s+/g, ""),
    t5_text: p.t5_text || undefined,
    clip_l_text: p.clip_l_text || undefined,
    image_guidance_scale:
      typeof p.image_guidance === "number" ? p.image_guidance : undefined,
    tiled_decoding: !!p.tiled_decoding,
    decoding_tile_width:
      typeof p.decoding_tile_width === "number"
        ? Math.max(1, Math.round(p.decoding_tile_width / 64))
        : undefined,
    decoding_tile_height:
      typeof p.decoding_tile_height === "number"
        ? Math.max(1, Math.round(p.decoding_tile_height / 64))
        : undefined,
    decoding_tile_overlap:
      typeof p.decoding_tile_overlap === "number"
        ? Math.max(0, Math.round(p.decoding_tile_overlap / 64))
        : undefined,
    compressionArtifacts: (p as any).compressionArtifacts ?? (p as any).compression_artifacts ?? undefined,
    compressionArtifactsQuality:
      typeof (p as any).compressionArtifactsQuality === "number"
        ? (p as any).compressionArtifactsQuality
        : typeof (p as any).compression_artifacts_quality === "number"
          ? (p as any).compression_artifacts_quality
          : undefined,
    num_frames: typeof p.num_frames === "number" ? p.num_frames : undefined,
  };

  // Cast: video and core may resolve flatbuffers from different node_modules
  // (file: symlink). At runtime there is only one copy — the cast is safe.
  const builder: any = new flatbuffers.Builder(2048);

  // Strings
  const modelOffset = p.model ? builder.createString(p.model) : 0;
  const clipLTextOffset = p.clip_l_text
    ? builder.createString(p.clip_l_text)
    : 0;
  const t5TextOffset = p.t5_text ? builder.createString(p.t5_text) : 0;
  const openClipGTextOffset = (p as any).open_clip_g_text
    ? builder.createString((p as any).open_clip_g_text)
    : 0;

  // LoRAs
  let lorasVector = 0;
  if (Array.isArray(p.loras) && p.loras.length) {
    const loff: number[] = [];
    for (const l of p.loras) {
      const f = l.file ? builder.createString(l.file) : 0;
      const wt = typeof l.weight === "number" ? l.weight : 1.0;
      loff.push(LoRA.createLoRA(builder, f, wt, LoRAMode.All));
    }
    lorasVector = GenerationConfiguration.createLorasVector(builder, loff);
  }

  // Start object
  GenerationConfiguration.startGenerationConfiguration(builder);
  GenerationConfiguration.addId(builder, BigInt(1));
  GenerationConfiguration.addStartWidth(builder, start_width);
  GenerationConfiguration.addStartHeight(builder, start_height);
  if (typeof p.seed === "number")
    GenerationConfiguration.addSeed(builder, p.seed);
  if (typeof p.steps === "number")
    GenerationConfiguration.addSteps(builder, p.steps);
  if (typeof p.guidance_scale === "number")
    GenerationConfiguration.addGuidanceScale(builder, p.guidance_scale);
  if (typeof p.strength === "number")
    GenerationConfiguration.addStrength(builder, p.strength);
  if (modelOffset) GenerationConfiguration.addModel(builder, modelOffset);
  GenerationConfiguration.addSampler(builder, samplerEnum);
  if (typeof (p as any).batch_count === "number") {
    GenerationConfiguration.addBatchCount(
      builder,
      (p as any).batch_count as number as any
    );
  }
  // Respect requested batch_size (clamped 1..4) to control variants; batch_count remains 1
  const batchSize =
    typeof (p as any).batch_size === "number"
      ? Math.max(1, Math.min(4, Math.round((p as any).batch_size)))
      : 1;
  GenerationConfiguration.addBatchSize(builder, batchSize as any);
  if (typeof p.image_guidance === "number")
    GenerationConfiguration.addImageGuidanceScale(builder, p.image_guidance);
  GenerationConfiguration.addSeedMode(builder, seedModeEnum);
  if (typeof p.clip_skip === "number")
    GenerationConfiguration.addClipSkip(builder, p.clip_skip);
  if (lorasVector) GenerationConfiguration.addLoras(builder, lorasVector);
  if (typeof p.mask_blur === "number")
    GenerationConfiguration.addMaskBlur(builder, p.mask_blur);
  if (typeof p.clip_weight === "number")
    GenerationConfiguration.addClipWeight(builder, p.clip_weight);
  if (typeof p.image_prior_steps === "number")
    GenerationConfiguration.addImagePriorSteps(builder, p.image_prior_steps);
  if (typeof p.negative_original_height === "number")
    GenerationConfiguration.addNegativeOriginalImageHeight(
      builder,
      p.negative_original_height
    );
  if (typeof p.negative_original_width === "number")
    GenerationConfiguration.addNegativeOriginalImageWidth(
      builder,
      p.negative_original_width
    );
  if (typeof p.fps === "number")
    GenerationConfiguration.addFpsId(builder, p.fps as any);
  if (typeof p.motion_scale === "number")
    GenerationConfiguration.addMotionBucketId(builder, p.motion_scale);
  if (typeof p.start_frame_guidance === "number")
    GenerationConfiguration.addStartFrameCfg(builder, p.start_frame_guidance);
  if (typeof p.num_frames === "number")
    GenerationConfiguration.addNumFrames(builder, p.num_frames);
  if (typeof p.sharpness === "number")
    GenerationConfiguration.addSharpness(builder, p.sharpness);
  if (typeof p.shift === "number")
    GenerationConfiguration.addShift(builder, p.shift);
  if (typeof p.tiled_decoding === "boolean")
    GenerationConfiguration.addTiledDecoding(builder, !!p.tiled_decoding);
  if (typeof p.decoding_tile_width === "number")
    GenerationConfiguration.addDecodingTileWidth(
      builder,
      Math.max(1, Math.round(p.decoding_tile_width / 64))
    );
  if (typeof p.decoding_tile_height === "number")
    GenerationConfiguration.addDecodingTileHeight(
      builder,
      Math.max(1, Math.round(p.decoding_tile_height / 64))
    );
  if (typeof p.decoding_tile_overlap === "number")
    GenerationConfiguration.addDecodingTileOverlap(
      builder,
      Math.max(0, Math.round(p.decoding_tile_overlap / 64))
    );
  if (typeof p.stochastic_sampling_gamma === "number")
    GenerationConfiguration.addStochasticSamplingGamma(
      builder,
      p.stochastic_sampling_gamma
    );
  if (typeof p.tiled_diffusion === "boolean")
    GenerationConfiguration.addTiledDiffusion(builder, !!p.tiled_diffusion);
  if (clipLTextOffset)
    GenerationConfiguration.addClipLText(builder, clipLTextOffset);
  if (openClipGTextOffset)
    GenerationConfiguration.addOpenClipGText(builder, openClipGTextOffset);
  if (typeof p.speed_up_with_guidance_embed === "boolean")
    GenerationConfiguration.addSpeedUpWithGuidanceEmbed(
      builder,
      !!p.speed_up_with_guidance_embed
    );
  if (typeof p.guidance_embed === "number")
    GenerationConfiguration.addGuidanceEmbed(builder, p.guidance_embed as any);
  if (typeof p.resolution_dependent_shift === "boolean")
    GenerationConfiguration.addResolutionDependentShift(
      builder,
      !!p.resolution_dependent_shift
    );
  if (typeof p.tea_cache_start === "number")
    GenerationConfiguration.addTeaCacheStart(builder, p.tea_cache_start);
  if (typeof p.tea_cache_end === "number")
    GenerationConfiguration.addTeaCacheEnd(builder, p.tea_cache_end);
  if (typeof p.tea_cache_threshold === "number")
    GenerationConfiguration.addTeaCacheThreshold(
      builder,
      p.tea_cache_threshold
    );
  if (typeof p.tea_cache === "boolean")
    GenerationConfiguration.addTeaCache(builder, !!p.tea_cache);
  if (typeof p.separate_t5 === "boolean")
    GenerationConfiguration.addSeparateT5(builder, !!p.separate_t5);
  if (t5TextOffset) GenerationConfiguration.addT5Text(builder, t5TextOffset);
  if (typeof p.tea_cache_max_skip_steps === "number")
    GenerationConfiguration.addTeaCacheMaxSkipSteps(
      builder,
      p.tea_cache_max_skip_steps
    );
  if (typeof p.causal_inference === "number")
    GenerationConfiguration.addCausalInference(builder, p.causal_inference);
  if (typeof p.causal_inference_pad === "number")
    GenerationConfiguration.addCausalInferencePad(
      builder,
      p.causal_inference_pad
    );
  if (typeof p.cfg_zero_star === "boolean")
    GenerationConfiguration.addCfgZeroStar(builder, !!p.cfg_zero_star);
  if (typeof p.cfg_zero_init_steps === "number")
    GenerationConfiguration.addCfgZeroInitSteps(builder, p.cfg_zero_init_steps);
  const caValue = (p as any).compressionArtifacts ?? (p as any).compression_artifacts;
  if (typeof caValue === "string") {
    const ca = caValue;
    const caEnum =
      ca === "h264" ? CompressionMethod.H264
      : ca === "h265" ? CompressionMethod.H265
      : ca === "jpeg" ? CompressionMethod.Jpeg
      : CompressionMethod.Disabled;
    GenerationConfiguration.addCompressionArtifacts(builder, caEnum);
  }
  const caQualityValue = (p as any).compressionArtifactsQuality ?? (p as any).compression_artifacts_quality;
  if (typeof caQualityValue === "number")
    GenerationConfiguration.addCompressionArtifactsQuality(builder, caQualityValue);

  const root = GenerationConfiguration.endGenerationConfiguration(builder);
  GenerationConfiguration.finishGenerationConfigurationBuffer(builder, root);
  const bytes = Buffer.from(builder.asUint8Array());

  return { view, bytes };
}