cpuai/Trellis.2.multiview

1import gradio as gr
2
3import os
4os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True"
5from datetime import datetime
6import shutil
7from typing import *
8import torch
9import numpy as np
10import trimesh
11from PIL import Image
12from trellis2.pipelines import Trellis2TexturingPipeline
13
14
15MAX_SEED = np.iinfo(np.int32).max
16TMP_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'tmp')
17
18
19def start_session(req: gr.Request):
20    user_dir = os.path.join(TMP_DIR, str(req.session_hash))
21    os.makedirs(user_dir, exist_ok=True)
22    
23    
24def end_session(req: gr.Request):
25    user_dir = os.path.join(TMP_DIR, str(req.session_hash))
26    shutil.rmtree(user_dir)
27
28
29def preprocess_image(image: Image.Image) -> Image.Image:
30    """
31    Preprocess the input image.
32
33    Args:
34        image (Image.Image): The input image.
35
36    Returns:
37        Image.Image: The preprocessed image.
38    """
39    processed_image = pipeline.preprocess_image(image)
40    return processed_image
41
42
43def get_seed(randomize_seed: bool, seed: int) -> int:
44    """
45    Get the random seed.
46    """
47    return np.random.randint(0, MAX_SEED) if randomize_seed else seed
48
49
50def shapeimage_to_tex(
51    mesh_file: str,
52    image: Image.Image,
53    seed: int,
54    resolution: str,
55    texture_size: int,
56    tex_slat_guidance_strength: float,
57    tex_slat_guidance_rescale: float,
58    tex_slat_sampling_steps: int,
59    tex_slat_rescale_t: float,
60    req: gr.Request,
61    progress=gr.Progress(track_tqdm=True),
62) -> str:
63    mesh = trimesh.load(mesh_file)
64    if isinstance(mesh, trimesh.Scene):
65        mesh = mesh.to_mesh()
66    output = pipeline.run(
67        mesh,
68        image,
69        seed=seed,
70        preprocess_image=False,
71        tex_slat_sampler_params={
72            "steps": tex_slat_sampling_steps,
73            "guidance_strength": tex_slat_guidance_strength,
74            "guidance_rescale": tex_slat_guidance_rescale,
75            "rescale_t": tex_slat_rescale_t,
76        },
77        resolution=int(resolution),
78        texture_size=texture_size,
79    )
80    now = datetime.now()
81    timestamp = now.strftime("%Y-%m-%dT%H%M%S") + f".{now.microsecond // 1000:03d}"
82    user_dir = os.path.join(TMP_DIR, str(req.session_hash))
83    os.makedirs(user_dir, exist_ok=True)
84    glb_path = os.path.join(user_dir, f'sample_{timestamp}.glb')
85    output.export(glb_path, extension_webp=True)
86    torch.cuda.empty_cache()
87    return glb_path, glb_path
88
89
90with gr.Blocks(delete_cache=(600, 600)) as demo:
91    gr.Markdown("""
92    ## Texturing a mesh with [TRELLIS.2](https://microsoft.github.io/TRELLIS.2)
93    * Upload a mesh and corresponding reference image (preferably with an alpha-masked foreground object) and click Generate to create a textured 3D asset.
94    """)
95    
96    with gr.Row():
97        with gr.Column(scale=1, min_width=360):
98            mesh_file = gr.File(label="Upload Mesh", file_types=[".ply", ".obj", ".glb", ".gltf"], file_count="single")
99            image_prompt = gr.Image(label="Image Prompt", format="png", image_mode="RGBA", type="pil", height=400)
100            
101            resolution = gr.Radio(["512", "1024", "1536"], label="Resolution", value="1024")
102            seed = gr.Slider(0, MAX_SEED, label="Seed", value=0, step=1)
103            randomize_seed = gr.Checkbox(label="Randomize Seed", value=True)
104            texture_size = gr.Slider(1024, 4096, label="Texture Size", value=2048, step=1024)
105            
106            generate_btn = gr.Button("Generate")
107                
108            with gr.Accordion(label="Advanced Settings", open=False):                
109                with gr.Row():
110                    tex_slat_guidance_strength = gr.Slider(1.0, 10.0, label="Guidance Strength", value=1.0, step=0.1)
111                    tex_slat_guidance_rescale = gr.Slider(0.0, 1.0, label="Guidance Rescale", value=0.0, step=0.01)
112                    tex_slat_sampling_steps = gr.Slider(1, 50, label="Sampling Steps", value=12, step=1)
113                    tex_slat_rescale_t = gr.Slider(1.0, 6.0, label="Rescale T", value=3.0, step=0.1)                
114
115        with gr.Column(scale=10):
116            glb_output = gr.Model3D(label="Extracted GLB", height=724, show_label=True, display_mode="solid", clear_color=(0.25, 0.25, 0.25, 1.0))
117            download_btn = gr.DownloadButton(label="Download GLB")
118                        
119
120    # Handlers
121    demo.load(start_session)
122    demo.unload(end_session)
123    
124    image_prompt.upload(
125        preprocess_image,
126        inputs=[image_prompt],
127        outputs=[image_prompt],
128    )
129
130    generate_btn.click(
131        get_seed,
132        inputs=[randomize_seed, seed],
133        outputs=[seed],
134    ).then(
135        shapeimage_to_tex,
136        inputs=[
137            mesh_file, image_prompt, seed, resolution, texture_size,
138            tex_slat_guidance_strength, tex_slat_guidance_rescale, tex_slat_sampling_steps, tex_slat_rescale_t,
139        ],
140        outputs=[glb_output, download_btn],
141    )
142        
143
144# Launch the Gradio app
145if __name__ == "__main__":
146    os.makedirs(TMP_DIR, exist_ok=True)
147
148    pipeline = Trellis2TexturingPipeline.from_pretrained('microsoft/TRELLIS.2-4B', config_file="texturing_pipeline.json")
149    pipeline.cuda()
150    
151    demo.launch()
152