Spaces:
Running on Zero
Running on Zero
| """Gradio demo for single-image 3D object generation and optional mesh export.""" | |
| import argparse | |
| import json | |
| from pathlib import Path | |
| import tempfile | |
| import threading | |
| import time | |
| import gradio as gr | |
| import numpy as np | |
| from PIL import Image, ImageOps | |
| import plotly.graph_objects as go | |
| EXAMPLE_DIRECTORY = Path(__file__).parent / "extra_files/test_cases" | |
| OBJECT_EXAMPLES = [EXAMPLE_DIRECTORY / name for name in ( | |
| "dp.png", "avocado.png", "doctor.png", "jige.png", "link.png", "monster.png", "popmart.png", | |
| )] | |
| def point_cloud_figure(xyz, colors, max_points=5000): | |
| indices = np.linspace(0, len(xyz) - 1, min(len(xyz), max_points), dtype=int) | |
| positions, rgb = xyz[indices], colors[indices] | |
| figure = go.Figure(go.Scatter3d( | |
| x=positions[:, 0].tolist(), y=positions[:, 1].tolist(), z=positions[:, 2].tolist(), | |
| mode="markers", marker=dict(size=2, color=[f"rgb({r},{g},{b})" for r, g, b in rgb]), | |
| hoverinfo="skip", showlegend=False, | |
| )) | |
| center = (xyz.max(axis=0) + xyz.min(axis=0)) / 2 | |
| radius = max(float(np.ptp(xyz, axis=0).max()), 0.1) * 0.65 | |
| axes = [dict(title=label, range=[float(value - radius), float(value + radius)]) | |
| for value, label in zip(center, ["X", "Y", "Z"])] | |
| figure.update_layout(height=420, margin=dict(l=0, r=0, t=0, b=0), | |
| scene=dict(aspectmode="cube", xaxis=axes[0], yaxis=axes[1], zaxis=axes[2], | |
| camera=dict(eye=dict(x=1.5, y=-1.8, z=1.0)))) | |
| return figure | |
| class ObjectDemo: | |
| def __init__(self, model_dir="CaiYuanhao/DiffusionGS", output_dir="outputs/object_demo", | |
| device="cuda:0"): | |
| self.model_dir = model_dir | |
| self.output_dir = Path(output_dir).resolve() | |
| self.device = device | |
| self.pipeline = None | |
| self.lock = threading.Lock() | |
| def generate(self, image, seed, foreground_ratio, force_remove_background, | |
| progress=gr.Progress()): | |
| if image is None: | |
| raise gr.Error("Upload an object image first.") | |
| if int(seed) != seed or not 0 <= seed <= 2**32 - 1: | |
| raise gr.Error("Seed must be an integer between 0 and 2**32 - 1.") | |
| if not 0 < foreground_ratio <= 1: | |
| raise gr.Error("Foreground ratio must be between 0 and 1.") | |
| image = ImageOps.exif_transpose(image).convert("RGBA") | |
| if image.getextrema()[3][1] == 0: | |
| raise gr.Error("The input image is fully transparent.") | |
| import torch | |
| from torchvision.utils import save_image | |
| from diffusionGS.pipline_obj import DiffusionGSPipeline | |
| from diffusionGS.models.gsrenderer.gs_core import SH2RGB | |
| if not torch.cuda.is_available() or torch.device(self.device).type != "cuda": | |
| raise gr.Error("Object generation requires a CUDA GPU.") | |
| with self.lock, torch.cuda.device(self.device): | |
| if self.pipeline is None: | |
| progress(0, desc="Loading object model") | |
| self.pipeline = DiffusionGSPipeline.from_pretrained( | |
| self.model_dir, device=self.device, torch_dtype=torch.float16, | |
| ) | |
| self.output_dir.mkdir(parents=True, exist_ok=True) | |
| directory = Path(tempfile.mkdtemp(prefix="object-", dir=self.output_dir)) | |
| image.save(directory / "input.png") | |
| start = time.perf_counter() | |
| def update_progress(step, total): | |
| progress(0.1 + 0.75 * step / total, desc=f"Generating object: {step}/{total}") | |
| output = self.pipeline( | |
| image, seed=int(seed), foreground_ratio=float(foreground_ratio), | |
| force_remove_background=force_remove_background, | |
| extract_mesh=False, progress_callback=update_progress, | |
| ) | |
| xyz = output.gaussians.get_xyz.detach().float().cpu().numpy() | |
| if len(xyz) == 0 or not np.isfinite(xyz).all() or not torch.isfinite(output.render_images).all(): | |
| raise RuntimeError("Generation produced empty or nonfinite outputs.") | |
| colors = (SH2RGB(output.gaussians._features_dc[:, 0]).detach().float() | |
| .clamp(0, 1).mul(255).byte().cpu().numpy()) | |
| progress(0.9, desc="Saving point cloud and rendered views") | |
| output.gaussians.save_ply(str(directory / "test.ply")) | |
| save_image(output.render_images, str(directory / "test.png")) | |
| views = [] | |
| for i, tensor in enumerate(output.render_images): | |
| pixels = tensor.detach().float().clamp(0, 1).mul(255).byte().permute(1, 2, 0).cpu().numpy() | |
| views.append((Image.fromarray(pixels), f"View {i + 1}")) | |
| elapsed = time.perf_counter() - start | |
| (directory / "result.json").write_text(json.dumps({ | |
| "status": "passed", "seed": int(seed), "foreground_ratio": float(foreground_ratio), | |
| "force_remove_background": bool(force_remove_background), | |
| "generation_seconds": elapsed, "gaussian_count": len(xyz), | |
| "resolution": 512, "mesh_exported": False, | |
| }, indent=2)) | |
| progress(1, desc="Object ready") | |
| return (views, point_cloud_figure(xyz, colors), str(directory / "test.ply"), | |
| str(directory / "test.png"), None, None, None, str(directory), | |
| f"**Done.** {len(xyz):,} Gaussians 路 {elapsed:.1f}s. Mesh export is available below.", | |
| "Mesh has not been extracted for this object.") | |
| def extract_mesh(self, result_directory, progress=gr.Progress()): | |
| if not result_directory: | |
| raise gr.Error("Generate an object first.") | |
| directory = Path(result_directory).resolve() | |
| if not directory.is_relative_to(self.output_dir) or not (directory / "test.ply").is_file(): | |
| raise gr.Error("The generated object is unavailable. Please generate it again.") | |
| destination = directory / "test.obj" | |
| with self.lock: | |
| if not destination.is_file() or destination.stat().st_size == 0: | |
| import torch | |
| from diffusionGS.models.gsrenderer.gs_core import GaussianModel | |
| progress(0.1, desc="Loading generated Gaussians for mesh extraction") | |
| start = time.perf_counter() | |
| with torch.inference_mode(), torch.cuda.device(self.device): | |
| gaussians = GaussianModel(sh_degree=0) | |
| gaussians.load_ply(str(directory / "test.ply")) | |
| gaussians.to(self.device) | |
| progress(0.2, desc="Extracting mesh; this can take several minutes") | |
| mesh = gaussians.extract_mesh() | |
| if not len(mesh.vertices) or not len(mesh.faces) or not np.isfinite(mesh.vertices).all(): | |
| raise RuntimeError("Mesh extraction produced invalid geometry.") | |
| mesh.export(destination) | |
| report = json.loads((directory / "result.json").read_text()) | |
| report.update(mesh_exported=True, mesh_seconds=time.perf_counter() - start, | |
| mesh_vertices=len(mesh.vertices), mesh_faces=len(mesh.faces)) | |
| (directory / "result.json").write_text(json.dumps(report, indent=2)) | |
| from diffusionGS.utils.mesh_preview import export_colored_mesh, PREVIEW_VERSION | |
| preview_path = directory / "mesh_preview.glb" | |
| report = json.loads((directory / "result.json").read_text()) | |
| if (not preview_path.is_file() | |
| or report.get("mesh_preview_version") != PREVIEW_VERSION | |
| or preview_path.stat().st_mtime < max(destination.stat().st_mtime, | |
| (directory / "test.ply").stat().st_mtime)): | |
| import trimesh | |
| progress(0.95, desc="Adding mesh colors") | |
| mesh = trimesh.load(destination, force="mesh", process=False) | |
| export_colored_mesh(mesh, directory / "test.ply", destination, preview_path) | |
| report.update(mesh_preview_version=PREVIEW_VERSION, colored_mesh=True) | |
| (directory / "result.json").write_text(json.dumps(report, indent=2)) | |
| progress(1, desc="Mesh ready") | |
| return (str(preview_path), str(destination), str(preview_path), | |
| "**Mesh ready.** Rotate the color preview or download the GLB / OBJ below.") | |
| def build_demo(model_dir="CaiYuanhao/DiffusionGS", output_dir="outputs/object_demo", device="cuda:0", *, | |
| service=None, resume_result=None, usage=None): | |
| service = service if service is not None else ObjectDemo(model_dir, output_dir, device) | |
| restored = None | |
| gallery_value, cloud_value = None, None | |
| if resume_result is not None: | |
| directory = Path(resume_result).resolve() | |
| # Restoring is explicit for a local launch, never shared between Space visitors. | |
| if not directory.is_relative_to(service.output_dir) or not (directory / "input.png").is_file(): | |
| raise ValueError("The restored result must be inside the demo output directory.") | |
| restored = service.extract_mesh(str(directory)) | |
| from diffusionGS.utils.mesh_preview import read_gaussian_colors | |
| xyz, rgb, _ = read_gaussian_colors(directory / "test.ply") | |
| cloud_value = point_cloud_figure(xyz, np.rint(rgb * 255).astype(np.uint8)) | |
| with Image.open(directory / "test.png") as grid: | |
| size = grid.height - 4 | |
| gallery_value = [(grid.crop((2 + i * (size + 2), 2, 2 + i * (size + 2) + size, 2 + size)), | |
| f"View {i + 1}") for i in range(4)] | |
| with gr.Blocks(title="DiffusionGS 路 Object Generation", theme=gr.themes.Soft()) as demo: | |
| result_directory = gr.State(str(directory) if restored else None) | |
| if usage is not None: | |
| usage.render_header("DiffusionGS 路 Object Generation") | |
| else: | |
| gr.Markdown("# DiffusionGS 路 Object Generation") | |
| gr.Markdown("Generate a 3D object from one image, then optionally extract its mesh.") | |
| with gr.Row(): | |
| with gr.Column(scale=1): | |
| image = gr.Image(value=str(directory / "input.png") if restored else str(OBJECT_EXAMPLES[0]), | |
| type="pil", image_mode="RGBA", | |
| label="Object image", height=300, elem_id="object-input") | |
| gr.Examples([[str(path)] for path in OBJECT_EXAMPLES], inputs=[image], | |
| examples_per_page=7, label="Object examples 路 click to select", | |
| cache_examples=False, api_name=False, elem_id="object-examples") | |
| seed = gr.Number(value=62, precision=0, minimum=0, maximum=2**32 - 1, label="Seed") | |
| ratio = gr.Slider(0.5, 1, value=0.825, step=0.005, label="Foreground ratio") | |
| force = gr.Checkbox(value=False, label="Remove background even when transparency is provided") | |
| generate = gr.Button("Generate object", variant="primary") | |
| gr.Markdown("Background removal runs automatically for opaque images. " | |
| "Transparent images can use their existing alpha mask.") | |
| with gr.Column(scale=2): | |
| views = gr.Gallery(value=gallery_value, label="Rendered object views", columns=4, height=270) | |
| cloud = gr.Plot(value=cloud_value, label="3D point cloud 路 drag to rotate / scroll to zoom") | |
| gr.Markdown("The preview shows up to 5,000 points; the PLY download contains the full object.") | |
| status = gr.Markdown("Previous object restored. Preview or download the result below." | |
| if restored else "Choose an image and generate an object.") | |
| with gr.Row(): | |
| ply = gr.DownloadButton("Download point cloud (.ply)", | |
| value=str(directory / "test.ply") if restored else None, | |
| interactive=bool(restored)) | |
| png = gr.DownloadButton("Download rendered views (.png)", | |
| value=str(directory / "test.png") if restored else None, | |
| interactive=bool(restored)) | |
| mesh_button = gr.Button("Extract mesh", interactive=bool(restored)) | |
| gr.Markdown("Mesh extraction is optional and can take several minutes. " | |
| "It reuses the generated object without another diffusion pass.") | |
| mesh_status = gr.Markdown(restored[3] if restored else None) | |
| mesh_preview = gr.Model3D(value=restored[0] if restored else None, | |
| label="Color mesh preview", height=450, interactive=False, | |
| display_mode="solid", clear_color=(1, 1, 1, 1), | |
| camera_position=(75, 75, None)) | |
| with gr.Row(): | |
| glb = gr.DownloadButton("Download color mesh (.glb)", | |
| value=restored[2] if restored else None, interactive=bool(restored)) | |
| obj = gr.DownloadButton("Download mesh (.obj)", | |
| value=restored[1] if restored else None, interactive=bool(restored)) | |
| generation = generate.click( | |
| service.generate, [image, seed, ratio, force], | |
| [views, cloud, ply, png, mesh_preview, obj, glb, result_directory, status, mesh_status], | |
| api_name="generate_object", concurrency_id="object_gpu", concurrency_limit=1, | |
| ) | |
| generation.success( | |
| lambda: (gr.update(interactive=True), gr.update(interactive=True), | |
| gr.update(interactive=True), gr.update(interactive=False), | |
| gr.update(interactive=False)), | |
| outputs=[ply, png, mesh_button, obj, glb], queue=False, api_name=False, | |
| ) | |
| extraction = mesh_button.click(service.extract_mesh, [result_directory], | |
| [mesh_preview, obj, glb, mesh_status], api_name="extract_mesh", | |
| concurrency_id="object_gpu", concurrency_limit=1) | |
| extraction.success(lambda: (gr.update(interactive=True), gr.update(interactive=True)), | |
| outputs=[obj, glb], queue=False, api_name=False) | |
| if usage is not None: | |
| usage.bind(demo, generation) | |
| return demo.queue(max_size=8, default_concurrency_limit=1) | |
| def main(): | |
| parser = argparse.ArgumentParser(description=__doc__) | |
| parser.add_argument("--model-dir", default="CaiYuanhao/DiffusionGS") | |
| parser.add_argument("--output-dir", default="outputs/object_demo") | |
| parser.add_argument("--device", default="cuda:0") | |
| parser.add_argument("--server-name", default="127.0.0.1") | |
| parser.add_argument("--server-port", type=int, default=7863) | |
| parser.add_argument("--share", action="store_true") | |
| parser.add_argument("--resume-result", help="Restore an existing local demo result without regenerating it") | |
| args = parser.parse_args() | |
| demo = build_demo(args.model_dir, args.output_dir, args.device, resume_result=args.resume_result) | |
| demo.launch(server_name=args.server_name, server_port=args.server_port, | |
| share=args.share, show_error=True, allowed_paths=[str(Path(args.output_dir).resolve())]) | |
| if __name__ == "__main__": | |
| main() | |