Move instantiate_from_config into module

Author: bfirsh

main.py

@@ -1,4 +1,4 @@
-import argparse, os, sys, datetime, glob, importlib
+import argparse, os, sys, datetime, glob
 from omegaconf import OmegaConf
 import numpy as np
 from PIL import Image
@@ -10,14 +10,7 @@
 from pytorch_lightning.trainer import Trainer
 from pytorch_lightning.callbacks import ModelCheckpoint, Callback, LearningRateMonitor
 from pytorch_lightning.utilities.distributed import rank_zero_only
-
-def get_obj_from_str(string, reload=False):
-    module, cls = string.rsplit(".", 1)
-    if reload:
-        module_imp = importlib.import_module(module)
-        importlib.reload(module_imp)
-    return getattr(importlib.import_module(module, package=None), cls)
-
+from taming.util import instantiate_from_config
 
 def get_parser(**parser_kwargs):
     def str2bool(v):
@@ -110,12 +103,6 @@ def nondefault_trainer_args(opt):
     return sorted(k for k in vars(args) if getattr(opt, k) != getattr(args, k))
 
 
-def instantiate_from_config(config):
-    if not "target" in config:
-        raise KeyError("Expected key `target` to instantiate.")
-    return get_obj_from_str(config["target"])(**config.get("params", dict()))
-
-
 class WrappedDataset(Dataset):
     """Wraps an arbitrary object with __len__ and __getitem__ into a pytorch dataset"""
     def __init__(self, dataset):

scripts/make_samples.py

@@ -3,7 +3,7 @@
 import numpy as np
 from omegaconf import OmegaConf
 from PIL import Image
-from main import instantiate_from_config, DataModuleFromConfig
+from taming.util import instantiate_from_config
 from torch.utils.data import DataLoader
 from torch.utils.data.dataloader import default_collate
 from tqdm import trange

scripts/sample_conditional.py

@@ -5,7 +5,7 @@
 import streamlit as st
 from streamlit import caching
 from PIL import Image
-from main import instantiate_from_config, DataModuleFromConfig
+from taming.util import instantiate_from_config
 from torch.utils.data import DataLoader
 from torch.utils.data.dataloader import default_collate
 

scripts/sample_fast.py

@@ -7,8 +7,8 @@
 from tqdm import tqdm, trange
 from einops import repeat
 
-from main import instantiate_from_config
 from taming.modules.transformer.mingpt import sample_with_past
+from taming.util import instantiate_from_config
 
 
 rescale = lambda x: (x + 1.) / 2.

taming/models/cond_transformer.py

@@ -3,8 +3,8 @@
 import torch.nn.functional as F
 import pytorch_lightning as pl
 
-from main import instantiate_from_config
 from taming.modules.util import SOSProvider
+from taming.util import instantiate_from_config
 
 
 def disabled_train(self, mode=True):

taming/models/vqgan.py

@@ -2,12 +2,11 @@
 import torch.nn.functional as F
 import pytorch_lightning as pl
 
-from main import instantiate_from_config
-
 from taming.modules.diffusionmodules.model import Encoder, Decoder
 from taming.modules.vqvae.quantize import VectorQuantizer2 as VectorQuantizer
 from taming.modules.vqvae.quantize import GumbelQuantize
-from taming.modules.vqvae.quantize import EMAVectorQuantizer
+from taming.util import instantiate_from_config
+
 
 class VQModel(pl.LightningModule):
     def __init__(self,
@@ -401,4 +400,4 @@ def configure_optimizers(self):
                                   lr=lr, betas=(0.5, 0.9))
         opt_disc = torch.optim.Adam(self.loss.discriminator.parameters(),
                                     lr=lr, betas=(0.5, 0.9))
-        return [opt_ae, opt_disc], []                                           
+        return [opt_ae, opt_disc], []                                           

taming/util.py

@@ -1,3 +1,4 @@
+import importlib
 import os, hashlib
 import requests
 from tqdm import tqdm
@@ -142,6 +143,20 @@ def retrieve(
         return list_or_dict, success
 
 
+def get_obj_from_str(string, reload=False):
+    module, cls = string.rsplit(".", 1)
+    if reload:
+        module_imp = importlib.import_module(module)
+        importlib.reload(module_imp)
+    return getattr(importlib.import_module(module, package=None), cls)
+
+
+def instantiate_from_config(config):
+    if not "target" in config:
+        raise KeyError("Expected key `target` to instantiate.")
+    return get_obj_from_str(config["target"])(**config.get("params", dict()))
+
+
 if __name__ == "__main__":
     config = {"keya": "a",
               "keyb": "b",
@@ -155,3 +170,4 @@ def retrieve(
     print(config)
     retrieve(config, "keya")
 
+