diffsynth-engine 0.5.1.dev4__py3-none-any.whl → 0.6.1.dev25__py3-none-any.whl

diffsynth_engine/pipelines/wan_s2v.py

@@ -239,7 +239,15 @@ class WanSpeech2VideoPipeline(WanVideoPipeline):
 
         return ref_latents, motion_latents, motion_frames
 
-    def encode_pose(self, pose_video: List[Image.Image], pose_video_fps: int, num_clips: int, num_frames_per_clip: int, height: int, width: int):
+    def encode_pose(
+        self,
+        pose_video: List[Image.Image],
+        pose_video_fps: int,
+        num_clips: int,
+        num_frames_per_clip: int,
+        height: int,
+        width: int,
+    ):
         self.load_models_to_device(["vae"])
         max_num_pose_frames = num_frames_per_clip * num_clips
         pose_video = read_n_frames(pose_video, pose_video_fps, max_num_pose_frames, target_fps=self.config.fps)
@@ -386,6 +394,7 @@ class WanSpeech2VideoPipeline(WanVideoPipeline):
         void_audio_input: torch.Tensor | None = None,
     ):
         latents = latents.to(dtype=self.config.model_dtype, device=self.device)
+        attn_kwargs = self.get_attn_kwargs(latents)
 
         noise_pred = model(
             x=latents,
@@ -400,6 +409,7 @@ class WanSpeech2VideoPipeline(WanVideoPipeline):
             drop_motion_frames=drop_motion_frames,
             audio_mask=audio_mask,
             void_audio_input=void_audio_input,
+            attn_kwargs=attn_kwargs,
         )
         return noise_pred
 
@@ -466,7 +476,9 @@ class WanSpeech2VideoPipeline(WanVideoPipeline):
                 dtype=self.dtype,
             ).to(self.device)
         if pose_video is not None:
-            pose_latents_all_clips = self.encode_pose(pose_video, pose_video_fps, num_clips, num_frames_per_clip, height, width)
+            pose_latents_all_clips = self.encode_pose(
+                pose_video, pose_video_fps, num_clips, num_frames_per_clip, height, width
+            )
 
         output_frames_all_clips = []
         for clip_idx in range(num_clips):
@@ -602,7 +614,9 @@ class WanSpeech2VideoPipeline(WanVideoPipeline):
         return cls.from_state_dict(state_dicts, config)
 
     @classmethod
