update CPU offload

app.py

@@ -105,7 +105,7 @@ def prepare_pipeline(model_version, enable_realism, enable_anti_blur):
     return pipeline
 
 
-@spaces.GPU
+@spaces.GPU(duration=120)
 def generate_image(
     input_image, 
     control_image, 
@@ -238,7 +238,7 @@ with gr.Blocks() as demo:
         inputs=[ui_id_image, ui_control_image, ui_prompt_text, ui_seed, ui_enable_realism, ui_enable_anti_blur, ui_model_version],
         outputs=[image_output],
         fn=generate_examples,
-        cache_examples=True,
+        cache_examples=True
     )
 
     ui_btn_generate.click(

pipelines/pipeline_flux_infusenet.py

@@ -261,6 +261,9 @@ class FluxInfuseNetPipeline(FluxControlNetPipeline):
             images.
         """
 
+        # CPU offload controlnet
+        self.controlnet.cpu()
+
         height = height or self.default_sample_size * self.vae_scale_factor
         width = width or self.default_sample_size * self.vae_scale_factor
 
@@ -487,6 +490,11 @@ class FluxInfuseNetPipeline(FluxControlNetPipeline):
             ]
             controlnet_keep.append(keeps[0] if isinstance(self.controlnet, FluxControlNetModel) else keeps)
 
+        # CPU offload T5, move back controlnet to GPU
+        self.text_encoder_2.cpu()
+        torch.cuda.empty_cache()
+        self.controlnet.to(device)
+
         # 7. Denoising loop
         with self.progress_bar(total=num_inference_steps) as progress_bar:
             for i, t in enumerate(timesteps):
@@ -591,6 +599,11 @@ class FluxInfuseNetPipeline(FluxControlNetPipeline):
 
                 if XLA_AVAILABLE:
                     xm.mark_step()
+        
+        # CPU offload controlnet, move back T5 to GPU
+        self.controlnet.cpu()
+        torch.cuda.empty_cache()
+        self.text_encoder_2.to(device)
 
         if output_type == "latent":
             image = latents

pipelines/pipeline_infu_flux.py

@@ -167,8 +167,11 @@ class InfUFluxPipeline:
                       'After that, run the code again. If you have downloaded it, please use `base_model_path` to specify the correct path.')
                 print('\nIf you are using other models, please download them to a local directory and use `base_model_path` to specify the correct path.')
                 exit()
-        # pipe.to('cuda', torch.bfloat16)
-        pipe.enable_model_cpu_offload()
+        pipe.to('cuda', torch.bfloat16)
+        # CPU offload controlnet in advance
+        pipe.controlnet.cpu()
+        torch.cuda.empty_cache()
+        # pipe.enable_model_cpu_offload()
         self.pipe = pipe
 
         # Load image proj model