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 AutoTokenizer

        self.tokenizer = AutoTokenizer.from_pretrained(
            str(Path(tokenizer_dir)),
            local_files_only=True,
            fix_mistral_regex=True,
        )

    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)