Add ffhq-uv: FFHQ-UV RGB fitting scripts (native env, like p3dm) + stevie example outputs

.gitattributes

@@ -38,3 +38,6 @@ outputs/scail2/scail2-test_00002_.mp4 filter=lfs diff=lfs merge=lfs -text
 outputs/qwen-edit-v4/mode1-4step_00001_.png filter=lfs diff=lfs merge=lfs -text
 outputs/qwen-edit-v4/mode2-8step_00001_.png filter=lfs diff=lfs merge=lfs -text
 outputs/qwen-edit-v4/mode3-quality-noref2_00001_.png filter=lfs diff=lfs merge=lfs -text
+outputs/ffhq-uv/stevie/stevie_uv.png filter=lfs diff=lfs merge=lfs -text
+outputs/ffhq-uv/stevie/stevie_vis.png filter=lfs diff=lfs merge=lfs -text
+outputs/ffhq-uv/stevie/stevie_input_vis.png filter=lfs diff=lfs merge=lfs -text

ffhq-uv/README.md

@@ -0,0 +1,77 @@
+# ffhq-uv — FFHQ-UV single-photo → 3D face mesh + UV texture
+
+One photo in, fitted 3D face mesh (HiFi3D++ topology) plus a 1024×1024 facial
+UV-texture map out, using [FFHQ-UV](https://github.com/csbhr/FFHQ-UV)
+(Bai et al., CVPR 2023) RGB fitting. Runs on the native python of the
+`plx1029/comfyui-qwen` RunPod image (python 3.12, torch 2.8+cu128) —
+**no venv, no conda, torch is not reinstalled**.
+
+## Install (once per pod, ~3 min + ~2.5 GB downloads)
+
+```bash
+./ffhq_uv_install.sh
+```
+
+Everything lands in `$FFHQUV_HOME` (default `/workspace/ffhq-uv`). Re-running
+is safe; completed steps are skipped. Checkpoints and topology assets are
+fetched from the **official Hugging Face mirror** (`csbhr/FFHQ-UV`) — no
+Baidu/OneDrive account needed, and no registration (unlike FLAME/p3dm).
+
+## Run (~2 min per photo on H100)
+
+```bash
+./ffhq_uv_run.sh photo.jpg                 # -> /workspace/outputs/ffhq-uv/photo/
+./ffhq_uv_run.sh photo.jpg -o /some/dir --force
+```
+
+Outputs: `<name>_uv.png` (1024×1024 UV-texture map — the headline output),
+`<name>_id.obj` (neutral/identity head mesh, HiFi3D++ topology, 20481 verts),
+`<name>_exp.obj` (mesh with the photo's expression), `<name>.mtl` (links the
+meshes to the texture — the obj/mtl/png triplet opens textured in Blender or
+any obj viewer), `<name>_vis.png` (input | skin mask | parsing | landmarks |
+fitted render | uv), `<name>_input_vis.png` (preprocessing check).
+
+Plus a `raw/` subfolder with the complete machine-readable fitting results:
+all three stages' meshes/coeffs (`stage*_coeffs.pt`: HiFi3D++ id/exp/tex
+coefficients, pose, lighting, predicted 68 landmarks; stages 2-3 also GAN
+latents `latents_z`/`latents_w`), per-step optimization visuals, fitting log,
+tensorboard events, the aligned 224×224 input tensors (`input_data.pt`,
+`preprocess/<name>_processed.pt`: image, skin mask, parse mask, landmarks,
+arcface warp matrix).
+
+Pipeline: MTCNN 5-pt + 68-landmark detection (TF frozen graphs, CPU) → skin
+mask + face parsing → similarity alignment to 224×224 → Deep3D (HiFi3D++
+basis) coefficient init → StyleGAN2-based texture-GAN latent search (spherical
+surface, fixed shape, 100 steps) → joint texture+shape optimization (200
+steps, arcface id loss + VGG/LPIPS + photometric + landmark losses, rendered
+with nvdiffrast CUDA rasterizer).
+
+## What the installer does differently (vs. upstream repo)
+
+Upstream targets python 3.7 / torch 1.7.1 / TF 1.15 + conda. To run on the
+native py3.12 / torch 2.8 / numpy 2.1 env:
+
+- **checkpoints from the official HF mirror** (`csbhr/FFHQ-UV` dataset repo)
+  instead of Baidu Netdisk/OneDrive; only the ~2.5 GB the RGB-fitting
+  pipeline needs (not the 300+ GB dataset, not the StyleGAN/e4e/DPR editing
+  checkpoints used only for dataset creation)
+- **tensorflow-cpu** (modern 2.x) instead of tensorflow-gpu 1.15; the two TF1
+  frozen-graph detectors (MTCNN, 68-landmarks) are run through the
+  `tf.compat.v1` API — the only source patch needed. CPU on purpose: seconds
+  of inference, and no CUDA/cuDNN conflicts with torch
+- **no pytorch3d needed** (RGB fitting only uses nvdiffrast, already present
+  from p3dm; the installer installs it if missing); stylegan2 custom CUDA ops
+  JIT-compile on first use and fall back to reference implementations if the
+  build fails
+- `TORCH_FORCE_NO_WEIGHTS_ONLY_LOAD=1` at runtime (repo predates torch 2.6
+  `weights_only` default)
+- pip runs against a constraints file pinning the already-installed
+  numpy/torch/opencv/kornia/… so the ComfyUI env is never disturbed; only two
+  new packages are actually installed (`dominate`, `tensorflow-cpu`)
+
+## Licenses
+
+FFHQ-UV code and checkpoints are for **non-commercial research** use (see the
+[upstream repo](https://github.com/csbhr/FFHQ-UV)); checkpoints are downloaded
+from the authors' official Hugging Face mirror at install time and are not
+redistributed here. HiFi3D++ topology © Tencent AI Lab (REALY).

ffhq-uv/ffhq_uv_install.sh

@@ -0,0 +1,151 @@
+#!/usr/bin/env bash
+# =============================================================================
+# ffhq_uv_install.sh — one-shot installer for FFHQ-UV RGB fitting
+#                      (single photo -> 3D face mesh + 1024x1024 UV texture)
+#
+# Installs https://github.com/csbhr/FFHQ-UV (RGB_Fitting pipeline, HiFi3D++
+# topology + GAN texture decoder) into the NATIVE python environment (no venv,
+# no conda, no torch reinstall). Tested on: Ubuntu 24.04, python 3.12,
+# torch 2.8.0+cu128, H100 (RunPod image plx1029/comfyui-qwen).
+#
+# Idempotent: safe to re-run; finished steps are skipped.
+#
+# Requirements:
+#   - torch with CUDA already installed (the script will not install torch)
+#   - ~2.5 GB disk for checkpoints (downloaded from the official HF mirror
+#     csbhr/FFHQ-UV — no Baidu/OneDrive account needed)
+#
+# Usage:
+#   ./ffhq_uv_install.sh
+#
+# Layout (override with FFHQUV_HOME):
+#   $FFHQUV_HOME/FFHQ-UV          code + checkpoints + topo assets
+#   $FFHQUV_HOME/staging          per-image input staging
+#   $FFHQUV_HOME/fitting_output   per-image fitting results
+# =============================================================================
+set -euo pipefail
+
+FFHQUV_HOME="${FFHQUV_HOME:-/workspace/ffhq-uv}"
+REPO="$FFHQUV_HOME/FFHQ-UV"
+
+log() { echo -e "\n\033[1;36m[ffhq-uv-install]\033[0m $*"; }
+
+mkdir -p "$FFHQUV_HOME" "$FFHQUV_HOME/staging" "$FFHQUV_HOME/fitting_output"
+cd "$FFHQUV_HOME"
+
+# ---------------------------------------------------------------- sanity checks
+log "checking python/torch"
+python3 - <<'EOF'
+import torch
+assert torch.cuda.is_available(), "CUDA torch required"
+print(f"  torch {torch.__version__} cuda {torch.version.cuda} ok")
+EOF
+
+# ---------------------------------------------------- constraints (protect env)
+# Pin the currently-installed versions of shared heavyweight packages so pip
+# cannot up/downgrade them while installing FFHQ-UV deps (keeps ComfyUI safe).
+log "generating pip constraints from current env"
+python3 - > "$FFHQUV_HOME/constraints.txt" <<'EOF'
+from importlib.metadata import version, PackageNotFoundError
+for p in ["numpy","torch","torchvision","pillow","scipy","scikit-image",
+          "opencv-python","kornia","tensorboard","huggingface_hub",
+          "transformers","safetensors","requests"]:
+    try: print(f"{p}=={version(p)}")
+    except PackageNotFoundError: pass
+EOF
+CONS="-c $FFHQUV_HOME/constraints.txt"
+
+# ------------------------------------------------------------------ pip deps
+# Almost everything (opencv, kornia, scikit-image, tensorboard, nvdiffrast) is
+# already in the native env. Only two new packages are needed:
+#   dominate        - tiny, HTML logging used by the fitting logger
+#   tensorflow-cpu  - runs the two frozen-graph detectors (MTCNN + 68 landmarks)
+#                     in step1 preprocessing; CPU build on purpose (a few sec of
+#                     inference, avoids CUDA/cuDNN clashes with torch)
+log "installing python dependencies (native env, constrained)"
+python3 -c "import dominate" 2>/dev/null || pip install -q $CONS dominate
+python3 -c "import tensorflow" 2>/dev/null || pip install -q $CONS tensorflow-cpu
+
+# nvdiffrast: normally already present (e.g. from p3dm_install.sh)
+python3 -c "import nvdiffrast" 2>/dev/null || {
+  log "installing nvdiffrast"
+  pip install -q $CONS ninja
+  pip install -q $CONS --no-build-isolation "git+https://github.com/NVlabs/nvdiffrast.git"
+}
+
+# ------------------------------------------------------------------ FFHQ-UV
+if [ ! -d "$REPO" ]; then
+  log "cloning FFHQ-UV"
+  git clone --depth 1 https://github.com/csbhr/FFHQ-UV.git "$REPO"
+fi
+
+# ------------------------------------------------------- numpy 2 compat patch
+# numpy 2 moved VisibleDeprecationWarning into np.exceptions, and no longer
+# builds ragged arrays implicitly (t[0]/t[1] are 1-element arrays)
+sed -i 's/category=np\.VisibleDeprecationWarning/category=np.exceptions.VisibleDeprecationWarning/;
+        s/trans_params = np\.array(\[w0, h0, s, t\[0\], t\[1\]\])/trans_params = np.array([w0, h0, s, float(t[0]), float(t[1])])/' \
+  "$REPO/RGB_Fitting/utils/preprocess_utils.py"
+
+# --------------------------------------------------- tf1 -> tf2 compat patches
+# The MTCNN + 68-landmark detectors are TF1 frozen graphs; modern tensorflow
+# runs them through the compat.v1 API with eager execution disabled.
+log "applying tensorflow compat patches"
+for f in "$REPO/RGB_Fitting/third_party/mtcnn/detect_face_with_mtcnn.py" \
+         "$REPO/RGB_Fitting/third_party/landmark68/detect_lm68.py"; do
+  grep -q "tensorflow.compat.v1" "$f" || \
+    sed -i 's/^import tensorflow as tf$/import tensorflow.compat.v1 as tf; tf.disable_eager_execution()/' "$f"
+  grep -q "tensorflow.compat.v1" "$f" || { echo "tf compat patch failed: $f"; exit 1; }
+done
+
+# ------------------------------------------------- checkpoints + topo assets
+# Official HF mirror of the project (csbhr/FFHQ-UV) hosts all checkpoints and
+# topology assets — only the files the RGB-fitting pipeline needs are fetched.
+log "downloading checkpoints + topo assets (~2.5 GB, official HF mirror)"
+REPO_DIR="$REPO" python3 - <<'EOF'
+import os
+from huggingface_hub import snapshot_download
+snapshot_download(
+    repo_id="csbhr/FFHQ-UV", repo_type="dataset",
+    local_dir=os.environ["REPO_DIR"],
+    allow_patterns=[
+        "checkpoints/lm_model/68lm_detector.pb",
+        "checkpoints/mtcnn_model/mtcnn_model.pb",
+        "checkpoints/parsing_model/79999_iter.pth",
+        "checkpoints/resnet_model/resnet18-5c106cde.pth",
+        "checkpoints/deep3d_model/epoch_latest.pth",
+        "checkpoints/arcface_model/ms1mv3_arcface_r50_fp16_backbone.pth",
+        "checkpoints/texgan_model/texgan_ffhq_uv.pth",
+        "checkpoints/vgg_model/vgg16.pt",
+        "topo_assets/*",
+    ],
+)
+print("  checkpoints ok")
+EOF
+
+# ------------------------------------------------------------------ verify
+log "verifying imports + files"
+REPO_DIR="$REPO" python3 - <<'EOF'
+import os, warnings
+warnings.filterwarnings("ignore")
+os.environ.setdefault("TF_CPP_MIN_LOG_LEVEL", "3")
+import nvdiffrast, kornia, dominate, cv2, skimage
+import tensorflow.compat.v1 as tf
+assert hasattr(tf, "GraphDef") and hasattr(tf, "InteractiveSession")
+repo = os.environ["REPO_DIR"]
+for f in ["checkpoints/lm_model/68lm_detector.pb",
+          "checkpoints/mtcnn_model/mtcnn_model.pb",
+          "checkpoints/parsing_model/79999_iter.pth",
+          "checkpoints/resnet_model/resnet18-5c106cde.pth",
+          "checkpoints/deep3d_model/epoch_latest.pth",
+          "checkpoints/arcface_model/ms1mv3_arcface_r50_fp16_backbone.pth",
+          "checkpoints/texgan_model/texgan_ffhq_uv.pth",
+          "checkpoints/vgg_model/vgg16.pt",
+          "topo_assets/hifi3dpp_model_info.mat",
+          "topo_assets/unwrap_1024_info.mat",
+          "topo_assets/similarity_Lm3D_all.mat"]:
+    p = os.path.join(repo, f)
+    assert os.path.getsize(p) > 0, p
+print("  all good")
+EOF
+
+log "install complete. run:  ./ffhq_uv_run.sh <photo.jpg>"

ffhq-uv/ffhq_uv_run.sh

@@ -0,0 +1,101 @@
+#!/usr/bin/env bash
+# =============================================================================
+# ffhq_uv_run.sh — single photo in -> 3D face mesh + UV texture out (FFHQ-UV)
+#
+# Usage:
+#   ./ffhq_uv_run.sh <photo.(jpg|png)> [-o <output_dir>] [--force]
+#
+#   -o DIR     output directory  (default: /workspace/outputs/ffhq-uv/<name>)
+#   --force    redo preprocessing/fitting even if cached results exist
+#
+# Outputs (in <output_dir>):
+#   <name>_uv.png         1024x1024 facial UV-texture map (headline output)
+#   <name>_id.obj         fitted head mesh, identity/neutral (HiFi3D++ topology)
+#   <name>_exp.obj        fitted head mesh with expression
+#   <name>.mtl            material linking the meshes to the UV texture
+#   <name>_vis.png        input | skin mask | parsing | landmarks | render | uv
+#   <name>_input_vis.png  preprocessing visualization (step1)
+#   raw/                  full fitting results (all stages, coeffs, latents,
+#                         logs, per-step visuals) + preprocessed input data
+#
+# Requires ffhq_uv_install.sh to have been run first.
+# =============================================================================
+set -euo pipefail
+
+FFHQUV_HOME="${FFHQUV_HOME:-/workspace/ffhq-uv}"
+REPO="$FFHQUV_HOME/FFHQ-UV"
+TEXGAN="texgan_ffhq_uv.pth"
+
+IMG="" ; OUT="" ; FORCE=0
+while [ $# -gt 0 ]; do
+  case "$1" in
+    -o)      OUT="$2"; shift 2;;
+    --force) FORCE=1; shift;;
+    -h|--help) grep '^#' "$0" | sed 's/^# \{0,1\}//'; exit 0;;
+    *)       IMG="$1"; shift;;
+  esac
+done
+[ -n "$IMG" ] && [ -f "$IMG" ] || { echo "usage: $0 <photo.jpg> [-o outdir] [--force]"; exit 1; }
+[ -d "$REPO" ] || { echo "FFHQ-UV not installed - run ffhq_uv_install.sh first"; exit 1; }
+
+# sanitized job name from the file stem (used as folder + file identifier)
+STEM="$(basename "$IMG")"; STEM="${STEM%.*}"
+NAME="$(echo "$STEM" | tr -c 'a-zA-Z0-9' '_' | sed 's/_*$//;s/^_*//' | cut -c1-60)"
+OUT="${OUT:-/workspace/outputs/ffhq-uv/$NAME}"
+
+export TORCH_FORCE_NO_WEIGHTS_ONLY_LOAD=1   # repo predates torch>=2.6 weights_only default
+export TF_CPP_MIN_LOG_LEVEL=3               # silence tensorflow banner
+if [ -z "${TORCH_CUDA_ARCH_LIST:-}" ]; then
+  export TORCH_CUDA_ARCH_LIST="$(python3 -c "import torch; print('%d.%d+PTX' % torch.cuda.get_device_capability())")"
+fi
+
+log() { echo -e "\033[1;35m[ffhq-uv]\033[0m $*"; }
+T0=$(date +%s)
+
+# stage the photo in a folder named after the job (step1 processes a folder)
+STAGE="$FFHQUV_HOME/staging/$NAME"
+PROC="$STAGE/processed_data"
+FITOUT="$FFHQUV_HOME/fitting_output/$NAME"
+if [ "$FORCE" = 1 ]; then rm -rf "$STAGE" "$FITOUT"; fi
+mkdir -p "$STAGE"
+cp -f "$IMG" "$STAGE/$NAME.${IMG##*.}"
+
+cd "$REPO/RGB_Fitting"
+
+log "1/2 preprocessing (MTCNN + 68 landmarks, skin mask, face parsing, alignment)"
+if [ ! -f "$PROC/$NAME.pt" ]; then
+  python3 step1_process_data.py --input_dir "$STAGE" --output_dir "$PROC" \
+      --checkpoints_dir ../checkpoints --topo_dir ../topo_assets > /dev/null 2>&1 \
+    || { echo "preprocessing failed, rerunning verbosely:"; \
+         python3 step1_process_data.py --input_dir "$STAGE" --output_dir "$PROC" \
+             --checkpoints_dir ../checkpoints --topo_dir ../topo_assets; exit 1; }
+fi
+[ -f "$PROC/$NAME.pt" ] || { echo "ERROR: no face detected / preprocessing failed for $IMG"; exit 1; }
+
+log "2/2 RGB fitting (deep3d init -> texture search -> joint tex+shape optimization)"
+if [ ! -f "$FITOUT/$NAME/stage3_mesh_id.obj" ]; then
+  python3 step2_fit_processed_data.py --input_dir "$PROC" --output_dir "$FITOUT" \
+      --checkpoints_dir ../checkpoints --topo_dir ../topo_assets \
+      --texgan_model_name "$TEXGAN" 2>&1 | grep -E "Stage|stage [0-9]|Fit image|Error|Traceback" || true
+fi
+RES="$FITOUT/$NAME"
+[ -f "$RES/stage3_mesh_id.obj" ] || { echo "ERROR: fitting produced no mesh"; exit 1; }
+
+# --------------------------------------------------------------- final outputs
+mkdir -p "$OUT" "$OUT/raw"
+cp "$RES/stage3_uv.png" "$OUT/${NAME}_uv.png"
+# meshes + material, renamed with internal references fixed up
+sed "s/stage3_mesh\.mlt/${NAME}.mtl/"   "$RES/stage3_mesh_id.obj"  > "$OUT/${NAME}_id.obj"
+sed "s/stage3_mesh\.mlt/${NAME}.mtl/"   "$RES/stage3_mesh_exp.obj" > "$OUT/${NAME}_exp.obj"
+sed "s/stage3_uv\.png/${NAME}_uv.png/"  "$RES/stage3_mesh.mlt"     > "$OUT/${NAME}.mtl"
+cp "$RES/stage3_vis.png" "$OUT/${NAME}_vis.png"
+[ -f "${PROC}_vis/$NAME.png" ] && cp "${PROC}_vis/$NAME.png" "$OUT/${NAME}_input_vis.png"
+
+# raw fitting + preprocessing data (all stages, coeffs, GAN latents, logs,
+# per-step visuals, aligned input tensors) for downstream use
+cp -r "$RES/." "$OUT/raw/"
+mkdir -p "$OUT/raw/preprocess"
+cp "$PROC/$NAME.pt" "$OUT/raw/preprocess/${NAME}_processed.pt"
+
+log "done in $(( $(date +%s) - T0 ))s -> $OUT"
+ls -la "$OUT"

outputs/ffhq-uv/stevie/stevie.mtl

@@ -0,0 +1,5 @@
+newmtl blinn1SG
+Ka 0.200000 0.200000 0.200000
+Kd 1.000000 1.000000 1.000000
+Ks 1.000000 1.000000 1.000000
+map_Kd stevie_uv.png