gen3c ¶

def forward(
    self,
    hidden_states: torch.Tensor,
    shift: torch.Tensor,
    scale: torch.Tensor,
) -> torch.Tensor:
    """
    Args:
        hidden_states: Input tensor
        shift: Shift parameter for modulation
        scale: Scale parameter for modulation

    Returns:
        normalized_hidden_states: Modulated normalized hidden states
    """
    # Apply layer norm and modulation
    hidden_states = self.norm(hidden_states)
    hidden_states = hidden_states * (1 + scale) + shift
    return hidden_states

def forward(
    self,
    hidden_states: torch.Tensor,
    encoder_hidden_states: torch.Tensor,
    attention_mask: torch.Tensor | None = None,
) -> torch.Tensor:
    """
    Args:
        hidden_states: (B, S, D)
        encoder_hidden_states: (B, N, D_text)
    """
    # Get QKV
    query, _ = self.to_q(hidden_states)
    key, _ = self.to_k(encoder_hidden_states)
    value, _ = self.to_v(encoder_hidden_states)

    # Reshape for multi-head attention
    query = query.unflatten(-1, (self.num_heads, self.head_dim))
    key = key.unflatten(-1, (self.num_heads, self.head_dim))
    value = value.unflatten(-1, (self.num_heads, self.head_dim))

    # Apply QK normalization
    query = self.norm_q(query)
    key = self.norm_k(key)

    attn_output = self.attn(query, key, value)
    attn_output = attn_output.flatten(-2, -1)

    # Output projection
    attn_output, _ = self.to_out(attn_output)
    return attn_output

def forward(
    self,
    timestep: torch.Tensor,
    dtype: torch.dtype,
) -> tuple[torch.Tensor, torch.Tensor | None]:
    """
    Args:
        timestep: (B,) tensor of timesteps
        dtype: Target dtype for embeddings

    Returns:
        embedded_timestep: Timestep embedding (B, D)
        adaln_lora: AdaLN-LoRA parameters (B, 3D) or None
    """
    # timestep should be 1D (B,)
    timestep = timestep.flatten()
    timesteps_proj = self.time_proj(timestep).to(dtype)  # (B, D)
    embedded_timestep, adaln_lora = self.t_embedder(timesteps_proj)

    return embedded_timestep, adaln_lora

def forward(
    self,
    hidden_states: torch.Tensor,
    affine_emb: torch.Tensor,
    adaln_lora: torch.Tensor | None = None,
) -> torch.Tensor:
    """
    Args:
        hidden_states: (B, T, H, W, D)
        affine_emb: (B, D)
        adaln_lora: (B, 3D) or None
    """
    # Generate modulation parameters
    modulation = _run_linear_stack(self.adaln_modulation, affine_emb)

    if self.use_adaln_lora and adaln_lora is not None:
        modulation = modulation + adaln_lora[..., : 2 * self.hidden_size]

    shift, scale = modulation.chunk(2, dim=-1)

    # Apply normalization and modulation
    hidden_states = self.norm(hidden_states)

    # Broadcast over all non-channel dimensions. Supports both:
    # - (B, T, H, W, D)
    # - (B, S, D) when sequence-parallel path is active.
    mod_shape = [hidden_states.shape[0]] + [1] * (hidden_states.dim() - 2) + [self.hidden_size]
    shift = shift.view(*mod_shape)
    scale = scale.view(*mod_shape)

    hidden_states = hidden_states * (1 + scale) + shift

    # Project to output
    hidden_states, _ = self.proj_out(hidden_states)

    return hidden_states

def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
    """
    Args:
        hidden_states: (B, T, H, W, D)
    Returns:
        pos_emb: (B, T, H, W, D)
    """
    B, T, H, W, D = hidden_states.shape

    emb_t = self.pos_emb_t[:T][None, :, None, None, :].repeat(B, 1, H, W, 1)
    emb_h = self.pos_emb_h[:H][None, None, :, None, :].repeat(B, T, 1, W, 1)
    emb_w = self.pos_emb_w[:W][None, None, None, :, :].repeat(B, T, H, 1, 1)

    emb = emb_t + emb_h + emb_w

    # Normalize
    norm = torch.linalg.vector_norm(emb, dim=-1, keepdim=True, dtype=torch.float32)
    norm = torch.add(self.eps, norm, alpha=np.sqrt(norm.numel() / emb.numel()))
    return (emb / norm).type_as(hidden_states)

