import os
import gc
import re
import ast
import json
import base64
import random
import uuid
import time
from io import BytesIO
from threading import Thread
import gradio as gr
import spaces
import torch
import numpy as np
from PIL import Image, ImageOps
import cv2
from transformers import (
Qwen2_5_VLForConditionalGeneration,
Qwen2VLForConditionalGeneration,
AutoProcessor,
TextIteratorStreamer,
)
MAX_MAX_NEW_TOKENS = 4096
DEFAULT_MAX_NEW_TOKENS = 1024
MAX_INPUT_TOKEN_LENGTH = int(os.getenv("MAX_INPUT_TOKEN_LENGTH", "4096"))
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
print("Using device:", device)
MODEL_ID_N = "prithivMLmods/DeepCaption-VLA-7B"
processor_n = AutoProcessor.from_pretrained(MODEL_ID_N, trust_remote_code=True)
model_n = Qwen2_5_VLForConditionalGeneration.from_pretrained(
MODEL_ID_N,
attn_implementation="kernels-community/flash-attn2",
trust_remote_code=True,
torch_dtype=torch.float16
).to(device).eval()
MODEL_ID_M = "Skywork/SkyCaptioner-V1"
processor_m = AutoProcessor.from_pretrained(MODEL_ID_M, trust_remote_code=True)
model_m = Qwen2_5_VLForConditionalGeneration.from_pretrained(
MODEL_ID_M,
attn_implementation="kernels-community/flash-attn2",
trust_remote_code=True,
torch_dtype=torch.float16
).to(device).eval()
MODEL_ID_Z = "remyxai/SpaceThinker-Qwen2.5VL-3B"
processor_z = AutoProcessor.from_pretrained(MODEL_ID_Z, trust_remote_code=True)
model_z = Qwen2_5_VLForConditionalGeneration.from_pretrained(
MODEL_ID_Z,
attn_implementation="kernels-community/flash-attn2",
trust_remote_code=True,
torch_dtype=torch.float16
).to(device).eval()
MODEL_ID_K = "prithivMLmods/coreOCR-7B-050325-preview"
processor_k = AutoProcessor.from_pretrained(MODEL_ID_K, trust_remote_code=True)
model_k = Qwen2VLForConditionalGeneration.from_pretrained(
MODEL_ID_K,
attn_implementation="kernels-community/flash-attn2",
trust_remote_code=True,
torch_dtype=torch.float16
).to(device).eval()
MODEL_ID_Y = "remyxai/SpaceOm"
processor_y = AutoProcessor.from_pretrained(MODEL_ID_Y, trust_remote_code=True)
model_y = Qwen2_5_VLForConditionalGeneration.from_pretrained(
MODEL_ID_Y,
attn_implementation="kernels-community/flash-attn2",
trust_remote_code=True,
torch_dtype=torch.float16
).to(device).eval()
MODEL_MAP = {
"DeepCaption-VLA-7B": (processor_n, model_n),
"SkyCaptioner-V1": (processor_m, model_m),
"SpaceThinker-3B": (processor_z, model_z),
"coreOCR-7B-050325-preview": (processor_k, model_k),
"SpaceOm-3B": (processor_y, model_y),
}
MODEL_CHOICES = list(MODEL_MAP.keys())
image_examples = [
{"query": "type out the messy hand-writing as accurately as you can.", "image": "images/1.jpg", "model": "coreOCR-7B-050325-preview"},
{"query": "count the number of birds and explain the scene in detail.", "image": "images/2.jpeg", "model": "DeepCaption-VLA-7B"},
{"query": "how far is the Goal from the penalty taker in this image?.", "image": "images/3.png", "model": "SpaceThinker-3B"},
{"query": "approximately how many meters apart are the chair and bookshelf?.", "image": "images/4.png", "model": "SkyCaptioner-V1"},
{"query": "how far is the man in the red hat from the pallet of boxes in feet?.", "image": "images/5.jpg", "model": "SpaceOm-3B"},
]
video_examples = [
{"query": "give the highlights of the movie scene video.", "video": "videos/1.mp4", "model": "DeepCaption-VLA-7B"},
{"query": "explain the advertisement in detail.", "video": "videos/2.mp4", "model": "SkyCaptioner-V1"},
]
def pil_to_data_url(img: Image.Image, fmt="PNG"):
buf = BytesIO()
img.save(buf, format=fmt)
data = base64.b64encode(buf.getvalue()).decode()
mime = "image/png" if fmt.upper() == "PNG" else "image/jpeg"
return f"data:{mime};base64,{data}"
def file_to_data_url(path):
if not os.path.exists(path):
return ""
ext = path.rsplit(".", 1)[-1].lower()
mime = {
"jpg": "image/jpeg",
"jpeg": "image/jpeg",
"png": "image/png",
"webp": "image/webp",
"mp4": "video/mp4",
"mov": "video/quicktime",
"webm": "video/webm",
"mkv": "video/x-matroska",
}.get(ext, "application/octet-stream")
with open(path, "rb") as f:
data = base64.b64encode(f.read()).decode()
return f"data:{mime};base64,{data}"
def make_thumb_b64(path, max_dim=240):
try:
img = Image.open(path).convert("RGB")
img.thumbnail((max_dim, max_dim))
return pil_to_data_url(img, "JPEG")
except Exception as e:
print("Thumbnail error:", e)
return ""
def make_video_thumb_b64(path, max_dim=240):
try:
cap = cv2.VideoCapture(path)
total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
target = max(0, total_frames // 2)
cap.set(cv2.CAP_PROP_POS_FRAMES, target)
success, frame = cap.read()
cap.release()
if not success:
return ""
frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
img = Image.fromarray(frame).convert("RGB")
img.thumbnail((max_dim, max_dim))
return pil_to_data_url(img, "JPEG")
except Exception as e:
print("Video thumbnail error:", e)
return ""
def build_example_cards_html():
cards = ""
for i, ex in enumerate(image_examples):
thumb = make_thumb_b64(ex["image"])
prompt_short = ex["query"][:72] + ("..." if len(ex["query"]) > 72 else "")
cards += f"""
{"" if thumb else "
Preview
"}
{ex["model"]}IMAGE
{prompt_short}
"""
for i, ex in enumerate(video_examples):
thumb = make_video_thumb_b64(ex["video"])
prompt_short = ex["query"][:72] + ("..." if len(ex["query"]) > 72 else "")
cards += f"""
![](" + thumb + ")
" if thumb else "