import os
import sys
import torch
from openai import OpenAI
from transformers import (
    LlavaNextProcessor, LlavaNextForConditionalGeneration, 
    Qwen2_5_VLForConditionalGeneration, AutoProcessor
)
## init device
device = "cuda"
torch_dtype = torch.float16


vlms_list = [
    {
        "type": "qwen2-vl",
        "name": "Qwen2.5-VL-7B-Instruct (Default)",
        "local_path": "models/vlms/Qwen/Qwen2.5-VL-7B-Instruct",
        "processor": AutoProcessor.from_pretrained(
            "models/vlms/Qwen/Qwen2.5-VL-7B-Instruct"
        ) if os.path.exists("models/vlms/Qwen/Qwen2.5-VL-7B-Instruct") else AutoProcessor.from_pretrained(
            "Qwen/Qwen2.5-VL-7B-Instruct"
        ),
        "model": Qwen2_5_VLForConditionalGeneration.from_pretrained(
            "models/vlms/Qwen/Qwen2.5-VL-7B-Instruct", torch_dtype=torch_dtype, device_map=device
        ).to(device) if os.path.exists("models/vlms/Qwen/Qwen2.5-VL-7B-Instruct") else 
            Qwen2_5_VLForConditionalGeneration.from_pretrained(
                "Qwen/Qwen2.5-VL-7B-Instruct", torch_dtype=torch_dtype, device_map=device
            ).to(device),
    },
    {
        "type": "openai",
        "name": "GPT4-o (Highly Recommended)",
        "local_path": "",
        "processor": "",
        "model": ""
    },
]

vlms_template = {k["name"]: (k["type"], k["local_path"], k["processor"], k["model"]) for k in vlms_list}