Speed UP!

app.py

@@ -3,6 +3,7 @@ import spaces
 from gradio_litmodel3d import LitModel3D
 
 import os
+os.environ['SPCONV_ALGO'] = 'native'
 from typing import *
 import torch
 import numpy as np
@@ -131,7 +132,7 @@ def extract_glb(state: dict, mesh_simplify: float, texture_size: int) -> Tuple[s
         str: The path to the extracted GLB file.
     """
     gs, mesh, model_id = unpack_state(state)
-    glb = postprocessing_utils.to_glb(gs, mesh, simplify=mesh_simplify, texture_size=texture_size)
+    glb = postprocessing_utils.to_glb(gs, mesh, simplify=mesh_simplify, texture_size=texture_size, verbose=False)
     glb_path = f"/tmp/Trellis-demo/{model_id}.glb"
     glb.export(glb_path)
     return glb_path, glb_path
@@ -161,12 +162,12 @@ with gr.Blocks() as demo:
                 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=5.0, step=0.1)
-                    ss_sampling_steps = gr.Slider(1, 50, label="Sampling Steps", value=25, step=1)
+                    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=12, 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=5.0, step=0.1)
-                    slat_sampling_steps = gr.Slider(1, 50, label="Sampling Steps", value=25, step=1)
+                    slat_guidance_strength = gr.Slider(0.0, 10.0, label="Guidance Strength", value=3.0, step=0.1)
+                    slat_sampling_steps = gr.Slider(1, 50, label="Sampling Steps", value=12, step=1)
 
             generate_btn = gr.Button("Generate")
             

trellis/modules/sparse/__init__.py

@@ -24,6 +24,8 @@ def __from_env():
     if env_sparse_attn is not None and env_sparse_attn in ['xformers', 'flash_attn']:
         ATTN = env_sparse_attn
         
+    print(f"[SPARSE] Backend: {BACKEND}, Attention: {ATTN}")
+        
 
 __from_env()
     

trellis/modules/sparse/conv/__init__.py

@@ -1,6 +1,21 @@
 from .. import BACKEND
 
+
+SPCONV_ALGO = 'auto'    # 'auto', 'implicit_gemm', 'native'
+
+def __from_env():
+    import os
+        
+    global SPCONV_ALGO
+    env_spconv_algo = os.environ.get('SPCONV_ALGO')
+    if env_spconv_algo is not None and env_spconv_algo in ['auto', 'implicit_gemm', 'native']:
+        SPCONV_ALGO = env_spconv_algo
+    print(f"[SPARSE][CONV] spconv algo: {SPCONV_ALGO}")
+        
+
+__from_env()
+
 if BACKEND == 'torchsparse':
     from .conv_torchsparse import *
 elif BACKEND == 'spconv':
-    from .conv_spconv import *
+    from .conv_spconv import *

trellis/modules/sparse/conv/conv_spconv.py

@@ -2,16 +2,22 @@ import torch
 import torch.nn as nn
 from .. import SparseTensor
 from .. import DEBUG
+from . import SPCONV_ALGO
 
 class SparseConv3d(nn.Module):
     def __init__(self, in_channels, out_channels, kernel_size, stride=1, dilation=1, padding=None, bias=True, indice_key=None):
         super(SparseConv3d, self).__init__()
         if 'spconv' not in globals():
             import spconv.pytorch as spconv
+        algo = None
+        if SPCONV_ALGO == 'native':
+            algo = spconv.ConvAlgo.Native
+        elif SPCONV_ALGO == 'implicit_gemm':
+            algo = spconv.ConvAlgo.MaskImplicitGemm
         if stride == 1 and (padding is None):
-            self.conv = spconv.SubMConv3d(in_channels, out_channels, kernel_size, dilation=dilation, bias=bias, indice_key=indice_key)
+            self.conv = spconv.SubMConv3d(in_channels, out_channels, kernel_size, dilation=dilation, bias=bias, indice_key=indice_key, algo=algo)
         else:
-            self.conv = spconv.SparseConv3d(in_channels, out_channels, kernel_size, stride=stride, dilation=dilation, padding=padding, bias=bias, indice_key=indice_key)
+            self.conv = spconv.SparseConv3d(in_channels, out_channels, kernel_size, stride=stride, dilation=dilation, padding=padding, bias=bias, indice_key=indice_key, algo=algo)
         self.stride = tuple(stride) if isinstance(stride, (list, tuple)) else (stride, stride, stride)
         self.padding = padding
 

trellis/utils/postprocessing_utils.py

@@ -448,7 +448,7 @@ def to_glb(
         observations, masks, extrinsics, intrinsics,
         texture_size=texture_size, mode='opt',
         lambda_tv=0.01,
-        verbose=True
+        verbose=verbose
     )
     texture = Image.fromarray(texture)