75 lines
2.8 KiB
Python
75 lines
2.8 KiB
Python
"""
|
|
Builds the STARFlow-V transformer skeleton (uninitialized/random weights)
|
|
from an STARFlowVConfig, matching the argument mapping used upstream in
|
|
utils/model_setup.py::setup_transformer for the generic (non named-factory)
|
|
model path, specialized to the exact starflow-v_7B_t2v_caus_480p.yaml
|
|
configuration (rope, sos, L2R, cond_top_only, learnable+conditional
|
|
denoiser, temporal_causal=2, shallow_block_local, denoiser_window=10,
|
|
top_block_channels=4096).
|
|
|
|
The optional ``dtype`` argument constructs the skeleton directly in that
|
|
floating dtype (e.g. bfloat16) via ``torch.set_default_dtype``, so the real
|
|
~7B model never needs a transient fp32 allocation before weights are
|
|
loaded. RoPE frequency tables are unaffected because they are built with an
|
|
explicit ``.float()`` upstream.
|
|
"""
|
|
import contextlib
|
|
|
|
import torch
|
|
|
|
from ..config import STARFlowVConfig
|
|
from .transformer_flow import Model
|
|
|
|
# FLAN-T5-XL encoder hidden size (google/flan-t5-xl config.d_model).
|
|
# This is architecturally baked into the checkpoint's proj_txt layers,
|
|
# so it is fixed here rather than derived from a loaded text encoder.
|
|
TEXT_HIDDEN_SIZE = 2048
|
|
|
|
|
|
@contextlib.contextmanager
|
|
def _default_dtype(dtype):
|
|
if dtype is None or not (isinstance(dtype, torch.dtype) and dtype.is_floating_point):
|
|
yield
|
|
return
|
|
prev = torch.get_default_dtype()
|
|
torch.set_default_dtype(dtype)
|
|
try:
|
|
yield
|
|
finally:
|
|
torch.set_default_dtype(prev)
|
|
|
|
|
|
def build_starflow_v_model(cfg: STARFlowVConfig = None, local_attn_window: int = None,
|
|
dtype: torch.dtype = None) -> Model:
|
|
cfg = cfg or STARFlowVConfig()
|
|
latent_img_size = cfg.img_size // cfg.vae_downsample_factor
|
|
|
|
with _default_dtype(dtype):
|
|
model = Model(
|
|
in_channels=cfg.channel_size,
|
|
img_size=latent_img_size,
|
|
patch_size=cfg.patch_size,
|
|
channels=cfg.channels,
|
|
num_blocks=cfg.blocks,
|
|
layers_per_block=list(cfg.layers_per_block),
|
|
head_dim=cfg.head_dim,
|
|
rope=cfg.rope,
|
|
pt_seq_len=cfg.pt_seq_len,
|
|
sos=cfg.sos,
|
|
txt_size=cfg.txt_size,
|
|
txt_dim=TEXT_HIDDEN_SIZE,
|
|
cond_top_only=cfg.cond_top_only,
|
|
use_softplus=cfg.use_softplus,
|
|
use_final_norm=cfg.use_final_norm,
|
|
seq_order=cfg.seq_order,
|
|
learnable_self_denoiser=cfg.learnable_self_denoiser,
|
|
conditional_denoiser=cfg.conditional_denoiser,
|
|
temporal_causal=cfg.temporal_causal,
|
|
top_block_channels=cfg.top_block_channels,
|
|
shallow_block_local=cfg.shallow_block_local,
|
|
denoiser_window=cfg.denoiser_window,
|
|
soft_clip=cfg.soft_clip,
|
|
local_attn_window=cfg.local_attn_window if local_attn_window is None else local_attn_window,
|
|
)
|
|
return model
|