"""Packed-input H3 transformer core, independent of ComfyUI node execution."""

import torch
from torch import nn

from .backbone import H3DenoiserBackbone
from .checkpoint import H3Checkpoint
from .final import H3FinalLayer


class H3PackedDenoiser(nn.Module):
    """Run the H3 transformer once its Ref2VA payload has been packed into hidden rows."""

    def __init__(self, backbone: H3DenoiserBackbone, final_layer: H3FinalLayer):
        super().__init__()
        self.backbone = backbone
        self.final_layer = final_layer

    @classmethod
    def from_checkpoint(cls, checkpoint: H3Checkpoint, *, output_dtype=torch.bfloat16, attention_backend: str = "sage2"):
        return cls(
            H3DenoiserBackbone.from_checkpoint(checkpoint, output_dtype=output_dtype, attention_backend=attention_backend),
            H3FinalLayer.from_checkpoint(checkpoint, output_dtype=output_dtype),
        )

    def forward(
        self,
        hidden: torch.Tensor,
        timesteps: torch.Tensor,
        position_ids: torch.Tensor,
        segments: list[tuple[int, int, int]],
        video_segment: tuple[int, int, int],
        audio_segment: tuple[int, int, int],
    ) -> tuple[torch.Tensor, torch.Tensor]:
        hidden = self.backbone(hidden, timesteps, position_ids, segments)
        return self.final_layer(hidden, timesteps, video_segment, audio_segment)