custom_nodes/h3_small_te/nodes.py

"""MiniMax H3 small text encoder: Qwen3-VL-4B + trained adapter -> 5120-dim conditioning.

Matches the training path in optimization/h3-shrink/scripts/train_te_adapter.py:
- H3 tokenizer (raw text, no chat template)
- Student final text-norm replaced with Identity (unnormalized hidden states)
- Adapter Linear(2560->4096)->GELU->Linear(4096->5120) in fp32
- minimax_token_tags = all-ones for pure text

Two nodes:
- H3SmallTELoader  -> CLIP  (plugs into MiniMaxH3ImageToVideo for T2V)
- H3SmallTextEncoder -> CONDITIONING  (direct encode of a prompt string)
"""

from __future__ import annotations

import importlib.util
import os
import sys
import threading
from typing import Any, Optional

import torch
import torch.nn as nn

import folder_paths

DEFAULT_STUDENT = "/home/bbear/Documents/OlympusServer/models/qwen3vl-4b-instruct"
DEFAULT_ADAPTER = (
    "/home/bbear/Documents/OlympusServer/optimization/h3-shrink/adapters/te_adapter_v1.safetensors"
)
DEFAULT_TOKENIZER = "/home/bbear/Documents/OlympusServer/optimization/h3-shrink/h3_tokenizer"
PAD_ID = 151643
EMBED_KEY = "qwen3vl_32b"  # keep teacher key so downstream nodes see a familiar dict shape


class Adapter(nn.Module):
    def __init__(self):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(2560, 4096),
            nn.GELU(),
            nn.Linear(4096, 5120),
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.net(x)


_CACHE_LOCK = threading.Lock()
_CACHE: dict[str, Any] = {}


def _pick_device(prefer: str = "auto") -> str:
    if prefer and prefer != "auto":
        return prefer
    if hasattr(torch, "xpu") and torch.xpu.is_available():
        # Prefer xpu:1 when two cards are present so the teacher TE dump
        # server can keep using xpu:0.
        n = torch.xpu.device_count()
        return f"xpu:{1 if n > 1 else 0}"
    if torch.cuda.is_available():
        return "cuda:0"
    return "cpu"


def _gguf_files() -> list[str]:
    files = folder_paths.get_filename_list("text_encoders")
    # ComfyUI-GGUF registers clip_gguf (text_encoders dirs, .gguf extension only).
    if "clip_gguf" in folder_paths.folder_names_and_paths:
        files = files + folder_paths.get_filename_list("clip_gguf")
    return sorted({f for f in files if f.endswith(".gguf")})


def _gguf_full_path(name: str) -> str:
    p = folder_paths.get_full_path("text_encoders", name)
    if p is None and "clip_gguf" in folder_paths.folder_names_and_paths:
        p = folder_paths.get_full_path("clip_gguf", name)
    if p is None:
        raise FileNotFoundError(f"h3_small_te: gguf not found: {name}")
    return p


def _import_gguf_backend():
    """Import the sibling ComfyUI-GGUF package (its dir name is not importable)."""
    name = "h3_small_te.gguf_backend"
    pkg = sys.modules.get(name)
    if pkg is None:
        pkg_dir = os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), "ComfyUI-GGUF")
        spec = importlib.util.spec_from_file_location(
            name, os.path.join(pkg_dir, "__init__.py"), submodule_search_locations=[pkg_dir]
        )
        pkg = importlib.util.module_from_spec(spec)
        sys.modules[name] = pkg
        spec.loader.exec_module(pkg)
    return pkg


