update app

app.py

@@ -7,10 +7,20 @@ import tempfile
 from typing import *
 from datetime import datetime
 from pathlib import Path
-
 from typing import Iterable
+
+# Rerun imports
+import rerun as rr
+from gradio_rerun import Rerun
+
+import gradio as gr
 from gradio.themes import Soft
 from gradio.themes.utils import colors, fonts, sizes
+from gradio_client import Client, handle_file
+import spaces
+from diffusers import ZImagePipeline
+from trellis2.pipelines import Trellis2ImageTo3DPipeline
+import o_voxel
 
 colors.orange_red = colors.Color(
     name="orange_red",
@@ -83,18 +93,9 @@ orange_red_theme = OrangeRedTheme()
 os.environ["OPENCV_IO_ENABLE_OPENEXR"] = '1'
 os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True"
 os.environ["ATTN_BACKEND"] = "flash_attn_3"
-
 os.environ["FLEX_GEMM_AUTOTUNE_CACHE_PATH"] = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'autotune_cache.json')
 os.environ["FLEX_GEMM_AUTOTUNER_VERBOSE"] = '1'
 
-import gradio as gr
-from gradio_client import Client, handle_file
-import spaces
-from diffusers import ZImagePipeline
-
-from trellis2.pipelines import Trellis2ImageTo3DPipeline
-import o_voxel
-
 MAX_SEED = np.iinfo(np.int32).max
 TMP_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'tmp')
 
@@ -145,7 +146,6 @@ def remove_background(input: Image.Image) -> Image.Image:
         return output
 
 def preprocess_image(input: Image.Image) -> Image.Image:
-    """Preprocess the input image: remove bg, crop, resize."""
     if input is None:
         return None
         
@@ -184,7 +184,6 @@ def get_seed(randomize_seed: bool, seed: int) -> int:
 
 @spaces.GPU
 def generate_txt2img(prompt, progress=gr.Progress(track_tqdm=True)):
-    """Generate Image using Z-Image Turbo"""
     if z_pipe is None:
         raise gr.Error("Z-Image-Turbo model failed to load.")
     if not prompt.strip():
@@ -268,15 +267,12 @@ def generate_3d(
             return_latent=True,
         )
         
-        # 2. Process Mesh
         progress(0.7, desc="Processing Mesh...")
         mesh = outputs[0]
-        mesh.simplify(16777216) # Simplify for processing limits
+        mesh.simplify(16777216) 
         
-        # 3. Export to GLB
         progress(0.9, desc="Baking Texture & Exporting GLB...")
         
-        # Note: We use the latent grid resolution from the pipeline output
         grid_size = latents[2] 
         
         glb = o_voxel.postprocess.to_glb(
@@ -299,10 +295,28 @@ def generate_3d(
         timestamp = now.strftime("%Y-%m-%dT%H%M%S")
         glb_path = os.path.join(user_dir, f'trellis_output_{timestamp}.glb')
         glb.export(glb_path, extension_webp=True)
+
+        # --- Rerun Visualization Setup ---
+        progress(0.95, desc="Creating Rerun Visualization...")
+        rr.init("Trellis_3D", spawn=False)
+        
+        # Log the 3D Asset
+        rr.log("generated_model", rr.Asset3D(path=glb_path))
+        
+        # Log Axis and View Coordinates (Z-Up)
+        rr.log("world", rr.ViewCoordinates.RIGHT_HAND_Z_UP, static=True)
+        rr.log("world/xyz", rr.Arrows3D(
+            vectors=[[1, 0, 0], [0, 1, 0], [0, 0, 1]], 
+            colors=[[255, 0, 0], [0, 255, 0], [0, 0, 255]],
+            labels=["X", "Y", "Z"],
+            origins=[[0, 0, 0], [0, 0, 0], [0, 0, 0]]
+        ))
+
+        rrd_path = os.path.join(user_dir, f'trellis_output_{timestamp}.rrd')
+        rr.save(rrd_path)
         
-        # Clean up
         torch.cuda.empty_cache()
-        return glb_path, glb_path
+        return rrd_path, glb_path
         
     except Exception as e:
         torch.cuda.empty_cache()
@@ -320,10 +334,11 @@ if __name__ == "__main__":
     os.makedirs(TMP_DIR, exist_ok=True)
 
     with gr.Blocks(delete_cache=(300, 300)) as demo:
-        gr.Markdown("# **TRELLIS.2 (Text-to-3D)**", elem_id="main-title")
+        gr.Markdown("# **TRELLIS.2 (Text-to-3D) with Rerun**", elem_id="main-title")
         gr.Markdown("""
         **Workflow:**
         Generate a 3D asset directly by converting Text-to-Image → 3D or Image-to-3D, powered by TRELLIS.2 and Z-Image-Turbo.
+        Visualization powered by Rerun.
         """)
 
         with gr.Row():
@@ -343,7 +358,6 @@ if __name__ == "__main__":
                     seed = gr.Slider(0, MAX_SEED, label="Seed", value=0, step=1)
                     randomize_seed = gr.Checkbox(label="Randomize Seed", value=True)
                     
-                #gr.Markdown("**Export Settings**")
                 decimation_target = gr.Slider(50000, 500000, label="Target Faces", value=150000, step=10000)
                 texture_size = gr.Slider(512, 4096, label="Texture Size", value=1024, step=512)
                 
@@ -369,13 +383,10 @@ if __name__ == "__main__":
                     tex_rescale_t = gr.Slider(1.0, 6.0, value=3.0, label="Rescale T")
 
             with gr.Column(scale=2):
-                gr.Markdown("### 3D Output")
-                glb_output = gr.Model3D(
-                    label="Generated GLB", 
-                    display_mode="solid", 
-                    clear_color=(0.2, 0.2, 0.2, 1.0),
-                    height=600,
-                    interactive=False  # Changed to False to hide upload area
+                gr.Markdown("### 3D Output (Rerun Viewer)")
+                rerun_output = Rerun(
+                    label="Rerun 3D Visualization",
+                    height=600
                 )
                 download_btn = gr.DownloadButton(label="3.Download GLB File", variant="primary")
 
@@ -561,7 +572,7 @@ if __name__ == "__main__":
                 shape_guidance, shape_rescale, shape_steps, shape_rescale_t,
                 tex_guidance, tex_rescale, tex_steps, tex_rescale_t,
             ],
-            outputs=[glb_output, download_btn],
+            outputs=[rerun_output, download_btn],
         )
 
     demo.launch(theme=orange_red_theme, css=css, mcp_server=True, ssr_mode=False, show_error=True)