nodes/hao-ai-lab-FastVideo
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

update rotary embedding and registry for gamecraft

Browse Source
This commit is contained in:
Mihir Jagtap
2026-02-07 14:04:26 -08:00
parent 66679d8ab3
commit 8efeee8da2
2 changed files with 62 additions and 24 deletions
Download Patch File Download Diff File Expand all files Collapse all files
fastvideo/layers/rotary_embedding.py
+59 -22
View File
@@ -109,26 +109,43 @@ def _apply_rotary_emb(
"""
Args:
x: [num_tokens, num_heads, head_size]
cos: [num_tokens, head_size // 2]
sin: [num_tokens, head_size // 2]
cos: [num_tokens, head_size] or [num_tokens, head_size // 2]
sin: [num_tokens, head_size] or [num_tokens, head_size // 2]
is_neox_style: Whether to use the Neox-style or GPT-J-style rotary
positional embeddings.
The function auto-detects whether cos/sin are full or half head_size:
- If cos/sin have head_size: use rotate_half style (for HunyuanVideo/GameCraft)
- If cos/sin have head_size // 2: use Neox/GPT-J style
"""
# cos = cos.unsqueeze(-2).to(x.dtype)
# sin = sin.unsqueeze(-2).to(x.dtype)
cos = cos.unsqueeze(-2)
sin = sin.unsqueeze(-2)
if is_neox_style:
x1, x2 = torch.chunk(x, 2, dim=-1)
head_size = x.shape[-1]
rope_dim = cos.shape[-1]
# Check if cos/sin are full head_dim (rotate_half style) or half (traditional style)
if rope_dim == head_size:
# Full head_dim - use rotate_half style (HunyuanVideo, GameCraft)
# x * cos + rotate_half(x) * sin
cos = cos.unsqueeze(-2) # [num_tokens, 1, head_size]
sin = sin.unsqueeze(-2) # [num_tokens, 1, head_size]
# rotate_half: split into pairs, negate and swap
x_real, x_imag = x.float().reshape(*x.shape[:-1], -1, 2).unbind(-1) # [B, H, D//2] each
x_rotated = torch.stack([-x_imag, x_real], dim=-1).flatten(-2) # [B, H, D]
return (x.float() * cos + x_rotated * sin).type_as(x)
else:
x1 = x[..., ::2]
x2 = x[..., 1::2]
o1 = (x1.float() * cos - x2.float() * sin).type_as(x)
o2 = (x2.float() * cos + x1.float() * sin).type_as(x)
if is_neox_style:
return torch.cat((o1, o2), dim=-1)
else:
return torch.stack((o1, o2), dim=-1).flatten(-2)
# Half head_dim - use traditional Neox/GPT-J style
cos = cos.unsqueeze(-2)
sin = sin.unsqueeze(-2)
if is_neox_style:
x1, x2 = torch.chunk(x, 2, dim=-1)
else:
x1 = x[..., ::2]
x2 = x[..., 1::2]
o1 = (x1.float() * cos - x2.float() * sin).type_as(x)
o2 = (x2.float() * cos + x1.float() * sin).type_as(x)
if is_neox_style:
return torch.cat((o1, o2), dim=-1)
else:
return torch.stack((o1, o2), dim=-1).flatten(-2)
@CustomOp.register("rotary_embedding")
@@ -278,6 +295,7 @@ def get_1d_rotary_pos_embed(
theta_rescale_factor: float = 1.0,
interpolation_factor: float = 1.0,
dtype: torch.dtype = torch.float32,
use_real: bool = True,
) -> tuple[torch.Tensor, torch.Tensor]:
"""
Precompute the frequency tensor for complex exponential (cis) with given dimensions.
@@ -292,9 +310,12 @@ def get_1d_rotary_pos_embed(
theta (float, optional): Scaling factor for frequency computation. Defaults to 10000.0.
theta_rescale_factor (float, optional): Rescale factor for theta. Defaults to 1.0.
interpolation_factor (float, optional): Factor to scale positions. Defaults to 1.0.
use_real (bool, optional): If True, output full head_dim with repeated cos/sin for
rotate_half style RoPE. If False, output half head_dim for complex style. Defaults to True.
Returns:
freqs_cos, freqs_sin: Precomputed frequency tensor with real and imaginary parts separately. [S, D]
freqs_cos, freqs_sin: Precomputed frequency tensor with real and imaginary parts separately.
Shape is [S, D] if use_real=True, [S, D/2] if use_real=False.
"""
if isinstance(pos, int):
pos = torch.arange(pos).float()
@@ -309,6 +330,16 @@ def get_1d_rotary_pos_embed(
freqs = torch.outer(pos * interpolation_factor, freqs) # [S, D/2]
freqs_cos = freqs.cos() # [S, D/2]
freqs_sin = freqs.sin() # [S, D/2]
if use_real:
# For rotate_half style RoPE (used by HunyuanVideo, GameCraft),
# we need to repeat cos/sin to full head_dim
# The rotate_half operation expects: x * cos + rotate_half(x) * sin
# where rotate_half splits x into [x1, x2] and returns [-x2, x1]
# So we need cos and sin to have shape [S, D] to broadcast correctly
freqs_cos = torch.cat([freqs_cos, freqs_cos], dim=-1) # [S, D]
freqs_sin = torch.cat([freqs_sin, freqs_sin], dim=-1) # [S, D]
return freqs_cos, freqs_sin
@@ -324,6 +355,7 @@ def get_nd_rotary_pos_embed(
sp_world_size: int = 1,
dtype: torch.dtype = torch.float32,
start_frame: int = 0,
use_real: bool = True,
) -> 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.
@@ -341,9 +373,10 @@ def get_nd_rotary_pos_embed(
shard_dim (int): Which dimension to shard for sequence parallelism. Defaults to 0.
sp_rank (int): Rank in the sequence parallel group. Defaults to 0.
sp_world_size (int): World size of the sequence parallel group. Defaults to 1.
use_real (bool): If True, output full head_dim for rotate_half style. Defaults to True.
Returns:
Tuple[torch.Tensor, torch.Tensor]: (cos, sin) tensors of shape [HW, D/2]
Tuple[torch.Tensor, torch.Tensor]: (cos, sin) tensors of shape [HW, D] if use_real, [HW, D/2] otherwise
"""
# Get the full grid
full_grid = get_meshgrid_nd(
@@ -412,11 +445,12 @@ def get_nd_rotary_pos_embed(
theta_rescale_factor=theta_rescale_factor[i],
interpolation_factor=interpolation_factor[i],
dtype=dtype,
) # 2 x [WHD, rope_dim_list[i]]
use_real=use_real,
) # 2 x [WHD, rope_dim_list[i]] or 2 x [WHD, rope_dim_list[i]*2] if use_real
embs.append(emb)
cos = torch.cat([emb[0] for emb in embs], dim=1) # (WHD, D/2)
sin = torch.cat([emb[1] for emb in embs], dim=1) # (WHD, D/2)
cos = torch.cat([emb[0] for emb in embs], dim=1) # (WHD, D) or (WHD, D/2)
sin = torch.cat([emb[1] for emb in embs], dim=1) # (WHD, D) or (WHD, D/2)
return cos, sin
@@ -432,6 +466,7 @@ def get_rotary_pos_embed(
do_sp_sharding: bool = False,
dtype: torch.dtype = torch.float32,
start_frame: int = 0,
use_real: bool = True,
) -> tuple[torch.Tensor, torch.Tensor]:
"""
Generate rotary positional embeddings for the given sizes.
@@ -446,9 +481,10 @@ def get_rotary_pos_embed(
interpolation_factor: Factor to scale positions. Defaults to 1.0
shard_dim: Which dimension to shard for sequence parallelism. Defaults to 0.
do_sp_sharding: Whether to shard the positional embeddings for sequence parallelism. Defaults to False.
use_real: If True, output full head_dim for rotate_half style RoPE. Defaults to True.
Returns:
Tuple of (cos, sin) tensors for rotary embeddings
Tuple of (cos, sin) tensors for rotary embeddings. Shape [S, D] if use_real, [S, D/2] otherwise.
"""
target_ndim = 3
@@ -481,6 +517,7 @@ def get_rotary_pos_embed(
sp_world_size=sp_world_size,
dtype=dtype,
start_frame=start_frame,
use_real=use_real,
)
return freqs_cos, freqs_sin
fastvideo/models/registry.py
+3 -2
View File
@@ -25,10 +25,10 @@ logger = init_logger(__name__)
_TEXT_TO_VIDEO_DIT_MODELS = {
"HunyuanVideoTransformer3DModel":
("dits", "hunyuanvideo", "HunyuanVideoTransformer3DModel"),
"HunyuanVideo15Transformer3DModel":
("dits", "hunyuanvideo15", "HunyuanVideo15Transformer3DModel"),
"HunyuanGameCraftTransformer3DModel":
("dits", "hunyuangamecraft", "HunyuanGameCraftTransformer3DModel"),
"HunyuanVideo15Transformer3DModel":
("dits", "hunyuanvideo15", "HunyuanVideo15Transformer3DModel"),
"HYWorldTransformer3DModel":
("dits", "hyworld", "HYWorldTransformer3DModel"),
"WanTransformer3DModel": ("dits", "wanvideo", "WanTransformer3DModel"),
@@ -73,6 +73,7 @@ _IMAGE_ENCODER_MODELS: dict[str, tuple] = {
_VAE_MODELS = {
"AutoencoderKLHunyuanVideo":
("vaes", "hunyuanvae", "AutoencoderKLHunyuanVideo"),
"AutoencoderKLCausal3D": ("vaes", "gamecraftvae", "GameCraftVAE"),
"AutoencoderKLHYWorld": ("vaes", "hyworldvae", "AutoencoderKLHYWorld"),
"AutoencoderKLHunyuanVideo15": ("vaes", "hunyuan15vae", "AutoencoderKLHunyuanVideo15"),
"AutoencoderKLWan": ("vaes", "wanvae", "AutoencoderKLWan"),