import gradio as gr
from huggingface_hub import InferenceClient
from transformers import LlavaProcessor, LlavaForConditionalGeneration, TextIteratorStreamer
from PIL import Image
from threading import Thread

# Initialize model and processor
model_id = "llava-hf/llava-interleave-qwen-0.5b-hf"
processor = LlavaProcessor.from_pretrained(model_id)
model = LlavaForConditionalGeneration.from_pretrained(model_id).to("cpu")

# Initialize inference clients
client_mistral = InferenceClient("mistralai/Mistral-7B-Instruct-v0.3")

def llava(inputs):
    """Processes an image and text input using Llava."""
    try:
        image = Image.open(inputs["files"][0]).convert("RGB")
        prompt = f"<|im_start|>user <image>\n{inputs['text']}<|im_end|>"
        processed = processor(prompt, image, return_tensors="pt").to("cpu")
        return processed
    except Exception as e:
        print(f"Error in llava function: {e}")
        return None

def respond(message, history):
    """Generate a response based on text or image input."""
    try:
        if "files" in message and message["files"]:
            # Handle image + text input
            inputs = llava(message)
            if inputs is None:
                raise ValueError("Failed to process image input")
            
            streamer = TextIteratorStreamer(skip_prompt=True, skip_special_tokens=True)
            thread = Thread(target=model.generate, kwargs=dict(inputs=inputs, max_new_tokens=512, streamer=streamer))
            thread.start()
            
            buffer = ""
            for new_text in streamer:
                buffer += new_text
                history[-1][1] = buffer
                yield history, history
        else:
            # Handle text-only input
            user_message = message["text"]
            history.append([user_message, None])
            prompt = [{"role": "user", "content": msg[0]} for msg in history if msg[0]]
            response = client_mistral.chat_completion(prompt, max_tokens=200)
            bot_message = response["choices"][0]["message"]["content"]
            history[-1][1] = bot_message
            yield history, history
    except Exception as e:
        print(f"Error in respond function: {e}")
        history[-1][1] = f"An error occurred: {str(e)}"
        yield history, history

# Set up Gradio interface
with gr.Blocks() as demo:
    chatbot = gr.Chatbot()
    with gr.Row():
        with gr.Column():
            text_input = gr.Textbox(placeholder="Enter your message...")
            file_input = gr.File(label="Upload an image")

    def handle_text(text, history=[]):
        """Handle text input and generate responses."""
        return respond({"text": text}, history)

    def handle_file_upload(files, history=[]):
        """Handle file uploads and generate responses."""
        return respond({"files": files, "text": "Describe this image."}, history)

    # Connect components to callbacks
    text_input.submit(handle_text, [text_input, chatbot], [chatbot, chatbot])
    file_input.change(handle_file_upload, [file_input, chatbot], [chatbot, chatbot])

# Launch the Gradio app
demo.launch()