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()