streaming

app.py

@@ -3,7 +3,8 @@ import random
 import torch
 import hashlib
 import gradio as gr
-from transformers import LlavaNextProcessor, LlavaNextForConditionalGeneration
+import threading
+from transformers import LlavaNextProcessor, LlavaNextForConditionalGeneration, TextIteratorStreamer
 
 model_id = "ibm-granite/granite-vision-3.1-2b-preview"
 processor = LlavaNextProcessor.from_pretrained(model_id, use_fast=True)
@@ -32,8 +33,8 @@ def chat_inference(image, text, temperature, top_p, top_k, max_tokens, conversat
         user_content.append({"type": "text", "text": text.strip()})
 
     if not user_content:
-        return conversation_display(conversation), conversation
-
+        return conversation_display(conversation), conversation, "", False
+    
     conversation.append({
         "role": "user",
         "content": user_content
@@ -63,17 +64,32 @@ def chat_inference(image, text, temperature, top_p, top_k, max_tokens, conversat
         generation_kwargs["temperature"] = temperature
         generation_kwargs["do_sample"] = True
 
-    output = model.generate(**inputs, **generation_kwargs)
-    raw_response = processor.decode(output[0], skip_special_tokens=True)
-    assistant_text = extract_answer(raw_response)
-
-    # Append the assistant's answer.
     conversation.append({
         "role": "assistant",
-        "content": [{"type": "text", "text": assistant_text}]
+        "content": [{"type": "text", "text": ""}]
     })
     
-    return conversation_display(conversation), conversation
+    yield conversation_display(conversation), conversation, "Processing...", True
+
+    streamer = TextIteratorStreamer(processor.tokenizer, skip_prompt=True, skip_special_tokens=True)
+    generation_kwargs["streamer"] = streamer
+
+    def generate_thread():
+        model.generate(**inputs, **generation_kwargs)
+        
+    thread = threading.Thread(target=generate_thread)
+    thread.start()
+
+    assistant_text = ""
+    for new_text in streamer:
+        assistant_text += new_text
+        conversation[-1]["content"][0]["text"] = extract_answer(assistant_text)
+        yield conversation_display(conversation), conversation, "Processing...", True
+
+    thread.join()
+    
+    yield conversation_display(conversation), conversation, "", False
+    return 
 
 def extract_answer(response):
     if "<|assistant|>" in response:
@@ -142,8 +158,11 @@ def conversation_display(conversation):
             })
     return chat_history
 
-def clear_chat():
-    return [], [], "", None
+def clear_chat(chat_history, conversation, text_value, image, is_generating):
+    if is_generating:
+        return chat_history, conversation, text_value, image, is_generating
+    else:
+        return [], [], "", None, is_generating
 
 with gr.Blocks(title="Granite Vision 3.1 2B", css="h1 { overflow: hidden; }") as demo:
     gr.Markdown("# [Granite Vision 3.1 2B](https://huggingface.co/ibm-granite/granite-vision-3.1-2b-preview)")
@@ -164,18 +183,19 @@ with gr.Blocks(title="Granite Vision 3.1 2B", css="h1 { overflow: hidden; }") as
                 send_button = gr.Button("Chat")
                 clear_button = gr.Button("Clear Chat")
     
-    state = gr.State([])
+    conversation_state = gr.State([])
+    is_generating = gr.State(False)
 
     send_button.click(
         chat_inference,
-        inputs=[image_input, text_input, temperature_input, top_p_input, top_k_input, max_tokens_input, state],
-        outputs=[chatbot, state]
+        inputs=[image_input, text_input, temperature_input, top_p_input, top_k_input, max_tokens_input, conversation_state],
+        outputs=[chatbot, conversation_state, text_input, is_generating]
     )
 
     clear_button.click(
         clear_chat,
-        inputs=None,
-        outputs=[chatbot, state, text_input, image_input]
+        inputs=[chatbot, conversation_state, text_input, image_input, is_generating],
+        outputs=[chatbot, conversation_state, text_input, image_input, is_generating]
     )
 
     gr.Examples(