Update app.py

app.py

@@ -48,7 +48,7 @@ class Demo:
         self.generating = False
         self.device = 'cuda'
         self.weight_dtype = torch.float16
-        self.pipe = StableDiffusionXLPipeline.from_pretrained('stabilityai/stable-diffusion-xl-base-1.0', torch_dtype=self.weight_dtype).to(self.device)
+        self.pipe = StableDiffusionXLPipeline.from_pretrained('stabilityai/stable-diffusion-xl-base-1.0', torch_dtype=self.weight_dtype).to('cpu')
         self.pipe.enable_xformers_memory_efficient_attention()
         with gr.Blocks() as demo:
             self.layout()
@@ -233,6 +233,7 @@ class Demo:
 #             positive_prompt = ''
 #         if negative_prompt is None:
 #             negative_prompt = ''
+
         if attributes_input == '':
             attributes_input = None
         print(target_concept, positive_prompt, negative_prompt, attributes_input, is_person)
@@ -267,7 +268,7 @@ class Demo:
         generator = torch.manual_seed(seed)
 
         model_path = model_map[model_name]
-        
+        self.pipe = self.pipe.to(self.device)
         unet = self.pipe.unet
         network_type = "c3lier"
         if 'full' in model_path:
@@ -310,7 +311,7 @@ class Demo:
         del unet, network
         unet = None
         network = None
-        pipe = None
+        self.pipe = self.pipe.to('cpu')
         torch.cuda.empty_cache()
         
         return edited_image, original_image