diff --git a/.github/workflows/pr-test.yml b/.github/workflows/pr-test.yml index 6aaa0abf9..afb450e66 100644 --- a/.github/workflows/pr-test.yml +++ b/.github/workflows/pr-test.yml @@ -248,13 +248,13 @@ jobs: JOB_ID: "ssim-test" RUNPOD_API_KEY: ${{ secrets.RUNPOD_API_KEY }} GITHUB_RUN_ID: ${{ github.run_id }} - timeout-minutes: 30 + timeout-minutes: 45 run: >- python .github/scripts/runpod_api.py --gpu-type "NVIDIA A40" --gpu-count 2 - --disk-size 100 - --volume-size 100 + --disk-size 200 + --volume-size 200 --test-command "pip install -e .[test] && pip install flash-attn==2.7.0.post2 --no-build-isolation && pytest ./fastvideo/v1/tests/ssim -vs" diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 310316204..40b86e93a 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -40,10 +40,10 @@ repos: - id: codespell additional_dependencies: ['tomli'] args: ['--toml', 'pyproject.toml'] -- repo: https://github.com/PyCQA/isort - rev: 0a0b7a830386ba6a31c2ec8316849ae4d1b8240d # 6.0.0 - hooks: - - id: isort +# - repo: https://github.com/PyCQA/isort +# rev: 0a0b7a830386ba6a31c2ec8316849ae4d1b8240d # 6.0.0 +# hooks: +# - id: isort - repo: https://github.com/jackdewinter/pymarkdown rev: v0.9.29 hooks: @@ -66,7 +66,7 @@ repos: entry: bash args: - -c - - 'git ls-files | grep -v "^fastvideo/v1/tests/ssim/reference_videos/" | grep " " && echo "Filenames should not contain spaces!" && exit 1 || exit 0' + - 'git ls-files | grep -v "^fastvideo/v1/tests/ssim/" | grep " " && echo "Filenames should not contain spaces!" && exit 1 || exit 0' language: system always_run: true pass_filenames: false diff --git a/docs/source/inference/wanvideo.md b/docs/source/inference/wanvideo.md new file mode 100644 index 000000000..e57eab188 --- /dev/null +++ b/docs/source/inference/wanvideo.md @@ -0,0 +1,44 @@ +(wanvideo)= + +# WanVideo +## Inference T2V with WanVideo +First, download the model: + +```bash +python scripts/huggingface/download_hf.py --repo_id=Wan-AI/Wan2.1-T2V-1.3B-Diffusers --local_dir=YOUR_LOCAL_DIR --repo_type=model +``` + +or + +```bash +python scripts/huggingface/download_hf.py --repo_id=Wan-AI/Wan2.1-T2V-14B-Diffusers --local_dir=YOUR_LOCAL_DIR --repo_type=model +``` + +Then run the inference using: + +```bash +sh scripts/inference/v1_inference_wan.sh +``` + +Remember to set `MODEL_BASE` and `num_gpus` accordingly. + +## Inference I2V with WanVideo +First, download the model: + +```bash +python scripts/huggingface/download_hf.py --repo_id=Wan-AI/Wan2.1-I2V-14B-480P-Diffusers --local_dir=YOUR_LOCAL_DIR --repo_type=model +``` + +or + +```bash +python scripts/huggingface/download_hf.py --repo_id=Wan-AI/Wan2.1-I2V-14B-720P-Diffusers --local_dir=YOUR_LOCAL_DIR --repo_type=model +``` + +Then run the inference using: + +```bash +sh scripts/inference/v1_inference_wan_i2v.sh +``` + +Remember to set `MODEL_BASE` and `num_gpus` accordingly. diff --git a/fastvideo/v1/inference_args.py b/fastvideo/v1/inference_args.py index 7ce58ee8a..f48427b14 100644 --- a/fastvideo/v1/inference_args.py +++ b/fastvideo/v1/inference_args.py @@ -31,7 +31,7 @@ class InferenceArgs: guidance_scale: float = 1.0 guidance_rescale: float = 0.0 embedded_cfg_scale: float = 6.0 - flow_shift: int = 7 + flow_shift: Optional[float] = None output_type: str = "pil" @@ -43,6 +43,9 @@ class InferenceArgs: vae_tiling: bool = True vae_sp: bool = False + # Image encoder configuration + image_encoder_precision: str = "fp32" + # Text encoder configuration text_encoder_precision: str = "fp16" text_len: int = 256 @@ -54,14 +57,14 @@ class InferenceArgs: # Flow Matching parameters flow_solver: str = "euler" - denoise_type: str = "flow" + denoise_type: str = "flow" # Deprecated. Will use scheduler_config.json # STA (Spatial-Temporal Attention) parameters mask_strategy_file_path: Optional[str] = None enable_torch_compile: bool = False # Scheduler options - scheduler_type: str = "euler" + scheduler_type: str = "euler" # Deprecated. Will use the param in scheduler_config.json neg_prompt: Optional[str] = None num_videos: int = 1 @@ -73,6 +76,7 @@ class InferenceArgs: log_level: str = "info" # Inference parameters + image_path: Optional[str] = None prompt: Optional[str] = None prompt_path: Optional[str] = None output_path: str = "outputs/" @@ -187,7 +191,7 @@ class InferenceArgs: parser.add_argument( "--flow-shift", "--shift", - type=int, + type=float, default=InferenceArgs.flow_shift, help="Flow shift parameter", ) @@ -240,6 +244,16 @@ class InferenceArgs: default=InferenceArgs.text_len, help="Maximum text length", ) + + # Image encoder config + parser.add_argument( + "--image-encoder-precision", + type=str, + default=InferenceArgs.image_encoder_precision, + choices=["fp32", "fp16", "bf16"], + help="Precision for image encoder", + ) + # Secondary text encoder parser.add_argument( @@ -343,6 +357,10 @@ class InferenceArgs: help="Path to a text file containing the prompt", ) + parser.add_argument("--image-path", + type=str, + help="Path to the image for I2V generation") + parser.add_argument( "--output-path", type=str, diff --git a/fastvideo/v1/inference_engine.py b/fastvideo/v1/inference_engine.py index 28e59f677..882021ca3 100644 --- a/fastvideo/v1/inference_engine.py +++ b/fastvideo/v1/inference_engine.py @@ -106,6 +106,7 @@ class InferenceEngine: guidance_scale = inference_args.guidance_scale flow_shift = inference_args.flow_shift embedded_guidance_scale = inference_args.embedded_cfg_scale + image_path = inference_args.image_path # ======================================================================== # Arguments: target_width, target_height, target_video_length @@ -162,6 +163,7 @@ class InferenceEngine: # local_rank = sp_group.rank device = torch.device(inference_args.device_str) batch = ForwardBatch( + image_path=image_path, prompt=prompt, negative_prompt=negative_prompt, num_videos_per_prompt=num_videos_per_prompt, diff --git a/fastvideo/v1/models/dits/base.py b/fastvideo/v1/models/dits/base.py index 2e0bcb430..2cf98697d 100644 --- a/fastvideo/v1/models/dits/base.py +++ b/fastvideo/v1/models/dits/base.py @@ -1,15 +1,46 @@ # SPDX-License-Identifier: Apache-2.0 +from abc import ABC, abstractmethod +from typing import List, Union +import torch from torch import nn # TODO -class BaseDiT(nn.Module): +class BaseDiT(nn.Module, ABC): _fsdp_shard_conditions: list = [] attention_head_dim: int | None = None + _param_names_mapping: dict + hidden_size: int + num_attention_heads: int + + def __init_subclass__(cls): + required_class_attrs = [ + "_fsdp_shard_conditions", "_param_names_mapping" + ] + super().__init_subclass__() + for attr in required_class_attrs: + if not hasattr(cls, attr): + raise AttributeError( + f"Subclasses of BaseDiT must define '{attr}' class variable" + ) def __init__(self, *args, **kwargs) -> None: super().__init__() - def forward(self, *args, **kwargs): + @abstractmethod + def forward(self, + hidden_states: torch.Tensor, + encoder_hidden_states: Union[torch.Tensor, List[torch.Tensor]], + timestep: torch.LongTensor, + guidance=None, + **kwargs) -> torch.Tensor: pass + + def __post_init__(self): + required_attrs = ["hidden_size", "num_attention_heads"] + for attr in required_attrs: + if not hasattr(self, attr): + raise AttributeError( + f"Subclasses of BaseDiT must define '{attr}' instance variable" + ) diff --git a/fastvideo/v1/models/dits/hunyuanvideo.py b/fastvideo/v1/models/dits/hunyuanvideo.py index c3b264e9d..c7ed836c1 100644 --- a/fastvideo/v1/models/dits/hunyuanvideo.py +++ b/fastvideo/v1/models/dits/hunyuanvideo.py @@ -648,6 +648,8 @@ class HunyuanVideoTransformer3DModel(BaseDiT): self.out_channels, dtype=dtype) + self.__post_init__() + # TODO: change the input the FORWAD_BACTCH Dict # TODO: change output to a dict def forward( diff --git a/fastvideo/v1/models/dits/wanvideo.py b/fastvideo/v1/models/dits/wanvideo.py index 7750d470e..37a9c0dd1 100644 --- a/fastvideo/v1/models/dits/wanvideo.py +++ b/fastvideo/v1/models/dits/wanvideo.py @@ -1,7 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 import math -from typing import Any, Dict, Optional, Tuple, Union +from typing import Any, Dict, List, Optional, Tuple, Union import torch import torch.nn as nn @@ -34,9 +34,10 @@ class WanImageEmbedding(torch.nn.Module): def forward(self, encoder_hidden_states_image: torch.Tensor) -> torch.Tensor: + dtype = encoder_hidden_states_image.dtype hidden_states = self.norm1(encoder_hidden_states_image) hidden_states = self.ff(hidden_states) - hidden_states = self.norm2(hidden_states) + hidden_states = self.norm2(hidden_states).to(dtype) return hidden_states @@ -52,10 +53,7 @@ class WanTimeTextImageEmbedding(nn.Module): super().__init__() self.time_embedder = TimestepEmbedder( - dim, - frequency_embedding_size=time_freq_dim, - act_layer="silu", - freq_dtype=torch.float64) + dim, frequency_embedding_size=time_freq_dim, act_layer="silu") self.time_modulation = ModulateProjection(dim, factor=6, act_layer="silu") @@ -75,9 +73,8 @@ class WanTimeTextImageEmbedding(nn.Module): encoder_hidden_states: torch.Tensor, encoder_hidden_states_image: Optional[torch.Tensor] = None, ): - with torch.cuda.amp.autocast(dtype=torch.float32): - temb = self.time_embedder(timestep.float()) - timestep_proj = self.time_modulation(temb) + temb = self.time_embedder(timestep) + timestep_proj = self.time_modulation(temb) encoder_hidden_states = self.text_embedder(encoder_hidden_states) if encoder_hidden_states_image is not None: @@ -116,7 +113,9 @@ class WanSelfAttention(nn.Module): self.norm_k = RMSNorm(dim, eps=eps) if qk_norm else nn.Identity() # Scaled dot product attention - self.attn = LocalAttention(dropout_rate=0, + self.attn = LocalAttention(num_heads=num_heads, + head_size=self.head_dim, + dropout_rate=0, softmax_scale=None, causal=False) @@ -284,16 +283,15 @@ class WanTransformerBlock(nn.Module): hidden_states = hidden_states.squeeze(1) bs, seq_length, _ = hidden_states.shape orig_dtype = hidden_states.dtype - assert temb.dtype == torch.float32 - with torch.cuda.amp.autocast(dtype=torch.float32): - e = self.scale_shift_table + temb - shift_msa, scale_msa, gate_msa, c_shift_msa, c_scale_msa, c_gate_msa = e.chunk( - 6, dim=1) + assert orig_dtype != torch.float32 + e = self.scale_shift_table + temb.float() + shift_msa, scale_msa, gate_msa, c_shift_msa, c_scale_msa, c_gate_msa = e.chunk( + 6, dim=1) assert shift_msa.dtype == torch.float32 # 1. Self-attention - norm_hidden_states = self.norm1(hidden_states.float()).to( - dtype=orig_dtype) * (1 + scale_msa) + shift_msa + norm_hidden_states = (self.norm1(hidden_states.float()) * + (1 + scale_msa) + shift_msa).to(orig_dtype) query, _ = self.to_q(norm_hidden_states) key, _ = self.to_k(norm_hidden_states) value, _ = self.to_v(norm_hidden_states) @@ -320,6 +318,8 @@ class WanTransformerBlock(nn.Module): null_shift = null_scale = torch.tensor([0], device=hidden_states.device) norm_hidden_states, hidden_states = self.self_attn_residual_norm( hidden_states, attn_output, gate_msa, null_shift, null_scale) + norm_hidden_states, hidden_states = norm_hidden_states.to( + orig_dtype), hidden_states.to(orig_dtype) # 2. Cross-attention attn_output = self.attn2(norm_hidden_states, @@ -327,10 +327,13 @@ class WanTransformerBlock(nn.Module): context_lens=None) norm_hidden_states, hidden_states = self.cross_attn_residual_norm( hidden_states, attn_output, 1, c_shift_msa, c_scale_msa) + norm_hidden_states, hidden_states = norm_hidden_states.to( + orig_dtype), hidden_states.to(orig_dtype) # 3. Feed-forward ff_output = self.ffn(norm_hidden_states) hidden_states = self.mlp_residual(hidden_states, ff_output, c_gate_msa) + hidden_states = hidden_states.to(orig_dtype) return hidden_states @@ -400,8 +403,9 @@ class WanTransformer3DModel(BaseDiT): super().__init__() inner_dim = num_attention_heads * attention_head_dim - self.inner_dim = inner_dim + self.hidden_size = inner_dim self.num_attention_heads = num_attention_heads + self.in_channels = in_channels self.out_channels = out_channels or in_channels self.patch_size = patch_size self.text_len = text_len @@ -440,19 +444,28 @@ class WanTransformer3DModel(BaseDiT): self.gradient_checkpointing = False + self.__post_init__() + def forward( self, hidden_states: torch.Tensor, + encoder_hidden_states: Union[torch.Tensor, List[torch.Tensor]], timestep: torch.LongTensor, - encoder_hidden_states: torch.Tensor, seq_len: Optional[int] = None, - encoder_hidden_states_image: Optional[torch.Tensor] = None, - y: Optional[torch.Tensor] = None, + encoder_hidden_states_image: Optional[Union[torch.Tensor, + List[torch.Tensor]]] = None, return_dict: bool = True, attention_kwargs: Optional[Dict[str, Any]] = None, - ) -> Union[torch.Tensor, Dict[str, torch.Tensor]]: - if y is not None: - hidden_states = torch.cat([hidden_states, y], dim=1) + guidance=None, + ) -> torch.Tensor: + orig_dtype = hidden_states.dtype + if not isinstance(encoder_hidden_states, torch.Tensor): + encoder_hidden_states = encoder_hidden_states[0] + if isinstance(encoder_hidden_states_image, + list) and len(encoder_hidden_states_image) > 0: + encoder_hidden_states_image = encoder_hidden_states_image[0] + else: + encoder_hidden_states_image = None batch_size, num_channels, num_frames, height, width = hidden_states.shape p_t, p_h, p_w = self.patch_size @@ -461,40 +474,23 @@ class WanTransformer3DModel(BaseDiT): post_patch_width = width // p_w # Get rotary embeddings - d = self.inner_dim // self.num_attention_heads + d = self.hidden_size // self.num_attention_heads rope_dim_list = [d - 4 * (d // 6), 2 * (d // 6), 2 * (d // 6)] freqs_cos, freqs_sin = get_rotary_pos_embed( (post_patch_num_frames * get_sequence_model_parallel_world_size(), post_patch_height, post_patch_width), - self.inner_dim, + self.hidden_size, self.num_attention_heads, rope_dim_list, dtype=torch.float64, rope_theta=10000) freqs_cos = freqs_cos.to(hidden_states.device) freqs_sin = freqs_sin.to(hidden_states.device) - freqs_cis = (freqs_cos, freqs_sin) if freqs_cos is not None else None + freqs_cis = (freqs_cos.float(), + freqs_sin.float()) if freqs_cos is not None else None hidden_states = self.patch_embedding(hidden_states) - grid_sizes = torch.stack( - [torch.tensor(hidden_states[0].shape[1:], dtype=torch.long)]) hidden_states = hidden_states.flatten(2).transpose(1, 2) - if seq_len is None: - seq_len = hidden_states.size(1) - hidden_states = torch.cat([ - hidden_states, - hidden_states.new_zeros(1, seq_len - hidden_states.size(1), - hidden_states.size(2)) - ], - dim=1) - - encoder_hidden_states = torch.cat([ - encoder_hidden_states, - encoder_hidden_states.new_zeros( - 1, self.text_len - encoder_hidden_states.size(1), - encoder_hidden_states.size(2)) - ], - dim=1) temb, timestep_proj, encoder_hidden_states, encoder_hidden_states_image = self.condition_embedder( timestep, encoder_hidden_states, encoder_hidden_states_image) @@ -504,6 +500,7 @@ class WanTransformer3DModel(BaseDiT): encoder_hidden_states = torch.concat( [encoder_hidden_states_image, encoder_hidden_states], dim=1) + assert encoder_hidden_states.dtype == orig_dtype # 4. Transformer blocks if torch.is_grad_enabled() and self.gradient_checkpointing: for block in self.blocks: @@ -516,39 +513,16 @@ class WanTransformer3DModel(BaseDiT): timestep_proj, freqs_cis) # 5. Output norm, projection & unpatchify - with torch.cuda.amp.autocast(dtype=torch.float32): - shift, scale = (self.scale_shift_table + temb.unsqueeze(1)).chunk( - 2, dim=1) - hidden_states = self.norm_out(hidden_states.float(), shift, scale) - hidden_states = self.proj_out(hidden_states) + shift, scale = (self.scale_shift_table + temb.unsqueeze(1)).chunk(2, + dim=1) + hidden_states = self.norm_out(hidden_states.float(), shift, scale) + hidden_states = self.proj_out(hidden_states) - output = self.unpatchify(hidden_states, grid_sizes) + hidden_states = hidden_states.reshape(batch_size, post_patch_num_frames, + post_patch_height, + post_patch_width, p_t, p_h, p_w, + -1) + hidden_states = hidden_states.permute(0, 7, 1, 4, 2, 5, 3, 6) + output = hidden_states.flatten(6, 7).flatten(4, 5).flatten(2, 3) - return output.float() - - def unpatchify(self, x, grid_sizes) -> torch.Tensor: - r""" - Reconstruct video tensors from patch embeddings. - - Args: - x (List[Tensor]): - List of patchified features, each with shape [L, C_out * prod(patch_size)] - grid_sizes (Tensor): - Original spatial-temporal grid dimensions before patching, - shape [B, 3] (3 dimensions correspond to F_patches, H_patches, W_patches) - - Returns: - Tensor: - Reconstructed video tensors with shape [B, C_out, F, H / 8, W / 8] - """ - - c = self.out_channels - out = [] - for u, v in zip(x, grid_sizes.tolist()): - u = u[:math.prod(v)].view(*v, *self.patch_size, c) - u = u.permute(6, 0, 3, 1, 4, 2, 5) - # u = torch.einsum('fhwpqrc->cfphqwr', u.contiguous()) - u = u.reshape(c, *[i * j for i, j in zip(v, self.patch_size)]) - out.append(u) - out = torch.cat(out, dim=0) - return out + return output diff --git a/fastvideo/v1/models/encoders/clip.py b/fastvideo/v1/models/encoders/clip.py index 5fbf77c16..a62bb64b9 100644 --- a/fastvideo/v1/models/encoders/clip.py +++ b/fastvideo/v1/models/encoders/clip.py @@ -598,6 +598,9 @@ class CLIPVisionTransformer(nn.Module): inputs_embeds=hidden_states, return_all_hidden_states=return_all_hidden_states) + if not return_all_hidden_states: + encoder_outputs = encoder_outputs[0] + # Handle post-norm (if applicable) and stacks feature layers if needed encoder_outputs = resolve_visual_encoder_outputs( encoder_outputs, feature_sample_layers, self.post_layernorm, @@ -654,6 +657,8 @@ class CLIPVisionModel(nn.Module, SupportsQuant): layer_count = len(self.vision_model.encoder.layers) for name, loaded_weight in weights: + if name.startswith("visual_projection"): + continue # post_layernorm is not needed in CLIPVisionModel if (name.startswith("vision_model.post_layernorm") and self.vision_model.post_layernorm is None): diff --git a/fastvideo/v1/models/encoders/t5.py b/fastvideo/v1/models/encoders/t5.py index 1422d16c3..85b6fb45a 100644 --- a/fastvideo/v1/models/encoders/t5.py +++ b/fastvideo/v1/models/encoders/t5.py @@ -28,11 +28,12 @@ import torch.nn.functional as F from torch import nn from transformers import T5Config -from fastvideo.v1.distributed import get_tensor_model_parallel_world_size +from fastvideo.v1.distributed import (get_tensor_model_parallel_rank, + get_tensor_model_parallel_world_size) from fastvideo.v1.layers.activation import get_act_fn from fastvideo.v1.layers.layernorm import RMSNorm -from fastvideo.v1.layers.linear import (ColumnParallelLinear, QKVParallelLinear, - RowParallelLinear) +from fastvideo.v1.layers.linear import (MergedColumnParallelLinear, + QKVParallelLinear, RowParallelLinear) from fastvideo.v1.layers.vocab_parallel_embedding import VocabParallelEmbedding from fastvideo.v1.models.loader.weight_utils import default_weight_loader @@ -67,7 +68,8 @@ class T5DenseActDense(nn.Module): config: T5Config, quant_config: Optional[QuantizationConfig] = None): super().__init__() - self.wi = ColumnParallelLinear(config.d_model, config.d_ff, bias=False) + self.wi = MergedColumnParallelLinear(config.d_model, [config.d_ff], + bias=False) self.wo = RowParallelLinear(config.d_ff, config.d_model, bias=False, @@ -87,14 +89,12 @@ class T5DenseGatedActDense(nn.Module): config: T5Config, quant_config: Optional[QuantizationConfig] = None): super().__init__() - self.wi_0 = ColumnParallelLinear(config.d_model, - config.d_ff, - bias=False, - quant_config=quant_config) - self.wi_1 = ColumnParallelLinear(config.d_model, - config.d_ff, - bias=False, - quant_config=quant_config) + self.wi_0 = MergedColumnParallelLinear(config.d_model, [config.d_ff], + bias=False, + quant_config=quant_config) + self.wi_1 = MergedColumnParallelLinear(config.d_model, [config.d_ff], + bias=False, + quant_config=quant_config) # Should not run in fp16 unless mixed-precision is used, # see https://github.com/huggingface/transformers/issues/20287. self.wo = RowParallelLinear(config.d_ff, @@ -170,6 +170,7 @@ class T5Attention(nn.Module): config.relative_attention_max_distance self.d_model = config.d_model self.key_value_proj_dim = config.d_kv + self.total_num_heads = self.total_num_kv_heads = config.num_heads # Partition heads across multiple tensor parallel GPUs. tp_world_size = get_tensor_model_parallel_world_size() @@ -178,29 +179,33 @@ class T5Attention(nn.Module): self.inner_dim = self.n_heads * self.key_value_proj_dim # No GQA in t5. - self.n_kv_heads = self.n_heads + # self.n_kv_heads = self.n_heads - self.qkv_proj = QKVParallelLinear(self.d_model, - self.d_model // self.n_heads, - self.n_heads, - self.n_kv_heads, - bias=False, - quant_config=quant_config) + self.qkv_proj = QKVParallelLinear( + self.d_model, + self.d_model // self.total_num_heads, + self.total_num_heads, + self.total_num_kv_heads, + bias=False, + quant_config=quant_config, + prefix=f"{prefix}.qkv_proj", + ) self.attn = T5MultiHeadAttention() if self.has_relative_attention_bias: self.relative_attention_bias = \ VocabParallelEmbedding(self.relative_attention_num_buckets, - self.n_heads, + self.total_num_heads, org_num_embeddings=self.relative_attention_num_buckets, padding_size=self.relative_attention_num_buckets, quant_config=quant_config) self.o = RowParallelLinear( - self.inner_dim, + self.d_model, self.d_model, bias=False, quant_config=quant_config, + prefix=f"{prefix}.o_proj", ) @staticmethod @@ -295,13 +300,13 @@ class T5Attention(nn.Module): ) -> torch.Tensor: bs, seq_len, _ = hidden_states.shape num_seqs = bs - n, c = self.n_heads, self.d_model // self.n_heads + n, c = self.n_heads, self.d_model // self.total_num_heads qkv, _ = self.qkv_proj(hidden_states) # Projection of 'own' hidden state (self-attention). No GQA here. q, k, v = qkv.split(self.inner_dim, dim=-1) - q = q.view(bs, -1, n, c) - k = k.view(bs, -1, n, c) - v = v.view(bs, -1, n, c) + q = q.reshape(bs, seq_len, n, c) + k = k.reshape(bs, seq_len, n, c) + v = v.reshape(bs, seq_len, n, c) assert attn_metadata is not None attn_bias = attn_metadata.attn_bias @@ -325,6 +330,11 @@ class T5Attention(nn.Module): -1) if attention_mask.ndim == 2 else attention_mask.unsqueeze(1) attn_bias.masked_fill_(attention_mask == 0, torch.finfo(q.dtype).min) + + if get_tensor_model_parallel_world_size() > 1: + rank = get_tensor_model_parallel_rank() + attn_bias = attn_bias[:, rank * self.n_heads:(rank + 1) * + self.n_heads, :, :] attn_output = self.attn(q, k, v, attn_bias) output, _ = self.o(attn_output) return output diff --git a/fastvideo/v1/models/hf_transformer_utils.py b/fastvideo/v1/models/hf_transformer_utils.py index feda42212..e2f39ea57 100644 --- a/fastvideo/v1/models/hf_transformer_utils.py +++ b/fastvideo/v1/models/hf_transformer_utils.py @@ -95,9 +95,12 @@ def get_diffusers_config( Returns: The loaded configuration. """ + config_name = "config.json" + if "scheduler" in model: + config_name = "scheduler_config.json" # Check if the model path exists if os.path.exists(model): - config_file = os.path.join(model, "config.json") + config_file = os.path.join(model, config_name) if os.path.exists(config_file): try: # Load the config directly from the file diff --git a/fastvideo/v1/models/loader/component_loader.py b/fastvideo/v1/models/loader/component_loader.py index bf7e1d09e..47e5f2fb4 100644 --- a/fastvideo/v1/models/loader/component_loader.py +++ b/fastvideo/v1/models/loader/component_loader.py @@ -10,7 +10,7 @@ from typing import Any, Generator, Iterable, List, Optional, Tuple, cast import torch import torch.nn as nn from safetensors.torch import load_file as safetensors_load_file -from transformers import AutoTokenizer, PretrainedConfig +from transformers import AutoImageProcessor, AutoTokenizer, PretrainedConfig from transformers.utils import SAFE_WEIGHTS_INDEX_NAME from fastvideo.v1.inference_args import InferenceArgs @@ -23,6 +23,7 @@ from fastvideo.v1.models.loader.weight_utils import ( filter_duplicate_safetensors_files, filter_files_not_needed_for_inference, pt_weights_iterator, safetensors_weights_iterator) from fastvideo.v1.models.registry import ModelRegistry +from fastvideo.v1.utils import PRECISION_TO_TYPE logger = init_logger(__name__) @@ -71,6 +72,8 @@ class ComponentLoader(ABC): "text_encoder_2": (TextEncoderLoader, "transformers"), "tokenizer": (TokenizerLoader, "transformers"), "tokenizer_2": (TokenizerLoader, "transformers"), + "image_processor": (ImageProcessorLoader, "transformers"), + "image_encoder": (ImageEncoderLoader, "transformers"), } if module_type in module_loaders: @@ -209,11 +212,15 @@ class TextEncoderLoader(ComponentLoader): target_device = torch.device(inference_args.device_str) # TODO(will): add support for other dtypes - return self.load_model(model_path, model_config, target_device) + return self.load_model(model_path, model_config, target_device, + inference_args.text_encoder_precision) - def load_model(self, model_path: str, model_config, - target_device: torch.device): - with set_default_torch_dtype(torch.float16): + def load_model(self, + model_path: str, + model_config, + target_device: torch.device, + dtype: str = "fp16"): + with set_default_torch_dtype(PRECISION_TO_TYPE[dtype]): with target_device: architectures = getattr(model_config, "architectures", []) model_cls, _ = ModelRegistry.resolve_model_cls(architectures) @@ -240,6 +247,40 @@ class TextEncoderLoader(ComponentLoader): return model.eval() +class ImageEncoderLoader(TextEncoderLoader): + + def load(self, model_path: str, architecture: str, + inference_args: InferenceArgs): + """Load the text encoders based on the model path, architecture, and inference args.""" + model_config: PretrainedConfig = get_hf_config( + model=model_path, + trust_remote_code=inference_args.trust_remote_code, + revision=inference_args.revision, + model_override_args=None, + inference_args=inference_args, + ) + logger.info("HF Model config: %s", model_config) + + target_device = torch.device(inference_args.device_str) + # TODO(will): add support for other dtypes + return self.load_model(model_path, model_config, target_device, + inference_args.image_encoder_precision) + + +class ImageProcessorLoader(ComponentLoader): + """Loader for image processor.""" + + def load(self, model_path: str, architecture: str, + inference_args: InferenceArgs): + """Load the image processor based on the model path, architecture, and inference args.""" + logger.info("Loading image processor from %s", model_path) + + image_processor = AutoImageProcessor.from_pretrained(model_path, ) + logger.info("Loaded image processor: %s", + image_processor.__class__.__name__) + return image_processor + + class TokenizerLoader(ComponentLoader): """Loader for tokenizers.""" @@ -265,8 +306,6 @@ class VAELoader(ComponentLoader): inference_args: InferenceArgs): """Load the VAE based on the model path, architecture, and inference args.""" # TODO(will): move this to a constants file - from fastvideo.v1.utils import PRECISION_TO_TYPE - config = get_diffusers_config(model=model_path) class_name = config.pop("_class_name") @@ -289,14 +328,6 @@ class VAELoader(ComponentLoader): dtype = PRECISION_TO_TYPE[inference_args.vae_precision] vae = vae.eval().to(dtype) - # TODO(will): should we define hunyuan vae config class? - vae_kwargs = { - "s_ratio": config["spatial_compression_ratio"], - "t_ratio": config["temporal_compression_ratio"], - } - - vae.kwargs = vae_kwargs - return vae @@ -326,6 +357,7 @@ class TransformerLoader(ComponentLoader): len(safetensors_list), model_path) # initialize_sequence_parallel_group(inference_args.sp_size) + default_dtype = PRECISION_TO_TYPE[inference_args.precision] # Load the model using FSDP loader logger.info("Loading model from %s", cls_name) @@ -333,12 +365,16 @@ class TransformerLoader(ComponentLoader): init_params=model_config, weight_dir_list=safetensors_list, device=inference_args.device, - cpu_offload=inference_args.use_cpu_offload) + cpu_offload=inference_args.use_cpu_offload, + default_dtype=default_dtype) total_params = sum(p.numel() for p in model.parameters()) logger.info("Loaded model with %.2fB parameters", total_params / 1e9) - model.eval() + dtypes = set(param.dtype for param in model.parameters()) + if len(dtypes) > 1: + model = model.to(default_dtype) + model = model.eval() return model @@ -348,21 +384,17 @@ class SchedulerLoader(ComponentLoader): def load(self, model_path: str, architecture: str, inference_args: InferenceArgs): """Load the scheduler based on the model path, architecture, and inference args.""" - if hasattr(inference_args, - 'denoise_type') and inference_args.denoise_type == "flow": - # TODO(will): add schedulers to register or create a new scheduler registry - # TODO(will): default to config file but allow override through - # inference args. Currently only uses inference args. - from fastvideo.v1.models.schedulers.scheduling_flow_match_euler_discrete import ( - FlowMatchDiscreteScheduler) - scheduler = FlowMatchDiscreteScheduler( - shift=inference_args.flow_shift, - solver=inference_args.flow_solver, - ) - logger.info("Scheduler loaded: %s", scheduler) - else: - raise ValueError( - f"Invalid denoise type: {inference_args.denoise_type}") + config = get_diffusers_config(model=model_path) + + class_name = config.pop("_class_name") + assert class_name is not None, "Model config does not contain a _class_name attribute. Only diffusers format is supported." + config.pop("_diffusers_version") + + scheduler_cls, _ = ModelRegistry.resolve_model_cls(class_name) + + scheduler = scheduler_cls(**config) + if inference_args.flow_shift is not None: + scheduler.set_shift(inference_args.flow_shift) return scheduler diff --git a/fastvideo/v1/models/registry.py b/fastvideo/v1/models/registry.py index e80e1e1c2..bc11f365b 100644 --- a/fastvideo/v1/models/registry.py +++ b/fastvideo/v1/models/registry.py @@ -38,6 +38,7 @@ _TEXT_ENCODER_MODELS = { _IMAGE_ENCODER_MODELS: dict[str, tuple] = { # "HunyuanVideoTransformer3DModel": ("image_encoder", "hunyuanvideo", "HunyuanVideoImageEncoder"), + "CLIPVisionModelWithProjection": ("encoders", "clip", "CLIPVisionModel"), } _VAE_MODELS = { @@ -46,12 +47,21 @@ _VAE_MODELS = { "AutoencoderKLWan": ("vaes", "wanvae", "AutoencoderKLWan"), } +_SCHEDULERS = { + "FlowMatchEulerDiscreteScheduler": + ("schedulers", "scheduling_flow_match_euler_discrete", + "FlowMatchDiscreteScheduler"), + "UniPCMultistepScheduler": + ("schedulers", "scheduling_unipc_multistep", "UniPCMultistepScheduler"), +} + _FAST_VIDEO_MODELS = { **_TEXT_TO_VIDEO_DIT_MODELS, **_IMAGE_TO_VIDEO_DIT_MODELS, **_TEXT_ENCODER_MODELS, **_IMAGE_ENCODER_MODELS, **_VAE_MODELS, + **_SCHEDULERS, } _SUBPROCESS_COMMAND = [ diff --git a/fastvideo/v1/models/schedulers/base.py b/fastvideo/v1/models/schedulers/base.py new file mode 100644 index 000000000..5333d07a9 --- /dev/null +++ b/fastvideo/v1/models/schedulers/base.py @@ -0,0 +1,47 @@ +# SPDX-License-Identifier: Apache-2.0 + +from abc import ABC, abstractmethod +from typing import Optional, Tuple, Union + +import torch +from diffusers.utils import BaseOutput + + +class BaseScheduler(ABC): + timesteps: torch.tensor + order: int + + def __init__(self, *args, **kwargs) -> None: + # Check if subclass has defined all required properties + required_attributes = ['timesteps', 'order'] + + for attr in required_attributes: + if not hasattr(self, attr): + raise AttributeError( + f"Subclasses of BaseScheduler must define '{attr}' property" + ) + + @abstractmethod + def set_shift(self, shift: float) -> None: + pass + + @abstractmethod + def set_timesteps(self, *args, **kwargs) -> None: + pass + + @abstractmethod + def scale_model_input(self, + sample: torch.Tensor, + timestep: Optional[int] = None) -> torch.Tensor: + pass + + @abstractmethod + def step( + self, + model_output: torch.FloatTensor, + timestep: Union[float, torch.FloatTensor], + sample: torch.FloatTensor, + return_dict: bool = True, + **kwargs, + ) -> Union[BaseOutput, Tuple]: + pass diff --git a/fastvideo/v1/models/schedulers/scheduling_flow_match_euler_discrete.py b/fastvideo/v1/models/schedulers/scheduling_flow_match_euler_discrete.py index 592af94e4..0e5f1e7d3 100644 --- a/fastvideo/v1/models/schedulers/scheduling_flow_match_euler_discrete.py +++ b/fastvideo/v1/models/schedulers/scheduling_flow_match_euler_discrete.py @@ -27,6 +27,8 @@ from diffusers.configuration_utils import ConfigMixin, register_to_config from diffusers.schedulers.scheduling_utils import SchedulerMixin from diffusers.utils import BaseOutput, logging +from fastvideo.v1.models.schedulers.base import BaseScheduler + logger = logging.get_logger(__name__) # pylint: disable=invalid-name @@ -44,7 +46,7 @@ class FlowMatchDiscreteSchedulerOutput(BaseOutput): prev_sample: torch.FloatTensor -class FlowMatchDiscreteScheduler(SchedulerMixin, ConfigMixin): +class FlowMatchDiscreteScheduler(SchedulerMixin, ConfigMixin, BaseScheduler): """ Euler scheduler. @@ -74,6 +76,7 @@ class FlowMatchDiscreteScheduler(SchedulerMixin, ConfigMixin): reverse: bool = True, solver: str = "euler", n_tokens: Optional[int] = None, + **kwargs, ): sigmas = torch.linspace(1, 0, num_train_timesteps + 1) @@ -94,6 +97,8 @@ class FlowMatchDiscreteScheduler(SchedulerMixin, ConfigMixin): f"Solver {solver} not supported. Supported solvers: {self.supported_solver}" ) + BaseScheduler.__init__(self) + @property def step_index(self): """ @@ -170,6 +175,9 @@ class FlowMatchDiscreteScheduler(SchedulerMixin, ConfigMixin): return idx + def set_shift(self, shift: float) -> None: + self.config.shift = shift + def _init_step_index(self, timestep) -> None: if self.begin_index is None: if isinstance(timestep, torch.Tensor): diff --git a/fastvideo/v1/models/schedulers/scheduling_unipc_multistep.py b/fastvideo/v1/models/schedulers/scheduling_unipc_multistep.py new file mode 100644 index 000000000..a81d95620 --- /dev/null +++ b/fastvideo/v1/models/schedulers/scheduling_unipc_multistep.py @@ -0,0 +1,1162 @@ +# SPDX-License-Identifier: Apache-2.0 + +# Copyright 2024 TSAIL Team and The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# DISCLAIMER: check https://arxiv.org/abs/2302.04867 and https://github.com/wl-zhao/UniPC for more info +# The codebase is modified based on https://github.com/huggingface/diffusers/blob/main/src/diffusers/schedulers/scheduling_dpmsolver_multistep.py +# ============================================================================== +# +# Modified from diffusers==0.33.0.dev0 +# +# ============================================================================== + +import math +from typing import List, Optional, Tuple, Union + +import numpy as np +import torch +from diffusers.configuration_utils import ConfigMixin, register_to_config +from diffusers.schedulers.scheduling_utils import (KarrasDiffusionSchedulers, + SchedulerMixin, + SchedulerOutput) +from diffusers.utils import deprecate, is_scipy_available + +from fastvideo.v1.models.schedulers.base import BaseScheduler + +if is_scipy_available(): + import scipy.stats + + +# Copied from diffusers.schedulers.scheduling_ddpm.betas_for_alpha_bar +def betas_for_alpha_bar( + num_diffusion_timesteps, + max_beta=0.999, + alpha_transform_type="cosine", +) -> torch.Tensor: + """ + Create a beta schedule that discretizes the given alpha_t_bar function, which defines the cumulative product of + (1-beta) over time from t = [0,1]. + + Contains a function alpha_bar that takes an argument t and transforms it to the cumulative product of (1-beta) up + to that part of the diffusion process. + + + Args: + num_diffusion_timesteps (`int`): the number of betas to produce. + max_beta (`float`): the maximum beta to use; use values lower than 1 to + prevent singularities. + alpha_transform_type (`str`, *optional*, default to `cosine`): the type of noise schedule for alpha_bar. + Choose from `cosine` or `exp` + + Returns: + betas (`np.ndarray`): the betas used by the scheduler to step the model outputs + """ + if alpha_transform_type == "cosine": + + def alpha_bar_fn(t: float) -> float: + return math.cos((t + 0.008) / 1.008 * math.pi / 2)**2 + + elif alpha_transform_type == "exp": + + def alpha_bar_fn(t: float) -> float: + return math.exp(t * -12.0) + + else: + raise ValueError( + f"Unsupported alpha_transform_type: {alpha_transform_type}") + + betas = [] + for i in range(num_diffusion_timesteps): + t1 = i / num_diffusion_timesteps + t2 = (i + 1) / num_diffusion_timesteps + betas.append(min(1 - alpha_bar_fn(t2) / alpha_bar_fn(t1), max_beta)) + return torch.tensor(betas, dtype=torch.float32) + + +# Copied from diffusers.schedulers.scheduling_ddim.rescale_zero_terminal_snr +def rescale_zero_terminal_snr(betas: torch.Tensor) -> torch.Tensor: + """ + Rescales betas to have zero terminal SNR Based on https://arxiv.org/pdf/2305.08891.pdf (Algorithm 1) + + + Args: + betas (`torch.Tensor`): + the betas that the scheduler is being initialized with. + + Returns: + `torch.Tensor`: rescaled betas with zero terminal SNR + """ + # Convert betas to alphas_bar_sqrt + alphas = 1.0 - betas + alphas_cumprod = torch.cumprod(alphas, dim=0) + alphas_bar_sqrt = alphas_cumprod.sqrt() + + # Store old values. + alphas_bar_sqrt_0 = alphas_bar_sqrt[0].clone() + alphas_bar_sqrt_T = alphas_bar_sqrt[-1].clone() + + # Shift so the last timestep is zero. + alphas_bar_sqrt -= alphas_bar_sqrt_T + + # Scale so the first timestep is back to the old value. + alphas_bar_sqrt *= alphas_bar_sqrt_0 / (alphas_bar_sqrt_0 - + alphas_bar_sqrt_T) + + # Convert alphas_bar_sqrt to betas + alphas_bar = alphas_bar_sqrt**2 # Revert sqrt + alphas = alphas_bar[1:] / alphas_bar[:-1] # Revert cumprod + alphas = torch.cat([alphas_bar[0:1], alphas]) + betas = 1 - alphas + + return betas + + +class UniPCMultistepScheduler(SchedulerMixin, ConfigMixin, BaseScheduler): + """ + `UniPCMultistepScheduler` is a training-free framework designed for the fast sampling of diffusion models. + + This model inherits from [`SchedulerMixin`] and [`ConfigMixin`]. Check the superclass documentation for the generic + methods the library implements for all schedulers such as loading and saving. + + Args: + num_train_timesteps (`int`, defaults to 1000): + The number of diffusion steps to train the model. + beta_start (`float`, defaults to 0.0001): + The starting `beta` value of inference. + beta_end (`float`, defaults to 0.02): + The final `beta` value. + beta_schedule (`str`, defaults to `"linear"`): + The beta schedule, a mapping from a beta range to a sequence of betas for stepping the model. Choose from + `linear`, `scaled_linear`, or `squaredcos_cap_v2`. + trained_betas (`np.ndarray`, *optional*): + Pass an array of betas directly to the constructor to bypass `beta_start` and `beta_end`. + solver_order (`int`, default `2`): + The UniPC order which can be any positive integer. The effective order of accuracy is `solver_order + 1` + due to the UniC. It is recommended to use `solver_order=2` for guided sampling, and `solver_order=3` for + unconditional sampling. + prediction_type (`str`, defaults to `epsilon`, *optional*): + Prediction type of the scheduler function; can be `epsilon` (predicts the noise of the diffusion process), + `sample` (directly predicts the noisy sample`) or `v_prediction` (see section 2.4 of [Imagen + Video](https://imagen.research.google/video/paper.pdf) paper). + thresholding (`bool`, defaults to `False`): + Whether to use the "dynamic thresholding" method. This is unsuitable for latent-space diffusion models such + as Stable Diffusion. + dynamic_thresholding_ratio (`float`, defaults to 0.995): + The ratio for the dynamic thresholding method. Valid only when `thresholding=True`. + sample_max_value (`float`, defaults to 1.0): + The threshold value for dynamic thresholding. Valid only when `thresholding=True` and `predict_x0=True`. + predict_x0 (`bool`, defaults to `True`): + Whether to use the updating algorithm on the predicted x0. + solver_type (`str`, default `bh2`): + Solver type for UniPC. It is recommended to use `bh1` for unconditional sampling when steps < 10, and `bh2` + otherwise. + lower_order_final (`bool`, default `True`): + Whether to use lower-order solvers in the final steps. Only valid for < 15 inference steps. This can + stabilize the sampling of DPMSolver for steps < 15, especially for steps <= 10. + disable_corrector (`list`, default `[]`): + Decides which step to disable the corrector to mitigate the misalignment between `epsilon_theta(x_t, c)` + and `epsilon_theta(x_t^c, c)` which can influence convergence for a large guidance scale. Corrector is + usually disabled during the first few steps. + solver_p (`SchedulerMixin`, default `None`): + Any other scheduler that if specified, the algorithm becomes `solver_p + UniC`. + use_karras_sigmas (`bool`, *optional*, defaults to `False`): + Whether to use Karras sigmas for step sizes in the noise schedule during the sampling process. If `True`, + the sigmas are determined according to a sequence of noise levels {σi}. + use_exponential_sigmas (`bool`, *optional*, defaults to `False`): + Whether to use exponential sigmas for step sizes in the noise schedule during the sampling process. + use_beta_sigmas (`bool`, *optional*, defaults to `False`): + Whether to use beta sigmas for step sizes in the noise schedule during the sampling process. Refer to [Beta + Sampling is All You Need](https://huggingface.co/papers/2407.12173) for more information. + timestep_spacing (`str`, defaults to `"linspace"`): + The way the timesteps should be scaled. Refer to Table 2 of the [Common Diffusion Noise Schedules and + Sample Steps are Flawed](https://huggingface.co/papers/2305.08891) for more information. + steps_offset (`int`, defaults to 0): + An offset added to the inference steps, as required by some model families. + final_sigmas_type (`str`, defaults to `"zero"`): + The final `sigma` value for the noise schedule during the sampling process. If `"sigma_min"`, the final + sigma is the same as the last sigma in the training schedule. If `zero`, the final sigma is set to 0. + rescale_betas_zero_snr (`bool`, defaults to `False`): + Whether to rescale the betas to have zero terminal SNR. This enables the model to generate very bright and + dark samples instead of limiting it to samples with medium brightness. Loosely related to + [`--offset_noise`](https://github.com/huggingface/diffusers/blob/74fd735eb073eb1d774b1ab4154a0876eb82f055/examples/dreambooth/train_dreambooth.py#L506). + """ + + _compatibles = [e.name for e in KarrasDiffusionSchedulers] + order = 1 + + @register_to_config + def __init__( + self, + num_train_timesteps: int = 1000, + beta_start: float = 0.0001, + beta_end: float = 0.02, + beta_schedule: str = "linear", + trained_betas: Optional[Union[np.ndarray, List[float]]] = None, + solver_order: int = 2, + prediction_type: str = "epsilon", + thresholding: bool = False, + dynamic_thresholding_ratio: float = 0.995, + sample_max_value: float = 1.0, + predict_x0: bool = True, + solver_type: str = "bh2", + lower_order_final: bool = True, + disable_corrector: Tuple[int, ...] = (), + solver_p: SchedulerMixin = None, + use_karras_sigmas: Optional[bool] = False, + use_exponential_sigmas: Optional[bool] = False, + use_beta_sigmas: Optional[bool] = False, + use_flow_sigmas: Optional[bool] = False, + flow_shift: Optional[float] = 1.0, + timestep_spacing: str = "linspace", + steps_offset: int = 0, + final_sigmas_type: Optional[str] = "zero", # "zero", "sigma_min" + rescale_betas_zero_snr: bool = False, + ): + if self.config.use_beta_sigmas and not is_scipy_available(): + raise ImportError( + "Make sure to install scipy if you want to use beta sigmas.") + if sum([ + self.config.use_beta_sigmas, self.config.use_exponential_sigmas, + self.config.use_karras_sigmas + ]) > 1: + raise ValueError( + "Only one of `config.use_beta_sigmas`, `config.use_exponential_sigmas`, `config.use_karras_sigmas` can be used." + ) + if trained_betas is not None: + self.betas = torch.tensor(trained_betas, dtype=torch.float32) + elif beta_schedule == "linear": + self.betas = torch.linspace(beta_start, + beta_end, + num_train_timesteps, + dtype=torch.float32) + elif beta_schedule == "scaled_linear": + # this schedule is very specific to the latent diffusion model. + self.betas = torch.linspace(beta_start**0.5, + beta_end**0.5, + num_train_timesteps, + dtype=torch.float32)**2 + elif beta_schedule == "squaredcos_cap_v2": + # Glide cosine schedule + self.betas = betas_for_alpha_bar(num_train_timesteps) + else: + raise NotImplementedError( + f"{beta_schedule} is not implemented for {self.__class__}") + + if rescale_betas_zero_snr: + self.betas = rescale_zero_terminal_snr(self.betas) + + self.alphas = 1.0 - self.betas + self.alphas_cumprod = torch.cumprod(self.alphas, dim=0) + + if rescale_betas_zero_snr: + # Close to 0 without being 0 so first sigma is not inf + # FP16 smallest positive subnormal works well here + self.alphas_cumprod[-1] = 2**-24 + + # Currently we only support VP-type noise schedule + self.alpha_t = torch.sqrt(self.alphas_cumprod) + self.sigma_t = torch.sqrt(1 - self.alphas_cumprod) + self.lambda_t = torch.log(self.alpha_t) - torch.log(self.sigma_t) + self.sigmas = ((1 - self.alphas_cumprod) / self.alphas_cumprod)**0.5 + + # standard deviation of the initial noise distribution + self.init_noise_sigma = 1.0 + + if solver_type not in ["bh1", "bh2"]: + if solver_type in ["midpoint", "heun", "logrho"]: + self.register_to_config(solver_type="bh2") + else: + raise NotImplementedError( + f"{solver_type} is not implemented for {self.__class__}") + + self.predict_x0 = predict_x0 + # setable values + self.num_inference_steps: Optional[int] = None + timesteps = np.linspace(0, + num_train_timesteps - 1, + num_train_timesteps, + dtype=np.float32)[::-1].copy() + self.timesteps = torch.from_numpy(timesteps) + self.model_outputs = [None] * solver_order + self.timestep_list: List[Union[int, + torch.Tensor]] = [None] * solver_order + self.lower_order_nums = 0 + self.disable_corrector = list(disable_corrector) + self.solver_p = solver_p + self.last_sample = None + self._step_index: Optional[int] = None + self._begin_index: Optional[int] = None + self.sigmas = self.sigmas.to( + "cpu") # to avoid too much CPU/GPU communication + + BaseScheduler.__init__(self) + + @property + def step_index(self): + """ + The index counter for current timestep. It will increase 1 after each scheduler step. + """ + return self._step_index + + @property + def begin_index(self): + """ + The index for the first timestep. It should be set from pipeline with `set_begin_index` method. + """ + return self._begin_index + + def set_shift(self, shift: float) -> None: + self.config.flow_shift = shift + + # Copied from diffusers.schedulers.scheduling_dpmsolver_multistep.DPMSolverMultistepScheduler.set_begin_index + def set_begin_index(self, begin_index: int = 0): + """ + Sets the begin index for the scheduler. This function should be run from pipeline before the inference. + + Args: + begin_index (`int`): + The begin index for the scheduler. + """ + self._begin_index = begin_index + + def set_timesteps(self, + num_inference_steps: int, + device: Union[str, torch.device] = None): + """ + Sets the discrete timesteps used for the diffusion chain (to be run before inference). + + Args: + num_inference_steps (`int`): + The number of diffusion steps used when generating samples with a pre-trained model. + device (`str` or `torch.device`, *optional*): + The device to which the timesteps should be moved to. If `None`, the timesteps are not moved. + """ + # "linspace", "leading", "trailing" corresponds to annotation of Table 2. of https://arxiv.org/abs/2305.08891 + if self.config.timestep_spacing == "linspace": + timesteps = (np.linspace( + 0, self.config.num_train_timesteps - 1, num_inference_steps + + 1).round()[::-1][:-1].copy().astype(np.int64)) + elif self.config.timestep_spacing == "leading": + step_ratio = self.config.num_train_timesteps // ( + num_inference_steps + 1) + # creates integer timesteps by multiplying by ratio + # casting to int to avoid issues when num_inference_step is power of 3 + timesteps = (np.arange(0, num_inference_steps + 1) * + step_ratio).round()[::-1][:-1].copy().astype(np.int64) + timesteps += self.config.steps_offset + elif self.config.timestep_spacing == "trailing": + step_ratio = self.config.num_train_timesteps / num_inference_steps + # creates integer timesteps by multiplying by ratio + # casting to int to avoid issues when num_inference_step is power of 3 + timesteps = np.arange(self.config.num_train_timesteps, 0, + -step_ratio).round().copy().astype(np.int64) + timesteps -= 1 + else: + raise ValueError( + f"{self.config.timestep_spacing} is not supported. Please make sure to choose one of 'linspace', 'leading' or 'trailing'." + ) + + sigmas = np.array( + ((1 - self.alphas_cumprod) / self.alphas_cumprod)**0.5) + if self.config.use_karras_sigmas: + log_sigmas = np.log(sigmas) + sigmas = np.flip(sigmas).copy() + sigmas = self._convert_to_karras( + in_sigmas=sigmas, num_inference_steps=num_inference_steps) + timesteps = np.array([ + self._sigma_to_t(sigma, log_sigmas) for sigma in sigmas + ]).round() + if self.config.final_sigmas_type == "sigma_min": + sigma_last = sigmas[-1] + elif self.config.final_sigmas_type == "zero": + sigma_last = 0 + else: + raise ValueError( + f"`final_sigmas_type` must be one of 'zero', or 'sigma_min', but got {self.config.final_sigmas_type}" + ) + sigmas = np.concatenate([sigmas, [sigma_last]]).astype(np.float32) + elif self.config.use_exponential_sigmas: + log_sigmas = np.log(sigmas) + sigmas = np.flip(sigmas).copy() + sigmas = self._convert_to_exponential( + in_sigmas=sigmas, num_inference_steps=num_inference_steps) + timesteps = np.array( + [self._sigma_to_t(sigma, log_sigmas) for sigma in sigmas]) + if self.config.final_sigmas_type == "sigma_min": + sigma_last = sigmas[-1] + elif self.config.final_sigmas_type == "zero": + sigma_last = 0 + else: + raise ValueError( + f"`final_sigmas_type` must be one of 'zero', or 'sigma_min', but got {self.config.final_sigmas_type}" + ) + sigmas = np.concatenate([sigmas, [sigma_last]]).astype(np.float32) + elif self.config.use_beta_sigmas: + log_sigmas = np.log(sigmas) + sigmas = np.flip(sigmas).copy() + sigmas = self._convert_to_beta( + in_sigmas=sigmas, num_inference_steps=num_inference_steps) + timesteps = np.array( + [self._sigma_to_t(sigma, log_sigmas) for sigma in sigmas]) + if self.config.final_sigmas_type == "sigma_min": + sigma_last = sigmas[-1] + elif self.config.final_sigmas_type == "zero": + sigma_last = 0 + else: + raise ValueError( + f"`final_sigmas_type` must be one of 'zero', or 'sigma_min', but got {self.config.final_sigmas_type}" + ) + sigmas = np.concatenate([sigmas, [sigma_last]]).astype(np.float32) + elif self.config.use_flow_sigmas: + alphas = np.linspace(1, 1 / self.config.num_train_timesteps, + num_inference_steps + 1) + sigmas = 1.0 - alphas + sigmas = np.flip( + self.config.flow_shift * sigmas / + (1 + (self.config.flow_shift - 1) * sigmas))[:-1].copy() + timesteps = (sigmas * self.config.num_train_timesteps).copy() + if self.config.final_sigmas_type == "sigma_min": + sigma_last = sigmas[-1] + elif self.config.final_sigmas_type == "zero": + sigma_last = 0 + else: + raise ValueError( + f"`final_sigmas_type` must be one of 'zero', or 'sigma_min', but got {self.config.final_sigmas_type}" + ) + sigmas = np.concatenate([sigmas, [sigma_last]]).astype(np.float32) + else: + sigmas = np.interp(timesteps, np.arange(0, len(sigmas)), sigmas) + if self.config.final_sigmas_type == "sigma_min": + sigma_last = ((1 - self.alphas_cumprod[0]) / + self.alphas_cumprod[0])**0.5 + elif self.config.final_sigmas_type == "zero": + sigma_last = 0 + else: + raise ValueError( + f"`final_sigmas_type` must be one of 'zero', or 'sigma_min', but got {self.config.final_sigmas_type}" + ) + sigmas = np.concatenate([sigmas, [sigma_last]]).astype(np.float32) + + self.sigmas = torch.from_numpy(sigmas) + self.timesteps = torch.from_numpy(timesteps).to(device=device, + dtype=torch.int64) + + self.num_inference_steps = len(timesteps) + + self.model_outputs = [ + None, + ] * self.config.solver_order + self.lower_order_nums = 0 + self.last_sample = None + if self.solver_p: + self.solver_p.set_timesteps(self.num_inference_steps, device=device) + + # add an index counter for schedulers that allow duplicated timesteps + self._step_index = None + self._begin_index = None + self.sigmas = self.sigmas.to( + "cpu") # to avoid too much CPU/GPU communication + + # Copied from diffusers.schedulers.scheduling_ddpm.DDPMScheduler._threshold_sample + def _threshold_sample(self, sample: torch.Tensor) -> torch.Tensor: + """ + "Dynamic thresholding: At each sampling step we set s to a certain percentile absolute pixel value in xt0 (the + prediction of x_0 at timestep t), and if s > 1, then we threshold xt0 to the range [-s, s] and then divide by + s. Dynamic thresholding pushes saturated pixels (those near -1 and 1) inwards, thereby actively preventing + pixels from saturation at each step. We find that dynamic thresholding results in significantly better + photorealism as well as better image-text alignment, especially when using very large guidance weights." + + https://arxiv.org/abs/2205.11487 + """ + dtype = sample.dtype + batch_size, channels, *remaining_dims = sample.shape + + if dtype not in (torch.float32, torch.float64): + sample = sample.float( + ) # upcast for quantile calculation, and clamp not implemented for cpu half + + # Flatten sample for doing quantile calculation along each image + sample = sample.reshape(batch_size, channels * np.prod(remaining_dims)) + + abs_sample = sample.abs() # "a certain percentile absolute pixel value" + + s = torch.quantile(abs_sample, + self.config.dynamic_thresholding_ratio, + dim=1) + s = torch.clamp( + s, min=1, max=self.config.sample_max_value + ) # When clamped to min=1, equivalent to standard clipping to [-1, 1] + s = s.unsqueeze( + 1) # (batch_size, 1) because clamp will broadcast along dim=0 + sample = torch.clamp( + sample, -s, s + ) / s # "we threshold xt0 to the range [-s, s] and then divide by s" + + sample = sample.reshape(batch_size, channels, *remaining_dims) + sample = sample.to(dtype) + + return sample + + # Copied from diffusers.schedulers.scheduling_euler_discrete.EulerDiscreteScheduler._sigma_to_t + def _sigma_to_t(self, sigma, log_sigmas) -> torch.Tensor: + # get log sigma + log_sigma = np.log(np.maximum(sigma, 1e-10)) + + # get distribution + dists = log_sigma - log_sigmas[:, np.newaxis] + + # get sigmas range + low_idx = np.cumsum( + (dists >= 0), + axis=0).argmax(axis=0).clip(max=log_sigmas.shape[0] - 2) + high_idx = low_idx + 1 + + low = log_sigmas[low_idx] + high = log_sigmas[high_idx] + + # interpolate sigmas + w = (low - log_sigma) / (low - high) + w = np.clip(w, 0, 1) + + # transform interpolation to time range + t = (1 - w) * low_idx + w * high_idx + t = t.reshape(sigma.shape) + return t + + # Copied from diffusers.schedulers.scheduling_dpmsolver_multistep.DPMSolverMultistepScheduler._sigma_to_alpha_sigma_t + def _sigma_to_alpha_sigma_t( + self, sigma: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: + if self.config.use_flow_sigmas: + alpha_t = 1 - sigma + sigma_t = sigma + else: + alpha_t = 1 / ((sigma**2 + 1)**0.5) + sigma_t = sigma * alpha_t + + return alpha_t, sigma_t + + # Copied from diffusers.schedulers.scheduling_euler_discrete.EulerDiscreteScheduler._convert_to_karras + def _convert_to_karras(self, in_sigmas: torch.Tensor, + num_inference_steps) -> torch.Tensor: + """Constructs the noise schedule of Karras et al. (2022).""" + + # Hack to make sure that other schedulers which copy this function don't break + # TODO: Add this logic to the other schedulers + if hasattr(self.config, "sigma_min"): + sigma_min = self.config.sigma_min + else: + sigma_min = None + + if hasattr(self.config, "sigma_max"): + sigma_max = self.config.sigma_max + else: + sigma_max = None + + sigma_min = sigma_min if sigma_min is not None else in_sigmas[-1].item() + sigma_max = sigma_max if sigma_max is not None else in_sigmas[0].item() + + rho = 7.0 # 7.0 is the value used in the paper + ramp = np.linspace(0, 1, num_inference_steps) + min_inv_rho = sigma_min**(1 / rho) + max_inv_rho = sigma_max**(1 / rho) + sigmas = (max_inv_rho + ramp * (min_inv_rho - max_inv_rho))**rho + return sigmas + + # Copied from diffusers.schedulers.scheduling_euler_discrete.EulerDiscreteScheduler._convert_to_exponential + def _convert_to_exponential(self, in_sigmas: torch.Tensor, + num_inference_steps: int) -> torch.Tensor: + """Constructs an exponential noise schedule.""" + + # Hack to make sure that other schedulers which copy this function don't break + # TODO: Add this logic to the other schedulers + if hasattr(self.config, "sigma_min"): + sigma_min = self.config.sigma_min + else: + sigma_min = None + + if hasattr(self.config, "sigma_max"): + sigma_max = self.config.sigma_max + else: + sigma_max = None + + sigma_min = sigma_min if sigma_min is not None else in_sigmas[-1].item() + sigma_max = sigma_max if sigma_max is not None else in_sigmas[0].item() + + sigmas = np.exp( + np.linspace(math.log(sigma_max), math.log(sigma_min), + num_inference_steps)) + return sigmas + + # Copied from diffusers.schedulers.scheduling_euler_discrete.EulerDiscreteScheduler._convert_to_beta + def _convert_to_beta(self, + in_sigmas: torch.Tensor, + num_inference_steps: int, + alpha: float = 0.6, + beta: float = 0.6) -> torch.Tensor: + """From "Beta Sampling is All You Need" [arXiv:2407.12173] (Lee et. al, 2024)""" + + # Hack to make sure that other schedulers which copy this function don't break + # TODO: Add this logic to the other schedulers + if hasattr(self.config, "sigma_min"): + sigma_min = self.config.sigma_min + else: + sigma_min = None + + if hasattr(self.config, "sigma_max"): + sigma_max = self.config.sigma_max + else: + sigma_max = None + + sigma_min = sigma_min if sigma_min is not None else in_sigmas[-1].item() + sigma_max = sigma_max if sigma_max is not None else in_sigmas[0].item() + + sigmas = np.array([ + sigma_min + (ppf * (sigma_max - sigma_min)) for ppf in [ + scipy.stats.beta.ppf(timestep, alpha, beta) + for timestep in 1 - np.linspace(0, 1, num_inference_steps) + ] + ]) + return sigmas + + def convert_model_output( + self, + model_output: torch.Tensor, + *args, + sample: torch.Tensor = None, + **kwargs, + ) -> torch.Tensor: + r""" + Convert the model output to the corresponding type the UniPC algorithm needs. + + Args: + model_output (`torch.Tensor`): + The direct output from the learned diffusion model. + timestep (`int`): + The current discrete timestep in the diffusion chain. + sample (`torch.Tensor`): + A current instance of a sample created by the diffusion process. + + Returns: + `torch.Tensor`: + The converted model output. + """ + timestep = args[0] if len(args) > 0 else kwargs.pop("timestep", None) + if sample is None: + if len(args) > 1: + sample = args[1] + else: + raise ValueError( + "missing `sample` as a required keyword argument") + if timestep is not None: + deprecate( + "timesteps", + "1.0.0", + "Passing `timesteps` is deprecated and has no effect as model output conversion is now handled via an internal counter `self.step_index`", + ) + + sigma = self.sigmas[self.step_index] + alpha_t, sigma_t = self._sigma_to_alpha_sigma_t(sigma) + + if self.predict_x0: + if self.config.prediction_type == "epsilon": + x0_pred = (sample - sigma_t * model_output) / alpha_t + elif self.config.prediction_type == "sample": + x0_pred = model_output + elif self.config.prediction_type == "v_prediction": + x0_pred = alpha_t * sample - sigma_t * model_output + elif self.config.prediction_type == "flow_prediction": + sigma_t = self.sigmas[self.step_index] + x0_pred = sample - sigma_t * model_output + else: + raise ValueError( + f"prediction_type given as {self.config.prediction_type} must be one of `epsilon`, `sample`, " + "`v_prediction`, or `flow_prediction` for the UniPCMultistepScheduler." + ) + + if self.config.thresholding: + x0_pred = self._threshold_sample(x0_pred) + + return x0_pred + else: + if self.config.prediction_type == "epsilon": + return model_output + elif self.config.prediction_type == "sample": + epsilon = (sample - alpha_t * model_output) / sigma_t + return epsilon + elif self.config.prediction_type == "v_prediction": + epsilon = alpha_t * model_output + sigma_t * sample + return epsilon + else: + raise ValueError( + f"prediction_type given as {self.config.prediction_type} must be one of `epsilon`, `sample`, or" + " `v_prediction` for the UniPCMultistepScheduler.") + + def multistep_uni_p_bh_update( + self, + model_output: torch.Tensor, + *args, + sample: torch.Tensor = None, + order: Optional[int] = None, + **kwargs, + ) -> torch.Tensor: + """ + One step for the UniP (B(h) version). Alternatively, `self.solver_p` is used if is specified. + + Args: + model_output (`torch.Tensor`): + The direct output from the learned diffusion model at the current timestep. + prev_timestep (`int`): + The previous discrete timestep in the diffusion chain. + sample (`torch.Tensor`): + A current instance of a sample created by the diffusion process. + order (`int`): + The order of UniP at this timestep (corresponds to the *p* in UniPC-p). + + Returns: + `torch.Tensor`: + The sample tensor at the previous timestep. + """ + prev_timestep = args[0] if len(args) > 0 else kwargs.pop( + "prev_timestep", None) + if sample is None: + if len(args) > 1: + sample = args[1] + else: + raise ValueError( + " missing `sample` as a required keyword argument") + if order is None: + if len(args) > 2: + order = args[2] + else: + raise ValueError( + " missing `order` as a required keyword argument") + if prev_timestep is not None: + deprecate( + "prev_timestep", + "1.0.0", + "Passing `prev_timestep` is deprecated and has no effect as model output conversion is now handled via an internal counter `self.step_index`", + ) + model_output_list = self.model_outputs + + s0 = self.timestep_list[-1] + m0 = model_output_list[-1] + x = sample + + if self.solver_p: + x_t = self.solver_p.step(model_output, s0, x).prev_sample + return x_t + + sigma_t, sigma_s0 = self.sigmas[self.step_index + + 1], self.sigmas[self.step_index] + alpha_t, sigma_t = self._sigma_to_alpha_sigma_t(sigma_t) + alpha_s0, sigma_s0 = self._sigma_to_alpha_sigma_t(sigma_s0) + + lambda_t = torch.log(alpha_t) - torch.log(sigma_t) + lambda_s0 = torch.log(alpha_s0) - torch.log(sigma_s0) + + h = lambda_t - lambda_s0 + device = sample.device + + rks = [] + D1s = [] + for i in range(1, order): + si = self.step_index - i + mi: torch.Tensor = model_output_list[-(i + 1)] + alpha_si, sigma_si = self._sigma_to_alpha_sigma_t(self.sigmas[si]) + lambda_si = torch.log(alpha_si) - torch.log(sigma_si) + rk = (lambda_si - lambda_s0) / h + rks.append(rk) + D1s.append((mi - m0) / rk) + + rks.append(1.0) + rks = torch.tensor(rks, device=device) + + R = [] + b = [] + + hh = -h if self.predict_x0 else h + h_phi_1 = torch.expm1(hh) # h\phi_1(h) = e^h - 1 + h_phi_k = h_phi_1 / hh - 1 + + factorial_i = 1 + + if self.config.solver_type == "bh1": + B_h = hh + elif self.config.solver_type == "bh2": + B_h = torch.expm1(hh) + else: + raise NotImplementedError() + + for i in range(1, order + 1): + R.append(torch.pow(rks, i - 1)) + b.append(h_phi_k * factorial_i / B_h) + factorial_i *= i + 1 + h_phi_k = h_phi_k / hh - 1 / factorial_i + + R_tensor: torch.Tensor = torch.stack(R) + b = torch.tensor(b, device=device) + + D1s_tensor: Optional[torch.Tensor] = None + if len(D1s) > 0: + D1s_tensor = torch.stack(D1s, dim=1) # (B, K) + # for order 2, we use a simplified version + if order == 2: + rhos_p = torch.tensor([0.5], dtype=x.dtype, device=device) + else: + rhos_p = torch.linalg.solve(R_tensor[:-1, :-1], + b[:-1]).to(device).to(x.dtype) + else: + D1s_tensor = None + + if self.predict_x0: + x_t_ = sigma_t / sigma_s0 * x - alpha_t * h_phi_1 * m0 + if D1s_tensor is not None: + pred_res = torch.einsum("k,bkc...->bc...", rhos_p, D1s_tensor) + else: + pred_res = 0 + x_t = x_t_ - alpha_t * B_h * pred_res + else: + x_t_ = alpha_t / alpha_s0 * x - sigma_t * h_phi_1 * m0 + if D1s_tensor is not None: + pred_res = torch.einsum("k,bkc...->bc...", rhos_p, D1s_tensor) + else: + pred_res = 0 + x_t = x_t_ - sigma_t * B_h * pred_res + + x_t = x_t.to(x.dtype) + return x_t + + def multistep_uni_c_bh_update( + self, + this_model_output: torch.Tensor, + *args, + last_sample: Optional[torch.Tensor] = None, + this_sample: Optional[torch.Tensor] = None, + order: Optional[int] = None, + **kwargs, + ) -> torch.Tensor: + """ + One step for the UniC (B(h) version). + + Args: + this_model_output (`torch.Tensor`): + The model outputs at `x_t`. + this_timestep (`int`): + The current timestep `t`. + last_sample (`torch.Tensor`): + The generated sample before the last predictor `x_{t-1}`. + this_sample (`torch.Tensor`): + The generated sample after the last predictor `x_{t}`. + order (`int`): + The `p` of UniC-p at this step. The effective order of accuracy should be `order + 1`. + + Returns: + `torch.Tensor`: + The corrected sample tensor at the current timestep. + """ + this_timestep = args[0] if len(args) > 0 else kwargs.pop( + "this_timestep", None) + if last_sample is None: + if len(args) > 1: + last_sample = args[1] + else: + raise ValueError( + " missing`last_sample` as a required keyword argument") + if this_sample is None: + if len(args) > 2: + this_sample = args[2] + else: + raise ValueError( + " missing`this_sample` as a required keyword argument") + if order is None: + if len(args) > 3: + order = args[3] + else: + raise ValueError( + " missing`order` as a required keyword argument") + if this_timestep is not None: + deprecate( + "this_timestep", + "1.0.0", + "Passing `this_timestep` is deprecated and has no effect as model output conversion is now handled via an internal counter `self.step_index`", + ) + + model_output_list = self.model_outputs + + m0 = model_output_list[-1] + x = last_sample + x_t = this_sample + model_t = this_model_output + + sigma_t, sigma_s0 = self.sigmas[self.step_index], self.sigmas[ + self.step_index - 1] + alpha_t, sigma_t = self._sigma_to_alpha_sigma_t(sigma_t) + alpha_s0, sigma_s0 = self._sigma_to_alpha_sigma_t(sigma_s0) + + lambda_t = torch.log(alpha_t) - torch.log(sigma_t) + lambda_s0 = torch.log(alpha_s0) - torch.log(sigma_s0) + + h = lambda_t - lambda_s0 + device = this_sample.device + + rks = [] + D1s = [] + for i in range(1, order): + si = self.step_index - (i + 1) + mi: torch.Tensor = model_output_list[-(i + 1)] + alpha_si, sigma_si = self._sigma_to_alpha_sigma_t(self.sigmas[si]) + lambda_si = torch.log(alpha_si) - torch.log(sigma_si) + rk = (lambda_si - lambda_s0) / h + rks.append(rk) + D1s.append((mi - m0) / rk) + + rks.append(1.0) + rks = torch.tensor(rks, device=device) + + R = [] + b = [] + + hh = -h if self.predict_x0 else h + h_phi_1 = torch.expm1(hh) # h\phi_1(h) = e^h - 1 + h_phi_k = h_phi_1 / hh - 1 + + factorial_i = 1 + + if self.config.solver_type == "bh1": + B_h = hh + elif self.config.solver_type == "bh2": + B_h = torch.expm1(hh) + else: + raise NotImplementedError() + + for i in range(1, order + 1): + R.append(torch.pow(rks, i - 1)) + b.append(h_phi_k * factorial_i / B_h) + factorial_i *= i + 1 + h_phi_k = h_phi_k / hh - 1 / factorial_i + + R = torch.stack(R) + b = torch.tensor(b, device=device) + + D1s_tensor: Optional[torch.Tensor] = torch.stack( + D1s, dim=1) if len(D1s) > 0 else None + + # for order 1, we use a simplified version + if order == 1: + rhos_c = torch.tensor([0.5], dtype=x.dtype, device=device) + else: + rhos_c = torch.linalg.solve(R, b).to(device).to(x.dtype) + + if self.predict_x0: + x_t_ = sigma_t / sigma_s0 * x - alpha_t * h_phi_1 * m0 + if D1s_tensor is not None: + corr_res = torch.einsum("k,bkc...->bc...", rhos_c[:-1], + D1s_tensor) + else: + corr_res = 0 + D1_t = model_t - m0 + x_t = x_t_ - alpha_t * B_h * (corr_res + rhos_c[-1] * D1_t) + else: + x_t_ = alpha_t / alpha_s0 * x - sigma_t * h_phi_1 * m0 + if D1s_tensor is not None: + corr_res = torch.einsum("k,bkc...->bc...", rhos_c[:-1], + D1s_tensor) + else: + corr_res = 0 + D1_t = model_t - m0 + x_t = x_t_ - sigma_t * B_h * (corr_res + rhos_c[-1] * D1_t) + x_t = x_t.to(x.dtype) + return x_t + + # Copied from diffusers.schedulers.scheduling_dpmsolver_multistep.DPMSolverMultistepScheduler.index_for_timestep + def index_for_timestep(self, timestep, schedule_timesteps=None) -> int: + if schedule_timesteps is None: + schedule_timesteps = self.timesteps + + index_candidates = (schedule_timesteps == timestep).nonzero() + + if len(index_candidates) == 0: + step_index = len(self.timesteps) - 1 + # The sigma index that is taken for the **very** first `step` + # is always the second index (or the last index if there is only 1) + # This way we can ensure we don't accidentally skip a sigma in + # case we start in the middle of the denoising schedule (e.g. for image-to-image) + elif len(index_candidates) > 1: + step_index = index_candidates[1].item() + else: + step_index = index_candidates[0].item() + + return step_index + + # Copied from diffusers.schedulers.scheduling_dpmsolver_multistep.DPMSolverMultistepScheduler._init_step_index + def _init_step_index(self, timestep) -> None: + """ + Initialize the step_index counter for the scheduler. + """ + + if self.begin_index is None: + if isinstance(timestep, torch.Tensor): + timestep = timestep.to(self.timesteps.device) + self._step_index = self.index_for_timestep(timestep) + else: + self._step_index = self._begin_index + + def step( + self, + model_output: torch.Tensor, + timestep: Union[int, torch.Tensor], + sample: torch.Tensor, + return_dict: bool = True, + ) -> Union[SchedulerOutput, Tuple]: + """ + Predict the sample from the previous timestep by reversing the SDE. This function propagates the sample with + the multistep UniPC. + + Args: + model_output (`torch.Tensor`): + The direct output from learned diffusion model. + timestep (`int`): + The current discrete timestep in the diffusion chain. + sample (`torch.Tensor`): + A current instance of a sample created by the diffusion process. + return_dict (`bool`): + Whether or not to return a [`~schedulers.scheduling_utils.SchedulerOutput`] or `tuple`. + + Returns: + [`~schedulers.scheduling_utils.SchedulerOutput`] or `tuple`: + If return_dict is `True`, [`~schedulers.scheduling_utils.SchedulerOutput`] is returned, otherwise a + tuple is returned where the first element is the sample tensor. + + """ + if self.num_inference_steps is None: + raise ValueError( + "Number of inference steps is 'None', you need to run 'set_timesteps' after creating the scheduler" + ) + + if self.step_index is None: + self._init_step_index(timestep) + + assert self.step_index is not None + use_corrector = (self.step_index > 0 + and self.step_index - 1 not in self.disable_corrector + and self.last_sample is not None) + + model_output_convert = self.convert_model_output(model_output, + sample=sample) + if use_corrector: + sample = self.multistep_uni_c_bh_update( + this_model_output=model_output_convert, + last_sample=self.last_sample, + this_sample=sample, + order=self.this_order, + ) + + for i in range(self.config.solver_order - 1): + self.model_outputs[i] = self.model_outputs[i + 1] + self.timestep_list[i] = self.timestep_list[i + 1] + + self.model_outputs[-1] = model_output_convert + self.timestep_list[-1] = timestep + + if self.config.lower_order_final: + this_order = min(self.config.solver_order, + len(self.timesteps) - self.step_index) + else: + this_order = self.config.solver_order + + self.this_order: int = min(this_order, self.lower_order_nums + + 1) # warmup for multistep + assert self.this_order > 0 + + self.last_sample = sample + prev_sample = self.multistep_uni_p_bh_update( + model_output= + model_output, # pass the original non-converted model output, in case solver-p is used + sample=sample, + order=self.this_order, + ) + + if self.lower_order_nums < self.config.solver_order: + self.lower_order_nums += 1 + + # upon completion increase step index by one + assert self._step_index is not None + self._step_index += 1 + + if not return_dict: + return (prev_sample, ) + + return SchedulerOutput(prev_sample=prev_sample) + + def scale_model_input(self, sample: torch.Tensor, *args, + **kwargs) -> torch.Tensor: + """ + Ensures interchangeability with schedulers that need to scale the denoising model input depending on the + current timestep. + + Args: + sample (`torch.Tensor`): + The input sample. + + Returns: + `torch.Tensor`: + A scaled input sample. + """ + return sample + + # Copied from diffusers.schedulers.scheduling_dpmsolver_multistep.DPMSolverMultistepScheduler.add_noise + def add_noise( + self, + original_samples: torch.Tensor, + noise: torch.Tensor, + timesteps: torch.IntTensor, + ) -> torch.Tensor: + # Make sure sigmas and timesteps have the same device and dtype as original_samples + sigmas = self.sigmas.to(device=original_samples.device, + dtype=original_samples.dtype) + if original_samples.device.type == "mps" and torch.is_floating_point( + timesteps): + # mps does not support float64 + schedule_timesteps = self.timesteps.to(original_samples.device, + dtype=torch.float32) + timesteps = timesteps.to(original_samples.device, + dtype=torch.float32) + else: + schedule_timesteps = self.timesteps.to(original_samples.device) + timesteps = timesteps.to(original_samples.device) + + # begin_index is None when the scheduler is used for training or pipeline does not implement set_begin_index + if self.begin_index is None: + step_indices = [ + self.index_for_timestep(t, schedule_timesteps) + for t in timesteps + ] + elif self.step_index is not None: + # add_noise is called after first denoising step (for inpainting) + step_indices = [self.step_index] * timesteps.shape[0] + else: + # add noise is called before first denoising step to create initial latent(img2img) + step_indices = [self.begin_index] * timesteps.shape[0] + + sigma = sigmas[step_indices].flatten() + while len(sigma.shape) < len(original_samples.shape): + sigma = sigma.unsqueeze(-1) + + alpha_t, sigma_t = self._sigma_to_alpha_sigma_t(sigma) + noisy_samples = alpha_t * original_samples + sigma_t * noise + return noisy_samples + + def __len__(self): + return self.config.num_train_timesteps diff --git a/fastvideo/v1/models/vaes/common.py b/fastvideo/v1/models/vaes/common.py index 8e9cc5920..c56d1454d 100644 --- a/fastvideo/v1/models/vaes/common.py +++ b/fastvideo/v1/models/vaes/common.py @@ -2,11 +2,12 @@ from abc import ABC, abstractmethod from math import prod -from typing import Iterator, Optional, Tuple +from typing import Iterator, Optional, Tuple, Union import numpy as np import torch import torch.distributed as dist +from diffusers.utils.torch_utils import randn_tensor from fastvideo.v1.distributed import (get_sequence_model_parallel_rank, get_sequence_model_parallel_world_size) @@ -20,8 +21,11 @@ class ParallelTiledVAE(ABC): tile_sample_stride_width: int tile_sample_stride_num_frames: int use_tiling: bool + use_temporal_tiling: bool + use_parallel_tiling: bool temporal_compression_ratio: int spatial_compression_ratio: int + scaling_factor: Union[float, torch.tensor] def __init__(self, *args, **kwargs) -> None: # Check if subclass has defined all required properties @@ -30,7 +34,8 @@ class ParallelTiledVAE(ABC): 'tile_sample_min_num_frames', 'tile_sample_stride_height', 'tile_sample_stride_width', 'tile_sample_stride_num_frames', 'spatial_compression_ratio', 'temporal_compression_ratio', - 'use_tiling' + 'use_tiling', 'use_temporal_tiling', 'use_parallel_tiling', + 'scaling_factor' ] for attr in required_attributes: @@ -52,13 +57,13 @@ class ParallelTiledVAE(ABC): latent_num_frames = (num_frames - 1) // self.temporal_compression_ratio + 1 - if self.use_tiling and num_frames > self.tile_sample_min_num_frames: + if self.use_tiling and self.use_temporal_tiling and num_frames > self.tile_sample_min_num_frames: latents = self.tiled_encode(x)[:, :, :latent_num_frames] elif self.use_tiling and (width > self.tile_sample_min_width or height > self.tile_sample_min_height): - latents = self.spatial_tiled_encode(x) + latents = self.spatial_tiled_encode(x)[:, :, :latent_num_frames] else: - latents = self._encode(x) + latents = self._encode(x)[:, :, :latent_num_frames] return DiagonalGaussianDistribution(latents) def decode(self, z: torch.Tensor) -> torch.Tensor: @@ -69,16 +74,17 @@ class ParallelTiledVAE(ABC): num_sample_frames = (num_frames - 1) * self.temporal_compression_ratio + 1 - if self.use_tiling and get_sequence_model_parallel_world_size() > 1: + if self.use_tiling and self.use_parallel_tiling and get_sequence_model_parallel_world_size( + ) > 1: return self.parallel_tiled_decode(z)[:, :, :num_sample_frames] - if self.use_tiling and num_frames > tile_latent_min_num_frames: + if self.use_tiling and self.use_temporal_tiling and num_frames > tile_latent_min_num_frames: return self.tiled_decode(z)[:, :, :num_sample_frames] if self.use_tiling and (width > tile_latent_min_width or height > tile_latent_min_height): - return self.spatial_tiled_decode(z) + return self.spatial_tiled_decode(z)[:, :, :num_sample_frames] - return self._decode(z) + return self._decode(z)[:, :, :num_sample_frames] def blend_v(self, a: torch.Tensor, b: torch.Tensor, blend_extent: int) -> torch.Tensor: @@ -462,7 +468,7 @@ class DiagonalGaussianDistribution: def sample(self, generator: Optional[torch.Generator] = None) -> torch.Tensor: # make sure sample is on the same device as the parameters and has same dtype - sample = torch.randn( + sample = randn_tensor( self.mean.shape, generator=generator, device=self.parameters.device, diff --git a/fastvideo/v1/models/vaes/hunyuanvae.py b/fastvideo/v1/models/vaes/hunyuanvae.py index 0d5fd4400..e68e5089b 100644 --- a/fastvideo/v1/models/vaes/hunyuanvae.py +++ b/fastvideo/v1/models/vaes/hunyuanvae.py @@ -847,6 +847,9 @@ class AutoencoderKLHunyuanVideo(nn.Module, ParallelTiledVAE): # frames spatially into smaller tiles and performing multiple forward passes for decoding, and then blending the # intermediate tiles together, the memory requirement can be lowered. self.use_tiling = True + self.use_temporal_tiling = True + self.use_parallel_tiling = True + self.scaling_factor = scaling_factor # The minimal tile height and width for spatial tiling to be used self.tile_sample_min_height = 256 diff --git a/fastvideo/v1/models/vaes/wanvae.py b/fastvideo/v1/models/vaes/wanvae.py index 094631875..691861837 100644 --- a/fastvideo/v1/models/vaes/wanvae.py +++ b/fastvideo/v1/models/vaes/wanvae.py @@ -1,3 +1,5 @@ +# SPDX-License-Identifier: Apache-2.0 + # Copyright 2025 The Wan Team and The HuggingFace Team. All rights reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); @@ -20,8 +22,8 @@ import torch.nn.functional as F import torch.utils.checkpoint from fastvideo.v1.layers.activation import get_act_fn +from fastvideo.v1.models.utils import auto_attributes from fastvideo.v1.models.vaes.common import ParallelTiledVAE -from fastvideo.v1.utils import auto_attributes CACHE_T = 2 @@ -647,6 +649,10 @@ class AutoencoderKLWan(nn.Module, ParallelTiledVAE): self.temperal_upsample = list(temperal_downsample)[::-1] self.latents_mean = list(latents_mean) self.latents_std = list(latents_std) + self.scaling_factor = 1.0 / torch.tensor(self.config.latents_std).view( + 1, self.config.z_dim, 1, 1, 1) + self.shift_factor = torch.tensor(self.config.latents_mean).view( + 1, self.config.z_dim, 1, 1, 1) if load_encoder: self.encoder = WanEncoder3d(base_dim, z_dim * 2, dim_mult, @@ -661,6 +667,8 @@ class AutoencoderKLWan(nn.Module, ParallelTiledVAE): self.temperal_upsample, dropout) self.use_tiling = True + self.use_temporal_tiling = False + self.use_parallel_tiling = False self.spatial_compression_ratio = 8 self.temporal_compression_ratio = 4 @@ -691,12 +699,16 @@ class AutoencoderKLWan(nn.Module, ParallelTiledVAE): enc = torch.cat([first_frame, enc], dim=2) return enc + def spatial_tiled_encode(self, x: torch.Tensor) -> torch.Tensor: + first_frame = x[:, :, 0, :, :].unsqueeze(2) + first_frame = self._encode(first_frame, first_frame=True) + + enc = ParallelTiledVAE.spatial_tiled_encode(self, x) + enc = enc[:, :, 1:] + enc = torch.cat([first_frame, enc], dim=2) + return enc + def _decode(self, z: torch.Tensor, first_frame=False) -> torch.Tensor: - latents_mean = (torch.tensor(self.latents_mean).view( - 1, self.z_dim, 1, 1, 1).to(z.device, z.dtype)) - latents_std = 1.0 / torch.tensor(self.latents_std).view( - 1, self.z_dim, 1, 1, 1).to(z.device, z.dtype) - z = z / latents_std + latents_mean x = self.post_quant_conv(z) out = self.decoder(x, first_frame=first_frame) @@ -711,6 +723,12 @@ class AutoencoderKLWan(nn.Module, ParallelTiledVAE): dec = dec[:, :, start_frame_idx:] return dec + def spatial_tiled_decode(self, z: torch.Tensor) -> torch.Tensor: + dec = ParallelTiledVAE.spatial_tiled_decode(self, z) + start_frame_idx = self.temporal_compression_ratio - 1 + dec = dec[:, :, start_frame_idx:] + return dec + def parallel_tiled_decode(self, z: torch.FloatTensor) -> torch.FloatTensor: self.blend_num_frames *= 2 dec = ParallelTiledVAE.parallel_tiled_decode(self, z) diff --git a/fastvideo/v1/models/vision_utils.py b/fastvideo/v1/models/vision_utils.py new file mode 100644 index 000000000..b86f7489a --- /dev/null +++ b/fastvideo/v1/models/vision_utils.py @@ -0,0 +1,220 @@ +# SPDX-License-Identifier: Apache-2.0 + +import os +from typing import Callable, List, Optional, Tuple, Union + +import numpy as np +import PIL.Image +import PIL.ImageOps +import requests +import torch +from packaging import version + +if version.parse(version.parse( + PIL.__version__).base_version) >= version.parse("9.1.0"): + PIL_INTERPOLATION = { + "linear": PIL.Image.Resampling.BILINEAR, + "bilinear": PIL.Image.Resampling.BILINEAR, + "bicubic": PIL.Image.Resampling.BICUBIC, + "lanczos": PIL.Image.Resampling.LANCZOS, + "nearest": PIL.Image.Resampling.NEAREST, + } +else: + PIL_INTERPOLATION = { + "linear": PIL.Image.LINEAR, + "bilinear": PIL.Image.BILINEAR, + "bicubic": PIL.Image.BICUBIC, + "lanczos": PIL.Image.LANCZOS, + "nearest": PIL.Image.NEAREST, + } + + +def pil_to_numpy( + images: Union[List[PIL.Image.Image], PIL.Image.Image]) -> np.ndarray: + r""" + Convert a PIL image or a list of PIL images to NumPy arrays. + + Args: + images (`PIL.Image.Image` or `List[PIL.Image.Image]`): + The PIL image or list of images to convert to NumPy format. + + Returns: + `np.ndarray`: + A NumPy array representation of the images. + """ + if not isinstance(images, list): + images = [images] + images = [np.array(image).astype(np.float32) / 255.0 for image in images] + images_arr: np.ndarray = np.stack(images, axis=0) + + return images_arr + + +def numpy_to_pt(images: np.ndarray) -> torch.Tensor: + r""" + Convert a NumPy image to a PyTorch tensor. + + Args: + images (`np.ndarray`): + The NumPy image array to convert to PyTorch format. + + Returns: + `torch.Tensor`: + A PyTorch tensor representation of the images. + """ + if images.ndim == 3: + images = images[..., None] + + images = torch.from_numpy(images.transpose(0, 3, 1, 2)) + return images + + +def normalize( + images: Union[np.ndarray, + torch.Tensor]) -> Union[np.ndarray, torch.Tensor]: + r""" + Normalize an image array to [-1,1]. + + Args: + images (`np.ndarray` or `torch.Tensor`): + The image array to normalize. + + Returns: + `np.ndarray` or `torch.Tensor`: + The normalized image array. + """ + return 2.0 * images - 1.0 + + +def load_image( + image: Union[str, PIL.Image.Image], + convert_method: Optional[Callable[[PIL.Image.Image], + PIL.Image.Image]] = None +) -> PIL.Image.Image: + """ + Loads `image` to a PIL Image. + + Args: + image (`str` or `PIL.Image.Image`): + The image to convert to the PIL Image format. + convert_method (Callable[[PIL.Image.Image], PIL.Image.Image], *optional*): + A conversion method to apply to the image after loading it. When set to `None` the image will be converted + "RGB". + + Returns: + `PIL.Image.Image`: + A PIL Image. + """ + if isinstance(image, str): + if image.startswith("http://") or image.startswith("https://"): + image = PIL.Image.open(requests.get(image, stream=True).raw) + elif os.path.isfile(image): + image = PIL.Image.open(image) + else: + raise ValueError( + f"Incorrect path or URL. URLs must start with `http://` or `https://`, and {image} is not a valid path." + ) + elif isinstance(image, PIL.Image.Image): + image = image + else: + raise ValueError( + "Incorrect format used for the image. Should be a URL linking to an image, a local path, or a PIL image." + ) + + image = PIL.ImageOps.exif_transpose(image) + + if convert_method is not None: + image = convert_method(image) + else: + image = image.convert("RGB") + + return image + + +def get_default_height_width( + image: Union[PIL.Image.Image, np.ndarray, torch.Tensor], + vae_scale_factor: int, + height: Optional[int] = None, + width: Optional[int] = None, +) -> Tuple[int, int]: + r""" + Returns the height and width of the image, downscaled to the next integer multiple of `vae_scale_factor`. + + Args: + image (`Union[PIL.Image.Image, np.ndarray, torch.Tensor]`): + The image input, which can be a PIL image, NumPy array, or PyTorch tensor. If it is a NumPy array, it + should have shape `[batch, height, width]` or `[batch, height, width, channels]`. If it is a PyTorch + tensor, it should have shape `[batch, channels, height, width]`. + height (`Optional[int]`, *optional*, defaults to `None`): + The height of the preprocessed image. If `None`, the height of the `image` input will be used. + width (`Optional[int]`, *optional*, defaults to `None`): + The width of the preprocessed image. If `None`, the width of the `image` input will be used. + + Returns: + `Tuple[int, int]`: + A tuple containing the height and width, both resized to the nearest integer multiple of + `vae_scale_factor`. + """ + + if height is None: + if isinstance(image, PIL.Image.Image): + height = image.height + elif isinstance(image, torch.Tensor): + height = image.shape[2] + else: + height = image.shape[1] + + if width is None: + if isinstance(image, PIL.Image.Image): + width = image.width + elif isinstance(image, torch.Tensor): + width = image.shape[3] + else: + width = image.shape[2] + + width, height = (x - x % vae_scale_factor for x in (width, height) + ) # resize to integer multiple of vae_scale_factor + + return height, width + + +def resize( + image: Union[PIL.Image.Image, np.ndarray, torch.Tensor], + height: int, + width: int, + resize_mode: str = "default", # "default", "fill", "crop" + resample: str = "lanczos", +) -> Union[PIL.Image.Image, np.ndarray, torch.Tensor]: + """ + Resize image. + + Args: + image (`PIL.Image.Image`, `np.ndarray` or `torch.Tensor`): + The image input, can be a PIL image, numpy array or pytorch tensor. + height (`int`): + The height to resize to. + width (`int`): + The width to resize to. + resize_mode (`str`, *optional*, defaults to `default`): + The resize mode to use, can be one of `default` or `fill`. If `default`, will resize the image to fit + within the specified width and height, and it may not maintaining the original aspect ratio. If `fill`, + will resize the image to fit within the specified width and height, maintaining the aspect ratio, and + then center the image within the dimensions, filling empty with data from image. If `crop`, will resize + the image to fit within the specified width and height, maintaining the aspect ratio, and then center + the image within the dimensions, cropping the excess. Note that resize_mode `fill` and `crop` are only + supported for PIL image input. + + Returns: + `PIL.Image.Image`, `np.ndarray` or `torch.Tensor`: + The resized image. + """ + if resize_mode != "default" and not isinstance(image, PIL.Image.Image): + raise ValueError( + f"Only PIL image input is supported for resize_mode {resize_mode}") + assert isinstance(image, PIL.Image.Image) + if resize_mode == "default": + image = image.resize((width, height), + resample=PIL_INTERPOLATION[resample]) + else: + raise ValueError(f"resize_mode {resize_mode} is not supported") + return image diff --git a/fastvideo/v1/pipelines/pipeline_batch_info.py b/fastvideo/v1/pipelines/pipeline_batch_info.py index 2c448ccd7..e89bdf75c 100644 --- a/fastvideo/v1/pipelines/pipeline_batch_info.py +++ b/fastvideo/v1/pipelines/pipeline_batch_info.py @@ -29,6 +29,10 @@ class ForwardBatch: generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None + # Image inputs + image_path: Optional[str] = None + image_embeds: List[torch.Tensor] = field(default_factory=list) + # Text inputs prompt: Optional[Union[str, List[str]]] = None negative_prompt: Optional[Union[str, List[str]]] = None @@ -55,6 +59,7 @@ class ForwardBatch: # Latent tensors latents: Optional[torch.Tensor] = None noise_pred: Optional[torch.Tensor] = None + image_latent: Optional[torch.Tensor] = None # Latent dimensions num_channels_latents: Optional[int] = None @@ -100,3 +105,4 @@ class ForwardBatch: # Set do_classifier_free_guidance based on guidance scale and negative prompt if self.guidance_scale > 1.0: self.do_classifier_free_guidance = True + self.negative_prompt_embeds = [] diff --git a/fastvideo/v1/pipelines/stages/__init__.py b/fastvideo/v1/pipelines/stages/__init__.py index 613ae46b1..432161dc2 100644 --- a/fastvideo/v1/pipelines/stages/__init__.py +++ b/fastvideo/v1/pipelines/stages/__init__.py @@ -7,15 +7,19 @@ complete diffusion pipelines. """ from fastvideo.v1.pipelines.stages.base import PipelineStage +from fastvideo.v1.pipelines.stages.clip_image_encoding import ( + CLIPImageEncodingStage) from fastvideo.v1.pipelines.stages.clip_text_encoding import ( CLIPTextEncodingStage) from fastvideo.v1.pipelines.stages.conditioning import ConditioningStage from fastvideo.v1.pipelines.stages.decoding import DecodingStage from fastvideo.v1.pipelines.stages.denoising import DenoisingStage +from fastvideo.v1.pipelines.stages.encoding import EncodingStage from fastvideo.v1.pipelines.stages.input_validation import InputValidationStage from fastvideo.v1.pipelines.stages.latent_preparation import ( LatentPreparationStage) from fastvideo.v1.pipelines.stages.llama_encoding import LlamaEncodingStage +from fastvideo.v1.pipelines.stages.t5_encoding import T5EncodingStage from fastvideo.v1.pipelines.stages.timestep_preparation import ( TimestepPreparationStage) @@ -26,7 +30,10 @@ __all__ = [ "LatentPreparationStage", "ConditioningStage", "DenoisingStage", + "EncodingStage", "DecodingStage", "LlamaEncodingStage", + "T5EncodingStage", "CLIPTextEncodingStage", + "CLIPImageEncodingStage", ] diff --git a/fastvideo/v1/pipelines/stages/clip_image_encoding.py b/fastvideo/v1/pipelines/stages/clip_image_encoding.py new file mode 100644 index 000000000..5b9d77807 --- /dev/null +++ b/fastvideo/v1/pipelines/stages/clip_image_encoding.py @@ -0,0 +1,71 @@ +# SPDX-License-Identifier: Apache-2.0 +""" +Image encoding stages for I2V diffusion pipelines. + +This module contains implementations of image encoding stages for diffusion pipelines. +""" + +import torch + +from fastvideo.v1.forward_context import set_forward_context +from fastvideo.v1.inference_args import InferenceArgs +from fastvideo.v1.logger import init_logger +from fastvideo.v1.models.vision_utils import load_image +from fastvideo.v1.pipelines.pipeline_batch_info import ForwardBatch +from fastvideo.v1.pipelines.stages.base import PipelineStage + +logger = init_logger(__name__) + + +class CLIPImageEncodingStage(PipelineStage): + """ + Stage for encoding image prompts into embeddings for diffusion models. + + This stage handles the encoding of image prompts into the embedding space + expected by the diffusion model. + """ + + def __init__(self, image_encoder, image_processor) -> None: + """ + Initialize the prompt encoding stage. + + Args: + enable_logging: Whether to enable logging for this stage. + is_secondary: Whether this is a secondary image encoder. + """ + super().__init__() + self.image_processor = image_processor + self.image_encoder = image_encoder + + def forward( + self, + batch: ForwardBatch, + inference_args: InferenceArgs, + ) -> ForwardBatch: + """ + Encode the prompt into image encoder hidden states. + + Args: + batch: The current batch information. + inference_args: The inference arguments. + + Returns: + The batch with encoded prompt embeddings. + """ + if inference_args.use_cpu_offload: + self.image_encoder = self.image_encoder.to(batch.device) + + image = load_image(batch.image_path) + + image_inputs = self.image_processor( + images=image, return_tensors="pt").to(batch.device) + with set_forward_context(current_timestep=0, attn_metadata=None): + image_embeds = self.image_encoder(**image_inputs) + + batch.image_embeds.append(image_embeds) + + if inference_args.use_cpu_offload: + self.image_encoder.to('cpu') + torch.cuda.empty_cache() + + return batch diff --git a/fastvideo/v1/pipelines/stages/clip_text_encoding.py b/fastvideo/v1/pipelines/stages/clip_text_encoding.py index 490f43f01..50be5aab6 100644 --- a/fastvideo/v1/pipelines/stages/clip_text_encoding.py +++ b/fastvideo/v1/pipelines/stages/clip_text_encoding.py @@ -5,6 +5,8 @@ Prompt encoding stages for diffusion pipelines. This module contains implementations of prompt encoding stages for diffusion pipelines. """ +import torch + from fastvideo.v1.forward_context import set_forward_context from fastvideo.v1.inference_args import InferenceArgs from fastvideo.v1.logger import init_logger @@ -49,6 +51,8 @@ class CLIPTextEncodingStage(PipelineStage): Returns: The batch with encoded prompt embeddings. """ + if inference_args.use_cpu_offload: + self.text_encoder = self.text_encoder.to(batch.device) text_inputs = self.tokenizer( batch.prompt, @@ -64,4 +68,24 @@ class CLIPTextEncodingStage(PipelineStage): batch.prompt_embeds.append(prompt_embeds) + if batch.do_classifier_free_guidance: + negative_text_inputs = self.tokenizer( + batch.negative_prompt, + truncation=True, + # better way to handle this? + max_length=77, + return_tensors="pt", + ) + with set_forward_context(current_timestep=0, attn_metadata=None): + negative_outputs = self.text_encoder( + input_ids=negative_text_inputs["input_ids"].to( + batch.device), ) + negative_prompt_embeds = negative_outputs["pooler_output"] + + batch.negative_prompt_embeds.append(negative_prompt_embeds) + + if inference_args.use_cpu_offload: + self.text_encoder.to('cpu') + torch.cuda.empty_cache() + return batch diff --git a/fastvideo/v1/pipelines/stages/conditioning.py b/fastvideo/v1/pipelines/stages/conditioning.py index 66a8e0fbe..fb88d33de 100644 --- a/fastvideo/v1/pipelines/stages/conditioning.py +++ b/fastvideo/v1/pipelines/stages/conditioning.py @@ -39,8 +39,7 @@ class ConditioningStage(PipelineStage): if not batch.do_classifier_free_guidance: return batch else: - raise NotImplementedError( - "Classifier-free guidance is not supported yet") + return batch logger.info("batch.negative_prompt_embeds: %s", batch.negative_prompt_embeds) diff --git a/fastvideo/v1/pipelines/stages/decoding.py b/fastvideo/v1/pipelines/stages/decoding.py index fb52e6c37..432f67282 100644 --- a/fastvideo/v1/pipelines/stages/decoding.py +++ b/fastvideo/v1/pipelines/stages/decoding.py @@ -54,13 +54,20 @@ class DecodingStage(PipelineStage): vae_autocast_enabled = (vae_dtype != torch.float32 ) and not inference_args.disable_autocast - # Apply scaling/shifting if needed - if (hasattr(self.vae.config, "shift_factor") - and self.vae.config.shift_factor): - latents = (latents / self.vae.config.scaling_factor + - self.vae.config.shift_factor) + if isinstance(self.vae.scaling_factor, torch.Tensor): + latents = latents / self.vae.scaling_factor.to( + latents.device, latents.dtype) else: - latents = latents / self.vae.config.scaling_factor + latents = latents / self.vae.scaling_factor + + # Apply shifting if needed + if (hasattr(self.vae, "shift_factor") + and self.vae.shift_factor is not None): + if isinstance(self.vae.shift_factor, torch.Tensor): + latents += self.vae.shift_factor.to(latents.device, + latents.dtype) + else: + latents += self.vae.shift_factor # Decode latents with torch.autocast(device_type="cuda", @@ -70,6 +77,8 @@ class DecodingStage(PipelineStage): self.vae.enable_tiling() # if inference_args.vae_sp: # self.vae.enable_parallel() + if not vae_autocast_enabled: + latents = latents.to(vae_dtype) image = self.vae.decode(latents) # Normalize image to [0, 1] range diff --git a/fastvideo/v1/pipelines/stages/denoising.py b/fastvideo/v1/pipelines/stages/denoising.py index 990c062dc..23d4ede92 100644 --- a/fastvideo/v1/pipelines/stages/denoising.py +++ b/fastvideo/v1/pipelines/stages/denoising.py @@ -53,6 +53,9 @@ class DenoisingStage(PipelineStage): Returns: The batch with denoised latents. """ + # If use cpu offload, need to load the model back into gpu again + if inference_args.use_cpu_offload: + self.transformer = self.transformer.to(batch.device) # Prepare extra step kwargs for scheduler extra_step_kwargs = self.prepare_extra_func_kwargs( self.scheduler.step, @@ -77,6 +80,12 @@ class DenoisingStage(PipelineStage): n=world_size).contiguous() latents = latents[:, :, rank, :, :, :] batch.latents = latents + if batch.image_latent is not None: + image_latent = rearrange(batch.image_latent, + "b t (n s) h w -> b t n s h w", + n=world_size).contiguous() + image_latent = image_latent[:, :, rank, :, :, :] + batch.image_latent = image_latent # Get timesteps and calculate warmup steps timesteps = batch.timesteps @@ -98,9 +107,28 @@ class DenoisingStage(PipelineStage): result[t][layer][h] = value return result + # Prepare image latents and embeddings for I2V generation + image_embeds = batch.image_embeds + if len(image_embeds) > 0: + assert torch.isnan(image_embeds[0]).sum() == 0 + image_embeds = [ + image_embed.to(target_dtype) for image_embed in image_embeds + ] + + image_kwargs = self.prepare_extra_func_kwargs( + self.transformer.forward, + { + "encoder_hidden_states_image": image_embeds, + }, + ) + # Get latents and embeddings latents = batch.latents prompt_embeds = batch.prompt_embeds + assert torch.isnan(prompt_embeds[0]).sum() == 0 + if batch.do_classifier_free_guidance: + neg_prompt_embeds = batch.negative_prompt_embeds + assert torch.isnan(neg_prompt_embeds[0]).sum() == 0 # Run denoising loop with self.progress_bar(total=num_inference_steps) as progress_bar: @@ -109,10 +137,13 @@ class DenoisingStage(PipelineStage): if hasattr(self, 'interrupt') and self.interrupt: continue - # Expand latents for classifier-free guidance - latent_model_input = (torch.cat( - [latents] * - 2) if batch.do_classifier_free_guidance else latents) + # Expand latents for I2V + latent_model_input = latents.to(target_dtype) + if batch.image_latent is not None: + latent_model_input = torch.cat( + [latent_model_input, batch.image_latent], + dim=1).to(target_dtype) + assert torch.isnan(latent_model_input).sum() == 0 latent_model_input = self.scheduler.scale_model_input( latent_model_input, t) @@ -160,8 +191,6 @@ class DenoisingStage(PipelineStage): inference_args=inference_args, ) assert attn_metadata is not None, "attn_metadata cannot be None" - else: - attn_metadata = None else: attn_metadata = None @@ -180,29 +209,43 @@ class DenoisingStage(PipelineStage): prompt_embeds, t_expand, guidance=guidance_expand, + **image_kwargs, ) - # Apply guidance - if batch.do_classifier_free_guidance: - noise_pred_uncond, noise_pred_text = noise_pred.chunk(2) - noise_pred = noise_pred_uncond + batch.guidance_scale * ( - noise_pred_text - noise_pred_uncond) + # Apply guidance + if batch.do_classifier_free_guidance: + with set_forward_context( + current_timestep=i, + attn_metadata=attn_metadata, + # inference_args=inference_args + ): + # Run transformer + noise_pred_uncond = self.transformer( + latent_model_input, + neg_prompt_embeds, + t_expand, + guidance=guidance_expand, + **image_kwargs, + ) + noise_pred_text = noise_pred + noise_pred = noise_pred_uncond + batch.guidance_scale * ( + noise_pred_text - noise_pred_uncond) - # Apply guidance rescale if needed - if batch.guidance_rescale > 0.0: - # Based on 3.4. in https://arxiv.org/pdf/2305.08891.pdf - noise_pred = self.rescale_noise_cfg( - noise_pred, - noise_pred_text, - guidance_rescale=batch.guidance_rescale, - ) + # Apply guidance rescale if needed + if batch.guidance_rescale > 0.0: + # Based on 3.4. in https://arxiv.org/pdf/2305.08891.pdf + noise_pred = self.rescale_noise_cfg( + noise_pred, + noise_pred_text, + guidance_rescale=batch.guidance_rescale, + ) - # Compute the previous noisy sample - latents = self.scheduler.step(noise_pred, - t, - latents, - **extra_step_kwargs, - return_dict=False)[0] + # Compute the previous noisy sample + latents = self.scheduler.step(noise_pred, + t, + latents, + **extra_step_kwargs, + return_dict=False)[0] # Update progress bar if i == len(timesteps) - 1 or ( @@ -218,11 +261,15 @@ class DenoisingStage(PipelineStage): # Update batch with final latents batch.latents = latents + if inference_args.use_cpu_offload: + self.transformer.to('cpu') + torch.cuda.empty_cache() + return batch def prepare_extra_func_kwargs(self, func, kwargs) -> Dict[str, Any]: """ - Prepare extra kwargs for the scheduler step. + Prepare extra kwargs for the scheduler step / denoise step. Args: func: The function to prepare kwargs for. diff --git a/fastvideo/v1/pipelines/stages/encoding.py b/fastvideo/v1/pipelines/stages/encoding.py new file mode 100644 index 000000000..e2f2d8ee9 --- /dev/null +++ b/fastvideo/v1/pipelines/stages/encoding.py @@ -0,0 +1,168 @@ +# SPDX-License-Identifier: Apache-2.0 +""" +Encoding stage for diffusion pipelines. +""" +from typing import Optional + +import PIL.Image +import torch + +from fastvideo.v1.inference_args import InferenceArgs +from fastvideo.v1.logger import init_logger +from fastvideo.v1.models.vision_utils import (get_default_height_width, + load_image, normalize, + numpy_to_pt, pil_to_numpy, resize) +from fastvideo.v1.pipelines.pipeline_batch_info import ForwardBatch +from fastvideo.v1.pipelines.stages.base import PipelineStage +from fastvideo.v1.utils import PRECISION_TO_TYPE + +logger = init_logger(__name__) + + +class EncodingStage(PipelineStage): + """ + Stage for encoding pixel representations into latent space. + + This stage handles the encoding of pixel representations into the final + input format (e.g., latents). + """ + + def __init__(self, vae) -> None: + self.vae = vae + + def forward( + self, + batch: ForwardBatch, + inference_args: InferenceArgs, + ) -> ForwardBatch: + """ + Encode pixel representations into latent space. + + Args: + batch: The current batch information. + inference_args: The inference arguments. + + Returns: + The batch with encoded outputs. + """ + image_path = batch.image_path + # TODO(will): remove this once we add input/output validation for stages + if image_path is None: + raise ValueError("Image Path must be provided") + latent_height = batch.height // self.vae.spatial_compression_ratio + latent_width = batch.width // self.vae.spatial_compression_ratio + + image = load_image(image_path) + image = self.preprocess( + image, + vae_scale_factor=self.vae.spatial_compression_ratio, + height=batch.height, + width=batch.width).to(batch.device, dtype=torch.float32) + image = image.unsqueeze(2) + video_condition = torch.cat([ + image, + image.new_zeros(image.shape[0], image.shape[1], + inference_args.num_frames - 1, batch.height, + batch.width) + ], + dim=2) + video_condition = video_condition.to(device=batch.device, + dtype=torch.float32) + + # Setup VAE precision + vae_dtype = PRECISION_TO_TYPE[inference_args.vae_precision] + vae_autocast_enabled = ( + vae_dtype != torch.float32) and not inference_args.disable_autocast + + # Encode Image + with torch.autocast(device_type="cuda", + dtype=vae_dtype, + enabled=vae_autocast_enabled): + if inference_args.vae_tiling: + self.vae.enable_tiling() + # if inference_args.vae_sp: + # self.vae.enable_parallel() + if not vae_autocast_enabled: + video_condition = video_condition.to(vae_dtype) + encoder_output = self.vae.encode(video_condition) + + generator = batch.generator + latent_condition = self.retrieve_latents(encoder_output, generator[0]) + + # Apply shifting if needed + if (hasattr(self.vae, "shift_factor") + and self.vae.shift_factor is not None): + if isinstance(self.vae.shift_factor, torch.Tensor): + latent_condition -= self.vae.shift_factor.to( + latent_condition.device, latent_condition.dtype) + else: + latent_condition -= self.vae.shift_factor + + if isinstance(self.vae.scaling_factor, torch.Tensor): + latent_condition = latent_condition * self.vae.scaling_factor.to( + latent_condition.device, latent_condition.dtype) + else: + latent_condition = latent_condition * self.vae.scaling_factor + + mask_lat_size = torch.ones(1, 1, inference_args.num_frames, + latent_height, latent_width) + mask_lat_size[:, :, list(range(1, inference_args.num_frames))] = 0 + first_frame_mask = mask_lat_size[:, :, 0:1] + first_frame_mask = torch.repeat_interleave( + first_frame_mask, + dim=2, + repeats=self.vae.temporal_compression_ratio) + mask_lat_size = torch.concat( + [first_frame_mask, mask_lat_size[:, :, 1:, :]], dim=2) + mask_lat_size = mask_lat_size.view(1, -1, + self.vae.temporal_compression_ratio, + latent_height, latent_width) + mask_lat_size = mask_lat_size.transpose(1, 2) + mask_lat_size = mask_lat_size.to(latent_condition.device) + + batch.image_latent = torch.concat([mask_lat_size, latent_condition], + dim=1) + + # Offload models if needed + if hasattr(self, 'maybe_free_model_hooks'): + self.maybe_free_model_hooks() + + return batch + + def retrieve_latents(self, + encoder_output: torch.Tensor, + generator: Optional[torch.Generator] = None, + sample_mode: str = "sample"): + if sample_mode == "sample": + return encoder_output.sample(generator) + elif sample_mode == "argmax": + return encoder_output.mode() + else: + raise AttributeError( + "Could not access latents of provided encoder_output") + + def preprocess( + self, + image: PIL.Image.Image, + vae_scale_factor: int, + height: Optional[int] = None, + width: Optional[int] = None, + resize_mode: str = "default", # "default", "fill", "crop" + ) -> torch.Tensor: + image = [image] + + height, width = get_default_height_width(image[0], vae_scale_factor, + height, width) + image = [ + resize(i, height, width, resize_mode=resize_mode) for i in image + ] + image = pil_to_numpy(image) # to np + image = numpy_to_pt(image) # to pt + + do_normalize = True + if image.min() < 0: + do_normalize = False + if do_normalize: + image = normalize(image) + + return image diff --git a/fastvideo/v1/pipelines/stages/latent_preparation.py b/fastvideo/v1/pipelines/stages/latent_preparation.py index 73beeef02..86854649b 100644 --- a/fastvideo/v1/pipelines/stages/latent_preparation.py +++ b/fastvideo/v1/pipelines/stages/latent_preparation.py @@ -6,6 +6,7 @@ from diffusers.utils.torch_utils import randn_tensor from fastvideo.v1.inference_args import InferenceArgs from fastvideo.v1.logger import init_logger +from fastvideo.v1.models.vaes.common import ParallelTiledVAE from fastvideo.v1.pipelines.pipeline_batch_info import ForwardBatch from fastvideo.v1.pipelines.stages.base import PipelineStage @@ -20,9 +21,10 @@ class LatentPreparationStage(PipelineStage): denoised during the diffusion process. """ - def __init__(self, scheduler) -> None: + def __init__(self, scheduler, vae=None) -> None: super().__init__() self.scheduler = scheduler + self.vae = vae def forward( self, @@ -42,7 +44,7 @@ class LatentPreparationStage(PipelineStage): # Adjust video length based on VAE version if needed if hasattr(self, 'adjust_video_length'): - batch = self.adjust_video_length(batch, inference_args) + batch = self.adjust_video_length(self.vae, batch, inference_args) # Determine batch size if isinstance(batch.prompt, list): batch_size = len(batch.prompt) @@ -101,7 +103,7 @@ class LatentPreparationStage(PipelineStage): return batch - def adjust_video_length(self, batch: ForwardBatch, + def adjust_video_length(self, vae: ParallelTiledVAE, batch: ForwardBatch, inference_args: InferenceArgs) -> ForwardBatch: """ Adjust video length based on VAE version. @@ -114,6 +116,7 @@ class LatentPreparationStage(PipelineStage): The batch with adjusted video length. """ video_length = batch.num_frames + temporal_scale_factor = vae.temporal_compression_ratio if vae is not None else 4 # TODO - batch.num_frames = (video_length - 1) // 4 + 1 + batch.num_frames = (video_length - 1) // temporal_scale_factor + 1 return batch diff --git a/fastvideo/v1/pipelines/stages/llama_encoding.py b/fastvideo/v1/pipelines/stages/llama_encoding.py index ef88e8216..c72dab181 100644 --- a/fastvideo/v1/pipelines/stages/llama_encoding.py +++ b/fastvideo/v1/pipelines/stages/llama_encoding.py @@ -7,6 +7,8 @@ This module contains implementations of prompt encoding stages for diffusion pip from typing import TypedDict +import torch + from fastvideo.v1.forward_context import set_forward_context from fastvideo.v1.inference_args import InferenceArgs from fastvideo.v1.logger import init_logger @@ -71,6 +73,8 @@ class LlamaEncodingStage(PipelineStage): Returns: The batch with encoded prompt embeddings. """ + if inference_args.use_cpu_offload: + self.text_encoder = self.text_encoder.to(batch.device) text = prompt_template_video["template"].format(batch.prompt) text_inputs = self.tokenizer( @@ -93,4 +97,33 @@ class LlamaEncodingStage(PipelineStage): last_hidden_state = last_hidden_state[:, crop_start:] batch.prompt_embeds.append(last_hidden_state) + if batch.do_classifier_free_guidance: + negative_text = prompt_template_video["template"].format( + batch.negative_prompt) + negative_text_inputs = self.tokenizer( + negative_text, + truncation=True, + # better way to handle this? + max_length=256, + return_tensors="pt", + ) + hidden_state_skip_layer = 2 + with set_forward_context(current_timestep=0, attn_metadata=None): + negative_outputs = self.text_encoder( + input_ids=negative_text_inputs["input_ids"].to( + batch.device), + output_hidden_states=hidden_state_skip_layer is not None, + ) + + negative_last_hidden_state = negative_outputs.hidden_states[-( + hidden_state_skip_layer + 1)] + crop_start = prompt_template_video.get("crop_start", -1) + negative_last_hidden_state = negative_last_hidden_state[:, + crop_start:] + batch.negative_prompt_embeds.append(negative_last_hidden_state) + + if inference_args.use_cpu_offload: + self.text_encoder.to('cpu') + torch.cuda.empty_cache() + return batch diff --git a/fastvideo/v1/pipelines/stages/t5_encoding.py b/fastvideo/v1/pipelines/stages/t5_encoding.py new file mode 100644 index 000000000..51ebec271 --- /dev/null +++ b/fastvideo/v1/pipelines/stages/t5_encoding.py @@ -0,0 +1,115 @@ +# SPDX-License-Identifier: Apache-2.0 +""" +Prompt encoding stages for diffusion pipelines. + +This module contains implementations of prompt encoding stages for diffusion pipelines. +""" +import torch + +from fastvideo.v1.forward_context import set_forward_context +from fastvideo.v1.inference_args import InferenceArgs +from fastvideo.v1.logger import init_logger +from fastvideo.v1.pipelines.pipeline_batch_info import ForwardBatch +from fastvideo.v1.pipelines.stages.base import PipelineStage + +logger = init_logger(__name__) + + +class T5EncodingStage(PipelineStage): + """ + Stage for encoding text prompts into embeddings for diffusion models. + + This stage handles the encoding of text prompts into the embedding space + expected by the diffusion model. + """ + + def __init__(self, text_encoder, tokenizer) -> None: + """ + Initialize the prompt encoding stage. + + Args: + enable_logging: Whether to enable logging for this stage. + is_secondary: Whether this is a secondary text encoder. + """ + super().__init__() + self.text_encoder = text_encoder + self.tokenizer = tokenizer + + def forward( + self, + batch: ForwardBatch, + inference_args: InferenceArgs, + ) -> ForwardBatch: + """ + Encode the prompt into text encoder hidden states. + + Args: + batch: The current batch information. + inference_args: The inference arguments. + + Returns: + The batch with encoded prompt embeddings. + """ + if inference_args.use_cpu_offload: + self.text_encoder = self.text_encoder.to(batch.device) + + text = batch.prompt + text_inputs = self.tokenizer( + text, + padding="max_length", + truncation=True, + max_length=512, + add_special_tokens=True, + return_attention_mask=True, + return_tensors="pt", + ).to(batch.device) + text_input_ids, mask = text_inputs.input_ids, text_inputs.attention_mask + seq_lens = mask.gt(0).sum(dim=1).long() + with set_forward_context(current_timestep=0, attn_metadata=None): + outputs = self.text_encoder( + input_ids=text_input_ids, + attention_mask=mask, + ) + assert torch.isnan(outputs).sum() == 0 + prompt_embeds = [u[:v] for u, v in zip(outputs, seq_lens)] + prompt_embeds = torch.stack([ + torch.cat([u, u.new_zeros(512 - u.size(0), u.size(1))]) + for u in prompt_embeds + ], + dim=0) + batch.prompt_embeds.append(prompt_embeds) + + if batch.do_classifier_free_guidance: + negative_text = batch.negative_prompt + negative_text_inputs = self.tokenizer( + negative_text, + padding="max_length", + truncation=True, + max_length=512, + add_special_tokens=True, + return_attention_mask=True, + return_tensors="pt", + ).to(batch.device) + text_input_ids, mask = negative_text_inputs.input_ids, negative_text_inputs.attention_mask + seq_lens = mask.gt(0).sum(dim=1).long() + with set_forward_context(current_timestep=0, attn_metadata=None): + negative_outputs = self.text_encoder( + input_ids=text_input_ids, + attention_mask=mask, + ) + assert torch.isnan(negative_outputs).sum() == 0 + neg_prompt_embeds = [ + u[:v] for u, v in zip(negative_outputs, seq_lens) + ] + neg_prompt_embeds = torch.stack([ + torch.cat([u, u.new_zeros(512 - u.size(0), u.size(1))]) + for u in neg_prompt_embeds + ], + dim=0) + batch.negative_prompt_embeds.append(neg_prompt_embeds) + + if inference_args.use_cpu_offload: + self.text_encoder.to('cpu') + torch.cuda.empty_cache() + + return batch diff --git a/fastvideo/v1/pipelines/wan/__init__.py b/fastvideo/v1/pipelines/wan/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/fastvideo/v1/pipelines/wan/wan_i2v_pipeline.py b/fastvideo/v1/pipelines/wan/wan_i2v_pipeline.py new file mode 100644 index 000000000..5ccd9f95a --- /dev/null +++ b/fastvideo/v1/pipelines/wan/wan_i2v_pipeline.py @@ -0,0 +1,81 @@ +# SPDX-License-Identifier: Apache-2.0 +""" +Wan video diffusion pipeline implementation. + +This module contains an implementation of the Wan video diffusion pipeline +using the modular pipeline architecture. +""" + +from fastvideo.v1.inference_args import InferenceArgs +from fastvideo.v1.logger import init_logger +from fastvideo.v1.pipelines.composed_pipeline_base import ComposedPipelineBase +from fastvideo.v1.pipelines.stages import ( + CLIPImageEncodingStage, ConditioningStage, DecodingStage, DenoisingStage, + EncodingStage, InputValidationStage, LatentPreparationStage, + T5EncodingStage, TimestepPreparationStage) + +# TODO(will): move PRECISION_TO_TYPE to better place + +logger = init_logger(__name__) + + +class WanImageToVideoPipeline(ComposedPipelineBase): + + _required_config_modules = [ + "text_encoder", "tokenizer", "vae", "transformer", "scheduler", \ + "image_encoder", "image_processor" + ] + + def create_pipeline_stages(self, inference_args: InferenceArgs): + """Set up pipeline stages with proper dependency injection.""" + + self.add_stage(stage_name="input_validation_stage", + stage=InputValidationStage()) + + self.add_stage(stage_name="prompt_encoding_stage", + stage=T5EncodingStage( + text_encoder=self.get_module("text_encoder"), + tokenizer=self.get_module("tokenizer"), + )) + + self.add_stage(stage_name="image_encoding_stage", + stage=CLIPImageEncodingStage( + image_encoder=self.get_module("image_encoder"), + image_processor=self.get_module("image_processor"), + )) + + 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"), + vae=self.get_module("vae"))) + + self.add_stage(stage_name="image_latent_preparation_stage", + stage=EncodingStage(vae=self.get_module("vae"))) + + self.add_stage(stage_name="denoising_stage", + stage=DenoisingStage( + transformer=self.get_module("transformer"), + scheduler=self.get_module("scheduler"))) + + self.add_stage(stage_name="decoding_stage", + stage=DecodingStage(vae=self.get_module("vae"))) + + def initialize_pipeline(self, inference_args: InferenceArgs): + """ + Initialize the pipeline. + """ + vae_scale_factor = self.get_module("vae").spatial_compression_ratio + inference_args.vae_scale_factor = vae_scale_factor + + num_channels_latents = self.get_module("transformer").out_channels + inference_args.num_channels_latents = num_channels_latents + + +EntryClass = WanImageToVideoPipeline diff --git a/fastvideo/v1/pipelines/wan/wan_pipeline.py b/fastvideo/v1/pipelines/wan/wan_pipeline.py new file mode 100644 index 000000000..e862588e8 --- /dev/null +++ b/fastvideo/v1/pipelines/wan/wan_pipeline.py @@ -0,0 +1,72 @@ +# SPDX-License-Identifier: Apache-2.0 +""" +Wan video diffusion pipeline implementation. + +This module contains an implementation of the Wan video diffusion pipeline +using the modular pipeline architecture. +""" + +from fastvideo.v1.inference_args import InferenceArgs +from fastvideo.v1.logger import init_logger +from fastvideo.v1.pipelines.composed_pipeline_base import ComposedPipelineBase +from fastvideo.v1.pipelines.stages import (ConditioningStage, DecodingStage, + DenoisingStage, InputValidationStage, + LatentPreparationStage, + T5EncodingStage, + TimestepPreparationStage) + +# TODO(will): move PRECISION_TO_TYPE to better place + +logger = init_logger(__name__) + + +class WanPipeline(ComposedPipelineBase): + + _required_config_modules = [ + "text_encoder", "tokenizer", "vae", "transformer", "scheduler" + ] + + def create_pipeline_stages(self, inference_args: InferenceArgs): + """Set up pipeline stages with proper dependency injection.""" + + self.add_stage(stage_name="input_validation_stage", + stage=InputValidationStage()) + + self.add_stage(stage_name="prompt_encoding_stage", + stage=T5EncodingStage( + text_encoder=self.get_module("text_encoder"), + tokenizer=self.get_module("tokenizer"), + )) + + 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"), + vae=self.get_module("vae"))) + + self.add_stage(stage_name="denoising_stage", + stage=DenoisingStage( + transformer=self.get_module("transformer"), + scheduler=self.get_module("scheduler"))) + + self.add_stage(stage_name="decoding_stage", + stage=DecodingStage(vae=self.get_module("vae"))) + + def initialize_pipeline(self, inference_args: InferenceArgs): + """ + Initialize the pipeline. + """ + vae_scale_factor = self.get_module("vae").spatial_compression_ratio + inference_args.vae_scale_factor = vae_scale_factor + + num_channels_latents = self.get_module("transformer").in_channels + inference_args.num_channels_latents = num_channels_latents + + +EntryClass = WanPipeline diff --git a/fastvideo/v1/tests/conftest.py b/fastvideo/v1/tests/conftest.py index 075afa886..a106e375b 100644 --- a/fastvideo/v1/tests/conftest.py +++ b/fastvideo/v1/tests/conftest.py @@ -1,4 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 +import pytest +import torch.distributed as dist import pytest import torch diff --git a/fastvideo/v1/tests/encoders/test_clip_encoder.py b/fastvideo/v1/tests/encoders/test_clip_encoder.py index 3c1df3296..14fc23bb8 100644 --- a/fastvideo/v1/tests/encoders/test_clip_encoder.py +++ b/fastvideo/v1/tests/encoders/test_clip_encoder.py @@ -1,5 +1,4 @@ # SPDX-License-Identifier: Apache-2.0 - # TODO: check if correct import os diff --git a/fastvideo/v1/tests/encoders/test_llama_encoder.py b/fastvideo/v1/tests/encoders/test_llama_encoder.py index 84a0eff4e..0a74c1473 100644 --- a/fastvideo/v1/tests/encoders/test_llama_encoder.py +++ b/fastvideo/v1/tests/encoders/test_llama_encoder.py @@ -1,5 +1,4 @@ # SPDX-License-Identifier: Apache-2.0 - import os import numpy as np diff --git a/fastvideo/v1/tests/encoders/test_t5_encoder.py b/fastvideo/v1/tests/encoders/test_t5_encoder.py new file mode 100644 index 000000000..e91e05354 --- /dev/null +++ b/fastvideo/v1/tests/encoders/test_t5_encoder.py @@ -0,0 +1,136 @@ +# SPDX-License-Identifier: Apache-2.0 +import os + +import numpy as np +import pytest +import torch +from transformers import AutoConfig, AutoTokenizer, UMT5EncoderModel + +from fastvideo.v1.forward_context import set_forward_context +from fastvideo.v1.logger import init_logger +from fastvideo.v1.models.loader.component_loader import TextEncoderLoader +from fastvideo.v1.utils import maybe_download_model + +logger = init_logger(__name__) + +os.environ["MASTER_ADDR"] = "localhost" +os.environ["MASTER_PORT"] = "29503" + +BASE_MODEL_PATH = "Wan-AI/Wan2.1-T2V-1.3B-Diffusers" +MODEL_PATH = maybe_download_model(BASE_MODEL_PATH, + local_dir=os.path.join( + 'data', BASE_MODEL_PATH)) +TEXT_ENCODER_PATH = os.path.join(MODEL_PATH, "text_encoder") +TOKENIZER_PATH = os.path.join(MODEL_PATH, "tokenizer") + + +@pytest.mark.usefixtures("distributed_setup") +def test_t5_encoder(): + # Initialize the two model implementations + hf_config = AutoConfig.from_pretrained(TEXT_ENCODER_PATH) + print(hf_config) + + device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") + precision = torch.float16 + model1 = UMT5EncoderModel.from_pretrained(TEXT_ENCODER_PATH).to( + precision).to(device).eval() + tokenizer = AutoTokenizer.from_pretrained(TOKENIZER_PATH) + + loader = TextEncoderLoader() + model2 = loader.load_model(TEXT_ENCODER_PATH, hf_config, device) + + # Convert to float16 and move to device + model2 = model2.to(precision) + model2 = model2.to(device) + model2.eval() + + # Sanity check weights between the two models + logger.info("Comparing model weights for sanity check...") + params1 = dict(model1.named_parameters()) + params2 = dict(model2.named_parameters()) + + # Check number of parameters + logger.info("Model1 has %s parameters", len(params1)) + logger.info("Model2 has %s parameters", len(params2)) + + weight_diffs = [] + # check if embed_tokens are the same + weights = ["encoder.block.{}.layer.0.layer_norm.weight", "encoder.block.{}.layer.0.SelfAttention.relative_attention_bias.weight", \ + "encoder.block.{}.layer.0.SelfAttention.o.weight", "encoder.block.{}.layer.1.DenseReluDense.wi_0.weight", "encoder.block.{}.layer.1.DenseReluDense.wi_1.weight",\ + "encoder.block.{}.layer.1.DenseReluDense.wo.weight", \ + "encoder.block.{}.layer.1.layer_norm.weight", "encoder.final_layer_norm.weight", "shared.weight"] + for idx in range(hf_config.num_hidden_layers): + for w in weights: + name1 = w.format(idx) + name2 = w.format(idx) + p1 = params1[name1] + p2 = params2[name2] + assert p1.dtype == p2.dtype + try: + logger.info("Parameter: %s vs %s", name1, name2) + max_diff = torch.max(torch.abs(p1 - p2)).item() + mean_diff = torch.mean(torch.abs(p1 - p2)).item() + weight_diffs.append((name1, name2, max_diff, mean_diff)) + logger.info(" Max diff: %s, Mean diff: %s", max_diff, + mean_diff) + except Exception as e: + logger.info("Error comparing %s and %s: %s", name1, name2, e) + + # Test with some sample prompts + prompts = [ + "Once upon a time", "The quick brown fox jumps over", + "In a galaxy far, far away" + ] + + logger.info("Testing T5 encoder with sample prompts") + + with torch.no_grad(): + for prompt in prompts: + logger.info("Testing prompt: %s", prompt) + + # Tokenize the prompt + tokens = tokenizer(prompt, + padding="max_length", + max_length=512, + truncation=True, + return_tensors="pt").to(device) + + # Get outputs from HuggingFace implementation + # filter out padding input_ids + # tokens.input_ids = tokens.input_ids[tokens.attention_mask==1] + # tokens.attention_mask = tokens.attention_mask[tokens.attention_mask==1] + outputs1 = model1(input_ids=tokens.input_ids, + attention_mask=tokens.attention_mask, + output_hidden_states=True).last_hidden_state + print("--------------------------------") + logger.info("Testing model2") + + # Get outputs from our implementation + with set_forward_context(current_timestep=0, attn_metadata=None): + outputs2 = model2( + input_ids=tokens.input_ids, + attention_mask=tokens.attention_mask, + ) + + # Compare last hidden states + last_hidden_state1 = outputs1[tokens.attention_mask == 1] + last_hidden_state2 = outputs2[tokens.attention_mask == 1] + + assert last_hidden_state1.shape == last_hidden_state2.shape, \ + f"Hidden state shapes don't match: {last_hidden_state1.shape} vs {last_hidden_state2.shape}" + + max_diff_hidden = torch.max( + torch.abs(last_hidden_state1 - last_hidden_state2)) + mean_diff_hidden = torch.mean( + torch.abs(last_hidden_state1 - last_hidden_state2)) + + logger.info("Maximum difference in last hidden states: %s", + max_diff_hidden.item()) + logger.info("Mean difference in last hidden states: %s", + mean_diff_hidden.item()) + + # Check if outputs are similar (allowing for small numerical differences) + assert mean_diff_hidden < 1e-4, \ + f"Hidden states differ significantly: mean diff = {mean_diff_hidden.item()}" + assert max_diff_hidden < 1e-4, \ + f"Hidden states differ significantly: max diff = {max_diff_hidden.item()}" \ No newline at end of file diff --git a/fastvideo/v1/tests/ssim/README.md b/fastvideo/v1/tests/ssim/README.md index f62ebb836..e32b2908c 100644 --- a/fastvideo/v1/tests/ssim/README.md +++ b/fastvideo/v1/tests/ssim/README.md @@ -1,7 +1,9 @@ The reference videos in the `reference_videos` directory are used as part of an e2e test to ensure consistency in video generation quality across code changes. `test_inference_similarity.py` compares newly generated videos against these references using Structural Similarity Index (SSIM) metrics to detect any regressions in visual quality across code changes. -`reference_videos/FLASH_ATTN/` videos were generated on commit `66107fd5b8469fed25972feb632cd48887dac451`. -`reference_videos/TORCH_SDPA/` videos were generated on commit `4ea008b8a16d7f5678a44b187ebdd7d9d0416ff1`. +`reference_videos/FastHunyuan-diffusers/FLASH_ATTN/` videos were generated on commit `66107fd5b8469fed25972feb632cd48887dac451`. +`reference_videos/FastHunyuan-diffusers/TORCH_SDPA/` videos were generated on commit `4ea008b8a16d7f5678a44b187ebdd7d9d0416ff1`. +`reference_videos/Wan2.1-T2V-1.3B-Diffusers` videos were generated on commit `d085770a70988c7b26632a0c3123c24a57f7ca77`. +`reference_videos/Wan2.1-I2V-14B-480P-Diffusers` videos were generated on commit `d085770a70988c7b26632a0c3123c24a57f7ca77`. ## Generation Details @@ -9,7 +11,7 @@ The reference videos in the `reference_videos` directory are used as part of an ## Generation Parameters -{ +FastHunyuan-diffusers: { "num_gpus": 2, "model_path": "data/FastHunyuan-diffusers", "height": 720, @@ -26,8 +28,51 @@ The reference videos in the `reference_videos` directory are used as part of an "fps": 24 } -### Prompts +Wan2.1-T2V-1.3B-Diffusers: { +"num_gpus": 2, +"model_path": "Wan-AI/Wan2.1-T2V-1.3B-Diffusers", +"height": 480, +"width": 832, +"num_frames": 45, +"num_inference_steps": 20, +"guidance_scale": 3, +"embedded_cfg_scale": 6, +"flow_shift": 7.0, +"seed": 1024, +"sp_size": 2, +"tp_size": 2, +"vae_sp": True, +"fps": 24, +"neg_prompt": "Bright tones, overexposed, static, blurred details, subtitles, style, works, paintings, images, static, overall gray, worst quality, low quality, JPEG compression residue, ugly, incomplete, extra fingers, poorly drawn hands, poorly drawn faces, deformed, disfigured, misshapen limbs, fused fingers, still picture, messy background, three legs, many people in the background, walking backwards", +"text-encoder-precision": "fp32" +} -1. Will Smith casually eats noodles, his relaxed demeanor contrasting with the energetic background of a bustling street food market. The scene captures a mix of humor and authenticity. Mid-shot framing, vibrant lighting." +Wan2.1-I2V-14B-480P-Diffusers: { +"num_gpus": 2, +"model_path": "Wan-AI/Wan2.1-I2V-14B-480P-Diffusers", +"height": 480, +"width": 832, +"num_frames": 45, +"num_inference_steps": 6, +"guidance_scale": 5.0, +"embedded_cfg_scale": 6, +"flow_shift": 7.0, +"seed": 1024, +"sp_size": 2, +"tp_size": 2, +"vae_sp": True, +"fps": 24, +"neg_prompt": "Bright tones, overexposed, static, blurred details, subtitles, style, works, paintings, images, static, overall gray, worst quality, low quality, JPEG compression residue, ugly, incomplete, extra fingers, poorly drawn hands, poorly drawn faces, deformed, disfigured, misshapen limbs, fused fingers, still picture, messy background, three legs, many people in the background, walking backwards", +"text-encoder-precision": "fp32" +} + +### Text-to-Video Prompts + +1. "Will Smith casually eats noodles, his relaxed demeanor contrasting with the energetic background of a bustling street food market. The scene captures a mix of humor and authenticity. Mid-shot framing, vibrant lighting." 2. "A lone hiker stands atop a towering cliff, silhouetted against the vast horizon. The rugged landscape stretches endlessly beneath, its earthy tones blending into the soft blues of the sky. The scene captures the spirit of exploration and human resilience. High angle, dynamic framing, with soft natural lighting emphasizing the grandeur of nature." + +### Image-to-Video Prompts + +1. "An astronaut hatching from an egg, on the surface of the moon, the darkness and depth of space realised in the background. High quality, ultrarealistic detail and breath-taking movie-like camera shot." + Image path: "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/astronaut.jpg" \ No newline at end of file diff --git a/fastvideo/v1/tests/ssim/compute_ssim.py b/fastvideo/v1/tests/ssim/compute_ssim.py index 943adf7f2..4910363e6 100644 --- a/fastvideo/v1/tests/ssim/compute_ssim.py +++ b/fastvideo/v1/tests/ssim/compute_ssim.py @@ -1,5 +1,4 @@ # SPDX-License-Identifier: Apache-2.0 - import argparse import os diff --git a/fastvideo/v1/tests/ssim/reference_videos/FLASH_ATTN/A lone hiker stands atop a towering cliff, silhouetted against the vast horizon. The rugged landscap.mp4 b/fastvideo/v1/tests/ssim/reference_videos/FastHunyuan-diffusers/FLASH_ATTN/A lone hiker stands atop a towering cliff, silhouetted against the vast horizon. The rugged landscap.mp4 similarity index 100% rename from fastvideo/v1/tests/ssim/reference_videos/FLASH_ATTN/A lone hiker stands atop a towering cliff, silhouetted against the vast horizon. The rugged landscap.mp4 rename to fastvideo/v1/tests/ssim/reference_videos/FastHunyuan-diffusers/FLASH_ATTN/A lone hiker stands atop a towering cliff, silhouetted against the vast horizon. The rugged landscap.mp4 diff --git a/fastvideo/v1/tests/ssim/reference_videos/FLASH_ATTN/Will Smith casually eats noodles, his relaxed demeanor contrasting with the energetic background of .mp4 b/fastvideo/v1/tests/ssim/reference_videos/FastHunyuan-diffusers/FLASH_ATTN/Will Smith casually eats noodles, his relaxed demeanor contrasting with the energetic background of .mp4 similarity index 100% rename from fastvideo/v1/tests/ssim/reference_videos/FLASH_ATTN/Will Smith casually eats noodles, his relaxed demeanor contrasting with the energetic background of .mp4 rename to fastvideo/v1/tests/ssim/reference_videos/FastHunyuan-diffusers/FLASH_ATTN/Will Smith casually eats noodles, his relaxed demeanor contrasting with the energetic background of .mp4 diff --git a/fastvideo/v1/tests/ssim/reference_videos/TORCH_SDPA/A lone hiker stands atop a towering cliff, silhouetted against the vast horizon. The rugged landscap.mp4 b/fastvideo/v1/tests/ssim/reference_videos/FastHunyuan-diffusers/TORCH_SDPA/A lone hiker stands atop a towering cliff, silhouetted against the vast horizon. The rugged landscap.mp4 similarity index 100% rename from fastvideo/v1/tests/ssim/reference_videos/TORCH_SDPA/A lone hiker stands atop a towering cliff, silhouetted against the vast horizon. The rugged landscap.mp4 rename to fastvideo/v1/tests/ssim/reference_videos/FastHunyuan-diffusers/TORCH_SDPA/A lone hiker stands atop a towering cliff, silhouetted against the vast horizon. The rugged landscap.mp4 diff --git a/fastvideo/v1/tests/ssim/reference_videos/TORCH_SDPA/Will Smith casually eats noodles, his relaxed demeanor contrasting with the energetic background of .mp4 b/fastvideo/v1/tests/ssim/reference_videos/FastHunyuan-diffusers/TORCH_SDPA/Will Smith casually eats noodles, his relaxed demeanor contrasting with the energetic background of .mp4 similarity index 100% rename from fastvideo/v1/tests/ssim/reference_videos/TORCH_SDPA/Will Smith casually eats noodles, his relaxed demeanor contrasting with the energetic background of .mp4 rename to fastvideo/v1/tests/ssim/reference_videos/FastHunyuan-diffusers/TORCH_SDPA/Will Smith casually eats noodles, his relaxed demeanor contrasting with the energetic background of .mp4 diff --git a/fastvideo/v1/tests/ssim/reference_videos/Wan2.1-I2V-14B-480P-Diffusers/FLASH_ATTN/An astronaut hatching from an egg, on the surface of the moon, the darkness and depth of space reali.mp4 b/fastvideo/v1/tests/ssim/reference_videos/Wan2.1-I2V-14B-480P-Diffusers/FLASH_ATTN/An astronaut hatching from an egg, on the surface of the moon, the darkness and depth of space reali.mp4 new file mode 100644 index 000000000..9cdddc1fe Binary files /dev/null and b/fastvideo/v1/tests/ssim/reference_videos/Wan2.1-I2V-14B-480P-Diffusers/FLASH_ATTN/An astronaut hatching from an egg, on the surface of the moon, the darkness and depth of space reali.mp4 differ diff --git a/fastvideo/v1/tests/ssim/reference_videos/Wan2.1-I2V-14B-480P-Diffusers/TORCH_SDPA/An astronaut hatching from an egg, on the surface of the moon, the darkness and depth of space reali.mp4 b/fastvideo/v1/tests/ssim/reference_videos/Wan2.1-I2V-14B-480P-Diffusers/TORCH_SDPA/An astronaut hatching from an egg, on the surface of the moon, the darkness and depth of space reali.mp4 new file mode 100644 index 000000000..942a2ba21 Binary files /dev/null and b/fastvideo/v1/tests/ssim/reference_videos/Wan2.1-I2V-14B-480P-Diffusers/TORCH_SDPA/An astronaut hatching from an egg, on the surface of the moon, the darkness and depth of space reali.mp4 differ diff --git a/fastvideo/v1/tests/ssim/reference_videos/Wan2.1-T2V-1.3B-Diffusers/FLASH_ATTN/A lone hiker stands atop a towering cliff, silhouetted against the vast horizon. The rugged landscap.mp4 b/fastvideo/v1/tests/ssim/reference_videos/Wan2.1-T2V-1.3B-Diffusers/FLASH_ATTN/A lone hiker stands atop a towering cliff, silhouetted against the vast horizon. The rugged landscap.mp4 new file mode 100644 index 000000000..7136d9ba9 Binary files /dev/null and b/fastvideo/v1/tests/ssim/reference_videos/Wan2.1-T2V-1.3B-Diffusers/FLASH_ATTN/A lone hiker stands atop a towering cliff, silhouetted against the vast horizon. The rugged landscap.mp4 differ diff --git a/fastvideo/v1/tests/ssim/reference_videos/Wan2.1-T2V-1.3B-Diffusers/FLASH_ATTN/Will Smith casually eats noodles, his relaxed demeanor contrasting with the energetic background of .mp4 b/fastvideo/v1/tests/ssim/reference_videos/Wan2.1-T2V-1.3B-Diffusers/FLASH_ATTN/Will Smith casually eats noodles, his relaxed demeanor contrasting with the energetic background of .mp4 new file mode 100644 index 000000000..cf05d79a6 Binary files /dev/null and b/fastvideo/v1/tests/ssim/reference_videos/Wan2.1-T2V-1.3B-Diffusers/FLASH_ATTN/Will Smith casually eats noodles, his relaxed demeanor contrasting with the energetic background of .mp4 differ diff --git a/fastvideo/v1/tests/ssim/reference_videos/Wan2.1-T2V-1.3B-Diffusers/TORCH_SDPA/A lone hiker stands atop a towering cliff, silhouetted against the vast horizon. The rugged landscap.mp4 b/fastvideo/v1/tests/ssim/reference_videos/Wan2.1-T2V-1.3B-Diffusers/TORCH_SDPA/A lone hiker stands atop a towering cliff, silhouetted against the vast horizon. The rugged landscap.mp4 new file mode 100644 index 000000000..27e7aba4b Binary files /dev/null and b/fastvideo/v1/tests/ssim/reference_videos/Wan2.1-T2V-1.3B-Diffusers/TORCH_SDPA/A lone hiker stands atop a towering cliff, silhouetted against the vast horizon. The rugged landscap.mp4 differ diff --git a/fastvideo/v1/tests/ssim/reference_videos/Wan2.1-T2V-1.3B-Diffusers/TORCH_SDPA/Will Smith casually eats noodles, his relaxed demeanor contrasting with the energetic background of .mp4 b/fastvideo/v1/tests/ssim/reference_videos/Wan2.1-T2V-1.3B-Diffusers/TORCH_SDPA/Will Smith casually eats noodles, his relaxed demeanor contrasting with the energetic background of .mp4 new file mode 100644 index 000000000..bbaf497b1 Binary files /dev/null and b/fastvideo/v1/tests/ssim/reference_videos/Wan2.1-T2V-1.3B-Diffusers/TORCH_SDPA/Will Smith casually eats noodles, his relaxed demeanor contrasting with the energetic background of .mp4 differ diff --git a/fastvideo/v1/tests/ssim/test_inference_similarity.py b/fastvideo/v1/tests/ssim/test_inference_similarity.py index 495c31205..9f9bb6751 100644 --- a/fastvideo/v1/tests/ssim/test_inference_similarity.py +++ b/fastvideo/v1/tests/ssim/test_inference_similarity.py @@ -1,5 +1,4 @@ # SPDX-License-Identifier: Apache-2.0 - import json import os @@ -12,7 +11,7 @@ from fastvideo.v1.tests.ssim.compute_ssim import compute_video_ssim_torchvision logger = init_logger(__name__) # Base parameters from the shell script -BASE_PARAMS = { +HUNYUAN_PARAMS = { "num_gpus": 2, "model_path": "FastVideo/FastHunyuan-diffusers", "height": 720, @@ -29,11 +28,66 @@ BASE_PARAMS = { "fps": 24, } +WAN_T2V_PARAMS = { + "num_gpus": 2, + "model_path": "Wan-AI/Wan2.1-T2V-1.3B-Diffusers", + "height": 480, + "width": 832, + "num_frames": 45, + "num_inference_steps": 20, + "guidance_scale": 3, + "embedded_cfg_scale": 6, + "flow_shift": 7.0, + "seed": 1024, + "sp_size": 2, + "tp_size": 2, + "vae_sp": True, + "fps": 24, + "neg_prompt": "Bright tones, overexposed, static, blurred details, subtitles, style, works, paintings, images, static, overall gray, worst quality, low quality, JPEG compression residue, ugly, incomplete, extra fingers, poorly drawn hands, poorly drawn faces, deformed, disfigured, misshapen limbs, fused fingers, still picture, messy background, three legs, many people in the background, walking backwards", + "text-encoder-precision": "fp32", +} + +WAN_I2V_PARAMS = { + "num_gpus": 2, + "model_path": "Wan-AI/Wan2.1-I2V-14B-480P-Diffusers", + "height": 480, + "width": 832, + "num_frames": 45, + "num_inference_steps": 6, + "guidance_scale": 5.0, + "embedded_cfg_scale": 6, + "flow_shift": 7.0, + "seed": 1024, + "sp_size": 2, + "tp_size": 2, + "vae_sp": True, + "fps": 24, + "neg_prompt": "Bright tones, overexposed, static, blurred details, subtitles, style, works, paintings, images, static, overall gray, worst quality, low quality, JPEG compression residue, ugly, incomplete, extra fingers, poorly drawn hands, poorly drawn faces, deformed, disfigured, misshapen limbs, fused fingers, still picture, messy background, three legs, many people in the background, walking backwards", + "text-encoder-precision": "fp32", +} + +MODEL_TO_PARAMS = { + "FastHunyuan-diffusers": HUNYUAN_PARAMS, + "Wan2.1-T2V-1.3B-Diffusers": WAN_T2V_PARAMS, +} + +I2V_MODEL_TO_PARAMS = { + "Wan2.1-I2V-14B-480P-Diffusers": WAN_I2V_PARAMS, +} + TEST_PROMPTS = [ "Will Smith casually eats noodles, his relaxed demeanor contrasting with the energetic background of a bustling street food market. The scene captures a mix of humor and authenticity. Mid-shot framing, vibrant lighting.", "A lone hiker stands atop a towering cliff, silhouetted against the vast horizon. The rugged landscape stretches endlessly beneath, its earthy tones blending into the soft blues of the sky. The scene captures the spirit of exploration and human resilience. High angle, dynamic framing, with soft natural lighting emphasizing the grandeur of nature." ] +I2V_TEST_PROMPTS = [ + "An astronaut hatching from an egg, on the surface of the moon, the darkness and depth of space realised in the background. High quality, ultrarealistic detail and breath-taking movie-like camera shot.", +] + +I2V_IMAGE_PATHS = [ + "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/astronaut.jpg", +] + def write_ssim_results(output_dir, ssim_values, reference_path, generated_path, num_inference_steps, prompt): @@ -74,24 +128,28 @@ def write_ssim_results(output_dir, ssim_values, reference_path, generated_path, return False -@pytest.mark.parametrize("num_inference_steps", [6]) -@pytest.mark.parametrize("prompt", TEST_PROMPTS) +@pytest.mark.parametrize("prompt", I2V_TEST_PROMPTS) @pytest.mark.parametrize("ATTENTION_BACKEND", ["FLASH_ATTN", "TORCH_SDPA"]) -def test_inference_similarity(num_inference_steps, prompt, ATTENTION_BACKEND): +@pytest.mark.parametrize("model_id", list(I2V_MODEL_TO_PARAMS.keys())) +def test_i2v_inference_similarity(prompt, ATTENTION_BACKEND, model_id): """ Test that runs inference with different parameters and compares the output to reference videos using SSIM. """ + assert len(I2V_TEST_PROMPTS) == len(I2V_IMAGE_PATHS), "Expect number of prompts equal to number of images" os.environ["FASTVIDEO_ATTENTION_BACKEND"] = ATTENTION_BACKEND script_dir = os.path.dirname(os.path.abspath(__file__)) - base_output_dir = os.path.join(script_dir, 'generated_videos') + base_output_dir = os.path.join(script_dir, 'generated_videos', model_id) output_dir = os.path.join(base_output_dir, ATTENTION_BACKEND) output_video_name = f"{prompt[:100]}.mp4" os.makedirs(output_dir, exist_ok=True) + BASE_PARAMS = I2V_MODEL_TO_PARAMS[model_id] + num_inference_steps = BASE_PARAMS["num_inference_steps"] + image_path = I2V_IMAGE_PATHS[I2V_TEST_PROMPTS.index(prompt)] launch_args = [ "--num-inference-steps", str(num_inference_steps), @@ -99,6 +157,8 @@ def test_inference_similarity(num_inference_steps, prompt, ATTENTION_BACKEND): prompt, "--output-path", output_dir, + "--image_path", + image_path, "--model-path", BASE_PARAMS["model_path"], "--height", @@ -125,13 +185,127 @@ def test_inference_similarity(num_inference_steps, prompt, ATTENTION_BACKEND): if BASE_PARAMS["vae_sp"]: launch_args.append("--vae-sp") + if "neg_prompt" in BASE_PARAMS.keys(): + launch_args.append("--neg_prompt") + launch_args.append(BASE_PARAMS["neg_prompt"]) + if "text-encoder-precision" in BASE_PARAMS.keys(): + launch_args.append("--text-encoder-precision") + launch_args.append(BASE_PARAMS["text-encoder-precision"]) launch_distributed(num_gpus=BASE_PARAMS["num_gpus"], args=launch_args) assert os.path.exists( output_dir), f"Output video was not generated at {output_dir}" - reference_folder = os.path.join(script_dir, 'reference_videos', ATTENTION_BACKEND) + reference_folder = os.path.join(script_dir, 'reference_videos', model_id, ATTENTION_BACKEND) + + if not os.path.exists(reference_folder): + logger.error("Reference folder missing") + raise FileNotFoundError( + f"Reference video folder does not exist: {reference_folder}") + + # Find the matching reference video based on the prompt + reference_video_name = None + + for filename in os.listdir(reference_folder): + if filename.endswith('.mp4') and prompt[:100] in filename: + reference_video_name = filename + break + + if not reference_video_name: + logger.error(f"Reference video not found for prompt: {prompt} with backend: {ATTENTION_BACKEND}") + raise FileNotFoundError(f"Reference video missing") + + reference_video_path = os.path.join(reference_folder, reference_video_name) + generated_video_path = os.path.join(output_dir, output_video_name) + + logger.info( + f"Computing SSIM between {reference_video_path} and {generated_video_path}" + ) + ssim_values = compute_video_ssim_torchvision(reference_video_path, + generated_video_path, + use_ms_ssim=True) + + mean_ssim = ssim_values[0] + logger.info(f"SSIM mean value: {mean_ssim}") + logger.info(f"Writing SSIM results to directory: {output_dir}") + + success = write_ssim_results(output_dir, ssim_values, reference_video_path, + generated_video_path, num_inference_steps, + prompt) + + if not success: + logger.error("Failed to write SSIM results to file") + + min_acceptable_ssim = 1 + assert mean_ssim >= min_acceptable_ssim, f"SSIM value {mean_ssim} is below threshold {min_acceptable_ssim}" + +@pytest.mark.parametrize("prompt", TEST_PROMPTS) +@pytest.mark.parametrize("ATTENTION_BACKEND", ["FLASH_ATTN", "TORCH_SDPA"]) +@pytest.mark.parametrize("model_id", list(MODEL_TO_PARAMS.keys())) +def test_inference_similarity(prompt, ATTENTION_BACKEND, model_id): + """ + Test that runs inference with different parameters and compares the output + to reference videos using SSIM. + """ + os.environ["FASTVIDEO_ATTENTION_BACKEND"] = ATTENTION_BACKEND + + script_dir = os.path.dirname(os.path.abspath(__file__)) + + base_output_dir = os.path.join(script_dir, 'generated_videos', model_id) + output_dir = os.path.join(base_output_dir, ATTENTION_BACKEND) + output_video_name = f"{prompt[:100]}.mp4" + + os.makedirs(output_dir, exist_ok=True) + + BASE_PARAMS = MODEL_TO_PARAMS[model_id] + num_inference_steps = BASE_PARAMS["num_inference_steps"] + launch_args = [ + "--num-inference-steps", + str(num_inference_steps), + "--prompt", + prompt, + "--output-path", + output_dir, + "--model-path", + BASE_PARAMS["model_path"], + "--height", + str(BASE_PARAMS["height"]), + "--width", + str(BASE_PARAMS["width"]), + "--num-frames", + str(BASE_PARAMS["num_frames"]), + "--guidance-scale", + str(BASE_PARAMS["guidance_scale"]), + "--embedded-cfg-scale", + str(BASE_PARAMS["embedded_cfg_scale"]), + "--flow-shift", + str(BASE_PARAMS["flow_shift"]), + "--seed", + str(BASE_PARAMS["seed"]), + "--sp-size", + str(BASE_PARAMS["sp_size"]), + "--tp-size", + str(BASE_PARAMS["tp_size"]), + "--fps", + str(BASE_PARAMS["fps"]), + ] + + if BASE_PARAMS["vae_sp"]: + launch_args.append("--vae-sp") + if "neg_prompt" in BASE_PARAMS.keys(): + launch_args.append("--neg_prompt") + launch_args.append(BASE_PARAMS["neg_prompt"]) + if "text-encoder-precision" in BASE_PARAMS.keys(): + launch_args.append("--text-encoder-precision") + launch_args.append(BASE_PARAMS["text-encoder-precision"]) + + launch_distributed(num_gpus=BASE_PARAMS["num_gpus"], args=launch_args) + + assert os.path.exists( + output_dir), f"Output video was not generated at {output_dir}" + + reference_folder = os.path.join(script_dir, 'reference_videos', model_id, ATTENTION_BACKEND) if not os.path.exists(reference_folder): logger.error("Reference folder missing") diff --git a/fastvideo/v1/tests/transformers/test_wanvideo.py b/fastvideo/v1/tests/transformers/test_wanvideo.py new file mode 100644 index 000000000..9f4342b63 --- /dev/null +++ b/fastvideo/v1/tests/transformers/test_wanvideo.py @@ -0,0 +1,118 @@ +# SPDX-License-Identifier: Apache-2.0 +import os + +import numpy as np +import pytest +import torch +from diffusers import WanTransformer3DModel + +from fastvideo.v1.forward_context import set_forward_context +from fastvideo.v1.inference_args import InferenceArgs +from fastvideo.v1.logger import init_logger +from fastvideo.v1.models.loader.component_loader import TransformerLoader +from fastvideo.v1.utils import maybe_download_model + +logger = init_logger(__name__) + +os.environ["MASTER_ADDR"] = "localhost" +os.environ["MASTER_PORT"] = "29503" + +BASE_MODEL_PATH = "Wan-AI/Wan2.1-T2V-1.3B-Diffusers" +MODEL_PATH = maybe_download_model(BASE_MODEL_PATH, + local_dir=os.path.join( + 'data', BASE_MODEL_PATH)) +TRANSFORMER_PATH = os.path.join(MODEL_PATH, "transformer") + + +@pytest.mark.usefixtures("distributed_setup") +def test_wan_transformer(): + device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") + precision = torch.bfloat16 + precision_str = "bf16" + args = InferenceArgs(model_path=TRANSFORMER_PATH, + use_cpu_offload=False, + precision=precision_str) + args.device = device + + loader = TransformerLoader() + model2 = loader.load(TRANSFORMER_PATH, "", args).to(device, dtype=precision) + + model1 = WanTransformer3DModel.from_pretrained( + TRANSFORMER_PATH, device=device, + torch_dtype=precision).to(device, dtype=precision).requires_grad_(False) + + total_params = sum(p.numel() for p in model1.parameters()) + # Calculate weight sum for model1 (converting to float64 to avoid overflow) + weight_sum_model1 = sum( + p.to(torch.float64).sum().item() for p in model1.parameters()) + # Also calculate mean for more stable comparison + weight_mean_model1 = weight_sum_model1 / total_params + logger.info("Model 1 weight sum: %s", weight_sum_model1) + logger.info("Model 1 weight mean: %s", weight_mean_model1) + + # Calculate weight sum for model2 (converting to float64 to avoid overflow) + total_params_model2 = sum(p.numel() for p in model2.parameters()) + weight_sum_model2 = sum( + p.to(torch.float64).sum().item() for p in model2.parameters()) + # Also calculate mean for more stable comparison + weight_mean_model2 = weight_sum_model2 / total_params_model2 + logger.info("Model 2 weight sum: %s", weight_sum_model2) + logger.info("Model 2 weight mean: %s", weight_mean_model2) + + weight_sum_diff = abs(weight_sum_model1 - weight_sum_model2) + logger.info("Weight sum difference: %s", weight_sum_diff) + weight_mean_diff = abs(weight_mean_model1 - weight_mean_model2) + logger.info("Weight mean difference: %s", weight_mean_diff) + + # Set both models to eval mode + model1 = model1.eval() + model2 = model2.eval() + + # Create identical inputs for both models + batch_size = 1 + seq_len = 30 + + # Video latents [B, C, T, H, W] + hidden_states = torch.randn(batch_size, + 16, + 21, + 160, + 90, + device=device, + dtype=precision) + + # Text embeddings [B, L, D] (including global token) + encoder_hidden_states = torch.randn(batch_size, + seq_len + 1, + 4096, + device=device, + dtype=precision) + + # Timestep + timestep = torch.tensor([500], device=device, dtype=precision) + + with torch.amp.autocast('cuda', dtype=precision): + output1 = model1( + hidden_states=hidden_states, + encoder_hidden_states=encoder_hidden_states, + timestep=timestep, + return_dict=False, + )[0] + with set_forward_context( + current_timestep=0, + attn_metadata=None, + ): + output2 = model2(hidden_states=hidden_states, + encoder_hidden_states=encoder_hidden_states, + timestep=timestep) + + # Check if outputs have the same shape + assert output1.shape == output2.shape, f"Output shapes don't match: {output1.shape} vs {output2.shape}" + assert output1.dtype == output2.dtype, f"Output dtype don't match: {output1.dtype} vs {output2.dtype}" + + # Check if outputs are similar (allowing for small numerical differences) + max_diff = torch.max(torch.abs(output1 - output2)) + mean_diff = torch.mean(torch.abs(output1 - output2)) + assert max_diff < 1e-1, f"Maximum difference between outputs: {max_diff.item()}" + # mean diff + assert mean_diff < 1e-2, f"Mean difference between outputs: {mean_diff.item()}" \ No newline at end of file diff --git a/fastvideo/v1/tests/vaes/test_wan_vae.py b/fastvideo/v1/tests/vaes/test_wan_vae.py new file mode 100644 index 000000000..80a45d9bf --- /dev/null +++ b/fastvideo/v1/tests/vaes/test_wan_vae.py @@ -0,0 +1,98 @@ +# SPDX-License-Identifier: Apache-2.0 +import os + +import numpy as np +import pytest +import torch +from diffusers import AutoencoderKLWan + +from fastvideo.v1.inference_args import InferenceArgs +from fastvideo.v1.logger import init_logger +from fastvideo.v1.models.loader.component_loader import VAELoader +from fastvideo.v1.utils import maybe_download_model + +logger = init_logger(__name__) + +os.environ["MASTER_ADDR"] = "localhost" +os.environ["MASTER_PORT"] = "29503" + +BASE_MODEL_PATH = "Wan-AI/Wan2.1-T2V-1.3B-Diffusers" +MODEL_PATH = maybe_download_model(BASE_MODEL_PATH, + local_dir=os.path.join( + 'data', BASE_MODEL_PATH)) +VAE_PATH = os.path.join(MODEL_PATH, "vae") + + +@pytest.mark.usefixtures("distributed_setup") +def test_wan_vae(): + device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") + precision = torch.bfloat16 + precision_str = "bf16" + args = InferenceArgs(model_path=VAE_PATH, vae_precision=precision_str) + args.device = device + + loader = VAELoader() + model2 = loader.load(VAE_PATH, "", args) + + model1 = AutoencoderKLWan.from_pretrained( + VAE_PATH, torch_dtype=precision).to(device).eval() + + # Create identical inputs for both models + batch_size = 1 + + # Video input [B, C, T, H, W] + input_tensor = torch.randn(batch_size, + 3, + 81, + 32, + 32, + device=device, + dtype=precision) + latent_tensor = torch.randn(batch_size, + 16, + 21, + 32, + 32, + device=device, + dtype=precision) + + # Disable gradients for inference + with torch.no_grad(): + # Test encoding + logger.info("Testing encoding...") + latent1 = model1.encode(input_tensor).latent_dist.mean + print("--------------------------------") + latent2 = model2.encode(input_tensor).mean + # Check if latents have the same shape + assert latent1.shape == latent2.shape, f"Latent shapes don't match: {latent1.shape} vs {latent2.shape}" + assert latent1.shape == latent2.shape, f"Latent shapes don't match: {latent1.shape} vs {latent2.shape}" + # Check if latents are similar + max_diff_encode = torch.max(torch.abs(latent1 - latent2)) + mean_diff_encode = torch.mean(torch.abs(latent1 - latent2)) + logger.info("Maximum difference between encoded latents: %s", + max_diff_encode.item()) + logger.info("Mean difference between encoded latents: %s", + mean_diff_encode.item()) + assert mean_diff_encode < 5e-1, f"Encoded latents differ significantly: mean diff = {mean_diff_encode.item()}" + # Test decoding + logger.info("Testing decoding...") + latents_mean = (torch.tensor(model1.config.latents_mean).view( + 1, model1.config.z_dim, 1, 1, 1).to(latent_tensor.device, + latent_tensor.dtype)) + latents_std = 1.0 / torch.tensor(model1.config.latents_std).view( + 1, model1.config.z_dim, 1, 1, 1).to(latent_tensor.device, + latent_tensor.dtype) + latent_tensor = latent_tensor / latents_std + latents_mean + output2 = model2.decode(latent_tensor) + output1 = model1.decode(latent_tensor).sample + # Check if outputs have the same shape + assert output1.shape == output2.shape, f"Output shapes don't match: {output1.shape} vs {output2.shape}" + + # Check if outputs are similar + max_diff_decode = torch.max(torch.abs(output1 - output2)) + mean_diff_decode = torch.mean(torch.abs(output1 - output2)) + logger.info("Maximum difference between decoded outputs: %s", + max_diff_decode.item()) + logger.info("Mean difference between decoded outputs: %s", + mean_diff_decode.item()) + assert mean_diff_decode < 1e-1, f"Decoded outputs differ significantly: mean diff = {mean_diff_decode.item()}" \ No newline at end of file diff --git a/pyproject.toml b/pyproject.toml index 4f1acf8ff..fba1d8536 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -19,7 +19,7 @@ dependencies = [ # Machine Learning & Transformers "transformers>=4.46.1", "tokenizers>=0.20.1", "sentencepiece==0.2.0", - "timm==1.0.11", "peft==0.13.2", "diffusers==0.32.0", "bitsandbytes", + "timm==1.0.11", "peft==0.13.2", "diffusers>=0.33.0", "bitsandbytes", "torch==2.5.1", "torchvision", # vLLM @@ -37,9 +37,8 @@ dependencies = [ # Miscellaneous Utilities "tqdm==4.66.5", "PyYAML==6.0.1", "idna==3.6", "protobuf==5.28.3", - "gradio==5.3.0", "huggingface_hub==0.26.1", "moviepy==1.0.3", "flask", + "gradio==5.3.0", "moviepy==1.0.3", "flask", "flask_restful", "aiohttp", "huggingface_hub", "cloudpickle", - # System & Monitoring Tools "gpustat", "watch", diff --git a/scripts/inference/v1_inference_wan.sh b/scripts/inference/v1_inference_wan.sh new file mode 100755 index 000000000..e4a97e78f --- /dev/null +++ b/scripts/inference/v1_inference_wan.sh @@ -0,0 +1,27 @@ +#!/bin/bash + +num_gpus=2 +export FASTVIDEO_ATTENTION_BACKEND= +export MODEL_BASE=/workspace/data/Wan2.1-T2V-1.3B-Diffusers +# export MODEL_BASE=hunyuanvideo-community/HunyuanVideo +# Note that the tp_size and sp_size should be the same and equal to the number +# of GPUs. They are used for different parallel groups. sp_size is used for +# dit model and tp_size is used for encoder models. +torchrun --nnodes=1 --nproc_per_node=$num_gpus --master_port 29503 \ + fastvideo/v1/sample/v1_fastvideo_inference.py \ + --sp_size $num_gpus \ + --tp_size $num_gpus \ + --height 480 \ + --width 832 \ + --num_frames 77 \ + --num_inference_steps 50 \ + --fps 16 \ + --guidance_scale 3.0 \ + --prompt_path ./assets/prompt.txt \ + --neg_prompt "Bright tones, overexposed, static, blurred details, subtitles, style, works, paintings, images, static, overall gray, worst quality, low quality, JPEG compression residue, ugly, incomplete, extra fingers, poorly drawn hands, poorly drawn faces, deformed, disfigured, misshapen limbs, fused fingers, still picture, messy background, three legs, many people in the background, walking backwards" \ + --seed 1024 \ + --output_path outputs_video/ \ + --model_path $MODEL_BASE \ + --vae-sp \ + --text-encoder-precision "fp32" \ + --use-cpu-offload \ No newline at end of file diff --git a/scripts/inference/v1_inference_wan_i2v.sh b/scripts/inference/v1_inference_wan_i2v.sh new file mode 100644 index 000000000..e52a81f69 --- /dev/null +++ b/scripts/inference/v1_inference_wan_i2v.sh @@ -0,0 +1,29 @@ +#!/bin/bash + +num_gpus=2 +export FASTVIDEO_ATTENTION_BACKEND= +export MODEL_BASE=/workspace/data/Wan2.1-I2V-14B-480P-Diffusers +# export MODEL_BASE=hunyuanvideo-community/HunyuanVideo +# Note that the tp_size and sp_size should be the same and equal to the number +# of GPUs. They are used for different parallel groups. sp_size is used for +# dit model and tp_size is used for encoder models. +torchrun --nnodes=1 --nproc_per_node=$num_gpus --master_port 29503 \ + fastvideo/v1/sample/v1_fastvideo_inference.py \ + --sp_size $num_gpus \ + --tp_size $num_gpus \ + --height 480 \ + --width 832 \ + --num_frames 77 \ + --num_inference_steps 40 \ + --fps 16 \ + --flow_shift 3.0 \ + --guidance_scale 5.0 \ + --image_path "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/astronaut.jpg" \ + --prompt "An astronaut hatching from an egg, on the surface of the moon, the darkness and depth of space realised in the background. High quality, ultrarealistic detail and breath-taking movie-like camera shot." \ + --neg_prompt "Bright tones, overexposed, static, blurred details, subtitles, style, works, paintings, images, static, overall gray, worst quality, low quality, JPEG compression residue, ugly, incomplete, extra fingers, poorly drawn hands, poorly drawn faces, deformed, disfigured, misshapen limbs, fused fingers, still picture, messy background, three legs, many people in the background, walking backwards" \ + --seed 1024 \ + --output_path outputs_i2v/ \ + --model_path $MODEL_BASE \ + --vae-sp \ + --text-encoder-precision "fp32" \ + --use-cpu-offload \ No newline at end of file