diff --git a/studio/backend/core/inference/diffusion_controlnet.py b/studio/backend/core/inference/diffusion_controlnet.py index 3c51aaa64d..8947e85bc6 100644 --- a/studio/backend/core/inference/diffusion_controlnet.py +++ b/studio/backend/core/inference/diffusion_controlnet.py @@ -250,7 +250,9 @@ def preprocess_control(image: Any, control_type: str) -> Any: mag = np.hypot(gx, gy) peak = float(mag.max()) if peak <= 1e-6: - return image # flat image -> nothing to trace + # Flat image -> no edges, which is an all-black map. Returning the source would + # instead condition the ControlNet on its raw luminance. + return Image.new("RGB", image.size, (0, 0, 0)) mag = mag / peak * 255.0 edges = (mag > 40.0).astype(np.uint8) * 255 # white edges on black (ControlNet convention) return Image.fromarray(edges).convert("RGB") diff --git a/studio/backend/core/inference/video_gallery.py b/studio/backend/core/inference/video_gallery.py index 860912fb3d..ead7f0c4d5 100644 --- a/studio/backend/core/inference/video_gallery.py +++ b/studio/backend/core/inference/video_gallery.py @@ -195,7 +195,8 @@ _REQUIRED_META = ( def _read_meta(sidecar: Path) -> Optional[dict[str, Any]]: try: raw = sidecar.read_text(encoding = "utf-8") - except OSError: + except (OSError, UnicodeError): + # Invalid UTF-8 is a corrupt sidecar, not a listing failure. return None try: meta = json.loads(raw) diff --git a/studio/backend/core/training/diffusion_train_common.py b/studio/backend/core/training/diffusion_train_common.py index 45566ce47c..0f3a1206c3 100644 --- a/studio/backend/core/training/diffusion_train_common.py +++ b/studio/backend/core/training/diffusion_train_common.py @@ -815,7 +815,12 @@ def discover_image_caption_pairs( sidecar = img.with_suffix(ext) if sidecar.is_file(): sidecar_present = True - caption = sidecar.read_text(encoding = "utf-8").strip() + try: + caption = sidecar.read_text(encoding = "utf-8").strip() + except (OSError, UnicodeError): + # Unreadable sidecar reads as the empty tombstone, so the + # instance_prompt fallback applies instead of a 500 preflight. + caption = "" break # 2. metadata row keyed by file name (basename or relative path; as_posix so a Windows # backslash path matches the jsonl's forward-slash keys). A sidecar, even empty, wins. diff --git a/studio/backend/models/training.py b/studio/backend/models/training.py index a80d0db91d..e857ba3793 100644 --- a/studio/backend/models/training.py +++ b/studio/backend/models/training.py @@ -803,6 +803,24 @@ class DiffusionTrainingStartRequest(BaseModel): "or auto (pick by free VRAM + GPU class). Dense modes need a non-prequant base." ), ) + # DiT-only levers the trainer implements. Undeclared, they were silently dropped by + # model_dump(); the defaults match DiffusionLoraConfig so an existing caller is unaffected. + ema_decay: float = Field( + 0.0, ge = 0.0, lt = 1.0, description = "EMA of the LoRA weights; 0 disables it" + ) + cfg_dropout: float = Field( + 0.0, ge = 0.0, le = 1.0, description = "Chance of dropping the caption to an empty prompt" + ) + weighting_scheme: Literal["none", "bell"] = Field( + "none", description = "Flow-matching timestep sampling: uniform, or logit-normal (bell)" + ) + flow_shift: Optional[float | Literal["auto"]] = Field( + None, + description = ( + "Flow-matching timestep shift. null uses the family default " + "(auto for qwen-image, 1.0 otherwise)." + ), + ) class DiffusionTrainingStopRequest(BaseModel): diff --git a/studio/backend/routes/training.py b/studio/backend/routes/training.py index 8aeea98d32..3589abf604 100644 --- a/studio/backend/routes/training.py +++ b/studio/backend/routes/training.py @@ -1557,7 +1557,8 @@ def _resolve_dataset_caption( if sidecar.is_file(): try: caption = sidecar.read_text(encoding = "utf-8").strip() - except OSError: + except (OSError, UnicodeError): + # Unreadable/invalid UTF-8 sidecar: no caption, not a 500. caption = None break if caption is None: diff --git a/studio/backend/tests/test_diffusion_controlnet.py b/studio/backend/tests/test_diffusion_controlnet.py index e7867045c3..36ce756773 100644 --- a/studio/backend/tests/test_diffusion_controlnet.py +++ b/studio/backend/tests/test_diffusion_controlnet.py @@ -146,11 +146,15 @@ def test_preprocess_control_passthrough_and_canny(): img = Image.new("RGB", (32, 24), (10, 20, 30)) # passthrough returns the same object. assert dc.preprocess_control(img, "passthrough") is img - # a flat image has no edges -> canny falls back to passthrough (no black map). - assert dc.preprocess_control(img, "canny") is img - # an image with structure yields an edge map: RGB, same size, some white pixels. + # a flat image has no edges, so the map is all black -- passing the source through would + # condition the ControlNet on its raw luminance, which is not an edge map at all. import numpy as np + flat = dc.preprocess_control(img, "canny") + assert flat.mode == "RGB" and flat.size == (32, 24) + assert np.asarray(flat).max() == 0 + # an image with structure yields an edge map: RGB, same size, some white pixels. + arr = np.zeros((24, 32, 3), np.uint8) arr[:, 16:, :] = 255 # a hard vertical edge edged = dc.preprocess_control(Image.fromarray(arr), "canny") diff --git a/studio/backend/tests/test_diffusion_lora_trainer.py b/studio/backend/tests/test_diffusion_lora_trainer.py index 20ac244d5b..4cad9388fa 100644 --- a/studio/backend/tests/test_diffusion_lora_trainer.py +++ b/studio/backend/tests/test_diffusion_lora_trainer.py @@ -87,6 +87,26 @@ def test_discover_empty_sidecar_without_instance_prompt_skips_image(tmp_path): assert pairs == {str(tmp_path / "cap.png"): "kept"} +def test_discover_reads_invalid_utf8_sidecar_as_tombstone(tmp_path): + # A sidecar with invalid UTF-8 raised UnicodeDecodeError out of the preflight (a 500 on + # /diffusion/start). It now reads as empty, so the instance prompt still applies. + _touch(tmp_path / "cat.png") + (tmp_path / "cat.txt").write_bytes(b"\xff\xfe not utf-8") + pairs = discover_image_caption_pairs(tmp_path, instance_prompt = "a photo of sks cat") + assert pairs == [(str(tmp_path / "cat.png"), "a photo of sks cat")] + + +def test_discover_null_metadata_caption_is_not_the_string_none(tmp_path): + # str(None) stored "None" as a real caption; a null row must fall through to the + # instance prompt instead. + _touch(tmp_path / "cat.png") + (tmp_path / "metadata.jsonl").write_text( + json.dumps({"file_name": "cat.png", "text": None}) + "\n", encoding = "utf-8" + ) + pairs = discover_image_caption_pairs(tmp_path, instance_prompt = "a photo of sks cat") + assert pairs == [(str(tmp_path / "cat.png"), "a photo of sks cat")] + + def test_discover_skips_uncaptioned_without_instance_prompt(tmp_path): _touch(tmp_path / "cap.png") _touch(tmp_path / "nocap.png") diff --git a/studio/backend/tests/test_diffusion_training.py b/studio/backend/tests/test_diffusion_training.py index 3c4143ba15..cb45f0304b 100644 --- a/studio/backend/tests/test_diffusion_training.py +++ b/studio/backend/tests/test_diffusion_training.py @@ -555,6 +555,39 @@ def test_route_start_forwards_num_epochs(client): assert client._fake.started_with["num_epochs"] == 8 +def test_route_start_forwards_dit_loss_knobs(client): + # The trainer implements these, but the request schema did not declare them, so + # model_dump() dropped them and the run silently used the defaults. + body = { + **_BODY, + "ema_decay": 0.99, + "cfg_dropout": 0.1, + "weighting_scheme": "bell", + "flow_shift": 3.0, + } + r = client.post("/api/train/diffusion/start", json = body) + assert r.status_code == 200, r.text + started = client._fake.started_with + assert started["ema_decay"] == 0.99 and started["cfg_dropout"] == 0.1 + assert started["weighting_scheme"] == "bell" and started["flow_shift"] == 3.0 + + +def test_request_model_dit_loss_knob_bounds(): + # Bounds mirror DiffusionLoraConfig.normalized(); flow_shift also accepts "auto". + from pydantic import ValidationError + + from models.training import DiffusionTrainingStartRequest + + base = {"base_model": "b", "data_dir": "d", "output_dir": "o"} + defaults = DiffusionTrainingStartRequest(**base) + assert (defaults.ema_decay, defaults.cfg_dropout) == (0.0, 0.0) + assert defaults.weighting_scheme == "none" and defaults.flow_shift is None + assert DiffusionTrainingStartRequest(**base, flow_shift = "auto").flow_shift == "auto" + for bad in ({"ema_decay": 1.0}, {"cfg_dropout": 1.5}, {"weighting_scheme": "bogus"}): + with pytest.raises(ValidationError): + DiffusionTrainingStartRequest(**base, **bad) + + def test_request_model_num_epochs_bounds(): # The request schema mirrors DiffusionLoraConfig's 0..1000 num_epochs range. from pydantic import ValidationError @@ -1136,6 +1169,20 @@ def test_diffusion_info_tolerates_invalid_utf8_jsonl(client, dataset_roots): assert summary["caption_count"] == 0 +def test_diffusion_info_tolerates_invalid_utf8_sidecar(client, dataset_roots): + # Same for a per-image .txt sidecar: read_text raises UnicodeDecodeError, which is not an + # OSError, so an unguarded read 500s the info endpoint after the upload already committed. + ds_root, _ = dataset_roots + folder = ds_root / "bad-utf8-sidecar" + folder.mkdir() + (folder / "a.png").write_bytes(b"x") + (folder / "a.txt").write_bytes(b"\xff\xfe not valid utf-8") + r = client.get("/api/train/diffusion/info") + assert r.status_code == 200, r.text + summary = next(d for d in r.json()["datasets"] if d["name"] == "bad-utf8-sidecar") + assert summary["caption_count"] == 0 + + def test_diffusion_dataset_mutations_blocked_while_training_active(client, dataset_roots): ds_root, _ = dataset_roots folder = ds_root / "locked" diff --git a/studio/backend/tests/test_video_gallery.py b/studio/backend/tests/test_video_gallery.py index d072dd44da..5f8f9a308c 100644 --- a/studio/backend/tests/test_video_gallery.py +++ b/studio/backend/tests/test_video_gallery.py @@ -244,6 +244,16 @@ def test_list_skips_corrupt_sidecar(): assert [r["prompt"] for r in listed] == ["ours"] +def test_list_skips_invalid_utf8_sidecar(): + # Invalid UTF-8 raises UnicodeDecodeError, which is not an OSError: one corrupt sidecar must + # be skipped like any other, not 500 the whole gallery listing. + directory = gallery.gallery_dir() + (directory / "badbytes.mp4").write_bytes(_mp4()) + (directory / "badbytes.json").write_bytes(b"\xff\xfe{}") + gallery.save(_mp4(), _meta(prompt = "ours")) + assert [r["prompt"] for r in gallery.list_videos()] == ["ours"] + + def test_clear_preserves_mp4_with_present_but_invalid_sidecar(): # A hand-dropped MP4 whose sidecar parses but lacks the required recipe keys (e.g. "{}") is # hidden by list_videos, so clear must not destroy it while removing the owned pair.