core/hair_transfer.py

"""
High-quality hair transfer with HairFastGAN (StyleGAN-based).

InsightFace swaps only the face; this transfers a reference hairstyle. It prefers
a LOCAL GPU install (vendored under external/HairFastGAN, run via its MSVC-env
launcher) and falls back to the hosted HairFastGAN Gradio Space. Set
HAIRFAST_LOCAL=0 to force the Space, or HAIRFAST_SPACE to your own duplicated
(HF Pro GPU) Space; HF_TOKEN lifts the public Space's ZeroGPU quota.

Either backend returns an FFHQ-aligned 1024 portrait (face + new hair, no
background), so callers paste the head back into the original scene themselves.
"""
import os
import tempfile

import cv2
import numpy as np

# Some environments (conda) set SSL_CERT_FILE to a path that doesn't exist,
# which makes httpx/gradio_client crash. Repair it from certifi before use.
_cf = os.environ.get("SSL_CERT_FILE")
if not _cf or not os.path.exists(_cf):
    try:
        import certifi
        os.environ["SSL_CERT_FILE"] = certifi.where()
        os.environ.setdefault("REQUESTS_CA_BUNDLE", certifi.where())
    except Exception:
        pass

HAIRFAST_SPACE = os.environ.get("HAIRFAST_SPACE", "AIRI-Institute/HairFastGAN")

# Local GPU HairFastGAN (vendored under external/HairFastGAN). When present with
# a CUDA GPU we run it via the MSVC-env launcher instead of the hosted Space —
# faster, private, no quota. Set HAIRFAST_LOCAL=0 to force the Space.
_HF_LOCAL_DIR = os.path.join("external", "HairFastGAN")
_HF_LOCAL_BAT = os.path.join(_HF_LOCAL_DIR, "run_hairfast.bat")

_client = None
_client_failed = False
_local_ok = None  # cached availability check


def _local_available() -> bool:
    global _local_ok
    if _local_ok is not None:
        return _local_ok
    _local_ok = False
    if os.environ.get("HAIRFAST_LOCAL", "1") not in ("1", "true", "on"):
        return False
    # Need the vendored repo with weights linked in (run scripts/setup_hairfast.py).
    if not os.path.isdir(os.path.join(_HF_LOCAL_DIR, "pretrained_models")):
        return False
    # Self-heal: drop the committed runner scripts into the repo if missing.
    if not os.path.isfile(_HF_LOCAL_BAT):
        src = os.path.join("scripts", "hairfast")
        try:
            import shutil
            for f in ("run_hairfast.py", "run_hairfast.bat"):
                s = os.path.join(src, f)
                if os.path.isfile(s):
                    shutil.copy(s, os.path.join(_HF_LOCAL_DIR, f))
        except Exception:
            pass
    if not os.path.isfile(_HF_LOCAL_BAT):
        return False
    try:
        import torch
        _local_ok = torch.cuda.is_available()
    except Exception:
        _local_ok = False
    if _local_ok:
        print("[hair_transfer] local GPU HairFastGAN available")
    return _local_ok


def _transfer_hair_local(face_bgr, shape_bgr, color_bgr):
    """Run the vendored HairFastGAN on the local GPU via its MSVC-env launcher."""
    import subprocess
    work = tempfile.mkdtemp(prefix="hairfast_")
    fp = os.path.join(work, "face.png")
    sp = os.path.join(work, "shape.png")
    cp = os.path.join(work, "color.png")
    op = os.path.join(work, "out.png")
    cv2.imwrite(fp, face_bgr)
    cv2.imwrite(sp, shape_bgr)
    cv2.imwrite(cp, color_bgr)
    try:
        # The .bat sets the VS build env + CUDA arch, cds into the repo, and runs
        # run_hairfast.py with absolute paths. ~18s warm; first ever call is slow
        # (one-time op compile + model downloads), hence the generous timeout.
        r = subprocess.run(
            ["cmd", "/c", os.path.abspath(_HF_LOCAL_BAT), fp, sp, cp, op],
            capture_output=True, text=True, timeout=900,
        )
        if os.path.exists(op):
            return cv2.imread(op)
        print(f"[hair_transfer] local run produced no output:\n{r.stdout[-500:]}\n{r.stderr[-500:]}")
        return None
    except subprocess.TimeoutExpired:
        print("[hair_transfer] local run timed out")
        return None
    except Exception as e:
        print(f"[hair_transfer] local run error: {e}")
        return None
    finally:
        for p in (fp, sp, cp, op):
            try:
                os.remove(p)
            except Exception:
                pass
        try:
            os.rmdir(work)
        except Exception:
            pass


