update app

app.py

@@ -7,26 +7,24 @@ import tempfile
 from typing import *
 from datetime import datetime
 from pathlib import Path
-from typing import Iterable
+import uuid
 
-# Gradio and Rerun Imports
-import gradio as gr
+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
 import spaces
+from diffusers import ZImagePipeline
+
+# Rerun Imports
 import rerun as rr
-from gradio_rerun import Rerun # Requires: pip install gradio-rerun
+from gradio_rerun import Rerun
 
-# Diffusers / Trellis Imports
-from diffusers import ZImagePipeline
 from trellis2.pipelines import Trellis2ImageTo3DPipeline
 import o_voxel
 
-# ==========================================
-# THEME CONFIGURATION
-# ==========================================
-
 colors.orange_red = colors.Color(
     name="orange_red",
     c50="#FFF0E5",
@@ -95,15 +93,9 @@ class OrangeRedTheme(Soft):
 
 orange_red_theme = OrangeRedTheme()
 
-# ==========================================
-# ENVIRONMENT & MODEL SETUP
-# ==========================================
-
 os.environ["OPENCV_IO_ENABLE_OPENEXR"] = '1'
 os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True"
 os.environ["ATTN_BACKEND"] = "flash_attn_3"
-
-# Adjust path if needed for your environment
 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'
 
@@ -139,10 +131,6 @@ except Exception as e:
 
 rmbg_client = Client("briaai/BRIA-RMBG-2.0")
 
-# ==========================================
-# UTILITY FUNCTIONS
-# ==========================================
-
 def start_session(req: gr.Request):
     user_dir = os.path.join(TMP_DIR, str(req.session_hash))
     os.makedirs(user_dir, exist_ok=True)
@@ -161,7 +149,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
         
@@ -198,13 +185,8 @@ def preprocess_image(input: Image.Image) -> Image.Image:
 def get_seed(randomize_seed: bool, seed: int) -> int:
     return np.random.randint(0, MAX_SEED) if randomize_seed else seed
 
-# ==========================================
-# GENERATION FUNCTIONS
-# ==========================================
-
 @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():
@@ -228,6 +210,30 @@ 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,
@@ -288,12 +294,10 @@ 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...")
         
         grid_size = latents[2] 
@@ -317,55 +321,17 @@ def generate_3d(
         now = datetime.now()
         timestamp = now.strftime("%Y-%m-%dT%H%M%S")
         glb_path = os.path.join(user_dir, f'trellis_output_{timestamp}.glb')
-        rrd_path = os.path.join(user_dir, f'trellis_output_{timestamp}.rrd')
-        
         glb.export(glb_path, extension_webp=True)
-
-        # ---------------------------------------------------------
-        # RERUN VISUALIZATION LOGIC
-        # ---------------------------------------------------------
-        progress(0.95, desc="Creating Rerun Visualization...")
         
-        # Initialize Rerun in standard mode (no spawn, just recording)
-        rr.init(f"trellis_session_{timestamp}", spawn=False)
+        rrd_path = create_rerun_rrd(glb_path)
         
-        # 1. Setup View Coordinates (Z-Up is standard for 3D web/GLB)
-        rr.log("world", rr.ViewCoordinates.RIGHT_HAND_Z_UP, static=True)
-
-        # 2. Log Axes (Red=X, Green=Y, Blue=Z)
-        rr.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]],
-                origins=[[0, 0, 0], [0, 0, 0], [0, 0, 0]],
-                labels=["X", "Y", "Z"],
-                radii=0.015
-            ),
-            static=True
-        )
-
-        # 3. Log the Model (Asset3D)
-        # Note: TRELLIS outputs are usually normalized.
-        rr.log("world/model", rr.Asset3D(path=glb_path))
-
-        # 4. Save to RRD for Gradio Viewer
-        rr.save(rrd_path)
-
-        # Clean up memory
         torch.cuda.empty_cache()
-        
-        # Return RRD path for visualization and GLB path for download
         return rrd_path, glb_path
         
     except Exception as e:
         torch.cuda.empty_cache()
         raise gr.Error(f"Generation failed: {str(e)}")
 
-# ==========================================
-# UI LAYOUT
-# ==========================================
-
 css="""
 #col-container {
     margin: 0 auto;
@@ -378,11 +344,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:** Output is displayed using Rerun.
+        Visualization powered by Rerun.
         """)
 
         with gr.Row():
@@ -428,13 +394,10 @@ if __name__ == "__main__":
 
             with gr.Column(scale=2):
                 gr.Markdown("### 3D Output (Rerun Viewer)")
-                
-                # REPLACED gr.Model3D with Rerun
                 rerun_output = Rerun(
-                    label="Rerun Visualization",
-                    height=600
+                    label="Rerun 3D Viewer",
+                    height=600,
                 )
-                
                 download_btn = gr.DownloadButton(label="3.Download GLB File", variant="primary")
 
                 gr.Examples(