@mcptoolshop/style-dataset-lab 3.4.0 → 3.5.0

package/lib/adapters/ai-toolkit.js

@@ -1,5 +1,5 @@
 /**
- * Ostris ai-toolkit adapter.
+ * Ostris ai-toolkit adapter (Flux + Qwen-Image target families, DESIGN-v3 §3.3, §3.5, G12).
  *
  * Consumes canonical rows (per D6) and emits a training package ai-toolkit
  * can consume directly:
@@ -10,14 +10,15 @@
  *     metadata/<partition>.jsonl                   (audit rows)
  *     ai-toolkit-config.yaml                       (training config)
  *
- * Precondition: profile.target_family must be 'flux'. ai-toolkit does support
- * SDXL, but this adapter only emits Flux-shaped configs — callers using SDXL
- * should use diffusers-lora instead. Rejecting here keeps config output honest.
+ * Precondition: profile.target_family must be 'flux' or 'qwen-image'.
+ * Callers using SDXL should use diffusers-lora instead.
  */
 
+import { createHash } from 'node:crypto';
 import { writeFile, mkdir, symlink, copyFile } from 'node:fs/promises';
-import { join, relative } from 'node:path';
-import { existsSync } from 'node:fs';
+import { join, relative, dirname } from 'node:path';
+import { existsSync, readFileSync } from 'node:fs';
+import { fileURLToPath } from 'node:url';
 import YAML from 'yaml';
 import { deriveStyleTrigger } from '../captions.js';
 import { groupRowsByPartition } from '../rows.js';
@@ -25,18 +26,89 @@ import { inputError } from '../errors.js';
 import { assertSafeAssetPath, assertSafeId } from '../export.js';
 import { warn } from '../log.js';
 