def _get_client():
    global _client, _client_failed
    if _client is not None:
        return _client
    if _client_failed:
        return None
    try:
        from gradio_client import Client
        # An HF token lifts the anonymous ZeroGPU quota that makes the public
        # Space error on every call. Required in practice; point HAIRFAST_SPACE
        # at your own duplicated (Pro GPU) Space for reliability.
        token = os.environ.get("HF_TOKEN") or os.environ.get("HUGGINGFACE_TOKEN")
        _client = Client(HAIRFAST_SPACE, token=token, verbose=False)
        print(f"[hair_transfer] connected to Space: {HAIRFAST_SPACE}"
              f"{' (authenticated)' if token else ' (anonymous — may hit GPU quota)'}")
        return _client
    except Exception as e:
        print(f"[hair_transfer] could not connect to {HAIRFAST_SPACE}: {e}")
        _client_failed = True
        return None


def _write_tmp(img_bgr: np.ndarray) -> str:
    fd, path = tempfile.mkstemp(suffix=".png")
    os.close(fd)
    cv2.imwrite(path, img_bgr)
    return path


def transfer_hair(face_bgr: np.ndarray,
                  shape_bgr: np.ndarray,
                  color_bgr=None):   # np.ndarray | None
    """
    Return an FFHQ-aligned 1024 portrait: the `face` identity wearing the
    hairstyle (shape) and hair colour (color) of the references. Returns a BGR
    image, or None on any failure (caller should fall back gracefully).

    face_bgr : whose face/identity to keep (e.g. the InsightFace swap result)
    shape_bgr: hairstyle-shape reference (e.g. the source person)
    color_bgr: hair-colour reference (defaults to shape)

    Uses the local GPU HairFastGAN when available; otherwise the hosted Space.
    """
    if color_bgr is None:
        color_bgr = shape_bgr

    # Prefer the local GPU model (faster, private, no quota).
    if _local_available():
        res = _transfer_hair_local(face_bgr, shape_bgr, color_bgr)
        if res is not None:
            return res
        print("[hair_transfer] local run failed — falling back to Space")

    client = _get_client()
    if client is None:
        return None

    from gradio_client import handle_file
    paths = []
    try:
        fp = _write_tmp(face_bgr);  paths.append(fp)
        sp = _write_tmp(shape_bgr); paths.append(sp)
        cp = _write_tmp(color_bgr); paths.append(cp)

        # 1. Align each photo to FFHQ (the Space's preprocessing step).
        align = ["Face", "Shape", "Color"]
        fa = client.predict(img=handle_file(fp), align=align, api_name="/resize_inner")
        sa = client.predict(img=handle_file(sp), align=align, api_name="/resize_inner_1")
        ca = client.predict(img=handle_file(cp), align=align, api_name="/resize_inner_2")

        # 2. Swap the hair.
        out = client.predict(
            face=handle_file(fa),
            shape=handle_file(sa),
            color=handle_file(ca),
            blending="Article",
            poisson_iters=0,
            poisson_erosion=15,
            api_name="/swap_hair",
        )
        result_path = out[0] if isinstance(out, (list, tuple)) else out
        err = out[1] if isinstance(out, (list, tuple)) and len(out) > 1 else ""
        if err:
            print(f"[hair_transfer] Space reported: {err}")
        if not result_path or not os.path.exists(result_path):
            print("[hair_transfer] no result image returned")
            return None
        res = cv2.imread(result_path)
        return res
    except Exception as e:
        print(f"[hair_transfer] transfer failed: {e}")
        return None
    finally:
        for p in paths:
            try:
                os.remove(p)
            except Exception:
                pass