Update app.py

app.py

@@ -55,13 +55,20 @@ def bot_streaming(message, history, max_new_tokens=250):
     else:
         inputs = processor(text=texts, images=images, return_tensors="pt").to("cuda")
         
-    generation_kwargs = dict(inputs, max_new_tokens=max_new_tokens)
-    
-    with torch.no_grad():
-        output = model.generate(**inputs, **generation_kwargs)
-    output_text = processor.decode(output[0][inputs['input_ids'].shape[1]:]).replace('<|eot_id|>', '')
+    streamer = TextIteratorStreamer(processor, skip_special_tokens=True, skip_prompt=True)
 
-    yield output_text
+    generation_kwargs = dict(inputs, streamer=streamer, max_new_tokens=max_new_tokens)
+    generated_text = ""
+    
+    thread = Thread(target=model.generate, kwargs=generation_kwargs)
+    thread.start()
+    buffer = ""
+    
+    for new_text in streamer:
+        buffer += new_text
+        generated_text_without_prompt = buffer
+        time.sleep(0.01)
+        yield buffer
 
 
 demo = gr.ChatInterface(fn=bot_streaming, title="LLaVA-CoT",