Update gradio_app.py

gradio_app.py

@@ -1,38 +1,190 @@
 import os
+import sys
+import uuid
+import logging
+import base64
+import shutil
+from typing import Optional, Tuple
+
+import gradio as gr
+import torch
+import cv2
+import numpy as np
+
+from huggingface_hub import snapshot_download
 
 
-# 替换原有的目录设置
+# -----------------------------------------------------------------------------
+# Environment for HF Spaces
+# -----------------------------------------------------------------------------
 os.environ.setdefault("GRADIO_TEMP_DIR", "/tmp/gradio")
 os.environ.setdefault("TMPDIR", "/tmp")
-os.makedirs("/tmp/gradio", exist_ok=True)
-os.makedirs("/tmp", exist_ok=True)
+os.makedirs(os.environ["GRADIO_TEMP_DIR"], exist_ok=True)
+os.makedirs(os.environ["TMPDIR"], exist_ok=True)
+
+
+# -----------------------------------------------------------------------------
+# Config via environment variables (set these in your Space settings)
+# -----------------------------------------------------------------------------
+# Required (you uploaded these as separate model repos on HF):
+#   - FFHQFACEALIGNMENT_REPO (e.g., "yourname/FFHQFaceAlignment")
+#   - HAIRMAPPER_REPO (e.g., "yourname/HairMapper")
+#   - SD15_REPO (e.g., "yourname/stable-diffusion-v1-5")
+# Optional:
+#   - TRAINED_MODEL_REPO (if you uploaded motion/control/ref ckpts as a repo)
+# If TRAINED_MODEL_REPO not provided, we will try to use local "./pretrain".
+FFHQFACEALIGNMENT_REPO = os.getenv("FFHQFACEALIGNMENT_REPO", "")
+HAIRMAPPER_REPO = os.getenv("HAIRMAPPER_REPO", "")
+SD15_REPO = os.getenv("SD15_REPO", "")
+TRAINED_MODEL_REPO = os.getenv("TRAINED_MODEL_REPO", "")
+
+
+# -----------------------------------------------------------------------------
+# Utilities
+# -----------------------------------------------------------------------------
+def _ensure_symlink(src_dir: str, dst_path: str) -> str:
+    """Create a directory symlink at dst_path pointing to src_dir if not exists.
+    If symlink creation is unavailable, fallback to copying a minimal structure.
+    Returns the final path that should be used by imports (dst_path if created, else src_dir).
+    """
+    try:
+        if os.path.islink(dst_path) or os.path.isdir(dst_path):
+            return dst_path
+        os.symlink(src_dir, dst_path, target_is_directory=True)
+        return dst_path
+    except Exception:
+        # Fallback: try to create the directory and copy only top-level python files/dirs needed
+        try:
+            if not os.path.exists(dst_path):
+                os.makedirs(dst_path, exist_ok=True)
+            # Last resort: shallow copy (can still be heavy; symlink is preferred on HF Linux)
+            for name in os.listdir(src_dir):
+                src = os.path.join(src_dir, name)
+                dst = os.path.join(dst_path, name)
+                if os.path.exists(dst):
+                    continue
+                if os.path.isdir(src):
+                    shutil.copytree(src, dst)
+                else:
+                    shutil.copy2(src, dst)
+            return dst_path
+        except Exception:
+            # Give up and return original source
+            return src_dir
+
+
+def _find_model_root(path: str) -> str:
+    """Given a snapshot path, return the directory containing model_index.json.
+    Handles repos that nest the folder (e.g., repo/stable-diffusion-v1-5/...).
+    """
+    if os.path.isfile(os.path.join(path, "model_index.json")):
+        return path
+    # Search one level deep for a folder with model_index.json
+    for name in os.listdir(path):
+        cand = os.path.join(path, name)
+        if os.path.isdir(cand) and os.path.isfile(os.path.join(cand, "model_index.json")):
+            return cand
+    # As a fallback, return original path
+    return path
+
+
+def _download_models() -> Tuple[Optional[str], Optional[str], Optional[str]]:
+    """Download HF model repos and prepare local paths.
+
+    Returns:
+        - sd15_path: path to the Stable Diffusion v1-5 folder (with model_index.json)
+        - hairmapper_dir: path to local HairMapper folder (import root)
+        - ffhq_dir: path to local FFHQFaceAlignment folder (import root)
+    """
+    cache_dir = os.getenv("HF_HUB_CACHE", None)
+
+    # 1) Stable Diffusion 1.5
+    sd15_path = None
+    if SD15_REPO:
+        sd_snap = snapshot_download(repo_id=SD15_REPO, local_files_only=False, cache_dir=cache_dir)
+        sd15_path = _find_model_root(sd_snap)
+
+    # 2) HairMapper
+    hairmapper_dir = None
+    if HAIRMAPPER_REPO:
+        hm_snap = snapshot_download(repo_id=HAIRMAPPER_REPO, local_files_only=False, cache_dir=cache_dir)
+        # Create a symlink so that imports like "from HairMapper..." work
+        hairmapper_dir = _ensure_symlink(hm_snap, os.path.abspath("HairMapper"))
+        if hairmapper_dir not in sys.path:
+            sys.path.insert(0, hairmapper_dir)
+
+    # 3) FFHQFaceAlignment
+    ffhq_dir = None
+    if FFHQFACEALIGNMENT_REPO:
+        fa_snap = snapshot_download(repo_id=FFHQFACEALIGNMENT_REPO, local_files_only=False, cache_dir=cache_dir)
+        # Create a symlink so that test_stablehairv2._maybe_align_image("./FFHQFaceAlignment") resolves
+        ffhq_dir = _ensure_symlink(fa_snap, os.path.abspath("FFHQFaceAlignment"))
+        if ffhq_dir not in sys.path:
+            sys.path.insert(0, ffhq_dir)
+
+    # 4) Optional: Trained model weights (motion/control/ref)
+    if TRAINED_MODEL_REPO:
+        tm_snap = snapshot_download(repo_id=TRAINED_MODEL_REPO, local_files_only=False, cache_dir=cache_dir)
+        # Symlink to ./trained_model so downstream code can load from there
+        _ = _ensure_symlink(tm_snap, os.path.abspath("trained_model"))
+
+    return sd15_path, hairmapper_dir, ffhq_dir
+
+
+# -----------------------------------------------------------------------------
+# Lazy imports that rely on downloaded models/paths
+# -----------------------------------------------------------------------------
+def _import_inference_bits():
+    from test_stablehairv2 import log_validation
+    from test_stablehairv2 import UNet3DConditionModel, ControlNetModel, CCProjection
+    from test_stablehairv2 import AutoTokenizer, CLIPVisionModelWithProjection, AutoencoderKL, UNet2DConditionModel
+    from test_stablehairv2 import _maybe_align_image
+    from HairMapper.hair_mapper_run import bald_head
+    return (
+        log_validation,
+        UNet3DConditionModel,
+        ControlNetModel,
+        CCProjection,
+        AutoTokenizer,
+        CLIPVisionModelWithProjection,
+        AutoencoderKL,
+        UNet2DConditionModel,
+        _maybe_align_image,
+        bald_head,
+    )
 
