update app

app.py

@@ -7,24 +7,23 @@ import tempfile
 from typing import *
 from datetime import datetime
 from pathlib import Path
-import uuid
-
 from typing import Iterable
-from gradio.themes import Soft
-from gradio.themes.utils import colors, fonts, sizes
 
 import gradio as gr
 from gradio_client import Client, handle_file
+from gradio.themes import Soft
+from gradio.themes.utils import colors, fonts, sizes
+
 import spaces
 from diffusers import ZImagePipeline
+from trellis2.pipelines import Trellis2ImageTo3DPipeline
+import o_voxel
 
-# Rerun Imports
+# Rerun imports
 import rerun as rr
 from gradio_rerun import Rerun
 
-from trellis2.pipelines import Trellis2ImageTo3DPipeline
-import o_voxel
-
+# Theme Configuration
 colors.orange_red = colors.Color(
     name="orange_red",
     c50="#FFF0E5",
@@ -93,6 +92,7 @@ class OrangeRedTheme(Soft):
 
 orange_red_theme = OrangeRedTheme()
 
+# Environment Config
 os.environ["OPENCV_IO_ENABLE_OPENEXR"] = '1'
 os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True"
 os.environ["ATTN_BACKEND"] = "flash_attn_3"
@@ -102,9 +102,8 @@ os.environ["FLEX_GEMM_AUTOTUNER_VERBOSE"] = '1'
 MAX_SEED = np.iinfo(np.int32).max
 TMP_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'tmp')
 
+# Model Loading
 print("Initializing models...")
-
-print("Loading Z-Image-Turbo...")
 try:
     z_pipe = ZImagePipeline.from_pretrained(
         "Tongyi-MAI/Z-Image-Turbo",
@@ -118,7 +117,6 @@ except Exception as e:
     print(f"Failed to load Z-Image-Turbo: {e}")
     z_pipe = None
 
-print("Loading TRELLIS.2...")
 try:
     trellis_pipeline = Trellis2ImageTo3DPipeline.from_pretrained('microsoft/TRELLIS.2-4B')
     trellis_pipeline.rembg_model = None
@@ -210,30 +208,6 @@ def generate_txt2img(prompt, progress=gr.Progress(track_tqdm=True)):
     except Exception as e:
         raise gr.Error(f"Z-Image Generation failed: {str(e)}")
 
-def create_rerun_rrd(glb_path: str) -> str:
-    """Creates a Rerun recording for the generated GLB."""
-    rrd_path = glb_path.replace(".glb", ".rrd")
-    
-    rec = rr.new_recording(application_id="trellis_app", recording_id=str(uuid.uuid4()))
-    
-    rec.log("world", rr.ViewCoordinates.RIGHT_HAND_Z_UP, static=True)
-    
-    rec.log(
-        "world/axes",
-        rr.Arrows3D(
-            vectors=[[0.5, 0, 0], [0, 0.5, 0], [0, 0, 0.5]],
-            colors=[[255, 0, 0], [0, 255, 0], [0, 0, 255]],
-            labels=["X", "Y", "Z"],
-            radii=0.02
-        ),
-        static=True
-    )
-
-    rec.log("world/model", rr.Asset3D(path=glb_path), static=True)
-
-    rec.save(rrd_path)
-    return rrd_path
-
 @spaces.GPU(duration=120)
 def generate_3d(
     image: Image.Image,
@@ -296,7 +270,7 @@ def generate_3d(
         
         progress(0.7, desc="Processing Mesh...")
         mesh = outputs[0]
-        mesh.simplify(16777216)
+        mesh.simplify(16777216) 
         
         progress(0.9, desc="Baking Texture & Exporting GLB...")
         
@@ -322,8 +296,29 @@ 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 LOGGING ---
+        progress(0.95, desc="Creating Rerun Visualization...")
+        rrd_path = os.path.join(user_dir, f'trellis_output_{timestamp}.rrd')
+        
+        # Initialize Rerun
+        rr.init(f"Trellis_Session_{timestamp}", spawn=False)
+        
+        # Setup View Coordinates (Z-UP is standard for Trellis)
+        rr.log("world", rr.ViewCoordinates.RIGHT_HAND_Z_UP, static=True)
+        
+        # Add Reference Axis
+        rr.log("world/axis", rr.Arrows3D(
+            vectors=[[1, 0, 0], [0, 1, 0], [0, 0, 1]],
+            colors=[[255, 0, 0], [0, 255, 0], [0, 0, 255]],
+            radii=0.02
+        ), static=True)
+        
+        # Log the generated 3D model
+        rr.log("world/model", rr.Asset3D(path=glb_path))
         
-        rrd_path = create_rerun_rrd(glb_path)
+        # Save to RRD file for Gradio
+        rr.save(rrd_path)
         
         torch.cuda.empty_cache()
         return rrd_path, glb_path
@@ -344,11 +339,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) with Rerun**", elem_id="main-title")
+        gr.Markdown("# **TRELLIS.2 (Text-to-3D)**", 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.
+        **Visualized with Rerun.**
         """)
 
         with gr.Row():
@@ -394,9 +389,10 @@ if __name__ == "__main__":
 
             with gr.Column(scale=2):
                 gr.Markdown("### 3D Output (Rerun Viewer)")
+                # Replaced Model3D with Rerun Component
                 rerun_output = Rerun(
-                    label="Rerun 3D Viewer",
-                    height=600,
+                    label="Rerun 3D Viewer", 
+                    height=600
                 )
                 download_btn = gr.DownloadButton(label="3.Download GLB File", variant="primary")