From 56b08cc15b1ddcb3d6f941ba590a5bd18264cd66 Mon Sep 17 00:00:00 2001
From: "pre-commit-ci[bot]"
 <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Date: Thu, 2 Jul 2026 16:08:56 +0000
Subject: [PATCH] [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci
---
 .../core/inference/diffusion_families.py      |   1 +
 .../core/training/diffusion_dit_trainer.py    | 123 ++++++++++++------
 .../core/training/diffusion_lora_trainer.py   |   1 -
 .../core/training/diffusion_train_common.py   |   4 +-
 studio/backend/routes/training.py             |  49 ++++---
 .../tests/test_diffusion_dataset_api.py       |  37 +++---
 .../tests/test_diffusion_lora_trainer.py      |   7 +-
 .../backend/tests/test_diffusion_training.py  |  29 ++++-
 8 files changed, 157 insertions(+), 94 deletions(-)

diff --git a/studio/backend/core/inference/diffusion_families.py b/studio/backend/core/inference/diffusion_families.py
index 7cff2c3fec..775be1c962 100644
--- a/studio/backend/core/inference/diffusion_families.py
+++ b/studio/backend/core/inference/diffusion_families.py
@@ -320,6 +320,7 @@ def trainable_family_names() -> tuple[str, ...]:
     """Names of families Studio can train a LoRA on, in registry order."""
     return tuple(fam.name for fam in _FAMILIES if fam.trainable)
 
+
 # Editing / inpaint checkpoints share an arch keyword but need a different
 # pipeline and an input image, which this text-to-image backend doesn't drive.
 # "layered" rejects Qwen-Image-Layered: its transformer sets additional_t_cond=True
diff --git a/studio/backend/core/training/diffusion_dit_trainer.py b/studio/backend/core/training/diffusion_dit_trainer.py
index 6a81b91804..82b65ee5ad 100644
--- a/studio/backend/core/training/diffusion_dit_trainer.py
+++ b/studio/backend/core/training/diffusion_dit_trainer.py
@@ -45,8 +45,14 @@ from core.training.diffusion_train_common import (
 # also carry added-kv projections; Z-Image is single-stream. Kept here (not in the generic
 # DEFAULT_LORA_TARGETS) because they are architecture-specific.
 _FLUX_TARGETS = (
-    "to_q", "to_k", "to_v", "to_out.0",
-    "add_q_proj", "add_k_proj", "add_v_proj", "to_add_out",
+    "to_q",
+    "to_k",
+    "to_v",
+    "to_out.0",
+    "add_q_proj",
+    "add_k_proj",
+    "add_v_proj",
+    "to_add_out",
 )
 _QWEN_TARGETS = _FLUX_TARGETS
 _ZIMAGE_TARGETS = ("to_q", "to_k", "to_v", "to_out.0")
@@ -128,7 +134,6 @@ def _encoders_to_device(pipe, device) -> None:
 def _bnb_4bit_config():
     from diffusers import BitsAndBytesConfig as DiffusersBnb
     import torch
-
     return DiffusersBnb(
         load_in_4bit = True,
         bnb_4bit_quant_type = "nf4",
@@ -147,7 +152,6 @@ def _repo_is_prequantized(base_model: str) -> bool:
 def _load_quantized_transformer(transformer_cls, cfg):
     """Load ``cfg.base_model``'s transformer subfolder as a trainable nf4 QLoRA module."""
     import torch
-
     return transformer_cls.from_pretrained(
         cfg.base_model,
         subfolder = "transformer",
@@ -171,7 +175,9 @@ def _flux_load(cfg, device, weight_dtype, qlora):
             token = cfg.hf_token,
         )
         pipe = FluxPipeline.from_pretrained(
-            cfg.base_model, transformer = transformer, torch_dtype = torch.bfloat16,
+            cfg.base_model,
+            transformer = transformer,
+            torch_dtype = torch.bfloat16,
             token = cfg.hf_token,
         )
     else:
@@ -191,7 +197,10 @@ def _flux_encode_prompts(pipe, captions, device):
     with torch.no_grad():
         for cap in captions:
             pe, pooled, text_ids = pipe.encode_prompt(
-                prompt = cap, prompt_2 = cap, device = device, num_images_per_prompt = 1,
+                prompt = cap,
+                prompt_2 = cap,
+                device = device,
+                num_images_per_prompt = 1,
                 max_sequence_length = 512,
             )
             out.append((pe.cpu(), pooled.cpu(), text_ids.cpu()))
@@ -233,7 +242,6 @@ def _flux_forward(transformer, noisy, timesteps, sigmas, embeds_batch, cfg, devi
 
 def _flux_save(pipe_cls, out_dir, transformer_lora_layers):
     from diffusers import FluxPipeline
-
     FluxPipeline.save_lora_weights(
         save_directory = out_dir,
         transformer_lora_layers = transformer_lora_layers,
@@ -265,7 +273,9 @@ def _qwen_encode_prompts(pipe, captions, device):
     with torch.no_grad():
         for cap in captions:
             pe, mask = pipe.encode_prompt(
-                prompt = cap, device = device, num_images_per_prompt = 1,
+                prompt = cap,
+                device = device,
+                num_images_per_prompt = 1,
                 max_sequence_length = 1024,
             )
             out.append((pe.cpu(), mask.cpu() if mask is not None else None))
@@ -311,7 +321,6 @@ def _qwen_forward(transformer, noisy, timesteps, sigmas, embeds_batch, cfg, devi
 
 def _qwen_save(pipe_cls, out_dir, transformer_lora_layers):
     from diffusers import QwenImagePipeline
-
     QwenImagePipeline.save_lora_weights(
         save_directory = out_dir,
         transformer_lora_layers = transformer_lora_layers,
@@ -342,7 +351,9 @@ def _zimage_encode_prompts(pipe, captions, device):
     with torch.no_grad():
         for cap in captions:
             pe, _neg = pipe.encode_prompt(
-                prompt = cap, device = device, do_classifier_free_guidance = False,
+                prompt = cap,
+                device = device,
+                do_classifier_free_guidance = False,
                 max_sequence_length = 512,
             )
             # pe is a list of one variable-length [seq, 2560] tensor per prompt.
@@ -353,7 +364,6 @@ def _zimage_encode_prompts(pipe, captions, device):
 
 def _zimage_encode_latents(vae, pixel_values):
     import torch
-
     with torch.no_grad():
         lat = vae.encode(pixel_values.to(torch.float32)).latent_dist.mode()
     return (lat - vae.config.shift_factor) * vae.config.scaling_factor
@@ -374,7 +384,6 @@ def _zimage_forward(transformer, noisy, timesteps, sigmas, embeds_batch, cfg, de
 
 def _zimage_save(pipe_cls, out_dir, transformer_lora_layers):
     from diffusers import ZImagePipeline
-
     ZImagePipeline.save_lora_weights(
         save_directory = out_dir,
         transformer_lora_layers = transformer_lora_layers,
@@ -384,19 +393,34 @@ def _zimage_save(pipe_cls, out_dir, transformer_lora_layers):
 
 _SPECS: dict[str, _FamilySpec] = {
     "flux.1": _FamilySpec(
-        family = "flux.1", lora_targets = _FLUX_TARGETS, force_bf16 = False,
-        load = _flux_load, encode_prompts = _flux_encode_prompts,
-        encode_latents = _flux_encode_latents, forward = _flux_forward, save = _flux_save,
+        family = "flux.1",
+        lora_targets = _FLUX_TARGETS,
+        force_bf16 = False,
+        load = _flux_load,
+        encode_prompts = _flux_encode_prompts,
+        encode_latents = _flux_encode_latents,
+        forward = _flux_forward,
+        save = _flux_save,
     ),
     "qwen-image": _FamilySpec(
-        family = "qwen-image", lora_targets = _QWEN_TARGETS, force_bf16 = True,
-        load = _qwen_load, encode_prompts = _qwen_encode_prompts,
-        encode_latents = _qwen_encode_latents, forward = _qwen_forward, save = _qwen_save,
+        family = "qwen-image",
+        lora_targets = _QWEN_TARGETS,
+        force_bf16 = True,
+        load = _qwen_load,
+        encode_prompts = _qwen_encode_prompts,
+        encode_latents = _qwen_encode_latents,
+        forward = _qwen_forward,
+        save = _qwen_save,
     ),
     "z-image": _FamilySpec(
-        family = "z-image", lora_targets = _ZIMAGE_TARGETS, force_bf16 = True,
-        load = _zimage_load, encode_prompts = _zimage_encode_prompts,
-        encode_latents = _zimage_encode_latents, forward = _zimage_forward, save = _zimage_save,
+        family = "z-image",
+        lora_targets = _ZIMAGE_TARGETS,
+        force_bf16 = True,
+        load = _zimage_load,
+        encode_prompts = _zimage_encode_prompts,
+        encode_latents = _zimage_encode_latents,
+        forward = _zimage_forward,
+        save = _zimage_save,
     ),
 }
 
@@ -499,8 +523,9 @@ def run_dit_lora_training(
     _emit(on_event, "model_load_started", num_images = len(pairs))
     if _check_stop():
         out_dir = Path(cfg.output_dir).expanduser()
-        _emit(on_event, "complete", output_dir = str(out_dir), lora_path = None,
-              stopped = True, steps_run = 0)
+        _emit(
+            on_event, "complete", output_dir = str(out_dir), lora_path = None, stopped = True, steps_run = 0
+        )
         return str(out_dir)
 
     # QLoRA by default for the big DiTs (nf4 transformer). The prequant Qwen/Z-Image repos
@@ -537,7 +562,6 @@ def run_dit_lora_training(
         # inputs do not require grad, which happens with a frozen 4-bit base).
         import functools
         import torch.utils.checkpoint as _ckpt
-
         transformer.enable_gradient_checkpointing(
             gradient_checkpointing_func = functools.partial(_ckpt.checkpoint, use_reentrant = False)
         )
@@ -562,9 +586,13 @@ def run_dit_lora_training(
         step_loss = 0.0
         for _ in range(cfg.gradient_accumulation_steps):
             i = rng.randrange(len(image_paths))
-            px = _load_pixel_tensor(
-                image_paths[i], cfg.resolution, cfg.center_crop, cfg.random_flip, rng
-            ).unsqueeze(0).to(device)
+            px = (
+                _load_pixel_tensor(
+                    image_paths[i], cfg.resolution, cfg.center_crop, cfg.random_flip, rng
+                )
+                .unsqueeze(0)
+                .to(device)
+            )
             latents = spec.encode_latents(vae, px).to(weight_dtype)
 
             noise = torch.randn_like(latents)
@@ -574,7 +602,8 @@ def run_dit_lora_training(
 
             emb = caption_embeds[captions[i]]
             emb_dev = tuple(
-                t.to(device = device, dtype = weight_dtype) if (t is not None and t.is_floating_point())
+                t.to(device = device, dtype = weight_dtype)
+                if (t is not None and t.is_floating_point())
                 else (t.to(device) if t is not None else None)
                 for t in emb
             )
@@ -607,12 +636,18 @@ def run_dit_lora_training(
                 peak_gb = round(torch.cuda.max_memory_allocated() / 1e9, 2)
             sps = round(
                 (done * cfg.train_batch_size * cfg.gradient_accumulation_steps)
-                / max(time.time() - t_start, 1e-6), 3,
+                / max(time.time() - t_start, 1e-6),
+                3,
             )
             _emit(
-                on_event, "progress", step = done, total_steps = cfg.train_steps,
-                loss = round(step_loss, 5), avg_loss = round(running_loss / done, 5),
-                learning_rate = cfg.learning_rate, samples_per_second = sps,
+                on_event,
+                "progress",
+                step = done,
+                total_steps = cfg.train_steps,
+                loss = round(step_loss, 5),
+                avg_loss = round(running_loss / done, 5),
+                learning_rate = cfg.learning_rate,
+                samples_per_second = sps,
                 peak_memory_gb = peak_gb or None,
             )
         if _check_stop():
@@ -629,9 +664,15 @@ def run_dit_lora_training(
         lora_path = str(out_dir / DEFAULT_LORA_FILENAME)
         catalog_path = _publish_to_lora_catalog(lora_path, cfg)
     _emit(
-        on_event, "complete", output_dir = str(out_dir), lora_path = lora_path,
-        catalog_path = catalog_path, family = cfg.resolved_family, base_model = cfg.base_model,
-        stopped = stopped, steps_run = done if cfg.train_steps else 0,
+        on_event,
+        "complete",
+        output_dir = str(out_dir),
+        lora_path = lora_path,
+        catalog_path = catalog_path,
+        family = cfg.resolved_family,
+        base_model = cfg.base_model,
+        stopped = stopped,
+        steps_run = done if cfg.train_steps else 0,
     )
     return str(out_dir)
 
@@ -640,10 +681,8 @@ def _make_optimizer(params, lr):
     """8-bit AdamW (bitsandbytes) when available -- half the optimizer state, no accuracy
     regression for LoRA -- else the torch AdamW fallback."""
     import torch
-
     try:
         import bitsandbytes as bnb
-
         return bnb.optim.AdamW8bit(params, lr = lr)
     except Exception:  # noqa: BLE001 -- bnb missing / no CUDA: fall back to torch AdamW
         return torch.optim.AdamW(params, lr = lr)
@@ -652,8 +691,14 @@ def _make_optimizer(params, lr):
 def _free_text_encoders(pipe) -> None:
     """Drop every text-encoder / tokenizer the pipeline holds, so the (large) encoders do
     not sit in VRAM during training. The embeddings are already precomputed."""
-    for attr in ("text_encoder", "text_encoder_2", "text_encoder_3", "tokenizer",
-                 "tokenizer_2", "tokenizer_3"):
+    for attr in (
+        "text_encoder",
+        "text_encoder_2",
+        "text_encoder_3",
+        "tokenizer",
+        "tokenizer_2",
+        "tokenizer_3",
+    ):
         if getattr(pipe, attr, None) is not None:
             try:
                 setattr(pipe, attr, None)
diff --git a/studio/backend/core/training/diffusion_lora_trainer.py b/studio/backend/core/training/diffusion_lora_trainer.py
index 59da0b6288..deef4d03ce 100644
--- a/studio/backend/core/training/diffusion_lora_trainer.py
+++ b/studio/backend/core/training/diffusion_lora_trainer.py
@@ -413,7 +413,6 @@ def _make_lora_optimizer(params: list, lr: float) -> Any:
     if os.environ.get("UNSLOTH_DIFFUSION_FP32_OPTIM", "") not in ("1", "true"):
         try:
             import bitsandbytes as bnb
-
             return bnb.optim.AdamW8bit(params, lr = lr)
         except Exception:  # noqa: BLE001 -- bnb missing / no CUDA: fall back to torch AdamW
             pass
diff --git a/studio/backend/core/training/diffusion_train_common.py b/studio/backend/core/training/diffusion_train_common.py
index f89b05ff80..9458c4e5c9 100644
--- a/studio/backend/core/training/diffusion_train_common.py
+++ b/studio/backend/core/training/diffusion_train_common.py
@@ -95,9 +95,7 @@ def resolve_trainable_family(base_model: str, model_family: Optional[str] = None
             known = ", ".join(supported_family_names())
             raise ValueError(f"Unknown model_family {model_family!r}. Known families: {known}.")
         if not fam.trainable:
-            raise ValueError(
-                f"'{fam.name}' models can't be trained yet. {_trainable_hint()}"
-            )
+            raise ValueError(f"'{fam.name}' models can't be trained yet. {_trainable_hint()}")
         return fam.name
 
     fam = detect_family_for_pick(base_model)
diff --git a/studio/backend/routes/training.py b/studio/backend/routes/training.py
index 2f493196fd..ff026b3814 100644
--- a/studio/backend/routes/training.py
+++ b/studio/backend/routes/training.py
@@ -1143,7 +1143,12 @@ def _preflight_gated_base(base_model: str, hf_token: Optional[str]) -> None:
 
     repo = (base_model or "").strip()
     # Only remote 'org/name' repos are gated; skip local paths and single-file names.
-    if not repo or repo.count("/") != 1 or repo.startswith((".", "/", "~")) or repo.endswith(".gguf"):
+    if (
+        not repo
+        or repo.count("/") != 1
+        or repo.startswith((".", "/", "~"))
+        or repo.endswith(".gguf")
+    ):
         return
     url = f"https://huggingface.co/{repo}/resolve/main/model_index.json"
     headers = {"Authorization": f"Bearer {hf_token}"} if hf_token else {}
@@ -1494,7 +1499,6 @@ def _image_record(
     width = height = 0
     try:
         from PIL import Image
-
         with Image.open(image_path) as im:
             width, height = im.size
     except Exception:  # noqa: BLE001 -- an unreadable image still lists (0x0) rather than 500
@@ -1509,9 +1513,7 @@ def _image_record(
     )
 
 
-@router.get(
-    "/diffusion/dataset/{name}/images", response_model = DiffusionDatasetImagesResponse
-)
+@router.get("/diffusion/dataset/{name}/images", response_model = DiffusionDatasetImagesResponse)
 async def list_diffusion_dataset_images(
     name: str, current_subject: str = Depends(get_current_subject)
 ):
@@ -1525,9 +1527,7 @@ async def list_diffusion_dataset_images(
         for p in sorted(folder.iterdir()):
             if p.is_file() and p.suffix.lower() in _DIFFUSION_DATASET_IMAGE_EXTS:
                 records.append(_image_record(folder, p, meta))
-        return DiffusionDatasetImagesResponse(
-            name = folder.name, path = str(folder), images = records
-        )
+        return DiffusionDatasetImagesResponse(name = folder.name, path = str(folder), images = records)
 
     return await asyncio.to_thread(scan)
 
@@ -1614,7 +1614,9 @@ async def set_diffusion_dataset_caption(
 
 @router.delete("/diffusion/dataset/{name}/image/{filename}")
 async def delete_diffusion_dataset_image(
-    name: str, filename: str, current_subject: str = Depends(get_current_subject)
+    name: str,
+    filename: str,
+    current_subject: str = Depends(get_current_subject),
 ):
     """Remove an image, its caption sidecars, and any cached thumbnails."""
     folder = _resolve_dataset_folder(name)
@@ -1696,9 +1698,7 @@ def _example_by_id(example_id: str) -> dict:
 
 
 @router.get("/diffusion/dataset-examples", response_model = DiffusionDatasetExamplesResponse)
-async def list_diffusion_dataset_examples(
-    current_subject: str = Depends(get_current_subject),
-):
+async def list_diffusion_dataset_examples(current_subject: str = Depends(get_current_subject)):
     """List the curated example datasets available for one-click import."""
     return DiffusionDatasetExamplesResponse(
         examples = [
@@ -1788,8 +1788,19 @@ def _materialize_imagefolder_jsonl(entry: dict, dest: Path, cap: int) -> int:
         snapshot_download(
             entry["repo"],
             repo_type = "dataset",
-            allow_patterns = ["*.jsonl", "*.jpg", "*.jpeg", "*.png", "*.webp", "*.bmp", "**/*.jpg",
-                              "**/*.jpeg", "**/*.png", "**/*.webp", "**/*.bmp"],
+            allow_patterns = [
+                "*.jsonl",
+                "*.jpg",
+                "*.jpeg",
+                "*.png",
+                "*.webp",
+                "*.bmp",
+                "**/*.jpg",
+                "**/*.jpeg",
+                "**/*.png",
+                "**/*.webp",
+                "**/*.bmp",
+            ],
         )
     )
     # Map basename -> caption from every jsonl carrying file_name + caption column.
@@ -1808,7 +1819,8 @@ def _materialize_imagefolder_jsonl(entry: dict, dest: Path, cap: int) -> int:
                 captions[Path(str(fn)).name] = str(row[caption_col])
     # Copy images (those with a caption first, so a cap keeps captioned pairs).
     images = sorted(
-        p for p in snap.rglob("*")
+        p
+        for p in snap.rglob("*")
         if p.is_file() and p.suffix.lower() in _DIFFUSION_DATASET_IMAGE_EXTS
     )
     images.sort(key = lambda p: (p.name not in captions, p.name))
@@ -1825,12 +1837,9 @@ def _materialize_imagefolder_jsonl(entry: dict, dest: Path, cap: int) -> int:
     return written
 
 
-@router.post(
-    "/diffusion/dataset/import-example", response_model = DiffusionDatasetImportResponse
-)
+@router.post("/diffusion/dataset/import-example", response_model = DiffusionDatasetImportResponse)
 async def import_diffusion_dataset_example(
-    body: DiffusionDatasetImportRequest,
-    current_subject: str = Depends(get_current_subject),
+    body: DiffusionDatasetImportRequest, current_subject: str = Depends(get_current_subject)
 ):
     """Materialize a curated example dataset into a Studio dataset folder (images + .txt
     captions), ready to train. Idempotent: a folder that already holds images is returned
diff --git a/studio/backend/tests/test_diffusion_dataset_api.py b/studio/backend/tests/test_diffusion_dataset_api.py
index f0f59d9a0d..7a07e8eb47 100644
--- a/studio/backend/tests/test_diffusion_dataset_api.py
+++ b/studio/backend/tests/test_diffusion_dataset_api.py
@@ -28,7 +28,11 @@ def _png_bytes(color = (200, 100, 50), size = (8, 8)) -> bytes:
     return buf.getvalue()
 
 
-def _write_png(path, color = (200, 100, 50), size = (8, 8)) -> None:
+def _write_png(
+    path,
+    color = (200, 100, 50),
+    size = (8, 8),
+) -> None:
     Image.new("RGB", size, color).save(path, format = "PNG")
 
 
@@ -104,9 +108,7 @@ def test_get_image_and_thumbnail_excluded_from_listing(client, ds_root):
 
 def test_get_image_missing_404(client, ds_root):
     (ds_root / "pics").mkdir()
-    assert (
-        client.get("/api/train/diffusion/dataset/pics/image/ghost.png").status_code == 404
-    )
+    assert client.get("/api/train/diffusion/dataset/pics/image/ghost.png").status_code == 404
 
 
 # ── caption write / clear ────────────────────────────────────────────────────
@@ -133,9 +135,7 @@ def test_put_caption_roundtrip_and_clear(client, ds_root):
 
 def test_put_caption_missing_image_404(client, ds_root):
     (ds_root / "cap").mkdir()
-    r = client.put(
-        "/api/train/diffusion/dataset/cap/caption/ghost.png", json = {"caption": "hi"}
-    )
+    r = client.put("/api/train/diffusion/dataset/cap/caption/ghost.png", json = {"caption": "hi"})
     assert r.status_code == 404
 
 
@@ -143,9 +143,7 @@ def test_put_caption_too_long_400(client, ds_root):
     folder = ds_root / "cap"
     folder.mkdir()
     _write_png(folder / "x.png")
-    r = client.put(
-        "/api/train/diffusion/dataset/cap/caption/x.png", json = {"caption": "z" * 2001}
-    )
+    r = client.put("/api/train/diffusion/dataset/cap/caption/x.png", json = {"caption": "z" * 2001})
     assert r.status_code == 400
 
 
@@ -188,7 +186,6 @@ def test_image_filename_validation_rejects_traversal():
 
 def test_clean_dataset_name_rejects_dotdot():
     from routes.training import _clean_diffusion_dataset_name
-
     for bad in ("../x", "a/b", "..", " "):
         with pytest.raises(HTTPException) as exc:
             _clean_diffusion_dataset_name(bad)
@@ -273,20 +270,20 @@ def test_import_example_writes_images_and_captions(client, ds_root, monkeypatch)
 
 def test_import_example_respects_cap(client, ds_root, monkeypatch):
     _install_fake_load_dataset(monkeypatch, n_rows = 5)
-    entry = next(e for e in __import__("routes.training", fromlist = ["_DATASET_EXAMPLES"])._DATASET_EXAMPLES if e["id"] == "tuxemon")
-    monkeypatch.setitem(entry, "image_cap", 2)
-    r = client.post(
-        "/api/train/diffusion/dataset/import-example", json = {"id": "tuxemon"}
+    entry = next(
+        e
+        for e in __import__("routes.training", fromlist = ["_DATASET_EXAMPLES"])._DATASET_EXAMPLES
+        if e["id"] == "tuxemon"
     )
+    monkeypatch.setitem(entry, "image_cap", 2)
+    r = client.post("/api/train/diffusion/dataset/import-example", json = {"id": "tuxemon"})
     assert r.status_code == 200, r.text
     assert r.json()["imported"] == 2
     assert r.json()["image_count"] == 2
 
 
 def test_import_example_unknown_id_404(client, ds_root):
-    r = client.post(
-        "/api/train/diffusion/dataset/import-example", json = {"id": "does-not-exist"}
-    )
+    r = client.post("/api/train/diffusion/dataset/import-example", json = {"id": "does-not-exist"})
     assert r.status_code == 404
 
 
@@ -297,8 +294,6 @@ def test_import_example_load_failure_maps_to_502(client, ds_root, monkeypatch):
         raise RuntimeError("network down")
 
     monkeypatch.setattr(datasets, "load_dataset", boom)
-    r = client.post(
-        "/api/train/diffusion/dataset/import-example", json = {"id": "tuxemon"}
-    )
+    r = client.post("/api/train/diffusion/dataset/import-example", json = {"id": "tuxemon"})
     assert r.status_code == 502
     assert "Could not import" in r.json()["detail"]
diff --git a/studio/backend/tests/test_diffusion_lora_trainer.py b/studio/backend/tests/test_diffusion_lora_trainer.py
index 8db45c588e..ff8f3ec89e 100644
--- a/studio/backend/tests/test_diffusion_lora_trainer.py
+++ b/studio/backend/tests/test_diffusion_lora_trainer.py
@@ -236,14 +236,12 @@ def test_config_accepts_sdxl_and_unknown_base_models():
 # ── trainer registry + family resolution + metadata sidecar (PR A platform) ──
 def test_get_trainer_resolves_sdxl():
     from core.training.diffusion_lora_trainer import get_trainer, run_diffusion_lora_training
-
     assert get_trainer("sdxl") is run_diffusion_lora_training
     assert get_trainer("SDXL") is run_diffusion_lora_training  # case-insensitive
 
 
 def test_get_trainer_unknown_family_raises():
     from core.training.diffusion_lora_trainer import get_trainer
-
     with pytest.raises(ValueError, match = "No trainer"):
         get_trainer("flux.2-dev")  # a real family with no registered trainer
 
@@ -251,7 +249,6 @@ def test_get_trainer_unknown_family_raises():
 def test_get_trainer_resolves_dit_families():
     from core.training.diffusion_dit_trainer import run_dit_lora_training
     from core.training.diffusion_lora_trainer import get_trainer
-
     for fam in ("flux.1", "qwen-image", "z-image"):
         assert get_trainer(fam) is run_dit_lora_training
 
@@ -261,7 +258,9 @@ def test_normalized_sets_resolved_family():
         base_model = "stabilityai/stable-diffusion-xl-base-1.0", data_dir = "d", output_dir = "o"
     ).normalized()
     assert cfg.resolved_family == "sdxl"
-    cfg2 = DiffusionLoraConfig(base_model = "my-custom-thing", data_dir = "d", output_dir = "o").normalized()
+    cfg2 = DiffusionLoraConfig(
+        base_model = "my-custom-thing", data_dir = "d", output_dir = "o"
+    ).normalized()
     assert cfg2.resolved_family == "sdxl"  # unknown -> default SDXL trainer
 
 
diff --git a/studio/backend/tests/test_diffusion_training.py b/studio/backend/tests/test_diffusion_training.py
index 6931418cde..886a3e5d4f 100644
--- a/studio/backend/tests/test_diffusion_training.py
+++ b/studio/backend/tests/test_diffusion_training.py
@@ -484,8 +484,15 @@ def test_route_start_refuses_non_sdxl_base_without_freeing_gpu(client, monkeypat
 def test_apply_event_records_metric_history_and_perf():
     svc = DiffusionTrainingService(ctx = _FakeCtx(), target = _happy_target)
     svc._apply_event(
-        {"type": "progress", "step": 1, "total_steps": 10, "loss": 0.5,
-         "learning_rate": 1e-4, "samples_per_second": 3.2, "peak_memory_gb": 7.1}
+        {
+            "type": "progress",
+            "step": 1,
+            "total_steps": 10,
+            "loss": 0.5,
+            "learning_rate": 1e-4,
+            "samples_per_second": 3.2,
+            "peak_memory_gb": 7.1,
+        }
     )
     svc._apply_event(
         {"type": "progress", "step": 2, "total_steps": 10, "loss": 0.4, "learning_rate": 9e-5}
@@ -531,8 +538,14 @@ def test_metric_history_decimates_at_cap():
 def test_complete_event_records_family_and_catalog():
     svc = DiffusionTrainingService(ctx = _FakeCtx(), target = _happy_target)
     svc._apply_event(
-        {"type": "complete", "output_dir": "/o", "lora_path": "/o/w.safetensors",
-         "catalog_path": "/loras/w.safetensors", "family": "sdxl", "base_model": "b"}
+        {
+            "type": "complete",
+            "output_dir": "/o",
+            "lora_path": "/o/w.safetensors",
+            "catalog_path": "/loras/w.safetensors",
+            "family": "sdxl",
+            "base_model": "b",
+        }
     )
     st = svc.status()
     assert st["status"] == "completed"
@@ -544,8 +557,12 @@ def test_complete_event_records_family_and_catalog():
 def test_status_route_nests_metric_history(client):
     # The status route folds the service's flat arrays into a nested metric_history object.
     client._fake.status_extra = {
-        "metric_steps": [1, 2], "metric_loss": [0.5, 0.4], "metric_lr": [1e-4, 9e-5],
-        "family": "sdxl", "samples_per_second": 2.0, "peak_memory_gb": 6.0,
+        "metric_steps": [1, 2],
+        "metric_loss": [0.5, 0.4],
+        "metric_lr": [1e-4, 9e-5],
+        "family": "sdxl",
+        "samples_per_second": 2.0,
+        "peak_memory_gb": 6.0,
     }
     r = client.get("/api/train/diffusion/status")
     assert r.status_code == 200, r.text