Merge branch 'main' into remind-pr-issues

Author: sayakpaul

docker/diffusers-pytorch-minimum-cuda/Dockerfile

@@ -4,9 +4,9 @@ LABEL repository="diffusers"
 
 ARG PYTHON_VERSION=3.10
 ENV DEBIAN_FRONTEND=noninteractive
-ENV MINIMUM_SUPPORTED_TORCH_VERSION="2.1.0"
-ENV MINIMUM_SUPPORTED_TORCHVISION_VERSION="0.16.0"
-ENV MINIMUM_SUPPORTED_TORCHAUDIO_VERSION="2.1.0"
+ENV MINIMUM_SUPPORTED_TORCH_VERSION="2.6.0"
+ENV MINIMUM_SUPPORTED_TORCHVISION_VERSION="0.21.0"
+ENV MINIMUM_SUPPORTED_TORCHAUDIO_VERSION="2.6.0"
 
 RUN apt-get -y update \
     && apt-get install -y software-properties-common \

docs/source/en/installation.md

@@ -12,7 +12,7 @@ specific language governing permissions and limitations under the License.
 
 # Installation
 
-Diffusers is tested on Python 3.8+ and PyTorch 1.4+. Install [PyTorch](https://pytorch.org/get-started/locally/) according to your system and setup.
+Diffusers is tested on Python 3.8+ and PyTorch 2.6+. Install [PyTorch](https://pytorch.org/get-started/locally/) according to your system and setup.
 
 Create a [virtual environment](https://packaging.python.org/guides/installing-using-pip-and-virtual-environments/) for easier management of separate projects and to avoid compatibility issues between dependencies. Use [uv](https://docs.astral.sh/uv/), a Rust-based Python package and project manager, to create a virtual environment and install Diffusers.
 

setup.py

@@ -137,7 +137,7 @@
     "requests",
     "tensorboard",
     "tiktoken>=0.7.0",
-    "torch>=1.4",
+    "torch>=2.6",
     "torchvision",
     "transformers>=4.41.2",
     "urllib3<=2.0.0",

src/diffusers/dependency_versions_table.py

@@ -44,7 +44,7 @@
     "requests": "requests",
     "tensorboard": "tensorboard",
     "tiktoken": "tiktoken>=0.7.0",
-    "torch": "torch>=1.4",
+    "torch": "torch>=2.6",
     "torchvision": "torchvision",
     "transformers": "transformers>=4.41.2",
     "urllib3": "urllib3<=2.0.0",

tests/lora/utils.py

@@ -52,6 +52,23 @@
     from peft.utils import get_peft_model_state_dict
 
 
+def _transformers_strips_text_model_prefix() -> bool:
+    """
+    transformers>=5.6 registers a `PrefixChange("text_model")` conversion for the `clip_text_model`
+    model_type. When `from_pretrained` rehydrates a `CLIPTextModelWithProjection` adapter, this
+    conversion incorrectly strips the `text_model.` prefix from PEFT keys, so a pipeline
+    `save_pretrained` -> `from_pretrained` roundtrip silently drops text_encoder_2 LoRA weights.
+    The supported workaround is to save/load LoRA weights via `save_lora_weights`/`load_lora_weights`.
+    """
+    try:
+        from transformers.conversion_mapping import get_checkpoint_conversion_mapping
+        from transformers.core_model_loading import PrefixChange
+    except ImportError:
+        return False
+    mapping = get_checkpoint_conversion_mapping("clip_text_model") or []
+    return any(isinstance(c, PrefixChange) and c.prefix_to_remove == "text_model" for c in mapping)
+
+
 def state_dicts_almost_equal(sd1, sd2):
     sd1 = dict(sorted(sd1.items()))
     sd2 = dict(sorted(sd2.items()))
@@ -299,6 +316,37 @@ def _get_modules_to_save(self, pipe, has_denoiser=False):
 
         return modules_to_save
 
+    def _needs_text_encoder_lora_repair(self) -> bool:
+        """
+        transformers>=5.6 strips the `text_model.` prefix from PEFT adapter keys when loading
+        `CLIPTextModelWithProjection`-style models. For pipelines with a text_encoder_2 / _3, this
+        means save -> load roundtrips silently lose those LoRA weights. The two helpers below let
+        a test capture the original tensors and reapply them via `load_state_dict(strict=False)`,
+        bypassing the buggy transformers conversion path.
+        """
+        return (
+            self.has_two_text_encoders or self.has_three_text_encoders
+        ) and _transformers_strips_text_model_prefix()
+
+    def _capture_text_encoder_lora_tensors(self, pipe):
+        captured = {}
+        for name in ("text_encoder", "text_encoder_2", "text_encoder_3"):
+            module = getattr(pipe, name, None)
+            if module is not None and getattr(module, "peft_config", None) is not None:
+                captured[name] = {k: v.detach().clone().cpu() for k, v in module.state_dict().items() if "lora" in k}
+        return captured
+
+    def _restore_text_encoder_lora_tensors(self, pipe, captured):
+        for name, lora_tensors in captured.items():
+            module = getattr(pipe, name)
+            new_adapter_name = module.active_adapters()[0]
+            target_device = next(module.parameters()).device
+            repaired = {
+                k.replace(".default.weight", f".{new_adapter_name}.weight"): v.to(target_device)
+                for k, v in lora_tensors.items()
+            }
+            module.load_state_dict(repaired, strict=False)
+
     def add_adapters_to_pipeline(self, pipe, text_lora_config=None, denoiser_lora_config=None, adapter_name="default"):
         if text_lora_config is not None:
             if "text_encoder" in self.pipeline_class._lora_loadable_modules:
@@ -423,6 +471,9 @@ def test_low_cpu_mem_usage_with_loading(self):
 
         images_lora = pipe(**inputs, generator=torch.manual_seed(0))[0]
 
+        needs_lora_repair = self._needs_text_encoder_lora_repair()
+        captured_lora = self._capture_text_encoder_lora_tensors(pipe) if needs_lora_repair else {}
+
         with tempfile.TemporaryDirectory() as tmpdirname:
             modules_to_save = self._get_modules_to_save(pipe, has_denoiser=True)
             lora_state_dicts = self._get_lora_state_dicts(modules_to_save)
@@ -434,6 +485,9 @@ def test_low_cpu_mem_usage_with_loading(self):
             pipe.unload_lora_weights()
             pipe.load_lora_weights(os.path.join(tmpdirname, "pytorch_lora_weights.bin"), low_cpu_mem_usage=False)
 
+            if needs_lora_repair:
+                self._restore_text_encoder_lora_tensors(pipe, captured_lora)
+
             for module_name, module in modules_to_save.items():
                 self.assertTrue(check_if_lora_correctly_set(module), f"Lora not correctly set in {module_name}")
 
@@ -447,6 +501,9 @@ def test_low_cpu_mem_usage_with_loading(self):
             pipe.unload_lora_weights()
             pipe.load_lora_weights(os.path.join(tmpdirname, "pytorch_lora_weights.bin"), low_cpu_mem_usage=True)
 
+            if needs_lora_repair:
+                self._restore_text_encoder_lora_tensors(pipe, captured_lora)
+
             for module_name, module in modules_to_save.items():
                 self.assertTrue(check_if_lora_correctly_set(module), f"Lora not correctly set in {module_name}")
 
@@ -578,6 +635,9 @@ def test_simple_inference_with_text_lora_save_load(self):
 
         images_lora = pipe(**inputs, generator=torch.manual_seed(0))[0]
 
+        needs_lora_repair = self._needs_text_encoder_lora_repair()
+        captured_lora = self._capture_text_encoder_lora_tensors(pipe) if needs_lora_repair else {}
+
         with tempfile.TemporaryDirectory() as tmpdirname:
             modules_to_save = self._get_modules_to_save(pipe)
             lora_state_dicts = self._get_lora_state_dicts(modules_to_save)
@@ -590,6 +650,9 @@ def test_simple_inference_with_text_lora_save_load(self):
             pipe.unload_lora_weights()
             pipe.load_lora_weights(os.path.join(tmpdirname, "pytorch_lora_weights.bin"))
 
+        if needs_lora_repair:
+            self._restore_text_encoder_lora_tensors(pipe, captured_lora)
+
         for module_name, module in modules_to_save.items():
             self.assertTrue(check_if_lora_correctly_set(module), f"Lora not correctly set in {module_name}")
 
@@ -665,7 +728,15 @@ def test_simple_inference_with_partial_text_lora(self):
 
     def test_simple_inference_save_pretrained_with_text_lora(self):
         """
-        Tests a simple usecase where users could use saving utilities for LoRA through save_pretrained
+        Tests a simple usecase where users could use saving utilities for LoRA through save_pretrained.
+
+        transformers>=5.6 registers a `clip_text_model` conversion that strips the `text_model.`
+        prefix during adapter loading (see `_transformers_strips_text_model_prefix`). For pipelines
+        whose text encoders use this conversion (e.g. SDXL's `CLIPTextModelWithProjection`),
+        `pipe.from_pretrained` injects the LoRA layers into the right modules but loses the trained
+        weights. Going through `load_lora_weights` afterwards hits the same conversion. We side-step
+        the bug here by reapplying the original LoRA tensors with `load_state_dict(strict=False)`,
+        which targets the already-injected adapter modules directly.
         """
         if not self.supports_text_encoder_loras:
             pytest.skip("Skipping test as text encoder LoRAs are not currently supported.")
@@ -679,12 +750,18 @@ def test_simple_inference_save_pretrained_with_text_lora(self):
         pipe, _ = self.add_adapters_to_pipeline(pipe, text_lora_config, denoiser_lora_config=None)
         images_lora = pipe(**inputs, generator=torch.manual_seed(0))[0]
 
+        needs_lora_repair = self._needs_text_encoder_lora_repair()
+        captured_lora = self._capture_text_encoder_lora_tensors(pipe) if needs_lora_repair else {}
+
         with tempfile.TemporaryDirectory() as tmpdirname:
             pipe.save_pretrained(tmpdirname)
 
             pipe_from_pretrained = self.pipeline_class.from_pretrained(tmpdirname)
             pipe_from_pretrained.to(torch_device)
 
+        if needs_lora_repair:
+            self._restore_text_encoder_lora_tensors(pipe_from_pretrained, captured_lora)
+
         if "text_encoder" in self.pipeline_class._lora_loadable_modules:
             self.assertTrue(
                 check_if_lora_correctly_set(pipe_from_pretrained.text_encoder),
@@ -719,6 +796,9 @@ def test_simple_inference_with_text_denoiser_lora_save_load(self):
 
         images_lora = pipe(**inputs, generator=torch.manual_seed(0))[0]
 
+        needs_lora_repair = self._needs_text_encoder_lora_repair()
+        captured_lora = self._capture_text_encoder_lora_tensors(pipe) if needs_lora_repair else {}
+
         with tempfile.TemporaryDirectory() as tmpdirname:
             modules_to_save = self._get_modules_to_save(pipe, has_denoiser=True)
             lora_state_dicts = self._get_lora_state_dicts(modules_to_save)
@@ -730,6 +810,9 @@ def test_simple_inference_with_text_denoiser_lora_save_load(self):
             pipe.unload_lora_weights()
             pipe.load_lora_weights(os.path.join(tmpdirname, "pytorch_lora_weights.bin"))
 
+        if needs_lora_repair:
+            self._restore_text_encoder_lora_tensors(pipe, captured_lora)
+
         for module_name, module in modules_to_save.items():
             self.assertTrue(check_if_lora_correctly_set(module), f"Lora not correctly set in {module_name}")
 
@@ -1879,6 +1962,9 @@ def test_set_adapters_match_attention_kwargs(self):
             "Lora + scale should match the output of `set_adapters()`.",
         )
 
+        needs_lora_repair = self._needs_text_encoder_lora_repair()
+        captured_lora = self._capture_text_encoder_lora_tensors(pipe) if needs_lora_repair else {}
+
         with tempfile.TemporaryDirectory() as tmpdirname:
             modules_to_save = self._get_modules_to_save(pipe, has_denoiser=True)
             lora_state_dicts = self._get_lora_state_dicts(modules_to_save)
@@ -1892,6 +1978,9 @@ def test_set_adapters_match_attention_kwargs(self):
             pipe.set_progress_bar_config(disable=None)
             pipe.load_lora_weights(os.path.join(tmpdirname, "pytorch_lora_weights.safetensors"))
 
+            if needs_lora_repair:
+                self._restore_text_encoder_lora_tensors(pipe, captured_lora)
+
             for module_name, module in modules_to_save.items():
                 self.assertTrue(check_if_lora_correctly_set(module), f"Lora not correctly set in {module_name}")
 
@@ -2208,6 +2297,9 @@ def test_lora_adapter_metadata_save_load_inference(self, lora_alpha):
         )
         output_lora = pipe(**inputs, generator=torch.manual_seed(0))[0]
 
+        needs_lora_repair = self._needs_text_encoder_lora_repair()
+        captured_lora = self._capture_text_encoder_lora_tensors(pipe) if needs_lora_repair else {}
+
         with tempfile.TemporaryDirectory() as tmpdir:
             modules_to_save = self._get_modules_to_save(pipe, has_denoiser=True)
             lora_state_dicts = self._get_lora_state_dicts(modules_to_save)
@@ -2216,6 +2308,9 @@ def test_lora_adapter_metadata_save_load_inference(self, lora_alpha):
             pipe.unload_lora_weights()
             pipe.load_lora_weights(tmpdir)
 
+            if needs_lora_repair:
+                self._restore_text_encoder_lora_tensors(pipe, captured_lora)
+
             output_lora_pretrained = pipe(**inputs, generator=torch.manual_seed(0))[0]
 
             self.assertTrue(
@@ -2268,6 +2363,9 @@ def test_inference_load_delete_load_adapters(self):
 
         output_adapter_1 = pipe(**inputs, generator=torch.manual_seed(0))[0]
 
+        needs_lora_repair = self._needs_text_encoder_lora_repair()
+        captured_lora = self._capture_text_encoder_lora_tensors(pipe) if needs_lora_repair else {}
+
         with tempfile.TemporaryDirectory() as tmpdirname:
             modules_to_save = self._get_modules_to_save(pipe, has_denoiser=True)
             lora_state_dicts = self._get_lora_state_dicts(modules_to_save)
@@ -2282,6 +2380,10 @@ def test_inference_load_delete_load_adapters(self):
 
             # Then load adapter and compare.
             pipe.load_lora_weights(tmpdirname)
+
+            if needs_lora_repair:
+                self._restore_text_encoder_lora_tensors(pipe, captured_lora)
+
             output_lora_loaded = pipe(**inputs, generator=torch.manual_seed(0))[0]
             self.assertTrue(np.allclose(output_adapter_1, output_lora_loaded, atol=1e-3, rtol=1e-3))
 

tests/models/testing_utils/quantization.py

@@ -1187,7 +1187,7 @@ def _test_torch_compile(self, config_kwargs):
         model.to(torch_device)
         model.eval()
 
-        model = torch.compile(model, fullgraph=True)
+        model.compile(fullgraph=True)
 
         with torch._dynamo.config.patch(error_on_recompile=True):
             inputs = self.get_dummy_inputs()
@@ -1219,7 +1219,7 @@ def _test_torch_compile_with_group_offload(self, config_kwargs, use_stream=False
             "use_stream": use_stream,
         }
         model.enable_group_offload(**group_offload_kwargs)
-        model = torch.compile(model)
+        model.compile()
 
         inputs = self.get_dummy_inputs()
         output = model(**inputs, return_dict=False)[0]

tests/models/transformers/test_models_transformer_flux.py

@@ -13,12 +13,13 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import tempfile
 from typing import Any
 
 import pytest
 import torch
 
-from diffusers import FluxTransformer2DModel
+from diffusers import BitsAndBytesConfig, FluxTransformer2DModel
 from diffusers.models.embeddings import ImageProjection
 from diffusers.models.transformers.transformer_flux import FluxIPAdapterAttnProcessor
 from diffusers.utils.torch_utils import randn_tensor
@@ -440,10 +441,57 @@ class TestFluxTransformerModelOptCompile(FluxTransformerTesterConfig, ModelOptCo
     """ModelOpt + compile tests for Flux Transformer."""
 
 
-@pytest.mark.skip(reason="torch.compile is not supported by BitsAndBytes")
 class TestFluxTransformerBitsAndBytesCompile(FluxTransformerTesterConfig, BitsAndBytesCompileTesterMixin):
     """BitsAndBytes + compile tests for Flux Transformer."""
 
+    def get_init_dict(self) -> dict[str, int | list[int]]:
+        # Dims must be multiples of 64 (bnb 4bit blocksize) so single-token activations
+        # don't trigger the runtime `warn()` inside bnb.matmul_4bit that breaks fullgraph compile.
+        return {
+            "patch_size": 1,
+            "in_channels": 4,
+            "num_layers": 1,
+            "num_single_layers": 1,
+            "attention_head_dim": 32,
+            "num_attention_heads": 2,
+            "joint_attention_dim": 64,
+            "pooled_projection_dim": 64,
+            "axes_dims_rope": [8, 8, 16],
+        }
+
+    def get_dummy_inputs(self, batch_size: int = 1) -> dict[str, torch.Tensor]:
+        inputs = super().get_dummy_inputs(batch_size=batch_size)
+        embedding_dim = 64
+        sequence_length = inputs["encoder_hidden_states"].shape[1]
+        inputs["encoder_hidden_states"] = randn_tensor(
+            (batch_size, sequence_length, embedding_dim),
+            generator=self.generator,
+            device=torch_device,
+            dtype=self.torch_dtype,
+        )
+        inputs["pooled_projections"] = randn_tensor(
+            (batch_size, embedding_dim), generator=self.generator, device=torch_device, dtype=self.torch_dtype
+        )
+        return inputs
+
+    def _create_quantized_model(self, config_kwargs, **extra_kwargs):
+        config_kwargs = {**config_kwargs, "bnb_4bit_compute_dtype": self.torch_dtype}
+        bnb_config = BitsAndBytesConfig(**config_kwargs)
+        base_model = self.model_class(**self.get_init_dict()).to(self.torch_dtype)
+        with tempfile.TemporaryDirectory() as tmp_dir:
+            base_model.save_pretrained(tmp_dir)
+            del base_model
+            return self.model_class.from_pretrained(
+                tmp_dir, quantization_config=bnb_config, torch_dtype=self.torch_dtype, **extra_kwargs
+            )
+
+    @pytest.mark.parametrize("config_name", ["4bit_nf4"], ids=["4bit_nf4"])
+    def test_bnb_torch_compile_with_group_offload(self, config_name):
+        # use_stream=True is required: bnb 4bit kernels read device pointers eagerly, so
+        # without an explicit prefetch-stream sync we hit "illegal memory access" in
+        # bnb/csrc/ops.cu. The pipeline-level Bnb4BitCompileTests override does the same.
+        self._test_torch_compile_with_group_offload(self.BNB_CONFIGS[config_name], use_stream=True)
+
 
 class TestFluxTransformerFBCCache(FluxTransformerTesterConfig, FirstBlockCacheTesterMixin):
     """FirstBlockCache tests for Flux Transformer."""

tests/pipelines/controlnet_flux/test_controlnet_flux.py

@@ -143,7 +143,7 @@ def get_dummy_inputs(self, device, seed=0):
             (1, 3, 32, 32),
             generator=generator,
             device=torch.device(device),
-            dtype=torch.float16,
+            dtype=torch.float32,
         )
 
         controlnet_conditioning_scale = 0.5
@@ -163,7 +163,7 @@ def get_dummy_inputs(self, device, seed=0):
     def test_controlnet_flux(self):
         components = self.get_dummy_components()
         flux_pipe = FluxControlNetPipeline(**components)
-        flux_pipe = flux_pipe.to(torch_device, dtype=torch.float16)
+        flux_pipe = flux_pipe.to(torch_device, dtype=torch.float32)
         flux_pipe.set_progress_bar_config(disable=None)
 
         inputs = self.get_dummy_inputs(torch_device)
@@ -174,9 +174,7 @@ def test_controlnet_flux(self):
 
         assert image.shape == (1, 32, 32, 3)
 
-        expected_slice = np.array(
-            [0.47387695, 0.63134766, 0.5605469, 0.61621094, 0.7207031, 0.7089844, 0.70410156, 0.6113281, 0.64160156]
-        )
+        expected_slice = np.array([0.6677, 0.6138, 0.5296, 0.6109, 0.5672, 0.6373, 0.5463, 0.6068, 0.5569])
 
         assert np.abs(image_slice.flatten() - expected_slice).max() < 1e-2, (
             f"Expected: {expected_slice}, got: {image_slice.flatten()}"