Update app.py

app.py

@@ -95,7 +95,9 @@ 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"
+# FIXED: Changed ATTN_BACKEND from 'flash_attn_3' to 'xformers' to avoid import error.
+# If 'xformers' also fails, try changing this to 'sdpa'
+os.environ["ATTN_BACKEND"] = "xformers" 
 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'
 
@@ -301,23 +303,19 @@ def generate_3d(
         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)
+        # 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))
         
-        # Save to RRD file for Gradio
         rr.save(rrd_path)
         
         torch.cuda.empty_cache()
@@ -325,7 +323,11 @@ def generate_3d(
         
     except Exception as e:
         torch.cuda.empty_cache()
-        raise gr.Error(f"Generation failed: {str(e)}")
+        # More descriptive error handling
+        error_msg = str(e)
+        if "No module named 'flash_attn'" in error_msg:
+            error_msg += " (Try changing ATTN_BACKEND to 'sdpa' in the code)"
+        raise gr.Error(f"Generation failed: {error_msg}")
 
 css="""
 #col-container {
@@ -389,7 +391,6 @@ 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