diff --git a/studio/backend/core/inference/diffusion.py b/studio/backend/core/inference/diffusion.py index da04a4ceee..b698dd71f4 100644 --- a/studio/backend/core/inference/diffusion.py +++ b/studio/backend/core/inference/diffusion.py @@ -2268,7 +2268,9 @@ class DiffusionBackend: @staticmethod def _align_vae_dtype( - pipe: Any, denoiser_attr: str = "transformer", vae_quant: Optional[str] = None + pipe: Any, + denoiser_attr: str = "transformer", + vae_quant: Optional[str] = None, ) -> None: """Cast the VAE to the denoiser's compute dtype before an image-conditioned call. The img2img/inpaint pipelines VAE-encode the input image at the text- diff --git a/studio/backend/core/inference/diffusion_vae_quant.py b/studio/backend/core/inference/diffusion_vae_quant.py index 92ccd0e985..b0a3cf685c 100644 --- a/studio/backend/core/inference/diffusion_vae_quant.py +++ b/studio/backend/core/inference/diffusion_vae_quant.py @@ -269,7 +269,9 @@ def _cast_vae_fp8_dynamic(vae: Any, target: Any) -> None: name = fqn.lower() if fqn else "" return not any(tok in name for tok in _VAE_KEEP_DENSE_TOKENS) - quantize_(vae, Float8DynamicActivationFloat8WeightConfig(granularity = PerTensor()), filter_fn = filter_fn) + quantize_( + vae, Float8DynamicActivationFloat8WeightConfig(granularity = PerTensor()), filter_fn = filter_fn + ) def _cast_vae_fp8(vae: Any, target: Any) -> None: diff --git a/studio/backend/tests/test_diffusion_vae_quant.py b/studio/backend/tests/test_diffusion_vae_quant.py index ba621d6526..a948a8454d 100644 --- a/studio/backend/tests/test_diffusion_vae_quant.py +++ b/studio/backend/tests/test_diffusion_vae_quant.py @@ -29,7 +29,12 @@ from core.inference.diffusion_vae_quant import ( ) -def _target(*, device = "cuda", dtype = "bfloat16", cc = (10, 0)): +def _target( + *, + device = "cuda", + dtype = "bfloat16", + cc = (10, 0), +): return types.SimpleNamespace(device = device, dtype = dtype, _cc = cc) @@ -43,7 +48,12 @@ class _Weight: return len(self.shape) -def _stub_torch(monkeypatch, *, with_fp8 = True, cc = (10, 0)): +def _stub_torch( + monkeypatch, + *, + with_fp8 = True, + cc = (10, 0), +): torch = types.ModuleType("torch") torch.bfloat16 = "bfloat16" torch.float16 = "float16" @@ -164,9 +174,7 @@ def test_select_offload_uses_layerwise_fp8(monkeypatch): monkeypatch.setattr( vq, "_vae_fp8_dynamic_probe", lambda device: pytest.fail("probe must not run under offload") ) - assert ( - select_vae_quant_scheme(_target(), "auto", offload_active = True) == VAE_QUANT_FP8 - ) + assert select_vae_quant_scheme(_target(), "auto", offload_active = True) == VAE_QUANT_FP8 def test_select_force_fp32_stays_dense(monkeypatch): @@ -258,6 +266,7 @@ def test_fp8_dynamic_conv_filter(monkeypatch): assert ff(_mod(nn.Conv2d, (128, 128, 3, 3)), "decoder.conv_norm_out") is False assert ff(_mod(nn.Conv2d, (128, 128, 3, 3)), "decoder.norm_out.conv") is False assert ff(_mod(nn.Linear, (512, 512)), "decoder.proj_out") is False + # A non-conv/linear module (e.g. a GroupNorm) is excluded outright. class _GroupNorm: pass