import gradio as gr
from transformers import pipeline
import torch

pipe = pipeline(
    "text-to-text",
    model="rahul7star/gemma-3bit",
    device="cpu",
    torch_dtype=torch.bfloat16
)

def chat(image, prompt):
    messages = [
        {"role": "user", "content": [{"type": "image", "image": image}, {"type": "text", "text": prompt}]}
    ]
    out = pipe(text=messages, max_new_tokens=200)
    return out[0]["generated_text"][-1]["content"]

demo = gr.Interface(
    fn=chat,
    inputs=[gr.Image(type="pil"), gr.Textbox(label="Prompt")],
    outputs=gr.Textbox(label="Response"),
    title="Gemma 3 4B‑IT Vision‑Text Chat",
    description="Upload an image and ask questions using Google’s gems!"
)

demo.launch()