Merge diffusion-krea2: grad norm reconciliation + review fixes

studio/backend/core/inference/diffusion_eager_patches.py

@@ -181,6 +181,11 @@ def install_compile_safe_patches() -> int:
     for cls, new_fn in _specs():
         if cls is None:
             continue
+        # torch < 2.4 has no F.rms_norm: leave diffusers' original RMSNorm.forward in
+        # place rather than installing a patch whose fast path would AttributeError.
+        if cls is _RMSNorm and not hasattr(F, "rms_norm"):
+            logger.info("eager-patch: skipping RMSNorm (this torch has no F.rms_norm)")
+            continue
         # Capture the live original BEFORE patching so the RMSNorm fast path can fall back
         # to it for the uncommon (NPU / bias / fp32-weight / tuple-dim) cases.
         if cls is _RMSNorm:

studio/backend/core/inference/diffusion_krea2.py

@@ -85,12 +85,20 @@ def load_krea2_text_encoder(
 
 def _load_model_index(repo_id: str, hf_token: Optional[str] = None) -> dict[str, Any]:
     """model_index.json as a dict, from a local path or the Hub cache."""
+    is_local_dir = False
     try:
-        local = Path(repo_id).expanduser() / "model_index.json"
+        root = Path(repo_id).expanduser()
+        is_local_dir = root.is_dir()
+        local = root / "model_index.json"
         if local.is_file():
             return json.loads(local.read_text())
     except OSError:
         pass
+    if is_local_dir:
+        # A local checkpoint dir without the file must fail clearly here: falling through
+        # to hf_hub_download with a filesystem path as the repo id would die with an
+        # opaque HFValidationError instead.
+        raise FileNotFoundError(f"model_index.json not found in local model dir {repo_id}")
     from huggingface_hub import hf_hub_download
 
     path = hf_hub_download(repo_id, "model_index.json", token = hf_token or None)

studio/backend/core/training/diffusion_train_common.py

@@ -420,6 +420,7 @@ class DiffusionLoraConfig:
             lora_target_modules = targets,
             max_grad_norm = float(self.max_grad_norm),
             hf_token = token or None,
+            num_epochs = int(self.num_epochs),
             cache_variants = int(self.cache_variants),
             compile_transformer = compile_transformer,
             base_precision = base_precision,
@@ -494,9 +495,10 @@ def discover_image_caption_pairs(
             if sidecar.is_file():
                 caption = sidecar.read_text(encoding = "utf-8").strip()
                 break
-        # 2. metadata row keyed by file name (basename or the name as written).
+        # 2. metadata row keyed by file name (basename or the relative path; as_posix so a
+        #    Windows backslash path still matches the jsonl's forward-slash keys).
         if caption is None:
-            caption = meta_caption.get(img.name) or meta_caption.get(str(img.relative_to(root)))
+            caption = meta_caption.get(img.name) or meta_caption.get(img.relative_to(root).as_posix())
         # 3. dreambooth instance prompt.
         if caption is None and instance_prompt:
             caption = instance_prompt

studio/backend/core/training/diffusion_training_service.py

@@ -99,7 +99,7 @@ def list_diffusion_runs(limit: int = 20) -> list[dict]:
     out: list[dict] = []
     for p in files[: max(0, int(limit))]:
         try:
-            rec = json.loads(p.read_text())
+            rec = json.loads(p.read_text(encoding = "utf-8"))
         except Exception:  # noqa: BLE001 -- a corrupt record never breaks the listing
             continue
         # A valid-JSON file with the wrong shape (an old or hand-edited record that is not a
@@ -124,7 +124,7 @@ def get_diffusion_run(job_id: str) -> Optional[dict]:
         return None
     p = _runs_dir() / f"{job_id}.json"
     try:
-        return json.loads(p.read_text())
+        return json.loads(p.read_text(encoding = "utf-8"))
     except Exception:  # noqa: BLE001 -- missing/corrupt record
         return None
 
@@ -403,7 +403,7 @@ class DiffusionTrainingService:
                 },
             }
             path = _runs_dir() / f"{s['job_id']}.json"
-            path.write_text(json.dumps(record))
+            path.write_text(json.dumps(record), encoding = "utf-8")
         except Exception:  # noqa: BLE001 -- persisting history must never break the run
             pass
 

studio/backend/routes/training.py

@@ -1549,7 +1549,15 @@ def _image_record(
                 caption = None
             break
     if caption is None:
+        # Basename first, then the relative path as written in the jsonl (as_posix so a
+        # Windows backslash path still matches forward-slash keys) -- the same lookup
+        # order discover_image_caption_pairs uses.
         meta = meta_captions.get(image_path.name)
+        if meta is None:
+            try:
+                meta = meta_captions.get(image_path.relative_to(folder).as_posix())
+            except ValueError:
+                meta = None
         if meta is not None:
             caption = meta
             source = "metadata"

studio/frontend/src/features/images/images-page.tsx

@@ -487,7 +487,7 @@ function AdvancedSelect({
   return (
     <div className="flex flex-col gap-1">
       <div className="flex items-center justify-between gap-2">
-        <span className="flex items-center gap-1 text-xs font-medium text-muted-foreground">
+        <span className="flex shrink-0 items-center gap-1 whitespace-nowrap text-xs font-medium text-muted-foreground">
           {label}
           {hint && <InfoHint>{hint}</InfoHint>}
         </span>
@@ -1878,7 +1878,7 @@ export function ImagesPage({ active = true }: { active?: boolean }) {
       {!status?.loaded || status.model_kind === "gguf" ? (
         <AdvancedSelect
           label="GGUF compute"
-          desc="Off runs the GGUF as-is. INT8/FP8/FP4 dequantise the transformer onto low-precision tensor cores for a faster step, at the cost of a larger download and more VRAM."
+          desc="Off runs the GGUF as-is. INT8/FP8/FP4 instead download the base model's bf16 transformer and quantise it directly onto low-precision tensor cores (the GGUF is not requantised): a faster step, at the cost of a larger download and more VRAM."
           hint="Optional speed-up for GGUF models. Off runs the GGUF as-is. FP8/INT8/FP4 instead load the FULL base model and quantise its transformer onto low-precision tensor cores: faster per step, but a larger download and more VRAM, and it falls back to the GGUF if it can't fit. Needs CUDA."
           value={transformerQuant}
           onValueChange={(v) => setTransformerQuant(v as typeof transformerQuant)}

studio/frontend/src/features/images/train/diffusion-charts.tsx

@@ -106,7 +106,6 @@ export function DiffusionCharts({
     () => buildYDomain(gradData.map((p) => p.displayGradNorm)),
     [gradData],
   );
-
   const avgRaw =
     lossItems.length > 0
       ? +(lossItems.reduce((s, p) => s + p.loss, 0) / lossItems.length).toFixed(4)