# app.py (IP-Adapter V2アップグレード版)

import gradio as gr
import torch
import numpy as np
import cv2
from PIL import Image
import time
import spaces
import traceback
import os

from diffusers import StableDiffusionPipeline, ControlNetModel
from insightface.app import FaceAnalysis

print("--- アプリケーションの初期化を開始 ---")


try:
    print("モデルのロードを開始...")
    is_cuda_available = torch.cuda.is_available()
    device = "cuda" if is_cuda_available else "cpu"
    dtype = torch.float16 if is_cuda_available else torch.float32
    print(f"実行デバイス: {device}, データ型: {dtype}")

    print("InsightFaceのプロバイダーとコンテキストIDを設定...")
    providers = ['CUDAExecutionProvider', 'CPUExecutionProvider'] if is_cuda_available else ['CPUExecutionProvider']
    ctx_id = 0 if is_cuda_available else -1
    print(f"使用するプロバイダー: {providers}, ctx_id: {ctx_id}")
    
    face_app = FaceAnalysis(name='buffalo_l', root='./', providers=providers)
    face_app.prepare(ctx_id=ctx_id, det_size=(640, 640))
    print("InsightFaceのロード完了。")
    
    source_repo = "InstantX/InstantID"
    
    print(f"ControlNetモデル ({source_repo}) のロード開始...")
    controlnet = ControlNetModel.from_pretrained(
        source_repo, 
        subfolder="ControlNetModel",
        torch_dtype=dtype
    )
    print("ControlNetモデルのロード完了。")
    
    civitai_model_url = "https://civitai.com/api/download/models/90505?type=Model&format=SafeTensor&size=full&fp=fp32"
    local_model_path = "civitai_model.safetensors"

    if not os.path.exists(local_model_path):
        print(f"{local_model_path} が存在しないため、ダウンロードを開始します...")
        os.system(f'wget -q -O {local_model_path} "{civitai_model_url}"')
    else:
        print(f"{local_model_path} は既に存在するため、ダウンロードをスキップします。")

    print(f"ローカルパス ({local_model_path}) からパイプラインをロード開始...")
    pipe = StableDiffusionPipeline.from_single_file(
        local_model_path,
        torch_dtype=dtype,
        safety_checker=None,
        use_safetensors=True,
    )
    pipe.controlnet = controlnet
    pipe.to(device)
    print("Stable Diffusionパイプラインのロード完了。")

    # ★★★★★ ここが最後の追加注文です ★★★★★
    # weight_nameをV2モデルに変更します。
    ip_adapter_repo = "h94/IP-Adapter"
    print(f"IP-Adapter ({ip_adapter_repo}) のロード開始...")
    pipe.load_ip_adapter(
        ip_adapter_repo,
        subfolder="models",
        weight_name="ip-adapter-plus_v2_sd15.bin" # V2の重みファイル名を指定
    )
    print("IP-Adapter V2のロード完了。")
    
    print("★★全てのモデルのロードが正常に完了しました★★")
    MODELS_LOADED = True
except Exception:
    error_details = traceback.format_exc()
    print("モデルのロード中に詳細なエラーが発生しました:")
    print(error_details)
    MODELS_LOADED = False

# --- 2. 画像生成関数 (変更なし) ---
@spaces.GPU(duration=300)
def generate_image(
    face_image, prompt, negative_prompt, guidance_scale, ip_adapter_scale, num_steps,
    progress=gr.Progress(track_tqdm=True)
):
    if not MODELS_LOADED:
        raise gr.Error("モデルがロードされていないため、画像を生成できません。ログを確認してください。")
    if face_image is None:
        raise gr.Error("顔画像をアップロードしてください。")
    if not prompt:
        raise gr.Error("プロンプトを入力してください。")

    face_image = Image.fromarray(face_image)
    face_info = face_app.get(cv2.cvtColor(np.array(face_image), cv2.COLOR_RGB2BGR))
    if not face_info:
        raise gr.Error("アップロードされた画像から顔を検出できませんでした。")
    
    face_info = sorted(face_info, key=lambda x: (x['bbox'][2] - x['bbox'][0]) * (x['bbox'][3] - x['bbox'][1]))[-1]
    face_emb = face_info['embedding']
    pipe.set_ip_adapter_scale(ip_adapter_scale)
    images = pipe(
        prompt=prompt, negative_prompt=negative_prompt, image_embeds=[face_emb], image=face_image,
        controlnet_conditioning_scale=ip_adapter_scale, num_inference_steps=int(num_steps), guidance_scale=guidance_scale,
    ).images
    return images[0]

# --- 3. GradioのUIとAPIの定義 (変更なし) ---
with gr.Blocks() as demo:
    gr.Markdown("# InstantID 画像生成アプリ")
    with gr.Row():
        with gr.Column():
            face_image_input = gr.Image(label="顔写真", type="numpy")
            prompt_input = gr.Textbox(label="プロンプト (例: 1girl, a photo of a cute girl in a suit)")
            negative_prompt_input = gr.Textbox(label="ネガティブプロンプト", value="(lowres, low quality, worst quality:1.2), ugly")
            with gr.Accordion("詳細設定", open=False):
                ip_adapter_scale_slider = gr.Slider(minimum=0, maximum=1.5, step=0.1, value=0.8, label="顔の忠実度 (IP Adapter Scale)")
                guidance_scale_slider = gr.Slider(minimum=1, maximum=10, step=0.5, value=5.0, label="プロンプトへの忠実度 (Guidance Scale)")
                num_steps_slider = gr.Slider(minimum=10, maximum=50, step=1, value=30, label="生成ステップ数 (Steps)")
            generate_button = gr.Button("画像を生成", variant="primary")
        with gr.Column():
            output_image = gr.Image(label="生成結果")
    generate_button.click(
        fn=generate_image,
        inputs=[
            face_image_input, prompt_input, negative_prompt_input, guidance_scale_slider,
            ip_adapter_scale_slider, num_steps_slider
        ],
        outputs=[output_image],
        api_name="generate"
    )

demo.queue().launch()
print("--- Gradioアプリの起動準備完了 ---")