"""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))