Load only our own weights and include necessary config and tokenizer files

cog.yaml

@@ -0,0 +1,21 @@
+build:
+  python_version: "3.8"
+
+  gpu: true
+
+  python_packages:
+    - 'huggingface_hub==0.8.1'
+    - 'torch==1.7.1'
+    - 'torchvision==0.8.2'
+    - 'tokenizers==0.10.2'
+    - 'pyflakes==2.2.0'
+    - 'tqdm==4.46.0'
+    - 'pytorch-lightning==1.5'
+    - 'einops==0.4.1'
+    - 'omegaconf==2.2.2'
+    - 'git+https://github.com/openai/CLIP.git'
+    - 'numpy==1.23.1'
+    - 'pillow==9.2.0'
+    - 'python-dotenv==0.20.0'
+
+predict: "predict.py:Predictor"

dalle/models/__init__.py

@@ -44,25 +44,15 @@ class Dalle(nn.Module):
     @classmethod
     def from_pretrained(cls,
                         path: str) -> nn.Module:
-        path = _MODELS[path] if path in _MODELS else path
-        path = utils.realpath_url_or_path(path, root=os.path.expanduser("~/.cache/minDALL-E"))
-
         config_base = get_base_config()
-        config_new = OmegaConf.load(os.path.join(path, 'config.yaml'))
+        config_new = OmegaConf.load('config.yaml')
         config_update = OmegaConf.merge(config_base, config_new)
 
         model = cls(config_update)
-        model.tokenizer = build_tokenizer(os.path.join(path, 'tokenizer'),
+        model.tokenizer = build_tokenizer('tokenizer',
                                           context_length=model.config_dataset.context_length,
                                           lowercase=True,
                                           dropout=None)
-        logging.info("Loading first stage")
-        #model.stage1.from_ckpt('last.ckpt')
-        model.stage1.from_ckpt(os.path.join(path, 'stage1_last.ckpt'))
-        logging.info("Loading second stage")
-        model.stage2.from_ckpt(os.path.join(path, 'stage2_last.ckpt'))
-        #model.stage2.from_ckpt(os.path.join(path, 'dalle_last.ckpt'))
-        #model.stage2.from_ckpt('dalle_last.ckpt')
         return model
 
     @torch.no_grad()

server.py

@@ -25,50 +25,9 @@ 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")
-# filename = os.path.basename(url)
-# pathname = filename[: -len(".tar.gz")]
-# download_target = os.path.join(root, filename)
-# result_path = os.path.join(root, pathname)
-# if not os.path.exists(result_path):
-#     result_path = download(url, root)
-
 
 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)
-
-
-# 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)
+model = Dalle.from_pretrained("minDALL-E/1.3B")
 
 # NEW METHOD
 model.load_state_dict(torch.load(full_dict_path, map_location=torch.device('cpu')))