import torch

from h3_blackwell_runtime.checkpoint import H3Checkpoint
from h3_blackwell_runtime.qwen3vl_text import Qwen3VLPromptConditioner
from h3_blackwell_runtime.token_refiner import H3TokenRefiner


states = Qwen3VLPromptConditioner("/text-encoders/qwen3vl_32b_minimax_h3_nvfp4_awq.safetensors", "/opt/h3-blackwell-runtime/qwen25_tokenizer")("A brass dragon flies through rain.")
refined = H3TokenRefiner(H3Checkpoint("/models/minimax_h3_fl2va_pruned_nvfp4.safetensors"))(states)
torch.cuda.synchronize()
print({"shape": tuple(refined.shape), "dtype": str(refined.dtype), "finite": torch.isfinite(refined).all().item()})