Update app.py

app.py

@@ -3,7 +3,7 @@
 import gradio as gr
 import torch
 from PIL import Image
-from transformers import AutoModelForImageTextToText, AutoProcessor
+from transformers import AutoModelForImageTextToText, AutoProcessor, BitsAndBytesConfig
 from gradio.events import SelectData
 import warnings
 import os
@@ -11,20 +11,48 @@ import requests
 
 warnings.filterwarnings("ignore", category=UserWarning, message="Overriding torch_dtype=None")
 
-# --- 1. Tải Model và Processor ---
+# --- 1. Tải Model và Processor (ĐÃ TỐI ƯU) ---
 MODEL_ID = "sunbv56/qwen2.5-vl-vqa-vibook"
-print(f"🚀 Đang tải model '{MODEL_ID}' và processor...")
+print(f"🚀 Đang tải model '{MODEL_ID}' và processor với các tối ưu hóa...")
+
+# ### THAY ĐỔI TỐI ƯU 1: Cấu hình Lượng tử hóa 4-bit (Quantization) ###
+# Sử dụng 4-bit quantization để tăng tốc độ inference và giảm VRAM
+quantization_config = BitsAndBytesConfig(
+    load_in_4bit=True,
+    bnb_4bit_quant_type="nf4",
+    bnb_4bit_compute_dtype=torch.bfloat16
+)
+
 try:
-    dtype = torch.bfloat16 if torch.cuda.is_available() and torch.cuda.is_bf16_supported() else torch.float16
-    model = AutoModelForImageTextToText.from_pretrained(MODEL_ID, torch_dtype=dtype, device_map="auto", trust_remote_code=True)
+    # ### THAY ĐỔI TỐI ƯU 2: Tải model với Quantization và Flash Attention 2 ###
+    model = AutoModelForImageTextToText.from_pretrained(
+        MODEL_ID, 
+        device_map="auto", 
+        trust_remote_code=True,
+        quantization_config=quantization_config,
+        # attn_implementation="flash_attention_2" # Bỏ comment dòng này nếu bạn có GPU tương thích (NVIDIA 30xx/40xx) và đã cài flash-attn thành công
+    )
     processor = AutoProcessor.from_pretrained(MODEL_ID, trust_remote_code=True, use_fast=True)
+    
+    # ### THAY ĐỔI TỐI ƯU 3: Biên dịch model với torch.compile ###
+    # Lần chạy đầu tiên sẽ mất chút thời gian để biên dịch, nhưng các lần sau sẽ rất nhanh.
+    # Chỉ hoạt động trên Linux/MacOS với PyTorch 2.0+ và GPU.
+    try:
+        print("🚀 Đang cố gắng biên dịch model với torch.compile()...")
+        model = torch.compile(model, mode="reduce-overhead", fullgraph=True)
+        print("✅ Biên dịch model thành công!")
+    except Exception as e:
+        print(f"⚠️ Không thể biên dịch model: {e}. Chạy ở chế độ thông thường.")
+
     model.eval()
-    print(f"✅ Model và processor đã được tải thành công!")
+    print(f"✅ Model và processor đã được tải và tối ưu thành công!")
+
 except Exception as e:
     print(f"❌ Lỗi khi tải model/processor: {e}")
     exit()
 
-# --- 2. Hàm Inference Cốt lõi ---
+# --- 2. Hàm Inference Cốt lõi (Không cần thay đổi) ---
+# Các tối ưu đã được áp dụng ở tầng model, nên hàm này sẽ tự động chạy nhanh hơn.
 def process_vqa(image: Image.Image, question: str):
     if image.mode != "RGB":
         image = image.convert("RGB")
@@ -32,22 +60,22 @@ def process_vqa(image: Image.Image, question: str):
     prompt_text = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
     model_inputs = processor(text=[prompt_text], images=[image], return_tensors="pt").to(model.device)
     
-    generated_ids = model.generate(
-        **model_inputs,
-        max_new_tokens=128,
-        do_sample=False,
-        temperature=1.0, 
-        eos_token_id=processor.tokenizer.eos_token_id,
-        pad_token_id=processor.tokenizer.pad_token_id
-    )
+    # Sử dụng torch.no_grad() để tắt việc tính toán gradient, giúp tiết kiệm bộ nhớ và tăng tốc độ
+    with torch.no_grad():
+        generated_ids = model.generate(
+            **model_inputs,
+            max_new_tokens=128,
+            do_sample=False,
+            temperature=1.0, 
+            eos_token_id=processor.tokenizer.eos_token_id,
+            pad_token_id=processor.tokenizer.pad_token_id
+        )
     
     generated_ids = generated_ids[:, model_inputs['input_ids'].shape[1]:]
     response = processor.tokenizer.decode(generated_ids[0], skip_special_tokens=True).strip()
     return response
 