def _load_text_stack_gguf(gguf_path: str):
    """Build the comfy-native Qwen3-VL-4B text stack from a GGUF file."""
    pkg = _import_gguf_backend()
    gguf_loader = sys.modules[pkg.__name__ + ".loader"]
    gguf_ops = sys.modules[pkg.__name__ + ".ops"]
    from comfy.text_encoders.llama import Llama2_, Qwen3VL_4BConfig

    print(f"[h3_small_te] loading student from {gguf_path} (gguf) ...", flush=True)
    sd = gguf_loader.gguf_clip_loader(gguf_path)
    sd = {k.removeprefix("model."): v for k, v in sd.items()}

    model = Llama2_(Qwen3VL_4BConfig(), device="cpu", dtype=torch.bfloat16, ops=gguf_ops.GGMLOps)
    missing, unexpected = model.load_state_dict(sd, strict=False)
    if missing or unexpected:
        raise RuntimeError(
            f"h3_small_te: gguf state dict mismatch: missing={missing} unexpected={unexpected}"
        )

    # gguf_clip_loader dequantizes token_embd to fp16; the safetensors path is bf16.
    emb = model.embed_tokens.weight
    model.embed_tokens.weight = nn.Parameter(emb.data.to(torch.bfloat16), requires_grad=False)
    model.eval()
    return model


def _load_stack(student_dir: str, adapter_path: str, tokenizer_dir: str, device: str,
                gguf_path: Optional[str] = None):
    """Load (and cache) student text stack + adapter + H3 tokenizer."""
    key = f"{gguf_path or student_dir}|{adapter_path}|{tokenizer_dir}|{device}"
    with _CACHE_LOCK:
        if key in _CACHE:
            return _CACHE[key]

        os.environ.setdefault("PYTORCH_ENABLE_XPU_FALLBACK", "1")
        os.environ.setdefault("ONEAPI_DEVICE_SELECTOR", "level_zero:*")

        from transformers import AutoTokenizer
        from safetensors.torch import load_file

        if gguf_path:
            text_model = _load_text_stack_gguf(gguf_path)
            path_used = f"gguf:{os.path.basename(gguf_path)}"
        else:
            from transformers import Qwen3VLForConditionalGeneration

            # Load on CPU first, peel the text stack, THEN move only that stack to
            # the target device. Moving the full VL (vision tower included) to XPU
            # has been observed to hang under concurrent SYCL loaders (llama-server
            # / Comfy GGUF TE). Training path is the same peel-then-run pattern.
            print(f"[h3_small_te] loading student from {student_dir} (cpu then {device}) ...", flush=True)
            model = Qwen3VLForConditionalGeneration.from_pretrained(
                student_dir, dtype=torch.bfloat16
            )
            model.eval()

            text_model = None
            path_used = None
            for cand in ("model.language_model", "language_model", "model.model.language_model"):
                obj = model
                ok = True
                for part in cand.split("."):
                    if hasattr(obj, part):
                        obj = getattr(obj, part)
                    else:
                        ok = False
                        break
                if ok and hasattr(obj, "norm") and hasattr(obj, "layers"):
                    text_model = obj
                    path_used = cand
                    break
            if text_model is None:
                raise RuntimeError("h3_small_te: could not locate student text stack")

        old_norm = text_model.norm
        text_model.norm = nn.Identity()
        print(
            f"[h3_small_te] text stack {path_used}; "
            f"{type(old_norm).__name__} -> Identity",
            flush=True,
        )
        for p in text_model.parameters():
            p.requires_grad_(False)

        # Detach text stack from the VL parent before device move so the vision
        # tower is not dragged onto the XPU.
        text_model = text_model.to(device)
        if not gguf_path:
            del model
            import gc
            gc.collect()
        if device.startswith("xpu") and hasattr(torch.xpu, "empty_cache"):
            torch.xpu.empty_cache()
        elif device.startswith("cuda") and hasattr(torch.cuda, "empty_cache"):
            torch.cuda.empty_cache()

        adapter = Adapter()
        if not os.path.isfile(adapter_path):
            raise FileNotFoundError(f"h3_small_te: adapter not found: {adapter_path}")
        sd = load_file(adapter_path)
        adapter.load_state_dict(sd, strict=True)
        adapter = adapter.to(device=device, dtype=torch.float32)
        adapter.eval()

        tokenizer = AutoTokenizer.from_pretrained(tokenizer_dir)

        bundle = {
            "text_model": text_model,
            "adapter": adapter,
            "tokenizer": tokenizer,
            "device": device,
        }
        _CACHE[key] = bundle
        print(f"[h3_small_te] ready on {device}", flush=True)
        return bundle


