fixes

Author: sayakpaul

src/diffusers/utils/testing_utils.py

@@ -291,6 +291,18 @@ def decorator(test_case):
     return decorator
 
 
+def require_torch_version_greater(torch_version):
+    """Decorator marking a test that requires torch with a specific version greater."""
+
+    def decorator(test_case):
+        correct_torch_version = is_torch_available() and is_torch_version(">", torch_version)
+        return unittest.skipUnless(
+            correct_torch_version, f"test requires torch with the version greater than {torch_version}"
+        )(test_case)
+
+    return decorator
+
+
 def require_torch_gpu(test_case):
     """Decorator marking a test that requires CUDA and PyTorch."""
     return unittest.skipUnless(is_torch_available() and torch_device == "cuda", "test requires PyTorch+CUDA")(

tests/quantization/bnb/test_4bit.py

@@ -28,9 +28,9 @@
     DiffusionPipeline,
     FluxControlPipeline,
     FluxTransformer2DModel,
-    PipelineQuantizationConfig,
     SD3Transformer2DModel,
 )
+from diffusers.quantizers import PipelineQuantizationConfig
 from diffusers.utils import is_accelerate_version, logging
 from diffusers.utils.testing_utils import (
     CaptureLogger,
@@ -46,7 +46,7 @@
     require_torch,
     require_torch_accelerator,
     require_torch_gpu,
-    require_torch_version_greater_equal,
+    require_torch_version_greater,
     require_transformers_version_greater,
     slow,
     torch_device,
@@ -875,17 +875,26 @@ def tearDown(self):
         backend_empty_cache(torch_device)
         torch.compiler.reset()
 
-    @require_torch_version_greater_equal("2.8")
+    @require_torch_version_greater("2.7.1")
     def test_torch_compile_4bit(self):
+        torch._dynamo.config.capture_dynamic_output_shape_ops = True
+
         quantization_config = PipelineQuantizationConfig(
             quant_backend="bitsandbytes_4bit",
-            quant_kwargs={"load_in_4bit": True},
+            quant_kwargs={
+                "load_in_4bit": True,
+                "bnb_4bit_quant_type": "nf4",
+                "bnb_4bit_compute_dtype": torch.bfloat16,
+            },
             components_to_quantize=["transformer"],
         )
         pipe = DiffusionPipeline.from_pretrained(
-            "hf-internal-testing/tiny-flux-pipe", quantization_config=quantization_config, torch_dtype=torch.bfloat16
+            "stabilityai/stable-diffusion-3-medium-diffusers",
+            quantization_config=quantization_config,
+            torch_dtype=torch.bfloat16,
         ).to("cuda")
         pipe.transformer.compile(fullgraph=True)
 
         for _ in range(2):
-            pipe("a dog", num_inference_steps=4, max_sequence_length=16)
+            # with torch._dynamo.config.patch(error_on_recompile=True):
+            pipe("a dog", num_inference_steps=4, max_sequence_length=16, height=256, width=256)