update app

app.py

@@ -12,18 +12,6 @@ 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
-
-from trellis2.pipelines import Trellis2ImageTo3DPipeline
-import o_voxel
-
 colors.orange_red = colors.Color(
     name="orange_red",
     c50="#FFF0E5",
@@ -99,10 +87,19 @@ 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')
 
 print("Initializing models...")
+
 print("Loading Z-Image-Turbo...")
 try:
     z_pipe = ZImagePipeline.from_pretrained(
@@ -148,6 +145,7 @@ 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
         
@@ -186,6 +184,7 @@ 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():
@@ -269,12 +268,15 @@ def generate_3d(
             return_latent=True,
         )
         
+        # 2. Process Mesh
         progress(0.7, desc="Processing Mesh...")
         mesh = outputs[0]
-        mesh.simplify(16777216)
+        mesh.simplify(16777216) # Simplify for processing limits
         
+        # 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(
@@ -297,28 +299,10 @@ 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 Logic
-        progress(0.95, desc="Preparing Rerun Visualization...")
-        rrd_path = os.path.join(user_dir, f'trellis_output_{timestamp}.rrd')
-        
-        rec = rr.new_recording(application_id="TRELLIS_Viewer", recording_id=str(req.session_hash))
-        
-        # Set view coordinates (Z-Up)
-        rec.log("world", rr.ViewCoordinates.RIGHT_HAND_Z_UP, static=True)
-
-        # Log visual reference axis
-        rec.log("world/axis/x", rr.Arrows3D(vectors=[[0.5, 0, 0]], colors=[[255, 0, 0]], radii=[0.02]), static=True)
-        rec.log("world/axis/y", rr.Arrows3D(vectors=[[0, 0.5, 0]], colors=[[0, 255, 0]], radii=[0.02]), static=True)
-        rec.log("world/axis/z", rr.Arrows3D(vectors=[[0, 0, 0.5]], colors=[[0, 0, 255]], radii=[0.02]), static=True)
-
-        # Log the 3D asset
-        rec.log("world/model", rr.Asset3D(path=glb_path), static=True)
-        
-        rec.save(rrd_path)
         
+        # Clean up
         torch.cuda.empty_cache()
-        return rrd_path, glb_path
+        return glb_path, glb_path
         
     except Exception as e:
         torch.cuda.empty_cache()
@@ -359,6 +343,7 @@ 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)
                 
@@ -385,13 +370,13 @@ if __name__ == "__main__":
 
             with gr.Column(scale=2):
                 gr.Markdown("### 3D Output")
-                
-                # Rerun Component
-                rerun_output = Rerun(
-                    label="Rerun 3D Viewer",
+                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
                 )
-                
                 download_btn = gr.DownloadButton(label="3.Download GLB File", variant="primary")
 
                 gr.Examples(
@@ -575,9 +560,8 @@ if __name__ == "__main__":
                 ss_guidance_strength, ss_guidance_rescale, ss_sampling_steps, ss_rescale_t,
                 shape_guidance, shape_rescale, shape_steps, shape_rescale_t,
                 tex_guidance, tex_rescale, tex_steps, tex_rescale_t,
-                
             ],
-            outputs=[rerun_output, download_btn],
+            outputs=[glb_output, download_btn],
         )
 
     demo.launch(theme=orange_red_theme, css=css, mcp_server=True, ssr_mode=False, show_error=True)