chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

diffusion: fix wan-2.2-TI2V and support sp (#12926)

Browse Source
This commit is contained in:
Mick
2025-11-10 14:37:57 +08:00
committed by GitHub
parent 90401cf7d2
commit e123648b36
12 changed files with 296 additions and 89 deletions
Download Patch File Download Diff File Expand all files Collapse all files

3
python/sglang/multimodal_gen/configs/pipelines/wan.py
View File

@@ -140,8 +140,7 @@ class Wan2_2_TI2V_5B_Config(WanT2V480PConfig):
vae_stride = self.vae_stride
oh = batch.height
ow = batch.width
shape = (z_dim, F, oh // vae_stride[1], ow // vae_stride[2])
shape = (batch_size, z_dim, F, oh // vae_stride[1], ow // vae_stride[2])
return shape
def __post_init__(self) -> None:

3
python/sglang/multimodal_gen/runtime/layers/rotary_embedding.py
View File

@@ -406,6 +406,9 @@ class NDRotaryEmbedding(torch.nn.Module):
start_frame: int = 0,
device: torch.device | str | None = None,
) -> tuple[torch.Tensor, torch.Tensor]:
"""
Handles sp internally
"""
# Caching wrapper: use grid parameters directly as the key.
# grid_tuple = _to_tuple(grid_size, dim=self.ndim)
device_str = str(device) if device is not None else "cpu"

6
python/sglang/multimodal_gen/runtime/models/dits/wanvideo.py
View File

@@ -690,7 +690,7 @@ class WanTransformer3DModel(CachableDiT):
d = self.hidden_size // self.num_attention_heads
self.rope_dim_list = [d - 4 * (d // 6), 2 * (d // 6), 2 * (d // 6)]
self.rope = NDRotaryEmbedding(
self.rotary_emb = NDRotaryEmbedding(
rope_dim_list=self.rope_dim_list,
rope_theta=10000,
dtype=torch.float32 if current_platform.is_mps() else torch.float64,
@@ -725,7 +725,8 @@ class WanTransformer3DModel(CachableDiT):
post_patch_height = height // p_h
post_patch_width = width // p_w
freqs_cos, freqs_sin = self.rope.forward_from_grid(
# The rotary embedding layer correctly handles SP offsets internally.
freqs_cos, freqs_sin = self.rotary_emb.forward_from_grid(
(
post_patch_num_frames * self.sp_size,
post_patch_height,
@@ -746,6 +747,7 @@ class WanTransformer3DModel(CachableDiT):
# timestep shape: batch_size, or batch_size, seq_len (wan 2.2 ti2v)
if timestep.dim() == 2:
# ti2v
ts_seq_len = timestep.shape[1]
timestep = timestep.flatten() # batch_size * seq_len
else:

228
python/sglang/multimodal_gen/runtime/pipelines/stages/denoising.py
View File

@@ -217,9 +217,6 @@ class DenoisingStage(PipelineStage):
target_dtype != torch.float32
) and not server_args.disable_autocast
# Handle sequence parallelism if enabled
self._preprocess_sp_latents(batch)
# Get timesteps and calculate warmup steps
timesteps = batch.timesteps
if timesteps is None:
@@ -263,8 +260,14 @@ class DenoisingStage(PipelineStage):
else:
boundary_timestep = None
# TI2V specific preparations
z, mask2, seq_len = None, None, None
# TI2V specific preparations - BEFORE SP sharding
z, z_sp, reserved_frames_masks, reserved_frames_mask_sp, seq_len = (
None,
None,
None,
None,
None,
)
# FIXME: should probably move to latent preparation stage, to handle with offload
if server_args.pipeline_config.ti2v_task and batch.pil_image is not None:
# Wan2.2 TI2V directly replaces the first frame of the latent with
@@ -285,11 +288,27 @@ class DenoisingStage(PipelineStage):
z = z * self.vae.scaling_factor.to(z.device, z.dtype)
else:
z = z * self.vae.scaling_factor
latent_model_input = latents.to(target_dtype).squeeze(0)
_, mask2 = masks_like([latent_model_input], zero=True)
# z: [B, C, 1, H, W]
latent_model_input = latents.to(target_dtype)
# Keep as [B, C, T, H, W] for proper broadcasting
assert latent_model_input.ndim == 5
latents = (1.0 - mask2[0]) * z + mask2[0] * latent_model_input
# Create mask with proper shape [B, C, T, H, W]
latent_for_mask = latent_model_input.squeeze(0) # [C, T, H, W]
_, reserved_frames_masks = masks_like([latent_for_mask], zero=True)
reserved_frames_mask = reserved_frames_masks[0].unsqueeze(
0
) # [1, C, T, H, W]
# replace GLOBAL first frame with image - proper broadcasting
# z: [B, C, 1, H, W], reserved_frames_mask: [1, C, T, H, W]
# Both will broadcast correctly
latents = (
1.0 - reserved_frames_mask
) * z + reserved_frames_mask * latent_model_input
assert latents.ndim == 5
latents = latents.to(get_local_torch_device())
batch.latents = latents
F = batch.num_frames
temporal_scale = (
@@ -309,6 +328,74 @@ class DenoisingStage(PipelineStage):
int(math.ceil(seq_len / get_sp_world_size())) * get_sp_world_size()
)
# Handle sequence parallelism AFTER TI2V processing
self._preprocess_sp_latents(batch)
latents = batch.latents
# Shard z and reserved_frames_mask for TI2V if SP is enabled
if (
server_args.pipeline_config.ti2v_task
and batch.pil_image is not None
and get_sp_world_size() > 1
):
sp_world_size = get_sp_world_size()
rank_in_sp_group = get_sp_parallel_rank()
if getattr(batch, "did_sp_shard_latents", False):
# Shard z (image latent) along time dimension
# z shape: [1, C, 1, H, W] - only first frame
# Only rank 0 has the first frame after sharding
if z.shape[2] == 1:
# z is single frame, only rank 0 needs it
if rank_in_sp_group == 0:
z_sp = z
else:
# Other ranks don't have the first frame
z_sp = None
else:
# Should not happen for TI2V
z_sp = z
# Shard reserved_frames_mask along time dimension to match sharded latents
# reserved_frames_mask is a list from masks_like, extract reserved_frames_mask[0] first
# reserved_frames_mask[0] shape: [C, T, H, W]
# All ranks need their portion of reserved_frames_mask for timestep calculation
if reserved_frames_masks is not None:
reserved_frames_mask = reserved_frames_masks[
0
] # Extract tensor from list
time_dim = reserved_frames_mask.shape[1] # [C, T, H, W]
if time_dim > 0 and time_dim % sp_world_size == 0:
reserved_frames_mask_sp_tensor = rearrange(
reserved_frames_mask,
"c (n t) h w -> c n t h w",
n=sp_world_size,
).contiguous()
reserved_frames_mask_sp_tensor = reserved_frames_mask_sp_tensor[
:, rank_in_sp_group, :, :, :
]
reserved_frames_mask_sp = (
reserved_frames_mask_sp_tensor # Store as tensor, not list
)
else:
reserved_frames_mask_sp = reserved_frames_mask
else:
reserved_frames_mask_sp = None
else:
# SP not enabled or latents not sharded
z_sp = z
reserved_frames_mask_sp = (
reserved_frames_masks[0]
if reserved_frames_masks is not None
else None
) # Extract tensor
else:
# TI2V not enabled or SP not enabled
z_sp = z
reserved_frames_mask_sp = (
reserved_frames_masks[0] if reserved_frames_masks is not None else None
) # Extract tensor
guidance = self.get_or_build_guidance(
# TODO: replace with raw_latent_shape?
latents.shape[0],
@@ -370,8 +457,9 @@ class DenoisingStage(PipelineStage):
"prompt_embeds": prompt_embeds,
"neg_prompt_embeds": neg_prompt_embeds,
"boundary_timestep": boundary_timestep,
"z": z,
"mask2": mask2,
"z": z_sp, # Use SP-sharded version
# ndim == 5
"reserved_frames_mask": reserved_frames_mask_sp, # Use SP-sharded version
"seq_len": seq_len,
"guidance": guidance,
}
@@ -582,6 +670,75 @@ class DenoisingStage(PipelineStage):
assert current_model is not None, "The model for the current step is not set."
return current_model, current_guidance_scale
def expand_timestep_before_forward(
self,
batch: Req,
server_args: ServerArgs,
t_device,
target_dtype,
seq_len,
reserved_frames_mask,
):
bsz = batch.raw_latent_shape[0]
# expand timestep
if server_args.pipeline_config.ti2v_task and batch.pil_image is not None:
# Explicitly cast t_device to the target float type at the beginning.
# This ensures any precision-based rounding (e.g., float32(999.0) -> bfloat16(1000.0))
# is applied consistently *before* it's used by any rank.
t_device_rounded = t_device.to(target_dtype)
local_seq_len = seq_len
if get_sp_world_size() > 1 and getattr(
batch, "did_sp_shard_latents", False
):
local_seq_len = seq_len // get_sp_world_size()
if get_sp_parallel_rank() == 0 and reserved_frames_mask is not None:
# Rank 0 has the first frame, create a special timestep tensor
# NOTE: The spatial downsampling in the next line is suspicious but kept
# to match original model's potential training configuration.
temp_ts = (
reserved_frames_mask[0][:, ::2, ::2] * t_device_rounded
).flatten()
# Pad to full local sequence length
temp_ts = torch.cat(
[
temp_ts,
temp_ts.new_ones(local_seq_len - temp_ts.size(0))
* t_device_rounded,
]
)
timestep = temp_ts.unsqueeze(0).repeat(bsz, 1)
else:
# Other ranks get a uniform timestep tensor of the correct shape [B, local_seq_len]
timestep = t_device.repeat(bsz, local_seq_len)
else:
timestep = t_device.repeat(bsz)
return timestep
def post_forward_for_ti2v_task(
self, batch: Req, server_args: ServerArgs, reserved_frames_mask, latents, z
):
"""
For Wan2.2 ti2v task, global first frame should be replaced with encoded image after each timestep
"""
if server_args.pipeline_config.ti2v_task and batch.pil_image is not None:
# Apply TI2V mask blending with SP-aware z and reserved_frames_mask.
# This ensures the first frame is always the condition image after each step.
# This is only applied on rank 0, where z is not None.
if z is not None and reserved_frames_mask is not None:
# z: [1, C, 1, H, W]
# latents: [1, C, T_local, H, W]
# reserved_frames_mask: [C, T_local, H, W]
# Unsqueeze mask to [1, C, T_local, H, W] for broadcasting.
# z will broadcast along the time dimension.
latents = (
1.0 - reserved_frames_mask.unsqueeze(0)
) * z + reserved_frames_mask.unsqueeze(0) * latents
return latents
@torch.no_grad()
def forward(
self,
@@ -613,7 +770,7 @@ class DenoisingStage(PipelineStage):
latents = prepared_vars["latents"]
boundary_timestep = prepared_vars["boundary_timestep"]
z = prepared_vars["z"]
mask2 = prepared_vars["mask2"]
reserved_frames_mask = prepared_vars["reserved_frames_mask"]
seq_len = prepared_vars["seq_len"]
guidance = prepared_vars["guidance"]
@@ -663,23 +820,14 @@ class DenoisingStage(PipelineStage):
[latent_model_input, batch.image_latent], dim=1
).to(target_dtype)
# expand timestep
if (
server_args.pipeline_config.ti2v_task
and batch.pil_image is not None
):
timestep = torch.stack([t_device]).to(get_local_torch_device())
temp_ts = (mask2[0][0][:, ::2, ::2] * timestep).flatten()
temp_ts = torch.cat(
[
temp_ts,
temp_ts.new_ones(seq_len - temp_ts.size(0)) * timestep,
]
)
timestep = temp_ts.unsqueeze(0)
t_expand = timestep.repeat(latent_model_input.shape[0], 1)
else:
t_expand = t_device.repeat(latent_model_input.shape[0])
timestep = self.expand_timestep_before_forward(
batch,
server_args,
t_device,
target_dtype,
seq_len,
reserved_frames_mask,
)
latent_model_input = self.scheduler.scale_model_input(
latent_model_input, t_device
@@ -690,7 +838,7 @@ class DenoisingStage(PipelineStage):
noise_pred = self._predict_noise_with_cfg(
current_model,
latent_model_input,
t_expand,
timestep,
batch,
i,
attn_metadata,
@@ -714,12 +862,10 @@ class DenoisingStage(PipelineStage):
**extra_step_kwargs,
return_dict=False,
)[0]
if (
server_args.pipeline_config.ti2v_task
and batch.pil_image is not None
):
latents = latents.squeeze(0)
latents = (1.0 - mask2[0]) * z + mask2[0] * latents
latents = self.post_forward_for_ti2v_task(
batch, server_args, reserved_frames_mask, latents, z
)
# save trajectory latents if needed
if batch.return_trajectory_latents:
@@ -876,7 +1022,7 @@ class DenoisingStage(PipelineStage):
self,
current_model,
latent_model_input,
t_expand,
timestep,
prompt_embeds,
target_dtype,
guidance: torch.Tensor,
@@ -885,7 +1031,7 @@ class DenoisingStage(PipelineStage):
return current_model(
hidden_states=latent_model_input,
encoder_hidden_states=prompt_embeds,
timestep=t_expand,
timestep=timestep,
guidance=guidance,
**kwargs,
)
@@ -894,7 +1040,7 @@ class DenoisingStage(PipelineStage):
self,
current_model: torch.nn.Module,
latent_model_input: torch.Tensor,
t_expand,
timestep,
batch,
timestep_index: int,
attn_metadata,
@@ -913,7 +1059,7 @@ class DenoisingStage(PipelineStage):
Args:
current_model: The transformer model to use for the current step.
latent_model_input: The input latents for the model.
t_expand: The expanded timestep tensor.
timestep: The expanded timestep tensor.
batch: The current batch information.
timestep_index: The current timestep index.
attn_metadata: Attention metadata for custom backends.
@@ -940,7 +1086,7 @@ class DenoisingStage(PipelineStage):
noise_pred_cond = self._predict_noise(
current_model=current_model,
latent_model_input=latent_model_input,
t_expand=t_expand,
timestep=timestep,
prompt_embeds=server_args.pipeline_config.get_pos_prompt_embeds(
batch
),
@@ -968,7 +1114,7 @@ class DenoisingStage(PipelineStage):
noise_pred_uncond = self._predict_noise(
current_model=current_model,
latent_model_input=latent_model_input,
t_expand=t_expand,
timestep=timestep,
prompt_embeds=server_args.pipeline_config.get_neg_prompt_embeds(
batch
),

2
python/sglang/multimodal_gen/runtime/utils/hf_diffusers_utils.py
View File

@@ -370,7 +370,7 @@ def maybe_download_model(
logger.info(
"Downloading model snapshot from HF Hub for %s...", model_name_or_path
)
with get_lock(model_name_or_path):
with get_lock(model_name_or_path).acquire(poll_interval=2):
local_path = snapshot_download(
repo_id=model_name_or_path,
ignore_patterns=["*.onnx", "*.msgpack"],

6
python/sglang/multimodal_gen/test/cli/test_generate_t2v_perf.py
View File

@@ -20,6 +20,10 @@ class TestFastWan2_1_T2V(TestGenerateBase):
"test_mixed": 15.0,
}
# disabled for vsa
def test_usp(self):
pass
class TestFastWan2_2_T2V(TestGenerateBase):
model_path = "FastVideo/FastWan2.2-TI2V-5B-FullAttn-Diffusers"
@@ -38,7 +42,7 @@ class TestWan2_1_T2V(TestGenerateBase):
extra_args = []
data_type: DataType = DataType.VIDEO
thresholds = {
"test_single_gpu": 76.0,
"test_single_gpu": 76.0 * 1.05,
"test_cfg_parallel": 46.5 * 1.05,
"test_usp": 22.5,
"test_mixed": 26.5,

32
python/sglang/multimodal_gen/test/cli/test_generate_ti2v_perf.py
View File

@@ -1,5 +1,3 @@
# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo
import unittest
from sglang.multimodal_gen.configs.sample.base import DataType
@@ -18,7 +16,8 @@ class TestGenerateTI2VBase(TestGenerateBase):
"sglang",
"generate",
f'--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."',
"--image-path=https://github.com/Wan-Video/Wan2.2/blob/990af50de458c19590c245151197326e208d7191/examples/i2v_input.JPG?raw=true",
"--image-path",
"https://github.com/Wan-Video/Wan2.2/blob/990af50de458c19590c245151197326e208d7191/examples/i2v_input.JPG?raw=true",
"--save-output",
"--log-level=debug",
f"--output-path={cls.output_path}",
@@ -36,12 +35,15 @@ class TestGenerateTI2VBase(TestGenerateBase):
class TestWan2_1_I2V_14B_480P(TestGenerateTI2VBase):
model_path = "Wan-AI/Wan2.1-I2V-14B-480P-Diffusers"
extra_args = ["--attention-backend=video_sparse_attn"]
thresholds = {
"test_single_gpu": 13.0,
"test_cfg_parallel": 191.7 * 1.05,
"test_usp": 15.0,
"test_mixed": 15.0,
"test_usp": 530.5 * 1.05,
}
class TestWan2_1_I2V_14B_720P(TestGenerateTI2VBase):
model_path = "Wan-AI/Wan2.1-I2V-14B-720P-Diffusers"
thresholds = {
"test_usp": 530.5 * 1.05,
}
@@ -50,13 +52,19 @@ class TestWan2_2_TI2V_5B(TestGenerateTI2VBase):
# FIXME: doesn't work with vsa at the moment
# extra_args = ["--attention-backend=video_sparse_attn"]
thresholds = {
"test_single_gpu": 13.0,
"test_cfg_parallel": 191.7 * 1.05,
"test_usp": 387.6 * 1.05,
"test_mixed": 15.0,
"test_usp": 82.3 * 1.05,
}
# OOM
# class TestWan2_2_I2V_A14B(TestGenerateTI2VBase):
# model_path = "Wan-AI/Wan2.2-I2V-A14B-Diffusers"
# # FIXME: doesn't work with vsa at the moment
# thresholds = {
# "test_usp": 66.3 * 1.05,
# }
if __name__ == "__main__":
del TestGenerateTI2VBase, TestGenerateBase
unittest.main()

2
python/sglang/multimodal_gen/test/test_offline_api.py
View File

@@ -19,7 +19,7 @@ class TestGeneratorAPIBase(unittest.TestCase):
server_kwargs = {}
# sampling
output_path: str = "outputs"
output_path: str = "test_outputs"
results = []

57
python/sglang/multimodal_gen/test/test_utils.py
View File

@@ -1,4 +1,5 @@
# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo
import dataclasses
import os
import shlex
import socket
@@ -6,6 +7,7 @@ import subprocess
import sys
import time
import unittest
from typing import Optional
from PIL import Image
@@ -15,7 +17,7 @@ from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
logger = init_logger(__name__)
def run_command(command):
def run_command(command) -> Optional[float]:
"""Runs a command and returns the execution time and status."""
print(f"Running command: {' '.join(command)}")
@@ -75,6 +77,18 @@ def check_image_size(ut, image, width, height):
ut.assertEqual(image.size, (width, height))
@dataclasses.dataclass
class TestResult:
name: str
key: str
duration: Optional[float]
succeed: bool
@property
def duration_str(self):
return f"{self.duration:.4f}" if self.duration else "NA"
class TestCLIBase(unittest.TestCase):
model_path: str = None
extra_args = []
@@ -84,7 +98,7 @@ class TestCLIBase(unittest.TestCase):
width: int = 720
height: int = 720
output_path: str = "outputs"
output_path: str = "test_outputs"
base_command = [
"sglang",
@@ -105,7 +119,7 @@ class TestCLIBase(unittest.TestCase):
def setUpClass(cls):
cls.results = []
def _run_command(self, name, model_path: str, test_key: str = "", args=[]):
def _run_command(self, name: str, model_path: str, test_key: str = "", args=[]):
command = (
self.base_command
+ [f"--model-path={model_path}"]
@@ -115,11 +129,10 @@ class TestCLIBase(unittest.TestCase):
)
duration = run_command(command)
status = "Success" if duration else "Failed"
succeed = duration is not None
duration_str = f"{duration:.4f}s" if duration else "NA"
self.__class__.results.append(
{"name": name, "key": test_key, "duration": duration_str, "status": status}
)
duration = float(duration) if succeed else None
self.results.append(TestResult(name, test_key, duration, succeed))
return name, duration, status
@@ -133,7 +146,7 @@ class TestGenerateBase(TestCLIBase):
width: int = 720
height: int = 720
output_path: str = "outputs"
output_path: str = "test_outputs"
image_path: str | None = None
prompt: str | None = "A curious raccoon"
@@ -150,7 +163,7 @@ class TestGenerateBase(TestCLIBase):
f"--output-path={output_path}",
]
results = []
results: list[TestResult] = []
@classmethod
def setUpClass(cls):
@@ -167,24 +180,28 @@ class TestGenerateBase(TestCLIBase):
test_key: order for order, test_key in enumerate(test_keys)
}
ordered_results: list[dict] = [{}] * len(test_keys)
ordered_results: list[TestResult] = [None] * len(test_keys)
for result in cls.results:
order = test_key_to_order[result["key"]]
order = test_key_to_order[result.key]
ordered_results[order] = result
for result in ordered_results:
if not result:
continue
status = (
result["status"] and result["duration"] <= cls.thresholds[result["key"]]
"Succeed"
if (
result.succeed
and float(result.duration) <= float(cls.thresholds[result.key])
)
else "Failed"
)
print(f"| {result['name']:<30} | {result['duration']:<8} | {status:<7} |")
print(f"| {result.name:<30} | {result.duration_str:<8} | {status:<7} |")
print()
durations = [result["duration"] for result in cls.results]
durations = [result.duration_str for result in cls.results]
print(" | ".join([""] + durations + [""]))
def _run_test(self, name, args, model_path: str, test_key: str):
def _run_test(self, name: str, args, model_path: str, test_key: str):
time_threshold = self.thresholds[test_key]
name, duration, status = self._run_command(
name, args=args, model_path=model_path, test_key=test_key
@@ -220,7 +237,7 @@ class TestGenerateBase(TestCLIBase):
def test_single_gpu(self):
"""single gpu"""
self._run_test(
name=f"{self.model_name()}, single gpu",
name=f"{self.model_name()}_single gpu",
args=None,
model_path=self.model_path,
test_key="test_single_gpu",
@@ -231,7 +248,7 @@ class TestGenerateBase(TestCLIBase):
if self.data_type == DataType.IMAGE:
return
self._run_test(
name=f"{self.model_name()}, cfg parallel",
name=f"{self.model_name()}_cfg parallel",
args="--num-gpus 2 --enable-cfg-parallel",
model_path=self.model_path,
test_key="test_cfg_parallel",
@@ -242,7 +259,7 @@ class TestGenerateBase(TestCLIBase):
if self.data_type == DataType.IMAGE:
return
self._run_test(
name=f"{self.model_name()}, usp",
name=f"{self.model_name()}_usp",
args="--num-gpus 4 --ulysses-degree=2 --ring-degree=2",
model_path=self.model_path,
test_key="test_usp",
@@ -253,7 +270,7 @@ class TestGenerateBase(TestCLIBase):
if self.data_type == DataType.IMAGE:
return
self._run_test(
name=f"{self.model_name()}, mixed",
name=f"{self.model_name()}_mixed",
args="--num-gpus 4 --ulysses-degree=2 --ring-degree=1 --enable-cfg-parallel",
model_path=self.model_path,
test_key="test_mixed",

34
python/sglang/multimodal_gen/utils.py
View File

@@ -698,12 +698,38 @@ def is_vmoba_available() -> bool:
# adapted from: https://github.com/Wan-Video/Wan2.2/blob/main/wan/utils/utils.py
def masks_like(
tensor, zero=False, generator=None, p=0.2
tensors, zero=False, generator=None, p=0.2
) -> tuple[list[torch.Tensor], list[torch.Tensor]]:
assert isinstance(tensor, list)
out1 = [torch.ones(u.shape, dtype=u.dtype, device=u.device) for u in tensor]
"""
Generate binary masks for Text-to-Image-to-Video (TI2V) tasks.
out2 = [torch.ones(u.shape, dtype=u.dtype, device=u.device) for u in tensor]
Creates masks to control which frames should be preserved vs replaced.
Primarily used to fix the first frame to the input image while generating other frames.
Args:
tensors: List of tensors with shape [C, T, H, W]
zero: If True, set first frame (dim 1, index 0) to zero. Default: False
generator: Optional random generator for stochastic masking
p: Probability of applying special noise when generator is provided. Default: 0.2
Returns:
Tuple of two lists of tensors:
- When zero=False: Both lists contain all-ones tensors
- When zero=True (no generator): First frame set to 0, others to 1
- When zero=True (with generator): First frame set to small random values with probability p
Example:
>>> latent = torch.randn(48, 69, 96, 160) # [C, T, H, W]
>>> _, mask = masks_like([latent], zero=True)
>>> # mask[0][:, 0] == 0 (first frame)
>>> # mask[0][:, 1:] == 1 (other frames)
>>> blended = (1.0 - mask[0]) * image + mask[0] * latent
>>> # Result: first frame = image, other frames = latent
"""
assert isinstance(tensors, list)
out1 = [torch.ones(u.shape, dtype=u.dtype, device=u.device) for u in tensors]
out2 = [torch.ones(u.shape, dtype=u.dtype, device=u.device) for u in tensors]
if zero:
if generator is not None: