# app.py — BRA v7 (AIGaming repo) × InstantID × ZeroGPU
# 2025-06-22

##############################################################################
# 0. diffusers-0.27 互換: cached_download() パッチ
##############################################################################
from huggingface_hub import hf_hub_download
import huggingface_hub as _hf
if not hasattr(_hf, "cached_download"):
    _hf.cached_download = hf_hub_download

##############################################################################
# 1. ライブラリ
##############################################################################
import os, io, base64, subprocess, traceback
from pathlib import Path
from typing import Optional
import numpy as np
import torch, gradio as gr, spaces
from fastapi import FastAPI, UploadFile, File, Form, HTTPException
from PIL import Image
from diffusers import (
    StableDiffusionControlNetPipeline,
    ControlNetModel,
    DPMSolverMultistepScheduler,
)
from diffusers.loaders import AttnProcsLayers
from insightface.app import FaceAnalysis
from realesrgan import RealESRGANer

##############################################################################
# 2. キャッシュパス
##############################################################################
ROOT      = Path("/data") if Path("/data").exists() else Path.home() / ".cache/instantid"
MODELS    = ROOT / "models"; LORA = ROOT / "lora"; UPSCALE = ROOT / "realesrgan"
for p in (MODELS, LORA, UPSCALE): p.mkdir(parents=True, exist_ok=True)

##############################################################################
# 3. モデル ID / ファイル
##############################################################################
# --- BRA v7 (公開) ---
BRA_REPO = "AIGaming/beautiful_realistic_asians"          # :contentReference[oaicite:0]
BRA_FILE = "beautifulRealistic_v7.safetensors"
BRA_REV  = "801a9b1999dd7018e58a1e2b432fdccd3d1d723d"     # 固定 revision

# --- IP-Adapter 本体 & LoRA ---
### 修正 ### 重みファイルへのフルパスを指定し、サブフォルダ定義を削除
IP_REPO,  IP_BIN  = "h94/IP-Adapter", "models/ip-adapter-plus-face_sd15.bin"
LORA_REPO,IP_LORA = "h94/IP-Adapter-FaceID", "ip-adapter-faceid-plusv2_sd15_lora.safetensors"

# --- ControlNet (MediaPipe Face) ---
CN_REPO, CN_SUBF = "CrucibleAI/ControlNetMediaPipeFace", "diffusion_sd15"

# --- Real-ESRGAN ---
ESRGAN_REPO, ESRGAN_FILE = "aimagelab/realesrgan", "RealESRGAN_x4plus.pth"

##############################################################################
# 4. HF Hub ダウンロード
##############################################################################
def dl(repo: str, file: str, sub: str | None = None, rev: str | None = None) -> Path:
    return Path(hf_hub_download(repo, file, subfolder=sub,
                                revision=rev, cache_dir=str(MODELS)))

##############################################################################
# 5. グローバル
##############################################################################
pipe: Optional[StableDiffusionControlNetPipeline] = None
face_analyser: Optional[FaceAnalysis] = None
upsampler: Optional[RealESRGANer] = None

##############################################################################
# 6. 初期化
##############################################################################
def init():
    global pipe, face_analyser, upsampler
    if pipe is not None:
        return
    print("[INIT] downloading models…")

    # 6-1 BRA v7
    bra_ckpt = dl(BRA_REPO, BRA_FILE, rev=BRA_REV)

    # 6-2 ControlNet
    cn = ControlNetModel.from_pretrained(
        CN_REPO, subfolder=CN_SUBF, torch_dtype=torch.float16,
        cache_dir=str(MODELS)
    )

    # 6-3 Pipeline from .safetensors + ControlNet
    pipe_ = StableDiffusionControlNetPipeline.from_single_file(
        bra_ckpt, controlnet=cn, torch_dtype=torch.float16,
        safety_checker=None
    )
    pipe_.scheduler = DPMSolverMultistepScheduler.from_config(pipe_.scheduler.config)

    # 6-4 IP-Adapter
    ### 修正 ### load_ip_adapter がダウンロードを処理するため、個別のdlは不要。
    ###        また、呼び出し方を修正し、リポジトリIDと重みのフルパスを渡す。
    ip_lora = dl(LORA_REPO, IP_LORA)
    pipe_.load_ip_adapter(IP_REPO, weight_name=IP_BIN, cache_dir=str(MODELS))
    AttnProcsLayers(pipe_.unet.attn_processors).load_lora_weights(
        ip_lora, adapter_name="ip_faceid", safe_load=True
    )
    pipe_.set_adapters(["ip_faceid"], adapter_weights=[0.6])
    pipe_.to("cuda"); pipe_ = pipe_

    pipe   = pipe_
    face_analyser = FaceAnalysis(
        name="buffalo_l", root=str(MODELS), providers=["CUDAExecutionProvider"]
    ); face_analyser.prepare(ctx_id=0, det_size=(640,640))

    esr = dl(ESRGAN_REPO, ESRGAN_FILE)
    upsampler = RealESRGANer(scale=4, model_path=str(esr), half=True,
                             tile=512, tile_pad=10, pre_pad=0, gpu_id=0)
    print("[INIT] ready.")

