File size: 15,306 Bytes
988efc8
 
 
 
 
8527388
988efc8
 
 
 
 
 
 
 
 
 
27271e8
 
 
ddff7d1
 
 
 
 
 
 
988efc8
 
 
 
 
ddff7d1
 
 
 
 
 
988efc8
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1eb7916
 
 
 
 
83d10c9
 
 
 
 
 
 
 
 
 
 
 
 
4d2a9aa
 
 
83d10c9
1eb7916
 
 
 
 
 
 
83d10c9
 
 
988efc8
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1eb7916
 
 
 
 
988efc8
 
 
1eb7916
 
 
 
 
988efc8
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1eb7916
 
 
 
 
988efc8
 
1eb7916
 
 
 
 
988efc8
 
 
 
ddff7d1
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
988efc8
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
9896238
988efc8
47a7175
 
 
988efc8
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
0b3b481
 
 
 
 
 
 
 
988efc8
0b3b481
 
 
 
 
 
 
 
c17a234
653000f
 
 
 
988efc8
 
ddff7d1
 
 
 
 
 
 
 
 
 
 
 
 
 
988efc8
ddff7d1
 
988efc8
 
ddff7d1
 
 
 
 
 
 
988efc8
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
import gradio as gr
import spaces

import os
import shutil
os.environ['TOKENIZERS_PARALLELISM'] = 'true'
os.environ['SPCONV_ALGO'] = 'native'
from typing import *
import torch
import numpy as np
import imageio
from easydict import EasyDict as edict
from trellis.pipelines import TrellisTextTo3DPipeline
from trellis.representations import Gaussian, MeshExtractResult
from trellis.utils import render_utils, postprocessing_utils

import traceback
import sys

# --- FastAPI / Threading Imports ---
import threading
import uvicorn
import logging
from fastapi import FastAPI, HTTPException
from pydantic import BaseModel


MAX_SEED = np.iinfo(np.int32).max
TMP_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'tmp')
os.makedirs(TMP_DIR, exist_ok=True)

# --- Global Pipeline Variable ---
pipeline = None

# --- Logging Setup ---
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

def start_session(req: gr.Request):
    user_dir = os.path.join(TMP_DIR, str(req.session_hash))
    os.makedirs(user_dir, exist_ok=True)
    
    
def end_session(req: gr.Request):
    user_dir = os.path.join(TMP_DIR, str(req.session_hash))
    shutil.rmtree(user_dir)


def pack_state(gs: Gaussian, mesh: MeshExtractResult) -> dict:
    return {
        'gaussian': {
            **gs.init_params,
            '_xyz': gs._xyz.cpu().numpy(),
            '_features_dc': gs._features_dc.cpu().numpy(),
            '_scaling': gs._scaling.cpu().numpy(),
            '_rotation': gs._rotation.cpu().numpy(),
            '_opacity': gs._opacity.cpu().numpy(),
        },
        'mesh': {
            'vertices': mesh.vertices.cpu().numpy(),
            'faces': mesh.faces.cpu().numpy(),
        },
    }
    
    
def unpack_state(state: dict) -> Tuple[Gaussian, edict, str]:
    gs = Gaussian(
        aabb=state['gaussian']['aabb'],
        sh_degree=state['gaussian']['sh_degree'],
        mininum_kernel_size=state['gaussian']['mininum_kernel_size'],
        scaling_bias=state['gaussian']['scaling_bias'],
        opacity_bias=state['gaussian']['opacity_bias'],
        scaling_activation=state['gaussian']['scaling_activation'],
    )
    gs._xyz = torch.tensor(state['gaussian']['_xyz'], device='cuda')
    gs._features_dc = torch.tensor(state['gaussian']['_features_dc'], device='cuda')
    gs._scaling = torch.tensor(state['gaussian']['_scaling'], device='cuda')
    gs._rotation = torch.tensor(state['gaussian']['_rotation'], device='cuda')
    gs._opacity = torch.tensor(state['gaussian']['_opacity'], device='cuda')
    
    mesh = edict(
        vertices=torch.tensor(state['mesh']['vertices'], device='cuda'),
        faces=torch.tensor(state['mesh']['faces'], device='cuda'),
    )
    
    return gs, mesh


