# custom_inference.py

import torch
from transformers import Qwen2VLForConditionalGeneration, AutoProcessor
from PIL import Image
from qwen_vl_utils import process_vision_info  # 必要に応じてインポートを調整
import os

def load_merged_model(merged_model_path):
    """
    マージ済みモデルとプロセッサのロード
    """
    print("マージ済みモデルをロード中...")
    model = Qwen2VLForConditionalGeneration.from_pretrained(
        merged_model_path, torch_dtype=torch.float16, device_map="auto"
    )
    processor = AutoProcessor.from_pretrained(merged_model_path)
    print("マージ済みモデルのロード完了.")
    return model, processor

def perform_inference(model, processor, image_path, prompt):
    """
    推論の実行
    """
    messages = [
        {
            "role": "user",
            "content": [
                {
                    "type": "image",
                    "image": image_path,
                },
                {"type": "text", "text": prompt},
            ],
        }
    ]

    # テキストの準備
    text = processor.apply_chat_template(
        messages, tokenize=False, add_generation_prompt=True
    )

    # 画像のロードと前処理
    image = Image.open(image_path).convert("RGB")

    # 画像の前処理
    image_inputs, video_inputs = process_vision_info(messages)

    # モデル入力の準備
    inputs = processor(
        text=[text],
        images=image_inputs,
        videos=video_inputs,
        padding=True,
        return_tensors="pt",
    )

    # デバイスへの転送
    device = "cuda" if torch.cuda.is_available() else "cpu"
    inputs = {k: v.to(device) for k, v in inputs.items()}
    model.to(device)

    # 推論
    with torch.no_grad():
        generated_ids = model.generate(**inputs, max_new_tokens=128)

    # 生成されたIDをトリム
    generated_ids_trimmed = [
        out_ids[len(in_ids):] for in_ids, out_ids in zip(inputs["input_ids"], generated_ids)
    ]

    # 結果のデコード
    output_text = processor.batch_decode(
        generated_ids_trimmed, skip_special_tokens=True, clean_up_tokenization_spaces=False
    )

    return output_text

def main():
    # マージ済みモデルのパス
    merged_model_path = "./checkpoint-merged"

    # 画像パスとプロンプトの設定
    image_path = "./images/0.jpg"  # ローカルの画像パス
    prompt = "<image>画像を見てシュールで面白いことを言ってください。空欄がある場合はそれを埋めるように答えてください。"

    # マージ済みモデルとプロセッサのロード
    model, processor = load_merged_model(merged_model_path)

    # 推論の実行
    print("推論を実行中...")
    output = perform_inference(model, processor, image_path, prompt)
    print("生成されたテキスト:", output)

if __name__ == "__main__":
    main()