[chore] Upgrade min Python version from 3.8 to 3.10 (#597)

comfyui/video_generator/load_image.py

@@ -1,6 +1,5 @@
 import hashlib
 import os
-from typing import List
 
 import folder_paths
 import numpy as np
@@ -39,8 +38,8 @@ class LoadImagePath:
 
         img = pillow(Image.open, image_path)
 
-        output_images: List[torch.Tensor] = []
-        output_masks: List[torch.Tensor] = []
+        output_images: list[torch.Tensor] = []
+        output_masks: list[torch.Tensor] = []
         w, h = None, None
 
         excluded_formats = ['MPO']

comfyui/video_generator/node_helpers.py

@@ -1,5 +1,6 @@
 import hashlib
-from typing import Any, Callable, Dict, List, Optional, Tuple, TypeVar
+from collections.abc import Callable
+from typing import Any, TypeVar
 
 import torch
 from comfy.cli_args import args
@@ -8,9 +9,8 @@ from PIL import ImageFile, UnidentifiedImageError
 T = TypeVar('T')
 
 
-def conditioning_set_values(
-        conditioning: List[Any],
-        values: Optional[Dict[str, Any]] = None) -> List[Any]:
+def conditioning_set_values(conditioning: list[Any],
+                            values: dict[str, Any] | None = None) -> list[Any]:
     if values is None:
         values = {}
     c = []
@@ -48,7 +48,7 @@ def hasher() -> Callable[[], Any]:
     return hashfuncs[args.default_hashing_function]
 
 
-def string_to_torch_dtype(string: str) -> Optional[torch.dtype]:
+def string_to_torch_dtype(string: str) -> torch.dtype | None:
     if string == "fp32":
         return torch.float32
     if string == "fp16":
@@ -59,7 +59,7 @@ def string_to_torch_dtype(string: str) -> Optional[torch.dtype]:
 
 
 def image_alpha_fix(destination: torch.Tensor,
-                    source: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
+                    source: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
     if destination.shape[-1] < source.shape[-1]:
         source = source[..., :destination.shape[-1]]
     elif destination.shape[-1] > source.shape[-1]:

docs/source/conf.py

@@ -18,7 +18,6 @@ import os
 import re
 import sys
 from pathlib import Path
-from typing import Optional
 
 import requests
 
@@ -168,8 +167,7 @@ _cached_base: str = ""
 _cached_branch: str = ""
 
 
-def get_repo_base_and_branch(
-        pr_number: str) -> tuple[Optional[str], Optional[str]]:
+def get_repo_base_and_branch(pr_number: str) -> tuple[str | None, str | None]:
     global _cached_base, _cached_branch
     if _cached_base and _cached_branch:
         return _cached_base, _cached_branch

docs/source/generate_examples.py

@@ -5,7 +5,6 @@ import itertools
 import re
 from dataclasses import dataclass, field
 from pathlib import Path
-from typing import Optional
 
 ROOT_DIR = Path(__file__).parent.parent.parent.resolve()
 ROOT_DIR_RELATIVE = '../../../..'
@@ -89,7 +88,7 @@ class Example:
         generate() -> str: Generates the documentation content.
     """ # noqa: E501
     path: Path
-    category: Optional[str] = None
+    category: str | None = None
     main_file: Path = field(init=False)
     other_files: list[Path] = field(init=False)
     title: str = field(init=False)

fastvideo/v1/STA_configuration.py

@@ -2,7 +2,7 @@
 import json
 import os
 from collections import defaultdict
-from typing import Any, Dict, List, Optional, Tuple
+from typing import Any
 
 import numpy as np
 
@@ -13,7 +13,7 @@ def configure_sta(mode: str = 'STA_searching',
                   layer_num: int = 40,
                   time_step_num: int = 50,
                   head_num: int = 40,
-                  **kwargs) -> List[List[List[Any]]]:
+                  **kwargs) -> list[list[list[Any]]]:
     """
     Configure Sliding Tile Attention (STA) parameters based on the specified mode.
     
@@ -53,22 +53,22 @@ def configure_sta(mode: str = 'STA_searching',
 
     if mode == 'STA_searching':
         # Get parameters with defaults
-        mask_candidates: Optional[List[str]] = kwargs.get('mask_candidates')
+        mask_candidates: list[str] | None = kwargs.get('mask_candidates')
         if mask_candidates is None:
             raise ValueError(
                 "mask_candidates is required for STA_searching mode")
-        mask_selected: List[int] = kwargs.get('mask_selected',
+        mask_selected: list[int] = kwargs.get('mask_selected',
                                               list(range(len(mask_candidates))))
 
         # Parse selected masks
-        selected_masks: List[List[int]] = []
+        selected_masks: list[list[int]] = []
         for index in mask_selected:
             mask = mask_candidates[index]
             masks_list = [int(x) for x in mask.split(',')]
             selected_masks.append(masks_list)
 
         # Create 3D mask structure with fixed dimensions (t=50, l=60)
-        masks_3d: List[List[List[List[int]]]] = []
+        masks_3d: list[list[list[list[int]]]] = []
         for i in range(time_step_num):  # Fixed t dimension = 50
             row = []
             for j in range(layer_num):  # Fixed l dimension = 60
@@ -79,25 +79,23 @@ def configure_sta(mode: str = 'STA_searching',
 
     elif mode == 'STA_tuning':
         # Get required parameters
-        mask_search_files_path: Optional[str] = kwargs.get(
+        mask_search_files_path: str | None = kwargs.get(
             'mask_search_files_path')
         if not mask_search_files_path:
             raise ValueError(
                 "mask_search_files_path is required for STA_tuning mode")
 
         # Get optional parameters with defaults
-        mask_candidates_tuning: Optional[List[str]] = kwargs.get(
-            'mask_candidates')
+        mask_candidates_tuning: list[str] | None = kwargs.get('mask_candidates')
         if mask_candidates_tuning is None:
             raise ValueError("mask_candidates is required for STA_tuning mode")
-        mask_selected_tuning: List[int] = kwargs.get(
+        mask_selected_tuning: list[int] = kwargs.get(
             'mask_selected', list(range(len(mask_candidates_tuning))))
-        skip_time_steps_tuning: Optional[int] = kwargs.get('skip_time_steps')
-        save_dir_tuning: Optional[str] = kwargs.get('save_dir',
-                                                    "mask_candidates")
+        skip_time_steps_tuning: int | None = kwargs.get('skip_time_steps')
+        save_dir_tuning: str | None = kwargs.get('save_dir', "mask_candidates")
 
         # Parse selected masks
-        selected_masks_tuning: List[List[int]] = []
+        selected_masks_tuning: list[list[int]] = []
         for index in mask_selected_tuning:
             mask = mask_candidates_tuning[index]
             masks_list = [int(x) for x in mask.split(',')]
@@ -108,7 +106,7 @@ def configure_sta(mode: str = 'STA_searching',
         averaged_results = average_head_losses(results, selected_masks_tuning)
 
         # Add full attention mask for specific cases
-        full_attention_mask_tuning: Optional[List[int]] = kwargs.get(
+        full_attention_mask_tuning: list[int] | None = kwargs.get(
             'full_attention_mask')
         if full_attention_mask_tuning is not None:
             selected_masks_tuning.append(full_attention_mask_tuning)
@@ -149,11 +147,11 @@ def configure_sta(mode: str = 'STA_searching',
         return mask_strategy_3d
     elif mode == 'STA_tuning_cfg':
         # Get required parameters for both positive and negative paths
-        mask_search_files_path_pos: Optional[str] = kwargs.get(
+        mask_search_files_path_pos: str | None = kwargs.get(
             'mask_search_files_path_pos')
-        mask_search_files_path_neg: Optional[str] = kwargs.get(
+        mask_search_files_path_neg: str | None = kwargs.get(
             'mask_search_files_path_neg')
-        save_dir_cfg: Optional[str] = kwargs.get('save_dir')
+        save_dir_cfg: str | None = kwargs.get('save_dir')
 
         if not mask_search_files_path_pos or not mask_search_files_path_neg or not save_dir_cfg:
             raise ValueError(
@@ -161,16 +159,16 @@ def configure_sta(mode: str = 'STA_searching',
             )
 
         # Get optional parameters with defaults
-        mask_candidates_cfg: Optional[List[str]] = kwargs.get('mask_candidates')
+        mask_candidates_cfg: list[str] | None = kwargs.get('mask_candidates')
         if mask_candidates_cfg is None:
             raise ValueError(
                 "mask_candidates is required for STA_tuning_cfg mode")
-        mask_selected_cfg: List[int] = kwargs.get(
+        mask_selected_cfg: list[int] = kwargs.get(
             'mask_selected', list(range(len(mask_candidates_cfg))))
-        skip_time_steps_cfg: Optional[int] = kwargs.get('skip_time_steps')
+        skip_time_steps_cfg: int | None = kwargs.get('skip_time_steps')
 
         # Parse selected masks
-        selected_masks_cfg: List[List[int]] = []
+        selected_masks_cfg: list[list[int]] = []
         for index in mask_selected_cfg:
             mask = mask_candidates_cfg[index]
             masks_list = [int(x) for x in mask.split(',')]
@@ -187,7 +185,7 @@ def configure_sta(mode: str = 'STA_searching',
                                                selected_masks_cfg)
 
         # Add full attention mask for specific cases
-        full_attention_mask_cfg: Optional[List[int]] = kwargs.get(
+        full_attention_mask_cfg: list[int] | None = kwargs.get(
             'full_attention_mask')
         if full_attention_mask_cfg is not None:
             selected_masks_cfg.append(full_attention_mask_cfg)
@@ -227,7 +225,7 @@ def configure_sta(mode: str = 'STA_searching',
 
     else:  # STA_inference
         # Get parameters with defaults
-        load_path: Optional[str] = kwargs.get(
+        load_path: str | None = kwargs.get(
             'load_path', "mask_candidates/mask_strategy.json")
         if load_path is None:
             raise ValueError("load_path is required for STA_inference mode")
@@ -248,9 +246,9 @@ def configure_sta(mode: str = 'STA_searching',
 # Helper functions
 
 
-def read_specific_json_files(folder_path: str) -> List[Dict[str, Any]]:
+def read_specific_json_files(folder_path: str) -> list[dict[str, Any]]:
     """Read and parse JSON files containing mask search results."""
-    json_contents: List[Dict[str, Any]] = []
+    json_contents: list[dict[str, Any]] = []
 
     # List files only in the current directory (no walk)
     files = os.listdir(folder_path)
@@ -268,11 +266,11 @@ def read_specific_json_files(folder_path: str) -> List[Dict[str, Any]]:
 
 
 def average_head_losses(
-        results: List[Dict[str, Any]],
-        selected_masks: List[List[int]]) -> Dict[str, Dict[str, np.ndarray]]:
+        results: list[dict[str, Any]],
+        selected_masks: list[list[int]]) -> dict[str, dict[str, np.ndarray]]:
     """Average losses across all prompts for each mask strategy."""
     # Initialize a dictionary to store the averaged results
-    averaged_losses: Dict[str, Dict[str, np.ndarray]] = {}
+    averaged_losses: dict[str, dict[str, np.ndarray]] = {}
     loss_type = 'L2_loss'
     # Get all loss types (e.g., 'L2_loss')
     averaged_losses[loss_type] = {}
@@ -294,14 +292,14 @@ def average_head_losses(
 
 
 def select_best_mask_strategy(
-        averaged_results: Dict[str, Dict[str, np.ndarray]],
-        selected_masks: List[List[int]],
+        averaged_results: dict[str, dict[str, np.ndarray]],
+        selected_masks: list[list[int]],
         skip_time_steps: int = 12,
         timesteps: int = 50,
         head_num: int = 40
-) -> Tuple[Dict[str, List[int]], float, Dict[str, int]]:
+) -> tuple[dict[str, list[int]], float, dict[str, int]]:
     """Select the best mask strategy for each head based on loss minimization."""
-    best_mask_strategy: Dict[str, List[int]] = {}
+    best_mask_strategy: dict[str, list[int]] = {}
     loss_type = 'L2_loss'
     # Get the shape of time steps and layers
     layers = len(averaged_results[loss_type][str(selected_masks[0])][0])
@@ -310,7 +308,7 @@ def select_best_mask_strategy(
     total_tokens = 0  # total number of masked tokens
     total_length = 0  # total sequence length
 
-    strategy_counts: Dict[str, int] = {
+    strategy_counts: dict[str, int] = {
         str(strategy): 0
         for strategy in selected_masks
     }
@@ -352,22 +350,22 @@ def select_best_mask_strategy(
 
 
 def save_mask_search_results(
-        mask_search_final_result: List[Dict[str, List[float]]],
+        mask_search_final_result: list[dict[str, list[float]]],
         prompt: str,
-        mask_strategies: List[str],
-        output_dir: str = 'output/mask_search_result/') -> Optional[str]:
+        mask_strategies: list[str],
+        output_dir: str = 'output/mask_search_result/') -> str | None:
     if not mask_search_final_result:
         print("No mask search results to save")
         return None
 
     # Create result dictionary with defaultdict for nested lists
-    mask_search_dict: Dict[str, Dict[str, List[List[float]]]] = {
+    mask_search_dict: dict[str, dict[str, list[list[float]]]] = {
         "L2_loss": defaultdict(list),
         "L1_loss": defaultdict(list)
     }
 
     mask_selected = list(range(len(mask_strategies)))
-    selected_masks: List[List[int]] = []
+    selected_masks: list[list[int]] = []
     for index in mask_selected:
         mask = mask_strategies[index]
         masks_list = [int(x) for x in mask.split(',')]
@@ -377,7 +375,7 @@ def save_mask_search_results(
     for i, mask_strategy in enumerate(selected_masks):
         mask_strategy_str = str(mask_strategy)
         # Process L2 loss
-        step_results: List[List[float]] = []
+        step_results: list[list[float]] = []
         for step_data in mask_search_final_result:
             if isinstance(step_data, dict) and "L2_loss" in step_data:
                 layer_losses = [float(loss) for loss in step_data["L2_loss"]]

fastvideo/v1/attention/backends/abstract.py

@@ -3,8 +3,7 @@
 
 from abc import ABC, abstractmethod
 from dataclasses import dataclass, fields
-from typing import (TYPE_CHECKING, Any, Dict, Generic, Optional, Protocol, Set,
-                    Type, TypeVar)
+from typing import TYPE_CHECKING, Any, Generic, Protocol, TypeVar
 
 if TYPE_CHECKING:
     from fastvideo.v1.fastvideo_args import FastVideoArgs
@@ -27,12 +26,12 @@ class AttentionBackend(ABC):
 
     @staticmethod
     @abstractmethod
-    def get_impl_cls() -> Type["AttentionImpl"]:
+    def get_impl_cls() -> type["AttentionImpl"]:
         raise NotImplementedError
 
     @staticmethod
     @abstractmethod
-    def get_metadata_cls() -> Type["AttentionMetadata"]:
+    def get_metadata_cls() -> type["AttentionMetadata"]:
         raise NotImplementedError
 
     # @staticmethod
@@ -46,7 +45,7 @@ class AttentionBackend(ABC):
 
     @staticmethod
     @abstractmethod
-    def get_builder_cls() -> Type["AttentionMetadataBuilder"]:
+    def get_builder_cls() -> type["AttentionMetadataBuilder"]:
         raise NotImplementedError
 
 
@@ -57,8 +56,7 @@ class AttentionMetadata:
     current_timestep: int
 
     def asdict_zerocopy(self,
-                        skip_fields: Optional[Set[str]] = None
-                        ) -> Dict[str, Any]:
+                        skip_fields: set[str] | None = None) -> dict[str, Any]:
         """Similar to dataclasses.asdict, but avoids deepcopying."""
         if skip_fields is None:
             skip_fields = set()
@@ -124,7 +122,7 @@ class AttentionImpl(ABC, Generic[T]):
         head_size: int,
         softmax_scale: float,
         causal: bool = False,
-        num_kv_heads: Optional[int] = None,
+        num_kv_heads: int | None = None,
         prefix: str = "",
         **extra_impl_args,
     ) -> None:

fastvideo/v1/attention/backends/flash_attn.py

@@ -1,7 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
 
-from typing import List, Optional, Type
-
 import torch
 from flash_attn import flash_attn_func as flash_attn_2_func
 
@@ -28,7 +26,7 @@ class FlashAttentionBackend(AttentionBackend):
     accept_output_buffer: bool = True
 
     @staticmethod
-    def get_supported_head_sizes() -> List[int]:
+    def get_supported_head_sizes() -> list[int]:
         return [32, 64, 96, 128, 160, 192, 224, 256]
 
     @staticmethod
@@ -36,15 +34,15 @@ class FlashAttentionBackend(AttentionBackend):
         return "FLASH_ATTN"
 
     @staticmethod
-    def get_impl_cls() -> Type["FlashAttentionImpl"]:
+    def get_impl_cls() -> type["FlashAttentionImpl"]:
         return FlashAttentionImpl
 
     @staticmethod
-    def get_metadata_cls() -> Type["AttentionMetadata"]:
+    def get_metadata_cls() -> type["AttentionMetadata"]:
         raise NotImplementedError
 
     @staticmethod
-    def get_builder_cls() -> Type["AttentionMetadataBuilder"]:
+    def get_builder_cls() -> type["AttentionMetadataBuilder"]:
         raise NotImplementedError
 
 
@@ -56,7 +54,7 @@ class FlashAttentionImpl(AttentionImpl):
         head_size: int,
         causal: bool,
         softmax_scale: float,
-        num_kv_heads: Optional[int] = None,
+        num_kv_heads: int | None = None,
         prefix: str = "",
         **extra_impl_args,
     ) -> None:

fastvideo/v1/attention/backends/sage_attn.py

@@ -1,5 +1,4 @@
 # SPDX-License-Identifier: Apache-2.0
-from typing import List, Optional, Type
 
 import torch
 from sageattention import sageattn
@@ -18,7 +17,7 @@ class SageAttentionBackend(AttentionBackend):
     accept_output_buffer: bool = True
 
     @staticmethod
-    def get_supported_head_sizes() -> List[int]:
+    def get_supported_head_sizes() -> list[int]:
         return [32, 64, 96, 128, 160, 192, 224, 256]
 
     @staticmethod
@@ -26,7 +25,7 @@ class SageAttentionBackend(AttentionBackend):
         return "SAGE_ATTN"
 
     @staticmethod
-    def get_impl_cls() -> Type["SageAttentionImpl"]:
+    def get_impl_cls() -> type["SageAttentionImpl"]:
         return SageAttentionImpl
 
     # @staticmethod
@@ -42,7 +41,7 @@ class SageAttentionImpl(AttentionImpl):
         head_size: int,
         causal: bool,
         softmax_scale: float,
-        num_kv_heads: Optional[int] = None,
+        num_kv_heads: int | None = None,
         prefix: str = "",
         **extra_impl_args,
     ) -> None:

fastvideo/v1/attention/backends/sdpa.py

@@ -1,5 +1,4 @@
 # SPDX-License-Identifier: Apache-2.0
-from typing import List, Optional, Type
 
 import torch
 
@@ -17,7 +16,7 @@ class SDPABackend(AttentionBackend):
     accept_output_buffer: bool = True
 
     @staticmethod
-    def get_supported_head_sizes() -> List[int]:
+    def get_supported_head_sizes() -> list[int]:
         return [32, 64, 96, 128, 160, 192, 224, 256]
 
     @staticmethod
@@ -25,7 +24,7 @@ class SDPABackend(AttentionBackend):
         return "SDPA"
 
     @staticmethod
-    def get_impl_cls() -> Type["SDPAImpl"]:
+    def get_impl_cls() -> type["SDPAImpl"]:
         return SDPAImpl
 
     # @staticmethod
@@ -41,7 +40,7 @@ class SDPAImpl(AttentionImpl):
         head_size: int,
         causal: bool,
         softmax_scale: float,
-        num_kv_heads: Optional[int] = None,
+        num_kv_heads: int | None = None,
         prefix: str = "",
         **extra_impl_args,
     ) -> None:

fastvideo/v1/attention/backends/sliding_tile_attn.py

@@ -1,7 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0
 import json
 from dataclasses import dataclass
-from typing import Any, List, Optional, Type
+from typing import Any
 
 import torch
 from einops import rearrange
@@ -40,7 +40,7 @@ class SlidingTileAttentionBackend(AttentionBackend):
     accept_output_buffer: bool = True
 
     @staticmethod
-    def get_supported_head_sizes() -> List[int]:
+    def get_supported_head_sizes() -> list[int]:
         # TODO(will-refactor): check this
         return [32, 64, 96, 128, 160, 192, 224, 256]
 
@@ -49,22 +49,22 @@ class SlidingTileAttentionBackend(AttentionBackend):
         return "SLIDING_TILE_ATTN"
 
     @staticmethod
-    def get_impl_cls() -> Type["SlidingTileAttentionImpl"]:
+    def get_impl_cls() -> type["SlidingTileAttentionImpl"]:
         return SlidingTileAttentionImpl
 
     @staticmethod
-    def get_metadata_cls() -> Type["SlidingTileAttentionMetadata"]:
+    def get_metadata_cls() -> type["SlidingTileAttentionMetadata"]:
         return SlidingTileAttentionMetadata
 
     @staticmethod
-    def get_builder_cls() -> Type["SlidingTileAttentionMetadataBuilder"]:
+    def get_builder_cls() -> type["SlidingTileAttentionMetadataBuilder"]:
         return SlidingTileAttentionMetadataBuilder
 
 
 @dataclass
 class SlidingTileAttentionMetadata(AttentionMetadata):
     current_timestep: int
-    STA_param: List[List[
+    STA_param: list[list[
         Any]]  # each timestep with one metadata, shape [num_layers, num_heads]
 
 
@@ -98,7 +98,7 @@ class SlidingTileAttentionImpl(AttentionImpl):
         head_size: int,
         causal: bool,
         softmax_scale: float,
-        num_kv_heads: Optional[int] = None,
+        num_kv_heads: int | None = None,
         prefix: str = "",
         **extra_impl_args,
     ) -> None:

fastvideo/v1/attention/backends/video_sparse_attn.py

@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 import math
 from dataclasses import dataclass
-from typing import List, Optional, Type
 
 import torch
 from einops import rearrange
@@ -28,7 +27,7 @@ class VideoSparseAttentionBackend(AttentionBackend):
     accept_output_buffer: bool = True
 
     @staticmethod
-    def get_supported_head_sizes() -> List[int]:
+    def get_supported_head_sizes() -> list[int]:
         return [64, 128]
 
     @staticmethod
@@ -36,22 +35,22 @@ class VideoSparseAttentionBackend(AttentionBackend):
         return "VIDEO_SPARSE_ATTN"
 
     @staticmethod
-    def get_impl_cls() -> Type["VideoSparseAttentionImpl"]:
+    def get_impl_cls() -> type["VideoSparseAttentionImpl"]:
         return VideoSparseAttentionImpl
 
     @staticmethod
-    def get_metadata_cls() -> Type["VideoSparseAttentionMetadata"]:
+    def get_metadata_cls() -> type["VideoSparseAttentionMetadata"]:
         return VideoSparseAttentionMetadata
 
     @staticmethod
-    def get_builder_cls() -> Type["VideoSparseAttentionMetadataBuilder"]:
+    def get_builder_cls() -> type["VideoSparseAttentionMetadataBuilder"]:
         return VideoSparseAttentionMetadataBuilder
 
 
 @dataclass
 class VideoSparseAttentionMetadata(AttentionMetadata):
     current_timestep: int
-    dit_seq_shape: List[int]
+    dit_seq_shape: list[int]
     VSA_sparsity: float
 
 
@@ -97,7 +96,7 @@ class VideoSparseAttentionImpl(AttentionImpl):
         head_size: int,
         causal: bool,
         softmax_scale: float,
-        num_kv_heads: Optional[int] = None,
+        num_kv_heads: int | None = None,
         prefix: str = "",
         **extra_impl_args,
     ) -> None:
@@ -105,8 +104,8 @@ class VideoSparseAttentionImpl(AttentionImpl):
         sp_group = get_sp_group()
         self.sp_size = sp_group.world_size
         self.VSA_base_tile_size = [4, 4, 4]
-        self.dit_seq_shape: List[int]
-        self.full_window_size: List[int]
+        self.dit_seq_shape: list[int]
+        self.full_window_size: list[int]
         self.img_seq_length: int
 
     def tile(self, x: torch.Tensor) -> torch.Tensor:

fastvideo/v1/attention/layer.py

@@ -1,7 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
 
-from typing import Optional, Tuple
-
 import torch
 import torch.nn as nn
 
@@ -23,11 +21,11 @@ class DistributedAttention(nn.Module):
     def __init__(self,
                  num_heads: int,
                  head_size: int,
-                 num_kv_heads: Optional[int] = None,
-                 softmax_scale: Optional[float] = None,
+                 num_kv_heads: int | None = None,
+                 softmax_scale: float | None = None,
                  causal: bool = False,
-                 supported_attention_backends: Optional[Tuple[
-                     AttentionBackendEnum, ...]] = None,
+                 supported_attention_backends: tuple[AttentionBackendEnum, ...]
+                 | None = None,
                  prefix: str = "",
                  **extra_impl_args) -> None:
         super().__init__()
@@ -63,10 +61,10 @@ class DistributedAttention(nn.Module):
         q: torch.Tensor,
         k: torch.Tensor,
         v: torch.Tensor,
-        replicated_q: Optional[torch.Tensor] = None,
-        replicated_k: Optional[torch.Tensor] = None,
-        replicated_v: Optional[torch.Tensor] = None,
-    ) -> tuple[torch.Tensor, Optional[torch.Tensor]]:
+        replicated_q: torch.Tensor | None = None,
+        replicated_k: torch.Tensor | None = None,
+        replicated_v: torch.Tensor | None = None,
+    ) -> tuple[torch.Tensor, torch.Tensor | None]:
         """Forward pass for distributed attention.
         
         Args:
@@ -144,11 +142,11 @@ class DistributedAttention_VSA(DistributedAttention):
         q: torch.Tensor,
         k: torch.Tensor,
         v: torch.Tensor,
-        replicated_q: Optional[torch.Tensor] = None,
-        replicated_k: Optional[torch.Tensor] = None,
-        replicated_v: Optional[torch.Tensor] = None,
-        gate_compress: Optional[torch.Tensor] = None,
-    ) -> tuple[torch.Tensor, Optional[torch.Tensor]]:
+        replicated_q: torch.Tensor | None = None,
+        replicated_k: torch.Tensor | None = None,
+        replicated_v: torch.Tensor | None = None,
+        gate_compress: torch.Tensor | None = None,
+    ) -> tuple[torch.Tensor, torch.Tensor | None]:
         """Forward pass for distributed attention.
         
         Args:
@@ -208,11 +206,11 @@ class LocalAttention(nn.Module):
     def __init__(self,
                  num_heads: int,
                  head_size: int,
-                 num_kv_heads: Optional[int] = None,
-                 softmax_scale: Optional[float] = None,
+                 num_kv_heads: int | None = None,
+                 softmax_scale: float | None = None,
                  causal: bool = False,
-                 supported_attention_backends: Optional[Tuple[
-                     AttentionBackendEnum, ...]] = None,
+                 supported_attention_backends: tuple[AttentionBackendEnum, ...]
+                 | None = None,
                  **extra_impl_args) -> None:
         super().__init__()
         if softmax_scale is None:

fastvideo/v1/attention/selector.py

@@ -2,9 +2,10 @@
 # Adapted from vllm: https://github.com/vllm-project/vllm/blob/v0.7.3/vllm/attention/selector.py
 
 import os
+from collections.abc import Generator
 from contextlib import contextmanager
 from functools import cache
-from typing import Generator, Optional, Tuple, Type, cast
+from typing import cast
 
 import torch
 
@@ -17,7 +18,7 @@ from fastvideo.v1.utils import STR_BACKEND_ENV_VAR, resolve_obj_by_qualname
 logger = init_logger(__name__)
 
 
-def backend_name_to_enum(backend_name: str) -> Optional[AttentionBackendEnum]:
+def backend_name_to_enum(backend_name: str) -> AttentionBackendEnum | None:
     """
     Convert a string backend name to a _Backend enum value.
 
@@ -31,7 +32,7 @@ def backend_name_to_enum(backend_name: str) -> Optional[AttentionBackendEnum]:
           None
 
 
-def get_env_variable_attn_backend() -> Optional[AttentionBackendEnum]:
+def get_env_variable_attn_backend() -> AttentionBackendEnum | None:
     '''
     Get the backend override specified by the FastVideo attention
     backend environment variable, if one is specified.
@@ -53,11 +54,11 @@ def get_env_variable_attn_backend() -> Optional[AttentionBackendEnum]:
 #
 # THIS SELECTION TAKES PRECEDENCE OVER THE
 # FASTVIDEO ATTENTION BACKEND ENVIRONMENT VARIABLE
-forced_attn_backend: Optional[AttentionBackendEnum] = None
+forced_attn_backend: AttentionBackendEnum | None = None
 
 
 def global_force_attn_backend(
-        attn_backend: Optional[AttentionBackendEnum]) -> None:
+        attn_backend: AttentionBackendEnum | None) -> None:
     '''
     Force all attention operations to use a specified backend.
 
@@ -72,7 +73,7 @@ def global_force_attn_backend(
     forced_attn_backend = attn_backend
 
 
-def get_global_forced_attn_backend() -> Optional[AttentionBackendEnum]:
+def get_global_forced_attn_backend() -> AttentionBackendEnum | None:
     '''
     Get the currently-forced choice of attention backend,
     or None if auto-selection is currently enabled.
@@ -83,9 +84,9 @@ def get_global_forced_attn_backend() -> Optional[AttentionBackendEnum]:
 def get_attn_backend(
     head_size: int,
     dtype: torch.dtype,
-    supported_attention_backends: Optional[Tuple[AttentionBackendEnum,
-                                                 ...]] = None,
-) -> Type[AttentionBackend]:
+    supported_attention_backends: tuple[AttentionBackendEnum, ...]
+    | None = None,
+) -> type[AttentionBackend]:
     return _cached_get_attn_backend(head_size, dtype,
                                     supported_attention_backends)
 
@@ -94,9 +95,9 @@ def get_attn_backend(
 def _cached_get_attn_backend(
     head_size: int,
     dtype: torch.dtype,
-    supported_attention_backends: Optional[Tuple[AttentionBackendEnum,
-                                                 ...]] = None,
-) -> Type[AttentionBackend]:
+    supported_attention_backends: tuple[AttentionBackendEnum, ...]
+    | None = None,
+) -> type[AttentionBackend]:
     # Check whether a particular choice of backend was
     # previously forced.
     #
@@ -105,13 +106,13 @@ def _cached_get_attn_backend(
     if not supported_attention_backends:
         raise ValueError("supported_attention_backends is empty")
     selected_backend = None
-    backend_by_global_setting: Optional[AttentionBackendEnum] = (
+    backend_by_global_setting: AttentionBackendEnum | None = (
         get_global_forced_attn_backend())
     if backend_by_global_setting is not None:
         selected_backend = backend_by_global_setting
     else:
         # Check the environment variable and override if specified
-        backend_by_env_var: Optional[str] = envs.FASTVIDEO_ATTENTION_BACKEND
+        backend_by_env_var: str | None = envs.FASTVIDEO_ATTENTION_BACKEND
         if backend_by_env_var is not None:
             selected_backend = backend_name_to_enum(backend_by_env_var)
 
@@ -123,7 +124,7 @@ def _cached_get_attn_backend(
     if not attention_cls:
         raise ValueError(
             f"Invalid attention backend for {current_platform.device_name}")
-    return cast(Type[AttentionBackend], resolve_obj_by_qualname(attention_cls))
+    return cast(type[AttentionBackend], resolve_obj_by_qualname(attention_cls))
 
 
 @contextmanager

fastvideo/v1/configs/models/base.py

@@ -1,6 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 from dataclasses import dataclass, field, fields
-from typing import Any, Dict, List, Tuple
+from typing import Any
 
 from fastvideo.v1.logger import init_logger
 
@@ -12,7 +12,7 @@ logger = init_logger(__name__)
 # 3. Any field in ArchConfig is fixed upon initialization, and should be hidden away from users
 @dataclass
 class ArchConfig:
-    stacked_params_mapping: List[Tuple[str, str, str]] = field(
+    stacked_params_mapping: list[tuple[str, str, str]] = field(
         default_factory=list
     )  # mapping from huggingface weight names to custom names
 
@@ -44,7 +44,7 @@ class ModelConfig:
         self.__dict__.update(state)
 
     # This should be used only when loading from transformers/diffusers
-    def update_model_arch(self, source_model_dict: Dict[str, Any]) -> None:
+    def update_model_arch(self, source_model_dict: dict[str, Any]) -> None:
         arch_config = self.arch_config
         valid_fields = {f.name for f in fields(arch_config)}
 
@@ -58,7 +58,7 @@ class ModelConfig:
         if hasattr(arch_config, "__post_init__"):
             arch_config.__post_init__()
 
-    def update_model_config(self, source_model_dict: Dict[str, Any]) -> None:
+    def update_model_config(self, source_model_dict: dict[str, Any]) -> None:
         assert "arch_config" not in source_model_dict, "Source model config shouldn't contain arch_config."
 
         valid_fields = {f.name for f in fields(self)}

fastvideo/v1/configs/models/dits/base.py

@@ -1,6 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 from dataclasses import dataclass, field
-from typing import Any, List, Optional, Tuple
+from typing import Any
 
 from fastvideo.v1.configs.models.base import ArchConfig, ModelConfig
 from fastvideo.v1.layers.quantization import QuantizationConfig
@@ -14,7 +14,7 @@ class DiTArchConfig(ArchConfig):
     _param_names_mapping: dict = field(default_factory=dict)
     _reverse_param_names_mapping: dict = field(default_factory=dict)
     _lora_param_names_mapping: dict = field(default_factory=dict)
-    _supported_attention_backends: Tuple[AttentionBackendEnum, ...] = (
+    _supported_attention_backends: tuple[AttentionBackendEnum, ...] = (
         AttentionBackendEnum.SLIDING_TILE_ATTN, AttentionBackendEnum.SAGE_ATTN,
         AttentionBackendEnum.FLASH_ATTN, AttentionBackendEnum.TORCH_SDPA,
         AttentionBackendEnum.VIDEO_SPARSE_ATTN)
@@ -22,7 +22,7 @@ class DiTArchConfig(ArchConfig):
     hidden_size: int = 0
     num_attention_heads: int = 0
     num_channels_latents: int = 0
-    exclude_lora_layers: List[str] = field(default_factory=list)
+    exclude_lora_layers: list[str] = field(default_factory=list)
 
     def __post_init__(self) -> None:
         if not self._compile_conditions:
@@ -35,7 +35,7 @@ class DiTConfig(ModelConfig):
 
     # FastVideoDiT-specific parameters
     prefix: str = ""
-    quant_config: Optional[QuantizationConfig] = None
+    quant_config: QuantizationConfig | None = None
 
     @staticmethod
     def add_cli_args(parser: Any, prefix: str = "dit-config") -> Any:

fastvideo/v1/configs/models/dits/hunyuanvideo.py

@@ -1,6 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
 from dataclasses import dataclass, field
-from typing import List, Optional, Tuple
 
 import torch
 
@@ -160,14 +159,14 @@ class HunyuanVideoArchConfig(DiTArchConfig):
     num_layers: int = 20
     num_single_layers: int = 40
     num_refiner_layers: int = 2
-    rope_axes_dim: Tuple[int, int, int] = (16, 56, 56)
+    rope_axes_dim: tuple[int, int, int] = (16, 56, 56)
     guidance_embeds: bool = False
-    dtype: Optional[torch.dtype] = None
+    dtype: torch.dtype | None = None
     text_embed_dim: int = 4096
     pooled_projection_dim: int = 768
     rope_theta: int = 256
     qk_norm: str = "rms_norm"
-    exclude_lora_layers: List[str] = field(
+    exclude_lora_layers: list[str] = field(
         default_factory=lambda: ["img_in", "txt_in", "time_in", "vector_in"])
 
     def __post_init__(self):

fastvideo/v1/configs/models/dits/stepvideo.py

@@ -1,6 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
 from dataclasses import dataclass, field
-from typing import List, Optional, Tuple, Union
 
 from fastvideo.v1.configs.models.dits.base import DiTArchConfig, DiTConfig
 
@@ -39,18 +38,18 @@ class StepVideoArchConfig(DiTArchConfig):
     num_attention_heads: int = 48
     attention_head_dim: int = 128
     in_channels: int = 64
-    out_channels: Optional[int] = 64
+    out_channels: int | None = 64
     num_layers: int = 48
     dropout: float = 0.0
     patch_size: int = 1
     norm_type: str = "ada_norm_single"
     norm_elementwise_affine: bool = False
     norm_eps: float = 1e-6
-    caption_channels: Optional[Union[int, List[int], Tuple[int, ...]]] = field(
+    caption_channels: int | list[int] | tuple[int, ...] | None = field(
         default_factory=lambda: [6144, 1024])
-    attention_type: Optional[str] = "torch"
-    use_additional_conditions: Optional[bool] = False
-    exclude_lora_layers: List[str] = field(default_factory=lambda: [])
+    attention_type: str | None = "torch"
+    use_additional_conditions: bool | None = False
+    exclude_lora_layers: list[str] = field(default_factory=lambda: [])
 
     def __post_init__(self):
         self.hidden_size = self.num_attention_heads * self.attention_head_dim

fastvideo/v1/configs/models/dits/wanvideo.py

@@ -1,6 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
 from dataclasses import dataclass, field
-from typing import List, Optional, Tuple
 
 from fastvideo.v1.configs.models.dits.base import DiTArchConfig, DiTConfig
 
@@ -74,7 +73,7 @@ class WanVideoArchConfig(DiTArchConfig):
             r"^blocks\.(\d+)\.ffn\.2\.(.*)$": r"blocks.\1.ffn.fc_out.\2",
         })
 
-    patch_size: Tuple[int, int, int] = (1, 2, 2)
+    patch_size: tuple[int, int, int] = (1, 2, 2)
     text_len = 512
     num_attention_heads: int = 40
     attention_head_dim: int = 128
@@ -87,10 +86,10 @@ class WanVideoArchConfig(DiTArchConfig):
     cross_attn_norm: bool = True
     qk_norm: str = "rms_norm_across_heads"
     eps: float = 1e-6
-    image_dim: Optional[int] = None
-    added_kv_proj_dim: Optional[int] = None
+    image_dim: int | None = None
+    added_kv_proj_dim: int | None = None
     rope_max_seq_len: int = 1024
-    exclude_lora_layers: List[str] = field(default_factory=lambda: ["embedder"])
+    exclude_lora_layers: list[str] = field(default_factory=lambda: ["embedder"])
 
     def __post_init__(self):
         super().__post_init__()

fastvideo/v1/configs/models/encoders/base.py

@@ -1,6 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 from dataclasses import dataclass, field
-from typing import Any, Dict, List, Optional, Tuple
+from typing import Any
 
 import torch
 
@@ -11,8 +11,8 @@ from fastvideo.v1.platforms import AttentionBackendEnum
 
 @dataclass
 class EncoderArchConfig(ArchConfig):
-    architectures: List[str] = field(default_factory=lambda: [])
-    _supported_attention_backends: Tuple[AttentionBackendEnum, ...] = (
+    architectures: list[str] = field(default_factory=lambda: [])
+    _supported_attention_backends: tuple[AttentionBackendEnum, ...] = (
         AttentionBackendEnum.FLASH_ATTN, AttentionBackendEnum.TORCH_SDPA)
     output_hidden_states: bool = False
     use_return_dict: bool = True
@@ -32,10 +32,10 @@ class TextEncoderArchConfig(EncoderArchConfig):
     output_past: bool = True
     scalable_attention: bool = True
     tie_word_embeddings: bool = False
-    stacked_params_mapping: List[Tuple[str, str, str]] = field(
+    stacked_params_mapping: list[tuple[str, str, str]] = field(
         default_factory=list
     )  # mapping from huggingface weight names to custom names
-    tokenizer_kwargs: Dict[str, Any] = field(default_factory=dict)
+    tokenizer_kwargs: dict[str, Any] = field(default_factory=dict)
     _fsdp_shard_conditions: list = field(default_factory=lambda: [])
 
     def __post_init__(self) -> None:
@@ -53,11 +53,11 @@ class ImageEncoderArchConfig(EncoderArchConfig):
 
 @dataclass
 class BaseEncoderOutput:
-    last_hidden_state: Optional[torch.FloatTensor] = None
-    pooler_output: Optional[torch.FloatTensor] = None
-    hidden_states: Optional[Tuple[torch.FloatTensor, ...]] = None
-    attentions: Optional[Tuple[torch.FloatTensor, ...]] = None
-    attention_mask: Optional[torch.Tensor] = None
+    last_hidden_state: torch.FloatTensor | None = None
+    pooler_output: torch.FloatTensor | None = None
+    hidden_states: tuple[torch.FloatTensor, ...] | None = None
+    attentions: tuple[torch.FloatTensor, ...] | None = None
+    attention_mask: torch.Tensor | None = None
 
 
 @dataclass
@@ -65,8 +65,8 @@ class EncoderConfig(ModelConfig):
     arch_config: ArchConfig = field(default_factory=EncoderArchConfig)
 
     prefix: str = ""
-    quant_config: Optional[QuantizationConfig] = None
-    lora_config: Optional[Any] = None
+    quant_config: QuantizationConfig | None = None
+    lora_config: Any | None = None
 
 
 @dataclass

fastvideo/v1/configs/models/encoders/clip.py

@@ -1,6 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
 from dataclasses import dataclass, field
-from typing import List, Optional, Tuple
 
 from fastvideo.v1.configs.models.encoders.base import (ImageEncoderArchConfig,
                                                        ImageEncoderConfig,
@@ -35,7 +34,7 @@ class CLIPTextArchConfig(TextEncoderArchConfig):
     bos_token_id: int = 49406
     eos_token_id: int = 49407
     text_len: int = 77
-    stacked_params_mapping: List[Tuple[str, str,
+    stacked_params_mapping: list[tuple[str, str,
                                        str]] = field(default_factory=lambda: [
                                            # (param_name, shard_name, shard_id)
                                            ("qkv_proj", "q_proj", "q"),
@@ -62,7 +61,7 @@ class CLIPVisionArchConfig(ImageEncoderArchConfig):
     attention_dropout: float = 0.0
     initializer_range: float = 0.02
     initializer_factor: float = 1.0
-    stacked_params_mapping: List[Tuple[str, str,
+    stacked_params_mapping: list[tuple[str, str,
                                        str]] = field(default_factory=lambda: [
                                            # (param_name, shard_name, shard_id)
                                            ("qkv_proj", "q_proj", "q"),
@@ -76,8 +75,8 @@ class CLIPTextConfig(TextEncoderConfig):
     arch_config: TextEncoderArchConfig = field(
         default_factory=CLIPTextArchConfig)
 
-    num_hidden_layers_override: Optional[int] = None
-    require_post_norm: Optional[bool] = None
+    num_hidden_layers_override: int | None = None
+    require_post_norm: bool | None = None
     prefix: str = "clip"
 
 
@@ -86,6 +85,6 @@ class CLIPVisionConfig(ImageEncoderConfig):
     arch_config: ImageEncoderArchConfig = field(
         default_factory=CLIPVisionArchConfig)
 
-    num_hidden_layers_override: Optional[int] = None
-    require_post_norm: Optional[bool] = None
+    num_hidden_layers_override: int | None = None
+    require_post_norm: bool | None = None
     prefix: str = "clip"

fastvideo/v1/configs/models/encoders/llama.py

@@ -1,6 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
 from dataclasses import dataclass, field
-from typing import List, Optional, Tuple
 
 from fastvideo.v1.configs.models.encoders.base import (TextEncoderArchConfig,
                                                        TextEncoderConfig)
@@ -25,7 +24,7 @@ class LlamaArchConfig(TextEncoderArchConfig):
     intermediate_size: int = 11008
     num_hidden_layers: int = 32
     num_attention_heads: int = 32
-    num_key_value_heads: Optional[int] = None
+    num_key_value_heads: int | None = None
     hidden_act: str = "silu"
     max_position_embeddings: int = 2048
     initializer_range: float = 0.02
@@ -37,14 +36,14 @@ class LlamaArchConfig(TextEncoderArchConfig):
     pretraining_tp: int = 1
     tie_word_embeddings: bool = False
     rope_theta: float = 10000.0
-    rope_scaling: Optional[float] = None
+    rope_scaling: float | None = None
     attention_bias: bool = False
     attention_dropout: float = 0.0
     mlp_bias: bool = False
-    head_dim: Optional[int] = None
+    head_dim: int | None = None
     hidden_state_skip_layer: int = 2
     text_len: int = 256
-    stacked_params_mapping: List[Tuple[str, str, str]] = field(
+    stacked_params_mapping: list[tuple[str, str, str]] = field(
         default_factory=lambda: [
             # (param_name, shard_name, shard_id)
             (".qkv_proj", ".q_proj", "q"),

fastvideo/v1/configs/models/encoders/t5.py

@@ -1,6 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
 from dataclasses import dataclass, field
-from typing import List, Optional, Tuple
 
 from fastvideo.v1.configs.models.encoders.base import (TextEncoderArchConfig,
                                                        TextEncoderConfig)
@@ -25,7 +24,7 @@ class T5ArchConfig(TextEncoderArchConfig):
     d_kv: int = 64
     d_ff: int = 2048
     num_layers: int = 6
-    num_decoder_layers: Optional[int] = None
+    num_decoder_layers: int | None = None
     num_heads: int = 8
     relative_attention_num_buckets: int = 32
     relative_attention_max_distance: int = 128
@@ -41,7 +40,7 @@ class T5ArchConfig(TextEncoderArchConfig):
     eos_token_id: int = 1
     classifier_dropout: float = 0.0
     text_len: int = 512
-    stacked_params_mapping: List[Tuple[str, str,
+    stacked_params_mapping: list[tuple[str, str,
                                        str]] = field(default_factory=lambda: [
                                            # (param_name, shard_name, shard_id)
                                            (".qkv_proj", ".q", "q"),

fastvideo/v1/configs/models/vaes/base.py

@@ -2,7 +2,7 @@
 import argparse
 import dataclasses
 from dataclasses import dataclass, field
-from typing import Any, Union
+from typing import Any
 
 import torch
 
@@ -12,7 +12,7 @@ from fastvideo.v1.utils import StoreBoolean
 
 @dataclass
 class VAEArchConfig(ArchConfig):
-    scaling_factor: Union[float, torch.tensor] = 0
+    scaling_factor: float | torch.Tensor = 0
 
     temporal_compression_ratio: int = 4
     spatial_compression_ratio: int = 8

fastvideo/v1/configs/models/vaes/hunyuanvae.py

@@ -1,6 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
 from dataclasses import dataclass, field
-from typing import Tuple
 
 from fastvideo.v1.configs.models.vaes.base import VAEArchConfig, VAEConfig
 
@@ -10,19 +9,19 @@ class HunyuanVAEArchConfig(VAEArchConfig):
     in_channels: int = 3
     out_channels: int = 3
     latent_channels: int = 16
-    down_block_types: Tuple[str, ...] = (
+    down_block_types: tuple[str, ...] = (
         "HunyuanVideoDownBlock3D",
         "HunyuanVideoDownBlock3D",
         "HunyuanVideoDownBlock3D",
         "HunyuanVideoDownBlock3D",
     )
-    up_block_types: Tuple[str, ...] = (
+    up_block_types: tuple[str, ...] = (
         "HunyuanVideoUpBlock3D",
         "HunyuanVideoUpBlock3D",
         "HunyuanVideoUpBlock3D",
         "HunyuanVideoUpBlock3D",
     )
-    block_out_channels: Tuple[int, ...] = (128, 256, 512, 512)
+    block_out_channels: tuple[int, ...] = (128, 256, 512, 512)
     layers_per_block: int = 2
     act_fn: str = "silu"
     norm_num_groups: int = 32

fastvideo/v1/configs/models/vaes/wanvae.py

@@ -1,6 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
 from dataclasses import dataclass, field
-from typing import Tuple
 
 import torch
 
@@ -11,12 +10,12 @@ from fastvideo.v1.configs.models.vaes.base import VAEArchConfig, VAEConfig
 class WanVAEArchConfig(VAEArchConfig):
     base_dim: int = 96
     z_dim: int = 16
-    dim_mult: Tuple[int, ...] = (1, 2, 4, 4)
+    dim_mult: tuple[int, ...] = (1, 2, 4, 4)
     num_res_blocks: int = 2
-    attn_scales: Tuple[float, ...] = ()
-    temperal_downsample: Tuple[bool, ...] = (False, True, True)
+    attn_scales: tuple[float, ...] = ()
+    temperal_downsample: tuple[bool, ...] = (False, True, True)
     dropout: float = 0.0
-    latents_mean: Tuple[float, ...] = (
+    latents_mean: tuple[float, ...] = (
         -0.7571,
         -0.7089,
         -0.9113,
@@ -34,7 +33,7 @@ class WanVAEArchConfig(VAEArchConfig):
         0.2503,
         -0.2921,
     )
-    latents_std: Tuple[float, ...] = (
+    latents_std: tuple[float, ...] = (
         2.8184,
         1.4541,
         2.3275,

fastvideo/v1/configs/pipelines/base.py

@@ -1,8 +1,9 @@
 # SPDX-License-Identifier: Apache-2.0
 import json
+from collections.abc import Callable
 from dataclasses import asdict, dataclass, field, fields
 from enum import Enum
-from typing import Any, Callable, Dict, List, Optional, Tuple, Union, cast
+from typing import Any, cast
 
 import torch
 
@@ -39,11 +40,11 @@ def postprocess_text(output: BaseEncoderOutput) -> torch.tensor:
 class PipelineConfig:
     """Base configuration for all pipeline architectures."""
     model_path: str = ""
-    pipeline_config_path: Optional[str] = None
+    pipeline_config_path: str | None = None
 
     # Video generation parameters
     embedded_cfg_scale: float = 6.0
-    flow_shift: Optional[float] = None
+    flow_shift: float | None = None
     disable_autocast: bool = False
 
     # Model configuration
@@ -62,30 +63,29 @@ class PipelineConfig:
 
     # Text encoder configuration
     DEFAULT_TEXT_ENCODER_PRECISIONS = ("fp16", )
-    text_encoder_configs: Tuple[EncoderConfig, ...] = field(
+    text_encoder_configs: tuple[EncoderConfig, ...] = field(
         default_factory=lambda: (EncoderConfig(), ))
-    text_encoder_precisions: Tuple[str, ...] = field(
+    text_encoder_precisions: tuple[str, ...] = field(
         default_factory=lambda: ("fp16", ))
-    preprocess_text_funcs: Tuple[Callable[[str], str], ...] = field(
+    preprocess_text_funcs: tuple[Callable[[str], str], ...] = field(
         default_factory=lambda: (preprocess_text, ))
-    postprocess_text_funcs: Tuple[Callable[[BaseEncoderOutput], torch.tensor],
+    postprocess_text_funcs: tuple[Callable[[BaseEncoderOutput], torch.tensor],
                                   ...] = field(default_factory=lambda:
                                                (postprocess_text, ))
 
     # LoRA parameters
-    lora_path: Optional[str] = None
-    lora_nickname: Optional[
-        str] = "default"  # for swapping adapters in the pipeline
-    lora_target_names: Optional[List[
-        str]] = None  # can restrict list of layers to adapt, e.g. ["q_proj"]
+    lora_path: str | None = None
+    lora_nickname: str | None = "default"  # for swapping adapters in the pipeline
+    lora_target_names: list[
+        str] | None = None  # can restrict list of layers to adapt, e.g. ["q_proj"]
 
     # StepVideo specific parameters
-    pos_magic: Optional[str] = None
-    neg_magic: Optional[str] = None
-    timesteps_scale: Optional[bool] = None
+    pos_magic: str | None = None
+    neg_magic: str | None = None
+    timesteps_scale: bool | None = None
 
     # STA (Sliding Tile Attention) parameters
-    mask_strategy_file_path: Optional[str] = None
+    mask_strategy_file_path: str | None = None
     STA_mode: STA_Mode = STA_Mode.STA_INFERENCE
     skip_time_steps: int = 15
 
@@ -217,7 +217,7 @@ class PipelineConfig:
         return parser
 
     def update_config_from_dict(self,
-                                args: Dict[str, Any],
+                                args: dict[str, Any],
                                 prefix: str = "") -> None:
         prefix_with_dot = f"{prefix}." if (prefix.strip() != "") else ""
         update_config_from_args(self, args, prefix, pop_args=True)
@@ -243,7 +243,7 @@ class PipelineConfig:
 
     @classmethod
     def from_kwargs(cls,
-                    kwargs: Dict[str, Any],
+                    kwargs: dict[str, Any],
                     config_cli_prefix: str = "") -> "PipelineConfig":
         """
         Load PipelineConfig from kwargs Dictionary.
@@ -255,11 +255,11 @@ class PipelineConfig:
 
         prefix_with_dot = f"{config_cli_prefix}." if (config_cli_prefix.strip()
                                                       != "") else ""
-        model_path: Optional[str] = kwargs.get(prefix_with_dot + 'model_path',
-                                               None) or kwargs.get('model_path')
-        pipeline_config_or_path: Optional[Union[str, PipelineConfig, Dict[
-            str, Any]]] = kwargs.get(prefix_with_dot + 'pipeline_config',
-                                     None) or kwargs.get('pipeline_config')
+        model_path: str | None = kwargs.get(prefix_with_dot + 'model_path',
+                                            None) or kwargs.get('model_path')
+        pipeline_config_or_path: str | PipelineConfig | dict[
+            str, Any] | None = kwargs.get(prefix_with_dot + 'pipeline_config',
+                                          None) or kwargs.get('pipeline_config')
         if model_path is None:
             raise ValueError("model_path is required in kwargs")
 
@@ -344,7 +344,7 @@ class PipelineConfig:
             input_pipeline_dict = json.load(f)
         self.update_pipeline_config(input_pipeline_dict)
 
-    def update_pipeline_config(self, source_pipeline_dict: Dict[str,
+    def update_pipeline_config(self, source_pipeline_dict: dict[str,
                                                                 Any]) -> None:
         for f in fields(self):
             key = f.name
@@ -360,8 +360,9 @@ class PipelineConfig:
                     assert len(current_value) == len(
                         new_value
                     ), "Users shouldn't delete or add text encoder config objects in your json"
-                    for target_config, source_config in zip(
-                            current_value, new_value):
+                    for target_config, source_config in zip(current_value,
+                                                            new_value,
+                                                            strict=True):
                         target_config.update_model_config(source_config)
                 else:
                     setattr(self, key, new_value)

fastvideo/v1/configs/pipelines/hunyuan.py

@@ -1,6 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0
+from collections.abc import Callable
 from dataclasses import dataclass, field
-from typing import Callable, Tuple, TypedDict
+from typing import TypedDict
 
 import torch
 
@@ -39,7 +40,7 @@ def llama_preprocess_text(prompt: str) -> str:
 def llama_postprocess_text(outputs: BaseEncoderOutput) -> torch.tensor:
     hidden_state_skip_layer = 2
     assert outputs.hidden_states is not None
-    hidden_states: Tuple[torch.Tensor, ...] = outputs.hidden_states
+    hidden_states: tuple[torch.Tensor, ...] = outputs.hidden_states
     last_hidden_state: torch.tensor = hidden_states[-(hidden_state_skip_layer +
                                                       1)]
     crop_start = prompt_template_video.get("crop_start", -1)
@@ -70,11 +71,11 @@ class HunyuanConfig(PipelineConfig):
     flow_shift: int = 7
 
     # Text encoding stage
-    text_encoder_configs: Tuple[EncoderConfig, ...] = field(
+    text_encoder_configs: tuple[EncoderConfig, ...] = field(
         default_factory=lambda: (LlamaConfig(), CLIPTextConfig()))
-    preprocess_text_funcs: Tuple[Callable[[str], str], ...] = field(
+    preprocess_text_funcs: tuple[Callable[[str], str], ...] = field(
         default_factory=lambda: (llama_preprocess_text, clip_preprocess_text))
-    postprocess_text_funcs: Tuple[
+    postprocess_text_funcs: tuple[
         Callable[[BaseEncoderOutput], torch.tensor],
         ...] = field(default_factory=lambda:
                      (llama_postprocess_text, clip_postprocess_text))
@@ -82,7 +83,7 @@ class HunyuanConfig(PipelineConfig):
     # Precision for each component
     dit_precision: str = "bf16"
     vae_precision: str = "fp16"
-    text_encoder_precisions: Tuple[str, ...] = field(
+    text_encoder_precisions: tuple[str, ...] = field(
         default_factory=lambda: ("fp16", "fp16"))
 
     def __post_init__(self):

fastvideo/v1/configs/pipelines/registry.py

@@ -2,7 +2,7 @@
 """Registry for pipeline weight-specific configurations."""
 
 import os
-from typing import Callable, Dict, Optional, Type
+from collections.abc import Callable
 
 from fastvideo.v1.configs.pipelines.base import PipelineConfig
 from fastvideo.v1.configs.pipelines.hunyuan import (FastHunyuanConfig,
@@ -19,7 +19,7 @@ from fastvideo.v1.utils import (maybe_download_model_index,
 logger = init_logger(__name__)
 
 # Registry maps specific model weights to their config classes
-PIPE_NAME_TO_CONFIG: Dict[str, Type[PipelineConfig]] = {
+PIPE_NAME_TO_CONFIG: dict[str, type[PipelineConfig]] = {
     "FastVideo/FastHunyuan-diffusers": FastHunyuanConfig,
     "hunyuanvideo-community/HunyuanVideo": HunyuanConfig,
     "Wan-AI/Wan2.1-T2V-1.3B-Diffusers": WanT2V480PConfig,
@@ -32,7 +32,7 @@ PIPE_NAME_TO_CONFIG: Dict[str, Type[PipelineConfig]] = {
 }
 
 # For determining pipeline type from model ID
-PIPELINE_DETECTOR: Dict[str, Callable[[str], bool]] = {
+PIPELINE_DETECTOR: dict[str, Callable[[str], bool]] = {
     "hunyuan": lambda id: "hunyuan" in id.lower(),
     "wanpipeline": lambda id: "wanpipeline" in id.lower(),
     "wanimagetovideo": lambda id: "wanimagetovideo" in id.lower(),
@@ -41,7 +41,7 @@ PIPELINE_DETECTOR: Dict[str, Callable[[str], bool]] = {
 }
 
 # Fallback configs when exact match isn't found but architecture is detected
-PIPELINE_FALLBACK_CONFIG: Dict[str, Type[PipelineConfig]] = {
+PIPELINE_FALLBACK_CONFIG: dict[str, type[PipelineConfig]] = {
     "hunyuan":
     HunyuanConfig,  # Base Hunyuan config as fallback for any Hunyuan variant
     "wanpipeline":
@@ -53,7 +53,7 @@ PIPELINE_FALLBACK_CONFIG: Dict[str, Type[PipelineConfig]] = {
 
 
 def get_pipeline_config_cls_from_name(
-        pipeline_name_or_path: str) -> Type[PipelineConfig]:
+        pipeline_name_or_path: str) -> type[PipelineConfig]:
     """Get the appropriate configuration class for a given pipeline name or path.
 
     This function implements a multi-step lookup process to find the most suitable
@@ -82,7 +82,7 @@ def get_pipeline_config_cls_from_name(
         - Warning messages are logged when falling back to less specific configurations
     """
 
-    pipeline_config_cls: Optional[Type[PipelineConfig]] = None
+    pipeline_config_cls: type[PipelineConfig] | None = None
 
     # First try exact match for specific weights
     if pipeline_name_or_path in PIPE_NAME_TO_CONFIG:

fastvideo/v1/configs/pipelines/wan.py

@@ -1,6 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
+from collections.abc import Callable
 from dataclasses import dataclass, field
-from typing import Callable, Tuple
 
 import torch
 
@@ -17,7 +17,7 @@ def t5_postprocess_text(outputs: BaseEncoderOutput) -> torch.tensor:
     hidden_state: torch.tensor = outputs.last_hidden_state
     seq_lens = mask.gt(0).sum(dim=1).long()
     assert torch.isnan(hidden_state).sum() == 0
-    prompt_embeds = [u[:v] for u, v in zip(hidden_state, seq_lens)]
+    prompt_embeds = [u[:v] for u, v in zip(hidden_state, seq_lens, strict=True)]
     prompt_embeds_tensor: torch.tensor = torch.stack([
         torch.cat([u, u.new_zeros(512 - u.size(0), u.size(1))])
         for u in prompt_embeds
@@ -42,16 +42,16 @@ class WanT2V480PConfig(PipelineConfig):
     flow_shift: int = 3
 
     # Text encoding stage
-    text_encoder_configs: Tuple[EncoderConfig, ...] = field(
+    text_encoder_configs: tuple[EncoderConfig, ...] = field(
         default_factory=lambda: (T5Config(), ))
-    postprocess_text_funcs: Tuple[Callable[[BaseEncoderOutput], torch.tensor],
+    postprocess_text_funcs: tuple[Callable[[BaseEncoderOutput], torch.tensor],
                                   ...] = field(default_factory=lambda:
                                                (t5_postprocess_text, ))
 
     # Precision for each component
     precision: str = "bf16"
     vae_precision: str = "fp32"
-    text_encoder_precisions: Tuple[str, ...] = field(
+    text_encoder_precisions: tuple[str, ...] = field(
         default_factory=lambda: ("fp32", ))
 
     # WanConfig-specific added parameters

fastvideo/v1/configs/sample/base.py

@@ -1,6 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 from dataclasses import dataclass
-from typing import Any, Dict, List, Optional, Union
+from typing import Any
 
 from fastvideo.v1.logger import init_logger
 
@@ -16,12 +16,12 @@ class SamplingParam:
     data_type: str = "video"
 
     # Image inputs
-    image_path: Optional[str] = None
+    image_path: str | None = None
 
     # Text inputs
-    prompt: Optional[Union[str, List[str]]] = None
-    negative_prompt: Optional[str] = None
-    prompt_path: Optional[str] = None
+    prompt: str | list[str] | None = None
+    negative_prompt: str | None = None
+    prompt_path: str | None = None
     output_path: str = "outputs/"
 
     # Batch info
@@ -54,7 +54,7 @@ class SamplingParam:
         if self.prompt_path and not self.prompt_path.endswith(".txt"):
             raise ValueError("prompt_path must be a txt file")
 
-    def update(self, source_dict: Dict[str, Any]) -> None:
+    def update(self, source_dict: dict[str, Any]) -> None:
         for key, value in source_dict.items():
             if hasattr(self, key):
                 setattr(self, key, value)

fastvideo/v1/configs/sample/registry.py

@@ -1,6 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0
 import os
-from typing import Any, Callable, Dict, Optional
+from collections.abc import Callable
+from typing import Any
 
 from fastvideo.v1.configs.sample.hunyuan import (FastHunyuanSamplingParam,
                                                  HunyuanSamplingParam)
@@ -15,7 +16,7 @@ from fastvideo.v1.utils import (maybe_download_model_index,
 
 logger = init_logger(__name__)
 # Registry maps specific model weights to their config classes
-SAMPLING_PARAM_REGISTRY: Dict[str, Any] = {
+SAMPLING_PARAM_REGISTRY: dict[str, Any] = {
     "FastVideo/FastHunyuan-diffusers": FastHunyuanSamplingParam,
     "hunyuanvideo-community/HunyuanVideo": HunyuanSamplingParam,
     "Wan-AI/Wan2.1-T2V-1.3B-Diffusers": WanT2V_1_3B_SamplingParam,
@@ -27,7 +28,7 @@ SAMPLING_PARAM_REGISTRY: Dict[str, Any] = {
 }
 
 # For determining pipeline type from model ID
-SAMPLING_PARAM_DETECTOR: Dict[str, Callable[[str], bool]] = {
+SAMPLING_PARAM_DETECTOR: dict[str, Callable[[str], bool]] = {
     "hunyuan": lambda id: "hunyuan" in id.lower(),
     "wanpipeline": lambda id: "wanpipeline" in id.lower(),
     "wanimagetovideo": lambda id: "wanimagetovideo" in id.lower(),
@@ -36,7 +37,7 @@ SAMPLING_PARAM_DETECTOR: Dict[str, Callable[[str], bool]] = {
 }
 
 # Fallback configs when exact match isn't found but architecture is detected
-SAMPLING_FALLBACK_PARAM: Dict[str, Any] = {
+SAMPLING_FALLBACK_PARAM: dict[str, Any] = {
     "hunyuan":
     HunyuanSamplingParam,  # Base Hunyuan config as fallback for any Hunyuan variant
     "wanpipeline":
@@ -47,8 +48,7 @@ SAMPLING_FALLBACK_PARAM: Dict[str, Any] = {
 }
 
 
-def get_sampling_param_cls_for_name(
-        pipeline_name_or_path: str) -> Optional[Any]:
+def get_sampling_param_cls_for_name(pipeline_name_or_path: str) -> Any | None:
     """Get the appropriate sampling param for specific pretrained weights."""
 
     if os.path.exists(pipeline_name_or_path):

fastvideo/v1/configs/utils.py

@@ -1,8 +1,8 @@
-from typing import Any, Dict
+from typing import Any
 
 
 def update_config_from_args(config: Any,
-                            args_dict: Dict[str, Any],
+                            args_dict: dict[str, Any],
                             prefix: str = "",
                             pop_args: bool = False) -> None:
     """

fastvideo/v1/dataset/benchmarks/benchmark_parquet_dataset_iterable_style.py

@@ -128,8 +128,8 @@ def main() -> None:
 
         # Verify all sums match
         if all(
-                abs(a - b) < 1e-6
-                for a, b in zip(first_pass_sums, second_pass_sums)):
+                abs(a - b) < 1e-6 for a, b in zip(
+                    first_pass_sums, second_pass_sums, strict=True)):
             logger.info(
                 "All latent sums match between passes - resume verification successful!"
             )

fastvideo/v1/dataset/benchmarks/benchmark_parquet_dataset_map_style.py

@@ -144,8 +144,8 @@ def main() -> None:
 
         # Verify all sums match
         if all(
-                abs(a - b) < 1e-6
-                for a, b in zip(first_pass_sums, second_pass_sums)):
+                abs(a - b) < 1e-6 for a, b in zip(
+                    first_pass_sums, second_pass_sums, strict=True)):
             logger.info(
                 "All latent sums match between passes - resume verification successful!"
             )

fastvideo/v1/dataset/latent_datasets.py

@@ -76,7 +76,7 @@ def latent_collate_function(batch):
     # latent_attn_mask: # b t h w
     # text_attn_mask: b 1 l
     # needs to check if the latent/prompt' size and apply padding & attn mask
-    latents, prompt_embeds, prompt_attention_masks = zip(*batch)
+    latents, prompt_embeds, prompt_attention_masks = zip(*batch, strict=True)
     # calculate max shape
     max_t = max([latent.shape[1] for latent in latents])
     max_h = max([latent.shape[2] for latent in latents])

fastvideo/v1/dataset/parquet_dataset_iterable_style.py

@@ -1,7 +1,6 @@
 import os
 import pickle
 import random
-from typing import Dict, List, Tuple
 
 import numpy as np
 import pyarrow as pa
@@ -139,7 +138,7 @@ def shard_parquet_files_across_sp_groups_and_workers(
     num_sp_groups: int,
     num_workers: int,
     seed: int = 42,
-) -> Tuple[List[List[str]], List[int], List[Dict[str, int]]]:
+) -> tuple[list[list[str]], list[int], list[dict[str, int]]]:
     """
     Shard parquet files across SP groups and workers in a balanced way.
     
@@ -217,7 +216,8 @@ def shard_parquet_files_across_sp_groups_and_workers(
         # Distribute files to shards using a greedy approach
         logger.info("Distributing files to shards...")
         for file, length in zip(reversed(sorted_files),
-                                reversed(sorted_lengths)):
+                                reversed(sorted_lengths),
+                                strict=True):
             # Find shard with minimum current length
             target_shard = np.argmin(shard_total_samples)
             shard_parquet_files[target_shard].append(file)
@@ -257,7 +257,7 @@ def build_parquet_iterable_style_dataloader(
     text_padding_length: int = 512,
     seed: int = 42,
     read_batch_size: int = 32
-) -> Tuple[LatentsParquetIterStyleDataset, StatefulDataLoader]:
+) -> tuple[LatentsParquetIterStyleDataset, StatefulDataLoader]:
     """Build a dataloader for the LatentsParquetIterStyleDataset."""
     dataset = LatentsParquetIterStyleDataset(
         path=path,

fastvideo/v1/dataset/parquet_dataset_map_style.py

@@ -1,7 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0
 import os
 import pickle
-from typing import Any, Dict, List, Tuple
+from typing import Any
 
 import pyarrow as pa
 import pyarrow.parquet as pq
@@ -20,7 +20,7 @@ from fastvideo.v1.logger import init_logger
 logger = init_logger(__name__)
 
 
-class DP_SP_BatchSampler(Sampler[List[int]]):
+class DP_SP_BatchSampler(Sampler[list[int]]):
     """
     A simple sequential batch sampler that yields batches of indices.
     """
@@ -121,8 +121,11 @@ def get_parquet_files_and_length(path: str):
             num_rows = pq.ParquetFile(file_path).metadata.num_rows
             lengths.append(num_rows)
         # sort according to file name to ensure all rank has the same order (in case os.walk is not sorted)
-        file_names_sorted, lengths_sorted = zip(
-            *sorted(zip(file_names, lengths), key=lambda x: x[0]))
+        file_names_sorted, lengths_sorted = zip(*sorted(zip(file_names,
+                                                            lengths,
+                                                            strict=True),
+                                                        key=lambda x: x[0]),
+                                                strict=True)
         assert len(
             file_names_sorted) != 0, "No parquet files found in the dataset"
 
@@ -138,8 +141,8 @@ def get_parquet_files_and_length(path: str):
     return get_parquet_files_and_length(path)
 
 
-def read_row_from_parquet_file(parquet_files: List[str], global_row_idx: int,
-                               lengths: List[int]) -> Dict[str, Any]:
+def read_row_from_parquet_file(parquet_files: list[str], global_row_idx: int,
+                               lengths: list[int]) -> dict[str, Any]:
     '''
     Read a row from a parquet file.
     Args:
@@ -250,7 +253,7 @@ class LatentsParquetMapStyleDataset(Dataset):
         return negative_prompt_embedding, negative_prompt_attention_mask, negative_prompt
 
     # PyTorch calls this ONLY because the batch_sampler yields a list
-    def __getitems__(self, indices: List[int]) -> Dict[str, Any]:
+    def __getitems__(self, indices: list[int]) -> dict[str, Any]:
         """
         Batch fetch using read_row_from_parquet_file for each index.
         """
@@ -285,7 +288,7 @@ def build_parquet_map_style_dataloader(
         drop_last=True,
         drop_first_row=False,
         text_padding_length=512,
-        seed=42) -> Tuple[LatentsParquetMapStyleDataset, StatefulDataLoader]:
+        seed=42) -> tuple[LatentsParquetMapStyleDataset, StatefulDataLoader]:
     dataset = LatentsParquetMapStyleDataset(
         path,
         batch_size,

fastvideo/v1/dataset/preprocessing_datasets.py

@@ -7,7 +7,7 @@ from abc import ABC, abstractmethod
 from collections import Counter
 from dataclasses import dataclass
 from os.path import join as opj
-from typing import Any, Dict, List, Optional, Union
+from typing import Any
 
 import numpy as np
 import torch
@@ -31,21 +31,21 @@ class PreprocessBatch:
     """
     # Raw metadata
     path: str
-    cap: Union[str, List[str]]
-    resolution: Optional[Dict] = None
-    fps: Optional[float] = None
-    duration: Optional[float] = None
+    cap: str | list[str]
+    resolution: dict | None = None
+    fps: float | None = None
+    duration: float | None = None
 
     # Processed metadata
-    num_frames: Optional[int] = None
-    sample_frame_index: Optional[List[int]] = None
-    sample_num_frames: Optional[int] = None
+    num_frames: int | None = None
+    sample_frame_index: list[int] | None = None
+    sample_num_frames: int | None = None
 
     # Processed data
-    pixel_values: Optional[torch.Tensor] = None
-    text: Optional[str] = None
-    input_ids: Optional[torch.Tensor] = None
-    cond_mask: Optional[torch.Tensor] = None
+    pixel_values: torch.Tensor | None = None
+    text: str | None = None
+    input_ids: torch.Tensor | None = None
+    cond_mask: torch.Tensor | None = None
 
     @property
     def is_video(self) -> bool:
@@ -488,7 +488,7 @@ class VideoCaptionMergedDataset(torch.utils.data.IterableDataset,
             text_max_length=args.text_max_length,
             cfg_rate=args.training_cfg_rate)
 
-    def _load_raw_data(self) -> List[Dict]:
+    def _load_raw_data(self) -> list[dict]:
         """Load raw data from JSON files."""
         # Read folder-annotation pairs
         with open(self.data_merge_path) as f:
@@ -501,7 +501,7 @@ class VideoCaptionMergedDataset(torch.utils.data.IterableDataset,
                    ) == 2, "Folder-annotation pair should have two elements"
         folder, annotation_file = folder_anno_pairs[0]
 
-        data_items: List[Dict] = []
+        data_items: list[dict] = []
         with open(annotation_file) as f:
             data_items = json.load(f)
 
@@ -511,7 +511,7 @@ class VideoCaptionMergedDataset(torch.utils.data.IterableDataset,
 
         return data_items
 
-    def _process_metadata(self) -> List[PreprocessBatch]:
+    def _process_metadata(self) -> list[PreprocessBatch]:
         """Process the raw metadata through all filtering stages."""
         raw_data = self._load_raw_data()
         processed_batches = []
@@ -522,7 +522,7 @@ class VideoCaptionMergedDataset(torch.utils.data.IterableDataset,
             "resolution_failed": 0,
             "frame_sampling_failed": 0
         }
-        sample_num_frames: List[int] = []
+        sample_num_frames: list[int] = []
 
         for item in raw_data:
             batch = PreprocessBatch(path=item["path"],
@@ -548,7 +548,7 @@ class VideoCaptionMergedDataset(torch.utils.data.IterableDataset,
         return processed_batches
 
     def _apply_filter_stages(self, batch: PreprocessBatch,
-                             filter_counts: Dict[str, int]) -> bool:
+                             filter_counts: dict[str, int]) -> bool:
         """Apply all filter stages and update counters. Returns True if batch should be kept."""
         if not self.validation_stage.should_keep(batch):
             filter_counts["validation_failed"] += 1
@@ -564,8 +564,8 @@ class VideoCaptionMergedDataset(torch.utils.data.IterableDataset,
 
         return True
 
-    def _log_filtering_stats(self, filter_counts: Dict[str, int],
-                             sample_num_frames: List[int], before_count: int,
+    def _log_filtering_stats(self, filter_counts: dict[str, int],
+                             sample_num_frames: list[int], before_count: int,
                              after_count: int):
         """Log filtering statistics."""
         logger.info(
@@ -584,7 +584,7 @@ class VideoCaptionMergedDataset(torch.utils.data.IterableDataset,
     def __len__(self):
         return len(self.processed_batches)
 
-    def _get_item(self, idx: int) -> Dict:
+    def _get_item(self, idx: int) -> dict:
         """Get a single processed data item."""
         batch = self.processed_batches[idx]
 
@@ -608,10 +608,10 @@ class VideoCaptionMergedDataset(torch.utils.data.IterableDataset,
 
         return result
 
-    def state_dict(self) -> Dict[str, Any]:
+    def state_dict(self) -> dict[str, Any]:
         """Return state dict for checkpointing."""
         return {"processed_batches": self.processed_batches}
 
-    def load_state_dict(self, state_dict: Dict[str, Any]) -> None:
+    def load_state_dict(self, state_dict: dict[str, Any]) -> None:
         """Load state dict from checkpoint."""
         self.processed_batches = state_dict["processed_batches"]

fastvideo/v1/dataset/utils.py

@@ -1,5 +1,5 @@
 import random
-from typing import Any, Dict, List, cast
+from typing import Any, cast
 
 import numpy as np
 import torch
@@ -21,7 +21,7 @@ def pad(t: torch.Tensor, padding_length: int) -> torch.Tensor:
         return t[:padding_length], torch.ones(padding_length)
 
 
-def get_torch_tensors_from_row_dict(row_dict, keys, cfg_rate) -> Dict[str, Any]:
+def get_torch_tensors_from_row_dict(row_dict, keys, cfg_rate) -> dict[str, Any]:
     """
     Get the latents and prompts from a row dictionary.
     """
@@ -61,7 +61,7 @@ def collate_latents_embs_masks(
         text_padding_length,
         keys,
         cfg_rate=0.0
-) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, List[str]]:
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, list[str]]:
     # Initialize tensors to hold padded embeddings and masks
     all_latents = []
     all_embs = []
@@ -95,7 +95,7 @@ def collate_latents_embs_masks(
 def collate_rows_from_parquet_schema(rows,
                                      parquet_schema,
                                      text_padding_length,
-                                     cfg_rate=0.0) -> Dict[str, Any]:
+                                     cfg_rate=0.0) -> dict[str, Any]:
     """
     Collate rows from parquet files based on the provided schema.
     Dynamically processes tensor fields based on schema and returns batched data.
@@ -108,10 +108,10 @@ def collate_rows_from_parquet_schema(rows,
         Dict containing batched tensors and metadata
     """
     if not rows:
-        return cast(Dict[str, Any], {})
+        return cast(dict[str, Any], {})
 
     # Initialize containers for different data types
-    batch_data: Dict[str, Any] = {}
+    batch_data: dict[str, Any] = {}
 
     # Get tensor and metadata field names from schema (fields ending with '_bytes')
     tensor_fields = []

fastvideo/v1/distributed/device_communicators/base_device_communicator.py

@@ -1,7 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0
 # Adapted from https://github.com/vllm-project/vllm/blob/v0.7.3/vllm/distributed/device_communicators/base_device_communicator.py
 
-from typing import Any, Optional, Tuple
+from typing import Any
 
 import torch
 import torch.distributed as dist
@@ -28,7 +28,7 @@ class DistributedAutograd:
         def forward(ctx: Any,
                     group: ProcessGroup,
                     input_: Tensor,
-                    op: Optional[dist.ReduceOp] = None) -> Tensor:
+                    op: dist.ReduceOp | None = None) -> Tensor:
             ctx.group = group
             ctx.op = op
             output = input_.clone()
@@ -37,7 +37,7 @@ class DistributedAutograd:
 
         @staticmethod
         def backward(ctx: Any,
-                     grad_output: Tensor) -> Tuple[None, Tensor, None]:
+                     grad_output: Tensor) -> tuple[None, Tensor, None]:
             grad_output = grad_output.clone()
             dist.all_reduce(grad_output, group=ctx.group, op=ctx.op)
             return None, grad_output, None
@@ -75,7 +75,7 @@ class DistributedAutograd:
 
         @staticmethod
         def backward(ctx: Any,
-                     grad_output: Tensor) -> Tuple[None, Tensor, None, None]:
+                     grad_output: Tensor) -> tuple[None, Tensor, None, None]:
             # Split the gradient tensor along the gathered dimension
             dim_size = grad_output.size(ctx.dim) // ctx.world_size
             grad_chunks = grad_output.reshape(grad_output.shape[:ctx.dim] +
@@ -163,7 +163,7 @@ class DistributedAutograd:
         @staticmethod
         def backward(
                 ctx: Any,
-                grad_output: Tensor) -> Tuple[None, Tensor, None, None, None]:
+                grad_output: Tensor) -> tuple[None, Tensor, None, None, None]:
             if ctx.world_size == 1:
                 return None, grad_output, None, None, None
 
@@ -184,8 +184,8 @@ class DeviceCommunicatorBase:
 
     def __init__(self,
                  cpu_group: ProcessGroup,
-                 device: Optional[torch.device] = None,
-                 device_group: Optional[ProcessGroup] = None,
+                 device: torch.device | None = None,
+                 device_group: ProcessGroup | None = None,
                  unique_name: str = ""):
         self.device = device or torch.device("cpu")
         self.cpu_group = cpu_group
@@ -201,7 +201,7 @@ class DeviceCommunicatorBase:
 
     def all_reduce(self,
                    input_: torch.Tensor,
-                   op: Optional[dist.ReduceOp] = ReduceOp.SUM) -> torch.Tensor:
+                   op: dist.ReduceOp | None = ReduceOp.SUM) -> torch.Tensor:
         """Performs an all_reduce operation with gradient support."""
         return DistributedAutograd.AllReduce.apply(self.device_group, input_,
                                                    op)
@@ -225,7 +225,7 @@ class DeviceCommunicatorBase:
     def gather(self,
                input_: torch.Tensor,
                dst: int = 0,
-               dim: int = -1) -> Optional[torch.Tensor]:
+               dim: int = -1) -> torch.Tensor | None:
         """
         NOTE: We assume that the input tensor is on the same device across
         all the ranks.
@@ -254,7 +254,7 @@ class DeviceCommunicatorBase:
             output_tensor = None
         return output_tensor
 
-    def send(self, tensor: torch.Tensor, dst: Optional[int] = None) -> None:
+    def send(self, tensor: torch.Tensor, dst: int | None = None) -> None:
         """Sends a tensor to the destination rank in a non-blocking way"""
         """NOTE: `dst` is the local rank of the destination rank."""
         if dst is None:
@@ -264,7 +264,7 @@ class DeviceCommunicatorBase:
     def recv(self,
              size: torch.Size,
              dtype: torch.dtype,
-             src: Optional[int] = None) -> torch.Tensor:
+             src: int | None = None) -> torch.Tensor:
         """Receives a tensor from the source rank."""
         """NOTE: `src` is the local rank of the source rank."""
         if src is None:

fastvideo/v1/distributed/device_communicators/cuda_communicator.py

@@ -1,8 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # Adapted from https://github.com/vllm-project/vllm/blob/v0.7.3/vllm/distributed/device_communicators/cuda_communicator.py
 
-from typing import Optional
-
 import torch
 from torch.distributed import ProcessGroup
 
@@ -14,24 +12,22 @@ class CudaCommunicator(DeviceCommunicatorBase):
 
     def __init__(self,
                  cpu_group: ProcessGroup,
-                 device: Optional[torch.device] = None,
-                 device_group: Optional[ProcessGroup] = None,
+                 device: torch.device | None = None,
+                 device_group: ProcessGroup | None = None,
                  unique_name: str = ""):
         super().__init__(cpu_group, device, device_group, unique_name)
 
         from fastvideo.v1.distributed.device_communicators.pynccl import (
             PyNcclCommunicator)
 
-        self.pynccl_comm: Optional[PyNcclCommunicator] = None
+        self.pynccl_comm: PyNcclCommunicator | None = None
         if self.world_size > 1:
             self.pynccl_comm = PyNcclCommunicator(
                 group=self.cpu_group,
                 device=self.device,
             )
 
-    def all_reduce(self,
-                   input_,
-                   op: Optional[torch.distributed.ReduceOp] = None):
+    def all_reduce(self, input_, op: torch.distributed.ReduceOp | None = None):
         pynccl_comm = self.pynccl_comm
         assert pynccl_comm is not None
         out = pynccl_comm.all_reduce(input_, op=op)
@@ -44,7 +40,7 @@ class CudaCommunicator(DeviceCommunicatorBase):
             torch.distributed.all_reduce(out, group=self.device_group, op=op)
         return out
 
-    def send(self, tensor: torch.Tensor, dst: Optional[int] = None) -> None:
+    def send(self, tensor: torch.Tensor, dst: int | None = None) -> None:
         """Sends a tensor to the destination rank in a non-blocking way"""
         """NOTE: `dst` is the local rank of the destination rank."""
         if dst is None:
@@ -59,7 +55,7 @@ class CudaCommunicator(DeviceCommunicatorBase):
     def recv(self,
              size: torch.Size,
              dtype: torch.dtype,
-             src: Optional[int] = None) -> torch.Tensor:
+             src: int | None = None) -> torch.Tensor:
         """Receives a tensor from the source rank."""
         """NOTE: `src` is the local rank of the source rank."""
         if src is None:

fastvideo/v1/distributed/device_communicators/pynccl.py

@@ -1,8 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # Adapted from https://github.com/vllm-project/vllm/blob/v0.7.3/vllm/distributed/device_communicators/pynccl.py
 
-from typing import Optional, Union
-
 # ===================== import region =====================
 import torch
 import torch.distributed as dist
@@ -22,9 +20,9 @@ class PyNcclCommunicator:
 
     def __init__(
         self,
-        group: Union[ProcessGroup, StatelessProcessGroup],
-        device: Union[int, str, torch.device],
-        library_path: Optional[str] = None,
+        group: ProcessGroup | StatelessProcessGroup,
+        device: int | str | torch.device,
+        library_path: str | None = None,
     ):
         """
         Args:

fastvideo/v1/distributed/device_communicators/pynccl_wrapper.py

@@ -27,7 +27,7 @@
 import ctypes
 import platform
 from dataclasses import dataclass
-from typing import Any, Dict, List, Optional
+from typing import Any
 
 import torch
 from torch.distributed import ReduceOp
@@ -124,7 +124,7 @@ class ncclRedOpTypeEnum:
 class Function:
     name: str
     restype: Any
-    argtypes: List[Any]
+    argtypes: list[Any]
 
 
 class NCCLLibrary:
@@ -212,13 +212,13 @@ class NCCLLibrary:
 
     # class attribute to store the mapping from the path to the library
     # to avoid loading the same library multiple times
-    path_to_library_cache: Dict[str, Any] = {}
+    path_to_library_cache: dict[str, Any] = {}
 
     # class attribute to store the mapping from library path
     #  to the corresponding dictionary
-    path_to_dict_mapping: Dict[str, Dict[str, Any]] = {}
+    path_to_dict_mapping: dict[str, dict[str, Any]] = {}
 
-    def __init__(self, so_file: Optional[str] = None):
+    def __init__(self, so_file: str | None = None):
 
         so_file = so_file or find_nccl_library()
 
@@ -240,7 +240,7 @@ class NCCLLibrary:
             raise e
 
         if so_file not in NCCLLibrary.path_to_dict_mapping:
-            _funcs: Dict[str, Any] = {}
+            _funcs: dict[str, Any] = {}
             for func in NCCLLibrary.exported_functions:
                 f = getattr(self.lib, func.name)
                 f.restype = func.restype

fastvideo/v1/distributed/parallel_state.py

@@ -28,10 +28,11 @@ import os
 import pickle
 import weakref
 from collections import namedtuple
+from collections.abc import Callable
 from contextlib import contextmanager
 from dataclasses import dataclass
 from multiprocessing import shared_memory
-from typing import Any, Callable, Dict, List, Optional, Tuple, Union
+from typing import Any, Optional
 from unittest.mock import patch
 
 import torch
@@ -58,15 +59,15 @@ TensorMetadata = namedtuple("TensorMetadata", ["device", "dtype", "size"])
 
 
 def _split_tensor_dict(
-    tensor_dict: Dict[str, Union[torch.Tensor, Any]]
-) -> Tuple[List[Tuple[str, Any]], List[torch.Tensor]]:
+    tensor_dict: dict[str, torch.Tensor | Any]
+) -> tuple[list[tuple[str, Any]], list[torch.Tensor]]:
     """Split the tensor dictionary into two parts:
     1. A list of (key, value) pairs. If the value is a tensor, it is replaced
          by its metadata.
     2. A list of tensors.
     """
-    metadata_list: List[Tuple[str, Any]] = []
-    tensor_list: List[torch.Tensor] = []
+    metadata_list: list[tuple[str, Any]] = []
+    tensor_list: list[torch.Tensor] = []
     for key, value in tensor_dict.items():
         if isinstance(value, torch.Tensor):
             # Note: we cannot use `value.device` here,
@@ -82,7 +83,7 @@ def _split_tensor_dict(
     return metadata_list, tensor_list
 
 
-_group_name_counter: Dict[str, int] = {}
+_group_name_counter: dict[str, int] = {}
 
 
 def _get_unique_name(name: str) -> str:
@@ -98,7 +99,7 @@ def _get_unique_name(name: str) -> str:
     return newname
 
 
-_groups: Dict[str, Callable[[], Optional["GroupCoordinator"]]] = {}
+_groups: dict[str, Callable[[], Optional["GroupCoordinator"]]] = {}
 
 
 def _register_group(group: "GroupCoordinator") -> None:
@@ -129,7 +130,7 @@ class GroupCoordinator:
 
     # available attributes:
     rank: int  # global rank
-    ranks: List[int]  # global ranks in the group
+    ranks: list[int]  # global ranks in the group
     world_size: int  # size of the group
     # difference between `local_rank` and `rank_in_group`:
     # if we have a group of size 4 across two nodes:
@@ -144,16 +145,16 @@ class GroupCoordinator:
     device_group: ProcessGroup  # group for device communication
     use_device_communicator: bool  # whether to use device communicator
     device_communicator: DeviceCommunicatorBase  # device communicator
-    mq_broadcaster: Optional[Any]  # shared memory broadcaster
+    mq_broadcaster: Any | None  # shared memory broadcaster
 
     def __init__(
         self,
-        group_ranks: List[List[int]],
+        group_ranks: list[list[int]],
         local_rank: int,
-        torch_distributed_backend: Union[str, Backend],
+        torch_distributed_backend: str | Backend,
         use_device_communicator: bool,
         use_message_queue_broadcaster: bool = False,
-        group_name: Optional[str] = None,
+        group_name: str | None = None,
     ):
         group_name = group_name or "anonymous"
         self.unique_name = _get_unique_name(group_name)
@@ -244,8 +245,8 @@ class GroupCoordinator:
         return self.ranks[(rank_in_group - 1) % world_size]
 
     @contextmanager
-    def graph_capture(
-            self, graph_capture_context: Optional[GraphCaptureContext] = None):
+    def graph_capture(self,
+                      graph_capture_context: GraphCaptureContext | None = None):
         if graph_capture_context is None:
             stream = torch.cuda.Stream()
             graph_capture_context = GraphCaptureContext(stream)
@@ -262,9 +263,9 @@ class GroupCoordinator:
             yield graph_capture_context
 
     def all_reduce(
-        self,
-        input_: torch.Tensor,
-        op: Optional[torch.distributed.ReduceOp] = ReduceOp.SUM
+            self,
+            input_: torch.Tensor,
+            op: torch.distributed.ReduceOp | None = ReduceOp.SUM
     ) -> torch.Tensor:
         """
         User-facing all-reduce function before we actually call the
@@ -291,9 +292,9 @@ class GroupCoordinator:
             return self._all_reduce_out_place(input_, op=op)
 
     def _all_reduce_out_place(
-        self,
-        input_: torch.Tensor,
-        op: Optional[torch.distributed.ReduceOp] = ReduceOp.SUM
+            self,
+            input_: torch.Tensor,
+            op: torch.distributed.ReduceOp | None = ReduceOp.SUM
     ) -> torch.Tensor:
         return self.device_communicator.all_reduce(input_, op=op)
 
@@ -310,7 +311,7 @@ class GroupCoordinator:
     def gather(self,
                input_: torch.Tensor,
                dst: int = 0,
-               dim: int = -1) -> Optional[torch.Tensor]:
+               dim: int = -1) -> torch.Tensor | None:
         """
         NOTE: We assume that the input tensor is on the same device across
         all the ranks.
@@ -346,7 +347,7 @@ class GroupCoordinator:
                                     group=self.device_group)
         return input_
 
-    def broadcast_object(self, obj: Optional[Any] = None, src: int = 0):
+    def broadcast_object(self, obj: Any | None = None, src: int = 0):
         """Broadcast the input object.
         NOTE: `src` is the local rank of the source rank.
         """
@@ -371,9 +372,9 @@ class GroupCoordinator:
             return recv[0]
 
     def broadcast_object_list(self,
-                              obj_list: List[Any],
+                              obj_list: list[Any],
                               src: int = 0,
-                              group: Optional[ProcessGroup] = None):
+                              group: ProcessGroup | None = None):
         """Broadcast the input object list.
         NOTE: `src` is the local rank of the source rank.
         """
@@ -453,11 +454,11 @@ class GroupCoordinator:
 
     def broadcast_tensor_dict(
         self,
-        tensor_dict: Optional[Dict[str, Union[torch.Tensor, Any]]] = None,
+        tensor_dict: dict[str, torch.Tensor | Any] | None = None,
         src: int = 0,
-        group: Optional[ProcessGroup] = None,
-        metadata_group: Optional[ProcessGroup] = None
-    ) -> Optional[Dict[str, Union[torch.Tensor, Any]]]:
+        group: ProcessGroup | None = None,
+        metadata_group: ProcessGroup | None = None
+    ) -> dict[str, torch.Tensor | Any] | None:
         """Broadcast the input tensor dictionary.
         NOTE: `src` is the local rank of the source rank.
         """
@@ -471,7 +472,7 @@ class GroupCoordinator:
 
         rank_in_group = self.rank_in_group
         if rank_in_group == src:
-            metadata_list: List[Tuple[Any, Any]] = []
+            metadata_list: list[tuple[Any, Any]] = []
             assert isinstance(
                 tensor_dict,
                 dict), (f"Expecting a dictionary, got {type(tensor_dict)}")
@@ -538,10 +539,10 @@ class GroupCoordinator:
 
     def send_tensor_dict(
         self,
-        tensor_dict: Dict[str, Union[torch.Tensor, Any]],
-        dst: Optional[int] = None,
+        tensor_dict: dict[str, torch.Tensor | Any],
+        dst: int | None = None,
         all_gather_group: Optional["GroupCoordinator"] = None,
-    ) -> Optional[Dict[str, Union[torch.Tensor, Any]]]:
+    ) -> dict[str, torch.Tensor | Any] | None:
         """Send the input tensor dictionary.
         NOTE: `dst` is the local rank of the source rank.
         """
@@ -561,7 +562,7 @@ class GroupCoordinator:
             dst = (self.rank_in_group + 1) % self.world_size
         assert dst < self.world_size, f"Invalid dst rank ({dst})"
 
-        metadata_list: List[Tuple[Any, Any]] = []
+        metadata_list: list[tuple[Any, Any]] = []
         assert isinstance(
             tensor_dict,
             dict), f"Expecting a dictionary, got {type(tensor_dict)}"
@@ -592,9 +593,9 @@ class GroupCoordinator:
 
     def recv_tensor_dict(
         self,
-        src: Optional[int] = None,
+        src: int | None = None,
         all_gather_group: Optional["GroupCoordinator"] = None,
-    ) -> Optional[Dict[str, Union[torch.Tensor, Any]]]:
+    ) -> dict[str, torch.Tensor | Any] | None:
         """Recv the input tensor dictionary.
         NOTE: `src` is the local rank of the source rank.
         """
@@ -615,7 +616,7 @@ class GroupCoordinator:
         assert src < self.world_size, f"Invalid src rank ({src})"
 
         recv_metadata_list = self.recv_object(src=src)
-        tensor_dict: Dict[str, Any] = {}
+        tensor_dict: dict[str, Any] = {}
         for key, value in recv_metadata_list:
             if isinstance(value, TensorMetadata):
                 tensor = torch.empty(value.size,
@@ -665,7 +666,7 @@ class GroupCoordinator:
         """
         torch.distributed.barrier(group=self.cpu_group)
 
-    def send(self, tensor: torch.Tensor, dst: Optional[int] = None) -> None:
+    def send(self, tensor: torch.Tensor, dst: int | None = None) -> None:
         """Sends a tensor to the destination rank in a non-blocking way"""
         """NOTE: `dst` is the local rank of the destination rank."""
         self.device_communicator.send(tensor, dst)
@@ -673,7 +674,7 @@ class GroupCoordinator:
     def recv(self,
              size: torch.Size,
              dtype: torch.dtype,
-             src: Optional[int] = None) -> torch.Tensor:
+             src: int | None = None) -> torch.Tensor:
         """Receives a tensor from the source rank."""
         """NOTE: `src` is the local rank of the source rank."""
         return self.device_communicator.recv(size, dtype, src)
@@ -691,7 +692,7 @@ class GroupCoordinator:
             self.mq_broadcaster = None
 
 
-_WORLD: Optional[GroupCoordinator] = None
+_WORLD: GroupCoordinator | None = None
 
 
 def get_world_group() -> GroupCoordinator:
@@ -699,7 +700,7 @@ def get_world_group() -> GroupCoordinator:
     return _WORLD
 
 
-def init_world_group(ranks: List[int], local_rank: int,
+def init_world_group(ranks: list[int], local_rank: int,
                      backend: str) -> GroupCoordinator:
     return GroupCoordinator(
         group_ranks=[ranks],
@@ -711,11 +712,11 @@ def init_world_group(ranks: List[int], local_rank: int,
 
 
 def init_model_parallel_group(
-    group_ranks: List[List[int]],
+    group_ranks: list[list[int]],
     local_rank: int,
     backend: str,
     use_message_queue_broadcaster: bool = False,
-    group_name: Optional[str] = None,
+    group_name: str | None = None,
 ) -> GroupCoordinator:
 
     return GroupCoordinator(
@@ -728,7 +729,7 @@ def init_model_parallel_group(
     )
 
 
-_TP: Optional[GroupCoordinator] = None
+_TP: GroupCoordinator | None = None
 
 
 def get_tp_group() -> GroupCoordinator:
@@ -750,7 +751,7 @@ def init_distributed_environment(
     distributed_init_method: str = "env://",
     local_rank: int = 0,
     backend: str = "nccl",
-    device_id: Optional[torch.device] = None,
+    device_id: torch.device | None = None,
 ):
     logger.debug(
         "world_size=%d rank=%d local_rank=%d "
@@ -786,7 +787,7 @@ def init_distributed_environment(
             "world group already initialized with a different world size")
 
 
-_SP: Optional[GroupCoordinator] = None
+_SP: GroupCoordinator | None = None
 
 
 def get_sp_group() -> GroupCoordinator:
@@ -794,7 +795,7 @@ def get_sp_group() -> GroupCoordinator:
     return _SP
 
 
-_DP: Optional[GroupCoordinator] = None
+_DP: GroupCoordinator | None = None
 
 
 def get_dp_group() -> GroupCoordinator:
@@ -805,7 +806,7 @@ def get_dp_group() -> GroupCoordinator:
 def initialize_model_parallel(
     tensor_model_parallel_size: int = 1,
     sequence_model_parallel_size: int = 1,
-    backend: Optional[str] = None,
+    backend: str | None = None,
 ) -> None:
     """
     Initialize model parallel groups.
@@ -1025,8 +1026,8 @@ def cleanup_dist_env_and_memory(shutdown_ray: bool = False):
             "torch._C._host_emptyCache() only available in Pytorch >=2.5")
 
 
-def in_the_same_node_as(pg: Union[ProcessGroup, StatelessProcessGroup],
-                        source_rank: int = 0) -> List[bool]:
+def in_the_same_node_as(pg: ProcessGroup | StatelessProcessGroup,
+                        source_rank: int = 0) -> list[bool]:
     """
     This is a collective operation that returns if each rank is in the same node
     as the source rank. It tests if processes are attached to the same
@@ -1112,7 +1113,7 @@ def in_the_same_node_as(pg: Union[ProcessGroup, StatelessProcessGroup],
 
 def initialize_tensor_parallel_group(
         tensor_model_parallel_size: int = 1,
-        backend: Optional[str] = None,
+        backend: str | None = None,
         group_name_suffix: str = "") -> GroupCoordinator:
     """Initialize a tensor parallel group for a specific model.
     
@@ -1176,7 +1177,7 @@ def initialize_tensor_parallel_group(
 
 def initialize_sequence_parallel_group(
         sequence_model_parallel_size: int = 1,
-        backend: Optional[str] = None,
+        backend: str | None = None,
         group_name_suffix: str = "") -> GroupCoordinator:
     """Initialize a sequence parallel group for a specific model.
     
@@ -1236,4 +1237,4 @@ def initialize_sequence_parallel_group(
                                          backend,
                                          group_name=group_name)
 
-    return sp_group
+    return sp_group

fastvideo/v1/distributed/utils.py

@@ -9,7 +9,8 @@ import dataclasses
 import pickle
 import time
 from collections import deque
-from typing import Any, Deque, Dict, Optional, Sequence, Tuple
+from collections.abc import Sequence
+from typing import Any
 
 import torch
 from torch.distributed import TCPStore
@@ -72,15 +73,15 @@ class StatelessProcessGroup:
     data_expiration_seconds: int = 3600  # 1 hour
 
     # dst rank -> counter
-    send_dst_counter: Dict[int, int] = dataclasses.field(default_factory=dict)
+    send_dst_counter: dict[int, int] = dataclasses.field(default_factory=dict)
     # src rank -> counter
-    recv_src_counter: Dict[int, int] = dataclasses.field(default_factory=dict)
+    recv_src_counter: dict[int, int] = dataclasses.field(default_factory=dict)
     broadcast_send_counter: int = 0
-    broadcast_recv_src_counter: Dict[int, int] = dataclasses.field(
+    broadcast_recv_src_counter: dict[int, int] = dataclasses.field(
         default_factory=dict)
 
     # A deque to store the data entries, with key and timestamp.
-    entries: Deque[Tuple[str, float]] = dataclasses.field(default_factory=deque)
+    entries: deque[tuple[str, float]] = dataclasses.field(default_factory=deque)
 
     def __post_init__(self):
         assert self.rank < self.world_size
@@ -114,7 +115,7 @@ class StatelessProcessGroup:
         self.recv_src_counter[src] += 1
         return obj
 
-    def broadcast_obj(self, obj: Optional[Any], src: int) -> Any:
+    def broadcast_obj(self, obj: Any | None, src: int) -> Any:
         """Broadcast an object from a source rank to all other ranks.
         It does not clean up after all ranks have received the object.
         Use it for limited times, e.g., for initialization.

fastvideo/v1/entrypoints/cli/generate.py

@@ -4,7 +4,7 @@
 import argparse
 import dataclasses
 import os
-from typing import List, cast
+from typing import cast
 
 from fastvideo import VideoGenerator
 from fastvideo.v1.configs.sample.base import SamplingParam
@@ -26,11 +26,11 @@ class GenerateSubcommand(CLISubcommand):
         self.init_arg_names = self._get_init_arg_names()
         self.generation_arg_names = self._get_generation_arg_names()
 
-    def _get_init_arg_names(self) -> List[str]:
+    def _get_init_arg_names(self) -> list[str]:
         """Get names of arguments for VideoGenerator initialization"""
         return ["num_gpus", "tp_size", "sp_size", "model_path"]
 
-    def _get_generation_arg_names(self) -> List[str]:
+    def _get_generation_arg_names(self) -> list[str]:
         """Get names of arguments for generate_video method"""
         return [field.name for field in dataclasses.fields(SamplingParam)]
 
@@ -120,5 +120,5 @@ class GenerateSubcommand(CLISubcommand):
         return cast(FlexibleArgumentParser, generate_parser)
 
 
-def cmd_init() -> List[CLISubcommand]:
+def cmd_init() -> list[CLISubcommand]:
     return [GenerateSubcommand()]

fastvideo/v1/entrypoints/cli/main.py

@@ -1,14 +1,12 @@
 # SPDX-License-Identifier: Apache-2.0
 # adapted from vllm: https://github.com/vllm-project/vllm/blob/v0.7.3/vllm/entrypoints/cli/main.py
 
-from typing import List
-
 from fastvideo.v1.entrypoints.cli.cli_types import CLISubcommand
 from fastvideo.v1.entrypoints.cli.generate import cmd_init as generate_cmd_init
 from fastvideo.v1.utils import FlexibleArgumentParser
 
 
-def cmd_init() -> List[CLISubcommand]:
+def cmd_init() -> list[CLISubcommand]:
     """Initialize all commands from separate modules"""
     commands = []
     commands.extend(generate_cmd_init())

fastvideo/v1/entrypoints/cli/utils.py

@@ -4,7 +4,6 @@ import argparse
 import os
 import subprocess
 import sys
-from typing import List, Optional
 
 from fastvideo.v1.logger import init_logger
 
@@ -19,8 +18,8 @@ class RaiseNotImplementedAction(argparse.Action):
 
 
 def launch_distributed(num_gpus: int,
-                       args: List[str],
-                       master_port: Optional[int] = None) -> int:
+                       args: list[str],
+                       master_port: int | None = None) -> int:
     """
     Launch a distributed job with the given arguments
     

fastvideo/v1/entrypoints/video_generator.py

@@ -10,7 +10,7 @@ import gc
 import math
 import os
 import time
-from typing import Any, Dict, List, Optional, Union
+from typing import Any
 
 import imageio
 import numpy as np
@@ -51,8 +51,8 @@ class VideoGenerator:
     @classmethod
     def from_pretrained(cls,
                         model_path: str,
-                        device: Optional[str] = None,
-                        torch_dtype: Optional[torch.dtype] = None,
+                        device: str | None = None,
+                        torch_dtype: torch.dtype | None = None,
                         **kwargs) -> "VideoGenerator":
         """
         Create a video generator from a pretrained model.
@@ -100,9 +100,9 @@ class VideoGenerator:
     def generate_video(
         self,
         prompt: str,
-        sampling_param: Optional[SamplingParam] = None,
+        sampling_param: SamplingParam | None = None,
         **kwargs,
-    ) -> Union[Dict[str, Any], List[np.ndarray]]:
+    ) -> dict[str, Any] | list[np.ndarray]:
         """
         Generate a video based on the given prompt.
         

fastvideo/v1/envs.py

@@ -2,28 +2,29 @@
 # Adapted from vllm: https://github.com/vllm-project/vllm/blob/v0.7.3/vllm/envs.py
 
 import os
-from typing import TYPE_CHECKING, Any, Callable, Dict, Optional
+from collections.abc import Callable
+from typing import TYPE_CHECKING, Any
 
 if TYPE_CHECKING:
     FASTVIDEO_RINGBUFFER_WARNING_INTERVAL: int = 60
-    FASTVIDEO_NCCL_SO_PATH: Optional[str] = None
-    LD_LIBRARY_PATH: Optional[str] = None
+    FASTVIDEO_NCCL_SO_PATH: str | None = None
+    LD_LIBRARY_PATH: str | None = None
     LOCAL_RANK: int = 0
-    CUDA_VISIBLE_DEVICES: Optional[str] = None
+    CUDA_VISIBLE_DEVICES: str | None = None
     FASTVIDEO_CACHE_ROOT: str = os.path.expanduser("~/.cache/fastvideo")
     FASTVIDEO_CONFIG_ROOT: str = os.path.expanduser("~/.config/fastvideo")
     FASTVIDEO_CONFIGURE_LOGGING: int = 1
     FASTVIDEO_LOGGING_LEVEL: str = "INFO"
     FASTVIDEO_LOGGING_PREFIX: str = ""
-    FASTVIDEO_LOGGING_CONFIG_PATH: Optional[str] = None
+    FASTVIDEO_LOGGING_CONFIG_PATH: str | None = None
     FASTVIDEO_TRACE_FUNCTION: int = 0
-    FASTVIDEO_ATTENTION_BACKEND: Optional[str] = None
-    FASTVIDEO_ATTENTION_CONFIG: Optional[str] = None
+    FASTVIDEO_ATTENTION_BACKEND: str | None = None
+    FASTVIDEO_ATTENTION_CONFIG: str | None = None
     FASTVIDEO_WORKER_MULTIPROC_METHOD: str = "fork"
     FASTVIDEO_TARGET_DEVICE: str = "cuda"
-    MAX_JOBS: Optional[str] = None
-    NVCC_THREADS: Optional[str] = None
-    CMAKE_BUILD_TYPE: Optional[str] = None
+    MAX_JOBS: str | None = None
+    NVCC_THREADS: str | None = None
+    CMAKE_BUILD_TYPE: str | None = None
     VERBOSE: bool = False
     FASTVIDEO_SERVER_DEV_MODE: bool = False
 
@@ -42,7 +43,7 @@ def get_default_config_root() -> str:
     )
 
 
-def maybe_convert_int(value: Optional[str]) -> Optional[int]:
+def maybe_convert_int(value: str | None) -> int | None:
     if value is None:
         return None
     return int(value)
@@ -53,7 +54,7 @@ def maybe_convert_int(value: Optional[str]) -> Optional[int]:
 
 # begin-env-vars-definition
 
-environment_variables: Dict[str, Callable[[], Any]] = {
+environment_variables: dict[str, Callable[[], Any]] = {
 
     # ================== Installation Time Env Vars ==================
 

fastvideo/v1/fastvideo_args.py

@@ -6,7 +6,7 @@ import argparse
 import dataclasses
 from contextlib import contextmanager
 from dataclasses import field
-from typing import Any, Dict, List, Optional
+from typing import Any
 
 from fastvideo.v1.configs.pipelines.base import PipelineConfig, STA_Mode
 from fastvideo.v1.logger import init_logger
@@ -15,7 +15,7 @@ from fastvideo.v1.utils import FlexibleArgumentParser, StoreBoolean
 logger = init_logger(__name__)
 
 
-def clean_cli_args(args: argparse.Namespace) -> Dict[str, Any]:
+def clean_cli_args(args: argparse.Namespace) -> dict[str, Any]:
     """
     Clean the arguments by removing the ones that not explicitly provided by the user.
     """
@@ -44,7 +44,7 @@ class FastVideoArgs:
 
     # HuggingFace specific parameters
     trust_remote_code: bool = False
-    revision: Optional[str] = None
+    revision: str | None = None
 
     # Parallelism
     num_gpus: int = 1
@@ -52,7 +52,7 @@ class FastVideoArgs:
     sp_size: int = -1
     hsdp_replicate_dim: int = 1
     hsdp_shard_dim: int = -1
-    dist_timeout: Optional[int] = None  # timeout for torch.distributed
+    dist_timeout: int | None = None  # timeout for torch.distributed
 
     pipeline_config: PipelineConfig = field(default_factory=PipelineConfig)
 
@@ -64,7 +64,7 @@ class FastVideoArgs:
     pin_cpu_memory: bool = True
 
     # STA (Sliding Tile Attention) parameters
-    mask_strategy_file_path: Optional[str] = None
+    mask_strategy_file_path: str | None = None
     STA_mode: STA_Mode = STA_Mode.STA_INFERENCE
     skip_time_steps: int = 15
 
@@ -280,7 +280,7 @@ class FastVideoArgs:
         return cls(**kwargs)  # type: ignore
 
     @classmethod
-    def from_kwargs(cls, kwargs: Dict[str, Any]) -> "FastVideoArgs":
+    def from_kwargs(cls, kwargs: dict[str, Any]) -> "FastVideoArgs":
         kwargs['pipeline_config'] = PipelineConfig.from_kwargs(kwargs)
         return cls(**kwargs)
 
@@ -325,7 +325,7 @@ class FastVideoArgs:
 _current_fastvideo_args = None
 
 
-def prepare_fastvideo_args(argv: List[str]) -> FastVideoArgs:
+def prepare_fastvideo_args(argv: list[str]) -> FastVideoArgs:
     """
     Prepare the inference arguments from the command line arguments.
 
@@ -410,7 +410,7 @@ class TrainingArgs(FastVideoArgs):
     log_validation: bool = False
     tracker_project_name: str = ""
     wandb_run_name: str = ""
-    seed: Optional[int] = None
+    seed: int | None = None
 
     # output
     output_dir: str = ""
@@ -427,7 +427,7 @@ class TrainingArgs(FastVideoArgs):
     lr_scheduler: str = "constant"
     lr_warmup_steps: int = 0
     max_grad_norm: float = 0.0
-    enable_gradient_checkpointing_type: Optional[str] = None
+    enable_gradient_checkpointing_type: str | None = None
     selective_checkpointing: float = 0.0
     allow_tf32: bool = False
     mixed_precision: str = ""

fastvideo/v1/forward_context.py

@@ -56,7 +56,7 @@ def get_forward_context() -> "ForwardContext":
 def set_forward_context(current_timestep,
                         attn_metadata,
                         forward_batch: Optional["ForwardBatch"] = None,
-                        fastvideo_args: Optional[FastVideoArgs] = None):
+                        fastvideo_args: FastVideoArgs | None = None):
     """A context manager that stores the current forward context,
     can be attention metadata, etc.
     Here we can inject common logic for every model forward pass.

fastvideo/v1/layers/custom_op.py

@@ -1,7 +1,8 @@
 # SPDX-License-Identifier: Apache-2.0
 # Adapted from vllm: https://github.com/vllm-project/vllm/blob/v0.7.3/vllm/model_executor/custom_op.py
 
-from typing import Any, Callable, Dict, Type
+from collections.abc import Callable
+from typing import Any
 
 import torch.nn as nn
 
@@ -81,7 +82,7 @@ class CustomOp(nn.Module):
     # Examples:
     # - MyOp.enabled()
     # - op_registry["my_op"].enabled()
-    op_registry: Dict[str, Type['CustomOp']] = {}
+    op_registry: dict[str, type['CustomOp']] = {}
 
     # Decorator to register custom ops.
     @classmethod

fastvideo/v1/layers/layernorm.py

@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # Adapted from vllm: https://github.com/vllm-project/vllm/blob/v0.7.3/vllm/model_executor/layers/layernorm.py
 """Custom normalization layers."""
-from typing import Optional, Tuple, Union
 
 import torch
 import torch.nn as nn
@@ -23,7 +22,7 @@ class RMSNorm(CustomOp):
         hidden_size: int,
         eps: float = 1e-6,
         dtype: torch.dtype = torch.float32,
-        var_hidden_size: Optional[int] = None,
+        var_hidden_size: int | None = None,
         has_weight: bool = True,
     ) -> None:
         super().__init__()
@@ -47,8 +46,8 @@ class RMSNorm(CustomOp):
     def forward_native(
         self,
         x: torch.Tensor,
-        residual: Optional[torch.Tensor] = None,
-    ) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:
+        residual: torch.Tensor | None = None,
+    ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
         """PyTorch-native implementation equivalent to forward()."""
         orig_dtype = x.dtype
         x = x.to(torch.float32)
@@ -159,7 +158,7 @@ class ScaleResidualLayerNormScaleShift(nn.Module):
 
     def forward(self, residual: torch.Tensor, x: torch.Tensor,
                 gate: torch.Tensor, shift: torch.Tensor,
-                scale: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
+                scale: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
         """
         Apply gated residual connection, followed by layernorm and 
         scale/shift in a single fused operation.

fastvideo/v1/layers/linear.py

@@ -2,7 +2,6 @@
 # Adapted from vllm: https://github.com/vllm-project/vllm/blob/v0.7.3/vllm/model_executor/layers/linear.py
 
 from abc import abstractmethod
-from typing import Optional, Union
 
 import torch
 import torch.nn.functional as F
@@ -39,7 +38,7 @@ WEIGHT_LOADER_V2_SUPPORTED = [
 
 def adjust_scalar_to_fused_array(
         param: torch.Tensor, loaded_weight: torch.Tensor,
-        shard_id: Union[str, int]) -> tuple[torch.Tensor, torch.Tensor]:
+        shard_id: str | int) -> tuple[torch.Tensor, torch.Tensor]:
     """For fused modules (QKV and MLP) we have an array of length
     N that holds 1 scale for each "logical" matrix. So the param
     is an array of length N. The loaded_weight corresponds to 
@@ -90,7 +89,7 @@ class LinearMethodBase(QuantizeMethodBase):
     def apply(self,
               layer: torch.nn.Module,
               x: torch.Tensor,
-              bias: Optional[torch.Tensor] = None) -> torch.Tensor:
+              bias: torch.Tensor | None = None) -> torch.Tensor:
         """Apply the weights in layer to the input tensor.
         Expects create_weights to have been called before on the layer."""
         raise NotImplementedError
@@ -115,7 +114,7 @@ class UnquantizedLinearMethod(LinearMethodBase):
     def apply(self,
               layer: torch.nn.Module,
               x: torch.Tensor,
-              bias: Optional[torch.Tensor] = None) -> torch.Tensor:
+              bias: torch.Tensor | None = None) -> torch.Tensor:
 
         return F.linear(x, layer.weight, bias)
 
@@ -137,8 +136,8 @@ class LinearBase(torch.nn.Module):
         input_size: int,
         output_size: int,
         skip_bias_add: bool = False,
-        params_dtype: Optional[torch.dtype] = None,
-        quant_config: Optional[QuantizationConfig] = None,
+        params_dtype: torch.dtype | None = None,
+        quant_config: QuantizationConfig | None = None,
         prefix: str = "",
     ):
         super().__init__()
@@ -151,14 +150,13 @@ class LinearBase(torch.nn.Module):
             params_dtype = torch.get_default_dtype()
         self.params_dtype = params_dtype
         if quant_config is None:
-            self.quant_method: Optional[
-                QuantizeMethodBase] = UnquantizedLinearMethod()
+            self.quant_method: QuantizeMethodBase | None = UnquantizedLinearMethod(
+            )
         else:
             self.quant_method = quant_config.get_quant_method(self,
                                                               prefix=prefix)
 
-    def forward(self,
-                x: torch.Tensor) -> tuple[torch.Tensor, Optional[Parameter]]:
+    def forward(self, x: torch.Tensor) -> tuple[torch.Tensor, Parameter | None]:
         raise NotImplementedError
 
 
@@ -181,8 +179,8 @@ class ReplicatedLinear(LinearBase):
                  output_size: int,
                  bias: bool = True,
                  skip_bias_add: bool = False,
-                 params_dtype: Optional[torch.dtype] = None,
-                 quant_config: Optional[QuantizationConfig] = None,
+                 params_dtype: torch.dtype | None = None,
+                 quant_config: QuantizationConfig | None = None,
                  prefix: str = ""):
         super().__init__(input_size,
                          output_size,
@@ -222,8 +220,7 @@ class ReplicatedLinear(LinearBase):
             f"to a parameter of size {param.size()}")
         param.data.copy_(loaded_weight)
 
-    def forward(self,
-                x: torch.Tensor) -> tuple[torch.Tensor, Optional[Parameter]]:
+    def forward(self, x: torch.Tensor) -> tuple[torch.Tensor, Parameter | None]:
         bias = self.bias if not self.skip_bias_add else None
         assert self.quant_method is not None
         output = self.quant_method.apply(self, x, bias)
@@ -267,9 +264,9 @@ class ColumnParallelLinear(LinearBase):
                  bias: bool = True,
                  gather_output: bool = False,
                  skip_bias_add: bool = False,
-                 params_dtype: Optional[torch.dtype] = None,
-                 quant_config: Optional[QuantizationConfig] = None,
-                 output_sizes: Optional[list[int]] = None,
+                 params_dtype: torch.dtype | None = None,
+                 quant_config: QuantizationConfig | None = None,
+                 output_sizes: list[int] | None = None,
                  prefix: str = ""):
         # Divide the weight matrix along the last dimension.
         self.tp_size = get_tp_world_size()
@@ -344,9 +341,8 @@ class ColumnParallelLinear(LinearBase):
             loaded_weight = loaded_weight.reshape(1)
         param.load_column_parallel_weight(loaded_weight=loaded_weight)
 
-    def forward(
-            self,
-            input_: torch.Tensor) -> tuple[torch.Tensor, Optional[Parameter]]:
+    def forward(self,
+                input_: torch.Tensor) -> tuple[torch.Tensor, Parameter | None]:
         bias = self.bias if not self.skip_bias_add else None
 
         # Matrix multiply.
@@ -398,8 +394,8 @@ class MergedColumnParallelLinear(ColumnParallelLinear):
                  bias: bool = True,
                  gather_output: bool = False,
                  skip_bias_add: bool = False,
-                 params_dtype: Optional[torch.dtype] = None,
-                 quant_config: Optional[QuantizationConfig] = None,
+                 params_dtype: torch.dtype | None = None,
+                 quant_config: QuantizationConfig | None = None,
                  prefix: str = ""):
         self.output_sizes = output_sizes
         tp_size = get_tp_world_size()
@@ -416,7 +412,7 @@ class MergedColumnParallelLinear(ColumnParallelLinear):
     def weight_loader(self,
                       param: Parameter,
                       loaded_weight: torch.Tensor,
-                      loaded_shard_id: Optional[int] = None) -> None:
+                      loaded_shard_id: int | None = None) -> None:
 
         param_data = param.data
         output_dim = getattr(param, "output_dim", None)
@@ -509,10 +505,8 @@ class MergedColumnParallelLinear(ColumnParallelLinear):
             # Special case for Quantization.
             # If quantized, we need to adjust the offset and size to account
             # for the packing.
-            if isinstance(
-                    param,
-                (PackedColumnParameter,
-                 PackedvLLMParameter)) and param.packed_dim == param.output_dim:
+            if isinstance(param, (PackedColumnParameter | PackedvLLMParameter
+                                  )) and param.packed_dim == param.output_dim:
                 shard_size, shard_offset = \
                     param.adjust_shard_indexes_for_packing(
                     shard_size=shard_size, shard_offset=shard_offset)
@@ -524,7 +518,7 @@ class MergedColumnParallelLinear(ColumnParallelLinear):
     def weight_loader_v2(self,
                          param: BasevLLMParameter,
                          loaded_weight: torch.Tensor,
-                         loaded_shard_id: Optional[int] = None) -> None:
+                         loaded_shard_id: int | None = None) -> None:
         if loaded_shard_id is None:
             if isinstance(param, PerTensorScaleParameter):
                 param.load_merged_column_weight(loaded_weight=loaded_weight,
@@ -597,11 +591,11 @@ class QKVParallelLinear(ColumnParallelLinear):
                  hidden_size: int,
                  head_size: int,
                  total_num_heads: int,
-                 total_num_kv_heads: Optional[int] = None,
+                 total_num_kv_heads: int | None = None,
                  bias: bool = True,
                  skip_bias_add: bool = False,
-                 params_dtype: Optional[torch.dtype] = None,
-                 quant_config: Optional[QuantizationConfig] = None,
+                 params_dtype: torch.dtype | None = None,
+                 quant_config: QuantizationConfig | None = None,
                  prefix: str = ""):
         self.hidden_size = hidden_size
         self.head_size = head_size
@@ -636,7 +630,7 @@ class QKVParallelLinear(ColumnParallelLinear):
                          quant_config=quant_config,
                          prefix=prefix)
 
-    def _get_shard_offset_mapping(self, loaded_shard_id: str) -> Optional[int]:
+    def _get_shard_offset_mapping(self, loaded_shard_id: str) -> int | None:
         shard_offset_mapping = {
             "q": 0,
             "k": self.num_heads * self.head_size,
@@ -645,7 +639,7 @@ class QKVParallelLinear(ColumnParallelLinear):
         }
         return shard_offset_mapping.get(loaded_shard_id)
 
-    def _get_shard_size_mapping(self, loaded_shard_id: str) -> Optional[int]:
+    def _get_shard_size_mapping(self, loaded_shard_id: str) -> int | None:
         shard_size_mapping = {
             "q": self.num_heads * self.head_size,
             "k": self.num_kv_heads * self.head_size,
@@ -678,10 +672,8 @@ class QKVParallelLinear(ColumnParallelLinear):
             # Special case for Quantization.
             # If quantized, we need to adjust the offset and size to account
             # for the packing.
-            if isinstance(
-                    param,
-                (PackedColumnParameter,
-                 PackedvLLMParameter)) and param.packed_dim == param.output_dim:
+            if isinstance(param, (PackedColumnParameter | PackedvLLMParameter
+                                  )) and param.packed_dim == param.output_dim:
                 shard_size, shard_offset = \
                     param.adjust_shard_indexes_for_packing(
                     shard_size=shard_size, shard_offset=shard_offset)
@@ -693,7 +685,7 @@ class QKVParallelLinear(ColumnParallelLinear):
     def weight_loader_v2(self,
                          param: BasevLLMParameter,
                          loaded_weight: torch.Tensor,
-                         loaded_shard_id: Optional[str] = None):
+                         loaded_shard_id: str | None = None):
         if loaded_shard_id is None:  # special case for certain models
             if isinstance(param, PerTensorScaleParameter):
                 param.load_qkv_weight(loaded_weight=loaded_weight, shard_id=0)
@@ -719,7 +711,7 @@ class QKVParallelLinear(ColumnParallelLinear):
     def weight_loader(self,
                       param: Parameter,
                       loaded_weight: torch.Tensor,
-                      loaded_shard_id: Optional[str] = None):
+                      loaded_shard_id: str | None = None):
 
         param_data = param.data
         output_dim = getattr(param, "output_dim", None)
@@ -844,9 +836,9 @@ class RowParallelLinear(LinearBase):
                  bias: bool = True,
                  input_is_parallel: bool = True,
                  skip_bias_add: bool = False,
-                 params_dtype: Optional[torch.dtype] = None,
+                 params_dtype: torch.dtype | None = None,
                  reduce_results: bool = True,
-                 quant_config: Optional[QuantizationConfig] = None,
+                 quant_config: QuantizationConfig | None = None,
                  prefix: str = ""):
         # Divide the weight matrix along the first dimension.
         self.tp_rank = get_tp_rank()
@@ -920,7 +912,7 @@ class RowParallelLinear(LinearBase):
 
         param.load_row_parallel_weight(loaded_weight=loaded_weight)
 
-    def forward(self, input_) -> tuple[torch.Tensor, Optional[Parameter]]:
+    def forward(self, input_) -> tuple[torch.Tensor, Parameter | None]:
         if self.input_is_parallel:
             input_parallel = input_
         else:

fastvideo/v1/layers/lora/linear.py

@@ -1,8 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # Code adapted from SGLang https://github.com/sgl-project/sglang/blob/main/python/sglang/srt/lora/layers.py
 
-from typing import Dict, List, Tuple, Type, Union
-
 import torch
 from torch import nn
 from torch.distributed.tensor import DTensor, distribute_tensor
@@ -201,7 +199,7 @@ class QKVParallelLinearWithLoRA(ColumnParallelLinearWithLoRA):
         return A
 
     def slice_lora_b_weights(
-            self, B: List[torch.Tensor]) -> Tuple[torch.Tensor, torch.Tensor]:
+            self, B: list[torch.Tensor]) -> tuple[torch.Tensor, torch.Tensor]:
         tp_rank = get_tp_rank()
         B_q, B_kv = B
         base_layer = self.base_layer
@@ -269,8 +267,8 @@ class RowParallelLinearWithLoRA(BaseLayerWithLoRA):
         return B
 
 
-def get_lora_layer(layer: nn.Module) -> Union[BaseLayerWithLoRA, None]:
-    supported_layer_types: Dict[Type[LinearBase], Type[BaseLayerWithLoRA]] = {
+def get_lora_layer(layer: nn.Module) -> BaseLayerWithLoRA | None:
+    supported_layer_types: dict[type[LinearBase], type[BaseLayerWithLoRA]] = {
         # the order matters
         # VocabParallelEmbedding: VocabParallelEmbeddingWithLoRA,
         QKVParallelLinear: QKVParallelLinearWithLoRA,

fastvideo/v1/layers/mlp.py

@@ -1,7 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
 
-from typing import Optional
-
 import torch
 import torch.nn as nn
 
@@ -18,10 +16,10 @@ class MLP(nn.Module):
         self,
         input_dim: int,
         mlp_hidden_dim: int,
-        output_dim: Optional[int] = None,
+        output_dim: int | None = None,
         bias: bool = True,
         act_type: str = "gelu_pytorch_tanh",
-        dtype: Optional[torch.dtype] = None,
+        dtype: torch.dtype | None = None,
         prefix: str = "",
     ):
         super().__init__()

fastvideo/v1/layers/quantization/base_config.py

@@ -3,7 +3,7 @@
 
 import inspect
 from abc import ABC, abstractmethod
-from typing import TYPE_CHECKING, Any, Optional
+from typing import TYPE_CHECKING, Any
 
 import torch
 from torch import nn
@@ -105,8 +105,8 @@ class QuantizationConfig(ABC):
         raise NotImplementedError
 
     @classmethod
-    def override_quantization_method(
-            cls, hf_quant_cfg, user_quant) -> Optional[QuantizationMethods]:
+    def override_quantization_method(cls, hf_quant_cfg,
+                                     user_quant) -> QuantizationMethods | None:
         """
            Detects if this quantization method can support a given checkpoint
            format by overriding the user specified quantization method -- 
@@ -135,7 +135,7 @@ class QuantizationConfig(ABC):
 
     @abstractmethod
     def get_quant_method(self, layer: torch.nn.Module,
-                         prefix: str) -> Optional[QuantizeMethodBase]:
+                         prefix: str) -> QuantizeMethodBase | None:
         """Get the quantize method to use for the quantized layer.
         
         Args:
@@ -147,5 +147,5 @@ class QuantizationConfig(ABC):
         """
         raise NotImplementedError
 
-    def get_cache_scale(self, name: str) -> Optional[str]:
-        return None
+    def get_cache_scale(self, name: str) -> str | None:
+        return None

fastvideo/v1/layers/rotary_embedding.py

@@ -23,7 +23,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 """Rotary Positional Embeddings."""
-from typing import Any, Dict, List, Optional, Tuple, Union
+from typing import Any
 
 import torch
 
@@ -84,7 +84,7 @@ class RotaryEmbedding(CustomOp):
         head_size: int,
         rotary_dim: int,
         max_position_embeddings: int,
-        base: Union[int, float],
+        base: int | float,
         is_neox_style: bool,
         dtype: torch.dtype,
     ) -> None:
@@ -101,7 +101,7 @@ class RotaryEmbedding(CustomOp):
         self.cos_sin_cache: torch.Tensor
         self.register_buffer("cos_sin_cache", cache, persistent=False)
 
-    def _compute_inv_freq(self, base: Union[int, float]) -> torch.Tensor:
+    def _compute_inv_freq(self, base: int | float) -> torch.Tensor:
         """Compute the inverse frequency."""
         # NOTE(woosuk): To exactly match the HF implementation, we need to
         # use CPU to compute the cache and then move it to GPU. However, we
@@ -127,8 +127,8 @@ class RotaryEmbedding(CustomOp):
         positions: torch.Tensor,
         query: torch.Tensor,
         key: torch.Tensor,
-        offsets: Optional[torch.Tensor] = None,
-    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        offsets: torch.Tensor | None = None,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
         """A PyTorch-native implementation of forward()."""
         if offsets is not None:
             positions = positions + offsets
@@ -159,7 +159,7 @@ class RotaryEmbedding(CustomOp):
         return s
 
 
-def _to_tuple(x: Union[int, Tuple[int, ...]], dim: int = 2) -> Tuple[int, ...]:
+def _to_tuple(x: int | tuple[int, ...], dim: int = 2) -> tuple[int, ...]:
     if isinstance(x, int):
         return (x, ) * dim
     elif len(x) == dim:
@@ -168,8 +168,8 @@ def _to_tuple(x: Union[int, Tuple[int, ...]], dim: int = 2) -> Tuple[int, ...]:
         raise ValueError(f"Expected length {dim} or int, but got {x}")
 
 
-def get_meshgrid_nd(start: Union[int, Tuple[int, ...]],
-                    *args: Union[int, Tuple[int, ...]],
+def get_meshgrid_nd(start: int | tuple[int, ...],
+                    *args: int | tuple[int, ...],
                     dim: int = 2) -> torch.Tensor:
     """
     Get n-D meshgrid with start, stop and num.
@@ -217,12 +217,12 @@ def get_meshgrid_nd(start: Union[int, Tuple[int, ...]],
 
 def get_1d_rotary_pos_embed(
     dim: int,
-    pos: Union[torch.FloatTensor, int],
+    pos: torch.FloatTensor | int,
     theta: float = 10000.0,
     theta_rescale_factor: float = 1.0,
     interpolation_factor: float = 1.0,
     dtype: torch.dtype = torch.float32,
-) -> Tuple[torch.Tensor, torch.Tensor]:
+) -> tuple[torch.Tensor, torch.Tensor]:
     """
     Precompute the frequency tensor for complex exponential (cis) with given dimensions.
     (Note: `cis` means `cos + i * sin`, where i is the imaginary unit.)
@@ -261,13 +261,13 @@ def get_nd_rotary_pos_embed(
     start,
     *args,
     theta=10000.0,
-    theta_rescale_factor: Union[float, List[float]] = 1.0,
-    interpolation_factor: Union[float, List[float]] = 1.0,
+    theta_rescale_factor: float | list[float] = 1.0,
+    interpolation_factor: float | list[float] = 1.0,
     shard_dim: int = 0,
     sp_rank: int = 0,
     sp_world_size: int = 1,
     dtype: torch.dtype = torch.float32,
-) -> Tuple[torch.Tensor, torch.Tensor]:
+) -> tuple[torch.Tensor, torch.Tensor]:
     """
     This is a n-d version of precompute_freqs_cis, which is a RoPE for tokens with n-d structure.
     Supports sequence parallelism by allowing sharding of a specific dimension.
@@ -324,7 +324,7 @@ def get_nd_rotary_pos_embed(
     else:
         grid = full_grid
 
-    if isinstance(theta_rescale_factor, (int, float)):
+    if isinstance(theta_rescale_factor, (int | float)):
         theta_rescale_factor = [theta_rescale_factor] * len(rope_dim_list)
     elif isinstance(theta_rescale_factor,
                     list) and len(theta_rescale_factor) == 1:
@@ -333,7 +333,7 @@ def get_nd_rotary_pos_embed(
         rope_dim_list
     ), "len(theta_rescale_factor) should equal to len(rope_dim_list)"
 
-    if isinstance(interpolation_factor, (int, float)):
+    if isinstance(interpolation_factor, (int | float)):
         interpolation_factor = [interpolation_factor] * len(rope_dim_list)
     elif isinstance(interpolation_factor,
                     list) and len(interpolation_factor) == 1:
@@ -370,7 +370,7 @@ def get_rotary_pos_embed(
     interpolation_factor=1.0,
     shard_dim: int = 0,
     dtype: torch.dtype = torch.float32,
-) -> Tuple[torch.Tensor, torch.Tensor]:
+) -> tuple[torch.Tensor, torch.Tensor]:
     """
     Generate rotary positional embeddings for the given sizes.
     
@@ -417,17 +417,17 @@ def get_rotary_pos_embed(
     return freqs_cos, freqs_sin
 
 
-_ROPE_DICT: Dict[Tuple, RotaryEmbedding] = {}
+_ROPE_DICT: dict[tuple, RotaryEmbedding] = {}
 
 
 def get_rope(
     head_size: int,
     rotary_dim: int,
     max_position: int,
-    base: Union[int, float],
+    base: int | float,
     is_neox_style: bool = True,
-    rope_scaling: Optional[Dict[str, Any]] = None,
-    dtype: Optional[torch.dtype] = None,
+    rope_scaling: dict[str, Any] | None = None,
+    dtype: torch.dtype | None = None,
     partial_rotary_factor: float = 1.0,
 ) -> RotaryEmbedding:
     if dtype is None:

fastvideo/v1/layers/utils.py

@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # Adapted from vllm: https://github.com/vllm-project/vllm/blob/v0.7.3/vllm/model_executor/layers/utils.py
 """Utility methods for model layers."""
-from typing import Tuple
 
 import torch
 
@@ -10,7 +9,7 @@ def get_token_bin_counts_and_mask(
     tokens: torch.Tensor,
     vocab_size: int,
     num_seqs: int,
-) -> Tuple[torch.Tensor, torch.Tensor]:
+) -> tuple[torch.Tensor, torch.Tensor]:
     # Compute the bin counts for the tokens.
     # vocab_size + 1 for padding.
     bin_counts = torch.zeros((num_seqs, vocab_size + 1),

fastvideo/v1/layers/visual_embedding.py

@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 
 import math
-from typing import Optional
 
 import torch
 import torch.nn as nn
@@ -36,7 +35,7 @@ class PatchEmbed(nn.Module):
                  prefix: str = ""):
         super().__init__()
         # Convert patch_size to 2-tuple
-        if isinstance(patch_size, (list, tuple)):
+        if isinstance(patch_size, (list | tuple)):
             if len(patch_size) == 1:
                 patch_size = (patch_size[0], patch_size[0])
         else:
@@ -133,7 +132,7 @@ class ModulateProjection(nn.Module):
         hidden_size: int,
         factor: int = 2,
         act_layer: str = "silu",
-        dtype: Optional[torch.dtype] = None,
+        dtype: torch.dtype | None = None,
         prefix: str = "",
     ):
         super().__init__()

fastvideo/v1/layers/vocab_parallel_embedding.py

@@ -1,7 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0
 
+from collections.abc import Sequence
 from dataclasses import dataclass
-from typing import List, Optional, Sequence, Tuple
 
 import torch
 import torch.nn.functional as F
@@ -23,7 +23,7 @@ class UnquantizedEmbeddingMethod(QuantizeMethodBase):
 
     def create_weights(self, layer: torch.nn.Module,
                        input_size_per_partition: int,
-                       output_partition_sizes: List[int], input_size: int,
+                       output_partition_sizes: list[int], input_size: int,
                        output_size: int, params_dtype: torch.dtype,
                        **extra_weight_attrs):
         """Create weights for embedding layer."""
@@ -38,7 +38,7 @@ class UnquantizedEmbeddingMethod(QuantizeMethodBase):
     def apply(self,
               layer: torch.nn.Module,
               x: torch.Tensor,
-              bias: Optional[torch.Tensor] = None) -> torch.Tensor:
+              bias: torch.Tensor | None = None) -> torch.Tensor:
         return F.linear(x, layer.weight, bias)
 
     def embedding(self, layer: torch.nn.Module,
@@ -138,7 +138,7 @@ def get_masked_input_and_mask(
         input_: torch.Tensor, org_vocab_start_index: int,
         org_vocab_end_index: int, num_org_vocab_padding: int,
         added_vocab_start_index: int,
-        added_vocab_end_index: int) -> Tuple[torch.Tensor, torch.Tensor]:
+        added_vocab_end_index: int) -> tuple[torch.Tensor, torch.Tensor]:
     # torch.compile will fuse all of the pointwise ops below
     # into a single kernel, making it very fast
     org_vocab_mask = (input_ >= org_vocab_start_index) & (input_
@@ -196,10 +196,10 @@ class VocabParallelEmbedding(torch.nn.Module):
     def __init__(self,
                  num_embeddings: int,
                  embedding_dim: int,
-                 params_dtype: Optional[torch.dtype] = None,
-                 org_num_embeddings: Optional[int] = None,
+                 params_dtype: torch.dtype | None = None,
+                 org_num_embeddings: int | None = None,
                  padding_size: int = DEFAULT_VOCAB_PADDING_SIZE,
-                 quant_config: Optional[QuantizationConfig] = None,
+                 quant_config: QuantizationConfig | None = None,
                  prefix: str = ""):
         super().__init__()
 
@@ -295,7 +295,7 @@ class VocabParallelEmbedding(torch.nn.Module):
             org_vocab_start_index, org_vocab_end_index, added_vocab_start_index,
             added_vocab_end_index)
 
-    def get_sharded_to_full_mapping(self) -> Optional[List[int]]:
+    def get_sharded_to_full_mapping(self) -> list[int] | None:
         """Get a mapping that can be used to reindex the gathered
         logits for sampling.
         
@@ -309,9 +309,9 @@ class VocabParallelEmbedding(torch.nn.Module):
         if self.tp_size < 2:
             return None
 
-        base_embeddings: List[int] = []
-        added_embeddings: List[int] = []
-        padding: List[int] = []
+        base_embeddings: list[int] = []
+        added_embeddings: list[int] = []
+        padding: list[int] = []
         for tp_rank in range(self.tp_size):
             shard_indices = self._get_indices(self.num_embeddings_padded,
                                               self.org_vocab_size_padded,

fastvideo/v1/logger.py

@@ -11,7 +11,7 @@ from logging import Logger
 from logging.config import dictConfig
 from os import path
 from types import MethodType
-from typing import Any, Optional, cast
+from typing import Any, cast
 
 import fastvideo.v1.envs as envs
 
@@ -278,8 +278,7 @@ def _trace_calls(log_path, root_dir, frame, event, arg=None):
     return partial(_trace_calls, log_path, root_dir)
 
 
-def enable_trace_function_call(log_file_path: str,
-                               root_dir: Optional[str] = None):
+def enable_trace_function_call(log_file_path: str, root_dir: str | None = None):
     """
     Enable tracing of every function call in code under `root_dir`.
     This is useful for debugging hangs or crashes.

fastvideo/v1/models/dits/base.py

@@ -1,6 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 from abc import ABC, abstractmethod
-from typing import Any, List, Optional, Tuple, Union
+from typing import Any
 
 import torch
 from torch import nn
@@ -19,7 +19,7 @@ class BaseDiT(nn.Module, ABC):
     num_attention_heads: int
     num_channels_latents: int
     # always supports torch_sdpa
-    _supported_attention_backends: Tuple[
+    _supported_attention_backends: tuple[
         AttentionBackendEnum, ...] = DiTConfig()._supported_attention_backends
 
     def __init_subclass__(cls) -> None:
@@ -47,10 +47,10 @@ class BaseDiT(nn.Module, ABC):
     @abstractmethod
     def forward(self,
                 hidden_states: torch.Tensor,
-                encoder_hidden_states: Union[torch.Tensor, List[torch.Tensor]],
+                encoder_hidden_states: torch.Tensor | list[torch.Tensor],
                 timestep: torch.LongTensor,
-                encoder_hidden_states_image: Optional[Union[
-                    torch.Tensor, List[torch.Tensor]]] = None,
+                encoder_hidden_states_image: torch.Tensor | list[torch.Tensor]
+                | None = None,
                 guidance=None,
                 **kwargs) -> torch.Tensor:
         pass
@@ -66,7 +66,7 @@ class BaseDiT(nn.Module, ABC):
                 )
 
     @property
-    def supported_attention_backends(self) -> Tuple[AttentionBackendEnum, ...]:
+    def supported_attention_backends(self) -> tuple[AttentionBackendEnum, ...]:
         return self._supported_attention_backends
 
 
@@ -86,7 +86,7 @@ class CachableDiT(BaseDiT):
     num_attention_heads: int
     num_channels_latents: int
     # always supports torch_sdpa
-    _supported_attention_backends: Tuple[
+    _supported_attention_backends: tuple[
         AttentionBackendEnum, ...] = DiTConfig()._supported_attention_backends
 
     def __init__(self, config: DiTConfig, **kwargs) -> None:

fastvideo/v1/models/dits/hunyuanvideo.py

@@ -1,6 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 
-from typing import Any, List, Optional, Tuple, Union
+from typing import Any
 
 import numpy as np
 import torch
@@ -95,9 +95,9 @@ class MMDoubleStreamBlock(nn.Module):
         hidden_size: int,
         num_attention_heads: int,
         mlp_ratio: float,
-        dtype: Optional[torch.dtype] = None,
-        supported_attention_backends: Optional[Tuple[AttentionBackendEnum,
-                                                     ...]] = None,
+        dtype: torch.dtype | None = None,
+        supported_attention_backends: tuple[AttentionBackendEnum, ...]
+        | None = None,
         prefix: str = "",
     ):
         super().__init__()
@@ -202,7 +202,7 @@ class MMDoubleStreamBlock(nn.Module):
         txt: torch.Tensor,
         vec: torch.Tensor,
         freqs_cis: tuple,
-    ) -> Tuple[torch.Tensor, torch.Tensor]:
+    ) -> tuple[torch.Tensor, torch.Tensor]:
         # Process modulation vectors
         img_mod_outputs = self.img_mod(vec)
         (
@@ -303,9 +303,9 @@ class MMSingleStreamBlock(nn.Module):
         hidden_size: int,
         num_attention_heads: int,
         mlp_ratio: float = 4.0,
-        dtype: Optional[torch.dtype] = None,
-        supported_attention_backends: Optional[Tuple[AttentionBackendEnum,
-                                                     ...]] = None,
+        dtype: torch.dtype | None = None,
+        supported_attention_backends: tuple[AttentionBackendEnum, ...]
+        | None = None,
         prefix: str = "",
     ):
         super().__init__()
@@ -367,7 +367,7 @@ class MMSingleStreamBlock(nn.Module):
         x: torch.Tensor,
         vec: torch.Tensor,
         txt_len: int,
-        freqs_cis: Tuple[torch.Tensor, torch.Tensor],
+        freqs_cis: tuple[torch.Tensor, torch.Tensor],
     ) -> torch.Tensor:
         # Process modulation
         mod_shift, mod_scale, mod_gate = self.modulation(vec).chunk(3, dim=-1)
@@ -546,10 +546,10 @@ class HunyuanVideoTransformer3DModel(CachableDiT):
     # TODO: change output to a dict
     def forward(self,
                 hidden_states: torch.Tensor,
-                encoder_hidden_states: Union[torch.Tensor, List[torch.Tensor]],
+                encoder_hidden_states: torch.Tensor | list[torch.Tensor],
                 timestep: torch.LongTensor,
-                encoder_hidden_states_image: Optional[Union[
-                    torch.Tensor, List[torch.Tensor]]] = None,
+                encoder_hidden_states_image: torch.Tensor | list[torch.Tensor]
+                | None = None,
                 guidance=None,
                 **kwargs):
         """

fastvideo/v1/models/dits/stepvideo.py

@@ -10,7 +10,7 @@
 # The above copyright notice and this permission notice shall be included in all
 # copies or substantial portions of the Software.
 # ==============================================================================
-from typing import Any, Dict, Optional, Tuple
+from typing import Any
 
 import torch
 from einops import rearrange, repeat
@@ -54,7 +54,7 @@ class PatchEmbed2D(nn.Module):
                  prefix: str = ""):
         super().__init__()
         # Convert patch_size to 2-tuple
-        if isinstance(patch_size, (list, tuple)):
+        if isinstance(patch_size, (list | tuple)):
             if len(patch_size) == 1:
                 patch_size = (patch_size[0], patch_size[0])
         else:
@@ -143,7 +143,7 @@ class SelfAttention(nn.Module):
         self,
         hidden_dim,
         head_dim,
-        rope_split: Tuple[int, int, int] = (64, 32, 32),
+        rope_split: tuple[int, int, int] = (64, 32, 32),
         bias: bool = False,
         with_rope: bool = True,
         with_qk_norm: bool = True,
@@ -190,8 +190,10 @@ class SelfAttention(nn.Module):
 
         outs = []
         idx = 0
-        for (chunk_size, cos_i, sin_i) in zip(self.rope_split, cos_splits,
-                                              sin_splits):
+        for (chunk_size, cos_i, sin_i) in zip(self.rope_split,
+                                              cos_splits,
+                                              sin_splits,
+                                              strict=True):
             # slice the corresponding channels
             x_chunk = x[..., idx:idx + chunk_size]  # [B,S,H,chunk_size]
             idx += chunk_size
@@ -332,8 +334,8 @@ class AdaLayerNormSingle(nn.Module):
     def forward(
         self,
         timestep: torch.Tensor,
-        added_cond_kwargs: Optional[Dict[str, torch.Tensor]] = None,
-    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        added_cond_kwargs: dict[str, torch.Tensor] | None = None,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
         embedded_timestep = self.emb(timestep * self.time_step_rescale)
 
         out, _ = self.linear(self.silu(embedded_timestep))
@@ -378,7 +380,7 @@ class StepVideoTransformerBlock(nn.Module):
                  dim: int,
                  attention_head_dim: int,
                  norm_eps: float = 1e-5,
-                 ff_inner_dim: Optional[int] = None,
+                 ff_inner_dim: int | None = None,
                  ff_bias: bool = False,
                  attention_type: str = 'torch'):
         super().__init__()
@@ -418,7 +420,7 @@ class StepVideoTransformerBlock(nn.Module):
                 kv: torch.Tensor,
                 t_expand: torch.LongTensor,
                 attn_mask=None,
-                rope_positions: Optional[list] = None,
+                rope_positions: list | None = None,
                 cos_sin=None,
                 mask_strategy=None) -> torch.Tensor:
 
@@ -540,7 +542,7 @@ class StepVideoModel(BaseDiT):
         return hidden_states
 
     def prepare_attn_mask(self, encoder_attention_mask, encoder_hidden_states,
-                          q_seqlen) -> Tuple[torch.Tensor, torch.Tensor]:
+                          q_seqlen) -> tuple[torch.Tensor, torch.Tensor]:
         kv_seqlens = encoder_attention_mask.sum(dim=1).int()
         mask = torch.zeros([len(kv_seqlens), q_seqlen,
                             max(kv_seqlens)],
@@ -594,12 +596,12 @@ class StepVideoModel(BaseDiT):
     def forward(
         self,
         hidden_states: torch.Tensor,
-        encoder_hidden_states: Optional[torch.Tensor] = None,
-        t_expand: Optional[torch.LongTensor] = None,
-        encoder_hidden_states_2: Optional[torch.Tensor] = None,
-        added_cond_kwargs: Optional[Dict[str, torch.Tensor]] = None,
-        encoder_attention_mask: Optional[torch.Tensor] = None,
-        fps: Optional[torch.Tensor] = None,
+        encoder_hidden_states: torch.Tensor | None = None,
+        t_expand: torch.LongTensor | None = None,
+        encoder_hidden_states_2: torch.Tensor | None = None,
+        added_cond_kwargs: dict[str, torch.Tensor] | None = None,
+        encoder_attention_mask: torch.Tensor | None = None,
+        fps: torch.Tensor | None = None,
         return_dict: bool = True,
         mask_strategy=None,
         guidance=None,

fastvideo/v1/models/dits/wanvideo.py

@@ -1,7 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0
 
 import math
-from typing import Any, List, Optional, Tuple, Union
+from typing import Any
 
 import numpy as np
 import torch
@@ -54,7 +54,7 @@ class WanTimeTextImageEmbedding(nn.Module):
         dim: int,
         time_freq_dim: int,
         text_embed_dim: int,
-        image_embed_dim: Optional[int] = None,
+        image_embed_dim: int | None = None,
     ):
         super().__init__()
 
@@ -77,7 +77,7 @@ class WanTimeTextImageEmbedding(nn.Module):
         self,
         timestep: torch.Tensor,
         encoder_hidden_states: torch.Tensor,
-        encoder_hidden_states_image: Optional[torch.Tensor] = None,
+        encoder_hidden_states_image: torch.Tensor | None = None,
     ):
         temb = self.time_embedder(timestep)
         timestep_proj = self.time_modulation(temb)
@@ -174,8 +174,8 @@ class WanI2VCrossAttention(WanSelfAttention):
         window_size=(-1, -1),
         qk_norm=True,
         eps=1e-6,
-        supported_attention_backends: Optional[Tuple[AttentionBackendEnum,
-                                                     ...]] = None
+        supported_attention_backends: tuple[AttentionBackendEnum, ...]
+        | None = None
     ) -> None:
         super().__init__(dim, num_heads, window_size, qk_norm, eps,
                          supported_attention_backends)
@@ -217,18 +217,17 @@ class WanI2VCrossAttention(WanSelfAttention):
 
 class WanTransformerBlock(nn.Module):
 
-    def __init__(
-            self,
-            dim: int,
-            ffn_dim: int,
-            num_heads: int,
-            qk_norm: str = "rms_norm_across_heads",
-            cross_attn_norm: bool = False,
-            eps: float = 1e-6,
-            added_kv_proj_dim: Optional[int] = None,
-            supported_attention_backends: Optional[Tuple[AttentionBackendEnum,
-                                                         ...]] = None,
-            prefix: str = ""):
+    def __init__(self,
+                 dim: int,
+                 ffn_dim: int,
+                 num_heads: int,
+                 qk_norm: str = "rms_norm_across_heads",
+                 cross_attn_norm: bool = False,
+                 eps: float = 1e-6,
+                 added_kv_proj_dim: int | None = None,
+                 supported_attention_backends: tuple[AttentionBackendEnum, ...]
+                 | None = None,
+                 prefix: str = ""):
         super().__init__()
 
         # 1. Self-attention
@@ -298,7 +297,7 @@ class WanTransformerBlock(nn.Module):
         hidden_states: torch.Tensor,
         encoder_hidden_states: torch.Tensor,
         temb: torch.Tensor,
-        freqs_cis: Tuple[torch.Tensor, torch.Tensor],
+        freqs_cis: tuple[torch.Tensor, torch.Tensor],
     ) -> torch.Tensor:
         if hidden_states.dim() == 4:
             hidden_states = hidden_states.squeeze(1)
@@ -362,18 +361,17 @@ class WanTransformerBlock(nn.Module):
 
 class WanTransformerBlock_VSA(nn.Module):
 
-    def __init__(
-            self,
-            dim: int,
-            ffn_dim: int,
-            num_heads: int,
-            qk_norm: str = "rms_norm_across_heads",
-            cross_attn_norm: bool = False,
-            eps: float = 1e-6,
-            added_kv_proj_dim: Optional[int] = None,
-            supported_attention_backends: Optional[Tuple[AttentionBackendEnum,
-                                                         ...]] = None,
-            prefix: str = ""):
+    def __init__(self,
+                 dim: int,
+                 ffn_dim: int,
+                 num_heads: int,
+                 qk_norm: str = "rms_norm_across_heads",
+                 cross_attn_norm: bool = False,
+                 eps: float = 1e-6,
+                 added_kv_proj_dim: int | None = None,
+                 supported_attention_backends: tuple[AttentionBackendEnum, ...]
+                 | None = None,
+                 prefix: str = ""):
         super().__init__()
 
         # 1. Self-attention
@@ -444,7 +442,7 @@ class WanTransformerBlock_VSA(nn.Module):
         hidden_states: torch.Tensor,
         encoder_hidden_states: torch.Tensor,
         temb: torch.Tensor,
-        freqs_cis: Tuple[torch.Tensor, torch.Tensor],
+        freqs_cis: tuple[torch.Tensor, torch.Tensor],
     ) -> torch.Tensor:
         if hidden_states.dim() == 4:
             hidden_states = hidden_states.squeeze(1)
@@ -593,10 +591,10 @@ class WanTransformer3DModel(CachableDiT):
 
     def forward(self,
                 hidden_states: torch.Tensor,
-                encoder_hidden_states: Union[torch.Tensor, List[torch.Tensor]],
+                encoder_hidden_states: torch.Tensor | list[torch.Tensor],
                 timestep: torch.LongTensor,
-                encoder_hidden_states_image: Optional[Union[
-                    torch.Tensor, List[torch.Tensor]]] = None,
+                encoder_hidden_states_image: torch.Tensor | list[torch.Tensor]
+                | None = None,
                 guidance=None,
                 **kwargs) -> torch.Tensor:
         forward_batch = get_forward_context().forward_batch

fastvideo/v1/models/encoders/base.py

@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 from abc import ABC, abstractmethod
 from dataclasses import field
-from typing import List, Optional, Tuple
 
 import torch
 from torch import nn
@@ -14,9 +13,9 @@ from fastvideo.v1.platforms import AttentionBackendEnum
 
 class TextEncoder(nn.Module, ABC):
     _fsdp_shard_conditions: list = field(default_factory=lambda: [])
-    _stacked_params_mapping: List[Tuple[str, str,
+    _stacked_params_mapping: list[tuple[str, str,
                                         str]] = field(default_factory=list)
-    _supported_attention_backends: Tuple[
+    _supported_attention_backends: tuple[
         AttentionBackendEnum,
         ...] = TextEncoderConfig()._supported_attention_backends
 
@@ -32,21 +31,21 @@ class TextEncoder(nn.Module, ABC):
 
     @abstractmethod
     def forward(self,
-                input_ids: Optional[torch.Tensor],
-                position_ids: Optional[torch.Tensor] = None,
-                attention_mask: Optional[torch.Tensor] = None,
-                inputs_embeds: Optional[torch.Tensor] = None,
-                output_hidden_states: Optional[bool] = None,
+                input_ids: torch.Tensor | None,
+                position_ids: torch.Tensor | None = None,
+                attention_mask: torch.Tensor | None = None,
+                inputs_embeds: torch.Tensor | None = None,
+                output_hidden_states: bool | None = None,
                 **kwargs) -> BaseEncoderOutput:
         pass
 
     @property
-    def supported_attention_backends(self) -> Tuple[AttentionBackendEnum, ...]:
+    def supported_attention_backends(self) -> tuple[AttentionBackendEnum, ...]:
         return self._supported_attention_backends
 
 
 class ImageEncoder(nn.Module, ABC):
-    _supported_attention_backends: Tuple[
+    _supported_attention_backends: tuple[
         AttentionBackendEnum,
         ...] = ImageEncoderConfig()._supported_attention_backends
 
@@ -64,5 +63,5 @@ class ImageEncoder(nn.Module, ABC):
         pass
 
     @property
-    def supported_attention_backends(self) -> Tuple[AttentionBackendEnum, ...]:
+    def supported_attention_backends(self) -> tuple[AttentionBackendEnum, ...]:
         return self._supported_attention_backends

fastvideo/v1/models/encoders/clip.py

@@ -3,7 +3,7 @@
 # Adapted from transformers: https://github.com/huggingface/transformers/blob/v4.39.0/src/transformers/models/clip/modeling_clip.py
 """Minimal implementation of CLIPVisionModel intended to be only used
 within a vision language model."""
-from typing import Iterable, Optional, Set, Tuple, Union
+from collections.abc import Iterable
 
 import torch
 import torch.nn as nn
@@ -90,9 +90,9 @@ class CLIPTextEmbeddings(nn.Module):
 
     def forward(
         self,
-        input_ids: Optional[torch.LongTensor] = None,
-        position_ids: Optional[torch.LongTensor] = None,
-        inputs_embeds: Optional[torch.FloatTensor] = None,
+        input_ids: torch.LongTensor | None = None,
+        position_ids: torch.LongTensor | None = None,
+        inputs_embeds: torch.FloatTensor | None = None,
     ) -> torch.Tensor:
         if input_ids is not None:
             seq_length = input_ids.shape[-1]
@@ -127,8 +127,8 @@ class CLIPAttention(nn.Module):
 
     def __init__(
         self,
-        config: Union[CLIPVisionConfig, CLIPTextConfig],
-        quant_config: Optional[QuantizationConfig] = None,
+        config: CLIPVisionConfig | CLIPTextConfig,
+        quant_config: QuantizationConfig | None = None,
         prefix: str = "",
     ):
         super().__init__()
@@ -208,8 +208,8 @@ class CLIPMLP(nn.Module):
 
     def __init__(
         self,
-        config: Union[CLIPVisionConfig, CLIPTextConfig],
-        quant_config: Optional[QuantizationConfig] = None,
+        config: CLIPVisionConfig | CLIPTextConfig,
+        quant_config: QuantizationConfig | None = None,
         prefix: str = "",
     ) -> None:
         super().__init__()
@@ -238,8 +238,8 @@ class CLIPEncoderLayer(nn.Module):
 
     def __init__(
         self,
-        config: Union[CLIPTextConfig, CLIPVisionConfig],
-        quant_config: Optional[QuantizationConfig] = None,
+        config: CLIPTextConfig | CLIPVisionConfig,
+        quant_config: QuantizationConfig | None = None,
         prefix: str = "",
     ) -> None:
         super().__init__()
@@ -283,9 +283,9 @@ class CLIPEncoder(nn.Module):
 
     def __init__(
         self,
-        config: Union[CLIPVisionConfig, CLIPTextConfig],
-        quant_config: Optional[QuantizationConfig] = None,
-        num_hidden_layers_override: Optional[int] = None,
+        config: CLIPVisionConfig | CLIPTextConfig,
+        quant_config: QuantizationConfig | None = None,
+        num_hidden_layers_override: int | None = None,
         prefix: str = "",
     ) -> None:
         super().__init__()
@@ -304,8 +304,8 @@ class CLIPEncoder(nn.Module):
         ])
 
     def forward(
-        self, inputs_embeds: torch.Tensor, return_all_hidden_states: bool
-    ) -> Union[torch.Tensor, list[torch.Tensor]]:
+            self, inputs_embeds: torch.Tensor, return_all_hidden_states: bool
+    ) -> torch.Tensor | list[torch.Tensor]:
         hidden_states_pool = [inputs_embeds]
         hidden_states = inputs_embeds
 
@@ -324,8 +324,8 @@ class CLIPTextTransformer(nn.Module):
 
     def __init__(self,
                  config: CLIPTextConfig,
-                 quant_config: Optional[QuantizationConfig] = None,
-                 num_hidden_layers_override: Optional[int] = None,
+                 quant_config: QuantizationConfig | None = None,
+                 num_hidden_layers_override: int | None = None,
                  prefix: str = ""):
         super().__init__()
         self.config = config
@@ -347,11 +347,11 @@ class CLIPTextTransformer(nn.Module):
 
     def forward(
         self,
-        input_ids: Optional[torch.Tensor],
-        position_ids: Optional[torch.Tensor] = None,
-        attention_mask: Optional[torch.Tensor] = None,
-        inputs_embeds: Optional[torch.Tensor] = None,
-        output_hidden_states: Optional[bool] = None,
+        input_ids: torch.Tensor | None,
+        position_ids: torch.Tensor | None = None,
+        attention_mask: torch.Tensor | None = None,
+        inputs_embeds: torch.Tensor | None = None,
+        output_hidden_states: bool | None = None,
     ) -> BaseEncoderOutput:
         r"""
         Returns:
@@ -439,11 +439,11 @@ class CLIPTextModel(TextEncoder):
 
     def forward(
         self,
-        input_ids: Optional[torch.Tensor],
-        position_ids: Optional[torch.Tensor] = None,
-        attention_mask: Optional[torch.Tensor] = None,
-        inputs_embeds: Optional[torch.Tensor] = None,
-        output_hidden_states: Optional[bool] = None,
+        input_ids: torch.Tensor | None,
+        position_ids: torch.Tensor | None = None,
+        attention_mask: torch.Tensor | None = None,
+        inputs_embeds: torch.Tensor | None = None,
+        output_hidden_states: bool | None = None,
         **kwargs,
     ) -> BaseEncoderOutput:
 
@@ -455,8 +455,8 @@ class CLIPTextModel(TextEncoder):
         )
         return outputs
 
-    def load_weights(self, weights: Iterable[Tuple[str,
-                                                   torch.Tensor]]) -> Set[str]:
+    def load_weights(self, weights: Iterable[tuple[str,
+                                                   torch.Tensor]]) -> set[str]:
 
         # Define mapping for stacked parameters
         stacked_params_mapping = [
@@ -466,7 +466,7 @@ class CLIPTextModel(TextEncoder):
             ("qkv_proj", "v_proj", "v"),
         ]
         params_dict = dict(self.named_parameters())
-        loaded_params: Set[str] = set()
+        loaded_params: set[str] = set()
         for name, loaded_weight in weights:
             # Handle q_proj, k_proj, v_proj -> qkv_proj mapping
             for param_name, weight_name, shard_id in stacked_params_mapping:
@@ -497,9 +497,9 @@ class CLIPVisionTransformer(nn.Module):
     def __init__(
         self,
         config: CLIPVisionConfig,
-        quant_config: Optional[QuantizationConfig] = None,
-        num_hidden_layers_override: Optional[int] = None,
-        require_post_norm: Optional[bool] = None,
+        quant_config: QuantizationConfig | None = None,
+        num_hidden_layers_override: int | None = None,
+        require_post_norm: bool | None = None,
         prefix: str = "",
     ) -> None:
         super().__init__()
@@ -539,7 +539,7 @@ class CLIPVisionTransformer(nn.Module):
     def forward(
         self,
         pixel_values: torch.Tensor,
-        feature_sample_layers: Optional[list[int]] = None,
+        feature_sample_layers: list[int] | None = None,
     ) -> torch.Tensor:
 
         hidden_states = self.embeddings(pixel_values)
@@ -581,7 +581,7 @@ class CLIPVisionModel(ImageEncoder):
     def forward(
         self,
         pixel_values: torch.Tensor,
-        feature_sample_layers: Optional[list[int]] = None,
+        feature_sample_layers: list[int] | None = None,
         **kwargs,
     ) -> BaseEncoderOutput:
         last_hidden_state = self.vision_model(pixel_values,
@@ -594,11 +594,11 @@ class CLIPVisionModel(ImageEncoder):
 
     # (TODO) Add prefix argument for filtering out weights to be loaded
     #        ref: https://github.com/vllm-project/vllm/pull/7186#discussion_r1734163986
-    def load_weights(self, weights: Iterable[Tuple[str,
-                                                   torch.Tensor]]) -> Set[str]:
+    def load_weights(self, weights: Iterable[tuple[str,
+                                                   torch.Tensor]]) -> set[str]:
 
         params_dict = dict(self.named_parameters())
-        loaded_params: Set[str] = set()
+        loaded_params: set[str] = set()
         layer_count = len(self.vision_model.encoder.layers)
 
         for name, loaded_weight in weights:

fastvideo/v1/models/encoders/llama.py

@@ -23,7 +23,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 """Inference-only LLaMA model compatible with HuggingFace weights."""
-from typing import Any, Dict, Iterable, Optional, Set, Tuple
+from collections.abc import Iterable
+from typing import Any
 
 import torch
 from torch import nn
@@ -52,7 +53,7 @@ class LlamaMLP(nn.Module):
         hidden_size: int,
         intermediate_size: int,
         hidden_act: str,
-        quant_config: Optional[QuantizationConfig] = None,
+        quant_config: QuantizationConfig | None = None,
         bias: bool = False,
         prefix: str = "",
     ) -> None:
@@ -92,9 +93,9 @@ class LlamaAttention(nn.Module):
                  num_heads: int,
                  num_kv_heads: int,
                  rope_theta: float = 10000,
-                 rope_scaling: Optional[Dict[str, Any]] = None,
+                 rope_scaling: dict[str, Any] | None = None,
                  max_position_embeddings: int = 8192,
-                 quant_config: Optional[QuantizationConfig] = None,
+                 quant_config: QuantizationConfig | None = None,
                  bias: bool = False,
                  bias_o_proj: bool = False,
                  prefix: str = "") -> None:
@@ -201,7 +202,7 @@ class LlamaDecoderLayer(nn.Module):
     def __init__(
         self,
         config: LlamaConfig,
-        quant_config: Optional[QuantizationConfig] = None,
+        quant_config: QuantizationConfig | None = None,
         prefix: str = "",
     ) -> None:
         super().__init__()
@@ -254,8 +255,8 @@ class LlamaDecoderLayer(nn.Module):
         self,
         positions: torch.Tensor,
         hidden_states: torch.Tensor,
-        residual: Optional[torch.Tensor],
-    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        residual: torch.Tensor | None,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
         # Self Attention
         if residual is None:
             residual = hidden_states
@@ -318,11 +319,11 @@ class LlamaModel(TextEncoder):
 
     def forward(
         self,
-        input_ids: Optional[torch.Tensor],
-        position_ids: Optional[torch.Tensor] = None,
-        attention_mask: Optional[torch.Tensor] = None,
-        inputs_embeds: Optional[torch.Tensor] = None,
-        output_hidden_states: Optional[bool] = None,
+        input_ids: torch.Tensor | None,
+        position_ids: torch.Tensor | None = None,
+        attention_mask: torch.Tensor | None = None,
+        inputs_embeds: torch.Tensor | None = None,
+        output_hidden_states: bool | None = None,
         **kwargs,
     ) -> BaseEncoderOutput:
         output_hidden_states = (output_hidden_states
@@ -339,7 +340,7 @@ class LlamaModel(TextEncoder):
                 0, hidden_states.shape[1],
                 device=hidden_states.device).unsqueeze(0)
 
-        all_hidden_states: Optional[Tuple[Any, ...]] = (
+        all_hidden_states: tuple[Any, ...] | None = (
         ) if output_hidden_states else None
         for layer in self.layers:
             if all_hidden_states is not None:
@@ -367,11 +368,11 @@ class LlamaModel(TextEncoder):
 
         return output
 
-    def load_weights(self, weights: Iterable[Tuple[str,
-                                                   torch.Tensor]]) -> Set[str]:
+    def load_weights(self, weights: Iterable[tuple[str,
+                                                   torch.Tensor]]) -> set[str]:
 
         params_dict = dict(self.named_parameters())
-        loaded_params: Set[str] = set()
+        loaded_params: set[str] = set()
         for name, loaded_weight in weights:
             if "rotary_emb.inv_freq" in name:
                 continue
@@ -393,7 +394,7 @@ class LlamaModel(TextEncoder):
             #     continue
             if "scale" in name:
                 # Remapping the name of FP8 kv-scale.
-                kv_scale_name: Optional[str] = maybe_remap_kv_scale_name(
+                kv_scale_name: str | None = maybe_remap_kv_scale_name(
                     name, params_dict)
                 if kv_scale_name is None:
                     continue

fastvideo/v1/models/encoders/stepllm.py

@@ -14,7 +14,6 @@
 # ==============================================================================
 import os
 from functools import wraps
-from typing import List, Optional
 
 import torch
 import torch.nn as nn
@@ -180,10 +179,10 @@ class StepChatTokenizer:
     def vocab_size(self):
         return self._tokenizer.vocab_size()
 
-    def tokenize(self, text: str) -> List[int]:
+    def tokenize(self, text: str) -> list[int]:
         return self._tokenizer.encode_as_ids(text)
 
-    def detokenize(self, token_ids: List[int]) -> str:
+    def detokenize(self, token_ids: list[int]) -> str:
         return self._tokenizer.decode_ids(token_ids)
 
 
@@ -348,9 +347,9 @@ class MultiQueryAttention(nn.Module):
     def forward(
         self,
         x: torch.Tensor,
-        mask: Optional[torch.Tensor],
-        cu_seqlens: Optional[torch.Tensor],
-        max_seq_len: Optional[torch.Tensor],
+        mask: torch.Tensor | None,
+        cu_seqlens: torch.Tensor | None,
+        max_seq_len: torch.Tensor | None,
     ):
         seqlen, bsz, dim = x.shape
         xqkv = self.wqkv(x)
@@ -472,9 +471,9 @@ class TransformerBlock(nn.Module):
     def forward(
         self,
         x: torch.Tensor,
-        mask: Optional[torch.Tensor],
-        cu_seqlens: Optional[torch.Tensor],
-        max_seq_len: Optional[torch.Tensor],
+        mask: torch.Tensor | None,
+        cu_seqlens: torch.Tensor | None,
+        max_seq_len: torch.Tensor | None,
     ):
         residual = self.attention.forward(self.attention_norm(x), mask,
                                           cu_seqlens, max_seq_len)

fastvideo/v1/models/encoders/t5.py

@@ -20,8 +20,8 @@
 """PyTorch T5 & UMT5 model."""
 
 import math
+from collections.abc import Iterable
 from dataclasses import dataclass
-from typing import Iterable, Optional, Set, Tuple
 
 import torch
 import torch.nn.functional as F
@@ -63,7 +63,7 @@ class T5DenseActDense(nn.Module):
 
     def __init__(self,
                  config: T5Config,
-                 quant_config: Optional[QuantizationConfig] = None):
+                 quant_config: QuantizationConfig | None = None):
         super().__init__()
         self.wi = MergedColumnParallelLinear(config.d_model, [config.d_ff],
                                              bias=False)
@@ -84,7 +84,7 @@ class T5DenseGatedActDense(nn.Module):
 
     def __init__(self,
                  config: T5Config,
-                 quant_config: Optional[QuantizationConfig] = None):
+                 quant_config: QuantizationConfig | None = None):
         super().__init__()
         self.wi_0 = MergedColumnParallelLinear(config.d_model, [config.d_ff],
                                                bias=False,
@@ -112,7 +112,7 @@ class T5LayerFF(nn.Module):
 
     def __init__(self,
                  config: T5Config,
-                 quant_config: Optional[QuantizationConfig] = None):
+                 quant_config: QuantizationConfig | None = None):
         super().__init__()
         if config.is_gated_act:
             self.DenseReluDense = T5DenseGatedActDense(
@@ -154,7 +154,7 @@ class T5Attention(nn.Module):
                  config: T5Config,
                  attn_type: str,
                  has_relative_attention_bias=False,
-                 quant_config: Optional[QuantizationConfig] = None,
+                 quant_config: QuantizationConfig | None = None,
                  prefix: str = ""):
         super().__init__()
         self.attn_type = attn_type
@@ -293,7 +293,7 @@ class T5Attention(nn.Module):
         self,
         hidden_states: torch.Tensor,  # (num_tokens, d_model)
         attention_mask: torch.Tensor,
-        attn_metadata: Optional[AttentionMetadata] = None,
+        attn_metadata: AttentionMetadata | None = None,
     ) -> torch.Tensor:
         bs, seq_len, _ = hidden_states.shape
         num_seqs = bs
@@ -343,7 +343,7 @@ class T5LayerSelfAttention(nn.Module):
         self,
         config,
         has_relative_attention_bias=False,
-        quant_config: Optional[QuantizationConfig] = None,
+        quant_config: QuantizationConfig | None = None,
         prefix: str = "",
     ):
         super().__init__()
@@ -360,7 +360,7 @@ class T5LayerSelfAttention(nn.Module):
         self,
         hidden_states: torch.Tensor,
         attention_mask: torch.Tensor,
-        attn_metadata: Optional[AttentionMetadata] = None,
+        attn_metadata: AttentionMetadata | None = None,
     ) -> torch.Tensor:
         normed_hidden_states = self.layer_norm(hidden_states)
         attention_output = self.SelfAttention(
@@ -376,7 +376,7 @@ class T5LayerCrossAttention(nn.Module):
 
     def __init__(self,
                  config,
-                 quant_config: Optional[QuantizationConfig] = None,
+                 quant_config: QuantizationConfig | None = None,
                  prefix: str = ""):
         super().__init__()
         self.EncDecAttention = T5Attention(config,
@@ -389,7 +389,7 @@ class T5LayerCrossAttention(nn.Module):
     def forward(
         self,
         hidden_states: torch.Tensor,
-        attn_metadata: Optional[AttentionMetadata] = None,
+        attn_metadata: AttentionMetadata | None = None,
     ) -> torch.Tensor:
         normed_hidden_states = self.layer_norm(hidden_states)
         attention_output = self.EncDecAttention(
@@ -406,7 +406,7 @@ class T5Block(nn.Module):
                  config: T5Config,
                  is_decoder: bool,
                  has_relative_attention_bias=False,
-                 quant_config: Optional[QuantizationConfig] = None,
+                 quant_config: QuantizationConfig | None = None,
                  prefix: str = ""):
         super().__init__()
         self.is_decoder = is_decoder
@@ -430,7 +430,7 @@ class T5Block(nn.Module):
         self,
         hidden_states: torch.Tensor,
         attention_mask: torch.Tensor,
-        attn_metadata: Optional[AttentionMetadata] = None,
+        attn_metadata: AttentionMetadata | None = None,
     ) -> torch.Tensor:
 
         hidden_states = self.layer[0](hidden_states=hidden_states,
@@ -454,7 +454,7 @@ class T5Stack(nn.Module):
                  is_decoder: bool,
                  n_layers: int,
                  embed_tokens=None,
-                 quant_config: Optional[QuantizationConfig] = None,
+                 quant_config: QuantizationConfig | None = None,
                  prefix: str = "",
                  is_umt5: bool = False):
         super().__init__()
@@ -523,11 +523,11 @@ class T5EncoderModel(TextEncoder):
 
     def forward(
         self,
-        input_ids: Optional[torch.Tensor],
-        position_ids: Optional[torch.Tensor] = None,
-        attention_mask: Optional[torch.Tensor] = None,
-        inputs_embeds: Optional[torch.Tensor] = None,
-        output_hidden_states: Optional[bool] = None,
+        input_ids: torch.Tensor | None,
+        position_ids: torch.Tensor | None = None,
+        attention_mask: torch.Tensor | None = None,
+        inputs_embeds: torch.Tensor | None = None,
+        output_hidden_states: bool | None = None,
         **kwargs,
     ) -> BaseEncoderOutput:
         attn_metadata = AttentionMetadata(None)
@@ -539,8 +539,8 @@ class T5EncoderModel(TextEncoder):
 
         return BaseEncoderOutput(last_hidden_state=hidden_states)
 
-    def load_weights(self, weights: Iterable[Tuple[str,
-                                                   torch.Tensor]]) -> Set[str]:
+    def load_weights(self, weights: Iterable[tuple[str,
+                                                   torch.Tensor]]) -> set[str]:
         stacked_params_mapping = [
             # (param_name, shard_name, shard_id)
             (".qkv_proj", ".q", "q"),
@@ -548,7 +548,7 @@ class T5EncoderModel(TextEncoder):
             (".qkv_proj", ".v", "v"),
         ]
         params_dict = dict(self.named_parameters())
-        loaded_params: Set[str] = set()
+        loaded_params: set[str] = set()
         for name, loaded_weight in weights:
             loaded = False
             if "decoder" in name or "lm_head" in name:
@@ -610,11 +610,11 @@ class UMT5EncoderModel(TextEncoder):
 
     def forward(
         self,
-        input_ids: Optional[torch.Tensor],
-        position_ids: Optional[torch.Tensor] = None,
-        attention_mask: Optional[torch.Tensor] = None,
-        inputs_embeds: Optional[torch.Tensor] = None,
-        output_hidden_states: Optional[bool] = None,
+        input_ids: torch.Tensor | None,
+        position_ids: torch.Tensor | None = None,
+        attention_mask: torch.Tensor | None = None,
+        inputs_embeds: torch.Tensor | None = None,
+        output_hidden_states: bool | None = None,
         **kwargs,
     ) -> BaseEncoderOutput:
         attn_metadata = AttentionMetadata(None)
@@ -629,10 +629,10 @@ class UMT5EncoderModel(TextEncoder):
             attention_mask=attention_mask,
         )
 
-    def load_weights(self, weights: Iterable[Tuple[str,
-                                                   torch.Tensor]]) -> Set[str]:
+    def load_weights(self, weights: Iterable[tuple[str,
+                                                   torch.Tensor]]) -> set[str]:
         params_dict = dict(self.named_parameters())
-        loaded_params: Set[str] = set()
+        loaded_params: set[str] = set()
         for name, loaded_weight in weights:
             loaded = False
             if "decoder" in name or "lm_head" in name:

fastvideo/v1/models/encoders/vision.py

@@ -2,7 +2,7 @@
 # Adapted from vllm: https://github.com/vllm-project/vllm/blob/v0.7.3/vllm/model_executor/models/vision.py
 
 from abc import ABC, abstractmethod
-from typing import Generic, Optional, TypeVar, Union
+from typing import Generic, TypeVar
 
 import torch
 from transformers import PretrainedConfig
@@ -48,9 +48,9 @@ class VisionEncoderInfo(ABC, Generic[_C]):
 
 
 def resolve_visual_encoder_outputs(
-    encoder_outputs: Union[torch.Tensor, list[torch.Tensor]],
-    feature_sample_layers: Optional[list[int]],
-    post_layer_norm: Optional[torch.nn.LayerNorm],
+    encoder_outputs: torch.Tensor | list[torch.Tensor],
+    feature_sample_layers: list[int] | None,
+    post_layer_norm: torch.nn.LayerNorm | None,
     max_possible_layers: int,
 ) -> torch.Tensor:
     """Given the outputs a visual encoder module that may correspond to the

fastvideo/v1/models/hf_transformer_utils.py

@@ -20,14 +20,14 @@ import contextlib
 import json
 import os
 from pathlib import Path
-from typing import Any, Dict, Optional, Type, Union
+from typing import Any
 
 from huggingface_hub import snapshot_download
 from transformers import AutoConfig, PretrainedConfig
 from transformers.models.auto.modeling_auto import (
     MODEL_FOR_CAUSAL_LM_MAPPING_NAMES)
 
-_CONFIG_REGISTRY: Dict[str, Type[PretrainedConfig]] = {
+_CONFIG_REGISTRY: dict[str, type[PretrainedConfig]] = {
     # ChatGLMConfig.model_type: ChatGLMConfig,
     # DbrxConfig.model_type: DbrxConfig,
     # ExaoneConfig.model_type: ExaoneConfig,
@@ -50,8 +50,8 @@ def download_from_hf(model_path: str):
 def get_hf_config(
     model: str,
     trust_remote_code: bool,
-    revision: Optional[str] = None,
-    model_override_args: Optional[dict] = None,
+    revision: str | None = None,
+    model_override_args: dict | None = None,
     **kwargs,
 ):
     is_gguf = check_gguf_file(model)
@@ -81,7 +81,7 @@ def get_hf_config(
     return config
 
 
-def get_diffusers_config(model: str, ) -> Dict[str, Any]:
+def get_diffusers_config(model: str, ) -> dict[str, Any]:
     """Gets a configuration for the given diffusers model.
     
     Args:
@@ -101,7 +101,7 @@ def get_diffusers_config(model: str, ) -> Dict[str, Any]:
             try:
                 # Load the config directly from the file
                 with open(config_file) as f:
-                    config_dict: Dict[str, Any] = json.load(f)
+                    config_dict: dict[str, Any] = json.load(f)
                 if "_diffusers_version" in config_dict:
                     config_dict.pop("_diffusers_version")
                 # TODO(will): apply any overrides from inference args
@@ -137,7 +137,7 @@ def attach_additional_stop_token_ids(tokenizer):
         tokenizer.additional_stop_token_ids = None
 
 
-def check_gguf_file(model: Union[str, os.PathLike]) -> bool:
+def check_gguf_file(model: str | os.PathLike) -> bool:
     """Check if the file is a GGUF model."""
     model = Path(model)
     if not model.is_file():

fastvideo/v1/models/loader/component_loader.py

@@ -6,8 +6,9 @@ import json
 import os
 import time
 from abc import ABC, abstractmethod
+from collections.abc import Generator, Iterable
 from copy import deepcopy
-from typing import Generator, Iterable, List, Optional, Tuple, cast
+from typing import cast
 
 import torch
 import torch.distributed as dist
@@ -109,7 +110,7 @@ class TextEncoderLoader(ComponentLoader):
         fall_back_to_pt: bool = True
         """Whether .pt weights can be used."""
 
-        allow_patterns_overrides: Optional[list[str]] = None
+        allow_patterns_overrides: list[str] | None = None
         """If defined, weights will load exclusively using these patterns."""
 
     counter_before_loading_weights: float = 0.0
@@ -119,8 +120,8 @@ class TextEncoderLoader(ComponentLoader):
         self,
         model_name_or_path: str,
         fall_back_to_pt: bool,
-        allow_patterns_overrides: Optional[list[str]],
-    ) -> Tuple[str, List[str], bool]:
+        allow_patterns_overrides: list[str] | None,
+    ) -> tuple[str, list[str], bool]:
         """Prepare weights for the model.
 
         If the model is not local, it will be downloaded."""
@@ -142,7 +143,7 @@ class TextEncoderLoader(ComponentLoader):
 
         hf_folder = model_name_or_path
 
-        hf_weights_files: List[str] = []
+        hf_weights_files: list[str] = []
         for pattern in allow_patterns:
             hf_weights_files += glob.glob(os.path.join(hf_folder, pattern))
             if len(hf_weights_files) > 0:
@@ -165,7 +166,7 @@ class TextEncoderLoader(ComponentLoader):
 
     def _get_weights_iterator(
             self, source: "Source",
-            to_cpu: bool) -> Generator[Tuple[str, torch.Tensor], None, None]:
+            to_cpu: bool) -> Generator[tuple[str, torch.Tensor], None, None]:
         """Get an iterator for the model weights based on the load format."""
         hf_folder, hf_weights_files, use_safetensors = self._prepare_weights(
             source.model_or_path, source.fall_back_to_pt,
@@ -188,7 +189,7 @@ class TextEncoderLoader(ComponentLoader):
         model: nn.Module,
         model_path: str,
         to_cpu: bool,
-    ) -> Generator[Tuple[str, torch.Tensor], None, None]:
+    ) -> Generator[tuple[str, torch.Tensor], None, None]:
         primary_weights = TextEncoderLoader.Source(
             model_path,
             prefix="",

fastvideo/v1/models/loader/fsdp_load.py

@@ -6,9 +6,9 @@
 
 import contextlib
 from collections import defaultdict
+from collections.abc import Callable, Generator
 from itertools import chain
-from typing import (Any, Callable, DefaultDict, Dict, Generator, List, Optional,
-                    Tuple, Type, Union)
+from typing import Any
 
 import torch
 from torch import nn
@@ -56,9 +56,9 @@ def set_default_dtype(dtype: torch.dtype) -> Generator[None, None, None]:
 
 # TODO(PY): add compile option
 def maybe_load_fsdp_model(
-    model_cls: Type[nn.Module],
-    init_params: Dict[str, Any],
-    weight_dir_list: List[str],
+    model_cls: type[nn.Module],
+    init_params: dict[str, Any],
+    weight_dir_list: list[str],
     device: torch.device,
     hsdp_replicate_dim: int,
     hsdp_shard_dim: int,
@@ -67,7 +67,7 @@ def maybe_load_fsdp_model(
     reduce_dtype: torch.dtype,
     cpu_offload: bool = False,
     fsdp_inference: bool = False,
-    output_dtype: Optional[torch.dtype] = None,
+    output_dtype: torch.dtype | None = None,
     training_mode: bool = True,
     pin_cpu_memory: bool = True,
 ) -> torch.nn.Module:
@@ -133,9 +133,9 @@ def shard_model(
     *,
     cpu_offload: bool,
     reshard_after_forward: bool = True,
-    mp_policy: Optional[MixedPrecisionPolicy] = MixedPrecisionPolicy(),  # noqa
-    mesh: Optional[DeviceMesh] = None,
-    fsdp_shard_conditions: List[Callable[[str, nn.Module], bool]] = [],  # noqa
+    mp_policy: MixedPrecisionPolicy | None = MixedPrecisionPolicy(),  # noqa
+    mesh: DeviceMesh | None = None,
+    fsdp_shard_conditions: list[Callable[[str, nn.Module], bool]] = [],  # noqa
     pin_cpu_memory: bool = True,
 ) -> None:
     """
@@ -203,13 +203,13 @@ def shard_model(
 
 # TODO(PY): device mesh for cfg parallel
 def load_model_from_full_model_state_dict(
-    model: Union[FSDPModule, torch.nn.Module],
-    full_sd_iterator: Generator[Tuple[str, torch.Tensor], None, None],
+    model: FSDPModule | torch.nn.Module,
+    full_sd_iterator: Generator[tuple[str, torch.Tensor], None, None],
     device: torch.device,
     param_dtype: torch.dtype,
     strict: bool = False,
     cpu_offload: bool = False,
-    param_names_mapping: Optional[Callable[[str], tuple[str, Any, Any]]] = None,
+    param_names_mapping: Callable[[str], tuple[str, Any, Any]] | None = None,
     training_mode: bool = True,
 ) -> _IncompatibleKeys:
     """
@@ -236,7 +236,7 @@ def load_model_from_full_model_state_dict(
     # Find new params
     used_keys = set()
     sharded_sd = {}
-    to_merge_params: DefaultDict[str, Dict[Any, Any]] = defaultdict(dict)
+    to_merge_params: defaultdict[str, dict[Any, Any]] = defaultdict(dict)
     reverse_param_names_mapping = {}
     assert param_names_mapping is not None
     for source_param_name, full_tensor in full_sd_iterator:

fastvideo/v1/models/loader/utils.py

@@ -2,7 +2,8 @@
 """Utilities for selecting and loading models."""
 import contextlib
 import re
-from typing import Any, Callable, Dict
+from collections.abc import Callable
+from typing import Any
 
 import torch
 
@@ -21,7 +22,7 @@ def set_default_torch_dtype(dtype: torch.dtype):
 
 
 def get_param_names_mapping(
-        mapping_dict: Dict[str, str]) -> Callable[[str], tuple[str, Any, Any]]:
+        mapping_dict: dict[str, str]) -> Callable[[str], tuple[str, Any, Any]]:
     """
     Creates a mapping function that transforms parameter names using regex patterns.
     

fastvideo/v1/models/loader/weight_utils.py

@@ -5,8 +5,8 @@ import hashlib
 import json
 import os
 import tempfile
+from collections.abc import Generator
 from pathlib import Path
-from typing import Generator, List, Optional, Tuple, Union
 
 import filelock
 import huggingface_hub.constants
@@ -47,8 +47,7 @@ class DisabledTqdm(tqdm):
         super().__init__(*args, **kwargs, disable=True)
 
 
-def get_lock(model_name_or_path: Union[str, Path],
-             cache_dir: Optional[str] = None):
+def get_lock(model_name_or_path: str | Path, cache_dir: str | None = None):
     lock_dir = cache_dir or temp_dir
     model_name_or_path = str(model_name_or_path)
     os.makedirs(os.path.dirname(lock_dir), exist_ok=True)
@@ -66,9 +65,9 @@ def get_lock(model_name_or_path: Union[str, Path],
 # Passing both of these to the weight loader functionality breaks.
 # So, we use the index_file to
 # look up which safetensors files should be used.
-def filter_duplicate_safetensors_files(hf_weights_files: List[str],
+def filter_duplicate_safetensors_files(hf_weights_files: list[str],
                                        hf_folder: str,
-                                       index_file: str) -> List[str]:
+                                       index_file: str) -> list[str]:
     # model.safetensors.index.json is a mapping from keys in the
     # torch state_dict to safetensors file holding that weight.
     index_file_name = os.path.join(hf_folder, index_file)
@@ -91,7 +90,7 @@ def filter_duplicate_safetensors_files(hf_weights_files: List[str],
 
 
 def filter_files_not_needed_for_inference(
-        hf_weights_files: List[str]) -> List[str]:
+        hf_weights_files: list[str]) -> list[str]:
     """
     Exclude files that are not needed for inference.
 
@@ -119,9 +118,9 @@ _BAR_FORMAT = "{desc}: {percentage:3.0f}% Completed | {n_fmt}/{total_fmt} [{elap
 
 
 def safetensors_weights_iterator(
-    hf_weights_files: List[str],
+    hf_weights_files: list[str],
     to_cpu: bool = True,
-) -> Generator[Tuple[str, torch.Tensor], None, None]:
+) -> Generator[tuple[str, torch.Tensor], None, None]:
     """Iterate over the weights in the model safetensor files."""
     enable_tqdm = not torch.distributed.is_initialized(
     ) or torch.distributed.get_rank() == 0
@@ -139,9 +138,9 @@ def safetensors_weights_iterator(
 
 
 def pt_weights_iterator(
-    hf_weights_files: List[str],
+    hf_weights_files: list[str],
     to_cpu: bool = True,
-) -> Generator[Tuple[str, torch.Tensor], None, None]:
+) -> Generator[tuple[str, torch.Tensor], None, None]:
     """Iterate over the weights in the model bin/pt files."""
     device = "cpu" if to_cpu else str(get_local_torch_device())
     enable_tqdm = not torch.distributed.is_initialized(
@@ -178,7 +177,7 @@ def default_weight_loader(param: torch.Tensor,
         raise
 
 
-def maybe_remap_kv_scale_name(name: str, params_dict: dict) -> Optional[str]:
+def maybe_remap_kv_scale_name(name: str, params_dict: dict) -> str | None:
     """Remap the name of FP8 k/v_scale parameters.
 
     This function handles the remapping of FP8 k/v_scale parameter names.

fastvideo/v1/models/parameter.py

@@ -1,8 +1,9 @@
 # SPDX-License-Identifier: Apache-2.0
 # Adapted from: https://github.com/vllm-project/vllm/blob/v0.7.3/vllm/model_executor/parameter.py
 
+from collections.abc import Callable
 from fractions import Fraction
-from typing import Any, Callable, Tuple, Union
+from typing import Any
 
 import torch
 from torch.nn import Parameter
@@ -113,8 +114,8 @@ class _ColumnvLLMParameter(BasevLLMParameter):
             raise ValueError("shard_offset and shard_size must be provided")
         if isinstance(
                 self,
-            (PackedColumnParameter,
-             PackedvLLMParameter)) and self.packed_dim == self.output_dim:
+            (PackedColumnParameter
+             | PackedvLLMParameter)) and self.packed_dim == self.output_dim:
             shard_size, shard_offset = self.adjust_shard_indexes_for_packing(
                 shard_offset=shard_offset, shard_size=shard_size)
 
@@ -142,8 +143,8 @@ class _ColumnvLLMParameter(BasevLLMParameter):
 
         if isinstance(
                 self,
-            (PackedColumnParameter,
-             PackedvLLMParameter)) and self.output_dim == self.packed_dim:
+            (PackedColumnParameter
+             | PackedvLLMParameter)) and self.output_dim == self.packed_dim:
             shard_size, shard_offset = self.adjust_shard_indexes_for_packing(
                 shard_offset=shard_offset, shard_size=shard_size)
 
@@ -230,7 +231,7 @@ class PerTensorScaleParameter(BasevLLMParameter):
         self.qkv_idxs = {"q": 0, "k": 1, "v": 2}
         super().__init__(**kwargs)
 
-    def _shard_id_as_int(self, shard_id: Union[str, int]) -> int:
+    def _shard_id_as_int(self, shard_id: str | int) -> int:
         if isinstance(shard_id, int):
             return shard_id
 
@@ -255,7 +256,7 @@ class PerTensorScaleParameter(BasevLLMParameter):
         super().load_row_parallel_weight(*args, **kwargs)
 
     def _load_into_shard_id(self, loaded_weight: torch.Tensor,
-                            shard_id: Union[str, int], **kwargs):
+                            shard_id: str | int, **kwargs):
         """
         Slice the parameter data based on the shard id for 
         loading.
@@ -282,7 +283,7 @@ class PackedColumnParameter(_ColumnvLLMParameter):
     for more details on the packed properties.
     """
 
-    def __init__(self, packed_factor: Union[int, Fraction], packed_dim: int,
+    def __init__(self, packed_factor: int | Fraction, packed_dim: int,
                  **kwargs):
         self._packed_factor = packed_factor
         self._packed_dim = packed_dim
@@ -297,7 +298,7 @@ class PackedColumnParameter(_ColumnvLLMParameter):
         return self._packed_factor
 
     def adjust_shard_indexes_for_packing(self, shard_size,
-                                         shard_offset) -> Tuple[Any, Any]:
+                                         shard_offset) -> tuple[Any, Any]:
         return _adjust_shard_indexes_for_packing(
             shard_size=shard_size,
             shard_offset=shard_offset,
@@ -315,7 +316,7 @@ class PackedvLLMParameter(ModelWeightParameter):
     by accounting for packing and optionally, marlin tile size.
     """
 
-    def __init__(self, packed_factor: Union[int, Fraction], packed_dim: int,
+    def __init__(self, packed_factor: int | Fraction, packed_dim: int,
                  **kwargs):
         self._packed_factor = packed_factor
         self._packed_dim = packed_dim
@@ -404,7 +405,7 @@ def permute_param_layout_(param: BasevLLMParameter, input_dim: int,
 
 
 def _adjust_shard_indexes_for_packing(shard_size, shard_offset,
-                                      packed_factor) -> Tuple[Any, Any]:
+                                      packed_factor) -> tuple[Any, Any]:
     shard_size = shard_size // packed_factor
     shard_offset = shard_offset // packed_factor
     return shard_size, shard_offset

fastvideo/v1/models/registry.py

@@ -8,10 +8,10 @@ import subprocess
 import sys
 import tempfile
 from abc import ABC, abstractmethod
+from collections.abc import Callable, Set
 from dataclasses import dataclass, field
 from functools import lru_cache
-from typing import (AbstractSet, Callable, Dict, List, NoReturn, Optional,
-                    Tuple, Type, TypeVar, Union, cast)
+from typing import NoReturn, TypeVar, cast
 
 import cloudpickle
 from torch import nn
@@ -80,7 +80,7 @@ class _ModelInfo:
     architecture: str
 
     @staticmethod
-    def from_model_cls(model: Type[nn.Module]) -> "_ModelInfo":
+    def from_model_cls(model: type[nn.Module]) -> "_ModelInfo":
         return _ModelInfo(architecture=model.__name__, )
 
 
@@ -91,7 +91,7 @@ class _BaseRegisteredModel(ABC):
         raise NotImplementedError
 
     @abstractmethod
-    def load_model_cls(self) -> Type[nn.Module]:
+    def load_model_cls(self) -> type[nn.Module]:
         raise NotImplementedError
 
 
@@ -102,10 +102,10 @@ class _RegisteredModel(_BaseRegisteredModel):
     """
 
     interfaces: _ModelInfo
-    model_cls: Type[nn.Module]
+    model_cls: type[nn.Module]
 
     @staticmethod
-    def from_model_cls(model_cls: Type[nn.Module]):
+    def from_model_cls(model_cls: type[nn.Module]):
         return _RegisteredModel(
             interfaces=_ModelInfo.from_model_cls(model_cls),
             model_cls=model_cls,
@@ -114,7 +114,7 @@ class _RegisteredModel(_BaseRegisteredModel):
     def inspect_model_cls(self) -> _ModelInfo:
         return self.interfaces
 
-    def load_model_cls(self) -> Type[nn.Module]:
+    def load_model_cls(self) -> type[nn.Module]:
         return self.model_cls
 
 
@@ -159,16 +159,16 @@ class _LazyRegisteredModel(_BaseRegisteredModel):
         return _run_in_subprocess(
             lambda: _ModelInfo.from_model_cls(self.load_model_cls()))
 
-    def load_model_cls(self) -> Type[nn.Module]:
+    def load_model_cls(self) -> type[nn.Module]:
         mod = importlib.import_module(self.module_name)
-        return cast(Type[nn.Module], getattr(mod, self.class_name))
+        return cast(type[nn.Module], getattr(mod, self.class_name))
 
 
 @lru_cache(maxsize=128)
 def _try_load_model_cls(
     model_arch: str,
     model: _BaseRegisteredModel,
-) -> Optional[Type[nn.Module]]:
+) -> type[nn.Module] | None:
     from fastvideo.v1.platforms import current_platform
     current_platform.verify_model_arch(model_arch)
     try:
@@ -182,7 +182,7 @@ def _try_load_model_cls(
 def _try_inspect_model_cls(
     model_arch: str,
     model: _BaseRegisteredModel,
-) -> Optional[_ModelInfo]:
+) -> _ModelInfo | None:
     try:
         return model.inspect_model_cls()
     except Exception:
@@ -194,15 +194,15 @@ def _try_inspect_model_cls(
 @dataclass
 class _ModelRegistry:
     # Keyed by model_arch
-    models: Dict[str, _BaseRegisteredModel] = field(default_factory=dict)
+    models: dict[str, _BaseRegisteredModel] = field(default_factory=dict)
 
-    def get_supported_archs(self) -> AbstractSet[str]:
+    def get_supported_archs(self) -> Set[str]:
         return self.models.keys()
 
     def register_model(
         self,
         model_arch: str,
-        model_cls: Union[Type[nn.Module], str],
+        model_cls: type[nn.Module] | str,
     ) -> None:
         """
         Register an external model to be used in vLLM.
@@ -232,7 +232,7 @@ class _ModelRegistry:
 
         self.models[model_arch] = model
 
-    def _raise_for_unsupported(self, architectures: List[str]) -> NoReturn:
+    def _raise_for_unsupported(self, architectures: list[str]) -> NoReturn:
         all_supported_archs = self.get_supported_archs()
 
         if any(arch in all_supported_archs for arch in architectures):
@@ -244,13 +244,13 @@ class _ModelRegistry:
             f"Model architectures {architectures} are not supported for now. "
             f"Supported architectures: {all_supported_archs}")
 
-    def _try_load_model_cls(self, model_arch: str) -> Optional[Type[nn.Module]]:
+    def _try_load_model_cls(self, model_arch: str) -> type[nn.Module] | None:
         if model_arch not in self.models:
             return None
 
         return _try_load_model_cls(model_arch, self.models[model_arch])
 
-    def _try_inspect_model_cls(self, model_arch: str) -> Optional[_ModelInfo]:
+    def _try_inspect_model_cls(self, model_arch: str) -> _ModelInfo | None:
         if model_arch not in self.models:
             return None
 
@@ -258,8 +258,8 @@ class _ModelRegistry:
 
     def _normalize_archs(
         self,
-        architectures: Union[str, List[str]],
-    ) -> List[str]:
+        architectures: str | list[str],
+    ) -> list[str]:
         if isinstance(architectures, str):
             architectures = [architectures]
         if not architectures:
@@ -274,8 +274,8 @@ class _ModelRegistry:
 
     def inspect_model_cls(
         self,
-        architectures: Union[str, List[str]],
-    ) -> Tuple[_ModelInfo, str]:
+        architectures: str | list[str],
+    ) -> tuple[_ModelInfo, str]:
         architectures = self._normalize_archs(architectures)
 
         for arch in architectures:
@@ -287,8 +287,8 @@ class _ModelRegistry:
 
     def resolve_model_cls(
         self,
-        architectures: Union[str, List[str]],
-    ) -> Tuple[Type[nn.Module], str]:
+        architectures: str | list[str],
+    ) -> tuple[type[nn.Module], str]:
         architectures = self._normalize_archs(architectures)
 
         for arch in architectures:

fastvideo/v1/models/schedulers/base.py

@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 
 from abc import ABC, abstractmethod
-from typing import Optional, Tuple, Union
 
 import torch
 from diffusers.utils import BaseOutput
@@ -32,15 +31,15 @@ class BaseScheduler(ABC):
     @abstractmethod
     def scale_model_input(self,
                           sample: torch.Tensor,
-                          timestep: Optional[int] = None) -> torch.Tensor:
+                          timestep: int | None = None) -> torch.Tensor:
         pass
 
     @abstractmethod
     def step(
         self,
         model_output: torch.Tensor,
-        timestep: Union[int, torch.Tensor],
+        timestep: int | torch.Tensor,
         sample: torch.Tensor,
         return_dict: bool = True,
-    ) -> Union[BaseOutput, Tuple]:
+    ) -> BaseOutput | tuple:
         pass

fastvideo/v1/models/schedulers/scheduling_flow_match_euler_discrete.py

@@ -20,7 +20,7 @@
 # ==============================================================================
 
 from dataclasses import dataclass
-from typing import Any, Optional, Tuple, Union
+from typing import Any
 
 import torch
 from diffusers.configuration_utils import ConfigMixin, register_to_config
@@ -76,7 +76,7 @@ class FlowMatchDiscreteScheduler(SchedulerMixin, ConfigMixin, BaseScheduler):
         shift: float = 1.0,
         reverse: bool = True,
         solver: str = "euler",
-        n_tokens: Optional[int] = None,
+        n_tokens: int | None = None,
         **kwargs,
     ):
         sigmas = torch.linspace(1, 0, num_train_timesteps + 1)
@@ -131,7 +131,7 @@ class FlowMatchDiscreteScheduler(SchedulerMixin, ConfigMixin, BaseScheduler):
     def set_timesteps(
         self,
         num_inference_steps: int,
-        device: Union[str, torch.device] = None,
+        device: str | torch.device = None,
         n_tokens: int = 0,
     ):
         """
@@ -194,7 +194,7 @@ class FlowMatchDiscreteScheduler(SchedulerMixin, ConfigMixin, BaseScheduler):
 
     def scale_model_input(self,
                           sample: torch.Tensor,
-                          timestep: Optional[int] = None) -> torch.Tensor:
+                          timestep: int | None = None) -> torch.Tensor:
         return sample
 
     def sd3_time_shift(self, t: torch.Tensor):
@@ -203,11 +203,11 @@ class FlowMatchDiscreteScheduler(SchedulerMixin, ConfigMixin, BaseScheduler):
     def step(
         self,
         model_output: torch.FloatTensor,
-        timestep: Union[float, torch.FloatTensor],
+        timestep: float | torch.FloatTensor,
         sample: torch.FloatTensor,
         return_dict: bool = True,
         **kwargs,
-    ) -> Union[FlowMatchDiscreteSchedulerOutput, Tuple]:
+    ) -> FlowMatchDiscreteSchedulerOutput | tuple:
         """
         Predict the sample from the previous timestep by reversing the SDE. This function propagates the diffusion
         process from the learned model outputs (most often the predicted noise).
@@ -233,7 +233,7 @@ class FlowMatchDiscreteScheduler(SchedulerMixin, ConfigMixin, BaseScheduler):
                 returned, otherwise a tuple is returned where the first element is the sample tensor.
         """
 
-        if isinstance(timestep, (int, torch.IntTensor, torch.LongTensor)):
+        if isinstance(timestep, (int | torch.IntTensor | torch.LongTensor)):
             raise ValueError((
                 "Passing integer indices (e.g. from `enumerate(timesteps)`) as timesteps to"
                 " `EulerDiscreteScheduler.step()` is not supported. Make sure to pass"

fastvideo/v1/models/schedulers/scheduling_flow_unipc_multistep.py

@@ -4,7 +4,7 @@
 # Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved.
 
 import math
-from typing import Any, List, Optional, Tuple, Union
+from typing import Any
 
 import numpy as np
 import torch
@@ -82,7 +82,7 @@ class FlowUniPCMultistepScheduler(SchedulerMixin, ConfigMixin, BaseScheduler):
             num_train_timesteps: int = 1000,
             solver_order: int = 2,
             prediction_type: str = "flow_prediction",
-            shift: Optional[float] = 1.0,
+            shift: float | None = 1.0,
             use_dynamic_shifting=False,
             thresholding: bool = False,
             dynamic_thresholding_ratio: float = 0.995,
@@ -90,11 +90,11 @@ class FlowUniPCMultistepScheduler(SchedulerMixin, ConfigMixin, BaseScheduler):
             predict_x0: bool = True,
             solver_type: str = "bh2",
             lower_order_final: bool = True,
-            disable_corrector: Tuple = (),
+            disable_corrector: tuple = (),
             solver_p: SchedulerMixin = None,
             timestep_spacing: str = "linspace",
             steps_offset: int = 0,
-            final_sigmas_type: Optional[str] = "zero",  # "zero", "sigma_min"
+            final_sigmas_type: str | None = "zero",  # "zero", "sigma_min"
             **kwargs):
 
         if solver_type not in ["bh1", "bh2"]:
@@ -106,7 +106,7 @@ class FlowUniPCMultistepScheduler(SchedulerMixin, ConfigMixin, BaseScheduler):
 
         self.predict_x0 = predict_x0
         # setable values
-        self.num_inference_steps: Optional[int] = None
+        self.num_inference_steps: int | None = None
         alphas = np.linspace(1, 1 / num_train_timesteps,
                              num_train_timesteps)[::-1].copy()
         sigmas = 1.0 - alphas
@@ -122,13 +122,13 @@ class FlowUniPCMultistepScheduler(SchedulerMixin, ConfigMixin, BaseScheduler):
         self.timesteps = sigmas * num_train_timesteps
 
         self.model_outputs = [None] * solver_order
-        self.timestep_list: List[Optional[Any]] = [None] * solver_order
+        self.timestep_list: list[Any | None] = [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._step_index: int | None = None
+        self._begin_index: int | None = None
 
         self.sigmas = self.sigmas.to(
             "cpu")  # to avoid too much CPU/GPU communication
@@ -166,11 +166,11 @@ class FlowUniPCMultistepScheduler(SchedulerMixin, ConfigMixin, BaseScheduler):
     # Modified from diffusers.schedulers.scheduling_flow_match_euler_discrete.FlowMatchEulerDiscreteScheduler.set_timesteps
     def set_timesteps(
         self,
-        num_inference_steps: Union[int, None] = None,
-        device: Union[str, torch.device] = None,
-        sigmas: Optional[List[float]] = None,
-        mu: Optional[Union[float, None]] = None,
-        shift: Optional[Union[float, None]] = None,
+        num_inference_steps: int | None = None,
+        device: str | torch.device = None,
+        sigmas: list[float] | None = None,
+        mu: float | None | None = None,
+        shift: float | None | None = None,
     ):
         """
         Sets the discrete timesteps used for the diffusion chain (to be run before inference).
@@ -280,7 +280,7 @@ class FlowUniPCMultistepScheduler(SchedulerMixin, ConfigMixin, BaseScheduler):
     def _sigma_to_t(self, sigma):
         return sigma * self.config.num_train_timesteps
 
-    def _sigma_to_alpha_sigma_t(self, sigma) -> Tuple[Any, Any]:
+    def _sigma_to_alpha_sigma_t(self, sigma) -> tuple[Any, Any]:
         return 1 - sigma, sigma
 
     # Copied from diffusers.schedulers.scheduling_flow_match_euler_discrete.set_timesteps
@@ -363,7 +363,7 @@ class FlowUniPCMultistepScheduler(SchedulerMixin, ConfigMixin, BaseScheduler):
         model_output: torch.Tensor,
         *args,
         sample: torch.Tensor = None,
-        order: Optional[int] = None,  # pyright: ignore
+        order: int | None = None,  # pyright: ignore
         **kwargs,
     ) -> torch.Tensor:
         """
@@ -425,7 +425,7 @@ class FlowUniPCMultistepScheduler(SchedulerMixin, ConfigMixin, BaseScheduler):
         device = sample.device
 
         rks = []
-        D1s: Optional[List[Any]] = []
+        D1s: list[Any] | None = []
         for i in range(1, order):
             si = self.step_index - i  # pyright: ignore
             mi = model_output_list[-(i + 1)]
@@ -502,7 +502,7 @@ class FlowUniPCMultistepScheduler(SchedulerMixin, ConfigMixin, BaseScheduler):
         *args,
         last_sample: torch.Tensor = None,
         this_sample: torch.Tensor = None,
-        order: Optional[int] = None,  # pyright: ignore
+        order: int | None = None,  # pyright: ignore
         **kwargs,
     ) -> torch.Tensor:
         """
@@ -570,7 +570,7 @@ class FlowUniPCMultistepScheduler(SchedulerMixin, ConfigMixin, BaseScheduler):
         device = this_sample.device
 
         rks = []
-        D1s: Optional[List[Any]] = []
+        D1s: list[Any] | None = []
         for i in range(1, order):
             si = self.step_index - (i + 1)  # pyright: ignore
             mi = model_output_list[-(i + 1)]
@@ -669,10 +669,10 @@ class FlowUniPCMultistepScheduler(SchedulerMixin, ConfigMixin, BaseScheduler):
 
     def step(self,
              model_output: torch.Tensor,
-             timestep: Union[int, torch.Tensor],
+             timestep: int | torch.Tensor,
              sample: torch.Tensor,
              return_dict: bool = True,
-             generator=None) -> Union[SchedulerOutput, Tuple]:
+             generator=None) -> SchedulerOutput | tuple:
         """
         Predict the sample from the previous timestep by reversing the SDE. This function propagates the sample with
         the multistep UniPC.

fastvideo/v1/models/schedulers/scheduling_unipc_multistep.py

@@ -23,7 +23,6 @@
 # ==============================================================================
 
 import math
-from typing import List, Optional, Tuple, Union
 
 import numpy as np
 import torch
@@ -203,7 +202,7 @@ class UniPCMultistepScheduler(SchedulerMixin, ConfigMixin, BaseScheduler):
         beta_start: float = 0.0001,
         beta_end: float = 0.02,
         beta_schedule: str = "linear",
-        trained_betas: Optional[Union[np.ndarray, List[float]]] = None,
+        trained_betas: np.ndarray | list[float] | None = None,
         solver_order: int = 2,
         prediction_type: str = "epsilon",
         thresholding: bool = False,
@@ -212,16 +211,16 @@ class UniPCMultistepScheduler(SchedulerMixin, ConfigMixin, BaseScheduler):
         predict_x0: bool = True,
         solver_type: str = "bh2",
         lower_order_final: bool = True,
-        disable_corrector: Tuple[int, ...] = (),
+        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,
+        use_karras_sigmas: bool | None = False,
+        use_exponential_sigmas: bool | None = False,
+        use_beta_sigmas: bool | None = False,
+        use_flow_sigmas: bool | None = False,
+        flow_shift: float | None = 1.0,
         timestep_spacing: str = "linspace",
         steps_offset: int = 0,
-        final_sigmas_type: Optional[str] = "zero",  # "zero", "sigma_min"
+        final_sigmas_type: str | None = "zero",  # "zero", "sigma_min"
         rescale_betas_zero_snr: bool = False,
     ):
         if self.config.use_beta_sigmas and not is_scipy_available():
@@ -283,21 +282,20 @@ class UniPCMultistepScheduler(SchedulerMixin, ConfigMixin, BaseScheduler):
 
         self.predict_x0 = predict_x0
         # setable values
-        self.num_inference_steps: Optional[int] = None
+        self.num_inference_steps: int | None = 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.timestep_list: list[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._step_index: int | None = None
+        self._begin_index: int | None = None
         self.sigmas = self.sigmas.to(
             "cpu")  # to avoid too much CPU/GPU communication
 
@@ -333,7 +331,7 @@ class UniPCMultistepScheduler(SchedulerMixin, ConfigMixin, BaseScheduler):
 
     def set_timesteps(self,
                       num_inference_steps: int,
-                      device: Union[str, torch.device] = None):
+                      device: str | torch.device = None):
         """
         Sets the discrete timesteps used for the diffusion chain (to be run before inference).
 
@@ -537,7 +535,7 @@ class UniPCMultistepScheduler(SchedulerMixin, ConfigMixin, BaseScheduler):
 
     # 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]:
+            self, sigma: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
         if self.config.use_flow_sigmas:
             alpha_t = 1 - sigma
             sigma_t = sigma
@@ -708,7 +706,7 @@ class UniPCMultistepScheduler(SchedulerMixin, ConfigMixin, BaseScheduler):
         model_output: torch.Tensor,
         *args,
         sample: torch.Tensor = None,
-        order: Optional[int] = None,
+        order: int | None = None,
         **kwargs,
     ) -> torch.Tensor:
         """
@@ -808,7 +806,7 @@ class UniPCMultistepScheduler(SchedulerMixin, ConfigMixin, BaseScheduler):
         R_tensor: torch.Tensor = torch.stack(R)
         b = torch.tensor(b, device=device)
 
-        D1s_tensor: Optional[torch.Tensor] = None
+        D1s_tensor: torch.Tensor | None = None
         if len(D1s) > 0:
             D1s_tensor = torch.stack(D1s, dim=1)  # (B, K)
             # for order 2, we use a simplified version
@@ -842,9 +840,9 @@ class UniPCMultistepScheduler(SchedulerMixin, ConfigMixin, BaseScheduler):
         self,
         this_model_output: torch.Tensor,
         *args,
-        last_sample: Optional[torch.Tensor] = None,
-        this_sample: Optional[torch.Tensor] = None,
-        order: Optional[int] = None,
+        last_sample: torch.Tensor | None = None,
+        this_sample: torch.Tensor | None = None,
+        order: int | None = None,
         **kwargs,
     ) -> torch.Tensor:
         """
@@ -950,7 +948,7 @@ class UniPCMultistepScheduler(SchedulerMixin, ConfigMixin, BaseScheduler):
         R = torch.stack(R)
         b = torch.tensor(b, device=device)
 
-        D1s_tensor: Optional[torch.Tensor] = torch.stack(
+        D1s_tensor: torch.Tensor | None = torch.stack(
             D1s, dim=1) if len(D1s) > 0 else None
 
         # for order 1, we use a simplified version
@@ -1016,10 +1014,10 @@ class UniPCMultistepScheduler(SchedulerMixin, ConfigMixin, BaseScheduler):
     def step(
         self,
         model_output: torch.Tensor,
-        timestep: Union[int, torch.Tensor],
+        timestep: int | torch.Tensor,
         sample: torch.Tensor,
         return_dict: bool = True,
-    ) -> Union[SchedulerOutput, Tuple]:
+    ) -> SchedulerOutput | tuple:
         """
         Predict the sample from the previous timestep by reversing the SDE. This function propagates the sample with
         the multistep UniPC.

fastvideo/v1/models/utils.py

@@ -1,7 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0
 # Adapted from: https://github.com/vllm-project/vllm/blob/v0.7.3/vllm/model_executor/utils.py
 """Utils for model executor."""
-from typing import Any, Dict, List, Optional
+from typing import Any
 
 import torch
 
@@ -53,7 +53,7 @@ def auto_attributes(init_func):
 
 def set_weight_attrs(
     weight: torch.Tensor,
-    weight_attrs: Optional[Dict[str, Any]],
+    weight_attrs: dict[str, Any] | None,
 ):
     """Set attributes on a weight tensor.
 
@@ -104,7 +104,7 @@ def extract_layer_index(layer_name: str) -> int:
     - "model.encoder.layers.0.sub.1" -> ValueError
     """
     subnames = layer_name.split(".")
-    int_vals: List[int] = []
+    int_vals: list[int] = []
     for subname in subnames:
         try:
             int_vals.append(int(subname))
@@ -116,8 +116,8 @@ def extract_layer_index(layer_name: str) -> int:
 
 
 def modulate(x: torch.Tensor,
-             shift: Optional[torch.Tensor] = None,
-             scale: Optional[torch.Tensor] = None) -> torch.Tensor:
+             shift: torch.Tensor | None = None,
+             scale: torch.Tensor | None = None) -> torch.Tensor:
     """modulate by shift and scale
 
     Args:

fastvideo/v1/models/vaes/common.py

@@ -1,8 +1,9 @@
 # SPDX-License-Identifier: Apache-2.0
 
 from abc import ABC, abstractmethod
+from collections.abc import Iterator
 from math import prod
-from typing import Iterator, Optional, Tuple, Union, cast
+from typing import Optional, cast
 
 import numpy as np
 import torch
@@ -50,8 +51,8 @@ class ParallelTiledVAE(ABC):
         return cast(int, self.config.spatial_compression_ratio)
 
     @property
-    def scaling_factor(self) -> Union[float, torch.tensor]:
-        return cast(Union[float, torch.tensor], self.config.scaling_factor)
+    def scaling_factor(self) -> float | torch.Tensor:
+        return cast(float | torch.Tensor, self.config.scaling_factor)
 
     @abstractmethod
     def _encode(self, *args, **kwargs) -> torch.Tensor:
@@ -159,7 +160,7 @@ class ParallelTiledVAE(ABC):
 
     def _parallel_data_generator(
             self, gathered_results,
-            gathered_dim_metadata) -> Iterator[Tuple[torch.Tensor, int]]:
+            gathered_dim_metadata) -> Iterator[tuple[torch.Tensor, int]]:
         global_idx = 0
         for i, per_rank_metadata in enumerate(gathered_dim_metadata):
             _start_shape = 0
@@ -410,16 +411,16 @@ class ParallelTiledVAE(ABC):
 
     def enable_tiling(
         self,
-        tile_sample_min_height: Optional[int] = None,
-        tile_sample_min_width: Optional[int] = None,
-        tile_sample_min_num_frames: Optional[int] = None,
-        tile_sample_stride_height: Optional[int] = None,
-        tile_sample_stride_width: Optional[int] = None,
-        tile_sample_stride_num_frames: Optional[int] = None,
-        blend_num_frames: Optional[int] = None,
-        use_tiling: Optional[bool] = None,
-        use_temporal_tiling: Optional[bool] = None,
-        use_parallel_tiling: Optional[bool] = None,
+        tile_sample_min_height: int | None = None,
+        tile_sample_min_width: int | None = None,
+        tile_sample_min_num_frames: int | None = None,
+        tile_sample_stride_height: int | None = None,
+        tile_sample_stride_width: int | None = None,
+        tile_sample_stride_num_frames: int | None = None,
+        blend_num_frames: int | None = None,
+        use_tiling: bool | None = None,
+        use_temporal_tiling: bool | None = None,
+        use_parallel_tiling: bool | None = None,
     ) -> None:
         r"""
         Enable tiled VAE decoding. When this option is enabled, the VAE will split the input tensor into tiles to
@@ -483,8 +484,7 @@ class DiagonalGaussianDistribution:
                 device=self.parameters.device,
                 dtype=self.parameters.dtype)
 
-    def sample(self,
-               generator: Optional[torch.Generator] = None) -> torch.Tensor:
+    def sample(self, generator: torch.Generator | None = None) -> torch.Tensor:
         # make sure sample is on the same device as the parameters and has same dtype
         sample = randn_tensor(
             self.mean.shape,
@@ -515,7 +515,7 @@ class DiagonalGaussianDistribution:
 
     def nll(
         self, sample: torch.Tensor,
-        dims: Tuple[int, ...] = (1, 2, 3)) -> torch.Tensor:
+        dims: tuple[int, ...] = (1, 2, 3)) -> torch.Tensor:
         if self.deterministic:
             return torch.Tensor([0.0])
         logtwopi = np.log(2.0 * np.pi)

fastvideo/v1/models/vaes/hunyuanvae.py

@@ -15,8 +15,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import Optional, Tuple, Union
-
 import numpy as np
 import torch
 import torch.nn as nn
@@ -32,7 +30,7 @@ def prepare_causal_attention_mask(
         height_width: int,
         dtype: torch.dtype,
         device: torch.device,
-        batch_size: Optional[int] = None) -> torch.Tensor:
+        batch_size: int | None = None) -> torch.Tensor:
     indices = torch.arange(1, num_frames + 1, dtype=torch.int32, device=device)
     indices_blocks = indices.repeat_interleave(height_width)
     x, y = torch.meshgrid(indices_blocks, indices_blocks, indexing="xy")
@@ -72,7 +70,7 @@ class HunyuanVAEAttention(nn.Module):
 
     def forward(self,
                 hidden_states: torch.Tensor,
-                attention_mask: Optional[torch.Tensor] = None) -> torch.Tensor:
+                attention_mask: torch.Tensor | None = None) -> torch.Tensor:
         residual = hidden_states
 
         batch_size, sequence_length, _ = hidden_states.shape
@@ -121,10 +119,10 @@ class HunyuanVideoCausalConv3d(nn.Module):
         self,
         in_channels: int,
         out_channels: int,
-        kernel_size: Union[int, Tuple[int, int, int]] = 3,
-        stride: Union[int, Tuple[int, int, int]] = 1,
-        padding: Union[int, Tuple[int, int, int]] = 0,
-        dilation: Union[int, Tuple[int, int, int]] = 1,
+        kernel_size: int | tuple[int, int, int] = 3,
+        stride: int | tuple[int, int, int] = 1,
+        padding: int | tuple[int, int, int] = 0,
+        dilation: int | tuple[int, int, int] = 1,
         bias: bool = True,
         pad_mode: str = "replicate",
     ) -> None:
@@ -163,11 +161,11 @@ class HunyuanVideoUpsampleCausal3D(nn.Module):
     def __init__(
             self,
             in_channels: int,
-            out_channels: Optional[int] = None,
+            out_channels: int | None = None,
             kernel_size: int = 3,
             stride: int = 1,
             bias: bool = True,
-            upsample_factor: Tuple[int, ...] = (2, 2, 2),
+            upsample_factor: tuple[int, ...] = (2, 2, 2),
     ) -> None:
         super().__init__()
 
@@ -213,7 +211,7 @@ class HunyuanVideoDownsampleCausal3D(nn.Module):
     def __init__(
         self,
         channels: int,
-        out_channels: Optional[int] = None,
+        out_channels: int | None = None,
         padding: int = 1,
         kernel_size: int = 3,
         bias: bool = True,
@@ -239,7 +237,7 @@ class HunyuanVideoResnetBlockCausal3D(nn.Module):
     def __init__(
         self,
         in_channels: int,
-        out_channels: Optional[int] = None,
+        out_channels: int | None = None,
         dropout: float = 0.0,
         groups: int = 32,
         eps: float = 1e-6,
@@ -313,7 +311,7 @@ class HunyuanVideoMidBlock3D(nn.Module):
                 non_linearity=resnet_act_fn,
             )
         ]
-        attentions: list[Optional[HunyuanVAEAttention]] = []
+        attentions: list[HunyuanVAEAttention | None] = []
 
         for _ in range(num_layers):
             if self.add_attention:
@@ -349,7 +347,9 @@ class HunyuanVideoMidBlock3D(nn.Module):
             hidden_states = self._gradient_checkpointing_func(
                 self.resnets[0], hidden_states)
 
-            for attn, resnet in zip(self.attentions, self.resnets[1:]):
+            for attn, resnet in zip(self.attentions,
+                                    self.resnets[1:],
+                                    strict=True):
                 if attn is not None:
                     batch_size, num_channels, num_frames, height, width = hidden_states.shape
                     hidden_states = hidden_states.permute(0, 2, 3, 4,
@@ -371,7 +371,9 @@ class HunyuanVideoMidBlock3D(nn.Module):
         else:
             hidden_states = self.resnets[0](hidden_states)
 
-            for attn, resnet in zip(self.attentions, self.resnets[1:]):
+            for attn, resnet in zip(self.attentions,
+                                    self.resnets[1:],
+                                    strict=True):
                 if attn is not None:
                     batch_size, num_channels, num_frames, height, width = hidden_states.shape
                     hidden_states = hidden_states.permute(0, 2, 3, 4,
@@ -404,7 +406,7 @@ class HunyuanVideoDownBlock3D(nn.Module):
         resnet_act_fn: str = "silu",
         resnet_groups: int = 32,
         add_downsample: bool = True,
-        downsample_stride: Tuple[int, ...] | int = 2,
+        downsample_stride: tuple[int, ...] | int = 2,
         downsample_padding: int = 1,
     ) -> None:
         super().__init__()
@@ -466,7 +468,7 @@ class HunyuanVideoUpBlock3D(nn.Module):
             resnet_act_fn: str = "silu",
             resnet_groups: int = 32,
             add_upsample: bool = True,
-            upsample_scale_factor: Tuple[int, ...] = (2, 2, 2),
+            upsample_scale_factor: tuple[int, ...] = (2, 2, 2),
     ) -> None:
         super().__init__()
         resnets = []
@@ -525,13 +527,13 @@ class HunyuanVideoEncoder3D(nn.Module):
         self,
         in_channels: int = 3,
         out_channels: int = 3,
-        down_block_types: Tuple[str, ...] = (
+        down_block_types: tuple[str, ...] = (
             "HunyuanVideoDownBlock3D",
             "HunyuanVideoDownBlock3D",
             "HunyuanVideoDownBlock3D",
             "HunyuanVideoDownBlock3D",
         ),
-        block_out_channels: Tuple[int, ...] = (128, 256, 512, 512),
+        block_out_channels: tuple[int, ...] = (128, 256, 512, 512),
         layers_per_block: int = 2,
         norm_num_groups: int = 32,
         act_fn: str = "silu",
@@ -546,7 +548,7 @@ class HunyuanVideoEncoder3D(nn.Module):
                                                 block_out_channels[0],
                                                 kernel_size=3,
                                                 stride=1)
-        self.mid_block: Optional[HunyuanVideoMidBlock3D] = None
+        self.mid_block: HunyuanVideoMidBlock3D | None = None
         self.down_blocks = nn.ModuleList([])
 
         output_channel = block_out_channels[0]
@@ -649,13 +651,13 @@ class HunyuanVideoDecoder3D(nn.Module):
         self,
         in_channels: int = 3,
         out_channels: int = 3,
-        up_block_types: Tuple[str, ...] = (
+        up_block_types: tuple[str, ...] = (
             "HunyuanVideoUpBlock3D",
             "HunyuanVideoUpBlock3D",
             "HunyuanVideoUpBlock3D",
             "HunyuanVideoUpBlock3D",
         ),
-        block_out_channels: Tuple[int, ...] = (128, 256, 512, 512),
+        block_out_channels: tuple[int, ...] = (128, 256, 512, 512),
         layers_per_block: int = 2,
         norm_num_groups: int = 32,
         act_fn: str = "silu",
@@ -832,7 +834,7 @@ class AutoencoderKLHunyuanVideo(nn.Module, ParallelTiledVAE):
         self,
         sample: torch.Tensor,
         sample_posterior: bool = False,
-        generator: Optional[torch.Generator] = None,
+        generator: torch.Generator | None = None,
     ) -> torch.Tensor:
         r"""
         Args:

fastvideo/v1/models/vaes/stepvideovae.py

@@ -11,7 +11,7 @@
 # The above copyright notice and this permission notice shall be included in all
 # copies or substantial portions of the Software.
 # ==============================================================================
-from typing import Any, List, Optional, Tuple
+from typing import Any
 
 import torch
 from einops import rearrange
@@ -103,7 +103,7 @@ def base_conv3d(x,
     return out
 
 
-def cal_outsize(input_sizes, kernel_sizes, stride, padding) -> List:
+def cal_outsize(input_sizes, kernel_sizes, stride, padding) -> list:
     stride_d, stride_h, stride_w = stride
     padding_d, padding_h, padding_w = padding
     dilation_d, dilation_h, dilation_w = 1, 1, 1
@@ -446,8 +446,8 @@ def base_group_norm_with_zero_pad(x,
 
 
 class CausalConvChannelLast(CausalConv):
-    time_causal_padding: Tuple[Any, ...]
-    time_uncausal_padding: Tuple[Any, ...]
+    time_causal_padding: tuple[Any, ...]
+    time_uncausal_padding: tuple[Any, ...]
 
     def __init__(self, chan_in, chan_out, kernel_size, **kwargs) -> None:
         super().__init__(chan_in, chan_out, kernel_size, **kwargs)
@@ -1122,7 +1122,7 @@ class AutoencoderKLStepvideo(nn.Module, ParallelTiledVAE):
         self,
         sample: torch.Tensor,
         sample_posterior: bool = False,
-        generator: Optional[torch.Generator] = None,
+        generator: torch.Generator | None = None,
     ) -> torch.Tensor:
         """
         Args:

fastvideo/v1/models/vaes/wanvae.py

@@ -16,7 +16,6 @@
 
 import contextvars
 from contextlib import contextmanager
-from typing import Optional, Tuple, Union
 
 import torch
 import torch.nn as nn
@@ -68,9 +67,9 @@ class WanCausalConv3d(nn.Conv3d):
         self,
         in_channels: int,
         out_channels: int,
-        kernel_size: Union[int, Tuple[int, int, int]],
-        stride: Union[int, Tuple[int, int, int]] = 1,
-        padding: Union[int, Tuple[int, int, int]] = 0,
+        kernel_size: int | tuple[int, int, int],
+        stride: int | tuple[int, int, int] = 1,
+        padding: int | tuple[int, int, int] = 0,
     ) -> None:
         super().__init__(
             in_channels=in_channels,
@@ -79,9 +78,9 @@ class WanCausalConv3d(nn.Conv3d):
             stride=stride,
             padding=padding,
         )
-        self.padding: Tuple[int, int, int]
+        self.padding: tuple[int, int, int]
         # Set up causal padding
-        self._padding: Tuple[int, ...] = (self.padding[2], self.padding[2],
+        self._padding: tuple[int, ...] = (self.padding[2], self.padding[2],
                                           self.padding[1], self.padding[1],
                                           2 * self.padding[0], 0)
         self.padding = (0, 0, 0)
@@ -434,7 +433,7 @@ class WanMidBlock(nn.Module):
         x = self.resnets[0](x)
 
         # Process through attention and residual blocks
-        for attn, resnet in zip(self.attentions, self.resnets[1:]):
+        for attn, resnet in zip(self.attentions, self.resnets[1:], strict=True):
             if attn is not None:
                 x = attn(x)
 
@@ -488,7 +487,8 @@ class WanEncoder3d(nn.Module):
 
         # downsample blocks
         self.down_blocks = nn.ModuleList([])
-        for i, (in_dim, out_dim) in enumerate(zip(dims[:-1], dims[1:])):
+        for i, (in_dim,
+                out_dim) in enumerate(zip(dims[:-1], dims[1:], strict=True)):
             # residual (+attention) blocks
             for _ in range(num_res_blocks):
                 self.down_blocks.append(
@@ -589,7 +589,7 @@ class WanUpBlock(nn.Module):
         out_dim: int,
         num_res_blocks: int,
         dropout: float = 0.0,
-        upsample_mode: Optional[str] = None,
+        upsample_mode: str | None = None,
         non_linearity: str = "silu",
     ):
         super().__init__()
@@ -687,7 +687,8 @@ class WanDecoder3d(nn.Module):
 
         # upsample blocks
         self.up_blocks = nn.ModuleList([])
-        for i, (in_dim, out_dim) in enumerate(zip(dims[:-1], dims[1:])):
+        for i, (in_dim,
+                out_dim) in enumerate(zip(dims[:-1], dims[1:], strict=True)):
             # residual (+attention) blocks
             if i > 0:
                 in_dim = in_dim // 2
@@ -944,7 +945,7 @@ class AutoencoderKLWan(nn.Module, ParallelTiledVAE):
         self,
         sample: torch.Tensor,
         sample_posterior: bool = False,
-        generator: Optional[torch.Generator] = None,
+        generator: torch.Generator | None = None,
     ) -> torch.Tensor:
         """
         Args:

fastvideo/v1/models/vision_utils.py

@@ -2,7 +2,7 @@
 
 import os
 import tempfile
-from typing import Callable, List, Optional, Tuple, Union
+from collections.abc import Callable
 from urllib.parse import unquote, urlparse
 
 import imageio
@@ -32,8 +32,7 @@ else:
     }
 
 
-def pil_to_numpy(
-        images: Union[List[PIL.Image.Image], PIL.Image.Image]) -> np.ndarray:
+def pil_to_numpy(images: list[PIL.Image.Image] | PIL.Image.Image) -> np.ndarray:
     r"""
     Convert a PIL image or a list of PIL images to NumPy arrays.
 
@@ -72,9 +71,7 @@ def numpy_to_pt(images: np.ndarray) -> torch.Tensor:
     return images
 
 
-def normalize(
-        images: Union[np.ndarray,
-                      torch.Tensor]) -> Union[np.ndarray, torch.Tensor]:
+def normalize(images: np.ndarray | torch.Tensor) -> np.ndarray | torch.Tensor:
     r"""
     Normalize an image array to [-1,1].
 
@@ -91,9 +88,8 @@ def normalize(
 
 # adapted from diffusers.utils import load_image
 def load_image(
-    image: Union[str, PIL.Image.Image],
-    convert_method: Optional[Callable[[PIL.Image.Image],
-                                      PIL.Image.Image]] = None
+    image: str | PIL.Image.Image,
+    convert_method: Callable[[PIL.Image.Image], PIL.Image.Image] | None = None
 ) -> PIL.Image.Image:
     """
     Loads `image` to a PIL Image.
@@ -138,9 +134,9 @@ def load_image(
 # adapted from diffusers.utils import load_video
 def load_video(
     video: str,
-    convert_method: Optional[Callable[[List[PIL.Image.Image]],
-                                      List[PIL.Image.Image]]] = None,
-) -> List[PIL.Image.Image]:
+    convert_method: Callable[[list[PIL.Image.Image]], list[PIL.Image.Image]]
+    | None = None,
+) -> list[PIL.Image.Image]:
     """
     Loads `video` to a list of PIL Image.
     Args:
@@ -215,11 +211,11 @@ def load_video(
 
 
 def get_default_height_width(
-    image: Union[PIL.Image.Image, np.ndarray, torch.Tensor],
+    image: PIL.Image.Image | np.ndarray | torch.Tensor,
     vae_scale_factor: int,
-    height: Optional[int] = None,
-    width: Optional[int] = None,
-) -> Tuple[int, int]:
+    height: int | None = None,
+    width: int | None = None,
+) -> tuple[int, int]:
     r"""
     Returns the height and width of the image, downscaled to the next integer multiple of `vae_scale_factor`.
 
@@ -262,12 +258,12 @@ def get_default_height_width(
 
 
 def resize(
-    image: Union[PIL.Image.Image, np.ndarray, torch.Tensor],
+    image: 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]:
+) -> PIL.Image.Image | np.ndarray | torch.Tensor:
     """
     Resize image.
 

fastvideo/v1/pipelines/composed_pipeline_base.py

@@ -8,7 +8,7 @@ This module defines the base class for pipelines that are composed of multiple s
 import argparse
 import os
 from abc import ABC, abstractmethod
-from typing import Any, Dict, List, Optional, Union, cast
+from typing import Any, cast
 
 import torch
 
@@ -36,17 +36,17 @@ class ComposedPipelineBase(ABC):
     """
 
     is_video_pipeline: bool = False  # To be overridden by video pipelines
-    _required_config_modules: List[str] = []
-    training_args: Optional[TrainingArgs] = None
-    fastvideo_args: Optional[FastVideoArgs] = None
-    modules: Dict[str, torch.nn.Module] = {}
+    _required_config_modules: list[str] = []
+    training_args: TrainingArgs | None = None
+    fastvideo_args: FastVideoArgs | None = None
+    modules: dict[str, torch.nn.Module] = {}
 
     # TODO(will): args should support both inference args and training args
     def __init__(self,
                  model_path: str,
-                 fastvideo_args: Union[FastVideoArgs, TrainingArgs],
-                 required_config_modules: Optional[List[str]] = None,
-                 loaded_modules: Optional[Dict[str, torch.nn.Module]] = None):
+                 fastvideo_args: FastVideoArgs | TrainingArgs,
+                 required_config_modules: list[str] | None = None,
+                 loaded_modules: dict[str, torch.nn.Module] | None = None):
         """
         Initialize the pipeline. After __init__, the pipeline should be ready to
         use. The pipeline should be stateless and not hold any batch state.
@@ -54,8 +54,8 @@ class ComposedPipelineBase(ABC):
         self.fastvideo_args = fastvideo_args
 
         self.model_path: str = model_path
-        self._stages: List[PipelineStage] = []
-        self._stage_name_mapping: Dict[str, PipelineStage] = {}
+        self._stages: list[PipelineStage] = []
+        self._stage_name_mapping: dict[str, PipelineStage] = {}
 
         if required_config_modules is not None:
             self._required_config_modules = required_config_modules
@@ -97,15 +97,13 @@ class ComposedPipelineBase(ABC):
     @classmethod
     def from_pretrained(cls,
                         model_path: str,
-                        device: Optional[str] = None,
-                        torch_dtype: Optional[torch.dtype] = None,
-                        pipeline_config: Optional[
-                            Union[str
-                                  | PipelineConfig]] = None,
-                        args: Optional[argparse.Namespace] = None,
-                        required_config_modules: Optional[List[str]] = None,
-                        loaded_modules: Optional[Dict[str,
-                                                      torch.nn.Module]] = None,
+                        device: str | None = None,
+                        torch_dtype: torch.dtype | None = None,
+                        pipeline_config: str | PipelineConfig | None = None,
+                        args: argparse.Namespace | None = None,
+                        required_config_modules: list[str] | None = None,
+                        loaded_modules: dict[str, torch.nn.Module]
+                        | None = None,
                         **kwargs) -> "ComposedPipelineBase":
         """
         Load a pipeline from a pretrained model.
@@ -150,16 +148,16 @@ class ComposedPipelineBase(ABC):
     def add_module(self, module_name: str, module: Any):
         self.modules[module_name] = module
 
-    def _load_config(self, model_path: str) -> Dict[str, Any]:
+    def _load_config(self, model_path: str) -> dict[str, Any]:
         model_path = maybe_download_model(self.model_path)
         self.model_path = model_path
         # fastvideo_args.downloaded_model_path = model_path
         logger.info("Model path: %s", model_path)
         config = verify_model_config_and_directory(model_path)
-        return cast(Dict[str, Any], config)
+        return cast(dict[str, Any], config)
 
     @property
-    def required_config_modules(self) -> List[str]:
+    def required_config_modules(self) -> list[str]:
         """
         List of modules that are required by the pipeline. The names should match
         the diffusers directory and model_index.json file. These modules will be
@@ -177,7 +175,7 @@ class ComposedPipelineBase(ABC):
         return self._required_config_modules
 
     @property
-    def stages(self) -> List[PipelineStage]:
+    def stages(self) -> list[PipelineStage]:
         """
         List of stages in the pipeline.
         """
@@ -205,8 +203,8 @@ class ComposedPipelineBase(ABC):
     def load_modules(
         self,
         fastvideo_args: FastVideoArgs,
-        loaded_modules: Optional[Dict[str, torch.nn.Module]] = None
-    ) -> Dict[str, Any]:
+        loaded_modules: dict[str, torch.nn.Module] | None = None
+    ) -> dict[str, Any]:
         """
         Load the modules from the config.
         loaded_modules: Optional[Dict[str, torch.nn.Module]] = None, 

fastvideo/v1/pipelines/lora_pipeline.py

@@ -1,6 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0
 from collections import defaultdict
-from typing import Any, DefaultDict, Dict, Hashable, List, Optional
+from collections.abc import Hashable
+from typing import Any
 
 import torch
 import torch.distributed as dist
@@ -22,12 +23,12 @@ class LoRAPipeline(ComposedPipelineBase):
     Pipeline that supports injecting LoRA adapters into the diffusion transformer.
     TODO: support training.
     """
-    lora_adapters: Dict[str, Dict[str, torch.Tensor]] = defaultdict(
+    lora_adapters: dict[str, dict[str, torch.Tensor]] = defaultdict(
         dict)  # state dicts of loaded lora adapters
     cur_adapter_name: str = ""
-    lora_layers: Dict[str, BaseLayerWithLoRA] = {}
+    lora_layers: dict[str, BaseLayerWithLoRA] = {}
     fastvideo_args: FastVideoArgs
-    exclude_lora_layers: List[str] = []
+    exclude_lora_layers: list[str] = []
     device: torch.device = torch.device(f"cuda:{torch.cuda.current_device()}")
 
     def __init__(self, *args, **kwargs):
@@ -72,7 +73,7 @@ class LoRAPipeline(ComposedPipelineBase):
 
     def set_lora_adapter(self,
                          lora_nickname: str,
-                         lora_path: Optional[str] = None):  # type: ignore
+                         lora_path: str | None = None):  # type: ignore
         """
         Loads a LoRA adapter into the pipeline and applies it to the transformer.
         Args:
@@ -95,8 +96,8 @@ class LoRAPipeline(ComposedPipelineBase):
             lora_param_names_mapping_fn = get_param_names_mapping(
                 self.modules["transformer"]._lora_param_names_mapping)
 
-            to_merge_params: DefaultDict[Hashable,
-                                         Dict[Any, Any]] = defaultdict(dict)
+            to_merge_params: defaultdict[Hashable,
+                                         dict[Any, Any]] = defaultdict(dict)
             for name, weight in lora_state_dict.items():
                 name = ".".join(
                     name.split(".")

fastvideo/v1/pipelines/pipeline_batch_info.py

@@ -9,7 +9,7 @@ in a functional manner, reducing the need for explicit parameter passing.
 
 import pprint
 from dataclasses import asdict, dataclass, field
-from typing import Any, Dict, List, Optional, Union
+from typing import Any
 
 import PIL.Image
 import torch
@@ -33,84 +33,84 @@ class ForwardBatch:
     # specific arguments.
     data_type: str
 
-    generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None
+    generator: torch.Generator | list[torch.Generator] | None = None
 
     # Image inputs
-    image_path: Optional[str] = None
-    image_embeds: List[torch.Tensor] = field(default_factory=list)
-    pil_image: Optional[PIL.Image.Image] = None
-    preprocessed_image: Optional[torch.Tensor] = None
+    image_path: str | None = None
+    image_embeds: list[torch.Tensor] = field(default_factory=list)
+    pil_image: PIL.Image.Image | None = None
+    preprocessed_image: torch.Tensor | None = None
 
     # Text inputs
-    prompt: Optional[Union[str, List[str]]] = None
-    negative_prompt: Optional[Union[str, List[str]]] = None
-    prompt_path: Optional[str] = None
+    prompt: str | list[str] | None = None
+    negative_prompt: str | list[str] | None = None
+    prompt_path: str | None = None
     output_path: str = "outputs/"
 
     # Primary encoder embeddings
-    prompt_embeds: List[torch.Tensor] = field(default_factory=list)
-    negative_prompt_embeds: Optional[List[torch.Tensor]] = None
-    prompt_attention_mask: Optional[List[torch.Tensor]] = None
-    negative_attention_mask: Optional[List[torch.Tensor]] = None
-    clip_embedding_pos: Optional[List[torch.Tensor]] = None
-    clip_embedding_neg: Optional[List[torch.Tensor]] = None
+    prompt_embeds: list[torch.Tensor] = field(default_factory=list)
+    negative_prompt_embeds: list[torch.Tensor] | None = None
+    prompt_attention_mask: list[torch.Tensor] | None = None
+    negative_attention_mask: list[torch.Tensor] | None = None
+    clip_embedding_pos: list[torch.Tensor] | None = None
+    clip_embedding_neg: list[torch.Tensor] | None = None
 
     # Additional text-related parameters
-    max_sequence_length: Optional[int] = None
-    prompt_template: Optional[Dict[str, Any]] = None
+    max_sequence_length: int | None = None
+    prompt_template: dict[str, Any] | None = None
     do_classifier_free_guidance: bool = False
 
     # Batch info
-    batch_size: Optional[int] = None
+    batch_size: int | None = None
     num_videos_per_prompt: int = 1
-    seed: Optional[int] = None
-    seeds: Optional[List[int]] = None
+    seed: int | None = None
+    seeds: list[int] | None = None
 
     # Tracking if embeddings are already processed
     is_prompt_processed: bool = False
 
     # Latent tensors
-    latents: Optional[torch.Tensor] = None
-    raw_latent_shape: Optional[torch.Tensor] = None
-    noise_pred: Optional[torch.Tensor] = None
-    image_latent: Optional[torch.Tensor] = None
+    latents: torch.Tensor | None = None
+    raw_latent_shape: torch.Tensor | None = None
+    noise_pred: torch.Tensor | None = None
+    image_latent: torch.Tensor | None = None
 
     # Latent dimensions
-    height_latents: Optional[int] = None
-    width_latents: Optional[int] = None
+    height_latents: int | None = None
+    width_latents: int | None = None
     num_frames: int = 1  # Default for image models
     num_frames_round_down: bool = False  # Whether to round down num_frames if it's not divisible by num_gpus
 
     # Original dimensions (before VAE scaling)
-    height: Optional[int] = None
-    width: Optional[int] = None
-    fps: Optional[int] = None
+    height: int | None = None
+    width: int | None = None
+    fps: int | None = None
 
     # Timesteps
-    timesteps: Optional[torch.Tensor] = None
-    timestep: Optional[Union[torch.Tensor, float, int]] = None
-    step_index: Optional[int] = None
+    timesteps: torch.Tensor | None = None
+    timestep: torch.Tensor | float | int | None = None
+    step_index: int | None = None
 
     # Scheduler parameters
     num_inference_steps: int = 50
     guidance_scale: float = 1.0
     guidance_rescale: float = 0.0
     eta: float = 0.0
-    sigmas: Optional[List[float]] = None
+    sigmas: list[float] | None = None
 
-    n_tokens: Optional[int] = None
+    n_tokens: int | None = None
 
     # Other parameters that may be needed by specific schedulers
-    extra_step_kwargs: Dict[str, Any] = field(default_factory=dict)
+    extra_step_kwargs: dict[str, Any] = field(default_factory=dict)
 
     # Component modules (populated by the pipeline)
-    modules: Dict[str, Any] = field(default_factory=dict)
+    modules: dict[str, Any] = field(default_factory=dict)
 
     # Final output (after pipeline completion)
     output: Any = None
 
     # Extra parameters that might be needed by specific pipeline implementations
-    extra: Dict[str, Any] = field(default_factory=dict)
+    extra: dict[str, Any] = field(default_factory=dict)
 
     # Misc
     save_video: bool = True
@@ -118,13 +118,13 @@ class ForwardBatch:
 
     # TeaCache parameters
     enable_teacache: bool = False
-    teacache_params: Optional[TeaCacheParams | WanTeaCacheParams] = None
+    teacache_params: TeaCacheParams | WanTeaCacheParams | None = None
 
     # STA parameters
-    STA_param: Optional[List] = None
+    STA_param: list | None = None
     is_cfg_negative: bool = False
-    mask_search_final_result_pos: Optional[List[List]] = None
-    mask_search_final_result_neg: Optional[List[List]] = None
+    mask_search_final_result_pos: list[list] | None = None
+    mask_search_final_result_neg: list[list] | None = None
 
     # VSA parameters
     VSA_sparsity: float = 0.0
@@ -148,25 +148,25 @@ class TrainingBatch:
     current_vsa_sparsity: float = 0.0
 
     # Dataloader batch outputs
-    latents: Optional[torch.Tensor] = None
-    encoder_hidden_states: Optional[torch.Tensor] = None
-    encoder_attention_mask: Optional[torch.Tensor] = None
+    latents: torch.Tensor | None = None
+    encoder_hidden_states: torch.Tensor | None = None
+    encoder_attention_mask: torch.Tensor | None = None
     # i2v
-    preprocessed_image: Optional[torch.Tensor] = None
-    image_embeds: Optional[torch.Tensor] = None
-    image_latents: Optional[torch.Tensor] = None
-    infos: Optional[List[Dict[str, Any]]] = None
+    preprocessed_image: torch.Tensor | None = None
+    image_embeds: torch.Tensor | None = None
+    image_latents: torch.Tensor | None = None
+    infos: list[dict[str, Any]] | None = None
 
     # Transformer inputs
-    noisy_model_input: Optional[torch.Tensor] = None
-    timesteps: Optional[torch.Tensor] = None
-    sigmas: Optional[torch.Tensor] = None
-    noise: Optional[torch.Tensor] = None
+    noisy_model_input: torch.Tensor | None = None
+    timesteps: torch.Tensor | None = None
+    sigmas: torch.Tensor | None = None
+    noise: torch.Tensor | None = None
 
-    attn_metadata: Optional[AttentionMetadata] = None
+    attn_metadata: AttentionMetadata | None = None
 
     # input kwargs
-    input_kwargs: Optional[Dict[str, Any]] = None
+    input_kwargs: dict[str, Any] | None = None
 
     # Training loss
     loss: torch.Tensor | None = None

fastvideo/v1/pipelines/pipeline_registry.py

@@ -4,9 +4,9 @@
 
 import importlib
 import pkgutil
+from collections.abc import Set
 from dataclasses import dataclass, field
 from functools import lru_cache
-from typing import AbstractSet, Dict, Optional, Tuple, Type, Union
 
 from fastvideo.v1.logger import init_logger
 from fastvideo.v1.pipelines.composed_pipeline_base import ComposedPipelineBase
@@ -18,14 +18,14 @@ logger = init_logger(__name__)
 @dataclass
 class _PipelineRegistry:
     # Keyed by pipeline_arch
-    pipelines: Dict[str, Optional[Type[ComposedPipelineBase]]] = field(
-        default_factory=dict)
+    pipelines: dict[str, type[ComposedPipelineBase]
+                    | None] = field(default_factory=dict)
 
-    def get_supported_archs(self) -> AbstractSet[str]:
+    def get_supported_archs(self) -> Set[str]:
         return self.pipelines.keys()
 
     def _try_load_pipeline_cls(
-            self, pipeline_arch: str) -> Optional[Type[ComposedPipelineBase]]:
+            self, pipeline_arch: str) -> type[ComposedPipelineBase] | None:
         if pipeline_arch not in self.pipelines:
             return None
 
@@ -34,7 +34,7 @@ class _PipelineRegistry:
     def resolve_pipeline_cls(
         self,
         architecture: str,
-    ) -> Tuple[Union[Type[ComposedPipelineBase], Type[LoRAPipeline]], str]:
+    ) -> tuple[type[ComposedPipelineBase] | type[LoRAPipeline], str]:
         if not architecture:
             logger.warning("No pipeline architecture is specified")
 

fastvideo/v1/pipelines/preprocess/preprocess_pipeline_base.py

@@ -5,7 +5,7 @@ import os
 from collections import defaultdict
 from concurrent.futures import ProcessPoolExecutor
 from itertools import chain
-from typing import Any, Dict, List, Optional
+from typing import Any
 
 import numpy as np
 import PIL.Image
@@ -47,24 +47,24 @@ class BasePreprocessPipeline(ComposedPipelineBase):
         args,
     ):
         # Initialize class variables for data sharing
-        self.video_data: Dict[str, Any] = {}  # Store video metadata and paths
-        self.latent_data: Dict[str, Any] = {}  # Store latent tensors
+        self.video_data: dict[str, Any] = {}  # Store video metadata and paths
+        self.latent_data: dict[str, Any] = {}  # Store latent tensors
         self.preprocess_validation(fastvideo_args, args)
         self.preprocess_video_and_text(fastvideo_args, args)
 
-    def get_extra_features(self, valid_data: Dict[str, Any],
-                           fastvideo_args: FastVideoArgs) -> Dict[str, Any]:
+    def get_extra_features(self, valid_data: dict[str, Any],
+                           fastvideo_args: FastVideoArgs) -> dict[str, Any]:
         """Get additional features specific to the pipeline type. Override in subclasses."""
         return {}
 
-    def get_schema_fields(self) -> List[str]:
+    def get_schema_fields(self) -> list[str]:
         """Get the schema fields for the pipeline type. Override in subclasses."""
         raise NotImplementedError
 
     def create_record_for_schema(self,
                                  preprocess_batch: PreprocessBatch,
                                  schema: pa.Schema,
-                                 strict: bool = False) -> Dict[str, Any]:
+                                 strict: bool = False) -> dict[str, Any]:
         """Create a record for the Parquet dataset using a generic schema-based approach.
         
         Args:
@@ -222,9 +222,9 @@ class BasePreprocessPipeline(ComposedPipelineBase):
             video_name: str,
             vae_latent: np.ndarray,
             text_embedding: np.ndarray,
-            valid_data: Dict[str, Any],
+            valid_data: dict[str, Any],
             idx: int,
-            extra_features: Optional[Dict[str, Any]] = None) -> Dict[str, Any]:
+            extra_features: dict[str, Any] | None = None) -> dict[str, Any]:
         """Create a record for the Parquet dataset."""
         record = {
             "id":
@@ -622,12 +622,12 @@ class BasePreprocessPipeline(ComposedPipelineBase):
             gc.collect()  # Force garbage collection
 
     def preprocess_image(self, image: PIL.Image.Image, height: int, width: int,
-                         fastvideo_args: FastVideoArgs) -> Dict[str, Any]:
+                         fastvideo_args: FastVideoArgs) -> dict[str, Any]:
         return {}
 
     def preprocess_video(self, video: list[PIL.Image.Image], height: int,
                          width: int,
-                         fastvideo_args: FastVideoArgs) -> Dict[str, Any]:
+                         fastvideo_args: FastVideoArgs) -> dict[str, Any]:
         return {}
 
     def _flush_tables(self, num_processed_samples: int, args,

fastvideo/v1/pipelines/preprocess/preprocess_pipeline_i2v.py

@@ -5,7 +5,7 @@ I2V Data Preprocessing pipeline implementation.
 This module contains an implementation of the I2V Data Preprocessing pipeline
 using the modular pipeline architecture.
 """
-from typing import Any, Dict, List, Optional
+from typing import Any
 
 import numpy as np
 import PIL
@@ -46,7 +46,7 @@ class PreprocessPipeline_I2V(BasePreprocessPipeline):
                        ))
 
     def preprocess_image(self, image: PIL.Image.Image, height: int, width: int,
-                         fastvideo_args: FastVideoArgs) -> Dict[str, Any]:
+                         fastvideo_args: FastVideoArgs) -> dict[str, Any]:
         assert hasattr(
             self,
             "image_encoding_stage"), "Image encoding stage must be created"
@@ -71,15 +71,15 @@ class PreprocessPipeline_I2V(BasePreprocessPipeline):
 
     def preprocess_video(self, video: list[PIL.Image.Image], height: int,
                          width: int,
-                         fastvideo_args: FastVideoArgs) -> Dict[str, Any]:
+                         fastvideo_args: FastVideoArgs) -> dict[str, Any]:
         return self.preprocess_image(video[0], height, width, fastvideo_args)
 
-    def get_schema_fields(self) -> List[str]:
+    def get_schema_fields(self) -> list[str]:
         """Get the schema fields for I2V pipeline."""
         return [f.name for f in pyarrow_schema_i2v]
 
-    def get_extra_features(self, valid_data: Dict[str, Any],
-                           fastvideo_args: FastVideoArgs) -> Dict[str, Any]:
+    def get_extra_features(self, valid_data: dict[str, Any],
+                           fastvideo_args: FastVideoArgs) -> dict[str, Any]:
 
         # TODO(will): move these to cpu at some point
         self.get_module("image_encoder").to(get_local_torch_device())
@@ -184,9 +184,9 @@ class PreprocessPipeline_I2V(BasePreprocessPipeline):
             video_name: str,
             vae_latent: np.ndarray,
             text_embedding: np.ndarray,
-            valid_data: Dict[str, Any],
+            valid_data: dict[str, Any],
             idx: int,
-            extra_features: Optional[Dict[str, Any]] = None) -> Dict[str, Any]:
+            extra_features: dict[str, Any] | None = None) -> dict[str, Any]:
         """Create a record for the Parquet dataset with CLIP features."""
         record = super().create_record(video_name=video_name,
                                        vae_latent=vae_latent,

fastvideo/v1/pipelines/stages/denoising.py

@@ -4,7 +4,8 @@ Denoising stage for diffusion pipelines.
 """
 
 import inspect
-from typing import Any, Dict, Iterable, Optional
+from collections.abc import Iterable
+from typing import Any
 
 import torch
 from einops import rearrange
@@ -302,7 +303,7 @@ class DenoisingStage(PipelineStage):
 
         return batch
 
-    def prepare_extra_func_kwargs(self, func, kwargs) -> Dict[str, Any]:
+    def prepare_extra_func_kwargs(self, func, kwargs) -> dict[str, Any]:
         """
         Prepare extra kwargs for the scheduler step / denoise step.
         
@@ -321,8 +322,8 @@ class DenoisingStage(PipelineStage):
         return extra_step_kwargs
 
     def progress_bar(self,
-                     iterable: Optional[Iterable] = None,
-                     total: Optional[int] = None) -> tqdm:
+                     iterable: Iterable | None = None,
+                     total: int | None = None) -> tqdm:
         """
         Create a progress bar for the denoising process.
         

fastvideo/v1/pipelines/stages/encoding.py

@@ -2,7 +2,6 @@
 """
 Encoding stage for diffusion pipelines.
 """
-from typing import Optional
 
 import PIL.Image
 import torch
@@ -143,7 +142,7 @@ class EncodingStage(PipelineStage):
 
     def retrieve_latents(self,
                          encoder_output: torch.Tensor,
-                         generator: Optional[torch.Generator] = None,
+                         generator: torch.Generator | None = None,
                          sample_mode: str = "sample"):
         if sample_mode == "sample":
             return encoder_output.sample(generator)
@@ -157,8 +156,8 @@ class EncodingStage(PipelineStage):
             self,
             image: PIL.Image.Image,
             vae_scale_factor: int,
-            height: Optional[int] = None,
-            width: Optional[int] = None,
+            height: int | None = None,
+            width: int | None = None,
             resize_mode: str = "default",  # "default", "fill", "crop"
     ) -> torch.Tensor:
         image = [image]

fastvideo/v1/pipelines/stages/text_encoding.py

@@ -56,12 +56,14 @@ class TextEncodingStage(PipelineStage):
             fastvideo_args.pipeline_config.text_encoder_configs)
 
         for tokenizer, text_encoder, encoder_config, preprocess_func, postprocess_func in zip(
-                self.tokenizers, self.text_encoders,
+                self.tokenizers,
+                self.text_encoders,
                 fastvideo_args.pipeline_config.text_encoder_configs,
                 fastvideo_args.pipeline_config.preprocess_text_funcs,
-                fastvideo_args.pipeline_config.postprocess_text_funcs):
+                fastvideo_args.pipeline_config.postprocess_text_funcs,
+                strict=True):
 
-            assert isinstance(batch.prompt, (str, list))
+            assert isinstance(batch.prompt, (str | list))
             if isinstance(batch.prompt, str):
                 batch.prompt = [batch.prompt]
             texts = []

fastvideo/v1/pipelines/stages/validators.py

@@ -6,7 +6,8 @@ This module provides reusable validation functions that can be used across
 all pipeline stages for input/output verification.
 """
 
-from typing import Any, Callable, Dict, List, Optional, Union
+from collections.abc import Callable
+from typing import Any
 
 import torch
 
@@ -27,12 +28,12 @@ class StageValidators:
     @staticmethod
     def positive_float(value: Any) -> bool:
         """Check if value is a positive float."""
-        return isinstance(value, (int, float)) and value > 0
+        return isinstance(value, (int | float)) and value > 0
 
     @staticmethod
     def non_negative_float(value: Any) -> bool:
         """Check if value is a non-negative float."""
-        return isinstance(value, (int, float)) and value >= 0
+        return isinstance(value, (int | float)) and value >= 0
 
     @staticmethod
     def divisible_by(value: Any, divisor: int) -> bool:
@@ -72,7 +73,7 @@ class StageValidators:
             return False
         if len(value.shape) != len(expected_shape):
             return False
-        for actual, expected in zip(value.shape, expected_shape):
+        for actual, expected in zip(value.shape, expected_shape, strict=True):
             if expected is not None and actual != expected:
                 return False
         return not torch.isnan(value).any().item()
@@ -268,8 +269,8 @@ class ValidationFailure:
     def __init__(self,
                  validator_name: str,
                  actual_value: Any,
-                 expected: Optional[str] = None,
-                 error_msg: Optional[str] = None):
+                 expected: str | None = None,
+                 error_msg: str | None = None):
         self.validator_name = validator_name
         self.actual_value = actual_value
         self.expected = expected
@@ -317,12 +318,12 @@ class VerificationResult:
     """Wrapper class for stage verification results."""
 
     def __init__(self) -> None:
-        self._checks: Dict[str, bool] = {}
-        self._failures: Dict[str, List[ValidationFailure]] = {}
+        self._checks: dict[str, bool] = {}
+        self._failures: dict[str, list[ValidationFailure]] = {}
 
     def add_check(
         self, field_name: str, value: Any,
-        validators: Union[Callable[[Any], bool], List[Callable[[Any], bool]]]
+        validators: Callable[[Any], bool] | list[Callable[[Any], bool]]
     ) -> 'VerificationResult':
         """
         Add a validation check for a field.
@@ -455,11 +456,11 @@ class VerificationResult:
         """Check if all validations passed."""
         return all(self._checks.values())
 
-    def get_failed_fields(self) -> List[str]:
+    def get_failed_fields(self) -> list[str]:
         """Get list of fields that failed validation."""
         return [field for field, passed in self._checks.items() if not passed]
 
-    def get_detailed_failures(self) -> Dict[str, List[ValidationFailure]]:
+    def get_detailed_failures(self) -> dict[str, list[ValidationFailure]]:
         """Get detailed failure information for each failed field."""
         return self._failures.copy()