aas

app.py

@@ -2,6 +2,8 @@ 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"
@@ -14,12 +16,23 @@ processor = AutoProcessor.from_pretrained(MODEL_ID)
 @spaces.GPU
 def tag_image(image):
     inputs = processor(images=image, text="<image>tag en", return_tensors="pt")
-    outputs = model.generate(**inputs)
-    return processor.decode(outputs[0], skip_special_tokens=True)
+    streamer = TextIteratorStreamer(tokenizer=processor.tokenizer, skip_prompt=True)
+    generation_kwargs = dict(inputs, streamer=streamer, return_full_text=False)
+    thread = Thread(target=model.generate, kwargs=generation_kwargs)
+    
+    text = ""
+    thread.start()
+    for new_text in streamer:
+        text += new_text
+        yield text
+    
+    thread.join()
+
+    return text
 
 gr.Interface(
     fn=tag_image,
-    inputs=gr.Image(type="pil"),
+    inputs=[gr.Image(type="pil")],
     outputs=gr.Textbox(type="text"),
     title=TITLE,
     description=DESCRIPTION,