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

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: diffsynth_engine
-Version: 0.4.3.dev3
+Version: 0.4.3.dev5
 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.dev5}/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.dev5}/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.dev5}/diffsynth_engine/utils/parallel.py

@@ -331,19 +331,14 @@ class ParallelWrapper:
         master_port: int = 29500,
         device: str = "cuda",
     ):
-        current_method = mp.get_start_method(allow_none=True)
-        if current_method is None or current_method != "spawn":
-            try:
-                mp.set_start_method("spawn")
-            except RuntimeError as e:
-                raise RuntimeError("Failed to set start method to spawn:", e)
         super().__init__()
         self.config = module.config if isinstance(module, BasePipeline) else None
         self._module_name = module.__class__.__name__
 
         self.world_size = cfg_degree * sp_ulysses_degree * sp_ring_degree * tp_degree
-        self.queue_in = mp.Queue()
-        self.queue_out = mp.Queue()
+        spawn_ctx = mp.get_context("spawn")
+        self.queue_in = spawn_ctx.Queue()
+        self.queue_out = spawn_ctx.Queue()
         self.ctx = mp.spawn(
             _worker_loop,
             args=(

{diffsynth_engine-0.4.3.dev3 → diffsynth_engine-0.4.3.dev5}/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.dev5
 Author: MuseAI x ModelScope
 Classifier: Programming Language :: Python :: 3
 Classifier: Operating System :: OS Independent