Update app.py

app.py

@@ -19,77 +19,54 @@ try:
 except Exception as e:
     print(f"Error detecting environment versions: {e}")
 
-# Install PyTorch3D from source
-print("Installing PyTorch3D from source...")
+# Install PyTorch3D with CPU support for ZeroGPU environment
+print("Installing PyTorch3D for CPU/ZeroGPU environment...")
 
 # First uninstall any existing PyTorch3D installation to avoid conflicts
 os.system("pip uninstall -y pytorch3d")
 
-# Install dependencies required for building PyTorch3D
-os.system("apt-get update && apt-get install -y git build-essential libglib2.0-0 libsm6 libxrender-dev libxext6 ninja-build")
-os.system("pip install 'imageio>=2.5.0' 'matplotlib>=3.1.2' 'numpy>=1.17.3' 'psutil>=5.6.5' 'scipy>=1.3.2' 'tqdm>=4.42.1' 'trimesh>=3.0.0'")
+# Install dependencies required for PyTorch3D
 os.system("pip install fvcore iopath")
 
-# Clone the PyTorch3D repository
-os.system("rm -rf pytorch3d")  # Remove any existing directory
-os.system("git clone https://github.com/facebookresearch/pytorch3d.git")
+# For ZeroGPU, we need a CPU-only version that still has the renderer module
+# Try installing from a pre-built wheel first (faster)
+os.system("pip install pytorch3d")
 
-# Try to build with CUDA support first
-print("Attempting to build PyTorch3D with CUDA support...")
-os.environ["FORCE_CUDA"] = "1"
-os.environ["TORCH_CUDA_ARCH_LIST"] = "3.5;5.0;6.0;6.1;7.0;7.5;8.0;8.6+PTX"
+# Verify if the renderer module is available
+import_result = os.popen('python -c "import pytorch3d; print(\'renderer\' in dir(pytorch3d))" 2>&1').read().strip()
+print(f"Renderer module check: {import_result}")
 
-# Use a specific commit that is known to be stable
-os.system("cd pytorch3d && git checkout v0.7.4")
-
-# Install with CUDA support
-cuda_build_result = os.system("cd pytorch3d && FORCE_CUDA=1 pip install -e .")
-
-# Verify installation and specifically check for renderer module
-import_result = os.popen('python -c "import pytorch3d; from pytorch3d import renderer; print(\'PyTorch3D and renderer successfully imported\')" 2>&1').read()
-print(import_result)
-
-# If CUDA build fails or import fails, try CPU-only build
-if cuda_build_result != 0 or 'No module named' in import_result or 'Error' in import_result:
-    print("CUDA build failed or import failed, trying CPU-only build...")
-    
-    # Uninstall any failed installation
-    os.system("pip uninstall -y pytorch3d")
-    
-    # Clean the repository
-    os.system("rm -rf pytorch3d")
-    os.system("git clone https://github.com/facebookresearch/pytorch3d.git")
-    os.system("cd pytorch3d && git checkout v0.7.4")
-    
-    # Set environment variables for CPU-only build
-    os.environ.pop("FORCE_CUDA", None)
-    
-    # Install without CUDA
-    os.system("cd pytorch3d && pip install -e .")
-    
-    # Verify again
-    import_result = os.popen('python -c "import pytorch3d; from pytorch3d import renderer; print(\'PyTorch3D and renderer successfully imported\')" 2>&1').read()
-    print(import_result)
-
-# Check specifically for the renderer module
-renderer_check = os.popen('python -c "import pytorch3d; print(\'renderer\' in dir(pytorch3d))" 2>&1').read().strip()
-print(f"Renderer module check: {renderer_check}")
-
-# If the renderer module is still missing, create a symbolic link to ensure it's found
-if renderer_check == "False" or "No module named" in import_result:
-    print("Renderer module is missing, creating a workaround...")
+# If the renderer module is not available, create it manually
+if import_result == "False" or "Error" in import_result:
+    print("Renderer module not found, creating it manually...")
     
     # Get the site-packages directory
     site_packages = os.popen('python -c "import site; print(site.getsitepackages()[0])"').read().strip()
     
     # Check if the pytorch3d directory exists
-    if os.path.exists(f"{site_packages}/pytorch3d"):
-        # Create the renderer directory if it doesn't exist
-        os.system(f"mkdir -p {site_packages}/pytorch3d/renderer")
+    pytorch3d_dir = f"{site_packages}/pytorch3d"
+    if not os.path.exists(pytorch3d_dir):
+        print(f"PyTorch3D directory not found at {pytorch3d_dir}, creating it...")
+        os.system(f"mkdir -p {pytorch3d_dir}")
         
-        # Create an __init__.py file in the renderer directory
-        with open(f"{site_packages}/pytorch3d/renderer/__init__.py", "w") as f:
+        # Create an __init__.py file in the pytorch3d directory
+        with open(f"{pytorch3d_dir}/__init__.py", "w") as f:
             f.write("""
+# PyTorch3D module
+import torch
+import warnings
+
+warnings.warn("Using custom PyTorch3D module")
+
+__version__ = "0.7.4"
+""")
+    
+    # Create the renderer directory if it doesn't exist
+    os.system(f"mkdir -p {pytorch3d_dir}/renderer")
+    
+    # Create an __init__.py file in the renderer directory
+    with open(f"{pytorch3d_dir}/renderer/__init__.py", "w") as f:
+        f.write("""
 # PyTorch3D renderer module
 import torch
 import warnings
@@ -156,10 +133,49 @@ class FoVPerspectiveCameras:
     def __init__(self, *args, **kwargs):
         pass
 """)
-        
-        print("Created custom renderer module")
-    else:
-        print(f"PyTorch3D directory not found at {site_packages}/pytorch3d")
+    
+    print("Created custom renderer module")
+
+# Patch the shap_e renderer to handle PyTorch3D renderer import error
+shap_e_renderer_path = "/usr/local/lib/python3.10/site-packages/shap_e/models/stf/renderer.py"
+if os.path.exists(shap_e_renderer_path):
+    print(f"Patching shap_e renderer at {shap_e_renderer_path}")
+    
+    # Read the current content
+    with open(shap_e_renderer_path, "r") as f:
+        content = f.read()
+    
+    # Create a backup
+    os.system(f"cp {shap_e_renderer_path} {shap_e_renderer_path}.bak")
+    
+    # Modify the content to handle the error more gracefully
+    modified_content = content
+    
+    # Replace the error message
+    if "exception rendering with PyTorch3D" in content:
+        modified_content = modified_content.replace(
+            'warnings.warn(f"exception rendering with PyTorch3D: {exc}")',
+            'warnings.warn("Using native PyTorch renderer")'
+        )
+    
+    # Replace the fallback warning
+    if "falling back on native PyTorch renderer" in modified_content:
+        modified_content = modified_content.replace(
+            'warnings.warn("falling back on native PyTorch renderer, which does not support full gradients")',
+            'warnings.warn("Using native PyTorch renderer")'
+        )
+    
+    # Write the modified content
+    with open(shap_e_renderer_path, "w") as f:
+        f.write(modified_content)
+    
+    print("Successfully patched shap_e renderer")
+else:
+    print(f"shap_e renderer not found at {shap_e_renderer_path}")
+
+# Verify the installation again
+import_result = os.popen('python -c "import pytorch3d; from pytorch3d import renderer; print(\'PyTorch3D and renderer successfully imported\')" 2>&1').read()
+print(import_result)
 
 import torch
 import torch.nn as nn