def get_seed(randomize_seed: bool, seed: int) -> int:
    """
    Get the random seed.
    """
    return np.random.randint(0, MAX_SEED) if randomize_seed else seed


@spaces.GPU
def text_to_3d(
    prompt: str,
    seed: int,
    ss_guidance_strength: float,
    ss_sampling_steps: int,
    slat_guidance_strength: float,
    slat_sampling_steps: int,
    req: gr.Request,
) -> Tuple[dict, str]:
    """
    Convert an text prompt to a 3D model.

    Args:
        prompt (str): The text prompt.
        seed (int): The random seed.
        ss_guidance_strength (float): The guidance strength for sparse structure generation.
        ss_sampling_steps (int): The number of sampling steps for sparse structure generation.
        slat_guidance_strength (float): The guidance strength for structured latent generation.
        slat_sampling_steps (int): The number of sampling steps for structured latent generation.

    Returns:
        dict: The information of the generated 3D model.
        str: The path to the video of the 3D model.
    """
    # --- Determine user_dir robustly ---
    session_hash_str = str(req.session_hash) if hasattr(req, 'session_hash') and req.session_hash else f"api_call_{np.random.randint(10000)}"
    user_dir = os.path.join(TMP_DIR, session_hash_str)
    os.makedirs(user_dir, exist_ok=True) # Ensure directory exists

    outputs = pipeline.run(
        prompt,
        seed=seed,
        formats=["gaussian", "mesh"],
        sparse_structure_sampler_params={
            "steps": ss_sampling_steps,
            "cfg_strength": ss_guidance_strength,
        },
        slat_sampler_params={
            "steps": slat_sampling_steps,
            "cfg_strength": slat_guidance_strength,
        },
    )
    video = render_utils.render_video(outputs['gaussian'][0], num_frames=120)['color']
    video_geo = render_utils.render_video(outputs['mesh'][0], num_frames=120)['normal']
    video = [np.concatenate([video[i], video_geo[i]], axis=1) for i in range(len(video))]
    video_path = os.path.join(user_dir, 'sample.mp4')
    try:
        imageio.mimsave(video_path, video, fps=15) # Now the directory should exist
    except FileNotFoundError:
         print(f"ERROR: Directory {user_dir} still not found before mimsave!", file=sys.stderr)
         # Decide if we should raise or return an error state?
         # Returning a dummy path might hide the error, so let's raise for now
         raise
    state = pack_state(outputs['gaussian'][0], outputs['mesh'][0])
    torch.cuda.empty_cache()
    return state, video_path


@spaces.GPU(duration=90)
def extract_glb(
    state: dict,
    mesh_simplify: float,
    texture_size: int,
    req: gr.Request,
) -> Tuple[str, str]:
    """
    Extract a GLB file from the 3D model.

    Args:
        state (dict): The state of the generated 3D model.
        mesh_simplify (float): The mesh simplification factor.
        texture_size (int): The texture resolution.

    Returns:
        str: The path to the extracted GLB file.
    """
    # --- Determine user_dir robustly ---
    session_hash_str = str(req.session_hash) if hasattr(req, 'session_hash') and req.session_hash else f"api_call_{np.random.randint(10000)}"
    user_dir = os.path.join(TMP_DIR, session_hash_str)
    os.makedirs(user_dir, exist_ok=True) # Ensure directory exists

    gs, mesh = unpack_state(state)
    glb = postprocessing_utils.to_glb(gs, mesh, simplify=mesh_simplify, texture_size=texture_size, verbose=False)
    glb_path = os.path.join(user_dir, 'sample.glb')
    try:
        glb.export(glb_path) # Now the directory should exist
    except FileNotFoundError:
        print(f"ERROR: Directory {user_dir} still not found before glb.export!", file=sys.stderr)
        raise
    torch.cuda.empty_cache()
    return glb_path, glb_path