+export const AI_TOOLKIT_UPSTREAM_ISSUE = 'https://github.com/ostris/ai-toolkit/issues/1087';
+
+export const SUPPORTED_AI_TOOLKIT_FAMILIES = Object.freeze(['flux', 'qwen-image']);
+
+const OVERRIDE_RELATIVE_PATH = 'scripts/colab/ai-toolkit-override-py315.txt';
+
+const HELD_DATA_TOKENS = Object.freeze([
+  'pls' + 'vessel',
+  'nano' + '_banana',
+  'nano' + '-banana',
+]);
+
+/**
+ * Load metadata and sha256 for `scripts/colab/ai-toolkit-override-py315.txt` (§3.5, G12).
+ *
+ * @returns {{ override_file: string, override_sha256: string, upstream_issue: string, packages: string[] }}
+ */
+export function getAiToolkitOverrideMetadata() {
+  const here = dirname(fileURLToPath(import.meta.url));
+  const fullPath = join(here, '..', '..', 'scripts', 'colab', 'ai-toolkit-override-py315.txt');
+  const raw = readFileSync(fullPath, 'utf-8');
+  const normalized = raw.replace(/\r\n/g, '\n');
+  const sha = 'sha256:' + createHash('sha256').update(normalized, 'utf-8').digest('hex');
+  const packages = normalized
+    .split('\n')
+    .map((l) => l.trim())
+    .filter((l) => l && !l.startsWith('#'));
+  return {
+    override_file: OVERRIDE_RELATIVE_PATH,
+    override_sha256: sha,
+    upstream_issue: AI_TOOLKIT_UPSTREAM_ISSUE,
+    packages,
+  };
+}
+
+/**
+ * Enforce that held Portlight training data (the 336 Nano Banana-derived pairs)
+ * is never packaged or trained on (§1.2, §5.2, G12).
+ *
+ * @param {Object} opts
+ * @param {Object} [opts.profile]
+ * @param {Array<Object>} [opts.rows]
+ * @param {Object} [opts.manifest]
+ */
+export function assertNotHeldTrainingData({ profile = {}, rows = [], manifest = {} } = {}) {
+  const haystacks = [
+    profile?.profile_id,
+    profile?.trigger_override,
+    profile?.label,
+    manifest?.source_export_id,
+    ...(Array.isArray(rows)
+      ? rows.flatMap((r) => [r?.entity_id, r?.caption, r?.asset_path])
+      : []),
+  ]
+    .filter(Boolean)
+    .map((s) => String(s).toLowerCase());
+
+  for (const text of haystacks) {
+    for (const token of HELD_DATA_TOKENS) {
+      if (text.includes(token)) {
+        throw inputError(
+          'AI_TOOLKIT_HELD_DATA_FORBIDDEN',
+          `AI_TOOLKIT_HELD_DATA_FORBIDDEN: Refusing to build training package or config referencing held dataset token "${token}".`
+        );
+      }
+    }
+  }
+}
+
 /**
  * Build training package in ai-toolkit format from canonical rows.
  */
 export async function buildPackage({ packageDir, rows, profile, manifest, config, projectRoot, copy }) {
-  if (profile?.target_family !== 'flux') {
+  if (!SUPPORTED_AI_TOOLKIT_FAMILIES.includes(profile?.target_family)) {
     throw inputError(
       'ADAPTER_TARGET_FAMILY_MISMATCH',
-      `ai-toolkit adapter requires profile.target_family === 'flux' (got "${profile?.target_family}")`,
-      'Use diffusers-lora for SDXL targets, or set target_family: "flux" on the profile.',
+      `ai-toolkit adapter requires profile.target_family === 'flux' or 'qwen-image' (got "${profile?.target_family}")`,
+      'Use diffusers-lora for SDXL targets, or set target_family: "flux" or "qwen-image" on the profile.',
     );
   }
 
+  assertNotHeldTrainingData({ profile, rows, manifest });
+
   const partitioned = groupRowsByPartition(rows || []);
 
   let imageCount = 0;
@@ -132,30 +204,121 @@ export async function buildPackage({ packageDir, rows, profile, manifest, config
   return { imageCount, imagesFailed };
 }
 
+/**
+ * Build the Qwen-Image ai-toolkit YAML config object (§3.3, §3.5, G12).
+ */
+export function buildQwenAiToolkitConfig({ profile, manifest }) {
+  assertNotHeldTrainingData({ profile, manifest });
+  const trigger = deriveStyleTrigger(profile) || profile.profile_id;
+  const baseModel = profile.base_model_recommendations?.[0] || 'Qwen/Qwen-Image';
+  const sourceExportId = manifest?.source_export_id || 'unknown';
+  const overrideMeta = getAiToolkitOverrideMetadata();
+
+  const hp = profile.training_hyperparameters || {};
+  const linear = hp.rank ?? 16;
+  const linearAlpha = hp.alpha ?? linear;
+  const trainSteps = hp.steps ?? 1000;
+  const lr = hp.lr ?? 1e-4;
+
+  return {
+    job: 'extension',
+    config: {
+      name: `${profile.profile_id}-${sourceExportId}`,
+      process: [
+        {
+          type: 'sd_trainer',
+          training_folder: 'output',
+          device: 'cuda:0',
+          trigger_word: trigger,
+          network: {
+            type: 'lora',
+            linear: linear,
+            linear_alpha: linearAlpha,
+          },
+          save: {
+            dtype: 'float16',
+            save_every: 250,
+            max_step_saves_to_keep: 4,
+          },
+          datasets: [
+            {
+              folder_path: './dataset/train',
+              caption_ext: 'txt',
+              caption_dropout_rate: 0.05,
+              shuffle_tokens: false,
+              cache_latents_to_disk: true,
+              resolution: [512, 768, 1024],
+            },
+          ],
+          train: {
+            batch_size: 1,
+            steps: trainSteps,
+            gradient_accumulation_steps: 1,
+            train_unet: true,
+            train_text_encoder: false,
+            gradient_checkpointing: true,
+            noise_scheduler: 'flowmatch',
+            optimizer: 'adamw8bit',
+            lr: lr,
+            ema_config: {
+              use_ema: true,
+              ema_decay: 0.99,
+            },
+            dtype: 'bf16',
+          },
+          model: {
+            name_or_path: baseModel,
+            arch: 'qwen_image',
+            quantize: true,
+            qtype: 'qfloat8',
+            quantize_te: true,
+            qtype_te: 'qfloat8',
+            low_vram: true,
+          },
+          is_style: profile.is_style_lora === true,
+          sample: {
+            sampler: 'flowmatch',
+            sample_every: 250,
+            width: 1024,
+            height: 1024,
+            prompts: [`${trigger} style sample prompt`],
+            neg: '',
+            seed: 42,
+            walk_seed: true,
+            guidance_scale: 3.5,
+            sample_steps: 22,
+          },
+        },
+      ],
+    },
+    meta: {
+      name: profile.profile_id,
+      version: '1.0',
+      target_family: 'qwen-image',
+      base_licence_id: 'Apache-2.0',
+      override_file: overrideMeta.override_file,
+      override_sha256: overrideMeta.override_sha256,
+      upstream_issue: overrideMeta.upstream_issue,
+    },
+  };
+}
+
 /**
  * Build the ai-toolkit YAML config object from a profile + manifest.
  *
- * Defaults follow ai-toolkit's published Flux examples and the 16GB-VRAM
- * ceiling documented in the slice-3 handoff. Users can override any field
- * by editing the emitted YAML before running ai-toolkit.
+ * Dispatches to `buildQwenAiToolkitConfig` when `profile.target_family === 'qwen-image'`;
+ * otherwise emits the Flux-shaped config.
  */
 export function buildAiToolkitConfig({ profile, manifest }) {
+  assertNotHeldTrainingData({ profile, manifest });
+  if (profile?.target_family === 'qwen-image') {
+    return buildQwenAiToolkitConfig({ profile, manifest });
+  }
+
   const trigger = deriveStyleTrigger(profile) || profile.profile_id;
   const baseModel = profile.base_model_recommendations?.[0] || 'black-forest-labs/FLUX.1-dev';
   const sourceExportId = manifest?.source_export_id || 'unknown';
 
-  // Per-profile training_hyperparameters override the World-LoRA-tuned defaults.
-  // World LoRAs: rank 32, steps 2500 (style absorption across a domain).
-  // Per-character LoRAs: rank 16, alpha 8, steps 2000 (identity on a single
-  // subject — higher rank just overfits face details). See the two-LoRA stack
-  // contract in the handbook + memory/two-lora-stack-research-2026-04-23.md.
-  //
-  // TODO (normalization): the implicit World-LoRA fallback is a transition-era
-  // convenience for existing profiles that pre-date training_hyperparameters.
-  // Once every shipping profile declares training_hyperparameters explicitly,
-  // the ?? fallbacks can be removed in favor of a required block — preventing
-  // the "is this World-tuned by omission or by intent?" ambiguity that grows
-  // once a second World project exists.
   const hp = profile.training_hyperparameters || {};
   const linear = hp.rank ?? 32;
   const linearAlpha = hp.alpha ?? linear;
@@ -211,12 +374,6 @@ export function buildAiToolkitConfig({ profile, manifest }) {
             is_flux: true,
             quantize: true,
           },
-          // ai-toolkit's is_style flag: true for world/style LoRAs that teach
-          // aesthetics broadly; false (default) for subject LoRAs that teach a
-          // specific character or object. Drives ai-toolkit's training
-          // regularization — wrong value silently mistrains. See the two-LoRA
-          // stack research (memory/two-lora-stack-research-2026-04-23.md) for
-          // the World LoRA vs per-character LoRA convention.
           is_style: profile.is_style_lora === true,
           sample: {
             sampler: 'flowmatch',

package/lib/adapters/comfyui-workflows.js

@@ -6,7 +6,6 @@
  * that into a ComfyUI node graph.
  */
 
-import { readFile } from 'node:fs/promises';
 import { join } from 'node:path';
 import { existsSync, readdirSync, readFileSync } from 'node:fs';
 import { getRuntimeDir } from '../paths.js';
@@ -72,6 +71,41 @@ export const QWEN_DEFAULTS = Object.freeze({
   height: 1024,
 });
 
+function normalizeLoraSpec(entry) {
+  if (!entry || typeof entry !== 'object') return null;
+  const name = entry.name || entry.lora_name;
+  if (typeof name !== 'string' || name.trim().length === 0) return null;
+  const weight = entry.weight ?? entry.strength_model ?? entry.strength ?? 1.0;
+  return { name: name.trim(), weight: Number(weight) };
+}
+
+/**
+ * Resolve the ordered Qwen-Image LoRA chain (G3) from project defaults, brief,
+ * runtime_plan, and explicit job overrides. Deduplicates by LoRA filename while
+ * preserving the latest weight override.
+ *
+ * @param {Object} [opts]
+ * @param {Object} [opts.brief]
+ * @param {Object} [opts.projectDefaults]
+ * @param {Array} [opts.loras]
+ * @returns {Array<{name: string, weight: number}>}
+ */
+export function resolveQwenLoras({ brief = {}, projectDefaults = {}, loras = [] } = {}) {
+  const sources = [
+    ...(Array.isArray(projectDefaults?.loras) ? projectDefaults.loras : []),
+    ...(Array.isArray(brief?.loras) ? brief.loras : []),
+    ...(Array.isArray(brief?.runtime_plan?.loras) ? brief.runtime_plan.loras : []),
+    ...(Array.isArray(loras) ? loras : []),
+  ];
+  const byName = new Map();
+  for (const raw of sources) {
+    const spec = normalizeLoraSpec(raw);
+    if (!spec) continue;
+    byName.set(spec.name, spec);
+  }
+  return [...byName.values()];
+}
+
 /**
  * Build the Qwen-Image ComfyUI prompt graph.
  *
@@ -118,13 +152,15 @@ export function buildQwenGraph({
   // DiT LoRA chain: 37 → 40 → 41 → … → ModelSamplingAuraFlow.
   let modelOut = ['37', 0];
   let loraId = 40;
-  for (const lora of loras) {
+  for (const rawLora of loras) {
+    const lora = normalizeLoraSpec(rawLora);
+    if (!lora) continue;
     const id = String(loraId++);
     nodes[id] = {
       class_type: 'LoraLoaderModelOnly',
       inputs: {
         lora_name: lora.name,
-        strength_model: lora.weight ?? 1.0,
+        strength_model: lora.weight,
         model: modelOut,
       },
     };
@@ -155,6 +191,43 @@ export function buildQwenGraph({
   return { nodes, saveNodeId: '9' };
 }
 
+/**
+ * Build the minimal 64x64 no-weights ComfyUI fixture graph for Stage 0 CPU
+ * stubs (EmptyImage -> ImageInvert -> SaveImage). Exercises ComfyUI --cpu
+ * execution, extra_pnginfo chunk writing, and /view download without loading
+ * multi-gigabyte diffusion weights.
+ *
+ * @param {Object} [opts]
+ * @param {number} [opts.seed=0]
+ * @param {number} [opts.width=64]
+ * @param {number} [opts.height=64]
+ * @param {string} [opts.filenamePrefix='sdl_stub']
+ * @returns {{ nodes: Object, saveNodeId: string }}
+ */
+export function buildStubFixtureGraph({
+  seed = 0,
+  width = 64,
+  height = 64,
+  filenamePrefix = 'sdl_stub',
+} = {}) {
+  const color = Math.abs(Number(seed) || 0) & 0xffffff;
+  const nodes = {
+    '1': {
+      class_type: 'EmptyImage',
+      inputs: { width, height, batch_size: 1, color },
+    },
+    '2': {
+      class_type: 'ImageInvert',
+      inputs: { image: ['1', 0] },
+    },
+    '3': {
+      class_type: 'SaveImage',
+      inputs: { filename_prefix: filenamePrefix, images: ['2', 0] },
+    },
+  };
+  return { nodes, saveNodeId: '3' };
+}
+
 // ─── Graph builder ───────────────────────────────────────────────
 
 /**
@@ -168,17 +241,19 @@ export function buildQwenGraph({
  * @param {Object} opts.brief — compiled brief from brief-compiler
  * @param {Object} opts.projectDefaults — project.json defaults block
  * @param {number} opts.seed — specific seed for this generation
+ * @param {Array} [opts.loras] — optional explicit LoRA overrides
  * @returns {{ nodes: Object, saveNodeId: string }}
  */
-export function buildWorkflowGraph({ brief, projectDefaults, seed }) {
-  const rp = brief.runtime_plan;
-  const defaults = projectDefaults;
+export function buildWorkflowGraph({ brief, projectDefaults, seed, loras }) {
+  const rp = brief.runtime_plan || {};
+  const defaults = projectDefaults || {};
 
   const prompt = brief.prompt;
   const negativePrompt = brief.negative_prompt;
 
   // ── Qwen-Image branch (non-anime studio base) ──
   if (String(defaults.base || '').toLowerCase() === 'qwen-image') {
+    const resolvedLoras = resolveQwenLoras({ brief, projectDefaults: defaults, loras });
     return buildQwenGraph({
       prompt,
       negativePrompt,
@@ -195,13 +270,13 @@ export function buildWorkflowGraph({ brief, projectDefaults, seed }) {
       vae: defaults.vae || QWEN_DEFAULTS.vae,
       clipType: defaults.clip_type || QWEN_DEFAULTS.clip_type,
       weightDtype: defaults.weight_dtype || QWEN_DEFAULTS.weight_dtype,
-      loras: defaults.loras || [],
+      loras: resolvedLoras,
     });
   }
 
   // ── SDXL branch (default) ──
   const checkpoint = defaults.checkpoint;
-  const loras = defaults.loras || [];
+  const sdxlLoras = loras || defaults.loras || [];
   const width = rp.width || defaults.width || 1024;
   const height = rp.height || defaults.height || 1024;
   const steps = rp.steps || defaults.steps || 30;
@@ -223,7 +298,7 @@ export function buildWorkflowGraph({ brief, projectDefaults, seed }) {
   let clipOut = [ckptId, 1];
 
   // LoRA loaders (chained)
-  for (const lora of loras) {
+  for (const lora of sdxlLoras) {
     const loraId = String(nextId++);
     nodes[loraId] = {
       class_type: 'LoraLoader',