app.py

"""Stable-Makeup: 妆容迁移 · ModelScope 创空间 xGPU

基于 sky24h/Stable-Makeup-unofficial 成功部署方案。
模型从 ModelScope 加载，预训练权重从 HuggingFace Hub 拉取。
"""

import os
import numpy as np
from PIL import Image
import torch
import gradio as gr

PROJECT_ROOT = os.path.dirname(os.path.abspath(__file__))
os.chdir(PROJECT_ROOT)

# ═══ 设备 · 精度自适应 ═══
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
DTYPE = torch.float16 if DEVICE == "cuda" else torch.float32
print(f"[Stable-Makeup] 设备: {DEVICE}, 精度: {DTYPE}")
if DEVICE == "cuda":
    print(f"[Stable-Makeup] GPU: {torch.cuda.get_device_name(0)}, 显存: {torch.cuda.get_device_properties(0).total_memory/1024**3:.0f}GB")

# ═══ 固定随机种子（与 sky24h 一致） ═══
torch.manual_seed(1024)
if torch.cuda.is_available():
    torch.cuda.manual_seed(1024)
    torch.cuda.manual_seed_all(1024)


# ═══════════════════════════════════════════════════════════
#  模型初始化（模块级加载）
# ═══════════════════════════════════════════════════════════

def _ms_download(repo_id, cache_dir="./models"):
    """从 ModelScope 下载模型"""
    from modelscope import snapshot_download
    return snapshot_download(repo_id, cache_dir=cache_dir)


def init_pipeline():
    """初始化推理管线——float16 GPU + DDIM"""
    from diffusers import UNet2DConditionModel as OriginalUNet2DConditionModel
    from diffusers import DDIMScheduler, ControlNetModel
    from pipeline_sd15 import StableDiffusionControlNetPipeline
    from detail_encoder.encoder_plus import detail_encoder

    # ── SD1.5 基础模型 ──
    sd15_dir = _ms_download("AI-ModelScope/stable-diffusion-v1-5")
    print(f"[Stable-Makeup] SD1.5: {sd15_dir}")

    Unet = OriginalUNet2DConditionModel.from_pretrained(
        sd15_dir, subfolder="unet", torch_dtype=DTYPE, local_files_only=True
    ).to(DEVICE)

    # ── 双路 ControlNet ──
    id_encoder = ControlNetModel.from_unet(Unet)
    pose_encoder = ControlNetModel.from_unet(Unet)

    # ── 妆容编码器（CLIP + SSR attention） ──
    clip_dir = _ms_download("AI-ModelScope/clip-vit-large-patch14")
    print(f"[Stable-Makeup] CLIP: {clip_dir}")
    makeup_encoder = detail_encoder(Unet, clip_dir, DEVICE, dtype=DTYPE)

    # ── 加载 Stable-Makeup 预训练权重（本地优先 → Google Drive 兜底） ──
    print("[Stable-Makeup] 加载预训练权重...")
    import gc

    REQUIRED = ["pytorch_model.bin", "pytorch_model_1.bin", "pytorch_model_2.bin"]
    WEIGHTS_DIR = os.path.join(PROJECT_ROOT, "weights")
    GDRIVE_FOLDER = "1397t27GrUyLPnj17qVpKWGwg93EcaFfg"

    # 检查缺失文件 → 自动从 Google Drive 下载（HF Spaces 海外服务器可直连）
    os.makedirs(WEIGHTS_DIR, exist_ok=True)
    missing = [f for f in REQUIRED if not os.path.exists(os.path.join(WEIGHTS_DIR, f))]
    if missing:
        print(f"[Stable-Makeup] 缺失: {missing}，从 Google Drive 下载...")
        try:
            import gdown
            gdown.download_folder(id=GDRIVE_FOLDER, output=WEIGHTS_DIR, quiet=False)
        except Exception as e:
            print(f"[Stable-Makeup] ⚠️ 自动下载失败: {e}")
            print(f"    请手动下载: https://drive.google.com/drive/folders/{GDRIVE_FOLDER}")
            print(f"    放到: {WEIGHTS_DIR}/")
            raise

    def _load_weight(filename):
        local_path = os.path.join(WEIGHTS_DIR, filename)
        print(f"  ✅ 加载: {local_path} ({os.path.getsize(local_path)/1024**3:.1f}GB)")
        return torch.load(local_path, map_location="cpu")

    id_encoder.load_state_dict(_load_weight("pytorch_model_1.bin"), strict=False)
    gc.collect()
    pose_encoder.load_state_dict(_load_weight("pytorch_model_2.bin"), strict=False)
    gc.collect()
    makeup_encoder.load_state_dict(_load_weight("pytorch_model.bin"), strict=False)
    gc.collect()

    id_encoder.to(device=DEVICE, dtype=DTYPE)
    pose_encoder.to(device=DEVICE, dtype=DTYPE)
    makeup_encoder.to(device=DEVICE, dtype=DTYPE)

    # ── 推理管线 ──
    pipe = StableDiffusionControlNetPipeline.from_pretrained(
        sd15_dir,
        safety_checker=None,
        unet=Unet,
        controlnet=[id_encoder, pose_encoder],
        torch_dtype=DTYPE,
        local_files_only=True,
    ).to(DEVICE)

    # DDIMScheduler（论文原版，与 sky24h 一致）
    pipe.scheduler = DDIMScheduler.from_config(pipe.scheduler.config)
    print("[Stable-Makeup] ✅ 模型加载完成 (float16 GPU + DDIM)")
    return pipe, makeup_encoder


