diffsynth-engine 0.4.3.dev3__tar.gz → 0.4.3.dev4__tar.gz

{diffsynth_engine-0.4.3.dev3 → diffsynth_engine-0.4.3.dev4}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: diffsynth_engine
-Version: 0.4.3.dev3
+Version: 0.4.3.dev4
 Author: MuseAI x ModelScope
 Classifier: Programming Language :: Python :: 3
 Classifier: Operating System :: OS Independent

{diffsynth_engine-0.4.3.dev3 → diffsynth_engine-0.4.3.dev4}/diffsynth_engine/__init__.py

@@ -17,7 +17,13 @@ from .pipelines import (
 from .models.flux import FluxControlNet, FluxIPAdapter, FluxRedux
 from .models.sd import SDControlNet
 from .models.sdxl import SDXLControlNetUnion
-from .utils.download import fetch_model, fetch_modelscope_model, fetch_civitai_model
+from .utils.download import (
+    fetch_model,
+    fetch_modelscope_model,
+    fetch_civitai_model,
+    register_fetch_modelscope_model,
+    reset_fetch_modelscope_model,
+)
 from .utils.video import load_video, save_video
 from .tools import (
     FluxInpaintingTool,
@@ -52,6 +58,8 @@ __all__ = [
     "ControlType",
     "fetch_model",
     "fetch_modelscope_model",
+    "register_fetch_modelscope_model",
+    "reset_fetch_modelscope_model",
     "fetch_civitai_model",
     "load_video",
     "save_video",

{diffsynth_engine-0.4.3.dev3 → diffsynth_engine-0.4.3.dev4}/diffsynth_engine/utils/download.py

@@ -20,6 +20,31 @@ logger = logging.get_logger(__name__)
 
 MODEL_SOURCES = ["modelscope", "civitai"]
 
+# Global registry for custom fetch function
+_CUSTOM_MODELSCOPE_FETCHER = None
+
+
+def register_fetch_modelscope_model(fetch_func):
+    """
+    Register a global custom fetch function for ModelScope models.
+
+    Args:
+        fetch_func (callable): Custom fetch function that should accept the same parameters
+                               as fetch_modelscope_model and return the model path(s)
+    """
+    global _CUSTOM_MODELSCOPE_FETCHER
+    _CUSTOM_MODELSCOPE_FETCHER = fetch_func
+    logger.info("Registered global custom ModelScope fetcher")
+
+
+def reset_fetch_modelscope_model():
+    """
+    Reset the global custom fetch function for ModelScope models.
+    """
+    global _CUSTOM_MODELSCOPE_FETCHER
+    _CUSTOM_MODELSCOPE_FETCHER = None
+    logger.info("Reset global custom ModelScope fetcher")
+
 
 def fetch_model(
     model_uri: str,
@@ -43,6 +68,11 @@ def fetch_modelscope_model(
     access_token: Optional[str] = None,
     fetch_safetensors: bool = True,
 ) -> str:
+    # Check if there's a global custom fetcher registered
+    if _CUSTOM_MODELSCOPE_FETCHER is not None:
+        logger.info(f"Using global custom fetcher for model: {model_id}")
+        return _CUSTOM_MODELSCOPE_FETCHER(model_id, revision, path, access_token, fetch_safetensors)
+
     lock_file_name = f"modelscope.{model_id.replace('/', '--')}.{revision if revision else '__version'}.lock"
     lock_file_path = os.path.join(DIFFSYNTH_FILELOCK_DIR, lock_file_name)
     ensure_directory_exists(lock_file_path)

{diffsynth_engine-0.4.3.dev3 → diffsynth_engine-0.4.3.dev4}/diffsynth_engine.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: diffsynth_engine
-Version: 0.4.3.dev3
+Version: 0.4.3.dev4
 Author: MuseAI x ModelScope
 Classifier: Programming Language :: Python :: 3
 Classifier: Operating System :: OS Independent