@mcptoolshop/style-dataset-lab 3.0.1 → 3.1.0

package/scripts/batch-generate.js

@@ -15,6 +15,7 @@ import { compileBatch } from '../lib/batch-compiler.js';
 import { getBatchMode } from '../lib/batch-modes.js';
 import {
   createBatchDir,
+  openBatchDirForResume,
   executeBatchRuns,
   saveBatchManifest,
 } from '../lib/batch-runs.js';
@@ -31,20 +32,23 @@ export async function run(argv = process.argv.slice(2)) {
       subject: { type: 'string' },
       theme: { type: 'string' },
       asset: { type: 'string' },
+      resume: { type: 'string' },
       'dry-run': { type: 'boolean' },
       json: { type: 'boolean' },
     },
     deprecated: { game: 'project' },
   });
 
-  if (!flags.mode) {
+  if (!flags.mode && !flags.resume) {
     console.log('Usage: sdlab batch generate --mode <id> [--subject <id>] [--theme <label>] [--project <name>]');
+    console.log('       sdlab batch generate --resume <batch_id> [--project <name>]');
     console.log('');
     console.log('Options:');
-    console.log('  --mode <id>        Batch mode (required)');
+    console.log('  --mode <id>        Batch mode (required unless --resume is given)');
     console.log('  --subject <id>     Subject for subject-driven modes');
     console.log('  --theme <label>    Theme label for environment modes');
     console.log('  --asset <id>       Training asset reference');
+    console.log('  --resume <id>      Resume an existing batch — re-runs only failed/missing slots');
     console.log('  --dry-run          Prepare batch without submitting to ComfyUI');
     console.log('  --json             Output manifest as JSON');
     return;
@@ -53,17 +57,47 @@ export async function run(argv = process.argv.slice(2)) {
   const projectName = flags.project || getProjectName(argv);
   const projectRoot = getProjectRoot(projectName);
   const dryRun = flags['dry-run'] || false;
+  const resumeId = flags.resume || null;
+
+  // Resume short-circuit: load prior manifest and reuse its mode_id + subject/theme
+  // so the user only needs to pass --resume <batch_id>.
+  let batchId;
+  let batchDir;
+  let priorResults = [];
+  let resumedMode = null;
+  let resumedSubject = null;
+  let resumedTheme = null;
+  let resumedAsset = null;
+  if (resumeId) {
+    const opened = openBatchDirForResume(projectRoot, resumeId);
+    batchId = opened.batchId;
+    batchDir = opened.batchDir;
+    priorResults = opened.priorResults;
+    resumedMode = opened.priorManifest.mode_id || null;
+    resumedSubject = opened.priorManifest.subject_id || null;
+    resumedTheme = opened.priorManifest.theme || null;
+    resumedAsset = opened.priorManifest.asset_ref || null;
+  }
+
+  const modeId = flags.mode || resumedMode;
+  const subjectId = flags.subject || resumedSubject;
+  const theme = flags.theme || resumedTheme;
+  const assetRef = flags.asset || resumedAsset;
 
   // Load mode for display
-  const mode = getBatchMode(projectRoot, flags.mode);
+  const mode = getBatchMode(projectRoot, modeId);
 
   console.log(`\x1b[1mstyle-dataset-lab\x1b[0m batch generate`);
   console.log(`  Project: ${projectName}`);
   console.log(`  Mode: ${mode.label} (${mode.mode_id})`);
   console.log(`  Type: ${mode.batch_type}`);
   console.log(`  Slots: ${mode.variant_plan.length}`);
-  if (flags.subject) console.log(`  Subject: ${flags.subject}`);
-  if (flags.theme) console.log(`  Theme: ${flags.theme}`);
+  if (subjectId) console.log(`  Subject: ${subjectId}`);
+  if (theme) console.log(`  Theme: ${theme}`);
+  if (resumeId) {
+    const completed = priorResults.filter(r => r && r.status === 'ok').length;
+    console.log(`  Resume: ${resumeId} (${completed}/${priorResults.length} slots already complete)`);
+  }
   if (dryRun) console.log('  ⚠ DRY RUN');
   console.log('');
 
@@ -71,14 +105,18 @@ export async function run(argv = process.argv.slice(2)) {
   const { baseBrief, slotBriefs } = await compileBatch({
     projectRoot,
     projectId: projectName,
-    modeId: flags.mode,
-    subjectId: flags.subject,
-    theme: flags.theme,
-    assetRef: flags.asset,
+    modeId,
+    subjectId,
+    theme,
+    assetRef,
   });
 
-  // 2. Create batch directory
-  const { batchId, batchDir } = await createBatchDir(projectRoot);
+  // 2. Create or reuse batch directory
+  if (!resumeId) {
+    const created = await createBatchDir(projectRoot);
+    batchId = created.batchId;
+    batchDir = created.batchDir;
+  }
   info('batch-generate', `Batch: ${batchId}`);
 
   // 3. Execute runs
@@ -90,6 +128,7 @@ export async function run(argv = process.argv.slice(2)) {
     mode,
     slotBriefs,
     dryRun,
+    priorResults,
   });
   const batchElapsedMs = Date.now() - batchStartMs;
   if (!dryRun && results.length > 0) {
@@ -111,7 +150,9 @@ export async function run(argv = process.argv.slice(2)) {
       label: r.label,
       brief_id: r.brief_id,
       run_id: r.run_id,
+      status: r.status,
       selected_output: r.selected_output,
+      ...(r.error ? { error: r.error } : {}),
     })),
     sheet_paths: [],
     summary: {
@@ -121,8 +162,10 @@ export async function run(argv = process.argv.slice(2)) {
     },
   };
 
-  if (flags.subject) manifest.subject_id = flags.subject;
-  if (flags.theme) manifest.theme = flags.theme;
+  if (subjectId) manifest.subject_id = subjectId;
+  if (theme) manifest.theme = theme;
+  if (assetRef) manifest.asset_ref = assetRef;
+  if (resumeId) manifest.resumed_from = resumeId;
 
   // 5. Render sheet
   const sheetHtml = renderSheetHTML({ manifest, mode, projectRoot, batchDir });

package/scripts/generate.js

@@ -11,12 +11,14 @@
  */
 
 import { writeFile, mkdir } from "node:fs/promises";
+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 { comfyHealth, submitAndWait, downloadImage } from "../lib/comfyui.js";
+import { pickOutputImage } from "../lib/comfyui-output.js";
 import { info, result } from "../lib/log.js";
 
 /**
@@ -185,11 +187,13 @@ export async function run(argv = process.argv.slice(2)) {
       project: { type: 'string' },
       game: { type: 'string' },
       'dry-run': { type: 'boolean' },
+      resume: { type: 'boolean' },
     },
     deprecated: { game: 'project' },
   });
   const packPath = parsed.positionals[0] || "inputs/prompts/rpg-icons-lane1.json";
   const dryRun = parsed.flags['dry-run'] === true;
+  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'] });
@@ -215,6 +219,7 @@ export async function run(argv = process.argv.slice(2)) {
   const d = pack.defaults;
   let generated = 0;
   let errors = 0;
+  let skipped = 0;
   const totalExpected = pack.subjects.length * pack.variations.length;
   const runStartMs = Date.now();
 
@@ -225,6 +230,20 @@ export async function run(argv = process.argv.slice(2)) {
       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})`);
 
       if (dryRun) {
@@ -233,7 +252,7 @@ export async function run(argv = process.argv.slice(2)) {
       }
 
       const startMs = Date.now();
-      const { nodes } = buildWorkflow(
+      const { nodes, saveNodeId } = buildWorkflow(
         fullPrompt, pack.negative,
         d.checkpoint, loras, seed,
         d.steps, d.cfg, d.sampler, d.scheduler,
@@ -241,29 +260,20 @@ export async function run(argv = process.argv.slice(2)) {
       );
 
       try {
-        const result = await submitAndWait(nodes, COMFY_URL);
+        const submitResult = await submitAndWait(nodes, COMFY_URL);
         const elapsed = Date.now() - startMs;
 
-        // Find output image
-        const outputs = result.outputs || {};
-        let imageFile = null;
-        let imageSubfolder = "";
-        for (const nodeOut of Object.values(outputs)) {
-          if (nodeOut.images && nodeOut.images.length > 0) {
-            imageFile = nodeOut.images[0].filename;
-            imageSubfolder = nodeOut.images[0].subfolder || "";
-            break;
-          }
-        }
-
-        if (!imageFile) {
+        // 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(imageFile, imageSubfolder, COMFY_URL);
+        const imgData = await downloadImage(picked.filename, picked.subfolder, COMFY_URL);
         const destPath = `outputs/candidates/${assetId}.png`;
         await writeFile(join(GAME_ROOT, destPath), imgData);
 
@@ -322,12 +332,14 @@ export async function run(argv = process.argv.slice(2)) {
   }
 
   console.log("");
-  const succeeded = generated - errors;
+  const succeeded = generated - errors - skipped;
   const totalElapsedMs = Date.now() - runStartMs;
-  console.log(`\x1b[32m✓\x1b[0m Generated ${succeeded} candidates (${errors} errors)`);
-  if (!dryRun && generated > 0) {
-    const avgMs = totalElapsedMs / generated;
-    console.log(`  Total: ${formatEta(totalElapsedMs)} — avg ${formatEta(avgMs)}/item`);
+  const skippedSuffix = skipped > 0 ? `, ${skipped} resumed` : '';
+  console.log(`\x1b[32m✓\x1b[0m Generated ${succeeded} candidates (${errors} errors${skippedSuffix})`);
+  if (!dryRun && generated > skipped) {
+    const ran = generated - skipped;
+    const avgMs = totalElapsedMs / ran;
+    console.log(`  Total: ${formatEta(totalElapsedMs)} — avg ${formatEta(avgMs)}/item over ${ran} new`);
   }
   if (dryRun) console.log("  (dry run — no images produced)");
   // Always print the output directory so scripted pipelines can parse it,