diff --git a/AGENTS.md b/AGENTS.md index 9035ad8..bdea4cc 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.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). +- **Status:** in active development; main is v0.18.1 (git tag `v0.18.1`). **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. +- **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 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; 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. diff --git a/app.py b/app.py index 3ac9885..6276a5d 100644 --- a/app.py +++ b/app.py @@ -55,7 +55,7 @@ ) from studio.jobs import JobControl from studio.shotplan import Shot, apply_prop_exclusion, apply_wardrobe, plan_for_type -from studio.trainer_configs import TRAINER_MODELS, TRAINERS +from studio.trainer_configs import TRAINER_MODELS, TRAINERS, optimizer_choices TRAINER_CHOICES = [(label, key) for key, label in TRAINERS.items()] @@ -1409,7 +1409,8 @@ 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.epochs, p.lr, p.batch_size) + p.resolution, p.rank, p.alpha, p.epochs, p.lr, p.batch_size, + gr.update(choices=optimizer_choices(trainer), value="adamw8bit")) def on_model_change(trainer: str, model_key: str): @@ -1515,7 +1516,7 @@ def do_generate_train_config(trainer: str, model_key: str, dataset_dir: str, multi_res: bool, dataset_type: str = "character", style_key: str = shot_style.MATCH, style_text: str = "", project_name: str = "", - repeats=0) -> str: + repeats=0, optimizer: str = "adamw8bit") -> 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 @@ -1545,13 +1546,15 @@ def do_generate_train_config(trainer: str, model_key: str, dataset_dir: str, preset = _preset(trainer, model_key) buckets = stats.buckets_for(int(resolution)) if multi_res else [] num_repeats = _repeats(trainer, stats, epochs, batch_size, repeats) + if optimizer not in {k for _, k in optimizer_choices(trainer)}: + optimizer = "adamw8bit" 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), epochs=max(1, int(epochs)), num_repeats=num_repeats, n_images=stats.n_images, - lr=float(lr), batch_size=int(batch_size), + lr=float(lr), optimizer=optimizer, 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()) @@ -1561,7 +1564,8 @@ def do_generate_train_config(trainer: str, model_key: str, dataset_dir: str, user_config.set_last_train_settings({ "trainer": trainer, "model": model_key, "resolution": int(resolution), "rank": int(rank), "alpha": int(alpha), "epochs": cfg.epochs, - "repeats": int(repeats or 0), "lr": float(lr), "batch_size": int(batch_size)}) + "repeats": int(repeats or 0), "lr": float(lr), "optimizer": optimizer, + "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") @@ -1569,6 +1573,12 @@ def do_generate_train_config(trainer: str, model_key: str, dataset_dir: str, 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 optimizer == "prodigy": + caveat += ("\n\nProdigy finds its own step size, so the config uses learning rate " + "1.0 and the Learning rate field is ignored.") + if trainer in ("musubi", "kohya"): + caveat += (f" {trainer} doesn't install Prodigy: run `pip install prodigyopt` in " + "its environment first.") if trainer in ("musubi", "fizgig"): caveat += (f"\n\n⚠️ {trainer} needs your local model paths — fill the " "<> placeholders in the command before running.") @@ -2247,8 +2257,14 @@ def _check_for_update(): 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.") + with gr.Row(): + tr_lr = gr.Number(value=_ai_presets[0].lr, label="Learning rate", + info="1e-4 is a common starting point.") + tr_optimizer = gr.Dropdown( + optimizer_choices("ai-toolkit"), value="adamw8bit", + label="Optimizer", + info="AdamW8bit is every trainer's tested default. Prodigy " + "picks its own step size and ignores the learning rate.") tr_exposure = gr.Markdown() tr_multi_res = gr.Checkbox( value=True, label="Multi-resolution buckets", @@ -2447,7 +2463,7 @@ def _cap_cost(key: str, model: str, selected: list[str]) -> str: 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, tr_path] + tr_hparams + [tr_optimizer]) 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]) @@ -2458,7 +2474,7 @@ def _cap_cost(key: str, model: str, selected: list[str]) -> str: 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, tr_repeats], + project_name, tr_repeats, tr_optimizer], [tr_result]) demo.load(_check_for_update, None, update_notice) diff --git a/docs/ARCHITECTURE.md b/docs/ARCHITECTURE.md index f91660a..a600c69 100644 --- a/docs/ARCHITECTURE.md +++ b/docs/ARCHITECTURE.md @@ -1,6 +1,6 @@ # Architecture -Version: 0.18.0 +Version: 0.18.1 ``` app.py Gradio UI — thin wiring over the stage functions (5 tabs). @@ -921,6 +921,15 @@ quietly gets someone else's value. 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. +- **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 + its README's diffusion advice (`weight_decay=0.01, use_bias_correction, safeguard_warmup`). + ai-toolkit ships `prodigyopt`; sd-scripts (`Prodigy`) and musubi (`prodigyopt.Prodigy`, its + module-path form) need `pip install prodigyopt`, which ⑤ says. Fizgig is AdamW8bit only: it + removed Prodigy because a self-tuning LR fights its adaptive LR. Adafactor, Lion and Automagic + are left out: for a LoRA the optimizer state is small, so their memory savings buy little, and + each needs its own LR recipe. - **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 `<>` @@ -1201,11 +1210,10 @@ 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. +- **Attention-only training (⑤, planned for 0.18.0, deferred).** A musubi/Fizgig toggle that + puts the LoRA on the attention projections only (via `--network_args "exclude_patterns=[…]"`), + leaving the feed-forward layers, which hold more of the textures, clothes and backgrounds, alone. + Evidence for a likeness gain is anecdotal and per-model; add when an A/B shows it. - **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 diff --git a/studio/__init__.py b/studio/__init__.py index 52fb9dd..22612cd 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.18.0" +__version__ = "0.18.1" diff --git a/studio/trainer_configs.py b/studio/trainer_configs.py index 8ce0049..9ec1cc2 100644 --- a/studio/trainer_configs.py +++ b/studio/trainer_configs.py @@ -34,6 +34,18 @@ # One seed for every validation sample, so epochs differ only by the LoRA. VALIDATION_SEED = 42 +# AdamW8bit is every trainer's own default and validated recipe. Prodigy finds +# its own step size (lr 1.0); these args are its README's advice for diffusion. +OPTIMIZERS = {"adamw8bit": "AdamW8bit (recommended)", + "prodigy": "Prodigy (sets its own learning rate)"} +PRODIGY_ARGS = {"weight_decay": 0.01, "use_bias_correction": True, "safeguard_warmup": True} + + +def optimizer_choices(trainer: str) -> list[tuple[str, str]]: + """Fizgig removed Prodigy: it fights Fizgig's adaptive LR (its docs/CLI.md).""" + keys = ["adamw8bit"] if trainer == "fizgig" else list(OPTIMIZERS) + return [(OPTIMIZERS[k], k) for k in keys] + def _yaml_str(value: str) -> str: """Render `value` as a safe single-line YAML scalar for interpolation. @@ -245,6 +257,7 @@ class TrainConfig(BaseModel): # Images in the dataset; 0 = count them when writing (`write_configs`). n_images: int = 0 lr: float = 1e-4 + optimizer: str = "adamw8bit" # a key of OPTIMIZERS; anything else renders as AdamW8bit batch_size: int = 1 # Multi-resolution buckets. Empty = single-bucket at `resolution` (the old # behaviour); populated from the dataset's real dimensions by the caller. @@ -345,7 +358,14 @@ def render_aitoolkit_yaml(cfg: TrainConfig) -> str: 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)) + if cfg.optimizer == "prodigy": + optimizer = (" optimizer: prodigy\n optimizer_params:\n" + + "".join(f" {k}: {str(v).lower()}\n" + for k, v in PRODIGY_ARGS.items()) + + " lr: 1.0\n") + else: + optimizer = f" optimizer: adamw8bit\n lr: {cfg.lr}\n" + 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" @@ -384,8 +404,7 @@ def render_aitoolkit_yaml(cfg: TrainConfig) -> str: " gradient_checkpointing: true\n" f" noise_scheduler: {m.noise_scheduler}\n" f"{timestep}" - f" optimizer: adamw8bit\n" - f" lr: {cfg.lr}\n" + f"{optimizer}" " dtype: bf16\n" " model:\n" f"{raw_note}" @@ -437,6 +456,14 @@ def aitoolkit_command(install_path: str, config_path: Path) -> str: return f'cd "{base}" && python run.py "{config_path.as_posix()}"' +def _optimizer_cli(cfg: TrainConfig, prodigy_type: str, adamw_type: str) -> str: + """Optimizer + learning-rate flags for musubi and sd-scripts.""" + if cfg.optimizer == "prodigy": + args = " ".join(f'"{k}={v}"' for k, v in PRODIGY_ARGS.items()) + return f"--optimizer_type {prodigy_type} --learning_rate 1.0 --optimizer_args {args}" + return f"--optimizer_type {adamw_type} --learning_rate {cfg.lr}" + + 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()}" ' @@ -474,7 +501,8 @@ def musubi_command(install_path: str, toml_path: Path, cfg: TrainConfig, f' --dataset_config "{toml}"{version} \\\n' " --sdpa --mixed_precision bf16 --gradient_checkpointing \\\n" f"{extra}" - f" --optimizer_type adamw8bit --learning_rate {cfg.lr} \\\n" + # musubi takes any optimizer as module.Class; Prodigy needs prodigyopt installed. + f" {_optimizer_cli(cfg, 'prodigyopt.Prodigy', 'adamw8bit')} \\\n" " --max_data_loader_n_workers 2 --persistent_data_loader_workers \\\n" f" --network_module {m.network_module} \\\n" f" --network_dim {cfg.rank} \\\n" @@ -532,8 +560,8 @@ def kohya_command(install_path: str, toml_path: Path, cfg: TrainConfig, 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_epochs {cfg.epochs} \\\n" - " --optimizer_type AdamW8bit --mixed_precision bf16 --sdpa \\\n" + f" {_optimizer_cli(cfg, 'Prodigy', 'AdamW8bit')} \\\n" + f" --max_train_epochs {cfg.epochs} --mixed_precision bf16 --sdpa \\\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." diff --git a/tests/test_provenance_heldout.py b/tests/test_provenance_heldout.py index 4785436..bcb877e 100644 --- a/tests/test_provenance_heldout.py +++ b/tests/test_provenance_heldout.py @@ -128,6 +128,25 @@ def test_do_export_reports_the_holdout(tmp_path: Path) -> None: assert "Held out 1" in result +def test_optimizer_reaches_the_config_and_fizgig_ignores_prodigy(tmp_path: Path) -> None: + folder = tmp_path / "d" + folder.mkdir() + for i in range(4): + Image.new("RGB", (64, 64)).save(folder / f"{i:02d}.png") + out = A.do_generate_train_config("kohya", "sdxl", str(folder), "", "lora", "trg", + 1024, 16, 16, 10, 1e-4, 1, False, optimizer="prodigy") + assert "--optimizer_type Prodigy" in out and "pip install prodigyopt" in out + out = A.do_generate_train_config("fizgig", "krea2", str(folder), "", "lora", "trg", + 1024, 8, 8, 30, 1e-4, 1, False, optimizer="prodigy") + assert "Prodigy" not in out and "--optimizer_type adamw8bit" in out + # The ⑤ dropdown is the last generate input and follows the trainer. + deps = [d for d in A.demo.fns.values() + if getattr(getattr(d, "fn", None), "__name__", "") == "do_generate_train_config"] + assert deps[0].inputs[-1].label == "Optimizer" + update = A.on_trainer_change("fizgig")[-1] + assert update["choices"] == [("AdamW8bit (recommended)", "adamw8bit")] + + def test_exposure_line_auto_repeats(tmp_path: Path) -> None: folder = tmp_path / "d" folder.mkdir() diff --git a/tests/test_trainer_configs.py b/tests/test_trainer_configs.py index a37e350..b5e09d9 100644 --- a/tests/test_trainer_configs.py +++ b/tests/test_trainer_configs.py @@ -224,3 +224,37 @@ def test_validation_guide_lists_held_out_photos(tmp_path: Path) -> None: 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 + + +# ---------- 0.18.1: optimizer choice ---------- + +def test_adamw8bit_is_the_default_everywhere(tmp_path: Path) -> None: + train = yaml.safe_load(render_aitoolkit_yaml(_cfg("ai-toolkit", tmp_path)))["config"][ + "process"][0]["train"] + assert (train["optimizer"], train["lr"]) == ("adamw8bit", 1e-4) + for trainer, flag in (("musubi", "adamw8bit"), ("kohya", "AdamW8bit"), ("fizgig", "adamw8bit")): + cfg = _cfg(trainer, tmp_path / trainer).model_copy( + update={"model": TRAINER_MODELS[trainer][0]}) + _, command = write_configs(cfg) + assert f"--optimizer_type {flag}" in command, trainer + + +def test_prodigy_runs_at_lr_1_with_its_diffusion_args(tmp_path: Path) -> None: + cfg = _cfg("ai-toolkit", tmp_path).model_copy(update={"optimizer": "prodigy"}) + train = yaml.safe_load(render_aitoolkit_yaml(cfg))["config"]["process"][0]["train"] + assert (train["optimizer"], train["lr"]) == ("prodigy", 1.0) + assert train["optimizer_params"] == {"weight_decay": 0.01, "use_bias_correction": True, + "safeguard_warmup": True} + for trainer, flag in (("musubi", "prodigyopt.Prodigy"), ("kohya", "Prodigy")): + cfg = _cfg(trainer, tmp_path / trainer).model_copy(update={"optimizer": "prodigy"}) + _, command = write_configs(cfg) + assert f"--optimizer_type {flag} --learning_rate 1.0" in command, trainer + # musubi and sd-scripts literal_eval each "key=value". + assert '"use_bias_correction=True" "safeguard_warmup=True"' in command + + +def test_fizgig_never_offers_prodigy() -> None: + from studio.trainer_configs import optimizer_choices + + assert [k for _, k in optimizer_choices("fizgig")] == ["adamw8bit"] + assert [k for _, k in optimizer_choices("kohya")] == ["adamw8bit", "prodigy"]