Spaces:
Running
on
Zero
Running
on
Zero
File size: 20,358 Bytes
bd9eb72 59ae2c2 e5c64d1 bd9eb72 dc5c57c bd9eb72 133857a bd9eb72 46cadc4 d8d23f5 839dee5 e5c64d1 59ae2c2 46cadc4 839dee5 46cadc4 839dee5 d8d23f5 839dee5 59ae2c2 839dee5 fa06469 e5c64d1 fa06469 e5c64d1 fa06469 e5c64d1 fa06469 45fe0fa fa06469 e5c64d1 fa06469 45fe0fa fa06469 e5c64d1 63cc8f4 e5c64d1 63cc8f4 e5c64d1 fa06469 3916713 63cc8f4 3916713 63cc8f4 fa06469 3916713 fa06469 e5c64d1 63cc8f4 e5c64d1 fa06469 e5c64d1 45fe0fa fa06469 63cc8f4 fa06469 e5c64d1 fa06469 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 289 290 291 292 293 294 295 296 297 298 299 300 301 302 303 304 305 306 307 308 309 310 311 312 313 314 315 316 317 318 319 320 321 322 323 324 325 326 327 328 329 330 331 332 333 334 335 336 337 338 339 340 341 342 343 344 345 346 347 348 349 350 351 352 353 354 355 356 357 358 359 360 361 362 363 364 365 366 367 368 369 370 371 372 373 374 375 376 377 378 379 380 381 382 383 384 385 386 387 388 389 390 391 392 393 394 395 396 397 398 399 400 401 402 403 404 405 406 407 408 409 410 411 412 413 414 415 416 417 418 419 420 421 422 423 424 425 426 427 428 429 430 431 432 433 434 435 436 437 438 439 440 441 442 443 444 445 446 447 448 449 450 451 452 453 454 455 456 457 458 459 460 461 462 463 464 465 466 467 468 469 470 471 472 473 474 475 476 477 478 |
# 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!")
@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()
# 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()
|