Update app.py

app.py

@@ -1,94 +1,91 @@
-# Copyright (c) Meta Platforms, Inc. and affiliates.
 import os
 import subprocess
 import sys
+import shutil
+
+# ============ Configuration ============
+REPO_URL = "https://github.com/facebookresearch/sam-3d-objects.git"
+REPO_DIR = "/home/user/app/sam-3d-objects"
 
 # ============ Install Dependencies & Setup ============
 def install_dependencies():
-    """Install all required dependencies."""
-    subprocess.run([sys.executable, "-m", "pip", "install", "--upgrade", "pip"], check=True)
-    
-    # Core dependencies
-    core_deps = [
-        "spaces",
-        "gradio>=4.0.0",
-        "numpy",
-        "Pillow",
-        "omegaconf",
-        "hydra-core",
-        "einops",
-        "timm",
-        "safetensors",
-        "accelerate",
-        "transformers",
-        "diffusers",
-        "trimesh",
-        "utils3d",
-    ]
-    subprocess.run([sys.executable, "-m", "pip", "install"] + core_deps, check=True)
-    
-    # PyTorch with CUDA 12.1
-    subprocess.run([
-        sys.executable, "-m", "pip", "install",
-        "torch>=2.1.0", "torchvision",
-        "--extra-index-url", "https://download.pytorch.org/whl/cu121"
-    ], check=True)
+    """
+    Installs dependencies using the official repo method (pip install -e .[extras])
+    instead of manual package listing.
+    """
+    print("Starting installation sequence...")
+
+    # 1. Clone Repository
+    if not os.path.exists(REPO_DIR):
+        print(f"Cloning repository to {REPO_DIR}...")
+        subprocess.run(["git", "clone", REPO_URL, REPO_DIR], check=True)
     
-    # PyTorch3D
-    subprocess.run([
-        sys.executable, "-m", "pip", "install",
-        "pytorch3d",
-        "--find-links", "https://dl.fbaipublicfiles.com/pytorch3d/packaging/wheels/py310_cu121_pyt251/download.html"
-    ], check=True)
+    # Switch working directory to repo for local installs
+    os.chdir(REPO_DIR)
+
+    # 2. Set Environment Variables for PIP
+    # As per instructions: export PIP_EXTRA_INDEX_URL and PIP_FIND_LINKS
+    env = os.environ.copy()
+    env["PIP_EXTRA_INDEX_URL"] = "https://pypi.ngc.nvidia.com https://download.pytorch.org/whl/cu121"
+    env["PIP_FIND_LINKS"] = "https://nvidia-kaolin.s3.us-east-2.amazonaws.com/torch-2.5.1_cu121.html"
     
-    # xformers (optional but recommended)
-    subprocess.run([sys.executable, "-m", "pip", "install", "xformers"], check=False)
+    # Upgrade pip first
+    subprocess.run([sys.executable, "-m", "pip", "install", "--upgrade", "pip"], env=env, check=True)
+
+    # 3. Install Dependencies via setup.py extras
+    # Step A: Install [dev]
+    print("Installing [dev] dependencies...")
+    subprocess.run([sys.executable, "-m", "pip", "install", "-e", ".[dev]"], env=env, check=True)
+
+    # Step B: Install [p3d] - The 2-step approach mentioned in instructions
+    print("Installing [p3d] dependencies...")
+    subprocess.run([sys.executable, "-m", "pip", "install", "-e", ".[p3d]"], env=env, check=True)
+
+    # Step C: Install [inference]
+    print("Installing [inference] dependencies...")
+    subprocess.run([sys.executable, "-m", "pip", "install", "-e", ".[inference]"], env=env, check=True)
+
+    # 4. Apply Patches
+    # Run ./patching/hydra
+    patch_script = os.path.join(REPO_DIR, "patching", "hydra")
+    if os.path.exists(patch_script):
+        print("Applying Hydra patch...")
+        subprocess.run(["chmod", "+x", patch_script], check=True)
+        subprocess.run([patch_script], check=True)
+    else:
+        print(f"Warning: Patch script not found at {patch_script}")
 
 # Run installation
-print("Installing dependencies...")
 install_dependencies()
 
-# Clone repo
-REPO_DIR = "/home/user/app/sam-3d-objects"
-if not os.path.exists(REPO_DIR):
-    print("Cloning sam-3d-objects repository...")
-    subprocess.run([
-        "git", "clone",
-        "https://github.com/facebookresearch/sam-3d-objects.git",
-        REPO_DIR
-    ], check=True)
-    
-    # Install sam3d-objects package
-    subprocess.run([
-        sys.executable, "-m", "pip", "install", "-e", REPO_DIR
-    ], check=True)
-
 # Add repo to Python path
 if REPO_DIR not in sys.path:
     sys.path.insert(0, REPO_DIR)
 
-# Set environment variables
+# Set environment variables required for runtime
 os.environ["CUDA_HOME"] = "/usr/local/cuda"
 os.environ["LIDRA_SKIP_INIT"] = "true"
