import gradio as gr
from transformers import AutoTokenizer, AutoModelForCausalLM
import torch

# ✅ LOAD YOUR MODEL
model_name = "rafiaislam/meme-caption-model"

tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(model_name)

# ✅ FUNCTION
def generate_joke(prompt):
    if not prompt:
        return "Enter prompt 😅"

    try:
        inputs = tokenizer(prompt, return_tensors="pt")

        outputs = model.generate(
            **inputs,
            max_length=60,
            do_sample=True,
            temperature=0.9,
            top_k=50,
            top_p=0.95
        )

        result = tokenizer.decode(outputs[0], skip_special_tokens=True)

        return result

    except Exception as e:
        return f"Error: {str(e)}"

# ✅ UI
demo = gr.Interface(
    fn=generate_joke,
    inputs=gr.Textbox(label="prompt"),
    outputs=gr.Textbox(label="output"),
    title="Meme Caption Generator"
)

demo.launch(server_name="0.0.0.0", server_port=7860)