Update app.py

app.py

@@ -1,147 +1,10 @@
 import os
-# this is a HF Spaces specific hack, as
-#  (i)  building pytorch3d with GPU support is a bit tricky here
-#  (ii) installing the wheel via requirements.txt breaks ZeroGPU
+# this is a HF Spaces specific hack for ZeroGPU
 import spaces
 
-# Use the dynamic approach from PyTorch3D documentation to determine the correct wheel
 import sys
 import torch
 
-# Print debug information about the environment
-try:
-    cuda_version = torch.version.cuda
-    torch_version = torch.__version__
-    python_version = f"{sys.version_info.major}.{sys.version_info.minor}"
-    print(f"CUDA Version: {cuda_version}")
-    print(f"PyTorch Version: {torch_version}")
-    print(f"Python Version: {python_version}")
-except Exception as e:
-    print(f"Error detecting environment versions: {e}")
-
-# Install PyTorch3D properly from source
-print("Installing PyTorch3D from source...")
-
-# First uninstall any existing PyTorch3D installation to avoid conflicts
-os.system("pip uninstall -y pytorch3d")
-
-# Install dependencies required for building PyTorch3D
-print("Installing build dependencies...")
-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'")
-os.system("pip install fvcore iopath")
-
-# Clone the PyTorch3D repository
-print("Cloning PyTorch3D repository...")
-os.system("rm -rf pytorch3d")  # Remove any existing directory
-clone_result = os.system("git clone https://github.com/facebookresearch/pytorch3d.git")
-if clone_result != 0:
-    print("Failed to clone PyTorch3D repository. Trying with git protocol...")
-    clone_result = os.system("git clone git://github.com/facebookresearch/pytorch3d.git")
-    if clone_result != 0:
-        print("Failed to clone PyTorch3D repository with both HTTPS and git protocols.")
-
-# Use a specific release tag that is known to be stable
-checkout_result = os.system("cd pytorch3d && git checkout v0.7.4")
-if checkout_result != 0:
-    print("Failed to checkout v0.7.4 tag. Trying with main branch...")
-    checkout_result = os.system("cd pytorch3d && git checkout main")
-    if checkout_result != 0:
-        print("Failed to checkout main branch.")
-
-# Install PyTorch3D from source with CPU support
-print("Building PyTorch3D from source...")
-build_result = os.system("cd pytorch3d && pip install -v -e .")
-if build_result != 0:
-    print("Failed to build PyTorch3D from source with default settings.")
-    
-    # Try with specific build flags
-    print("Trying with specific build flags...")
-    os.environ["FORCE_CUDA"] = "0"  # Explicitly disable CUDA for build
-    build_result = os.system("cd pytorch3d && pip install -v -e .")
-    
-    if build_result != 0:
-        print("Failed to build PyTorch3D from source with specific build flags.")
-        
-        # Try with setup.py directly
-        print("Trying with setup.py directly...")
-        build_result = os.system("cd pytorch3d && python setup.py install")
-        
-        if build_result != 0:
-            print("All attempts to build from source failed.")
-
-# Verify the installation
-import_result = os.popen('python -c "import pytorch3d; print(\'pytorch3d imported successfully\'); try: from pytorch3d import renderer; print(\'renderer module imported successfully\'); except ImportError as e: print(f\'Error importing renderer: {e}\');" 2>&1').read()
-print(import_result)
-
-# If the installation fails, try a different approach with wheels from PyPI
-if "Error importing renderer" in import_result or "No module named" in import_result:
-    print("Source installation failed to provide renderer module, trying with PyPI...")
-    os.system("pip uninstall -y pytorch3d")
-    
-    # Try with PyPI version first (which might be CPU-only but should have renderer)
-    os.system("pip install pytorch3d")
-    
-    # Verify again
-    import_result = os.popen('python -c "import pytorch3d; print(\'pytorch3d imported successfully\'); try: from pytorch3d import renderer; print(\'renderer module imported successfully\'); except ImportError as e: print(f\'Error importing renderer: {e}\');" 2>&1').read()
-    print(import_result)
-
-# Patch the shap_e renderer to handle PyTorch3D renderer import error if needed
-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}")
-
-# Add a helper function to ensure PyTorch3D works with ZeroGPU
-def ensure_pytorch3d_cuda_compatibility():
-    """
-    This function ensures PyTorch3D works correctly with CUDA in ZeroGPU environments.
-    It should be called at the beginning of any @spaces.GPU decorated function.
-    """
-    try:
-        import pytorch3d
-        if torch.cuda.is_available():
-            # Check if we can access the renderer module
-            from pytorch3d import renderer
-            print("PyTorch3D renderer module is available with CUDA")
-        else:
-            print("CUDA is not available, using CPU version of PyTorch3D")
-    except ImportError as e:
-        print(f"Error importing PyTorch3D: {e}")
-    except Exception as e:
-        print(f"Unexpected error with PyTorch3D: {e}")
-
 import torch
 import torch.nn as nn
 import gradio as gr
@@ -342,9 +205,6 @@ def load_models():
 @spaces.GPU(duration=20)
 def process_images(input_images, prompt, steps=75, guidance_scale=7.5, pipeline=None):
     """Process input images and run refinement"""
-    # Ensure PyTorch3D works with CUDA
-    ensure_pytorch3d_cuda_compatibility()
-    
     device = pipeline.device
     
     if isinstance(input_images, list):
@@ -415,9 +275,6 @@ def process_images(input_images, prompt, steps=75, guidance_scale=7.5, pipeline=
 @spaces.GPU(duration=20)
 def create_mesh(refined_image, model, infer_config):
     """Generate mesh from refined image"""
-    # Ensure PyTorch3D works with CUDA
-    ensure_pytorch3d_cuda_compatibility()
-    
     # Convert PIL image to tensor
     image = np.array(refined_image) / 255.0
     image = torch.from_numpy(image).float().permute(2, 0, 1)
@@ -680,9 +537,6 @@ def create_demo():
         @spaces.GPU(duration=20)  # Reduced duration to 20 seconds
         def generate(prompt, guidance_scale, num_steps):
             try:
-                # Ensure PyTorch3D works with CUDA
-                ensure_pytorch3d_cuda_compatibility()
-                
                 torch.cuda.empty_cache()  # Clear GPU memory before starting
                 with torch.no_grad():
                     layout, _ = shap_e.generate_views(prompt, guidance_scale, num_steps)
@@ -696,9 +550,6 @@ def create_demo():
         @spaces.GPU(duration=20)  # Reduced duration to 20 seconds
         def refine(input_image, prompt, steps, guidance_scale):
             try:
-                # Ensure PyTorch3D works with CUDA
-                ensure_pytorch3d_cuda_compatibility()
-                
                 torch.cuda.empty_cache()  # Clear GPU memory before starting
                 refined_img, mesh_path = refiner.refine_model(
                     input_image,