diff --git a/fastvideo/attention/backends/flash_attn.py b/fastvideo/attention/backends/flash_attn.py index 16fbb9d44..5c22d705c 100644 --- a/fastvideo/attention/backends/flash_attn.py +++ b/fastvideo/attention/backends/flash_attn.py @@ -46,11 +46,13 @@ class FlashAttentionBackend(AttentionBackend): def get_builder_cls() -> type["AttentionMetadataBuilder"]: raise NotImplementedError + @dataclass class FlashAttnMetadata(AttentionMetadata): current_timestep: int attn_mask: torch.Tensor | None = None + class FlashAttnMetadataBuilder(AttentionMetadataBuilder): def __init__(self): @@ -64,7 +66,8 @@ class FlashAttnMetadataBuilder(AttentionMetadataBuilder): current_timestep: int, attn_mask: torch.Tensor, ) -> FlashAttnMetadata: - return FlashAttnMetadata(current_timestep=current_timestep, attn_mask=attn_mask) + return FlashAttnMetadata(current_timestep=current_timestep, + attn_mask=attn_mask) class FlashAttentionImpl(AttentionImpl): @@ -94,14 +97,13 @@ class FlashAttentionImpl(AttentionImpl): attn_mask = attn_metadata.attn_mask qkv = torch.stack([query, key, value], dim=2) - attn_mask = F.pad(attn_mask, (qkv.shape[1] - attn_mask.shape[1], 0), value=True) - output = flash_attn_no_pad( - qkv, - attn_mask, - causal=False, - dropout_p=0, - softmax_scale=None - ) + attn_mask = F.pad(attn_mask, (qkv.shape[1] - attn_mask.shape[1], 0), + value=True) + output = flash_attn_no_pad(qkv, + attn_mask, + causal=False, + dropout_p=0, + softmax_scale=None) else: output = flash_attn_func( query, # type: ignore[no-untyped-call] diff --git a/fastvideo/attention/backends/sdpa.py b/fastvideo/attention/backends/sdpa.py index b7746e59b..00a0cee60 100644 --- a/fastvideo/attention/backends/sdpa.py +++ b/fastvideo/attention/backends/sdpa.py @@ -3,7 +3,8 @@ import torch from dataclasses import dataclass from fastvideo.attention.backends.abstract import ( # FlashAttentionMetadata, - AttentionBackend, AttentionImpl, AttentionMetadata, AttentionMetadataBuilder) + AttentionBackend, AttentionImpl, AttentionMetadata, + AttentionMetadataBuilder) from fastvideo.logger import init_logger logger = init_logger(__name__) @@ -29,11 +30,13 @@ class SDPABackend(AttentionBackend): # def get_metadata_cls() -> Type["AttentionMetadata"]: # return FlashAttentionMetadata + @dataclass class SDPAMetadata(AttentionMetadata): current_timestep: int attn_mask: torch.Tensor | None = None + class SDPAMetadataBuilder(AttentionMetadataBuilder): def __init__(self): @@ -47,7 +50,9 @@ class SDPAMetadataBuilder(AttentionMetadataBuilder): current_timestep: int, attn_mask: torch.Tensor, ) -> SDPAMetadata: - return SDPAMetadata(current_timestep=current_timestep, attn_mask=attn_mask) + return SDPAMetadata(current_timestep=current_timestep, + attn_mask=attn_mask) + class SDPAImpl(AttentionImpl): diff --git a/fastvideo/attention/utils/flash_attn_no_pad.py b/fastvideo/attention/utils/flash_attn_no_pad.py index 6a6fcbbfb..1983f8d38 100644 --- a/fastvideo/attention/utils/flash_attn_no_pad.py +++ b/fastvideo/attention/utils/flash_attn_no_pad.py @@ -5,7 +5,7 @@ # https://github.com/Tencent-Hunyuan/HunyuanVideo-1.5/blob/main/LICENSE # # Unless and only to the extent required by applicable law, the Tencent Hunyuan works and any -# output and results therefrom are provided "AS IS" without any express or implied warranties of +# output and results there from are provided "AS IS" without any express or implied warranties of # any kind including any warranties of title, merchantability, noninfringement, course of dealing, # usage of trade, or fitness for a particular purpose. You are solely responsible for determining the # appropriateness of using, reproducing, modifying, performing, displaying or distributing any of @@ -17,9 +17,12 @@ from einops import rearrange -def flash_attn_no_pad( - qkv, key_padding_mask, causal=False, dropout_p=0.0, softmax_scale=None, deterministic=False -): +def flash_attn_no_pad(qkv, + key_padding_mask, + causal=False, + dropout_p=0.0, + softmax_scale=None, + deterministic=False): from flash_attn import flash_attn_varlen_qkvpacked_func from flash_attn.bert_padding import pad_input, unpad_input batch_size = qkv.shape[0] @@ -27,10 +30,12 @@ def flash_attn_no_pad( nheads = qkv.shape[-2] x = rearrange(qkv, "b s three h d -> b s (three h d)") x_unpad, indices, cu_seqlens, max_s, used_seqlens_in_batch = unpad_input( - x, key_padding_mask - ) + x, key_padding_mask) - x_unpad = rearrange(x_unpad, "nnz (three h d) -> nnz three h d", three=3, h=nheads) + x_unpad = rearrange(x_unpad, + "nnz (three h d) -> nnz three h d", + three=3, + h=nheads) output_unpad = flash_attn_varlen_qkvpacked_func( x_unpad, cu_seqlens, @@ -41,52 +46,54 @@ def flash_attn_no_pad( deterministic=deterministic, ) output = rearrange( - pad_input( - rearrange(output_unpad, "nnz h d -> nnz (h d)"), indices, batch_size, seqlen - ), + pad_input(rearrange(output_unpad, "nnz h d -> nnz (h d)"), indices, + batch_size, seqlen), "b s (h d) -> b s h d", h=nheads, ) return output -def flash_attn_no_pad_v3( - qkv, key_padding_mask, causal=False, dropout_p=0.0, softmax_scale=None, deterministic=False -): - from flash_attn import flash_attn_varlen_qkvpacked_func + +def flash_attn_no_pad_v3(qkv, + key_padding_mask, + causal=False, + dropout_p=0.0, + softmax_scale=None, + deterministic=False): from flash_attn.bert_padding import pad_input, unpad_input from flash_attn_interface import flash_attn_varlen_func as flash_attn_varlen_func_v3 if flash_attn_varlen_func_v3 is None: raise ImportError("FlashAttention V3 backend not available") - + batch_size, seqlen, _, nheads, head_dim = qkv.shape query, key, value = qkv.unbind(dim=2) - + query_unpad, indices, cu_seqlens_q, max_seqlen_q, _ = unpad_input( - rearrange(query, "b s h d -> b s (h d)"), key_padding_mask - ) + rearrange(query, "b s h d -> b s (h d)"), key_padding_mask) key_unpad, _, cu_seqlens_k, _, _ = unpad_input( - rearrange(key, "b s h d -> b s (h d)"), key_padding_mask - ) + rearrange(key, "b s h d -> b s (h d)"), key_padding_mask) value_unpad, _, _, _, _ = unpad_input( - rearrange(value, "b s h d -> b s (h d)"), key_padding_mask - ) - + rearrange(value, "b s h d -> b s (h d)"), key_padding_mask) + query_unpad = rearrange(query_unpad, "nnz (h d) -> nnz h d", h=nheads) key_unpad = rearrange(key_unpad, "nnz (h d) -> nnz h d", h=nheads) value_unpad = rearrange(value_unpad, "nnz (h d) -> nnz h d", h=nheads) - - output_unpad = flash_attn_varlen_func_v3( - query_unpad, key_unpad, value_unpad, - cu_seqlens_q, cu_seqlens_k, - max_seqlen_q, max_seqlen_q, - softmax_scale=softmax_scale, - causal=causal, - deterministic=deterministic - ) - - output = rearrange( - pad_input(rearrange(output_unpad, "nnz h d -> nnz (h d)"), indices, batch_size, seqlen), - "b s (h d) -> b s h d", h=nheads - ) + + output_unpad = flash_attn_varlen_func_v3(query_unpad, + key_unpad, + value_unpad, + cu_seqlens_q, + cu_seqlens_k, + max_seqlen_q, + max_seqlen_q, + softmax_scale=softmax_scale, + causal=causal, + deterministic=deterministic) + + output = rearrange(pad_input( + rearrange(output_unpad, "nnz h d -> nnz (h d)"), indices, batch_size, + seqlen), + "b s (h d) -> b s h d", + h=nheads) return output diff --git a/fastvideo/configs/models/dits/hunyuanvideo15.py b/fastvideo/configs/models/dits/hunyuanvideo15.py index 666065b6f..6d7f77208 100644 --- a/fastvideo/configs/models/dits/hunyuanvideo15.py +++ b/fastvideo/configs/models/dits/hunyuanvideo15.py @@ -1,8 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 from dataclasses import dataclass, field -import torch - from fastvideo.configs.models.dits.base import DiTArchConfig, DiTConfig @@ -10,7 +8,7 @@ def is_double_block(n: str, m) -> bool: return "double_blocks" in n and str.isdigit(n.split(".")[-1]) -def is_refiner_block(n: str, m) -> bool: +def is_refiner_block(n: str, m) -> bool: return "refiner" in n and str.isdigit(n.split(".")[-1]) @@ -21,8 +19,7 @@ def is_txt_in(n: str, m) -> bool: @dataclass class HunyuanVideo15ArchConfig(DiTArchConfig): _fsdp_shard_conditions: list = field( - default_factory=lambda: - [is_double_block, is_refiner_block]) + default_factory=lambda: [is_double_block, is_refiner_block]) _compile_conditions: list = field( default_factory=lambda: [is_double_block, is_refiner_block, is_txt_in]) diff --git a/fastvideo/configs/models/encoders/__init__.py b/fastvideo/configs/models/encoders/__init__.py index 04afb17e8..86a3e3334 100644 --- a/fastvideo/configs/models/encoders/__init__.py +++ b/fastvideo/configs/models/encoders/__init__.py @@ -11,5 +11,6 @@ from fastvideo.configs.models.encoders.qwen2_5 import Qwen2_5_VLConfig __all__ = [ "EncoderConfig", "TextEncoderConfig", "ImageEncoderConfig", "BaseEncoderOutput", "CLIPTextConfig", "CLIPVisionConfig", - "WAN2_1ControlCLIPVisionConfig", "LlamaConfig", "T5Config", "T5LargeConfig", "Qwen2_5_VLConfig" + "WAN2_1ControlCLIPVisionConfig", "LlamaConfig", "T5Config", "T5LargeConfig", + "Qwen2_5_VLConfig" ] diff --git a/fastvideo/configs/models/encoders/qwen2_5.py b/fastvideo/configs/models/encoders/qwen2_5.py index dd471fdb5..3af6a6687 100644 --- a/fastvideo/configs/models/encoders/qwen2_5.py +++ b/fastvideo/configs/models/encoders/qwen2_5.py @@ -1,10 +1,10 @@ # SPDX-License-Identifier: Apache-2.0 from dataclasses import dataclass, field -from typing import Optional from fastvideo.configs.models.encoders.base import (TextEncoderArchConfig, TextEncoderConfig) + def _is_transformer_layer(n: str, m) -> bool: return "layers" in n and str.isdigit(n.split(".")[-1]) @@ -33,26 +33,26 @@ class Qwen2_5_VLArchConfig(TextEncoderArchConfig): tie_word_embeddings: bool = False rope_theta: float = 1000000.0 use_sliding_window: bool = False - sliding_window: int = 4096 + sliding_window: int | None = 4096 max_window_layers: int = 80 layer_types: list = field(default_factory=list) attention_dropout: float = 0.0 - rope_scaling: Optional[dict] = None - bos_token_id: Optional[int] = None - eos_token_id: Optional[int] = None - pad_token_id: Optional[int] = None + rope_scaling: dict | None = None + bos_token_id: int | None = None + eos_token_id: int | None = None + pad_token_id: int | None = None vision_token_id: int = 151654 model_type: str = "qwen2_5_vl_text" dtype: str = "bfloat16" - - stacked_params_mapping: list[tuple[str, str, str]] = field( - default_factory=lambda: [ - (".qkv_proj", ".q_proj", "q"), - (".qkv_proj", ".k_proj", "k"), - (".qkv_proj", ".v_proj", "v"), - (".gate_up_proj", ".gate_proj", 0), - (".gate_up_proj", ".up_proj", 1), - ]) + + stacked_params_mapping: list[tuple[str, str, str + | int]] = field(default_factory=lambda: [ + (".qkv_proj", ".q_proj", "q"), + (".qkv_proj", ".k_proj", "k"), + (".qkv_proj", ".v_proj", "v"), + (".gate_up_proj", ".gate_proj", 0), + (".gate_up_proj", ".up_proj", 1), + ]) _fsdp_shard_conditions: list = field( default_factory=lambda: [_is_transformer_layer, _is_embeddings, _is_final_norm]) @@ -65,16 +65,15 @@ class Qwen2_5_VLArchConfig(TextEncoderArchConfig): self.num_key_value_heads = self.num_attention_heads if self.layer_types is None: self.layer_types = [ - "sliding_attention" - if self.sliding_window is not None and i >= self.max_window_layers - else "full_attention" + "sliding_attention" if self.sliding_window is not None + and i >= self.max_window_layers else "full_attention" for i in range(self.num_hidden_layers) ] if self.rope_scaling is not None and "type" in self.rope_scaling: if self.rope_scaling["type"] == "mrope": self.rope_scaling["type"] = "default" self.rope_scaling["rope_type"] = self.rope_scaling["type"] - + self.tokenizer_kwargs = { "add_generation_prompt": True, "tokenize": True, @@ -85,8 +84,10 @@ class Qwen2_5_VLArchConfig(TextEncoderArchConfig): "return_tensors": "pt", } + @dataclass class Qwen2_5_VLConfig(TextEncoderConfig): - arch_config: TextEncoderArchConfig = field(default_factory=Qwen2_5_VLArchConfig) + arch_config: TextEncoderArchConfig = field( + default_factory=Qwen2_5_VLArchConfig) prefix: str = "qwen2_5_vl" - is_chat_model: bool = True \ No newline at end of file + is_chat_model: bool = True diff --git a/fastvideo/configs/pipelines/__init__.py b/fastvideo/configs/pipelines/__init__.py index 41432d1a5..530f9352e 100644 --- a/fastvideo/configs/pipelines/__init__.py +++ b/fastvideo/configs/pipelines/__init__.py @@ -11,9 +11,9 @@ from fastvideo.configs.pipelines.wan import (SelfForcingWanT2V480PConfig, WanT2V480PConfig, WanT2V720PConfig) __all__ = [ - "HunyuanConfig", "FastHunyuanConfig", "PipelineConfig", "Hunyuan15T2V480PConfig", - "SlidingTileAttnConfig", "WanT2V480PConfig", "WanI2V480PConfig", - "WanT2V720PConfig", "WanI2V720PConfig", "StepVideoT2VConfig", - "SelfForcingWanT2V480PConfig", "CosmosConfig", + "HunyuanConfig", "FastHunyuanConfig", "PipelineConfig", + "Hunyuan15T2V480PConfig", "SlidingTileAttnConfig", "WanT2V480PConfig", + "WanI2V480PConfig", "WanT2V720PConfig", "WanI2V720PConfig", + "StepVideoT2VConfig", "SelfForcingWanT2V480PConfig", "CosmosConfig", "get_pipeline_config_cls_from_name" ] diff --git a/fastvideo/configs/pipelines/hunyuan15.py b/fastvideo/configs/pipelines/hunyuan15.py index be97051f4..98233a302 100644 --- a/fastvideo/configs/pipelines/hunyuan15.py +++ b/fastvideo/configs/pipelines/hunyuan15.py @@ -1,7 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 from collections.abc import Callable from dataclasses import dataclass, field -from typing import TypedDict, List, Dict, Any +from typing import Any import re import torch @@ -22,7 +22,8 @@ PROMPT_TEMPLATE_ENCODE_VIDEO = "You are a helpful assistant. Describe the video 4. background environment, light, style and atmosphere. \ 5. camera angles, movements, and transitions used in the video." -def extract_glyph_texts(prompt: str) -> List[str]: + +def extract_glyph_texts(prompt: str) -> str | None: """ Extract glyph texts from prompt using regex pattern. @@ -38,13 +39,15 @@ def extract_glyph_texts(prompt: str) -> List[str]: result = list(dict.fromkeys(result)) if len(result) > 1 else result if result: - formatted_result = ". ".join([f'Text "{text}"' for text in result]) + ". " + formatted_result = ". ".join([f'Text "{text}"' + for text in result]) + ". " else: formatted_result = None return formatted_result -def format_text_input(prompt: str, system_message: str) -> List[Dict[str, Any]]: + +def format_text_input(prompt: str, system_message: str) -> list[dict[str, Any]]: """ Apply text to template. @@ -56,17 +59,25 @@ def format_text_input(prompt: str, system_message: str) -> List[Dict[str, Any]]: List[Dict[str, Any]]: List of chat conversation. """ - template = [{"role": "system", "content": system_message}, {"role": "user", "content": prompt if prompt else " "}] + template = [{ + "role": "system", + "content": system_message + }, { + "role": "user", + "content": prompt if prompt else " " + }] return template -def qwen_preprocess_text(prompt: str) -> str: - prompt = format_text_input(prompt, PROMPT_TEMPLATE_ENCODE_VIDEO) - return prompt +def qwen_preprocess_text(prompt: str) -> list[dict[str, Any]]: + output = format_text_input(prompt, PROMPT_TEMPLATE_ENCODE_VIDEO) + return output -def qwen_postprocess_text(outputs: BaseEncoderOutput, mask: torch.tensor) -> tuple[torch.tensor, torch.tensor]: +def qwen_postprocess_text( + outputs: BaseEncoderOutput, + mask: torch.tensor) -> tuple[torch.tensor, torch.tensor]: assert outputs.hidden_states is not None output = outputs.hidden_states[-3] output = output[:, PROMPT_TEMPLATE_TOKEN_LENGTH:] @@ -74,9 +85,9 @@ def qwen_postprocess_text(outputs: BaseEncoderOutput, mask: torch.tensor) -> tup return output, mask -def byt5_preprocess_text(prompt: str) -> str: - prompt = [prompt] if isinstance(prompt, str) else prompt - glyph_texts = [extract_glyph_texts(p) for p in prompt] +def byt5_preprocess_text(prompt: str) -> str | None: + prompts = [prompt] if isinstance(prompt, str) else prompt + glyph_texts = [extract_glyph_texts(p) for p in prompts] return glyph_texts[0] @@ -99,12 +110,10 @@ class Hunyuan15T2V480PConfig(PipelineConfig): # Text encoding stage text_encoder_configs: tuple[EncoderConfig, ...] = field( default_factory=lambda: (Qwen2_5_VLConfig(), T5Config())) - preprocess_text_funcs: tuple[Callable[[str], str], ...] = field( + preprocess_text_funcs: tuple[Callable[[Any], Any], ...] = field( default_factory=lambda: (qwen_preprocess_text, byt5_preprocess_text)) - postprocess_text_funcs: tuple[ - Callable[[BaseEncoderOutput], torch.tensor], - ...] = field(default_factory=lambda: - (qwen_postprocess_text, byt5_postprocess_text)) + postprocess_text_funcs: tuple[Callable[..., Any], ...] = field( + default_factory=lambda: (qwen_postprocess_text, byt5_postprocess_text)) # Precision for each component dit_precision: str = "bf16" diff --git a/fastvideo/configs/pipelines/registry.py b/fastvideo/configs/pipelines/registry.py index 383d4cba1..0d70c6be2 100644 --- a/fastvideo/configs/pipelines/registry.py +++ b/fastvideo/configs/pipelines/registry.py @@ -28,7 +28,8 @@ logger = init_logger(__name__) PIPE_NAME_TO_CONFIG: dict[str, type[PipelineConfig]] = { "FastVideo/FastHunyuan-diffusers": FastHunyuanConfig, "hunyuanvideo-community/HunyuanVideo": HunyuanConfig, - "hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v": Hunyuan15T2V480PConfig, + "hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v": + Hunyuan15T2V480PConfig, "Wan-AI/Wan2.1-T2V-1.3B-Diffusers": WanT2V480PConfig, "weizhou03/Wan2.1-Fun-1.3B-InP-Diffusers": WanI2V480PConfig, "IRMChen/Wan2.1-Fun-1.3B-Control-Diffusers": WANV2VConfig, diff --git a/fastvideo/configs/sample/hunyuan15.py b/fastvideo/configs/sample/hunyuan15.py index 4807a53f6..37a83647d 100644 --- a/fastvideo/configs/sample/hunyuan15.py +++ b/fastvideo/configs/sample/hunyuan15.py @@ -17,4 +17,5 @@ class Hunyuan15_480P_SamplingParam(SamplingParam): guidance_scale: float = 6.0 prompt_attention_mask: list = field(default_factory=list) negative_attention_mask: list = field(default_factory=list) - sigmas: list[float] | None = field(default_factory=lambda: list(np.linspace(1.0, 0.0, 50 + 1)[:-1])) \ No newline at end of file + sigmas: list[float] | None = field( + default_factory=lambda: list(np.linspace(1.0, 0.0, 50 + 1)[:-1])) diff --git a/fastvideo/models/loader/component_loader.py b/fastvideo/models/loader/component_loader.py index 778cc4cc1..e8bbd5453 100644 --- a/fastvideo/models/loader/component_loader.py +++ b/fastvideo/models/loader/component_loader.py @@ -249,7 +249,6 @@ class TextEncoderLoader(ComponentLoader): dtype: str = "fp16"): use_cpu_offload = fastvideo_args.text_encoder_cpu_offload and len( getattr(model_config, "_fsdp_shard_conditions", [])) > 0 - logger.info("Text encoder use_cpu_offload: %s", use_cpu_offload) from fastvideo.platforms import current_platform @@ -481,8 +480,6 @@ class TransformerLoader(ComponentLoader): fastvideo_args.pipeline_config.dit_precision] # Load the model using FSDP loader - logger.info("Loading model from %s, default_dtype: %s, cpu_offload: %s", cls_name, - default_dtype, fastvideo_args.dit_cpu_offload) assert fastvideo_args.hsdp_shard_dim is not None model = maybe_load_fsdp_model( model_cls=model_cls, diff --git a/fastvideo/pipelines/basic/hunyuan15/hunyuan15_pipeline.py b/fastvideo/pipelines/basic/hunyuan15/hunyuan15_pipeline.py index 62d6b8aca..87a58c692 100644 --- a/fastvideo/pipelines/basic/hunyuan15/hunyuan15_pipeline.py +++ b/fastvideo/pipelines/basic/hunyuan15/hunyuan15_pipeline.py @@ -57,11 +57,10 @@ class HunyuanVideo15Pipeline(ComposedPipelineBase): stage=LatentPreparationStage( scheduler=self.get_module("scheduler"), transformer=self.get_module("transformer"))) - + self.add_stage(stage_name="image_encoding_stage", - stage=Hy15ImageEncodingStage( - image_encoder=None, - image_processor=None)) + stage=Hy15ImageEncodingStage(image_encoder=None, + image_processor=None)) self.add_stage(stage_name="denoising_stage", stage=DenoisingStage( diff --git a/fastvideo/pipelines/stages/denoising.py b/fastvideo/pipelines/stages/denoising.py index 084ed36b6..4ab4e1b76 100644 --- a/fastvideo/pipelines/stages/denoising.py +++ b/fastvideo/pipelines/stages/denoising.py @@ -3,7 +3,6 @@ Denoising stage for diffusion pipelines. """ -import gc import inspect import weakref from collections.abc import Iterable @@ -475,11 +474,6 @@ class DenoisingStage(PipelineStage): logger.info("Memory after deallocating transformer: %s", torch.mps.current_allocated_memory()) - if fastvideo_args.dit_cpu_offload: - self.transformer = self.transformer.to('cpu') - gc.collect() - torch.cuda.empty_cache() - return batch def prepare_extra_func_kwargs(self, func, kwargs) -> dict[str, Any]: diff --git a/fastvideo/pipelines/stages/image_encoding.py b/fastvideo/pipelines/stages/image_encoding.py index 743bf2af3..8bccd29e9 100644 --- a/fastvideo/pipelines/stages/image_encoding.py +++ b/fastvideo/pipelines/stages/image_encoding.py @@ -99,16 +99,19 @@ class ImageEncodingStage(PipelineStage): V.list_of_tensors_dims(3)) return result + class Hy15ImageEncodingStage(ImageEncodingStage): """ Stage for encoding image prompts into embeddings for HunyuanVideo1.5 models. """ + def verify_input(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> VerificationResult: """Verify image encoding stage inputs.""" return VerificationResult() - def forward(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> ForwardBatch: + def forward(self, batch: ForwardBatch, + fastvideo_args: FastVideoArgs) -> ForwardBatch: """ Encode the prompt into image encoder hidden states. """ @@ -119,7 +122,8 @@ class Hy15ImageEncodingStage(ImageEncodingStage): raw_latent_shape = list(batch.raw_latent_shape) raw_latent_shape[1] = 1 - batch.video_latent = torch.zeros(tuple(raw_latent_shape), device=get_local_torch_device()) + batch.video_latent = torch.zeros(tuple(raw_latent_shape), + device=get_local_torch_device()) return batch diff --git a/fastvideo/pipelines/stages/text_encoding.py b/fastvideo/pipelines/stages/text_encoding.py index 19c7c19b2..612749bdc 100644 --- a/fastvideo/pipelines/stages/text_encoding.py +++ b/fastvideo/pipelines/stages/text_encoding.py @@ -6,6 +6,7 @@ This module contains implementations of prompt encoding stages for diffusion pip """ import torch +from typing import Any from fastvideo.distributed import get_local_torch_device from fastvideo.fastvideo_args import FastVideoArgs @@ -69,17 +70,12 @@ class TextEncodingStage(PipelineStage): encoder_index=all_indices, return_attention_mask=True, ) - logger.info(f"prompt: {prompt_text}") - logger.info(f"pos text 1 sum: {prompt_embeds_list[0].sum()}, pos text 2 sum: {prompt_embeds_list[1].sum()}") - logger.info(f"pos text 1 mask sum: {prompt_masks_list[0].sum()}, pos text 2 mask sum: {prompt_masks_list[1].sum()}") for pe in prompt_embeds_list: batch.prompt_embeds.append(pe) if batch.prompt_attention_mask is not None: for am in prompt_masks_list: batch.prompt_attention_mask.append(am) - else: - raise ValueError("prompt_attention_mask is required") # Encode negative prompt if CFG is enabled if batch.do_classifier_free_guidance: @@ -90,9 +86,7 @@ class TextEncodingStage(PipelineStage): encoder_index=all_indices, return_attention_mask=True, ) - logger.info(f"neg text: {batch.negative_prompt}") - logger.info(f"neg text 1 sum: {neg_embeds_list[0].sum()}, neg text 2 sum: {neg_embeds_list[1].sum()}") - logger.info(f"neg text 1 mask sum: {neg_masks_list[0].sum()}, neg text 2 mask sum: {neg_masks_list[1].sum()}") + assert batch.negative_prompt_embeds is not None for ne in neg_embeds_list: batch.negative_prompt_embeds.append(ne) @@ -213,8 +207,11 @@ class TextEncodingStage(PipelineStage): tok_kwargs = dict(encoder_config.tokenizer_kwargs) if max_length is not None: tok_kwargs["max_length"] = max_length - elif hasattr(fastvideo_args.pipeline_config, "text_encoder_max_lengths"): - tok_kwargs["max_length"] = fastvideo_args.pipeline_config.text_encoder_max_lengths[i] + elif hasattr(fastvideo_args.pipeline_config, + "text_encoder_max_lengths"): + tok_kwargs[ + "max_length"] = fastvideo_args.pipeline_config.text_encoder_max_lengths[ + i] if truncation is not None: tok_kwargs["truncation"] = truncation @@ -228,18 +225,21 @@ class TextEncodingStage(PipelineStage): processed_texts.append(processed_text) else: # Assuming batch_size = 1 - prompt_embeds = torch.zeros( - (1, tok_kwargs["max_length"], encoder_config.hidden_size), device=target_device - ) - attention_mask = torch.zeros((1, tok_kwargs["max_length"]), device=target_device, dtype=torch.int64) + prompt_embeds = torch.zeros((1, tok_kwargs["max_length"], + encoder_config.hidden_size), + device=target_device) + attention_mask = torch.zeros((1, tok_kwargs["max_length"]), + device=target_device, + dtype=torch.int64) embeds_list.append(prompt_embeds) attn_masks_list.append(attention_mask) - return self.return_embeds(embeds_list, attn_masks_list, return_type, return_attention_mask, indices) + return self.return_embeds(embeds_list, attn_masks_list, + return_type, + return_attention_mask, indices) if encoder_config.is_chat_model: - logger.info(f"processed_texts: {processed_texts}") - logger.info(f"tok_kwargs: {tok_kwargs}") - text_inputs = tokenizer.apply_chat_template(processed_texts, **tok_kwargs).to(target_device) + text_inputs = tokenizer.apply_chat_template( + processed_texts, **tok_kwargs).to(target_device) else: text_inputs = tokenizer(processed_texts, **tok_kwargs).to(target_device) @@ -256,8 +256,9 @@ class TextEncodingStage(PipelineStage): try: prompt_embeds = postprocess_func(outputs) - except Exception as e: - prompt_embeds, attention_mask = postprocess_func(outputs, attention_mask) + except Exception: + prompt_embeds, attention_mask = postprocess_func( + outputs, attention_mask) if dtype is not None: prompt_embeds = prompt_embeds.to(dtype=dtype) @@ -265,15 +266,17 @@ class TextEncodingStage(PipelineStage): if return_attention_mask: attn_masks_list.append(attention_mask) - return self.return_embeds(embeds_list, attn_masks_list, return_type, return_attention_mask, indices) - - def return_embeds(self, - embeds_list: list[torch.Tensor], - attn_masks_list: list[torch.Tensor], + return self.return_embeds(embeds_list, attn_masks_list, return_type, + return_attention_mask, indices) + + def return_embeds( + self, + embeds_list: list[torch.Tensor], + attn_masks_list: list[torch.Tensor], return_type: str = "list", return_attention_mask: bool = False, - indices: list[int] = None, - ) -> tuple[list[torch.Tensor], list[torch.Tensor]]: + indices: list[int] | None = None, + ) -> Any: # Shape results according to return_type if return_type == "list": if return_attention_mask: