diffsynth-engine 0.6.1.dev28__py3-none-any.whl → 0.6.1.dev29__py3-none-any.whl

diffsynth_engine/models/basic/video_sparse_attention.py

@@ -3,10 +3,15 @@ import math
 import functools
 
 from diffsynth_engine.utils.flag import VIDEO_SPARSE_ATTN_AVAILABLE
-from diffsynth_engine.utils.parallel import get_sp_ulysses_group, get_sp_ring_world_size
+from diffsynth_engine.utils.process_group import get_sp_ulysses_group, get_sp_ring_world_size
 
+
+vsa_core = None
 if VIDEO_SPARSE_ATTN_AVAILABLE:
-    from vsa import video_sparse_attn as vsa_core
+    try:
+        from vsa import video_sparse_attn as vsa_core
+    except Exception:
+        vsa_core = None
 
 VSA_TILE_SIZE = (4, 4, 4)
 
@@ -171,6 +176,12 @@ def video_sparse_attn(
     variable_block_sizes: torch.LongTensor,
     non_pad_index: torch.LongTensor,
 ):
+    if vsa_core is None:
+        raise RuntimeError(
+            "Video sparse attention (VSA) is not available. "
+            "Please install the 'vsa' package and ensure all its dependencies (including pytest) are installed."
+        )
+
     q = tile(q, num_tiles, tile_partition_indices, non_pad_index)
     k = tile(k, num_tiles, tile_partition_indices, non_pad_index)
     v = tile(v, num_tiles, tile_partition_indices, non_pad_index)
@@ -212,7 +223,8 @@ def distributed_video_sparse_attn(
 ):
     from yunchang.comm.all_to_all import SeqAllToAll4D
 
-    assert get_sp_ring_world_size() == 1, "distributed video sparse attention requires ring degree to be 1"
+    ring_world_size = get_sp_ring_world_size()
+    assert ring_world_size == 1, "distributed video sparse attention requires ring degree to be 1"
     sp_ulysses_group = get_sp_ulysses_group()
 
     q = SeqAllToAll4D.apply(sp_ulysses_group, q, scatter_idx, gather_idx)

diffsynth_engine/pipelines/wan_video.py

@@ -650,7 +650,7 @@ class WanVideoPipeline(BasePipeline):
                 dit_type = "wan2.2-i2v-a14b"
             elif model_state_dict["high_noise_model"]["patch_embedding.weight"].shape[1] == 16:
                 dit_type = "wan2.2-t2v-a14b"
-        elif model_state_dict["patch_embedding.weight"].shape[1] == 48:
+        elif has_any_key("patch_embedding.weight") and model_state_dict["patch_embedding.weight"].shape[1] == 48:
             dit_type = "wan2.2-ti2v-5b"
         elif has_any_key("img_emb.emb_pos", "condition_embedder.image_embedder.pos_embed"):
             dit_type = "wan2.1-flf2v-14b"
@@ -680,6 +680,30 @@ class WanVideoPipeline(BasePipeline):
             if config.attn_params is None:
                 config.attn_params = VideoSparseAttentionParams(sparsity=0.9)
 
+    def update_weights(self, state_dicts: WanStateDicts) -> None:
+        is_dual_model_state_dict = (isinstance(state_dicts.model, dict) and
+                                     ("high_noise_model" in state_dicts.model or "low_noise_model" in state_dicts.model))
+        is_dual_model_pipeline = self.dit2 is not None
+
+        if is_dual_model_state_dict != is_dual_model_pipeline:
+            raise ValueError(
+                f"Model structure mismatch: pipeline has {'dual' if is_dual_model_pipeline else 'single'} model "
+                f"but state_dict is for {'dual' if is_dual_model_state_dict else 'single'} model. "
+                f"Cannot update weights between WAN 2.1 (single model) and WAN 2.2 (dual model)."
+            )
+
+        if is_dual_model_state_dict:
+            if "high_noise_model" in state_dicts.model:
+                self.update_component(self.dit, state_dicts.model["high_noise_model"], self.config.device, self.config.model_dtype)
+            if "low_noise_model" in state_dicts.model:
+                self.update_component(self.dit2, state_dicts.model["low_noise_model"], self.config.device, self.config.model_dtype)
+        else:
+            self.update_component(self.dit, state_dicts.model, self.config.device, self.config.model_dtype)
+
+        self.update_component(self.text_encoder, state_dicts.t5, self.config.device, self.config.t5_dtype)
+        self.update_component(self.vae, state_dicts.vae, self.config.device, self.config.vae_dtype)
+        self.update_component(self.image_encoder, state_dicts.image_encoder, self.config.device, self.config.image_encoder_dtype)
+
     def compile(self):
         self.dit.compile_repeated_blocks()
         if self.dit2 is not None:

