"""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()