updated

.gitignore

@@ -1 +1,7 @@
-venv/
+venv/
+AdaBins/
+k-diffusion/
+MiDaS/
+models/*
+pytorch3d-lite/
+stable-diffusion/

requirements.txt

@@ -1,11 +1,11 @@
 --extra-index-url https://download.pytorch.org/whl/cu113
 torch==1.12.1+cu113 
 torchvision==0.13.1+cu113
+torchtext==0.13.1
 omegaconf==2.2.3
 einops==0.4.1
 pytorch-lightning==1.7.4
 torchmetrics==0.9.3
-torchtext==0.13.1
 transformers==4.21.2
 kornia==0.6.7
 git+https://github.com/CompVis/taming-transformers.git@master#egg=taming-transformers
@@ -17,4 +17,5 @@ matplotlib
 resize-right
 timm
 torchdiffeq
-#opencv-python
+opencv-python
+scikit-image

sd.py

@@ -16,12 +16,17 @@ from torch import autocast
 
 from PIL import Image
 import os
+import requests
 
-from huggingface_hub import snapshot_download
+#from huggingface_hub import snapshot_download
 from omegaconf import OmegaConf
 
 # 1. Download stable diffusion repository and set path
 os.system("git clone https://github.com/deforum/stable-diffusion")
+os.system("git clone https://github.com/shariqfarooq123/AdaBins.git")
+os.system("git clone https://github.com/isl-org/MiDaS.git")
+os.system("git clone https://github.com/MSFTserver/pytorch3d-lite.git")
+
 os.system("git clone https://github.com/deforum/k-diffusion/")
 with open('k-diffusion/k_diffusion/__init__.py', 'w') as f:
         f.write('')
@@ -31,6 +36,9 @@ sys.path.extend([
     'src/clip',
     'stable-diffusion/',
     'k-diffusion',
+    'pytorch3d-lite',
+    'AdaBins',
+    'MiDaS',
 ])
 from helpers import sampler_fn
 from k_diffusion.external import CompVisDenoiser
@@ -45,9 +53,9 @@ def load_model_from_config(config, ckpt, verbose=False, device='cuda', half_prec
     pl_sd = torch.load(ckpt, map_location=map_location)
     if "global_step" in pl_sd:
         print(f"Global Step: {pl_sd['global_step']}")
-    sd = pl_sd["state_dict"]
+    sd_var = pl_sd["state_dict"]
     model = instantiate_from_config(config.model)
-    m, u = model.load_state_dict(sd, strict=False)
+    m, u = model.load_state_dict(sd_var, strict=False)
     if len(m) > 0 and verbose:
         print("missing keys:")
         print(m)
@@ -62,16 +70,31 @@ def load_model_from_config(config, ckpt, verbose=False, device='cuda', half_prec
     model.eval()
     return model
 
+print('Model loading...')
+
+models_path = "models"
+model_checkpoint =  "sd-v1-4.ckpt"
+ckpt_path = os.path.join(models_path, model_checkpoint)
+#ckpt_path = snapshot_download(repo_id="CompVis/stable-diffusion-v-1-4-original", filename="sd-v1-4.ckpt")
+
 ckpt_config_path = "./stable-diffusion/configs/stable-diffusion/v1-inference.yaml"
-ckpt_path = snapshot_download(repo_id="CompVis/stable-diffusion-v-1-4-original", filename="sd-v1-4.ckpt")
+
+if os.path.exists(ckpt_path):
+    print(f"{ckpt_path} exists")
+else:   
+    print(f"Attempting to download {model_checkpoint}...this may take a while")
+    url = 'https://kael558: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@huggingface.co/CompVis/stable-diffusion-v-1-4-original/resolve/main/sd-v1-4.ckpt'
+    ckpt_request = requests.get(url)
+
+    print('Model downloaded.')
+    with open(os.path.join(models_path, model_checkpoint), 'wb') as model_file:
+        model_file.write(ckpt_request.content)
 
 local_config = OmegaConf.load(f"{ckpt_config_path}")
 model = load_model_from_config(local_config, f"{ckpt_path}")
 device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
 model = model.to(device)
-
-
-
+print('Model saved.')
 
 class DeformAnimKeys():
     def __init__(self, anim_args):