MOSS-Video-Preview-Real-Time-SFT

import os
import queue
import threading

import torch
from transformers import AutoModelForCausalLM, AutoProcessor

checkpoint = "fnlp-vision/moss-video-preview-realtime-sft"
video_path = "data/example_video.mp4"
prompt = "Describe the video."

max_new_tokens = 1024
temperature = 1.0
top_k = 50
top_p = 1.0
repetition_penalty = 1.0

video_fps = 1.0
video_minlen = 8
video_maxlen = 256


def load_model(checkpoint: str):
    processor = AutoProcessor.from_pretrained(
        checkpoint, trust_remote_code=True, frame_extract_num_threads=1
    )
    model = AutoModelForCausalLM.from_pretrained(
        checkpoint,
        trust_remote_code=True,
        device_map="auto",
        torch_dtype=torch.bfloat16,
        attn_implementation="flash_attention_2",
    )
    return model, processor


if not checkpoint:
    raise ValueError("Missing `checkpoint`.")
if not video_path:
    raise ValueError("Missing `video_path`.")
if not os.path.isfile(video_path):
    raise FileNotFoundError(f"Video not found: {video_path}")

model, processor = load_model(checkpoint)
new_queries: "queue.Queue[dict]" = queue.Queue()
output_text_queue: "queue.Queue[str]" = queue.Queue()

new_queries.put(
    {
        "prompt": f"\n{prompt}",
        "images": [],
        "videos": [video_path],
        "media_kwargs": {
            "video_fps": video_fps,
            "video_minlen": video_minlen,
            "video_maxlen": video_maxlen,
        },
        "thinking_mode": "no_thinking",
        "system_prompt_type": "video",
        "generate_kwargs": {
            "temperature": temperature,
            "top_k": top_k,
            "top_p": top_p,
            "max_new_tokens": max_new_tokens,
            "repetition_penalty": repetition_penalty,
        },
        "stop_offline_generate": False,
    }
)
new_queries.put({"stop_offline_generate": True})


def drain_output():
    while True:
        tok = output_text_queue.get()
        if tok == "<|round_end|>":
            break
        print(tok, end="", flush=True)


t = threading.Thread(target=drain_output, daemon=True)
t.start()
with torch.no_grad():
    model.offline_generate(processor, new_queries, output_text_queue, vision_chunked_length=64)
t.join(timeout=5.0)

import os, queue, threading, torch
from PIL import Image
from transformers import AutoModelForCausalLM, AutoProcessor
checkpoint = "fnlp-vision/moss-video-preview-realtime-sft"
image_path = "data/example_image.jpg"
prompt = "Describe this image."
if not os.path.isfile(image_path):
    raise FileNotFoundError(image_path)

processor = AutoProcessor.from_pretrained(
    checkpoint, trust_remote_code=True, frame_extract_num_threads=1
)
model = AutoModelForCausalLM.from_pretrained(
    checkpoint, trust_remote_code=True, device_map="auto", torch_dtype=torch.bfloat16, attn_implementation="flash_attention_2"
)

new_q, out_q = queue.Queue(), queue.Queue()
new_q.put(
    {
        "prompt": f"\n{prompt}",
        "images": [Image.open(image_path).convert("RGB")],
        "videos": [],
        "system_prompt_type": "text_image",
        "thinking_mode": "no_thinking",
        "generate_kwargs": {"temperature": 1.0, "top_k": 50, "top_p": 1.0, "max_new_tokens": 256, "repetition_penalty": 1.0},
        "stop_offline_generate": False,
    }
)
new_q.put({"stop_offline_generate": True})

threading.Thread(
    target=lambda: (lambda: [print(t, end="", flush=True) for t in iter(out_q.get, "<|round_end|>")])(),
    daemon=True,
).start()

with torch.no_grad():
    model.offline_generate(processor, new_q, out_q, vision_chunked_length=64)