directly place on cuda.

Author: sayakpaul

src/diffusers/pipelines/flux/pipeline_flux.py

@@ -777,7 +777,7 @@ def __call__(
 
         # Offload all models
         self.maybe_free_model_hooks()
-        
+
         if not return_dict:
             return (image,)
 

tests/lora/test_lora_layers_flux.py

@@ -202,7 +202,7 @@ def test_flux_the_last_ben(self):
         self.pipeline.load_lora_weights("TheLastBen/Jon_Snow_Flux_LoRA", weight_name="jon_snow.safetensors")
         self.pipeline.fuse_lora()
         self.pipeline.unload_lora_weights()
-        self.pipeline.enable_model_cpu_offload()
+        self.pipeline = self.pipeline.to("cuda")
 
         prompt = "jon snow eating pizza with ketchup"
 
@@ -225,7 +225,10 @@ def test_flux_kohya(self):
         self.pipeline.load_lora_weights("Norod78/brain-slug-flux")
         self.pipeline.fuse_lora()
         self.pipeline.unload_lora_weights()
-        self.pipeline.enable_model_cpu_offload()
+        # Instead of calling `enable_model_cpu_offload()`, we do a cuda placement here because the CI
+        # run supports it. We have about 34GB RAM in the CI runner which kills the test when run with
+        # `enable_model_cpu_offload()`.
+        self.pipeline = self.pipeline.to("cuda")
 
         prompt = "The cat with a brain slug earring"
         out = self.pipeline(
@@ -248,7 +251,7 @@ def test_flux_kohya_with_text_encoder(self):
         self.pipeline.load_lora_weights("cocktailpeanut/optimus", weight_name="optimus.safetensors")
         self.pipeline.fuse_lora()
         self.pipeline.unload_lora_weights()
-        self.pipeline.enable_model_cpu_offload()
+        self.pipeline = self.pipeline.to("cuda")
 
         prompt = "optimus is cleaning the house with broomstick"
         out = self.pipeline(
@@ -271,7 +274,7 @@ def test_flux_xlabs(self):
         self.pipeline.load_lora_weights("XLabs-AI/flux-lora-collection", weight_name="disney_lora.safetensors")
         self.pipeline.fuse_lora()
         self.pipeline.unload_lora_weights()
-        self.pipeline.enable_model_cpu_offload()
+        self.pipeline = self.pipeline.to("cuda")
 
         prompt = "A blue jay standing on a large basket of rainbow macarons, disney style"
 

tests/lora/test_lora_layers_sd3.py

@@ -176,7 +176,9 @@ def get_inputs(self, device, seed=0):
     def test_sd3_img2img_lora(self):
         pipe = self.pipeline_class.from_pretrained(self.repo_id, torch_dtype=torch.float16)
         pipe.load_lora_weights("nerijs/pixel-art-xl", weight_name="pixel-art-xl.safetensors")
-        pipe.enable_sequential_cpu_offload()
+        pipe.fuse_lora()
+        pipe.unload_lora_weights()
+        pipe = pipe.to("cuda")
 
         inputs = self.get_inputs(torch_device)
 

tests/pipelines/flux/test_pipeline_flux.py

@@ -236,8 +236,7 @@ def get_inputs(self, device, seed=0):
     def test_flux_inference(self):
         pipe = self.pipeline_class.from_pretrained(
             self.repo_id, torch_dtype=torch.bfloat16, text_encoder=None, text_encoder_2=None
-        )
-        pipe.enable_model_cpu_offload()
+        ).to("cuda")
 
         inputs = self.get_inputs(torch_device)