Fix invalid-UTF-8 500s, the flat Canny map and the dropped DiT knobs

read_text raises UnicodeDecodeError, which is not an OSError, so one bad caption
sidecar or video sidecar 500d the info, upload and gallery routes. A flat image
now yields the all-black edge map instead of its own luminance, and the four DiT
loss knobs the trainer implements are declared so model_dump keeps them.

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")

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)

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.

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):

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:

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")

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")

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"

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.