Update app.py

app.py

@@ -18,7 +18,7 @@ class Demo:
 
         self.training = False
 
-        
+        self.diffuser = StableDiffuser(scheduler='DDIM', seed=42).to('cuda').eval().half()
 
         with gr.Blocks() as demo:
             self.layout()
@@ -155,7 +155,10 @@ class Demo:
         # self.diffuser = StableDiffuser(scheduler='DDIM', seed=42).to('cuda').eval().half()
         if self.training:
             return [gr.update(interactive=True, value='Train'), gr.update(value='Someone else is training... Try again soon'), None, gr.update()]
-
+        # clear the diffusers
+        del self.diffuser
+        torch.cuda.empty_cache()
+        
         if train_method == 'ESD-x':
 
             modules = ".*attn2$"
@@ -186,12 +189,13 @@ class Demo:
         model_map['Custom'] = save_path
 
         # del self.diffuser
-        # torch.cuda.empty_cache()
+        torch.cuda.empty_cache()
+        self.diffuser = StableDiffuser(scheduler='DDIM', seed=42).to('cuda').eval().half()
         return [gr.update(interactive=True, value='Train'), gr.update(value='Done Training'), save_path, gr.Dropdown.update(choices=list(model_map.keys()), value='Custom')]
 
 
     def inference(self, prompt, seed, model_name, pbar = gr.Progress(track_tqdm=True)):
-        self.diffuser = StableDiffuser(scheduler='DDIM', seed=42).to('cuda').eval().half()
+        
         self.diffuser._seed = seed or 42
 
         model_path = model_map[model_name]
@@ -222,7 +226,7 @@ class Demo:
 
         edited_image = images[0][0]
 
-        del self.finetuner, self.diffuser
+        del self.finetuner
         torch.cuda.empty_cache()
 
         return edited_image, orig_image