-# --- 3. Logic Chatbot ---
-# ### THAY ĐỔI MỚI 1: Định nghĩa HTML và CSS cho hiệu ứng động ###
-# HTML cho hiệu ứng "đang gõ"
+# --- 3. Logic Chatbot (Giữ nguyên) ---
 THINKING_HTML = """
 <div class="typing-indicator">
     <span></span>
@@ -55,7 +83,6 @@ THINKING_HTML = """
     <span></span>
 </div>
 """
-# CSS để tạo hiệu ứng
 CUSTOM_CSS = """
 @keyframes blink {
     0% { opacity: .2; }
@@ -65,14 +92,14 @@ CUSTOM_CSS = """
 .typing-indicator {
     display: flex;
     align-items: center;
-    justify-content: flex-start; /* Căn trái */
-    padding: 8px 0; /* Thêm chút khoảng đệm */
+    justify-content: flex-start;
+    padding: 8px 0;
 }
 .typing-indicator span {
     height: 10px;
     width: 10px;
     margin: 0 2px;
-    background-color: #9E9E9E; /* Màu xám */
+    background-color: #9E9E9E;
     border-radius: 50%;
     animation: blink 1.4s infinite both;
 }
@@ -84,7 +111,6 @@ CUSTOM_CSS = """
 }
 """
 
-# Hàm dành cho việc người dùng tự nhập câu hỏi
 def manual_chat_responder(user_question: str, chat_history: list, uploaded_image: Image.Image):
     if uploaded_image is None:
         gr.Warning("Vui lòng tải ảnh lên trước để đặt câu hỏi về nó.")
@@ -94,7 +120,6 @@ def manual_chat_responder(user_question: str, chat_history: list, uploaded_image
         return "", chat_history
     
     chat_history.append({"role": "user", "content": user_question})
-    # ### THAY ĐỔI MỚI 2: Sử dụng HTML động thay cho text tĩnh ###
     chat_history.append({"role": "assistant", "content": THINKING_HTML})
     yield "", chat_history
 
@@ -103,14 +128,12 @@ def manual_chat_responder(user_question: str, chat_history: list, uploaded_image
     chat_history[-1]["content"] = bot_response
     yield "", chat_history
 
-# Hàm dành riêng cho việc xử lý khi nhấn vào ví dụ
 def run_example(evt: SelectData):
     selected_example = example_list[evt.index]
     image_path, question = selected_example
     gr.Info(f"Đang chạy ví dụ: \"{question}\"")
     image = Image.open(image_path).convert("RGB")
     
-    # ### THAY ĐỔI MỚI 3: Sử dụng HTML động thay cho text tĩnh ###
     chat_history = [
         {"role": "user", "content": question},
         {"role": "assistant", "content": THINKING_HTML}
@@ -125,8 +148,7 @@ def run_example(evt: SelectData):
 def clear_chat():
     return []
 
-# --- 4. Định nghĩa Giao diện Người dùng Gradio ---
-# ### THAY ĐỔI MỚI 4: Thêm CSS vào Blocks ###
+# --- 4. Giao diện Người dùng Gradio (Giữ nguyên) ---
 with gr.Blocks(theme=gr.themes.Soft(primary_hue="blue", secondary_hue="sky"), title="Vibook VQA Chatbot", css=CUSTOM_CSS) as demo:
     gr.Markdown("# 🤖 Vibook VQA Chatbot")
     
@@ -147,15 +169,12 @@ with gr.Blocks(theme=gr.themes.Soft(primary_hue="blue", secondary_hue="sky"), ti
             chatbot = gr.Chatbot(label="Cuộc trò chuyện", height=600, avatar_images=(None, "https://huggingface.co/datasets/huggingface/brand-assets/resolve/main/hf-logo.png"), type="messages", value=[])
             question_input = gr.Textbox(label="Hoặc nhập câu hỏi về ảnh đã tải lên", placeholder="Nhập câu hỏi và nhấn Enter...", container=False, scale=7)
 
-    # --- 5. Xử lý Sự kiện ---
     question_input.submit(fn=manual_chat_responder, inputs=[question_input, chatbot, image_input], outputs=[question_input, chatbot])
-    
     example_dataset.select(fn=run_example, inputs=None, outputs=[image_input, question_input, chatbot], show_progress="full")
-    
     image_input.upload(fn=clear_chat, inputs=None, outputs=[chatbot])
     image_input.clear(fn=clear_chat, inputs=None, outputs=[chatbot])
 
-# --- Phần cuối ---
+# --- Phần cuối (Giữ nguyên) ---
 if __name__ == "__main__":
     ASSETS_DIR = "assets"
     if not os.path.exists(ASSETS_DIR):