-os.environ["PYTORCH3D_NO_CUDA_CHECK"] = "1"
+# Often required to prevent Pytorch3D checks in certain container environments
+os.environ["PYTORCH3D_NO_CUDA_CHECK"] = "1" 
 
+# ============ Imports ============
 import spaces
+import builtins
 from typing import Optional, List, Callable
+from copy import deepcopy
+import gradio as gr
 import numpy as np
 from PIL import Image
-from omegaconf import OmegaConf, DictConfig, ListConfig
-from hydra.utils import instantiate, get_method
 import torch
 import math
-import shutil
-import builtins
-from copy import deepcopy
-import gradio as gr
+from omegaconf import OmegaConf, DictConfig, ListConfig
+from hydra.utils import instantiate, get_method
 
-# Lazy imports - loaded after GPU allocation
+# Lazy imports placehoder
 _sam3d_imported = False
 _pipeline = None
 
+# ============ Security / Config Filters ============
 WHITELIST_FILTERS = [
     lambda target: target.split(".", 1)[0] in {"sam3d_objects", "torch", "torchvision", "moge"},
 ]
@@ -124,6 +121,7 @@ def check_hydra_safety(config: DictConfig, whitelist_filters: List[Callable], bl
         elif isinstance(node, ListConfig):
             to_check.extend(list(node))
 
+# ============ Lazy Loading & Model Logic ============
 def lazy_import_sam3d():
     """Import sam3d modules lazily after GPU is available."""
     global _sam3d_imported
@@ -131,25 +129,30 @@ def lazy_import_sam3d():
         global utils3d, sam3d_objects, InferencePipelinePointMap, render_utils, SceneVisualizer
         global quaternion_multiply, quaternion_invert
         
-        import utils3d as _utils3d
-        utils3d = _utils3d
-        
-        import sam3d_objects as _sam3d_objects
-        sam3d_objects = _sam3d_objects
-        
-        from sam3d_objects.pipeline.inference_pipeline_pointmap import InferencePipelinePointMap as _IPP
-        InferencePipelinePointMap = _IPP
-        
-        from sam3d_objects.model.backbone.tdfy_dit.utils import render_utils as _ru
-        render_utils = _ru
-        
-        from sam3d_objects.utils.visualization import SceneVisualizer as _SV
-        SceneVisualizer = _SV
-        
-        from pytorch3d.transforms import quaternion_multiply as _qm, quaternion_invert as _qi
-        quaternion_multiply, quaternion_invert = _qm, _qi
-        
-        _sam3d_imported = True
+        try:
+            import utils3d as _utils3d
+            utils3d = _utils3d
+            
+            import sam3d_objects as _sam3d_objects
+            sam3d_objects = _sam3d_objects
+            
+            from sam3d_objects.pipeline.inference_pipeline_pointmap import InferencePipelinePointMap as _IPP
+            InferencePipelinePointMap = _IPP
+            
+            from sam3d_objects.model.backbone.tdfy_dit.utils import render_utils as _ru
+            render_utils = _ru
+            
+            from sam3d_objects.utils.visualization import SceneVisualizer as _SV
+            SceneVisualizer = _SV
+            
+            from pytorch3d.transforms import quaternion_multiply as _qm, quaternion_invert as _qi
+            quaternion_multiply, quaternion_invert = _qm, _qi
+            
+            _sam3d_imported = True
+        except ImportError as e:
+            print(f"Failed to import SAM 3D modules: {e}")
+            print("Ensure the installation step completed successfully.")
+            raise
 
 def load_pipeline(config_file: str):
     """Load the inference pipeline (call inside GPU context)."""
@@ -173,11 +176,14 @@ def merge_mask_to_rgba(image, mask):
 def run_inference(image: np.ndarray, mask: np.ndarray, config_file: str, seed: Optional[int] = None, pointmap=None) -> dict:
     """GPU-decorated inference function for ZeroGPU."""
     global _pipeline
+    
+    # Ensure pipeline is loaded
     _pipeline = load_pipeline(config_file)
     if hasattr(_pipeline, 'to'):
         _pipeline.to('cuda')
     
     rgba_image = merge_mask_to_rgba(image, mask)
+    
     return _pipeline.run(
         rgba_image, None, seed,
         stage1_only=False,
@@ -189,6 +195,8 @@ def run_inference(image: np.ndarray, mask: np.ndarray, config_file: str, seed: O
         pointmap=pointmap,
     )
 
+# ============ Rendering Helpers ============
+# (Retained from original script logic for rendering frames)
 def _yaw_pitch_r_fov_to_extrinsics_intrinsics(yaws, pitchs, rs, fovs):
     lazy_import_sam3d()
     is_list = isinstance(yaws, list)
@@ -225,72 +233,6 @@ def render_video_gpu(sample, resolution=512, bg_color=(0,0,0), num_frames=300, r
     extr, intr = _yaw_pitch_r_fov_to_extrinsics_intrinsics(yaws, pitch, r, fov)
     return render_utils.render_frames(sample, extr, intr, {"resolution": resolution, "bg_color": bg_color, "backend": "gsplat"}, **kwargs)
 
-def normalized_gaussian(scene_gs, in_place=False, outlier_percentile=None):
-    if not in_place:
-        scene_gs = deepcopy(scene_gs)
-    orig_xyz, orig_scale = scene_gs.get_xyz, scene_gs.get_scaling
-    active_mask = (scene_gs.get_opacity > 0.9).squeeze()
-    inv_scale = (orig_xyz[active_mask].max(dim=0)[0] - orig_xyz[active_mask].min(dim=0)[0]).max()
-    norm_scale, norm_xyz = orig_scale / inv_scale, orig_xyz / inv_scale
-    if outlier_percentile is None:
-        lower = torch.min(norm_xyz[active_mask], dim=0)[0]
-        upper = torch.max(norm_xyz[active_mask], dim=0)[0]
-    else:
-        lower = torch.quantile(norm_xyz[active_mask], outlier_percentile, dim=0)
-        upper = torch.quantile(norm_xyz[active_mask], 1.0 - outlier_percentile, dim=0)
-    scene_gs.from_xyz(norm_xyz - (lower + upper) / 2)
-    scene_gs.mininum_kernel_size /= inv_scale.item()
-    scene_gs.from_scaling(norm_scale)
-    return scene_gs
-
-def _fix_gaussian_alignment(scene_gs, in_place=False):
-    if not in_place:
-        scene_gs = deepcopy(scene_gs)
-    device, dtype = scene_gs._xyz.device, scene_gs._xyz.dtype
-    scene_gs._xyz = scene_gs._xyz @ torch.tensor([[-1,0,0],[0,0,1],[0,1,0]], device=device, dtype=dtype).T
-    return scene_gs
-
-def ready_gaussian_for_video_rendering(scene_gs, in_place=False, fix_alignment=False):
-    if fix_alignment:
-        scene_gs = _fix_gaussian_alignment(scene_gs, in_place=in_place)
-    return normalized_gaussian(scene_gs, in_place=fix_alignment)
-
-def make_scene(*outputs, in_place=False):
-    lazy_import_sam3d()
-    if not in_place:
-        outputs = [deepcopy(o) for o in outputs]
-    all_outs, min_kernel = [], float("inf")
-    for output in outputs:
-        PC = SceneVisualizer.object_pointcloud(
-            points_local=output["gaussian"][0].get_xyz.unsqueeze(0),
-            quat_l2c=output["rotation"], trans_l2c=output["translation"], scale_l2c=output["scale"])
-        output["gaussian"][0].from_xyz(PC.points_list()[0])
-        output["gaussian"][0].from_rotation(quaternion_multiply(quaternion_invert(output["rotation"]), output["gaussian"][0].get_rotation))
-        scale = output["gaussian"][0].get_scaling * output["scale"]
-        assert output["scale"][0,0].item() == output["scale"][0,1].item() == output["scale"][0,2].item()
-        output["gaussian"][0].mininum_kernel_size *= output["scale"][0,0].item()
-        scale = torch.maximum(scale, torch.tensor(output["gaussian"][0].mininum_kernel_size * 1.1, device=scale.device))
-        output["gaussian"][0].from_scaling(scale)
-        min_kernel = min(min_kernel, output["gaussian"][0].mininum_kernel_size)
-        all_outs.append(output)
-    scene_gs = all_outs[0]["gaussian"][0]
-    scene_gs.mininum_kernel_size = min_kernel
-    for out in all_outs[1:]:
-        gs = out["gaussian"][0]
-        scene_gs._xyz = torch.cat([scene_gs._xyz, gs._xyz], dim=0)
-        scene_gs._features_dc = torch.cat([scene_gs._features_dc, gs._features_dc], dim=0)
-        scene_gs._scaling = torch.cat([scene_gs._scaling, gs._scaling], dim=0)
-        scene_gs._rotation = torch.cat([scene_gs._rotation, gs._rotation], dim=0)
-        scene_gs._opacity = torch.cat([scene_gs._opacity, gs._opacity], dim=0)
-    return scene_gs
-
-def load_image(path):
-    return np.array(Image.open(path)).astype(np.uint8)
-
-def load_mask(path):
-    mask = load_image(path) > 0
-    return mask[..., -1] if mask.ndim == 3 else mask
-
 # ============ Gradio Interface ============
 CONFIG_FILE = os.path.join(REPO_DIR, "configs/inference.yaml")
 
@@ -312,6 +254,8 @@ def process_image(input_image, input_mask, seed):
             return ply_path, "✅ Inference complete!"
         return None, "⚠️ No 3D output generated"
     except Exception as e:
+        import traceback
+        traceback.print_exc()
         return None, f"❌ Error: {str(e)}"
 
 with gr.Blocks(title="SAM 3D Objects", theme=gr.themes.Soft()) as demo: