Spaces:
Running on Zero
Running on Zero
| import spaces # Must precede torch and CUDA-touching project imports. | |
| import logging | |
| import os | |
| from pathlib import Path | |
| import time | |
| from urllib.parse import urlencode | |
| import gradio as gr | |
| import numpy as np | |
| from PIL import Image, ImageOps | |
| import torch | |
| from huggingface_hub import hf_hub_download | |
| from bootstrap import prepare_solver | |
| ROOT = Path(__file__).resolve().parent | |
| MODEL_REVISION = "0a74d1d328ed6f47da5ec3632991c02a0266adca" | |
| prepare_solver() | |
| from engine import ReconstructionEngine | |
| def weight(name: str) -> str: | |
| return hf_hub_download("LinzhouLi/RealDenseFace", filename=name, revision=MODEL_REVISION, token=False) | |
| def flame_model() -> str: | |
| token = os.environ.get("FLAME_READ_TOKEN") | |
| if not token: | |
| raise RuntimeError("Configure the FLAME_READ_TOKEN Space secret with read access to the private demo assets repository.") | |
| return hf_hub_download( | |
| os.environ.get("FLAME_REPO_ID", "LinzhouLi/RealDenseFace-Demo-Assets"), | |
| filename="flame2023.pkl", revision=os.environ.get("FLAME_REVISION", "main"), | |
| token=token, | |
| ) | |
| print("[startup] Loading ViT-B and FaceBox", flush=True) | |
| FLAME_MODEL_PATH = flame_model() | |
| ENGINE = ReconstructionEngine( | |
| ROOT, weight("weights/realdenseface/vitb.pth"), weight("weights/facebox/face_box.pth"), | |
| weight("weights/flame/flame_assets.npz"), FLAME_MODEL_PATH, | |
| ) | |
| print("[startup] Networks loaded; native contexts will be created in GPU requests", flush=True) | |
| def reconstruct(image: Image.Image): | |
| """Reconstruct a single portrait and return dense priors, overlay and mesh visualizations.""" | |
| started = time.perf_counter() | |
| if image is None: | |
| raise gr.Error("Upload a portrait or choose an example first.") | |
| if image.width * image.height > 40_000_000: | |
| raise gr.Error("Please use an image with fewer than 40 million pixels.") | |
| image = ImageOps.exif_transpose(image).convert("RGB") | |
| image.thumbnail((1280, 1280), Image.Resampling.LANCZOS) | |
| if min(image.size) < 32: | |
| raise gr.Error("The image is too small. Please upload a larger portrait.") | |
| try: | |
| dense, overlay, mesh, report = ENGINE.reconstruct(np.ascontiguousarray(np.asarray(image))) | |
| except ValueError as exc: | |
| raise gr.Error(str(exc)) from exc | |
| except Exception: | |
| logging.exception("Reconstruction failed") | |
| raise gr.Error("Reconstruction failed. Please try a clearer portrait.") | |
| report["callback_seconds"] = round(time.perf_counter() - started, 3) | |
| print(f"[request] {report}", flush=True) | |
| return dense, overlay, mesh, f"Inference: {report['seconds']:.2f} s (excludes queue and transfer)." | |
| with gr.Blocks(title="RealDenseFace", delete_cache=(3600, 3600)) as demo: | |
| gr.Markdown("""# RealDenseFace | |
| ### Real-time Monocular 3D Face Reconstruction from Dense UV-space Priors | |
| [Paper](https://arxiv.org/abs/2608.09238) 路 [Project](https://gapszju.github.io/RealDenseFace/) 路 [Code](https://github.com/LinzhouLi/RealDenseFace) 路 [Models](https://huggingface.co/LinzhouLi/RealDenseFace) | |
| Upload a portrait or select an example, then click **Reconstruct**. Powered by RealDenseFace ViT-B. | |
| """) | |
| with gr.Row(): | |
| with gr.Column(scale=1, min_width=300): | |
| with gr.Tabs(): | |
| with gr.Tab("Input"): | |
| image_input = gr.Image(type="pil", label="Input portrait", sources=["upload"], height=420, format="webp") | |
| with gr.Column(scale=1, min_width=300): | |
| with gr.Tabs(): | |
| with gr.Tab("Mesh overlay"): | |
| overlay = gr.Image(label="Mesh overlay", interactive=False, height=420, format="webp") | |
| with gr.Tab("Reconstructed mesh"): | |
| mesh = gr.Image(label="Reconstructed mesh", interactive=False, height=420, format="webp") | |
| with gr.Tab("Dense priors"): | |
| dense = gr.Image(label="Dense priors", interactive=False, height=420, format="webp") | |
| with gr.Row(): | |
| run = gr.Button("Reconstruct", variant="primary", scale=1) | |
| status = gr.Markdown(scale=1) | |
| run.click(reconstruct, inputs=image_input, outputs=[dense, overlay, mesh, status], | |
| api_name="reconstruct", concurrency_limit=1, concurrency_id="gpu") | |
| gr.Examples(examples=[[str(p)] for p in sorted((ROOT / "examples").glob("*.jpg"))], | |
| inputs=[image_input], label="Examples", examples_per_page=10, cache_examples=False) | |
| gr.Markdown("Please use an image with one clearly visible face. ZeroGPU requests may queue before processing.") | |
| issue_url = "https://github.com/LinzhouLi/RealDenseFace/issues/new?" + urlencode({ | |
| "title": "[Demo] Poor reconstruction result", | |
| "body": "### Problem\nPlease describe the problem.\n\n" | |
| "### Input image\nUpload your input image here.\n\n" | |
| "### Result image\nUpload your result image here.\n", | |
| }) | |
| gr.Markdown(f"[Report a poor result]({issue_url})") | |
| image_input.change(lambda: (None, None, None, ""), | |
| inputs=[], outputs=[dense, overlay, mesh, status], | |
| queue=False, api_name=False) | |
| if __name__ == "__main__": | |
| demo.queue(max_size=12).launch(show_error=False, blocked_paths=[FLAME_MODEL_PATH, str(Path(FLAME_MODEL_PATH).resolve())]) | |