Fix lint

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]

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):
 

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

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])

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"
 ]

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
+    is_chat_model: bool = True

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"
 ]

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"

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,

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]))
+    sigmas: list[float] | None = field(
+        default_factory=lambda: list(np.linspace(1.0, 0.0, 50 + 1)[:-1]))

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,

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(

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]:

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
 
 

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: