from torch import nn

from .latent_encoder import LatentEncoder


def add_vista4d_modules(model):
    model.vista4d = True
    model.latent_encoder = LatentEncoder(
        source_init_mode="wan_patch_embed",
        point_cloud_init_mode="wan_patch_embed",
        mask_init_mode="zero_init",
        use_source_masks=True,
        use_point_cloud_masks=True,
        wan_patch_embedding=model.patch_embedding,
        rgb_in_channels=model.in_dim,
        mask_in_channels=2 * 4 * 8 * 8,
    )
    hidden_dim = model.blocks[0].self_attn.q.weight.shape[0]
    for block in model.blocks:
        block.cam_encoder = nn.Linear(6, hidden_dim)
        block.projector = nn.Linear(hidden_dim, hidden_dim)
        if not block.cam_encoder.weight.is_meta:
            block.cam_encoder.weight.data.zero_()
            block.cam_encoder.bias.data.zero_()
        if not block.projector.weight.is_meta:
            nn.init.eye_(block.projector.weight)
            nn.init.zeros_(block.projector.bias)