app.py

# app.py — BRA v7 (AIGaming repo) × InstantID × ZeroGPU
# 2025-06-22
##############################################################################
# torchvision 0.17+ 互換パッチ（functional_tensor → functional）
##############################################################################
import sys, types
try:
    import torchvision.transforms.functional as F
    if "torchvision.transforms.functional_tensor" not in sys.modules:
        faux = types.ModuleType("torchvision.transforms.functional_tensor")
        # 必要最低限だけ持たせる
        faux.rgb_to_grayscale = getattr(F, "rgb_to_grayscale", None)
        sys.modules["torchvision.transforms.functional_tensor"] = faux
except Exception as e:
    print("[WARN] torchvision compatibility patch failed:", e)

##############################################################################
# 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"
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
    ip_lora = dl(LORA_REPO, IP_LORA)
    
    ### 最終修正 ### subfolder引数に空文字列""を渡し、TypeErrorを回避する
    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)