Update app.py

app.py

@@ -9,7 +9,7 @@ import torch
 from PIL import Image
 import gradio as gr
 import spaces
-from transformers import AutoModel, AutoTokenizer, TextIteratorStreamer
+from transformers import AutoModel, AutoTokenizer
 import os
 
 
@@ -35,9 +35,8 @@ CSS = """
 
 model = AutoModel.from_pretrained(
     MODEL_ID,
-    torch_dtype=torch.float16,
     trust_remote_code=True
-).to(0)
+)
 tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, trust_remote_code=True)
 model.eval()
 
@@ -71,12 +70,18 @@ def stream_chat(message, history: list, temperature: float, max_new_tokens: int)
         temperature=temperature,
         sampling=True,
         tokenizer=tokenizer,
+        stream=True
     )
     if temperature == 0:
         generate_kwargs["sampling"] = False
 
     response = model.chat(**generate_kwargs)
-    return response
+
+    generated_text = ""
+    for new_text in response:
+        generated_text += new_text
+        yeild(new_text, flush=True, end='')
+
 
 
 chatbot = gr.Chatbot(height=450)