app.py

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

# 模型路徑
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) ====
bnb_config = BitsAndBytesConfig(
    load_in_4bit=True,
    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()

# ==== 單張圖片處理 + 減少 max_new_tokens ====
def generate_answer(image, text):
    conversation = [
        {"role": "User", "content": "<image_placeholder>" + text, "images": [image]},
        {"role": "Assistant", "content": ""}
    ]

    pil_images = load_pil_images(conversation)
    prepare_inputs = vl_chat_processor(
        conversations=conversation,
        images=pil_images,
        force_batchify=True
    ).to(vl_gpt.device)

    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,
        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
    )

    answer = tokenizer.decode(outputs[0].cpu().tolist(), skip_special_tokens=True)
    return f"{prepare_inputs['sft_format'][0]} {answer}"

# ==== 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()