Update app.py

app.py

@@ -3,7 +3,8 @@ import torch
 from PIL import Image
 import gradio as gr
 import spaces
-from transformers import AutoModelForCausalLM, AutoTokenizer, AutoConfig, TextIteratorStreamer
+from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
+from huggingface_hub.inference._generated.types import TextGenerationStreamOutput, TextGenerationStreamOutputToken
 import os
 from huggingface_hub import hf_hub_download
 
@@ -93,13 +94,21 @@ def stream_chat(message, history: list, system: str, temperature: float, max_new
     if message["files"]:
         image = Image.open(message["files"][0]).convert('RGB')
         # Process the conversation text
-        inputs = model.build_conversation_input_ids(tokenizer, query=message['text'], image=image, image_processor=image_processor)
+        inputs = model.build_conversation_input_ids(
+            tokenizer, 
+            query=message['text'], 
+            image=image, 
+            image_processor=image_processor,
+        )
         input_ids = inputs["input_ids"].to(device='cuda', non_blocking=True)
         images = inputs["image"].to(dtype=torch.float16, device='cuda', non_blocking=True)
     else:
-        input_ids = tokenizer.apply_chat_template(conversation, add_generation_prompt=True, return_tensors="pt").to(model.device)
+        input_ids = tokenizer.apply_chat_template(
+            conversation, 
+            add_generation_prompt=True, 
+            return_tensors="pt"
+        ).to(model.device)
         images = None
-
     streamer = TextIteratorStreamer(tokenizer, timeout=10.0, skip_prompt=True, skip_special_tokens=True)
 
     generate_kwargs = dict(
@@ -116,10 +125,19 @@ def stream_chat(message, history: list, system: str, temperature: float, max_new
 
     t = Thread(target=model.generate, kwargs=generate_kwargs)
     t.start()
+    input_token_len = input_ids.shape[1]
+    
     output = ""
     for new_token in streamer:
-        output += new_token
-        yield output
+        yield TextGenerationStreamOutput(
+            index=0,
+            token=TextGenerationStreamOutputToken(
+                id=0,
+                logprob=0,
+                text=next_text,
+                special=False,
+            )
+        )
 
 
 chatbot = gr.Chatbot(height=450)