ComfyUI/tests/language/test_language.py
2025-12-03 14:40:52 -08:00

15 lines
522 B
Python

import torch
from comfy_extras.nodes.nodes_language import TransformersLoader, OneShotInstructTokenize
def test_integration_transformers_loader_and_tokenize():
loader = TransformersLoader()
tokenize = OneShotInstructTokenize()
model, = loader.execute("llava-hf/llava-v1.6-mistral-7b-hf", "")
tokens, = tokenize.execute(model, "Describe this image:", torch.rand((1, 224, 224, 3)), "llava-v1.6-mistral-7b-hf", )
assert isinstance(tokens, dict)
assert "input_ids" in tokens or "inputs" in tokens