[pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

studio/backend/core/training/diffusion_dit_trainer.py

@@ -258,7 +258,12 @@ def _flux_encode_latent_stats(vae, pixel_values):
     return (dist.mean - vae.config.shift_factor) * scale, dist.std * scale
 
 
-def _flux_collate(entries, device, weight_dtype, pad_to = None):
+def _flux_collate(
+    entries,
+    device,
+    weight_dtype,
+    pad_to = None,
+):
     import torch
 
     # FLUX embeds are fixed-length (encode_prompt pads to max_sequence_length), so a plain
@@ -381,7 +386,12 @@ def _qwen_encode_latent_stats(vae, pixel_values):
     return (dist.mean - mean) / std, dist.std / std
 
 
-def _qwen_collate(entries, device, weight_dtype, pad_to = None):
+def _qwen_collate(
+    entries,
+    device,
+    weight_dtype,
+    pad_to = None,
+):
     import torch
     import torch.nn.functional as F
 
@@ -482,7 +492,12 @@ def _zimage_encode_latent_stats(vae, pixel_values):
     return _zimage_encode_latents(vae, pixel_values), None
 
 
-def _zimage_collate(entries, device, weight_dtype, pad_to = None):
+def _zimage_collate(
+    entries,
+    device,
+    weight_dtype,
+    pad_to = None,
+):
     caps = [e[0].to(device = device, dtype = weight_dtype) for e in entries]
     return (caps,)
 
@@ -647,7 +662,7 @@ def _build_latent_cache(spec, vae, image_paths, cfg, device, weight_dtype, on_ev
     total = len(image_paths)
     for i, path in enumerate(image_paths):
         variants = []
-        for (u_left, u_top, flip) in plan[i]:
+        for u_left, u_top, flip in plan[i]:
             px = (
                 _load_pixel_tensor_planned(
                     path, cfg.resolution, cfg.center_crop, u_left, u_top, flip
@@ -706,7 +721,9 @@ def _maybe_compile_transformer(transformer, cfg, base_is_bnb, device, on_event)
 
     fn = getattr(transformer, "compile_repeated_blocks", None)
     if not callable(fn):
-        _emit(on_event, "warning", message = "torch.compile unavailable for this model; running eager.")
+        _emit(
+            on_event, "warning", message = "torch.compile unavailable for this model; running eager."
+        )
         return False
     try:
         dynamo_cfg = getattr(getattr(torch, "_dynamo", None), "config", None)
@@ -792,7 +809,14 @@ def run_dit_lora_training(
     perf_snap = _apply_perf_flags(cfg, device)
     try:
         return _train_dit(
-            cfg, spec, pairs, rng, device, weight_dtype, on_event, _check_stop,
+            cfg,
+            spec,
+            pairs,
+            rng,
+            device,
+            weight_dtype,
+            on_event,
+            _check_stop,
             lambda: save_on_stop,
         )
     finally:
@@ -841,8 +865,12 @@ def _train_dit(cfg, spec, pairs, rng, device, weight_dtype, on_event, _check_sto
         )
         if latent_cache is None:  # stopped during the cache build; nothing trained yet
             _emit(
-                on_event, "complete", output_dir = str(out_dir), lora_path = None,
-                stopped = True, steps_run = 0,
+                on_event,
+                "complete",
+                output_dir = str(out_dir),
+                lora_path = None,
+                stopped = True,
+                steps_run = 0,
             )
             return str(out_dir)
         try:
@@ -952,7 +980,9 @@ def _train_dit(cfg, spec, pairs, rng, device, weight_dtype, on_event, _check_sto
             noisy = (1.0 - sigmas) * latents + sigmas * noise
 
             embeds = spec.collate(
-                [caption_embeds[captions[i]] for i in idxs], device, weight_dtype,
+                [caption_embeds[captions[i]] for i in idxs],
+                device,
+                weight_dtype,
                 pad_to = qwen_pad_to,
             )
             with autocast:
@@ -1027,6 +1057,7 @@ def _make_optimizer(params, lr):
     regression for LoRA -- else torch AdamW, fused on CUDA (with a fallback when this
     build/device lacks the fused kernel)."""
     import torch
+
     try:
         import bitsandbytes as bnb
         return bnb.optim.AdamW8bit(params, lr = lr)

studio/backend/core/training/diffusion_lora_trainer.py

@@ -201,7 +201,7 @@ def _build_sdxl_latent_cache(
     total = len(image_paths)
     for i, path in enumerate(image_paths):
         variants = []
-        for (u_left, u_top, flip) in plan[i]:
+        for u_left, u_top, flip in plan[i]:
             tensor, time_ids = _load_image_tensor_planned(
                 path, cfg.resolution, cfg.center_crop, u_left, u_top, flip
             )
@@ -390,7 +390,13 @@ def run_diffusion_lora_training(
         latent_cache = None
         if use_cache:
             latent_cache = _build_sdxl_latent_cache(
-                vae, vae_scale, [p for p, _ in pairs], cfg, device, weight_dtype, on_event,
+                vae,
+                vae_scale,
+                [p for p, _ in pairs],
+                cfg,
+                device,
+                weight_dtype,
+                on_event,
                 _check_stop,
             )
             if latent_cache is None:  # stopped during the cache build; nothing trained yet

studio/backend/core/training/diffusion_train_common.py

@@ -376,11 +376,7 @@ def _emit(on_event: Optional[EventCb], type_: str, **kw: Any) -> None:
 
 
 def _plan_cache_variants(
-    num_images: int,
-    cache_variants: int,
-    center_crop: bool,
-    random_flip: bool,
-    seed: int,
+    num_images: int, cache_variants: int, center_crop: bool, random_flip: bool, seed: int
 ) -> list[list[tuple[float, float, bool]]]:
     """Seed-deterministic crop/flip plan for the latent cache: per image, up to
     ``cache_variants`` draws of (u_left, u_top, flip) with the crop as unit fractions the
@@ -405,7 +401,9 @@ def _plan_cache_variants(
 
 
 def _apply_perf_flags(
-    cfg: "DiffusionLoraConfig", device: str, cudnn_benchmark: bool = False
+    cfg: "DiffusionLoraConfig",
+    device: str,
+    cudnn_benchmark: bool = False,
 ) -> dict:
     """Set the run-scoped torch backend knobs: TF32 matmuls + high fp32 matmul precision
     (under ``cfg.enable_tf32``), plus cudnn autotuning when the caller opts in. Autotune is
@@ -444,7 +442,6 @@ def _restore_perf_flags(snap: Optional[dict]) -> None:
     if snap.get("matmul_precision"):
         try:
             import torch
-
             torch.set_float32_matmul_precision(snap["matmul_precision"])
         except Exception:  # noqa: BLE001 -- best-effort restore
             pass

studio/backend/tests/test_diffusion_train_perf.py

@@ -82,9 +82,7 @@ def test_plan_cache_variants_deterministic_and_deduped():
 # ── per-family collate fns ────────────────────────────────────────────────────
 def test_flux_collate_shapes():
     # FLUX embeds are fixed length: 3 entries batch by a plain cat; text_ids are shared.
-    entries = [
-        (torch.randn(1, 512, 32), torch.randn(1, 16), torch.randn(512, 3)) for _ in range(3)
-    ]
+    entries = [(torch.randn(1, 512, 32), torch.randn(1, 16), torch.randn(512, 3)) for _ in range(3)]
     pe, pooled, text_ids = _flux_collate(entries, "cpu", torch.float32)
     assert pe.shape == (3, 512, 32)
     assert pooled.shape == (3, 16)
@@ -239,9 +237,7 @@ def test_service_stop_save_flag():
 # ── preparing / warning events + stopped completion messages ──────────────────
 def test_apply_event_preparing_and_warning():
     svc = DiffusionTrainingService()
-    svc._apply_event(
-        {"type": "preparing", "stage": "cache_latents", "done": 4, "total": 8}
-    )
+    svc._apply_event({"type": "preparing", "stage": "cache_latents", "done": 4, "total": 8})
     st = svc.status()
     assert st["status"] == "running"
     assert st["in_model_load"] is True