import base64 import html import math import os from pathlib import Path import tempfile from typing import Any, List, Optional, Tuple, Union import gradio as gr import numpy as np import spaces import torch from media_loader import load_media_views from zipsplat import ZipSplat, viz from zipsplat.camera import Camera from zipsplat.gaussians import Gaussians # ---------------------------------------------------------------------------- # Viewer HTML Template Loader # ---------------------------------------------------------------------------- VIEWER_HTML_PATH = Path(__file__).parent / "static" / "viewer" / "viewer.html" def get_viewer_template() -> str: if VIEWER_HTML_PATH.exists(): return VIEWER_HTML_PATH.read_text(encoding="utf-8") return "" def build_viewer_iframe(ply_path: str, cx: float = 0, cy: float = 0, cz: float = 0, r: float = 2) -> str: template = get_viewer_template() if not template: return ( '
' "Error: Viewer template file static/viewer/viewer.html not found on server.
" ) try: with open(ply_path, "rb") as f: b64_ply = base64.b64encode(f.read()).decode("utf-8") except Exception as e: return ( f'
' f"Error reading PLY file for viewer: {e}
" ) config_script = f""" """ if "" in template: full_html = template.replace("", f"{config_script}\n", 1) else: full_html = config_script + template escaped_html = html.escape(full_html, quote=True) return f""" """ def render_stroke_turntable( gaussians: Gaussians, path: str, gaussians_per_stroke: int = 32, num_frames: int = 180, fps: int = 30, render_size: int = 512, bg: Tuple[float, float, float] = (1.0, 1.0, 1.0), ) -> str: """Render a video showing the 3D scene painted stroke-by-stroke during orbit. As the camera orbits, strokes (groups of `gaussians_per_stroke` Gaussians) are progressively added frame by frame, revealing the 3D structure from initial strokes to full detail. """ device = gaussians.device fov = torch.tensor(math.radians(55.0)) camera = Camera.from_fov(fov, w=render_size, h=render_size).to(device) center, radius = viz.scene_center_radius(gaussians) poses = viz.orbit_poses(center, radius, num_frames, sweep_deg=360.0) cameras = Camera(camera.data_.unsqueeze(0).expand(num_frames, -1).clone()).to(device) poses = poses.to(device) total_gaussians = gaussians.num_gaussians total_strokes = math.ceil(total_gaussians / gaussians_per_stroke) reveal_frames = int(num_frames * 0.8) bg_t = torch.tensor(bg, device=device, dtype=torch.float32) frames = [] for f in range(num_frames): if f >= reveal_frames: stroke_count = total_strokes else: stroke_count = max(1, math.ceil(((f + 1) / reveal_frames) * total_strokes)) num_visible = min(stroke_count * gaussians_per_stroke, total_gaussians) sub_g = Gaussians(gaussians.data_[:num_visible]) rgb, _ = sub_g.render( cameras[f : f + 1], poses[f : f + 1], mode="RGB", backgrounds=bg_t.expand(1, 3) ) rgb = rgb.float().clamp(0, 1).cpu() frame_np = (rgb.permute(0, 2, 3, 1).numpy() * 255).astype(np.uint8)[0] frames.append(frame_np) frames_np = np.stack(frames, axis=0) viz.save_video(frames_np, path, fps=fps) return path def normalize_uploaded_paths(files: Any) -> List[Path]: """Normalize Gradio return types (None, str, Path, list, dict, file objects) to a list of Path objects.""" if files is None: return [] if isinstance(files, (str, Path)): return [Path(files)] if isinstance(files, dict): p = files.get("name") or files.get("path") orig = files.get("orig_name") if p: path_obj = Path(p) if orig and not path_obj.suffix: orig_suffix = Path(orig).suffix if orig_suffix: new_path = path_obj.with_name(path_obj.name + orig_suffix) if not new_path.exists() and path_obj.exists(): try: os.symlink(path_obj, new_path) return [new_path] except Exception: pass return [path_obj] return [] paths = [] if isinstance(files, (list, tuple)): for item in files: paths.extend(normalize_uploaded_paths(item)) elif hasattr(files, "name"): p = files.name orig = getattr(files, "orig_name", None) path_obj = Path(p) if orig and not path_obj.suffix: orig_suffix = Path(orig).suffix if orig_suffix: new_path = path_obj.with_name(path_obj.name + orig_suffix) if not new_path.exists() and path_obj.exists(): try: os.symlink(path_obj, new_path) return [new_path] except Exception: pass paths.append(path_obj) return paths # ---------------------------------------------------------------------------- # ZipSplat Model Initialization (loaded CPU-side, ZeroGPU moves to GPU) # ---------------------------------------------------------------------------- model = ZipSplat(weights="zipsplat").eval() @spaces.GPU def generate(files: Any): paths = normalize_uploaded_paths(files) if not paths: default_placeholder = ( '
' "Please upload images, videos, or Live Photo pairs.
" ) return None, None, None, default_placeholder, "No files uploaded." try: result = load_media_views(paths, max_views=500, frames_per_video=8) except Exception as e: err_msg = f"Error loading media: {e}" err_placeholder = ( f'
' f"{err_msg}
" ) return None, None, None, err_placeholder, err_msg if not result.images: status = "Error: No valid views could be loaded." if result.warnings: status += "\n\nWarnings:\n" + "\n".join(result.warnings) err_placeholder = ( f'
' f"{status}
" ) return None, None, None, err_placeholder, status import traceback try: device = torch.device("cuda" if torch.cuda.is_available() else "cpu") model.to(device) gaussians = model(result.images)[0] video_path = tempfile.mktemp(suffix=".mp4") stroke_video_path = tempfile.mktemp(suffix="_stroke.mp4") ply_path = tempfile.mktemp(suffix=".ply") viz.turntable(gaussians, video_path, sweep_deg=None) render_stroke_turntable(gaussians, stroke_video_path) gaussians.save_ply(ply_path) center, radius = viz.scene_center_radius(gaussians) cx, cy, cz = center.tolist() viewer_html = build_viewer_iframe(ply_path, cx=cx, cy=cy, cz=cz, r=radius) total_g = gaussians.num_gaussians total_s = math.ceil(total_g / 32) status = ( f"Loaded {result.image_count} still image(s) and selected " f"{result.selected_video_frames} of {result.decoded_video_frames} decoded video frame(s).\n" f"Total ZipSplat views: {len(result.images)}.\n" f"Generated 3D scene with {total_g:,} Gaussians across {total_s:,} strokes (32 Gaussians/stroke)." ) if result.warnings: status += "\n\nWarnings:\n" + "\n".join(result.warnings) return video_path, stroke_video_path, ply_path, viewer_html, status except Exception as e: err_msg = f"Backend Crash during generation:\n{traceback.format_exc()}" err_html = ( f'
' f"{err_msg.replace(chr(10), '
')}
" ) return None, None, None, err_html, err_msg # ---------------------------------------------------------------------------- # SEO Head Metadata & Social OpenGraph Tags # ---------------------------------------------------------------------------- SEO_HEAD_HTML = """ """ # ---------------------------------------------------------------------------- # Gradio UI Layout # ---------------------------------------------------------------------------- with gr.Blocks( title="ZipSplatPlus: Real-time 3D Gaussian Splatting & Stroke Playback Viewer", head=SEO_HEAD_HTML, ) as demo: gr.Markdown("# 🔥 ZipSplatPlus: Real-time 3D Gaussian Splatting & Stroke Playback") gr.Markdown( "ZeroGPU Powered Demo for [veichta/zipsplat](https://huggingface.co/veichta/zipsplat). " "Upload images (JPEG, PNG, WebP, HEIC/HEIF), videos (MOV, MP4, M4V), or Apple Live Photo pairs " "(HEIC + MOV) to generate a 3D Gaussian Splatting scene and **play back each stroke step-by-step** in real time!" ) with gr.Row(): with gr.Column(scale=1): media_in = gr.File( label="Upload Images, Videos, or Live Photo Pairs", file_count="multiple", type="filepath", file_types=[ ".jpg", ".jpeg", ".png", ".webp", ".heic", ".heif", ".mov", ".mp4", ".m4v", ], ) btn = gr.Button("🚀 Generate 3D Scene", variant="primary") with gr.Column(scale=2): viewer_out = gr.HTML( label="Interactive 3D Gaussian & Stroke Viewer (Spark.js)", value='
Upload media and click "Generate 3D Scene" to view interactive 3D Gaussian Splat with stroke-by-stroke playback
', ) with gr.Tabs(): with gr.TabItem("🖌️ Stroke-by-Stroke Playback Video"): stroke_video_out = gr.Video(label="Stroke-by-Stroke Playback Video") with gr.TabItem("🎬 360° Turntable Video"): video_out = gr.Video(label="360° Turntable Video Render") with gr.Row(): file_out = gr.File(label="Download 3D Model (.ply)") status_out = gr.Textbox(label="Status & Warnings", interactive=False) btn.click( fn=generate, inputs=[media_in], outputs=[video_out, stroke_video_out, file_out, viewer_out, status_out], js="""function(files) { if (!files || (Array.isArray(files) && files.length === 0)) { alert("Please upload at least one image or video before generating."); throw new Error("No files uploaded"); } return [files]; }""", ) if __name__ == "__main__": demo.launch(allowed_paths=[tempfile.gettempdir()])