From 3367b963bb3f3922f6ff7f353f2dd6cc3d6c75ec Mon Sep 17 00:00:00 2001 From: EnragedAntelope Date: Thu, 8 Oct 2026 11:56:04 -0400 Subject: [PATCH] 0.18.0: epochs, validation pack, held-out photos, Fizgig target MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - ⑤ counts epochs (one checkpoint + one sample set each); Repeats 0 = auto, sized to n x 75 target steps; Fizgig auto = 1 - ⑤ live exposure line (images x repeats / batch = steps/epoch x epochs) with low/high warnings - every config writes validation/: 8 fixed prompts at seed 42, validation.md, per-epoch score sheet - ai-toolkit saves/samples once per epoch, keeps every epoch, seed 42 with walk_seed off - musubi/kohya/Fizgig sample every epoch from the prompt file; musubi Krea 2 samples via --turbo_dit - Fizgig trainer target (Krea 2, Qwen Image 2.1 at 704, Klein) with its cache/train commands - ① writes a dd_source provenance PNG chunk; ⑤ names upscaled sources; ④ records provenance in metadata.json - ④ "Hold out N reference photos" (+ cli export --holdout) copies the largest real photos to -heldout/ - dataset naming skips names whose -heldout sibling exists - version 0.18.0; AGENTS.md, ARCHITECTURE.md, README updated --- AGENTS.md | 8 +- README.md | 12 +- app.py | 157 ++++++++---- cli.py | 6 +- docs/ARCHITECTURE.md | 92 +++++-- studio/__init__.py | 2 +- studio/dataset_stats.py | 76 ++++-- studio/package.py | 62 ++++- studio/preprocess.py | 14 +- studio/trainer_configs.py | 365 ++++++++++++++++++++++------ tests/test_dataset_stats.py | 14 +- tests/test_identity_sync.py | 2 +- tests/test_provenance_heldout.py | 145 +++++++++++ tests/test_review_ideas.py | 8 +- tests/test_shot_style.py | 12 +- tests/test_trainer_configs.py | 78 +++++- tests/test_trainer_configs_0_4_0.py | 22 +- tests/test_trainer_configs_kohya.py | 11 +- tests/test_ui_wiring_0_16_0.py | 7 +- 19 files changed, 879 insertions(+), 214 deletions(-) create mode 100644 tests/test_provenance_heldout.py diff --git a/AGENTS.md b/AGENTS.md index 846cd80..9035ad8 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -8,9 +8,9 @@ Turn a character, style, or concept into a ready-to-train LoRA dataset. One refe _Last verified: 2026-10-08_ -- **Status:** in active development; main is v0.17.3 (git tag `v0.17.3`). **Every version bump must ship a GitHub Release** or the in-app update check never fires. The project was renamed twice (lora-dataset-studio → lora-distillery → dataset-deviser); older references under the previous names are stale. -- **Works:** all five stages end to end (preprocess → generate & curate → caption → export → train config); the three dataset types (character, style, concept) with their own shot plans and caption framing; local and cloud paths for every stage; gallery curation with selection carried forward between stages, including shift-click range select; advisory dedupe, quality flags and caption lint; trainer configs for ai-toolkit (incl. SDXL), kohya sd-scripts and musubi-tuner; opt-in private Hugging Face publish. Since 0.15.0: per-image fault isolation in ① (one bad source never ends the batch), a cooperative ⏹ Stop across ①/②/③ with documented resume paths, translated ComfyUI failures (missing node, bad input, unreachable server) and an in-app 🩺 setup check. Since 0.15.1: a ② **shot style** (default = keep the reference's own medium; eight presets + custom text) threaded into both prompt builders, the CLI, ④ metadata and ⑤'s sample prompt, a final-prompt preview, and a prev/next/save-and-next caption editor with the image on screen. Since 0.16.0: **one name box and one trigger box in the header** that ②/③/④/⑤ all read (they used to be per-tab and went stale), ④ stating the pair it stamped, an opt-in ④ **hand-off to [Idiot LoRa Builder](https://github.com/Fablestarexpanse/Idiot-Lora-Builder)** (writes its `.lora-studio/ratings.json` so its grid opens pre-triaged — nothing is launched), a ⑤ advisory when the dataset can't fill the training resolution, every bundled ComfyUI model filename overridable from `.env`, and a `cli build` preflight for the local engine. Since 0.17.0: the local ② engine is **Qwen-Image 2.1** (plain-English prompts shared with Gemini, ~26 s/shot, `LDS_QWEN21_*` filenames) instead of Qwen-Image-Edit 2511 + the Multiple-Angles LoRA, whose graph stays in `studio/comfy_workflows/legacy/`; a dead ComfyUI mid-batch now fails one shot, not the run. Since 0.17.1: the UI port is chosen by `_pick_port` in `app.py` — 7861 if bindable, else an OS-assigned free port (Windows reserves a ~500-port block around 7861), or `LDS_PORT` to pin one. Since 0.17.2: `comfy_api.upload_image` downscales a source over 90 MB before upload (ComfyUI answers 413 above its 100 MB `--max-upload-size`). Since 0.17.3: the cloud image model defaults to **`auto`** (Google shut down the `-preview` ids on 2026-06-25; `resolve_image_model` picks from the cached live list, Nano Banana 2.1 first), Groq captions use `qwen/qwen3.8-27b` (key `groq-qwen`, old key aliased), every API image goes through `config.api_image` (true mime, size cap), musubi commands include the cache steps and the per-architecture network module, ai-toolkit gains Qwen-Image 2.1 and a correct Krea 2 (Raw) preset and keeps every checkpoint, ① applies EXIF rotation, no longer 4x-upscales large JPEGs, and can jump a busy ComfyUI queue like ② (`--front` on `preprocess`/`build`). -- **In progress:** a phased LoRA best-practice update. 0.18.0 is next: epochs as the ⑤ unit with an exposure line, a validation pack, held-out references, a Fizgig trainer target, source-size provenance in a PNG chunk. Then 0.19.0 (training copies keep their background, identity policy, framing tiers, coverage panel, local multi-reference) and 0.20.0 (a Quick build tab). Nothing from those is half-built on main. +- **Status:** in active development; main is v0.18.0 (git tag `v0.18.0`). **Every version bump must ship a GitHub Release** or the in-app update check never fires. The project was renamed twice (lora-dataset-studio → lora-distillery → dataset-deviser); older references under the previous names are stale. +- **Works:** all five stages end to end (preprocess → generate & curate → caption → export → train config); the three dataset types (character, style, concept) with their own shot plans and caption framing; local and cloud paths for every stage; gallery curation with selection carried forward between stages, including shift-click range select; advisory dedupe, quality flags and caption lint; trainer configs for ai-toolkit (incl. SDXL), kohya sd-scripts, musubi-tuner and Fizgig; opt-in private Hugging Face publish. Since 0.15.0: per-image fault isolation in ① (one bad source never ends the batch), a cooperative ⏹ Stop across ①/②/③ with documented resume paths, translated ComfyUI failures (missing node, bad input, unreachable server) and an in-app 🩺 setup check. Since 0.15.1: a ② **shot style** (default = keep the reference's own medium; eight presets + custom text) threaded into both prompt builders, the CLI, ④ metadata and ⑤'s sample prompt, a final-prompt preview, and a prev/next/save-and-next caption editor with the image on screen. Since 0.16.0: **one name box and one trigger box in the header** that ②/③/④/⑤ all read (they used to be per-tab and went stale), ④ stating the pair it stamped, an opt-in ④ **hand-off to [Idiot LoRa Builder](https://github.com/Fablestarexpanse/Idiot-Lora-Builder)** (writes its `.lora-studio/ratings.json` so its grid opens pre-triaged — nothing is launched), a ⑤ advisory when the dataset can't fill the training resolution, every bundled ComfyUI model filename overridable from `.env`, and a `cli build` preflight for the local engine. Since 0.17.0: the local ② engine is **Qwen-Image 2.1** (plain-English prompts shared with Gemini, ~26 s/shot, `LDS_QWEN21_*` filenames) instead of Qwen-Image-Edit 2511 + the Multiple-Angles LoRA, whose graph stays in `studio/comfy_workflows/legacy/`; a dead ComfyUI mid-batch now fails one shot, not the run. Since 0.17.1: the UI port is chosen by `_pick_port` in `app.py` — 7861 if bindable, else an OS-assigned free port (Windows reserves a ~500-port block around 7861), or `LDS_PORT` to pin one. Since 0.17.2: `comfy_api.upload_image` downscales a source over 90 MB before upload (ComfyUI answers 413 above its 100 MB `--max-upload-size`). Since 0.17.3: the cloud image model defaults to **`auto`** (Google shut down the `-preview` ids on 2026-06-25; `resolve_image_model` picks from the cached live list, Nano Banana 2.1 first), Groq captions use `qwen/qwen3.8-27b` (key `groq-qwen`, old key aliased), every API image goes through `config.api_image` (true mime, size cap), musubi commands include the cache steps and the per-architecture network module, ai-toolkit gains Qwen-Image 2.1 and a correct Krea 2 (Raw) preset and keeps every checkpoint, ① applies EXIF rotation, no longer 4x-upscales large JPEGs, and can jump a busy ComfyUI queue like ② (`--front` on `preprocess`/`build`). Since 0.18.0: ⑤ counts **epochs** (one checkpoint and one sample set each) with auto repeats sized to a step target and a live exposure line; every config ships a `validation/` pack (8 fixed prompts at seed 42, `validation.md`, a per-epoch score sheet); **Fizgig** is a ⑤ target (Krea 2 / Qwen Image 2.1 / Klein); ① stamps a `dd_source` provenance PNG chunk so ⑤ names upscaled sources and ④ knows real photos; ④ can hold out N real photos in a `-heldout` sibling as a likeness test. +- **In progress:** a phased LoRA best-practice update. 0.19.0 is next: training copies keep their background, identity policy (outfits vary by default), framing tiers, a coverage panel, a full-body warning, near-duplicate/burst flags, an anchor shot, local multi-reference. Then 0.20.0 (a Quick build tab). Nothing from those is half-built on main. The 0.18.0 optimizer dropdown and attention-only toggle were deferred (ARCHITECTURE.md → Future ideas). - **Known gaps / next steps:** Auto's model order (`IMAGE_MODEL_PREFERENCE`, Nano Banana 2.1 then Pro) rests on one A/B subject — see ARCHITECTURE.md; local ② still sends only the first reference. Otherwise pick up from `docs/ARCHITECTURE.md` → "Future ideas" / "Deferred", and read its *Maintainer principles* before changing anything; `gradio` is pinned `<6` because of a real stuck-loading-overlay regression — unpinning needs that verified fixed upstream; training is never launched for you and is not planned to be; the default export `output_root` writes into the repo root, so a smoke-test export leaves an untracked dataset folder to delete before committing. - **Deep docs:** `docs/ARCHITECTURE.md` (module map, stage and backend detail, gotchas, backlog — the deep reference), `docs/comfyui-setup.md`. Worklogs under `docs/` are gitignored and local-only by design. @@ -21,7 +21,7 @@ _Last verified: 2026-10-08_ - **Local OR cloud per stage.** Every step has a free/private local path AND a no-GPU cloud path. Mix and match: generate on cloud, caption on GPU, or reverse. - **Multiple backends:** ComfyUI (generation, preprocessing), built-in SAM3 (isolation), ONNX booru tagger (WD + e621), Qwen3-VL-8B / JoyCaption / NSFW finetune (local captioning), Gemini Flash / Groq (cloud captioning), LM Studio / Ollama / any OpenAI-compatible endpoint. - **Curate by clicking.** Toggle shots from the gallery with ✅/⬜ marks. Selection carries through stages: what you keep in ② is what ③ captions, and what ③ captions is what ④ preselects. -- **Trainer configs:** Generated for ai-toolkit (incl. SDXL), kohya sd-scripts, musubi-tuner. Training is NEVER launched for you — configs are shown, you run them. +- **Trainer configs:** Generated for ai-toolkit (incl. SDXL), kohya sd-scripts, musubi-tuner, Fizgig, plus a validation pack. Training is NEVER launched for you — configs are shown, you run them. - **Honest by design.** No telemetry, no upsell. Cloud calls bill YOUR key; this tool takes no cut. ## Layout diff --git a/README.md b/README.md index 3caeeea..a4848be 100644 --- a/README.md +++ b/README.md @@ -6,7 +6,7 @@ settings; for a **concept**, 18 shots around one object — a turnaround plus fr For a **style**, bring your own images and get smart, correctly-framed captions. Either way the output is a flat folder that drops straight into any trainer (**ai-toolkit / kohya / OneTrainer / …**), plus a ready-to-edit **training config** for -ai-toolkit, kohya or musubi. +ai-toolkit, kohya, musubi or Fizgig, and a validation pack for picking the best checkpoint. ![Generate & curate tab](docs/images/ui-generate.png) @@ -48,7 +48,7 @@ Run them in order (each step auto-fills the next) or jump straight to the one yo | ② Generate shots *(character + concept)* | ComfyUI: Qwen-Image 2.1 *(Qwen Research License)* | Gemini (Nano Banana) | | ③ Caption | Qwen3-VL-8B, JoyCaption, NSFW finetune, **WD + e621 taggers**, LM Studio / Ollama / any OpenAI endpoint | Gemini Flash, Groq free tier | | ④ Export | always local (+ optional **.zip** and **Hugging Face** publish) | — | -| ⑤ Train config | ai-toolkit (incl. SDXL) / **kohya sd-scripts** / musubi-tuner | — | +| ⑤ Train config | ai-toolkit (incl. SDXL) / **kohya sd-scripts** / musubi-tuner / **Fizgig** | — | ## Dataset types @@ -230,6 +230,14 @@ Both have a switch in ②, because both bit us: *is* the character). Best fix: **isolate the source in ①** — SAM3 keeps bags and held objects out of the subject mask. The "Exclude props" toggle in ② helps too, especially on Gemini. +## Picking the best checkpoint + +⑤ sets training in **epochs**: every epoch saves a checkpoint and renders the same 8 test +prompts at the same seed. Leave **Repeats** at 0 and it sizes each epoch for you; the exposure +line shows the total. The `validation/` folder it writes holds the prompts, a short guide and a +score sheet. To test likeness fairly, set **Hold out N reference photos** in ④: those photos +stay out of training, in a `-heldout` folder beside the dataset, to compare samples against. + ## CLI Every stage is a standalone subcommand; `--help` shows all options. diff --git a/app.py b/app.py index 2da456b..3ac9885 100644 --- a/app.py +++ b/app.py @@ -1215,7 +1215,7 @@ def load_export_preview(folders_text: str, dup_distance: float = 5, carry=None): def do_export(selected: list[str], name: str, trigger: str, output_root: str, make_zip: bool = False, dataset_type: str = "character", style_key: str = shot_style.MATCH, style_text: str = "", - ilb_handoff: bool = False): + ilb_handoff: bool = False, holdout=0): if not selected: raise gr.Error("Click '📂 Load & preview', then keep at least one image checked.") from studio.package import package_dataset, resolve_export_items @@ -1238,7 +1238,8 @@ def do_export(selected: list[str], name: str, trigger: str, output_root: str, "skipped_empty_caption": res.empties} out_root = _validate_out_dir(output_root) try: - ds = package_dataset(res.items, out_root, name, trigger, metadata) + ds = package_dataset(res.items, out_root, name, trigger, metadata, + holdout=int(holdout or 0)) except OSError as e: raise gr.Error(f"Couldn't write the dataset to '{out_root}': {e}. Check the " f"output folder path (valid drive, no forbidden characters, writable).") from e @@ -1282,9 +1283,16 @@ def do_export(selected: list[str], name: str, trigger: str, output_root: str, identity += ("\n⚠️ No trigger word — captions have nothing to teach the LoRA " "to respond to. Set one at the top and re-caption if that wasn't " "deliberate.") - result = (f"✅ Dataset ready: {ds} ({len(res.items)} image/caption pairs from the " - f"{checked} image(s) you checked)" - f"{identity}{skipped}{empty_note}{zip_note}{ilb_note}{sample_block}") + held_dir = ds.parent / f"{ds.name}-heldout" + held = list_images(held_dir) + held_note = (f"\n🔒 Held out {len(held)} photo(s) as a likeness test, never trained " + f"on: {held_dir}" if held else "") + if int(holdout or 0) and not held: + held_note = ("\n⚠️ Nothing held out — only ①'s photos qualify, and none of the " + "checked images came from ①.") + result = (f"✅ Dataset ready: {ds} ({len(res.items) - len(held)} image/caption pairs " + f"from the {checked} image(s) you checked)" + f"{identity}{held_note}{skipped}{empty_note}{zip_note}{ilb_note}{sample_block}") # ds path auto-fills the ⑤ Train tab AND the HF-publish box below. return result, str(ds), str(ds) @@ -1401,12 +1409,47 @@ def on_trainer_change(trainer: str): p = TRAINER_MODELS[trainer][0] return (_model_dropdown(trainer), user_config.get_trainer_path(trainer), - p.resolution, p.rank, p.alpha, p.steps, p.lr, p.batch_size) + p.resolution, p.rank, p.alpha, p.epochs, p.lr, p.batch_size) def on_model_change(trainer: str, model_key: str): p = _preset(trainer, model_key) - return p.resolution, p.rank, p.alpha, p.steps, p.lr, p.batch_size + return p.resolution, p.rank, p.alpha, p.epochs, p.lr, p.batch_size + + +def _repeats(trainer: str, stats, epochs, batch_size, repeats) -> int: + """The typed repeats, or (0 = auto) the count that reaches the target steps.""" + if int(repeats or 0) > 0: + return int(repeats) + if trainer == "fizgig": # Fizgig's own guidance: repeats 1, more epochs + return 1 + return stats.suggested_repeats(int(epochs or 1), int(batch_size or 1)) + + +def exposure_line(trainer: str, dataset_dir: str, epochs, repeats, batch_size) -> str: + """Live "how much training is this" line under ⑤'s epochs/repeats fields.""" + from studio.dataset_stats import inspect + from studio.trainer_configs import exposure + + ds = Path(dataset_dir.strip()) if dataset_dir.strip() else None + if not ds or not ds.is_dir(): + return "" + stats = inspect(ds) + if not stats.n_images: + return "" + reps = _repeats(trainer, stats, epochs, batch_size, repeats) + batch, n_epochs = max(int(batch_size or 1), 1), max(int(epochs or 1), 1) + per_epoch, total = exposure(stats.n_images, reps, batch, n_epochs) + auto = " (auto)" if not int(repeats or 0) else "" + line = (f"**Exposure:** {stats.n_images} images × {reps} repeats{auto} ÷ batch {batch} " + f"= {per_epoch} steps/epoch × {n_epochs} epochs = **{total} steps**, saving " + f"{n_epochs} checkpoints with samples.") + # Fizgig's adaptive LR is tuned for small sets at repeats 1, so few steps is its norm. + if total < 400 and trainer != "fizgig": + line += " ⚠️ Few steps — likely undertrained; raise epochs or repeats." + elif total > 6000: + line += " ⚠️ Many steps — slow, and late epochs will likely overfit." + return line def save_trainer_path(trainer: str, path: str) -> str: @@ -1448,31 +1491,31 @@ def dataset_type_note(ds: Path, selected_type: str) -> str: return line -def inspect_dataset(dataset_dir: str, dataset_type: str = "character") -> tuple[str, gr.Number]: - """Read the dataset and suggest a step count derived from its image count.""" +def inspect_dataset(dataset_dir: str, dataset_type: str = "character") -> str: + """Summarize the dataset: count, sizes, captions and the target step count.""" from studio.dataset_stats import inspect if not dataset_dir.strip(): - return "", gr.Number() + return "" ds = Path(dataset_dir.strip()) if not ds.is_dir(): - return f"⚠️ Folder not found: {ds}", gr.Number() + return f"⚠️ Folder not found: {ds}" try: stats = inspect(ds) except Exception as e: # unreadable/corrupt image headers — report, never crash - return f"⚠️ Couldn't inspect {ds}: {e}", gr.Number() + return f"⚠️ Couldn't inspect {ds}: {e}" if not stats.n_images: - return f"⚠️ No images in {ds}", gr.Number() - return stats.summary() + dataset_type_note(ds, dataset_type), \ - gr.Number(value=stats.suggested_steps) + return f"⚠️ No images in {ds}" + return stats.summary() + dataset_type_note(ds, dataset_type) def do_generate_train_config(trainer: str, model_key: str, dataset_dir: str, install_path: str, name: str, trigger: str, - resolution, rank, alpha, steps, lr, batch_size, + resolution, rank, alpha, epochs, lr, batch_size, multi_res: bool, dataset_type: str = "character", style_key: str = shot_style.MATCH, - style_text: str = "", project_name: str = "") -> str: + style_text: str = "", project_name: str = "", + repeats=0) -> str: # ⑤'s "LoRA name" is the trained file's name, not the subject's — the one # identity-ish field that is legitimately its own (people want "-v2"). Blank # means "follow the header name", which is why it is not auto-filled: an @@ -1501,33 +1544,36 @@ def do_generate_train_config(trainer: str, model_key: str, dataset_dir: str, raise gr.Error(f"No images found in {ds} — export a dataset first (④).") preset = _preset(trainer, model_key) buckets = stats.buckets_for(int(resolution)) if multi_res else [] + num_repeats = _repeats(trainer, stats, epochs, batch_size, repeats) cfg = TrainConfig( trainer=trainer, model=preset, dataset_dir=ds, trigger=trigger.strip(), name=(name.strip() or "lora"), dataset_type=dataset_type, resolution=int(resolution), rank=int(rank), alpha=int(alpha), - steps=int(steps), lr=float(lr), batch_size=int(batch_size), + epochs=max(1, int(epochs)), num_repeats=num_repeats, n_images=stats.n_images, + lr=float(lr), batch_size=int(batch_size), buckets=buckets, shot_style=style_key, shot_style_text=style_text) try: - written, command = write_configs(cfg, install_path.strip(), - num_repeats=max(1, round(400 / stats.n_images))) + written, command = write_configs(cfg, install_path.strip()) except OSError as e: raise gr.Error(f"Couldn't write the config into {ds}: {e}. Check the dataset " f"folder is writable.") from e user_config.set_last_train_settings({ "trainer": trainer, "model": model_key, "resolution": int(resolution), - "rank": int(rank), "alpha": int(alpha), "steps": int(steps), - "lr": float(lr), "batch_size": int(batch_size)}) + "rank": int(rank), "alpha": int(alpha), "epochs": cfg.epochs, + "repeats": int(repeats or 0), "lr": float(lr), "batch_size": int(batch_size)}) files = "\n".join(str(p) for p in written) bucket_note = (f"\nBuckets: {buckets} (from the dataset's actual sizes)" if buckets else f"\nSingle bucket at {int(resolution)}px") bucket_note += stats.upscale_note(int(resolution)) - caveat = "" - if trainer == "musubi": - caveat = ("\n\n⚠️ musubi needs your local DiT / VAE / text-encoder paths — " - "fill the <> placeholders in the command before running.") + exposure = exposure_line(trainer, str(ds), cfg.epochs, num_repeats, batch_size) + caveat = ("\n\n📋 validation/validation.md says how to pick the best epoch from " + "the per-epoch samples.") + if trainer in ("musubi", "fizgig"): + caveat += (f"\n\n⚠️ {trainer} needs your local model paths — fill the " + "<> placeholders in the command before running.") elif trainer == "kohya": - caveat = ("\n\n⚠️ kohya sd-scripts: SDXL base runs from the HF id shown; for a " + caveat += ("\n\n⚠️ kohya sd-scripts: SDXL base runs from the HF id shown; for a " "Pony / Illustrious / NoobAI checkpoint, replace the <> pretrained " "path. Verify flags against the sd-scripts docs before a long run.") # Advisory ④→⑤ sanity check: do the dataset's captions fit this base model? @@ -1538,7 +1584,8 @@ def do_generate_train_config(trainer: str, model_key: str, dataset_dir: str, if mismatch: caveat += f"\n\n{mismatch}" return (f"✅ Wrote:\n{files}\n\nDataset: {stats.n_images} images, " - f"{stats.min_long_side}-{stats.max_long_side}px long side{bucket_note}\n\n" + f"{stats.min_long_side}-{stats.max_long_side}px long side{bucket_note}\n" + f"{exposure.replace('**', '')}\n\n" f"Run it with:\n{command}{caveat}\n\n" f"⚠️ Configs are generated, not test-trained — verify keys against your " f"trainer's own docs before a long run.") @@ -2118,6 +2165,11 @@ def _check_for_update(): info="Writes a ratings sidecar into the dataset folder so Idiot LoRa " "Builder's grid opens pre-triaged — blurry, over/under-exposed and " "near-duplicate shots marked 'needs edit'. Nothing is launched.") + exp_holdout = gr.Number( + value=0, precision=0, minimum=0, label="Hold out N reference photos", + info="Keeps your N largest ① photos OUT of training, in a " + "'-heldout' folder beside the dataset — a fair test of whether " + "the LoRA learned the face. 0 trains on everything.") btn_export = gr.Button("④ Export dataset", variant="primary") exp_result = gr.Textbox(label="Result", lines=8) with gr.Accordion("Publish to Hugging Face (optional)", open=False): @@ -2145,17 +2197,20 @@ def _check_for_update(): gr.Markdown( "Generate a ready-to-edit LoRA training config for your dataset. " "**ai-toolkit** produces a one-command `config.yaml` (`python run.py …`); " - "**musubi-tuner** produces a `dataset.toml` plus a command template where " - "you fill in your local model paths. Nothing is launched or executed here — " - "the config is written into the dataset folder and the run command is shown.") + "**musubi-tuner**, **kohya** and **Fizgig** produce a `dataset.toml` plus a " + "command template where you fill in your local model paths. Every trainer " + "saves a checkpoint and renders fixed validation prompts each epoch; " + "`validation/validation.md` explains how to pick the best one. Nothing is " + "launched here — the files are written into the dataset folder and the run " + "command is shown.") from studio import user_config as _uc _ai_presets = TRAINER_MODELS["ai-toolkit"] with gr.Row(): with gr.Column(scale=1): tr_trainer = gr.Radio(TRAINER_CHOICES, value="ai-toolkit", label="Trainer", - info="ai-toolkit is one-command; musubi/kohya emit a " - "config plus a run-command template.") + info="ai-toolkit is one-command; musubi/kohya/Fizgig " + "emit a config plus a run-command template.") tr_path = gr.Textbox(label="Trainer install path (saved on this machine)", value=_uc.get_trainer_path("ai-toolkit"), placeholder=r"C:\ai-toolkit", @@ -2186,11 +2241,15 @@ def _check_for_update(): tr_alpha = gr.Number(value=_ai_presets[0].alpha, precision=0, label="Alpha", info="Usually equal to rank.") with gr.Row(): - tr_steps = gr.Number(value=_ai_presets[0].steps, precision=0, - label="Steps", - info="Auto-suggested from image count on Inspect.") - tr_lr = gr.Number(value=_ai_presets[0].lr, label="Learning rate", - info="1e-4 is a common starting point.") + tr_epochs = gr.Number(value=_ai_presets[0].epochs, precision=0, + label="Epochs", + info="One checkpoint + one sample set each.") + tr_repeats = gr.Number(value=0, precision=0, minimum=0, + label="Repeats (0 = auto)", + info="Auto sizes epochs to the target steps.") + tr_lr = gr.Number(value=_ai_presets[0].lr, label="Learning rate", + info="1e-4 is a common starting point.") + tr_exposure = gr.Markdown() tr_multi_res = gr.Checkbox( value=True, label="Multi-resolution buckets", info="Bucket by the dataset's real aspect ratios instead of " @@ -2198,10 +2257,10 @@ def _check_for_update(): with gr.Column(scale=2): tr_dataset = gr.Textbox( label="Dataset folder (auto-filled by ④ Export)", - info="The config file is written INTO this folder, and the " - "step count and buckets are derived from the images in " - "it. Works on any dataset folder, not just ④'s.") - btn_inspect = gr.Button("🔍 Inspect dataset & suggest steps") + info="The config and a validation/ pack are written INTO this " + "folder; repeats and buckets are derived from its images. " + "Works on any dataset folder, not just ④'s.") + btn_inspect = gr.Button("🔍 Inspect dataset") tr_stats = gr.Markdown() tr_gen = gr.Button("⑤ Generate training config", variant="primary") tr_result = gr.Textbox(label="Result / run command", lines=14) @@ -2380,22 +2439,26 @@ def _cap_cost(key: str, model: str, selected: list[str]) -> str: [exp_rows, exp_gallery, exp_select, exp_preview_note]) btn_export.click(do_export, [exp_select, project_name, project_trigger, output_root, exp_zip, - dataset_type, gen_style, gen_style_text, exp_ilb], + dataset_type, gen_style, gen_style_text, exp_ilb, exp_holdout], [exp_result, tr_dataset, exp_ds_dir]) \ - .then(inspect_dataset, [tr_dataset, dataset_type], [tr_stats, tr_steps]) + .then(inspect_dataset, [tr_dataset, dataset_type], [tr_stats]) btn_publish_hf.click(do_publish_hf, [exp_ds_dir, exp_hf_repo, exp_hf_private], [exp_hf_note]) - tr_hparams = [tr_res, tr_rank, tr_alpha, tr_steps, tr_lr, tr_batch] + tr_hparams = [tr_res, tr_rank, tr_alpha, tr_epochs, tr_lr, tr_batch] tr_trainer.change(on_trainer_change, [tr_trainer], [tr_model, tr_path] + tr_hparams) tr_model.change(on_model_change, [tr_trainer, tr_model], tr_hparams) tr_save_path.click(save_trainer_path, [tr_trainer, tr_path], [tr_path_note]) - btn_inspect.click(inspect_dataset, [tr_dataset, dataset_type], [tr_stats, tr_steps]) + btn_inspect.click(inspect_dataset, [tr_dataset, dataset_type], [tr_stats]) + # Epochs change with the trainer/model presets, so this follows them too. + for field in (tr_trainer, tr_dataset, tr_epochs, tr_repeats, tr_batch): + field.change(exposure_line, [tr_trainer, tr_dataset, tr_epochs, tr_repeats, tr_batch], + [tr_exposure]) tr_gen.click(do_generate_train_config, [tr_trainer, tr_model, tr_dataset, tr_path, tr_name, project_trigger] + tr_hparams + [tr_multi_res, dataset_type, gen_style, gen_style_text, - project_name], + project_name, tr_repeats], [tr_result]) demo.load(_check_for_update, None, update_notice) diff --git a/cli.py b/cli.py index 1817f6b..7d7efa2 100644 --- a/cli.py +++ b/cli.py @@ -370,6 +370,10 @@ def export( help="Publish to the HF Hub as this dataset id (needs a write HF_TOKEN)"), hf_private: bool = typer.Option( True, "--hf-private/--hf-public", help="Visibility of the published HF dataset"), + holdout: int = typer.Option( + 0, "--holdout", + help="Keep the N largest preprocessed photos out of training, in a sibling " + "-heldout folder (a likeness test)"), ): """Package captioned folders into a flat NN.png/NN.txt dataset (standalone).""" from studio.package import package_dataset, resolve_export_items @@ -387,7 +391,7 @@ def export( metadata = {"character_name": name, "trigger": trigger, "dataset_type": dtype, "source_folders": [str(f) for f in folders], "skipped_uncaptioned": res.missing, "skipped_empty_caption": res.empties} - ds = package_dataset(res.items, output_root, name, trigger, metadata) + ds = package_dataset(res.items, output_root, name, trigger, metadata, holdout=holdout) typer.echo(f"Dataset written to {ds}") if zip_: from studio.package import zip_dataset diff --git a/docs/ARCHITECTURE.md b/docs/ARCHITECTURE.md index f0fba63..f91660a 100644 --- a/docs/ARCHITECTURE.md +++ b/docs/ARCHITECTURE.md @@ -1,6 +1,6 @@ # Architecture -Version: 0.17.3 +Version: 0.18.0 ``` app.py Gradio UI — thin wiring over the stage functions (5 tabs). @@ -54,7 +54,8 @@ studio/ flag) + should_stop_now(). Checked BETWEEN items by the three batch loops; UI-free so the CLI and tests drive it directly preprocess.py Restore (comfyui|basic|auto) + isolate + optional tighten-crop - + resize. Atomic: any failure deletes the partial output. + + resize, stamping a `dd_source` provenance PNG chunk. + Atomic: any failure deletes the partial output. failed_report() turns a failure into data, not an exception isolate.py Subject isolation: builtin SAM3 (transformers) | comfyui. crop_to_content() tightens the isolated (subject-on-white, or @@ -95,8 +96,10 @@ studio/ gotcha; the gemini backend retries transient statuses with backoff (GEMINI_RETRIES / GEMINI_BACKOFF_S) package.py Dataset export (NN./NN.txt + metadata.json + metadata.jsonl - + README.txt); records dataset_type + a detected caption_style - in metadata.json; zip_dataset() bundles a folder into a .zip; + + README.txt); records dataset_type, a detected caption_style, + ① provenance and held-out photos in metadata.json; `holdout=N` + copies the N largest real photos to `-heldout/`; + zip_dataset() bundles a folder into a .zip; resolve_export_items() classifies candidates by caption sidecar state (ready/empty/missing) — shared by the UI gate and the CLI shot_style.py The medium ②'s prompts ask for: SHOT_STYLES (match | @@ -117,15 +120,19 @@ studio/ wardrobe.py Compositional unisex outfit pool for the outfit column (colour x garment; one pool, no gender picker) plan_io.py Save/load shot plans as YAML (user-editable prompt libraries) - dataset_stats.py Inspect a dataset folder (count/sizes/aspects) -> suggested - steps + bucket ladder, so ⑤'s numbers are derived not guessed + dataset_stats.py Inspect a dataset folder (count/sizes/aspects/① provenance) -> + target steps, auto repeats + bucket ladder, so ⑤'s numbers are + derived not guessed. provenance() reads ①'s `dd_source` PNG chunk trainer_configs.py Emit LoRA-trainer configs (ai-toolkit config.yaml / kohya - sd-scripts kohya-dataset.toml / musubi dataset.toml) + - run-command builders + model registry. ModelPreset carries - per-arch train/sample knobs (noise_scheduler/sample_guidance/ - sample_steps) and expects_tags (drives the caption/model - sanity check). _sample_prompt varies by TrainConfig.dataset_type; - caption_mismatch_warning() is the ④→⑤ advisory + sd-scripts kohya-dataset.toml / musubi dataset.toml / Fizgig + fizgig-dataset.toml) + run-command builders + model registry + + the validation pack (validation/: prompts, validation.md, scores + CSV). ModelPreset carries per-arch train/sample knobs + (noise_scheduler/sample_guidance/sample_steps/epochs) and + expects_tags (drives the caption/model sanity check). + validation_prompts() varies by TrainConfig.dataset_type; + exposure() is the steps/epoch math; caption_mismatch_warning() + is the ④→⑤ advisory user_config.py Persist trainer install paths, last training settings, the header dataset type, and the custom captioner endpoint (URL/model/ key-env-NAME/spacing) to .cache/user_settings.json (no secrets; @@ -881,10 +888,45 @@ quietly gets someone else's value. "Nothing was preprocessed"; a model-name mismatch only *warns*, because that check covers every bundled template including ones the run may never submit. - **Trainer configs are derived from the dataset, not from constants.** `dataset_stats.inspect()` - reads image count/dimensions (Pillow header parse only) so steps scale with the set - (`STEPS_PER_IMAGE`, clamped to 1000–4000) and buckets come from the images that actually exist. - Only config keys attested in each trainer's own examples are emitted — inventing plausible keys - produces configs that fail hours into a run. + reads image count/dimensions (Pillow header parse only) so training scales with the set and + buckets come from the images that actually exist. Only config keys attested in each trainer's + own examples are emitted — inventing plausible keys produces configs that fail hours into a run. +- **⑤ counts epochs, not steps (0.18.0).** An epoch is the unit you compare: every trainer saves + one checkpoint and renders one sample set per epoch, so "pick the best checkpoint" means "pick + an epoch". Each preset carries `epochs` (16 for ai-toolkit/musubi/kohya, 30–55 for Fizgig per + its docs). Plain epochs undertrain a small set (24 images × 16 epochs = 384 steps), so + **Repeats = 0 means auto**: `suggested_repeats()` picks the repeats that bring the run near + `target_steps` (n × 75, clamped 1000–4000 — the old steps rule, now a target). Fizgig's auto + is always 1 — its guidance is more epochs, not repeats. ai-toolkit counts steps, so its YAML + gets `steps = epochs × steps_per_epoch` and saves/samples every `steps_per_epoch`. + `exposure()` is the one formula; ⑤'s live exposure line shows it and warns below 400 or above + 6000 total steps. +- **Every config ships a validation pack (0.18.0).** `write_configs` writes + `validation/validation_prompts.txt` (8 fixed prompts per dataset type), `validation.md` (how to + pick the epoch, inference settings, held-out photos) and `validation_scores.csv` (a row per + epoch × prompt), all sampled at `VALIDATION_SEED` so between epochs only the LoRA changes. They + live in a **subfolder** because every trainer pairs `NN.txt` with `NN.png` in the dataset + folder; a `.txt` beside the images is one bad glob away from being read as a caption. musubi + and sd-scripts take per-line `--w --h --d` (+ the preset's `sample_line_args`); Fizgig and + ai-toolkit take plain prompts. musubi Krea 2 samples through `--turbo_dit` at `--l 1 --s 8`, + which needs the text encoder in the train command — it only renders samples. +- **Provenance rides in the PNG, not a side file (0.18.0).** ① writes a `dd_source` tEXt chunk + (`name, w, h, captured, restored, isolated`). It survives ③ (sidecars only) and ④ (`copy2`), + and sits before IDAT, so `inspect()` reads it from the header without decoding pixels. That is + how ⑤ can name the photos that were upscaled at ① (`photo.jpg (580px)`) and how ④ tells + a real photo from a generated shot. Generated shots carry none and count as native size. +- **Held-out photos sit beside the dataset, never in it (0.18.0).** ④'s "Hold out N" copies the + N largest real photos (by original area — the best likeness yardstick) to `-heldout/` + with numbered names, and records them in `metadata.json["heldout"]`; ⑤'s `validation.md` lists + them. Only photos with provenance qualify (a generated shot is no test of likeness) and at + least one image always stays in training. The dataset naming loop skips a name whose + `-heldout` sibling exists, so a stale held-out folder never pairs with a new dataset. +- **Fizgig is a ⑤ target, not a launch (0.18.0).** It reads musubi's `dataset.toml` shape, so + `render_musubi_toml` serves both; `arch` holds its `--family` (`krea2` / `qwen_image21` / + `klein`). The command is its three `cache`/`train` steps from `docs/CLI.md` with `<>` + model paths, `--precision auto --blocks_to_swap -1` (what its GUI does) and 1024 samples (it + warns smaller previews undersell a checkpoint). Preset knobs — adaptive LR, EMA, rank 8, + `--speed_lora`/`--training_adapter` — follow its per-family presets. - **`musubi_command()` takes the whole `TrainConfig`, not just the `ModelPreset`.** It previously took the preset alone and hardcoded `--network_dim 16` / `--max_train_steps 2000`, so every ⑤-tab slider was silently discarded. Regression-tested in `tests/test_trainer_configs_0_4_0.py`. @@ -901,7 +943,8 @@ quietly gets someone else's value. version flag, the text-encoder flags and the timestep args, copied from each model's `docs/.md` example in musubi-tuner; `tests/test_trainer_configs.py` pins them. - **ai-toolkit keeps every checkpoint.** `max_step_saves_to_keep` was 4, so a long run deleted the - early checkpoints, which are often the best likeness. It is now `ceil(steps / save_every)`. + early checkpoints, which are often the best likeness. Since 0.18.0 it is `epochs` (one save per + epoch). - **Train LoRAs on the base checkpoint, never a Turbo/distilled one.** Krea 2 trains on `krea/Krea-2-Raw` (arch `krea2`; `krea` was never an ai-toolkit arch). A test fails if any preset path contains "turbo". @@ -1031,10 +1074,12 @@ grouped by stage. Milestone versions are noted only where they explain a design - **④ Export** — flat NN.png/NN.txt + metadata.json (records dataset_type + detected caption_style) + metadata.jsonl (HF imagefolder) + README.txt; click-to-pick selection gate preseeded with what ③ captioned, with near-duplicate and caption-health advisories; optional .zip; opt-in - private-by-default Hugging Face publish. + private-by-default Hugging Face publish; hold out N real photos as a likeness test. - **⑤ Train config** — ai-toolkit (one-command, incl. correct SDXL knobs), kohya sd-scripts SDXL, - musubi-tuner; steps + multi-resolution buckets derived from the dataset; caption/model sanity - check; type-aware sample prompt. Nothing is ever launched — configs + the run command are shown. + musubi-tuner, Fizgig; epochs + auto repeats + multi-resolution buckets derived from the dataset + with a live exposure line; caption/model sanity check; a validation pack (8 type-aware prompts + at one seed, a guide, a score sheet). Nothing is ever launched — configs + the run command are + shown. The **CLI** mirrors every stage (`preprocess`/`generate`/`caption`/`lint`/`export`/`build`) with the same options and the same shared seams (`resolve_captioner_config`, `merge_tagger_overrides`, @@ -1156,6 +1201,13 @@ ordered by benefit-to-cost. ⑤'s bucket line (idea borrowed from Idiot LoRa Builder's crop tool, which shows the same verdict per crop). The *transforming* half (pad-to-square / center-crop) is still open and still wants its own design pass: it mutates pixels, and the advisory covers the case where knowing is enough. +- **⑤ optimizer choice and attention-only training (planned for 0.18.0, deferred).** An + "Advanced" accordion with an optimizer dropdown (AdamW8bit default; Prodigy at lr 1.0, hidden + for Fizgig, which removed it; Automagic where the trainer has it) and a musubi/Fizgig + attention-only toggle via `--network_args "exclude_patterns=[…]"`. Low value for the target + user, who should not need either; add when someone asks. +- **Fizgig look-score report (⑤).** Fizgig writes `fizgig_look_scores.json` per run; reading it + back into a per-epoch report beside `validation_scores.csv` would close the pick-the-epoch loop. - **Exact CLIP token count (③/④).** The 77-token warning is a tokenizer-free estimate; loading the real CLIP tokenizer would make it exact. Only worth it if a user needs high accuracy — the estimate errs safe and is fine for an advisory. diff --git a/studio/__init__.py b/studio/__init__.py index 8ce5787..52fb9dd 100644 --- a/studio/__init__.py +++ b/studio/__init__.py @@ -1,3 +1,3 @@ """Dataset Deviser: turn a character, style, or concept into a ready-to-train LoRA dataset.""" -__version__ = "0.17.3" +__version__ = "0.18.0" diff --git a/studio/dataset_stats.py b/studio/dataset_stats.py index f346d68..ee2f68f 100644 --- a/studio/dataset_stats.py +++ b/studio/dataset_stats.py @@ -7,6 +7,7 @@ from __future__ import annotations +import json from collections import Counter from dataclasses import dataclass, field from pathlib import Path @@ -22,6 +23,20 @@ MIN_STEPS = 1000 MAX_STEPS = 4000 +# PNG text chunk ① writes into every prepped image: the source photo's name, +# original size and capture time. ③ only writes sidecars and ④ copies bytes, +# so it reaches the dataset — where an upscaled source otherwise looks native. +PROVENANCE_KEY = "dd_source" + + +def provenance(im: Image.Image) -> dict: + """The ① provenance of an opened image, or {} (generated shots carry none).""" + try: + data = json.loads(im.info.get(PROVENANCE_KEY) or "{}") + except (TypeError, ValueError): + return {} + return data if isinstance(data, dict) else {} + # ai-toolkit's established multi-resolution idiom: list the buckets and it sorts # images into them. Only keys attested in ai-toolkit's own examples are emitted. BUCKET_LADDER = [512, 768, 1024] @@ -33,6 +48,8 @@ class DatasetStats: n_captioned: int sizes: list[tuple[int, int]] = field(default_factory=list) aspect_counts: Counter = field(default_factory=Counter) + # (source name, original long side) for images ① enlarged from a smaller source. + upscaled_sources: list[tuple[str, int]] = field(default_factory=list) @property def min_long_side(self) -> int: @@ -43,11 +60,23 @@ def max_long_side(self) -> int: return max((max(s) for s in self.sizes), default=0) @property - def suggested_steps(self) -> int: + def target_steps(self) -> int: if not self.n_images: return MIN_STEPS return max(MIN_STEPS, min(MAX_STEPS, self.n_images * STEPS_PER_IMAGE)) + def suggested_repeats(self, epochs: int, batch_size: int = 1) -> int: + """Repeats that bring `epochs` passes near `target_steps`. + + Epochs set how many checkpoints there are to compare; repeats set how + long each one trains. 16 epochs of a 24-image set at repeats 1 is only + 384 steps — too few to learn a face. + """ + if not self.n_images: + return 1 + per_repeat = self.n_images * max(epochs, 1) / max(batch_size, 1) + return max(1, round(self.target_steps / per_repeat)) + def buckets_for(self, resolution: int) -> list[int]: """Bucket ladder capped at `resolution` — never upscale past the training resolution, and never past what the images actually contain.""" @@ -79,15 +108,24 @@ def upscale_note(self, resolution: int) -> str: config alone — the images have to be measured. Empty string when every image already clears the bar, so a clean dataset says nothing. """ + note = "" small = self.undersized(resolution) - if not small: - return "" - factor = self.max_upscale(resolution) - return (f"\n⚠️ {len(small)} of {self.n_images} image(s) are below " - f"{resolution}px on the long side — up to {factor:.1f}× upscale to " - f"fill it. The kohya/musubi configs set `bucket_no_upscale`, so " - f"those land in a smaller bucket instead; for ai-toolkit, either " - f"lower the resolution or restore them at ① first.") + if small: + factor = self.max_upscale(resolution) + note = (f"\n⚠️ {len(small)} of {self.n_images} image(s) are below " + f"{resolution}px on the long side — up to {factor:.1f}× upscale to " + f"fill it. The kohya/musubi/Fizgig configs set `bucket_no_upscale`, " + f"so those land in a smaller bucket instead; for ai-toolkit, either " + f"lower the resolution or restore them at ① first.") + grown = self.upscaled_sources + if grown: + names = ", ".join(f"{n} ({side}px)" for n, side in grown[:5]) + more = f" and {len(grown) - 5} more" if len(grown) > 5 else "" + note += (f"\n⚠️ {len(grown)} of {self.n_images} image(s) were upscaled at ① " + f"from smaller sources: {names}{more}. They look full-size but hold " + f"only the source's detail — consider leaving them out at ④ when " + f"generated shots cover the same views.") + return note def summary(self) -> str: if not self.n_images: @@ -97,9 +135,9 @@ def summary(self) -> str: f"**{self.n_images} images** ({self.n_captioned} with captions)", f"Long side: {self.min_long_side}–{self.max_long_side}px", f"Aspect ratios: {shapes}", - f"Suggested steps: **{self.suggested_steps}** " + f"Target: **~{self.target_steps} steps** " f"({self.n_images} images × {STEPS_PER_IMAGE}, clamped to " - f"{MIN_STEPS}–{MAX_STEPS})", + f"{MIN_STEPS}–{MAX_STEPS}); Repeats = 0 sizes each epoch to reach it", ] if self.n_captioned < self.n_images: lines.append(f"⚠️ {self.n_images - self.n_captioned} image(s) have no " @@ -118,11 +156,12 @@ def _aspect_label(w: int, h: int) -> str: def inspect(dataset_dir: Path) -> DatasetStats: - """Read image dimensions + caption coverage for `dataset_dir`. + """Read image dimensions, caption coverage and ① provenance for `dataset_dir`. - Pillow parses headers lazily, so this never decodes pixel data. Unreadable - files are skipped rather than raising — a stray non-image shouldn't break - config generation. + Pillow parses headers lazily, so this never decodes pixel data — and a PNG's + text chunks sit before its image data, so provenance comes with the header. + Unreadable files are skipped rather than raising — a stray non-image + shouldn't break config generation. """ images = list_images(dataset_dir) stats = DatasetStats(n_images=0, n_captioned=0) @@ -130,8 +169,15 @@ def inspect(dataset_dir: Path) -> DatasetStats: try: with Image.open(path) as im: size = im.size + source = provenance(im) except Exception: continue + try: + original = max(int(source.get("w") or 0), int(source.get("h") or 0)) + except (TypeError, ValueError): # a malformed chunk is no provenance + original = 0 + if 0 < original < max(size): + stats.upscaled_sources.append((str(source.get("name", path.name)), original)) stats.n_images += 1 stats.sizes.append(size) stats.aspect_counts[_aspect_label(*size)] += 1 diff --git a/studio/package.py b/studio/package.py index 831cd8b..0a99e15 100644 --- a/studio/package.py +++ b/studio/package.py @@ -66,22 +66,78 @@ def resolve_export_items(image_paths: Iterable[Path]) -> ExportResolution: return ExportResolution(items=items, empties=empties, missing=missing) +def _provenance(img: Path) -> dict: + """① provenance of an image file, or {} for generated shots / unreadables.""" + from PIL import Image + + from studio.dataset_stats import provenance + + try: + with Image.open(img) as im: + return provenance(im) + except Exception: # never fail an export over metadata + return {} + + +def _heldout_dir(ds_dir: Path) -> Path: + return ds_dir.parent / f"{ds_dir.name}-heldout" + + +def _area(origin: dict) -> int: + try: + return int(origin.get("w", 0)) * int(origin.get("h", 0)) + except (TypeError, ValueError): + return 0 + + +def _hold_out(items: list[tuple[Path, str]], holdout: int, origins: dict[Path, dict]): + """Split off up to `holdout` real photos (largest originals first). + + Only ①'s photos qualify — a generated shot is no test of likeness — and at + least one image always stays in the training set. + """ + real = sorted((it for it in items if origins[it[0]]), key=lambda it: -_area(origins[it[0]])) + held = real[:max(0, min(holdout, len(items) - 1))] + return [it for it in items if it not in held], held + + def package_dataset( items: list[tuple[Path, str]], # (image_path, final_caption) output_root: Path, character_name: str, trigger: str, metadata: dict, + holdout: int = 0, ) -> Path: + """Write the dataset folder; `holdout` real photos go to a sibling instead. + + Held-out photos are copied to `-heldout/`, outside the folder any + trainer reads, so they stay a fair likeness test for the finished LoRA. + """ ds_name = f"{slugify(character_name or trigger or 'character')}-dataset" ds_dir = output_root / ds_name n = 1 - while ds_dir.exists(): # never clobber an existing dataset + # Never clobber an existing dataset, nor a held-out folder whose dataset was deleted. + while ds_dir.exists() or _heldout_dir(ds_dir).exists(): n += 1 ds_dir = output_root / f"{ds_name}-{n}" ds_dir.mkdir(parents=True) + origins = {img: _provenance(img) for img, _ in items} + items, held = _hold_out(items, holdout, origins) if holdout else (items, []) + if held: + held_dir = _heldout_dir(ds_dir) + held_dir.mkdir() + names = [] + for k, (img, caption) in enumerate(held, start=1): + name = f"{k:02d}-{img.name}" # numbered: sources may share a name + shutil.copy2(img, held_dir / name) + (held_dir / name).with_suffix(".txt").write_text(caption, encoding="utf-8") + names.append(name) + metadata = {**metadata, "heldout": {"dir": str(held_dir), "files": names}} + jsonl_lines: list[str] = [] + sources: list[dict] = [] for i, (img, caption) in enumerate(items, start=1): stem = f"{i:02d}" filename = f"{stem}{normalize_suffix(img.suffix)}" @@ -91,7 +147,11 @@ def package_dataset( # directly via datasets.load_dataset("imagefolder", ...) with captions. # file_name is resolved literally, so it must be the name actually written. jsonl_lines.append(json.dumps({"file_name": filename, "text": caption})) + if origins[img]: + sources.append({"file": filename, **origins[img]}) (ds_dir / "metadata.jsonl").write_text("\n".join(jsonl_lines) + "\n", encoding="utf-8") + if sources: # which files are real photos, and how big they started + metadata = {**metadata, "provenance": sources} # Sources are copied verbatim, so a mixed-format folder exports mixed # extensions. ".png", or ".png/.jpg" when they differ. extensions = sorted({normalize_suffix(img.suffix) for img, _ in items}) diff --git a/studio/preprocess.py b/studio/preprocess.py index bd31516..e3c6423 100644 --- a/studio/preprocess.py +++ b/studio/preprocess.py @@ -8,6 +8,7 @@ from __future__ import annotations +import json import tempfile from collections.abc import Callable from dataclasses import dataclass @@ -15,8 +16,10 @@ import numpy as np from PIL import Image, ImageOps +from PIL.PngImagePlugin import PngInfo from studio.config import settings +from studio.dataset_stats import PROVENANCE_KEY from studio.isolate import isolate_subject BLUR_THRESHOLD = 120.0 # Laplacian variance below this = soft/degraded image @@ -156,9 +159,12 @@ def preprocess( restore_backend = restore_backend or settings.restore_backend work_dir.mkdir(parents=True, exist_ok=True) with Image.open(source) as raw: + exif = raw.getexif() # Phone photos store portrait shots landscape plus an EXIF rotate tag; # ignoring it put sideways people in the dataset. - rotated = raw.getexif().get(0x0112, 1) != 1 # EXIF Orientation + rotated = exif.get(0x0112, 1) != 1 # EXIF Orientation + # DateTimeOriginal (Exif IFD), else DateTime: groups burst shots later. + captured = str(exif.get_ifd(0x8769).get(0x9003) or exif.get(0x0132) or "") oriented = ImageOps.exif_transpose(raw) img = oriented.convert("RGB") original_size = img.size @@ -218,7 +224,11 @@ def preprocess( img = crop_to_content(img) img = _resize_to_target(img, target) - img.save(out_path, "PNG") + info = PngInfo() + info.add_text(PROVENANCE_KEY, json.dumps({ + "name": source.name, "w": original_size[0], "h": original_size[1], + "captured": captured, "restored": restore, "isolated": isolate})) + img.save(out_path, "PNG", pnginfo=info) except BaseException: out_path.unlink(missing_ok=True) raise diff --git a/studio/trainer_configs.py b/studio/trainer_configs.py index 6d8b03e..8ce0049 100644 --- a/studio/trainer_configs.py +++ b/studio/trainer_configs.py @@ -1,15 +1,19 @@ """Generate LoRA-trainer config files for a finished dataset folder. -Two trainers are supported: +Four trainers are supported: - **ostris ai-toolkit** — one self-contained `config.yaml`, launched with a single `python run.py config.yaml`. Genuinely one-command: the model is a HF id and every hyperparameter lives in the file. -- **kohya-ss musubi-tuner** — the standard image `dataset.toml`. musubi's - *training* invocation additionally needs the user's local DiT / VAE / text- - encoder paths, which this tool cannot know, so `musubi_command()` returns a - template with clearly-marked `<>` placeholders. It is deliberately - NOT presented as one-click. +- **kohya-ss musubi-tuner**, **kohya-ss sd-scripts** and **Fizgig** — a + `dataset.toml` plus a command template. Their training invocations need the + user's local DiT / VAE / text-encoder paths, which this tool cannot know, so + the commands carry clearly-marked `<>` placeholders. They are + deliberately NOT presented as one-click. + +Every trainer also gets a validation pack in `validation/`: fixed prompts the +trainer samples after each epoch, plus a guide and a score sheet for picking +the best checkpoint. Configs are hand-templated (not serialized) so inline comments, sample prompts, and placeholders are preserved. No secrets are ever written — only the model id @@ -18,14 +22,17 @@ from __future__ import annotations +import csv +import io +import json import math from pathlib import Path import yaml from pydantic import BaseModel -# ai-toolkit checkpoint/sample interval, in steps. -SAVE_EVERY = 250 +# One seed for every validation sample, so epochs differ only by the LoRA. +VALIDATION_SEED = 42 def _yaml_str(value: str) -> str: @@ -47,6 +54,7 @@ def _yaml_str(value: str) -> str: "ai-toolkit": "ostris ai-toolkit (one-command: python run.py config.yaml)", "musubi": "kohya-ss musubi-tuner (dataset.toml + command template)", "kohya": "kohya-ss sd-scripts (SDXL LoRA — dataset.toml + command)", + "fizgig": "Fizgig (Krea 2 / Qwen Image 2.1 / Klein — dataset.toml + command)", } @@ -72,6 +80,9 @@ class ModelPreset(BaseModel): noise_scheduler: str = "flowmatch" sample_guidance: float = 4.0 sample_steps: int = 20 + # musubi / sd-scripts read sample settings per prompt line (`--l 7 --s 25`); + # appended after the size and seed. Empty = the trainer's per-arch defaults. + sample_line_args: str = "" # Caption style this base model expects: tag-trained checkpoints (SDXL / # Illustrious / NoobAI / Pony) learn from comma tags, prose models (Flux / # Qwen-Image / Z-Image / Krea) from natural language. Drives the ④→⑤ advisory @@ -83,17 +94,22 @@ class ModelPreset(BaseModel): # Every musubi arch has its own LoRA module; plain networks.lora is wrong # for all of them. network_module: str = "networks.lora" - musubi_args: str = "" # timestep sampling etc. for the train command + # Extra train-command flags (musubi and Fizgig): timestep sampling, Fizgig's + # recipe, sample-only model paths. + musubi_args: str = "" musubi_version: str = "" # --model_version, passed to all three commands + # The train command always gets these: sampling each epoch needs them, even + # where training itself reads cached text-encoder outputs (Krea 2). musubi_text_encoders: list[str] = ["text_encoder"] - musubi_train_te: bool = True # Krea 2 trains from cached TE outputs alone # kohya-ss sd-scripts training script (SDXL uses sdxl_train_network.py). kohya_script: str = "<>" # per-model defaults (UI pre-fills these; the user can override) resolution: int = 1024 rank: int = 16 alpha: int = 16 - steps: int = 2000 + # One checkpoint and one sample set per epoch. Repeats are derived from the + # dataset size (⑤) so the total lands near the target step count. + epochs: int = 16 lr: float = 1e-4 batch_size: int = 1 @@ -161,12 +177,15 @@ class ModelPreset(BaseModel): ModelPreset(key="zimage", label="Z-Image", arch="zimage", musubi_script="zimage_train_network.py", network_module="networks.lora_zimage", musubi_args=_SHIFT.format(2.0)), - # docs/krea2.md: --dit is the Raw checkpoint; training reads the cached - # text-encoder outputs, so the train command takes no --text_encoder. + # docs/krea2.md: --dit is the Raw checkpoint. Samples render on Turbo + # (how the LoRA is used): CFG off, 8 steps. ModelPreset(key="krea2", label="Krea 2 (Raw)", arch="krea2", musubi_script="krea2_train_network.py", - network_module="networks.lora_krea2", musubi_args=_SHIFT.format(2.5), - musubi_train_te=False, rank=32, alpha=32), + network_module="networks.lora_krea2", + musubi_args=_SHIFT.format(2.5) + + " --turbo_dit <>", + sample_line_args="--l 1 --s 8", sample_guidance=1.0, sample_steps=8, + rank=32, alpha=32), ], # kohya-ss sd-scripts is the standard SDXL LoRA trainer — the natural home for # the Danbooru/e621 tag captions (③). SDXL base is a runnable HF id; a family @@ -174,11 +193,38 @@ class ModelPreset(BaseModel): "kohya": [ ModelPreset(key="sdxl", label="SDXL 1.0 (base)", name_or_path="stabilityai/stable-diffusion-xl-base-1.0", - arch="sdxl", kohya_script="sdxl_train_network.py", expects_tags=True), + arch="sdxl", kohya_script="sdxl_train_network.py", expects_tags=True, + sample_guidance=7.0, sample_steps=25, sample_line_args="--l 7 --s 25"), ModelPreset(key="sdxl-custom", label="SDXL-family checkpoint — Pony / Illustrious / NoobAI (set path)", name_or_path="<>", - arch="sdxl", kohya_script="sdxl_train_network.py", expects_tags=True), + arch="sdxl", kohya_script="sdxl_train_network.py", expects_tags=True, + sample_guidance=7.0, sample_steps=25, sample_line_args="--l 7 --s 25"), + ], + # Fizgig (shootthesound/Fizgig, docs/CLI.md): `arch` is its --family. Each + # preset is the GUI's default recipe — adaptive LR (which ignores + # --learning_rate), EMA, and the speed LoRA / adapter it renders and trains + # with. Fizgig wants repeats 1 and more epochs. + "fizgig": [ + ModelPreset(key="krea2", label="Krea 2 (Raw)", arch="krea2", + musubi_args="--adaptive_lr --adaptive_lr_min 2e-4 --adaptive_lr_max 4e-4 " + "--ema_decay 0.98 --speed_lora <> " + "--sample_steps 8", + sample_guidance=1.0, sample_steps=8, rank=8, alpha=8, epochs=30), + # 0.5 MP: Fizgig's Qwen presets train at [704, 704]. + ModelPreset(key="qwen_image21", label="Qwen Image 2.1", arch="qwen_image21", + musubi_args="--adaptive_lr --adaptive_lr_min 2e-4 --adaptive_lr_max 4e-4 " + "--ema_decay 0.98 --training_adapter <> --speed_lora <>", + # Previews render on the turbo LoRA at 6 steps (docs/CLI.md). + sample_guidance=1.0, sample_steps=6, resolution=704, rank=8, alpha=8, + epochs=30), + ModelPreset(key="klein", label="FLUX.2 Klein 9B (Base)", arch="klein", + musubi_args="--adaptive_lr --adaptive_lr_min 5e-5 --adaptive_lr_max 4e-4 " + "--preview_checkpoint <>", + # Previews render on the distilled DiT at 4 steps (docs/CLI.md). + sample_guidance=1.0, sample_steps=4, epochs=55), ], } @@ -194,7 +240,10 @@ class TrainConfig(BaseModel): resolution: int = 1024 rank: int = 16 alpha: int = 16 - steps: int = 2000 + epochs: int = 16 + num_repeats: int = 1 + # Images in the dataset; 0 = count them when writing (`write_configs`). + n_images: int = 0 lr: float = 1e-4 batch_size: int = 1 # Multi-resolution buckets. Empty = single-bucket at `resolution` (the old @@ -206,26 +255,55 @@ class TrainConfig(BaseModel): shot_style_text: str = "" -def _sample_prompt(cfg: TrainConfig) -> str: - """The sample prompt written into the trainer config. +def validation_prompts(cfg: TrainConfig) -> list[str]: + """Eight fixed prompts the trainer samples every epoch, same seed each time. + + They probe what a LoRA gets wrong: a tight face (likeness), full body and a + back view (physique, the views a few references rarely show), an outfit and + a setting the dataset never had (does the trigger carry the subject or the + training images?), another medium (flexibility), and action. House rules + as ②: no negation, never "photorealistic". The medium follows the ② shot style: "a photo of " is wrong for a - dataset of anime shots, and the sample images are how a user judges the run. - `match` keeps the historical "a photo of" wording — with an unknown source - medium there is nothing better to say, and changing it would alter every - existing character config for no gain. + dataset of anime shots. `match` keeps "a photo of" — with an unknown source + medium there is nothing better to say. """ from studio.shot_style import resolve who = cfg.trigger or cfg.name or "the subject" if cfg.dataset_type == "style": - # The trigger is an aesthetic; the prompt names content it renders. - return f"{who}, a mountain landscape at sunset" + # The trigger is an aesthetic; the prompts name content it renders. + return [f"{who}, {content}" for content in ( + "a mountain landscape at sunset", "a portrait of an old fisherman", + "a bowl of fruit on a wooden table", "a busy city street at night", + "a cat asleep on a windowsill", "a castle on a hill under storm clouds", + "a woman reading in a cafe", "a spaceship over a desert")] lead = resolve(cfg.shot_style, cfg.shot_style_text).sample_lead subject = f"{lead} {who}" if lead.endswith(" of") else f"{lead}{who}" if cfg.dataset_type == "concept": - return subject - return f"{subject}, standing outdoors in daylight" + return [subject, f"{subject}, close-up", f"{subject} on a wooden table", + f"{subject} outdoors in daylight", f"{subject} at night under neon light", + f"{subject} in a snowy forest", f"a watercolor painting of {who}", + f"{subject}, seen from above"] + return [f"{subject}, standing outdoors in daylight", + f"{subject}, a tight close-up of the face, soft window light", + f"{subject}, full body, standing facing the camera, plain studio backdrop", + f"{subject}, seen from behind over the shoulder, walking down a city " + f"street at dusk", + f"{subject}, sitting at a cafe table wearing a yellow raincoat", + f"{subject}, standing on the deck of a sailing ship in a storm", + f"a charcoal sketch of {who}, waist-up portrait", + f"{subject}, jumping mid-air on a beach, arms raised"] + + +def exposure(n_images: int, num_repeats: int, batch_size: int, epochs: int) -> tuple[int, int]: + """(steps per epoch, total steps). An empty folder counts as one image.""" + per_epoch = max(1, math.ceil(max(n_images, 1) * max(num_repeats, 1) / max(batch_size, 1))) + return per_epoch, per_epoch * max(epochs, 1) + + +def _exposure(cfg: TrainConfig) -> tuple[int, int]: + return exposure(cfg.n_images, cfg.num_repeats, cfg.batch_size, cfg.epochs) def caption_mismatch_warning(preset: ModelPreset, caption_kind: str) -> str: @@ -265,6 +343,9 @@ def render_aitoolkit_yaml(cfg: TrainConfig) -> str: raw_note = (" # Train on Raw. Turbo is distilled: a LoRA trained on it fights the\n" " # distillation. Use the LoRA with Turbo at inference instead.\n" if m.arch.startswith("krea2") else "") + # ai-toolkit counts steps, so an "epoch" is the steps one pass takes. + per_epoch, total = _exposure(cfg) + prompts = "".join(f" - {_yaml_str(p)}\n" for p in validation_prompts(cfg)) return ( "# ai-toolkit LoRA config generated by Dataset Deviser.\n" "# Run from your ai-toolkit install: python run.py \n" @@ -283,10 +364,10 @@ def render_aitoolkit_yaml(cfg: TrainConfig) -> str: f" linear_alpha: {cfg.alpha}\n" " save:\n" " dtype: float16\n" - f" save_every: {SAVE_EVERY}\n" + f" save_every: {per_epoch}\n" # Keep every save: the best checkpoint is often an early one, and # pruning to the last few deleted it before anyone could compare. - f" max_step_saves_to_keep: {max(1, math.ceil(cfg.steps / SAVE_EVERY))}\n" + f" max_step_saves_to_keep: {cfg.epochs}\n" " datasets:\n" f" - folder_path: \"{cfg.dataset_dir.as_posix()}\"\n" " caption_ext: \"txt\"\n" @@ -296,7 +377,7 @@ def render_aitoolkit_yaml(cfg: TrainConfig) -> str: f" resolution: {_resolution_list(cfg)}\n" " train:\n" f" batch_size: {cfg.batch_size}\n" - f" steps: {cfg.steps}\n" + f" steps: {total} # {cfg.epochs} epochs x {per_epoch} steps\n" " gradient_accumulation_steps: 1\n" " train_unet: true\n" " train_text_encoder: false\n" @@ -313,11 +394,13 @@ def render_aitoolkit_yaml(cfg: TrainConfig) -> str: f" quantize: {str(m.quantize).lower()}\n" f"{model_extras}" " sample:\n" - f" sample_every: {SAVE_EVERY}\n" + f" sample_every: {per_epoch}\n" f" width: {cfg.resolution}\n" f" height: {cfg.resolution}\n" + f" seed: {VALIDATION_SEED}\n" + " walk_seed: false\n" " prompts:\n" - f" - {_yaml_str(_sample_prompt(cfg))}\n" + f"{prompts}" f" guidance_scale: {m.sample_guidance:g}\n" f" sample_steps: {m.sample_steps}\n" "meta:\n" @@ -326,11 +409,11 @@ def render_aitoolkit_yaml(cfg: TrainConfig) -> str: ) -def render_musubi_toml(cfg: TrainConfig, num_repeats: int = 1) -> str: - """The standard musubi-tuner image dataset.toml.""" +def render_musubi_toml(cfg: TrainConfig) -> str: + """The standard musubi-tuner image dataset.toml (Fizgig reads the same shape).""" cache = (cfg.dataset_dir / "cache").as_posix() return ( - "# musubi-tuner dataset config generated by Dataset Deviser.\n" + f"# {cfg.trainer} dataset config generated by Dataset Deviser.\n" "# Pass to training with: --dataset_config \n" "# (the training command also needs your DiT/VAE/text-encoder paths.)\n" "\n" @@ -345,7 +428,7 @@ def render_musubi_toml(cfg: TrainConfig, num_repeats: int = 1) -> str: "[[datasets]]\n" f'image_directory = "{cfg.dataset_dir.as_posix()}"\n' f'cache_directory = "{cache}"\n' - f"num_repeats = {num_repeats}\n" + f"num_repeats = {cfg.num_repeats}\n" ) @@ -354,10 +437,17 @@ def aitoolkit_command(install_path: str, config_path: Path) -> str: return f'cd "{base}" && python run.py "{config_path.as_posix()}"' -def musubi_command(install_path: str, toml_path: Path, cfg: TrainConfig) -> str: +def _sampling(prompts_path: Path) -> str: + """Per-epoch sampling flags, shared by musubi, sd-scripts and Fizgig.""" + return (f' --sample_prompts "{prompts_path.as_posix()}" ' + "--sample_every_n_epochs 1 --sample_at_first \\\n") + + +def musubi_command(install_path: str, toml_path: Path, cfg: TrainConfig, + prompts_path: Path) -> str: """Build the musubi run command from `cfg`. - Takes the whole TrainConfig, not just the preset: rank/alpha/steps/lr are + Takes the whole TrainConfig, not just the preset: rank/alpha/epochs/lr are user-tunable in the UI and must actually reach the command line. """ m = cfg.model @@ -367,11 +457,10 @@ def musubi_command(install_path: str, toml_path: Path, cfg: TrainConfig) -> str: version = f" --model_version {m.musubi_version}" if m.musubi_version else "" te = " ".join(f"--{f} <>" for f in m.musubi_text_encoders) - train_te = f" {te} \\\n" if m.musubi_train_te else "" extra = f" {m.musubi_args} \\\n" if m.musubi_args else "" return ( "# 1-2: cache latents and text-encoder outputs (re-run after changing the\n" - "# images or captions). 3: train, saving a LoRA every epoch to compare.\n" + "# images or captions). 3: train, saving a LoRA and samples every epoch.\n" f'cd "{base}"\n' f'python src/musubi_tuner/{prefix}_cache_latents.py --dataset_config "{toml}" ' f"--vae <>{version}\n" @@ -381,7 +470,7 @@ def musubi_command(install_path: str, toml_path: Path, cfg: TrainConfig) -> str: f"src/musubi_tuner/{m.musubi_script} \\\n" " --dit <> \\\n" " --vae <> \\\n" - f"{train_te}" + f" {te} \\\n" f' --dataset_config "{toml}"{version} \\\n' " --sdpa --mixed_precision bf16 --gradient_checkpointing \\\n" f"{extra}" @@ -390,14 +479,20 @@ def musubi_command(install_path: str, toml_path: Path, cfg: TrainConfig) -> str: f" --network_module {m.network_module} \\\n" f" --network_dim {cfg.rank} \\\n" f" --network_alpha {cfg.alpha} \\\n" - f" --max_train_steps {cfg.steps} --save_every_n_epochs 1 --seed 42 \\\n" + f" --max_train_epochs {cfg.epochs} --save_every_n_epochs 1 " + f"--seed {VALIDATION_SEED} \\\n" + f"{_sampling(prompts_path)}" f' --output_dir output --output_name "{cfg.name}"\n' - f"# Flags follow musubi-tuner's docs/{prefix}.md example; check " - "it if a run complains." + "# The text encoder in step 3 only renders the samples. " + + ("--turbo_dit can't be combined\n# with --blocks_to_swap: drop it to sample " + "on Raw (edit the prompt lines to --l 5.5 --s 28).\n" + if "--turbo_dit" in m.musubi_args else "\n") + + f"# Flags follow musubi-tuner's docs/{prefix}.md example; check it if a " + "run complains." ) -def render_kohya_toml(cfg: TrainConfig, num_repeats: int = 1) -> str: +def render_kohya_toml(cfg: TrainConfig) -> str: """A kohya-ss sd-scripts image `dataset.toml` (subsets layout).""" return ( "# kohya-ss sd-scripts dataset config generated by Dataset Deviser.\n" @@ -416,11 +511,12 @@ def render_kohya_toml(cfg: TrainConfig, num_repeats: int = 1) -> str: "\n" " [[datasets.subsets]]\n" f' image_dir = "{cfg.dataset_dir.as_posix()}"\n' - f" num_repeats = {num_repeats}\n" + f" num_repeats = {cfg.num_repeats}\n" ) -def kohya_command(install_path: str, toml_path: Path, cfg: TrainConfig) -> str: +def kohya_command(install_path: str, toml_path: Path, cfg: TrainConfig, + prompts_path: Path) -> str: """Build the kohya sd-scripts run command from `cfg`. The pretrained model is the preset's `name_or_path` (SDXL base is a runnable @@ -436,13 +532,52 @@ def kohya_command(install_path: str, toml_path: Path, cfg: TrainConfig) -> str: f' --output_dir output --output_name "{cfg.name}" \\\n' " --network_module networks.lora \\\n" f" --network_dim {cfg.rank} --network_alpha {cfg.alpha} \\\n" - f" --learning_rate {cfg.lr} --max_train_steps {cfg.steps} \\\n" + f" --learning_rate {cfg.lr} --max_train_epochs {cfg.epochs} \\\n" " --optimizer_type AdamW8bit --mixed_precision bf16 --sdpa \\\n" - " --gradient_checkpointing --save_model_as safetensors --save_every_n_steps 250\n" + f"{_sampling(prompts_path)}" + " --gradient_checkpointing --save_model_as safetensors --save_every_n_epochs 1\n" "# SDXL uses sdxl_train_network.py; verify flags against the kohya-ss/sd-scripts docs." ) +def fizgig_command(install_path: str, toml_path: Path, cfg: TrainConfig, + prompts_path: Path) -> str: + """Fizgig's three steps, per its docs/CLI.md: cache latents, cache text, train. + + `arch` is the --family. `--precision auto --blocks_to_swap -1` is the GUI's + behaviour (the CLI defaults to bf16). Samples render at 1024: Fizgig warns + that smaller previews undersell the checkpoint. + """ + m = cfg.model + base = install_path.strip() or "<>" + toml = toml_path.as_posix() + fam = f"--family {m.arch}" + cache = f'python src/fizgig/families/cache.py {fam} --dataset_config "{toml}"' + return ( + "# 1-2: cache latents and text-encoder outputs (re-run after changing the\n" + "# images or captions). 3: train, saving a LoRA and samples every epoch.\n" + "# On Windows use venv\\Scripts\\python.exe and put each command on one line.\n" + f'cd "{base}"\n' + f"{cache} --stage latents --model <>\n" + f"{cache} --stage text --model <>\n" + f"python src/fizgig/families/train.py {fam} \\\n" + f' --dataset_config "{toml}" \\\n' + " --dit <> \\\n" + " --vae <> --text_encoder <> \\\n" + f' --output_dir output --output_name "{cfg.name}" \\\n' + f" --network_dim {cfg.rank} --network_alpha {cfg.alpha} " + f"--learning_rate {cfg.lr} \\\n" + f" --max_train_epochs {cfg.epochs} --save_every_n_epochs 1 " + f"--seed {VALIDATION_SEED} \\\n" + " --optimizer_type adamw8bit --precision auto --blocks_to_swap -1 \\\n" + f" {m.musubi_args} \\\n" + f"{_sampling(prompts_path)}" + f" --sample_width 1024 --sample_height 1024 --sample_seed {VALIDATION_SEED}\n" + "# --adaptive_lr sets the learning rate itself (--learning_rate is ignored).\n" + "# Flags follow Fizgig's docs/CLI.md; run train.py --help if one is refused." + ) + + def _nonclobber(path: Path) -> Path: if not path.exists(): return path @@ -454,33 +589,113 @@ def _nonclobber(path: Path) -> Path: n += 1 -def write_configs(cfg: TrainConfig, install_path: str = "", - num_repeats: int = 1) -> tuple[list[Path], str]: - """Write the trainer's config file(s) into the dataset folder. +def _write(path: Path, text: str) -> Path: + path = _nonclobber(path) + path.write_text(text, encoding="utf-8") + return path + + +def _prompt_file(cfg: TrainConfig, prompts: list[str]) -> str: + """The trainer's --sample_prompts file. musubi and sd-scripts take size, seed + and sampler settings per line; Fizgig and ai-toolkit take plain prompts.""" + suffix = "" + if cfg.trainer in ("musubi", "kohya"): + suffix = (f" --w {cfg.resolution} --h {cfg.resolution} --d {VALIDATION_SEED} " + f"{cfg.model.sample_line_args}").rstrip() + return ("# Validation prompts generated by Dataset Deviser, sampled every epoch " + "at one seed.\n# One prompt per line; lines starting with # are comments.\n" + + "".join(f"{p}{suffix}\n" for p in prompts)) + + +def _heldout(dataset_dir: Path) -> dict: + """The held-out photos ④ recorded in metadata.json, or {}.""" + try: + meta = json.loads((dataset_dir / "metadata.json").read_text(encoding="utf-8")) + except (OSError, ValueError): + return {} + held = meta.get("heldout") if isinstance(meta, dict) else None + return held if isinstance(held, dict) else {} + - Returns (written_paths, run_command). Never clobbers existing files — - collisions get a `.N` suffix. `install_path` is only used to compose the - displayed run command; it is never written into any file. +def _guide(cfg: TrainConfig, prompts: list[str]) -> str: + m = cfg.model + if m.arch.startswith("krea2"): + settings = ("Use the LoRA on **Krea 2 Turbo**: 8 steps, CFG 1. On Raw: 28 steps, " + "CFG 5.5.") + else: + settings = (f"{m.label}: {m.sample_steps} steps, CFG {m.sample_guidance:g} — the " + f"settings the samples use.") + held = _heldout(cfg.dataset_dir) + held_block = "" + if held.get("files"): + held_block = ( + "\n## Held-out photos\n\n" + f"These real photos were kept out of training: `{held.get('dir', '')}`\n" + + "".join(f"- {f}\n" for f in held["files"]) + + "\nThe LoRA never saw them, so they are the fair likeness test: compare " + "each epoch's close-up and full-body samples against them.\n") + numbered = "".join(f"{i}. {p}\n" for i, p in enumerate(prompts, 1)) + return ( + f"# Validating {cfg.name}\n\n" + f"{TRAINERS[cfg.trainer].split(' (')[0]} · {m.label} · {cfg.epochs} epochs. " + f"Every epoch saves a checkpoint and renders the prompts below at seed " + f"{VALIDATION_SEED}, so between epochs only the LoRA changes.\n\n" + "## Pick the checkpoint\n\n" + "1. Open the sample folder inside the trainer's output folder.\n" + "2. Score epochs in `validation_scores.csv`, 1–5: **likeness** (the face), " + "**physique** (build, height, proportions), **adherence** (did it wear, do and " + "go where the prompt said), **unwanted** (something from the training images " + "that appears unasked: a prop, a pose, a background, a colour cast).\n" + "3. Keep the earliest epoch where likeness stops improving. When adherence " + "drops or unwanted traits climb, later epochs are overfitting.\n" + f"{held_block}\n" + f"## Inference settings\n\n{settings}\n\n" + f"## Prompts (seed {VALIDATION_SEED})\n\n{numbered}") + + +def _score_sheet(cfg: TrainConfig, n_prompts: int) -> str: + out = io.StringIO() + writer = csv.writer(out, lineterminator="\n") + writer.writerow(["epoch", "prompt", "likeness", "physique", "adherence", "unwanted", + "notes"]) + for epoch in range(1, cfg.epochs + 1): + for prompt in range(1, n_prompts + 1): + writer.writerow([epoch, prompt, "", "", "", "", ""]) + return out.getvalue() + + +def write_configs(cfg: TrainConfig, install_path: str = "") -> tuple[list[Path], str]: + """Write the trainer's config file(s) and the validation pack. + + Config files go into the dataset folder; the validation pack into its + `validation/` subfolder, where no trainer reads a `.txt` as a caption. + Returns (written_paths, run_command), the config first. Never clobbers + existing files — collisions get a `.N` suffix. `install_path` is only used + to compose the displayed run command; it is never written into any file. """ - cfg.dataset_dir.mkdir(parents=True, exist_ok=True) - written: list[Path] = [] + if cfg.trainer not in TRAINERS: + raise ValueError(f"Unknown trainer: {cfg.trainer}") + ds = cfg.dataset_dir + ds.mkdir(parents=True, exist_ok=True) + if not cfg.n_images: + from studio.config import list_images + + cfg = cfg.model_copy(update={"n_images": len(list_images(ds))}) + prompts = validation_prompts(cfg) + val = ds / "validation" + val.mkdir(exist_ok=True) + prompts_path = _write(val / "validation_prompts.txt", _prompt_file(cfg, prompts)) if cfg.trainer == "ai-toolkit": - path = _nonclobber(cfg.dataset_dir / "ai-toolkit.yaml") - path.write_text(render_aitoolkit_yaml(cfg), encoding="utf-8") - written.append(path) + path = _write(ds / "ai-toolkit.yaml", render_aitoolkit_yaml(cfg)) command = aitoolkit_command(install_path, path) - elif cfg.trainer == "musubi": - path = _nonclobber(cfg.dataset_dir / "dataset.toml") - path.write_text(render_musubi_toml(cfg, num_repeats=num_repeats), - encoding="utf-8") - written.append(path) - command = musubi_command(install_path, path, cfg) elif cfg.trainer == "kohya": - path = _nonclobber(cfg.dataset_dir / "kohya-dataset.toml") - path.write_text(render_kohya_toml(cfg, num_repeats=num_repeats), - encoding="utf-8") - written.append(path) - command = kohya_command(install_path, path, cfg) - else: - raise ValueError(f"Unknown trainer: {cfg.trainer}") - return written, command + path = _write(ds / "kohya-dataset.toml", render_kohya_toml(cfg)) + command = kohya_command(install_path, path, cfg, prompts_path) + elif cfg.trainer == "musubi": + path = _write(ds / "dataset.toml", render_musubi_toml(cfg)) + command = musubi_command(install_path, path, cfg, prompts_path) + else: # fizgig reads musubi's dataset.toml shape + path = _write(ds / "fizgig-dataset.toml", render_musubi_toml(cfg)) + command = fizgig_command(install_path, path, cfg, prompts_path) + return [path, prompts_path, _write(val / "validation.md", _guide(cfg, prompts)), + _write(val / "validation_scores.csv", _score_sheet(cfg, len(prompts)))], command diff --git a/tests/test_dataset_stats.py b/tests/test_dataset_stats.py index 86ed290..223fb9a 100644 --- a/tests/test_dataset_stats.py +++ b/tests/test_dataset_stats.py @@ -38,16 +38,16 @@ def test_steps_scale_with_image_count(tmp_path: Path) -> None: # The bug being fixed: a flat 2000 steps regardless of dataset size. small = ds.inspect(_make(tmp_path / "s", [(1024, 1024)] * 24)) large = ds.inspect(_make(tmp_path / "l", [(1024, 1024)] * 40)) - assert small.suggested_steps == 24 * ds.STEPS_PER_IMAGE - assert large.suggested_steps == 40 * ds.STEPS_PER_IMAGE - assert large.suggested_steps > small.suggested_steps + assert small.target_steps == 24 * ds.STEPS_PER_IMAGE + assert large.target_steps == 40 * ds.STEPS_PER_IMAGE + assert large.target_steps > small.target_steps def test_steps_are_clamped_at_both_ends(tmp_path: Path) -> None: tiny = ds.inspect(_make(tmp_path / "t", [(1024, 1024)] * 2)) huge = ds.inspect(_make(tmp_path / "h", [(1024, 1024)] * 200)) - assert tiny.suggested_steps == ds.MIN_STEPS - assert huge.suggested_steps == ds.MAX_STEPS + assert tiny.target_steps == ds.MIN_STEPS + assert huge.target_steps == ds.MAX_STEPS def test_aspect_ratios_are_labelled(tmp_path: Path) -> None: @@ -129,7 +129,7 @@ def test_the_train_tab_actually_shows_the_advisory(tmp_path: Path) -> None: folder = _make(tmp_path / "small", [(512, 512)] * 4, captions=4) out = A.do_generate_train_config( "ai-toolkit", "flux-dev", str(folder), "", "lora", "sks", - 1024, 16, 16, 2000, 1e-4, 1, True) + 1024, 16, 16, 16, 1e-4, 1, True) assert "below 1024px" in out assert "2.0×" in out @@ -140,5 +140,5 @@ def test_the_train_tab_stays_quiet_when_the_dataset_is_big_enough(tmp_path: Path folder = _make(tmp_path / "big", [(1024, 1024)] * 4, captions=4) out = A.do_generate_train_config( "ai-toolkit", "flux-dev", str(folder), "", "lora", "sks", - 1024, 16, 16, 2000, 1e-4, 1, True) + 1024, 16, 16, 16, 1e-4, 1, True) assert "below 1024px" not in out diff --git a/tests/test_identity_sync.py b/tests/test_identity_sync.py index 35c6ebf..f6b0bc8 100644 --- a/tests/test_identity_sync.py +++ b/tests/test_identity_sync.py @@ -122,7 +122,7 @@ def _dataset(folder: Path, n: int = 2) -> Path: def _train(tmp_path: Path, lora_name: str, project: str) -> str: return A.do_generate_train_config( "ai-toolkit", "flux-dev", str(_dataset(tmp_path / "ds")), "", lora_name, - "sysnootles", 512, 16, 16, 2000, 1e-4, 1, False, "character", + "sysnootles", 512, 16, 16, 16, 1e-4, 1, False, "character", A.shot_style.MATCH, "", project) diff --git a/tests/test_provenance_heldout.py b/tests/test_provenance_heldout.py new file mode 100644 index 0000000..4785436 --- /dev/null +++ b/tests/test_provenance_heldout.py @@ -0,0 +1,145 @@ +"""0.18.0: ① provenance, ④ held-out photos, and ⑤'s exposure line. + +Provenance is a PNG text chunk ① writes, so a later stage can tell a real photo +(and how small it started) from a generated shot without any side file. +""" + +from __future__ import annotations + +import json +from pathlib import Path + +import numpy as np +from PIL import Image + +import app as A +from studio import dataset_stats as ds +from studio.package import package_dataset +from studio.preprocess import preprocess + + +def _photo(path: Path, size: tuple[int, int], seed: int = 1) -> Path: + rng = np.random.default_rng(seed=seed) + Image.fromarray(rng.integers(48, 208, size=(size[1], size[0], 3), dtype=np.uint8)).save(path) + return path + + +def _prep(tmp_path: Path, sizes: list[tuple[int, int]]) -> list: + src = tmp_path / "src" + src.mkdir() + return [preprocess(_photo(src / f"p{i}.png", s, i), tmp_path / "prepped", target=256, + force_restore=False, isolate=False).output + for i, s in enumerate(sizes, 1)] + + +def test_preprocess_records_the_source(tmp_path: Path) -> None: + out = _prep(tmp_path, [(120, 80)])[0] + with Image.open(out) as im: + origin = ds.provenance(im) + assert origin["name"] == "p1.png" + assert (origin["w"], origin["h"]) == (120, 80) + assert origin["restored"] is False and origin["isolated"] is False + + +def test_provenance_reads_from_the_header_without_decoding(tmp_path: Path) -> None: + """The chunk sits before IDAT, so a truncated pixel stream still yields it.""" + out = _prep(tmp_path, [(120, 80)])[0] + data = out.read_bytes() + cut = tmp_path / "cut.png" + cut.write_bytes(data[:data.index(b"IDAT") + 16]) + with Image.open(cut) as im: + assert ds.provenance(im)["name"] == "p1.png" + + +def test_generated_and_corrupt_chunks_read_as_empty(tmp_path: Path) -> None: + plain = _photo(tmp_path / "gen.png", (64, 64)) + with Image.open(plain) as im: + assert ds.provenance(im) == {} + from PIL.PngImagePlugin import PngInfo + + info = PngInfo() + info.add_text(ds.PROVENANCE_KEY, "{not json") + bad = tmp_path / "bad.png" + Image.new("RGB", (8, 8)).save(bad, pnginfo=info) + with Image.open(bad) as im: + assert ds.provenance(im) == {} + + +def test_upscaled_sources_are_named_at_train_time(tmp_path: Path) -> None: + _prep(tmp_path, [(120, 80), (400, 300)]) # target 256: only the first grew + stats = ds.inspect(tmp_path / "prepped") + assert stats.upscaled_sources == [("p1.png", 120)] + note = stats.upscale_note(256) + assert "p1.png (120px)" in note and "p2.png" not in note + + +def test_export_carries_provenance_into_metadata(tmp_path: Path) -> None: + prepped = _prep(tmp_path, [(120, 80)]) + gen = _photo(tmp_path / "prepped" / "shot.png", (256, 256), 9) + items = [(prepped[0], "a photo of trg"), (gen, "trg, side view")] + out = package_dataset(items, tmp_path / "out", "Sy", "trg", {}) + meta = json.loads((out / "metadata.json").read_text(encoding="utf-8")) + assert meta["provenance"] == [{"file": "01.png", **{ + "name": "p1.png", "w": 120, "h": 80, "captured": "", + "restored": False, "isolated": False}}] + with Image.open(out / "01.png") as im: # copy2 kept the chunk + assert ds.provenance(im)["name"] == "p1.png" + + +def test_held_out_photos_leave_the_training_folder(tmp_path: Path) -> None: + prepped = _prep(tmp_path, [(300, 300), (900, 600), (500, 500)]) + gen = _photo(tmp_path / "prepped" / "shot.png", (256, 256), 9) + items = [(p, f"caption {i}") for i, p in enumerate(prepped)] + [(gen, "shot")] + out = package_dataset(items, tmp_path / "out", "Sy", "trg", {}, holdout=2) + + held = out.parent / f"{out.name}-heldout" + # Largest originals first; numbered, with their captions alongside. + assert sorted(p.name for p in held.iterdir()) == [ + "01-p2_prepped.png", "01-p2_prepped.txt", "02-p3_prepped.png", "02-p3_prepped.txt"] + assert len(list(out.glob("*.png"))) == 2 # p1 + the generated shot + meta = json.loads((out / "metadata.json").read_text(encoding="utf-8")) + assert meta["heldout"] == {"dir": str(held), "files": ["01-p2_prepped.png", + "02-p3_prepped.png"]} + + +def test_holdout_never_empties_the_dataset_or_takes_generated_shots(tmp_path: Path) -> None: + prepped = _prep(tmp_path, [(300, 300)]) + out = package_dataset([(prepped[0], "c")], tmp_path / "out", "Sy", "trg", {}, holdout=3) + assert not (out.parent / f"{out.name}-heldout").exists() + gen = _photo(tmp_path / "g.png", (64, 64)) + out2 = package_dataset([(gen, "c"), (gen, "d")], tmp_path / "out2", "Sy", "trg", {}, + holdout=1) + assert not (out2.parent / f"{out2.name}-heldout").exists() + + +def test_a_stale_heldout_folder_is_never_reused(tmp_path: Path) -> None: + (tmp_path / "out" / "sy-dataset-heldout").mkdir(parents=True) + out = package_dataset([(_photo(tmp_path / "g.png", (64, 64)), "c")], + tmp_path / "out", "Sy", "trg", {}) + assert out.name == "sy-dataset-2" + + +def test_do_export_reports_the_holdout(tmp_path: Path) -> None: + prepped = _prep(tmp_path, [(300, 300), (500, 500)]) + for p in prepped: + p.with_suffix(".txt").write_text("a caption", encoding="utf-8") + result, _ds, _hf = A.do_export([str(p) for p in prepped], "Sy", "trg", + str(tmp_path / "out"), holdout=1) + assert "Held out 1" in result + + +def test_exposure_line_auto_repeats(tmp_path: Path) -> None: + folder = tmp_path / "d" + folder.mkdir() + for i in range(20): + Image.new("RGB", (64, 64)).save(folder / f"{i:02d}.png") + # 20 x 75 = 1500 target steps; 16 epochs of 20 images -> ~5 repeats. + line = A.exposure_line("ai-toolkit", str(folder), 16, 0, 1) + assert "20 images × 5 repeats (auto)" in line and "**1600 steps**" in line + assert "⚠️" not in line + # Fizgig's guidance is repeats 1; a typed value always wins. + assert "× 1 repeats (auto)" in A.exposure_line("fizgig", str(folder), 30, 0, 1) + assert "× 2 repeats ÷" in A.exposure_line("fizgig", str(folder), 30, 2, 1) + assert "Few steps" in A.exposure_line("ai-toolkit", str(folder), 4, 1, 4) + assert "Few steps" not in A.exposure_line("fizgig", str(folder), 4, 1, 4) + assert A.exposure_line("ai-toolkit", "", 16, 0, 1) == "" diff --git a/tests/test_review_ideas.py b/tests/test_review_ideas.py index 19462fe..ad3274b 100644 --- a/tests/test_review_ideas.py +++ b/tests/test_review_ideas.py @@ -18,8 +18,8 @@ from studio.trainer_configs import ( TRAINER_MODELS, TrainConfig, - _sample_prompt, caption_mismatch_warning, + validation_prompts, ) # ---------- CLIP token estimate ---------- @@ -91,9 +91,9 @@ def cfg(dt: str) -> TrainConfig: return TrainConfig(trainer="ai-toolkit", model=TRAINER_MODELS["ai-toolkit"][0], dataset_dir=tmp_path, trigger="tok", dataset_type=dt) - char = _sample_prompt(cfg("character")) - style = _sample_prompt(cfg("style")) - concept = _sample_prompt(cfg("concept")) + char = validation_prompts(cfg("character"))[0] + style = validation_prompts(cfg("style"))[0] + concept = validation_prompts(cfg("concept"))[0] assert char != style != concept assert char.startswith("a photo of tok") assert style.startswith("tok,") # style names content the trigger renders diff --git a/tests/test_shot_style.py b/tests/test_shot_style.py index d7f7f15..c8a195d 100644 --- a/tests/test_shot_style.py +++ b/tests/test_shot_style.py @@ -145,15 +145,15 @@ def test_prompts_have_no_double_punctuation_or_dangling_commas() -> None: # ---------- ⑤ sample prompt ---------- def test_sample_prompt_names_the_right_medium() -> None: - from studio.trainer_configs import TRAINER_MODELS, TrainConfig, _sample_prompt + from studio.trainer_configs import TRAINER_MODELS, TrainConfig, validation_prompts preset = TRAINER_MODELS["ai-toolkit"][0] def prompt(style: str, text: str = "", dtype: str = "character") -> str: - return _sample_prompt(TrainConfig( + return validation_prompts(TrainConfig( trainer="ai-toolkit", model=preset, dataset_dir=Path("."), trigger="trg", name="n", dataset_type=dtype, - shot_style=style, shot_style_text=text)) + shot_style=style, shot_style_text=text))[0] # match keeps the historical wording — no regression for existing configs. assert prompt(MATCH) == "a photo of trg, standing outdoors in daylight" @@ -165,12 +165,12 @@ def prompt(style: str, text: str = "", dtype: str = "character") -> str: def test_sample_prompt_handles_a_custom_style() -> None: - from studio.trainer_configs import TRAINER_MODELS, TrainConfig, _sample_prompt + from studio.trainer_configs import TRAINER_MODELS, TrainConfig, validation_prompts - out = _sample_prompt(TrainConfig( + out = validation_prompts(TrainConfig( trainer="ai-toolkit", model=TRAINER_MODELS["ai-toolkit"][0], dataset_dir=Path("."), trigger="trg", name="n", - shot_style=CUSTOM, shot_style_text="a woodcut print")) + shot_style=CUSTOM, shot_style_text="a woodcut print"))[0] assert "woodcut print" in out and "trg" in out diff --git a/tests/test_trainer_configs.py b/tests/test_trainer_configs.py index 64a7041..a37e350 100644 --- a/tests/test_trainer_configs.py +++ b/tests/test_trainer_configs.py @@ -24,7 +24,8 @@ def _cfg(trainer: str, tmp_path: Path) -> TrainConfig: return TrainConfig(trainer=trainer, model=TRAINER_MODELS[trainer][0], dataset_dir=tmp_path, trigger="sysnootles", name="sy-lora", - resolution=1024, rank=16, alpha=16, steps=1500, lr=1e-4) + resolution=1024, rank=16, alpha=16, epochs=10, num_repeats=10, + n_images=15, lr=1e-4) def test_aitoolkit_yaml_parses_and_has_keys(tmp_path: Path) -> None: @@ -32,7 +33,7 @@ def test_aitoolkit_yaml_parses_and_has_keys(tmp_path: Path) -> None: proc = doc["config"]["process"][0] assert proc["type"] == "sd_trainer" assert proc["network"]["linear"] == 16 - assert proc["train"]["steps"] == 1500 + assert proc["train"]["steps"] == 1500 # 15 images x 10 repeats x 10 epochs assert proc["datasets"][0]["folder_path"] == tmp_path.as_posix() assert proc["model"]["name_or_path"] # non-empty @@ -157,14 +158,69 @@ def test_musubi_krea2_matches_its_reference_recipe(tmp_path: Path) -> None: preset = _preset("musubi", "krea2") assert (preset.rank, preset.alpha) == (32, 32) cfg = _cfg("musubi", tmp_path).model_copy(update={"model": preset}) - _, command = write_configs(cfg) + written, command = write_configs(cfg) assert "--discrete_flow_shift 2.5" in command train = command.split("accelerate launch", 1)[1] - # Krea 2 trains from cached text-encoder outputs alone. - assert "--text_encoder" not in train - - -def test_aitoolkit_keeps_every_checkpoint(tmp_path: Path) -> None: - cfg = _cfg("ai-toolkit", tmp_path).model_copy(update={"steps": 3000}) - save = yaml.safe_load(render_aitoolkit_yaml(cfg))["config"]["process"][0]["save"] - assert save["max_step_saves_to_keep"] * save["save_every"] >= 3000 + # Training reads cached text outputs; the encoder and Turbo DiT only render samples. + assert "--text_encoder" in train and "--turbo_dit" in train + assert "--l 1 --s 8" in written[1].read_text(encoding="utf-8") + + +def test_aitoolkit_saves_and_samples_once_per_epoch(tmp_path: Path) -> None: + proc = yaml.safe_load(render_aitoolkit_yaml(_cfg("ai-toolkit", tmp_path)))["config"]["process"][0] + per_epoch = 150 # 15 images x 10 repeats / batch 1 + assert proc["save"]["save_every"] == per_epoch + assert proc["save"]["max_step_saves_to_keep"] == 10 # every epoch kept + assert proc["sample"]["sample_every"] == per_epoch + assert proc["sample"]["seed"] == 42 and proc["sample"]["walk_seed"] is False + assert len(proc["sample"]["prompts"]) == 8 + + +@pytest.mark.parametrize("trainer", ["ai-toolkit", "musubi", "kohya", "fizgig"]) +def test_every_trainer_gets_the_validation_pack(trainer: str, tmp_path: Path) -> None: + written, _ = write_configs(_cfg(trainer, tmp_path)) + names = [p.relative_to(tmp_path).as_posix() for p in written[1:]] + assert names == ["validation/validation_prompts.txt", "validation/validation.md", + "validation/validation_scores.csv"] + prompts = [ln for ln in written[1].read_text(encoding="utf-8").splitlines() + if ln and not ln.startswith("#")] + assert len(prompts) == 8 and all("sysnootles" in p for p in prompts) + rows = written[3].read_text(encoding="utf-8").splitlines() + assert len(rows) == 1 + 10 * 8 # header + epochs x prompts + # No trainer may read a validation prompt as a caption. + assert not list(tmp_path.glob("*.txt")) + + +def test_exposure_math() -> None: + from studio.trainer_configs import exposure + + assert exposure(10, 3, 2, 12) == (15, 180) + assert exposure(7, 1, 2, 4) == (4, 16) # a partial batch still costs a step + assert exposure(0, 0, 0, 5) == (1, 5) # never zero + + +@pytest.mark.parametrize("key,res", [("krea2", 1024), ("qwen_image21", 704), ("klein", 1024)]) +def test_fizgig_preset_toml_and_commands(key: str, res: int, tmp_path: Path) -> None: + preset = _preset("fizgig", key) + assert preset.resolution == res + cfg = _cfg("fizgig", tmp_path).model_copy( + update={"model": preset, "resolution": preset.resolution}) + written, command = write_configs(cfg, install_path="C:/Fizgig") + assert written[0].name == "fizgig-dataset.toml" + assert tomllib.loads(written[0].read_text(encoding="utf-8"))["general"]["resolution"] == [res, res] + assert command.count(f"--family {preset.arch}") == 3 + assert "--stage latents" in command and "--stage text" in command + assert "--max_train_epochs 10 --save_every_n_epochs 1" in command + assert "--sample_prompts" in command and "--sample_every_n_epochs 1" in command + # Fizgig's prompt file is plain prompts — no musubi per-line options. + assert " --w " not in written[1].read_text(encoding="utf-8") + + +def test_validation_guide_lists_held_out_photos(tmp_path: Path) -> None: + import json + + (tmp_path / "metadata.json").write_text(json.dumps( + {"heldout": {"dir": "X-heldout", "files": ["01-a.png"]}}), encoding="utf-8") + written, _ = write_configs(_cfg("ai-toolkit", tmp_path)) + guide = written[2].read_text(encoding="utf-8") + assert "X-heldout" in guide and "01-a.png" in guide diff --git a/tests/test_trainer_configs_0_4_0.py b/tests/test_trainer_configs_0_4_0.py index 1a82fce..610b54b 100644 --- a/tests/test_trainer_configs_0_4_0.py +++ b/tests/test_trainer_configs_0_4_0.py @@ -21,30 +21,34 @@ def _cfg(trainer: str, tmp_path: Path, **kw) -> TrainConfig: preset = TRAINER_MODELS[trainer][0] base = {"trainer": trainer, "model": preset, "dataset_dir": tmp_path, "name": "sy-lora", "trigger": "sysnootles", "resolution": 1024, - "rank": 32, "alpha": 64, "steps": 1800, "lr": 8e-5, "batch_size": 2} + "rank": 32, "alpha": 64, "epochs": 12, "num_repeats": 3, "n_images": 10, + "lr": 8e-5, "batch_size": 2} base.update(kw) return TrainConfig(**base) +def _musubi(cfg: TrainConfig, tmp_path: Path) -> str: + return musubi_command("C:/musubi", tmp_path / "dataset.toml", cfg, tmp_path / "p.txt") + + def test_musubi_command_honors_every_hyperparameter(tmp_path: Path) -> None: """The 0.3.1 bug: musubi_command() took only the ModelPreset, so rank/steps were hardcoded and every ⑤-tab slider was silently discarded.""" - cfg = _cfg("musubi", tmp_path) - cmd = musubi_command("C:/musubi", tmp_path / "dataset.toml", cfg) + cmd = _musubi(_cfg("musubi", tmp_path), tmp_path) assert "--network_dim 32" in cmd assert "--network_alpha 64" in cmd - assert "--max_train_steps 1800" in cmd + assert "--max_train_epochs 12" in cmd assert "--learning_rate 8e-05" in cmd assert "sy-lora" in cmd # The old code hardcoded these; make sure they can't creep back. assert "--network_dim 16" not in cmd - assert "--max_train_steps 2000" not in cmd + assert "--max_train_steps" not in cmd def test_musubi_command_has_no_unrendered_fstring_braces(tmp_path: Path) -> None: """The old code emitted the literal `{16}` from a broken f-string.""" - cmd = musubi_command("C:/musubi", tmp_path / "dataset.toml", _cfg("musubi", tmp_path)) + cmd = _musubi(_cfg("musubi", tmp_path), tmp_path) # <> placeholders are intentional; bare braces are not. assert not re.search(r"\{\d", cmd) assert "{16}" not in cmd @@ -53,11 +57,11 @@ def test_musubi_command_has_no_unrendered_fstring_braces(tmp_path: Path) -> None def test_write_configs_threads_config_into_musubi_command(tmp_path: Path) -> None: _written, command = write_configs(_cfg("musubi", tmp_path), "C:/musubi") assert "--network_dim 32" in command - assert "--max_train_steps 1800" in command + assert "--max_train_epochs 12" in command def test_musubi_toml_scales_repeats(tmp_path: Path) -> None: - toml = render_musubi_toml(_cfg("musubi", tmp_path), num_repeats=17) + toml = render_musubi_toml(_cfg("musubi", tmp_path, num_repeats=17)) assert "num_repeats = 17" in toml @@ -82,7 +86,7 @@ def test_aitoolkit_respects_sliders(tmp_path: Path) -> None: yaml = render_aitoolkit_yaml(_cfg("ai-toolkit", tmp_path)) assert "linear: 32" in yaml assert "linear_alpha: 64" in yaml - assert "steps: 1800" in yaml + assert "steps: 180 " in yaml # 10 images x 3 repeats / batch 2 = 15/epoch x 12 assert "batch_size: 2" in yaml diff --git a/tests/test_trainer_configs_kohya.py b/tests/test_trainer_configs_kohya.py index dfce024..1989605 100644 --- a/tests/test_trainer_configs_kohya.py +++ b/tests/test_trainer_configs_kohya.py @@ -21,11 +21,11 @@ def _cfg(tmp_path: Path, key: str = "sdxl") -> TrainConfig: preset = next(p for p in TRAINER_MODELS["kohya"] if p.key == key) return TrainConfig(trainer="kohya", model=preset, dataset_dir=tmp_path, name="sy-lora", trigger="trig", resolution=1024, rank=16, - alpha=16, steps=1500, lr=1e-4, batch_size=2) + alpha=16, epochs=10, lr=1e-4, batch_size=2) def test_kohya_toml_parses_with_subsets(tmp_path: Path) -> None: - doc = tomllib.loads(render_kohya_toml(_cfg(tmp_path), num_repeats=5)) + doc = tomllib.loads(render_kohya_toml(_cfg(tmp_path).model_copy(update={"num_repeats": 5}))) assert doc["general"]["caption_extension"] == ".txt" ds = doc["datasets"][0] assert ds["resolution"] == [1024, 1024] @@ -36,17 +36,18 @@ def test_kohya_toml_parses_with_subsets(tmp_path: Path) -> None: def test_kohya_command_threads_hparams_and_script(tmp_path: Path) -> None: - cmd = kohya_command("C:/sd-scripts", tmp_path / "kohya-dataset.toml", _cfg(tmp_path)) + cmd = kohya_command("C:/sd-scripts", tmp_path / "kohya-dataset.toml", _cfg(tmp_path), + tmp_path / "p.txt") assert "sdxl_train_network.py" in cmd assert "--network_dim 16" in cmd assert "--network_alpha 16" in cmd - assert "--max_train_steps 1500" in cmd + assert "--max_train_epochs 10" in cmd assert "stabilityai/stable-diffusion-xl-base-1.0" in cmd assert "C:/sd-scripts" in cmd def test_kohya_custom_preset_keeps_fill(tmp_path: Path) -> None: - cmd = kohya_command("", tmp_path / "d.toml", _cfg(tmp_path, "sdxl-custom")) + cmd = kohya_command("", tmp_path / "d.toml", _cfg(tmp_path, "sdxl-custom"), tmp_path / "p.txt") assert "< None: "the hand-off checkbox is in the layout but reaches do_export as nothing" -def test_handoff_checkbox_is_the_last_input_matching_the_signature() -> None: - """`ilb_handoff` is do_export's last parameter, so it must be the last input.""" - assert _export_inputs()[-1] == "Prepare for Idiot LoRa Builder" +def test_export_inputs_end_in_signature_order() -> None: + """`ilb_handoff, holdout` end do_export's signature, so they end the inputs.""" + assert _export_inputs()[-2:] == ["Prepare for Idiot LoRa Builder", + "Hold out N reference photos"] def test_export_with_the_box_ticked_writes_the_sidecar(tmp_path: Path) -> None: