Merge branch 'main' into speedup-model-loading

Author: a-r-r-o-w

examples/server/requirements.in

@@ -1,4 +1,4 @@
-torch~=2.4.0
+torch~=2.7.0
 transformers==4.46.1
 sentencepiece
 aiohttp

examples/server/requirements.txt

@@ -63,36 +63,42 @@ networkx==3.2.1
     # via torch
 numpy==2.0.2
     # via transformers
-nvidia-cublas-cu12==12.1.3.1
+nvidia-cublas-cu12==12.6.4.1
     # via
     #   nvidia-cudnn-cu12
     #   nvidia-cusolver-cu12
     #   torch
-nvidia-cuda-cupti-cu12==12.1.105
+nvidia-cuda-cupti-cu12==12.6.80
     # via torch
-nvidia-cuda-nvrtc-cu12==12.1.105
+nvidia-cuda-nvrtc-cu12==12.6.77
     # via torch
-nvidia-cuda-runtime-cu12==12.1.105
+nvidia-cuda-runtime-cu12==12.6.77
     # via torch
-nvidia-cudnn-cu12==9.1.0.70
+nvidia-cudnn-cu12==9.5.1.17
     # via torch
-nvidia-cufft-cu12==11.0.2.54
+nvidia-cufft-cu12==11.3.0.4
     # via torch
-nvidia-curand-cu12==10.3.2.106
+nvidia-cufile-cu12==1.11.1.6
     # via torch
-nvidia-cusolver-cu12==11.4.5.107
+nvidia-curand-cu12==10.3.7.77
     # via torch
-nvidia-cusparse-cu12==12.1.0.106
+nvidia-cusolver-cu12==11.7.1.2
+    # via torch
+nvidia-cusparse-cu12==12.5.4.2
     # via
     #   nvidia-cusolver-cu12
     #   torch
-nvidia-nccl-cu12==2.20.5
+nvidia-cusparselt-cu12==0.6.3
+    # via torch
+nvidia-nccl-cu12==2.26.2
     # via torch
-nvidia-nvjitlink-cu12==12.9.86
+nvidia-nvjitlink-cu12==12.6.85
     # via
+    #   nvidia-cufft-cu12
     #   nvidia-cusolver-cu12
     #   nvidia-cusparse-cu12
-nvidia-nvtx-cu12==12.1.105
+    #   torch
+nvidia-nvtx-cu12==12.6.77
     # via torch
 packaging==24.1
     # via
@@ -137,20 +143,19 @@ sympy==1.13.3
     # via torch
 tokenizers==0.20.1
     # via transformers
-torch==2.4.1
+torch==2.7.0
     # via -r requirements.in
 tqdm==4.66.5
     # via
     #   huggingface-hub
     #   transformers
 transformers==4.46.1
     # via -r requirements.in
-triton==3.0.0
+triton==3.3.0
     # via torch
 typing-extensions==4.12.2
     # via
     #   anyio
-    #   exceptiongroup
     #   fastapi
     #   huggingface-hub
     #   multidict

src/diffusers/loaders/lora_base.py

@@ -470,7 +470,7 @@ def _func_optionally_disable_offloading(_pipeline):
             for _, component in _pipeline.components.items():
                 if not isinstance(component, nn.Module) or not hasattr(component, "_hf_hook"):
                     continue
-            remove_hook_from_module(component, recurse=is_sequential_cpu_offload)
+                remove_hook_from_module(component, recurse=is_sequential_cpu_offload)
 
     return (is_model_cpu_offload, is_sequential_cpu_offload, is_group_offload)
 

src/diffusers/utils/torch_utils.py

@@ -175,6 +175,8 @@ def get_device():
         return "npu"
     elif hasattr(torch, "xpu") and torch.xpu.is_available():
         return "xpu"
+    elif torch.backends.mps.is_available():
+        return "mps"
     else:
         return "cpu"
 

tests/lora/utils.py

@@ -2510,3 +2510,34 @@ def test_group_offloading_inference_denoiser(self, offload_type, use_stream):
                 # materializes the test methods on invocation which cannot be overridden.
                 return
         self._test_group_offloading_inference_denoiser(offload_type, use_stream)
+
+    @require_torch_accelerator
+    def test_lora_loading_model_cpu_offload(self):
+        components, _, denoiser_lora_config = self.get_dummy_components(self.scheduler_classes[0])
+        _, _, inputs = self.get_dummy_inputs(with_generator=False)
+        pipe = self.pipeline_class(**components)
+        pipe = pipe.to(torch_device)
+        pipe.set_progress_bar_config(disable=None)
+
+        denoiser = pipe.transformer if self.unet_kwargs is None else pipe.unet
+        denoiser.add_adapter(denoiser_lora_config)
+        self.assertTrue(check_if_lora_correctly_set(denoiser), "Lora not correctly set in denoiser.")
+
+        output_lora = pipe(**inputs, generator=torch.manual_seed(0))[0]
+
+        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)
+            self.pipeline_class.save_lora_weights(
+                save_directory=tmpdirname, safe_serialization=True, **lora_state_dicts
+            )
+            # reinitialize the pipeline to mimic the inference workflow.
+            components, _, denoiser_lora_config = self.get_dummy_components(self.scheduler_classes[0])
+            pipe = self.pipeline_class(**components)
+            pipe.enable_model_cpu_offload(device=torch_device)
+            pipe.load_lora_weights(tmpdirname)
+            denoiser = pipe.transformer if self.unet_kwargs is None else pipe.unet
+            self.assertTrue(check_if_lora_correctly_set(denoiser), "Lora not correctly set in denoiser.")
+
+        output_lora_loaded = pipe(**inputs, generator=torch.manual_seed(0))[0]
+        self.assertTrue(np.allclose(output_lora, output_lora_loaded, atol=1e-3, rtol=1e-3))