Update app.py

app.py

@@ -1,6 +1,7 @@
 import gradio as gr
 from transformers import AutoProcessor, Idefics3ForConditionalGeneration, TextIteratorStreamer
 from threading import Thread
+from queue import Queue
 import re
 import time
 from PIL import Image
@@ -24,13 +25,13 @@ def model_inference(
         input_dict, history, decoding_strategy, temperature, max_new_tokens,
         repetition_penalty, top_p
 ):
-\
+    print(input_dict)
     text = input_dict["text"]
    
     if len(input_dict["files"]) > 1:
-        images = [Image.open(image).convert("RGB") for image in input_dict["files"]]
+        images = [Image.open(image).convert("RGB") for image in input_dict["files"]["path"]]
     elif len(input_dict["files"]) == 1:
-        images = [Image.open(input_dict["files"][0]).convert("RGB")]
+        images = [Image.open(input_dict["files"][0]["path"]).convert("RGB")]
     else:
         images = []
 
@@ -78,8 +79,9 @@ def model_inference(
     # start token id = argmax last logit
     start_token_id = int(np.argmax(prefill_out["logits"][:, -1, :], axis=-1)[0])
 
-
-    generated_text = ""
+    # 创建输出队列用于线程间通信
+    output_queue = Queue()
+    
     generation_args = {
         "llm_session" : llm_session,
         "embed_tokens_session": embed_tokens_session,
@@ -89,14 +91,32 @@ def model_inference(
         "freqs_cos": freqs_cos,
         "freqs_sin": freqs_sin,
         "attention_mask": attention_mask.numpy(),
-        "max_new_tokens": 128,
+        "max_new_tokens": max_new_tokens,
         "eos_token_id": 2,
-        "start_pos": seqlen
+        "start_pos": seqlen,
+        "output_queue": output_queue
     }
-        
     
+    # 在后台线程启动生成
     thread = Thread(target=generate_autoregressive, kwargs=generation_args)
     thread.start()
+    
+    # 从队列中读取生成的文本并 yield
+    yield "..."
+    buffer = ""
+    
+    while True:
+        text_chunk = output_queue.get()  # 阻塞等待队列中的数据
+        
+        if text_chunk is None:  # 生成完成信号
+            break
+        
+        buffer += text_chunk
+        time.sleep(0.01)
+        yield buffer
+    
+    # 等待线程完成
+    thread.join()
 
 
 examples = [