Update app.py

app.py

@@ -19,124 +19,44 @@ try:
 except Exception as e:
     print(f"Error detecting environment versions: {e}")
 
-# Install PyTorch3D with CPU support for ZeroGPU environment
-print("Installing PyTorch3D for CPU/ZeroGPU environment...")
+# 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 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'")
 os.system("pip install fvcore iopath")
 
-# 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")
+# Clone the PyTorch3D repository
+os.system("rm -rf pytorch3d")  # Remove any existing directory
+os.system("git clone https://github.com/facebookresearch/pytorch3d.git")
 
-# 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 release tag that is known to be stable
+os.system("cd pytorch3d && git checkout v0.7.4")
 
-# 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
-    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 pytorch3d directory
-        with open(f"{pytorch3d_dir}/__init__.py", "w") as f:
-            f.write("""
-# PyTorch3D module
-import torch
-import warnings
+# Install PyTorch3D from source with CPU support
+os.system("cd pytorch3d && pip install -e .")
 
-warnings.warn("Using custom PyTorch3D module")
+# Verify the installation
+import_result = os.popen('python -c "import pytorch3d; from pytorch3d import renderer; print(\'PyTorch3D and renderer successfully imported\')" 2>&1').read()
+print(import_result)
 
-__version__ = "0.7.4"
-""")
+# If the installation fails, try a different approach with a specific wheel
+if "No module named" in import_result or "Error" in import_result:
+    print("Source installation failed, trying with a specific wheel...")
+    os.system("pip uninstall -y pytorch3d")
     
-    # Create the renderer directory if it doesn't exist
-    os.system(f"mkdir -p {pytorch3d_dir}/renderer")
+    # Try with a specific wheel that's known to work
+    os.system("pip install https://dl.fbaipublicfiles.com/pytorch3d/packaging/wheels/py310_cpu_pyt201/pytorch3d-0.7.4-cp310-cp310-linux_x86_64.whl")
     
-    # 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
-
-warnings.warn("Using custom PyTorch3D renderer module")
-
-# Basic renderer components
-class TexturesBase:
-    def __init__(self, *args, **kwargs):
-        pass
-
-class TexturesVertex(TexturesBase):
-    def __init__(self, *args, **kwargs):
-        super().__init__(*args, **kwargs)
-
-class TexturesUV(TexturesBase):
-    def __init__(self, *args, **kwargs):
-        super().__init__(*args, **kwargs)
-
-class TexturesAtlas(TexturesBase):
-    def __init__(self, *args, **kwargs):
-        super().__init__(*args, **kwargs)
-
-class RasterizationSettings:
-    def __init__(self, *args, **kwargs):
-        pass
-
-class MeshRasterizer:
-    def __init__(self, *args, **kwargs):
-        pass
-
-class PointLights:
-    def __init__(self, *args, **kwargs):
-        pass
-
-class Materials:
-    def __init__(self, *args, **kwargs):
-        pass
+    # 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)
 
-class MeshRenderer:
-    def __init__(self, *args, **kwargs):
-        pass
-
-class SoftPhongShader:
-    def __init__(self, *args, **kwargs):
-        pass
-
-class HardPhongShader:
-    def __init__(self, *args, **kwargs):
-        pass
-
-class SoftSilhouetteShader:
-    def __init__(self, *args, **kwargs):
-        pass
-
-class BlendParams:
-    def __init__(self, *args, **kwargs):
-        pass
-
-def look_at_view_transform(*args, **kwargs):
-    return torch.eye(4), torch.eye(4)
-
-class FoVPerspectiveCameras:
-    def __init__(self, *args, **kwargs):
-        pass
-""")
-    
-    print("Created custom renderer module")
-
-# Patch the shap_e renderer to handle PyTorch3D renderer import error
+# 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}")
@@ -173,9 +93,24 @@ if os.path.exists(shap_e_renderer_path):
 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)
+# 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
@@ -354,6 +289,9 @@ 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):
@@ -424,6 +362,9 @@ 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)
@@ -686,6 +627,9 @@ 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)
@@ -699,11 +643,14 @@ 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, 
-                    prompt, 
-                    steps, 
+                    prompt,
+                    steps,
                     guidance_scale
                 )
                 return refined_img, mesh_path, None  # Return None for error message