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

studio/backend/core/inference/diffusion.py

@@ -1471,13 +1471,19 @@ class DiffusionBackend:
         if denoiser is None or vae is None:
             return
         try:
-            # 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
+            # Read the dtype from the parameters (not denoiser.dtype): a plain nn.Module
+            # has no .dtype, and a torch.compile'd / wrapped denoiser can obscure it. Take
+            # the first FLOATING dtype: a GGUF-quantized transformer's leading params are
+            # packed uint8 storage, and nn.Module.to() rejects integer dtypes outright.
+            target_dtype = next(
+                (p.dtype for p in denoiser.parameters() if p.dtype.is_floating_point),
+                None,
+            )
+            if target_dtype is None:
+                return
             if next(vae.parameters()).dtype != target_dtype:
                 vae.to(dtype = target_dtype)
-        except (StopIteration, AttributeError, RuntimeError):
+        except (StopIteration, AttributeError, RuntimeError, TypeError):
             pass
 
     def _apply_loras(

studio/backend/tests/test_diffusion_sdxl.py

@@ -105,26 +105,58 @@ 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. The
     # dtype is read from a parameter (denoiser has no .dtype), so use a _FakeVae denoiser.
-    vae = _FakeVae(dtype = "float32")
-    unet = _FakeVae(dtype = "bfloat16")
+    import torch
+
+    vae = _FakeVae(dtype = torch.float32)
+    unet = _FakeVae(dtype = torch.bfloat16)
     pipe = types.SimpleNamespace(unet = unet, vae = vae)
     DiffusionBackend._align_vae_dtype(pipe, "unet")
-    assert vae.moved_to == "bfloat16"
+    assert vae.moved_to == torch.bfloat16
 
 
 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")
-    transformer = _FakeVae(dtype = "bfloat16")
+    import torch
+
+    vae = _FakeVae(dtype = torch.float32)
+    transformer = _FakeVae(dtype = torch.bfloat16)
     pipe = types.SimpleNamespace(transformer = transformer, vae = vae)
     DiffusionBackend._align_vae_dtype(pipe)
-    assert vae.moved_to == "bfloat16"
+    assert vae.moved_to == torch.bfloat16
     # No denoiser attribute -> no-op (does not raise, does not move the VAE).
-    vae2 = _FakeVae(dtype = "float32")
+    vae2 = _FakeVae(dtype = torch.float32)
     DiffusionBackend._align_vae_dtype(types.SimpleNamespace(vae = vae2), "unet")
     assert vae2.moved_to is None
 
 
+def test_align_vae_dtype_skips_gguf_packed_uint8_params():
+    # A GGUF-quantized transformer's leading parameters are packed uint8 storage; the
+    # dtype probe must skip them and use the first FLOATING dtype, or nn.Module.to()
+    # rejects the integer dtype and an Edit/img2img call 500s (regression: Qwen-Image-
+    # Edit GGUF). All-integer params (no floating dtype at all) must be a clean no-op.
+    import torch
+
+    class _GgufDenoiser:
+        def parameters(self):
+            yield types.SimpleNamespace(dtype = torch.uint8)   # packed GGUF block
+            yield types.SimpleNamespace(dtype = torch.bfloat16)  # compute dtype
+
+    vae = _FakeVae(dtype = torch.float32)
+    pipe = types.SimpleNamespace(transformer = _GgufDenoiser(), vae = vae)
+    DiffusionBackend._align_vae_dtype(pipe)
+    assert vae.moved_to == torch.bfloat16
+
+    class _AllPacked:
+        def parameters(self):
+            yield types.SimpleNamespace(dtype = torch.uint8)
+
+    vae2 = _FakeVae(dtype = torch.float32)
+    DiffusionBackend._align_vae_dtype(
+        types.SimpleNamespace(transformer = _AllPacked(), vae = vae2)
+    )
+    assert vae2.moved_to is None
+
+
 def test_sdxl_lora_supported_on_diffusers():
     # SDXL is bf16/bnb-4bit on diffusers -> LoRA is allowed (unlike GGUF-via-diffusers).
     assert diffusion_lora.supports_lora(