# Any4D DiT built on the pure Wan 2.1 T2V-1.3B base (text -> image/video).
#
# Reference: the canonical Wan2.1-T2V-1.3B text-to-video model
#   https://github.com/modelscope/DiffSynth-Studio/blob/main/examples/wanvideo/
#       model_inference/Wan2.1-T2V-1.3B.py
# Loads diffusion_pytorch_model*.safetensors from Wan-AI/Wan2.1-T2V-1.3B.
#
# Differences from the Fun-InP variant in ``a4d_network_wan.py``:
#   - patch_embedding in_dim = 16 (not 36). The T2V DiT input is JUST the noisy RGB
#     latent: there is no 4-ch input-frame mask slot and no 16-ch reference latent
#     (those are Fun-InP / I2V conditioning channels). Conditioning is text-only.
#   - has_image_input = False. The Fun-InP/Fun-Control checkpoints ship a CLIP image
#     cross-attention branch (img_emb / k_img / v_img); the pure T2V checkpoint does
#     not, so we don't build those modules and the checkpoint loads cleanly.
#   - out_dim = 16 (unchanged — still predicts the 16-ch latent).
#
# Per-view stream layout (18 channels — Fun-InP minus the reference slot):
#     [0:16]   noisy RGB latent  (the only thing the DiT consumes)
#     [16:17]  input_mask        (Any4D bookkeeping; broadcast read-through)
#     [17:18]  output_mask       (Any4D bookkeeping)
# The reshape helper drops the mask channels and feeds only [0:16] to the 16-ch
# patch_embedding. The head pads its 16-ch prediction back to 18 ch so the
# pipeline's mask/ref arithmetic is shape-consistent (the extra channels are
# conditioning-only and carry supervise_mask=0).
#
# Everything else (forward(), view embeddings, low-dim streams, FSDP sharding,
# learnable structured-noise projection ``sn_P_raw``) is inherited unchanged from
# ``Any4DWanDiT``. ``forward`` calls ``self._to_wan_36ch`` via MRO, so we keep that
# name (the actual output is 16 ch here).

from typing import Any

import torch
from custom.wan.a4d_network_wan import Any4DWanDiT


class Any4DWanT2VDiT(Any4DWanDiT):
    """Any4DWanDiT specialized for the pure Wan 2.1 T2V-1.3B base (in_dim=16).

    Inherits ``__init__`` and ``forward()`` from ``Any4DWanDiT``. Two deltas:
        * ``WAN_1_3B_INP_KWARGS`` flips ``in_dim`` 36 -> 16 and ``has_image_input``
          True -> False (read by the parent ``__init__`` when building ``WanModel``).
        * ``_to_wan_36ch`` is overridden to slice the per-view stream down to the
          16-ch noisy RGB latent. The method name is kept for MRO compatibility with
          the parent's ``forward``; the actual output is 16 ch.
    """

    # Canonical Wan 2.1 T2V-1.3B kwargs (from Wan-AI/Wan2.1-T2V-1.3B config.json:
    # dim=1536, in_dim=16, ffn_dim=8960, out_dim=16, text_dim=4096, freq_dim=256,
    # eps=1e-6, patch_size=[1,2,2], num_heads=12, num_layers=30, model_type=t2v).
    WAN_1_3B_INP_KWARGS: dict[str, Any] = dict(
        **{**Any4DWanDiT.WAN_1_3B_INP_KWARGS, 'in_dim': 16, 'has_image_input': False},
    )

    def _to_wan_36ch(self, x_v: torch.Tensor) -> torch.Tensor:
        """Slice the Any4D per-view stream down to Wan T2V's 16-ch input.

        Input layout (per-view stream after pack_streams + assemble_network_inputs):
            x_v[:, 0:16]  = noisy rgb VAE latent (noisy at all frames — Any4DWanPipeline
                            ._rewrite_masks_for_wan sets the rgb slot input_mask=0)
            x_v[:, 16:17] = input_mask  (Any4D bookkeeping; not consumed by the T2V DiT)
            x_v[:, 17:18] = output_mask (Any4D bookkeeping; not consumed)

        Output: the 16-ch noisy RGB latent only — no mask/ref slots (pure text->image).
        """
        return x_v[:, 0:16]


__all__ = ['Any4DWanT2VDiT']