# Copyright © 2025, Adobe Inc. and its licensors. All rights reserved. # # This file is licensed under the Adobe Research License. You may obtain a copy # of the license at https://raw.githubusercontent.com/adobe-research/FaceLift/main/LICENSE.md """ FaceLift: Single Image 3D Face Reconstruction Generates 3D head models from single images using multi-view diffusion and GS-LRM. Note: To enable the interactive 3D viewer, this Space needs write access to wlyu/FaceLift_demo. Set the HF_TOKEN environment variable in Space settings with a token that has write access. """ # Disable HF fast transfer if hf_transfer is not installed # This MUST be done before importing huggingface_hub import os if os.environ.get("HF_HUB_ENABLE_HF_TRANSFER") == "1": try: import hf_transfer except ImportError: print("⚠️ hf_transfer not available, disabling fast download") os.environ["HF_HUB_ENABLE_HF_TRANSFER"] = "0" import json from pathlib import Path from datetime import datetime import random import gradio as gr import numpy as np import torch print(f"[GPU-DIAG] torch {torch.__version__} / CUDA {torch.version.cuda}") import yaml from easydict import EasyDict as edict from einops import rearrange from PIL import Image from huggingface_hub import snapshot_download, HfApi import spaces # Install diff-gaussian-rasterization at runtime (requires GPU) import subprocess import sys # Outputs directory for generated files (ensures folder exists even if CWD differs) OUTPUTS_DIR = Path.cwd() / "outputs" OUTPUTS_DIR.mkdir(exist_ok=True) def _log_viewer_file(ply_path: Path): """Print a concise JSON line about the viewer file so users can debug from Space logs.""" info = { "ply_path": str(Path(ply_path).absolute()), "exists": Path(ply_path).exists(), "size_bytes": (Path(ply_path).stat().st_size if Path(ply_path).exists() else None) } print("[VIEWER-RETURN]", json.dumps(info)) def upload_ply_to_hf(ply_path: Path, repo_id: str = "wlyu/FaceLift_demo") -> str: """Upload PLY file to HuggingFace and return the public URL.""" try: # Get HF token from environment (automatically available in HF Spaces) hf_token = os.environ.get("HF_TOKEN") or os.environ.get("HUGGING_FACE_HUB_TOKEN") if not hf_token: print("⚠️ No HF_TOKEN found in environment, skipping upload") return None api = HfApi(token=hf_token) ply_filename = ply_path.name # Upload to tmp_ply folder path_in_repo = f"tmp_ply/{ply_filename}" print(f"Uploading {ply_filename} to HuggingFace...") api.upload_file( path_or_fileobj=str(ply_path), path_in_repo=path_in_repo, repo_id=repo_id, repo_type="model", token=hf_token, ) # Return the public URL hf_url = f"https://huggingface.co/{repo_id}/resolve/main/{path_in_repo}" print(f"✓ Uploaded to: {hf_url}") return hf_url except Exception as e: print(f"⚠️ Failed to upload to HuggingFace: {e}") print(" Make sure the Space has write access to the repository") return None # ----------------------------- # Ensure diff-gaussian-rasterization builds for current GPU # ----------------------------- try: import diff_gaussian_rasterization # noqa: F401 except ImportError: print("Installing diff-gaussian-rasterization (compiling for detected CUDA arch)...") env = os.environ.copy() # ZeroGPU ships a cu130 torch and a matching CUDA 13.0 toolkit, so compile # against the system toolkit (no version mismatch). Include Blackwell sm_120 # (the ZeroGPU GPU is capability 12.0); the build stage may not see a GPU, so # use a fixed cross-arch set rather than get_device_capability. env["TORCH_CUDA_ARCH_LIST"] = "8.0;8.6;8.9;9.0;12.0+PTX" # --no-build-isolation: the package's setup.py imports torch to detect the # CUDA arch, but pip's isolated build env has no torch. Reuse the main env # (torch + ninja are already installed) so the extension can compile. subprocess.check_call( [sys.executable, "-m", "pip", "install", "--no-build-isolation", "git+https://github.com/graphdeco-inria/diff-gaussian-rasterization"], env=env, ) import diff_gaussian_rasterization # noqa: F401 from gslrm.model.gaussians_renderer import render_turntable, imageseq2video from mvdiffusion.pipelines.pipeline_mvdiffusion_unclip import StableUnCLIPImg2ImgPipeline from utils_folder.face_utils import preprocess_image, preprocess_image_without_cropping # HuggingFace repository configuration HF_REPO_ID = "wlyu/OpenFaceLift" def download_weights_from_hf() -> Path: """Download model weights from HuggingFace if not already present. Returns: Path to the downloaded repository """ workspace_dir = Path(__file__).parent # Check if weights already exist locally mvdiffusion_path = workspace_dir / "checkpoints/mvdiffusion/pipeckpts" gslrm_path = workspace_dir / "checkpoints/gslrm/ckpt_0000000000021125.pt" if mvdiffusion_path.exists() and gslrm_path.exists(): print("Using local model weights") return workspace_dir print(f"Downloading model weights from HuggingFace: {HF_REPO_ID}") print("This may take a few minutes on first run...") # Download to local directory snapshot_download( repo_id=HF_REPO_ID, local_dir=str(workspace_dir / "checkpoints"), local_dir_use_symlinks=False, ) print("Model weights downloaded successfully!") return workspace_dir class FaceLiftPipeline: """Pipeline for FaceLift 3D head generation from single images.""" def __init__(self): # Download weights from HuggingFace if needed workspace_dir = download_weights_from_hf() # Setup paths self.output_dir = workspace_dir / "outputs" self.examples_dir = workspace_dir / "examples" self.output_dir.mkdir(exist_ok=True) # Parameters self.device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") self.image_size = 512 self.camera_indices = [2, 1, 0, 5, 4, 3] # Load models (keep on CPU for ZeroGPU compatibility) print("Loading models... (gradio", getattr(gr, "__version__", "unknown"), ")") try: self.mvdiffusion_pipeline = StableUnCLIPImg2ImgPipeline.from_pretrained( str(workspace_dir / "checkpoints/mvdiffusion/pipeckpts"), torch_dtype=torch.float16, ) # Don't move to device or enable xformers here - will be done in GPU-decorated function self._models_on_gpu = False with open(workspace_dir / "configs/gslrm.yaml", "r") as f: config = edict(yaml.safe_load(f)) module_name, class_name = config.model.class_name.rsplit(".", 1) module = __import__(module_name, fromlist=[class_name]) ModelClass = getattr(module, class_name) self.gs_lrm_model = ModelClass(config) checkpoint = torch.load( workspace_dir / "checkpoints/gslrm/ckpt_0000000000021125.pt", map_location="cpu" ) # Filter out loss_calculator weights (training-only, not needed for inference) state_dict = {k: v for k, v in checkpoint["model"].items() if not k.startswith("loss_calculator.")} self.gs_lrm_model.load_state_dict(state_dict) # Keep on CPU initially - will move to GPU in decorated function self.color_prompt_embedding = torch.load( workspace_dir / "mvdiffusion/fixed_prompt_embeds_6view/clr_embeds.pt", map_location="cpu" ) with open(workspace_dir / "utils_folder/opencv_cameras.json", 'r') as f: self.cameras_data = json.load(f)["frames"] print("Models loaded successfully!") except Exception as e: print(f"Error loading models: {e}") import traceback traceback.print_exc() raise def _move_models_to_gpu(self): """Move models to GPU and enable optimizations. Called within @spaces.GPU context.""" if not self._models_on_gpu and torch.cuda.is_available(): print("Moving models to GPU...") self.device = torch.device("cuda:0") self.mvdiffusion_pipeline.to(self.device) # NOTE: no enable_xformers_memory_efficient_attention() here. ZeroGPU's # Blackwell GPU (compute capability 12.0 / sm_120) is too new for # xformers' prebuilt kernels (fa3F/cutlassF cap out at 9.0), so the call # raises NotImplementedError. diffusers falls back to PyTorch's built-in # SDPA, which supports sm_120 and is the standard efficient-attention path. self.gs_lrm_model.to(self.device) self.gs_lrm_model.eval() # Set to eval mode self.color_prompt_embedding = self.color_prompt_embedding.to(self.device) self._models_on_gpu = True torch.cuda.empty_cache() # Clear cache after moving models print("Models on GPU (using PyTorch SDPA attention)!") @spaces.GPU(duration=120) def generate_3d_head(self, image_path, auto_crop=True, guidance_scale=3.0, random_seed=4, num_steps=50): """Generate 3D head from single image.""" try: # Move models to GPU now that we're in the GPU context self._move_models_to_gpu() # Diagnostic: confirm torch build and actual VRAM available inside the # ZeroGPU worker. Distinguishes a real memory cap from a spurious malloc # failure (the "NVML_SUCCESS" assert is just PyTorch's OOM reporter). try: free_b, total_b = torch.cuda.mem_get_info() print(f"[GPU-DIAG] torch {torch.__version__} / CUDA {torch.version.cuda} | " f"free {free_b/2**30:.2f} GiB / total {total_b/2**30:.2f} GiB | " f"reserved {torch.cuda.memory_reserved()/2**30:.2f} GiB") except Exception as _diag_e: print(f"[GPU-DIAG] mem_get_info failed: {_diag_e}") # Setup output directory timestamp = datetime.now().strftime("%Y%m%d_%H%M%S") output_dir = self.output_dir / timestamp output_dir.mkdir(exist_ok=True) # Preprocess input original_img = np.array(Image.open(image_path)) input_image = preprocess_image(original_img) if auto_crop else \ preprocess_image_without_cropping(original_img) if input_image.size != (self.image_size, self.image_size): input_image = input_image.resize((self.image_size, self.image_size)) input_path = output_dir / "input.png" input_image.save(input_path) # Generate multi-view images generator = torch.Generator(device=self.mvdiffusion_pipeline.unet.device) generator.manual_seed(random_seed) result = self.mvdiffusion_pipeline( input_image, None, prompt_embeds=self.color_prompt_embedding, height=self.image_size, width=self.image_size, guidance_scale=guidance_scale, num_images_per_prompt=1, num_inference_steps=num_steps, generator=generator, eta=1.0, ) selected_views = result.images[:6] # Save multi-view composite multiview_image = Image.new("RGB", (self.image_size * 6, self.image_size)) for i, view in enumerate(selected_views): multiview_image.paste(view, (self.image_size * i, 0)) multiview_path = output_dir / "multiview.png" multiview_image.save(multiview_path) # Move diffusion model to CPU to free GPU memory for GS-LRM print("Moving diffusion model to CPU to free memory...") self.mvdiffusion_pipeline.to("cpu") # Delete intermediate variables to free memory del result, generator torch.cuda.empty_cache() torch.cuda.synchronize() # Prepare 3D reconstruction input view_arrays = [np.array(view) for view in selected_views] lrm_input = torch.from_numpy(np.stack(view_arrays, axis=0)).float() lrm_input = lrm_input[None].to(self.device) / 255.0 lrm_input = rearrange(lrm_input, "b v h w c -> b v c h w") # Prepare camera parameters selected_cameras = [self.cameras_data[i] for i in self.camera_indices] fxfycxcy_list = [[c["fx"], c["fy"], c["cx"], c["cy"]] for c in selected_cameras] c2w_list = [np.linalg.inv(np.array(c["w2c"])) for c in selected_cameras] fxfycxcy = torch.from_numpy(np.stack(fxfycxcy_list, axis=0).astype(np.float32)) c2w = torch.from_numpy(np.stack(c2w_list, axis=0).astype(np.float32)) fxfycxcy = fxfycxcy[None].to(self.device) c2w = c2w[None].to(self.device) batch_indices = torch.stack([ torch.zeros(lrm_input.size(1)).long(), torch.arange(lrm_input.size(1)).long(), ], dim=-1)[None].to(self.device) batch = edict({ "image": lrm_input, "c2w": c2w, "fxfycxcy": fxfycxcy, "index": batch_indices, }) # Ensure GS-LRM model is on GPU if next(self.gs_lrm_model.parameters()).device.type == "cpu": print("Moving GS-LRM model to GPU...") self.gs_lrm_model.to(self.device) torch.cuda.empty_cache() # Final memory cleanup before reconstruction torch.cuda.empty_cache() # Run 3D reconstruction with torch.no_grad(), torch.autocast(enabled=True, device_type="cuda", dtype=torch.float16): result = self.gs_lrm_model.forward(batch, create_visual=False, split_data=True) comp_image = result.render[0].unsqueeze(0).detach() gaussians = result.gaussians[0] # Clear CUDA cache after reconstruction torch.cuda.empty_cache() # Save filtered gaussians filtered_gaussians = gaussians.apply_all_filters( cam_origins=None, opacity_thres=0.04, scaling_thres=0.2, floater_thres=0.75, crop_bbx=[-0.91, 0.91, -0.91, 0.91, -1.0, 1.0], nearfar_percent=(0.0001, 1.0), ) # Generate random filename to support multiple concurrent users random_id = random.randint(0, 999) ply_filename = f"gaussians_{random_id:03d}.ply" ply_path = output_dir / ply_filename filtered_gaussians.save_ply(str(ply_path)) # Save output image comp_image = rearrange(comp_image, "x v c h w -> (x h) (v w) c") comp_image = (comp_image.cpu().numpy() * 255.0).clip(0, 255).astype(np.uint8) output_path = output_dir / "output.png" Image.fromarray(comp_image).save(output_path) # Generate turntable video turntable_resolution = 512 num_turntable_views = 180 turntable_frames = render_turntable(gaussians, rendering_resolution=turntable_resolution, num_views=num_turntable_views) turntable_frames = rearrange(turntable_frames, "h (v w) c -> v h w c", v=num_turntable_views) turntable_frames = np.ascontiguousarray(turntable_frames) turntable_path = output_dir / "turntable.mp4" imageseq2video(turntable_frames, str(turntable_path), fps=30) # Log the viewer file for quick debugging _log_viewer_file(ply_path) # Upload PLY to HuggingFace for public access hf_ply_url = upload_ply_to_hf(ply_path) # Final CUDA cache clear torch.cuda.empty_cache() # Create viewer HTML if hf_ply_url: # Successfully uploaded - show iframe viewer viewer_url = f"https://www.wlyu.me/FaceLift/splat/index.html?url={hf_ply_url}" viewer_html = f"""
🎮 Open Interactive Viewer in New Tab

