Update app.py

app.py

@@ -1,18 +1,31 @@
 import torch
-from transformers import AutoModelForCausalLM
-
-from deepseek_vl.models import VLChatProcessor, MultiModalityCausalLM
+from transformers import AutoModelForCausalLM, BitsAndBytesConfig
+from deepseek_vl.models import VLChatProcessor
 from deepseek_vl.utils.io import load_pil_images
 
-
-# specify the path to the model
+# 模型路徑
 model_path = "deepseek-ai/deepseek-vl-7b-chat"
+
+# 讀取 processor
 vl_chat_processor: VLChatProcessor = VLChatProcessor.from_pretrained(model_path)
 tokenizer = vl_chat_processor.tokenizer
 
-vl_gpt: MultiModalityCausalLM = AutoModelForCausalLM.from_pretrained(model_path, trust_remote_code=True)
-vl_gpt = vl_gpt.to(torch.bfloat16).cuda().eval()
+# ==== 量化模型設定 ====
+bnb_config = BitsAndBytesConfig(
+    load_in_4bit=True,                 # 4-bit 量化
+    bnb_4bit_compute_dtype=torch.float16,
+    bnb_4bit_use_double_quant=True
+)
+
+vl_gpt: AutoModelForCausalLM = AutoModelForCausalLM.from_pretrained(
+    model_path,
+    trust_remote_code=True,
+    device_map="auto",
+    quantization_config=bnb_config
+)
+vl_gpt.eval()
 
+# 範例對話
 conversation = [
     {
         "role": "User",
@@ -25,28 +38,33 @@ conversation = [
     }
 ]
 
-# load images and prepare for inputs
-pil_images = load_pil_images(conversation)
-prepare_inputs = vl_chat_processor(
-    conversations=conversation,
-    images=pil_images,
-    force_batchify=True
-).to(vl_gpt.device)
-
-# run image encoder to get the image embeddings
-inputs_embeds = vl_gpt.prepare_inputs_embeds(**prepare_inputs)
-
-# run the model to get the response
-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,
-    do_sample=False,
-    use_cache=True
-)
+# ==== 逐張圖片處理，降低 VRAM 使用 ====
+answers = []
+for conv in conversation:
+    pil_images = load_pil_images([conv])
+    prepare_inputs = vl_chat_processor(
+        conversations=[conv],
+        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
+        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}")
 
-answer = tokenizer.decode(outputs[0].cpu().tolist(), skip_special_tokens=True)
-print(f"{prepare_inputs['sft_format'][0]}", answer)
+# 輸出結果
+for ans in answers:
+    print(ans)