update wan22 ti2v i2v parallel

lightx2v/models/networks/wan/infer/pre_infer.py

@@ -42,7 +42,7 @@ class WanPreInfer:
         else:
             timestep = self.scheduler.timesteps[self.scheduler.step_index]
             t = torch.stack([timestep])
-            if hasattr(self.scheduler, "mask"):
+            if self.config["model_cls"] == "wan2.2" and self.config["task"] == "i2v":
                 t = (self.scheduler.mask[0][:, ::2, ::2] * t).flatten()
 
         if positive:

lightx2v/models/networks/wan/model.py

@@ -380,41 +380,37 @@ class WanModel:
 
     @torch.no_grad()
     def _seq_parallel_pre_process(self, pre_infer_out):
-        embed, x, embed0 = pre_infer_out.embed, pre_infer_out.x, pre_infer_out.embed0
-
+        x = pre_infer_out.x
         world_size = dist.get_world_size(self.seq_p_group)
         cur_rank = dist.get_rank(self.seq_p_group)
 
         padding_size = (world_size - (x.shape[0] % world_size)) % world_size
-
         if padding_size > 0:
-            # 使用 F.pad 填充第一维
-            x = F.pad(x, (0, 0, 0, padding_size))  # (后维度填充, 前维度填充)
+            x = F.pad(x, (0, 0, 0, padding_size))
 
         x = torch.chunk(x, world_size, dim=0)[cur_rank]
-        # if self.config["model_cls"] == "wan2.2":
-        #     padding_size = (world_size - (embed0.shape[0] % world_size)) % world_size
-        #     if padding_size > 0:
-        #         embed0 = F.pad(embed0, (0, 0, 0, 0, 0, padding_size))  # (后维度填充, 前维度填充)
-        #         embed = F.pad(embed, (0, 0, 0, padding_size))
 
-        pre_infer_out.x = x
+        if self.config["model_cls"] == "wan2.2" and self.config["task"] == "i2v":
+            embed, embed0 = pre_infer_out.embed, pre_infer_out.embed0
+
+            padding_size = (world_size - (embed.shape[0] % world_size)) % world_size
+            if padding_size > 0:
+                embed = F.pad(embed, (0, 0, 0, padding_size))
+                embed0 = F.pad(embed0, (0, 0, 0, 0, 0, padding_size))
+
+            embed = torch.chunk(embed, world_size, dim=0)[cur_rank]
+            embed0 = torch.chunk(embed0, world_size, dim=0)[cur_rank]
             pre_infer_out.embed = embed
             pre_infer_out.embed0 = embed0
 
+        pre_infer_out.x = x
+
         return pre_infer_out
 
     @torch.no_grad()
     def _seq_parallel_post_process(self, x):
         world_size = dist.get_world_size(self.seq_p_group)
-
-        # 创建一个列表，用于存储所有进程的输出
         gathered_x = [torch.empty_like(x) for _ in range(world_size)]
-
-        # 收集所有进程的输出
         dist.all_gather(gathered_x, x, group=self.seq_p_group)
-
-        # 在指定的维度上合并所有进程的输出
         combined_output = torch.cat(gathered_x, dim=0)
-
-        return combined_output  # 返回合并后的输出
+        return combined_output

scripts/dist_infer/run_wan22_ti2v_i2v_cfg.sh

@@ -11,9 +11,10 @@ source ${lightx2v_path}/scripts/base/base.sh
 
 torchrun --nproc_per_node=2 -m lightx2v.infer \
 --model_cls wan2.2 \
---task t2v \
+--task i2v \
 --model_path $model_path \
 --config_json ${lightx2v_path}/configs/dist_infer/wan22_ti2v_i2v_cfg.json \
---prompt "Two anthropomorphic cats in comfy boxing gear and bright gloves fight intensely on a spotlighted stage" \
+--prompt "Summer beach vacation style, a white cat wearing sunglasses sits on a surfboard. The fluffy-furred feline gazes directly at the camera with a relaxed expression. Blurred beach scenery forms the background featuring crystal-clear waters, distant green hills, and a blue sky dotted with white clouds. The cat assumes a naturally relaxed posture, as if savoring the sea breeze and warm sunlight. A close-up shot highlights the feline's intricate details and the refreshing atmosphere of the seaside." \
 --negative_prompt "色调艳丽，过曝，静态，细节模糊不清，字幕，风格，作品，画作，画面，静止，整体发灰，最差质量，低质量，JPEG压缩残留，丑陋的，残缺的，多余的手指，画得不好的手部，画得不好的脸部，畸形的，毁容的，形态畸形的肢体，手指融合，静止不动的画面，杂乱的背景，三条腿，背景人很多，倒着走" \
