Update app.py

app.py

@@ -60,43 +60,50 @@ def bot_streaming(message, history):
     
     # Generate the prompt for the model
     prompt = message['text']
+    streamer = TextStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
     
-    # 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=args,
-            image_file=image_path,
-            text=prompt,
-            tokenizer=tokenizer,
-            model=llava_model,
-            image_processor=image_processor,
-            context_len=context_len
-        )
-        for new_text in output:
-            streamer.append(new_text)
+    # Set up the generation arguments, including the streamer
+    generation_kwargs = dict(
+        args=args,
+        image_file=image_path,
+        text=prompt,
+        tokenizer=tokenizer,
+        model=llava_model,
+        streamer=streamer
+        image_processor=image_processor,  # todo: input model name or path
+        context_len=context_len)
+    
+    # Define the function to call `chat_llava` with the given arguments
+    def generate_output(generation_kwargs):
+        chat_llava(**generation_kwargs)
     
     # Start the generation in a separate thread
-    thread = Thread(target=generate_output)
+    thread = Thread(target=generate_output, kwargs=generation_kwargs)
     thread.start()
     
-    # Stream the output
+    # Initialize a buffer to accumulate the generated text
     buffer = ""
-    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)
+    # Allow the generation to start
+    time.sleep(0.5)
+    
+    # Iterate over the streamer to handle the incoming text in chunks
+    for new_text in streamer:
+        # Look for the end of text token and remove it
+        if "<|eot_id|>" in new_text:
+            new_text = new_text.split("<|eot_id|>")[0]
+        
+        # Add the new text to the buffer
         buffer += new_text
-        yield buffer
-
+    
+        # Remove the prompt from the generated text (if necessary)
+        generated_text_without_prompt = buffer[len(prompt):]
+    
+        # Simulate processing time (optional)
+        time.sleep(0.06)
+    
+        # Yield the current generated text for further processing or display
+        yield generated_text_without_prompt
 
 chatbot = gr.Chatbot(scale=1)
 chat_input = gr.MultimodalTextbox(interactive=True, file_types=["image"], placeholder="Enter message or upload file...", show_label=False)