hunyuanvideo15 ¶

def __init__(self,
             hidden_size,
             patch_size,
             out_channels,
             dtype=None,
             prefix: str = "") -> None:
    super().__init__()

    # Normalization
    self.norm_final = nn.LayerNorm(hidden_size,
                                   eps=1e-6,
                                   elementwise_affine=False,
                                   dtype=dtype)

    output_dim = patch_size[0] * patch_size[1] * patch_size[2] * out_channels

    self.linear = ReplicatedLinear(hidden_size,
                                   output_dim,
                                   bias=True,
                                   params_dtype=dtype,
                                   prefix=f"{prefix}.linear")

    # Modulation
    self.adaLN_modulation = ModulateProjection(
        hidden_size,
        factor=2,
        act_layer="silu",
        dtype=dtype,
        prefix=f"{prefix}.adaLN_modulation")

def __init__(
    self,
    dim: int,
    elementwise_affine=True,
    eps: float = 1e-6,
    device=None,
    dtype=None,
):
    """
    Initialize the RMSNorm normalization layer.

    Args:
        dim (int): The dimension of the input tensor.
        eps (float, optional): A small value added to the denominator for numerical stability. Default is 1e-6.

    Attributes:
        eps (float): A small value added to the denominator for numerical stability.
        weight (nn.Parameter): Learnable scaling parameter.

    """
    factory_kwargs = {"device": device, "dtype": dtype}
    super().__init__()
    self.eps = eps
    if elementwise_affine:
        self.weight = nn.Parameter(torch.ones(dim, **factory_kwargs))

def __init__(self, embedding_dim: int, use_meanflow: bool = False):
    super().__init__()

    self.timestep_embedder = TimestepEmbedder(hidden_size=embedding_dim)

    self.use_meanflow = use_meanflow
    self.time_proj_r = None
    self.timestep_embedder_r = None
    if use_meanflow:
        self.timestep_embedder_r = TimestepEmbedder(hidden_size=embedding_dim)

def __init__(
    self,
    config: HunyuanVideo15Config,
    hf_config: dict[str, Any],
) -> None:
    super().__init__(config=config, hf_config=hf_config)

    self.hidden_size = config.hidden_size
    self.num_attention_heads = config.num_attention_heads
    self.num_channels_latents = config.num_channels_latents
    self.out_channels = config.out_channels or config.in_channels
    self.patch_size = (config.patch_size_t, config.patch_size, config.patch_size)

    # 1. Latent and condition embedders
    self.img_in = PatchEmbed(self.patch_size, 
                                config.in_channels, 
                                self.hidden_size,
                                prefix=f"{config.prefix}.img_in")
    self.image_embedder = HunyuanVideo15ImageProjection(config.image_embed_dim, self.hidden_size)

    self.txt_in = SingleTokenRefiner(config.text_embed_dim,
                                               self.hidden_size,
                                               config.num_attention_heads,
                                               depth=config.num_refiner_layers,
                                               dtype=None,
                                               prefix=f"{config.prefix}.txt_in")

    self.txt_in_2 = HunyuanVideo15ByT5TextProjection(config.text_embed_2_dim, 2048, self.hidden_size)

    self.time_in = HunyuanVideo15TimeEmbedding(self.hidden_size, use_meanflow=config.use_meanflow)

    self.cond_type_embed = nn.Embedding(3, self.hidden_size)

    # 3. Dual stream transformer blocks

    self.double_blocks = nn.ModuleList(
        [
            MMDoubleStreamBlock(
                hidden_size=self.hidden_size,
                num_attention_heads=config.num_attention_heads,
                mlp_ratio=config.mlp_ratio,
                dtype=None,
                supported_attention_backends=self._supported_attention_backends,
                prefix=f"{config.prefix}.double_blocks.{i}"
            )
            for i in range(config.num_layers)
        ]
    )

    # 5. Output projection
    self.final_layer = FinalLayer(self.hidden_size,
                            self.patch_size,
                            self.out_channels,
                            prefix=f"{config.prefix}.final_layer")

    self.gradient_checkpointing = False

    self.__post_init__()

def __init__(
    self,
    hidden_size,
    num_attention_heads,
    mlp_ratio=4.0,
    qkv_bias=True,
    dtype=None,
    prefix: str = "",
) -> None:
    super().__init__()
    self.num_attention_heads = num_attention_heads
    mlp_hidden_dim = int(hidden_size * mlp_ratio)

    # Normalization and attention
    self.norm1 = nn.LayerNorm(hidden_size,
                              eps=1e-6,
                              elementwise_affine=True,
                              dtype=dtype)

    self.self_attn_qkv = ReplicatedLinear(hidden_size,
                                          hidden_size * 3,
                                          bias=qkv_bias,
                                          params_dtype=dtype,
                                          prefix=f"{prefix}.self_attn_qkv")

    self.self_attn_proj = ReplicatedLinear(
        hidden_size,
        hidden_size,
        bias=qkv_bias,
        params_dtype=dtype,
        prefix=f"{prefix}.self_attn_proj")

    # MLP
    self.norm2 = nn.LayerNorm(hidden_size,
                              eps=1e-6,
                              elementwise_affine=True,
                              dtype=dtype)
    self.mlp = MLP(hidden_size,
                   mlp_hidden_dim,
                   bias=True,
                   act_type="silu",
                   dtype=dtype,
                   prefix=f"{prefix}.mlp")

    # Modulation
    self.adaLN_modulation = ModulateProjection(
        hidden_size,
        factor=2,
        act_layer="silu",
        dtype=dtype,
        prefix=f"{prefix}.adaLN_modulation")

    # Scaled dot product attention
    self.attn = LocalAttention(
        num_heads=num_attention_heads,
        head_size=hidden_size // num_attention_heads,
        # TODO: remove hardcode
        supported_attention_backends=(AttentionBackendEnum.FLASH_ATTN,
                                      AttentionBackendEnum.TORCH_SDPA),
    )