-    def from_state_dict(cls, state_dicts: WanS2VStateDicts, config: WanSpeech2VideoPipelineConfig) -> "WanSpeech2VideoPipeline":
+    def from_state_dict(
+        cls, state_dicts: WanS2VStateDicts, config: WanSpeech2VideoPipelineConfig
+    ) -> "WanSpeech2VideoPipeline":
         if config.parallelism > 1:
             pipe = ParallelWrapper(
                 cfg_degree=config.cfg_degree,
@@ -617,7 +631,9 @@ class WanSpeech2VideoPipeline(WanVideoPipeline):
         return pipe
 
     @classmethod
-    def _from_state_dict(cls, state_dicts: WanS2VStateDicts, config: WanSpeech2VideoPipelineConfig) -> "WanSpeech2VideoPipeline":
+    def _from_state_dict(
+        cls, state_dicts: WanS2VStateDicts, config: WanSpeech2VideoPipelineConfig
+    ) -> "WanSpeech2VideoPipeline":
         # default params from model config
         vae_type = "wan2.1-vae"
         dit_type = "wan2.2-s2v-14b"
@@ -632,25 +648,20 @@ class WanSpeech2VideoPipeline(WanVideoPipeline):
         init_device = "cpu" if config.offload_mode is not None else config.device
         tokenizer = WanT5Tokenizer(WAN_TOKENIZER_CONF_PATH, seq_len=512, clean="whitespace")
         text_encoder = WanTextEncoder.from_state_dict(state_dicts.t5, device=init_device, dtype=config.t5_dtype)
-        vae = WanVideoVAE.from_state_dict(state_dicts.vae, config=vae_config, device=init_device, dtype=config.vae_dtype)
+        vae = WanVideoVAE.from_state_dict(
+            state_dicts.vae, config=vae_config, device=init_device, dtype=config.vae_dtype
+        )
         audio_encoder = Wav2Vec2Model.from_state_dict(
             state_dicts.audio_encoder, config=Wav2Vec2Config(), device=init_device, dtype=config.audio_encoder_dtype
         )
 
         with LoRAContext():
-            attn_kwargs = {
-                "attn_impl": config.dit_attn_impl,
-                "sparge_smooth_k": config.sparge_smooth_k,
-                "sparge_cdfthreshd": config.sparge_cdfthreshd,
-                "sparge_simthreshd1": config.sparge_simthreshd1,
-                "sparge_pvthreshd": config.sparge_pvthreshd,
-            }
             dit = WanS2VDiT.from_state_dict(
                 state_dicts.model,
                 config=model_config,
-                device=init_device,
+                device=("cpu" if config.use_fsdp else init_device),
                 dtype=config.model_dtype,
-                attn_kwargs=attn_kwargs,
+                use_vsa=(config.dit_attn_impl.value == "vsa"),
             )
             if config.use_fp8_linear:
                 enable_fp8_linear(dit)

diffsynth_engine/pipelines/wan_video.py

@@ -95,8 +95,14 @@ class WanLoRAConverter(LoRAStateDictConverter):
         return state_dict
 
 
+class WanLowNoiseLoRAConverter(WanLoRAConverter):
+    def convert(self, state_dict):
+        return {"dit2": super().convert(state_dict)["dit"]}
+
+
 class WanVideoPipeline(BasePipeline):
     lora_converter = WanLoRAConverter()
+    low_noise_lora_converter = WanLowNoiseLoRAConverter()
 
     def __init__(
         self,
@@ -133,7 +139,13 @@ class WanVideoPipeline(BasePipeline):
         self.image_encoder = image_encoder
         self.model_names = ["text_encoder", "dit", "dit2", "vae", "image_encoder"]
 
-    def load_loras(self, lora_list: List[Tuple[str, float]], fused: bool = True, save_original_weight: bool = False):
+    def load_loras(
+        self,
+        lora_list: List[Tuple[str, float]],
+        fused: bool = True,
+        save_original_weight: bool = False,
+        lora_converter: Optional[WanLoRAConverter] = None,
+    ):
         assert self.config.tp_degree is None or self.config.tp_degree == 1, (
             "load LoRA is not allowed when tensor parallel is enabled; "
             "set tp_degree=None or tp_degree=1 during pipeline initialization"
@@ -142,10 +154,24 @@ class WanVideoPipeline(BasePipeline):
             "load fused LoRA is not allowed when fully sharded data parallel is enabled; "
             "either load LoRA with fused=False or set use_fsdp=False during pipeline initialization"
         )
-        super().load_loras(lora_list, fused, save_original_weight)
+        super().load_loras(lora_list, fused, save_original_weight, lora_converter)
+
+    def load_loras_low_noise(
+        self, lora_list: List[Tuple[str, float]], fused: bool = True, save_original_weight: bool = False
+    ):
+        assert self.dit2 is not None, "low noise LoRA can only be applied to Wan2.2"
+        self.load_loras(lora_list, fused, save_original_weight, self.low_noise_lora_converter)
+
+    def load_loras_high_noise(
+        self, lora_list: List[Tuple[str, float]], fused: bool = True, save_original_weight: bool = False
+    ):
+        assert self.dit2 is not None, "high noise LoRA can only be applied to Wan2.2"
+        self.load_loras(lora_list, fused, save_original_weight)
 
     def unload_loras(self):
         self.dit.unload_loras()
+        if self.dit2 is not None:
+            self.dit2.unload_loras()
         self.text_encoder.unload_loras()
 
     def get_default_fps(self) -> int:
@@ -301,6 +327,7 @@ class WanVideoPipeline(BasePipeline):
 
     def predict_noise(self, model, latents, image_clip_feature, image_y, timestep, context):
         latents = latents.to(dtype=self.config.model_dtype, device=self.device)
+        attn_kwargs = self.get_attn_kwargs(latents)
 
         noise_pred = model(
             x=latents,
@@ -308,6 +335,7 @@ class WanVideoPipeline(BasePipeline):
             context=context,
             clip_feature=image_clip_feature,
             y=image_y,
+            attn_kwargs=attn_kwargs,
         )
         return noise_pred
 
@@ -556,19 +584,12 @@ class WanVideoPipeline(BasePipeline):
             dit_state_dict = state_dicts.model
 
         with LoRAContext():
-            attn_kwargs = {
-                "attn_impl": config.dit_attn_impl,
-                "sparge_smooth_k": config.sparge_smooth_k,
-                "sparge_cdfthreshd": config.sparge_cdfthreshd,
-                "sparge_simthreshd1": config.sparge_simthreshd1,
-                "sparge_pvthreshd": config.sparge_pvthreshd,
-            }
             dit = WanDiT.from_state_dict(
                 dit_state_dict,
                 config=dit_config,
-                device=init_device,
+                device=("cpu" if config.use_fsdp else init_device),
                 dtype=config.model_dtype,
-                attn_kwargs=attn_kwargs,
+                use_vsa=(config.dit_attn_impl.value == "vsa"),
             )
             if config.use_fp8_linear:
                 enable_fp8_linear(dit)
@@ -578,9 +599,9 @@ class WanVideoPipeline(BasePipeline):
                 dit2 = WanDiT.from_state_dict(
                     dit2_state_dict,
                     config=dit_config,
-                    device=init_device,
+                    device=("cpu" if config.use_fsdp else init_device),
                     dtype=config.model_dtype,
-                    attn_kwargs=attn_kwargs,
+                    use_vsa=(config.dit_attn_impl.value == "vsa"),
                 )
                 if config.use_fp8_linear:
                     enable_fp8_linear(dit2)
@@ -618,19 +639,22 @@ class WanVideoPipeline(BasePipeline):
     @staticmethod
     def _get_dit_type(model_state_dict: Dict[str, torch.Tensor] | Dict[str, Dict[str, torch.Tensor]]) -> str:
         # determine wan dit type by model params
+        def has_any_key(*xs):
+            return any(x in model_state_dict for x in xs)
+
         dit_type = None
-        if "high_noise_model" in model_state_dict and "low_noise_model" in model_state_dict:
+        if has_any_key("high_noise_model"):
             if model_state_dict["high_noise_model"]["patch_embedding.weight"].shape[1] == 36:
                 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:
             dit_type = "wan2.2-ti2v-5b"
-        elif "img_emb.emb_pos" in model_state_dict:
+        elif has_any_key("img_emb.emb_pos", "condition_embedder.image_embedder.pos_embed"):
             dit_type = "wan2.1-flf2v-14b"
-        elif "img_emb.proj.0.weight" in model_state_dict:
+        elif has_any_key("img_emb.proj.0.weight", "condition_embedder.image_embedder.norm1"):
             dit_type = "wan2.1-i2v-14b"
-        elif "blocks.39.self_attn.norm_q.weight" in model_state_dict:
+        elif has_any_key("blocks.39.self_attn.norm_q.weight", "blocks.39.attn1.norm_q.weight"):
             dit_type = "wan2.1-t2v-14b"
         else:
             dit_type = "wan2.1-t2v-1.3b"
@@ -645,6 +669,6 @@ class WanVideoPipeline(BasePipeline):
         return vae_type
 
     def compile(self):
-        self.dit.compile_repeated_blocks(dynamic=True)
+        self.dit.compile_repeated_blocks()
         if self.dit2 is not None:
-            self.dit2.compile_repeated_blocks(dynamic=True)
+            self.dit2.compile_repeated_blocks()

diffsynth_engine/tokenizers/base.py

@@ -1,10 +1,16 @@
 # Modified from transformers.tokenization_utils_base
 from typing import Dict, List, Union, overload
+from enum import Enum
 
 
 TOKENIZER_CONFIG_FILE = "tokenizer_config.json"
 
 
+class PaddingStrategy(str, Enum):
+    LONGEST = "longest"
+    MAX_LENGTH = "max_length"
+
+
 class BaseTokenizer:
     SPECIAL_TOKENS_ATTRIBUTES = [
         "bos_token",

diffsynth_engine/tokenizers/qwen2.py

@@ -4,7 +4,7 @@ import torch
 from typing import Dict, List, Union, Optional
 from tokenizers import Tokenizer as TokenizerFast, AddedToken
 
-from diffsynth_engine.tokenizers.base import BaseTokenizer, TOKENIZER_CONFIG_FILE
+from diffsynth_engine.tokenizers.base import BaseTokenizer, PaddingStrategy, TOKENIZER_CONFIG_FILE
 
 
 VOCAB_FILES_NAMES = {"vocab_file": "vocab.json", "merges_file": "merges.txt", "tokenizer_file": "tokenizer.json"}
@@ -165,22 +165,28 @@ class Qwen2TokenizerFast(BaseTokenizer):
         texts: Union[str, List[str]],
         max_length: Optional[int] = None,
         padding_side: Optional[str] = None,
+        padding_strategy: Union[PaddingStrategy, str] = "longest",
         **kwargs,
     ) -> Dict[str, "torch.Tensor"]:
         """
         Tokenize text and prepare for model inputs.
 
         Args:
-            text (`str`, `List[str]`, *optional*):
+            texts (`str`, `List[str]`):
                 The sequence or batch of sequences to be encoded.
 
             max_length (`int`, *optional*):
-                Each encoded sequence will be truncated or padded to max_length.
+                Maximum length of the encoded sequences.
 
             padding_side (`str`, *optional*):
                 The side on which the padding should be applied. Should be selected between `"right"` and `"left"`.
                 Defaults to `"right"`.
 
+            padding_strategy (`PaddingStrategy`, `str`, *optional*):
+                If `"longest"`, will pad the sequences to the longest sequence in the batch.
+                If `"max_length"`, will pad the sequences to the `max_length` argument.
+                Defaults to `"longest"`.
+
         Returns:
             `Dict[str, "torch.Tensor"]`: tensor dict compatible with model_input_names.
         """
@@ -190,7 +196,9 @@ class Qwen2TokenizerFast(BaseTokenizer):
 
         batch_ids = self.batch_encode(texts)
         ids_lens = [len(ids_) for ids_ in batch_ids]
-        max_length = max_length if max_length is not None else min(max(ids_lens), self.model_max_length)
+        max_length = max_length if max_length is not None else self.model_max_length
+        if padding_strategy == PaddingStrategy.LONGEST:
+            max_length = min(max(ids_lens), max_length)
         padding_side = padding_side if padding_side is not None else self.padding_side
 
         encoded = torch.zeros(len(texts), max_length, dtype=torch.long)

diffsynth_engine/utils/constants.py

@@ -27,18 +27,19 @@ SD3_TEXT_ENCODER_CONFIG_FILE = os.path.join(CONF_PATH, "models", "sd3", "sd3_tex
 SDXL_TEXT_ENCODER_CONFIG_FILE = os.path.join(CONF_PATH, "models", "sdxl", "sdxl_text_encoder.json")
 SDXL_UNET_CONFIG_FILE = os.path.join(CONF_PATH, "models", "sdxl", "sdxl_unet.json")
 
-WAN2_1_DIT_T2V_1_3B_CONFIG_FILE = os.path.join(CONF_PATH, "models", "wan", "dit", "wan2.1-t2v-1.3b.json")
-WAN2_1_DIT_T2V_14B_CONFIG_FILE = os.path.join(CONF_PATH, "models", "wan", "dit", "wan2.1-t2v-14b.json")
-WAN2_1_DIT_I2V_14B_CONFIG_FILE = os.path.join(CONF_PATH, "models", "wan", "dit", "wan2.1-i2v-14b.json")
-WAN2_1_DIT_FLF2V_14B_CONFIG_FILE = os.path.join(CONF_PATH, "models", "wan", "dit", "wan2.1-flf2v-14b.json")
-WAN2_2_DIT_TI2V_5B_CONFIG_FILE = os.path.join(CONF_PATH, "models", "wan", "dit", "wan2.2-ti2v-5b.json")
-WAN2_2_DIT_T2V_A14B_CONFIG_FILE = os.path.join(CONF_PATH, "models", "wan", "dit", "wan2.2-t2v-a14b.json")
-WAN2_2_DIT_I2V_A14B_CONFIG_FILE = os.path.join(CONF_PATH, "models", "wan", "dit", "wan2.2-i2v-a14b.json")
-WAN2_2_DIT_S2V_14B_CONFIG_FILE = os.path.join(CONF_PATH, "models", "wan", "dit", "wan2.2-s2v-14b.json")
-
-WAN2_1_VAE_CONFIG_FILE = os.path.join(CONF_PATH, "models", "wan", "vae", "wan2.1-vae.json")
-WAN2_2_VAE_CONFIG_FILE = os.path.join(CONF_PATH, "models", "wan", "vae", "wan2.2-vae.json")
-WAN_VAE_KEYMAP_FILE = os.path.join(CONF_PATH, "models", "wan", "vae", "wan-vae-keymap.json")
+WAN2_1_DIT_T2V_1_3B_CONFIG_FILE = os.path.join(CONF_PATH, "models", "wan", "dit", "wan2.1_t2v_1.3b.json")
+WAN2_1_DIT_T2V_14B_CONFIG_FILE = os.path.join(CONF_PATH, "models", "wan", "dit", "wan2.1_t2v_14b.json")
+WAN2_1_DIT_I2V_14B_CONFIG_FILE = os.path.join(CONF_PATH, "models", "wan", "dit", "wan2.1_i2v_14b.json")
+WAN2_1_DIT_FLF2V_14B_CONFIG_FILE = os.path.join(CONF_PATH, "models", "wan", "dit", "wan2.1_flf2v_14b.json")
+WAN2_2_DIT_TI2V_5B_CONFIG_FILE = os.path.join(CONF_PATH, "models", "wan", "dit", "wan2.2_ti2v_5b.json")
+WAN2_2_DIT_T2V_A14B_CONFIG_FILE = os.path.join(CONF_PATH, "models", "wan", "dit", "wan2.2_t2v_a14b.json")
+WAN2_2_DIT_I2V_A14B_CONFIG_FILE = os.path.join(CONF_PATH, "models", "wan", "dit", "wan2.2_i2v_a14b.json")
+WAN2_2_DIT_S2V_14B_CONFIG_FILE = os.path.join(CONF_PATH, "models", "wan", "dit", "wan2.2_s2v_14b.json")
+WAN_DIT_KEYMAP_FILE = os.path.join(CONF_PATH, "models", "wan", "dit", "wan_dit_keymap.json")
+
+WAN2_1_VAE_CONFIG_FILE = os.path.join(CONF_PATH, "models", "wan", "vae", "wan2.1_vae.json")
+WAN2_2_VAE_CONFIG_FILE = os.path.join(CONF_PATH, "models", "wan", "vae", "wan2.2_vae.json")
+WAN_VAE_KEYMAP_FILE = os.path.join(CONF_PATH, "models", "wan", "vae", "wan_vae_keymap.json")
 
 QWEN_IMAGE_CONFIG_FILE = os.path.join(CONF_PATH, "models", "qwen_image", "qwen2_5_vl_config.json")
 QWEN_IMAGE_VISION_CONFIG_FILE = os.path.join(CONF_PATH, "models", "qwen_image", "qwen2_5_vl_vision_config.json")

diffsynth_engine/utils/download.py

@@ -12,7 +12,7 @@ from modelscope import snapshot_download
 from modelscope.hub.api import HubApi
 from diffsynth_engine.utils import logging
 from diffsynth_engine.utils.lock import HeartbeatFileLock
-from diffsynth_engine.utils.env import DIFFSYNTH_FILELOCK_DIR, DIFFSYNTH_CACHE
+from diffsynth_engine.utils.env import DIFFSYNTH_FILELOCK_DIR, DIFFSYNTH_CACHE, MS_HUB_OFFLINE
 from diffsynth_engine.utils.constants import MB
 
 logger = logging.get_logger(__name__)
@@ -81,7 +81,9 @@ def fetch_modelscope_model(
         api.login(access_token)
     with HeartbeatFileLock(lock_file_path):
         directory = os.path.join(DIFFSYNTH_CACHE, "modelscope", model_id, revision if revision else "__version")
-        dirpath = snapshot_download(model_id, revision=revision, local_dir=directory, allow_patterns=path)
+        dirpath = snapshot_download(
+            model_id, revision=revision, local_dir=directory, allow_patterns=path, local_files_only=MS_HUB_OFFLINE
+        )
 
     if isinstance(path, str):
         path = glob.glob(os.path.join(dirpath, path))

diffsynth_engine/utils/env.py

@@ -8,3 +8,5 @@ DIFFSYNTH_CACHE = os.environ.get("DIFFSYNTH_CACHE", os.path.join(HOME, ".cache",
 DIFFSYNTH_FILELOCK_DIR = os.environ.get(
     "DIFFSYNTH_FILELOCK_DIR", os.path.join(HOME, ".cache", "diffsynth", "filelocks")
 )
+
+MS_HUB_OFFLINE = os.getenv("MS_HUB_OFFLINE", "0").lower() in ("1", "true", "yes")

diffsynth_engine/utils/flag.py

@@ -44,3 +44,9 @@ if SPARGE_ATTN_AVAILABLE:
     logger.info("Sparge attention is available")
 else:
     logger.info("Sparge attention is not available")
+
+VIDEO_SPARSE_ATTN_AVAILABLE = importlib.util.find_spec("vsa") is not None
+if VIDEO_SPARSE_ATTN_AVAILABLE:
+    logger.info("Video sparse attention is available")
+else:
+    logger.info("Video sparse attention is not available")

diffsynth_engine/utils/loader.py

@@ -9,12 +9,10 @@ try:
 
     use_fast_safetensors = True
 except ImportError:
-    from safetensors.torch import load_file as _load_file
-
     use_fast_safetensors = False
 
 
-def load_file(path: str | os.PathLike, device: str = "cpu"):
+def load_file(path: str | os.PathLike, device: str = "cpu", need_metadata: bool = False):
     if use_fast_safetensors:
         logger.info(f"FastSafetensors load model from {path}")
         start_time = time.time()
@@ -24,13 +22,34 @@ def load_file(path: str | os.PathLike, device: str = "cpu"):
             direct_io=(os.environ.get("FAST_SAFETENSORS_DIRECT_IO", "False").upper() == "TRUE"),
         )
         logger.info(f"FastSafetensors Load Model End. Time: {time.time() - start_time:.2f}s")
-        return {k: v.to(device) for k, v in result.items()}
+        state_dict = {k: v.to(device) for k, v in result.items()}
+
+        if need_metadata:
+            # FastSafetensors不直接支持metadata，需要用标准safetensors获取
+            from safetensors import safe_open
+
+            with safe_open(str(path), framework="pt", device="cpu") as f:
+                metadata = f.metadata()
+            return state_dict, metadata
+        else:
+            return state_dict
     else:
         logger.info(f"Safetensors load model from {path}")
         start_time = time.time()
-        result = _load_file(path, device=device)
+
+        from safetensors import safe_open
+
+        with safe_open(path, framework="pt", device="cpu") as f:
+            state_dict = {k: f.get_tensor(k).to(device) for k in f.keys()}
+            if need_metadata:
+                metadata = f.metadata()
+
         logger.info(f"Safetensors Load Model End. Time: {time.time() - start_time:.2f}s")
-        return result
+
+        if need_metadata:
+            return state_dict, metadata
+        else:
+            return state_dict
 
 
 save_file = _save_file

diffsynth_engine/utils/parallel.py

@@ -8,19 +8,17 @@ import torch.multiprocessing as mp
 import torch.distributed as dist
 from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
 from torch.distributed.fsdp import ShardingStrategy
-from torch.distributed.fsdp.wrap import lambda_auto_wrap_policy
+from torch.distributed.fsdp.wrap import transformer_auto_wrap_policy
 from torch.distributed.device_mesh import DeviceMesh
 from torch.distributed.tensor.parallel.style import ParallelStyle
 from torch.distributed.tensor.parallel._utils import _validate_tp_mesh_dim
 from contextlib import contextmanager
 from datetime import timedelta
 from functools import partial
-from typing import Dict, List, Union, Optional
+from typing import Dict, List, Set, Type, Union, Optional
 from queue import Empty
 
 import diffsynth_engine.models.basic.attention as attention_ops
-from diffsynth_engine.models import PreTrainedModel
-from diffsynth_engine.pipelines import BasePipeline
 from diffsynth_engine.utils.platform import empty_cache
 from diffsynth_engine.utils import logging
 
@@ -40,10 +38,14 @@ 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] = []
 
 
@@ -82,6 +84,38 @@ 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
 
@@ -127,23 +161,32 @@ def init_parallel_pgs(
     blocks = [list(range(world_size))]
     cfg_groups, cfg_blocks = make_parallel_groups(blocks, cfg_degree)
     for cfg_ranks in cfg_groups:
-        cfg_group = dist.new_group(cfg_ranks)
         if rank in cfg_ranks:
-            PROCESS_GROUP.CFG_GROUP = cfg_group
+            PROCESS_GROUP.CFG_GROUP = dist.new_group(cfg_ranks)
             PROCESS_GROUP.CFG_RANKS = cfg_ranks
 
     sp_groups, sp_blocks = make_parallel_groups(cfg_blocks, sp_degree)
     for sp_ranks in sp_groups:
-        group = dist.new_group(sp_ranks)
         if rank in sp_ranks:
-            PROCESS_GROUP.SP_GROUP = group
+            PROCESS_GROUP.SP_GROUP = dist.new_group(sp_ranks)
             PROCESS_GROUP.SP_RANKS = sp_ranks
 
+    sp_ulysses_groups, sp_ulysses_blocks = make_parallel_groups(cfg_blocks, sp_ulysses_degree)
+    for sp_ulysses_ranks in sp_ulysses_groups:
+        if rank in sp_ulysses_ranks:
+            PROCESS_GROUP.SP_ULYSSUES_GROUP = dist.new_group(sp_ulysses_ranks)
+            PROCESS_GROUP.SP_ULYSSUES_RANKS = sp_ulysses_ranks
+
+    sp_ring_groups, _ = make_parallel_groups(sp_ulysses_blocks, sp_ring_degree)
+    for sp_ring_ranks in sp_ring_groups:
+        if rank in sp_ring_ranks:
+            PROCESS_GROUP.SP_RING_GROUP = dist.new_group(sp_ring_ranks)
+            PROCESS_GROUP.SP_RING_RANKS = sp_ring_ranks
+
     tp_groups, _ = make_parallel_groups(sp_blocks, tp_degree)
     for tp_ranks in tp_groups:
-        group = dist.new_group(tp_ranks)
         if rank in tp_ranks:
-            PROCESS_GROUP.TP_GROUP = group
+            PROCESS_GROUP.TP_GROUP = dist.new_group(tp_ranks)
             PROCESS_GROUP.TP_RANKS = tp_ranks
 
     set_seq_parallel_pg(sp_ulysses_degree, sp_ring_degree, rank, world_size)
@@ -174,25 +217,14 @@ def to_device(data, device):
 def shard_model(
     module: nn.Module,
     device_id: int | torch.device,
+    wrap_module_cls: Set[Type[nn.Module]],
     sharding_strategy: ShardingStrategy = ShardingStrategy.FULL_SHARD,
-    wrap_module_names: Optional[List[str]] = None,
 ):
-    wrap_module_names = wrap_module_names or []
-
-    def wrap_fn(m):
-        for name in wrap_module_names:
-            submodule = getattr(module, name)
-            if isinstance(submodule, nn.ModuleList) and m in submodule:
-                return True
-            elif not isinstance(submodule, nn.ModuleList) and m is submodule:
-                return True
-        return False
-
     return FSDP(
         module,
         device_id=device_id,
         sharding_strategy=sharding_strategy,
-        auto_wrap_policy=partial(lambda_auto_wrap_policy, lambda_fn=wrap_fn),
+        auto_wrap_policy=partial(transformer_auto_wrap_policy, transformer_layer_cls=wrap_module_cls),
     )
 
 
@@ -266,14 +298,15 @@ def _worker_loop(
             world_size=world_size,
         )
 
-        def wrap_for_parallel(module: Union[PreTrainedModel, BasePipeline]):
-            if isinstance(module, BasePipeline):
-                for model_name in module.model_names:
-                    if isinstance(submodule := getattr(module, model_name), PreTrainedModel):
+        def wrap_for_parallel(module):
+            if hasattr(module, "model_names"):
+                for model_name in getattr(module, "model_names"):
+                    submodule = getattr(module, model_name)
+                    if getattr(submodule, "_supports_parallelization", False):
                         setattr(module, model_name, wrap_for_parallel(submodule))
                 return module
 
-            if not module._supports_parallelization:
+            if not getattr(module, "_supports_parallelization", False):
                 return module
 
             if tp_degree > 1:
@@ -283,7 +316,7 @@ def _worker_loop(
                     parallelize_plan=module.get_tp_plan(),
                 )
             elif use_fsdp:
-                module = shard_model(module, device_id=device, wrap_module_names=module.get_fsdp_modules())
+                module = shard_model(module, device_id=device, wrap_module_cls=module.get_fsdp_module_cls())
             return module
 
         module = None

diffsynth_engine/utils/video.py

@@ -41,7 +41,9 @@ def save_video(frames, save_path, fps=15):
         writer.write(frames, fps=fps, codec=codec)
 
 
-def read_n_frames(frames: List[Image.Image], original_fps: int, n_frames: int, target_fps: int = 16) -> List[Image.Image]:
+def read_n_frames(
+    frames: List[Image.Image], original_fps: int, n_frames: int, target_fps: int = 16
+) -> List[Image.Image]:
     num_frames = len(frames)
     interval = max(1, round(original_fps / target_fps))
     sampled_frames: List[Image.Image] = []

{diffsynth_engine-0.5.1.dev4.dist-info → diffsynth_engine-0.6.1.dev25.dist-info}/METADATA

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