Spaces:
Running on Zero
Running on Zero
| from __future__ import annotations | |
| from dataclasses import dataclass, field | |
| from typing import Optional | |
| class DiTConfig: | |
| dim: int = 1536 | |
| in_channels: int = 16 | |
| out_channels: int = 16 | |
| text_dim: int = 4096 | |
| freq_dim: int = 256 | |
| ffn_dim: int = 8960 | |
| num_heads: int = 12 | |
| num_layers: int = 30 | |
| patch_size: tuple[int, int, int] = (1, 2, 2) | |
| eps: float = 1e-6 | |
| def validate(self) -> None: | |
| if self.dim % self.num_heads: | |
| raise ValueError("dim must be divisible by num_heads") | |
| if any(value <= 0 for value in self.patch_size): | |
| raise ValueError("patch_size values must be positive") | |
| if self.num_layers <= 0: | |
| raise ValueError("num_layers must be positive") | |
| class CausalAttentionConfig: | |
| block_size_frames: int = 1 | |
| include_current_block: bool = True | |
| pad_to_block: bool = True | |
| def validate(self) -> None: | |
| if self.block_size_frames <= 0: | |
| raise ValueError("causal block_size_frames must be positive") | |
| class SparseAttentionConfig: | |
| window_size: tuple[int, int, int] = (2, 8, 8) | |
| topk: Optional[int] = None | |
| topk_ratio: Optional[float] = 2.5 | |
| local_blocks: Optional[int] = None | |
| local_range: Optional[int] = None | |
| def validate(self) -> None: | |
| if len(self.window_size) != 3 or any( | |
| value <= 0 for value in self.window_size | |
| ): | |
| raise ValueError( | |
| "sparse.window_size must contain three positive values" | |
| ) | |
| class SlidingAttentionConfig: | |
| local_range: int = 3 | |
| temporal_block_size: int = 2 | |
| temporal_mode: str = "local" | |
| spatial_mode: str = "local" | |
| cache_temporal_windows: int = 1 | |
| def validate(self) -> None: | |
| self.temporal_mode = self.temporal_mode.lower() | |
| self.spatial_mode = self.spatial_mode.lower() | |
| if self.local_range <= 0: | |
| raise ValueError("sliding.local_range must be positive") | |
| if self.temporal_block_size <= 0: | |
| raise ValueError("sliding.temporal_block_size must be positive") | |
| if self.cache_temporal_windows < 0: | |
| raise ValueError( | |
| "sliding.cache_temporal_windows cannot be negative" | |
| ) | |
| if self.temporal_mode not in {"local", "causal", "short_causal"}: | |
| raise ValueError( | |
| "sliding.temporal_mode must be 'local', 'causal', or " | |
| "'short_causal'" | |
| ) | |
| if ( | |
| self.temporal_mode == "short_causal" | |
| and self.cache_temporal_windows == 0 | |
| ): | |
| raise ValueError( | |
| "sliding short_causal mode requires at least one cached " | |
| "temporal window" | |
| ) | |
| if self.spatial_mode not in {"local", "base"}: | |
| raise ValueError("sliding.spatial_mode must be 'local' or 'base'") | |
| class AttentionConfig: | |
| backend: str = "dense" | |
| kernel: str = "auto" | |
| is_causal: bool = False | |
| is_sparse: bool = False | |
| is_hybrid_sliding: bool = False | |
| causal: CausalAttentionConfig = field(default_factory=CausalAttentionConfig) | |
| sparse: SparseAttentionConfig = field(default_factory=SparseAttentionConfig) | |
| sliding: SlidingAttentionConfig = field(default_factory=SlidingAttentionConfig) | |
| def validate(self) -> None: | |
| self.backend = self.backend.lower() | |
| if self.backend not in {"dense", "sparse", "sliding_local", "hybrid_sliding"}: | |
| raise ValueError(f"Unknown attention backend: {self.backend}") | |
| self.kernel = self.kernel.lower().replace("-", "_") | |
| aliases = { | |
| "fa2": "flash_attn2", | |
| "fa3": "flash_attn3", | |
| "flash2": "flash_attn2", | |
| "flash3": "flash_attn3", | |
| "triton_sparse": "triton_block_sparse", | |
| } | |
| self.kernel = aliases.get(self.kernel, self.kernel) | |
| if self.kernel not in { | |
| "auto", | |
| "sdpa", | |
| "flash_attn2", | |
| "flash_attn3", | |
| "block_sparse", | |
| "triton_block_sparse", | |
| }: | |
| raise ValueError(f"Unknown attention kernel: {self.kernel}") | |
| self.causal.validate() | |
| self.sparse.validate() | |
| self.sliding.validate() | |
| if ( | |
| self.is_hybrid_sliding | |
| and not self.is_sparse | |
| and self.sparse.local_range is None | |
| ): | |
| raise ValueError( | |
| "hybrid_sliding requires draft attention or a fixed local_range" | |
| ) | |
| if ( | |
| self.is_hybrid_sliding | |
| and self.sliding.temporal_mode in {"causal", "short_causal"} | |
| and not self.is_causal | |
| ): | |
| raise ValueError( | |
| "hybrid_sliding causal temporal modes require causal attention" | |
| ) | |
| if ( | |
| self.is_hybrid_sliding | |
| and self.is_causal | |
| and self.causal.block_size_frames != self.sliding.temporal_block_size | |
| ): | |
| raise ValueError( | |
| "hybrid_sliding requires causal.block_size_frames to match " | |
| "sliding.temporal_block_size" | |
| ) | |
| if not self.is_causal and self.causal.block_size_frames != 1: | |
| raise ValueError("block_size_frames is only meaningful when is_causal=True") | |
| if self.is_sparse and self.backend not in {"sparse", "sliding_local", "hybrid_sliding"}: | |
| raise ValueError("is_sparse=True requires a sparse attention backend") | |
| if self.sparse.topk is not None and self.sparse.topk <= 0: | |
| raise ValueError("sparse.topk must be positive") | |
| if self.sparse.topk_ratio is not None and self.sparse.topk_ratio <= 0: | |
| raise ValueError("sparse.topk_ratio must be positive") | |
| if self.sparse.local_blocks is not None and self.sparse.local_blocks <= 0: | |
| raise ValueError("sparse.local_blocks must be positive") | |
| if self.sparse.local_range is not None and self.sparse.local_range <= 0: | |
| raise ValueError("sparse.local_range must be positive") | |
| sparse_masked = self.backend in {"sparse", "sliding_local", "hybrid_sliding"} and ( | |
| self.is_sparse or self.sparse.local_range is not None | |
| ) | |
| if self.kernel in {"flash_attn2", "flash_attn3"} and ( | |
| self.is_causal or sparse_masked | |
| ): | |
| raise ValueError( | |
| f"{self.kernel} cannot represent block causal/sparse masks; " | |
| "use kernel='block_sparse', 'triton_block_sparse', 'sdpa', or 'auto'" | |
| ) | |
| if self.kernel in {"block_sparse", "triton_block_sparse"} and not sparse_masked: | |
| raise ValueError( | |
| "sparse attention kernels require draft sparse attention or local_range" | |
| ) | |
| if self.kernel in {"block_sparse", "triton_block_sparse"}: | |
| block_tokens = ( | |
| self.sparse.window_size[0] | |
| * self.sparse.window_size[1] | |
| * self.sparse.window_size[2] | |
| ) | |
| if block_tokens != 128: | |
| raise ValueError( | |
| "kernel='block_sparse' requires sparse.window_size to " | |
| f"contain 128 tokens, got {self.sparse.window_size} " | |
| f"({block_tokens} tokens)" | |
| ) | |
| class LQProjectionConfig: | |
| enabled: bool = True | |
| num_injection_layers: int = 30 | |
| input_channels: int = 3 | |
| spatial_factor: int = 16 | |
| hidden_dim: int = 2048 | |
| hidden_dim_2: int = 3072 | |
| zero_init: bool = True | |
| pad_last_frames: int = 0 | |
| def validate(self, dit: DiTConfig) -> None: | |
| if self.num_injection_layers < 0: | |
| raise ValueError("num_injection_layers cannot be negative") | |
| if self.pad_last_frames < 0: | |
| raise ValueError("pad_last_frames cannot be negative") | |
| if self.enabled and self.num_injection_layers > dit.num_layers: | |
| raise ValueError("LQ injection layers cannot exceed DiT layers") | |
| class BlockReplacementConfig: | |
| layers: tuple[int, ...] = () | |
| def validate(self) -> None: | |
| if len(set(self.layers)) != len(self.layers): | |
| raise ValueError("replacement.layers cannot contain duplicates") | |
| class FlashVSRModelConfig: | |
| dit: DiTConfig = field(default_factory=DiTConfig) | |
| attention: AttentionConfig = field(default_factory=AttentionConfig) | |
| lq_projection: LQProjectionConfig = field(default_factory=LQProjectionConfig) | |
| replacement: BlockReplacementConfig = field(default_factory=BlockReplacementConfig) | |
| def validate(self) -> None: | |
| self.dit.validate() | |
| self.attention.validate() | |
| self.lq_projection.validate(self.dit) | |
| self.replacement.validate() | |
| if self.replacement.layers and not self.attention.is_hybrid_sliding: | |
| raise ValueError("replacement layers require hybrid_sliding attention") | |
| invalid = [index for index in self.replacement.layers if not 0 <= index < self.dit.num_layers] | |
| if invalid: | |
| raise ValueError(f"Replacement layer indices out of range: {invalid}") | |