update app

app.py

@@ -4,6 +4,7 @@ import torch
 import numpy as np
 from PIL import Image
 import tempfile
+import uuid
 from typing import *
 from datetime import datetime
 from pathlib import Path
@@ -12,6 +13,10 @@ from typing import Iterable
 from gradio.themes import Soft
 from gradio.themes.utils import colors, fonts, sizes
 
+import rerun as rr
+from gradio_rerun import Rerun
+
+# --- Theme Configuration ---
 colors.orange_red = colors.Color(
     name="orange_red",
     c50="#FFF0E5",
@@ -80,10 +85,10 @@ class OrangeRedTheme(Soft):
 
 orange_red_theme = OrangeRedTheme()
 
+# --- Environment Setup ---
 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'
 
@@ -91,7 +96,6 @@ 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
 
@@ -300,9 +304,32 @@ def generate_3d(
         glb_path = os.path.join(user_dir, f'trellis_output_{timestamp}.glb')
         glb.export(glb_path, extension_webp=True)
         
+        # --- Rerun Visualization Logic ---
+        progress(0.95, desc="Creating Rerun Visualization...")
+        
+        run_id = str(uuid.uuid4())
+        
+        # Create a new Rerun recording
+        rec = rr.new_recording(application_id="TRELLIS-3D-Viewer", recording_id=run_id)
+        
+        # Read the GLB data to embed it in the RRD
+        # This ensures the web viewer can see it without needing local file access
+        with open(glb_path, 'rb') as f:
+            glb_data = f.read()
+            
+        # Log the GLB asset to Rerun
+        rec.log("generated_model", rr.Asset3D(blob=glb_data, media_type=rr.MediaType.GLB), static=True)
+        
+        # Also log the input image for context
+        rec.log("input_image", rr.Image(np.array(image)), static=True)
+        
+        # Save the Rerun recording (.rrd)
+        rrd_path = os.path.join(user_dir, f'trellis_output_{timestamp}.rrd')
+        rec.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()
@@ -343,7 +370,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)
                 
@@ -370,13 +396,13 @@ if __name__ == "__main__":
 
             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
+                
+                # Replaced gr.Model3D with Rerun component
+                rerun_output = Rerun(
+                    label="Rerun 3D Viewer",
+                    height=600
                 )
+                
                 download_btn = gr.DownloadButton(label="3.Download GLB File", variant="primary")
 
                 gr.Examples(
@@ -386,72 +412,6 @@ if __name__ == "__main__":
                         ["example-images/A (3).webp"],
                         ["example-images/A (4).webp"],
                         ["example-images/A (5).webp"],
-                        ["example-images/A (6).webp"],
-                        ["example-images/A (7).webp"],
-                        ["example-images/A (8).webp"],
-                        ["example-images/A (9).webp"],
-                        ["example-images/A (10).webp"],
-                        ["example-images/A (11).webp"],
-                        ["example-images/A (12).webp"],
-                        ["example-images/A (13).webp"],
-                        ["example-images/A (14).webp"],
-                        ["example-images/A (15).webp"],
-                        ["example-images/A (16).webp"],
-                        ["example-images/A (17).webp"],
-                        ["example-images/A (18).webp"],
-                        ["example-images/A (19).webp"],
-                        ["example-images/A (20).webp"],
-                        ["example-images/A (21).webp"],
-                        ["example-images/A (22).webp"],
-                        ["example-images/A (23).webp"],
-                        ["example-images/A (24).webp"],
-                        ["example-images/A (25).webp"],
-                        ["example-images/A (26).webp"],
-                        ["example-images/A (27).webp"],
-                        ["example-images/A (28).webp"],
-                        ["example-images/A (29).webp"],
-                        ["example-images/A (30).webp"],
-                        ["example-images/A (31).webp"],
-                        ["example-images/A (32).webp"],
-                        ["example-images/A (33).webp"],
-                        ["example-images/A (34).webp"],
-                        ["example-images/A (35).webp"],
-                        ["example-images/A (36).webp"],
-                        ["example-images/A (37).webp"],
-                        ["example-images/A (38).webp"],
-                        ["example-images/A (39).webp"],
-                        ["example-images/A (40).webp"],
-                        ["example-images/A (41).webp"],
-                        ["example-images/A (42).webp"],
-                        ["example-images/A (43).webp"],
-                        ["example-images/A (44).webp"],
-                        ["example-images/A (45).webp"],
-                        ["example-images/A (46).webp"],
-                        ["example-images/A (47).webp"],
-                        ["example-images/A (48).webp"],
-                        ["example-images/A (49).webp"],
-                        ["example-images/A (50).webp"],
-                        ["example-images/A (51).webp"],
-                        ["example-images/A (52).webp"],
-                        ["example-images/A (53).webp"],
-                        ["example-images/A (54).webp"],
-                        ["example-images/A (55).webp"],
-                        ["example-images/A (56).webp"],
-                        ["example-images/A (57).webp"],
-                        ["example-images/A (58).webp"],
-                        ["example-images/A (59).webp"],
-                        ["example-images/A (60).webp"],
-                        ["example-images/A (61).webp"],
-                        ["example-images/A (62).webp"],
-                        ["example-images/A (63).webp"],
-                        ["example-images/A (64).webp"],
-                        ["example-images/A (65).webp"],
-                        ["example-images/A (66).webp"],
-                        ["example-images/A (67).webp"],
-                        ["example-images/A (68).webp"],
-                        ["example-images/A (69).webp"],
-                        ["example-images/A (70).webp"],
-                        ["example-images/A (71).webp"],
                     ],
                     inputs=[image_prompt],
                     label="Image Examples [image-to-3d]"
@@ -460,70 +420,9 @@ if __name__ == "__main__":
                 gr.Examples(
                     examples=[
                         ["A Cat 3D model"],
-                        ["A realistic Cat 3D model"],
-                        ["A cartoon Cat 3D model"],
-                        ["A low poly Cat 3D"],
-                        ["A cyberpunk Cat 3D"],
-                        ["A robotic Cat 3D"],
-                        ["A fluffy Cat 3D"],
-                        ["A fantasy Cat 3D creature"],
-                        ["A stylized Cat 3D"],
-                        ["A Cat 3D sculpture"],
-                
-                        ["A Plane 3D model"],
-                        ["A commercial Plane 3D"],
-                        ["A fighter jet Plane 3D"],
-                        ["A low poly Plane 3D"],
-                        ["A vintage Plane 3D"],
-                        ["A futuristic Plane 3D"],
-                        ["A cargo Plane 3D"],
-                        ["A private jet Plane 3D"],
-                        ["A toy Plane 3D"],
                         ["A realistic Plane 3D"],
-                
-                        ["A Car 3D model"],
-                        ["A sports Car 3D"],
-                        ["A luxury Car 3D"],
-                        ["A low poly Car 3D"],
-                        ["A racing Car 3D"],
                         ["A cyberpunk Car 3D"],
-                        ["A vintage Car 3D"],
-                        ["A futuristic Car 3D"],
-                        ["A SUV Car 3D"],
-                        ["A electric Car 3D"],
-                
-                        ["A Shoe 3D model"],
-                        ["A sneaker Shoe 3D"],
-                        ["A running Shoe 3D"],
-                        ["A leather Shoe 3D"],
-                        ["A high heel Shoe 3D"],
-                        ["A boot Shoe 3D"],
-                        ["A low poly Shoe 3D"],
                         ["A futuristic Shoe 3D"],
-                        ["A sports Shoe 3D"],
-                        ["A casual Shoe 3D"],
-                
-                        ["A Chair 3D model"],
-                        ["A Table 3D model"],
-                        ["A Sofa 3D model"],
-                        ["A Lamp 3D model"],
-                        ["A Watch 3D model"],
-                        ["A Backpack 3D model"],
-                        ["A Drone 3D model"],
-                        ["A Robot 3D model"],
-                        ["A Smartphone 3D model"],
-                        ["A Headphones 3D model"],
-                
-                        ["A House 3D model"],
-                        ["A Skyscraper 3D model"],
-                        ["A Bridge 3D model"],
-                        ["A Castle 3D model"],
-                        ["A Spaceship 3D model"],
-                        ["A Rocket 3D model"],
-                        ["A Satellite 3D model"],
-                        ["A Tank 3D model"],
-                        ["A Motorcycle 3D model"],
-                        ["A Bicycle 3D model"]
                     ],
                     inputs=[txt_prompt],
                     label="3D Prompt Examples [text-to-3d]"
@@ -561,7 +460,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)