94 lines
2.8 KiB
Python
94 lines
2.8 KiB
Python
"""
|
|
Exact STARFlow-V T2V/I2V (7B, 480p causal video) architecture and inference
|
|
configuration, transcribed from Apple ml-starflow
|
|
configs/starflow-v_7B_t2v_caus_480p.yaml (commit 20984834c1a68ee679faa8ed9cda2842f02b042d).
|
|
|
|
Only the fields needed for inference (no training/optimizer/dataloader knobs)
|
|
are kept here.
|
|
"""
|
|
from dataclasses import dataclass, field
|
|
from typing import List, Optional
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class STARFlowVConfig:
|
|
# ---- image / video geometry ----
|
|
img_size: int = 640
|
|
txt_size: int = 256
|
|
vid_frames: int = 81
|
|
vid_fps: int = 16
|
|
fps_cond: bool = True
|
|
|
|
# ---- latent / patchify ----
|
|
channel_size: int = 48 # VAE latent channels (Wan2.2 TI2V C=48)
|
|
patch_size: int = 1
|
|
|
|
# ---- transformer ----
|
|
channels: int = 3072
|
|
top_block_channels: int = 4096
|
|
blocks: int = 6
|
|
layers_per_block: List[int] = field(default_factory=lambda: [2, 2, 2, 2, 2, 24])
|
|
head_dim: int = 64
|
|
|
|
# ---- flow / attention structure ----
|
|
rope: bool = True
|
|
adaln: bool = False
|
|
sos: bool = True
|
|
seq_order: str = "L2R"
|
|
pt_seq_len: int = 32
|
|
use_softplus: bool = True
|
|
cond_top_only: bool = True
|
|
use_final_norm: bool = True
|
|
|
|
# ---- denoiser ----
|
|
learnable_self_denoiser: bool = True
|
|
conditional_denoiser: bool = True
|
|
denoiser_window: int = 10
|
|
cond_noise_level: bool = True
|
|
|
|
# ---- temporal / locality ----
|
|
temporal_causal: int = 2
|
|
shallow_block_local: bool = True
|
|
local_attn_window: Optional[int] = 21 # training-length latent window (81 frames -> 21 latents)
|
|
|
|
# ---- diffusion-adjacent knobs ----
|
|
noise_std: float = 0.5
|
|
soft_clip: float = 4.0
|
|
|
|
# ---- VAE ----
|
|
vae_name: str = "Wan-AI/Wan2.2-TI2V-5B-Diffusers"
|
|
vae_temporal_scale: float = 0.6
|
|
vae_downsample_factor: int = 16
|
|
vae_temporal_downsample_factor: int = 4
|
|
|
|
# ---- text ----
|
|
text_encoder_name: str = "google/flan-t5-xl"
|
|
|
|
# ---- sampling defaults ----
|
|
cfg: float = 2.5
|
|
|
|
# ---- capability flags (this suite supports T2V + first-frame I2V only) ----
|
|
supports_t2v: bool = True
|
|
supports_i2v_first_frame: bool = True
|
|
supports_v2v: bool = False
|
|
|
|
|
|
STARFLOW_V_T2V = STARFlowVConfig()
|
|
|
|
|
|
def frames_valid(num_frames: int) -> bool:
|
|
"""Video frame counts must be 4n+1 (VAE causal temporal compression)."""
|
|
return num_frames >= 1 and (num_frames - 1) % 4 == 0
|
|
|
|
|
|
def latent_temporal_length(num_frames: int, temporal_downsample_factor: int = 4) -> int:
|
|
assert frames_valid(num_frames), f"num_frames must be 4n+1, got {num_frames}"
|
|
return (num_frames - 1) // temporal_downsample_factor + 1
|
|
|
|
|
|
def spatial_latent_size(pixel_size: int, downsample_factor: int = 16) -> int:
|
|
assert pixel_size % downsample_factor == 0, (
|
|
f"spatial size {pixel_size} must be divisible by downsample_factor {downsample_factor}"
|
|
)
|
|
return pixel_size // downsample_factor
|