def __init__(
    self,
    hidden_size: int,
    num_attention_heads: int,
    mlp_ratio: float,
    dtype: torch.dtype | None = None,
    supported_attention_backends: tuple[AttentionBackendEnum, ...]
    | None = None,
    prefix: str = "",
):
    super().__init__()

    self.deterministic = False
    self.num_attention_heads = num_attention_heads
    head_dim = hidden_size // num_attention_heads
    mlp_hidden_dim = int(hidden_size * mlp_ratio)

    # Image modulation components
    self.img_mod = ModulateProjection(
        hidden_size,
        factor=6,
        act_layer="silu",
        dtype=dtype,
        prefix=f"{prefix}.img_mod",
    )

    # Fused operations for image stream
    self.img_attn_norm = LayerNormScaleShift(hidden_size,
                                             norm_type="layer",
                                             elementwise_affine=False,
                                             dtype=dtype)
    self.img_attn_residual_mlp_norm = ScaleResidualLayerNormScaleShift(
        hidden_size,
        norm_type="layer",
        elementwise_affine=False,
        dtype=dtype)
    self.img_mlp_residual = ScaleResidual()

    # Image attention components
    self.img_attn_qkv = ReplicatedLinear(hidden_size,
                                         hidden_size * 3,
                                         bias=True,
                                         params_dtype=dtype,
                                         prefix=f"{prefix}.img_attn_qkv")

    self.img_attn_q_norm = HunyuanRMSNorm(head_dim, eps=1e-6, dtype=dtype)
    self.img_attn_k_norm = HunyuanRMSNorm(head_dim, eps=1e-6, dtype=dtype)

    self.img_attn_proj = ReplicatedLinear(hidden_size,
                                          hidden_size,
                                          bias=True,
                                          params_dtype=dtype,
                                          prefix=f"{prefix}.img_attn_proj")

    self.img_mlp = MLP(hidden_size,
                       mlp_hidden_dim,
                       bias=True,
                       dtype=dtype,
                       prefix=f"{prefix}.img_mlp")

    # Text modulation components
    self.txt_mod = ModulateProjection(
        hidden_size,
        factor=6,
        act_layer="silu",
        dtype=dtype,
        prefix=f"{prefix}.txt_mod",
    )

    # Fused operations for text stream
    self.txt_attn_norm = LayerNormScaleShift(hidden_size,
                                             norm_type="layer",
                                             elementwise_affine=False,
                                             dtype=dtype)
    self.txt_attn_residual_mlp_norm = ScaleResidualLayerNormScaleShift(
        hidden_size,
        norm_type="layer",
        elementwise_affine=False,
        dtype=dtype)
    self.txt_mlp_residual = ScaleResidual()

    # Text attention components
    self.txt_attn_qkv = ReplicatedLinear(hidden_size,
                                         hidden_size * 3,
                                         bias=True,
                                         params_dtype=dtype)

    # QK norm layers for text
    self.txt_attn_q_norm = HunyuanRMSNorm(head_dim, eps=1e-6, dtype=dtype)
    self.txt_attn_k_norm = HunyuanRMSNorm(head_dim, eps=1e-6, dtype=dtype)

    self.txt_attn_proj = ReplicatedLinear(hidden_size,
                                          hidden_size,
                                          bias=True,
                                          params_dtype=dtype)

    self.txt_mlp = MLP(hidden_size, mlp_hidden_dim, bias=True, dtype=dtype)

    # Distributed attention
    self.attn = DistributedAttention(
        num_heads=num_attention_heads,
        head_size=head_dim,
        causal=False,
        supported_attention_backends=supported_attention_backends,
        prefix=f"{prefix}.attn")

def __init__(
    self,
    in_channels,
    hidden_size,
    num_attention_heads,
    depth=2,
    qkv_bias=True,
    dtype=None,
    prefix: str = "",
) -> None:
    super().__init__()

    # Input projection
    # self.input_embedder = ReplicatedLinear(
    #     in_channels,
    #     hidden_size,
    #     bias=True,
    #     params_dtype=dtype,
    #     prefix=f"{prefix}.input_embedder")
    self.input_embedder = nn.Linear(in_channels, hidden_size, bias=True)

    # Timestep embedding
    self.t_embedder = TimestepEmbedder(hidden_size,
                                       act_layer="silu",
                                       dtype=dtype,
                                       prefix=f"{prefix}.t_embedder")

    # Context embedding
    self.c_embedder = MLP(in_channels,
                          hidden_size,
                          hidden_size,
                          act_type="silu",
                          dtype=dtype,
                          prefix=f"{prefix}.c_embedder")

    # Refiner blocks
    self.refiner_blocks = nn.ModuleList([
        IndividualTokenRefinerBlock(
            hidden_size,
            num_attention_heads,
            qkv_bias=qkv_bias,
            dtype=dtype,
            prefix=f"{prefix}.refiner_blocks.{i}",
        ) for i in range(depth)
    ])