Update app.py

app.py

@@ -1,4 +1,5 @@
 import torch
+import gradio as gr
 from transformers import AutoModelForCausalLM, BitsAndBytesConfig
 from deepseek_vl.models import VLChatProcessor
 from deepseek_vl.utils.io import load_pil_images
@@ -6,13 +7,13 @@ from deepseek_vl.utils.io import load_pil_images
 # 模型路徑
 model_path = "deepseek-ai/deepseek-vl-7b-chat"
 
-# 讀取 processor
+# 載入 processor 和 tokenizer
 vl_chat_processor: VLChatProcessor = VLChatProcessor.from_pretrained(model_path)
 tokenizer = vl_chat_processor.tokenizer
 
-# ==== 量化模型設定 ====
+# ==== 量化模型設定 (4-bit) ====
 bnb_config = BitsAndBytesConfig(
-    load_in_4bit=True,                 # 4-bit 量化
+    load_in_4bit=True,
     bnb_4bit_compute_dtype=torch.float16,
     bnb_4bit_use_double_quant=True
 )
@@ -25,46 +26,44 @@ vl_gpt: AutoModelForCausalLM = AutoModelForCausalLM.from_pretrained(
 )
 vl_gpt.eval()
 
-# 範例對話
-conversation = [
-    {
-        "role": "User",
-        "content": "<image_placeholder>Describe each stage of this image.",
-        "images": ["./images/training_pipelines.png"]
-    },
-    {
-        "role": "Assistant",
-        "content": ""
-    }
-]
+# ==== 單張圖片處理 + 減少 max_new_tokens ====
+def generate_answer(image, text):
+    conversation = [
+        {"role": "User", "content": "<image_placeholder>" + text, "images": [image]},
+        {"role": "Assistant", "content": ""}
+    ]
 
-# ==== 逐張圖片處理，降低 VRAM 使用 ====
-answers = []
-for conv in conversation:
-    pil_images = load_pil_images([conv])
+    pil_images = load_pil_images(conversation)
     prepare_inputs = vl_chat_processor(
-        conversations=[conv],
+        conversations=conversation,
         images=pil_images,
         force_batchify=True
     ).to(vl_gpt.device)
 
     inputs_embeds = vl_gpt.prepare_inputs_embeds(**prepare_inputs)
 
-    # 減少生成長度 max_new_tokens
     outputs = vl_gpt.language_model.generate(
         inputs_embeds=inputs_embeds,
         attention_mask=prepare_inputs.attention_mask,
         pad_token_id=tokenizer.eos_token_id,
         bos_token_id=tokenizer.bos_token_id,
         eos_token_id=tokenizer.eos_token_id,
-        max_new_tokens=128,   # 原本 512 → 128
+        max_new_tokens=128,  # 降低生成長度
         do_sample=False,
         use_cache=True
     )
 
     answer = tokenizer.decode(outputs[0].cpu().tolist(), skip_special_tokens=True)
-    answers.append(f"{prepare_inputs['sft_format'][0]} {answer}")
+    return f"{prepare_inputs['sft_format'][0]} {answer}"
 
-# 輸出結果
-for ans in answers:
-    print(ans)
+# ==== Gradio Web UI ====
+demo = gr.Interface(
+    fn=generate_answer,
+    inputs=[gr.Image(type="pil", label="Upload Image"), gr.Textbox(label="Question")],
+    outputs="text",
+    title="DeepSeek-VL-7B Chat Demo",
+    description="上傳圖片並輸入問題，模型會生成與圖片相關的回答（4-bit 量化，低記憶體模式）"
+)
+
+if __name__ == "__main__":
+    demo.launch()