diff --git a/studio/backend/core/inference/diffusion.py b/studio/backend/core/inference/diffusion.py index 0eef14de04..43dc906d0d 100644 --- a/studio/backend/core/inference/diffusion.py +++ b/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( diff --git a/studio/backend/tests/test_diffusion_sdxl.py b/studio/backend/tests/test_diffusion_sdxl.py index fa3001e4f1..2ebe7f2013 100644 --- a/studio/backend/tests/test_diffusion_sdxl.py +++ b/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(