Upload imagedream/model_zoo.py with huggingface_hub

imagedream/model_zoo.py

@@ -1,64 +1,64 @@
-""" Utiliy functions to load pre-trained models more easily """
-import os
-import pkg_resources
-from omegaconf import OmegaConf
-
-import torch
-from huggingface_hub import hf_hub_download
-
-from imagedream.ldm.util import instantiate_from_config
-
-
-PRETRAINED_MODELS = {
-    "sd-v2.1-base-4view-ipmv": {
-        "config": "sd_v2_base_ipmv.yaml",
-        "repo_id": "Peng-Wang/ImageDream",
-        "filename": "sd-v2.1-base-4view-ipmv.pt",
-    },
-    "sd-v2.1-base-4view-ipmv-local": {
-        "config": "sd_v2_base_ipmv_local.yaml",
-        "repo_id": "Peng-Wang/ImageDream",
-        "filename": "sd-v2.1-base-4view-ipmv-local.pt",
-    },
-}
-
-
-def get_config_file(config_path):
-    cfg_file = pkg_resources.resource_filename(
-        "imagedream", os.path.join("configs", config_path)
-    )
-    if not os.path.exists(cfg_file):
-        raise RuntimeError(f"Config {config_path} not available!")
-    return cfg_file
-
-
-def build_model(model_name, config_path=None, ckpt_path=None, cache_dir=None):
-    if (config_path is not None) and (ckpt_path is not None):
-        config = OmegaConf.load(config_path)
-        model = instantiate_from_config(config.model)
-        model.load_state_dict(torch.load(ckpt_path, map_location="cpu"), strict=False)
-        return model
-        
-    if not model_name in PRETRAINED_MODELS:
-        raise RuntimeError(
-            f"Model name {model_name} is not a pre-trained model. Available models are:\n- "
-            + "\n- ".join(PRETRAINED_MODELS.keys())
-        )
-    model_info = PRETRAINED_MODELS[model_name]
-
-    # Instiantiate the model
-    print(f"Loading model from config: {model_info['config']}")
-    config_file = get_config_file(model_info["config"])
-    config = OmegaConf.load(config_file)
-    model = instantiate_from_config(config.model)
-
-    # Load pre-trained checkpoint from huggingface
-    if not ckpt_path:
-        ckpt_path = hf_hub_download(
-            repo_id=model_info["repo_id"],
-            filename=model_info["filename"],
-            cache_dir=cache_dir,
-        )
-        print(f"Loading model from cache file: {ckpt_path}")
-    model.load_state_dict(torch.load(ckpt_path, map_location="cpu"), strict=False)
-    return model
+""" Utiliy functions to load pre-trained models more easily """
+import os
+import pkg_resources
+from omegaconf import OmegaConf
+
+import torch
+from huggingface_hub import hf_hub_download
+
+from imagedream.ldm.util import instantiate_from_config
+
+
+PRETRAINED_MODELS = {
+    "sd-v2.1-base-4view-ipmv": {
+        "config": "sd_v2_base_ipmv.yaml",
+        "repo_id": "Peng-Wang/ImageDream",
+        "filename": "sd-v2.1-base-4view-ipmv.pt",
+    },
+    "sd-v2.1-base-4view-ipmv-local": {
+        "config": "sd_v2_base_ipmv_local.yaml",
+        "repo_id": "Peng-Wang/ImageDream",
+        "filename": "sd-v2.1-base-4view-ipmv-local.pt",
+    },
+}
+
+
+def get_config_file(config_path):
+    cfg_file = pkg_resources.resource_filename(
+        "imagedream", os.path.join("configs", config_path)
+    )
+    if not os.path.exists(cfg_file):
+        raise RuntimeError(f"Config {config_path} not available!")
+    return cfg_file
+
+
+def build_model(model_name, config_path=None, ckpt_path=None, cache_dir=None):
+    if (config_path is not None) and (ckpt_path is not None):
+        config = OmegaConf.load(config_path)
+        model = instantiate_from_config(config.model)
+        model.load_state_dict(torch.load(ckpt_path, map_location="cpu"), strict=False)
+        return model
+        
+    if not model_name in PRETRAINED_MODELS:
+        raise RuntimeError(
+            f"Model name {model_name} is not a pre-trained model. Available models are:\n- "
+            + "\n- ".join(PRETRAINED_MODELS.keys())
+        )
+    model_info = PRETRAINED_MODELS[model_name]
+
+    # Instiantiate the model
+    print(f"Loading model from config: {model_info['config']}")
+    config_file = get_config_file(model_info["config"])
+    config = OmegaConf.load(config_file)
+    model = instantiate_from_config(config.model)
+
+    # Load pre-trained checkpoint from huggingface
+    if not ckpt_path:
+        ckpt_path = hf_hub_download(
+            repo_id=model_info["repo_id"],
+            filename=model_info["filename"],
+            cache_dir=cache_dir,
+        )
+        print(f"Loading model from cache file: {ckpt_path}")
+    model.load_state_dict(torch.load(ckpt_path, map_location="cpu"), strict=False)
+    return model