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

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

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

{diffsynth_engine-0.4.3.dev2 → 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.dev2 → 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.dev2 → diffsynth_engine-0.4.3.dev4}/diffsynth_engine/utils/offload.py

@@ -1,7 +1,7 @@
 import torch
 import torch.nn as nn
 from typing import Dict
-
+import platform
 
 def enable_sequential_cpu_offload(module: nn.Module, device: str = "cuda"):
     module = module.to("cpu")
@@ -26,13 +26,14 @@ def add_cpu_offload_hook(module: nn.Module, device: str = "cuda", recurse: bool
             for name, buffer in module.named_buffers(recurse=recurse):
                 buffer.data = buffer.data.to(device=device)
             return tuple(x.to(device=device) if isinstance(x, torch.Tensor) else x for x in input_)
-
-        for name, param in module.named_parameters(recurse=recurse):
-            param.data = param.data.pin_memory()
+        for name, param in module.named_parameters(recurse=recurse):            
+            if platform.system()  == 'Linux':
+                param.data = param.data.pin_memory()
             offload_param_dict[name] = param.data
             param.data = param.data.to(device=device)
         for name, buffer in module.named_buffers(recurse=recurse):
-            buffer.data = buffer.data.pin_memory()
+            if platform.system()  == 'Linux':            
+                buffer.data = buffer.data.pin_memory()
             offload_param_dict[name] = buffer.data
             buffer.data = buffer.data.to(device=device)
         setattr(module, "_offload_param_dict", offload_param_dict)
@@ -58,10 +59,12 @@ def offload_model_to_dict(module: nn.Module) -> Dict[str, torch.Tensor]:
     module = module.to("cpu")
     offload_param_dict = {}
     for name, param in module.named_parameters(recurse=True):
-        param.data = param.data.pin_memory()
+        if platform.system()  == 'Linux':        
+            param.data = param.data.pin_memory()
         offload_param_dict[name] = param.data
     for name, buffer in module.named_buffers(recurse=True):
-        buffer.data = buffer.data.pin_memory()
+        if platform.system()  == 'Linux':        
+            buffer.data = buffer.data.pin_memory()
         offload_param_dict[name] = buffer.data
     return offload_param_dict
 

{diffsynth_engine-0.4.3.dev2 → 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.dev2
+Version: 0.4.3.dev4
 Author: MuseAI x ModelScope
 Classifier: Programming Language :: Python :: 3
 Classifier: Operating System :: OS Independent