+--image_path ${lightx2v_path}/assets/inputs/imgs/img_0.jpg \
 --save_video_path ${lightx2v_path}/save_results/output_lightx2v_wan22_ti2v_i2v_parallel_cfg.mp4

scripts/dist_infer/run_wan22_ti2v_i2v_cfg_ulysses.sh

@@ -11,9 +11,10 @@ source ${lightx2v_path}/scripts/base/base.sh
 
 torchrun --nproc_per_node=8 -m lightx2v.infer \
 --model_cls wan2.2 \
---task t2v \
+--task i2v \
 --model_path $model_path \
 --config_json ${lightx2v_path}/configs/dist_infer/wan22_ti2v_i2v_cfg_ulysses.json \
---prompt "Two anthropomorphic cats in comfy boxing gear and bright gloves fight intensely on a spotlighted stage" \
+--prompt "Summer beach vacation style, a white cat wearing sunglasses sits on a surfboard. The fluffy-furred feline gazes directly at the camera with a relaxed expression. Blurred beach scenery forms the background featuring crystal-clear waters, distant green hills, and a blue sky dotted with white clouds. The cat assumes a naturally relaxed posture, as if savoring the sea breeze and warm sunlight. A close-up shot highlights the feline's intricate details and the refreshing atmosphere of the seaside." \
 --negative_prompt "色调艳丽，过曝，静态，细节模糊不清，字幕，风格，作品，画作，画面，静止，整体发灰，最差质量，低质量，JPEG压缩残留，丑陋的，残缺的，多余的手指，画得不好的手部，画得不好的脸部，畸形的，毁容的，形态畸形的肢体，手指融合，静止不动的画面，杂乱的背景，三条腿，背景人很多，倒着走" \
+--image_path ${lightx2v_path}/assets/inputs/imgs/img_0.jpg \
 --save_video_path ${lightx2v_path}/save_results/output_lightx2v_wan22_ti2v_i2v_parallel_cfg_ulysses.mp4

scripts/dist_infer/run_wan22_ti2v_i2v_ulysses.sh

@@ -11,9 +11,10 @@ source ${lightx2v_path}/scripts/base/base.sh
 
 torchrun --nproc_per_node=4 -m lightx2v.infer \
 --model_cls wan2.2 \
---task t2v \
+--task i2v \
 --model_path $model_path \
 --config_json ${lightx2v_path}/configs/dist_infer/wan22_ti2v_i2v_ulysses.json \
---prompt "Two anthropomorphic cats in comfy boxing gear and bright gloves fight intensely on a spotlighted stage" \
+--prompt "Summer beach vacation style, a white cat wearing sunglasses sits on a surfboard. The fluffy-furred feline gazes directly at the camera with a relaxed expression. Blurred beach scenery forms the background featuring crystal-clear waters, distant green hills, and a blue sky dotted with white clouds. The cat assumes a naturally relaxed posture, as if savoring the sea breeze and warm sunlight. A close-up shot highlights the feline's intricate details and the refreshing atmosphere of the seaside." \
 --negative_prompt "色调艳丽，过曝，静态，细节模糊不清，字幕，风格，作品，画作，画面，静止，整体发灰，最差质量，低质量，JPEG压缩残留，丑陋的，残缺的，多余的手指，画得不好的手部，画得不好的脸部，畸形的，毁容的，形态畸形的肢体，手指融合，静止不动的画面，杂乱的背景，三条腿，背景人很多，倒着走" \
---save_video_path ${lightx2v_path}/save_results/output_lightx2v_wan22_ti2v_i2v_ulysses.mp4
+--image_path ${lightx2v_path}/assets/inputs/imgs/img_0.jpg \
+--save_video_path ${lightx2v_path}/save_results/output_lightx2v_wan22_ti2v_i2v_parallel_ulysses.mp4