-# 同时修改你的输出目录为相对路径
-os.makedirs("gradio_inputs", exist_ok=True)
-os.makedirs("gradio_outputs", exist_ok=True)
 
+# -----------------------------------------------------------------------------
+# Prepare models on startup
+# -----------------------------------------------------------------------------
+SD15_PATH, _, _ = _download_models()
 
-import logging
-import gradio as gr
-import torch
-import os
-import uuid
-from test_stablehairv2 import log_validation
-from test_stablehairv2 import UNet3DConditionModel, ControlNetModel, CCProjection
-from test_stablehairv2 import AutoTokenizer, CLIPVisionModelWithProjection, AutoencoderKL, UNet2DConditionModel
-from omegaconf import OmegaConf
-import numpy as np
-import cv2
-from test_stablehairv2 import _maybe_align_image
-from HairMapper.hair_mapper_run import bald_head
-
-import base64
 
+# -----------------------------------------------------------------------------
+# Gradio inference
+# -----------------------------------------------------------------------------
 with open("imgs/background.jpg", "rb") as f:
-    b64_img = base64.b64encode(f.read()).decode()
+    _b64_bg = base64.b64encode(f.read()).decode()
 
 
 def inference(id_image, hair_image):
+    # Require GPU (HairMapper currently uses CUDA explicitly)
+    if not torch.cuda.is_available():
+        raise RuntimeError("This demo requires a GPU Space. Please enable a GPU in this Space.")
+
+    (
+        log_validation,
+        UNet3DConditionModel,
+        ControlNetModel,
+        CCProjection,
+        AutoTokenizer,
+        CLIPVisionModelWithProjection,
+        AutoencoderKL,
+        UNet2DConditionModel,
+        _maybe_align_image,
+        bald_head,
+    ) = _import_inference_bits()
+
     os.makedirs("gradio_inputs", exist_ok=True)
     os.makedirs("gradio_outputs", exist_ok=True)
 