# ── 启动时加载模型 ──
pipeline, makeup_encoder = init_pipeline()


# ═══════════════════════════════════════════════════════════
#  推理函数
# ═══════════════════════════════════════════════════════════

def get_draw(pil_img, size):
    """生成人脸结构控制图（PIL → cv2 BGR → SPIGA → 骨架图）"""
    import cv2
    from spiga_draw import spiga_process, spiga_segmentation
    cv2_img = cv2.cvtColor(np.array(pil_img), cv2.COLOR_RGB2BGR)
    spigas = spiga_process(cv2_img)
    if spigas is False:
        width, height = pil_img.size
        return Image.new("RGB", (width, height), color=(0, 0, 0))
    return spiga_segmentation(spigas, size=size)


def makeup_transfer(id_image, makeup_image, guidance_scale=1.6):
    """妆容迁移——与 sky24h 推理逻辑完全一致"""
    size = 512
    id_image_resized = id_image.resize((size, size))
    makeup_image_resized = makeup_image.resize((size, size))

    pose_image = get_draw(id_image_resized, size=size)
    result = makeup_encoder.generate(
        id_image=[id_image_resized, pose_image],
        makeup_image=makeup_image_resized,
        guidance_scale=guidance_scale,
        pipe=pipeline,
    )
    return result


# ═══════════════════════════════════════════════════════════
#  Gradio UI
# ═══════════════════════════════════════════════════════════

with gr.Blocks(title="Stable-Makeup 妆容迁移") as demo:
    gr.Markdown("""
    # 💄 Stable-Makeup · 妆容迁移
    上传素颜照 + 参考妆容图，AI 将妆容迁移到你的照片上。
    基于 [Stable-Makeup](https://arxiv.org/abs/2403.07764) (arXiv 2403.07764)
    """)
    with gr.Row():
        with gr.Column():
            id_img = gr.Image(label="素颜照", type="pil", height=400)
        with gr.Column():
            makeup_img = gr.Image(label="参考妆容", type="pil", height=400)

    guidance = gr.Slider(
        minimum=1.01, maximum=3.0, value=1.6, step=0.05,
        label="妆容浓度 (guidance_scale)",
        info="淡妆建议 1.05-1.15，浓妆建议 2.0"
    )
    btn = gr.Button("开始试妆", variant="primary")
    output = gr.Image(label="试妆结果", type="pil")

    btn.click(
        fn=makeup_transfer,
        inputs=[id_img, makeup_img, guidance],
        outputs=output,
    )

if __name__ == "__main__":
    demo.queue(max_size=10).launch()