@spaces.GPU
def extract_gaussian(state: dict, req: gr.Request) -> Tuple[str, str]:
    """
    Extract a Gaussian file from the 3D model.

    Args:
        state (dict): The state of the generated 3D model.

    Returns:
        str: The path to the extracted Gaussian file.
    """
    # --- Determine user_dir robustly ---
    session_hash_str = str(req.session_hash) if hasattr(req, 'session_hash') and req.session_hash else f"api_call_{np.random.randint(10000)}"
    user_dir = os.path.join(TMP_DIR, session_hash_str)
    os.makedirs(user_dir, exist_ok=True) # Ensure directory exists

    gs, _ = unpack_state(state)
    gaussian_path = os.path.join(user_dir, 'sample.ply')
    try:
        gs.save_ply(gaussian_path) # Now the directory should exist
    except FileNotFoundError:
        print(f"ERROR: Directory {user_dir} still not found before gs.save_ply!", file=sys.stderr)
        raise
    torch.cuda.empty_cache()
    return gaussian_path, gaussian_path


# --- FastAPI App Setup ---
api_app = FastAPI()

class GenerateRequest(BaseModel):
    prompt: str
    seed: int = 0 # Default seed
    mesh_simplify: float = 0.95 # Default simplify factor
    texture_size: int = 1024 # Default texture size
    # Add other generation parameters if needed

@api_app.post("/api/generate-sync")
@spaces.GPU(duration=300) # Allow longer duration for API calls
async def generate_sync_api(request_data: GenerateRequest):
    global pipeline # Access the globally initialized pipeline
    if pipeline is None:
        logger.error("API Error: Pipeline not initialized")
        raise HTTPException(status_code=503, detail="Pipeline not ready")

    prompt = request_data.prompt
    seed = request_data.seed
    mesh_simplify = request_data.mesh_simplify
    texture_size = request_data.texture_size
    # Extract other params if added to GenerateRequest

    logger.info(f"API /generate-sync received prompt: {prompt}")

    try:
        # --- Determine a unique temporary directory for this API call ---
        api_call_hash = f"api_sync_{np.random.randint(100000)}"
        user_dir = os.path.join(TMP_DIR, api_call_hash)
        os.makedirs(user_dir, exist_ok=True)
        logger.info(f"API using temp dir: {user_dir}")

        # --- Stage 1: Run the text-to-3D pipeline ---
        logger.info("API running pipeline...")
        # Use default values for parameters not exposed in the simple API for now
        outputs = pipeline.run(
            prompt,
            seed=seed,
            formats=["gaussian", "mesh"],
            sparse_structure_sampler_params={"steps": 25, "cfg_strength": 7.5},
            slat_sampler_params={"steps": 25, "cfg_strength": 7.5},
        )
        gs = outputs['gaussian'][0]
        mesh = outputs['mesh'][0]
        logger.info("API pipeline finished.")
        torch.cuda.empty_cache()

        # --- Stage 2: Extract GLB ---
        logger.info("API extracting GLB...")
        glb = postprocessing_utils.to_glb(gs, mesh, simplify=mesh_simplify, texture_size=texture_size, verbose=False)
        glb_path = os.path.join(user_dir, 'generated_sync.glb')
        glb.export(glb_path)
        logger.info(f"API GLB exported to: {glb_path}")
        torch.cuda.empty_cache()

        # Return the absolute path within the container
        return {"status": "success", "glb_path": os.path.abspath(glb_path)}

    except Exception as e:
        logger.error(f"API /generate-sync error: {str(e)}", exc_info=True)
        # Clean up temp dir on error if it exists
        if os.path.exists(user_dir):
            try:
                shutil.rmtree(user_dir)
            except Exception as cleanup_e:
                logger.error(f"API Error cleaning up dir {user_dir}: {cleanup_e}")
        raise HTTPException(status_code=500, detail=f"Generation failed: {str(e)}")
    # Note: We don't automatically clean up the user_dir on success,
    # as the file needs to be accessible for download by the calling server.
    # A separate cleanup mechanism might be needed eventually.


