Re-enable full loading

dalle/models/__init__.py

@@ -57,10 +57,12 @@ class Dalle(nn.Module):
                                           lowercase=True,
                                           dropout=None)
         logging.info("Loading first stage")
-        model.stage1.from_ckpt('last.ckpt')
+        #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')
+        #model.stage2.from_ckpt('dalle_last.ckpt')
         return model
 
     @torch.no_grad()

server.py

@@ -34,30 +34,30 @@ 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_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)
 
 # model_clip, preprocess_clip = clip.load("ViT-B/32", device=device)
 # model_clip.to(device=device)