@@ -41,40 +193,46 @@ def inference(id_image, hair_image):
     id_image.save(id_path)
     hair_image.save(hair_path)
 
-    # ===== 图像对齐 =====
+    # Align
     aligned_id = _maybe_align_image(id_path, output_size=1024, prefer_cuda=True)
     aligned_hair = _maybe_align_image(hair_path, output_size=1024, prefer_cuda=True)
 
-    # 保存对齐结果（方便 Gradio 输出）
     aligned_id_path = "gradio_outputs/aligned_id.png"
     aligned_hair_path = "gradio_outputs/aligned_hair.png"
     cv2.imwrite(aligned_id_path, cv2.cvtColor(aligned_id, cv2.COLOR_RGB2BGR))
     cv2.imwrite(aligned_hair_path, cv2.cvtColor(aligned_hair, cv2.COLOR_RGB2BGR))
 
-    # ===== 调用 HairMapper 秃头化 =====
+    # Balding
     bald_id_path = "gradio_outputs/bald_id.png"
     cv2.imwrite(bald_id_path, cv2.cvtColor(aligned_id, cv2.COLOR_RGB2BGR))
     bald_head(bald_id_path, bald_id_path)
 
-    # ===== 原本的 Args =====
+    # Resolve trained model dir
+    trained_model_dir = os.path.abspath("trained_model") if os.path.isdir("trained_model") else None
+    if trained_model_dir is None and os.path.isdir("pretrain"):
+        trained_model_dir = os.path.abspath("pretrain")
+    if trained_model_dir is None:
+        raise RuntimeError("Missing trained model weights. Provide TRAINED_MODEL_REPO or include ./pretrain.")
+
     class Args:
-        pretrained_model_name_or_path = "./stable-diffusion-v1-5/stable-diffusion-v1-5"
-        model_path = "./trained_model"
+        pretrained_model_name_or_path = SD15_PATH or os.path.abspath("stable-diffusion-v1-5/stable-diffusion-v1-5")
+        model_path = trained_model_dir
         image_encoder = "openai/clip-vit-large-patch14"
         controlnet_model_name_or_path = None
         revision = None
         output_dir = "gradio_outputs"
         seed = 42
         num_validation_images = 1
-        validation_ids = [aligned_id_path]  # 用对齐后的图像
-        validation_hairs = [aligned_hair_path]  # 用对齐后的图像
+        validation_ids = [aligned_id_path]
+        validation_hairs = [aligned_hair_path]
         use_fp16 = False
+        align_before_infer = True
+        align_size = 1024
 
     args = Args()
 
     device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
-    # 初始化 logger
     logging.basicConfig(
         format="%(asctime)s - %(levelname)s - %(name)s - %(message)s",
         datefmt="%m/%d/%Y %H:%M:%S",
@@ -82,15 +240,17 @@ def inference(id_image, hair_image):
     )
     logger = logging.getLogger(__name__)
 
-    # ===== 模型加载（和 main() 对齐） =====
+    # Load tokenizer/encoders/vae
     tokenizer = AutoTokenizer.from_pretrained(args.pretrained_model_name_or_path, subfolder="tokenizer",
                                               revision=args.revision)
     image_encoder = CLIPVisionModelWithProjection.from_pretrained(args.image_encoder, revision=args.revision).to(device)
-    vae = AutoencoderKL.from_pretrained(args.pretrained_model_name_or_path, subfolder="vae", revision=args.revision).to(
-        device, dtype=torch.float32)
+    vae = AutoencoderKL.from_pretrained(args.pretrained_model_name_or_path, subfolder="vae",
+                                        revision=args.revision).to(device, dtype=torch.float32)
 
+    from omegaconf import OmegaConf
     infer_config = OmegaConf.load('./configs/inference/inference_v2.yaml')
 
+    # UNet2D with 8-channel conv_in
     unet2 = UNet2DConditionModel.from_pretrained(
         args.pretrained_model_name_or_path, subfolder="unet", revision=args.revision, torch_dtype=torch.float32
     ).to(device)
