# Version: 1.1.2 - Removed torch_dtype from from_pretrained call # Applied: # - Removed unsupported inputs/outputs kwargs on demo.load/unload # - Converted NumPy arrays to lists in pack_state for JSON safety # - Fixed indentation in Blocks event-handlers # - Verified clear() callbacks use only callback + outputs # - Removed `torch_dtype` arg from TrellisTextTo3DPipeline.from_pretrained # - Bumped version, added comments at change sites 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 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) def start_session(req: gr.Request): user_dir = os.path.join(TMP_DIR, str(req.session_hash)) os.makedirs(user_dir, exist_ok=True) print(f"Started session, created directory: {user_dir}") def end_session(req: gr.Request): user_dir = os.path.join(TMP_DIR, str(req.session_hash)) if os.path.exists(user_dir): try: shutil.rmtree(user_dir) print(f"Ended session, removed directory: {user_dir}") except OSError as e: print(f"Error removing tmp directory {user_dir}: {e.strerror}", file=sys.stderr) else: print(f"Ended session, directory already removed: {user_dir}") def pack_state(gs: Gaussian, mesh: MeshExtractResult) -> dict: """Packs Gaussian and Mesh data into a JSON-serializable dictionary.""" packed_data = { 'gaussian': { **{k: v for k, v in gs.init_params.items()}, # FIX: convert arrays to lists for JSON '_xyz': gs._xyz.detach().cpu().numpy().tolist(), '_features_dc': gs._features_dc.detach().cpu().numpy().tolist(), '_scaling': gs._scaling.detach().cpu().numpy().tolist(), '_rotation': gs._rotation.detach().cpu().numpy().tolist(), '_opacity': gs._opacity.detach().cpu().numpy().tolist(), }, 'mesh': { 'vertices': mesh.vertices.detach().cpu().numpy().tolist(), 'faces': mesh.faces.detach().cpu().numpy().tolist(), }, } return packed_data def unpack_state(state_dict: dict) -> Tuple[Gaussian, edict]: print("[unpack_state] Unpacking state from dictionary... ") device = 'cuda' if torch.cuda.is_available() else 'cpu' gauss_data = state_dict['gaussian'] mesh_data = state_dict['mesh'] gs = Gaussian( aabb=gauss_data.get('aabb'), sh_degree=gauss_data.get('sh_degree'), mininum_kernel_size=gauss_data.get('mininum_kernel_size'), scaling_bias=gauss_data.get('scaling_bias'), opacity_bias=gauss_data.get('opacity_bias'), scaling_activation=gauss_data.get('scaling_activation'), ) gs._xyz = torch.tensor(np.array(gauss_data['_xyz']), device=device, dtype=torch.float32) gs._features_dc = torch.tensor(np.array(gauss_data['_features_dc']), device=device, dtype=torch.float32) gs._scaling = torch.tensor(np.array(gauss_data['_scaling']), device=device, dtype=torch.float32) gs._rotation = torch.tensor(np.array(gauss_data['_rotation']), device=device, dtype=torch.float32) gs._opacity = torch.tensor(np.array(gauss_data['_opacity']), device=device, dtype=torch.float32) mesh = edict( vertices=torch.tensor(np.array(mesh_data['vertices']), device=device, dtype=torch.float32), faces=torch.tensor(np.array(mesh_data['faces']), device=device, dtype=torch.int64), ) return gs, mesh def get_seed(randomize_seed: bool, seed: int) -> int: new_seed = np.random.randint(0, MAX_SEED) if randomize_seed else seed return int(new_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]: outputs = pipeline.run( prompt, seed=seed, formats=["gaussian", "mesh"], sparse_structure_sampler_params={"steps": int(ss_sampling_steps), "cfg_strength": float(ss_guidance_strength)}, slat_sampler_params={"steps": int(slat_sampling_steps), "cfg_strength": float(slat_guidance_strength)}, ) state_dict = pack_state(outputs['gaussian'][0], outputs['mesh'][0]) 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_combined = [np.concatenate([v.astype(np.uint8), vg.astype(np.uint8)], axis=1) for v, vg in zip(video, video_geo)] user_dir = os.path.join(TMP_DIR, str(req.session_hash)) os.makedirs(user_dir, exist_ok=True) video_path = os.path.join(user_dir, 'sample.mp4') imageio.mimsave(video_path, video_combined, fps=15, quality=8) if torch.cuda.is_available(): torch.cuda.empty_cache() return state_dict, video_path @spaces.GPU(duration=120) def extract_glb( state_dict: dict, mesh_simplify: float, texture_size: int, req: gr.Request, ) -> Tuple[str, str]: gs, mesh = unpack_state(state_dict) user_dir = os.path.join(TMP_DIR, str(req.session_hash)) os.makedirs(user_dir, exist_ok=True) glb = postprocessing_utils.to_glb(gs, mesh, simplify=float(mesh_simplify), texture_size=int(texture_size), verbose=True) glb_path = os.path.join(user_dir, 'sample.glb') glb.export(glb_path) if torch.cuda.is_available(): torch.cuda.empty_cache() return glb_path, glb_path @spaces.GPU def extract_gaussian( state_dict: dict, req: gr.Request ) -> Tuple[str, str]: gs, _ = unpack_state(state_dict) user_dir = os.path.join(TMP_DIR, str(req.session_hash)) os.makedirs(user_dir, exist_ok=True) gaussian_path = os.path.join(user_dir, 'sample.ply') gs.save_ply(gaussian_path) if torch.cuda.is_available(): torch.cuda.empty_cache() return gaussian_path, gaussian_path # --- Gradio UI Definition --- with gr.Blocks(delete_cache=(600, 600), title="TRELLIS Text-to-3D") as demo: gr.Markdown(""" # Text to 3D Asset with [TRELLIS](https://trellis3d.github.io/) """) # State buffer output_buf = gr.State() with gr.Row(): with gr.Column(scale=1): 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("---\n**Stage 1**") ss_guidance_strength = gr.Slider(0.0, 15.0, label="Guidance Strength", value=7.5, step=0.1) ss_sampling_steps = gr.Slider(10, 50, label="Sampling Steps", value=25, step=1) gr.Markdown("---\n**Stage 2**") slat_guidance_strength = gr.Slider(0.0, 15.0, label="Guidance Strength", value=7.5, step=0.1) slat_sampling_steps = gr.Slider(10, 50, label="Sampling Steps", value=25, step=1) generate_btn = gr.Button("Generate 3D Preview", variant="primary") with gr.Accordion(label="GLB Extraction Settings", open=True): mesh_simplify = gr.Slider(0.9, 0.99, label="Simplify Factor", value=0.95, step=0.01) texture_size = gr.Slider(512, 2048, label="Texture Size", value=1024, step=512) extract_glb_btn = gr.Button("Extract GLB", interactive=False) extract_gs_btn = gr.Button("Extract Gaussian (PLY)", interactive=False) download_glb = gr.DownloadButton(label="Download GLB", interactive=False) download_gs = gr.DownloadButton(label="Download Gaussian (PLY)", interactive=False) with gr.Column(scale=1): video_output = gr.Video(label="3D Preview", autoplay=True, loop=True) model_output = gr.Model3D(label="Extracted Model Preview") # --- Event handlers --- demo.load(start_session) # FIX: remove inputs/outputs kwargs demo.unload(end_session) # FIX: remove inputs/outputs kwargs # Align indentation to one level under Blocks generate_event = 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: (extract_glb_btn.update(interactive=True), extract_gs_btn.update(interactive=True)), outputs=[extract_glb_btn, extract_gs_btn], ) extract_glb_event = extract_glb_btn.click( extract_glb, inputs=[output_buf, mesh_simplify, texture_size], outputs=[model_output, download_glb], ).then( lambda: download_glb.update(interactive=True), outputs=[download_glb], ) extract_gs_event = extract_gs_btn.click( extract_gaussian, inputs=[output_buf], outputs=[model_output, download_gs], ).then( lambda: download_gaussian.update(interactive=True), outputs=[download_gs], ) # Clear callbacks model_output.clear( lambda: (download_glb.update(interactive=False), download_gs.update(interactive=False)), outputs=[download_glb, download_gs], ) video_output.clear( lambda: (extract_glb_btn.update(interactive=False), extract_gs_btn.update(interactive=False), download_glb.update(interactive=False), download_gs.update(interactive=False)), outputs=[extract_glb_btn, extract_gs_btn, download_glb, download_gs], ) if __name__ == "__main__": # Removed torch_dtype argument to match current API pipeline = TrellisTextTo3DPipeline.from_pretrained( "JeffreyXiang/TRELLIS-text-xlarge" ) if torch.cuda.is_available(): pipeline = pipeline.to("cuda") demo.queue().launch(debug=True)