Update app.py

app.py

@@ -40,8 +40,10 @@ vlm.to(DEVICE, dtype=DTYPE)
 
 prompt_builder = vlm.get_prompt_builder()
 
-@spaces.GPU
+@spaces.GPU(duration=20)
 def bot_streaming(message, history, temperature, top_k, max_new_tokens):
+    streamer = TextIteratorStreamer(processor.tokenizer, skip_special_tokens=True)
+    
     if len(history) == 0:
         prompt_builder.prompt, prompt_builder.turn_count = "", 0
 
@@ -61,22 +63,36 @@ def bot_streaming(message, history, temperature, top_k, max_new_tokens):
     
     prompt_builder.add_turn(role="human", message=message['text'])
     prompt_text = prompt_builder.get_prompt()
+
+    generation_kwargs = {
+        "pixel_values": image,
+        "input_ids": prompt_text,
+        "streamer": streamer,
+        "max_new_tokens": max_new_tokens,
+        "use_cache": True,
+        "temperature": temperature,
+        "do_sample": True,
+        "top_k": top_k,
+    }
     
-    # Generate from the VLM
-    with torch.no_grad():
-        generated_text = vlm.generate(
-            image,
-            prompt_text,
-            use_cache=True,
-            do_sample=True,
-            temperature=temperature,
-            top_k=top_k,
-            max_new_tokens=max_new_tokens,
-        )
-    prompt_builder.add_turn(role="gpt", message=generated_text)
+    # Generate from the VLM  
+    thread = Thread(target=vlm.generate, kwargs=generation_kwargs)
+    thread.start()
 
-    time.sleep(0.04)
-    yield generated_text
+    buffer = ""
+    output_started = False
+    for new_text in streamer:
+        if not output_started:
+            if "<|assistant|>\n" in new_text:
+                output_started = True
+            continue
+        buffer += new_text
+        if len(buffer) > 1:
+            yield buffer
+
+    prompt_builder.add_turn(role="gpt", message=buffer)
+    
+    return buffer
     
     
 demo = gr.ChatInterface(fn=bot_streaming,