import os import random import tempfile from typing import Any, List, Union import spaces import gradio as gr import numpy as np import torch # from gradio_image_prompter import ImagePrompter # from gradio_litmodel3d import LitModel3D from huggingface_hub import snapshot_download from PIL import Image import trimesh from skimage import measure from detailgen3d.pipelines.pipeline_detailgen3d import DetailGen3DPipeline from detailgen3d.inference_utils import generate_dense_grid_points import sys sys.path.append(os.path.dirname(os.path.abspath(__file__))) # Constants MAX_SEED = np.iinfo(np.int32).max TMP_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), "tmp") DTYPE = torch.bfloat16 DEVICE = "cuda" if torch.cuda.is_available() else "cpu" MARKDOWN = """ ## Generating geometry details guided by reference image with [DetailGen3D](https://detailgen3d.github.io/DetailGen3D/) 1. Upload a detailed image of the frontal view and a coarse model. Then clik "Generate Details" to generate the refined result. 2. If you find the generated 3D scene satisfactory, download it by clicking the "Download GLB" button. 3. If you want the refine result to be more consistent with the image, please manually increase the CFG strength. """ EXAMPLES = [ [ "assets/image/503d193a-1b9b-4685-b05f-00ac82f93d7b.png", "assets/model/503d193a-1b9b-4685-b05f-00ac82f93d7b.glb", 42, False, ], [ "assets/image/34933195-9c2c-4271-8d31-a28bc5348b7a.png", "assets/model/34933195-9c2c-4271-8d31-a28bc5348b7a.glb", 2131379184, False, ], [ "assets/image/a5d09c66-1617-465c-aec9-431f48d9a7e1.png", "assets/model/a5d09c66-1617-465c-aec9-431f48d9a7e1.glb", 42, False, ], [ "assets/image/cb7e6c4a-b4dd-483c-9789-3d4887ee7434.png", "assets/model/cb7e6c4a-b4dd-483c-9789-3d4887ee7434.glb", 42, False, ], [ "assets/image/e799e6b4-3b47-40e0-befb-b156af8758ad.png", "assets/model/e799e6b4-3b47-40e0-befb-b156af8758ad.glb", 42, False, ], [ "assets/image/100.png", "assets/model/100.glb", 42, False, ], ] os.makedirs(TMP_DIR, exist_ok=True) local_dir = "pretrained_weights/DetailGen3D" snapshot_download(repo_id="VAST-AI/DetailGen3D", local_dir=local_dir) pipeline = DetailGen3DPipeline.from_pretrained( local_dir ).to(DEVICE, dtype=DTYPE) def load_mesh(mesh_path, num_pc=20480): mesh = trimesh.load(mesh_path,force="mesh") center = mesh.bounding_box.centroid mesh.apply_translation(-center) scale = max(mesh.bounding_box.extents) mesh.apply_scale(1.9 / scale) surface, face_indices = trimesh.sample.sample_surface(mesh, 1000000,) normal = mesh.face_normals[face_indices] rng = np.random.default_rng() ind = rng.choice(surface.shape[0], num_pc, replace=False) surface = torch.FloatTensor(surface[ind]) normal = torch.FloatTensor(normal[ind]) surface = torch.cat([surface, normal], dim=-1).unsqueeze(0).cuda() return surface @torch.no_grad() @torch.autocast(device_type=DEVICE) def run_detailgen3d( pipeline, image, mesh, seed, num_inference_steps, guidance_scale, ): surface = load_mesh(mesh) # image = Image.open(image).convert("RGB") batch_size = 1 # sample query points for decoding box_min = np.array([-1.005, -1.005, -1.005]) box_max = np.array([1.005, 1.005, 1.005]) sampled_points, grid_size, bbox_size = generate_dense_grid_points( bbox_min=box_min, bbox_max=box_max, octree_depth=8, indexing="ij" ) sampled_points = torch.FloatTensor(sampled_points).to(DEVICE, dtype=DTYPE) sampled_points = sampled_points.unsqueeze(0).repeat(batch_size, 1, 1) # inference pipeline sample = pipeline.vae.encode(surface).latent_dist.sample() occ = pipeline(image, latents=sample, sampled_points=sampled_points, guidance_scale=guidance_scale, noise_aug_level=0, num_inference_steps=num_inference_steps).samples[0] # marching cubes grid_logits = occ.view(grid_size).cpu().numpy() vertices, faces, normals, _ = measure.marching_cubes( grid_logits, 0, method="lewiner" ) vertices = vertices / grid_size * bbox_size + box_min mesh = trimesh.Trimesh(vertices.astype(np.float32), np.ascontiguousarray(faces)) return mesh @spaces.GPU(duration=180) @torch.no_grad() @torch.autocast(device_type=DEVICE) def run_refinement( rgb_image: Any, mesh: Any, seed: int, randomize_seed: bool = False, num_inference_steps: int = 50, guidance_scale: float = 4.0, ): if randomize_seed: seed = random.randint(0, MAX_SEED) scene = run_detailgen3d( pipeline, rgb_image, mesh, seed, num_inference_steps, guidance_scale, ) _, tmp_path = tempfile.mkstemp(suffix=".glb", prefix="detailgen3d_", dir=TMP_DIR) scene.export(tmp_path) torch.cuda.empty_cache() return tmp_path, tmp_path, seed # Demo with gr.Blocks() as demo: gr.Markdown(MARKDOWN) with gr.Row(): with gr.Column(): with gr.Row(): # image_prompts = ImagePrompter(label="Input Image", type="pil") image_prompts = gr.Image(label="Example Image", type="pil") with gr.Accordion("Generation Settings", open=False): seed = gr.Slider( label="Seed", minimum=0, maximum=MAX_SEED, step=1, value=0, ) randomize_seed = gr.Checkbox(label="Randomize seed", value=True) num_inference_steps = gr.Slider( label="Number of inference steps", minimum=1, maximum=50, step=1, value=50, ) guidance_scale = gr.Slider( label="CFG scale", minimum=0.0, maximum=50.0, step=0.1, value=10.0, ) gen_button = gr.Button("Generate Details", variant="primary") with gr.Column(): mesh = gr.Model3D(label="Input Coarse Model",camera_position=(90,90,3)) # model_output = LitModel3D(label="Generated GLB", exposure=1.0, height=500,camera_position=(90,90,3)) model_output = gr.Model3D(label="Generated GLB", camera_position=(90,90,3)) download_glb = gr.DownloadButton(label="Download GLB", interactive=False) with gr.Row(): gr.Examples( examples=EXAMPLES, fn=run_refinement, inputs=[image_prompts, mesh, seed, randomize_seed], outputs=[model_output, download_glb, seed], cache_examples=False, ) gen_button.click( run_refinement, inputs=[ image_prompts, mesh, seed, randomize_seed, num_inference_steps, guidance_scale, ], outputs=[model_output, download_glb, seed], ).then(lambda: gr.Button(interactive=True), outputs=[download_glb]) demo.launch()