# --- Gradio Blocks Definition ---
with gr.Blocks(delete_cache=(600, 600)) as demo:
    gr.Markdown("""
    ## Text to 3D Asset with [TRELLIS](https://trellis3d.github.io/)
    * Type a text prompt and click "Generate" to create a 3D asset.
    * If you find the generated 3D asset satisfactory, click "Extract GLB" to extract the GLB file and download it.
    """)
    
    with gr.Row():
        with gr.Column():
            text_prompt = gr.Textbox(label="Text Prompt", lines=5)
            
            with gr.Accordion(label="Generation Settings", open=False):
                seed = gr.Slider(0, MAX_SEED, label="Seed", value=0, step=1)
                randomize_seed = gr.Checkbox(label="Randomize Seed", value=True)
                gr.Markdown("Stage 1: Sparse Structure Generation")
                with gr.Row():
                    ss_guidance_strength = gr.Slider(0.0, 10.0, label="Guidance Strength", value=7.5, step=0.1)
                    ss_sampling_steps = gr.Slider(1, 50, label="Sampling Steps", value=25, step=1)
                gr.Markdown("Stage 2: Structured Latent Generation")
                with gr.Row():
                    slat_guidance_strength = gr.Slider(0.0, 10.0, label="Guidance Strength", value=7.5, step=0.1)
                    slat_sampling_steps = gr.Slider(1, 50, label="Sampling Steps", value=25, step=1)

            generate_btn = gr.Button("Generate")
            
            with gr.Accordion(label="GLB Extraction Settings", open=False):
                mesh_simplify = gr.Slider(0.9, 0.98, label="Simplify", value=0.95, step=0.01)
                texture_size = gr.Slider(512, 2048, label="Texture Size", value=1024, step=512)
            
            with gr.Row():
                extract_glb_btn = gr.Button("Extract GLB", interactive=False)
                extract_gs_btn = gr.Button("Extract Gaussian", interactive=False)
            gr.Markdown("""
                        *NOTE: Gaussian file can be very large (~50MB), it will take a while to display and download.*
                        """)

        with gr.Column():
            video_output = gr.Video(label="Generated 3D Asset", autoplay=True, loop=True, height=300)
            model_output = gr.Model3D(label="Extracted GLB/Gaussian", height=300)
            
            with gr.Row():
                download_glb = gr.DownloadButton(label="Download GLB", interactive=False)
                download_gs = gr.DownloadButton(label="Download Gaussian", interactive=False)  
    
    output_buf = gr.State()

    # Handlers
    demo.load(start_session)
    demo.unload(end_session)

    generate_btn.click(
        get_seed,
        inputs=[randomize_seed, seed],
        outputs=[seed],
    ).then(
        text_to_3d,
        inputs=[text_prompt, seed, ss_guidance_strength, ss_sampling_steps, slat_guidance_strength, slat_sampling_steps],
        outputs=[output_buf, video_output],
    ).then(
        lambda: tuple([gr.Button(interactive=True), gr.Button(interactive=True)]),
        outputs=[extract_glb_btn, extract_gs_btn],
    )

    video_output.clear(
        lambda: tuple([gr.Button(interactive=False), gr.Button(interactive=False)]),
        outputs=[extract_glb_btn, extract_gs_btn],
    )

    extract_glb_btn.click(
        extract_glb,
        inputs=[output_buf, mesh_simplify, texture_size],
        outputs=[model_output, download_glb],
    ).then(
        lambda: gr.Button(interactive=True),
        outputs=[download_glb],
    )
    
    extract_gs_btn.click(
        extract_gaussian,
        inputs=[output_buf],
        outputs=[model_output, download_gs],
    ).then(
        lambda: gr.Button(interactive=True),
        outputs=[download_gs],
    )

    model_output.clear(
        lambda: gr.Button(interactive=False),
        outputs=[download_glb],
    )
    

# --- Functions to Run FastAPI in Background ---
def run_api():
    """Run the FastAPI server."""
    uvicorn.run(api_app, host="0.0.0.0", port=8000) # Run on port 8000

def start_api_thread():
    """Start the API server in a background thread."""
    api_thread = threading.Thread(target=run_api, daemon=True)
    api_thread.start()
    logger.info("Started FastAPI server thread on port 8000")
    return api_thread


# Launch the Gradio app and FastAPI server
if __name__ == "__main__":
    logger.info("Initializing Trellis Pipeline...")
    # Make pipeline global so API endpoint can access it
    pipeline = TrellisTextTo3DPipeline.from_pretrained("JeffreyXiang/TRELLIS-text-xlarge")
    pipeline.cuda()
    logger.info("Trellis Pipeline Initialized.")

    # Start the background API server
    start_api_thread()

    # Launch the Gradio interface (blocking call)
    logger.info("Launching Gradio Demo...")
    demo.launch()