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

"""Prompt-only conditioning primitives independent of ComfyUI's node API."""

from pathlib import Path

import torch


class H3PromptTokenizer:
    """Tokenize raw H3 prompt text without Qwen chat-template tokens."""

    def __init__(self, tokenizer_dir: str | Path):
        from transformers import Qwen2Tokenizer

        tokenizer_dir = Path(tokenizer_dir)
        self.tokenizer = Qwen2Tokenizer(
            vocab_file=str(tokenizer_dir / "vocab.json"),
            merges_file=str(tokenizer_dir / "merges.txt"),
        )

    def __call__(self, prompt: str, *, device: torch.device | str = "cuda") -> torch.Tensor:
        if not prompt:
            prompt = " "
        encoded = self.tokenizer(prompt, add_special_tokens=False, return_tensors="pt")
        return encoded.input_ids.to(device)