def _encode_ids(text_model, adapter, ids: list[int], device: str) -> torch.Tensor:
    """Return (1, L, 5120) fp32 conditioning tensor."""
    if not ids:
        ids = [PAD_ID]
    input_ids = torch.tensor([ids], dtype=torch.long, device=device)
    attention_mask = torch.ones_like(input_ids)
    with torch.no_grad():
        # Positional ids: HF takes input_ids first; comfy Llama2_ takes x (ids) first.
        out = text_model(input_ids, attention_mask=attention_mask)
        hidden = out.last_hidden_state if hasattr(out, "last_hidden_state") else out[0]
        cond = adapter(hidden.float())  # (1, L, 5120)
    return cond


def _token_ids_from_text(tokenizer, text: str) -> list[int]:
    return list(tokenizer.encode(text, add_special_tokens=False))


def _token_ids_from_clip_tokens(tokens) -> list[int]:
    """Extract flat token id list from a comfy-style tokenize() result."""
    if isinstance(tokens, dict):
        batches = next(iter(tokens.values()))
    else:
        batches = tokens
    if not batches:
        return [PAD_ID]
    entries = batches[0]
    ids = []
    for entry in entries:
        tid = entry[0] if isinstance(entry, (tuple, list)) else entry
        if isinstance(tid, dict):
            # Vision embed — Phase 2 is T2V-only; refuse silently-wrong paths.
            raise RuntimeError(
                "h3_small_te: vision/image tokens are not supported yet "
                "(adapter is text-only). Use pure T2V prompts."
            )
        ids.append(int(tid))
    return ids if ids else [PAD_ID]


class H3SmallCLIP:
    """Duck-typed CLIP for MiniMaxH3ImageToVideo (T2V pure-text path)."""

    def __init__(self, student_dir: str, adapter_path: str, tokenizer_dir: str, device: str,
                 gguf_path: Optional[str] = None):
        self.student_dir = student_dir
        self.adapter_path = adapter_path
        self.tokenizer_dir = tokenizer_dir
        self.device = device
        self.gguf_path = gguf_path
        self._bundle: Optional[dict] = None

    def _ensure(self):
        if self._bundle is None:
            self._bundle = _load_stack(
                self.student_dir, self.adapter_path, self.tokenizer_dir, self.device,
                gguf_path=self.gguf_path,
            )
        return self._bundle

    def tokenize(self, text, return_word_ids=False, images=None, minimax_ref_items=None, **kwargs):
        if images:
            raise RuntimeError(
                "h3_small_te: FL2VA image conditioning not supported yet "
                "(student adapter is text-only). Use T2V (no first/last frame)."
            )
        if minimax_ref_items:
            raise RuntimeError(
                "h3_small_te: ref2va not supported yet (student adapter is text-only)."
            )
        b = self._ensure()
        ids = _token_ids_from_text(b["tokenizer"], text)
        entries = [(tid, 1.0) for tid in ids] or [(PAD_ID, 1.0)]
        if return_word_ids:
            entries = [t + (0,) for t in entries]
        return {EMBED_KEY: [entries]}

    def encode_from_tokens_scheduled(self, tokens, unprojected=False, add_dict=None, show_pbar=True):
        add_dict = add_dict or {}
        b = self._ensure()
        ids = _token_ids_from_clip_tokens(tokens)
        cond = _encode_ids(b["text_model"], b["adapter"], ids, b["device"])
        # Match comfy TE output placement (usually CPU / model management device).
        cond = cond.cpu()
        tags = torch.ones(cond.shape[1], dtype=torch.long)
        pooled = {
            "pooled_output": cond[:, -1, :].clone(),
            "minimax_token_tags": tags,
        }
        pooled.update(add_dict)
        return [[cond, pooled]]

    def encode_from_tokens(self, tokens, return_pooled=False, return_dict=False):
        scheduled = self.encode_from_tokens_scheduled(tokens)
        cond, pooled = scheduled[0]
        if return_dict:
            out = {"cond": cond, "pooled_output": pooled.get("pooled_output")}
            for k, v in pooled.items():
                if k != "pooled_output":
                    out[k] = v
            return out
        if return_pooled:
            return cond, pooled.get("pooled_output")
        return cond

    def encode(self, text):
        return self.encode_from_tokens(self.tokenize(text))


