Update app.py

app.py

@@ -62,22 +62,20 @@ from ldm.models.diffusion.plms import PLMSSampler
 from open_clip import tokenizer
 import open_clip
 
-def load_model_from_config(config, ckpt, verbose=False):
-    print(f"Loading model from {ckpt}")
-    pl_sd = torch.load(ckpt, map_location="cuda")
-    sd = pl_sd["state_dict"]
-    model = instantiate_from_config(config.model)
-    m, u = model.load_state_dict(sd, strict=False)
-    if len(m) > 0 and verbose:
-        print("missing keys:")
-        print(m)
-    if len(u) > 0 and verbose:
-        print("unexpected keys:")
-        print(u)
-
-    model = model.half().cuda()
-    model.eval()
-    return model
+config = OmegaConf.load("latent-diffusion/configs/latent-diffusion/txt2img-1p4B-eval.yaml")
+print(f"Loading model from {ckpt}")
+pl_sd = torch.load(ckpt, map_location="cuda")
+sd = pl_sd["state_dict"]
+model = instantiate_from_config(config.model)
+m, u = model.load_state_dict(sd, strict=False)
+if len(m) > 0 and verbose:
+    print("missing keys:")
+    print(m)
+if len(u) > 0 and verbose:
+    print("unexpected keys:")
+    print(u)
+
+model.half().to("cuda")
 
 def load_safety_model(clip_model):
     """load the safety model"""
@@ -127,10 +125,8 @@ def is_unsafe(safety_model, embeddings, threshold=0.5):
     x = np.array([e[0] for e in nsfw_values])
     return True if x > threshold else False
 
-config = OmegaConf.load("latent-diffusion/configs/latent-diffusion/txt2img-1p4B-eval.yaml")
-model = load_model_from_config(config,model_path_e)
+# model = load_model_from_config(config,model_path_e)
 device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
-model = model.to(device)
 
 #NSFW CLIP Filter
 safety_model = load_safety_model("ViT-B/32")