import gradio as gr
import spaces
from PIL import Image
from transformers import PaliGemmaForConditionalGeneration, AutoProcessor
from transformers import TextIteratorStreamer
from threading import Thread

TITLE = "E621 Tagger"
DESCRIPTION = "Tag images with E621 tags"

MODEL_ID = "estrogen/paligemma2-3b-e621-224"

model = PaliGemmaForConditionalGeneration.from_pretrained(MODEL_ID)
model.to("cuda")
processor = AutoProcessor.from_pretrained(MODEL_ID)

@spaces.GPU
def tag_image(image, max_new_tokens=128, temperature=1, top_p=1, min_p=0):
    inputs = processor(images=image, text="<image>tag en", return_tensors="pt").to("cuda")
    streamer = TextIteratorStreamer(tokenizer=processor.tokenizer, skip_prompt=True)
    generation_kwargs = dict(inputs, streamer=streamer, max_new_tokens=max_new_tokens, use_cache=True, cache_implementation="hybrid", do_sample=True, temperature=temperature, top_p=top_p, min_p=min_p)
    thread = Thread(target=model.generate, kwargs=generation_kwargs)
    
    text = ""
    thread.start()
    for new_text in streamer:
        text += new_text
        yield text

    return text

gr.Interface(
    fn=tag_image,
    inputs=[gr.Image(type="pil"), gr.Slider(label="Max new tokens", minimum=1, maximum=1024, value=128), gr.Slider(label="Temperature", minimum=0, maximum=2, value=1), gr.Slider(label="Top p", minimum=0, maximum=1, value=1), gr.Slider(label="Min p", minimum=0, maximum=1, value=0)],
    outputs=gr.Textbox(type="text"),
    title=TITLE,
    description=DESCRIPTION,
).launch()