@@ -126,10 +286,10 @@ def inference(id_image, hair_image):
         args.pretrained_model_name_or_path, subfolder="unet", revision=args.revision, low_cpu_mem_usage=False,
         device_map=None, ignore_mismatched_sizes=True
     ).to(device)
-    state_dict2 = torch.load(os.path.join(args.model_path, "pytorch_model_2.bin"), map_location="cpu")
-    Hair_Encoder.load_state_dict(state_dict2, strict=False)
+    state_dict4 = torch.load(os.path.join(args.model_path, "pytorch_model_2.bin"), map_location="cpu")
+    Hair_Encoder.load_state_dict(state_dict4, strict=False)
 
-    # 推理
+    # Run inference
     log_validation(
         vae, tokenizer, image_encoder, denoising_unet,
         args, device, logger,
@@ -138,7 +298,7 @@ def inference(id_image, hair_image):
 
     output_video = os.path.join(args.output_dir, "validation", "generated_video_0.mp4")
 
-    # 提取视频帧用于可拖动预览
+    # Extract frames for slider preview
     frames_dir = os.path.join(args.output_dir, "frames", uuid.uuid4().hex)
     os.makedirs(frames_dir, exist_ok=True)
     cap = cv2.VideoCapture(output_video)
@@ -157,34 +317,21 @@ def inference(id_image, hair_image):
     max_frames = len(frames_list) if frames_list else 1
     first_frame = frames_list[0] if frames_list else None
 
-    return aligned_id_path, aligned_hair_path, bald_id_path, output_video, frames_list, gr.update(minimum=1,
-                                                                                                  maximum=max_frames,
-                                                                                                  value=1,
-                                                                                                  step=1), first_frame
-
-
-# Gradio 前端
-# 原 Interface 版本（保留以便回退）
-# demo = gr.Interface(
-#     fn=inference,
-#     inputs=[
-#         gr.Image(type="pil", label="上传身份图（ID Image）"),
-#         gr.Image(type="pil", label="上传发型图（Hair Reference Image）")
-#     ],
-#     outputs=[
-#         gr.Image(type="filepath", label="对齐后的身份图"),
-#         gr.Image(type="filepath", label="对齐后的发型图"),
-#         gr.Image(type="filepath", label="秃头化后的身份图"),
-#         gr.Video(label="生成的视频")
-#     ],
-#     title="StableHairV2 多视角发型迁移",
-#     description="上传身份图和发型参考图，查看对齐结果并生成多视角视频"
-# )
-# if __name__ == "__main__":
-#     demo.launch(server_name="0.0.0.0", server_port=7860)
-
-# Blocks 美化版
-css = f"""
+    return (
+        aligned_id_path,
+        aligned_hair_path,
+        bald_id_path,
+        output_video,
+        frames_list,
+        gr.update(minimum=1, maximum=max_frames, value=1, step=1),
+        first_frame,
+    )
+
+
+# -----------------------------------------------------------------------------
+# UI (Blocks)
+# -----------------------------------------------------------------------------
+CSS = f"""
     html, body {{
         height: 100%;
         margin: 0;
@@ -195,10 +342,10 @@ css = f"""
         height: 100% !important;
         margin: 0 !important;
         padding: 0 !important;
-        background-image: url("data:image/jpeg;base64,{b64_img}");
+        background-image: url("data:image/jpeg;base64,{_b64_bg}");
         background-size: cover;
         background-position: center;
-        background-attachment: fixed; /* 背景固定 */
+        background-attachment: fixed;
     }}
     #title-card {{
         background: rgba(255, 255, 255, 0.8);
@@ -226,7 +373,6 @@ css = f"""
     }}
     .left-pane {{min-width: 360px}}
     .right-pane {{min-width: 680px}}
-    /* Tabs 美化 */
     .tabs {{
         background: rgba(255,255,255,0.88);
         border-radius: 12px;
@@ -240,31 +386,11 @@ css = f"""
         border-bottom: 1px solid #e5e7eb;
         padding-bottom: 6px;
     }}
-    .tab-nav button {{
-        background: rgba(255,255,255,0.7);
-        border: 1px solid #e5e7eb;
-        backdrop-filter: blur(6px);
-        border-radius: 8px;
-        padding: 6px 12px;
-        color: #111827;
-        transition: all .2s ease;
-    }}
-    .tab-nav button:hover {{
-        transform: translateY(-1px);
-        box-shadow: 0 4px 10px rgba(0,0,0,0.06);
-    }}
-    .tab-nav button[aria-selected="true"] {{
-        background: #4f46e5;
-        color: #fff;
-        border-color: #4f46e5;
-        box-shadow: 0 6px 14px rgba(79,70,229,0.25);
-    }}
     .tabitem {{
         background: rgba(255,255,255,0.88);
         border-radius: 10px;
         padding: 8px;
     }}
-    /* 发型库滚动限制容器：固定260px高度，内部可滚动 */
     #hair_gallery_wrap {{
         height: 260px !important;
         overflow-y: scroll !important;
@@ -274,17 +400,13 @@ css = f"""
         height: 100% !important;
         overflow-y: scroll !important;
     }}
-    /* 确保画廊本体占满容器高度，避免滚动条落到页面底部 */
     #hair_gallery {{
         height: 100% !important;
     }}
 """
 
-with gr.Blocks(
-        theme=gr.themes.Soft(primary_hue="indigo", neutral_hue="slate"),
-        css=css
-) as demo:
-    # ==== 顶部 Panel ====
+
+with gr.Blocks(theme=gr.themes.Soft(primary_hue="indigo", neutral_hue="slate"), css=CSS) as demo:
     with gr.Group(elem_id="title-card"):
         gr.Markdown("""
         <h2 id='title'>StableHairV2 多视角发型迁移</h2>