diffsynth_engine/utils/parallel.py

@@ -21,117 +21,33 @@ from queue import Empty
 import diffsynth_engine.models.basic.attention as attention_ops
 from diffsynth_engine.utils.platform import empty_cache
 from diffsynth_engine.utils import logging
+from diffsynth_engine.utils.process_group import (
+    PROCESS_GROUP,
+    get_cfg_group,
+    get_cfg_world_size,
+    get_cfg_rank,
+    get_cfg_ranks,
+    get_sp_group,
+    get_sp_world_size,
+    get_sp_rank,
+    get_sp_ranks,
+    get_sp_ulysses_group,
+    get_sp_ulysses_world_size,
+    get_sp_ulysses_rank,
+    get_sp_ulysses_ranks,
+    get_sp_ring_group,
+    get_sp_ring_world_size,
+    get_sp_ring_rank,
+    get_sp_ring_ranks,
+    get_tp_group,
+    get_tp_world_size,
+    get_tp_rank,
+    get_tp_ranks,
+)
 
 logger = logging.get_logger(__name__)
 
 
-class Singleton:
-    _instance = None
-
-    def __new__(cls, *args, **kwargs):
-        if not cls._instance:
-            cls._instance = super(Singleton, cls).__new__(cls, *args, **kwargs)
-        return cls._instance
-
-
-class ProcessGroupSingleton(Singleton):
-    def __init__(self):
-        self.CFG_GROUP: Optional[dist.ProcessGroup] = None
-        self.SP_GROUP: Optional[dist.ProcessGroup] = None
-        self.SP_ULYSSUES_GROUP: Optional[dist.ProcessGroup] = None
-        self.SP_RING_GROUP: Optional[dist.ProcessGroup] = None
-        self.TP_GROUP: Optional[dist.ProcessGroup] = None
-
-        self.CFG_RANKS: List[int] = []
-        self.SP_RANKS: List[int] = []
-        self.SP_ULYSSUES_RANKS: List[int] = []
-        self.SP_RING_RANKS: List[int] = []
-        self.TP_RANKS: List[int] = []
-
-
-PROCESS_GROUP = ProcessGroupSingleton()
-
-
-def get_cfg_group():
-    return PROCESS_GROUP.CFG_GROUP
-
-
-def get_cfg_world_size():
-    return PROCESS_GROUP.CFG_GROUP.size() if PROCESS_GROUP.CFG_GROUP is not None else 1
-
-
-def get_cfg_rank():
-    return PROCESS_GROUP.CFG_GROUP.rank() if PROCESS_GROUP.CFG_GROUP is not None else 0
-
-
-def get_cfg_ranks():
-    return PROCESS_GROUP.CFG_RANKS
-
-
-def get_sp_group():
-    return PROCESS_GROUP.SP_GROUP
-
-
-def get_sp_world_size():
-    return PROCESS_GROUP.SP_GROUP.size() if PROCESS_GROUP.SP_GROUP is not None else 1
-
-
-def get_sp_rank():
-    return PROCESS_GROUP.SP_GROUP.rank() if PROCESS_GROUP.SP_GROUP is not None else 0
-
-
-def get_sp_ranks():
-    return PROCESS_GROUP.SP_RANKS
-
-
-def get_sp_ulysses_group():
-    return PROCESS_GROUP.SP_ULYSSUES_GROUP
-
-
-def get_sp_ulysses_world_size():
-    return PROCESS_GROUP.SP_ULYSSUES_GROUP.size() if PROCESS_GROUP.SP_ULYSSUES_GROUP is not None else 1
-
-
-def get_sp_ulysses_rank():
-    return PROCESS_GROUP.SP_ULYSSUES_GROUP.rank() if PROCESS_GROUP.SP_ULYSSUES_GROUP is not None else 0
-
-
-def get_sp_ulysses_ranks():
-    return PROCESS_GROUP.SP_ULYSSUES_RANKS
-
-
-def get_sp_ring_group():
-    return PROCESS_GROUP.SP_RING_GROUP
-
-
-def get_sp_ring_world_size():
-    return PROCESS_GROUP.SP_RING_GROUP.size() if PROCESS_GROUP.SP_RING_GROUP is not None else 1
-
-
-def get_sp_ring_rank():
-    return PROCESS_GROUP.SP_RING_GROUP.rank() if PROCESS_GROUP.SP_RING_GROUP is not None else 0
-
-
-def get_sp_ring_ranks():
-    return PROCESS_GROUP.SP_RING_RANKS
-
-
-def get_tp_group():
-    return PROCESS_GROUP.TP_GROUP
-
-
-def get_tp_world_size():
-    return PROCESS_GROUP.TP_GROUP.size() if PROCESS_GROUP.TP_GROUP is not None else 1
-
-
-def get_tp_rank():
-    return PROCESS_GROUP.TP_GROUP.rank() if PROCESS_GROUP.TP_GROUP is not None else 0
-
-
-def get_tp_ranks():
-    return PROCESS_GROUP.TP_RANKS
-
-
 def init_parallel_pgs(
     cfg_degree: int = 1,
     sp_ulysses_degree: int = 1,

diffsynth_engine/utils/process_group.py

@@ -0,0 +1,149 @@
+"""
+Process group management for distributed training.
+
+This module provides singleton-based process group management for distributed training,
+including support for CFG parallelism, sequence parallelism (Ulysses + Ring), and tensor parallelism.
+"""
+
+import torch.distributed as dist
+from typing import Optional, List
+
+
+class Singleton:
+    _instance = None
+
+    def __new__(cls, *args, **kwargs):
+        if not cls._instance:
+            cls._instance = super(Singleton, cls).__new__(cls, *args, **kwargs)
+        return cls._instance
+
+
+class ProcessGroupSingleton(Singleton):
+    def __init__(self):
+        if not hasattr(self, 'initialized'):
+            self.CFG_GROUP: Optional[dist.ProcessGroup] = None
+            self.SP_GROUP: Optional[dist.ProcessGroup] = None
+            self.SP_ULYSSUES_GROUP: Optional[dist.ProcessGroup] = None
+            self.SP_RING_GROUP: Optional[dist.ProcessGroup] = None
+            self.TP_GROUP: Optional[dist.ProcessGroup] = None
+
+            self.CFG_RANKS: List[int] = []
+            self.SP_RANKS: List[int] = []
+            self.SP_ULYSSUES_RANKS: List[int] = []
+            self.SP_RING_RANKS: List[int] = []
+            self.TP_RANKS: List[int] = []
+
+            self.initialized = True
+
+
+PROCESS_GROUP = ProcessGroupSingleton()
+
+
+# CFG parallel group functions
+def get_cfg_group():
+    return PROCESS_GROUP.CFG_GROUP
+
+
+def get_cfg_world_size():
+    return PROCESS_GROUP.CFG_GROUP.size() if PROCESS_GROUP.CFG_GROUP is not None else 1
+
+
+def get_cfg_rank():
+    return PROCESS_GROUP.CFG_GROUP.rank() if PROCESS_GROUP.CFG_GROUP is not None else 0
+
+
+def get_cfg_ranks():
+    return PROCESS_GROUP.CFG_RANKS
+
+
+# Sequence parallel group functions
+def get_sp_group():
+    return PROCESS_GROUP.SP_GROUP
+
+
+def get_sp_world_size():
+    return PROCESS_GROUP.SP_GROUP.size() if PROCESS_GROUP.SP_GROUP is not None else 1
+
+
+def get_sp_rank():
+    return PROCESS_GROUP.SP_GROUP.rank() if PROCESS_GROUP.SP_GROUP is not None else 0
+
+
+def get_sp_ranks():
+    return PROCESS_GROUP.SP_RANKS
+
+
+# Sequence parallel Ulysses group functions
+def get_sp_ulysses_group():
+    return PROCESS_GROUP.SP_ULYSSUES_GROUP
+
+
+def get_sp_ulysses_world_size():
+    return PROCESS_GROUP.SP_ULYSSUES_GROUP.size() if PROCESS_GROUP.SP_ULYSSUES_GROUP is not None else 1
+
+
+def get_sp_ulysses_rank():
+    return PROCESS_GROUP.SP_ULYSSUES_GROUP.rank() if PROCESS_GROUP.SP_ULYSSUES_GROUP is not None else 0
+
+
+def get_sp_ulysses_ranks():
+    return PROCESS_GROUP.SP_ULYSSUES_RANKS
+
+
+# Sequence parallel Ring group functions
+def get_sp_ring_group():
+    return PROCESS_GROUP.SP_RING_GROUP
+
+
+def get_sp_ring_world_size():
+    return PROCESS_GROUP.SP_RING_GROUP.size() if PROCESS_GROUP.SP_RING_GROUP is not None else 1
+
+
+def get_sp_ring_rank():
+    return PROCESS_GROUP.SP_RING_GROUP.rank() if PROCESS_GROUP.SP_RING_GROUP is not None else 0
+
+
+def get_sp_ring_ranks():
+    return PROCESS_GROUP.SP_RING_RANKS
+
+
+# Tensor parallel group functions
+def get_tp_group():
+    return PROCESS_GROUP.TP_GROUP
+
+
+def get_tp_world_size():
+    return PROCESS_GROUP.TP_GROUP.size() if PROCESS_GROUP.TP_GROUP is not None else 1
+
+
+def get_tp_rank():
+    return PROCESS_GROUP.TP_GROUP.rank() if PROCESS_GROUP.TP_GROUP is not None else 0
+
+
+def get_tp_ranks():
+    return PROCESS_GROUP.TP_RANKS
+
+
+__all__ = [
+    "PROCESS_GROUP",
+    "get_cfg_group",
+    "get_cfg_world_size",
+    "get_cfg_rank",
+    "get_cfg_ranks",
+    "get_sp_group",
+    "get_sp_world_size",
+    "get_sp_rank",
+    "get_sp_ranks",
+    "get_sp_ulysses_group",
+    "get_sp_ulysses_world_size",
+    "get_sp_ulysses_rank",
+    "get_sp_ulysses_ranks",
+    "get_sp_ring_group",
+    "get_sp_ring_world_size",
+    "get_sp_ring_rank",
+    "get_sp_ring_ranks",
+    "get_tp_group",
+    "get_tp_world_size",
+    "get_tp_rank",
+    "get_tp_ranks",
+]

{diffsynth_engine-0.6.1.dev28.dist-info → diffsynth_engine-0.6.1.dev29.dist-info}/METADATA

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

{diffsynth_engine-0.6.1.dev28.dist-info → diffsynth_engine-0.6.1.dev29.dist-info}/RECORD

@@ -93,7 +93,7 @@ diffsynth_engine/models/basic/relative_position_emb.py,sha256=rCXOweZMcayVnNUVvB
 diffsynth_engine/models/basic/timestep.py,sha256=WJODYqkSXEM0wcS42YkkfrGwxWt0e60zMTkDdUBQqBw,2810
 diffsynth_engine/models/basic/transformer_helper.py,sha256=6K7A5bVnN2bOoq6I0IQf7RJBhSZUP4jNf1n7NPGu8zA,5287
 diffsynth_engine/models/basic/unet_helper.py,sha256=4lN6F80Ubm6ip4dkLVmB-Og5-Y25Wduhs9Q8qjyzK6E,9044
-diffsynth_engine/models/basic/video_sparse_attention.py,sha256=iXA3sHDLWk1ns1lVCNbZdiaDu94kBIsw-9vrCGAll7g,7843
+diffsynth_engine/models/basic/video_sparse_attention.py,sha256=GxDN6PTpA1rCoQaXUwSPgH4708bEzVI1qsD48WVDXLA,8201
 diffsynth_engine/models/flux/__init__.py,sha256=x0JoxL0CdiiVrY0BjkIrGinud7mcXecLleGO0km91XQ,686
 diffsynth_engine/models/flux/flux_controlnet.py,sha256=NvFKQIx0NldX5uUxdmYwuS2s-xaFRlKotiE6lr3-HRY,8018
 diffsynth_engine/models/flux/flux_dit.py,sha256=7sdV8KFQiHcK-8aqyvXBgC7E_-D9rcgBcnMXUq_AybI,23403
@@ -151,7 +151,7 @@ diffsynth_engine/pipelines/sd_image.py,sha256=nr-Nhsnomq8CsUqhTM3i2l2zG01YjwXdfR
 diffsynth_engine/pipelines/sdxl_image.py,sha256=v7ZACGPb6EcBunL6e5E9jynSQjE7GQx8etEV-ZLP91g,21704
 diffsynth_engine/pipelines/utils.py,sha256=HZbJHErNJS1DhlwJKvZ9dY7Kh8Zdlsw3zE2e88TYGRY,2277
 diffsynth_engine/pipelines/wan_s2v.py,sha256=QHlCLMqlmnp55iYm2mzg4qCq4jceRAP3Zt5Mubz3mAM,29384
-diffsynth_engine/pipelines/wan_video.py,sha256=rJq60LiaCoLq1PkqUzzrdvFkp6h73fc-ZUu0MiMQC-c,29668
+diffsynth_engine/pipelines/wan_video.py,sha256=9xjSvQ4mlVEDdaL6QuUURj4iyxhJ2xABBphQjkfzK8s,31323
 diffsynth_engine/processor/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 diffsynth_engine/processor/canny_processor.py,sha256=hV30NlblTkEFUAmF_O-LJrNlGVM2SFrqq6okfF8VpOo,602
 diffsynth_engine/processor/depth_processor.py,sha256=dQvs3JsnyMbz4dyI9QoR8oO-mMFBFAgNvgqeCoaU5jk,1532
@@ -182,15 +182,16 @@ diffsynth_engine/utils/lock.py,sha256=1Ipgst9eEFfFdViAvD5bxdB6HnHHBcqWYOb__fGaPU
 diffsynth_engine/utils/logging.py,sha256=XB0xTT8PBN6btkOjFtOvjlrOCRVgDGT8PFAp1vmse28,467
 diffsynth_engine/utils/offload.py,sha256=94og79TIkxldwYUgZT3L4OVu1WBlE7gfVPvO2MRhm6c,3551
 diffsynth_engine/utils/onnx.py,sha256=jeWUudJHnESjuiEAHyUZYUZz7dCj34O9aGjHCe8yjWo,1149
-diffsynth_engine/utils/parallel.py,sha256=6T8oCTp-7Gb3qsgNRB2Bp3DF4eyx1FzvS6pFnEJbsek,19789
+diffsynth_engine/utils/parallel.py,sha256=OBGsAK-3ncArRyMU1lea7tbYgxSdCucQvXheL3Ssl5M,17653
 diffsynth_engine/utils/platform.py,sha256=nbpG-XHJFRmYY6u_e7IBQ9Q6GyItrIkKf3VKuBPTUpY,627
+diffsynth_engine/utils/process_group.py,sha256=P-X04a--Zb4M4kjc3DddmusrxCKqv8wiDGhXG4Al-rE,3783
 diffsynth_engine/utils/prompt.py,sha256=YItMchoVzsG6y-LB4vzzDUWrkhKRVlt1HfVhxZjSxMQ,280
 diffsynth_engine/utils/video.py,sha256=8FCaeqIdUsWMgWI_6SO9SPynsToGcLCQAVYFTc4CDhg,2200
 diffsynth_engine/utils/memory/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 diffsynth_engine/utils/memory/linear_regression.py,sha256=oW_EQEw13oPoyUrxiL8A7Ksa5AuJ2ynI2qhCbfAuZbg,3930
 diffsynth_engine/utils/memory/memory_predcit_model.py,sha256=EXprSl_zlVjgfMWNXP-iw83Ot3hyMcgYaRPv-dvyL84,3943
-diffsynth_engine-0.6.1.dev28.dist-info/licenses/LICENSE,sha256=x7aBqQuVI0IYnftgoTPI_A0I_rjdjPPQkjnU6N2nikM,11346
-diffsynth_engine-0.6.1.dev28.dist-info/METADATA,sha256=2LB9DNq9Pf8-sOj9A7_7EzP88Fh--XGF-0hDSP94DPE,1164
-diffsynth_engine-0.6.1.dev28.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-diffsynth_engine-0.6.1.dev28.dist-info/top_level.txt,sha256=6zgbiIzEHLbhgDKRyX0uBJOV3F6VnGGBRIQvSiYYn6w,17
-diffsynth_engine-0.6.1.dev28.dist-info/RECORD,,
+diffsynth_engine-0.6.1.dev29.dist-info/licenses/LICENSE,sha256=x7aBqQuVI0IYnftgoTPI_A0I_rjdjPPQkjnU6N2nikM,11346
+diffsynth_engine-0.6.1.dev29.dist-info/METADATA,sha256=8A5q0qhRMxeJi7IOvP3dcqk58BsgIBxy16ndlnDM_6I,1164
+diffsynth_engine-0.6.1.dev29.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+diffsynth_engine-0.6.1.dev29.dist-info/top_level.txt,sha256=6zgbiIzEHLbhgDKRyX0uBJOV3F6VnGGBRIQvSiYYn6w,17
+diffsynth_engine-0.6.1.dev29.dist-info/RECORD,,