import gradio as gr
import os

# Set the CUDA_VISIBLE_DEVICES environment variable
os.environ["CUDA_VISIBLE_DEVICES"] = "0"

# Import the demo script
from mllm.demo.web_demo import NextChatInference

# Initialize the model
model_path = "AoZhang/nextchat-7b-336"
vit_path = "openai/clip-vit-large-patch14-336"
model = NextChatInference(model_path, vit_path, 576)

# Gradio Interface
iface = gr.Interface(
    fn=model,
    inputs=["text", "image"],
    outputs=["text", "image"],
    live=True,
    capture_session=True,
    interpretation="default",
)

# Launch the interface
iface.launch()