add model cpu offload

app.py

@@ -105,7 +105,7 @@ def prepare_pipeline(model_version, enable_realism, enable_anti_blur):
     return pipeline
 
 
-# @spaces.GPU
+@spaces.GPU
 def generate_image(
     input_image, 
     control_image, 

pipelines/pipeline_infu_flux.py

@@ -167,7 +167,8 @@ 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.to('cuda', torch.bfloat16)
+        pipe.enable_model_cpu_offload()
         self.pipe = pipe
 
         # Load image proj model