Add usage for both ckpts

dalle/models/__init__.py

@@ -6,6 +6,7 @@
 
 import os
 import torch
+import logging
 import torch.nn as nn
 import pytorch_lightning as pl
 from typing import Optional, Tuple
@@ -55,8 +56,11 @@ class Dalle(nn.Module):
                                           context_length=model.config_dataset.context_length,
                                           lowercase=True,
                                           dropout=None)
-        model.stage1.from_ckpt(os.path.join(path, 'stage1_last.ckpt'))
-        model.stage2.from_ckpt(os.path.join(path, 'stage2_last.ckpt'))
+        logging.info("Loading first stage")
+        model.stage1.from_ckpt('last.ckpt')
+        logging.info("Loading second stage")
+        #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

@@ -34,17 +34,17 @@ model = Dalle.from_pretrained("minDALL-E/1.3B")  # This will automatically downl
 model.to(device=device)
 
 # -----------------------------------------------------------
-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_ = 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()