Update gradio_app.py

gradio_app.py

@@ -6,8 +6,8 @@ import base64
 import shutil
 from typing import Optional, Tuple
 
-import spaces
 import gradio as gr
+import spaces
 import torch
 import cv2
 import numpy as np
@@ -121,10 +121,18 @@ def _download_models() -> Tuple[Optional[str], Optional[str], Optional[str]]:
             cache_dir=cache_dir,
             token=HF_AUTH_TOKEN,
         )
+        # If repo root contains a nested "HairMapper" folder, link to that subfolder.
+        hm_src = hm_snap
+        nested_hm = os.path.join(hm_snap, "HairMapper")
+        if os.path.isdir(nested_hm) and (
+            os.path.isfile(os.path.join(nested_hm, "hair_mapper_run.py")) or
+            os.path.isdir(os.path.join(nested_hm, "mapper"))
+        ):
+            hm_src = nested_hm
         # Create a symlink so that imports like "from HairMapper..." work
-        hairmapper_dir = _ensure_symlink(hm_snap, os.path.abspath("HairMapper"))
+        hairmapper_dir = _ensure_symlink(hm_src, os.path.abspath("HairMapper"))
         if hairmapper_dir not in sys.path:
-            sys.path.insert(0, hairmapper_dir)
+            sys.path.insert(0, os.path.dirname(hairmapper_dir))
 
     # 3) FFHQFaceAlignment
     ffhq_dir = None
@@ -135,10 +143,18 @@ def _download_models() -> Tuple[Optional[str], Optional[str], Optional[str]]:
             cache_dir=cache_dir,
             token=HF_AUTH_TOKEN,
         )
-        # Create a symlink so that test_stablehairv2._maybe_align_image("./FFHQFaceAlignment") resolves
-        ffhq_dir = _ensure_symlink(fa_snap, os.path.abspath("FFHQFaceAlignment"))
+        # If repo root contains a nested "FFHQFaceAlignment" folder, link to that subfolder.
+        fa_src = fa_snap
+        nested_fa = os.path.join(fa_snap, "FFHQFaceAlignment")
+        if os.path.isdir(nested_fa) and (
+            os.path.isfile(os.path.join(nested_fa, "align.py")) or
+            os.path.isdir(os.path.join(nested_fa, "lib"))
+        ):
+            fa_src = nested_fa
+        # Create a symlink so that _maybe_align_image can import modules relatively
+        ffhq_dir = _ensure_symlink(fa_src, os.path.abspath("FFHQFaceAlignment"))
         if ffhq_dir not in sys.path:
-            sys.path.insert(0, ffhq_dir)
+            sys.path.insert(0, os.path.dirname(ffhq_dir))
 
     # 4) Optional: Trained model weights (motion/control/ref)
     if TRAINED_MODEL_REPO:
@@ -186,9 +202,10 @@ SD15_PATH, _, _ = _download_models()
 # -----------------------------------------------------------------------------
 # Gradio inference
 # -----------------------------------------------------------------------------
-with open("imgs/background.png", "rb") as f:
+with open("imgs/background.jpg", "rb") as f:
     _b64_bg = base64.b64encode(f.read()).decode()
 
+
 @spaces.GPU
 def inference(id_image, hair_image):
     # Require GPU (HairMapper currently uses CUDA explicitly)