Merge branch 'diffusion-auto-badges' into diffusion-more-families

studio/backend/core/inference/diffusion.py

@@ -71,6 +71,7 @@ from .diffusion_speed import (
 from .diffusion_attention import (
     apply_attention_backend,
     select_attention_backend,
+    _ensure_attention_backend_installed,
 )
 from . import diffusion_compile_cache as compile_cache
 from . import diffusion_gguf_compile as gguf_compile
@@ -960,6 +961,23 @@ class DiffusionBackend:
 
         import diffusers
 
+        # Pre-install the optional attention kernel BEFORE taking the load locks. The
+        # wheel-only pip install can run up to 600s, and doing it under _lock /
+        # _generate_lock (as the in-lock apply_attention_backend otherwise would) blocks
+        # unload() and cancellation for that whole window. Only an explicit backend pulls
+        # a package -- auto resolves to cuDNN / native, which ship with torch -- and an
+        # explicit backend's resolution ignores the speed tier, so it can run here without
+        # effective_speed. Best-effort: the authoritative resolve + set still happens under
+        # the lock, where the now-satisfied install call is a fast no-op.
+        try:
+            preinstall_backend = select_attention_backend(
+                target, attention_backend, speed_active = True
+            )
+            if preinstall_backend is not None:
+                _ensure_attention_backend_installed(preinstall_backend, logger)
+        except Exception:  # noqa: BLE001 — the locked path re-resolves and validates
+            pass
+
         # Signal an in-flight denoise to abort, then take _generate_lock to WAIT for
         # it to actually exit before allocating the replacement: a load is about to
         # claim VRAM, so unlike unload() it must not overlap a still-live pipeline.

studio/backend/core/inference/diffusion_attention.py

@@ -213,11 +213,24 @@ def _ensure_attention_backend_installed(backend: str, logger: Any = None) -> Non
         )
     except Exception as exc:  # noqa: BLE001 — no wheel / no network -> native fallback
         if logger is not None:
-            logger.warning(
-                "diffusion.attention: could not install %s (%s); falling back to default",
-                package,
-                exc,
-            )
+            # A failed pip install raises CalledProcessError whose str() shows only the
+            # exit code and command; the real reason (no matching wheel, resolver error)
+            # is in exc.stderr. Surface it so a fallback to native is diagnosable.
+            stderr = getattr(exc, "stderr", None)
+            if stderr:
+                if isinstance(stderr, bytes):
+                    stderr = stderr.decode("utf-8", errors = "replace")
+                logger.warning(
+                    "diffusion.attention: could not install %s; pip failed with: %s",
+                    package,
+                    stderr.strip() or str(exc),
+                )
+            else:
+                logger.warning(
+                    "diffusion.attention: could not install %s (%s); falling back to default",
+                    package,
+                    exc,
+                )
 
 
 def apply_attention_backend(

studio/backend/core/inference/diffusion_speed.py

@@ -218,7 +218,7 @@ def apply_speed_optims(
     # and the UNSLOTH_DISABLE_FP16_ACCUM kill switch.
     if on_cuda:
         applied["fp16_accum"] = _enable_fp16_accumulation(
-            family, logger, dtype = getattr(target, "dtype", None), speed_mode = speed_mode
+            family, logger, dtype = getattr(target, "dtype", None), speed_mode = mode
         )
 
     # --- the compile lever, remapped per tier ----------------------------------------

studio/backend/core/training/diffusion_train_common.py

@@ -704,6 +704,15 @@ def _coerce_gradient_checkpointing(value: Any) -> bool:
     return bool(value)
 
 
+def _coerce_bool(value: Any) -> bool:
+    """Coerce a flag that may arrive as a string through the generic Studio config path
+    (e.g. "false" / "0" / "off"). A non-empty string like "false" is otherwise truthy, so
+    an opt-out would silently no-op. A real bool passes through."""
+    if isinstance(value, str):
+        return value.strip().lower() not in ("", "none", "false", "0", "no", "off")
+    return bool(value)
+
+
 def _config_from_dict(config: dict) -> DiffusionLoraConfig:
     """Build a DiffusionLoraConfig from a plain dict. Unknown keys are ignored so a richer
     request payload (UI form) does not break construction; a small set of generic Studio
@@ -734,4 +743,7 @@ def _config_from_dict(config: dict) -> DiffusionLoraConfig:
         kwargs["gradient_checkpointing"] = _coerce_gradient_checkpointing(
             kwargs["gradient_checkpointing"]
         )
+    for flag in ("cache_latents", "enable_tf32"):
+        if flag in kwargs:
+            kwargs[flag] = _coerce_bool(kwargs[flag])
     return DiffusionLoraConfig(**kwargs)

studio/backend/tests/test_diffusion_attention.py

@@ -290,6 +290,35 @@ def test_install_never_attempted_for_builtin_backends(monkeypatch):
     assert run.calls == []
 
 
+def test_install_failure_logs_pip_stderr(monkeypatch):
+    # A CalledProcessError's str() hides the pip reason; the warning must surface the
+    # captured stderr (decoding bytes) so a fallback to native is diagnosable.
+    monkeypatch.setenv("UNSLOTH_DIFFUSION_ATTENTION_INSTALL", "auto")
+    import importlib.util
+    import subprocess as sp
+
+    monkeypatch.setattr(importlib.util, "find_spec", lambda name: None)
+
+    def _boom(cmd, **kwargs):
+        raise sp.CalledProcessError(
+            returncode = 1, cmd = cmd, stderr = b"ERROR: No matching distribution found"
+        )
+
+    _stub_subprocess(monkeypatch, _boom)
+
+    warnings: list[str] = []
+
+    class _Logger:
+        def info(self, *a, **k):
+            pass
+
+        def warning(self, msg, *args):
+            warnings.append(msg % args if args else msg)
+
+    att._ensure_attention_backend_installed("sage", _Logger())
+    assert warnings and "No matching distribution found" in warnings[-1]
+
+
 def test_install_failure_falls_back_to_native(monkeypatch):
     # pip failing (no wheel for this platform) must not break the load: the apply
     # path proceeds, set_attention_backend raises on the missing package, and the

studio/backend/tests/test_diffusion_speed.py

@@ -485,7 +485,7 @@ def test_fp16_accum_allowed_on_fp16_dtype_under_max(monkeypatch):
         _target(dtype = "float16"),
         is_gguf = True,
         family = _family(),
-        speed_mode = "max",
+        speed_mode = "MAX",
     )
     assert applied["fp16_accum"] is True
     assert torch.backends.cuda.matmul.allow_fp16_accumulation is True

studio/backend/tests/test_diffusion_train_perf.py

@@ -186,6 +186,20 @@ def test_config_validates_new_fields():
     assert cfg.enable_tf32 is False
     assert cfg.cache_latents is False
 
+    # String flags from the generic Studio dict path are coerced: "false" is otherwise a
+    # non-empty (truthy) string, so an opt-out would silently no-op.
+    cfg = _config_from_dict(
+        {
+            "base_model": _SDXL,
+            "data_dir": "d",
+            "output_dir": "o",
+            "enable_tf32": "false",
+            "cache_latents": "0",
+        }
+    )
+    assert cfg.enable_tf32 is False
+    assert cfg.cache_latents is False
+
 
 # ── torch.compile policy ──────────────────────────────────────────────────────
 def test_should_compile_policy():