feat: enhance Auto Class to support remote and local path

.gitignore

@@ -170,3 +170,5 @@ cython_debug/
 
 *.mp4
 *.mkv
+
+/cccv/cache_models/

README.md

@@ -20,7 +20,7 @@ pip install cccv
 
 ### Start
 
-#### cv2
+#### Load a registered model in cccv
 
 a simple example to use the SISR (Single Image Super-Resolution) model to process an image
 
@@ -37,6 +37,22 @@ img = model.inference_image(img)
 cv2.imwrite("test_out.jpg", img)
 ```
 
+#### Load a custom model from remote repository or local path
+
+a simple example to use [remote repository](https://github.com/EutropicAI/cccv_demo_remote_model), auto register then load
+
+```python
+import cv2
+import numpy as np
+
+from cccv import AutoModel, SRBaseModel
+
+# remote repo
+model: SRBaseModel = AutoModel.from_pretrained("https://github.com/EutropicAI/cccv_demo_remote_model")
+# local path
+model: SRBaseModel = AutoModel.from_pretrained("/path/to/cccv_demo_model")
+```
+
 #### VapourSynth
 
 a simple example to use the VapourSynth to process a video

cccv/__init__.py

@@ -26,6 +26,6 @@
 
 from cccv.arch import ARCH_REGISTRY
 from cccv.auto import AutoConfig, AutoModel
-from cccv.config import CONFIG_REGISTRY, BaseConfig, SRBaseConfig, VFIBaseConfig, VSRBaseConfig
+from cccv.config import CONFIG_REGISTRY, AutoBaseConfig, BaseConfig, SRBaseConfig, VFIBaseConfig, VSRBaseConfig
 from cccv.model import MODEL_REGISTRY, AuxiliaryBaseModel, CCBaseModel, SRBaseModel, VFIBaseModel, VSRBaseModel
 from cccv.type import ArchType, BaseModelInterface, ConfigType, ModelType

cccv/auto/config.py

@@ -1,32 +1,81 @@
+import importlib.util
+import json
+from pathlib import Path
 from typing import Any, Optional, Union
 
-from cccv.config import CONFIG_REGISTRY, BaseConfig
+from cccv.config import CONFIG_REGISTRY, AutoBaseConfig
 from cccv.type import ConfigType
+from cccv.util.remote import git_clone
 
 
 class AutoConfig:
     @staticmethod
     def from_pretrained(
-        pretrained_model_name: Union[ConfigType, str],
+        pretrained_model_name_or_path: Union[ConfigType, str, Path],
+        *,
+        model_dir: Optional[Union[Path, str]] = None,
         **kwargs: Any,
     ) -> Any:
         """
-        Get a config instance of a pretrained model configuration.
+        Get a config instance of a pretrained model configuration, can be a registered config name or a local path or a git url.
 
-        :param pretrained_model_name: The name of the pretrained model configuration
+        :param pretrained_model_name_or_path: The name or path of the pretrained model configuration
+        :param model_dir: The path to cache the downloaded model configuration. Should be a full path. If None, use default cache path.
         :return:
         """
-        return CONFIG_REGISTRY.get(pretrained_model_name)
+        if "pretrained_model_name" in kwargs:
+            print(
+                "[CCCV] warning: 'pretrained_model_name' is deprecated, please use 'pretrained_model_name_or_path' instead."
+            )
+            pretrained_model_name_or_path = kwargs.pop("pretrained_model_name")
 
-    @staticmethod
-    def register(config: Union[BaseConfig, Any], name: Optional[str] = None) -> None:
-        """
-        Register the given config class instance under the name BaseConfig.name or the given name.
-        Can be used as a function call. See docstring of this class for usage.
+        # 1. check if it's a registered config name, early return if found
+        if isinstance(pretrained_model_name_or_path, ConfigType):
+            pretrained_model_name_or_path = pretrained_model_name_or_path.value
+        if str(pretrained_model_name_or_path) in CONFIG_REGISTRY:
+            return CONFIG_REGISTRY.get(str(pretrained_model_name_or_path))
 
-        :param config: The config class instance to register
-        :param name: The name to register the config class instance under. If None, use BaseConfig.name
-        :return:
-        """
-        # used as a function call
-        CONFIG_REGISTRY.register(obj=config, name=name)
+        # 2. check is a url or not, if it's a url, git clone it to model_dir then replace pretrained_model_name_or_path with the local path (Path)
+        if str(pretrained_model_name_or_path).startswith("http"):
+            pretrained_model_name_or_path = git_clone(
+                git_url=str(pretrained_model_name_or_path),
+                model_dir=model_dir,
+                **kwargs,
+            )
+
+        # 3. check if it's a real path
+        dir_path = Path(str(pretrained_model_name_or_path))
+
+        if not dir_path.exists() or not dir_path.is_dir():
+            raise ValueError(f"[CCCV] model configuration '{dir_path}' is not a valid config name or path")
+
+        # load config,json from the directory
+        config_path = dir_path / "config.json"
+        # check if config.json exists
+        if not config_path.exists():
+            raise FileNotFoundError(f"[CCCV] no valid config.json not found in {dir_path}")
+
+        with open(config_path, "r", encoding="utf-8") as f:
+            config_dict = json.load(f)
+
+        for k in ["arch", "model", "name"]:
+            if k not in config_dict:
+                raise KeyError(
+                    f"[CCCV] no key '{k}' in config.json in {dir_path}, you should provide a valid config.json contain a key '{k}'"
+                )
+
+        # auto import all .py files in the directory to register the arch, model and config
+        for py_file in dir_path.glob("*.py"):
+            spec = importlib.util.spec_from_file_location(py_file.stem, py_file)
+            if spec is None or spec.loader is None:
+                continue
+            module = importlib.util.module_from_spec(spec)
+            spec.loader.exec_module(module)
+
+        if "path" not in config_dict or config_dict["path"] is None or config_dict["path"] == "":
+            # add the path to the config_dict
+            config_dict["path"] = str(dir_path / config_dict["name"])
+
+        # convert config_dict to pydantic model
+        cfg = AutoBaseConfig.model_validate(config_dict)
+        return cfg

cccv/auto/model.py

@@ -1,31 +1,34 @@
+from pathlib import Path
 from typing import Any, Optional, Tuple, Union
 
 import torch
 
-from cccv.config import CONFIG_REGISTRY, BaseConfig
+from cccv.auto.config import AutoConfig
+from cccv.config import BaseConfig
 from cccv.model import MODEL_REGISTRY
 from cccv.type import ConfigType
 
 
 class AutoModel:
     @staticmethod
     def from_pretrained(
-        pretrained_model_name: Union[ConfigType, str],
+        pretrained_model_name_or_path: Union[ConfigType, str, Path],
+        *,
         device: Optional[torch.device] = None,
         fp16: bool = True,
         compile: bool = False,
         compile_backend: Optional[str] = None,
         tile: Optional[Tuple[int, int]] = (128, 128),
         tile_pad: int = 8,
         pad_img: Optional[Tuple[int, int]] = None,
-        model_dir: Optional[str] = None,
+        model_dir: Optional[Union[Path, str]] = None,
         gh_proxy: Optional[str] = None,
         **kwargs: Any,
     ) -> Any:
         """
-        Get a model instance from a pretrained model name.
+        Get a model instance from a pretrained model name or path.
 
-        :param pretrained_model_name: The name of the pretrained model. It should be registered in CONFIG_REGISTRY.
+        :param pretrained_model_name_or_path: The name or path of the pretrained model. It should be registered in CONFIG_REGISTRY.
         :param device: inference device
         :param fp16: use fp16 precision or not
         :param compile: use torch.compile or not
@@ -37,8 +40,8 @@ def from_pretrained(
         :param gh_proxy: The proxy for downloading from github release. Example: https://github.abskoop.workers.dev/
         :return:
         """
+        config = AutoConfig.from_pretrained(pretrained_model_name_or_path, model_dir=model_dir, **kwargs)
 
-        config = CONFIG_REGISTRY.get(pretrained_model_name)
         return AutoModel.from_config(
             config=config,
             device=device,
@@ -56,14 +59,15 @@ def from_pretrained(
     @staticmethod
     def from_config(
         config: Union[BaseConfig, Any],
+        *,
         device: Optional[torch.device] = None,
         fp16: bool = True,
         compile: bool = False,
         compile_backend: Optional[str] = None,
         tile: Optional[Tuple[int, int]] = (128, 128),
         tile_pad: int = 8,
         pad_img: Optional[Tuple[int, int]] = None,
-        model_dir: Optional[str] = None,
+        model_dir: Optional[Union[Path, str]] = None,
         gh_proxy: Optional[str] = None,
         **kwargs: Any,
     ) -> Any:
@@ -99,29 +103,3 @@ def from_config(
         )
 
         return model
-
-    @staticmethod
-    def register(obj: Optional[Any] = None, name: Optional[str] = None) -> Any:
-        """
-        Register the given object under the name `obj.__name__` or the given name.
-        Can be used as either a decorator or not. See docstring of this class for usage.
-
-        :param obj: The object to register. If None, this is being used as a decorator.
-        :param name: The name to register the object under. If None, use `obj.__name__`.
-        :return:
-        """
-        if obj is None:
-            # used as a decorator
-            def deco(func_or_class: Any) -> Any:
-                _name = name
-                if _name is None:
-                    _name = func_or_class.__name__
-                MODEL_REGISTRY.register(obj=func_or_class, name=_name)
-                return func_or_class
-
-            return deco
-
-        # used as a function call
-        if name is None:
-            name = obj.__name__
-        MODEL_REGISTRY.register(obj=obj, name=name)

cccv/config/__init__.py

@@ -3,7 +3,7 @@
 
 CONFIG_REGISTRY: RegistryConfigInstance = RegistryConfigInstance("CONFIG")
 
-from cccv.config.base_config import BaseConfig, SRBaseConfig, VSRBaseConfig, VFIBaseConfig
+from cccv.config.base_config import BaseConfig, SRBaseConfig, VSRBaseConfig, VFIBaseConfig, AutoBaseConfig
 
 # Auxiliary Network
 

cccv/config/base_config.py

@@ -1,6 +1,6 @@
 from typing import Optional, Union
 
-from pydantic import BaseModel, FilePath, HttpUrl
+from pydantic import BaseModel, ConfigDict, FilePath, HttpUrl
 
 from cccv.type.arch import ArchType
 from cccv.type.model import ModelType
@@ -15,6 +15,10 @@ class BaseConfig(BaseModel):
     model: Union[ModelType, str]
 
 
+class AutoBaseConfig(BaseConfig):
+    model_config = ConfigDict(extra="allow")
+
+
 class AuxiliaryBaseConfig(BaseConfig):
     pass
 

cccv/config/sr/srcnn_config.py

@@ -6,7 +6,7 @@
 
 
 class SRCNNConfig(SRBaseConfig):
-    arch: ArchType = ArchType.SRCNN
+    arch: Union[ArchType, str] = ArchType.SRCNN
     model: Union[ModelType, str] = ModelType.SRBaseModel
     scale: int = 2
     num_channels: int = 1

cccv/model/base_model.py

@@ -1,14 +1,15 @@
 import sys
 from inspect import signature
-from typing import Any, Optional, Tuple
+from pathlib import Path
+from typing import Any, Optional, Tuple, Union
 
 import torch
 
 from cccv.arch import ARCH_REGISTRY
-from cccv.cache_models import load_file_from_url
 from cccv.config import BaseConfig
 from cccv.type import BaseModelInterface
 from cccv.util.device import DEFAULT_DEVICE
+from cccv.util.remote import load_file_from_url
 
 
 class CCBaseModel(BaseModelInterface):
@@ -37,7 +38,7 @@ def __init__(
         tile: Optional[Tuple[int, int]] = (128, 128),
         tile_pad: int = 8,
         pad_img: Optional[Tuple[int, int]] = None,
-        model_dir: Optional[str] = None,
+        model_dir: Optional[Union[Path, str]] = None,
         gh_proxy: Optional[str] = None,
         **kwargs: Any,
     ) -> None:
@@ -56,7 +57,7 @@ def __init__(
         self.tile: Optional[Tuple[int, int]] = tile
         self.tile_pad: int = tile_pad
         self.pad_img: Optional[Tuple[int, int]] = pad_img
-        self.model_dir: Optional[str] = model_dir
+        self.model_dir: Optional[Union[Path, str]] = model_dir
         self.gh_proxy: Optional[str] = gh_proxy
 
         # post-hook: edit parameters here if needed
@@ -105,7 +106,7 @@ def get_state_dict(self) -> Any:
         cfg: BaseConfig = self.config
 
         if cfg.path is not None:
-            state_dict_path = str(cfg.path)
+            state_dict_path = cfg.path
         else:
             try:
                 state_dict_path = load_file_from_url(

cccv/util/registry.py

@@ -41,7 +41,7 @@ def __init__(self, name: str) -> None:
 
     def _do_register(self, name: str, obj: Any) -> None:
         if name in self._obj_map:
-            print("[CCCV] An object named '{}' was already registered in '{}' registry!".format(name, self._name))
+            raise KeyError(f"[CCCV] An object named '{name}' was already registered in '{self._name}' registry!")
         else:
             self._obj_map[name] = obj
 
@@ -69,7 +69,7 @@ def deco(func_or_class: Any) -> Any:
     def get(self, name: str) -> Any:
         ret = self._obj_map.get(name)
         if ret is None:
-            raise KeyError("[CCCV] No object named '{}' found in '{}' registry!".format(name, self._name))
+            raise KeyError(f"[CCCV] No object named '{name}' found in '{self._name}' registry!")
         return ret
 
     def __contains__(self, name: str) -> bool: