Update app.py

app.py

@@ -24,7 +24,6 @@ tokenizer, llava_model, image_processor, context_len = load_pretrained_model(
     device=device
 )
 
-
 @spaces.GPU
 def bot_streaming(message, history):
     print(message)
@@ -53,20 +52,39 @@ def bot_streaming(message, history):
     # Generate the prompt for the model
     prompt = message['text']
     
-    # Call the chat_llava function to generate the output
-    output = chat_llava(
-        args=None,
-        image_file=image,
-        text=prompt,
-        tokenizer=tokenizer,
-        model=llava_model,
-        image_processor=image_processor,
-        context_len=context_len
-    )
+    # Use a streamer to generate the output in a streaming fashion
+    streamer = []
+
+    # Define a function to call chat_llava in a separate thread
+    def generate_output():
+        output = chat_llava(
+            args=None,
+            image_file=image,
+            text=prompt,
+            tokenizer=tokenizer,
+            model=llava_model,
+            image_processor=image_processor,
+            context_len=context_len
+        )
+        for new_text in output:
+            streamer.append(new_text)
+    
+    # Start the generation in a separate thread
+    thread = Thread(target=generate_output)
+    thread.start()
     
     # Stream the output
     buffer = ""
-    for new_text in output:
+    while thread.is_alive() or streamer:
+        while streamer:
+            new_text = streamer.pop(0)
+            buffer += new_text
+            yield buffer
+        time.sleep(0.1)
+    
+    # Ensure any remaining text is yielded after the thread completes
+    while streamer:
+        new_text = streamer.pop(0)
         buffer += new_text
         yield buffer