hunyuangamecraft ¶

def __init__(
    self,
    in_channels: int = 6,
    downscale_coef: int = 8,
    out_channels: int = 16,
    patch_size: List[int] = [1, 2, 2],
    hidden_size: int = 3072,
    dtype: Optional[torch.dtype] = None,
    prefix: str = "",
):
    super().__init__()
    _ = prefix  # Unused

    start_channels = in_channels * (downscale_coef ** 2)
    input_channels = [start_channels, start_channels // 2, start_channels // 4]
    self.input_channels = input_channels

    self.unshuffle = nn.PixelUnshuffle(downscale_coef)

    self.encode_first = nn.Sequential(
        nn.Conv2d(input_channels[0], input_channels[1], kernel_size=1, stride=1, padding=0),
        nn.GroupNorm(2, input_channels[1]),
        nn.ReLU(),
    )
    self._initialize_weights(self.encode_first)

    self.encode_second = nn.Sequential(
        nn.Conv2d(input_channels[1], input_channels[2], kernel_size=1, stride=1, padding=0),
        nn.GroupNorm(2, input_channels[2]),
        nn.ReLU(),
    )
    self._initialize_weights(self.encode_second)

    self.final_proj = nn.Conv2d(input_channels[2], out_channels, kernel_size=1)
    self._zeros_init_linear(self.final_proj)

    self.scale = nn.Parameter(torch.ones(1))

    self.camera_in = PatchEmbed(
        patch_size=patch_size,
        in_chans=out_channels,
        embed_dim=hidden_size,
    )

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

    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")

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

def __init__(self, config: HunyuanGameCraftConfig, hf_config: dict[str, Any]):
    super().__init__(config=config, hf_config=hf_config)

    arch = config.arch_config

    if isinstance(arch.patch_size, (list, tuple)):
        self.patch_size = list(arch.patch_size)
    else:
        self.patch_size = [arch.patch_size_t, arch.patch_size, arch.patch_size]

    self.in_channels = arch.in_channels
    self.out_channels = arch.out_channels
    self.unpatchify_channels = self.out_channels
    self.num_channels_latents = self.out_channels  # Alias for latent_preparation stage
    self.hidden_size = arch.hidden_size
    self.num_heads = arch.num_attention_heads
    self.num_attention_heads = arch.num_attention_heads  # Alias for compatibility
    self.guidance_embeds = arch.guidance_embeds
    self.rope_dim_list = list(arch.rope_axes_dim)
    self.rope_theta = arch.rope_theta
    self.text_states_dim = arch.text_embed_dim
    self.text_states_dim_2 = arch.pooled_projection_dim
    self.dtype = arch.dtype

    pe_dim = self.hidden_size // self.num_heads
    if sum(self.rope_dim_list) != pe_dim:
        raise ValueError(f"rope_axes_dim sum {sum(self.rope_dim_list)} != {pe_dim}")

    factory_kwargs = {'dtype': self.dtype}

    self.img_in = PatchEmbed(
        patch_size=self.patch_size,
        in_chans=self.in_channels,
        embed_dim=self.hidden_size,
        **factory_kwargs,
    )

    self.txt_in = SingleTokenRefiner(
        self.text_states_dim,
        self.hidden_size,
        self.num_heads,
        depth=arch.num_refiner_layers,
        **factory_kwargs,
    )

    self.time_in = TimestepEmbedder(self.hidden_size, **factory_kwargs)

    self.vector_in = MLP(
        self.text_states_dim_2,
        self.hidden_size,
        self.hidden_size,
        act_type="silu",
        **factory_kwargs,
    )

    self.guidance_in = (
        TimestepEmbedder(self.hidden_size, **factory_kwargs)
        if self.guidance_embeds else None
    )

    self.double_blocks = nn.ModuleList([
        MMDoubleStreamBlock(
            hidden_size=self.hidden_size,
            num_attention_heads=self.num_heads,
            mlp_ratio=arch.mlp_ratio,
            supported_attention_backends=self._supported_attention_backends,
            **factory_kwargs,
        )
        for _ in range(arch.num_layers)
    ])

    self.single_blocks = nn.ModuleList([
        MMSingleStreamBlock(
            hidden_size=self.hidden_size,
            num_attention_heads=self.num_heads,
            mlp_ratio=arch.mlp_ratio,
            supported_attention_backends=self._supported_attention_backends,
            **factory_kwargs,
        )
        for _ in range(arch.num_single_layers)
    ])

    self.final_layer = GameCraftFinalLayer(
        self.hidden_size,
        self.patch_size,
        self.out_channels,
        **factory_kwargs,
    )

    self.camera_net = CameraNet(
        in_channels=arch.camera_in_channels,
        out_channels=16,
        downscale_coef=arch.camera_downscale_coef,
        patch_size=self.patch_size,
        hidden_size=self.hidden_size,
    )