@@ -300,13 +422,10 @@ with gr.Blocks(
                 run_btn = gr.Button("开始生成", variant="primary")
                 clear_btn = gr.Button("清空")
 
-            # ========= 发型库（点击即填充到“发型参考图”） =========
             def _list_imgs(dir_path: str):
                 exts = (".png", ".jpg", ".jpeg", ".webp")
-                # exts = (".jpg")
                 try:
-                    files = [os.path.join(dir_path, f) for f in sorted(os.listdir(dir_path))
-                             if f.lower().endswith(exts)]
+                    files = [os.path.join(dir_path, f) for f in sorted(os.listdir(dir_path)) if f.lower().endswith(exts)]
                     return files
                 except Exception:
                     return []
@@ -315,11 +434,8 @@ with gr.Blocks(
 
             with gr.Accordion("发型库（点击选择后自动填充）", open=True):
                 with gr.Group(elem_id="hair_gallery_wrap"):
-                    gallery = gr.Gallery(
-                        value=hair_list,
-                        columns=4, rows=2, allow_preview=True, label="发型库",
-                        elem_id="hair_gallery"
-                    )
+                    gallery = gr.Gallery(value=hair_list, columns=4, rows=2, allow_preview=True, label="发型库",
+                                         elem_id="hair_gallery")
 
             def _pick_hair(evt: gr.SelectData):  # type: ignore[name-defined]
                 i = evt.index if hasattr(evt, 'index') else 0
@@ -350,12 +466,11 @@ with gr.Blocks(
                     with gr.Group(elem_classes=["out-card"]):
                         bald_id_out = gr.Image(type="filepath", label="秃头化后的身份图", height=260)
 
-    # 逻辑保持不变
-    run_btn.click(fn=inference,
-                  inputs=[id_input, hair_input],
-                  outputs=[aligned_id_out, aligned_hair_out, bald_id_out,
-                           video_out, frames_state, frame_slider, frame_preview])
-
+    run_btn.click(
+        fn=inference,
+        inputs=[id_input, hair_input],
+        outputs=[aligned_id_out, aligned_hair_out, bald_id_out, video_out, frames_state, frame_slider, frame_preview],
+    )
 
     def _on_slide(frames, idx):
         if not frames:
@@ -364,20 +479,15 @@ with gr.Blocks(
         i = max(0, min(i, len(frames) - 1))
         return gr.update(value=frames[i])
 
-
     frame_slider.change(_on_slide, inputs=[frames_state, frame_slider], outputs=frame_preview)
 
-
     def _clear():
         return None, None, None, None, None
 
+    clear_btn.click(_clear, None, [id_input, hair_input, aligned_id_out, aligned_hair_out, bald_id_out])
 
-    clear_btn.click(_clear, None,
-                    [id_input, hair_input, aligned_id_out, aligned_hair_out, bald_id_out])
 
 if __name__ == "__main__":
     demo.queue().launch(server_name="0.0.0.0", server_port=7860)
 
 
-
-