small fix

Author: patrickvonplaten

src/diffusers/pipelines/alt_diffusion/pipeline_alt_diffusion.py

@@ -236,7 +236,7 @@ def enable_sequential_cpu_offload(self, gpu_id=0):
         if self.safety_checker is not None:
             # TODO(Patrick) - there is currently a bug with cpu offload of nn.Parameter in accelerate
             # fix by only offloading self.safety_checker for now
-            cpu_offload(self.safety_checker.vision_model)
+            cpu_offload(self.safety_checker.vision_model, device)
 
     @property
     def _execution_device(self):

src/diffusers/pipelines/alt_diffusion/pipeline_alt_diffusion_img2img.py

@@ -231,7 +231,7 @@ def enable_sequential_cpu_offload(self, gpu_id=0):
         if self.safety_checker is not None:
             # TODO(Patrick) - there is currently a bug with cpu offload of nn.Parameter in accelerate
             # fix by only offloading self.safety_checker for now
-            cpu_offload(self.safety_checker.vision_model)
+            cpu_offload(self.safety_checker.vision_model, device)
 
     @property
     def _execution_device(self):

src/diffusers/pipelines/stable_diffusion/pipeline_cycle_diffusion.py

@@ -264,7 +264,7 @@ def enable_sequential_cpu_offload(self, gpu_id=0):
         if self.safety_checker is not None:
             # TODO(Patrick) - there is currently a bug with cpu offload of nn.Parameter in accelerate
             # fix by only offloading self.safety_checker for now
-            cpu_offload(self.safety_checker.vision_model)
+            cpu_offload(self.safety_checker.vision_model, device)
 
     @property
     # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline._execution_device

src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion.py

@@ -235,7 +235,7 @@ def enable_sequential_cpu_offload(self, gpu_id=0):
         if self.safety_checker is not None:
             # TODO(Patrick) - there is currently a bug with cpu offload of nn.Parameter in accelerate
             # fix by only offloading self.safety_checker for now
-            cpu_offload(self.safety_checker.vision_model)
+            cpu_offload(self.safety_checker.vision_model, device)
 
     @property
     def _execution_device(self):

src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_img2img.py

@@ -233,7 +233,7 @@ def enable_sequential_cpu_offload(self, gpu_id=0):
         if self.safety_checker is not None:
             # TODO(Patrick) - there is currently a bug with cpu offload of nn.Parameter in accelerate
             # fix by only offloading self.safety_checker for now
-            cpu_offload(self.safety_checker.vision_model)
+            cpu_offload(self.safety_checker.vision_model, device)
 
     @property
     # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline._execution_device

src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_inpaint.py

@@ -298,7 +298,7 @@ def enable_sequential_cpu_offload(self, gpu_id=0):
         if self.safety_checker is not None:
             # TODO(Patrick) - there is currently a bug with cpu offload of nn.Parameter in accelerate
             # fix by only offloading self.safety_checker for now
-            cpu_offload(self.safety_checker.vision_model)
+            cpu_offload(self.safety_checker.vision_model, device)
 
     # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.enable_xformers_memory_efficient_attention
     def enable_xformers_memory_efficient_attention(self):

src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_inpaint_legacy.py

@@ -246,7 +246,7 @@ def enable_sequential_cpu_offload(self, gpu_id=0):
         if self.safety_checker is not None:
             # TODO(Patrick) - there is currently a bug with cpu offload of nn.Parameter in accelerate
             # fix by only offloading self.safety_checker for now
-            cpu_offload(self.safety_checker.vision_model)
+            cpu_offload(self.safety_checker.vision_model, device)
 
     # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.enable_xformers_memory_efficient_attention
     def enable_xformers_memory_efficient_attention(self):

tests/pipelines/versatile_diffusion/test_versatile_diffusion_image_variation.py

@@ -54,5 +54,6 @@ def test_inference_image_variations(self):
         image_slice = image[0, 253:256, 253:256, -1]
 
         assert image.shape == (1, 512, 512, 3)
+        print(torch.from_numpy(image_slice.flatten()))
         expected_slice = np.array([0.0113, 0.2241, 0.4024, 0.0839, 0.0871, 0.2725, 0.2581, 0.0, 0.1096])
         assert np.abs(image_slice.flatten() - expected_slice).max() < 1e-2