Drag to rotate • Scroll to zoom • Right-click to pan

""" else: # Upload failed - provide instructions to use viewer manually viewer_base_url = "https://www.wlyu.me/FaceLift/splat/index.html" viewer_html = f"""
🎮

Interactive 3D Viewer

Download the PLY file below, then drag and drop it into the viewer
or use the viewer with a public URL

🔗 Open Interactive Viewer

Controls: Drag to rotate • Scroll to zoom • Right-click to pan

""" return ( viewer_html, # Viewer HTML (top) str(output_path), # Reconstruction grid str(turntable_path), # Turntable video str(ply_path), # Download file ) except Exception as e: import traceback error_details = traceback.format_exc() print(f"Error details:\n{error_details}") raise gr.Error(f"Generation failed: {str(e)}") def main(): """Run the FaceLift application.""" pipeline = FaceLiftPipeline() # Prepare examples (same as before) examples = [] if pipeline.examples_dir.exists(): examples = [[str(f), True, 3.0, 4, 50] for f in sorted(pipeline.examples_dir.iterdir()) if f.suffix.lower() in {'.png', '.jpg', '.jpeg'}] with gr.Blocks(title="FaceLift: Single Image 3D Face Reconstruction") as demo: gr.Markdown("## [ICCV 2025] FaceLift: Learning Generalizable Single Image 3D Face Reconstruction from Synthetic Heads") gr.Markdown(""" ### 💡 Tips for Best Results - Works best with near-frontal portrait images. - The provided checkpoints were not trained with accessories (glasses, hats, etc.). Portraits containing accessories may produce suboptimal results. - If face detection fails, try disabling auto-cropping and manually crop to square. - Inference complete when the turntable video is generated, the interactive 3D gaussian might take several seconds to load. """) with gr.Row(): with gr.Column(scale=1): in_image = gr.Image(type="filepath", label="Input Portrait Image") auto_crop = gr.Checkbox(value=True, label="Auto Cropping") guidance = gr.Slider(1.0, 10.0, 3.0, step=0.1, label="Guidance Scale") seed = gr.Number(value=4, label="Random Seed") steps = gr.Slider(10, 100, 50, step=5, label="Generation Steps") run_btn = gr.Button("Generate 3D Head", variant="primary") # Examples (match input signature) if examples: gr.Examples( examples=examples, inputs=[in_image, auto_crop, guidance, seed, steps], examples_per_page=10, ) with gr.Column(scale=1): out_viewer = gr.HTML(label="🎮 Interactive 3D Viewer") out_recon = gr.Image(label="3D Reconstruction Views") out_video = gr.PlayableVideo(label="Turntable Animation (360° View)", height=600) out_ply = gr.File(label="Download 3D Model (.ply)") # Wrapper: call the pipeline and forward outputs in the exact order expected def _generate_and_filter_outputs(image_path, auto_crop, guidance_scale, random_seed, num_steps): return pipeline.generate_3d_head(image_path, auto_crop, guidance_scale, random_seed, num_steps) # Run generation and display all outputs run_btn.click( fn=_generate_and_filter_outputs, inputs=[in_image, auto_crop, guidance, seed, steps], outputs=[out_viewer, out_recon, out_video, out_ply], ) demo.queue(max_size=10) # ssr_mode=False: Gradio's SSR front-server forks/shares the CUDA context in a # way that breaks NVML in the ZeroGPU worker, causing the caching allocator to # assert (NVML_SUCCESS == r) during inference. share=True is unsupported on Spaces. demo.launch(server_name="0.0.0.0", server_port=7860, show_error=True, ssr_mode=False) if __name__ == "__main__": main()