Update app.py

app.py

@@ -1,79 +1,92 @@
 import torch
 import gradio as gr
 from transformers import AutoModelForCausalLM, BitsAndBytesConfig
-from deepseek_vl.models import VLChatProcessor
+from deepseek_vl.models import VLChatProcessor, MultiModalityCausalLM
+from deepseek_vl.utils.io import load_pil_images
 
 # 模型路徑
 model_path = "deepseek-ai/deepseek-vl-7b-chat"
 
-# 載入 processor 和 tokenizer
-vl_chat_processor: VLChatProcessor = VLChatProcessor.from_pretrained(model_path)
-tokenizer = vl_chat_processor.tokenizer
-
-# ==== 量化模型設定 (4-bit) ====
+# ==== BitsAndBytes 4-bit 量化設定 ====
 bnb_config = BitsAndBytesConfig(
     load_in_4bit=True,
-    bnb_4bit_compute_dtype=torch.float16,
+    bnb_4bit_compute_dtype=torch.float16,  # 強制 float16
     bnb_4bit_use_double_quant=True
 )
 
-vl_gpt: AutoModelForCausalLM = AutoModelForCausalLM.from_pretrained(
+# 載入 processor 和 tokenizer
+vl_chat_processor: VLChatProcessor = VLChatProcessor.from_pretrained(model_path)
+tokenizer = vl_chat_processor.tokenizer
+
+# 載入模型 (4-bit 量化 + float16)
+vl_gpt: MultiModalityCausalLM = AutoModelForCausalLM.from_pretrained(
     model_path,
-    trust_remote_code=True,
+    quantization_config=bnb_config,
     device_map="auto",
-    quantization_config=bnb_config
-)
-vl_gpt.eval()
+    trust_remote_code=True
+).eval()
 
-# ==== 單張圖片處理 + 減少 max_new_tokens ====
-def generate_answer(image, text):
+# ==== 單張圖片推理函式 ====
+def chat_with_image(image, user_message):
     try:
-        # 將圖片與文字組合成對話格式
+        # 建立對話格式
         conversation = [
-            {"role": "User", "content": "<image_placeholder>" + text, "images": [image]},
+            {"role": "User", "content": "<image_placeholder>" + user_message, "images": [image]},
             {"role": "Assistant", "content": ""}
         ]
 
-        # 🚨 這裡直接傳入 [image]，因為已經是 PIL.Image，不需要 load_pil_images
+        # 輸入處理
+        pil_images = load_pil_images(conversation)
         prepare_inputs = vl_chat_processor(
             conversations=conversation,
-            images=[image],
+            images=pil_images,
             force_batchify=True
         ).to(vl_gpt.device)
 
-        prepare_inputs = dict(prepare_inputs)
-        prepare_inputs = {k: (v.to(torch.float16) if torch.is_tensor(v) else v) 
-                  for k, v in prepare_inputs.items()}
-        
-        # 轉換成 embeddings
+        # 🚨 正確 dtype 處理
+        # 只將需要的 tensor 轉 float16，input_ids 必須是 long
+        new_inputs = {}
+        for k, v in prepare_inputs.items():
+            if torch.is_tensor(v):
+                if k in ["input_ids", "labels"]:
+                    new_inputs[k] = v.to(torch.long)
+                else:
+                    new_inputs[k] = v.to(torch.float16)
+            else:
+                new_inputs[k] = v
+        prepare_inputs = new_inputs
+
+        # 取得輸入 embeddings
         inputs_embeds = vl_gpt.prepare_inputs_embeds(**prepare_inputs)
 
-        # 產生回答
+        # 生成回答
         outputs = vl_gpt.language_model.generate(
             inputs_embeds=inputs_embeds,
-            attention_mask=prepare_inputs.attention_mask,
+            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,   # 降低生成長度以減少記憶體
+            max_new_tokens=128,  # 降低生成長度以減少記憶體
             do_sample=False,
             use_cache=True
         )
 
+        # 解碼
         answer = tokenizer.decode(outputs[0].cpu().tolist(), skip_special_tokens=True)
-        return f"{prepare_inputs['sft_format'][0]} {answer}"
+        return answer
 
     except Exception as e:
         return f"Error: {str(e)}"
 
 # ==== Gradio Web UI ====
 demo = gr.Interface(
-    fn=generate_answer,
-    inputs=[gr.Image(type="pil", label="Upload Image"), gr.Textbox(label="Question")],
+    fn=chat_with_image,
+    inputs=[gr.Image(type="pil", label="Upload Image"),
+            gr.Textbox(lines=2, placeholder="Ask about the image...")],
     outputs="text",
-    title="DeepSeek-VL-7B Chat Demo",
-    description="上傳圖片並輸入問題，模型會生成與圖片相關的回答（4-bit 量化，低記憶體模式）"
+    title="DeepSeek-VL-7B-Chat Demo (4-bit, float16)",
+    description="上傳圖片並輸入問題，模型會生成與圖片相關的回答"
 )
 
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch()