| from extend3d import Extend3D
|
| from trellis.utils import render_utils, postprocessing_utils
|
|
|
| import imageio
|
| import random
|
| import uuid
|
| from pathlib import Path
|
|
|
| import numpy as np
|
| import torch
|
| import gradio as gr
|
|
|
| MODEL_ID = "microsoft/TRELLIS-image-large"
|
| DEFAULT_OUTPUT_DIR = "./output"
|
|
|
|
|
|
|
|
|
|
|
| PIPELINE: Extend3D = Extend3D.from_pretrained(MODEL_ID).cuda() |
|
|
|
|
|
|
|
|
|
|
| def run_extend3d( |
| image_pil,
|
| seed: int,
|
| randomize_seed: bool,
|
| width: int,
|
| length: int,
|
| div: int,
|
| ss_optim: bool,
|
| ss_iterations: int,
|
| ss_steps: int,
|
| ss_rescale_t: float,
|
| ss_t_noise: float,
|
| ss_t_start: float,
|
| ss_cfg_strength: float,
|
| ss_alpha: float,
|
| ss_batch_size: int,
|
| slat_optim: bool,
|
| slat_steps: int,
|
| slat_rescale_t: float,
|
| slat_cfg_strength: float,
|
| slat_batch_size: int,
|
| progress=gr.Progress(),
|
| ):
|
| if randomize_seed:
|
| seed = random.randint(0, 2147483647)
|
|
|
| torch.manual_seed(seed)
|
| np.random.seed(seed)
|
| random.seed(seed)
|
|
|
| pipe = PIPELINE |
|
|
| output = pipe.run(
|
| image_pil,
|
| width, length, div,
|
| ss_optim, ss_iterations, ss_steps,
|
| ss_rescale_t, ss_t_noise, ss_t_start,
|
| ss_cfg_strength, ss_alpha, ss_batch_size,
|
| slat_optim, slat_steps, slat_rescale_t,
|
| slat_cfg_strength, slat_batch_size,
|
| progress_callback=lambda frac, desc: progress(frac, desc=desc),
|
| )
|
|
|
| gaussian = output["gaussian"][0]
|
| mesh = output["mesh"][0]
|
|
|
| out_dir = Path(DEFAULT_OUTPUT_DIR)
|
| out_dir.mkdir(parents=True, exist_ok=True)
|
| run_id = uuid.uuid4().hex
|
|
|
|
|
| progress(0, desc="Rendering video...")
|
| color_frames = render_utils.render_video(gaussian, r=1.6, resolution=1024)["color"]
|
| progress(0.5, desc="Rendering video...")
|
| normal_frames = render_utils.render_video(mesh, r=1.6, resolution=1024)["normal"]
|
| progress(1.0, desc="Rendering video...")
|
| video_frames = [
|
| np.concatenate([c, n], axis=1)
|
| for c, n in zip(color_frames, normal_frames)
|
| ]
|
| video_path = str(out_dir / f"preview_{run_id}.mp4")
|
| imageio.mimsave(video_path, video_frames, fps=30)
|
|
|
|
|
| progress(0, desc="Exporting GLB...")
|
| glb = postprocessing_utils.to_glb(gaussian, mesh, simplify=0.98, texture_size=1024)
|
| glb.visual.material.metallicFactor = 0.0
|
| glb_path = str(out_dir / f"preview_{run_id}.glb")
|
| glb.export(glb_path)
|
| progress(1.0, desc="Done!")
|
|
|
| return video_path, glb_path, seed
|
|
|
|
|
|
|
|
|
|
|
|
|
| css = """
|
| #examples_gallery .gallery-item {
|
| width: 160px !important;
|
| height: 160px !important;
|
| min-width: 160px !important;
|
| }
|
| #examples_gallery img {
|
| width: 100% !important;
|
| height: 100% !important;
|
| object-fit: cover;
|
| }
|
| #examples_gallery .gallery {
|
| width: 100% !important;
|
| height: 100% !important;
|
| object-fit: cover;
|
| max-width: none !important;
|
| justify-content: center;
|
| }
|
| """
|
|
|
| with gr.Blocks(title="Extend3D Demo", css=css) as demo:
|
| gr.Markdown("# Extend3D: Town-scale 3D Generation")
|
| gr.Markdown("[Project Page](https://seungwoo-yoon.github.io/extend3d-page/) | [Code](https://github.com/Seungwoo-Yoon/Extend3D) | [Paper](#)")
|
|
|
| with gr.Row():
|
|
|
| with gr.Column(scale=4, min_width=420):
|
| gr.Markdown("### Input")
|
| image_in = gr.Image(label="Input Image", type="pil")
|
| run_btn = gr.Button("Run", variant="primary")
|
|
|
| with gr.Accordion("Settings", open=True):
|
| seed = gr.Slider(0, 2147483647, value=42, step=1, label="seed")
|
| randomize_seed = gr.Checkbox(value=True, label="randomize_seed")
|
| with gr.Row():
|
| width = gr.Slider(1, 8, value=2, step=1, label="width")
|
| length = gr.Slider(1, 8, value=2, step=1, label="length")
|
| div = gr.Slider(1, 8, value=4, step=1, label="div")
|
|
|
| with gr.Accordion("Sparse Structure Settings", open=False):
|
| ss_optim = gr.Checkbox(value=True, label="optimize")
|
| with gr.Row():
|
| ss_iterations = gr.Slider(1, 10, value=3, step=1, label="iterations")
|
| ss_steps = gr.Slider(1, 100, value=25, step=1, label="steps")
|
| with gr.Row():
|
| ss_rescale_t = gr.Slider(1, 10, value=3.0, step=0.1, label="rescale_t")
|
| ss_cfg_strength = gr.Slider(1, 10, value=7.5, step=0.1, label="cfg_strength")
|
| with gr.Row():
|
| ss_t_noise = gr.Slider(0, 1, value=0.6, step=0.1, label="t_noise")
|
| ss_t_start = gr.Slider(0, 1, value=0.8, step=0.1, label="t_start")
|
| ss_alpha = gr.Slider(1, 10, value=5.0, step=0.1, label="alpha")
|
| ss_batch_size = gr.Slider(1, 16, value=1, step=1, label="batch_size")
|
|
|
| with gr.Accordion("SLAT Settings", open=False):
|
| slat_optim = gr.Checkbox(value=True, label="optimize")
|
| with gr.Row():
|
| slat_steps = gr.Slider(1, 100, value=25, step=1, label="steps")
|
| with gr.Row():
|
| slat_rescale_t = gr.Slider(1, 10, value=3.0, step=0.1, label="rescale_t")
|
| slat_cfg_strength = gr.Slider(1, 10, value=3.0, step=0.1, label="cfg_strength")
|
| slat_batch_size = gr.Slider(1, 16, value=1, step=1, label="batch_size")
|
|
|
|
|
| with gr.Column(scale=5, min_width=420):
|
| gr.Markdown("### Output")
|
| preview_video = gr.Video(label="3D Preview (Video)", value=None, autoplay=True, loop=True)
|
| preview_glb = gr.Model3D(label="3D Preview (GLB)", value=None)
|
|
|
| gr.Examples(
|
| examples=[
|
| "assets/examples/0.png",
|
| "assets/examples/1.png",
|
| "assets/examples/2.png",
|
| "assets/examples/3.png",
|
| "assets/examples/4.png",
|
| "assets/examples/5.webp",
|
| ],
|
| inputs=[image_in],
|
| label="Examples",
|
| examples_per_page=6,
|
| elem_id="examples_gallery",
|
| )
|
|
|
| run_btn.click(
|
| fn=run_extend3d,
|
| inputs=[
|
| image_in,
|
| seed, randomize_seed,
|
| width, length, div,
|
| ss_optim, ss_iterations, ss_steps, ss_rescale_t, ss_t_noise, ss_t_start,
|
| ss_cfg_strength, ss_alpha, ss_batch_size,
|
| slat_optim, slat_steps, slat_rescale_t, slat_cfg_strength, slat_batch_size,
|
| ],
|
| outputs=[preview_video, preview_glb, seed],
|
| )
|
|
|
| if __name__ == "__main__": |
| demo.launch() |
|
|
