Gpu mem stuff

app.py

@@ -6,6 +6,7 @@ from convertModels import convert_ldm_unet_checkpoint, create_unet_diffusers_con
 from omegaconf import OmegaConf
 from StableDiffuser import StableDiffuser
 from diffusers import UNet2DConditionModel
+import torch
 
 ckpt_path = "stable_diffusion/models/ldm/sd-v1-4-full-ema.ckpt"
 config_path = "stable_diffusion/configs/stable-diffusion/v1-inference.yaml"
@@ -18,10 +19,16 @@ class Demo:
 
         self.training = False
         self.generating = False
+        self.model_edited_sd = None
+        self.model_orig_sd = None
+
+        self.diffuser = StableDiffuser(42)
+        self.diffuser.to('cpu')
+        self.diffuser = self.diffuser.half()
 
         with gr.Blocks() as demo:
             self.layout()
-            demo.queue(concurrency_count=10).launch()
+            demo.queue(concurrency_count=1).launch()
 
     def disable(self):
         return [gr.update(interactive=False), gr.update(interactive=False)]
@@ -131,6 +138,8 @@ class Demo:
         else:
             self.training = True
 
+        self.diffuser.to('cpu')
+
         model_orig, model_edited = train_esd(prompt,
                 train_method,
                 3,
@@ -146,8 +155,16 @@ class Demo:
         original_config = OmegaConf.load(config_path)
         original_config["model"]["params"]["unet_config"]["params"]["in_channels"] = 4
         unet_config = create_unet_diffusers_config(original_config, image_size=512)
-        model_edited_sd = convert_ldm_unet_checkpoint(model_edited.state_dict(), unet_config)
-        model_orig_sd = convert_ldm_unet_checkpoint(model_orig.state_dict(), unet_config)
+        _model_edited_sd = convert_ldm_unet_checkpoint(model_edited.state_dict(), unet_config)
+        _model_orig_sd = convert_ldm_unet_checkpoint(model_orig.state_dict(), unet_config)
+        
+        model_edited_sd = {key: value.cpu() for key, value in _model_edited_sd.items()}
+        model_orig_sd = {key: value.cpu() for key, value in _model_orig_sd.items()}
+
+        del model_orig, _model_orig_sd
+        del model_edited, _model_edited_sd
+
+        torch.cuda.empty_cache()
 
         self.init_inference(model_edited_sd, model_orig_sd, unet_config)
 
@@ -155,16 +172,17 @@ class Demo:
 
     def init_inference(self, model_edited_sd, model_orig_sd, unet_config):
 
+        del self.model_edited_sd, self.model_orig_sd
+
+        torch.cuda.empty_cache()
+
         self.model_edited_sd = model_edited_sd
         self.model_orig_sd = model_orig_sd
 
-        self.diffuser = StableDiffuser(42)
-
-        self.diffuser.unet = UNet2DConditionModel(**unet_config)
         self.diffuser.to('cuda')
 
         self.training = False
-
+ 
 
     def inference(self, prompt, seed, pbar = gr.Progress(track_tqdm=True)):
 
@@ -185,6 +203,8 @@ class Demo:
 
         orig_image = images[0][0]
 
+        torch.cuda.empty_cache()
+
         self.diffuser.unet.load_state_dict(self.model_edited_sd)
 
         images = self.diffuser(
@@ -197,6 +217,8 @@ class Demo:
 
         self.generating = False
 
+        torch.cuda.empty_cache()
+
         return edited_image, orig_image