Update app.py

app.py

@@ -264,6 +264,73 @@ def generate_video(model_name: str, text: str, video_path: str,
         yield buffer, buffer
 
 
+@spaces.GPU
+# def generate_pdf(model_name: str, text: str, state: Dict[str, Any], 
+#                  max_new_tokens: int = 2048, 
+#                  temperature: float = 0.6, 
+#                  top_p: float = 0.9, 
+#                  top_k: int = 50, 
+#                  repetition_penalty: float = 1.2):
+    
+#     # if model_name == "Qwen2.5-VL-7B-Instruct":
+#     #     processor, model = processor_m, model_m
+#     # elif model_name == "Qwen2.5-VL-3B-Instruct":
+#     #     processor, model = processor_x, model_x
+#     if model_name == "Qwen3-VL-4B-Instruct":
+#         processor, model = processor_q, model_q
+#     elif model_name == "Qwen3-VL-8B-Instruct":
+#         processor, model = processor_y, model_y
+#     # elif model_name == "Qwen3-VL-8B-Thinking":
+#     #     processor, model = processor_z, model_z
+#     elif model_name == "Qwen3-VL-4B-Thinking":
+#         processor, model = processor_t, model_t
+#     elif model_name == "Qwen3-VL-2B-Instruct":
+#         processor, model = processor_l, model_l
+#     elif model_name == "Qwen3-VL-2B-Thinking":
+#         processor, model = processor_j, model_j
+#     else:
+#         yield "Invalid model selected.", "Invalid model selected."
+#         return
+
+#     if not state or not state["pages"]:
+#         yield "Please upload a PDF file first.", "Please upload a PDF file first."
+#         return
+    
+#     page_images = state["pages"]
+#     full_response = ""
+#     for i, image in enumerate(page_images):
+#         page_header = f"--- Page {i+1}/{len(page_images)} ---\n"
+#         yield full_response + page_header, full_response + page_header
+        
+#         messages = [{"role": "user", "content": [{"type": "image"}, {"type": "text", "text": text}]}]
+#         # Sử dụng processor đã chọn
+#         prompt_full = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
+#         inputs = processor(text=[prompt_full], images=[image], return_tensors="pt", padding=True).to(device)
+#         streamer = TextIteratorStreamer(processor, skip_prompt=True, skip_special_tokens=True)
+        
+#         generation_kwargs = {
+#             **inputs, 
+#             "streamer": streamer, 
+#             "max_new_tokens": max_new_tokens,
+#             # "do_sample": True, 
+#             # "temperature": temperature, 
+#             # "top_p": top_p, 
+#             # "top_k": top_k, 
+#             # "repetition_penalty": repetition_penalty
+#         }
+        
+#         # Sử dụng model đã chọn
+#         thread = Thread(target=model.generate, kwargs=generation_kwargs)
+#         thread.start()
+        
+#         page_buffer = ""
+#         for new_text in streamer:
+#             page_buffer += new_text
+#             yield full_response + page_header + page_buffer, full_response + page_header + page_buffer
+#             time.sleep(0.01)
+        
+#         full_response += page_header + page_buffer + "\n\n"
+
 @spaces.GPU
 def generate_pdf(model_name: str, text: str, state: Dict[str, Any], 
                  max_new_tokens: int = 2048, 
@@ -272,16 +339,10 @@ def generate_pdf(model_name: str, text: str, state: Dict[str, Any],
                  top_k: int = 50, 
                  repetition_penalty: float = 1.2):
     
-    # if model_name == "Qwen2.5-VL-7B-Instruct":
-    #     processor, model = processor_m, model_m
-    # elif model_name == "Qwen2.5-VL-3B-Instruct":
-    #     processor, model = processor_x, model_x
     if model_name == "Qwen3-VL-4B-Instruct":
         processor, model = processor_q, model_q
     elif model_name == "Qwen3-VL-8B-Instruct":
         processor, model = processor_y, model_y
-    # elif model_name == "Qwen3-VL-8B-Thinking":
-    #     processor, model = processor_z, model_z
     elif model_name == "Qwen3-VL-4B-Thinking":
         processor, model = processor_t, model_t
     elif model_name == "Qwen3-VL-2B-Instruct":
@@ -297,39 +358,44 @@ def generate_pdf(model_name: str, text: str, state: Dict[str, Any],
         return
     
     page_images = state["pages"]
-    full_response = ""
-    for i, image in enumerate(page_images):
-        page_header = f"--- Page {i+1}/{len(page_images)} ---\n"
-        yield full_response + page_header, full_response + page_header
-        
-        messages = [{"role": "user", "content": [{"type": "image"}, {"type": "text", "text": text}]}]
-        # Sử dụng processor đã chọn
-        prompt_full = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
-        inputs = processor(text=[prompt_full], images=[image], return_tensors="pt", padding=True).to(device)
-        streamer = TextIteratorStreamer(processor, skip_prompt=True, skip_special_tokens=True)
-        
-        generation_kwargs = {
-            **inputs, 
-            "streamer": streamer, 
-            "max_new_tokens": max_new_tokens,
-            # "do_sample": True, 
-            # "temperature": temperature, 
-            # "top_p": top_p, 
-            # "top_k": top_k, 
-            # "repetition_penalty": repetition_penalty
-        }
-        
-        # Sử dụng model đã chọn
-        thread = Thread(target=model.generate, kwargs=generation_kwargs)
-        thread.start()
-        
-        page_buffer = ""
-        for new_text in streamer:
-            page_buffer += new_text
-            yield full_response + page_header + page_buffer, full_response + page_header + page_buffer
-            time.sleep(0.01)
-        
-        full_response += page_header + page_buffer + "\n\n"
+
+    messages = [{"role": "user", "content": [{"type": "text", "text": text}]}]
+    images_for_processor = []
+    for frame in page_images:
+        messages[0]["content"].append({"type": "image"})
+        images_for_processor.append(frame)
+
+    prompt_full = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
+    
+    inputs = processor(
+        text=[prompt_full], 
+        images=images_for_processor,  # Truyền cả list ảnh
+        return_tensors="pt", 
+        padding=True
+    ).to(device)
+    
+    streamer = TextIteratorStreamer(processor, skip_prompt=True, skip_special_tokens=True)
+    
+    generation_kwargs = {
+        **inputs, 
+        "streamer": streamer, 
+        "max_new_tokens": max_new_tokens,
+        "do_sample": True, 
+        "temperature": temperature, 
+        "top_p": top_p, 
+        "top_k": top_k, 
+        "repetition_penalty": repetition_penalty
+    }
+    
+    thread = Thread(target=model.generate, kwargs=generation_kwargs)
+    thread.start()
+    
+    buffer = ""
+    for new_text in streamer:
+        buffer += new_text
+        buffer = buffer.replace("<|im_end|>", "") # Thêm dòng này giống video
+        yield buffer, buffer
+        time.sleep(0.01)
 
 image_examples = [
     ["Explain the content in detail.", "images/force.jpg"],