Update app.py

app.py

@@ -127,15 +127,17 @@ class process:
             ckpt = os.path.join("models", opt.base_model)
             pl_sd = read_state_dict(ckpt)
             if "state_dict" in pl_sd:
-                st = pl_sd["state_dict"]
+                pl_sd = pl_sd["state_dict"]
             else:
-                st = pl_sd
-            self.sd_model.load_state_dict(st, strict=False)
+                pl_sd = pl_sd
+            self.sd_model.load_state_dict(pl_sd, strict=False)
+            del pl_sd
             self.base_model = opt.base_model
             if self.base_model!='v1-5-pruned-emaonly.ckpt' and self.base_model!='sd-v1-4.ckpt':
                 vae_sd = torch.load(os.path.join('models', 'anything-v4.0.vae.pt'), map_location="cuda")
                 st = vae_sd["state_dict"]
                 self.sd_model.first_stage_model.load_state_dict(st, strict=False)
+                del st
 
         with torch.inference_mode(), \
                 self.sd_model.ema_scope(), \
@@ -304,7 +306,7 @@ with gr.Blocks(css='style.css') as demo:
                         label=f"Input type for drawing",
                         interactive=True,
                         value="Nothing")
-                    im1 = gr.Image(source='canvas', tool='color-sketch', label='Pay attention to adjusting stylus thickness!', visible=False) 
+                    im1 = gr.Image(source='canvas', tool='color-sketch', label='Pay attention to adjusting stylus thickness!', visible=False)
                     im2 = im1
                     cond_weight = gr.Slider(
                         label="Condition weight",