h3-blackwell-runtime/src/h3_blackwell_runtime/qwen3vl_text.py

"""Direct prompt-only Qwen3-VL-32B conditioning for MiniMax H3.

The MiniMax checkpoint contains the first 50 Qwen3-VL decoder layers.  H3
uses their unnormalized final output, not the language-model head or final
RMSNorm.  This module intentionally has no ComfyUI imports and does not load
the Qwen vision encoder; image/video prompt construction remains a separate
feature.
"""

from __future__ import annotations

from dataclasses import dataclass
from pathlib import Path

import torch
import torch.nn.functional as F
from torch import nn

from .checkpoint import H3Checkpoint


@dataclass(frozen=True)
class Qwen3VL32BTextConfig:
    vocab_size: int = 151936
    hidden_size: int = 5120
    intermediate_size: int = 25600
    num_layers: int = 50
    num_attention_heads: int = 64
    num_key_value_heads: int = 8
    head_dim: int = 128
    rms_norm_eps: float = 1e-6
    rope_theta: float = 5_000_000.0


class _RMSNorm(nn.Module):
    def __init__(self, weight: torch.Tensor, eps: float):
        super().__init__()
        self.eps = eps
        self.register_buffer("weight", weight, persistent=False)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        variance = x.float().square().mean(dim=-1, keepdim=True)
        return (x * torch.rsqrt(variance + self.eps)).to(x.dtype) * self.weight.to(x.dtype)


def _rope(query: torch.Tensor, key: torch.Tensor, theta: float) -> tuple[torch.Tensor, torch.Tensor]:
    """Apply Qwen's split-half rotary embedding to [B, H, S, D] Q/K tensors."""
    positions = torch.arange(query.shape[-2], device=query.device, dtype=torch.float32)
    dimensions = torch.arange(0, query.shape[-1], 2, device=query.device, dtype=torch.float32)
    frequencies = positions[:, None] / theta ** (dimensions / query.shape[-1])
    angles = torch.cat((frequencies, frequencies), dim=-1)
    cos = angles.cos()[None, None].to(query.dtype)
    sin = angles.sin()[None, None].to(query.dtype)

    def rotate_half(value: torch.Tensor) -> torch.Tensor:
        first, second = value.chunk(2, dim=-1)
        return torch.cat((-second, first), dim=-1)

    return query * cos + rotate_half(query) * sin, key * cos + rotate_half(key) * sin


class _Qwen3VLBlock(nn.Module):
    def __init__(self, checkpoint: H3Checkpoint, prefix: str, config: Qwen3VL32BTextConfig, dtype: torch.dtype):
        super().__init__()
        self.config = config
        self.input_layernorm = _RMSNorm(checkpoint.tensor(f"{prefix}.input_layernorm.weight", dtype=dtype), config.rms_norm_eps)
        self.post_attention_layernorm = _RMSNorm(checkpoint.tensor(f"{prefix}.post_attention_layernorm.weight", dtype=dtype), config.rms_norm_eps)
        self.q_norm = _RMSNorm(checkpoint.tensor(f"{prefix}.self_attn.q_norm.weight", dtype=dtype), config.rms_norm_eps)
        self.k_norm = _RMSNorm(checkpoint.tensor(f"{prefix}.self_attn.k_norm.weight", dtype=dtype), config.rms_norm_eps)
        self.q_proj = checkpoint.nvfp4_linear(f"{prefix}.self_attn.q_proj", output_dtype=dtype)
        self.k_proj = checkpoint.nvfp4_linear(f"{prefix}.self_attn.k_proj", output_dtype=dtype)
        self.v_proj = checkpoint.nvfp4_linear(f"{prefix}.self_attn.v_proj", output_dtype=dtype)
        self.o_proj = checkpoint.nvfp4_linear(f"{prefix}.self_attn.o_proj", output_dtype=dtype)
        self.gate_proj = checkpoint.nvfp4_linear(f"{prefix}.mlp.gate_proj", output_dtype=dtype)
        self.up_proj = checkpoint.nvfp4_linear(f"{prefix}.mlp.up_proj", output_dtype=dtype)
        self.down_proj = checkpoint.nvfp4_linear(f"{prefix}.mlp.down_proj", output_dtype=dtype)

    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
        residual = hidden_states
        x = self.input_layernorm(hidden_states)
        batch, sequence, _ = x.shape
        query = self.q_proj(x).view(batch, sequence, self.config.num_attention_heads, self.config.head_dim).transpose(1, 2)
        key = self.k_proj(x).view(batch, sequence, self.config.num_key_value_heads, self.config.head_dim).transpose(1, 2)
        value = self.v_proj(x).view(batch, sequence, self.config.num_key_value_heads, self.config.head_dim).transpose(1, 2)
        query = self.q_norm(query)
        key = self.k_norm(key)
        query, key = _rope(query, key, self.config.rope_theta)
        attention = F.scaled_dot_product_attention(query, key, value, is_causal=True, enable_gqa=True)
        hidden_states = residual + self.o_proj(attention.transpose(1, 2).reshape(batch, sequence, -1))
        residual = hidden_states
        x = self.post_attention_layernorm(hidden_states)
        return residual + self.down_proj(F.silu(self.gate_proj(x)) * self.up_proj(x))


