import os
import shutil
import subprocess

# Clone the repository if not already present
if not os.path.exists("edge_vlm"):
    subprocess.run(["git", "clone", "https://huggingface.co/irotem98/edge_vlm"])

# Install the required dependencies
subprocess.run(["pip", "install", "-r", "edge_vlm/requirements.txt"])

# Copy all files from edge_vlm to the current directory
source_dir = "edge_vlm"
destination_dir = "."

for filename in os.listdir(source_dir):
    source_file = os.path.join(source_dir, filename)
    destination_file = os.path.join(destination_dir, filename)

    # Copy files, skipping directories like .git
    if os.path.isfile(source_file):
        shutil.copy(source_file, destination_file)

# Now import the model from the copied files
from model import MoondreamModel
import torch
import gradio as gr

# Load the model and tokenizer
model = MoondreamModel.load_model()
tokenizer = MoondreamModel.load_tokenizer()

# Define the default question
default_question = "Describe the image."

# Function to handle image and return generated caption
def generate_caption_with_default(image):
    # Preprocess the image
    preprocessed_image = MoondreamModel.preprocess_image(image)
    
    # Generate caption
    caption = MoondreamModel.generate_caption(model, preprocessed_image, tokenizer)
    
    return caption

# Create Gradio interface
interface = gr.Interface(
    fn=generate_caption_with_default,
    inputs=gr.inputs.Image(type="pil", label="Upload an Image"),
    outputs="text",
    title="Image Caption Generator",
    description=f"The default question is: '{default_question}'. Upload an image to generate a description."
)

# Launch the interface
interface.launch()