Add h3_small_te custom node (H3SmallTELoader + H3SmallTextEncoder)

custom_nodes/h3_small_te/__init__.py

@@ -0,0 +1,3 @@
+from .nodes import NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS
+
+__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS"]

custom_nodes/h3_small_te/nodes.py

@@ -0,0 +1,410 @@
+"""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)",
+}