diff --git a/.env.example b/.env.example index abad4b6..418ccfb 100644 --- a/.env.example +++ b/.env.example @@ -19,7 +19,8 @@ GROQ_API_KEY= # Built-in SAM3 subject isolation (gated model): accept the license at # https://huggingface.co/facebook/sam3 then create a READ token at -# https://huggingface.co/settings/tokens +# https://huggingface.co/settings/tokens (HUGGINGFACE_API_KEY in the environment works too, when +# there is no HF_TOKEN and no `hf auth login`) HF_TOKEN= # --- Optional overrides (defaults shown) --- diff --git a/AGENTS.md b/AGENTS.md index 5cb8270..f47d044 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -6,11 +6,11 @@ Turn a character, style, or concept into a ready-to-train LoRA dataset. One refe ## Current state -_Last verified: 2026-10-08_ +_Last verified: 2026-10-09_ -- **Status:** in active development; main is v0.19.0 (git tag `v0.19.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. Since 0.18.1: a ⑤ optimizer dropdown, AdamW8bit by default, Prodigy (lr 1.0, its diffusion args) for ai-toolkit/musubi/kohya; Fizgig stays AdamW8bit. Since 0.19.0: ① training copies **keep their background** and the isolated copy goes to `prepped/refs/` for ② only (angle re-isolation is its own ② option, off); a header **identity policy** for characters (Identity only varies outfits and captions clothing; Signature costume keeps the outfit out of captions); four framing tiers in the shot plan and a ② coverage line with a dominance warning; a cut-off-reference warning; ② near-duplicate finder and primary-reference preview, ④ burst grouping; local ② sends several references (`LDS_QWEN21_MAX_REFS`, default 1) at `LDS_QWEN21_RESOLUTION` (1536); Gemini gets a per-shot aspect and a reference cap; an opt-in **anchor shot**; SAM3 load errors name a stale `HF_TOKEN`. -- **In progress:** a phased LoRA best-practice update. 0.20.0 is next: a ⚡ Quick build tab and `cli build` on shared `pipeline.build_start`/`build_finish`. Nothing from it is half-built on main. The attention-only toggle planned for 0.18.0 was deferred (ARCHITECTURE.md → Future ideas). +- **Status:** in active development; main is v0.20.0 (git tag `v0.20.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. Since 0.18.1: a ⑤ optimizer dropdown, AdamW8bit by default, Prodigy (lr 1.0, its diffusion args) for ai-toolkit/musubi/kohya; Fizgig stays AdamW8bit. Since 0.19.0: ① training copies **keep their background** and the isolated copy goes to `prepped/refs/` for ② only (angle re-isolation is its own ② option, off); a header **identity policy** for characters (Identity only varies outfits and captions clothing; Signature costume keeps the outfit out of captions); four framing tiers in the shot plan and a ② coverage line with a dominance warning; a cut-off-reference warning; ② near-duplicate finder and primary-reference preview, ④ burst grouping; local ② sends several references (`LDS_QWEN21_MAX_REFS`, default 1) at `LDS_QWEN21_RESOLUTION` (1536); Gemini gets a per-shot aspect and a reference cap; an opt-in **anchor shot**; SAM3 load errors name a stale `HF_TOKEN`. Since 0.20.0: a **⚡ Quick build** tab (first tab: drop images, name, Cloud or Local, trainer; Build runs ①→②, untick, Finish runs ③→④→⑤) and `cli build` share `pipeline.build_start`/`build_finish`, and `cli build` writes ⑤ configs with `--trainer/--model`; a failed isolation keeps the photo and gives ② the whole image; builtin SAM3 retries `person` on real photos; a ComfyUI job's timeout no longer runs while it waits in the queue; `HUGGINGFACE_API_KEY` stands in when there is neither an `HF_TOKEN` nor an `hf auth login`; a caption that opens with the name in another case (`Ann` for trigger `ann`) now opens with the trigger's exact casing; `cli build` and Quick build Finish show ⑤'s undersized/upscaled warnings. +- **In progress:** nothing half-built on main. The phased LoRA best-practice update (0.17.3–0.20.0) is complete; the optional local face model (drift flag, verified coverage) awaits a go-ahead. The attention-only toggle planned for 0.18.0 was 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; the local defaults (one reference, 1536, anchor off) likewise rest on one A/B subject. 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. @@ -28,7 +28,7 @@ _Last verified: 2026-10-08_ | File / Directory | Purpose | |------------------|---------| -| `app.py` | Gradio UI — thin wiring over stage functions (5 tabs) | +| `app.py` | Gradio UI — thin wiring over stage functions (⚡ Quick build + 5 stage tabs) | | `cli.py` | Typer CLI — one subcommand per stage + `build`, `doctor`, `keys`, `custom-endpoint` | | `setup.bat` / `setup.sh` | One-time install: Python gate, venv, GPU/CPU torch, requirements, ONNX, keys, doctor | | `start.bat` / `start.sh` | Launch the UI | diff --git a/README.md b/README.md index e823d59..779a5c1 100644 --- a/README.md +++ b/README.md @@ -39,7 +39,17 @@ ai-toolkit, kohya, musubi or Fizgig, and a validation pack for picking the best **① Preprocess → ② Generate & curate → ③ Caption → ④ Export → ⑤ Train config** -Run them in order (each step auto-fills the next) or jump straight to the one you need. +New to this? Use the **⚡ Quick build** tab: drop your images, type a name at the top, pick Cloud +or Local and a trainer, press **Build**. Untick any shot that looks wrong and press **Finish**: you +get the captioned dataset, the trainer config with its run command, and the validation pack. +Name anything the subject is holding (*"cup, plate"*) so it isn't redrawn in every shot. If a +shadow or object crosses the face in your best photo, pick Cloud: Local tends to copy it into +the generated shots, Cloud draws a clean face. Photos +① had to enlarge a lot start unticked (they still guide generation). The trigger is made from +the name if you leave it blank. + +The numbered tabs run the same steps one at a time, with every option. Run them in order (each +step auto-fills the next) or jump straight to the one you need. | Stage | Local | Cloud | |---|---|---| @@ -257,6 +267,8 @@ python cli.py generate ./prepped/refs --name "Sy Snootles" --engine comfyui python cli.py generate ./prepped/refs --name "Sy Snootles" --identity costume # outfit is the character python cli.py generate ./prepped/refs --name "Sy Snootles" --shot-style anime # keep an illustrated look python cli.py generate ./prepped/refs --name "brass compass" --dataset-type concept # 18-shot object set +python cli.py build ./photos --name "Ann Lee" --trigger annlee --exclude-prompt "cup" \ + --trainer ai-toolkit --model krea # everything, plus ⑤ python cli.py caption ./folder --trigger sysnootles # prose sidecars python cli.py caption ./folder --trigger sysnootles --caption-style tags # Danbooru tags python cli.py caption ./folder --trigger mystyle --dataset-type style # style-framed @@ -278,7 +290,7 @@ Set or change any of them with `python cli.py keys` — they live only in a giti - **Gemini** — image generation and/or captioning: - **Groq** — free-tier captioning: -- **Hugging Face** (`HF_TOKEN`) — two uses: +- **Hugging Face** (`HF_TOKEN`, or `HUGGINGFACE_API_KEY` in the environment when you have no `hf auth login`) — two uses: - **Built-in SAM3 isolation.** `facebook/sam3` is **gated**: accept the licence on the [model page](https://huggingface.co/facebook/sam3), wait for approval, then add a **read** token. Weights (~3.4 GB) download once. diff --git a/app.py b/app.py index 147f678..31aadda 100644 --- a/app.py +++ b/app.py @@ -428,6 +428,7 @@ def _picker_mark(rows: list[tuple[str, str, str]], selected): ("dd-gallery-gen", "dd-picks-gen", "dd-zoom-gen"), ("dd-gallery-cap", "dd-picks-cap", "dd-zoom-cap"), ("dd-gallery-exp", "dd-picks-exp", "dd-zoom-exp"), + ("dd-gallery-quick", "dd-picks-quick", "dd-zoom-quick"), ] # Clicking a thumbnail must toggle it, and Gradio's own `Gallery.select` event cannot @@ -537,6 +538,21 @@ def val(row, k): return [Shot(**{k: val(row, k) for k in cols}) for _, row in df.iterrows() if val(row, "id").strip()] +def _flagged(path: Path, label: str) -> str: + """`label` plus this image's advisory quality flags (blurry, dark, …).""" + try: + from studio.quality import composition_flags, is_blurry + + flags: list[str] = [] + blurry, score = is_blurry(path) + if blurry: + flags.append(f"blurry ({score:.0f})") + flags += composition_flags(path) + except Exception: + return label # quality checks are advisory — never block the gallery on them + return f"{label} ⚠ {', '.join(flags)}" if flags else label + + def _gen_gallery(results: list[pipeline.GenResult], selected=None): """Picker rows + gallery + CheckboxGroup for ②'s kept-shot list. @@ -544,22 +560,7 @@ def _gen_gallery(results: list[pipeline.GenResult], selected=None): existing pick across a re-sync instead of silently re-checking rejected shots. """ ok = [r for r in results if r.path and r.path.exists()] - rows: list[tuple[str, str, str]] = [] - for r in ok: - label = r.shot.id - try: - from studio.quality import composition_flags, is_blurry - - flags: list[str] = [] - blurry, score = is_blurry(r.path) - if blurry: - flags.append(f"blurry ({score:.0f})") - flags += composition_flags(r.path) - if flags: - label = f"{r.shot.id} ⚠ {', '.join(flags)}" - except Exception: - pass # quality checks are advisory — never block the gallery on them - rows.append((str(r.path), r.shot.id, label)) + rows = [(str(r.path), r.shot.id, _flagged(r.path, r.shot.id)) for r in ok] ids = [r.shot.id for r in ok] keep = ids if selected is None else _picker_order(rows, selected) return rows, _picker_gallery(rows, keep), gr.CheckboxGroup(choices=ids, value=keep) @@ -621,6 +622,11 @@ def _preprocess_note(reports, out_dir: Path, alpha_cutout: bool) -> str: elif cut: head += (f"\n\nCut off at the bottom: {', '.join(cut)} — the other " f"reference(s) show more of the body.") + whole = [r.source.name for r in ok if r.reference and not r.isolated] + if whole: + head += (f"\n\n⚠️ **Not isolated:** {', '.join(whole)} — isolation failed (the Log " + f"says why), so ② uses the whole photo, background and props included. " + f"The training copy is unaffected.") if not failed: return head lines = "\n".join(f"- `{r.source.name}` — {r.error} \n → {_failure_hint(r.error)}" @@ -1749,6 +1755,166 @@ def _check_for_update(): text = "" return gr.Markdown(value=text, visible=bool(text)) +# ---------- ⚡ Quick build ---------- + +QUICK_ENGINES = [("Cloud — Gemini, billed to your Google key", "gemini"), + ("Local — ComfyUI, free", "comfyui")] + + +def quick_trigger(name: str, trigger: str) -> str: + """The header trigger, or one made from the name ("Sy Snootles" → "sysnootles").""" + from studio.package import slugify + + return trigger.strip() or slugify(name).replace("-", "") + + +def quick_engine_note(engine: str, dataset_type: str) -> str: + """What Build will cost (cloud) or whether it can run at all (local).""" + if dataset_type == "style": + return "Style datasets generate nothing: Build prepares and Finish captions your own images." + n = len(plan_for_type(dataset_type, "")) + if engine == "gemini": + if not settings.resolved_gemini_key(): + return ("⚠️ No Gemini key yet: run `python cli.py keys --set GEMINI_API_KEY` " + "and restart, or pick Local.") + model, price = image_price(AUTO_MODEL) + cost = f"~${n * price:.2f}" if price else "price unknown" + return f"Auto → `{model}`: {cost} for {n} shots (estimate, billed to your Google key)." + from studio.doctor import check_comfyui, check_comfyui_models + + up = check_comfyui() + if up.warn: + return f"⚠️ ComfyUI: {up.detail}" + models = check_comfyui_models() + if models and not models.ok: + return f"⚠️ {models.detail}" + return f"✅ ComfyUI is ready: {n} shots, $0." + + +def _quick_rows(start: pipeline.BuildStart) -> tuple[list[tuple[str, str, str]], list[str]]: + """Picker rows (your photos first, then the shots) and what starts ticked: + everything except a photo ① enlarged a lot — a LoRA learns upscaling blur as + the subject's texture. That photo still served as a reference.""" + rows: list[tuple[str, str, str]] = [] + keep: list[str] = [] + for r in start.reports: + if not r.output: + continue + label = f"📷 {r.source.name}" + factor = pipeline.upscale_factor(r) + if factor > pipeline.UPSCALE_UNTICK: + label += f" ⚠ upscaled ×{factor:.1f} from {max(r.original_size)}px" + else: + keep.append(str(r.output)) + rows.append((str(r.output), str(r.output), _flagged(r.output, label))) + for g in start.results: + if g.path and g.path.exists(): + rows.append((str(g.path), str(g.path), _flagged(g.path, g.shot.id))) + keep.append(str(g.path)) + return rows, keep + + +def _quick_note(start: pipeline.BuildStart) -> str: + note = _preprocess_note(start.reports, start.run_dir / "prepped", alpha_cutout=False) + failed = [g for g in start.results if not g.path] + if failed: + note += ("\n\n❌ **Failed shots** (the rest carried on; the ② tab can regenerate " + "them):\n" + "\n".join(f"- `{g.shot.id}`: {g.error}" for g in failed)) + big = [r.source.name for r in start.reports + if r.output and pipeline.upscale_factor(r) > pipeline.UPSCALE_UNTICK] + if big: + note += (f"\n\nUnticked: {', '.join(big)}. ① enlarged them, and a LoRA " + f"learns the upscaling blur as texture. They were still used as " + f"references; tick them to train on them anyway.") + if start.results: + note += "\n\nUntick anything that looks wrong, then press **Finish**." + return note + + +def quick_review(start, keep: list[str]) -> str: + """Coverage of the kept shots and any near-duplicates among everything kept.""" + if not start: + return "" + kept = set(keep or []) + shots = [g.shot for g in start.results if g.path and str(g.path) in kept] + note = coverage(shots) if shots else "" + try: + from studio.dedupe import find_near_duplicate_groups + + groups = find_near_duplicate_groups(sorted(Path(k) for k in kept), max_distance=5) + except Exception: + groups = [] # advisory + if groups: + note += ("\n\n" if note else "") + ( + f"🔁 **{len(groups)} near-duplicate group(s)**, keep one of each: " + f"{_groups_text(groups)}") + return note + + +def do_quick_build(files: list[str], name: str, trigger: str, dataset_type: str, + identity: str, engine: str, held: str, front: bool, + progress=gr.Progress()): + if not files: + raise gr.Error("Drop one or more images of your subject first.") + if not name.strip(): + raise gr.Error("Type a name at the top first: it names the dataset and goes into " + "every prompt.") + trigger = quick_trigger(name, trigger) + log: list[str] = [] + total = len(files) + (0 if dataset_type == "style" else + len(plan_for_type(dataset_type, ""))) + 2 + + def report(msg: str): + log.append(msg) + progress((min(len(log), total), total), desc=msg) + + JOB.start() + try: + start = pipeline.build_start( + [Path(f) for f in files], name.strip(), engine, dataset_type=dataset_type, + exclude_prompt=held.strip(), identity=identity, front=front, should_stop=JOB, + progress=report) + except Exception as e: + raise gr.Error(f"Build failed: {e}") from e + rows, keep = _quick_rows(start) + boxes = gr.CheckboxGroup(choices=[(Path(v).name, v) for _, v, _ in rows], value=keep) + return (start, rows, _picker_gallery(rows, keep), boxes, _quick_note(start), + "\n".join(log), trigger) + + +def do_quick_finish(start, keep: list[str], name: str, trigger: str, dataset_type: str, + identity: str, engine: str, captioner: str, trainer: str, + model_key: str, progress=gr.Progress()): + if not start: + raise gr.Error("Press Build first.") + images = [Path(k) for k in keep or [] if Path(k).exists()] + if not images: + raise gr.Error("Tick at least one image to keep.") + trigger = quick_trigger(name, trigger) + log: list[str] = [] + + def report(msg: str): + log.append(msg) + progress((min(len(log), len(images) + 2), len(images) + 2), desc=msg) + + # A tag-trained base model (SDXL family) wants tag captions. + style = "tags" if _preset(trainer, model_key).expects_tags else "prose" + try: + ds, command = pipeline.build_finish( + start, images, name=name.strip(), trigger=trigger, captioner=captioner, + output_root=settings.output_root, dataset_type=dataset_type, identity=identity, + engine_key=engine, caption_style=style, trainer=trainer, model_key=model_key, + progress=report) + except Exception as e: + raise gr.Error(f"Finish failed: {e}") from e + warnings_ = "".join(f"{m}\n\n" for m in log if m.startswith("⚠️")) + note = (f"✅ **Dataset ready:** `{ds}` ({len(images)} images, trigger `{trigger}`).\n\n" + f"{warnings_}Train with:\n```\n{command}\n```\n`validation/validation.md` in the dataset " + f"says how to pick the best epoch.\n\nConfigs are generated, not " + f"test-trained. Fill any `<>` placeholders first; ⑤ has every setting.") + return note, "\n".join(log), str(ds) + + # ---------- layout ---------- # Gradio 5 warns that `head=` moves to `launch()` in Gradio 6 — but `launch()` does not @@ -1863,6 +2029,51 @@ def _check_for_update(): cap_edit_names = gr.State([]) with gr.Tabs() as tabs: + with gr.Tab("⚡ Quick build", id="quick"): + gr.Markdown( + "Drop a few images, type a name at the top, press **Build**. Untick " + "anything that looks wrong, then **Finish** captions the dataset and " + "writes the trainer config. The numbered tabs run the same steps one " + "at a time, with every option.") + with gr.Row(): + with gr.Column(scale=1): + qb_files = gr.File(label="Images of your subject", file_count="multiple", + file_types=["image"], height=160) + qb_held = gr.Textbox( + label="Anything held to remove?", placeholder="cup, plate", + info="Name objects the subject is holding, so they are cut out " + "and not redrawn in every shot.") + qb_engine = gr.Radio(QUICK_ENGINES, value=settings.default_engine, + label="Generate with") + qb_engine_note = gr.Markdown() + with gr.Row(): + qb_trainer = gr.Dropdown(TRAINER_CHOICES, value="ai-toolkit", + label="Trainer") + qb_model = gr.Dropdown( + [(m.label, m.key) for m in TRAINER_MODELS["ai-toolkit"]], + value=TRAINER_MODELS["ai-toolkit"][0].key, label="Base model") + qb_captioner = gr.Dropdown(CAPTIONER_CHOICES, + value=settings.default_captioner, + label="Captioner") + qb_front = gr.Checkbox(value=False, + label="Prioritize this app's ComfyUI jobs") + btn_qb_build = gr.Button("⚡ Build", variant="primary") + with gr.Column(scale=2): + qb_note = gr.Markdown() + qb_gallery = gr.Gallery( + label="Click a thumbnail to untick/tick it", columns=6, height=420, + allow_preview=False, elem_id="dd-gallery-quick") + qb_zoom = gr.Checkbox(value=False, label="🔍 Zoom on click", + elem_id="dd-zoom-quick") + qb_keep = gr.CheckboxGroup(label="✅ Kept: UNCHECK to leave out", + choices=[], elem_id="dd-picks-quick") + qb_review = gr.Markdown() + btn_qb_finish = gr.Button("✅ Finish: caption, export, write configs", + variant="primary") + qb_result = gr.Markdown() + qb_state = gr.State(None) + qb_rows = gr.State([]) + with gr.Tab("① Preprocess (optional)", id="preprocess"): gr.Markdown("Restore / upscale / isolate source images. Skip this tab entirely " "if your images are already clean.") @@ -1994,7 +2205,8 @@ def _check_for_update(): "other shot's references with it. Helps when your sources " "are partial or poor (face in shadow, body cut off). It " "copies whatever the front view gets wrong (a held object, " - "an outfit) into every shot, so check that view first.") + "an outfit) into every shot, so check that view first, and " + "check full-body shots did not just repeat its stance.") gen_isolate = gr.Checkbox(value=False, label="Isolate generated angle shots (white background)", info="Cut generated angle shots onto white too. " @@ -2502,9 +2714,25 @@ def _force_refresh(current: str): # button or a reload — re-marks the gallery labels. for _gallery, _rows, _boxes, _zoom in ((gen_gallery, gen_rows, keep, gen_zoom), (cap_gallery, cap_rows, cap_select, cap_zoom), - (exp_gallery, exp_rows, exp_select, exp_zoom)): + (exp_gallery, exp_rows, exp_select, exp_zoom), + (qb_gallery, qb_rows, qb_keep, qb_zoom)): _boxes.change(_picker_mark, [_rows, _boxes], [_gallery]) _zoom.change(_set_zoom, [_zoom], [_gallery]) + btn_qb_build.click( + do_quick_build, + [qb_files, project_name, project_trigger, dataset_type, identity_policy, qb_engine, + qb_held, qb_front], + [qb_state, qb_rows, qb_gallery, qb_keep, qb_note, log_box, project_trigger]) + qb_keep.change(quick_review, [qb_state, qb_keep], [qb_review]) + btn_qb_finish.click( + do_quick_finish, + [qb_state, qb_keep, project_name, project_trigger, dataset_type, identity_policy, + qb_engine, qb_captioner, qb_trainer, qb_model], + [qb_result, log_box, tr_dataset]) + qb_trainer.change(_model_dropdown, [qb_trainer], [qb_model]) + for _event in (qb_engine.change, dataset_type.change): + _event(quick_engine_note, [qb_engine, dataset_type], [qb_engine_note]) + demo.load(quick_engine_note, [qb_engine, dataset_type], [qb_engine_note]) for _event in (keep.change, gen_dominance.change): _event(coverage_note, [results_state, keep, gen_dominance], [gen_coverage]) for _event in (gen_files.change, gen_src_folder.change): diff --git a/cli.py b/cli.py index 6e0919d..fb1f21a 100644 --- a/cli.py +++ b/cli.py @@ -24,6 +24,7 @@ from studio import env_keys, pipeline from studio.config import CAPTIONERS, list_images, settings from studio.shotplan import plan_for_type +from studio.trainer_configs import TRAINER_MODELS app = typer.Typer(add_completion=False, help=__doc__) @@ -507,90 +508,63 @@ def build( help="Comma-separated tags to strip from tag captions (e.g. 'watermark, signature')"), anchor: bool = typer.Option(False, "--anchor/--no-anchor", help=_ANCHOR_HELP), front: bool = typer.Option(False, help="Jump ComfyUI's pending queue"), + holdout: int = typer.Option( + 0, help="Keep N real photos out of training, in a -heldout folder, as a likeness test"), + trainer: str = typer.Option( + "", help="Also write ⑤ configs for this trainer (ai-toolkit | kohya | musubi | fizgig)"), + model: str = typer.Option( + "", help="Base model key for --trainer (default: its first). " + + "; ".join(f"{t}: " + ", ".join(m.key for m in ms) + for t, ms in TRAINER_MODELS.items())), ): - """Full pipeline: preprocess -> generate -> caption -> export.""" - from studio.captioner import caption_images - from studio.package import package_dataset - + """Full pipeline: preprocess -> generate -> caption -> export (-> train config).""" style = _check_caption_style(caption_style) dtype = _check_dataset_type(dataset_type) - identity = _check_identity(identity) + identity = "identity" if randomize_outfits else _check_identity(identity) + if trainer and trainer not in TRAINER_MODELS: + raise typer.BadParameter(f"--trainer must be one of {list(TRAINER_MODELS)}") # Check the generation backend before spending a preprocess pass on images # it would then refuse to use. Style never generates, so it never needs one. if dtype != "style": _preflight_comfyui(engine) + n = len(_plan_for(dtype, name, _check_shot_style(shot_style), shot_style_text)) + _echo_cloud_estimate(engine, cloud_model, min(n, max_shots) if max_shots else n) run_dir = pipeline.new_run_dir(name or trigger) typer.echo(f"Run dir: {run_dir}") - # Isolation defaults follow the dataset type (a style is whole-image), the - # same rule the ① tab applies when the header type changes. - do_isolate = (dtype != "style") if isolate is None else isolate - reports = pipeline.preprocess_sources( - _expand(images), run_dir / "prepped", target=target, force_restore=restore, - isolate=do_isolate, subject_prompt=subject_prompt, exclude_prompt=exclude_prompt, - tighten_crop=tighten, front=front, progress=typer.echo) - _echo_preprocess_failures(reports) + start = pipeline.build_start( + _expand(images), name, engine, dataset_type=dtype, run_dir=run_dir, + cloud_model=cloud_model, target=target, restore=restore, isolate=isolate, + subject_prompt=subject_prompt, exclude_prompt=exclude_prompt, tighten=tighten, + isolate_angles=isolate_angles, shot_style=_check_shot_style(shot_style), + shot_style_text=shot_style_text, max_shots=max_shots, identity=identity, + exclude_props=exclude_props, anchor=anchor, front=front, progress=typer.echo) + _echo_preprocess_failures(start.reports) # A per-image failure is reported, not raised, so `output` can be None — - # everything downstream must run on the images that actually exist. - prepped = _report_outputs(reports) - if not prepped: + # everything downstream runs on the images that actually exist. + if not start.prepped: typer.echo("No source image could be preprocessed; aborting.") raise typer.Exit(1) + kept = [r.path for r in start.results if r.path] + if dtype != "style" and not kept: + typer.echo("No shots succeeded; aborting before captioning.") + raise typer.Exit(1) - # Style has no synthetic generation — go straight from the preprocessed - # sources to captioning instead of running (and billing) a character - # turnaround the user can't use. - results = [] - kept: list[Path] = [] - if dtype == "style": - typer.echo("Style dataset: skipping ② generation — captioning your own images.") - else: - shots = _plan_for(dtype, name, _check_shot_style(shot_style), shot_style_text) - if max_shots: - shots = shots[:max_shots] - if _wants_outfits(identity, randomize_outfits, dtype): - shots = _dress(shots, dtype) - _echo_cloud_estimate(engine, cloud_model, len(shots)) - - # Generate from the isolated references; caption the training copies. - refs = [r.reference or r.output for r in reports if r.output] - results = pipeline.generate_shots( - refs, shots, engine, run_dir / "generated", - cloud_model=cloud_model, isolate_angles=isolate_angles, subject_prompt=subject_prompt, - exclude_prompt=exclude_prompt, - exclude_props=_props_default(exclude_props, dtype), front=front, - anchor=anchor, progress=typer.echo) - kept = [r.path for r in results if r.path] - if not kept: - typer.echo("No shots succeeded; aborting before captioning.") - raise typer.Exit(1) - - all_images = prepped + kept - items = caption_images(all_images, captioner, name, trigger, progress=typer.echo, - style=style, prefix=prefix, suffix=suffix, blacklist=drop_tags, - dataset_type=dtype, sparse=sparse, identity=identity) - metadata = { - "character_name": name, - "trigger": trigger, - "dataset_type": dtype, - "shot_style": shot_style, - "shot_style_text": shot_style_text, - "captioner": captioner, - "caption_style": style, - "sources": [str(s) for s in images], - } - if dtype == "character": - metadata["identity"] = identity - if results: # generation ran (character/concept) - metadata["engine"] = engine - metadata["shots"] = [{"id": r.shot.id, "seed": r.seed, "error": r.error} - for r in results] - ds = package_dataset(items, output_root, name, trigger, metadata) + # Caption the training copies (not the isolated refs/) and the generated shots. + ds, command = pipeline.build_finish( + start, start.prepped + kept, name=name, trigger=trigger, captioner=captioner, + output_root=output_root, dataset_type=dtype, identity=identity, engine_key=engine, + shot_style=shot_style, shot_style_text=shot_style_text, caption_style=style, + prefix=prefix, suffix=suffix, drop_tags=drop_tags, sparse=sparse, holdout=holdout, + trainer=trainer, model_key=model, progress=typer.echo) if zip_: from studio.package import zip_dataset typer.echo(f"Zipped: {zip_dataset(ds)}") typer.echo(f"\nDone: {ds}") + if command: + typer.echo(f"\nTrain with ({ds / 'validation' / 'validation.md'} says how to pick " + f"the best epoch):\n{command}") # --- setup / diagnostics --------------------------------------------------- diff --git a/docs/ARCHITECTURE.md b/docs/ARCHITECTURE.md index da1b620..af3bbfd 100644 --- a/docs/ARCHITECTURE.md +++ b/docs/ARCHITECTURE.md @@ -1,9 +1,10 @@ # Architecture -Version: 0.19.0 +Version: 0.20.0 ``` -app.py Gradio UI — thin wiring over the stage functions (5 tabs). +app.py Gradio UI — thin wiring over the stage functions (⚡ Quick + build + the 5 stage tabs). Owns the click-to-pick gallery helpers (_picker_*, _PICKER_SCRIPT/PICKER_IDS — the only browser-side code in the project) and the selection carried ②→③→④; see the @@ -49,7 +50,9 @@ studio/ the SDK's misleading dual-key warning, see the Gotcha) pipeline.py Stage orchestration: preprocess_sources(), generate_shots(). Both isolate per-item failures (one bad image never ends the - batch) and honour the cooperative `should_stop` token + batch) and honour the cooperative `should_stop` token. + build_start() (①→②) and build_finish() (③→④→⑤) are the one + full-build path behind `cli build` and ⚡ Quick build jobs.py Cooperative cancellation: JobControl (a threading.Event stop flag) + should_stop_now(). Checked BETWEEN items by the three batch loops; UI-free so the CLI and tests drive it directly @@ -641,7 +644,10 @@ quietly gets someone else's value. `prompt_id`, so another client's jobs are never mistaken for ours. Every stage that queues ComfyUI work carries the flag: ① and ② each have a "Prioritize" checkbox, and `preprocess`/`generate`/`build` take `--front`. Until 0.17.3 ① had none, so a busy - queue refused every ComfyUI restore/isolation with advice that only worked on ②. + queue refused every ComfyUI restore/isolation with advice that only worked on ②. A job's + timeout restarts while it is still pending (checked every 10 s, for up to `QUEUE_WAIT_S`, an + hour): a front job waits out whatever is running, and before 0.20.0 a restore queued behind + a long user job timed out before it started. - **The bundled workflows use ONLY core nodes.** `isolate_*.json` previously needed `MaskPreview+` from the third-party `ComfyUI_essentials` pack — used purely as a hack to render a white backdrop — which silently broke the workflows for anyone who didn't @@ -937,6 +943,14 @@ quietly gets someone else's value. an isolated copy to `prepped/refs/` that ② generates from. `list_images` is non-recursive, so ③, ④, the ILB hand-off and dataset stats never see `refs/` (a test pins this). ②'s "isolate generated angle shots" is a separate option, default off, for the same reason. +- **A failed isolation costs the reference its cut-out, never the photo (0.20.0).** Before, + SAM3 finding no subject skipped the image outright, training copy included; on the + acceptance pair it dropped the sharpest photo. Now ① writes the training copy as usual and + puts the whole photo in `refs/`, so ② still uses it (background and props included), and ①'s + and ⚡'s notes name it. Other errors (a crash, ComfyUI down) still fail the image. The builtin + backend also retries `person` when the default `character` finds nothing: SAM3 reads + "character" as a drawn figure and scored 0.0 and 0.06 on two real photos where "person" + scored 0.98. - **Identity policy decides whether clothing is part of the subject (0.19.0).** "Identity only" (default) dresses each character shot in a different random outfit and tells the captioner to describe clothing, so the trigger learns the face and body, not one outfit. "Signature @@ -982,6 +996,28 @@ quietly gets someone else's value. as part of the subject. Naming it in ①'s exclude prompt is the fix. - Neither input can supply what it never showed: everything below the hips and the shadowed eye are invented or copied, whatever the settings. + - Local Qwen-Image 2.1 copies a shadow across the face: 10 of a 24-shot acceptance + set (0.20.0) carried the spoon's shadow over one eye, where the 0.17.3 cloud runs on the same + photos showed it in none. Two follow-up probes, same seeds: a positive phrase ("her face is + evenly and softly lit, with both eyes clearly visible") cleared it on one of two shots, so it + was not added to the prompts; anchoring on that clean front view cleared it on a close-up + but collapsed all three full-body shots onto the front view's stance over a garbled + background (the 0.19.0 anchored angle-low did the same). So the advice for an obscured face + is Cloud, not a setting. The 0.20.0 cloud acceptance run (Auto → Nano Banana 2.1, same two + photos, identity policy, "plate, spoon" held) bore it out: 24 of 24 generated, no shadow, + plate or spoon in any shot, glasses on every visible face, back and low views recognisable, + all four framing tiers present. Its one weakness: close-ups get no outfit by design + (`wardrobe.OUTFIT_SHOT_KINDS`), and all seven came back in the same black top. The captions + name the top, which keeps it off the trigger. +- **Quick build is the stage functions with defaults, not a second pipeline (0.20.0).** + `pipeline.build_start` (① then ②) and `build_finish` (③, ④, and ⑤ through + `trainer_configs.default_config`) hold the whole orchestration; `cli build` and the ⚡ tab + both call them, so a fix lands in both. The split is the review point: the user unticks bad + shots between the two calls. Defaults chosen for poor inputs: the highest-resolution source + leads ②'s references (it sets the local output shape); a held-object field feeds ①'s exclude + prompt; a photo ① enlarged more than 1.33x starts unticked (`UPSCALE_UNTICK`, a LoRA learns + upscaling blur as texture) but still serves as a reference; a tag-trained base model gets tag + captions; a blank trigger is made from the name and written back to the header. - **The optimizer defaults to AdamW8bit; Prodigy is the one alternative (0.18.1).** AdamW8bit is the default and the validated recipe in ai-toolkit's UI, musubi's examples, sd-scripts and Fizgig. Prodigy is offered because it removes the learning-rate guess: it runs at lr 1.0 with diff --git a/studio/__init__.py b/studio/__init__.py index fc86838..464bd31 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.19.0" +__version__ = "0.20.0" diff --git a/studio/captioner.py b/studio/captioner.py index a8162b5..99fc269 100644 --- a/studio/captioner.py +++ b/studio/captioner.py @@ -473,7 +473,12 @@ def finalize_caption(raw: str, trigger: str, character_name: str, aliases: list[ # Replace generic nouns the VLM used for the subject with the real name for alias in aliases: caption = re.sub(rf"\bthe {re.escape(alias)}\b", character_name, caption, flags=re.I) - if trigger and not caption.lower().startswith(trigger.lower()): + lead = re.match(rf"{re.escape(trigger)}(?!\w)", caption, flags=re.I) if trigger else None + if lead: + # The trigger's exact casing is the token the LoRA learns and the validation + # prompts use: name "Ann" with trigger "ann" must not lead with "Ann". + caption = trigger + caption[lead.end():] + elif trigger: keep_name = (dataset_type == "character" and character_name and caption.startswith(character_name.split()[0])) if caption and not keep_name: diff --git a/studio/comfy_api.py b/studio/comfy_api.py index 12d18c7..5a99216 100644 --- a/studio/comfy_api.py +++ b/studio/comfy_api.py @@ -22,6 +22,8 @@ class ComfyError(Exception): # Consecutive failed /history polls before a running job is given up on (each poll # waits up to 30 s, so a stalled server gets minutes, not seconds). MAX_STALLED_POLLS = 5 +# Longest a queued job may wait for other jobs before its own timeout counts down. +QUEUE_WAIT_S = 3600 # Model filenames inside the bundled templates, remapped to whatever the user @@ -261,6 +263,15 @@ def queue_backlog() -> int: return 0 +def _is_pending(prompt_id: str) -> bool: + """True while `prompt_id` is still waiting in ComfyUI's queue (not yet running).""" + try: + q = httpx.get(f"{settings.comfy_url}/queue", timeout=10).json() + except Exception: + return False + return any(len(j) > 1 and j[1] == prompt_id for j in q.get("queue_pending", [])) + + def run_prompt(graph: dict, timeout: float = 600.0, front: bool = False) -> list[dict]: """Queue an API-format graph, wait for completion, return output image refs. @@ -296,7 +307,9 @@ def run_prompt(graph: dict, timeout: float = 600.0, front: bool = False) -> list raise ComfyError(describe_rejection(r.text)) prompt_id = r.json()["prompt_id"] - deadline = time.monotonic() + timeout + queued_at = time.monotonic() + deadline = queued_at + timeout + next_queue_check = queued_at stalled = 0 while time.monotonic() < deadline: # ComfyUI stops answering HTTP for tens of seconds while it swaps models in @@ -339,6 +352,16 @@ def run_prompt(graph: dict, timeout: float = 600.0, front: bool = False) -> list return images if status.get("completed"): raise ComfyError("run completed but produced no images") + else: + # Waiting in line behind someone else's job is not ours being slow: a + # front job sat behind a long user job and timed out before it started + # (0.20.0 acceptance). The clock restarts while ours is still pending, + # for up to QUEUE_WAIT_S in line. + now = time.monotonic() + if now >= next_queue_check and now - queued_at < QUEUE_WAIT_S: + next_queue_check = now + 10 + if _is_pending(prompt_id): + deadline = now + timeout time.sleep(1.5) raise ComfyError( f"ComfyUI did not finish within {timeout:g}s (prompt {prompt_id}). It may " diff --git a/studio/config.py b/studio/config.py index c3198f2..b784c35 100644 --- a/studio/config.py +++ b/studio/config.py @@ -8,6 +8,7 @@ import io import json +import os from datetime import datetime, timedelta, timezone from pathlib import Path @@ -21,6 +22,21 @@ # GROQ_API_KEY, HF_TOKEN for the gated SAM3 download) work from there too. load_dotenv(REPO_ROOT / ".env") + +def _alias_hf_token() -> None: + """Hugging Face's libraries read HF_TOKEN, then the `hf auth login` file. Take the + other common name only when neither exists: an env token overrides the login, so a + stale HUGGINGFACE_API_KEY would otherwise break a login that works.""" + if os.environ.get("HF_TOKEN") or not os.environ.get("HUGGINGFACE_API_KEY"): + return + from huggingface_hub import get_token + + if not get_token(): + os.environ["HF_TOKEN"] = os.environ["HUGGINGFACE_API_KEY"] + + +_alias_hf_token() + # Base captioning instruction, tuned for LoRA training captions: describe what # VARIES between images (pose/angle/setting/lighting), not the subject's fixed # identity — identity is what the trigger token learns. diff --git a/studio/isolate.py b/studio/isolate.py index bd737f7..0428d8c 100644 --- a/studio/isolate.py +++ b/studio/isolate.py @@ -55,7 +55,8 @@ def _load_sam3(): except Exception as e: # An HF_TOKEN (from .env or the environment) wins over `hf auth login`, # so a stale one 401s even when the saved login would work. - hint = (" An HF_TOKEN is set (in .env or the environment) and is used " + hint = (" An HF_TOKEN is set (in .env or the environment, or taken from " + "HUGGINGFACE_API_KEY when there is no login) and is used " "instead of `hf auth login` — replace it (`python cli.py keys --set " "HF_TOKEN`) or remove it." if os.environ.get("HF_TOKEN") else "") raise IsolationError( @@ -149,6 +150,10 @@ def isolate_builtin(image_path: Path, out_path: Path, subject_prompt: str = "cha alpha_cutout: bool = False, label: str = "") -> Path: image = Image.open(image_path).convert("RGB") subject = _segment(image, subject_prompt) + if not subject.any() and subject_prompt == "character": + # SAM3 reads "character" as a drawn figure: on two real photos it scored + # 0.0 and 0.06 where "person" scored 0.98 (0.20.0 acceptance). + subject = _segment(image, "person") if not subject.any(): # `label` names the file the USER recognises. By the time preprocess calls # this, `image_path` can be a restored intermediate ("cat_prepped.png") diff --git a/studio/pipeline.py b/studio/pipeline.py index 7568ef1..a66649d 100644 --- a/studio/pipeline.py +++ b/studio/pipeline.py @@ -89,7 +89,7 @@ def preprocess_sources( progress(f" SKIPPED {src.name}: {e}") reports.append(failed_report(src, str(e))) continue - extra = f", isolated reference in {REFS_DIR}/" if rep.reference else "" + extra = f", isolated reference in {REFS_DIR}/" if rep.isolated else "" progress( f" {src.name}: {rep.original_size[0]}x{rep.original_size[1]} -> " f"{rep.final_size[0]}x{rep.final_size[1]} ({rep.reason}{extra})" @@ -191,3 +191,159 @@ def generate_shots( ok = sum(1 for r in results if r.path) progress(f"Generation done: {ok} succeeded, {len(results) - ok} failed.") return results + + +# ---------- full build: ①→② then ③→④(→⑤), shared by `cli build` and Quick build ---------- + +# A source ① enlarged more than this starts unticked in Quick build: upscaling +# invents detail, and the LoRA learns that as the subject's texture. It still +# serves as a reference. +UPSCALE_UNTICK = 1.33 + + +def upscale_factor(report: PreprocessReport) -> float: + """How much ① enlarged this source (1.0 = not at all, or shrunk).""" + return max(report.final_size) / max(1, *report.original_size) + + +@dataclass +class BuildStart: + """What ①→② made: the run folder, ①'s reports, ②'s references and results.""" + run_dir: Path + reports: list[PreprocessReport] + refs: list[Path] + results: list[GenResult] + + @property + def prepped(self) -> list[Path]: + """The training copies ① wrote (a skipped source has none).""" + return [r.output for r in self.reports if r.output] + + +def build_start( + images: list[Path], + name: str, + engine_key: str, + *, + dataset_type: str = "character", + run_dir: Path | None = None, + cloud_model: str = "", + target: int | None = None, + restore: bool | None = None, + isolate: bool | None = None, + subject_prompt: str = "character", + exclude_prompt: str = "", + tighten: bool = False, + isolate_angles: bool = False, + shot_style: str = "match", + shot_style_text: str = "", + max_shots: int = 0, + identity: str = "identity", + exclude_props: bool | None = None, + anchor: bool = False, + front: bool = False, + should_stop: ShouldStop | None = None, + progress: ProgressFn = print, +) -> BuildStart: + """① preprocess, then ② generate from the isolated references. + + The highest-resolution source leads the references: the local engine takes + its output aspect from the first one, and at one reference it is the only one. + Style datasets stop after ① (an aesthetic can't be generated from a reference). + """ + from studio.shotplan import plan_for_type + from studio.wardrobe import dress + + run_dir = run_dir or new_run_dir(name) + reports = preprocess_sources( + images, run_dir / "prepped", target=target, force_restore=restore, + isolate=(dataset_type != "style") if isolate is None else isolate, + subject_prompt=subject_prompt, exclude_prompt=exclude_prompt, tighten_crop=tighten, + front=front, should_stop=should_stop, progress=progress) + ok = sorted((r for r in reports if r.output), + key=lambda r: r.original_size[0] * r.original_size[1], reverse=True) + start = BuildStart(run_dir, reports, [r.reference or r.output for r in ok], []) + if dataset_type == "style": + progress("Style dataset: skipping ② generation — captioning your own images.") + return start + if not start.refs or should_stop_now(should_stop): + return start + shots = plan_for_type(dataset_type, name, shot_style, shot_style_text) + if max_shots: + shots = shots[:max_shots] + if identity == "identity" and dataset_type == "character": + shots = dress(shots) + start.results = generate_shots( + start.refs, shots, engine_key, run_dir / "generated", cloud_model=cloud_model, + isolate_angles=isolate_angles, subject_prompt=subject_prompt, + exclude_prompt=exclude_prompt, + exclude_props=(dataset_type == "character") if exclude_props is None else exclude_props, + front=front, anchor=anchor, should_stop=should_stop, progress=progress) + return start + + +def build_finish( + start: BuildStart, + images: list[Path], + *, + name: str, + trigger: str, + captioner: str, + output_root: Path, + dataset_type: str = "character", + identity: str = "identity", + engine_key: str = "", + shot_style: str = "match", + shot_style_text: str = "", + caption_style: str = "prose", + prefix: str = "", + suffix: str = "", + drop_tags: str = "", + sparse: bool = False, + holdout: int = 0, + trainer: str = "", + model_key: str = "", + progress: ProgressFn = print, +) -> tuple[Path, str]: + """③ caption `images`, ④ export them, and (with `trainer`) ⑤ write its configs. + + Returns the dataset folder and the train command ("" without a trainer). + """ + from studio.captioner import caption_images + from studio.package import package_dataset + + items = caption_images(images, captioner, name, trigger, progress=progress, + style=caption_style, prefix=prefix, suffix=suffix, + blacklist=drop_tags, dataset_type=dataset_type, sparse=sparse, + identity=identity) + metadata: dict = { + "character_name": name, + "trigger": trigger, + "dataset_type": dataset_type, + "shot_style": shot_style, + "shot_style_text": shot_style_text, + "captioner": captioner, + "caption_style": caption_style, + "sources": [str(r.source) for r in start.reports], + } + if dataset_type == "character": + metadata["identity"] = identity + if start.results: # generation ran (character/concept) + metadata["engine"] = engine_key + metadata["shots"] = [{"id": r.shot.id, "seed": r.seed, "error": r.error} + for r in start.results] + ds = package_dataset(items, output_root, name, trigger, metadata, holdout=holdout) + command = "" + if trainer: + from studio.dataset_stats import inspect + from studio.trainer_configs import TRAINER_MODELS, default_config + + _, command = default_config(ds, trainer, model_key, name=name, trigger=trigger, + dataset_type=dataset_type, shot_style=shot_style, + shot_style_text=shot_style_text) + preset = next((p for p in TRAINER_MODELS[trainer] if p.key == model_key), + TRAINER_MODELS[trainer][0]) + for line in inspect(ds).upscale_note(preset.resolution).splitlines(): + if line: + progress(line) # the ⑤ tab's undersized/upscaled warnings + return ds, command diff --git a/studio/preprocess.py b/studio/preprocess.py index bc3a05c..984c78f 100644 --- a/studio/preprocess.py +++ b/studio/preprocess.py @@ -20,7 +20,7 @@ from studio.config import settings from studio.dataset_stats import PROVENANCE_KEY -from studio.isolate import isolate_subject +from studio.isolate import IsolationError, isolate_subject BLUR_THRESHOLD = 120.0 # Laplacian variance below this = soft/degraded image # Subfolder of ①'s output holding the isolated copies ② uses as references. @@ -225,13 +225,23 @@ def save(img: Image.Image, path: Path, isolated: bool) -> None: "captured": captured, "restored": restore, "isolated": isolated})) img.save(path, "PNG", pnginfo=info) + isolated = False if ref_path: # The isolated copy is ②'s reference only: it keeps props and the - # old background out of the generated shots. + # old background out of the generated shots. So a failed isolation + # costs the reference its cut-out, never the photo: ② falls back to + # the whole image and the training copy is written as usual. refs_dir.mkdir(exist_ok=True) - isolate_subject(stage_path, ref_path, subject_prompt, exclude_prompt, - backend=isolation_backend, progress=progress, - alpha_cutout=alpha_cutout, label=source.name, front=front) + try: + isolate_subject(stage_path, ref_path, subject_prompt, exclude_prompt, + backend=isolation_backend, progress=progress, + alpha_cutout=alpha_cutout, label=source.name, front=front) + isolated = True + except IsolationError as e: + reason = (f"{reason}; " if reason else "") + ( + f"NOT isolated ({e}) — ② uses the whole photo, background and props " + f"included") + if ref_path and isolated: with Image.open(ref_path) as im: ref = im.copy() if alpha_cutout else im.convert("RGB") # keep RGBA from studio.isolate import touches_bottom @@ -248,6 +258,8 @@ def save(img: Image.Image, path: Path, isolated: bool) -> None: with Image.open(stage_path) as im: img = _resize_to_target(im.convert("RGB"), target) save(img, out_path, False) + if ref_path and not isolated: + save(img, ref_path, False) except BaseException: out_path.unlink(missing_ok=True) if ref_path: @@ -263,7 +275,7 @@ def save(img: Image.Image, path: Path, isolated: bool) -> None: final_size=img.size, restored=restore, reason=reason or "clean source, resize only", - isolated=isolate, + isolated=isolated, reference=ref_path, cut_off=cut_off, ) diff --git a/studio/trainer_configs.py b/studio/trainer_configs.py index 9ec1cc2..ba1a45d 100644 --- a/studio/trainer_configs.py +++ b/studio/trainer_configs.py @@ -727,3 +727,32 @@ def write_configs(cfg: TrainConfig, install_path: str = "") -> tuple[list[Path], 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 + + +def default_config(dataset_dir: Path, trainer: str, model_key: str, *, name: str, + trigger: str, dataset_type: str = "character", + shot_style: str = "match", shot_style_text: str = "", + install_path: str = "") -> tuple[list[Path], str]: + """⑤ with every field at the preset's default: Quick build and `cli build --trainer`. + + Same derivation ⑤ makes on its own: buckets from the dataset's real sizes, + repeats sized to the step target (Fizgig: 1, its own guidance). + """ + from studio.dataset_stats import inspect + from studio.package import slugify + + if trainer not in TRAINER_MODELS: + raise ValueError(f"Unknown trainer: {trainer}") + preset = next((p for p in TRAINER_MODELS[trainer] if p.key == model_key), + TRAINER_MODELS[trainer][0]) + stats = inspect(dataset_dir) + repeats = 1 if trainer == "fizgig" else stats.suggested_repeats(preset.epochs, + preset.batch_size) + cfg = TrainConfig( + trainer=trainer, model=preset, dataset_dir=dataset_dir, trigger=trigger, + name=slugify(name) or "lora", dataset_type=dataset_type, + resolution=preset.resolution, rank=preset.rank, alpha=preset.alpha, + epochs=preset.epochs, num_repeats=repeats, n_images=stats.n_images, lr=preset.lr, + batch_size=preset.batch_size, buckets=stats.buckets_for(preset.resolution), + shot_style=shot_style, shot_style_text=shot_style_text) + return write_configs(cfg, install_path) diff --git a/tests/test_caption_style.py b/tests/test_caption_style.py index a0ae7cb..1aee550 100644 --- a/tests/test_caption_style.py +++ b/tests/test_caption_style.py @@ -110,6 +110,15 @@ def test_finalize_prose_still_prepends_trigger_and_names() -> None: assert "Sy Snootles" in out # alias replacement still runs in prose mode +def test_finalize_prose_leads_with_the_triggers_exact_casing() -> None: + # Quick build makes the trigger from the name, so a one-word name differs only in case. + out = finalize_caption("The woman sits at a table.", "ann", "Ann", SUBJECT_ALIASES) + assert out == "ann sits at a table." + # A longer word that merely starts with the trigger is not the trigger. + assert finalize_caption("Annika sits.", "ann", "Annika", SUBJECT_ALIASES) == \ + "ann, Annika sits." + + def test_finalize_default_style_is_prose() -> None: # No style arg → prose behaviour: the sentence's first letter is lowercased # after the trigger. The tags path would instead lowercase the whole thing. diff --git a/tests/test_cli_build_types.py b/tests/test_cli_build_types.py index 4f88ce2..0801e62 100644 --- a/tests/test_cli_build_types.py +++ b/tests/test_cli_build_types.py @@ -57,7 +57,7 @@ def fake_caption_images(images, *a, **kw): calls["captioned"] = list(images) return [(p, "caption") for p in images] - def fake_package(items, output_root, name, trigger, metadata): + def fake_package(items, output_root, name, trigger, metadata, holdout=0): calls["metadata"] = metadata ds = Path(output_root) / "ds" ds.mkdir(parents=True, exist_ok=True) @@ -145,3 +145,54 @@ def test_generate_command_builds_a_concept_set(stubs: dict, tmp_path: Path) -> N "--out", str(tmp_path / "gen")]) assert result.exit_code == 0, result.output assert len(stubs["shots"]) == 18 + + +def test_build_orders_references_largest_first_and_writes_train_config( + stubs: dict, tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: + """The highest-resolution source leads ②'s references; --trainer writes ⑤.""" + from studio import trainer_configs + + small, big = tmp_path / "small.png", tmp_path / "big.png" + for p in (small, big): + p.write_bytes(b"x") + + def fake_preprocess(sources, out_dir, **kw): + out_dir.mkdir(parents=True, exist_ok=True) + reps = [] + for src, size in zip(sources, [(580, 580), (3000, 4000)], strict=True): + out = out_dir / f"{src.stem}_prepped.png" + out.write_bytes(b"x") + reps.append(PreprocessReport(source=src, output=out, original_size=size, + final_size=(1024, 1024), restored=False, + reason="ok", reference=out_dir / "refs" / out.name)) + return reps + + written: dict = {} + + def fake_default_config(ds, trainer, model_key, **kw): + written.update(trainer=trainer, model=model_key, **kw) + return [], "train.py --go" + + monkeypatch.setattr(cli.pipeline, "preprocess_sources", fake_preprocess) + monkeypatch.setattr(trainer_configs, "default_config", fake_default_config) + from studio import dataset_stats + + class _Stats: + def upscale_note(self, resolution: int) -> str: + return f"\n⚠️ small_prepped was upscaled (bar {resolution})" + + monkeypatch.setattr(dataset_stats, "inspect", lambda ds: _Stats()) + result = runner.invoke(cli.app, [ + "build", str(small), str(big), "--output-root", str(tmp_path / "out"), + "--name", "Ann", "--trainer", "fizgig", "--model", "krea2"]) + + assert result.exit_code == 0, result.output + assert [p.name for p in stubs["sources"]] == ["big_prepped.png", "small_prepped.png"] + assert written["trainer"] == "fizgig" and written["model"] == "krea2" + assert "train.py --go" in result.output + assert "⚠️ small_prepped was upscaled" in result.output # ⑤'s warning reaches the CLI + + +def test_build_rejects_an_unknown_trainer(stubs: dict, tmp_path: Path) -> None: + result = runner.invoke(cli.app, ["build", str(stubs["src"]), "--trainer", "nope"]) + assert result.exit_code != 0 and stubs["generated"] == 0 diff --git a/tests/test_comfy_diagnostics.py b/tests/test_comfy_diagnostics.py index f0db2cf..97538a2 100644 --- a/tests/test_comfy_diagnostics.py +++ b/tests/test_comfy_diagnostics.py @@ -177,6 +177,32 @@ def history(*a, **k): assert comfy_api.run_prompt({}) == [{"filename": "o.png"}] +def test_time_waiting_behind_another_job_does_not_count( + monkeypatch: pytest.MonkeyPatch) -> None: + """A front job sat pending behind a long user job and timed out before it ran.""" + clock = {"t": 0.0} + monkeypatch.setattr(comfy_api.time, "monotonic", lambda: clock["t"]) + polls = {"n": 0} + + def get(url, *a, **k): + req = httpx.Request("GET", url) + if url.endswith("/queue"): # pending for the first 100 s, then running + pending = [[-1, "p1", {}]] if clock["t"] < 100 else [] + return httpx.Response(200, json={"queue_pending": pending}, request=req) + polls["n"] += 1 + clock["t"] += 5 + done = {"p1": {"status": {"completed": True}, + "outputs": {"8": {"images": [{"filename": "o.png"}]}}}} + return httpx.Response(200, json=done if clock["t"] >= 130 else {}, request=req) + + _queue_stubs(monkeypatch, get) + assert comfy_api.run_prompt({}, timeout=60) == [{"filename": "o.png"}] + with pytest.raises(comfy_api.ComfyError, match="did not finish within 60s"): + clock["t"] = 0.0 + monkeypatch.setattr(comfy_api, "_is_pending", lambda pid: False) + comfy_api.run_prompt({}, timeout=60) + + def test_a_server_that_stays_silent_is_given_up_on(monkeypatch: pytest.MonkeyPatch) -> None: calls = {"n": 0} diff --git a/tests/test_env_keys.py b/tests/test_env_keys.py index e959839..8ddcef5 100644 --- a/tests/test_env_keys.py +++ b/tests/test_env_keys.py @@ -8,6 +8,7 @@ from __future__ import annotations +import os from pathlib import Path import pytest @@ -251,3 +252,20 @@ def test_settings_reads_an_lds_prefixed_hf_token(tmp_path: Path) -> None: s = Settings(_env_file=None, hf_token="hf_aliased") assert s.resolved_key("HF_TOKEN") == "hf_aliased" + + +# --- HUGGINGFACE_API_KEY alias --------------------------------------------- + +@pytest.mark.parametrize("login, expected", [(None, "alias"), ("saved", None)]) +def test_hf_alias_never_overrides_a_saved_login(monkeypatch: pytest.MonkeyPatch, + login: str | None, + expected: str | None) -> None: + import huggingface_hub + + from studio import config + + monkeypatch.delenv("HF_TOKEN", raising=False) + monkeypatch.setenv("HUGGINGFACE_API_KEY", "alias") + monkeypatch.setattr(huggingface_hub, "get_token", lambda: login) + config._alias_hf_token() + assert os.environ.get("HF_TOKEN") == expected diff --git a/tests/test_isolate.py b/tests/test_isolate.py index 0e8a573..37fa92c 100644 --- a/tests/test_isolate.py +++ b/tests/test_isolate.py @@ -132,6 +132,22 @@ def test_missing_prop_reports_rather_than_silently_passing( assert any("not found" in m for m in log) +def test_a_photo_with_no_character_falls_back_to_person( + scene: Path, tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + lookup = _masks() + asked: list[str] = [] + + def segment(i, p, threshold=0.5): + asked.append(p) + return np.zeros((40, 40), dtype=bool) if p == "character" else lookup["character"] + + monkeypatch.setattr(iso, "_segment", segment) + out = tmp_path / "out.png" + iso.isolate_builtin(scene, out, subject_prompt="character") + assert asked == ["character", "person"] and _kept(out) == 800 + + def test_no_subject_raises_actionable_error(scene: Path, tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: monkeypatch.setattr(iso, "_segment", diff --git a/tests/test_preprocess.py b/tests/test_preprocess.py index 099c0a2..2969a03 100644 --- a/tests/test_preprocess.py +++ b/tests/test_preprocess.py @@ -5,6 +5,7 @@ from pathlib import Path import numpy as np +import pytest from PIL import Image from studio.preprocess import preprocess @@ -124,13 +125,12 @@ def fake_preprocess(source, work_dir, **kwargs): # ---------- a failed stage must leave nothing behind ---------- -def test_preprocess_removes_the_partial_output_when_isolation_fails( +def test_preprocess_removes_the_partial_output_when_isolation_crashes( tmp_path: Path, monkeypatch) -> None: """Restoration writes out_path BEFORE isolation runs. When isolation then - fails, that restored-but-not-isolated, not-resized file must not survive — + crashes, that restored-but-not-resized file must not survive — `list_images` would serve it to ②/③ as a finished source.""" import studio.preprocess as pp - from studio.isolate import IsolationError src_dir = tmp_path / "src" src_dir.mkdir() @@ -143,26 +143,23 @@ def fake_restore(source, out_path, upscale=True, front=False): return out_path def fake_isolate(image_path, out_path, *a, **kw): - raise IsolationError("SAM3 found no 'character'") + raise RuntimeError("CUDA out of memory") monkeypatch.setattr(pp, "_restore_comfyui", fake_restore) monkeypatch.setattr(pp, "isolate_subject", fake_isolate) - try: + with pytest.raises(RuntimeError, match="out of memory"): pp.preprocess(src, out, target=256, force_restore=True, isolate=True, restore_backend="comfyui") - except IsolationError: - pass - else: - raise AssertionError("expected the isolation failure to propagate") assert list(out.glob("*.png")) == [], "a half-processed file was left behind" -def test_isolation_error_names_the_users_file_not_the_intermediate( +def test_no_subject_keeps_the_photo_and_uses_it_whole_as_the_reference( tmp_path: Path, monkeypatch) -> None: - """preprocess passes the source name through: naming the deleted - 'cat_prepped.png' intermediate sent users hunting for a file that is gone.""" + """Isolation only makes ②'s reference, so SAM3 finding no subject must not + drop the photo (0.20.0 acceptance: it dropped the sharpest source). The + note names the user's file, not the 'cat_prepped.png' intermediate.""" import studio.preprocess as pp from studio.isolate import IsolationError @@ -175,13 +172,13 @@ def fake_isolate(image_path, out_path, subject_prompt="character", *a, label="", raise IsolationError(f"SAM3 found no '{subject_prompt}' in {label}") monkeypatch.setattr(pp, "isolate_subject", fake_isolate) - try: - pp.preprocess(src, tmp_path / "out", target=64, force_restore=False, isolate=True) - except IsolationError as e: - assert "cat.png" in str(e) - assert "prepped" not in str(e) - else: - raise AssertionError("expected the isolation failure to propagate") + rep = pp.preprocess(src, tmp_path / "out", target=64, force_restore=False, isolate=True) + assert rep.output.exists() and rep.reference and rep.reference.exists() + assert not rep.isolated + assert "NOT isolated" in rep.reason and "cat.png" in rep.reason + assert "prepped" not in rep.reason + with Image.open(rep.reference) as ref, Image.open(rep.output) as out: + assert ref.size == out.size and ref.tobytes() == out.tobytes() # ---------- 0.17.3: phone photos (size, JPEG, EXIF rotation) ---------- diff --git a/tests/test_quick_build.py b/tests/test_quick_build.py new file mode 100644 index 0000000..97ed321 --- /dev/null +++ b/tests/test_quick_build.py @@ -0,0 +1,111 @@ +"""⚡ Quick build: one tab that runs ①→② then ③→④→⑤ with defaults. Every stage stubbed.""" + +from __future__ import annotations + +from pathlib import Path + +import pytest + +import app as A +from studio import pipeline +from studio.preprocess import PreprocessReport +from studio.shotplan import default_plan + + +def _start(tmp_path: Path) -> pipeline.BuildStart: + def png(name: str) -> Path: + p = tmp_path / name + p.write_bytes(b"x") + return p + + big = PreprocessReport(source=tmp_path / "big.jpg", output=png("big_prepped.png"), + original_size=(3000, 4000), final_size=(768, 1024), + restored=False, reason="ok") + small = PreprocessReport(source=tmp_path / "small.jpg", output=png("small_prepped.png"), + original_size=(580, 580), final_size=(1024, 1024), + restored=True, reason="upscaled") + shots = default_plan()[:2] + results = [pipeline.GenResult(shots[0], png("angle-front.png"), 1), + pipeline.GenResult(shots[1], None, 2, error="refused")] + return pipeline.BuildStart(tmp_path, [big, small], [big.output], results) + + +def test_trigger_comes_from_the_name_when_blank() -> None: + assert A.quick_trigger("Sy Snootles", "") == "sysnootles" + assert A.quick_trigger("Sy Snootles", " sy_s ") == "sy_s" + + +def test_an_upscaled_photo_starts_unticked_and_failures_are_named(tmp_path: Path) -> None: + start = _start(tmp_path) + rows, keep = A._quick_rows(start) + assert [Path(v).name for _, v, _ in rows] == ["big_prepped.png", "small_prepped.png", + "angle-front.png"] + assert [Path(k).name for k in keep] == ["big_prepped.png", "angle-front.png"] + assert "upscaled ×1.8 from 580px" in rows[1][2] + note = A._quick_note(start) + assert "refused" in note and "Unticked: small.jpg" in note + + +def test_a_failed_isolation_is_named_in_the_note(tmp_path: Path) -> None: + start = _start(tmp_path) + start.reports[0].reference = start.reports[0].output + start.reports[0].isolated = False + assert "Not isolated:** big.jpg" in A._quick_note(start) + + +def test_build_derives_the_trigger_and_returns_the_picker( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: + seen: dict = {} + + def fake_start(images, name, engine, **kw): + seen.update(images=images, name=name, engine=engine, **kw) + return _start(tmp_path) + + monkeypatch.setattr(pipeline, "build_start", fake_start) + out = A.do_quick_build([str(tmp_path / "a.jpg")], "Ann Lee", "", "character", + "identity", "comfyui", " plate ", False, progress=lambda *a, **k: None) + state, rows, gallery, boxes, note, log, trigger = out + assert trigger == "annlee" and seen["exclude_prompt"] == "plate" + assert seen["identity"] == "identity" and state.refs + assert len(gallery) == 3 and boxes.value == [r for _, r, _ in rows if "big" in r + or "angle" in r] + + +def test_build_needs_images_and_a_name() -> None: + import gradio as gr + + with pytest.raises(gr.Error, match="Drop"): + A.do_quick_build([], "Ann", "", "character", "identity", "gemini", "", False) + with pytest.raises(gr.Error, match="name"): + A.do_quick_build(["x.png"], " ", "", "character", "identity", "gemini", "", False) + + +def test_finish_captions_tags_for_a_tag_model_and_writes_the_config( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: + seen: dict = {} + + def fake_finish(start, images, **kw): + seen.update(images=images, **kw) + return tmp_path / "ds", "train --go" + + monkeypatch.setattr(pipeline, "build_finish", fake_finish) + start = _start(tmp_path) + keep = [str(start.reports[0].output)] + note, _log, ds = A.do_quick_finish(start, keep, "Ann", "", "character", "identity", + "comfyui", "qwen3vl", "kohya", "sdxl", + progress=lambda *a, **k: None) + assert seen["caption_style"] == "tags" and seen["trainer"] == "kohya" + assert seen["trigger"] == "ann" and [p.name for p in seen["images"]] == ["big_prepped.png"] + assert "train --go" in note and ds.endswith("ds") + + +def test_the_quick_tab_is_first_and_wired() -> None: + names = [getattr(d.fn, "__name__", "") for d in A.demo.fns.values()] + assert "do_quick_build" in names and "do_quick_finish" in names + assert ("dd-gallery-quick", "dd-picks-quick", "dd-zoom-quick") in A.PICKER_IDS + + +def test_engine_note_prices_cloud_and_skips_style(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(A.settings, "gemini_api_key", "k") + assert "$" in A.quick_engine_note("gemini", "character") + assert "generate nothing" in A.quick_engine_note("gemini", "style") diff --git a/tests/test_trainer_configs.py b/tests/test_trainer_configs.py index b5e09d9..fb32f1e 100644 --- a/tests/test_trainer_configs.py +++ b/tests/test_trainer_configs.py @@ -258,3 +258,17 @@ def test_fizgig_never_offers_prodigy() -> None: assert [k for _, k in optimizer_choices("fizgig")] == ["adamw8bit"] assert [k for _, k in optimizer_choices("kohya")] == ["adamw8bit", "prodigy"] + + +def test_default_config_writes_the_presets_defaults(tmp_path: Path) -> None: + from PIL import Image + + from studio.trainer_configs import default_config + + for i in range(3): + Image.new("RGB", (64, 96)).save(tmp_path / f"{i:02d}.png") + written, command = default_config(tmp_path, "fizgig", "krea2", name="Ann Lee", + trigger="annlee") + names = {p.name for p in written} + assert "fizgig-dataset.toml" in names and "validation_prompts.txt" in names + assert "--max_train_epochs 30" in command and "ann-lee" in command