diff --git a/fastvideo/configs/pipelines/wan.py b/fastvideo/configs/pipelines/wan.py index 2d90a98f3..19b83c1a9 100644 --- a/fastvideo/configs/pipelines/wan.py +++ b/fastvideo/configs/pipelines/wan.py @@ -12,13 +12,13 @@ from fastvideo.configs.models.vaes import WanVAEConfig from fastvideo.configs.pipelines.base import PipelineConfig -def t5_postprocess_text(outputs: BaseEncoderOutput) -> torch.tensor: - mask: torch.tensor = outputs.attention_mask - hidden_state: torch.tensor = outputs.last_hidden_state +def t5_postprocess_text(outputs: BaseEncoderOutput) -> torch.Tensor: + mask: torch.Tensor = outputs.attention_mask + hidden_state: torch.Tensor = outputs.last_hidden_state seq_lens = mask.gt(0).sum(dim=1).long() assert torch.isnan(hidden_state).sum() == 0 prompt_embeds = [u[:v] for u, v in zip(hidden_state, seq_lens, strict=True)] - prompt_embeds_tensor: torch.tensor = torch.stack([ + prompt_embeds_tensor: torch.Tensor = torch.stack([ torch.cat([u, u.new_zeros(512 - u.size(0), u.size(1))]) for u in prompt_embeds ], @@ -39,12 +39,12 @@ class WanT2V480PConfig(PipelineConfig): vae_sp: bool = False # Denoising stage - flow_shift: int = 3 + flow_shift: float | None = 3.0 # Text encoding stage text_encoder_configs: tuple[EncoderConfig, ...] = field( default_factory=lambda: (T5Config(), )) - postprocess_text_funcs: tuple[Callable[[BaseEncoderOutput], torch.tensor], + postprocess_text_funcs: tuple[Callable[[BaseEncoderOutput], torch.Tensor], ...] = field(default_factory=lambda: (t5_postprocess_text, )) @@ -68,7 +68,7 @@ class WanT2V720PConfig(WanT2V480PConfig): # WanConfig-specific parameters with defaults # Denoising stage - flow_shift: int = 5 + flow_shift: float | None = 5.0 @dataclass @@ -94,7 +94,7 @@ class WanI2V720PConfig(WanI2V480PConfig): # WanConfig-specific parameters with defaults # Denoising stage - flow_shift: int = 5 + flow_shift: float | None = 5.0 @dataclass @@ -104,7 +104,7 @@ class FastWan2_1_T2V_480P_Config(WanT2V480PConfig): # WanConfig-specific parameters with defaults # Denoising stage - flow_shift: int = 8 + flow_shift: float | None = 8.0 dmd_denoising_steps: list[int] | None = field( default_factory=lambda: [1000, 757, 522]) @@ -115,7 +115,7 @@ class FastWan2_1_T2V_480P_Config(WanT2V480PConfig): @dataclass class Wan2_2_TI2V_5B_Config(WanT2V480PConfig): - flow_shift: int = 5 + flow_shift: float | None = 5.0 ti2v_task: bool = True def __post_init__(self) -> None: @@ -125,7 +125,7 @@ class Wan2_2_TI2V_5B_Config(WanT2V480PConfig): @dataclass class FastWan2_2_TI2V_5B_Config(Wan2_2_TI2V_5B_Config): - flow_shift: int = 5 + flow_shift: float | None = 5.0 dmd_denoising_steps: list[int] | None = field( default_factory=lambda: [1000, 757, 522]) @@ -146,5 +146,7 @@ class Wan2_2_I2V_A14B_Config(WanT2V480PConfig): @dataclass class SelfForcingWanT2V480PConfig(WanT2V480PConfig): is_causal: bool = True + flow_shift: float | None = 5.0 dmd_denoising_steps: list[int] | None = field( default_factory=lambda: [1000, 750, 500, 250]) + warp_denoising_step: bool = True diff --git a/fastvideo/layers/layernorm.py b/fastvideo/layers/layernorm.py index fd81db6d6..d93e1d966 100644 --- a/fastvideo/layers/layernorm.py +++ b/fastvideo/layers/layernorm.py @@ -100,7 +100,16 @@ class ScaleResidual(nn.Module): def forward(self, residual: torch.Tensor, x: torch.Tensor, gate: torch.Tensor) -> torch.Tensor: """Apply gated residual connection.""" - return residual + x * gate + # x.shape: [batch_size, seq_len, inner_dim] + if gate.dim() == 4: + # gate.shape: [batch_size, num_frames, 1, inner_dim] + num_frames = gate.shape[1] + frame_seqlen = x.shape[1] // num_frames + return residual + (x.unflatten( + dim=1, sizes=(num_frames, frame_seqlen)) * gate).flatten(1, 2) + else: + # gate.shape: [batch_size, 1, inner_dim] + return residual + x * gate # adapted from Diffusers: https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/normalization.py @@ -159,7 +168,7 @@ class ScaleResidualLayerNormScaleShift(nn.Module): raise NotImplementedError(f"Norm type {norm_type} not implemented") def forward(self, residual: torch.Tensor, x: torch.Tensor, - gate: torch.Tensor, shift: torch.Tensor, + gate: torch.Tensor | int, shift: torch.Tensor, scale: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: """ Apply gated residual connection, followed by layernorm and @@ -171,12 +180,41 @@ class ScaleResidualLayerNormScaleShift(nn.Module): - residual value (value after residual connection but before normalization) """ + # x.shape: [batch_size, seq_len, inner_dim] # Apply residual connection with gating - residual_output = residual + x * gate + if isinstance(gate, int): + # used by cross-attention, should be 1 + assert gate == 1 + residual_output = residual + x + elif isinstance(gate, torch.Tensor): + if gate.dim() == 4: + # gate.shape: [batch_size, num_frames, 1, inner_dim] + num_frames = gate.shape[1] + frame_seqlen = x.shape[1] // num_frames + residual_output = residual + ( + x.unflatten(dim=1, sizes=(num_frames, frame_seqlen)) * + gate).flatten(1, 2) + else: + # used by bidirectional self attention + # gate.shape: [batch_size, 1, inner_dim] + residual_output = residual + x * gate + else: + raise ValueError(f"Gate type {type(gate)} not supported") + # residual_output.shape: [batch_size, seq_len, inner_dim] + # Apply normalization normalized = self.norm(residual_output) # Apply scale and shift - modulated = normalized * (1.0 + scale) + shift + if isinstance(scale, torch.Tensor) and scale.dim() == 4: + # scale.shape: [batch_size, num_frames, 1, inner_dim] + # shift.shape: [batch_size, num_frames, 1, inner_dim] + num_frames = scale.shape[1] + frame_seqlen = normalized.shape[1] // num_frames + modulated = ( + normalized.unflatten(dim=1, sizes=(num_frames, frame_seqlen)) * + (1.0 + scale) + shift).flatten(1, 2) + else: + modulated = normalized * (1.0 + scale) + shift return modulated, residual_output @@ -218,8 +256,24 @@ class LayerNormScaleShift(nn.Module): def forward(self, x: torch.Tensor, shift: torch.Tensor, scale: torch.Tensor) -> torch.Tensor: """Apply ln followed by scale and shift in a single fused operation.""" + # x.shape: [batch_size, seq_len, inner_dim] normalized = self.norm(x) if self.compute_dtype == torch.float32: - return (normalized.float() * (1.0 + scale) + shift).to(x.dtype) + normalized = normalized.float() + + if scale.dim() == 4: + # scale.shape: [batch_size, num_frames, 1, inner_dim] + num_frames = scale.shape[1] + frame_seqlen = normalized.shape[1] // num_frames + output = ( + normalized.unflatten(dim=1, sizes=(num_frames, frame_seqlen)) * + (1.0 + scale) + shift).flatten(1, 2) else: - return normalized * (1.0 + scale) + shift + # scale.shape: [batch_size, 1, inner_dim] + # shift.shape: [batch_size, 1, inner_dim] + output = normalized * (1.0 + scale) + shift + + if self.compute_dtype == torch.float32: + output = output.to(x.dtype) + + return output diff --git a/fastvideo/models/dits/causal_wanvideo.py b/fastvideo/models/dits/causal_wanvideo.py index 993a8044c..a936485b6 100644 --- a/fastvideo/models/dits/causal_wanvideo.py +++ b/fastvideo/models/dits/causal_wanvideo.py @@ -244,19 +244,26 @@ class CausalWanTransformerBlock(nn.Module): current_start: int = 0, cache_start: int | None = None, ) -> torch.Tensor: + # hidden_states.shape: [batch_size, seq_length, inner_dim] + # temb.shape: [batch_size, num_frames, 6, inner_dim] if hidden_states.dim() == 4: hidden_states = hidden_states.squeeze(1) + num_frames = temb.shape[1] + frame_seqlen = hidden_states.shape[1] // num_frames bs, seq_length, _ = hidden_states.shape orig_dtype = hidden_states.dtype # assert orig_dtype != torch.float32 e = self.scale_shift_table + temb.float() + # e.shape: [batch_size, num_frames, 6, inner_dim] + assert e.shape == (bs, num_frames, 6, self.hidden_dim) shift_msa, scale_msa, gate_msa, c_shift_msa, c_scale_msa, c_gate_msa = e.chunk( - 6, dim=1) + 6, dim=2) + # *_msa.shape: [batch_size, num_frames, 1, inner_dim] assert shift_msa.dtype == torch.float32 # 1. Self-attention - norm_hidden_states = (self.norm1(hidden_states.float()) * - (1 + scale_msa) + shift_msa).to(orig_dtype) + norm_hidden_states = (self.norm1(hidden_states.float()).unflatten(dim=1, sizes=(num_frames, frame_seqlen)) * + (1 + scale_msa) + shift_msa).flatten(1, 2).to(orig_dtype) query, _ = self.to_q(norm_hidden_states) key, _ = self.to_k(norm_hidden_states) value, _ = self.to_v(norm_hidden_states) @@ -487,8 +494,8 @@ class CausalWanTransformer3DModel(BaseDiT): hidden_states = hidden_states.flatten(2).transpose(1, 2) temb, timestep_proj, encoder_hidden_states, encoder_hidden_states_image = self.condition_embedder( - timestep, encoder_hidden_states, encoder_hidden_states_image) - timestep_proj = timestep_proj.unflatten(1, (6, -1)) + timestep.flatten(), encoder_hidden_states, encoder_hidden_states_image) + timestep_proj = timestep_proj.unflatten(1, (6, self.hidden_size)).unflatten(dim=0, sizes=timestep.shape) if encoder_hidden_states_image is not None: encoder_hidden_states = torch.concat( @@ -526,8 +533,9 @@ class CausalWanTransformer3DModel(BaseDiT): **causal_kwargs) # 5. Output norm, projection & unpatchify - shift, scale = (self.scale_shift_table + temb.unsqueeze(1)).chunk(2, - dim=1) + temb = temb.unflatten(dim=0, sizes=timestep.shape).unsqueeze(2) + shift, scale = (self.scale_shift_table.unsqueeze(1) + temb).chunk(2, + dim=2) hidden_states = self.norm_out(hidden_states, shift, scale) hidden_states = self.proj_out(hidden_states) @@ -596,8 +604,8 @@ class CausalWanTransformer3DModel(BaseDiT): hidden_states = hidden_states.flatten(2).transpose(1, 2) temb, timestep_proj, encoder_hidden_states, encoder_hidden_states_image = self.condition_embedder( - timestep, encoder_hidden_states, encoder_hidden_states_image) - timestep_proj = timestep_proj.unflatten(1, (6, -1)) + timestep.flatten(), encoder_hidden_states, encoder_hidden_states_image) + timestep_proj = timestep_proj.unflatten(1, (6, self.hidden_size)).unflatten(dim=0, sizes=timestep.shape) if encoder_hidden_states_image is not None: encoder_hidden_states = torch.concat( @@ -623,8 +631,9 @@ class CausalWanTransformer3DModel(BaseDiT): block_mask=self.block_mask) # 5. Output norm, projection & unpatchify - shift, scale = (self.scale_shift_table + temb.unsqueeze(1)).chunk(2, - dim=1) + temb = temb.unflatten(dim=0, sizes=timestep.shape).unsqueeze(2) + shift, scale = (self.scale_shift_table.unsqueeze(1) + temb).chunk(2, + dim=2) hidden_states = self.norm_out(hidden_states, shift, scale) hidden_states = self.proj_out(hidden_states) diff --git a/fastvideo/models/registry.py b/fastvideo/models/registry.py index 68be9bc13..2b919627f 100644 --- a/fastvideo/models/registry.py +++ b/fastvideo/models/registry.py @@ -61,6 +61,9 @@ _SCHEDULERS = { "FlowMatchEulerDiscreteScheduler"), "UniPCMultistepScheduler": ("schedulers", "scheduling_unipc_multistep", "UniPCMultistepScheduler"), + "SelfForcingFlowMatchScheduler": + ("schedulers", "scheduling_self_forcing_flow_match", + "SelfForcingFlowMatchScheduler"), } _FAST_VIDEO_MODELS = { diff --git a/fastvideo/models/schedulers/scheduling_self_forcing_flow_match.py b/fastvideo/models/schedulers/scheduling_self_forcing_flow_match.py new file mode 100644 index 000000000..1d76c14ae --- /dev/null +++ b/fastvideo/models/schedulers/scheduling_self_forcing_flow_match.py @@ -0,0 +1,122 @@ +# SPDX-License-Identifier: Apache-2.0 +from diffusers.configuration_utils import ConfigMixin, register_to_config +from diffusers.schedulers.scheduling_utils import SchedulerMixin +from diffusers.utils import BaseOutput +import torch + +from fastvideo.logger import init_logger +from fastvideo.models.schedulers.base import BaseScheduler + +logger = init_logger(__name__) + + +class SelfForcingFlowMatchSchedulerOutput(BaseOutput): + """ + Output class for the scheduler's `step` function output. + + Args: + prev_sample (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)` for images): + Computed sample `(x_{t-1})` of previous timestep. `prev_sample` should be used as next model input in the + denoising loop. + """ + prev_sample: torch.FloatTensor + +class SelfForcingFlowMatchScheduler(BaseScheduler, ConfigMixin, SchedulerMixin): + + order = 1 + def __init__(self, num_inference_steps=100, num_train_timesteps=1000, shift=3.0, sigma_max=1.0, sigma_min=0.003 / 1.002, inverse_timesteps=False, extra_one_step=False, reverse_sigmas=False, training=False): + self.num_train_timesteps = num_train_timesteps + self.shift = shift + self.sigma_max = sigma_max + self.sigma_min = sigma_min + self.inverse_timesteps = inverse_timesteps + self.extra_one_step = extra_one_step + self.reverse_sigmas = reverse_sigmas + self.set_timesteps(num_inference_steps, training=training) + + def set_timesteps(self, num_inference_steps=100, denoising_strength=1.0, training=False, return_dict=False, **kwargs): + sigma_start = self.sigma_min + \ + (self.sigma_max - self.sigma_min) * denoising_strength + if self.extra_one_step: + self.sigmas = torch.linspace( + sigma_start, self.sigma_min, num_inference_steps + 1)[:-1] + else: + self.sigmas = torch.linspace( + sigma_start, self.sigma_min, num_inference_steps) + if self.inverse_timesteps: + self.sigmas = torch.flip(self.sigmas, dims=[0]) + self.sigmas = self.shift * self.sigmas / \ + (1 + (self.shift - 1) * self.sigmas) + if self.reverse_sigmas: + self.sigmas = 1 - self.sigmas + self.timesteps = self.sigmas * self.num_train_timesteps + if training: + x = self.timesteps + y = torch.exp(-2 * ((x - num_inference_steps / 2) / + num_inference_steps) ** 2) + y_shifted = y - y.min() + bsmntw_weighing = y_shifted * \ + (num_inference_steps / y_shifted.sum()) + self.linear_timesteps_weights = bsmntw_weighing + + def step(self, model_output: torch.FloatTensor, timestep: torch.FloatTensor, sample: torch.FloatTensor, to_final=False, return_dict=False, **kwargs): + if timestep.ndim == 2: + timestep = timestep.flatten(0, 1) + self.sigmas = self.sigmas.to(model_output.device) + self.timesteps = self.timesteps.to(model_output.device) + timestep_id = torch.argmin( + (self.timesteps.unsqueeze(0) - timestep.unsqueeze(1)).abs(), dim=1) + sigma = self.sigmas[timestep_id].reshape(-1, 1, 1, 1) + if to_final or (timestep_id + 1 >= len(self.timesteps)).any(): + sigma_ = 1 if ( + self.inverse_timesteps or self.reverse_sigmas) else 0 + else: + sigma_ = self.sigmas[timestep_id + 1].reshape(-1, 1, 1, 1) + prev_sample = sample + model_output * (sigma_ - sigma) + if isinstance(prev_sample, torch.Tensor | float) and not return_dict: + return (prev_sample, ) + return SelfForcingFlowMatchSchedulerOutput(prev_sample=prev_sample) + + def add_noise(self, original_samples, noise, timestep): + """ + Diffusion forward corruption process. + Input: + - clean_latent: the clean latent with shape [B*T, C, H, W] + - noise: the noise with shape [B*T, C, H, W] + - timestep: the timestep with shape [B*T] + Output: the corrupted latent with shape [B*T, C, H, W] + """ + if timestep.ndim == 2: + timestep = timestep.flatten(0, 1) + self.sigmas = self.sigmas.to(noise.device) + self.timesteps = self.timesteps.to(noise.device) + timestep_id = torch.argmin( + (self.timesteps.unsqueeze(0) - timestep.unsqueeze(1)).abs(), dim=1) + sigma = self.sigmas[timestep_id].reshape(-1, 1, 1, 1) + sample = (1 - sigma) * original_samples + sigma * noise + return sample.type_as(noise) + + def training_target(self, sample, noise, timestep): + target = noise - sample + return target + + def training_weight(self, timestep): + """ + Input: + - timestep: the timestep with shape [B*T] + Output: the corresponding weighting [B*T] + """ + if timestep.ndim == 2: + timestep = timestep.flatten(0, 1) + self.linear_timesteps_weights = self.linear_timesteps_weights.to(timestep.device) + timestep_id = torch.argmin( + (self.timesteps.unsqueeze(1) - timestep.unsqueeze(0)).abs(), dim=0) + weights = self.linear_timesteps_weights[timestep_id] + return weights + + def scale_model_input(self, sample: torch.Tensor, timestep: int | None = None) -> torch.Tensor: + return sample + + def set_shift(self, shift: float) -> None: + self.shift = shift + \ No newline at end of file diff --git a/fastvideo/models/utils.py b/fastvideo/models/utils.py index 57cb7a74f..bc5df4c36 100644 --- a/fastvideo/models/utils.py +++ b/fastvideo/models/utils.py @@ -145,8 +145,30 @@ def pred_noise_to_pred_video(pred_noise: torch.Tensor, scheduler: Any) -> torch.Tensor: """ Convert predicted noise to clean latent. + + Args: + pred_noise: the predicted noise with shape [B, C, H, W] + where B is batch_size or batch_size * num_frames + noise_input_latent: the noisy latent with shape [B, C, H, W], + timestep: the timestep with shape [1] or [bs * num_frames] or [bs, num_frames] + scheduler: the scheduler + + Returns: + the predicted video with shape [B, C, H, W] """ - timestep = timestep.expand(noise_input_latent.shape[0]) + # If timestep is [bs, num_frames] + if timestep.ndim == 2: + timestep = timestep.flatten(0, 1) + assert timestep.numel() == noise_input_latent.shape[0] + elif timestep.ndim == 1: + # If timestep is [1] + if timestep.shape[0] == 1: + timestep = timestep.expand(noise_input_latent.shape[0]) + else: + assert timestep.numel() == noise_input_latent.shape[0] + else: + raise ValueError(f"[pred_noise_to_pred_video] Invalid timestep shape: {timestep.shape}") + # timestep shape should be [B] dtype = pred_noise.dtype device = pred_noise.device pred_noise = pred_noise.float().to(device) diff --git a/fastvideo/pipelines/basic/wan/wan_causal_dmd_pipeline.py b/fastvideo/pipelines/basic/wan/wan_causal_dmd_pipeline.py index 97426420d..42a8baa9b 100644 --- a/fastvideo/pipelines/basic/wan/wan_causal_dmd_pipeline.py +++ b/fastvideo/pipelines/basic/wan/wan_causal_dmd_pipeline.py @@ -16,8 +16,7 @@ from fastvideo.pipelines.stages import (ConditioningStage, DecodingStage, CausalDMDDenosingStage, InputValidationStage, LatentPreparationStage, - TextEncodingStage, - TimestepPreparationStage) + TextEncodingStage) # isort: on logger = init_logger(__name__) @@ -48,10 +47,6 @@ class WanCausalDMDPipeline(LoRAPipeline, ComposedPipelineBase): self.add_stage(stage_name="conditioning_stage", stage=ConditioningStage()) - self.add_stage(stage_name="timestep_preparation_stage", - stage=TimestepPreparationStage( - scheduler=self.get_module("scheduler"))) - self.add_stage(stage_name="latent_preparation_stage", stage=LatentPreparationStage( scheduler=self.get_module("scheduler"), diff --git a/fastvideo/pipelines/stages/causal_denoising.py b/fastvideo/pipelines/stages/causal_denoising.py index e25b62cf0..925c570f4 100644 --- a/fastvideo/pipelines/stages/causal_denoising.py +++ b/fastvideo/pipelines/stages/causal_denoising.py @@ -4,11 +4,11 @@ from fastvideo.distributed import get_local_torch_device from fastvideo.fastvideo_args import FastVideoArgs from fastvideo.forward_context import set_forward_context from fastvideo.logger import init_logger -from fastvideo.models.schedulers.scheduling_flow_match_euler_discrete import ( - FlowMatchEulerDiscreteScheduler) from fastvideo.models.utils import pred_noise_to_pred_video from fastvideo.pipelines.pipeline_batch_info import ForwardBatch from fastvideo.pipelines.stages.denoising import DenoisingStage +from fastvideo.pipelines.stages.validators import StageValidators as V +from fastvideo.pipelines.stages.validators import VerificationResult try: from fastvideo.attention.backends.sliding_tile_attn import ( @@ -36,7 +36,6 @@ class CausalDMDDenosingStage(DenoisingStage): def __init__(self, transformer, scheduler) -> None: super().__init__(transformer, scheduler) - self.scheduler = FlowMatchEulerDiscreteScheduler(shift=8.0) # KV and cross-attention cache state (initialized on first forward) self.kv_cache1: list | None = None self.crossattn_cache: list | None = None @@ -71,8 +70,16 @@ class CausalDMDDenosingStage(DenoisingStage): # Timesteps for DMD timesteps = torch.tensor( fastvideo_args.pipeline_config.dmd_denoising_steps, - dtype=torch.long, - device=get_local_torch_device()) + dtype=torch.long).cpu() + + if fastvideo_args.pipeline_config.warp_denoising_step: + logger.info("Warping timesteps...") + scheduler_timesteps = torch.cat((self.scheduler.timesteps.cpu(), + torch.tensor([0], + dtype=torch.float32))) + timesteps = scheduler_timesteps[1000 - timesteps] + timesteps = timesteps.to(get_local_torch_device()) + logger.info("Using timesteps: %s", timesteps) # Image kwargs (kept empty unless caller provides compatible args) image_kwargs: dict = {} @@ -262,10 +269,14 @@ class CausalDMDDenosingStage(DenoisingStage): attn_metadata=attn_metadata, forward_batch=batch): # Run transformer; follow DMD stage pattern + t_expanded_noise = t_cur * torch.ones( + (latent_model_input.shape[0], 1), + device=latent_model_input.device, + dtype=torch.long) pred_noise_btchw = self.transformer( latent_model_input, prompt_embeds, - t_expand, + t_expanded_noise, kv_cache=self.kv_cache1, crossattn_cache=self.crossattn_cache, current_start=(pos_start_base + start_index) * @@ -326,10 +337,11 @@ class CausalDMDDenosingStage(DenoisingStage): set_forward_context(current_timestep=0, attn_metadata=attn_metadata, forward_batch=batch): + t_expanded_context = t_context.unsqueeze(1) _ = self.transformer( context_bcthw, prompt_embeds, - t_context, + t_expanded_context, kv_cache=self.kv_cache1, crossattn_cache=self.crossattn_cache, current_start=(pos_start_base + start_index) * @@ -407,3 +419,27 @@ class CausalDMDDenosingStage(DenoisingStage): False, }) self.crossattn_cache = crossattn_cache + + def verify_input(self, batch: ForwardBatch, + fastvideo_args: FastVideoArgs) -> VerificationResult: + """Verify denoising stage inputs.""" + result = VerificationResult() + result.add_check("latents", batch.latents, + [V.is_tensor, V.with_dims(5)]) + result.add_check("prompt_embeds", batch.prompt_embeds, V.list_not_empty) + result.add_check("image_embeds", batch.image_embeds, V.is_list) + result.add_check("image_latent", batch.image_latent, + V.none_or_tensor_with_dims(5)) + result.add_check("num_inference_steps", batch.num_inference_steps, + V.positive_int) + result.add_check("guidance_scale", batch.guidance_scale, + V.positive_float) + result.add_check("eta", batch.eta, V.non_negative_float) + result.add_check("generator", batch.generator, + V.generator_or_list_generators) + result.add_check("do_classifier_free_guidance", + batch.do_classifier_free_guidance, V.bool_value) + result.add_check( + "negative_prompt_embeds", batch.negative_prompt_embeds, lambda x: + not batch.do_classifier_free_guidance or V.list_not_empty(x)) + return result diff --git a/fastvideo/tests/ssim/L40S_reference_videos/SFWan2.1-T2V-1.3B-Diffusers/FLASH_ATTN/Will Smith casually eats noodles, his relaxed demeanor contrasting with the energetic background of .mp4 b/fastvideo/tests/ssim/L40S_reference_videos/SFWan2.1-T2V-1.3B-Diffusers/FLASH_ATTN/Will Smith casually eats noodles, his relaxed demeanor contrasting with the energetic background of .mp4 index b5f2b535e..6a55fe10e 100644 Binary files a/fastvideo/tests/ssim/L40S_reference_videos/SFWan2.1-T2V-1.3B-Diffusers/FLASH_ATTN/Will Smith casually eats noodles, his relaxed demeanor contrasting with the energetic background of .mp4 and b/fastvideo/tests/ssim/L40S_reference_videos/SFWan2.1-T2V-1.3B-Diffusers/FLASH_ATTN/Will Smith casually eats noodles, his relaxed demeanor contrasting with the energetic background of .mp4 differ