diffsynth-engine 0.4.3.dev2__py3-none-any.whl → 0.4.3.dev4__py3-none-any.whl

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/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/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.dist-info → diffsynth_engine-0.4.3.dev4.dist-info}/METADATA

@@ -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.dist-info → diffsynth_engine-0.4.3.dev4.dist-info}/RECORD

@@ -1,4 +1,4 @@
-diffsynth_engine/__init__.py,sha256=bY_ekp0--kriKjSDdXmppjOahdoL1DlAc22_kDh7b7g,1423
+diffsynth_engine/__init__.py,sha256=tXaLuKje4NQ3zARAvqBUdj1pGLjP0ttkXKE6ysuzsOc,1586
 diffsynth_engine/algorithm/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 diffsynth_engine/algorithm/noise_scheduler/__init__.py,sha256=YvcwE2tCNua-OAX9GEPm0EXsINNWH4XvJMNZb-uaZMM,745
 diffsynth_engine/algorithm/noise_scheduler/base_scheduler.py,sha256=WICrLEh7b2TdZMMEN14NqiYydj7dxXT6RolXymKiMk8,188
@@ -153,7 +153,7 @@ diffsynth_engine/tools/flux_replace_tool.py,sha256=M_q8KnsBEwNi4w8NOK-F2Bmj7cKUN
 diffsynth_engine/utils/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 diffsynth_engine/utils/cache.py,sha256=Ivef22pCuhEq-4H00gSvkLS8ceVZoGis7OSitYL6gH4,2101
 diffsynth_engine/utils/constants.py,sha256=W9znMP8O7td3RF_2ykeIARCWVtIxfro9ytmlXR_GG2c,3121
-diffsynth_engine/utils/download.py,sha256=NCgfL9tUca-sOhT41k6w4o__Ktbw-1aDwFTR4JDkT28,5639
+diffsynth_engine/utils/download.py,sha256=8O56zQr2taY4BnCBEKVk3YGorTqTTCqTblf7tI9otXA,6721
 diffsynth_engine/utils/env.py,sha256=43x-kBjt5zI2cwZ9G4BOeTbedi2k6TuBzHGOBeFbFvU,280
 diffsynth_engine/utils/flag.py,sha256=6zQLnoEaU69pBEyhavCgydQfP0khw5ppCU7sue4yRqg,1370
 diffsynth_engine/utils/fp8_linear.py,sha256=NosnWMoAr_IpFcLn-OYbAx-vXySphjxutDZqmXLNjJI,4064
@@ -162,14 +162,14 @@ diffsynth_engine/utils/image.py,sha256=_46CVs1Qe7GdZNulWWJISnR_Y6FotC2tZGLKtr04g
 diffsynth_engine/utils/loader.py,sha256=Z5v1WNDWFY0OrVubB70j5VU3zeaAfEK_j8c1KrGI4yM,1240
 diffsynth_engine/utils/lock.py,sha256=1Ipgst9eEFfFdViAvD5bxdB6HnHHBcqWYOb__fGaPUI,1601
 diffsynth_engine/utils/logging.py,sha256=XB0xTT8PBN6btkOjFtOvjlrOCRVgDGT8PFAp1vmse28,467
-diffsynth_engine/utils/offload.py,sha256=qEiqbeMQqeV1DtHF-6OuUO8Akdr1enfpPYwF0OBpx98,3500
+diffsynth_engine/utils/offload.py,sha256=8dHMEzAQcE-GNbFYotGQg2YBBB0eywoRTIffVmK4A0k,3746
 diffsynth_engine/utils/onnx.py,sha256=jeWUudJHnESjuiEAHyUZYUZz7dCj34O9aGjHCe8yjWo,1149
 diffsynth_engine/utils/parallel.py,sha256=TtB6FzP2qo4VQqSenVnV-ZaKmp9xHaIWJ8D1ZjHtukE,17064
 diffsynth_engine/utils/platform.py,sha256=2lXdw6YkqcRONCeT98n4cyg1Ii8Ybbyj2Ns72Se9tlk,496
 diffsynth_engine/utils/prompt.py,sha256=YItMchoVzsG6y-LB4vzzDUWrkhKRVlt1HfVhxZjSxMQ,280
 diffsynth_engine/utils/video.py,sha256=Ne0rd2lb59UT1q5EotpjlY7OT8F9oTCFDyo1ST77uoQ,1004
-diffsynth_engine-0.4.3.dev2.dist-info/licenses/LICENSE,sha256=x7aBqQuVI0IYnftgoTPI_A0I_rjdjPPQkjnU6N2nikM,11346
-diffsynth_engine-0.4.3.dev2.dist-info/METADATA,sha256=wv13KwkzgW7iGL29CfOU1C2CsLrLM_2_tNOxhBIXwFQ,1110
-diffsynth_engine-0.4.3.dev2.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-diffsynth_engine-0.4.3.dev2.dist-info/top_level.txt,sha256=6zgbiIzEHLbhgDKRyX0uBJOV3F6VnGGBRIQvSiYYn6w,17
-diffsynth_engine-0.4.3.dev2.dist-info/RECORD,,
+diffsynth_engine-0.4.3.dev4.dist-info/licenses/LICENSE,sha256=x7aBqQuVI0IYnftgoTPI_A0I_rjdjPPQkjnU6N2nikM,11346
+diffsynth_engine-0.4.3.dev4.dist-info/METADATA,sha256=4D5ubLDy1tWA_n94e0QPdgTHQHOjFxEp-FM7mzS_ClE,1110
+diffsynth_engine-0.4.3.dev4.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+diffsynth_engine-0.4.3.dev4.dist-info/top_level.txt,sha256=6zgbiIzEHLbhgDKRyX0uBJOV3F6VnGGBRIQvSiYYn6w,17
+diffsynth_engine-0.4.3.dev4.dist-info/RECORD,,