mirror of
https://github.com/comfyanonymous/ComfyUI.git
synced 2026-01-10 06:10:50 +08:00
40 lines
1.7 KiB
Python
40 lines
1.7 KiB
Python
import pytest
|
|
from comfy_execution.graph_utils import GraphBuilder
|
|
from comfy.client.embedded_comfy_client import Comfy
|
|
from comfy.api.components.schema.prompt import Prompt
|
|
|
|
class TestQwenVLVideo:
|
|
@pytest.mark.asyncio
|
|
async def test_qwenvl_video_loading(self):
|
|
graph = GraphBuilder()
|
|
|
|
# Load QwenVL model (using a small one as requested)
|
|
# Qwen/Qwen2-VL-2B-Instruct is a good candidate for a "small" QwenVL model
|
|
model_loader = graph.node("TransformersLoader1", ckpt_name="Qwen/Qwen2-VL-2B-Instruct")
|
|
|
|
# Load video from URL with frame cap to avoid OOM
|
|
video_url = "https://upload.wikimedia.org/wikipedia/commons/f/f7/2024-04-05_Luisenpark_MA_Ziegen_2.webm"
|
|
load_video = graph.node("LoadVideoFromURL", value=video_url, frame_load_cap=16, select_every_nth=10)
|
|
|
|
# Tokenize with video
|
|
# OneShotInstructTokenize has optional 'videos' input
|
|
tokenizer = graph.node("OneShotInstructTokenize", model=model_loader.out(0), prompt="Describe this video.", videos=load_video.out(0), chat_template="default")
|
|
|
|
# Generate
|
|
generation = graph.node("TransformersGenerate", model=model_loader.out(0), tokens=tokenizer.out(0), max_new_tokens=50, seed=42)
|
|
|
|
# OmitThink (as requested)
|
|
omit_think = graph.node("OmitThink", value=generation.out(0))
|
|
|
|
# Save output
|
|
graph.node("SaveString", value=omit_think.out(0), filename_prefix="qwenvl_video_test")
|
|
|
|
workflow = graph.finalize()
|
|
prompt = Prompt.validate(workflow)
|
|
|
|
async with Comfy() as client:
|
|
outputs = await client.queue_prompt(prompt)
|
|
|
|
# We expect it to fail before this, but if it succeeds, we should check the output
|
|
assert len(outputs) > 0
|