import gradio as gr
from gradio_client import Client, handle_file
from huggingface_hub import InferenceClient

moondream_client = Client("vikhyatk/moondream2")
qwq_client = InferenceClient("Qwen/QwQ-32B-Preview")

def describe_image(image, user_message):
    result = moondream_client.predict(
        img=handle_file(image),
        prompt="Describe this image.",
        api_name="/answer_question"
    )
    
    description = result

    user_message = description + "\n" + user_message 

    qwq_result = qwq_client.chat_completion(
        messages=[{"role": "user", "content": user_message}],
        max_tokens=512,
        temperature=0.7,
        top_p=0.95
    )
    
    return qwq_result['choices'][0]['message']['content']

def chat_or_image(image, user_message):
    if image:
        return describe_image(image, user_message)
    else:
        qwq_result = qwq_client.chat_completion(
            messages=[{"role": "user", "content": user_message}],
            max_tokens=512,
            temperature=0.7,
            top_p=0.95
        )
        return qwq_result['choices'][0]['message']['content']

demo = gr.Interface(
    fn=chat_or_image,
    inputs=[
        gr.Image(type="filepath", label="Upload image (Optional)"),
        gr.Textbox(label="Ask anything", placeholder="Ask...", lines=2)
    ],
    outputs="text",
)

if __name__ == "__main__":
    demo.launch(show_error=True)