class H3SmallTELoader:
    """Load Qwen3-VL-4B + adapter as a CLIP substitute for MiniMax H3 T2V."""

    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {},
            "optional": {
                "gguf_name": (["none"] + _gguf_files(),),
                "student_dir": ("STRING", {"default": DEFAULT_STUDENT}),
                "adapter_path": ("STRING", {"default": DEFAULT_ADAPTER}),
                "tokenizer_dir": ("STRING", {"default": DEFAULT_TOKENIZER}),
                "device": ("STRING", {"default": "auto"}),
            },
        }

    RETURN_TYPES = ("CLIP",)
    FUNCTION = "load"
    CATEGORY = "h3"
    TITLE = "H3 Small TE Loader (4B+adapter)"

    def load(self, gguf_name="none", student_dir=DEFAULT_STUDENT, adapter_path=DEFAULT_ADAPTER,
             tokenizer_dir=DEFAULT_TOKENIZER, device="auto"):
        dev = _pick_device(device)
        gguf_path = None
        if gguf_name != "none":
            gguf_path = _gguf_full_path(gguf_name)
        clip = H3SmallCLIP(student_dir, adapter_path, tokenizer_dir, dev, gguf_path=gguf_path)
        # Eager-load so the first workflow step surfaces errors immediately.
        clip._ensure()
        return (clip,)


class H3SmallTextEncoder:
    """Encode a prompt with Qwen3-VL-4B + adapter -> CONDITIONING (1, L, 5120)."""

    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "text": ("STRING", {"multiline": True, "dynamicPrompts": True}),
            },
            "optional": {
                "gguf_name": (["none"] + _gguf_files(),),
                "student_dir": ("STRING", {"default": DEFAULT_STUDENT}),
                "adapter_path": ("STRING", {"default": DEFAULT_ADAPTER}),
                "tokenizer_dir": ("STRING", {"default": DEFAULT_TOKENIZER}),
                "device": ("STRING", {"default": "auto"}),
            },
        }

    RETURN_TYPES = ("CONDITIONING",)
    FUNCTION = "encode"
    CATEGORY = "h3"
    TITLE = "H3 Small Text Encoder (4B+adapter)"
    OUTPUT_NODE = True

    def encode(self, text, gguf_name="none", student_dir=DEFAULT_STUDENT, adapter_path=DEFAULT_ADAPTER,
               tokenizer_dir=DEFAULT_TOKENIZER, device="auto"):
        dev = _pick_device(device)
        gguf_path = None
        if gguf_name != "none":
            gguf_path = _gguf_full_path(gguf_name)
        b = _load_stack(student_dir, adapter_path, tokenizer_dir, dev, gguf_path=gguf_path)
        ids = _token_ids_from_text(b["tokenizer"], text)
        cond = _encode_ids(b["text_model"], b["adapter"], ids, b["device"]).cpu()
        tags = torch.ones(cond.shape[1], dtype=torch.long)
        pooled = {
            "pooled_output": cond[:, -1, :].clone(),
            "minimax_token_tags": tags,
        }
        print(
            f"[h3_small_te] encoded L={cond.shape[1]} dim={cond.shape[2]} "
            f"mean||h||={cond[0].norm(dim=-1).mean().item():.1f}",
            flush=True,
        )
        return ([[cond, pooled]],)


NODE_CLASS_MAPPINGS = {
    "H3SmallTELoader": H3SmallTELoader,
    "H3SmallTextEncoder": H3SmallTextEncoder,
}

NODE_DISPLAY_NAME_MAPPINGS = {
    "H3SmallTELoader": "H3 Small TE Loader (4B+adapter)",
    "H3SmallTextEncoder": "H3 Small Text Encoder (4B+adapter)",
}