diff --git a/fastvideo/layers/rotary_embedding.py b/fastvideo/layers/rotary_embedding.py index 19b02b71d..1a5c60d61 100644 --- a/fastvideo/layers/rotary_embedding.py +++ b/fastvideo/layers/rotary_embedding.py @@ -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 diff --git a/fastvideo/models/registry.py b/fastvideo/models/registry.py index 93bbfe151..885cc23a7 100644 --- a/fastvideo/models/registry.py +++ b/fastvideo/models/registry.py @@ -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"),