back (#19)

- Update app.py (f6532a09ea9fe76356273dc20ccd2b1b1010e969)
- Update requirements.txt (485f8955612562c987019e40ac347e02cb2948ca)

app.py

@@ -1,38 +1,3 @@
-# Add this code at the very beginning of your script, before any other imports
-
-import sys
-import types
-
-# Create a fake torch._six module with string_classes
-torch_six = types.ModuleType('torch._six')
-torch_six.string_classes = (str,)  # In Python 3, string_classes is just (str,)
-
-# Create torch module if it doesn't exist in sys.modules
-if 'torch' not in sys.modules:
-    import torch
-    
-# Add the _six submodule to torch
-sys.modules['torch._six'] = torch_six
-torch._six = torch_six
-
-# Monkey patch for pytorch_lightning.utilities.distributed
-try:
-    from pytorch_lightning.utilities.rank_zero import rank_zero_only
-    # Create the old module path
-    pl_utils_dist = types.ModuleType('pytorch_lightning.utilities.distributed')
-    pl_utils_dist.rank_zero_only = rank_zero_only
-    sys.modules['pytorch_lightning.utilities.distributed'] = pl_utils_dist
-except ImportError:
-    # If even the new import fails, create a dummy decorator
-    def rank_zero_only(fn):
-        """Dummy decorator that just returns the function as-is"""
-        return fn
-    
-    pl_utils_dist = types.ModuleType('pytorch_lightning.utilities.distributed')
-    pl_utils_dist.rank_zero_only = rank_zero_only
-    sys.modules['pytorch_lightning.utilities.distributed'] = pl_utils_dist
-
-# Now continue with your original imports
 from pydoc import describe
 import gradio as gr
 import torch
@@ -46,9 +11,9 @@ from ldm.util import instantiate_from_config
 from huggingface_hub import hf_hub_download
 import spaces
 
-# Rest of your code continues here...
 model_path_e = hf_hub_download(repo_id="multimodalart/compvis-latent-diffusion-text2img-large", filename="txt2img-f8-large.ckpt")
 
+#@title Import stuff
 import argparse, os, sys, glob
 import numpy as np
 from PIL import Image
@@ -62,14 +27,22 @@ from ldm.models.diffusion.plms import PLMSSampler
 from open_clip import tokenizer
 import open_clip
 
-config = OmegaConf.load("latent-diffusion/configs/latent-diffusion/txt2img-1p4B-eval.yaml")
-
-pl_sd = torch.load(model_path_e, map_location="cuda")
-sd = pl_sd["state_dict"]
-model = instantiate_from_config(config.model)
-m, u = model.load_state_dict(sd, strict=False)
-
-model.half().to("cuda")
+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
 
 def load_safety_model(clip_model):
     """load the safety model"""
@@ -119,8 +92,10 @@ 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
 
-# model = load_model_from_config(config,model_path_e)
+config = OmegaConf.load("latent-diffusion/configs/latent-diffusion/txt2img-1p4B-eval.yaml")
+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")

requirements.txt

@@ -3,11 +3,12 @@ ftfy
 regex 
 tqdm 
 omegaconf 
-pytorch-lightning
+pytorch-lightning==1.7.7
 torch-fidelity
 transformers   
 einops
 gradio
+torch==1.13.1
 open_clip_torch
 numpy
 tqdm