@mcptoolshop/style-dataset-lab 3.1.0 → 3.2.0

package/scripts/generate.js

@@ -15,10 +15,11 @@ import { existsSync } from "node:fs";
 import { join } from "node:path";
 import { parseArgs, getProjectName } from "../lib/args.js";
 import { REPO_ROOT, resolveSafeProjectPath } from "../lib/paths.js";
-import { readJsonFile } from "../lib/config.js";
-import { runtimeError, handleCliError } from "../lib/errors.js";
+import { readJsonFile, loadProjectMeta, resolveGpuModel } from "../lib/config.js";
+import { runtimeError, inputError, handleCliError } from "../lib/errors.js";
 import { comfyHealth, submitAndWait, downloadImage } from "../lib/comfyui.js";
 import { pickOutputImage } from "../lib/comfyui-output.js";
+import { buildQwenGraph, QWEN_DEFAULTS } from "../lib/adapters/comfyui-workflows.js";
 import { info, result } from "../lib/log.js";
 
 /**
@@ -177,6 +178,89 @@ function buildWorkflow(prompt, negativePrompt, checkpoint, loras, seed, steps, c
   return { nodes, saveNodeId: saveId };
 }
 
+/**
+ * Detect the input schema and normalize to a unified shape:
+ *   { lane, style_prefix, negative, defaults, base, jobs:[{assetId, subjectId, variationId, prompt, seed, loras}] }
+ *
+ * Two accepted schemas (both produce jobs in subject-major order, so the legacy
+ * seed sequence stays byte-identical):
+ *
+ *   • Legacy pack — top-level `lane` + `variations` (array) applied to every
+ *     subject. seed = base_seed + jobIndex + (variation.seed_offset||0);
+ *     assetId = `${subject.id}_${variation.id}`; `negative` is top-level.
+ *
+ *   • Wave (what `sdlab init` scaffolds, e.g. example-wave.json) — each subject
+ *     carries `variations` as an integer count; `negative` lives in `defaults`.
+ *     seed = base_seed + jobIndex; assetId = `${subject.id}_v${v}`. This matches
+ *     scripts/qwen_generate.py's seed/filename scheme.
+ *
+ * `base` resolves: input.base → input.defaults.base → project.json defaults.base
+ * → "sdxl". `"qwen-image"` routes to the non-anime Qwen graph.
+ */
+function normalizeGenerationInput(raw, projectMeta) {
+  if (!Array.isArray(raw.subjects) || raw.subjects.length === 0) {
+    throw inputError(
+      'INPUT_MISSING_FIELDS',
+      'Prompt pack/wave must have a non-empty "subjects" array.',
+      'See inputs/prompts/example-wave.json for the expected shape.'
+    );
+  }
+  const defaults = raw.defaults || {};
+  const base = String(
+    raw.base || defaults.base || projectMeta?.defaults?.base || 'sdxl'
+  ).toLowerCase();
+  const baseSeed = Number.isFinite(defaults.base_seed) ? defaults.base_seed : 0;
+  const stylePrefix = raw.style_prefix || '';
+  const compose = (p) => (stylePrefix ? `${stylePrefix}, ${p}` : p);
+
+  const jobs = [];
+  let k = 0;
+
+  if (Array.isArray(raw.variations)) {
+    // Legacy pack schema
+    if (!raw.lane) {
+      throw inputError(
+        'INPUT_MISSING_FIELDS',
+        'Legacy pack (top-level "variations" array) requires a "lane".',
+        'Add a "lane" string, or use the wave schema (per-subject integer "variations").'
+      );
+    }
+    const negative = raw.negative ?? '';
+    for (const subject of raw.subjects) {
+      for (const variation of raw.variations) {
+        jobs.push({
+          assetId: `${subject.id}_${variation.id}`,
+          subjectId: subject.id,
+          variationId: variation.id,
+          prompt: compose(subject.prompt),
+          seed: baseSeed + k + (variation.seed_offset || 0),
+          loras: variation.loras || [],
+        });
+        k++;
+      }
+    }
+    return { lane: raw.lane, style_prefix: stylePrefix, negative, defaults, base, jobs, schema: 'pack' };
+  }
+
+  // Wave schema
+  const negative = raw.negative ?? defaults.negative ?? '';
+  for (const subject of raw.subjects) {
+    const nvar = Number.isFinite(subject.variations) ? subject.variations : 1;
+    for (let v = 0; v < nvar; v++) {
+      jobs.push({
+        assetId: `${subject.id}_v${v}`,
+        subjectId: subject.id,
+        variationId: `v${v}`,
+        prompt: compose(subject.prompt),
+        seed: baseSeed + k,
+        loras: [],
+      });
+      k++;
+    }
+  }
+  return { lane: raw.wave || raw.lane || 'default', style_prefix: stylePrefix, negative, defaults, base, jobs, schema: 'wave' };
+}
+
 export async function run(argv = process.argv.slice(2)) {
   const projectName = getProjectName(argv);
   const GAME_ROOT = join(REPO_ROOT, 'projects', projectName);
@@ -196,13 +280,40 @@ export async function run(argv = process.argv.slice(2)) {
   const resume = parsed.flags['resume'] === true;
 
   const fullPackPath = resolveSafeProjectPath(GAME_ROOT, packPath, { flagName: 'pack-path' });
-  const pack = await readJsonFile(fullPackPath, { requiredFields: ['lane', 'subjects', 'variations', 'defaults'] });
+  const raw = await readJsonFile(fullPackPath);
+  const projectMeta = loadProjectMeta(GAME_ROOT);
+  const gpuModel = resolveGpuModel(projectMeta?.defaults);
+  const pack = normalizeGenerationInput(raw, projectMeta);
+
+  const d = pack.defaults;
+  const base = pack.base;
+  const isQwen = base === 'qwen-image';
+
+  // Resolve effective sampler params once. Qwen applies its own defaults when a
+  // wave omits them; SDXL uses the pack defaults exactly as authored (unchanged).
+  const resolved = isQwen
+    ? {
+        width: d.width || QWEN_DEFAULTS.width,
+        height: d.height || QWEN_DEFAULTS.height,
+        steps: d.steps || QWEN_DEFAULTS.steps,
+        cfg: d.cfg || QWEN_DEFAULTS.cfg,
+        sampler: d.sampler || QWEN_DEFAULTS.sampler,
+        scheduler: d.scheduler || QWEN_DEFAULTS.scheduler,
+      }
+    : {
+        width: d.width,
+        height: d.height,
+        steps: d.steps,
+        cfg: d.cfg,
+        sampler: d.sampler,
+        scheduler: d.scheduler,
+      };
 
   console.log(`\x1b[1mstyle-dataset-lab\x1b[0m generate`);
   console.log(`  Lane: ${pack.lane}`);
-  console.log(`  Subjects: ${pack.subjects.length}`);
-  console.log(`  Variations: ${pack.variations.length}`);
-  console.log(`  Total candidates: ${pack.subjects.length * pack.variations.length}`);
+  console.log(`  Base: ${base}${isQwen ? ' \x1b[2m(non-anime)\x1b[0m' : ''}`);
+  console.log(`  Subjects: ${raw.subjects.length}`);
+  console.log(`  Total candidates: ${pack.jobs.length}`);
   console.log("");
 
   if (!dryRun) {
@@ -216,118 +327,143 @@ export async function run(argv = process.argv.slice(2)) {
   await mkdir(join(GAME_ROOT, "outputs/candidates"), { recursive: true });
   await mkdir(join(GAME_ROOT, "records"), { recursive: true });
 
-  const d = pack.defaults;
   let generated = 0;
   let errors = 0;
   let skipped = 0;
-  const totalExpected = pack.subjects.length * pack.variations.length;
+  const totalExpected = pack.jobs.length;
   const runStartMs = Date.now();
 
-  for (const subject of pack.subjects) {
-    for (const variation of pack.variations) {
-      const assetId = `${subject.id}_${variation.id}`;
-      const seed = (d.base_seed + generated) + (variation.seed_offset || 0);
-      const fullPrompt = `${pack.style_prefix}, ${subject.prompt}`;
-      const loras = variation.loras || [];
-
-      // --resume: skip subjects whose record + image are both already on disk.
-      // Seeds are deterministic per (base_seed, generated, variation), so
-      // increment `generated` for skipped slots to keep downstream seeds stable.
-      if (resume && !dryRun) {
-        const recordPath = join(GAME_ROOT, `records/${assetId}.json`);
-        const imagePath = join(GAME_ROOT, `outputs/candidates/${assetId}.png`);
-        if (existsSync(recordPath) && existsSync(imagePath)) {
-          console.log(`  [${generated + 1}/${totalExpected}] ${assetId} \x1b[2m(resumed — skipped)\x1b[0m`);
-          generated++;
-          skipped++;
-          continue;
-        }
-      }
-
-      console.log(`  [${generated + 1}/${totalExpected}] ${assetId} (seed: ${seed}, loras: ${loras.length})`);
+  for (const job of pack.jobs) {
+    const { assetId, prompt: fullPrompt, seed, loras } = job;
 
-      if (dryRun) {
+    // --resume: skip jobs whose record + image are both already on disk.
+    // Seeds are precomputed per job index, so skipping never shifts later seeds.
+    if (resume && !dryRun) {
+      const recordPath = join(GAME_ROOT, `records/${assetId}.json`);
+      const imagePath = join(GAME_ROOT, `outputs/candidates/${assetId}.png`);
+      if (existsSync(recordPath) && existsSync(imagePath)) {
+        console.log(`  [${generated + 1}/${totalExpected}] ${assetId} \x1b[2m(resumed — skipped)\x1b[0m`);
         generated++;
+        skipped++;
         continue;
       }
+    }
 
-      const startMs = Date.now();
-      const { nodes, saveNodeId } = buildWorkflow(
-        fullPrompt, pack.negative,
-        d.checkpoint, loras, seed,
-        d.steps, d.cfg, d.sampler, d.scheduler,
-        d.width, d.height,
-      );
+    console.log(`  [${generated + 1}/${totalExpected}] ${assetId} (seed: ${seed}, loras: ${isQwen ? (d.loras || []).length : loras.length})`);
 
-      try {
-        const submitResult = await submitAndWait(nodes, COMFY_URL);
-        const elapsed = Date.now() - startMs;
-
-        // Prefer the explicit SaveImage node from buildWorkflow; fall back
-        // to highest numeric node id (see lib/comfyui-output.js).
-        const picked = pickOutputImage(submitResult.outputs, { preferNodeId: saveNodeId });
-        if (!picked) {
-          console.log(`    \x1b[33m⚠\x1b[0m No output image found`);
-          generated++;
-          continue;
-        }
-
-        // Download and save
-        const imgData = await downloadImage(picked.filename, picked.subfolder, COMFY_URL);
-        const destPath = `outputs/candidates/${assetId}.png`;
-        await writeFile(join(GAME_ROOT, destPath), imgData);
-
-        // Write provenance record
-        const record = {
-          id: assetId,
-          schema_version: "1.0.0",
-          created_at: new Date().toISOString(),
-          asset_path: destPath,
-          image: { format: "png", width: d.width, height: d.height, bytes: imgData.length },
-          provenance: {
-            workflow_id: pack.lane,
-            workflow_version: "1.0.0",
-            checkpoint: d.checkpoint,
-            loras,
-            prompt: fullPrompt,
-            negative_prompt: pack.negative,
-            seed,
-            steps: d.steps,
-            cfg: d.cfg,
-            sampler: d.sampler,
-            scheduler: d.scheduler,
-            width: d.width,
-            height: d.height,
-            generation_time_ms: elapsed,
-            gpu_model: "RTX 5080",
-            batch_index: generated,
-          },
-          judgment: null,
-          canon: null,
-          iteration: null,
-        };
-
-        await writeFile(
-          join(GAME_ROOT, `records/${assetId}.json`),
-          JSON.stringify(record, null, 2),
+    if (dryRun) {
+      generated++;
+      continue;
+    }
+
+    const startMs = Date.now();
+    const { nodes, saveNodeId } = isQwen
+      ? buildQwenGraph({
+          prompt: fullPrompt,
+          negativePrompt: pack.negative,
+          seed,
+          width: resolved.width,
+          height: resolved.height,
+          steps: resolved.steps,
+          cfg: resolved.cfg,
+          sampler: resolved.sampler,
+          scheduler: resolved.scheduler,
+          shift: d.shift ?? QWEN_DEFAULTS.shift,
+          unet: d.unet || QWEN_DEFAULTS.unet,
+          clip: d.clip || QWEN_DEFAULTS.clip,
+          vae: d.vae || QWEN_DEFAULTS.vae,
+          clipType: d.clip_type || QWEN_DEFAULTS.clip_type,
+          weightDtype: d.weight_dtype || QWEN_DEFAULTS.weight_dtype,
+          loras: d.loras || [],
+          filenamePrefix: 'sdl',
+        })
+      : buildWorkflow(
+          fullPrompt, pack.negative,
+          d.checkpoint, loras, seed,
+          resolved.steps, resolved.cfg, resolved.sampler, resolved.scheduler,
+          resolved.width, resolved.height,
         );
 
-        console.log(`    \x1b[32m✓\x1b[0m ${destPath} (${elapsed}ms, ${imgData.length} bytes)`);
-      } catch (err) {
-        console.log(`    \x1b[31m✗\x1b[0m ${err.message}`);
-        errors++;
-      }
+    try {
+      const submitResult = await submitAndWait(nodes, COMFY_URL);
+      const elapsed = Date.now() - startMs;
 
-      generated++;
+      // Prefer the explicit SaveImage node from the builder; fall back
+      // to highest numeric node id (see lib/comfyui-output.js).
+      const picked = pickOutputImage(submitResult.outputs, { preferNodeId: saveNodeId });
+      if (!picked) {
+        console.log(`    \x1b[33m⚠\x1b[0m No output image found`);
+        generated++;
+        continue;
+      }
 
-      // ETA every 5 items (not every item — too noisy for long runs).
-      if (!dryRun && generated > 0 && generated % 5 === 0 && generated < totalExpected) {
-        const elapsedTotal = Date.now() - runStartMs;
-        const avgMs = elapsedTotal / generated;
-        const remaining = totalExpected - generated;
-        const etaMs = avgMs * remaining;
-        info('generate', `progress ${generated}/${totalExpected} — avg ${formatEta(avgMs)}/item, ETA ~${formatEta(etaMs)}`);
+      // Download and save
+      const imgData = await downloadImage(picked.filename, picked.subfolder, COMFY_URL);
+      const destPath = `outputs/candidates/${assetId}.png`;
+      await writeFile(join(GAME_ROOT, destPath), imgData);
+
+      // Write provenance record (base-aware: Qwen records unet/clip/vae/shift,
+      // SDXL records checkpoint/loras).
+      const provenance = {
+        workflow_id: pack.lane,
+        workflow_version: "1.0.0",
+        base,
+        prompt: fullPrompt,
+        negative_prompt: pack.negative,
+        seed,
+        steps: resolved.steps,
+        cfg: resolved.cfg,
+        sampler: resolved.sampler,
+        scheduler: resolved.scheduler,
+        width: resolved.width,
+        height: resolved.height,
+        generation_time_ms: elapsed,
+        gpu_model: gpuModel,
+        batch_index: generated,
+      };
+      if (isQwen) {
+        provenance.unet = d.unet || QWEN_DEFAULTS.unet;
+        provenance.clip = d.clip || QWEN_DEFAULTS.clip;
+        provenance.vae = d.vae || QWEN_DEFAULTS.vae;
+        provenance.shift = d.shift ?? QWEN_DEFAULTS.shift;
+        provenance.loras = d.loras || [];
+      } else {
+        provenance.checkpoint = d.checkpoint;
+        provenance.loras = loras;
       }
+
+      const record = {
+        id: assetId,
+        schema_version: "1.0.0",
+        created_at: new Date().toISOString(),
+        asset_path: destPath,
+        image: { format: "png", width: resolved.width, height: resolved.height, bytes: imgData.length },
+        provenance,
+        judgment: null,
+        canon: null,
+        iteration: null,
+      };
+
+      await writeFile(
+        join(GAME_ROOT, `records/${assetId}.json`),
+        JSON.stringify(record, null, 2),
+      );
+
+      console.log(`    \x1b[32m✓\x1b[0m ${destPath} (${elapsed}ms, ${imgData.length} bytes)`);
+    } catch (err) {
+      console.log(`    \x1b[31m✗\x1b[0m ${err.message}`);
+      errors++;
+    }
+
+    generated++;
+
+    // ETA every 5 items (not every item — too noisy for long runs).
+    if (!dryRun && generated > 0 && generated % 5 === 0 && generated < totalExpected) {
+      const elapsedTotal = Date.now() - runStartMs;
+      const avgMs = elapsedTotal / generated;
+      const remaining = totalExpected - generated;
+      const etaMs = avgMs * remaining;
+      info('generate', `progress ${generated}/${totalExpected} — avg ${formatEta(avgMs)}/item, ETA ~${formatEta(etaMs)}`);
     }
   }
 

package/scripts/qwen_generate.py

@@ -0,0 +1,124 @@
+#!/usr/bin/env python3
+"""Studio NON-ANIME generation on Qwen-Image (the verified base) via ComfyUI.
+
+Reads a canon-derived WAVE json ({style_prefix, defaults, subjects:[{id,prompt,variations}]}) and
+generates each subject on Qwen-Image, saving images + a provenance receipt (PIN_PER_STEP) to an
+output dir. This is the bridge generator until sdlab's buildWorkflowGraph() gains native Qwen
+support (tracked follow-on). It is the studio's NON-ANIME path — see memory/feedback_no_anime.md.
+
+All Qwen-Image components are on the rig (diffusion_models/text_encoders/vae). Run with ComfyUI's
+embedded python (has PIL): start the watchdog (E:/AI/training/_watchdog.ps1) before GPU runs.
+
+    python qwen_generate.py --wave <wave.json> --out <dir> [--variations N]
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+import time
+import urllib.request
+import uuid
+from pathlib import Path
+
+from PIL import Image
+
+COMFY = "http://127.0.0.1:8188"
+COMFY_OUT = Path(r"E:/AI-Models/ComfyUI_windows_portable/ComfyUI/output")
+
+# Qwen-Image native components (on rig). NON-ANIME base — never an anime checkpoint.
+QWEN_UNET = "qwen_image_fp8_e4m3fn.safetensors"
+QWEN_CLIP = "qwen_2.5_vl_7b_fp8_scaled.safetensors"
+QWEN_VAE = "qwen_image_vae.safetensors"
+
+
+def qwen_graph(pos: str, neg: str, seed: int, w: int, h: int, steps: int, cfg: float, prefix: str,
+               loras: list | None = None) -> dict:
+    """loras: [{name, weight}] chained as DiT-only LoraLoaderModelOnly nodes (ids from
+    "40", between UNETLoader and ModelSamplingAuraFlow) — mirrors sdlab buildQwenGraph.
+    Model-only because studio Qwen style LoRAs train the transformer with the text
+    encoder frozen."""
+    nodes = {
+        "37": {"class_type": "UNETLoader", "inputs": {"unet_name": QWEN_UNET, "weight_dtype": "default"}},
+        "38": {"class_type": "CLIPLoader", "inputs": {"clip_name": QWEN_CLIP, "type": "qwen_image"}},
+        "39": {"class_type": "VAELoader", "inputs": {"vae_name": QWEN_VAE}},
+        "6": {"class_type": "CLIPTextEncode", "inputs": {"text": pos, "clip": ["38", 0]}},
+        "7": {"class_type": "CLIPTextEncode", "inputs": {"text": neg, "clip": ["38", 0]}},
+    }
+    model_out = ["37", 0]
+    for i, lora in enumerate(loras or []):
+        lid = str(40 + i)
+        nodes[lid] = {"class_type": "LoraLoaderModelOnly", "inputs": {
+            "lora_name": lora["name"], "strength_model": lora.get("weight", 1.0), "model": model_out}}
+        model_out = [lid, 0]
+    nodes.update({
+        "66": {"class_type": "ModelSamplingAuraFlow", "inputs": {"shift": 3.1, "model": model_out}},
+        "5": {"class_type": "EmptySD3LatentImage", "inputs": {"width": w, "height": h, "batch_size": 1}},
+        "3": {"class_type": "KSampler", "inputs": {"seed": seed, "steps": steps, "cfg": cfg, "sampler_name": "euler",
+              "scheduler": "simple", "denoise": 1.0, "model": ["66", 0], "positive": ["6", 0], "negative": ["7", 0], "latent_image": ["5", 0]}},
+        "8": {"class_type": "VAEDecode", "inputs": {"samples": ["3", 0], "vae": ["39", 0]}},
+        "9": {"class_type": "SaveImage", "inputs": {"filename_prefix": prefix, "images": ["8", 0]}},
+    })
+    return nodes
+
+
+def _post(path: str, payload: dict):
+    return json.loads(urllib.request.urlopen(urllib.request.Request(
+        f"{COMFY}{path}", json.dumps(payload).encode(), {"Content-Type": "application/json"}), timeout=90).read())
+
+
+def generate(pos: str, neg: str, seed: int, dst: Path, **kw) -> Path:
+    g = qwen_graph(pos, neg, seed, kw["w"], kw["h"], kw["steps"], kw["cfg"], dst.stem, loras=kw.get("loras"))
+    pid = _post("/prompt", {"prompt": g, "client_id": uuid.uuid4().hex})["prompt_id"]
+    for _ in range(1200):
+        hist = json.loads(urllib.request.urlopen(f"{COMFY}/history/{pid}", timeout=30).read())
+        if pid in hist and hist[pid].get("outputs"):
+            img = hist[pid]["outputs"]["9"]["images"][0]
+            src = COMFY_OUT / img["subfolder"] / img["filename"]
+            for _w in range(120):
+                if src.exists():
+                    Image.open(src).convert("RGB").save(dst)
+                    return dst
+                time.sleep(0.5)
+        time.sleep(0.5)
+    raise RuntimeError("ComfyUI generation timed out")
+
+
+def main() -> None:
+    ap = argparse.ArgumentParser(description=__doc__)
+    ap.add_argument("--wave", required=True, type=Path)
+    ap.add_argument("--out", required=True, type=Path)
+    ap.add_argument("--variations", type=int, default=None, help="override per-subject variation count")
+    args = ap.parse_args()
+
+    wave = json.loads(args.wave.read_text(encoding="utf-8"))
+    prefix = wave.get("style_prefix", "")
+    d = wave.get("defaults", {})
+    neg = d.get("negative", "")
+    w, h = d.get("width", 1024), d.get("height", 1024)
+    steps, cfg = d.get("steps", 22), d.get("cfg", 3.5)
+    base_seed = d.get("base_seed", 1000)
+    loras = d.get("loras", [])
+    args.out.mkdir(parents=True, exist_ok=True)
+
+    receipt = {"wave": wave.get("wave"), "base": "qwen-image", "unet": QWEN_UNET, "clip": QWEN_CLIP,
+               "vae": QWEN_VAE, "sampler": "euler", "scheduler": "simple", "shift": 3.1,
+               "steps": steps, "cfg": cfg, "size": [w, h], "style_prefix": prefix, "negative": neg,
+               "loras": loras, "items": []}
+    seed = base_seed
+    for subj in wave["subjects"]:
+        nvar = args.variations if args.variations is not None else subj.get("variations", 1)
+        pos = f"{prefix}, {subj['prompt']}" if prefix else subj["prompt"]
+        for v in range(nvar):
+            dst = args.out / f"{subj['id']}_v{v}.png"
+            generate(pos, neg, seed, dst, w=w, h=h, steps=steps, cfg=cfg, loras=loras)
+            print(f"  {subj['id']} v{v} (seed {seed}) -> {dst.name}", flush=True)
+            receipt["items"].append({"id": subj["id"], "variation": v, "seed": seed, "file": dst.name, "prompt": pos})
+            seed += 1
+
+    (args.out / "generation.json").write_text(json.dumps(receipt, indent=2), encoding="utf-8")
+    print(f"\n{len(receipt['items'])} images -> {args.out}  (+ generation.json provenance)")
+
+
+if __name__ == "__main__":
+    main()