videox_fun/models/creator_audio.py

"""Creator audio diffusion transformer used by the release inference path."""

import glob
import json
import logging
import os
from typing import Optional

import torch
import torch.nn as nn
from diffusers.configuration_utils import ConfigMixin, register_to_config
from diffusers.loaders.single_file_model import FromOriginalModelMixin
from diffusers.models.modeling_utils import ModelMixin
from einops import rearrange

from .creator.creator_video_dit import DiTBlock
from .creator.creator_audio_dit import (
    Head,
    MLP,
    sinusoidal_embedding_1d,
    precompute_freqs_cis_1d,
)


class CreatorAudioModel(ModelMixin, ConfigMixin, FromOriginalModelMixin):
    """Creator Audio DiT with wan_audio2-compatible forward interface."""

    @register_to_config
    def __init__(
        self,
        patch_size=(1,),
        text_len=512,
        in_dim=128,
        dim=2048,
        ffn_dim=8192,
        freq_dim=256,
        text_dim=4096,
        out_dim=128,
        num_heads=16,
        num_layers=32,
        eps=1e-6,
        has_image_input=False,
        has_image_pos_emb=False,
        has_ref_conv=False,
        vae_type="dac",
        **kwargs
    ):
        super().__init__()

        self.patch_size = tuple(patch_size) if isinstance(patch_size, (list, tuple)) else (patch_size,)
        self.text_len = text_len
        self.in_dim = in_dim
        self.dim = dim
        self.ffn_dim = ffn_dim
        self.freq_dim = freq_dim
        self.text_dim = text_dim
        self.out_dim = out_dim
        self.num_heads = num_heads
        self.num_layers = num_layers
        self.eps = eps
        self.has_image_input = has_image_input
        self.vae_type = vae_type

        # Patch embedding (1D conv)
        self.patch_embedding = nn.Conv1d(
            in_dim, dim, kernel_size=self.patch_size[0], stride=self.patch_size[0],
        )
        self.text_embedding = nn.Sequential(
            nn.Linear(text_dim, dim),
            nn.GELU(approximate="tanh"),
            nn.Linear(dim, dim),
        )
        self.time_embedding = nn.Sequential(
            nn.Linear(freq_dim, dim),
            nn.SiLU(),
            nn.Linear(dim, dim),
        )
        self.time_projection = nn.Sequential(
            nn.SiLU(), nn.Linear(dim, dim * 6),
        )

        # DiTBlock stack
        self.blocks = nn.ModuleList([
            DiTBlock(has_image_input, dim, num_heads, ffn_dim, eps)
            for _ in range(num_layers)
        ])

        self.head = Head(dim, out_dim, self.patch_size, eps)

        # RoPE precomputation
        head_dim = dim // num_heads
        self.freqs = precompute_freqs_cis_1d(head_dim)

        # Optional image / reference embeddings
        if has_image_input:
            self.img_emb = MLP(1280, dim, has_pos_emb=has_image_pos_emb)
        if has_ref_conv:
            self.ref_conv = nn.Conv2d(16, dim, kernel_size=(2, 2), stride=(2, 2))
        self.has_ref_conv = has_ref_conv

    # -----------------------------------------------------------------
    # Unpatchify
    # -----------------------------------------------------------------

    def unpatchify(self, x, grid_sizes, output_shapes):
        """Unpatchify tokens back to audio latent shape.

        Args:
            x: [B, seq_len, out_dim * patch_size]
            grid_sizes: [B, 1] token lengths
            output_shapes: list of original audio shapes [(C, T), ...]

        Returns:
            list of restored tensors matching output_shapes.
        """
        output = []
        for i, (grid_size, original_shape) in enumerate(zip(grid_sizes, output_shapes)):
            f = int(grid_size[0].item())
            restored = rearrange(
                x[i, :f].unsqueeze(0),
                "b f (p c) -> b c (f p)",
                f=f, p=self.patch_size[0],
            )
            # Trim to original time length
            orig_T = original_shape[-1]
            restored = restored[:, :, :orig_T].squeeze(0)
            output.append(restored)
        return output

    # -----------------------------------------------------------------
    # RoPE frequencies
    # -----------------------------------------------------------------
    def _build_freqs(self, seq_len: int, device: torch.device) -> torch.Tensor:
        """Build RoPE complex frequencies [seq_len, 1, rope_dim]."""
        freqs = torch.cat([
            self.freqs[0][:seq_len].view(seq_len, -1),
            self.freqs[1][:seq_len].view(seq_len, -1),
            self.freqs[2][:seq_len].view(seq_len, -1),
        ], dim=-1).reshape(seq_len, 1, -1).to(device)
        return freqs

    # -----------------------------------------------------------------
    # Weight initialisation
    # -----------------------------------------------------------------
    def init_weights(self):
        for module in self.modules():
            if isinstance(module, nn.Linear):
                nn.init.xavier_uniform_(module.weight)
                if module.bias is not None:
                    nn.init.zeros_(module.bias)

        nn.init.xavier_uniform_(self.patch_embedding.weight.flatten(1))
        if self.patch_embedding.bias is not None:
            nn.init.zeros_(self.patch_embedding.bias)
        for module in self.text_embedding.modules():
            if isinstance(module, nn.Linear):
                nn.init.normal_(module.weight, std=0.02)
        for module in self.time_embedding.modules():
            if isinstance(module, nn.Linear):
                nn.init.normal_(module.weight, std=0.02)

        nn.init.zeros_(self.head.head.weight)
        if self.head.head.bias is not None:
            nn.init.zeros_(self.head.head.bias)

    # -----------------------------------------------------------------
    # forward (matches wan_audio2.WanAudioModel interface)
    # -----------------------------------------------------------------
    def forward(
        self,
        x,
        t,
        context,
        seq_len,
        clip_fea=None,
        y=None,
        dtype=torch.bfloat16,
        **kwargs,
    ):
        """
        Args:
            x: list of [C, T] audio latent samples, or [B, C, T] tensor.
            t: [B] or [B, T] timesteps.
            context: list of [S, text_dim] or [B, S, text_dim] tensor.
            seq_len: max sequence length after patching.
            clip_fea: optional CLIP features for image conditioning.
            y: optional conditioning latent (e.g. for i2a), same format as x.
        """
        # --- Normalise inputs to lists ---
        if isinstance(x, torch.Tensor):
            x = [sample for sample in x]
        if isinstance(context, torch.Tensor):
            context = [sample for sample in context]
        if y is not None and isinstance(y, torch.Tensor):
            y = [sample for sample in y]

        device = self.patch_embedding.weight.device
        dtype = x[0].dtype if len(x) > 0 else self.patch_embedding.weight.dtype
        batch_size = len(x)

        # Concatenate conditioning y (e.g. for i2v/i2a)
        if y is not None:
            x = [torch.cat([u, v], dim=0) for u, v in zip(x, y)]

        # Remember original shapes for unpatchify
        original_audio_shapes = [tuple(sample.shape) for sample in x]

        # --- Patchify per sample, then pad tokens ---
        patchified_samples = []
        grid_sizes_list = []
        for sample in x:
            tokens_i = self.patch_embedding(sample.unsqueeze(0).to(device))  # [1, dim, T']
            tokens_i = rearrange(tokens_i, "1 c f -> f c").contiguous()
            patchified_samples.append(tokens_i)
            grid_sizes_list.append(tokens_i.shape[0])

        grid_sizes = torch.tensor(
            [[g] for g in grid_sizes_list], dtype=torch.long, device=device,
        )
        seq_lens = grid_sizes[:, 0]
        assert seq_lens.max() <= seq_len, (
            f"Max token length {seq_lens.max().item()} exceeds seq_len {seq_len}"
        )

        # Pad each sample's tokens to seq_len and stack
        x = torch.stack([
            torch.cat([u, u.new_zeros(seq_len - u.size(0), u.size(1))], dim=0)
            for u in patchified_samples
        ])

        # --- Time embedding ---
        with torch.amp.autocast("cuda", dtype=torch.float32):
            if t.dim() != 1:
                # Per-token timesteps [B, T]
                if t.size(1) < seq_len:
                    pad_size = seq_len - t.size(1)
                    last_elements = t[:, -1].unsqueeze(1)
                    t = torch.cat([t, last_elements.repeat(1, pad_size)], dim=1)
                bt = t.size(0)
                e = self.time_embedding(
                    sinusoidal_embedding_1d(self.freq_dim, t.flatten())
                    .unflatten(0, (bt, seq_len)).float()
                )
                e0 = self.time_projection(e).unflatten(2, (6, self.dim))
            else:
                e = self.time_embedding(
                    sinusoidal_embedding_1d(self.freq_dim, t).float()
                )
                e0 = self.time_projection(e).unflatten(1, (6, self.dim))

        # --- Context embedding ---
        context = self.text_embedding(
            torch.stack([
                torch.cat([u, u.new_zeros(self.text_len - u.size(0), u.size(1))])
                for u in context
            ])
        )

        # Image embedding
        if self.has_image_input and clip_fea is not None:
            clip_embedding = self.img_emb(clip_fea)
            context = torch.cat([clip_embedding, context], dim=1)

        # --- RoPE frequencies ---
        freqs = self._build_freqs(seq_len, device)

        for block in self.blocks:
            x = block(x, context, e0, freqs, seq_lens=seq_lens)

        x = self.head(x, e)
        x = self.unpatchify(x, grid_sizes, original_audio_shapes)
        return x

    # -----------------------------------------------------------------
    # from_pretrained
    # -----------------------------------------------------------------
    @classmethod
    def from_pretrained(
        cls,
        pretrained_model_path,
        subfolder=None,
        transformer_additional_kwargs=None,
        low_cpu_mem_usage=False,
        in_dim=None,
        out_dim=None,
        patch_size=None,
        torch_dtype=torch.bfloat16,
    ):
        transformer_additional_kwargs = dict(transformer_additional_kwargs or {})
        if subfolder is not None:
            pretrained_model_path = os.path.join(pretrained_model_path, subfolder)

        config_file = os.path.join(pretrained_model_path, "config.json")
        if not os.path.isfile(config_file):
            raise RuntimeError(f"{config_file} does not exist")

        with open(config_file, "r") as fp:
            config = json.load(fp)

        from diffusers.utils import WEIGHTS_NAME

        model_file = os.path.join(pretrained_model_path, WEIGHTS_NAME)
        model_file_safetensors = model_file.replace(".bin", ".safetensors")

        # Explicit architecture overrides
        if in_dim is not None:
            transformer_additional_kwargs["in_dim"] = in_dim
        if out_dim is not None:
            transformer_additional_kwargs["out_dim"] = out_dim
        if patch_size is not None:
            transformer_additional_kwargs["patch_size"] = patch_size

        if "dict_mapping" in transformer_additional_kwargs:
            for key, value in transformer_additional_kwargs["dict_mapping"].items():
                if value not in transformer_additional_kwargs:
                    transformer_additional_kwargs[value] = config[key]

        # Merge overrides into config
        model_config = dict(config)
        model_config.update(transformer_additional_kwargs)

        model = cls.from_config(model_config, **transformer_additional_kwargs)

        # --- Load state dict ---
        if os.path.exists(model_file):
            state_dict = torch.load(model_file, map_location="cpu")
        elif os.path.exists(model_file_safetensors):
            from safetensors.torch import load_file
            state_dict = load_file(model_file_safetensors)
        else:
            from safetensors.torch import load_file
            state_dict = {}
            for shard in glob.glob(os.path.join(pretrained_model_path, "*.safetensors")):
                state_dict.update(load_file(shard))

        # Filter by shape match
        model_sd = model.state_dict()
        filtered_state_dict = {}
        for key, value in state_dict.items():
            if key in model_sd and model_sd[key].shape == value.shape:
                filtered_state_dict[key] = value
            else:
                logging.info("Skipping key %s due to size mismatch or absence.", key)

        missing, unexpected = model.load_state_dict(filtered_state_dict, strict=False)
        logging.info(
            "CreatorAudioModel missing keys: %d, unexpected keys: %d",
            len(missing), len(unexpected),
        )
        return model.to(torch_dtype)

# Convenience aliases
CreatorAudioTransformerModel = CreatorAudioModel