def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
    """
    Args:
        hidden_states: (B, C, T, H, W)
    Returns:
        (B, T', H', W', D) where T'=T//pt, H'=H//ph, W'=W//pw
    """
    batch_size, num_channels, num_frames, height, width = hidden_states.shape
    p_t, p_h, p_w = self.patch_size

    # Rearrange: b c (t pt) (h ph) (w pw) -> b t h w (c pt ph pw)
    hidden_states = hidden_states.reshape(
        batch_size, num_channels,
        num_frames // p_t, p_t,
        height // p_h, p_h,
        width // p_w, p_w
    )
    hidden_states = hidden_states.permute(0, 2, 4, 6, 1, 3, 5, 7)
    hidden_states = hidden_states.flatten(4, 7)  # Flatten patch dimensions

    # Project to model dimension
    hidden_states, _ = self.proj(hidden_states)
    return hidden_states

def forward(
    self, hidden_states: torch.Tensor, fps: int | None = None
) -> tuple[torch.Tensor, torch.Tensor]:
    """
    Generate 3D RoPE embeddings.
    """
    batch_size, T, H, W, input_dim = hidden_states.shape
    device = hidden_states.device

    # Generate frequency scales with NTK
    h_theta = 10000.0 * self.h_ntk_factor
    w_theta = 10000.0 * self.w_ntk_factor
    t_theta = 10000.0 * self.t_ntk_factor

    seq = torch.arange(max(self.max_size), device=device, dtype=torch.float32)

    dim_h_range = torch.arange(0, self.dim_h, 2, device=device, dtype=torch.float32)[: (self.dim_h // 2)] / self.dim_h
    dim_w_range = torch.arange(0, self.dim_w, 2, device=device, dtype=torch.float32)[: (self.dim_w // 2)] / self.dim_w
    dim_t_range = torch.arange(0, self.dim_t, 2, device=device, dtype=torch.float32)[: (self.dim_t // 2)] / self.dim_t

    h_spatial_freqs = 1.0 / (h_theta ** dim_h_range)
    w_spatial_freqs = 1.0 / (w_theta ** dim_w_range)
    temporal_freqs = 1.0 / (t_theta ** dim_t_range)

    # Generate positional embeddings
    half_emb_h = torch.outer(seq[:H], h_spatial_freqs)
    half_emb_w = torch.outer(seq[:W], w_spatial_freqs)

    if self.enable_fps_modulation and fps is not None:
        half_emb_t = torch.outer(seq[:T] / fps * self.base_fps, temporal_freqs)
    else:
        half_emb_t = torch.outer(seq[:T], temporal_freqs)

    # Broadcast and concatenate embeddings
    emb_t = half_emb_t[:, None, None, :].repeat(1, H, W, 1)
    emb_h = half_emb_h[None, :, None, :].repeat(T, 1, W, 1)
    emb_w = half_emb_w[None, None, :, :].repeat(T, H, 1, 1)

    # Concatenate [t, h, w, t, h, w] for sin/cos pairs
    freqs = torch.cat([emb_t, emb_h, emb_w] * 2, dim=-1)
    freqs = freqs.flatten(0, 2).float()  # (THW, D)

    cos = torch.cos(freqs)
    sin = torch.sin(freqs)

    return cos, sin

def forward(
    self,
    hidden_states: torch.Tensor,
    rope_emb: tuple[torch.Tensor, torch.Tensor] | None = None,
    original_seq_len: int | None = None,
) -> torch.Tensor:
    """
    Args:
        hidden_states: (B, S, D) where S = T*H*W
        rope_emb: Tuple of (cos, sin) for RoPE
        original_seq_len: Original (unpadded) full sequence length
    """
    # Get QKV
    query, _ = self.to_q(hidden_states)
    key, _ = self.to_k(hidden_states)
    value, _ = self.to_v(hidden_states)

    # Reshape for multi-head attention: (B, S, D) -> (B, S, H, D_h) -> (B, H, S, D_h)
    query = query.unflatten(-1, (self.num_heads, self.head_dim)).transpose(1, 2)
    key = key.unflatten(-1, (self.num_heads, self.head_dim)).transpose(1, 2)
    value = value.unflatten(-1, (self.num_heads, self.head_dim)).transpose(1, 2)

    # Apply QK normalization
    query = self.norm_q(query)
    key = self.norm_k(key)

    # Apply RoPE if provided (query/key are now in (B, H, S, D_h) format)
    if rope_emb is not None:
        cos, sin = rope_emb
        query = apply_rotary_emb(query, (cos, sin), use_real=True, use_real_unbind_dim=-2)
        key = apply_rotary_emb(key, (cos, sin), use_real=True, use_real_unbind_dim=-2)

    # Attention computation
    query = query.transpose(1, 2)  # (B, H, S, D_h) -> (B, S, H, D_h)
    key = key.transpose(1, 2)
    value = value.transpose(1, 2)

    attn_output, _ = self.attn(
        query, key, value, original_seq_len=original_seq_len
    )
    attn_output = attn_output.flatten(-2, -1)  # (B, S, H*D_h)

    # Output projection
    attn_output, _ = self.to_out(attn_output)
    return attn_output

def forward(self, sample: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor | None]:
    """
    Returns:
        emb: Standard embedding (B, D) - the original sinusoidal input
        adaln_lora: AdaLN-LoRA parameters (B, 3D) or None

    Note: When use_adaln_lora=True, the standard embedding is the INPUT
    (sinusoidal timestep embedding), not the processed output. The processed
    output (linear_2) is used exclusively for AdaLN-LoRA parameters.
    This matches the official GEN3C implementation.
    """
    emb, _ = self.linear_1(sample)
    emb = self.activation(emb)
    emb, _ = self.linear_2(emb)

    if self.use_adaln_lora:
        adaln_lora = emb  # (B, 3D) - full processed embedding for LoRA
        emb_standard = sample  # (B, D) - original input as the standard embedding
    else:
        emb_standard = emb
        adaln_lora = None

    return emb_standard, adaln_lora

def forward(
    self,
    hidden_states: torch.Tensor,
    timestep: torch.Tensor,
    encoder_hidden_states: torch.Tensor | list[torch.Tensor],
    attention_mask: torch.Tensor | None = None,
    fps: int | None = None,
    condition_video_input_mask: torch.Tensor | None = None,
    condition_video_pose: torch.Tensor | None = None,
    condition_video_augment_sigma: torch.Tensor | None = None,
    padding_mask: torch.Tensor | None = None,
    **kwargs,
) -> torch.Tensor:
    """
    Args:
        hidden_states: (B, C, T, H, W) latent video (16 channels)
        timestep: (B,) diffusion timesteps
        encoder_hidden_states: (B, N, D_text) text embeddings
        attention_mask: Optional attention mask
        fps: Frames per second
        condition_video_input_mask: (B, 1, T, H, W) conditioning mask for video frames
        condition_video_pose: (B, buffer_channels, T, H, W) VAE-encoded 3D cache buffers
        condition_video_augment_sigma: (B,) sigma for conditioning noise augmentation
        padding_mask: (B, 1, H, W) padding mask
    """
    if not isinstance(encoder_hidden_states, torch.Tensor):
        encoder_hidden_states = encoder_hidden_states[0]

    batch_size, num_channels, num_frames, height, width = hidden_states.shape

    # 1. Concatenate condition_video_input_mask
    if condition_video_input_mask is not None:
        hidden_states = torch.cat([hidden_states, condition_video_input_mask], dim=1)
    else:
        # Default: zeros mask (no conditioning)
        condition_video_input_mask = torch.zeros(
            batch_size, 1, num_frames, height, width, 
            device=hidden_states.device, dtype=hidden_states.dtype
        )
        hidden_states = torch.cat([hidden_states, condition_video_input_mask], dim=1)

    # 2. Concatenate condition_video_pose (3D cache buffers)
    if condition_video_pose is not None:
        hidden_states = torch.cat([hidden_states, condition_video_pose], dim=1)
    else:
        # Default: zeros for pose buffers
        pose_zeros = torch.zeros(
            batch_size, self.buffer_channels, num_frames, height, width,
            device=hidden_states.device, dtype=hidden_states.dtype
        )
        hidden_states = torch.cat([hidden_states, pose_zeros], dim=1)

    # 3. Concatenate padding mask if needed
    if self.concat_padding_mask:
        if padding_mask is None:
            padding_mask = torch.ones(
                batch_size, 1, height, width,
                device=hidden_states.device, dtype=hidden_states.dtype
            )
        padding_mask = transforms.functional.resize(
            padding_mask,
            list(hidden_states.shape[-2:]),
            interpolation=transforms.InterpolationMode.NEAREST,
        )
        hidden_states = torch.cat(
            [hidden_states, padding_mask.unsqueeze(2).repeat(1, 1, num_frames, 1, 1)],
            dim=1,
        )

    # 4. Patchify input
    p_t, p_h, p_w = self.patch_size
    hidden_states = self.patch_embed(hidden_states)  # (B, T', H', W', D)
    post_patch_num_frames, post_patch_height, post_patch_width = hidden_states.shape[1:4]

    # 5. Generate RoPE embeddings
    rope_emb = self.rope(hidden_states, fps=fps)

    # 6. Generate learnable positional embeddings
    extra_pos_emb = None
    if self.learnable_pos_embed is not None:
        extra_pos_emb = self.learnable_pos_embed(hidden_states)

    # Flatten to sequence representation for transformer blocks.
    hidden_states = hidden_states.flatten(1, 3)  # (B, S, D)
    if extra_pos_emb is not None:
        extra_pos_emb = extra_pos_emb.flatten(1, 3)

    # Sequence parallel sharding with optional padding.
    # This is a no-op when SP world size is 1.
    sp_world_size = get_sp_world_size()
    original_seq_len = hidden_states.shape[1]
    if sp_world_size > 1:
        hidden_states, original_seq_len = sequence_model_parallel_shard(
            hidden_states, dim=1
        )
        if extra_pos_emb is not None:
            extra_pos_emb, _ = sequence_model_parallel_shard(extra_pos_emb, dim=1)
        rope_cos, _ = sequence_model_parallel_shard(rope_emb[0], dim=0)
        rope_sin, _ = sequence_model_parallel_shard(rope_emb[1], dim=0)
        rope_emb = (rope_cos, rope_sin)

    # 7. Timestep embeddings
    affine_emb, adaln_lora = self.time_embed(timestep, hidden_states.dtype)

    # 8. Augment sigma embedding (GEN3C-specific)
    if self.add_augment_sigma_embedding:
        if condition_video_augment_sigma is None:
            condition_video_augment_sigma = torch.zeros_like(timestep)
        augment_emb, _ = self.augment_sigma_embed(condition_video_augment_sigma, hidden_states.dtype)
        affine_emb = affine_emb + augment_emb

    # 9. Apply affine normalization
    affine_emb = self.affine_norm(affine_emb)

    # 10. Transformer blocks
    if torch.is_grad_enabled() and self.gradient_checkpointing:
        for block in self.transformer_blocks:
            hidden_states = self._gradient_checkpointing_func(
                block,
                hidden_states,
                encoder_hidden_states,
                affine_emb,
                adaln_lora,
                rope_emb,
                extra_pos_emb,
                original_seq_len,
            )
    else:
        for block in self.transformer_blocks:
            hidden_states = block(
                hidden_states=hidden_states,
                encoder_hidden_states=encoder_hidden_states,
                affine_emb=affine_emb,
                adaln_lora=adaln_lora,
                rope_emb=rope_emb,
                extra_pos_emb=extra_pos_emb,
                original_seq_len=original_seq_len,
            )

    # 11. Final layer
    hidden_states = self.final_layer(hidden_states, affine_emb, adaln_lora)

    if sp_world_size > 1:
        hidden_states = sequence_model_parallel_all_gather_with_unpad(
            hidden_states, original_seq_len, dim=1
        )

    # 12. Unpatchify: (B, S, P) -> (B, T', H', W', P) -> (B, C, T, H, W)
    hidden_states = hidden_states.unflatten(
        1, (post_patch_num_frames, post_patch_height, post_patch_width)
    )
    hidden_states = hidden_states.unflatten(-1, (p_t, p_h, p_w, self.out_channels))
    hidden_states = hidden_states.permute(0, 7, 1, 4, 2, 5, 3, 6)
    hidden_states = hidden_states.flatten(2, 3).flatten(3, 4).flatten(4, 5)

    return hidden_states