class Qwen3VL32BTextEncoder(nn.Module):
    """Mounted-checkpoint Qwen3-VL prompt conditioner returning layer-50 states."""
    config = Qwen3VL32BTextConfig()

    def __init__(self, checkpoint_path: str | Path, *, device: str | torch.device = "cuda", dtype: torch.dtype = torch.bfloat16):
        super().__init__()
        self.checkpoint_path = str(checkpoint_path)
        self.device_name = str(device)
        self.dtype = dtype
        checkpoint = H3Checkpoint(checkpoint_path, device=device)
        self._validate_checkpoint(checkpoint)
        self.register_buffer("embed_tokens", checkpoint.tensor("model.embed_tokens.weight", dtype=dtype), persistent=False)
        self.layers = nn.ModuleList(
            _Qwen3VLBlock(checkpoint, f"model.layers.{index}", self.config, dtype)
            for index in range(self.config.num_layers)
        )

    @classmethod
    def _validate_checkpoint(cls, checkpoint: H3Checkpoint) -> None:
        from safetensors import safe_open

        required = {"model.embed_tokens.weight"}
        linear_names = (
            "self_attn.q_proj", "self_attn.k_proj", "self_attn.v_proj", "self_attn.o_proj",
            "mlp.gate_proj", "mlp.up_proj", "mlp.down_proj",
        )
        for index in range(cls.config.num_layers):
            prefix = f"model.layers.{index}"
            required.update({
                f"{prefix}.input_layernorm.weight",
                f"{prefix}.post_attention_layernorm.weight",
                f"{prefix}.self_attn.q_norm.weight",
                f"{prefix}.self_attn.k_norm.weight",
            })
            for linear in linear_names:
                required.update(
                    f"{prefix}.{linear}.{suffix}"
                    for suffix in ("comfy_quant", "weight", "weight_scale", "weight_scale_2")
                )
        with safe_open(checkpoint.path, framework="pt", device="cpu") as file:
            names = set(file.keys())
        missing = sorted(required - names)
        if missing:
            raise ValueError(
                "Not a supported MiniMax Qwen3-VL-32B NVFP4 checkpoint; missing "
                + ", ".join(missing)
            )

    @torch.inference_mode()
    def forward(self, input_ids: torch.Tensor) -> torch.Tensor:
        """Return unnormalized `[batch, tokens, 5120]` output after decoder layer 50."""
        if input_ids.ndim != 2:
            raise ValueError(f"input_ids must have shape [batch, tokens], got {tuple(input_ids.shape)}")
        if input_ids.numel() == 0:
            raise ValueError("input_ids must contain at least one token")
        hidden_states = F.embedding(input_ids.to(self.embed_tokens.device), self.embed_tokens).to(self.dtype)
        for layer in self.layers:
            hidden_states = layer(hidden_states)
        return hidden_states


class Qwen3VLPromptConditioner:
    """Tokenize raw H3 prompt text and produce Qwen layer-50 conditioning."""
    def __init__(self, checkpoint_path: str | Path, tokenizer_dir: str | Path, *, device: str | torch.device = "cuda", dtype: torch.dtype = torch.bfloat16):
        from .conditioning import H3PromptTokenizer

        self.tokenizer = H3PromptTokenizer(tokenizer_dir)
        self.encoder = Qwen3VL32BTextEncoder(checkpoint_path, device=device, dtype=dtype)
        self.device = device

    def __call__(self, prompt: str) -> torch.Tensor:
        return self.encoder(self.tokenizer(prompt, device=self.device))