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/arch/sr/dat_arch.py

@@ -365,7 +365,7 @@ def __init__(
         elif idx == 1:
             W_sp, H_sp = self.split_size[0], self.split_size[1]
         else:
-            print("ERROR MODE", idx)
+            print("[CCCV] ERROR MODE", idx)
             exit(0)
         self.H_sp = H_sp
         self.W_sp = W_sp

cccv/arch/sr/upcunet_arch.py

@@ -376,7 +376,7 @@ def forward(self, x, tile_mode, cache_mode, alpha, pro):
             t2 = tile_mode * 2
             crop_size = (((h0 - 1) // t2 * t2 + t2) // tile_mode, ((w0 - 1) // t2 * t2 + t2) // tile_mode)
         else:
-            print("tile_mode config error")
+            print("[CCCV] tile_mode config error")
             os._exit(233)
 
         ph = ((h0 - 1) // crop_size[0] + 1) * crop_size[0]
@@ -526,7 +526,7 @@ def forward_gap_sync(self, x, tile_mode, alpha, pro):
             t2 = tile_mode * 2
             crop_size = (((h0 - 1) // t2 * t2 + t2) // tile_mode, ((w0 - 1) // t2 * t2 + t2) // tile_mode)
         else:
-            print("tile_mode config error")
+            print("[CCCV] tile_mode config error")
             os._exit(233)
         ph = ((h0 - 1) // crop_size[0] + 1) * crop_size[0]
         pw = ((w0 - 1) // crop_size[1] + 1) * crop_size[1]
@@ -767,7 +767,7 @@ def forward(self, x, tile_mode, cache_mode, alpha, pro):
             t4 = tile_mode * 4
             crop_size = (((h0 - 1) // t4 * t4 + t4) // tile_mode, ((w0 - 1) // t4 * t4 + t4) // tile_mode)
         else:
-            print("tile_mode config error")
+            print("[CCCV] tile_mode config error")
             os._exit(233)
         ph = ((h0 - 1) // crop_size[0] + 1) * crop_size[0]
         pw = ((w0 - 1) // crop_size[1] + 1) * crop_size[1]
@@ -916,7 +916,7 @@ def forward_gap_sync(self, x, tile_mode, alpha, pro):
             t4 = tile_mode * 4
             crop_size = (((h0 - 1) // t4 * t4 + t4) // tile_mode, ((w0 - 1) // t4 * t4 + t4) // tile_mode)
         else:
-            print("tile_mode config error")
+            print("[CCCV] tile_mode config error")
             os._exit(233)
         ph = ((h0 - 1) // crop_size[0] + 1) * crop_size[0]
         pw = ((w0 - 1) // crop_size[1] + 1) * crop_size[1]
@@ -1162,7 +1162,7 @@ def forward(self, x, tile_mode, cache_mode, alpha, pro):
             t2 = tile_mode * 2
             crop_size = (((h0 - 1) // t2 * t2 + t2) // tile_mode, ((w0 - 1) // t2 * t2 + t2) // tile_mode)
         else:
-            print("tile_mode config error")
+            print("[CCCV] tile_mode config error")
             os._exit(233)
         ph = ((h0 - 1) // crop_size[0] + 1) * crop_size[0]
         pw = ((w0 - 1) // crop_size[1] + 1) * crop_size[1]
@@ -1323,7 +1323,7 @@ def forward_gap_sync(self, x, tile_mode, alpha, pro):
             t2 = tile_mode * 2
             crop_size = (((h0 - 1) // t2 * t2 + t2) // tile_mode, ((w0 - 1) // t2 * t2 + t2) // tile_mode)  # 5.6G
         else:
-            print("tile_mode config error")
+            print("[CCCV] tile_mode config error")
             os._exit(233)
         ph = ((h0 - 1) // crop_size[0] + 1) * crop_size[0]
         pw = ((w0 - 1) // crop_size[1] + 1) * crop_size[1]

cccv/arch/vfi/drba_arch.py

@@ -1,4 +1,6 @@
 # type: ignore
+import warnings
+
 import numpy as np
 import torch
 import torch.nn as nn
@@ -61,7 +63,7 @@ def inference(self, x, timestep=0.5, scale_list=None, fastmode=True, ensemble=Fa
                     torch.cat((img0[:, :3], img1[:, :3], f0, f1, timestep), 1), None, scale=scale_list[i]
                 )
                 if ensemble:
-                    print("warning: ensemble is not supported since RIFEv4.21")
+                    warnings.warn("[CCCV] ensemble is not supported since RIFEv4.21", stacklevel=2)
             else:
                 wf0 = warp(f0, flow[:, :2])
                 wf1 = warp(f1, flow[:, 2:4])
@@ -71,7 +73,7 @@ def inference(self, x, timestep=0.5, scale_list=None, fastmode=True, ensemble=Fa
                     scale=scale_list[i],
                 )
                 if ensemble:
-                    print("warning: ensemble is not supported since RIFEv4.21")
+                    warnings.warn("[CCCV] ensemble is not supported since RIFEv4.21", stacklevel=2)
                 else:
                     mask = m0
                 flow = flow + fd
@@ -83,7 +85,7 @@ def inference(self, x, timestep=0.5, scale_list=None, fastmode=True, ensemble=Fa
         mask = torch.sigmoid(mask)
         merged[4] = warped_img0 * mask + warped_img1 * (1 - mask)
         if not fastmode:
-            print("contextnet is removed")
+            warnings.warn("[CCCV] contextnet is removed", stacklevel=2)
             """
             c0 = self.contextnet(img0, flow[:, :2])
             c1 = self.contextnet(img1, flow[:, 2:4])

cccv/arch/vfi/ifnet_arch.py

@@ -1,4 +1,6 @@
 # type: ignore
+import warnings
+
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
@@ -43,7 +45,7 @@ def forward(self, x, timestep=0.5, scale_list=None, fastmode=True, ensemble=Fals
                     torch.cat((img0[:, :3], img1[:, :3], f0, f1, timestep), 1), None, scale=scale_list[i]
                 )
                 if ensemble:
-                    print("warning: ensemble is not supported since RIFEv4.21")
+                    warnings.warn("[CCCV] ensemble is not supported since RIFEv4.21", stacklevel=2)
             else:
                 wf0 = warp(f0, flow[:, :2])
                 wf1 = warp(f1, flow[:, 2:4])
@@ -53,7 +55,7 @@ def forward(self, x, timestep=0.5, scale_list=None, fastmode=True, ensemble=Fals
                     scale=scale_list[i],
                 )
                 if ensemble:
-                    print("warning: ensemble is not supported since RIFEv4.21")
+                    warnings.warn("[CCCV] ensemble is not supported since RIFEv4.21", stacklevel=2)
                 else:
                     mask = m0
                 flow = flow + fd
@@ -65,7 +67,7 @@ def forward(self, x, timestep=0.5, scale_list=None, fastmode=True, ensemble=Fals
         mask = torch.sigmoid(mask)
         merged[4] = warped_img0 * mask + warped_img1 * (1 - mask)
         if not fastmode:
-            print("contextnet is removed")
+            warnings.warn("[CCCV] contextnet is removed", stacklevel=2)
             """
             c0 = self.contextnet(img0, flow[:, :2])
             c1 = self.contextnet(img1, flow[:, 2:4])

cccv/arch/vfi/vfi_utils/softsplat.py

@@ -61,7 +61,7 @@ def cuda_kernel(strFunction: str, strKernel: str, objVariables: typing.Dict):
             strKey += str(objValue.stride())
 
         elif True:
-            print(strVariable, type(objValue))
+            print(f"[CCCV] {strVariable}, {type(objValue)}")
 
         # end
     # end
@@ -106,10 +106,10 @@ def cuda_kernel(strFunction: str, strKernel: str, objVariables: typing.Dict):
                 strKernel = strKernel.replace("{{type}}", "long")
 
             elif isinstance(objValue, torch.Tensor):
-                print(strVariable, objValue.dtype)
+                print(f"[CCCV] {strVariable}, {objValue.dtype}")
 
             elif True:
-                print(strVariable, type(objValue))
+                print(f"[CCCV] {strVariable}, {type(objValue)}")
 
             # end
         # end

cccv/auto/config.py

@@ -1,32 +1,84 @@
+import importlib.util
+import json
+import warnings
+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:
+            warnings.warn(
+                "[CCCV] 'pretrained_model_name' is deprecated, please use 'pretrained_model_name_or_path' instead.",
+                DeprecationWarning,
+                stacklevel=2,
+            )
+            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)