Load complete ckpt from hub with secret

app.py

@@ -9,11 +9,8 @@ from PIL import Image
 
 import requests
 import logging
-from huggingface_hub import hf_hub_download
 
-logging.info("Start downloading")
-hf_hub_download(repo_id="MatthiasC/dall-e-logo", filename="README.md")
-logging.info("End downloading")
+
 def start_server():
     os.system("uvicorn server:app --port 8080 --host 0.0.0.0 --workers 1")
 

server.py

@@ -13,11 +13,20 @@ from PIL import Image
 #import clip
 from dalle.models import Dalle
 import logging
+import streamlit as st
 from dalle.utils.utils import clip_score, download
 
 print("Loading models...")
 app = FastAPI()
 
+from huggingface_hub import hf_hub_download
+
+logging.info("Start downloading")
+full_dict_path = hf_hub_download(repo_id="MatthiasC/dall-e-logo", filename="full_dict_new.ckpt",
+                                 use_auth_token=st.secrets["model_hub"])
+logging.info("End downloading")
+logging.info(full_dict_path)
+
 
 # url = "https://arena.kakaocdn.net/brainrepo/models/minDALL-E/57b008f02ceaa02b779c8b7463143315/1.3B.tar.gz"
 # root = os.path.expanduser("~/.cache/minDALLE")
@@ -31,33 +40,39 @@ app = FastAPI()
 
 device = "cuda" if torch.cuda.is_available() else "cpu"
 model = Dalle.from_pretrained("minDALL-E/1.3B")  # This will automatically download the pretrained model.
-model.to(device=device)
+#model.to(device=device)
 
+
+# OLD CODE
 # -----------------------------------------------------------
-state_dict_ = torch.load('last.ckpt', map_location='cpu')
-vqgan_stage_dict = model.stage1.state_dict()
-
-for name, param in state_dict_['state_dict'].items():
-    if name not in model.stage1.state_dict().keys():
-        continue
-    if isinstance(param, nn.parameter.Parameter):
-        param = param.data
-    vqgan_stage_dict[name].copy_(param)
-
-model.stage1.load_state_dict(vqgan_stage_dict)
-#---------------------------------------------------------
-state_dict_dalle = torch.load('dalle_last.ckpt', map_location='cpu')
-dalle_stage_dict = model.stage2.state_dict()
-
-for name, param in state_dict_dalle['state_dict'].items():
-    if name[6:] not in model.stage2.state_dict().keys():
-        print(name)
-        continue
-    if isinstance(param, nn.parameter.Parameter):
-        param = param.data
-    dalle_stage_dict[name[6:]].copy_(param)
-
-model.stage2.load_state_dict(dalle_stage_dict)
+# state_dict_ = torch.load('last.ckpt', map_location='cpu')
+# vqgan_stage_dict = model.stage1.state_dict()
+#
+# for name, param in state_dict_['state_dict'].items():
+#     if name not in model.stage1.state_dict().keys():
+#         continue
+#     if isinstance(param, nn.parameter.Parameter):
+#         param = param.data
+#     vqgan_stage_dict[name].copy_(param)
+#
+# model.stage1.load_state_dict(vqgan_stage_dict)
+# #---------------------------------------------------------
+# state_dict_dalle = torch.load('dalle_last.ckpt', map_location='cpu')
+# dalle_stage_dict = model.stage2.state_dict()
+#
+# for name, param in state_dict_dalle['state_dict'].items():
+#     if name[6:] not in model.stage2.state_dict().keys():
+#         print(name)
+#         continue
+#     if isinstance(param, nn.parameter.Parameter):
+#         param = param.data
+#     dalle_stage_dict[name[6:]].copy_(param)
+#
+# model.stage2.load_state_dict(dalle_stage_dict)
+
+# NEW METHOD
+model.load_state_dict(torch.load(full_dict_path))
+model.to(device=device)
 
 # model_clip, preprocess_clip = clip.load("ViT-B/32", device=device)
 # model_clip.to(device=device)