diffsynth 2.1.12__tar.gz → 2.2.0__tar.gz

{diffsynth-2.1.12 → diffsynth-2.2.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: diffsynth
-Version: 2.1.12
+Version: 2.2.0
 Summary: Enjoy the magic of Diffusion models!
 Author: ModelScope Team
 License-Expression: Apache-2.0

diffsynth-2.2.0/diffsynth/core/attention/__init__.py

@@ -0,0 +1,2 @@
+from .attention import attention_forward, list_attention_implementations
+from .flex_attention import FLEX_ATTN_AVAILABLE

diffsynth-2.2.0/diffsynth/core/attention/attention.py

@@ -0,0 +1,51 @@
+import os
+from .registry import (
+    get_attention_implementation,
+    list_attention_implementations,
+    register_attention_implementation,
+)
+from .flash_attention import CUSTOMIZED_FA_KERNEL_AVAILABLE, FLASH_ATTN_4_AVAILABLE, FLASH_ATTN_3_AVAILABLE, FLASH_ATTN_2_AVAILABLE
+from .sage_attention import SAGE_ATTN_AVAILABLE
+from .xformers_attention import XFORMERS_AVAILABLE
+from .torch_attention import torch_sdpa, torch_sdpa_sliding_window
+from .flex_attention import flex_attention
+
+
+def initialize_attention_priority():
+    if os.environ.get('DIFFSYNTH_ATTENTION_IMPLEMENTATION') is not None:
+        return os.environ.get('DIFFSYNTH_ATTENTION_IMPLEMENTATION').strip().lower()
+    else:
+        for name in ["customized_fa_kernel", "flash_attention_4", "flash_attention_3", "flash_attention_2", "sage_attention", "xformers"]:
+            if name in list_attention_implementations():
+                return name
+        return "torch"
+
+
+ATTENTION_IMPLEMENTATION = initialize_attention_priority()
+
+
+def attention_forward(q, k, v, q_pattern="b n s d", k_pattern="b n s d", v_pattern="b n s d", out_pattern="b n s d", dims=None, attn_mask=None, scale=None, is_causal=False, compatibility_mode=False, window_size=None, use_flex=False, score_mod=None):
+    if use_flex or score_mod is not None:
+        return flex_attention(q, k, v, q_pattern, k_pattern, v_pattern, out_pattern, dims, attn_mask=attn_mask, scale=scale, score_mod=score_mod)
+
+    if compatibility_mode:
+        return torch_sdpa(q, k, v, q_pattern, k_pattern, v_pattern, out_pattern, dims, attn_mask=attn_mask, scale=scale, is_causal=is_causal, window_size=window_size)
+
+    backend = get_attention_implementation(ATTENTION_IMPLEMENTATION)
+
+    if attn_mask is not None and not backend.supports_mask:
+        return torch_sdpa(q, k, v, q_pattern, k_pattern, v_pattern, out_pattern, dims, attn_mask=attn_mask, scale=scale, is_causal=is_causal)
+    if window_size is not None and not backend.supports_window:
+        assert is_causal == False and attn_mask is None
+        return torch_sdpa_sliding_window(q, k, v, window_size, q_pattern, k_pattern, v_pattern, out_pattern, dims, scale=scale)
+    if is_causal and not backend.supports_causal:
+        return torch_sdpa(q, k, v, q_pattern, k_pattern, v_pattern, out_pattern, dims, attn_mask=attn_mask, scale=scale, is_causal=is_causal)
+
+    kwargs = {"scale": scale}
+    if backend.supports_causal:
+        kwargs["is_causal"] = is_causal
+    if backend.supports_mask:
+        kwargs["attn_mask"] = attn_mask
+    if backend.supports_window:
+        kwargs["window_size"] = window_size
+    return backend.func(q, k, v, q_pattern, k_pattern, v_pattern, out_pattern, dims, **kwargs)

diffsynth-2.2.0/diffsynth/core/attention/common.py

@@ -0,0 +1,27 @@
+import inspect
+
+import torch
+from einops import rearrange
+
+try:
+    TORCH_SUPPORT_GQA = "enable_gqa" in inspect.signature(torch.nn.functional.scaled_dot_product_attention).parameters
+except Exception:
+    TORCH_SUPPORT_GQA = False
+
+
+def rearrange_qkv(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, q_pattern="b n s d", k_pattern="b n s d", v_pattern="b n s d", required_in_pattern="b n s d", dims=None):
+    dims = {} if dims is None else dims
+    if q_pattern != required_in_pattern:
+        q = rearrange(q, f"{q_pattern} -> {required_in_pattern}", **dims)
+    if k_pattern != required_in_pattern:
+        k = rearrange(k, f"{k_pattern} -> {required_in_pattern}", **dims)
+    if v_pattern != required_in_pattern:
+        v = rearrange(v, f"{v_pattern} -> {required_in_pattern}", **dims)
+    return q, k, v
+
+
+def rearrange_out(out: torch.Tensor, out_pattern="b n s d", required_out_pattern="b n s d", dims=None):
+    dims = {} if dims is None else dims
+    if out_pattern != required_out_pattern:
+        out = rearrange(out, f"{required_out_pattern} -> {out_pattern}", **dims)
+    return out

diffsynth-2.2.0/diffsynth/core/attention/flash_attention.py

@@ -0,0 +1,118 @@
+import os
+
+from .common import rearrange_qkv, rearrange_out
+from .registry import register_attention_implementation
+
+if os.environ.get("DIFFSYNTH_FLASH_ATTN_KERNEL_REPO_ID") is not None:
+    try:
+        from kernels import get_kernel
+        kernel_module = get_kernel(os.environ.get("DIFFSYNTH_FLASH_ATTN_KERNEL_REPO_ID"), version=int(os.environ.get("DIFFSYNTH_FLASH_ATTN_KERNEL_VERSION")))
+        customized_flash_attn_func = kernel_module.flash_attn_func
+        CUSTOMIZED_FA_KERNEL_AVAILABLE = True
+    except Exception:
+        CUSTOMIZED_FA_KERNEL_AVAILABLE = False
+else:
+    CUSTOMIZED_FA_KERNEL_AVAILABLE = False
+
+try:
+    from flash_attn.cute import flash_attn_func as flash_attn_func_cute
+    FLASH_ATTN_4_AVAILABLE = True
+except Exception:
+    FLASH_ATTN_4_AVAILABLE = False
+
+try:
+    from flash_attn_3 import flash_attn_interface
+    FLASH_ATTN_3_AVAILABLE = True
+except Exception:
+    FLASH_ATTN_3_AVAILABLE = False
+
+try:
+    import flash_attn
+    FLASH_ATTN_2_AVAILABLE = hasattr(flash_attn, "flash_attn_func")
+except Exception:
+    FLASH_ATTN_2_AVAILABLE = False
+
+
+def _require(flag, name):
+    if not flag:
+        raise RuntimeError(f"{name} is not available.")
+
+
+def flash_attention_4(q, k, v, q_pattern="b n s d", k_pattern="b n s d", v_pattern="b n s d", out_pattern="b n s d", dims=None, scale=None, is_causal=False, window_size=None):
+    _require(FLASH_ATTN_4_AVAILABLE, "FlashAttention-4")
+    required_in_pattern, required_out_pattern = "b s n d", "b s n d"
+    q, k, v = rearrange_qkv(q, k, v, q_pattern, k_pattern, v_pattern, required_in_pattern, dims)
+    window_size = (window_size, window_size) if window_size is not None else (-1, -1)
+    out = flash_attn_func_cute(q, k, v, softmax_scale=scale, causal=is_causal, window_size=window_size)
+    if isinstance(out, tuple):
+        out = out[0]
+    out = rearrange_out(out, out_pattern, required_out_pattern, dims)
+    return out
+
+
+def flash_attention_3(q, k, v, q_pattern="b n s d", k_pattern="b n s d", v_pattern="b n s d", out_pattern="b n s d", dims=None, scale=None, is_causal=False, window_size=None):
+    _require(FLASH_ATTN_3_AVAILABLE, "FlashAttention-3")
+    required_in_pattern, required_out_pattern = "b s n d", "b s n d"
+    q, k, v = rearrange_qkv(q, k, v, q_pattern, k_pattern, v_pattern, required_in_pattern, dims)
+    window_size = (window_size, window_size) if window_size is not None else (-1, -1)
+    out = flash_attn_interface.flash_attn_func(q, k, v, softmax_scale=scale, causal=is_causal, window_size=window_size)
+    if isinstance(out, tuple):
+        out = out[0]
+    out = rearrange_out(out, out_pattern, required_out_pattern, dims)
+    return out
+
+
+def flash_attention_2(q, k, v, q_pattern="b n s d", k_pattern="b n s d", v_pattern="b n s d", out_pattern="b n s d", dims=None, scale=None, is_causal=False, window_size=None):
+    _require(FLASH_ATTN_2_AVAILABLE, "FlashAttention-2")
+    required_in_pattern, required_out_pattern = "b s n d", "b s n d"
+    q, k, v = rearrange_qkv(q, k, v, q_pattern, k_pattern, v_pattern, required_in_pattern, dims)
+    window_size = (window_size, window_size) if window_size is not None else (-1, -1)
+    out = flash_attn.flash_attn_func(q, k, v, softmax_scale=scale, causal=is_causal, window_size=window_size)
+    out = rearrange_out(out, out_pattern, required_out_pattern, dims)
+    return out
+
+
+def flash_attention_customized_kernel(q, k, v, q_pattern="b n s d", k_pattern="b n s d", v_pattern="b n s d", out_pattern="b n s d", dims=None, scale=None, is_causal=False, window_size=None):
+    _require(CUSTOMIZED_FA_KERNEL_AVAILABLE, "customized FlashAttention kernel")
+    required_in_pattern, required_out_pattern = "b s n d", "b s n d"
+    q, k, v = rearrange_qkv(q, k, v, q_pattern, k_pattern, v_pattern, required_in_pattern, dims)
+    window_size = (window_size, window_size) if window_size is not None else (-1, -1)
+    out = customized_flash_attn_func(q, k, v, softmax_scale=scale, causal=is_causal, window_size=window_size)
+    if isinstance(out, tuple):
+        out = out[0]
+    out = rearrange_out(out, out_pattern, required_out_pattern, dims)
+    return out
+
+
+if CUSTOMIZED_FA_KERNEL_AVAILABLE:
+    register_attention_implementation(
+        "customized_fa_kernel",
+        flash_attention_customized_kernel,
+        supports_causal=True,
+        supports_window=True,
+        supports_gqa=True,
+    )
+if FLASH_ATTN_4_AVAILABLE:
+    register_attention_implementation(
+        "flash_attention_4",
+        flash_attention_4,
+        supports_causal=True,
+        supports_window=True,
+        supports_gqa=True,
+    )
+if FLASH_ATTN_3_AVAILABLE:
+    register_attention_implementation(
+        "flash_attention_3",
+        flash_attention_3,
+        supports_causal=True,
+        supports_window=True,
+        supports_gqa=True,
+    )
+if FLASH_ATTN_2_AVAILABLE:
+    register_attention_implementation(
+        "flash_attention_2",
+        flash_attention_2,
+        supports_causal=True,
+        supports_window=True,
+        supports_gqa=True,
+    )

diffsynth-2.2.0/diffsynth/core/attention/flex_attention.py

@@ -0,0 +1,28 @@
+import torch
+
+from .common import rearrange_qkv, rearrange_out
+
+try:
+    from torch.nn.attention.flex_attention import flex_attention as _flex_attention_func
+    FLEX_ATTN_AVAILABLE = True
+except Exception:
+    FLEX_ATTN_AVAILABLE = False
+
+# Compile flex attention lazily
+_compiled_flex_attention_func = None
+
+
+def get_flex_attention_func():
+    global _compiled_flex_attention_func
+    if _compiled_flex_attention_func is None:
+        _compiled_flex_attention_func = torch.compile(_flex_attention_func)
+    return _compiled_flex_attention_func
+
+
+def flex_attention(q, k, v, q_pattern="b n s d", k_pattern="b n s d", v_pattern="b n s d", out_pattern="b n s d", dims=None, attn_mask=None, scale=None, score_mod=None):
+    assert FLEX_ATTN_AVAILABLE, "Flex Attention is not available. Please upgrade torch to 2.5.0 or later."
+    required_in_pattern, required_out_pattern = "b n s d", "b n s d"
+    q, k, v = rearrange_qkv(q, k, v, q_pattern, k_pattern, v_pattern, required_in_pattern, dims)
+    out = get_flex_attention_func()(query=q, key=k, value=v, block_mask=attn_mask, scale=scale, score_mod=score_mod)
+    out = rearrange_out(out, out_pattern, required_out_pattern, dims)
+    return out

diffsynth-2.2.0/diffsynth/core/attention/registry.py

@@ -0,0 +1,48 @@
+from dataclasses import dataclass
+from typing import Callable, Optional
+
+
+@dataclass(frozen=True)
+class AttentionBackend:
+    func: Callable
+    supports_causal: bool = False
+    supports_window: bool = False
+    supports_mask: bool = False
+    supports_gqa: bool = False
+
+
+_ATTENTION_BACKENDS = {}
+
+
+def register_attention_implementation(
+    name: str,
+    func: Optional[Callable] = None,
+    supports_causal: bool = False,
+    supports_window: bool = False,
+    supports_mask: bool = False,
+    supports_gqa: bool = False,
+):
+    key = name.strip().lower()
+    _ATTENTION_BACKENDS[key] = AttentionBackend(
+        func=func,
+        supports_causal=supports_causal,
+        supports_window=supports_window,
+        supports_mask=supports_mask,
+        supports_gqa=supports_gqa,
+    )
+
+
+def get_attention_implementation(name: str) -> AttentionBackend:
+    key = name.strip().lower()
+    if key not in _ATTENTION_BACKENDS:
+        available = ", ".join(list_attention_implementations()) or "none"
+        raise ValueError(
+            f"Unknown attention implementation: {name!r}. "
+            f"Available implementations: {available}. "
+            f"Please check the environment variable `DIFFSYNTH_ATTENTION_IMPLEMENTATION`."
+        )
+    return _ATTENTION_BACKENDS[key]
+
+
+def list_attention_implementations():
+    return sorted(_ATTENTION_BACKENDS)

diffsynth-2.2.0/diffsynth/core/attention/sage_attention.py

@@ -0,0 +1,20 @@
+from .common import rearrange_qkv, rearrange_out
+from .registry import register_attention_implementation
+
+try:
+    from sageattention import sageattn
+    SAGE_ATTN_AVAILABLE = True
+except Exception:
+    SAGE_ATTN_AVAILABLE = False
+
+def sage_attention(q, k, v, q_pattern="b n s d", k_pattern="b n s d", v_pattern="b n s d", out_pattern="b n s d", dims=None, scale=None):
+    if not SAGE_ATTN_AVAILABLE:
+        raise RuntimeError("SageAttention is not available.")
+    required_in_pattern, required_out_pattern = "b n s d", "b n s d"
+    q, k, v = rearrange_qkv(q, k, v, q_pattern, k_pattern, v_pattern, required_in_pattern, dims)
+    out = sageattn(q, k, v, sm_scale=scale)
+    out = rearrange_out(out, out_pattern, required_out_pattern, dims)
+    return out
+
+if SAGE_ATTN_AVAILABLE:
+    register_attention_implementation("sage_attention", sage_attention)

diffsynth-2.2.0/diffsynth/core/attention/torch_attention.py

@@ -0,0 +1,82 @@
+import torch
+from einops import repeat
+
+from .common import TORCH_SUPPORT_GQA, rearrange_qkv, rearrange_out
+from .registry import register_attention_implementation
+
+
+def torch_sdpa(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, q_pattern="b n s d", k_pattern="b n s d", v_pattern="b n s d", out_pattern="b n s d", dims=None, attn_mask=None, scale=None, is_causal=False, window_size=None):
+    if window_size is not None:
+        if is_causal or attn_mask is not None:
+            raise ValueError("Sliding window attention does not support `is_causal=True` or `attn_mask`.")
+        return torch_sdpa_sliding_window(q, k, v, window_size, q_pattern, k_pattern, v_pattern, out_pattern, dims, scale=scale)
+
+    required_in_pattern, required_out_pattern = "b n s d", "b n s d"
+    q, k, v = rearrange_qkv(q, k, v, q_pattern, k_pattern, v_pattern, required_in_pattern, dims)
+    if q.shape[1] != k.shape[1] or q.shape[1] != v.shape[1]:
+        # Grouped Query Attention
+        if TORCH_SUPPORT_GQA:
+            out = torch.nn.functional.scaled_dot_product_attention(q, k, v, attn_mask, scale=scale, is_causal=is_causal, enable_gqa=True)
+        else:
+            # In low-version torch, `enable_gqa` is not supported.
+            k = repeat(k, "b n s d -> b (n m) s d", m=q.shape[1]//k.shape[1])
+            v = repeat(v, "b n s d -> b (n m) s d", m=q.shape[1]//v.shape[1])
+            out = torch.nn.functional.scaled_dot_product_attention(q, k, v, attn_mask, scale=scale, is_causal=is_causal)
+    else:
+        out = torch.nn.functional.scaled_dot_product_attention(q, k, v, attn_mask, scale=scale, is_causal=is_causal)
+    out = rearrange_out(out, out_pattern, required_out_pattern, dims)
+    return out
+
+
+def torch_sdpa_sliding_window(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, sliding_window: int, q_pattern="b n s d", k_pattern="b n s d", v_pattern="b n s d", out_pattern="b n s d", dims=None, scale=None):
+    required_in_pattern, required_out_pattern = "b n s d", "b n s d"
+    q, k, v = rearrange_qkv(q, k, v, q_pattern, k_pattern, v_pattern, required_in_pattern, dims)
+
+    B, N, S, D = q.shape
+    W = sliding_window
+    chunk_size = W
+    num_chunks = (S + chunk_size - 1) // chunk_size
+
+    output = torch.empty_like(q)
+    dtype = q.dtype
+    device = q.device
+    min_val = torch.finfo(dtype).min
+
+    for i in range(num_chunks):
+        q_start = i * chunk_size
+        q_end = min(q_start + chunk_size, S)
+        actual_chunk_size = q_end - q_start
+
+        kv_start = max(0, q_start - W)
+        kv_end = min(S, q_end + W)
+
+        q_chunk = q[:, :, q_start:q_end, :]
+        k_chunk = k[:, :, kv_start:kv_end, :]
+        v_chunk = v[:, :, kv_start:kv_end, :]
+
+        q_indices = torch.arange(q_start, q_end, device=device)
+        k_indices = torch.arange(kv_start, kv_end, device=device)
+        diff = q_indices.unsqueeze(1) - k_indices.unsqueeze(0)
+        valid = diff.abs() <= W
+
+        local_mask = torch.zeros(actual_chunk_size, kv_end - kv_start, dtype=dtype, device=device)
+        local_mask.masked_fill_(~valid, min_val)
+        local_mask = local_mask.unsqueeze(0).unsqueeze(0)
+
+        out_chunk = torch_sdpa(
+            q_chunk, k_chunk, v_chunk, attn_mask=local_mask, scale=scale
+        )
+        output[:, :, q_start:q_end, :] = out_chunk
+
+    output = rearrange_out(output, out_pattern, required_out_pattern, dims)
+    return output
+
+
+register_attention_implementation(
+    "torch",
+    torch_sdpa,
+    supports_window=True,
+    supports_causal=True,
+    supports_mask=True,
+    supports_gqa=True,
+)

diffsynth-2.2.0/diffsynth/core/attention/xformers_attention.py

@@ -0,0 +1,22 @@
+from .common import rearrange_qkv, rearrange_out
+from .registry import register_attention_implementation
+
+try:
+    import xformers.ops as xops
+    XFORMERS_AVAILABLE = True
+except Exception:
+    XFORMERS_AVAILABLE = False
+
+
+def xformers_attention(q, k, v, q_pattern="b n s d", k_pattern="b n s d", v_pattern="b n s d", out_pattern="b n s d", dims=None, scale=None):
+    if not XFORMERS_AVAILABLE:
+        raise RuntimeError("xformers is not available.")
+    required_in_pattern, required_out_pattern = "b s n d", "b s n d"
+    q, k, v = rearrange_qkv(q, k, v, q_pattern, k_pattern, v_pattern, required_in_pattern, dims)
+    out = xops.memory_efficient_attention(q, k, v, scale=scale)
+    out = rearrange_out(out, out_pattern, required_out_pattern, dims)
+    return out
+
+
+if XFORMERS_AVAILABLE:
+    register_attention_implementation("xformers", xformers_attention)

{diffsynth-2.1.12 → diffsynth-2.2.0}/diffsynth/core/data/operators.py

@@ -1,6 +1,5 @@
 import math, warnings, time
-import torch, torchvision, imageio, os
-import imageio.v3 as iio
+import torch, os
 import urllib.request
 from io import BytesIO
 from PIL import Image
@@ -100,16 +99,20 @@ class ImageCropAndResize(DataProcessingOperator):
         self.max_pixels = max_pixels
         self.height_division_factor = height_division_factor
         self.width_division_factor = width_division_factor
+        import torchvision
+        self.resize = torchvision.transforms.functional.resize
+        self.center_crop = torchvision.transforms.functional.center_crop
+        self.interpolation_mode = torchvision.transforms.InterpolationMode.BILINEAR
 
     def crop_and_resize(self, image, target_height, target_width):
         width, height = image.size
         scale = max(target_width / width, target_height / height)
-        image = torchvision.transforms.functional.resize(
+        image = self.resize(
             image,
             (round(height*scale), round(width*scale)),
-            interpolation=torchvision.transforms.InterpolationMode.BILINEAR
+            interpolation=self.interpolation_mode
         )
-        image = torchvision.transforms.functional.center_crop(image, (target_height, target_width))
+        image = self.center_crop(image, (target_height, target_width))
         return image
     
     def get_height_width(self, image):
@@ -143,6 +146,7 @@ class FrameSamplerByRateMixin:
         self.fix_frame_rate = fix_frame_rate
 
     def get_reader(self, data: str):
+        import imageio
         return imageio.get_reader(data)
 
     def get_available_num_frames(self, reader):
@@ -213,6 +217,7 @@ class LoadGIF(DataProcessingOperator):
 
     def get_num_frames(self, path):
         num_frames = self.num_frames
+        import imageio.v3 as iio
         images = iio.imread(path, mode="RGB")
         if len(images) < num_frames:
             num_frames = len(images)
@@ -223,6 +228,7 @@ class LoadGIF(DataProcessingOperator):
     def __call__(self, data: str):
         num_frames = self.get_num_frames(data)
         frames = []
+        import imageio.v3 as iio
         images = iio.imread(data, mode="RGB")
         for img in images:
             frame = Image.fromarray(img)

{diffsynth-2.1.12 → diffsynth-2.2.0}/diffsynth/core/data/unified_dataset.py

@@ -1,5 +1,5 @@
 from .operators import *
-import torch, json, pandas
+import torch, json, os
 
 
 class UnifiedDataset(torch.utils.data.Dataset):
@@ -84,6 +84,7 @@ class UnifiedDataset(torch.utils.data.Dataset):
                     metadata.append(json.loads(line.strip()))
             self.data = metadata
         else:
+            import pandas
             metadata = pandas.read_csv(metadata_path)
             self.data = [metadata.iloc[i].to_dict() for i in range(len(metadata))]
 

{diffsynth-2.1.12 → diffsynth-2.2.0}/diffsynth/core/gradient/gradient_checkpoint.py

@@ -1,11 +1,9 @@
 import torch
+import sys
 
 
-try:
-    import deepspeed
-    _HAS_DEEPSPEED = True
-except ModuleNotFoundError:
-    _HAS_DEEPSPEED = False
+def get_deepspeed():
+    return sys.modules.get("deepspeed")
 
 
 def create_custom_forward(module):
@@ -36,7 +34,8 @@ def gradient_checkpoint_forward(
     *args,
     **kwargs,
 ):
-    if use_gradient_checkpointing and _HAS_DEEPSPEED and deepspeed.checkpointing.is_configured():
+    deepspeed = get_deepspeed() if use_gradient_checkpointing else None
+    if use_gradient_checkpointing and deepspeed is not None and deepspeed.checkpointing.is_configured():
         kwarg_keys = tuple(kwargs)
         all_args = args + tuple(kwargs.values())
         if not judge_args_requires_grad(*all_args):

{diffsynth-2.1.12 → diffsynth-2.2.0}/diffsynth/core/loader/config.py

@@ -1,8 +1,6 @@
 import torch, glob, os, sys
 from typing import Optional, Union, Dict
 from dataclasses import dataclass
-from modelscope import snapshot_download
-from huggingface_hub import snapshot_download as hf_snapshot_download
 from typing import Optional
 from ..quant import QuantizeConfig, MixedQuantizeConfig
 
@@ -98,6 +96,7 @@ class ModelConfig:
         downloaded_files = glob.glob(origin_file_pattern, root_dir=os.path.join(self.local_model_path, self.model_id))
         download_source = self.parse_download_source()
         if download_source.lower() == "modelscope":
+            from modelscope import snapshot_download
             snapshot_download(
                 self.model_id,
                 local_dir=os.path.join(self.local_model_path, self.model_id),
@@ -106,6 +105,7 @@ class ModelConfig:
                 local_files_only=False
             )
         elif download_source.lower() == "huggingface":
+            from huggingface_hub import snapshot_download as hf_snapshot_download
             hf_snapshot_download(
                 self.model_id,
                 local_dir=os.path.join(self.local_model_path, self.model_id),

{diffsynth-2.1.12 → diffsynth-2.2.0}/diffsynth/core/loader/model.py

@@ -3,14 +3,22 @@ from ..vram.disk_map import DiskMap
 from ..vram.layers import enable_vram_management
 from .file import load_state_dict, load_metadata_from_safetensors
 import torch
-from contextlib import contextmanager
-from transformers.integrations import is_deepspeed_zero3_enabled
-from transformers.utils import ContextManagers
+import sys
+from contextlib import ExitStack
+
+
+def is_deepspeed_zero3_enabled():
+    if "transformers" not in sys.modules:
+        return False
+    from transformers.integrations import is_deepspeed_zero3_enabled as _is_deepspeed_zero3_enabled
+    return _is_deepspeed_zero3_enabled()
 
 
 def load_model(model_class, path, config=None, torch_dtype=torch.bfloat16, device="cpu", state_dict_converter=None, use_disk_map=False, module_map=None, vram_config=None, vram_limit=None, state_dict=None, quantize=None):
     config = {} if config is None else config
-    with ContextManagers(get_init_context(torch_dtype=torch_dtype, device=device)):
+    with ExitStack() as stack:
+        for init_context in get_init_context(torch_dtype=torch_dtype, device=device):
+            stack.enter_context(init_context)
         model = model_class(**config)
     # What is `module_map`?
     # This is a module mapping table for VRAM management.

{diffsynth-2.1.12 → diffsynth-2.2.0}/diffsynth/core/quant/backends/__init__.py

@@ -5,6 +5,8 @@ _LAZY_BACKENDS = {
     "torchao": ".torchao",
     "comfy_kitchen": ".comfy_kitchen",
     "entropack": ".entropack",
+    "deep_gemm": ".deep_gemm",
+    "deep_gemm_fp4": ".deep_gemm",
 }
 _loaded = set()