keep pipe_inference in device except for the inversion phase

app.py

@@ -37,7 +37,7 @@ scheduler_class = MyEulerAncestralDiscreteScheduler
 
 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)
@@ -109,6 +109,7 @@ with gr.Blocks(css="style.css") as demo:
                            edit_guidance_scale=edit_guidance_scale,
                            inversion_max_step=inversion_max_step)
         if device == 'cuda':
+            pipe_inference.to('cpu')
             torch.cuda.empty_cache()
             # pipe_inversion = pipe_inversion.to(device)
             # if image_editor is not None:
@@ -122,6 +123,7 @@ with gr.Blocks(css="style.css") as demo:
         if device == 'cuda':
             pipe_inversion.to('cpu')
             torch.cuda.empty_cache()
+            pipe_inference.to(device)
             # pipe_inversion = pipe_inversion.to('cpu')
         print(f"#### 3 #### pipe_inversion.device: {pipe_inversion.device}")
         print(f"#### 4 #### pipe_inference.device: {pipe_inference.device}")
@@ -149,18 +151,18 @@ with gr.Blocks(css="style.css") as demo:
         #       image = editor.to(device).edit(target_prompt)
         #  else:
 
-        if device == 'cuda':
-            torch.cuda.empty_cache()
+        # if device == 'cuda':
+        #     torch.cuda.empty_cache()
 
         print(f"#### 5 #### pipe_inversion.device: {pipe_inversion.device}")
         print(f"#### 6 #### pipe_inference.device: {pipe_inference.device}")
 
-        image = ImageEditorDemo.edit(pipe_inference.to(device), target_prompt, inversion_state['latent'], inversion_state['noise'],
+        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.to('cpu')
-            torch.cuda.empty_cache()
+        # if device == 'cuda':
+            # pipe_inference.to('cpu')
+            # torch.cuda.empty_cache()
 
 
         print(f"#### 7 #### pipe_inversion.device: {pipe_inversion.device}")