diff --git a/.env.example b/.env.example index 8c8bf01..581ba94 100644 --- a/.env.example +++ b/.env.example @@ -36,7 +36,7 @@ HF_TOKEN= # reserved by Windows (Hyper-V/WSL/Docker). Set it to pin one exactly. # LDS_PORT=7861 # LDS_OUTPUT_ROOT=./datasets -# LDS_GEMINI_IMAGE_MODEL=gemini-3-pro-image-preview +# LDS_GEMINI_IMAGE_MODEL=auto # auto = best image model your key can use; or pin an id from ② refresh # LDS_UPDATE_CHECK_ENABLED=true # checks GitHub releases on launch (no data sent); set false to disable # ComfyUI model filenames, if yours are named differently (see docs/comfyui-setup.md) diff --git a/AGENTS.md b/AGENTS.md index ee50917..846cd80 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -6,12 +6,12 @@ Turn a character, style, or concept into a ready-to-train LoRA dataset. One refe ## Current state -_Last verified: 2026-10-03_ +_Last verified: 2026-10-08_ -- **Status:** in active development, released at v0.17.2 (git tag `v0.17.2`). **Every version bump must ship a GitHub Release** or the in-app update check never fires. The project was renamed twice (lora-dataset-studio → lora-distillery → dataset-deviser); older references under the previous names are stale. -- **Works:** all five stages end to end (preprocess → generate & curate → caption → export → train config); the three dataset types (character, style, concept) with their own shot plans and caption framing; local and cloud paths for every stage; gallery curation with selection carried forward between stages, including shift-click range select; advisory dedupe, quality flags and caption lint; trainer configs for ai-toolkit (incl. SDXL), kohya sd-scripts and musubi-tuner; opt-in private Hugging Face publish. Since 0.15.0: per-image fault isolation in ① (one bad source never ends the batch), a cooperative ⏹ Stop across ①/②/③ with documented resume paths, translated ComfyUI failures (missing node, bad input, unreachable server) and an in-app 🩺 setup check. Since 0.15.1: a ② **shot style** (default = keep the reference's own medium; eight presets + custom text) threaded into both prompt builders, the CLI, ④ metadata and ⑤'s sample prompt, a final-prompt preview, and a prev/next/save-and-next caption editor with the image on screen. Since 0.16.0: **one name box and one trigger box in the header** that ②/③/④/⑤ all read (they used to be per-tab and went stale), ④ stating the pair it stamped, an opt-in ④ **hand-off to [Idiot LoRa Builder](https://github.com/Fablestarexpanse/Idiot-Lora-Builder)** (writes its `.lora-studio/ratings.json` so its grid opens pre-triaged — nothing is launched), a ⑤ advisory when the dataset can't fill the training resolution, every bundled ComfyUI model filename overridable from `.env`, and a `cli build` preflight for the local engine. Since 0.17.0: the local ② engine is **Qwen-Image 2.1** (plain-English prompts shared with Gemini, ~26 s/shot, `LDS_QWEN21_*` filenames) instead of Qwen-Image-Edit 2511 + the Multiple-Angles LoRA, whose graph stays in `studio/comfy_workflows/legacy/`; a dead ComfyUI mid-batch now fails one shot, not the run. Since 0.17.1: the UI port is chosen by `_pick_port` in `app.py` — 7861 if bindable, else an OS-assigned free port (Windows reserves a ~500-port block around 7861), or `LDS_PORT` to pin one. Since 0.17.2: `comfy_api.upload_image` downscales a source over 90 MB before upload (ComfyUI answers 413 above its 100 MB `--max-upload-size`). -- **In progress:** nothing half-built — 0.16.0 closed out the hand-off, cross-tab-identity and config-hardening work. `docs/ARCHITECTURE.md` carries the "Future ideas" and "Deferred" sections that hold the real backlog. -- **Known gaps / next steps:** 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. +- **Status:** in active development; main is v0.17.3 (git tag `v0.17.3`). **Every version bump must ship a GitHub Release** or the in-app update check never fires. The project was renamed twice (lora-dataset-studio → lora-distillery → dataset-deviser); older references under the previous names are stale. +- **Works:** all five stages end to end (preprocess → generate & curate → caption → export → train config); the three dataset types (character, style, concept) with their own shot plans and caption framing; local and cloud paths for every stage; gallery curation with selection carried forward between stages, including shift-click range select; advisory dedupe, quality flags and caption lint; trainer configs for ai-toolkit (incl. SDXL), kohya sd-scripts and musubi-tuner; opt-in private Hugging Face publish. Since 0.15.0: per-image fault isolation in ① (one bad source never ends the batch), a cooperative ⏹ Stop across ①/②/③ with documented resume paths, translated ComfyUI failures (missing node, bad input, unreachable server) and an in-app 🩺 setup check. Since 0.15.1: a ② **shot style** (default = keep the reference's own medium; eight presets + custom text) threaded into both prompt builders, the CLI, ④ metadata and ⑤'s sample prompt, a final-prompt preview, and a prev/next/save-and-next caption editor with the image on screen. Since 0.16.0: **one name box and one trigger box in the header** that ②/③/④/⑤ all read (they used to be per-tab and went stale), ④ stating the pair it stamped, an opt-in ④ **hand-off to [Idiot LoRa Builder](https://github.com/Fablestarexpanse/Idiot-Lora-Builder)** (writes its `.lora-studio/ratings.json` so its grid opens pre-triaged — nothing is launched), a ⑤ advisory when the dataset can't fill the training resolution, every bundled ComfyUI model filename overridable from `.env`, and a `cli build` preflight for the local engine. Since 0.17.0: the local ② engine is **Qwen-Image 2.1** (plain-English prompts shared with Gemini, ~26 s/shot, `LDS_QWEN21_*` filenames) instead of Qwen-Image-Edit 2511 + the Multiple-Angles LoRA, whose graph stays in `studio/comfy_workflows/legacy/`; a dead ComfyUI mid-batch now fails one shot, not the run. Since 0.17.1: the UI port is chosen by `_pick_port` in `app.py` — 7861 if bindable, else an OS-assigned free port (Windows reserves a ~500-port block around 7861), or `LDS_PORT` to pin one. Since 0.17.2: `comfy_api.upload_image` downscales a source over 90 MB before upload (ComfyUI answers 413 above its 100 MB `--max-upload-size`). Since 0.17.3: the cloud image model defaults to **`auto`** (Google shut down the `-preview` ids on 2026-06-25; `resolve_image_model` picks from the cached live list, Nano Banana 2.1 first), Groq captions use `qwen/qwen3.8-27b` (key `groq-qwen`, old key aliased), every API image goes through `config.api_image` (true mime, size cap), musubi commands include the cache steps and the per-architecture network module, ai-toolkit gains Qwen-Image 2.1 and a correct Krea 2 (Raw) preset and keeps every checkpoint, ① applies EXIF rotation, no longer 4x-upscales large JPEGs, and can jump a busy ComfyUI queue like ② (`--front` on `preprocess`/`build`). +- **In progress:** a phased LoRA best-practice update. 0.18.0 is next: epochs as the ⑤ unit with an exposure line, a validation pack, held-out references, a Fizgig trainer target, source-size provenance in a PNG chunk. Then 0.19.0 (training copies keep their background, identity policy, framing tiers, coverage panel, local multi-reference) and 0.20.0 (a Quick build tab). Nothing from those is half-built on main. +- **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. ## Architecture in 60 seconds diff --git a/README.md b/README.md index c5b9940..3caeeea 100644 --- a/README.md +++ b/README.md @@ -201,8 +201,8 @@ drop-list. It's advisory: nothing is blocked or rewritten. | Qwen3-VL-8B NSFW-Caption | your GPU, ~17 GB | explicit-dataset specialist | | **WD EVA02 / ViT v3** *(tagger)* | GPU or CPU, ONNX | canonical **Danbooru** tags | | **Z3D e621 ConvNeXt** *(tagger)* | GPU or CPU, ONNX | canonical **e621/furry** tags | -| Gemini Flash | Google API | SFW, ~$0.0007/img est., billed to your key | -| Groq Qwen3.6 27B | Groq API | SFW, free tier, rate-limited | +| Gemini Flash | Google API | SFW, ~$0.002/img est., billed to your key | +| Groq Qwen3.8 27B | Groq API | SFW, free or paid plan, rate-limited | | LM Studio / Ollama / custom | your choice | any OpenAI-compatible endpoint | ## Shot style diff --git a/app.py b/app.py index 9752faa..2da456b 100644 --- a/app.py +++ b/app.py @@ -40,13 +40,19 @@ from studio.config import ( CAPTIONERS, CAPTIONERS_BY_KEY, - CLOUD_IMAGE_PRICES, friendly_api_error, list_images, load_caption_model_cache, read_caption, settings, ) +from studio.engines.gemini import ( + AUTO_MODEL, + image_model_choices, + image_price, + known_image_models, + resolve_image_model, +) 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 @@ -57,7 +63,17 @@ ("Cloud — Gemini image model (best identity fidelity, SFW only)", "gemini"), ("Local — ComfyUI Qwen Image 2.1 (free, private, uncensored)", "comfyui"), ] -CLOUD_MODEL_CHOICES = [(f"{m} (~${p:.3f}/img est.)", m) for m, p in CLOUD_IMAGE_PRICES.items()] +CLOUD_MODEL_CHOICES = image_model_choices(known_image_models()) + + +def _cloud_model_default(pinned: str, choices: list[tuple[str, str]]) -> str: + """The ② dropdown's start value: the .env pin (retired ids mapped) if listed, else Auto.""" + if pinned != AUTO_MODEL: + pinned = resolve_image_model(pinned) + return pinned if pinned in {m for _, m in choices} else AUTO_MODEL + + +CLOUD_MODEL_DEFAULT = _cloud_model_default(settings.gemini_image_model, CLOUD_MODEL_CHOICES) CAPTIONER_CHOICES = [(c.label, c.key) for c in CAPTIONERS] # Gemini caption-model dropdown seed: use the local cache if present, else a @@ -581,7 +597,8 @@ def _preprocess_note(reports, out_dir: Path, alpha_cutout: bool) -> str: def do_preprocess(files: list[str], folder: str, target: int, restore_mode: str, restore_backend: str, isolate: bool, isolation_backend: str, subject_prompt: str, exclude_prompt: str, tighten: bool = False, - alpha_cutout: bool = False, progress=gr.Progress()): + alpha_cutout: bool = False, front: bool = False, + progress=gr.Progress()): sources = _inputs(files, folder) out_dir = _stamped("prepped") force = {"Auto (only if needed)": None, "Always": True, "Never": False}[restore_mode] @@ -598,7 +615,7 @@ def report(msg: str): subject_prompt=subject_prompt or "character", exclude_prompt=exclude_prompt or "", restore_backend=restore_backend, isolation_backend=isolation_backend, tighten_crop=tighten, - alpha_cutout=alpha_cutout, should_stop=JOB, progress=report) + alpha_cutout=alpha_cutout, front=front, should_stop=JOB, progress=report) except OSError as e: raise gr.Error(f"Couldn't write to '{out_dir}': {e}. Check the output folder " f"(valid drive, writable, enough space).") from e @@ -1358,17 +1375,10 @@ def do_load_plan(plan_name: str): def estimate_cost(engine: str, cloud_model: str, df: pd.DataFrame) -> str: n = len(df) if engine == "gemini": - from studio.config import CLOUD_IMAGE_PRICES, load_cloud_model_cache - - price = CLOUD_IMAGE_PRICES.get(cloud_model) - cached = load_cloud_model_cache() or [] - for m in cached: - if m.get("model_id") == cloud_model and m.get("price") is not None: - price = m["price"] - break + model, price = image_price(cloud_model) if price is None: - return f"**Cost:** {n} images on `{cloud_model}` (price unknown — billed to your API key)" - return (f"**Cost:** ~${n * price:.2f} for {n} images on `{cloud_model}` " + return f"**Cost:** {n} images on `{model}` (price unknown — billed to your API key)" + return (f"**Cost:** ~${n * price:.2f} for {n} images on `{model}` " f"(estimate at build time — billed to your own Google API key)") return f"**Cost:** {n} images, $0 (local generation)" @@ -1534,15 +1544,16 @@ def do_generate_train_config(trainer: str, model_key: str, dataset_dir: str, f"trainer's own docs before a long run.") -def refresh_cloud_models(force: bool = False): +def refresh_cloud_models(current: str = "", force: bool = False): + """New choices from the API; keeps the user's pick when it is still offered.""" from studio.engines.gemini import list_image_models try: - models = list_image_models(force_refresh=force) + choices = image_model_choices(list_image_models(force_refresh=force)) except Exception as e: raise gr.Error(f"Could not list models: {e}") from e - return gr.Dropdown(choices=models, - value=models[0][1] if models else settings.gemini_image_model) + keep = current if current in {m for _, m in choices} else AUTO_MODEL + return gr.Dropdown(choices=choices, value=keep) def refresh_caption_models(): @@ -1749,6 +1760,10 @@ def _check_for_update(): "own compositing workflows. Builtin SAM3 backend only. Leave off " "(default) if you're continuing to ② Generate — it expects a white " "background reference. Needs isolation on.") + pre_front = gr.Checkbox( + value=False, label="Prioritize this app's ComfyUI jobs", + info="Puts ComfyUI restore/isolation jobs at the head of its " + "pending queue. Does not interrupt a job already running.") btn_pre = gr.Button("① Preprocess", variant="primary") with gr.Column(scale=2): pre_note = gr.Markdown() @@ -1793,8 +1808,10 @@ def _check_for_update(): label="Generation engine", info="Cloud Gemini needs no GPU (best identity, SFW); " "local ComfyUI is free, private, uncensored.") + # A pinned id the list lacks (a retired preview in .env) would + # make Gradio reject every event that reads this dropdown. cloud_model = gr.Dropdown(CLOUD_MODEL_CHOICES, - value=settings.gemini_image_model, + value=CLOUD_MODEL_DEFAULT, label="Cloud image model", info="Only used by the Cloud engine. Prices are " "build-time estimates.") @@ -2200,7 +2217,7 @@ def _check_for_update(): do_preprocess, [pre_files, pre_folder, target, restore_mode, restore_backend, isolate, isolation_backend, subject_prompt, exclude_prompt, pre_tighten, - pre_alpha_cutout], + pre_alpha_cutout, pre_front], [prep_gallery, pre_note, log_box, gen_src_folder, cap_folder]) \ .then(lambda s, e: (s, e), [subject_prompt, exclude_prompt], [gen_subject, gen_exclude]) @@ -2249,12 +2266,12 @@ def _check_for_update(): btn_outfits_clear.click(clear_outfits, [plan], [plan, plan_note]) btn_save_plan.click(do_save_plan, [plan, plan_name], [plan_note]) btn_load_plan.click(do_load_plan, [plan_name], [plan, plan_note]) - refresh_models.click(refresh_cloud_models, [], [cloud_model]) + refresh_models.click(refresh_cloud_models, [cloud_model], [cloud_model]) - def _force_refresh(): - return refresh_cloud_models(force=True) + def _force_refresh(current: str): + return refresh_cloud_models(current, force=True) - force_refresh_models.click(_force_refresh, [], [cloud_model]) + force_refresh_models.click(_force_refresh, [cloud_model], [cloud_model]) engine.change(estimate_cost, [engine, cloud_model, plan], [cost]) cloud_model.change(estimate_cost, [engine, cloud_model, plan], [cost]) plan.change(estimate_cost, [engine, cloud_model, plan], [cost]) diff --git a/cli.py b/cli.py index f23566f..1817f6b 100644 --- a/cli.py +++ b/cli.py @@ -22,7 +22,7 @@ import typer from studio import env_keys, pipeline -from studio.config import CAPTIONERS_BY_KEY, list_images, settings +from studio.config import CAPTIONERS, list_images, settings from studio.shotplan import plan_for_type app = typer.Typer(add_completion=False, help=__doc__) @@ -63,17 +63,13 @@ def _echo_cloud_estimate(engine: str, cloud_model: str, n_shots: int) -> None: """Warn what a cloud run will cost before it starts billing the user.""" if engine != "gemini": return - from studio.config import CLOUD_IMAGE_PRICES, load_cloud_model_cache - - model_id = cloud_model or settings.gemini_image_model - price = CLOUD_IMAGE_PRICES.get(model_id) - for m in load_cloud_model_cache() or []: - if m.get("model_id") == model_id and m.get("price") is not None: - price = m["price"] - break + from studio.engines.gemini import image_price + + model_id, price = image_price(cloud_model) if price: typer.echo(f"Cloud engine: ~${n_shots * price:.2f} estimated for {n_shots} " - f"images (build-time estimate, billed to your Google API key).") + f"images on {model_id} (build-time estimate, billed to your Google " + f"API key).") def _preflight_comfyui(engine: str) -> None: @@ -203,6 +199,7 @@ def preprocess( False, "--alpha-cutout/--no-alpha-cutout", help="Export on a transparent background instead of white (builtin backend " "only; for your own compositing workflows — not meant to feed 'generate')"), + front: bool = typer.Option(False, help="Jump ComfyUI's pending queue"), ): """Restore/upscale/isolate images (standalone).""" out = out or pipeline.new_run_dir("prepped") @@ -210,7 +207,7 @@ def preprocess( _expand(inputs), out, target=target, force_restore=restore, isolate=isolate, subject_prompt=subject_prompt, exclude_prompt=exclude_prompt, restore_backend=restore_backend, isolation_backend=isolation_backend, - tighten_crop=tighten, alpha_cutout=alpha_cutout, progress=typer.echo) + tighten_crop=tighten, alpha_cutout=alpha_cutout, front=front, progress=typer.echo) _echo_preprocess_failures(reports) # Exit non-zero only when nothing was written: a partial run still produced # a usable folder, and failing the whole command would hide that. @@ -279,7 +276,7 @@ def caption( folder: Path = typer.Argument(..., exists=True, file_okay=False, help="Folder of images to tag"), captioner: str = typer.Option(settings.default_captioner, - help=f"one of {list(CAPTIONERS_BY_KEY)}"), + help=f"one of {[c.key for c in CAPTIONERS]}"), name: str = typer.Option("", help="Character name used in captions"), trigger: str = typer.Option("", help="Trigger word placed first in every caption"), model: str = typer.Option("", help="Model id (gemini captioner; blank = default)"), @@ -421,7 +418,7 @@ def build( help="gemini (cloud, billed to your key) or comfyui (local, free — needs " "the models in docs/comfyui-setup.md; `doctor` says if it is ready)"), captioner: str = typer.Option(settings.default_captioner, - help=f"one of {list(CAPTIONERS_BY_KEY)}"), + help=f"one of {[c.key for c in CAPTIONERS]}"), caption_style: str = typer.Option( "prose", "--caption-style", help="prose (natural language), tags (Danbooru: SDXL/Illustrious) or e621 (furry/Pony)"), @@ -467,6 +464,7 @@ def build( drop_tags: str = typer.Option( "", "--drop-tags", help="Comma-separated tags to strip from tag captions (e.g. 'watermark, signature')"), + front: bool = typer.Option(False, help="Jump ComfyUI's pending queue"), ): """Full pipeline: preprocess -> generate -> caption -> export.""" from studio.captioner import caption_images @@ -487,7 +485,7 @@ def build( 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, progress=typer.echo) + tighten_crop=tighten, front=front, progress=typer.echo) _echo_preprocess_failures(reports) # A per-image failure is reported, not raised, so `output` can be None — # everything downstream must run on the images that actually exist. @@ -515,7 +513,8 @@ def build( prepped, shots, engine, run_dir / "generated", cloud_model=cloud_model, isolate_angles=do_isolate, subject_prompt=subject_prompt, exclude_prompt=exclude_prompt, - exclude_props=_props_default(exclude_props, dtype), progress=typer.echo) + exclude_props=_props_default(exclude_props, dtype), front=front, + progress=typer.echo) kept = [r.path for r in results if r.path] if not kept: typer.echo("No shots succeeded; aborting before captioning.") @@ -709,4 +708,10 @@ def custom_endpoint( if __name__ == "__main__": + import sys + + # Redirected to a file on Windows, stdout is cp1252 and a "→" in a progress + # line raised UnicodeEncodeError, ending the whole batch mid-run. + for stream in (sys.stdout, sys.stderr): + stream.reconfigure(errors="replace") app() diff --git a/docs/ARCHITECTURE.md b/docs/ARCHITECTURE.md index 7aa022c..f0fba63 100644 --- a/docs/ARCHITECTURE.md +++ b/docs/ARCHITECTURE.md @@ -1,6 +1,6 @@ # Architecture -Version: 0.17.2 +Version: 0.17.3 ``` app.py Gradio UI — thin wiring over the stage functions (5 tabs). @@ -571,6 +571,16 @@ quietly gets someone else's value. white, so this is backend-agnostic and needs no mask handle. It runs only when isolation is on and the toggle is set (default off → the Character path is byte-identical), before the resize. An effectively all-white image (nothing found) is returned unchanged. +- **Restoration is judged at the size the dataset holds (0.17.3).** Every JPEG used to be flagged + "lossy" and sent through DeJPG + the 4x upscale at full size: a 3000x4000 phone photo became a + 12000x16000 PNG, fetched only to be shrunk to 1024. Now "lossy" applies only below 2x the + target, sharpness is measured after resizing to the target, the upload is pre-shrunk to 2x the + target, and `_restore_comfyui(upscale=False)` rewires the template's SaveImage to the DeJPG + output when the source already reaches the target. +- **EXIF orientation is applied at ①.** Phone portraits are stored landscape plus a rotate tag; + `Image.open().convert()` ignores it, so they entered the dataset sideways. `preprocess()` + transposes first and hands restore/isolation an upright temp copy (`_stage_copy`) whenever + the file on disk is rotated or larger than 2x the target. - **Preprocess output names must not clobber.** `list_images` admits several extensions, so two sources sharing a stem (`cat.jpg` + `cat.png`, in one folder or across merged input folders) both map to `cat_prepped.png`. `preprocess()` makes the output non-clobbering @@ -611,7 +621,10 @@ quietly gets someone else's value. ComfyUI to put the job at the head of the pending queue (`/prompt` accepts `front` and negates the job's priority number). `front` does **not** interrupt the job already running, so the UI must not promise "immediate". Polling is always by our own - `prompt_id`, so another client's jobs are never mistaken for ours. + `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 ②. - **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 @@ -652,6 +665,30 @@ quietly gets someone else's value. to `.cache/gemini_caption_models.json`, both with a 24-hour TTL. The UI seeds each dropdown from cache (never a network call on startup); a refresh button force-pulls. Both fall back to a static list if the API is unreachable or no key is set. +- **Gemini image models have no `-latest` alias, so the default is `auto` (0.17.3).** The old + default `gemini-3-pro-image-preview` and `gemini-3.1-flash-image-preview` were shut down on + 2026-06-25, so every fresh install 404'd on its first ② shot. `resolve_image_model()` walks + `IMAGE_MODEL_PREFERENCE` against the cached live list (never a network call), maps a pinned + retired preview id to its GA id, and the cost line names what Auto resolved to. A model 404 + now says to refresh and pick another. The live + list filter must match `nano-banana` ids too — `gemini-nano-banana-2.1` has no "image" in it. + `models.list` still returns the shut-down `-preview` ids, so their presence proves nothing. +- **Auto prefers Nano Banana 2.1 over Pro (0.17.3 A/B).** One real person, two poor phone + photos (a cropped, occluded 12 MP shot and a 580 px heavy-JPEG one), both isolated onto white + by ①, the 24-shot character plan on each model with identical prompts: + - **2.1** (~$0.82, ~6 min): 24/24, every shot a consistent photograph in a real setting; one + near-duplicate pair of profiles and one very dark shot. + - **Pro** (~$3.22, ~7 min): 24/24, comparable likeness, more outfit variety drawn from both + references, but six close-ups/poses came back on the reference's white background with + cut-out edges and one drifted to a painted look. + Neither showed the removed props or the face shadow. One subject and one run each, so treat + the order as a default to revisit, not a law; the white-background leak should shrink once + ② references keep their background (0.19). +- **Every image sent to an API goes through `config.api_image()`.** Cloud generation, Gemini and + OpenAI-compatible captioning used to send raw file bytes labelled `image/png` whatever they + were, with no size cap. `api_image` passes PNG/JPEG/WebP within the limit through untouched, + otherwise downscales and encodes PNG (alpha) or JPEG q95, and returns the true mime. Limits: + 2048 for generation references, 1536 for cloud captions; local VLMs get 1024 (the training size). - **Update check is best-effort by design.** `update_check.py` hits the GitHub releases API on `demo.load` (not at process start, so it can't slow down launch), caches the result 24h in `.cache/update_check.json`, and falls back to a stale cache on a failed @@ -681,11 +718,13 @@ quietly gets someone else's value. `gemini-flash-latest` so it doesn't 404 when a pinned version is decommissioned. The Caption tab can refresh and pick a specific model; `Captioner(model_override=...)` applies it only when the backend is Gemini. -- **Groq Qwen3.6 is a reasoning model.** Its spec sets `extra_params={"reasoning_effort": +- **Groq Qwen3.8 is a reasoning model.** Its spec sets `extra_params={"reasoning_effort": "none"}` to switch off the `` scratchpad (otherwise the token budget is spent - on reasoning and the caption comes back empty/truncated), plus a higher `max_tokens` - as insurance and `_clean()` strips any residual `` tags. 429s are retried with - backoff; `min_interval_s` spaces requests for the free tier. + on reasoning and the caption comes back empty/truncated), and `_clean()` strips any + residual `` tags. 429s are retried with backoff. Groq meters tokens per minute + (8K on the base tier) and a captioned image costs ~2.6K, so `min_interval_s` is 20 s and + `max_tokens` stays ~400. The key is `groq-qwen`; the pre-0.17.3 `groq-qwen3.6` (model shut + down 2026-09-14) stays as an alias so a saved `.env` keeps working. - **Custom OpenAI-compatible captioner.** The `custom` captioner carries no endpoint in config; either the Caption tab or `cli.py custom-endpoint` (0.12.2 — previously UI-only, a CLI-only user had no way to set this at all) collects base URL / model / key-env-NAME / @@ -853,8 +892,19 @@ quietly gets someone else's value. run.** The ⑤ tab says so. Don't imply otherwise. - **ai-toolkit configs are one-command; musubi configs are not.** ai-toolkit emits a fully runnable `config.yaml` (model is a HF id → `python run.py config.yaml`). musubi's `dataset.toml` - is complete, but its *training* invocation needs the user's local DiT/VAE/text-encoder paths, so - the musubi run command is a `<>` template. The UI says so; it is not faked as one-click. + is complete, but its invocation needs the user's local DiT/VAE/text-encoder paths, so the + command carries `<>` paths. The UI says so; it is not faked as one-click. +- **musubi needs three commands and a per-architecture network module (0.17.3).** Before 0.17.3 + the command skipped the latent and text-encoder cache steps (training fails without them) and + passed `networks.lora` for every model; each architecture has its own (`lora_qwen_image`, + `lora_flux_2`, `lora_zimage`, `lora_krea2`, …). `ModelPreset` now carries the module, the + version flag, the text-encoder flags and the timestep args, copied from each model's + `docs/.md` example in musubi-tuner; `tests/test_trainer_configs.py` pins them. +- **ai-toolkit keeps every checkpoint.** `max_step_saves_to_keep` was 4, so a long run deleted the + early checkpoints, which are often the best likeness. It is now `ceil(steps / save_every)`. +- **Train LoRAs on the base checkpoint, never a Turbo/distilled one.** Krea 2 trains on + `krea/Krea-2-Raw` (arch `krea2`; `krea` was never an ai-toolkit arch). A test fails if any + preset path contains "turbo". - **Trainer model registry is curated, not exhaustive.** Where a model's canonical HF id or musubi script isn't something we can guarantee, the preset carries a `<>` placeholder so the emitted config is honest rather than silently wrong. @@ -963,7 +1013,7 @@ grouped by stage. Milestone versions are noted only where they explain a design - **① Preprocess** — restore (ComfyUI models / basic Lanczos / auto), SAM3 subject isolation (built-in transformers or ComfyUI, with measured over-cut fixes), optional tighten-to-subject - crop, resize to target long-side. Optional alpha (RGBA) cutout output (builtin backend + crop, resize to target long-side, EXIF rotation applied. Optional alpha (RGBA) cutout output (builtin backend only) for external compositing workflows. - **② Generate & curate** — curated 24-shot Character plan (angles/poses/emotions/settings) or 18-shot Concept plan (turnaround/framing/context), Qwen-Image 2.1 (local, 0.17.0 — was Qwen-Image-Edit diff --git a/requirements.txt b/requirements.txt index 6d11014..ab4288f 100644 --- a/requirements.txt +++ b/requirements.txt @@ -10,7 +10,7 @@ gradio>=5.0,<6 transformers>=4.57 accelerate>=1.0 google-genai>=1.0 -pillow>=10.0 +pillow>=10.1 # has_transparency_data httpx>=0.27 typer>=0.12 pydantic>=2.7 diff --git a/studio/__init__.py b/studio/__init__.py index a490d35..8ce5787 100644 --- a/studio/__init__.py +++ b/studio/__init__.py @@ -1,3 +1,3 @@ """Dataset Deviser: turn a character, style, or concept into a ready-to-train LoRA dataset.""" -__version__ = "0.17.2" +__version__ = "0.17.3" diff --git a/studio/captioner.py b/studio/captioner.py index 9d41e60..08f149d 100644 --- a/studio/captioner.py +++ b/studio/captioner.py @@ -35,6 +35,12 @@ _JOYCAPTION_SYSTEM = "You are a helpful image captioner." +# Long-side caps for the image a captioner sees. Cloud: big enough for detail, +# small enough that Groq's 2,048-token image budget and Gemini's tiling don't +# bill for pixels the caption can't use. Local: the training size. +CLOUD_CAPTION_MAX_SIDE = 1536 +LOCAL_VLM_MAX_SIDE = 1024 + # Cloud captioning retries: 4 attempts at 2s/4s/8s covers the seconds-long Gemini # demand spikes that produce a 503 without leaving a user staring at a stuck batch. GEMINI_RETRIES = 4 @@ -221,7 +227,10 @@ def _caption_transformers(self, image_path: Path, instruction: str) -> str: import torch self._load_transformers() + # Capped at the training size: a full-res phone photo otherwise becomes + # thousands of vision tokens and a VRAM spike for no better caption. image = Image.open(image_path).convert("RGB") + image.thumbnail((LOCAL_VLM_MAX_SIDE, LOCAL_VLM_MAX_SIDE), Image.LANCZOS) messages = [] if self.spec.prompt_style == "llava": # JoyCaption expects its system prompt messages.append({"role": "system", "content": _JOYCAPTION_SYSTEM}) @@ -260,8 +269,9 @@ def _caption_gemini(self, image_path: Path, instruction: str) -> str: from google.genai import types client = config.gemini_client(key) + data, mime = config.api_image(image_path, CLOUD_CAPTION_MAX_SIDE) contents = [ - types.Part.from_bytes(data=image_path.read_bytes(), mime_type="image/png"), + types.Part.from_bytes(data=data, mime_type=mime), instruction, ] # Retry transient statuses with backoff, like the OpenAI backend already did. @@ -304,7 +314,8 @@ def _caption_openai(self, image_path: Path, instruction: str) -> str: headers["Authorization"] = f"Bearer {key}" model = self.model or self._first_served_model(headers) - b64 = base64.b64encode(image_path.read_bytes()).decode() + data, mime = config.api_image(image_path, CLOUD_CAPTION_MAX_SIDE) + b64 = base64.b64encode(data).decode() payload = { "model": model, "max_tokens": self.spec.max_tokens, @@ -314,7 +325,7 @@ def _caption_openai(self, image_path: Path, instruction: str) -> str: "role": "user", "content": [ {"type": "image_url", - "image_url": {"url": f"data:image/png;base64,{b64}"}}, + "image_url": {"url": f"data:{mime};base64,{b64}"}}, {"type": "text", "text": instruction}, ], } diff --git a/studio/comfy_api.py b/studio/comfy_api.py index df4d44b..12d18c7 100644 --- a/studio/comfy_api.py +++ b/studio/comfy_api.py @@ -274,8 +274,8 @@ def run_prompt(graph: dict, timeout: float = 600.0, front: bool = False) -> list if backlog > 10 and not front: raise ComfyError( f"ComfyUI queue is busy: {backlog} jobs already pending. This app's jobs " - f"would wait behind them. Either enable 'Prioritize this app's ComfyUI jobs' " - f"to jump the pending queue, or open ComfyUI ({settings.comfy_url}) and clear " + f"would wait behind them. Either tick 'Prioritize this app's ComfyUI jobs' " + f"(① and ② each have one; `--front` in the CLI) to jump the pending queue, or open ComfyUI ({settings.comfy_url}) and clear " f"it (Queue panel → Clear, or the Manager's 'Clear Queue'), then retry. To " f"generate without ComfyUI, switch the engine to the Cloud (Gemini) option, " f"or set the isolation/restore backend to Built-in/Basic." diff --git a/studio/config.py b/studio/config.py index 1a4bf1b..89d6989 100644 --- a/studio/config.py +++ b/studio/config.py @@ -6,12 +6,13 @@ from __future__ import annotations +import io import json from datetime import datetime, timedelta, timezone from pathlib import Path from dotenv import load_dotenv -from pydantic import BaseModel +from pydantic import BaseModel, field_validator from pydantic_settings import BaseSettings, SettingsConfigDict REPO_ROOT = Path(__file__).resolve().parent.parent @@ -309,22 +310,25 @@ def _compose_prompt(self, style: str, dataset_type: str, sparse: bool) -> str: model="gemini-flash-latest", api_key_env="GEMINI_API_KEY", nsfw_capable=False, - cost_note="billed by Google to your key (~$0.001/img est. — check current pricing)", + cost_note="billed by Google to your key (~$0.002/img est. — check current pricing)", ), CaptionerSpec( - key="groq-qwen3.6", - label="Cloud: Groq Qwen3.6 27B (free tier, SFW)", + # Versionless key: Groq retires models every few months (qwen3.6-27b went + # on 2026-09-14) and the key is what .env / scripts name. + key="groq-qwen", + label="Cloud: Groq Qwen3.8 27B (SFW, rate-limited)", backend="openai", base_url="https://api.groq.com/openai/v1", - model="qwen/qwen3.6-27b", + model="qwen/qwen3.8-27b", api_key_env="GROQ_API_KEY", - min_interval_s=3.0, # respect the free-tier 8K TPM limit + # 8K tokens/min on the base tier, and one call is ~2.6K (an image costs up + # to 2,048 tokens). 3 s spacing ran straight into 429s; 20 s stays under. + min_interval_s=20.0, nsfw_capable=False, - cost_note="free tier (rate-limited; 8K TPM)", - # Qwen3.6 is a reasoning model; "none" disables its scratchpad - # so the response is just the caption. Keep a slightly higher token - # budget as insurance if a future model ignores the flag. - max_tokens=800, + cost_note="Groq free or paid plan (~$0.003/img est. paid; 8K tokens/min)", + # A reasoning model; "none" disables its scratchpad so the response + # is just the caption. Kept low because Groq counts it against the minute. + max_tokens=400, extra_params={"reasoning_effort": "none"}, ), CaptionerSpec( @@ -350,29 +354,37 @@ def _compose_prompt(self, style: str, dataset_type: str, sparse: bool) -> str: ] CAPTIONERS_BY_KEY = {c.key: c for c in CAPTIONERS} +# Retired keys still named in someone's .env (LDS_DEFAULT_CAPTIONER) or scripts. +CAPTIONERS_BY_KEY["groq-qwen3.6"] = CAPTIONERS_BY_KEY["groq-qwen"] # Known Gemini image models -> USD per standard-resolution (1K) image. # ESTIMATES captured at build time — actual costs are billed by Google to the # user's own API key; the UI can live-pull the current model list. +# Captured 2026-10; the -preview ids of Pro and Nano Banana 2 were shut down on +# 2026-06-25. Order is the dropdown order when the API can't be listed. CLOUD_IMAGE_PRICES = { - "gemini-3-pro-image-preview": 0.134, # Nano Banana Pro (1K-2K) - "gemini-3.1-flash-image-preview": 0.067, # Nano Banana 2 (1K) + "gemini-3-pro-image": 0.134, # Nano Banana Pro (1K-2K), 5 character refs + "gemini-nano-banana-2.1": 0.034, # Nano Banana 2.1 (1K), 4 character refs + "gemini-3.1-flash-image": 0.067, # Nano Banana 2 (1K), 4 character refs + "gemini-3.1-flash-lite-image": 0.034, # Nano Banana 2 Lite (1K only) "gemini-2.5-flash-image": 0.039, # Nano Banana (1K) } +# Still served, but Google has announced a shutdown; labelled in the dropdown. +DEPRECATED_IMAGE_MODELS = {"gemini-2.5-flash-image"} # Known Gemini caption models -> USD per captioned image. ESTIMATES captured at -# build time, derived from published token pricing: ~1290 tokens for a 1K image -# in, ~120 tokens of caption out. Actual costs are billed by Google to the -# user's own key — always check current pricing. +# build time (2026-10), derived from published token pricing: ~1.5K tokens for +# a 1K image plus instruction in, ~120 tokens of caption out. Gemini 3.8 Flash +# (what flash-latest serves) doubles its price on 2027-01-01. Actual costs are +# billed by Google to the user's own key — always check current pricing. CAPTION_IMAGE_PRICES = { - "gemini-flash-latest": 0.0007, - "gemini-2.5-flash": 0.0007, - "gemini-flash-lite-latest": 0.0002, - "gemini-2.5-flash-lite": 0.0002, + "gemini-flash-latest": 0.0016, + "gemini-flash-lite-latest": 0.0008, + "gemini-3.5-flash-lite": 0.0008, } # Applied to a Gemini caption model we have no specific price for, so the # estimate degrades to "about a Flash" rather than silently reporting $0. -DEFAULT_CAPTION_PRICE = 0.0007 +DEFAULT_CAPTION_PRICE = 0.0016 # Local cache for the live Gemini model list so the dropdown can be populated # without an API call on every UI load. 24-hour TTL; falls back to stale cache, @@ -443,7 +455,10 @@ class Settings(BaseSettings): # `LDS_HF_TOKEN=…` in .env was read by nothing and silently ignored, while the # Gemini/Groq aliases worked — an inconsistency users had no way to see. hf_token: str = "" # also read from HF_TOKEN if unset - gemini_image_model: str = "gemini-3-pro-image-preview" + # "auto" = the best image model the API currently offers (see + # IMAGE_MODEL_PREFERENCE in studio/engines/gemini.py), so a retired id + # can't break the default again. + gemini_image_model: str = "auto" # Checks GitHub releases for a newer version at UI launch (cached 24h, # best-effort, never blocks). Set false to disable entirely (no network call). @@ -452,6 +467,14 @@ class Settings(BaseSettings): target_long_side: int = 1024 default_engine: str = "gemini" default_captioner: str = "qwen3vl" + + @field_validator("default_captioner") + @classmethod + def _current_captioner_key(cls, v: str) -> str: + # A retired alias in .env would otherwise reach the ③ dropdown as a + # value that isn't one of its choices. + spec = CAPTIONERS_BY_KEY.get(v) + return spec.key if spec else v # "builtin" = SAM3 via transformers in-process; "comfyui" = SAM3 workflow isolation_backend: str = "builtin" sam3_hf_id: str = "facebook/sam3" # gated: accept license + `hf auth login` @@ -538,6 +561,37 @@ def list_images(folder: Path) -> list[Path]: if p.is_file() and p.suffix.lower() in IMAGE_EXTS) +_PASSTHROUGH_FORMATS = {"PNG": "image/png", "JPEG": "image/jpeg", "WEBP": "image/webp"} + + +def api_image(path: Path, max_side: int) -> tuple[bytes, str]: + """`path` as (bytes, true mime) for a cloud or VLM call, long side <= `max_side`. + + Every API path used to send raw file bytes labelled image/png: a JPEG went out + mislabelled and a 12 MP phone photo went out whole. A file already small enough + in a format every provider reads is sent untouched; anything else is re-encoded + (PNG when it has alpha, else JPEG q95) after an EXIF-aware Lanczos downscale. + """ + from PIL import Image, ImageOps + + with Image.open(path) as im: + mime = _PASSTHROUGH_FORMATS.get(im.format or "") + if mime and max(im.size) <= max_side: + return path.read_bytes(), mime + img = ImageOps.exif_transpose(im) + if img.mode.startswith("I;16"): # 16-bit grey: convert() would clip it to white + img = img.convert("I").point(lambda v: v / 256, "L") + # has_transparency_data also catches a palette image's tRNS transparency. + img = img.convert("RGBA" if img.has_transparency_data else "RGB") + img.thumbnail((max_side, max_side), Image.LANCZOS) + buf = io.BytesIO() + if img.mode == "RGBA": + img.save(buf, "PNG") + return buf.getvalue(), "image/png" + img.save(buf, "JPEG", quality=95) + return buf.getvalue(), "image/jpeg" + + # HTTP statuses a cloud provider returns for "try again", mapped to plain English. # 503 is by far the common one — a Gemini demand spike, not anything the user did. _TRANSIENT_STATUS = { @@ -572,9 +626,15 @@ def friendly_api_error(exc: BaseException) -> str: named cause and the two things actually worth doing about them; everything else degrades to a trimmed message so no failure is ever swallowed. """ - reason = _TRANSIENT_STATUS.get(api_status_code(exc) or 0, "") + status = api_status_code(exc) or 0 + reason = _TRANSIENT_STATUS.get(status, "") if reason: return f"{reason}. {_RETRY_ADVICE}" + if status == 404 and "model" in str(exc).lower(): + # What a retired model id looks like (gemini-3-pro-image-preview, 2026-06). + return ("this model isn't available to your key — the provider may have " + "retired it. Refresh the model list and pick another (for ② images, " + "'Auto (recommended)').") text = " ".join(str(exc).split()) return text[:300] + ("…" if len(text) > 300 else "") or exc.__class__.__name__ diff --git a/studio/engines/gemini.py b/studio/engines/gemini.py index cbcf605..8527b71 100644 --- a/studio/engines/gemini.py +++ b/studio/engines/gemini.py @@ -12,7 +12,9 @@ from studio.config import ( CLOUD_IMAGE_PRICES, + DEPRECATED_IMAGE_MODELS, MODEL_CACHE_FILE, + api_image, friendly_api_error, gemini_client, is_transient_api_error, @@ -26,6 +28,21 @@ from studio.shotplan import Shot MAX_REFERENCE_IMAGES = 14 +# References are capped here before upload. ① already outputs 1024 px, but ② +# also takes files straight from the user (a pasted 12 MP photo), and every +# current Gemini image model works at 1-2K. +REFERENCE_MAX_SIDE = 2048 + +# What the "auto" model setting resolves to: the first of these the model list +# offers. Image models have no "-latest" alias, so this list is the one place a +# release moves the default. Order set by the 0.17.3 A/B (docs/ARCHITECTURE.md): +# 2.1 matched Pro's likeness at a quarter of the price and drifted less. +AUTO_MODEL = "auto" +IMAGE_MODEL_PREFERENCE = [ + "gemini-nano-banana-2.1", + "gemini-3-pro-image", # Nano Banana Pro: 5 character refs + "gemini-3.1-flash-image", # Nano Banana 2 +] # 3 attempts at 2s/4s. Image generation is slow and billed per call, so this stays # shorter than the captioner's ladder — enough to ride out a spike, not enough to # quietly burn a user's budget retrying a model that is genuinely down. @@ -46,15 +63,73 @@ def _save_model_cache(models: list[dict]) -> None: save_cloud_model_cache(models) -def _model_label(name: str, price: float | None) -> str: +# Shut-down ids still named in someone's .env (LDS_GEMINI_IMAGE_MODEL) or a saved +# CLI script, mapped to the GA release that replaced them. +_RETIRED_IMAGE_MODELS = { + "gemini-3-pro-image-preview": "gemini-3-pro-image", + "gemini-3.1-flash-image-preview": "gemini-3.1-flash-image", +} + + +def _model_label(name: str) -> str: + # Priced from the table at display time, never from the cache: a cached + # price outlives the release that corrected it. + price = CLOUD_IMAGE_PRICES.get(name) if price is None: return f"{name} (price unknown)" - return f"{name} (${price:.3f}/img est.)" + note = ", deprecated" if name in DEPRECATED_IMAGE_MODELS else "" + return f"{name} (~${price:.3f}/img est.{note})" + + +def _labelled(ids: list[str]) -> list[tuple[str, str]]: + return [(_model_label(m), m) for m in ids] def _fallback_models() -> list[tuple[str, str]]: """Static fallback when live listing is impossible.""" - return [(_model_label(m, p), m) for m, p in CLOUD_IMAGE_PRICES.items()] + return _labelled(list(CLOUD_IMAGE_PRICES)) + + +def _known_ids() -> list[str]: + """Model ids from the fresh cache, else the price table. Never a network call.""" + cached = _load_model_cache() + return [m["model_id"] for m in cached] if cached else list(CLOUD_IMAGE_PRICES) + + +def known_image_models() -> list[tuple[str, str]]: + """`list_image_models` without the network: for building the UI at startup.""" + return _labelled(_known_ids()) + + +def resolve_image_model(model: str = "", available: list[str] | None = None) -> str: + """The concrete model id for `model` ("" or "auto" = best available). + + Resolves against `available`, else the cached/static list, so cost estimates + can call it on every keystroke without touching the network. + """ + model = model or settings.gemini_image_model + if model != AUTO_MODEL: + return _RETIRED_IMAGE_MODELS.get(model, model) + ids = available if available is not None else _known_ids() + return next((m for m in IMAGE_MODEL_PREFERENCE if m in ids), IMAGE_MODEL_PREFERENCE[0]) + + +def image_price(model: str = "") -> tuple[str, float | None]: + """(resolved model id, estimated USD per 1K image or None if unknown).""" + resolved = resolve_image_model(model) + return resolved, CLOUD_IMAGE_PRICES.get(resolved) + + +def image_model_choices(models: list[tuple[str, str]]) -> list[tuple[str, str]]: + """Dropdown choices: "Auto" first, labelled with what it resolves to.""" + auto = resolve_image_model(AUTO_MODEL, [m for _, m in models]) + return [(f"Auto (recommended) → {_model_label(auto)}", AUTO_MODEL), *models] + + +def _is_image_model(name: str) -> bool: + # imagen = text-to-image only, no reference editing. "nano-banana" ids carry + # no "image" in the name (gemini-nano-banana-2.1). + return ("image" in name or "nano-banana" in name) and "imagen" not in name def list_image_models(force_refresh: bool = False) -> list[tuple[str, str]]: @@ -67,10 +142,7 @@ def list_image_models(force_refresh: bool = False) -> list[tuple[str, str]]: if not force_refresh: cached = _load_model_cache() if cached: - return [ - (_model_label(m["model_id"], m.get("price")), m["model_id"]) - for m in cached - ] + return _labelled([m["model_id"] for m in cached]) key = settings.resolved_gemini_key() if not key: @@ -81,35 +153,33 @@ def list_image_models(force_refresh: bool = False) -> list[tuple[str, str]]: found: list[dict] = [] for m in client.models.list(): name = m.name.removeprefix("models/") - if "image" not in name or "imagen" in name: # imagen = t2i only, no reference edit + if not _is_image_model(name): continue # Skip deprecated / shutdown models when the API lists them. if any(tag in name for tag in ("-shut-down", "deprecated", "-experimental")): continue - price = CLOUD_IMAGE_PRICES.get(name) found.append( { "model_id": name, "display_name": getattr(m, "display_name", name), - "price": price, "cached_at": datetime.now(tz=timezone.utc).isoformat(), } ) + # A "-preview" id the API still lists next to its GA release is the + # retired one (gemini-3-pro-image-preview, shut down 2026-06-25). + ids = {f["model_id"] for f in found} + found = [f for f in found + if not (f["model_id"].endswith("-preview") + and f["model_id"].removesuffix("-preview") in ids)] if found: _save_model_cache(found) - return [ - (_model_label(f["model_id"], f["price"]), f["model_id"]) - for f in found - ] + return _labelled([f["model_id"] for f in found]) except Exception: # Live pull failed; try stale cache as a last resort before falling # back to the static dict. stale = _load_model_cache() if stale: - return [ - (_model_label(m["model_id"], m.get("price")), m["model_id"]) - for m in stale - ] + return _labelled([m["model_id"] for m in stale]) return _fallback_models() @@ -120,8 +190,9 @@ def list_image_models(force_refresh: bool = False) -> list[tuple[str, str]]: _CAPTION_FALLBACK = [ "gemini-flash-latest", "gemini-flash-lite-latest", - "gemini-2.5-flash", - "gemini-2.5-flash-lite", + # Pinned fallbacks. Not 2.5: new projects no longer get access to it. + "gemini-3.5-flash", + "gemini-3.5-flash-lite", ] @@ -132,17 +203,17 @@ def list_caption_models(force_refresh: bool = False) -> list[tuple[str, str]]: filtered to text/vision models that support `generateContent` (excludes image-generation, embedding, TTS, and audio/live models). """ - def _labelled(ids: list[str]) -> list[tuple[str, str]]: + def _plain(ids: list[str]) -> list[tuple[str, str]]: return [(m, m) for m in ids] if not force_refresh: cached = load_caption_model_cache() if cached: - return _labelled([m["model_id"] for m in cached]) + return _plain([m["model_id"] for m in cached]) key = settings.resolved_gemini_key() if not key: - return _labelled(_CAPTION_FALLBACK) + return _plain(_CAPTION_FALLBACK) try: client = gemini_client(key) @@ -167,13 +238,13 @@ def _labelled(ids: list[str]) -> list[tuple[str, str]]: found.sort(key=lambda f: (0 if f["model_id"].endswith("latest") else 1, f["model_id"])) save_caption_model_cache(found) - return _labelled([f["model_id"] for f in found]) + return _plain([f["model_id"] for f in found]) except Exception: stale = load_caption_model_cache() if stale: - return _labelled([m["model_id"] for m in stale]) + return _plain([m["model_id"] for m in stale]) - return _labelled(_CAPTION_FALLBACK) + return _plain(_CAPTION_FALLBACK) class GeminiEngine: @@ -190,15 +261,15 @@ def __init__(self, model: str = "") -> None: # gemini_client defers the google-genai import, so local-only installs never # need it configured. self._client = gemini_client(key) - self._model = model or settings.gemini_image_model + self._model = resolve_image_model(model) def generate(self, sources: list[Path], shot: Shot, out_path: Path, seed: int) -> Path: from google.genai import types - parts: list = [ - types.Part.from_bytes(data=p.read_bytes(), mime_type="image/png") - for p in sources[:MAX_REFERENCE_IMAGES] - ] + parts: list = [] + for p in sources[:MAX_REFERENCE_IMAGES]: + data, mime = api_image(p, REFERENCE_MAX_SIDE) + parts.append(types.Part.from_bytes(data=data, mime_type=mime)) parts.append(shot.cloud_prompt) last_err: Exception | None = None diff --git a/studio/pipeline.py b/studio/pipeline.py index 610681f..1177448 100644 --- a/studio/pipeline.py +++ b/studio/pipeline.py @@ -60,6 +60,7 @@ def preprocess_sources( isolation_backend: str = "", tighten_crop: bool = False, alpha_cutout: bool = False, + front: bool = False, should_stop: ShouldStop | None = None, progress: ProgressFn = print, ) -> list[PreprocessReport]: @@ -82,7 +83,7 @@ def preprocess_sources( isolate=isolate, subject_prompt=subject_prompt, exclude_prompt=exclude_prompt, restore_backend=restore_backend, isolation_backend=isolation_backend, tighten_crop=tighten_crop, - alpha_cutout=alpha_cutout, progress=progress) + alpha_cutout=alpha_cutout, front=front, progress=progress) except Exception as e: progress(f" SKIPPED {src.name}: {e}") reports.append(failed_report(src, str(e))) diff --git a/studio/preprocess.py b/studio/preprocess.py index 4a3fc3a..bd31516 100644 --- a/studio/preprocess.py +++ b/studio/preprocess.py @@ -8,12 +8,13 @@ from __future__ import annotations +import tempfile from collections.abc import Callable from dataclasses import dataclass from pathlib import Path import numpy as np -from PIL import Image +from PIL import Image, ImageOps from studio.config import settings from studio.isolate import isolate_subject @@ -60,13 +61,19 @@ def _laplacian_variance(img: Image.Image) -> float: def _needs_restoration(img: Image.Image, path: Path, target: int) -> str: - """Return a human-readable reason, or '' if the image is fine as-is.""" + """Return a human-readable reason, or '' if the image is fine as-is. + + Judged at the size the dataset will hold. A lossy photo at 2x the target or + more loses its block artefacts in the downscale, and sharpness measured on + 12 MP of sensor noise says nothing about the 1024 px result — flagging every + phone JPEG sent it through a 4x upscale for nothing. + """ long_side = max(img.size) if long_side < target: return f"long side {long_side}px < target {target}px" - if path.suffix.lower() in (".jpg", ".jpeg", ".webp"): + if path.suffix.lower() in (".jpg", ".jpeg", ".webp") and long_side < 2 * target: return "lossy source format" - if _laplacian_variance(img) < BLUR_THRESHOLD: + if _laplacian_variance(_resize_to_target(img, target)) < BLUR_THRESHOLD: return "low sharpness (blur/grain)" return "" @@ -80,16 +87,44 @@ def _resize_to_target(img: Image.Image, target: int) -> Image.Image: return img.resize(new_size, Image.LANCZOS) -def _restore_comfyui(source: Path, out_path: Path) -> Path: +def _restore_comfyui(source: Path, out_path: Path, upscale: bool = True, + front: bool = False) -> Path: + """DeJPG, then the 4x photo upscale only when `upscale` (source below target). + + Running the 4x model on a full-size photo turned a 3000x4000 JPEG into a + 12000x16000 PNG that was saved, fetched and shrunk straight back to 1024. + """ from studio import comfy_api uploaded = comfy_api.upload_image(source) graph = comfy_api.load_template("restore_upscale") graph["1"]["inputs"]["image"] = uploaded - refs = comfy_api.run_prompt(graph, timeout=420) + if not upscale: + graph["6"]["inputs"]["images"] = ["3", 0] # save the DeJPG output + del graph["4"], graph["5"] + refs = comfy_api.run_prompt(graph, timeout=420, front=front) return comfy_api.fetch_image(refs[0], out_path) +def _stage_copy(oriented: Image.Image, changed: bool, max_side: int | None) -> Path | None: + """A temp copy of the source for restore/isolation, or None to use the file. + + Needed when EXIF rotation changed the pixels (the file on disk is sideways) + or the source exceeds `max_side` (restore/isolation cost scales with pixels). + None = keep full size. + """ + if not changed and (max_side is None or max(oriented.size) <= max_side): + return None + # convert(): a CMYK JPEG cannot be written as PNG. + staged = oriented.convert("RGBA" if oriented.has_transparency_data else "RGB") + if max_side: + staged.thumbnail((max_side, max_side), Image.LANCZOS) + with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as tmp: + path = Path(tmp.name) + staged.save(path, "PNG") + return path + + def preprocess( source: Path, work_dir: Path, @@ -102,6 +137,7 @@ def preprocess( isolation_backend: str = "", tighten_crop: bool = False, alpha_cutout: bool = False, + front: bool = False, progress: Callable[[str], None] | None = None, ) -> PreprocessReport: """Copy + clean one source image into `work_dir` at target resolution. @@ -119,7 +155,12 @@ def preprocess( target = target or settings.target_long_side restore_backend = restore_backend or settings.restore_backend work_dir.mkdir(parents=True, exist_ok=True) - img = Image.open(source).convert("RGB") + with Image.open(source) as raw: + # Phone photos store portrait shots landscape plus an EXIF rotate tag; + # ignoring it put sideways people in the dataset. + rotated = raw.getexif().get(0x0112, 1) != 1 # EXIF Orientation + oriented = ImageOps.exif_transpose(raw) + img = oriented.convert("RGB") original_size = img.size reason = _needs_restoration(img, source, target) @@ -141,15 +182,20 @@ def preprocess( # isolated, not resized) image behind, where `list_images` happily served it # to ②/③ as a finished source — a silent half-processed file in the dataset. # The stage is therefore atomic: complete output, or none at all. + # 2x the target is all a resize needs, but tighten-crop keeps only the + # subject's box: shrinking first would upscale a small subject back up. + staged = _stage_copy(oriented, rotated, + None if isolate and tighten_crop else 2 * target) try: - stage_path = source + stage_path = staged or source if restore: if restore_backend == "auto": from studio import comfy_api restore_backend = "comfyui" if comfy_api.is_up() else "basic" if restore_backend == "comfyui": - _restore_comfyui(stage_path, out_path) + _restore_comfyui(stage_path, out_path, + upscale=max(original_size) < target, front=front) stage_path = out_path else: # Basic path: Lanczos handles resolution; sharpness/compression @@ -159,7 +205,7 @@ def preprocess( if isolate: isolate_subject(stage_path, out_path, subject_prompt, exclude_prompt, backend=isolation_backend, progress=progress, - alpha_cutout=alpha_cutout, label=source.name) + alpha_cutout=alpha_cutout, label=source.name, front=front) stage_path = out_path if isolate and alpha_cutout: @@ -176,6 +222,9 @@ def preprocess( except BaseException: out_path.unlink(missing_ok=True) raise + finally: + if staged: + staged.unlink(missing_ok=True) return PreprocessReport( source=source, output=out_path, diff --git a/studio/trainer_configs.py b/studio/trainer_configs.py index 10993d2..6d8b03e 100644 --- a/studio/trainer_configs.py +++ b/studio/trainer_configs.py @@ -18,11 +18,15 @@ from __future__ import annotations +import math from pathlib import Path import yaml from pydantic import BaseModel +# ai-toolkit checkpoint/sample interval, in steps. +SAVE_EVERY = 250 + def _yaml_str(value: str) -> str: """Render `value` as a safe single-line YAML scalar for interpolation. @@ -57,6 +61,11 @@ class ModelPreset(BaseModel): arch: str = "" is_flux: bool = False quantize: bool = True + # ai-toolkit: extra `model:` keys (quantize_te, low_vram, qtype…) and the + # train.timestep_type, copied from the arch's defaults in ai-toolkit's + # extensions_built_in/diffusion_models/ui.tsx. + model_extras: dict[str, bool | str] = {} + timestep_type: str = "" # ai-toolkit train/sample knobs that vary by architecture family. Defaults # match the flow-matching models (Flux / Qwen-Image / Z-Image); SDXL, which # is not flow-matching, overrides them (ddpm scheduler, higher CFG). @@ -68,8 +77,16 @@ class ModelPreset(BaseModel): # Qwen-Image / Z-Image / Krea) from natural language. Drives the ④→⑤ advisory # that warns when the dataset's captions don't match — nothing else. expects_tags: bool = False - # musubi training script (…_train_network.py). Placeholder when unverified. + # musubi training script (…_train_network.py). Its prefix also names the two + # cache scripts. Everything below is from the arch's page in musubi's docs/. musubi_script: str = "<>" + # Every musubi arch has its own LoRA module; plain networks.lora is wrong + # for all of them. + network_module: str = "networks.lora" + musubi_args: str = "" # timestep sampling etc. for the train command + musubi_version: str = "" # --model_version, passed to all three commands + musubi_text_encoders: list[str] = ["text_encoder"] + musubi_train_te: bool = True # Krea 2 trains from cached TE outputs alone # kohya-ss sd-scripts training script (SDXL uses sdxl_train_network.py). kohya_script: str = "<>" # per-model defaults (UI pre-fills these; the user can override) @@ -81,6 +98,8 @@ class ModelPreset(BaseModel): batch_size: int = 1 +_SHIFT = "--timestep_sampling shift --weighting_scheme none --discrete_flow_shift {}" + # Curated, extensible — not exhaustive. Where a model's canonical HF id or # musubi script is not something we can guarantee, it is a <> placeholder # so the emitted config is honest rather than silently wrong. @@ -105,24 +124,49 @@ class ModelPreset(BaseModel): name_or_path="<>", arch="sdxl", quantize=False, noise_scheduler="ddpm", sample_guidance=7.0, sample_steps=25, expects_tags=True), + ModelPreset(key="qwen-image-2.1", label="Qwen-Image 2.1", + name_or_path="Comfy-Org/Qwen-Image-2.1", arch="qwen_image_2", + timestep_type="shift", sample_guidance=3.0, + # The Comfy-Org weights are pre-quantized int8 convrot. + model_extras={"quantize_te": True, "low_vram": True, + "qtype": "convrot8", "qtype_te": "convrot8"}), ModelPreset(key="zimage", label="Z-Image", - name_or_path="<>", - arch="zimage"), - ModelPreset(key="krea", label="Krea 2", - name_or_path="<>", - arch="krea"), + name_or_path="Tongyi-MAI/Z-Image", arch="zimage", + timestep_type="weighted", sample_steps=30, + model_extras={"quantize_te": True, "low_vram": True, + "qtype": "qfloat8"}), + # Key stays "krea" so saved ⑤ settings still find it. Train on Raw: + # Turbo is distilled, and samples use Raw's own 28 steps / CFG 5.5. + ModelPreset(key="krea", label="Krea 2 (Raw)", + name_or_path="krea/Krea-2-Raw", arch="krea2", + timestep_type="linear", sample_guidance=5.5, sample_steps=28, + model_extras={"quantize_te": True, "low_vram": True}), ModelPreset(key="custom", label="Custom (edit name_or_path below)", name_or_path="<>", arch="flux"), ], "musubi": [ ModelPreset(key="qwen-image", label="Qwen-Image", arch="qwen_image", - musubi_script="qwen_image_train_network.py"), + musubi_script="qwen_image_train_network.py", + network_module="networks.lora_qwen_image", musubi_version="original", + musubi_args=_SHIFT.format(2.2)), ModelPreset(key="flux-kontext", label="FLUX.1 Kontext", arch="flux_kontext", - musubi_script="<>"), + musubi_script="flux_kontext_train_network.py", + network_module="networks.lora_flux", + musubi_args="--timestep_sampling flux_shift --weighting_scheme none", + musubi_text_encoders=["text_encoder1", "text_encoder2"]), ModelPreset(key="flux2", label="FLUX.2", arch="flux2", - musubi_script="<>"), + musubi_script="flux_2_train_network.py", + network_module="networks.lora_flux_2", musubi_version="dev", + musubi_args="--timestep_sampling flux2_shift --weighting_scheme none"), ModelPreset(key="zimage", label="Z-Image", arch="zimage", - musubi_script="<>"), + musubi_script="zimage_train_network.py", + network_module="networks.lora_zimage", musubi_args=_SHIFT.format(2.0)), + # docs/krea2.md: --dit is the Raw checkpoint; training reads the cached + # text-encoder outputs, so the train command takes no --text_encoder. + ModelPreset(key="krea2", label="Krea 2 (Raw)", arch="krea2", + musubi_script="krea2_train_network.py", + network_module="networks.lora_krea2", musubi_args=_SHIFT.format(2.5), + musubi_train_te=False, rank=32, alpha=32), ], # kohya-ss sd-scripts is the standard SDXL LoRA trainer — the natural home for # the Danbooru/e621 tag captions (③). SDXL base is a runnable HF id; a family @@ -214,6 +258,13 @@ def render_aitoolkit_yaml(cfg: TrainConfig) -> str: """A complete, runnable ai-toolkit config.yaml.""" m = cfg.model arch_line = " is_flux: true\n" if m.is_flux else f' arch: "{m.arch}"\n' + model_extras = "".join( + f" {k}: {str(v).lower() if isinstance(v, bool) else _yaml_str(v)}\n" + for k, v in m.model_extras.items()) + timestep = f" timestep_type: {m.timestep_type}\n" if m.timestep_type else "" + raw_note = (" # Train on Raw. Turbo is distilled: a LoRA trained on it fights the\n" + " # distillation. Use the LoRA with Turbo at inference instead.\n" + if m.arch.startswith("krea2") else "") return ( "# ai-toolkit LoRA config generated by Dataset Deviser.\n" "# Run from your ai-toolkit install: python run.py \n" @@ -232,8 +283,10 @@ def render_aitoolkit_yaml(cfg: TrainConfig) -> str: f" linear_alpha: {cfg.alpha}\n" " save:\n" " dtype: float16\n" - " save_every: 250\n" - " max_step_saves_to_keep: 4\n" + f" save_every: {SAVE_EVERY}\n" + # Keep every save: the best checkpoint is often an early one, and + # pruning to the last few deleted it before anyone could compare. + f" max_step_saves_to_keep: {max(1, math.ceil(cfg.steps / SAVE_EVERY))}\n" " datasets:\n" f" - folder_path: \"{cfg.dataset_dir.as_posix()}\"\n" " caption_ext: \"txt\"\n" @@ -249,15 +302,18 @@ def render_aitoolkit_yaml(cfg: TrainConfig) -> str: " train_text_encoder: false\n" " gradient_checkpointing: true\n" f" noise_scheduler: {m.noise_scheduler}\n" + f"{timestep}" f" optimizer: adamw8bit\n" f" lr: {cfg.lr}\n" " dtype: bf16\n" " model:\n" + f"{raw_note}" f" name_or_path: {_yaml_str(m.name_or_path)}\n" f"{arch_line}" f" quantize: {str(m.quantize).lower()}\n" + f"{model_extras}" " sample:\n" - " sample_every: 250\n" + f" sample_every: {SAVE_EVERY}\n" f" width: {cfg.resolution}\n" f" height: {cfg.resolution}\n" " prompts:\n" @@ -304,20 +360,40 @@ def musubi_command(install_path: str, toml_path: Path, cfg: TrainConfig) -> str: Takes the whole TrainConfig, not just the preset: rank/alpha/steps/lr are user-tunable in the UI and must actually reach the command line. """ + m = cfg.model base = install_path.strip() or "<>" + prefix = m.musubi_script.removesuffix("_train_network.py") + toml = toml_path.as_posix() + version = f" --model_version {m.musubi_version}" if m.musubi_version else "" + te = " ".join(f"--{f} <>" + for f in m.musubi_text_encoders) + train_te = f" {te} \\\n" if m.musubi_train_te else "" + extra = f" {m.musubi_args} \\\n" if m.musubi_args else "" return ( - f'cd "{base}" && accelerate launch src/musubi_tuner/{cfg.model.musubi_script} \\\n' - f' --dataset_config "{toml_path.as_posix()}" \\\n' + "# 1-2: cache latents and text-encoder outputs (re-run after changing the\n" + "# images or captions). 3: train, saving a LoRA every epoch to compare.\n" + f'cd "{base}"\n' + f'python src/musubi_tuner/{prefix}_cache_latents.py --dataset_config "{toml}" ' + f"--vae <>{version}\n" + f'python src/musubi_tuner/{prefix}_cache_text_encoder_outputs.py ' + f'--dataset_config "{toml}" {te} --batch_size 1{version}\n' + "accelerate launch --num_cpu_threads_per_process 1 --mixed_precision bf16 " + f"src/musubi_tuner/{m.musubi_script} \\\n" " --dit <> \\\n" " --vae <> \\\n" - " --text_encoder <> \\\n" - " --network_module networks.lora \\\n" + f"{train_te}" + f' --dataset_config "{toml}"{version} \\\n' + " --sdpa --mixed_precision bf16 --gradient_checkpointing \\\n" + f"{extra}" + f" --optimizer_type adamw8bit --learning_rate {cfg.lr} \\\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" f" --network_alpha {cfg.alpha} \\\n" - f" --learning_rate {cfg.lr} \\\n" - f' --output_dir output --output_name "{cfg.name}" \\\n' - f" --max_train_steps {cfg.steps} --mixed_precision bf16\n" - "# Consult the musubi-tuner arch-specific doc for the exact required flags." + f" --max_train_steps {cfg.steps} --save_every_n_epochs 1 --seed 42 \\\n" + f' --output_dir output --output_name "{cfg.name}"\n' + f"# Flags follow musubi-tuner's docs/{prefix}.md example; check " + "it if a run complains." ) diff --git a/tests/test_caption_resilience.py b/tests/test_caption_resilience.py index a0b61ad..b4a841f 100644 --- a/tests/test_caption_resilience.py +++ b/tests/test_caption_resilience.py @@ -10,11 +10,17 @@ from pathlib import Path import pytest +from PIL import Image from studio import captioner as C from studio import config +def _png(path: Path) -> Path: + Image.new("RGB", (8, 8)).save(path) + return path + + class _Boom(Exception): def __init__(self, code: int = 503) -> None: super().__init__(f"{code} UNAVAILABLE") @@ -131,8 +137,7 @@ def generate_content(self, **kw): def test_gemini_retries_a_503_then_succeeds(tmp_path: Path, monkeypatch) -> None: - img = tmp_path / "a.png" - img.write_bytes(b"") + img = _png(tmp_path / "a.png") ok = type("R", (), {"text": "a caption"})() cap, attempts = _gemini_captioner(monkeypatch, [_Boom(503), _Boom(503), ok]) assert cap._caption_gemini(img, "describe") == "a caption" @@ -140,8 +145,7 @@ def test_gemini_retries_a_503_then_succeeds(tmp_path: Path, monkeypatch) -> None def test_gemini_gives_up_after_the_last_attempt(tmp_path: Path, monkeypatch) -> None: - img = tmp_path / "a.png" - img.write_bytes(b"") + img = _png(tmp_path / "a.png") cap, attempts = _gemini_captioner(monkeypatch, [_Boom(503)] * C.GEMINI_RETRIES) with pytest.raises(_Boom): cap._caption_gemini(img, "describe") @@ -149,8 +153,7 @@ def test_gemini_gives_up_after_the_last_attempt(tmp_path: Path, monkeypatch) -> def test_gemini_does_not_retry_a_bad_request(tmp_path: Path, monkeypatch) -> None: - img = tmp_path / "a.png" - img.write_bytes(b"") + img = _png(tmp_path / "a.png") cap, attempts = _gemini_captioner(monkeypatch, [_Boom(400)]) with pytest.raises(_Boom): cap._caption_gemini(img, "describe") diff --git a/tests/test_captioner.py b/tests/test_captioner.py index b4629d4..56899fb 100644 --- a/tests/test_captioner.py +++ b/tests/test_captioner.py @@ -6,17 +6,28 @@ from studio.config import CAPTIONERS_BY_KEY -def test_groq_qwen3_6_spec_exists() -> None: - spec = CAPTIONERS_BY_KEY["groq-qwen3.6"] - assert spec.model == "qwen/qwen3.6-27b" +def test_groq_qwen_spec_exists() -> None: + spec = CAPTIONERS_BY_KEY["groq-qwen"] + # qwen3.6-27b was shut down 2026-09-14. + assert spec.model == "qwen/qwen3.8-27b" assert spec.backend == "openai" assert spec.base_url == "https://api.groq.com/openai/v1" assert spec.api_key_env == "GROQ_API_KEY" -def test_groq_qwen3_6_rate_interval_is_conservative() -> None: - spec = CAPTIONERS_BY_KEY["groq-qwen3.6"] - assert spec.min_interval_s >= 3.0 +def test_retired_groq_key_still_resolves() -> None: + # An .env naming the old key must keep working, and the ③ dropdown must get + # the current key (the alias is not one of its choices). + from studio.config import Settings + + assert CAPTIONERS_BY_KEY["groq-qwen3.6"].key == "groq-qwen" + assert Settings(default_captioner="groq-qwen3.6").default_captioner == "groq-qwen" + + +def test_groq_qwen_rate_interval_fits_8k_tpm() -> None: + # ~2.6K tokens a call against 8K tokens/minute: anything under ~20 s 429s. + spec = CAPTIONERS_BY_KEY["groq-qwen"] + assert spec.min_interval_s >= 20.0 def test_groq_llama4_scout_removed() -> None: @@ -24,12 +35,12 @@ def test_groq_llama4_scout_removed() -> None: assert "groq-llama4-scout" not in CAPTIONERS_BY_KEY -def test_groq_qwen3_6_disables_reasoning() -> None: - # Qwen3.6 is a thinking model; we switch its scratchpad off so the - # response is just the caption, and keep headroom on the token budget. - spec = CAPTIONERS_BY_KEY["groq-qwen3.6"] +def test_groq_qwen_disables_reasoning() -> None: + # A thinking model; its scratchpad is switched off so the response is just + # the caption, and the budget stays small because Groq bills it per minute. + spec = CAPTIONERS_BY_KEY["groq-qwen"] assert spec.extra_params.get("reasoning_effort") == "none" - assert spec.max_tokens >= 800 + assert 300 <= spec.max_tokens <= 400 def test_gemini_caption_default_is_rolling_alias() -> None: diff --git a/tests/test_gemini.py b/tests/test_gemini.py index 9c4a2a7..bc03377 100644 --- a/tests/test_gemini.py +++ b/tests/test_gemini.py @@ -100,3 +100,151 @@ def test_list_image_models_fallback_without_key( models = gemini.list_image_models() ids = {m[1] for m in models} assert ids == set(CLOUD_IMAGE_PRICES) + + +# ---------- model resolution (0.17.3: the -preview defaults were shut down) ---------- + +def _point_cache_at(temp_cache_dir: Path, monkeypatch: pytest.MonkeyPatch) -> Path: + from studio import config as config_mod + from studio.engines import gemini + + cache_file = temp_cache_dir / "gemini_image_models.json" + monkeypatch.setattr(config_mod, "MODEL_CACHE_FILE", cache_file) + monkeypatch.setattr(gemini, "MODEL_CACHE_FILE", cache_file) + return cache_file + + +def test_auto_picks_first_preferred_model_on_offer() -> None: + from studio.engines import gemini + + pref = gemini.IMAGE_MODEL_PREFERENCE + assert gemini.resolve_image_model("auto", [pref[1], "other"]) == pref[1] + assert gemini.resolve_image_model("auto", ["other"]) == pref[0] + assert gemini.resolve_image_model("gemini-x", []) == "gemini-x" + + +def test_retired_preview_ids_map_to_their_ga_release() -> None: + from studio.engines import gemini + + assert gemini.resolve_image_model("gemini-3-pro-image-preview") == "gemini-3-pro-image" + + +def test_every_preferred_model_has_a_price() -> None: + from studio.engines import gemini + + assert set(gemini.IMAGE_MODEL_PREFERENCE) <= set(CLOUD_IMAGE_PRICES) + + +def test_default_model_is_auto_and_choices_lead_with_it( + temp_cache_dir: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + from studio.config import Settings + from studio.engines import gemini + + _point_cache_at(temp_cache_dir, monkeypatch) + assert Settings.model_fields["gemini_image_model"].default == gemini.AUTO_MODEL + choices = gemini.image_model_choices(gemini.known_image_models()) + assert choices[0][1] == gemini.AUTO_MODEL + assert gemini.IMAGE_MODEL_PREFERENCE[0] in choices[0][0] + + +def test_live_list_keeps_nano_banana_and_drops_retired_previews( + temp_cache_dir: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + from studio import config as config_mod + from studio.engines import gemini + + _point_cache_at(temp_cache_dir, monkeypatch) + monkeypatch.setattr(config_mod.settings, "gemini_api_key", "k") + names = ["gemini-nano-banana-2.1", "gemini-3-pro-image", "gemini-3-pro-image-preview", + "imagen-4.0-generate", "gemini-flash-latest"] + listed = [type("M", (), {"name": f"models/{n}"})() for n in names] + client = type("C", (), {"models": type("Ms", (), {"list": lambda self: listed})()})() + monkeypatch.setattr(gemini, "gemini_client", lambda key: client) + + ids = [m for _, m in gemini.list_image_models(force_refresh=True)] + assert ids == ["gemini-nano-banana-2.1", "gemini-3-pro-image"] + + +def test_refresh_keeps_the_users_pick(monkeypatch: pytest.MonkeyPatch) -> None: + import app + from studio.engines import gemini + + live = [("a", "gemini-3-pro-image"), ("b", "gemini-nano-banana-2.1")] + monkeypatch.setattr(gemini, "list_image_models", lambda force_refresh=False: live) + assert app.refresh_cloud_models("gemini-nano-banana-2.1").value == "gemini-nano-banana-2.1" + # A pick the API no longer offers falls back to Auto, never to whatever is first. + assert app.refresh_cloud_models("gone-model").value == gemini.AUTO_MODEL + + +# ---------- api_image: what every cloud/VLM call sends ---------- + +def test_api_image_passes_small_files_through_with_their_real_mime(tmp_path: Path) -> None: + from PIL import Image + + from studio.config import api_image + + jpg = tmp_path / "a.png" # wrong extension on purpose: the format decides + Image.new("RGB", (64, 32)).save(jpg, "JPEG") + data, mime = api_image(jpg, 1024) + assert mime == "image/jpeg" + assert data == jpg.read_bytes() + + +def test_api_image_downscales_and_keeps_alpha(tmp_path: Path) -> None: + import io + + from PIL import Image + + from studio.config import api_image + + big = tmp_path / "big.png" + Image.new("RGBA", (3000, 4000)).save(big) + data, mime = api_image(big, 2048) + assert mime == "image/png" + with Image.open(io.BytesIO(data)) as im: + assert max(im.size) == 2048 and im.mode == "RGBA" + + photo = tmp_path / "photo.tif" + Image.new("RGB", (3000, 4000)).save(photo) + data, mime = api_image(photo, 1536) + assert mime == "image/jpeg" + with Image.open(io.BytesIO(data)) as im: + assert im.size == (1152, 1536) + + +def test_a_pinned_retired_id_opens_on_its_ga_release_not_auto() -> None: + # Falling back to Auto (Pro) silently doubled a Flash pin's cost. + import app + + choices = [("a", "auto"), ("b", "gemini-3-pro-image"), ("c", "gemini-3.1-flash-image")] + assert app._cloud_model_default("gemini-3.1-flash-image-preview", choices) == \ + "gemini-3.1-flash-image" + assert app._cloud_model_default("auto", choices) == "auto" + assert app._cloud_model_default("some-unlisted-id", choices) == "auto" + + +def test_api_image_keeps_16_bit_grey_and_palette_transparency(tmp_path: Path) -> None: + import io + + import numpy as np + from PIL import Image + + from studio.config import api_image + + grey = tmp_path / "grey16.png" + Image.fromarray(np.full((3000, 3000), 32768, dtype=np.uint16)).save(grey) + data, _ = api_image(grey, 1536) + with Image.open(io.BytesIO(data)) as im: + assert 120 <= im.convert("L").getpixel((10, 10)) <= 136 # mid-grey, not clipped white + + pal = tmp_path / "pal.png" + p = Image.new("P", (3000, 3000), 1) + p.putpalette([255, 0, 0, 0, 0, 255] + [0] * 762) + p.info["transparency"] = 0 + p.paste(0, (0, 0, 100, 100)) # transparent corner + p.save(pal, transparency=0) + data, mime = api_image(pal, 1536) + assert mime == "image/png" + with Image.open(io.BytesIO(data)) as im: + assert im.mode == "RGBA" and im.getpixel((2, 2))[3] == 0 diff --git a/tests/test_preprocess.py b/tests/test_preprocess.py index 99b65e8..0374210 100644 --- a/tests/test_preprocess.py +++ b/tests/test_preprocess.py @@ -111,7 +111,7 @@ def test_preprocess_removes_the_partial_output_when_isolation_fails( _img(src, size=(300, 300)) out = tmp_path / "out" - def fake_restore(source, out_path): + def fake_restore(source, out_path, upscale=True, front=False): Image.new("RGB", (512, 512), (10, 20, 30)).save(out_path, "PNG") return out_path @@ -157,6 +157,138 @@ def fake_isolate(image_path, out_path, subject_prompt="character", *a, label="", raise AssertionError("expected the isolation failure to propagate") +# ---------- 0.17.3: phone photos (size, JPEG, EXIF rotation) ---------- + +def _noisy_jpeg(path: Path, size: tuple[int, int], exif=None) -> Path: + # Noise keeps the sharpness check from firing on a flat test image. + rng = np.random.default_rng(0) + arr = rng.integers(0, 255, (size[1], size[0], 3), dtype=np.uint8) + Image.fromarray(arr).save(path, "JPEG", quality=90, exif=exif or Image.Exif()) + return path + + +def _restore_calls(monkeypatch) -> list[tuple[tuple[int, int], bool]]: + import studio.preprocess as pp + + calls = [] + + def fake_restore(source, out_path, upscale=True, front=False): + with Image.open(source) as im: + calls.append((im.size, upscale)) + im.convert("RGB").save(out_path, "PNG") + return out_path + + monkeypatch.setattr(pp, "_restore_comfyui", fake_restore) + return calls + + +def test_preprocess_can_jump_the_comfyui_queue(tmp_path: Path, monkeypatch) -> None: + # ① had no way to prioritize, so a busy ComfyUI queue refused every ① job. + import studio.preprocess as pp + from studio import comfy_api + + fronts = [] + + def fake_run(graph, timeout=600.0, front=False): + fronts.append(front) + return ["ref"] + + def fake_isolate(image_path, out_path, *a, front=False, **kw): + fronts.append(front) + Image.open(image_path).convert("RGB").save(out_path, "PNG") + + monkeypatch.setattr(comfy_api, "upload_image", lambda p: "up.png") + monkeypatch.setattr(comfy_api, "run_prompt", fake_run) + monkeypatch.setattr(comfy_api, "fetch_image", + lambda ref, out: Image.new("RGB", (64, 64)).save(out) or out) + monkeypatch.setattr(pp, "isolate_subject", fake_isolate) + src = _noisy_jpeg(tmp_path / "small.jpg", (64, 64)) + from studio import pipeline + + pipeline.preprocess_sources([src], tmp_path / "out", target=128, front=True, + restore_backend="comfyui", progress=lambda _m: None) + assert fronts == [True, True] # restore, then isolation + + +def test_a_large_phone_jpeg_is_not_restored(tmp_path: Path, monkeypatch) -> None: + # The 4x upscale turned a 3000x4000 photo into 12000x16000 for nothing. + calls = _restore_calls(monkeypatch) + src = _noisy_jpeg(tmp_path / "phone.jpg", (3000, 4000)) + r = preprocess(src, tmp_path / "out", target=1024, isolate=False, + restore_backend="comfyui") + assert not r.restored and calls == [] + assert max(r.final_size) == 1024 + + +def test_a_small_jpeg_gets_both_restore_passes(tmp_path: Path, monkeypatch) -> None: + calls = _restore_calls(monkeypatch) + src = _noisy_jpeg(tmp_path / "small.jpg", (580, 580)) + r = preprocess(src, tmp_path / "out", target=1024, isolate=False, + restore_backend="comfyui") + assert r.restored and calls == [((580, 580), True)] + + +def test_a_forced_restore_of_a_big_photo_skips_the_4x_pass(tmp_path: Path, + monkeypatch) -> None: + calls = _restore_calls(monkeypatch) + src = _noisy_jpeg(tmp_path / "phone.jpg", (3000, 4000)) + preprocess(src, tmp_path / "out", target=1024, force_restore=True, isolate=False, + restore_backend="comfyui") + # Uploaded pre-shrunk to 2x the target, de-JPEG only. + assert calls == [((1536, 2048), False)] + + +def test_restore_without_upscale_saves_the_dejpg_output(tmp_path: Path, + monkeypatch) -> None: + import studio.preprocess as pp + from studio import comfy_api + + sent = {} + monkeypatch.setattr(comfy_api, "upload_image", lambda p: "up.png") + monkeypatch.setattr(comfy_api, "run_prompt", + lambda g, timeout, front=False: sent.setdefault("graph", g) and ["ref"]) + monkeypatch.setattr(comfy_api, "fetch_image", lambda ref, out: out) + pp._restore_comfyui(tmp_path / "a.png", tmp_path / "b.png", upscale=False) + graph = sent["graph"] + assert graph["6"]["inputs"]["images"] == ["3", 0] + linked = {v[0] for n in graph.values() for v in n["inputs"].values() + if isinstance(v, list)} + assert linked <= set(graph) # no input points at a deleted node + assert not any("4x" in str(n["inputs"]) for n in graph.values()) + + +def test_tighten_crop_isolates_the_full_size_source(tmp_path: Path, monkeypatch) -> None: + # Cropping a 2x-target copy to a small subject would upscale it back up. + import studio.preprocess as pp + + seen = [] + + def fake_isolate(image_path, out_path, *a, **kw): + with Image.open(image_path) as im: + seen.append(im.size) + im.convert("RGB").save(out_path, "PNG") + + monkeypatch.setattr(pp, "isolate_subject", fake_isolate) + src = _noisy_jpeg(tmp_path / "big.jpg", (3000, 4000)) + pp.preprocess(src, tmp_path / "out", target=1024, isolate=True, tighten_crop=True, + restore_backend="comfyui") + pp.preprocess(src, tmp_path / "out2", target=1024, isolate=True, + restore_backend="comfyui") + assert seen == [(3000, 4000), (1536, 2048)] + + +def test_exif_rotation_is_applied(tmp_path: Path, monkeypatch) -> None: + calls = _restore_calls(monkeypatch) + exif = Image.Exif() + exif[0x0112] = 6 # stored landscape, shown portrait + src = _noisy_jpeg(tmp_path / "rot.jpg", (400, 300), exif=exif) + r = preprocess(src, tmp_path / "out", target=256, force_restore=True, + isolate=False, restore_backend="comfyui") + assert r.original_size == (300, 400) + assert calls[0][0] == (300, 400) # restore sees the upright copy too + assert r.final_size == (192, 256) + + # ---------- batch resilience: one bad image must not kill the run ---------- # # A real report: SAM3 found no subject in image 3 of a batch, IsolationError diff --git a/tests/test_prop_exclusion.py b/tests/test_prop_exclusion.py index 514c380..23dd114 100644 --- a/tests/test_prop_exclusion.py +++ b/tests/test_prop_exclusion.py @@ -63,7 +63,7 @@ def test_gemini_cost_scales_with_image_count() -> None: many = estimate_caption_cost("gemini-flash", "gemini-flash-latest", 100) assert "1 image(s)" in one assert "100 image(s)" in many - assert "$0.07" in many # 100 * 0.0007 + assert "$0.16" in many # 100 * 0.0016 def test_gemini_cost_differs_by_model() -> None: diff --git a/tests/test_trainer_configs.py b/tests/test_trainer_configs.py index f4ac419..64a7041 100644 --- a/tests/test_trainer_configs.py +++ b/tests/test_trainer_configs.py @@ -92,3 +92,79 @@ def test_unknown_trainer_raises(tmp_path: Path) -> None: cfg.trainer = "nope" with pytest.raises(ValueError): write_configs(cfg) + + +# ---------- 0.17.3: arch strings pinned to what the trainers actually accept ---------- + +def _preset(trainer: str, key: str) -> ModelPreset: + return next(p for p in TRAINER_MODELS[trainer] if p.key == key) + + +# ai-toolkit: extensions_built_in/diffusion_models/ui.tsx. "krea" was never an arch. +AITOOLKIT_ARCHS = { + "flux2": ("flux2", "black-forest-labs/FLUX.2-dev"), + "qwen-image": ("qwen_image", "Qwen/Qwen-Image"), + "qwen-image-2.1": ("qwen_image_2", "Comfy-Org/Qwen-Image-2.1"), + "zimage": ("zimage", "Tongyi-MAI/Z-Image"), + "krea": ("krea2", "krea/Krea-2-Raw"), +} + +# musubi-tuner: src/musubi_tuner/*_train_network.py and networks/lora_*.py. +MUSUBI_SCRIPTS = { + "qwen-image": ("qwen_image_train_network.py", "networks.lora_qwen_image"), + "flux-kontext": ("flux_kontext_train_network.py", "networks.lora_flux"), + "flux2": ("flux_2_train_network.py", "networks.lora_flux_2"), + "zimage": ("zimage_train_network.py", "networks.lora_zimage"), + "krea2": ("krea2_train_network.py", "networks.lora_krea2"), +} + + +@pytest.mark.parametrize("key", AITOOLKIT_ARCHS) +def test_aitoolkit_arch_and_model_are_pinned(key: str, tmp_path: Path) -> None: + preset = _preset("ai-toolkit", key) + assert (preset.arch, preset.name_or_path) == AITOOLKIT_ARCHS[key] + cfg = _cfg("ai-toolkit", tmp_path).model_copy(update={"model": preset}) + model = yaml.safe_load(render_aitoolkit_yaml(cfg))["config"]["process"][0]["model"] + assert model["arch"] == preset.arch + for k, v in preset.model_extras.items(): + assert model[k] == v + + +def test_no_preset_trains_on_a_turbo_checkpoint() -> None: + # Turbo checkpoints are distilled; a LoRA must be trained on the base/Raw one. + for presets in TRAINER_MODELS.values(): + for p in presets: + assert "turbo" not in p.name_or_path.lower(), p.key + + +@pytest.mark.parametrize("key", MUSUBI_SCRIPTS) +def test_musubi_script_module_and_cache_steps(key: str, tmp_path: Path) -> None: + script, module = MUSUBI_SCRIPTS[key] + preset = _preset("musubi", key) + assert (preset.musubi_script, preset.network_module) == (script, module) + cfg = _cfg("musubi", tmp_path).model_copy(update={"model": preset}) + _, command = write_configs(cfg, install_path="/opt/musubi") + prefix = script.removesuffix("_train_network.py") + assert f"{prefix}_cache_latents.py" in command + assert f"{prefix}_cache_text_encoder_outputs.py" in command + assert f"--network_module {module}" in command + assert "--network_module networks.lora " not in command + assert "--save_every_n_epochs 1" in command + assert "< None: + preset = _preset("musubi", "krea2") + assert (preset.rank, preset.alpha) == (32, 32) + cfg = _cfg("musubi", tmp_path).model_copy(update={"model": preset}) + _, command = write_configs(cfg) + assert "--discrete_flow_shift 2.5" in command + train = command.split("accelerate launch", 1)[1] + # Krea 2 trains from cached text-encoder outputs alone. + assert "--text_encoder" not in train + + +def test_aitoolkit_keeps_every_checkpoint(tmp_path: Path) -> None: + cfg = _cfg("ai-toolkit", tmp_path).model_copy(update={"steps": 3000}) + save = yaml.safe_load(render_aitoolkit_yaml(cfg))["config"]["process"][0]["save"] + assert save["max_step_saves_to_keep"] * save["save_every"] >= 3000