updated model/_init_.py

dalle/models/__init__.py

@@ -43,20 +43,24 @@ 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"))
+        #path = _MODELS[path] if path in _MODELS else path
+        #path = utils.realpath_url_or_path(path, root=os.path.expanduser(".cache/minDALL-E"))
+        path = ''
 
         config_base = get_base_config()
-        config_new = OmegaConf.load(os.path.join(path, 'config.yaml'))
+        config_new = OmegaConf.load(os.path.join(path, '.cache/minDALL-E/1.3B/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('.cache/minDALL-E/1.3B/tokenizer',
                                           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'))
+        model.stage1.from_ckpt('.cache/minDALL-E/1.3B/stage1_last.ckpt')
+        model.stage2.from_ckpt('.cache/minDALL-E/1.3B/stage2_last.ckpt')
+        #model.stage1.from_ckpt('https://utexas.box.com/shared/static/rpt9miyj2kikogyekpqnkd6y115xp51i.ckpt')
+        #model.stage2.from_ckpt('https://utexas.box.com/shared/static/54jc9fw0bious5nx6wvayeqaskcrdgv4.ckpt')
+        
         return model
 
     @torch.no_grad()
@@ -199,4 +203,4 @@ class ImageGPT(pl.LightningModule):
         self.log("lr", self.lr_schedulers().get_last_lr()[0], on_step=True, on_epoch=False, prog_bar=True, logger=True)
 
     def on_epoch_start(self):
-        self.stage1.eval()
+        self.stage1.eval()