keep pipes in cpu until using them

app.py

@@ -35,9 +35,9 @@ if device == "cuda":
 
 scheduler_class = MyEulerAncestralDiscreteScheduler
 
-pipe_inversion = SDXLDDIMPipeline.from_pretrained("stabilityai/sdxl-turbo", use_safetensors=True).to(device)
+pipe_inversion = SDXLDDIMPipeline.from_pretrained("stabilityai/sdxl-turbo", use_safetensors=True)#.to(device)
 pipe_inference = AutoPipelineForImage2Image.from_pretrained("stabilityai/sdxl-turbo",
-                                                            use_safetensors=True).to(device)
+                                                            use_safetensors=True)#.to(device)
 pipe_inference.scheduler = scheduler_class.from_config(pipe_inference.scheduler.config)
 pipe_inversion.scheduler = scheduler_class.from_config(pipe_inversion.scheduler.config)
 pipe_inversion.scheduler_inference = scheduler_class.from_config(pipe_inference.scheduler.config)
@@ -92,11 +92,6 @@ with gr.Blocks(css="style.css") as demo:
     # @spaces.GPU
     def set_pipe(inversion_state, input_image, description_prompt, edit_guidance_scale, num_inference_steps=4,
                  num_inversion_steps=4, inversion_max_step=0.6, rnri_iterations=2, rnri_alpha=0.1, rnri_lr=0.2):
-        if device == 'cuda':
-            # if image_editor is not None:
-            #     image_editor = image_editor.to('cpu')
-
-            torch.cuda.empty_cache()
 
         if input_image is None or not description_prompt:
             return None, "Please set all inputs."
@@ -113,8 +108,17 @@ with gr.Blocks(css="style.css") as demo:
                            num_inversion_steps=num_inversion_steps,
                            edit_guidance_scale=edit_guidance_scale,
                            inversion_max_step=inversion_max_step)
+        if device == 'cuda':
+            torch.cuda.empty_cache()
+            pipe_inversion = pipe_inversion.to(device)
+            # if image_editor is not None:
+            #     image_editor = image_editor.to('cpu')
+
         inversion_state = ImageEditorDemo.invert(pipe_inversion, input_image, description_prompt, config,
                                                  [rnri_iterations, rnri_alpha, rnri_lr], device)
+        if device == 'cuda':
+            torch.cuda.empty_cache()
+            pipe_inversion = pipe_inversion.to('cpu')
         return inversion_state, "Input has set!"
 
 
@@ -122,6 +126,7 @@ with gr.Blocks(css="style.css") as demo:
     def edit1(editor, target_prompt):
         if editor is None:
             raise gr.Error("Set inputs before editing.")
+
         # if device == "cuda":
         #       image = editor.to(device).edit(target_prompt)
         #  else:
@@ -136,8 +141,16 @@ with gr.Blocks(css="style.css") as demo:
         # if device == "cuda":
         #       image = editor.to(device).edit(target_prompt)
         #  else:
+
+        if device == 'cuda':
+            torch.cuda.empty_cache()
+            pipe_inference = pipe_inference.to(device)
         image = ImageEditorDemo.edit(pipe_inference, target_prompt, inversion_state['latent'], inversion_state['noise'],
                                      inversion_state['cfg'], inversion_state['cfg'].edit_guidance_scale)
+
+        if device == 'cuda':
+            pipe_inference = pipe_inference.to('cpu')
+            torch.cuda.empty_cache()
         return image