From af638f98c3d037077d31903aca95cdbb49b5e025 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Thu, 2 Jul 2026 01:06:42 +0000 Subject: [PATCH 01/11] Diffusion LoRA training: harden config handling, cancellation, SDXL conditioning, and safety Addresses review findings on the SDXL LoRA trainer: - Gate the base model with the same trust check as inference (unsloth/*, allowlisted official bases, or a local path) before from_pretrained, so an untrusted remote repo is never fetched or deserialised. - Check the stop signal before the (slow) model load, not only between steps, so a cancel during download is honoured; a stop may carry save=False to cancel without leaving a partial adapter. - Per-sample SDXL add_time_ids from the actual crop (original size + crop offset, with the offset mirrored on horizontal flip) instead of a fixed uncropped-square tensor. - Apply EXIF orientation before resize/crop so rotated photos train upright. - Skip gradient clipping when max_grad_norm <= 0 (the Studio 'disable' value) instead of scaling every gradient to zero. - Coerce Studio config strings/blanks: learning_rate string to float, blank hf_token to anonymous, gradient_checkpointing 'none'/'true'/'unsloth' to bool; reject a zero/negative lora_alpha or learning_rate. - Alias the generic Studio training payload keys (model_name/max_steps/batch_size/lora_r/ lr_scheduler_type/random_seed) onto the diffusion field names. - Mirror the trained adapter into loras/diffusion so the Images LoRA picker discovers it. - Report worker exceptions in both message and error keys so the failure is not lost. Adds regression tests for the config coercion/validation and aliasing. --- .../core/training/diffusion_lora_trainer.py | 259 ++++++++++++++---- .../tests/test_diffusion_lora_trainer.py | 78 ++++++ 2 files changed, 287 insertions(+), 50 deletions(-) diff --git a/studio/backend/core/training/diffusion_lora_trainer.py b/studio/backend/core/training/diffusion_lora_trainer.py index ee095e63d2..fba01d0b14 100644 --- a/studio/backend/core/training/diffusion_lora_trainer.py +++ b/studio/backend/core/training/diffusion_lora_trainer.py @@ -43,7 +43,9 @@ _CAPTION_EXTS = (".txt", ".caption") DEFAULT_LORA_FILENAME = "pytorch_lora_weights.safetensors" EventCb = Callable[[dict[str, Any]], None] -StopCb = Callable[[], bool] +# 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. +StopCb = Callable[[], Any] @dataclass @@ -83,7 +85,10 @@ class DiffusionLoraConfig: def normalized(self) -> "DiffusionLoraConfig": """Return a copy with derived/validated fields filled in. Raises ValueError on a - request that cannot train (bad numbers, or no caption source).""" + request that cannot train (bad numbers, or no caption source). + + 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 "").""" if self.train_steps < 1: raise ValueError("train_steps must be >= 1") if self.train_batch_size < 1: @@ -92,13 +97,35 @@ class DiffusionLoraConfig: raise ValueError("gradient_accumulation_steps must be >= 1") if self.lora_rank < 1: raise ValueError("lora_rank must be >= 1") + if self.lora_alpha is not None and self.lora_alpha < 1: + raise ValueError( + "lora_alpha must be >= 1 (a zero/negative alpha scales the adapter to nothing)" + ) if self.resolution < 64 or self.resolution % 8 != 0: raise ValueError("resolution must be a multiple of 8 and >= 64") if self.mixed_precision not in ("bf16", "fp16", "no"): raise ValueError("mixed_precision must be one of bf16 / fp16 / no") + # learning_rate can arrive as a string ("1e-4") from the Studio config path, which + # preserves it as a string after validation; coerce so AdamW receives a float. + try: + learning_rate = float(self.learning_rate) + except (TypeError, ValueError) as exc: + raise ValueError(f"learning_rate must be a number, got {self.learning_rate!r}") from exc + if learning_rate <= 0: + raise ValueError("learning_rate must be > 0") alpha = self.lora_alpha if self.lora_alpha is not None else self.lora_rank targets = tuple(self.lora_target_modules) or DEFAULT_LORA_TARGETS - return replace(self, lora_alpha = alpha, lora_target_modules = targets) + # A blank Hub token (the Studio default when none is configured) must load + # anonymously, not as an explicit empty credential. + token = self.hf_token.strip() if isinstance(self.hf_token, str) else self.hf_token + return replace( + self, + learning_rate = learning_rate, + lora_alpha = alpha, + lora_target_modules = targets, + max_grad_norm = float(self.max_grad_norm), + hf_token = token or None, + ) def discover_image_caption_pairs( @@ -172,7 +199,8 @@ def discover_image_caption_pairs( def compute_sdxl_add_time_ids(resolution: int) -> tuple[int, int, int, int, int, int]: """SDXL micro-conditioning ``add_time_ids`` for a square ``resolution`` train crop: (original_h, original_w, crop_top, crop_left, target_h, target_w). Pure; the trainer - turns it into a tensor. No crop offset is applied (top-left = 0).""" + turns it into a tensor. No crop offset is applied (top-left = 0). The training loop + derives per-image time-ids from the actual crop instead; this is the square default.""" return (resolution, resolution, 0, 0, resolution, resolution) @@ -183,30 +211,42 @@ def _emit(on_event: Optional[EventCb], type_: str, **kw: Any) -> None: def _load_image_tensor( path: str, resolution: int, center_crop: bool, random_flip: bool, rng: random.Random -) -> Any: +) -> tuple[Any, tuple[int, int, int, int, int, int]]: """Load an image to a normalised CxHxW tensor in [-1, 1] (resize shorter side to - ``resolution``, crop to a square, optional horizontal flip). No torchvision.""" + ``resolution``, crop to a square, optional horizontal flip). No torchvision. + + Returns ``(tensor, add_time_ids)`` where add_time_ids is the SDXL micro-conditioning + (original_h, original_w, crop_top, crop_left, target_h, target_w) for THIS sample, so + the U-Net is told the real original size and crop offset (not a fixed uncropped + square). EXIF orientation is applied first so rotated phone photos train upright.""" import numpy as np import torch - from PIL import Image + from PIL import Image, ImageOps - img = Image.open(path).convert("RGB") - w, h = img.size - scale = resolution / min(w, h) - img = img.resize( - (max(resolution, round(w * scale)), max(resolution, round(h * scale))), Image.LANCZOS - ) - w, h = img.size + # Honour EXIF orientation before any geometry, or rotated camera/phone photos would + # train in their stored (sideways) orientation, mismatched to their captions. + img = ImageOps.exif_transpose(Image.open(path)).convert("RGB") + original_w, original_h = img.size + scale = resolution / min(original_w, original_h) + resized_w = max(resolution, round(original_w * scale)) + resized_h = max(resolution, round(original_h * scale)) + img = img.resize((resized_w, resized_h), Image.LANCZOS) if center_crop: - left, top = (w - resolution) // 2, (h - resolution) // 2 + left, top = (resized_w - resolution) // 2, (resized_h - resolution) // 2 else: - left = rng.randint(0, max(0, w - resolution)) - top = rng.randint(0, max(0, h - resolution)) + left = rng.randint(0, max(0, resized_w - resolution)) + top = rng.randint(0, max(0, resized_h - resolution)) img = img.crop((left, top, left + resolution, top + resolution)) + crop_left = left if random_flip and rng.random() < 0.5: img = img.transpose(Image.FLIP_LEFT_RIGHT) + # A horizontal flip mirrors the crop's left origin, so report the mirrored offset + # (diffusers' SDXL training scripts do the same) to keep the conditioning honest. + crop_left = max(0, resized_w - resolution - left) arr = np.asarray(img, dtype = np.float32) / 255.0 - return torch.from_numpy(arr).permute(2, 0, 1) * 2.0 - 1.0 + tensor = torch.from_numpy(arr).permute(2, 0, 1) * 2.0 - 1.0 + time_ids = (original_h, original_w, top, crop_left, resolution, resolution) + return tensor, time_ids def _encode_sdxl_prompts( @@ -236,6 +276,20 @@ def _encode_sdxl_prompts( return prompt_embeds, pooled +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 + BEFORE ``from_pretrained`` so an untrusted remote repo (which could ship pickle weights) + is never fetched or deserialised.""" + from core.inference.diffusion import _is_trusted_diffusion_repo + + if not _is_trusted_diffusion_repo(base_model): + raise ValueError( + f"Refusing to train from untrusted base model '{base_model}'. Use a local path or " + f"a trusted repo (an unsloth/* repo or an official SDXL base)." + ) + + def run_diffusion_lora_training( config: DiffusionLoraConfig, *, @@ -246,7 +300,8 @@ def run_diffusion_lora_training( Emits ``model_load_started`` / ``model_load_completed`` / ``progress`` (step, loss) / ``complete`` (output_dir, lora_path) events via ``on_event``; ``error`` is emitted by - the process adapter. Honours ``should_stop`` between optimizer steps (partial save).""" + the process adapter. Honours ``should_stop`` (checked before model load and between + optimizer steps); a stop saves a partial adapter unless it carries ``save=False``.""" import torch import torch.nn.functional as F from diffusers import DDPMScheduler, StableDiffusionXLPipeline @@ -260,6 +315,21 @@ def run_diffusion_lora_training( rng = random.Random(cfg.seed) torch.manual_seed(cfg.seed) + # A stop signal may be a bare truthy value or a dict carrying save=False (cancel without + # saving a partial adapter). ``save_on_stop`` records that decision for the export step. + save_on_stop = True + + def _check_stop() -> bool: + nonlocal save_on_stop + if should_stop is None: + return False + sig = should_stop() + if not sig: + return False + if isinstance(sig, dict) and sig.get("save") is False: + save_on_stop = False + return True + device = "cuda" if torch.cuda.is_available() else "cpu" precision = cfg.mixed_precision if device == "cuda" else "no" if precision == "bf16" and device == "cuda" and not torch.cuda.is_bf16_supported(): @@ -268,11 +338,24 @@ def run_diffusion_lora_training( precision = "fp16" weight_dtype = {"bf16": torch.bfloat16, "fp16": torch.float16, "no": torch.float32}[precision] + # Preflight the base model against the same trust gate as inference, before any fetch. + _assert_trusted_base_model(cfg.base_model) + pairs = discover_image_caption_pairs( cfg.data_dir, instance_prompt = cfg.instance_prompt, caption_column = cfg.caption_column ) _emit(on_event, "model_load_started", num_images = len(pairs)) + # Honour a stop requested before the (potentially large / slow) base model loads, the + # same way the LLM training worker checks its stop thread around model load. + 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, + ) + return str(out_dir) + pipe = StableDiffusionXLPipeline.from_pretrained( cfg.base_model, torch_dtype = weight_dtype, token = cfg.hf_token, add_watermarker = False ) @@ -312,9 +395,6 @@ def run_diffusion_lora_training( num_training_steps = cfg.train_steps * cfg.gradient_accumulation_steps, ) - add_time_ids = torch.tensor( - [compute_sdxl_add_time_ids(cfg.resolution)], device = device, dtype = weight_dtype - ) vae_scale = vae.config.scaling_factor prediction_type = noise_scheduler.config.prediction_type @@ -334,12 +414,15 @@ def run_diffusion_lora_training( step_loss = 0.0 for _ in range(cfg.gradient_accumulation_steps): img_paths, captions = _next_batch() - pixel_values = torch.stack( - [ - _load_image_tensor(p, cfg.resolution, cfg.center_crop, cfg.random_flip, rng) - for p in img_paths - ] - ).to(device, dtype = torch.float32) + loaded = [ + _load_image_tensor(p, cfg.resolution, cfg.center_crop, cfg.random_flip, rng) + for p in img_paths + ] + pixel_values = torch.stack([t for t, _ in loaded]).to(device, dtype = torch.float32) + # Per-sample SDXL micro-conditioning from the actual crop (original size + offset). + batch_time_ids = torch.tensor( + [tid for _, tid in loaded], device = device, dtype = weight_dtype + ) with torch.no_grad(): latents = vae.encode(pixel_values).latent_dist.sample() * vae_scale @@ -357,7 +440,7 @@ def run_diffusion_lora_training( ) prompt_embeds = prompt_embeds.to(dtype = weight_dtype) pooled = pooled.to(dtype = weight_dtype) - added = {"text_embeds": pooled, "time_ids": add_time_ids.repeat(bsz, 1)} + added = {"text_embeds": pooled, "time_ids": batch_time_ids} model_pred = unet( noisy, timesteps, prompt_embeds, added_cond_kwargs = added, return_dict = False @@ -384,7 +467,10 @@ def run_diffusion_lora_training( step_loss += float(loss.detach()) / cfg.gradient_accumulation_steps micro += 1 - torch.nn.utils.clip_grad_norm_(lora_params, cfg.max_grad_norm) + # max_grad_norm <= 0 means "disable clipping" (the Studio payload sends 0.0 for that); + # passing 0.0 to clip_grad_norm_ would scale every gradient to zero (no learning). + if cfg.max_grad_norm and cfg.max_grad_norm > 0: + torch.nn.utils.clip_grad_norm_(lora_params, cfg.max_grad_norm) optimizer.step() lr_sched.step() @@ -404,32 +490,62 @@ def run_diffusion_lora_training( learning_rate = lr_sched.get_last_lr()[0], ) - if should_stop is not None and should_stop(): + if _check_stop(): stopped = True break - # Export the trained LoRA in diffusers format (loadable via load_lora_weights). + # Export the trained LoRA in diffusers format (loadable via load_lora_weights), unless + # the run was cancelled with save disabled -- then leave no partial adapter behind. out_dir = Path(cfg.output_dir).expanduser() - out_dir.mkdir(parents = True, exist_ok = True) - unet_lora = convert_state_dict_to_diffusers(get_peft_model_state_dict(unet)) - StableDiffusionXLPipeline.save_lora_weights( - save_directory = str(out_dir), - unet_lora_layers = unet_lora, - safe_serialization = True, - weight_name = DEFAULT_LORA_FILENAME, - ) - lora_path = str(out_dir / DEFAULT_LORA_FILENAME) + lora_path: Optional[str] = None + catalog_path: Optional[str] = None + if not (stopped and not save_on_stop): + out_dir.mkdir(parents = True, exist_ok = True) + unet_lora = convert_state_dict_to_diffusers(get_peft_model_state_dict(unet)) + StableDiffusionXLPipeline.save_lora_weights( + save_directory = str(out_dir), + unet_lora_layers = unet_lora, + safe_serialization = True, + weight_name = DEFAULT_LORA_FILENAME, + ) + lora_path = str(out_dir / DEFAULT_LORA_FILENAME) + # Mirror into the Studio diffusion LoRA directory so the Images picker discovers it + # (its scan lists only files directly under loras/diffusion, not subdirectories). + 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, stopped = stopped, steps_run = done if cfg.train_steps else 0, ) return str(out_dir) +def _publish_to_lora_catalog(lora_path: str, cfg: DiffusionLoraConfig) -> Optional[str]: + """Best-effort copy of the trained adapter into the Studio diffusion LoRA directory so + the Images LoRA picker (which scans only files directly under ``loras/diffusion``) finds + it without the user moving files. Returns the published path, or None on any failure.""" + try: + import shutil + + from core.inference.diffusion_lora import loras_dir, sanitize_alias + + base = ( + cfg.adapter_name + if cfg.adapter_name and cfg.adapter_name != "default" + else Path(cfg.output_dir).name + ) + dest = loras_dir() / f"{sanitize_alias(base)}.safetensors" + if Path(lora_path).resolve() != dest.resolve(): + shutil.copy2(lora_path, dest) + return str(dest) + except Exception: # noqa: BLE001 -- the catalog mirror is best-effort, never fatal + return None + + def run_diffusion_training_process(*, event_queue: Any, stop_queue: Any, config: dict) -> None: """mp.Queue subprocess adapter: run a training job, translating the ``on_event`` callback to ``event_queue`` and a ``stop_queue`` poll to ``should_stop``. Any @@ -439,29 +555,72 @@ def run_diffusion_training_process(*, event_queue: Any, stop_queue: Any, config: def on_event(ev: dict) -> None: event_queue.put(ev) - def should_stop() -> bool: + def should_stop() -> Any: + # Drain the queue and return the last stop message (bool True, or a dict that may + # carry save=False for cancel-without-save); False when nothing was requested. + got: Any = None + saw = False try: while not stop_queue.empty(): - stop_queue.get_nowait() - return True + got = stop_queue.get_nowait() + saw = True except Exception: # noqa: BLE001 -- an empty/closed queue just means "keep going" pass - return False + return got if saw else False try: cfg = _config_from_dict(config) run_diffusion_lora_training(cfg, on_event = on_event, should_stop = should_stop) except Exception as exc: # noqa: BLE001 -- surfaced to the parent as an error event - event_queue.put({"type": "error", "message": str(exc), "ts": time.time()}) + # Emit both keys: the diffusion service reads ``message``, but the generic Studio + # training worker reads ``error``; carrying both keeps the real failure visible on + # either path instead of surfacing as "Unknown error". + event_queue.put( + {"type": "error", "message": str(exc), "error": str(exc), "ts": time.time()} + ) + + +# Aliases from the generic Studio training payload onto DiffusionLoraConfig fields, so the +# diffusion trainer can also be driven by the shared training request shape (not only its +# own request model whose keys already match). +_CONFIG_ALIASES = { + "model_name": "base_model", + "max_steps": "train_steps", + "batch_size": "train_batch_size", + "lora_r": "lora_rank", + "lr_scheduler_type": "lr_scheduler", + "random_seed": "seed", + "lr": "learning_rate", +} + + +def _coerce_gradient_checkpointing(value: Any) -> bool: + """Studio sends gradient_checkpointing as a string ("none" / "true" / "unsloth"); the + disable words are False, anything else truthy is True. 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, ignoring unknown keys so a richer - request payload (UI form) does not break construction.""" + """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 + training keys are aliased onto the diffusion field names, and string flags are coerced.""" valid = DiffusionLoraConfig.__dataclass_fields__.keys() - kwargs = {k: v for k, v in config.items() if k in valid} - if "lora_target_modules" in kwargs and kwargs["lora_target_modules"]: + kwargs: dict[str, Any] = {} + # Aliases first (lowest priority); a canonical key present in the payload overrides. + for src, dst in _CONFIG_ALIASES.items(): + if src in config and config[src] is not None and dst in valid: + kwargs[dst] = config[src] + for k, v in config.items(): + if k in valid: + kwargs[k] = v + if kwargs.get("lora_target_modules"): kwargs["lora_target_modules"] = tuple(kwargs["lora_target_modules"]) + if "gradient_checkpointing" in kwargs: + kwargs["gradient_checkpointing"] = _coerce_gradient_checkpointing( + kwargs["gradient_checkpointing"] + ) return DiffusionLoraConfig(**kwargs) diff --git a/studio/backend/tests/test_diffusion_lora_trainer.py b/studio/backend/tests/test_diffusion_lora_trainer.py index 0d8f2fa9e7..d35e1b3f1a 100644 --- a/studio/backend/tests/test_diffusion_lora_trainer.py +++ b/studio/backend/tests/test_diffusion_lora_trainer.py @@ -16,6 +16,7 @@ import pytest from core.training.diffusion_lora_trainer import ( DEFAULT_LORA_TARGETS, DiffusionLoraConfig, + _coerce_gradient_checkpointing, _config_from_dict, compute_sdxl_add_time_ids, discover_image_caption_pairs, @@ -117,3 +118,80 @@ def test_config_from_dict_ignores_unknown_and_tuples_targets(): ) assert cfg.lora_target_modules == ("to_q", "to_v") assert not hasattr(cfg, "unknown_field") + + +def test_config_rejects_zero_lora_alpha(): + # An explicit zero alpha would scale the adapter to nothing; reject it. + with pytest.raises(ValueError, match = "lora_alpha"): + DiffusionLoraConfig( + base_model = "b", data_dir = "d", output_dir = "o", lora_alpha = 0 + ).normalized() + + +def test_config_coerces_string_learning_rate(): + # The Studio config path preserves learning_rate as a string; normalize to float. + cfg = DiffusionLoraConfig( + base_model = "b", data_dir = "d", output_dir = "o", learning_rate = "1e-4" + ).normalized() + assert cfg.learning_rate == 1e-4 + with pytest.raises(ValueError, match = "learning_rate"): + DiffusionLoraConfig( + base_model = "b", data_dir = "d", output_dir = "o", learning_rate = "abc" + ).normalized() + + +def test_config_blank_hf_token_is_anonymous(): + cfg = DiffusionLoraConfig( + base_model = "b", data_dir = "d", output_dir = "o", hf_token = " " + ).normalized() + assert cfg.hf_token is None + + +def test_config_from_dict_aliases_generic_studio_keys(): + # The generic Studio training payload uses different key names; alias them. + cfg = _config_from_dict( + { + "model_name": "b", + "data_dir": "d", + "output_dir": "o", + "max_steps": 25, + "batch_size": 3, + "lora_r": 8, + "lr_scheduler_type": "cosine", + "random_seed": 7, + } + ) + assert cfg.base_model == "b" + assert cfg.train_steps == 25 + assert cfg.train_batch_size == 3 + assert cfg.lora_rank == 8 + assert cfg.lr_scheduler == "cosine" + assert cfg.seed == 7 + + +def test_config_from_dict_canonical_key_beats_alias(): + cfg = _config_from_dict( + {"base_model": "canon", "model_name": "alias", "data_dir": "d", "output_dir": "o"} + ) + assert cfg.base_model == "canon" + + +def test_gradient_checkpointing_string_coercion(): + # Studio sends a string; the disable words are False, everything else truthy True. + for off in ("none", "None", "false", "0", "no", "off", ""): + assert _coerce_gradient_checkpointing(off) is False + for on in ("true", "unsloth", "yes"): + assert _coerce_gradient_checkpointing(on) is True + assert _coerce_gradient_checkpointing(True) is True + assert _coerce_gradient_checkpointing(False) is False + cfg = _config_from_dict( + {"base_model": "b", "data_dir": "d", "output_dir": "o", "gradient_checkpointing": "none"} + ) + assert cfg.gradient_checkpointing is False + + +def test_config_rejects_nonpositive_learning_rate(): + with pytest.raises(ValueError, match = "learning_rate"): + DiffusionLoraConfig( + base_model = "b", data_dir = "d", output_dir = "o", learning_rate = 0 + ).normalized() From f28be14639e37db890050d479d9a511ce0a40704 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 01:07:36 +0000 Subject: [PATCH 02/11] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- studio/backend/core/training/diffusion_lora_trainer.py | 9 ++++++--- studio/backend/tests/test_diffusion_lora_trainer.py | 4 +--- 2 files changed, 7 insertions(+), 6 deletions(-) diff --git a/studio/backend/core/training/diffusion_lora_trainer.py b/studio/backend/core/training/diffusion_lora_trainer.py index fba01d0b14..70f7c8aedf 100644 --- a/studio/backend/core/training/diffusion_lora_trainer.py +++ b/studio/backend/core/training/diffusion_lora_trainer.py @@ -282,7 +282,6 @@ def _assert_trusted_base_model(base_model: str) -> None: BEFORE ``from_pretrained`` so an untrusted remote repo (which could ship pickle weights) is never fetched or deserialised.""" from core.inference.diffusion import _is_trusted_diffusion_repo - if not _is_trusted_diffusion_repo(base_model): raise ValueError( f"Refusing to train from untrusted base model '{base_model}'. Use a local path or " @@ -351,8 +350,12 @@ def run_diffusion_lora_training( 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, + on_event, + "complete", + output_dir = str(out_dir), + lora_path = None, + stopped = True, + steps_run = 0, ) return str(out_dir) diff --git a/studio/backend/tests/test_diffusion_lora_trainer.py b/studio/backend/tests/test_diffusion_lora_trainer.py index d35e1b3f1a..2f83eee1d3 100644 --- a/studio/backend/tests/test_diffusion_lora_trainer.py +++ b/studio/backend/tests/test_diffusion_lora_trainer.py @@ -123,9 +123,7 @@ def test_config_from_dict_ignores_unknown_and_tuples_targets(): def test_config_rejects_zero_lora_alpha(): # An explicit zero alpha would scale the adapter to nothing; reject it. with pytest.raises(ValueError, match = "lora_alpha"): - DiffusionLoraConfig( - base_model = "b", data_dir = "d", output_dir = "o", lora_alpha = 0 - ).normalized() + DiffusionLoraConfig(base_model = "b", data_dir = "d", output_dir = "o", lora_alpha = 0).normalized() def test_config_coerces_string_learning_rate(): From cdfc6f0f56946efd1c8aa92ea6b7f5cb6b918b28 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Thu, 2 Jul 2026 01:08:02 +0000 Subject: [PATCH 03/11] ControlNet: address review findings on the diffusers path - resolve_controlnet enforces catalog family compatibility so a direct API call cannot load a ControlNet built for another family through the wrong pipeline. - Unknown ControlNet ids now surface as a 400 (call site maps FileNotFoundError to ValueError) instead of a generic 500. - strength 0 disables ControlNet entirely, so a no-op selection never pays the download / VRAM cost; the control image is decoded and validated BEFORE the ControlNet is resolved or built, so a malformed image fails fast for the same reason. - ControlNet loads use the base compute dtype (state.dtype is a display string, not a torch.dtype, so it silently fell back to float32) and honor the base offload policy via group offloading instead of forcing the module resident. - Empty/malformed HF token coerced to anonymous access. - Flux Union ControlNet control_mode mapped from the selected control type. - resolve_controlnet drops the unused hf_token/cancel_event params. - ControlNetSpec validates guidance_start <= guidance_end (clean 422). - Images UI ControlNet Select shows its placeholder when nothing is selected. Adds regression tests for family enforcement and the union control-mode map. --- studio/backend/core/inference/diffusion.py | 142 +++++++++++++----- .../core/inference/diffusion_controlnet.py | 43 +++++- studio/backend/models/inference.py | 8 + .../tests/test_diffusion_controlnet.py | 20 +++ .../src/features/images/images-page.tsx | 2 +- 5 files changed, 173 insertions(+), 42 deletions(-) diff --git a/studio/backend/core/inference/diffusion.py b/studio/backend/core/inference/diffusion.py index fa88706e25..bb9da0d59e 100644 --- a/studio/backend/core/inference/diffusion.py +++ b/studio/backend/core/inference/diffusion.py @@ -1216,11 +1216,30 @@ class DiffusionBackend: if cn_model is None: if cancel.is_set(): raise RuntimeError(DIFFUSION_CANCELLED_MSG) - cn_model = ( - getattr(diffusers, model_cls_name) - .from_pretrained(resolved_cn.path, torch_dtype = state.dtype, token = state.hf_token) - .to(state.device) + import torch + + # state.dtype is the display string saved at load ("bfloat16"), NOT a + # torch.dtype; pass the real dtype so diffusers loads the ControlNet at the + # base compute dtype instead of silently defaulting to float32 (extra VRAM). + cn_dtype = getattr(torch, str(state.dtype).replace("torch.", ""), None) + cn_model = getattr(diffusers, model_cls_name).from_pretrained( + resolved_cn.path, + torch_dtype = cn_dtype, + # An empty / malformed token means anonymous access; the HF client can + # raise on a blank credential instead of falling back, so coerce to None. + token = state.hf_token or None, ) + # Placement must follow the base model's offload policy. A resident base moves + # the ControlNet resident too; an offloaded (low-VRAM) base streams it through + # the device with group offloading instead of forcing the whole module onto the + # GPU, which would defeat the offload and risk an OOM. Best-effort: any failure + # falls back to the resident placement (the prior behaviour). + if getattr(state, "offload_policy", OFFLOAD_NONE) != OFFLOAD_NONE and ( + _offload_controlnet_module(cn_model, state.device, logger) + ): + pass + else: + cn_model = cn_model.to(state.device) if cancel.is_set(): # An unload raced the blocking download above and already cleared the # ControlNet caches; caching now would pin the module past the unload. @@ -1401,7 +1420,7 @@ class DiffusionBackend: pipe = state.pipe init_pil = mask_pil = None control_pil = None - cn_scale = cn_gstart = cn_gend = None + cn_scale = cn_gstart = cn_gend = cn_mode = None ref_extra: list = [] if getattr(state.family, "edit", False): # Instruction editing: the loaded pipe is the edit pipeline. It always @@ -1472,40 +1491,55 @@ class DiffusionBackend: if controlnet is not None: from core.inference import diffusion_controlnet - if workflow != "txt2img": - raise ValueError( - "ControlNet currently combines with plain text-to-image only, not the " - f"{workflow} workflow." - ) - if not diffusion_controlnet.supports_controlnet( - engine = "diffusers", - family = state.family.name, - has_controlnet_pipeline = bool( - getattr(state.family, "controlnet_pipeline_class", None) - ), - model_kind = state.kind, - transformer_quant = state.transformer_quant, - ): - raise ValueError( - "ControlNet is not supported for this model/quantisation on the " - "diffusers engine (needs a bf16 or bnb-4bit load of a family with a " - "ControlNet pipeline; not GGUF-via-diffusers or torchao fp8/int8)." - ) cn_id, cn_image_b64, cn_type, cn_strength, cn_gs, cn_ge = controlnet - resolved_cn = diffusion_controlnet.resolve_controlnet( - cn_id, - family = state.family.name, - hf_token = state.hf_token, - cancel_event = cancel, - ) - pipe = self._controlnet_pipe(state, resolved_cn, cancel) - workflow = "controlnet" - src = _decode_b64_image(cn_image_b64, mode = "RGB") - # Control map at the OUTPUT size so it aligns with the generated latents. - control_pil = diffusion_controlnet.preprocess_control(src, cn_type).resize( - (width, height), Image.LANCZOS - ) - cn_scale, cn_gstart, cn_gend = cn_strength, cn_gs, cn_ge + # strength 0 disables ControlNet (documented on the request model, and the + # frontend slider allows it): skip the whole path so a no-op selection never + # pays the multi-GB ControlNet download / VRAM cost. + if cn_strength in (None, 0, 0.0): + controlnet = None + else: + if workflow != "txt2img": + raise ValueError( + "ControlNet currently combines with plain text-to-image only, not " + f"the {workflow} workflow." + ) + if not diffusion_controlnet.supports_controlnet( + engine = "diffusers", + family = state.family.name, + has_controlnet_pipeline = bool( + getattr(state.family, "controlnet_pipeline_class", None) + ), + model_kind = state.kind, + transformer_quant = state.transformer_quant, + ): + raise ValueError( + "ControlNet is not supported for this model/quantisation on the " + "diffusers engine (needs a bf16 or bnb-4bit load of a family with a " + "ControlNet pipeline; not GGUF-via-diffusers or torchao fp8/int8)." + ) + # Decode + preprocess the control image FIRST so a malformed / unsupported + # image fails as a clean 400 BEFORE any ControlNet download or pipe build, + # rather than after paying that cost. Control map at the OUTPUT size so it + # aligns with the generated latents. + src = _decode_b64_image(cn_image_b64, mode = "RGB") + control_pil = diffusion_controlnet.preprocess_control(src, cn_type).resize( + (width, height), Image.LANCZOS + ) + try: + resolved_cn = diffusion_controlnet.resolve_controlnet( + cn_id, family = state.family.name + ) + except FileNotFoundError as exc: + # An unknown / missing ControlNet id is a bad selection -> 400, not a + # generic 500 (the route maps ValueError, not FileNotFoundError). + raise ValueError(str(exc)) from exc + pipe = self._controlnet_pipe(state, resolved_cn, cancel) + workflow = "controlnet" + cn_scale, cn_gstart, cn_gend = cn_strength, cn_gs, cn_ge + # Flux Union ControlNet selects the active mode by an integer + # ``control_mode`` (canny/depth/pose/...); map the chosen control type so + # the union model applies the right head instead of a default/wrong one. + cn_mode = diffusion_controlnet.union_control_mode(cn_id, cn_type) # Auto-resize odd-sized inputs to a multiple of 16 for the workflows whose # OUTPUT size is taken from the input image (img2img / inpaint / extend / edit), # so an upload like 186px tall no longer fails the pipeline's divisibility check. @@ -1580,6 +1614,10 @@ class DiffusionBackend: kwargs["control_guidance_start"] = cn_gstart if "control_guidance_end" in call_params and cn_gend is not None: kwargs["control_guidance_end"] = cn_gend + # Union ControlNet mode index (Flux); only when the pipe accepts it and the + # selected control type maps to a known mode. + if "control_mode" in call_params and cn_mode is not None: + kwargs["control_mode"] = cn_mode gen = _GenState(total_steps = steps) @@ -1821,6 +1859,34 @@ def _hf_base_model(repo_id: str, hf_token: Optional[str]) -> Optional[str]: return base if isinstance(base, str) and base.strip() else None +def _offload_controlnet_module(cn_model: Any, device: str, logger: Any) -> bool: + """Stream a ControlNet module through ``device`` via diffusers group offloading. + + Used when the base model was loaded with an offload policy: forcing the ControlNet + fully resident with ``.to(device)`` would defeat that low-VRAM placement and can OOM. + Group offloading is applied to this single module (it does not touch the base pipe's + existing hooks), so it is isolated and reversible. Returns True on success; on any + failure the caller falls back to a resident placement, so this never blocks a load.""" + try: + import torch + from diffusers.hooks import apply_group_offloading + + onload = torch.device(device) + apply_group_offloading( + cn_model, + onload_device = onload, + offload_device = torch.device("cpu"), + offload_type = "block_level", + num_blocks_per_group = 1, + use_stream = onload.type == "cuda", + ) + return True + except Exception as exc: # noqa: BLE001 — offload is best-effort; resident is the fallback + if logger is not None: + logger.warning("diffusion.controlnet: group offload failed (%s); loading resident", exc) + return False + + def _base_file_downloaded(rfilename: str) -> bool: """True for base-repo files ``from_pretrained`` actually fetches. diff --git a/studio/backend/core/inference/diffusion_controlnet.py b/studio/backend/core/inference/diffusion_controlnet.py index bf103a9206..3ad0468c1a 100644 --- a/studio/backend/core/inference/diffusion_controlnet.py +++ b/studio/backend/core/inference/diffusion_controlnet.py @@ -19,7 +19,6 @@ backend read an arbitrary location. from __future__ import annotations import re -import threading from dataclasses import dataclass from pathlib import Path from typing import Any, Optional @@ -142,17 +141,28 @@ def resolve_controlnet( spec_id: str, *, family: Optional[str] = None, - hf_token: Optional[str] = None, - cancel_event: Optional[threading.Event] = None, ) -> ResolvedControlNet: """Resolve a ControlNet id to a loadable repo id / local dir. Accepts a catalog/local id, or a bare public HF repo id (``owner/name``). The backend loads the result with ``ControlNetModelClass.from_pretrained(path)`` (download + cache handled there, like the base pipeline). Raises on an unknown id -> the caller maps to 400. + + ``family`` (the loaded base family) enforces catalog compatibility: a ControlNet is + architecture-specific, so a catalog entry tagged for another family is rejected here + with a clear error rather than being loaded through the wrong pipeline class later. """ entry = _catalog_by_id().get(spec_id) if entry is not None: + # A curated/local entry may declare the families it is built for. A client that + # bypasses the UI filter (direct API call) could send an entry for another family; + # reject it before any download so it never reaches the wrong ControlNet pipeline. + fam = (family or "").strip().lower() + if entry.families and fam and fam not in {f.lower() for f in entry.families}: + raise ValueError( + f"ControlNet '{spec_id}' is for {', '.join(entry.families)}, not the loaded " + f"'{family}' model; pick a ControlNet built for this family." + ) if entry.source == "local": path = entry.local_path or "" if not path or not Path(path).is_dir(): @@ -174,6 +184,33 @@ def resolve_controlnet( ) +# Union ControlNet mode indices. A single "union" model covers several control modes and +# selects the active one via an integer ``control_mode`` argument; these are the standard +# indices used by the FLUX.1 / Qwen-Image union ControlNets. "passthrough" (an already-made +# map) carries no intrinsic mode, so it maps to nothing (the caller omits control_mode). +_UNION_CONTROL_MODES: dict[str, int] = { + "canny": 0, + "tile": 1, + "depth": 2, + "blur": 3, + "pose": 4, + "gray": 5, + "lq": 6, +} + + +def union_control_mode(spec_id: str, control_type: str) -> Optional[int]: + """The integer ``control_mode`` for a union ControlNet, or None. + + Returns a mode only for a curated *union* catalog entry AND a control type that maps to a + known index; otherwise None so the caller omits the kwarg (a non-union ControlNet has a + single fixed mode, and 'passthrough' does not name one). Pure lookup, no network.""" + entry = _catalog_by_id().get(spec_id) + if entry is None or not entry.is_union: + return None + return _UNION_CONTROL_MODES.get((control_type or "").strip().lower()) + + def preprocess_control(image: Any, control_type: str) -> Any: """Turn a source image into a control map. diff --git a/studio/backend/models/inference.py b/studio/backend/models/inference.py index 580545b766..979dad1252 100644 --- a/studio/backend/models/inference.py +++ b/studio/backend/models/inference.py @@ -1856,6 +1856,14 @@ class ControlNetSpec(BaseModel): 1.0, ge = 0.0, le = 1.0, description = "Fraction of steps at which ControlNet ends" ) + @model_validator(mode = "after") + def _check_guidance_range(self) -> "ControlNetSpec": + # An inverted range (start > end) means "act over no steps"; reject it as a clean + # 422 instead of letting the diffusers pipeline raise a 500 deep in the denoise. + if self.guidance_start > self.guidance_end: + raise ValueError("guidance_start must be <= guidance_end") + return self + class DiffusionGenerateRequest(BaseModel): """Request to generate one image from the loaded diffusion model.""" diff --git a/studio/backend/tests/test_diffusion_controlnet.py b/studio/backend/tests/test_diffusion_controlnet.py index f9f3d5bd36..e91abad053 100644 --- a/studio/backend/tests/test_diffusion_controlnet.py +++ b/studio/backend/tests/test_diffusion_controlnet.py @@ -45,6 +45,26 @@ def test_resolve_controlnet_rejects_filesystem_like_ids(): dc.resolve_controlnet(bad) +def test_resolve_controlnet_enforces_family_match(): + # A curated entry tagged for another family must be rejected before download so it + # never reaches the wrong ControlNet pipeline class. + with pytest.raises(ValueError, match = "not the"): + dc.resolve_controlnet("qwen-union", family = "flux.1") + # The matching family resolves fine, and no family (unfiltered) is permissive. + assert dc.resolve_controlnet("qwen-union", family = "qwen-image").path + assert dc.resolve_controlnet("qwen-union").path + + +def test_union_control_mode_maps_only_union_entries(): + # Union entries map a known control type to its integer mode; passthrough / unknown + # types and non-union ids return None so the caller omits control_mode. + assert dc.union_control_mode("flux-union-pro", "canny") == 0 + assert dc.union_control_mode("flux-union-pro", "depth") == 2 + assert dc.union_control_mode("flux-union-pro", "pose") == 4 + assert dc.union_control_mode("flux-union-pro", "passthrough") is None + assert dc.union_control_mode("some/bare-repo", "canny") is None + + def test_resolve_controlnet_local(tmp_path, monkeypatch): d = tmp_path / "controlnets" d.mkdir() diff --git a/studio/frontend/src/features/images/images-page.tsx b/studio/frontend/src/features/images/images-page.tsx index f58c1a565e..a51f75795a 100644 --- a/studio/frontend/src/features/images/images-page.tsx +++ b/studio/frontend/src/features/images/images-page.tsx @@ -2095,7 +2095,7 @@ export function ImagesPage({ active = true }: { active?: boolean }) { hint="Condition the image on a control map (edges / depth / pose). Union models cover many types. Use 'Canny' to trace edges from your image, or 'Passthrough' if it is already a control map." >
- From c8984736206711c1272cc9b7d7c26c9312c26594 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 01:08:43 +0000 Subject: [PATCH 04/11] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- studio/backend/core/inference/diffusion.py | 1 - studio/backend/core/inference/diffusion_controlnet.py | 6 +----- 2 files changed, 1 insertion(+), 6 deletions(-) diff --git a/studio/backend/core/inference/diffusion.py b/studio/backend/core/inference/diffusion.py index bb9da0d59e..9fe3a5c136 100644 --- a/studio/backend/core/inference/diffusion.py +++ b/studio/backend/core/inference/diffusion.py @@ -1490,7 +1490,6 @@ class DiffusionBackend: # and passes a control map. v1 conditions txt2img only (not img2img/inpaint/edit). if controlnet is not None: from core.inference import diffusion_controlnet - cn_id, cn_image_b64, cn_type, cn_strength, cn_gs, cn_ge = controlnet # strength 0 disables ControlNet (documented on the request model, and the # frontend slider allows it): skip the whole path so a no-op selection never diff --git a/studio/backend/core/inference/diffusion_controlnet.py b/studio/backend/core/inference/diffusion_controlnet.py index 3ad0468c1a..6b7cdec9fb 100644 --- a/studio/backend/core/inference/diffusion_controlnet.py +++ b/studio/backend/core/inference/diffusion_controlnet.py @@ -137,11 +137,7 @@ def _catalog_by_id() -> dict[str, ControlNetCatalogEntry]: return {e.id: e for e in (list(_CURATED) + _scan_local())} -def resolve_controlnet( - spec_id: str, - *, - family: Optional[str] = None, -) -> ResolvedControlNet: +def resolve_controlnet(spec_id: str, *, family: Optional[str] = None) -> ResolvedControlNet: """Resolve a ControlNet id to a loadable repo id / local dir. Accepts a catalog/local id, or a bare public HF repo id (``owner/name``). The backend From 50a313f93db17047648f1c2eff9c4bdcc1f57032 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Thu, 2 Jul 2026 01:11:07 +0000 Subject: [PATCH 05/11] Diffusion LoRA: harden resolution, native tag precedence, and diffusers teardown Address review findings on the LoRA path: - resolve_one: normalise a blank/whitespace hf_token to None (anonymous access) and reject a client-supplied weight file with traversal / absolute path. - resolve_specs: convert FileNotFoundError from an unknown/stale id to ValueError so the route returns 400 instead of a generic 500. - _scan_local: disambiguate local adapters that share a stem (foo.safetensors vs foo.gguf) so each is uniquely addressable. - inject_prompt_tags: the backend-validated weight now wins over a user-typed for a selected adapter; unselected user tags are left alone. - diffusers _apply_loras: reject a .gguf adapter with a clear error before touching the pipe (diffusers loads safetensors only). - _unload_locked: drop the explicit unload_lora_weights() on teardown; the pipe is dropped wholesale (freeing adapters), so the previous call could race an in-flight denoise on the same pipe. - Images page: use a stable LoRA key and clear the selection (not just the options) when the catalog refresh fails. --- studio/backend/core/inference/diffusion.py | 23 +++-- .../backend/core/inference/diffusion_lora.py | 85 +++++++++++++------ studio/backend/tests/test_diffusion_lora.py | 68 ++++++++++++++- .../src/features/images/images-page.tsx | 9 +- 4 files changed, 149 insertions(+), 36 deletions(-) diff --git a/studio/backend/core/inference/diffusion.py b/studio/backend/core/inference/diffusion.py index a281edea69..79b8a50b11 100644 --- a/studio/backend/core/inference/diffusion.py +++ b/studio/backend/core/inference/diffusion.py @@ -1249,6 +1249,15 @@ class DiffusionBackend: ) resolved = diffusion_lora.resolve_specs(specs, hf_token = state.hf_token, cancel_event = cancel) + # The shared catalog scans both .safetensors and .gguf, but diffusers' + # load_lora_weights only takes safetensors; a .gguf adapter would otherwise fail + # deep in generation. Reject it here as a clean 400 before touching the pipe. + bad = [r.id for r in resolved if r.fmt != "safetensors"] + if bad: + raise ValueError( + "GGUF LoRA adapters are not supported on the diffusers engine " + f"({', '.join(bad)}); use a .safetensors adapter, or the native engine." + ) # Unique adapter names (diffusers requires distinct names; sanitized stems can collide). uniq: list[tuple[str, str, float]] = [] seen: set[str] = set() @@ -1579,13 +1588,13 @@ class DiffusionBackend: if state.eager_patched: uninstall_patches() uninstall_arch_patches() - # Drop any LoRA adapters applied to the pipe so a later reference-path load is - # bit-identical and the freed transformer carries no adapter layers. Idempotent. - try: - if getattr(state.pipe, "_unsloth_loras", ()): - state.pipe.unload_lora_weights() - except Exception: # noqa: BLE001 -- best-effort cleanup on teardown - pass + # NOTE: we deliberately do NOT call state.pipe.unload_lora_weights() here. unload() + # sets the cancel event but does not take _generate_lock, so a LoRA-backed denoise + # can still be running on this same pipe for up to one more callback; mutating its + # adapter layers now would race that in-flight generation. The whole pipe is dropped + # just below (self._state = None; del state; clear_gpu_cache()), so the adapter + # tensors are freed with it -- no explicit unload is needed for memory or for a + # later load (which builds a fresh pipe). # Drop the workflow pipes built around this load's modules so they don't pin the # freed pipeline (they only re-wire its components, but holding the wrappers # would keep the modules alive past unload). diff --git a/studio/backend/core/inference/diffusion_lora.py b/studio/backend/core/inference/diffusion_lora.py index 630857323d..35c52e5abb 100644 --- a/studio/backend/core/inference/diffusion_lora.py +++ b/studio/backend/core/inference/diffusion_lora.py @@ -95,26 +95,31 @@ def sanitize_alias(raw: str) -> str: def _scan_local() -> list[LoraCatalogEntry]: - entries: list[LoraCatalogEntry] = [] root = loras_dir() try: children = sorted(root.iterdir()) except OSError: - return entries - for p in children: - if not p.is_file(): - continue + return [] + files = [p for p in children if p.is_file() and p.suffix.lower() in _ALL_EXTS] + # Two files that share a stem but differ in extension (foo.safetensors + foo.gguf) + # would collide on id (== stem), so the frontend select value and resolve_one's + # id->entry lookup could only ever address one of them. Disambiguate a colliding + # stem by keeping the full filename as the id; a unique stem stays the clean stem. + stem_counts: dict[str, int] = {} + for p in files: + stem_counts[p.stem] = stem_counts.get(p.stem, 0) + 1 + entries: list[LoraCatalogEntry] = [] + for p in files: ext = p.suffix.lower() - if ext not in _ALL_EXTS: - continue try: size = p.stat().st_size except OSError: size = 0 + entry_id = p.name if stem_counts.get(p.stem, 0) > 1 else p.stem entries.append( LoraCatalogEntry( - id = p.stem, - display_name = p.stem, + id = entry_id, + display_name = entry_id, source = "local", fmt = "gguf" if ext == ".gguf" else "safetensors", local_path = str(p), @@ -157,6 +162,9 @@ def resolve_one( shared xet-fallback helper. Raises FileNotFoundError/ValueError on an unresolvable or unsupported id -- the caller maps that to a clear 400. """ + # An empty / whitespace token sent verbatim to HfApi triggers an auth error instead + # of falling back to anonymous access; normalise it to None. + hf_token = hf_token.strip() if hf_token and hf_token.strip() else None entry = _catalog_by_id().get(spec_id) if entry is not None: if entry.source == "local": @@ -176,6 +184,17 @@ def resolve_one( if "/" in spec_id: repo_id, _, weight_name = spec_id.partition(":") weight_name = weight_name or None + if weight_name is not None: + # A client-supplied weight file must stay a plain filename inside the repo: + # reject traversal / absolute paths so it can never resolve outside the HF + # cache dir once handed to the downloader. + if ( + ".." in weight_name + or weight_name.startswith(("/", "\\", "~")) + or "\\" in weight_name + or os.path.isabs(weight_name) + ): + raise ValueError(f"invalid LoRA weight file path '{weight_name}'") if weight_name is None: weight_name = _pick_repo_weight_file(repo_id, hf_token) ext = os.path.splitext(weight_name)[1].lower() @@ -218,12 +237,21 @@ def resolve_specs( hf_token: Optional[str] = None, cancel_event: Optional[threading.Event] = None, ) -> list[ResolvedLora]: - """Resolve request (id, weight) pairs, dropping zero-weight entries.""" + """Resolve request (id, weight) pairs, dropping zero-weight entries. + + A stale / unknown id raises FileNotFoundError inside resolve_one; convert it to + ValueError so the route (which maps only ValueError to a 400) reports bad client + input instead of a generic 500.""" out: list[ResolvedLora] = [] - for spec_id, weight in specs: - if weight == 0: - continue - out.append(resolve_one(spec_id, weight, hf_token = hf_token, cancel_event = cancel_event)) + try: + for spec_id, weight in specs: + if weight == 0: + continue + out.append( + resolve_one(spec_id, weight, hf_token = hf_token, cancel_event = cancel_event) + ) + except FileNotFoundError as exc: + raise ValueError(str(exc)) from exc return out @@ -265,20 +293,29 @@ _TAG_RE = re.compile(r"]+):([^>]+)>") def inject_prompt_tags(prompt: str, resolved: list[ResolvedLora]) -> str: - """Append `` tags to the prompt, skipping any the user already typed. + """Append `` tags for the selected adapters, using the backend- + validated weights. sd-cli strips these tags before they reach the model, so appending them is safe and - deterministic. Duplicate protection: if the prompt already contains a tag for the same - alias, we don't add a second one. + deterministic. A selected adapter's weight is validated (0-2) and recorded in the + request/gallery, so the injected tag must WIN over any `` the user + typed for that same alias: strip a user tag whose alias matches a selected adapter, + then append the validated one. Tags for aliases the user typed that are NOT selected + are left untouched (free-form use). """ - existing = {m.group(1) for m in _TAG_RE.finditer(prompt)} - tags = [ - f"" for r in resolved if r.alias not in existing - ] + selected = {r.alias for r in resolved} + # Drop any user-typed tag whose alias is one of the selected adapters, so the typed + # weight can't override the validated weight (or slip outside the 0-2 bounds). + cleaned = _TAG_RE.sub( + lambda m: "" if m.group(1) in selected else m.group(0), prompt + ) + # Collapse whitespace left by stripped tags without disturbing the user's text. + cleaned = re.sub(r"[ \t]{2,}", " ", cleaned).strip() + tags = [f"" for r in resolved] if not tags: - return prompt - sep = "" if not prompt or prompt.endswith(" ") else " " - return f"{prompt}{sep}{' '.join(tags)}" + return cleaned + sep = "" if not cleaned or cleaned.endswith(" ") else " " + return f"{cleaned}{sep}{' '.join(tags)}" def _fmt_weight(w: float) -> str: diff --git a/studio/backend/tests/test_diffusion_lora.py b/studio/backend/tests/test_diffusion_lora.py index de5aefefcc..e7df3da936 100644 --- a/studio/backend/tests/test_diffusion_lora.py +++ b/studio/backend/tests/test_diffusion_lora.py @@ -37,10 +37,19 @@ def test_inject_prompt_tags_appends_with_spacing(): assert dl.inject_prompt_tags("x", [r1]) == "x " -def test_inject_prompt_tags_dedupes_user_typed_tag(): +def test_inject_prompt_tags_validated_weight_overrides_user_typed(): r = dl.ResolvedLora("id", "style", "/p", "safetensors", 0.8) - # user already wrote a tag for the same alias -> not duplicated - assert dl.inject_prompt_tags("a cat ", [r]) == "a cat " + # A user-typed tag for a SELECTED adapter is replaced by the backend-validated weight + # (so the recorded/validated 0-2 weight wins over whatever was typed), not duplicated. + assert dl.inject_prompt_tags("a cat ", [r]) == "a cat " + + +def test_inject_prompt_tags_keeps_unselected_user_tags(): + r = dl.ResolvedLora("id", "style", "/p", "safetensors", 0.8) + # A user tag for an alias that is NOT one of the selected adapters is left untouched. + out = dl.inject_prompt_tags("a cat ", [r]) + assert "" in out + assert "" in out def test_inject_prompt_tags_empty_returns_prompt(): @@ -129,6 +138,41 @@ def test_resolve_specs_drops_zero_weight(tmp_path, monkeypatch): assert len(out) == 1 and out[0].weight == 1.0 +def test_resolve_specs_maps_unknown_id_to_valueerror(tmp_path, monkeypatch): + # An unknown / stale id raises FileNotFoundError in resolve_one; resolve_specs must + # surface it as ValueError so the route returns 400, not a generic 500. + d = tmp_path / "loras" + d.mkdir() + monkeypatch.setattr(dl, "loras_dir", lambda: d) + with pytest.raises(ValueError): + dl.resolve_specs([("nope", 1.0)]) + + +def test_scan_local_disambiguates_identical_stems(tmp_path, monkeypatch): + # foo.safetensors and foo.gguf must get distinct ids so each is addressable; a + # unique stem keeps its clean stem id. + d = tmp_path / "loras" + d.mkdir() + (d / "foo.safetensors").write_bytes(b"x") + (d / "foo.gguf").write_bytes(b"y") + (d / "solo.safetensors").write_bytes(b"z") + monkeypatch.setattr(dl, "loras_dir", lambda: d) + by_id = {e.id: e for e in dl.list_loras()} + assert "foo.safetensors" in by_id and "foo.gguf" in by_id + assert by_id["foo.safetensors"].fmt == "safetensors" + assert by_id["foo.gguf"].fmt == "gguf" + assert "solo" in by_id # unique stem is untouched + + +def test_resolve_one_rejects_traversal_weight_name(tmp_path, monkeypatch): + # A client-supplied weight file with traversal / absolute path is rejected before it + # can reach the downloader (it must stay a plain filename inside the repo). + monkeypatch.setattr(dl, "loras_dir", lambda: tmp_path) + for bad in ("owner/name:../secret.safetensors", "owner/name:/etc/x.safetensors"): + with pytest.raises(ValueError): + dl.resolve_one(bad, 1.0) + + # ── Request-model validation ──────────────────────────────────────────────── @@ -264,3 +308,21 @@ def test_diffusers_apply_rejects_unsupported_quant(): [("styleA", 1.0)], threading.Event(), ) + + +def test_diffusers_apply_rejects_gguf_adapter(monkeypatch): + # A .gguf adapter (discoverable in the shared catalog) cannot load on the diffusers + # engine; it must be rejected as a clean 400 before touching the pipe. + import threading + + monkeypatch.setattr( + dl, + "resolve_specs", + lambda specs, **_: [ + dl.ResolvedLora(i, dl.sanitize_alias(i), f"/{i}.gguf", "gguf", w) for i, w in specs + ], + ) + pipe = _FakePipe() + with pytest.raises(ValueError, match = "GGUF LoRA"): + _backend()._apply_loras(_fake_state(pipe), [("styleA", 1.0)], threading.Event()) + assert pipe.loaded == [] # never touched the pipe diff --git a/studio/frontend/src/features/images/images-page.tsx b/studio/frontend/src/features/images/images-page.tsx index 594daafded..27fb48c1eb 100644 --- a/studio/frontend/src/features/images/images-page.tsx +++ b/studio/frontend/src/features/images/images-page.tsx @@ -975,7 +975,12 @@ export function ImagesPage({ active = true }: { active?: boolean }) { setLoras((prev) => prev.filter((s) => ids.has(s.id))); }) .catch(() => { - if (!cancelled) setAvailableLoras([]); + if (cancelled) return; + // Clear the SELECTED adapters too, not just the options: leaving a stale `loras` + // selection in state (with the picker now hidden/empty) would still be posted by + // handleGenerate and could apply adapters from the previous model, or fail. + setAvailableLoras([]); + setLoras([]); }); return () => { cancelled = true; @@ -1968,7 +1973,7 @@ export function ImagesPage({ active = true }: { active?: boolean }) {
{loras.map((sel, i) => (
From 1d4461b23b3452a29f02e5c5de2db20fc69de812 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 01:12:12 +0000 Subject: [PATCH 06/11] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- studio/backend/core/inference/diffusion_lora.py | 8 ++------ 1 file changed, 2 insertions(+), 6 deletions(-) diff --git a/studio/backend/core/inference/diffusion_lora.py b/studio/backend/core/inference/diffusion_lora.py index 35c52e5abb..20b40aaf55 100644 --- a/studio/backend/core/inference/diffusion_lora.py +++ b/studio/backend/core/inference/diffusion_lora.py @@ -247,9 +247,7 @@ def resolve_specs( for spec_id, weight in specs: if weight == 0: continue - out.append( - resolve_one(spec_id, weight, hf_token = hf_token, cancel_event = cancel_event) - ) + out.append(resolve_one(spec_id, weight, hf_token = hf_token, cancel_event = cancel_event)) except FileNotFoundError as exc: raise ValueError(str(exc)) from exc return out @@ -306,9 +304,7 @@ def inject_prompt_tags(prompt: str, resolved: list[ResolvedLora]) -> str: selected = {r.alias for r in resolved} # Drop any user-typed tag whose alias is one of the selected adapters, so the typed # weight can't override the validated weight (or slip outside the 0-2 bounds). - cleaned = _TAG_RE.sub( - lambda m: "" if m.group(1) in selected else m.group(0), prompt - ) + cleaned = _TAG_RE.sub(lambda m: "" if m.group(1) in selected else m.group(0), prompt) # Collapse whitespace left by stripped tags without disturbing the user's text. cleaned = re.sub(r"[ \t]{2,}", " ", cleaned).strip() tags = [f"" for r in resolved] From 94d74e8bbe896118960794d74baac457aa5c5241 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Thu, 2 Jul 2026 01:13:04 +0000 Subject: [PATCH 07/11] Diffusion: guard trust check against OSError and validate conditioning inputs - _is_trusted_diffusion_repo: wrap Path.exists() so a repo id with invalid characters (or a bare owner/name id) can't raise OSError; treat any failure as not-a-local-path and fall through to the unsloth/ allowlist. validate_load_request still raises the clear FileNotFoundError for a genuinely missing local pick. - generate(): reject mask_image / upscale / reference_images supplied without an input image, and reject reference_images on a family that does not support reference conditioning, instead of silently degrading to txt2img / img2img. --- studio/backend/core/inference/diffusion.py | 31 +++++++++++++++++-- .../backend/tests/test_diffusion_backend.py | 30 ++++++++++++++++++ 2 files changed, 58 insertions(+), 3 deletions(-) diff --git a/studio/backend/core/inference/diffusion.py b/studio/backend/core/inference/diffusion.py index af8b14df0d..a221021358 100644 --- a/studio/backend/core/inference/diffusion.py +++ b/studio/backend/core/inference/diffusion.py @@ -191,9 +191,18 @@ def _is_trusted_diffusion_repo(repo_id: str) -> bool: on an arbitrary repo, which fetches and deserialises third-party weights. So the non-GGUF paths are gated to the ``unsloth/*`` org (the curated safetensors models) and to local paths the user explicitly pointed at (already on their disk). The GGUF path - is unchanged and stays open to any repo, as before.""" - if Path(repo_id).expanduser().exists(): - return True + is unchanged and stays open to any repo, as before. + + A bare ``owner/name`` HF id is never a real filesystem path, and an id with invalid + characters makes ``Path.exists()`` raise OSError; treat any such failure as "not a + local path" so the trust decision falls through to the unsloth/ check (the loader's + validate_load_request raises the clear FileNotFoundError for a genuinely missing + local pick).""" + try: + if Path(repo_id).expanduser().exists(): + return True + except OSError: + pass return repo_id.strip().lower().startswith("unsloth/") @@ -1273,6 +1282,22 @@ class DiffusionBackend: pipe = state.pipe init_pil = mask_pil = None ref_extra: list = [] + # Validate parameter dependencies up front: mask / upscale / reference all + # need an input image, and reference conditioning needs a family that + # supports it. Without these guards an unsupported combination would be + # silently ignored and quietly fall back to txt2img / img2img. + if init_image is None: + if mask_image is not None: + raise ValueError("mask_image requires an input image (init_image).") + if upscale is not None and upscale > 1.0: + raise ValueError("upscale requires an input image (init_image).") + if reference_images: + raise ValueError("reference_images require an input image (init_image).") + if reference_images and not getattr(state.family, "reference", False): + raise ValueError( + f"Reference images are not supported for the '{state.family.name}' " + "model family." + ) if getattr(state.family, "edit", False): # Instruction editing: the loaded pipe is the edit pipeline. It always # needs an input image; the prompt is the edit instruction. No mask, no diff --git a/studio/backend/tests/test_diffusion_backend.py b/studio/backend/tests/test_diffusion_backend.py index 6e88968316..947d9e7ae3 100644 --- a/studio/backend/tests/test_diffusion_backend.py +++ b/studio/backend/tests/test_diffusion_backend.py @@ -468,6 +468,36 @@ def test_generate_img2img_unsupported_family_raises(fake_runtime, tmp_path, monk backend.generate(prompt = "x", steps = 4, init_image = _tiny_png_b64()) +def test_generate_rejects_conditioning_without_init_image(fake_runtime, tmp_path): + """mask / upscale / reference all need an input image; without one they must raise a + clear ValueError rather than silently degrading to txt2img.""" + (tmp_path / "model.gguf").write_bytes(b"x") + backend = DiffusionBackend() + backend.load_pipeline( + str(tmp_path), gguf_filename = "model.gguf", base_repo = "base/repo", family_override = "z-image" + ) + with pytest.raises(ValueError, match = "mask_image requires"): + backend.generate(prompt = "x", steps = 4, mask_image = _mask_b64(64)) + with pytest.raises(ValueError, match = "upscale requires"): + backend.generate(prompt = "x", steps = 4, upscale = 2.0) + with pytest.raises(ValueError, match = "reference_images require"): + backend.generate(prompt = "x", steps = 4, reference_images = [_tiny_png_b64()]) + + +def test_generate_rejects_reference_on_unsupported_family(fake_runtime, tmp_path): + """A non-reference family rejects reference_images instead of silently dropping them.""" + (tmp_path / "model.gguf").write_bytes(b"x") + backend = DiffusionBackend() + backend.load_pipeline( + str(tmp_path), gguf_filename = "model.gguf", base_repo = "base/repo", family_override = "z-image" + ) + with pytest.raises(ValueError, match = "Reference images are not supported"): + backend.generate( + prompt = "x", steps = 4, init_image = _tiny_png_b64(), + reference_images = [_tiny_png_b64()], + ) + + def test_generate_upscale_enlarges_and_low_strength(fake_runtime, tmp_path): """An init_image + upscale factor routes generate() through the family's img2img pipeline (hires fix): the source is enlarged to size*factor (rounded to /16) before the From a4fb348f9fcf1eec92ba2db1ee4fe7f6d68b18f1 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Thu, 2 Jul 2026 01:13:18 +0000 Subject: [PATCH 08/11] SDXL: reject GGUF up front, skip unused base weights, drop refiner, and harden helpers Addresses review findings on the SDXL family: - Reject a GGUF load for single_file_is_pipeline families (SDXL) in validate_load_request, before the route evicts the current model; SDXL has no transformer-only GGUF variant. - Skip base-repo weight files when a whole-pipeline single file is loaded: from_single_file (config=base) needs only the base config/tokenizer/scheduler, so a local .safetensors no longer triggers a multi-GB base download. - Remove the SDXL refiner from the non-GGUF trust allowlist: it is an img2img-only pipeline but this backend loads every sdxl repo as the base txt2img pipeline. - Normalize a blank/whitespace hf_token to None once in load_pipeline so every load branch degrades to anonymous instead of erroring on a malformed token. - Read the denoiser dtype from a parameter (compile-wrapped modules may lack .dtype) and access state.family.denoiser_attr directly. Adds/updates regression tests for the trust allowlist, GGUF rejection, and base-config filter. --- studio/backend/core/inference/diffusion.py | 74 +++++++++++++++++++-- studio/backend/tests/test_diffusion_sdxl.py | 50 ++++++++++++-- 2 files changed, 113 insertions(+), 11 deletions(-) diff --git a/studio/backend/core/inference/diffusion.py b/studio/backend/core/inference/diffusion.py index addf528819..d89f2d9220 100644 --- a/studio/backend/core/inference/diffusion.py +++ b/studio/backend/core/inference/diffusion.py @@ -191,10 +191,13 @@ def _snap_to_multiple(img: Any, multiple: int = 16) -> Any: # no unsloth-hosted GGUF, so without this its curated catalog entry could not load. # Exact-match, lowercased, so it cannot be widened by a typo-squat. Extend # deliberately, and never add a repo that carries pickled weights or remote code. +# The SDXL refiner is intentionally NOT here: it is an img2img-only refiner pipeline +# (StableDiffusionXLImg2ImgPipeline), but this backend loads every ``sdxl`` repo as the +# base txt2img StableDiffusionXLPipeline and advertises txt2img, so allowlisting the +# refiner would surface the wrong workflow and call it without its required input image. _TRUSTED_NON_GGUF_REPOS = frozenset( { "stabilityai/stable-diffusion-xl-base-1.0", - "stabilityai/stable-diffusion-xl-refiner-1.0", "stabilityai/sdxl-turbo", } ) @@ -434,6 +437,16 @@ class DiffusionBackend: f"pass family_override with that family name. (Video models and image models " f"whose diffusers transformer has no single-file loader are not supported.)" ) + # A GGUF load builds a transformer-only file via the generic GGUF branch + # (UNet2DConditionModel.from_single_file(subfolder="transformer", GGUFQuantizationConfig)). + # Families whose single file IS the whole pipeline (SDXL) have no transformer-only + # GGUF path, so reject GGUF here -- before the route evicts the current model and + # the background load fails deep in from_single_file. + if kind == "gguf" and fam.single_file_is_pipeline: + raise ValueError( + f"'{fam.name}' checkpoints are whole-pipeline single files and have no GGUF " + f"transformer variant; load the .safetensors pipeline instead of a GGUF." + ) # Non-GGUF loads (a single-file safetensors transformer, or a full pipeline) # are gated to the unsloth org or a local path -- they fetch + deserialise # weights, so an arbitrary remote repo is rejected here, before any work. @@ -568,6 +581,7 @@ class DiffusionBackend: base, kwargs.get("hf_token"), kind = kind, + single_file_is_pipeline = bool(fam and fam.single_file_is_pipeline), ) with self._lock: # Stamp progress only if this load is still current; a superseding @@ -634,6 +648,7 @@ class DiffusionBackend: hf_token: Optional[str], *, kind: str = "gguf", + single_file_is_pipeline: bool = False, ) -> tuple[int, list[str]]: """Total download size for the progress bar, plus the base-repo files to fetch (the prefetch reuses this list, so the base is listed only once). @@ -641,7 +656,9 @@ class DiffusionBackend: For a ``pipeline`` load the whole repo IS the pipeline (``base_repo`` is the repo itself), so the transformer/ subfolder is INCLUDED -- unlike the GGUF / single-file paths, where the transformer is the single file and the base repo - supplies only the companions.""" + supplies only the companions. For a ``single_file_is_pipeline`` family (SDXL) the + single file is the WHOLE pipeline, so the base repo supplies only config/tokenizer + (no weights) and its weight files are skipped.""" from huggingface_hub import HfApi api = HfApi() @@ -671,9 +688,16 @@ class DiffusionBackend: if gguf_filename and not Path(repo_id).expanduser().exists(): info = api.model_info(repo_id, files_metadata = True, token = hf_token) total += sum(s.size or 0 for s in info.siblings if s.rfilename == gguf_filename) + # A whole-pipeline single file (SDXL) needs only the base repo's config/tokenizer, + # not its (unused, multi-GB) weight files. + base_filter = ( + _base_config_file_downloaded + if (kind == "single_file" and single_file_is_pipeline) + else _base_file_downloaded + ) base_info = api.model_info(base_repo, files_metadata = True, token = hf_token) for s in base_info.siblings: - if _base_file_downloaded(s.rfilename): + if base_filter(s.rfilename): base_files.append(s.rfilename) total += s.size or 0 except Exception as exc: # noqa: BLE001 — estimate is best-effort @@ -719,6 +743,13 @@ class DiffusionBackend: model_kind: Optional[str] = None, _load_token: Optional[int] = None, ) -> dict[str, Any]: + # A blank / whitespace-only token must degrade to anonymous access, not be passed + # as an explicit credential (from_single_file / from_pretrained / the Hub client + # can error on a malformed token instead of falling back). Normalize once here so + # every load branch and the size estimate below use a real token or None. + hf_token = hf_token.strip() if isinstance(hf_token, str) else hf_token + hf_token = hf_token or None + # Validate first (cheap, no torch/diffusers) so a direct call with a bad # family fails with ValueError even in a no-diffusers runtime. fam = self.validate_load_request( @@ -1288,7 +1319,10 @@ class DiffusionBackend: if denoiser is None or vae is None: return try: - target_dtype = denoiser.dtype + # Read the dtype from a parameter (not denoiser.dtype): a plain nn.Module has no + # .dtype, and a torch.compile'd/ wrapped denoiser can obscure it; this also + # matches how the VAE dtype is read on the next line. + target_dtype = next(denoiser.parameters()).dtype if next(vae.parameters()).dtype != target_dtype: vae.to(dtype = target_dtype) except (StopIteration, AttributeError, RuntimeError): @@ -1556,9 +1590,8 @@ class DiffusionBackend: mask_pil = mask_pil.resize(init_pil.size, _PILImage.NEAREST) if init_pil is not None: # Keep the VAE encode dtype consistent with the input image. - self._align_vae_dtype( - pipe, getattr(state.family, "denoiser_attr", "transformer") - ) + # state.family is always a DiffusionFamily, which defines denoiser_attr. + self._align_vae_dtype(pipe, state.family.denoiser_attr) # Pipelines vary in which kwargs they accept (img2img derives size from the # input image and may reject width/height; a distilled pipe may take no @@ -1876,6 +1909,33 @@ def _base_file_downloaded(rfilename: str) -> bool: return not rfilename.startswith("assets/") +# Weight file extensions the base repo need NOT supply when the single file is the whole +# pipeline (SDXL): from_single_file(config=base) reads only the base repo's structure +# (config/tokenizer/scheduler) and takes the weights from the single file. +_BASE_WEIGHT_EXTS = ( + ".safetensors", + ".bin", + ".ckpt", + ".pt", + ".pth", + ".gguf", + ".onnx", + ".onnx_data", + ".msgpack", + ".h5", + ".pb", +) + + +def _base_config_file_downloaded(rfilename: str) -> bool: + """True for base-repo files needed to BUILD a pipeline structure around a whole-pipeline + single file WITHOUT its weights: config / tokenizer / scheduler JSON, but no weight + tensors (the single file supplies those). Used for ``single_file_is_pipeline`` families.""" + if not _base_file_downloaded(rfilename): + return False + return not rfilename.lower().endswith(_BASE_WEIGHT_EXTS) + + def _pipeline_file_downloaded(rfilename: str) -> bool: """True for files a full-pipeline ``from_pretrained`` fetches. diff --git a/studio/backend/tests/test_diffusion_sdxl.py b/studio/backend/tests/test_diffusion_sdxl.py index 2a78c0c7de..fa3001e4f1 100644 --- a/studio/backend/tests/test_diffusion_sdxl.py +++ b/studio/backend/tests/test_diffusion_sdxl.py @@ -15,6 +15,8 @@ from __future__ import annotations import types +import pytest + from core.inference import diffusion_lora from core.inference.diffusion import ( DiffusionBackend, @@ -70,7 +72,8 @@ def test_sdxl_base_repos_are_trusted_non_gguf(): # Official safetensors-only base repos are allowlisted so their catalog entries load. assert _is_trusted_diffusion_repo("stabilityai/stable-diffusion-xl-base-1.0") assert _is_trusted_diffusion_repo("stabilityai/sdxl-turbo") - assert _is_trusted_diffusion_repo("stabilityai/stable-diffusion-xl-refiner-1.0") + # The refiner is img2img-only and is intentionally NOT allowlisted (see + # test_sdxl_refiner_not_trusted). # Case-insensitive match. assert _is_trusted_diffusion_repo("StabilityAI/SDXL-Turbo") # A random repo (even one that detects as SDXL) is NOT trusted for a non-GGUF load. @@ -100,9 +103,11 @@ class _FakeVae: def test_align_vae_dtype_uses_unet_denoiser(): # For SDXL the denoiser lives at pipe.unet; _align_vae_dtype must read it (a pipe - # with only .unet and no .transformer) and cast the VAE to the U-Net's dtype. + # with only .unet and no .transformer) and cast the VAE to the U-Net's dtype. The + # dtype is read from a parameter (denoiser has no .dtype), so use a _FakeVae denoiser. vae = _FakeVae(dtype = "float32") - pipe = types.SimpleNamespace(unet = types.SimpleNamespace(dtype = "bfloat16"), vae = vae) + unet = _FakeVae(dtype = "bfloat16") + pipe = types.SimpleNamespace(unet = unet, vae = vae) DiffusionBackend._align_vae_dtype(pipe, "unet") assert vae.moved_to == "bfloat16" @@ -110,7 +115,8 @@ def test_align_vae_dtype_uses_unet_denoiser(): def test_align_vae_dtype_transformer_default_unchanged(): # DiT default: reads pipe.transformer; a pipe with no transformer is a safe no-op. vae = _FakeVae(dtype = "float32") - pipe = types.SimpleNamespace(transformer = types.SimpleNamespace(dtype = "bfloat16"), vae = vae) + transformer = _FakeVae(dtype = "bfloat16") + pipe = types.SimpleNamespace(transformer = transformer, vae = vae) DiffusionBackend._align_vae_dtype(pipe) assert vae.moved_to == "bfloat16" # No denoiser attribute -> no-op (does not raise, does not move the VAE). @@ -147,3 +153,39 @@ def test_pipeline_prefetch_skips_non_torch_artifacts(): assert not keep("unet/flax_model.msgpack") assert not keep("vae_decoder/model.onnx_data") assert not keep("assets/preview.png") + + +def test_sdxl_refiner_not_trusted(): + # The refiner is an img2img-only pipeline; the sdxl family loads every repo as the + # base txt2img pipeline, so the refiner must NOT be allowlisted for a non-GGUF load. + assert not _is_trusted_diffusion_repo("stabilityai/stable-diffusion-xl-refiner-1.0") + # The base and turbo remain trusted. + assert _is_trusted_diffusion_repo("stabilityai/stable-diffusion-xl-base-1.0") + assert _is_trusted_diffusion_repo("stabilityai/sdxl-turbo") + + +def test_sdxl_gguf_load_rejected_up_front(): + # SDXL has no transformer-only GGUF variant (its single file is the whole pipeline), + # so a GGUF request must fail cheap validation before the GPU handoff. + backend = DiffusionBackend() + with pytest.raises(ValueError, match = "no GGUF"): + backend.validate_load_request( + "some-org/my-sdxl.gguf", gguf_filename = "my-sdxl.gguf", family_override = "sdxl" + ) + + +def test_base_config_filter_skips_weights(): + # For a whole-pipeline single file, the base repo supplies only config/tokenizer, not + # its (unused) weight tensors. + from core.inference.diffusion import _base_config_file_downloaded as keep + + assert keep("model_index.json") + assert keep("text_encoder/config.json") + assert keep("tokenizer/vocab.json") + assert keep("scheduler/scheduler_config.json") + assert not keep("unet/diffusion_pytorch_model.safetensors") + assert not keep("vae/diffusion_pytorch_model.bin") + assert not keep("text_encoder/model.onnx") + # transformer/ and assets/ stay excluded (inherited from _base_file_downloaded). + assert not keep("transformer/config.json") + assert not keep("assets/x.png") From e05c9cc947e9b894e733b4d73f3d091b7dcd2de6 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Thu, 2 Jul 2026 01:13:39 +0000 Subject: [PATCH 09/11] Images: preserve restored LoRAs through model load and never send hidden LoRAs - The LoRA effect cleared the selection on every load->capable transition, which wiped adapters restored from a gallery recipe before the model finished loading. Track the previously-loaded family in a ref and clear only on a real family swap; keep the selection on the initial load and on unload. - Gate the generate payload's loras on loraCapable so a restored selection that is hidden (loaded model does not support LoRA) is never sent to the backend. --- .../src/features/images/images-page.tsx | 24 +++++++++++++++---- 1 file changed, 20 insertions(+), 4 deletions(-) diff --git a/studio/frontend/src/features/images/images-page.tsx b/studio/frontend/src/features/images/images-page.tsx index eacfea5c15..b8d676df3d 100644 --- a/studio/frontend/src/features/images/images-page.tsx +++ b/studio/frontend/src/features/images/images-page.tsx @@ -985,18 +985,30 @@ export function ImagesPage({ active = true }: { active?: boolean }) { }, [images, hasMore, selectedId, quant]); // Refresh the LoRA picker's suggestions when the loaded model (family) changes. A LoRA is - // trained for a specific base family, so a model swap invalidates the current selection -- - // clear it (the user re-adds a suggestion or types a Hub repo id for the new family). We do + // trained for a specific base family, so a real model SWAP invalidates the current selection + // -- clear it then (the user re-adds a suggestion or types a Hub repo id for the new family). + // But do NOT clear on the first load or an unload: a user can restore a saved recipe (which + // sets loras) BEFORE the model finishes loading, and clearing on that load->capable + // transition would silently drop the restored adapters. We track the previously-loaded + // family in a ref and clear only when it changes to a different loaded family. We also do // NOT filter the selection against the discovered catalog: a valid pick can be a free-text // Hugging Face repo id that is not in the (often empty) curated list. const loraCapable = Boolean(status?.loaded && status?.supports_lora); + const prevLoraFamilyRef = useRef(undefined); useEffect(() => { if (!loraCapable) { + // Options are gone with the model, but keep the selection: it may have just been + // restored while the model is (re)loading. It is only SENT when loraCapable, and a + // real family swap below clears it. setAvailableLoras([]); - setLoras([]); return; } - setLoras([]); + const fam = status?.family ?? null; + const prev = prevLoraFamilyRef.current; + if (prev != null && prev !== fam) { + setLoras([]); + } + prevLoraFamilyRef.current = fam; let cancelled = false; listDiffusionLoras(status?.family ?? undefined) .then((list) => { @@ -1564,7 +1576,11 @@ export function ImagesPage({ active = true }: { active?: boolean }) { reference_images: condRefImages, // Drop empty (no id typed yet) and zero-weight rows, and trim hand-typed repo ids, // so the recipe records only adapters that actually applied. Empty -> omit entirely. + // Gate on loraCapable: a restore can leave adapters in state while the loaded model + // does not support LoRA (picker hidden), and sending them would fail generation with + // no visible row to remove. loras: (() => { + if (!loraCapable) return undefined; const active = loras .map((l) => ({ id: l.id.trim(), weight: l.weight })) .filter((l) => l.id && l.weight > 0); From dcf6cf511742e21d747d9ea7b21bc6e1ea787479 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 01:14:19 +0000 Subject: [PATCH 10/11] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- studio/backend/tests/test_diffusion_backend.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/studio/backend/tests/test_diffusion_backend.py b/studio/backend/tests/test_diffusion_backend.py index 947d9e7ae3..5aa8c9a71e 100644 --- a/studio/backend/tests/test_diffusion_backend.py +++ b/studio/backend/tests/test_diffusion_backend.py @@ -493,7 +493,9 @@ def test_generate_rejects_reference_on_unsupported_family(fake_runtime, tmp_path ) with pytest.raises(ValueError, match = "Reference images are not supported"): backend.generate( - prompt = "x", steps = 4, init_image = _tiny_png_b64(), + prompt = "x", + steps = 4, + init_image = _tiny_png_b64(), reference_images = [_tiny_png_b64()], ) From 76b5ee354d9f692de9a1751bec290a809b541211 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Thu, 2 Jul 2026 01:23:15 +0000 Subject: [PATCH 11/11] Merge diffusion-sdxl into diffusion-lora-ux; keep options-only LoRA catch The catalog-refresh .catch from the lower branch clears the selected adapters too, which is right for its catalog-only picker but wrong here: this picker holds free-text HF repo ids that are valid without being in the catalog, so a transient refresh failure must not wipe them. Family swaps still clear the selection and hidden LoRAs are never sent. --- studio/frontend/src/features/images/images-page.tsx | 13 +++++++------ 1 file changed, 7 insertions(+), 6 deletions(-) diff --git a/studio/frontend/src/features/images/images-page.tsx b/studio/frontend/src/features/images/images-page.tsx index 2f31492122..d71c35e37d 100644 --- a/studio/frontend/src/features/images/images-page.tsx +++ b/studio/frontend/src/features/images/images-page.tsx @@ -1015,12 +1015,13 @@ export function ImagesPage({ active = true }: { active?: boolean }) { if (!cancelled) setAvailableLoras(list); }) .catch(() => { - if (cancelled) return; - // Clear the SELECTED adapters too, not just the options: leaving a stale `loras` - // selection in state (with the picker now hidden/empty) would still be posted by - // handleGenerate and could apply adapters from the previous model, or fail. - setAvailableLoras([]); - setLoras([]); + // Clear only the OPTIONS on a failed catalog refresh. Unlike the catalog-only + // picker below the stack, this free-text picker holds selections (bare HF repo + // ids) that are valid without being in the catalog; a transient refresh failure + // must not wipe them. Stale cross-family selections are already cleared by the + // family-swap check above, and hidden LoRAs are never sent (handleGenerate is + // gated on loraCapable). + if (!cancelled) setAvailableLoras([]); }); return () => { cancelled = true;