Spaces:
Running
on
Zero
Running
on
Zero
| # Copyright (C) 2025, FaceLift Research Group | |
| # https://github.com/weijielyu/FaceLift | |
| # | |
| # This software is free for non-commercial, research and evaluation use | |
| # under the terms of the LICENSE.md file. | |
| # | |
| # For inquiries contact: wlyu3@ucmerced.edu | |
| """ | |
| FaceLift: Single Image 3D Face Reconstruction | |
| Generates 3D head models from single images using multi-view diffusion and GS-LRM. | |
| """ | |
| import json | |
| from pathlib import Path | |
| from datetime import datetime | |
| import uuid | |
| import time | |
| import shutil | |
| import zipfile | |
| import gradio as gr | |
| import numpy as np | |
| import torch | |
| import yaml | |
| from easydict import EasyDict as edict | |
| from einops import rearrange | |
| from PIL import Image | |
| from huggingface_hub import snapshot_download | |
| import spaces | |
| # Install diff-gaussian-rasterization at runtime (requires GPU) | |
| import subprocess | |
| import sys | |
| import os | |
| # No custom viewer - using standard Gradio outputs only | |
| # ----------------------------- | |
| # 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() | |
| try: | |
| import torch as _torch | |
| if _torch.cuda.is_available(): | |
| maj, minr = _torch.cuda.get_device_capability() | |
| arch = f"{maj}.{minr}" # e.g., "9.0" on H100/H200, "8.0" on A100 | |
| env["TORCH_CUDA_ARCH_LIST"] = f"{arch}+PTX" | |
| else: | |
| # Build stage may not see a GPU on HF Spaces: compile a cross-arch set | |
| env["TORCH_CUDA_ARCH_LIST"] = "8.0;8.6;8.9;9.0+PTX" | |
| except Exception: | |
| env["TORCH_CUDA_ARCH_LIST"] = "8.0;8.6;8.9;9.0+PTX" | |
| # (Optional) side-step allocator+NVML quirks in restrictive containers | |
| env.setdefault("PYTORCH_NO_CUDA_MEMORY_CACHING", "1") | |
| subprocess.check_call( | |
| [sys.executable, "-m", "pip", "install", | |
| "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 | |
| def create_viewer_html(ply_path: str, output_dir: Path) -> str: | |
| """Create a standalone HTML viewer with embedded PLY path.""" | |
| viewer_html_path = output_dir / "viewer.html" | |
| ply_filename = Path(ply_path).name | |
| html_content = f"""<!DOCTYPE html> | |
| <html lang="en"> | |
| <head> | |
| <meta charset="UTF-8"> | |
| <meta name="viewport" content="width=device-width, initial-scale=1.0"> | |
| <title>FaceLift 3D Viewer</title> | |
| <style> | |
| body {{ margin: 0; overflow: hidden; background: #000; }} | |
| #canvas {{ width: 100vw; height: 100vh; display: block; }} | |
| #info {{ position: absolute; top: 10px; left: 10px; color: white; font-family: Arial; | |
| background: rgba(0,0,0,0.7); padding: 10px; border-radius: 5px; }} | |
| #fps {{ position: absolute; top: 10px; right: 10px; color: white; font-family: monospace; | |
| background: rgba(0,0,0,0.7); padding: 10px; border-radius: 5px; }} | |
| </style> | |
| </head> | |
| <body> | |
| <div id="info"> | |
| <strong>FaceLift 3D Gaussian Splat Viewer</strong><br> | |
| Controls: Drag to rotate | Scroll to zoom | Right-drag to pan | |
| </div> | |
| <div id="fps">Loading...</div> | |
| <canvas id="canvas"></canvas> | |
| <script src="viewer.js"></script> | |
| <script> | |
| // Auto-load the PLY file | |
| fetch('{ply_filename}') | |
| .then(response => response.arrayBuffer()) | |
| .then(buffer => {{ | |
| const event = new MessageEvent('message', {{ data: {{ ply: buffer }} }}); | |
| if (window.worker) window.worker.dispatchEvent(event); | |
| }}); | |
| </script> | |
| </body> | |
| </html>""" | |
| viewer_html_path.write_text(html_content) | |
| return str(viewer_html_path) | |
| # 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...") | |
| 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) | |
| self.mvdiffusion_pipeline.unet.enable_xformers_memory_efficient_attention() | |
| 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, xformers enabled!") | |
| 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() | |
| # 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), | |
| ) | |
| ply_path = output_dir / "gaussians.ply" | |
| 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) | |
| # Create standalone HTML viewer | |
| viewer_js_src = Path(__file__).parent / "viewer.js" | |
| if viewer_js_src.exists(): | |
| shutil.copy2(viewer_js_src, output_dir / "viewer.js") | |
| viewer_html_path = create_viewer_html(str(ply_path), output_dir) | |
| else: | |
| viewer_html_path = None | |
| # Final CUDA cache clear | |
| torch.cuda.empty_cache() | |
| return str(input_path), str(multiview_path), str(output_path), \ | |
| str(turntable_path), str(ply_path), viewer_html_path | |
| 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)}") | |
| # No custom head needed - using Gradio's built-in Model3D component | |
| 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: | |
| # Wrapper to return outputs for display | |
| def _generate_and_filter_outputs(image_path, auto_crop, guidance_scale, random_seed, num_steps): | |
| input_path, multiview_path, output_path, turntable_path, ply_path, viewer_html = \ | |
| pipeline.generate_3d_head(image_path, auto_crop, guidance_scale, random_seed, num_steps) | |
| return turntable_path, output_path, viewer_html, ply_path | |
| gr.Markdown("## FaceLift: Single Image 3D Face Reconstruction.") | |
| 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. | |
| """) | |
| 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_video = gr.PlayableVideo(label="Turntable Animation (360° View)", height=600) | |
| out_recon = gr.Image(label="3D Reconstruction Views") | |
| out_viewer = gr.File(label="📦 Download Interactive 3D Viewer (HTML + PLY)") | |
| out_ply = gr.File(label="Download 3D Model (.ply only)") | |
| gr.Markdown(""" | |
| **💡 How to View the 3D Model:** | |
| **Option 1 (Recommended):** Download the "Interactive 3D Viewer" package above, extract it, and open `viewer.html` in your browser for full interactive exploration! | |
| **Option 2:** Download just the `.ply` file and view it online: | |
| - [antimatter15/splat](https://antimatter15.com/splat/) - Drag & drop viewer (based on [GitHub repo](https://github.com/antimatter15/splat)) | |
| - [SuperSplat Viewer](https://playcanvas.com/supersplat/editor) - Advanced editor | |
| - [Polycam](https://poly.cam/) or [Luma AI](https://lumalabs.ai/) - Mobile apps | |
| """) | |
| # Function to package viewer as zip | |
| def package_viewer(viewer_html_path): | |
| if not viewer_html_path or not Path(viewer_html_path).exists(): | |
| return None | |
| output_dir = Path(viewer_html_path).parent | |
| zip_path = output_dir / "facelift_3d_viewer.zip" | |
| with zipfile.ZipFile(zip_path, 'w') as zipf: | |
| # Add README | |
| readme_src = Path(__file__).parent / "VIEWER_README.md" | |
| if readme_src.exists(): | |
| zipf.write(readme_src, "README.txt") | |
| # Add HTML viewer | |
| zipf.write(viewer_html_path, "viewer.html") | |
| # Add JS | |
| js_path = output_dir / "viewer.js" | |
| if js_path.exists(): | |
| zipf.write(js_path, "viewer.js") | |
| # Add PLY | |
| ply_path = output_dir / "gaussians.ply" | |
| if ply_path.exists(): | |
| zipf.write(ply_path, "gaussians.ply") | |
| return str(zip_path) | |
| # Run generation and display all outputs | |
| run_btn.click( | |
| fn=_generate_and_filter_outputs, | |
| inputs=[in_image, auto_crop, guidance, seed, steps], | |
| outputs=[out_video, out_recon, out_viewer, out_ply], | |
| ).then( | |
| fn=package_viewer, | |
| inputs=[out_viewer], | |
| outputs=[out_viewer], | |
| ) | |
| demo.queue(max_size=10) | |
| demo.launch(share=True, server_name="0.0.0.0", server_port=7860, show_error=True) | |
| if __name__ == "__main__": | |
| main() | |