Update app.py

app.py

@@ -3,9 +3,9 @@ import torch
 from PIL import Image
 import gradio as gr
 import spaces
-from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
-from huggingface_hub.inference._generated.types import TextGenerationStreamOutput, TextGenerationStreamOutputToken
+from transformers import AutoModelForCausalLM, AutoTokenizer
 import os
+import time
 from huggingface_hub import hf_hub_download
 
 
@@ -109,35 +109,28 @@ def stream_chat(message, history: list, system: str, temperature: float, max_new
             return_tensors="pt"
         ).to(model.device)
         images = None
-    streamer = TextIteratorStreamer(tokenizer, timeout=10.0, skip_prompt=True, skip_special_tokens=True)
 
     generate_kwargs = dict(
         input_ids=input_ids,
-        streamer=streamer,
         max_new_tokens=max_new_tokens,
         temperature=temperature,
         do_sample=True,
+        num_beams=1,
         eos_token_id=terminators,
         images=images
     )
     if temperature == 0:
         generate_kwargs["do_sample"] = False
 
-    t = Thread(target=model.generate, kwargs=generate_kwargs)
-    t.start()
+    output_ids=model.generate(**generate_kwargs)
     input_token_len = input_ids.shape[1]
-    
-    output = ""
-    for next_text in streamer:
-        yield TextGenerationStreamOutput(
-            index=0,
-            token=TextGenerationStreamOutputToken(
-                id=0,
-                logprob=0,
-                text=next_text,
-                special=False,
-            )
-        )
+    outputs = tokenizer.batch_decode(output_ids[:, input_token_len:], skip_special_tokens=True)[0]
+    outputs = outputs.strip()
+
+    for i in range(len(outputs)):
+        time.sleep(0.05)
+        yield outputs[: i + 1]
+
 
 
 chatbot = gr.Chatbot(height=450)