Merge remote-tracking branch 'origin/diffusion-lora-training' into diffusion-lora-training-api

studio/backend/core/training/diffusion_lora_trainer.py

@@ -28,6 +28,7 @@ import json
 import math
 import os
 import random
+import re
 import time
 from dataclasses import dataclass, field, replace
 from pathlib import Path
@@ -42,6 +43,18 @@ _CAPTION_EXTS = (".txt", ".caption")
 # diffusers' canonical single-file LoRA name, so load_lora_weights(dir) finds it.
 DEFAULT_LORA_FILENAME = "pytorch_lora_weights.safetensors"
 
+# Families Studio can LOAD but not train (DiT architectures). A base-model name that
+# clearly belongs to one is refused in normalized(), so a wrong pick fails at start
+# (an instant HTTP 400 through the API) instead of minutes later inside
+# StableDiffusionXLPipeline.from_pretrained. Single tokens match on word boundaries;
+# hyphenated markers match as substrings of the hyphen-condensed name.
+_NON_SDXL_TOKENS = frozenset({"flux", "sd3", "kontext", "pixart", "sana", "lumina", "cogview"})
+_NON_SDXL_PHRASES = ("qwen-image", "z-image", "stable-diffusion-3", "hunyuan-dit")
+_ONLY_SDXL_HINT = (
+    "Only SDXL bases can be trained right now (e.g. stabilityai/stable-diffusion-xl-base-1.0 "
+    "or stabilityai/sdxl-turbo). Other families can load LoRAs but not train them yet."
+)
+
 EventCb = Callable[[dict[str, Any]], None]
 # Returns a falsy value to keep training, or a truthy stop signal: bare True, or a dict
 # that may carry ``save=False`` to cancel without saving a partial adapter.
@@ -89,6 +102,7 @@ class DiffusionLoraConfig:
 
         Also coerces values that arrive as strings/blanks through the Studio config path
         (``learning_rate`` is preserved as a string there; ``hf_token`` defaults to "")."""
+        assert_trainable_base_model(self.base_model)
         if self.train_steps < 1:
             raise ValueError("train_steps must be >= 1")
         if self.train_batch_size < 1:
@@ -276,6 +290,32 @@ def _encode_sdxl_prompts(
     return prompt_embeds, pooled
 
 
+def assert_trainable_base_model(base_model: str) -> None:
+    """Refuse base models that are recognisably not SDXL, before anything is downloaded.
+
+    Purely name-based: a GGUF filename or a known DiT-family name (FLUX / Qwen-Image /
+    Z-Image / SD3 / ...) can never train on the SDXL U-Net trainer, so failing here turns
+    a confusing mid-run crash into an immediate, actionable error. Names this cannot
+    classify pass through; from_pretrained still fails cleanly on a genuinely wrong pick."""
+    name = str(base_model or "").strip().lower()
+    if name.endswith(".gguf"):
+        raise ValueError(
+            f"'{base_model}' is a GGUF checkpoint, which can't be trained. {_ONLY_SDXL_HINT}"
+        )
+    condensed = re.sub(r"[^a-z0-9]+", "-", name)
+    hit = next(
+        (p for p in _NON_SDXL_PHRASES if p in condensed),
+        None,
+    ) or next(
+        (t for t in condensed.split("-") if t in _NON_SDXL_TOKENS),
+        None,
+    )
+    if hit:
+        raise ValueError(
+            f"'{base_model}' looks like a {hit} model, which isn't trainable. {_ONLY_SDXL_HINT}"
+        )
+
+
 def _assert_trusted_base_model(base_model: str) -> None:
     """Gate the training base model the same way the inference backend gates non-GGUF loads:
     a local path or a trusted repo (``unsloth/*`` or an allowlisted official base). This runs

studio/backend/tests/test_diffusion_lora_trainer.py

@@ -193,3 +193,34 @@ def test_config_rejects_nonpositive_learning_rate():
         DiffusionLoraConfig(
             base_model = "b", data_dir = "d", output_dir = "o", learning_rate = 0
         ).normalized()
+
+
+def test_config_rejects_known_non_sdxl_base_models():
+    # Known DiT families and GGUF checkpoints must fail at normalise time (an instant
+    # 400 via the API) instead of minutes later inside from_pretrained.
+    for bad in (
+        "unsloth/FLUX.1-dev-GGUF",
+        "black-forest-labs/FLUX.1-schnell",
+        "unsloth/Qwen-Image-2512-unsloth-bnb-4bit",
+        "Tongyi-MAI/Z-Image-Turbo",
+        "stabilityai/stable-diffusion-3-medium",
+        "unsloth/FLUX.1-Kontext-dev",
+        "z-image-turbo-Q4_K_M.gguf",
+    ):
+        with pytest.raises(ValueError, match = "SDXL"):
+            DiffusionLoraConfig(
+                base_model = bad, data_dir = "d", output_dir = "o"
+            ).normalized()
+
+
+def test_config_accepts_sdxl_and_unknown_base_models():
+    # SDXL names and unclassifiable custom names/paths must pass the guard (a wrong
+    # custom pick still fails cleanly in from_pretrained).
+    for ok in (
+        "stabilityai/stable-diffusion-xl-base-1.0",
+        "stabilityai/sdxl-turbo",
+        "/data/checkpoints/my-custom-sdxl",
+        "my-finetune",
+    ):
+        cfg = DiffusionLoraConfig(base_model = ok, data_dir = "d", output_dir = "o").normalized()
+        assert cfg.base_model == ok