##############################################################################
# 7. プロンプト
##############################################################################
BASE = "(masterpiece:1.2), best quality, ultra-realistic, RAW photo, 8k, cinematic lighting, textured skin, "
NEG  = "verybadimagenegative_v1.3, ng_deepnegative_v1_75t, (worst quality:2), (low quality:2), lowres, blurry, bad anatomy, bad hands, extra digits, watermark, signature"

##############################################################################
# 8. 生成コア
##############################################################################
@spaces.GPU(duration=60)
def generate(face: Image.Image, subj: str, add: str, neg: str,
             cfg: float, ipw: float, steps: int, w: int, h: int,
             up: bool, upf: int, progress=gr.Progress(track_tqdm=True)):
    if pipe is None:
        init()
    if len(face_analyser.get(np.array(face))) == 0:
        raise ValueError("顔が検出できません。他の画像でお試しください。")
    pipe.set_adapters(["ip_faceid"], adapter_weights=[ipw])
    img = pipe(prompt=BASE+subj+", "+add,
               negative_prompt=NEG+", "+neg,
               num_inference_steps=steps, guidance_scale=cfg,
               image=face, width=w, height=h).images[0]
    if up:
        upsampler.scale = int(upf)
        img, _ = upsampler.enhance(np.array(img)); img = Image.fromarray(img)
    return img

##############################################################################
# 9. Gradio UI
##############################################################################
with gr.Blocks(title="BRA v7 × InstantID (ZeroGPU)") as demo:
    gr.Markdown("## BRA v7 × InstantID")
    with gr.Row():
        f = gr.Image(type="pil", label="Face ID"); s = gr.Textbox(label="被写体説明")
    ap = gr.Textbox(label="追加プロンプト"); ng = gr.Textbox(label="追加ネガ")
    with gr.Row():
        cf = gr.Slider(1,20,7.5,0.5,"CFG"); ip = gr.Slider(0.1,1.0,0.6,0.05,"IP-Adapter Weight")
    with gr.Row():
        st = gr.Slider(10,50,30,1,"Steps"); W = gr.Slider(512,1024,768,64,"W"); H = gr.Slider(512,1024,768,64,"H")
    with gr.Row():
        up = gr.Checkbox(label="Real-ESRGAN"); upf = gr.Radio([4,8], value=4, label="アップスケール")
    btn = gr.Button("Generate"); out = gr.Image(type="pil", label="Result")
    btn.click(generate, [f,s,ap,ng,cf,ip,st,W,H,up,upf], out, show_progress=True)

##############################################################################
# 10. FastAPI
##############################################################################
app = FastAPI()

@app.post("/api/generate")
async def api_gen(subj: str=Form(...), cfg: float=Form(7.5), stp: int=Form(30),
                  ipw: float=Form(0.6), W: int=Form(768), H: int=Form(768),
                  file: UploadFile=File(...)):
    img = Image.open(io.BytesIO(await file.read())).convert("RGB")
    res = generate(img, subj, "", "", cfg, ipw, stp, W, H, False, 4)
    buf = io.BytesIO(); res.save(buf,"PNG")
    return {"image":"data:image/png;base64,"+base64.b64encode(buf.getvalue()).decode()}

##############################################################################
# 11. Launch
##############################################################################
demo.queue(default_concurrency_limit=2).launch(share=False)