Update app.py

app.py

@@ -4,13 +4,11 @@ import sys
 import gradio as gr
 import numpy as np
 from PIL import Image
-import io
 import tempfile
 from pathlib import Path
 import subprocess
-import shutil
 
-# Add notebook directory to path for inference code
+# Repository configuration
 REPO_DIR = "sam-3d-objects"
 REPO_URL = "https://github.com/facebookresearch/sam-3d-objects"
 
@@ -42,63 +40,11 @@ def ensure_repository():
     
     return True
 
-# Ensure repository is available
-if not ensure_repository():
-    print("Warning: Could not clone repository. Running in limited mode.")
+# Ensure repository is available at startup
+ensure_repository()
 
-# Print current sys.path for debugging
-print("Current sys.path:")
-for p in sys.path[:5]:
-    print(f"  {p}")
-
-# Set CUDA_HOME environment variable before importing inference
-# The inference.py tries to do: os.environ["CUDA_HOME"] = os.environ["CONDA_PREFIX"]
-# which fails in non-conda environments
-if "CUDA_HOME" not in os.environ:
-    # Try common CUDA locations
-    cuda_paths = [
-        "/usr/local/cuda",
-        "/usr/local/cuda-11",
-        "/usr/local/cuda-12",
-        "/opt/cuda",
-    ]
-    for cuda_path in cuda_paths:
-        if os.path.exists(cuda_path):
-            os.environ["CUDA_HOME"] = cuda_path
-            print(f"Set CUDA_HOME to {cuda_path}")
-            break
-    else:
-        # Set a dummy path if no CUDA found - some features may not work
-        os.environ["CUDA_HOME"] = "/usr/local/cuda"
-        print("Warning: CUDA not found, set CUDA_HOME to /usr/local/cuda")
-
-# Also set CONDA_PREFIX if not set (in case other parts of the code need it)
-if "CONDA_PREFIX" not in os.environ:
-    os.environ["CONDA_PREFIX"] = os.environ.get("CUDA_HOME", "/usr/local/cuda")
-    print(f"Set CONDA_PREFIX to {os.environ['CONDA_PREFIX']}")
-
-# Import inference code with error handling
-try:
-    from inference import Inference, load_image, load_single_mask
-    INFERENCE_AVAILABLE = True
-    print("Inference module loaded successfully")
-except ImportError as e:
-    print(f"Warning: Could not import inference module: {e}")
-    print("Checking for inference.py location...")
-    
-    # Debug: look for inference.py
-    for root, dirs, files in os.walk(REPO_DIR):
-        if "inference.py" in files:
-            print(f"  Found inference.py at: {os.path.join(root, 'inference.py')}")
-    
-    print("Running in demo mode with mock functionality")
-    INFERENCE_AVAILABLE = False
-except Exception as e:
-    print(f"Warning: Error importing inference module: {e}")
-    import traceback
-    traceback.print_exc()
-    print("Running in demo mode with mock functionality")
-    INFERENCE_AVAILABLE = False
+# Global variable to cache the inference model
+_cached_inference = None
 
 def create_demo_3d_output():
     """Create a demo 3D file for demonstration purposes"""
@@ -117,7 +63,6 @@ def create_demo_3d_output():
         "property uchar blue",
         "end_header"
     ]
-    # Add some demo vertices
     for i in range(1000):
         x, y, z = np.random.normal(0, 1, 3)
         nx, ny, nz = np.random.normal(0, 1, 3)
@@ -126,33 +71,31 @@ def create_demo_3d_output():
     
     return "\n".join(lines).encode('utf-8')
 
-def load_and_validate_image(image_path):
-    """Load and validate image file"""
-    try:
-        img = Image.open(image_path)
-        img = img.convert('RGB')
-        return np.array(img)
-    except Exception as e:
-        raise ValueError(f"Error loading image: {str(e)}")
-
-@spaces.GPU()
+@spaces.GPU(duration=120)
 def process_image_to_3d(image, mask=None, seed=42, model_tag="hf"):
-    """Process image to 3D model"""
+    """Process image to 3D model - runs on GPU"""
+    global _cached_inference
+    
     try:
-        if not INFERENCE_AVAILABLE:
-            # Demo mode - return mock output
-            demo_content = create_demo_3d_output()
-            return {
-                "status": "demo",
-                "message": "Demo mode - inference module not available. This is a sample 3D model file.",
-                "file_content": demo_content,
-                "filename": "demo_splat.ply"
-            }
+        # Set environment variables for CUDA (needed by inference.py)
+        if "CUDA_HOME" not in os.environ:
+            cuda_paths = ["/usr/local/cuda", "/usr/local/cuda-11", "/usr/local/cuda-12"]
+            for cuda_path in cuda_paths:
+                if os.path.exists(cuda_path):
+                    os.environ["CUDA_HOME"] = cuda_path
+                    break
+            else:
+                os.environ["CUDA_HOME"] = "/usr/local/cuda"
+        
+        if "CONDA_PREFIX" not in os.environ:
+            os.environ["CONDA_PREFIX"] = os.environ.get("CUDA_HOME", "/usr/local/cuda")
         
-        # Check if model directory exists (inside the cloned repo)
+        # Import inside the GPU function where CUDA is available
+        from inference import Inference, load_image, load_single_mask
+        
+        # Find model checkpoint
         model_dir = os.path.join(REPO_DIR, "checkpoints", model_tag)
         if not os.path.exists(model_dir):
-            # Try alternative location at root level
             model_dir = os.path.join("checkpoints", model_tag)
         
         if not os.path.exists(model_dir):
@@ -160,7 +103,6 @@ def process_image_to_3d(image, mask=None, seed=42, model_tag="hf"):
                 "status": "error",
                 "message": f"Model checkpoint not found. Please ensure the model is downloaded to checkpoints/{model_tag}/",
                 "file_content": None,
-                "filename": None
             }
         
         config_path = os.path.join(model_dir, "pipeline.yaml")
@@ -170,17 +112,23 @@ def process_image_to_3d(image, mask=None, seed=42, model_tag="hf"):
         if not os.path.exists(config_path):
             return {
                 "status": "error",
-                "message": f"Pipeline configuration not found. Expected at {config_path}",
+                "message": f"Pipeline configuration not found at {config_path}",
                 "file_content": None,
-                "filename": None
             }
         
-        # Create temporary files for the uploaded image and mask
+        # Load or reuse cached model
+        if _cached_inference is None:
+            print("Loading inference model...")
+            _cached_inference = Inference(config_path, compile=False)
+            print("Model loaded successfully")
+        
+        # Save uploaded image to temp file
         with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as img_temp:
             img = Image.fromarray(image)
             img.save(img_temp.name)
             temp_image_path = img_temp.name
         
+        # Save mask if provided
         temp_mask_path = None
         if mask is not None:
             with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as mask_temp:
@@ -188,30 +136,28 @@ def process_image_to_3d(image, mask=None, seed=42, model_tag="hf"):
                 mask_img.save(mask_temp.name)
                 temp_mask_path = mask_temp.name
         
-        # Load the model using the API from docs
-        inference = Inference(config_path, compile=False)
-        
-        # Load image and mask using the helper functions
+        # Load image and mask using helper functions
         loaded_image = load_image(temp_image_path)
         loaded_mask = load_single_mask(temp_mask_path) if temp_mask_path else None
         
         # Run inference
-        output = inference(loaded_image, loaded_mask, seed=seed)
+        print(f"Running inference with seed={seed}...")
+        output = _cached_inference(loaded_image, loaded_mask, seed=seed)
         
-        # Export gaussian splat to temp file
-        output_path = tempfile.NamedTemporaryFile(suffix=".ply", delete=False).name
-        output["gs"].save_ply(output_path)
+        # Export gaussian splat
+        output_ply_path = tempfile.NamedTemporaryFile(suffix=".ply", delete=False).name
+        output["gs"].save_ply(output_ply_path)
         
         # Read the generated file
-        with open(output_path, "rb") as f:
+        with open(output_ply_path, "rb") as f:
             file_content = f.read()
         
-        # Clean up temporary files
+        # Cleanup temp files
         try:
             os.unlink(temp_image_path)
             if temp_mask_path:
                 os.unlink(temp_mask_path)
-            os.unlink(output_path)
+            os.unlink(output_ply_path)
         except:
             pass
         
@@ -219,7 +165,6 @@ def process_image_to_3d(image, mask=None, seed=42, model_tag="hf"):
             "status": "success",
             "message": "3D model generated successfully!",
             "file_content": file_content,
-            "filename": f"splat_{seed}.ply"
         }
         
     except Exception as e:
@@ -229,7 +174,6 @@ def process_image_to_3d(image, mask=None, seed=42, model_tag="hf"):
             "status": "error",
             "message": f"Error processing image: {str(e)}",
             "file_content": None,
-            "filename": None
         }
 
 def update_mask_status(mask_image):
@@ -244,19 +188,17 @@ def process_wrapper(image, mask, seed, model_tag):
     if image is None:
         return "Please upload an image first", None
     
-    # Show processing status
-    yield "Processing image to 3D model...", None
+    yield "Processing image to 3D model... (this may take 1-2 minutes)", None
     
     result = process_image_to_3d(image, mask, seed, model_tag)
     
-    if result["status"] == "success" or result["status"] == "demo":
+    if result["status"] == "success":
         # Write content to a temporary file for gr.File component
         if result["file_content"] is not None:
-            temp_path = tempfile.NamedTemporaryFile(suffix=".ply", delete=False)
-            temp_path.write(result["file_content"])
-            temp_path.close()
-            msg = result["message"] if result["status"] == "success" else "Demo: " + result["message"]
-            yield msg, temp_path.name
+            temp_file = tempfile.NamedTemporaryFile(suffix=".ply", delete=False)
+            temp_file.write(result["file_content"])
+            temp_file.close()
+            yield result["message"], temp_file.name
         else:
             yield result["message"], None
     else:
@@ -265,68 +207,26 @@ def process_wrapper(image, mask, seed, model_tag):
 def create_interface():
     """Create the Gradio interface"""
     
-    # Custom CSS for better styling
     css = """
-    .gradio-container {
-        max-width: 1200px !important;
-        margin: auto !important;
-    }
-    .upload-section {
-        border: 2px dashed #ccc;
-        padding: 20px;
-        border-radius: 10px;
-        background-color: #f9f9f9;
-    }
-    .status-message {
-        padding: 10px;
-        border-radius: 5px;
-        margin: 10px 0;
-    }
-    .success {
-        background-color: #d4edda;
-        color: #155724;
-        border: 1px solid #c3e6cb;
-    }
-    .error {
-        background-color: #f8d7da;
-        color: #721c24;
-        border: 1px solid #f5c6cb;
-    }
-    .upload-area {
-        border: 2px dashed #4CAF50 !important;
-        border-radius: 10px !important;
-        padding: 10px !important;
-    }
-    .mask-status {
-        background-color: #e3f2fd;
-        border: 1px solid #2196F3;
-        padding: 5px;
-        border-radius: 5px;
-        font-weight: bold;
-    }
+    .gradio-container { max-width: 1200px !important; margin: auto !important; }
+    .upload-section { border: 2px dashed #ccc; padding: 20px; border-radius: 10px; background-color: #f9f9f9; }
+    .status-message { padding: 10px; border-radius: 5px; margin: 10px 0; }
+    .upload-area { border: 2px dashed #4CAF50 !important; border-radius: 10px !important; padding: 10px !important; }
+    .mask-status { background-color: #e3f2fd; border: 1px solid #2196F3; padding: 5px; border-radius: 5px; font-weight: bold; }
     """
     
     with gr.Blocks(css=css, title="Image to 3D Converter") as demo:
         
-        # Header
         gr.HTML("""
         <div style="text-align: center; padding: 20px; background: linear-gradient(135deg, #667eea 0%, #764ba2 100%); color: white; border-radius: 10px; margin-bottom: 30px;">
             <h1 style="margin: 0; font-size: 2.5em;">🎨 Image to 3D Converter</h1>
-            <p style="margin: 10px 0 0 0; font-size: 1.2em;">Transform your 2D images into stunning 3D models</p>
-            <div style="margin-top: 15px; font-size: 0.9em;">
-                <a href="https://huggingface.co/spaces/akhaliq/anycoder" target="_blank" style="color: #fff; text-decoration: none; border: 1px solid rgba(255,255,255,0.5); padding: 5px 15px; border-radius: 20px;">Built with anycoder</a>
-            </div>
+            <p style="margin: 10px 0 0 0; font-size: 1.2em;">Transform your 2D images into stunning 3D models using SAM 3D Objects</p>
         </div>
         """)
         
         with gr.Row():
             with gr.Column(scale=1):
-                gr.HTML("""
-                <div class="upload-section">
-                    <h3>📤 Upload Image</h3>
-                    <p>Upload the image you want to convert to 3D</p>
-                </div>
-                """)
+                gr.HTML("""<div class="upload-section"><h3>📤 Upload Image</h3><p>Upload the image you want to convert to 3D</p></div>""")
                 
                 image_input = gr.Image(
                     label="Input Image",
@@ -335,20 +235,14 @@ def create_interface():
                     elem_classes=["upload-area"],
                 )
                 
-                gr.HTML("""
-                <div class="upload-section" style="margin-top: 15px;">
-                    <h3>🎭 Optional Mask</h3>
-                    <p>Upload a mask to focus on specific areas (optional: binary mask to segment the object)</p>
-                </div>
-                """)
+                gr.HTML("""<div class="upload-section" style="margin-top: 15px;"><h3>🎭 Optional Mask</h3><p>Upload a binary mask to focus on specific objects</p></div>""")
                 
-                with gr.Row():
-                    mask_upload = gr.Image(
-                        label="Segmentation Mask (Optional)",
-                        type="numpy", 
-                        image_mode="L",
-                        elem_classes=["upload-area"],
-                    )
+                mask_upload = gr.Image(
+                    label="Segmentation Mask (Optional)",
+                    type="numpy", 
+                    image_mode="L",
+                    elem_classes=["upload-area"],
+                )
                 
                 mask_status = gr.Textbox(
                     label="Mask Status",
@@ -358,21 +252,15 @@ def create_interface():
                 )
             
             with gr.Column(scale=1):
-                gr.HTML("""
-                <div style="background: #f0f8ff; padding: 20px; border-radius: 10px; margin-bottom: 20px; border: 1px solid #2196F3;">
-                    <h3>⚙️ Configuration</h3>
-                    <p>Fine-tune the 3D generation parameters</p>
-                </div>
-                """)
+                gr.HTML("""<div style="background: #f0f8ff; padding: 20px; border-radius: 10px; margin-bottom: 20px; border: 1px solid #2196F3;"><h3>⚙️ Configuration</h3><p>Fine-tune the 3D generation parameters</p></div>""")
                 
-                with gr.Row():
-                    seed = gr.Slider(
-                        minimum=0,
-                        maximum=999999,
-                        value=42,
-                        step=1,
-                        label="Random Seed",
-                    )
+                seed = gr.Slider(
+                    minimum=0,
+                    maximum=999999,
+                    value=42,
+                    step=1,
+                    label="Random Seed",
+                )
                 
                 model_tag = gr.Dropdown(
                     choices=["hf"],
@@ -380,42 +268,26 @@ def create_interface():
                     label="Model Configuration",
                 )
                 
-                gr.HTML("""
-                <div style="margin-top: 20px; text-align: center;">
-                    <p><strong>Generation Status:</strong></p>
-                    <p style="color: #666; font-size: 0.9em;">This process may take several minutes depending on image complexity</p>
-                </div>
-                """)
+                gr.HTML("""<div style="margin-top: 20px; text-align: center;"><p style="color: #666; font-size: 0.9em;">⏱️ Processing typically takes 1-2 minutes on ZeroGPU</p></div>""")
                 
-                run_button = gr.Button(
-                    "🚀 Generate 3D Model",
-                    variant="primary",
-                    size="lg"
-                )
+                run_button = gr.Button("🚀 Generate 3D Model", variant="primary", size="lg")
         
         with gr.Row():
-            with gr.Column():
-                status_output = gr.Textbox(
-                    label="Generation Status",
-                    max_lines=5,
-                    interactive=False,
-                    elem_classes=["status-message"],
-                )
+            status_output = gr.Textbox(
+                label="Generation Status",
+                max_lines=5,
+                interactive=False,
+                elem_classes=["status-message"],
+            )
         
         with gr.Row():
-            with gr.Column():
-                output_file = gr.File(
-                    label="📥 Download 3D Model",
-                    file_types=[".ply"],
-                    elem_classes=["download-section"],
-                )
+            output_file = gr.File(
+                label="📥 Download 3D Model (.ply)",
+                file_types=[".ply"],
+            )
         
-        # Wire up the interface
-        mask_upload.upload(
-            fn=update_mask_status,
-            inputs=[mask_upload],
-            outputs=[mask_status]
-        )
+        # Event handlers
+        mask_upload.upload(fn=update_mask_status, inputs=[mask_upload], outputs=[mask_status])
         
         run_button.click(
             fn=process_wrapper,
@@ -423,71 +295,38 @@ def create_interface():
             outputs=[status_output, output_file]
         )
         
-        # Examples and instructions section
         gr.HTML("""
         <div style="margin-top: 40px; text-align: center; background: #f8f9fa; padding: 30px; border-radius: 15px;">
             <h3>📖 How to Use</h3>
             <div style="display: flex; justify-content: space-around; margin-top: 20px; flex-wrap: wrap; gap: 20px;">
-                <div style="max-width: 280px; padding: 20px; background: white; border-radius: 10px; margin: 10px; box-shadow: 0 2px 10px rgba(0,0,0,0.1);">
+                <div style="max-width: 280px; padding: 20px; background: white; border-radius: 10px; box-shadow: 0 2px 10px rgba(0,0,0,0.1);">
                     <h4>1. Upload Image</h4>
-                    <p>Choose a clear, well-lit image for best results. The image should show the object you want to convert to 3D.</p>
+                    <p>Choose a clear, well-lit image showing the object you want to convert to 3D.</p>
                 </div>
-                <div style="max-width: 280px; padding: 20px; background: white; border-radius: 10px; margin: 10px; box-shadow: 0 2px 10px rgba(0,0,0,0.1);">
+                <div style="max-width: 280px; padding: 20px; background: white; border-radius: 10px; box-shadow: 0 2px 10px rgba(0,0,0,0.1);">
                     <h4>2. Add Mask (Optional)</h4>
-                    <p>Upload a mask to focus on specific areas. This helps the model understand what part of the image to convert.</p>
+                    <p>Upload a binary mask to isolate specific objects in the scene.</p>
                 </div>
-                <div style="max-width: 280px; padding: 20px; background: white; border-radius: 10px; margin: 10px; box-shadow: 0 2px 10px rgba(0,0,0,0.1);">
+                <div style="max-width: 280px; padding: 20px; background: white; border-radius: 10px; box-shadow: 0 2px 10px rgba(0,0,0,0.1);">
                     <h4>3. Generate</h4>
-                    <p>Click generate and wait for your 3D model. The process typically takes 1-5 minutes.</p>
+                    <p>Click generate and download your 3D model in PLY format.</p>
                 </div>
             </div>
-            
-            <div style="margin-top: 30px; padding: 20px; background: white; border-radius: 10px; text-align: left; max-width: 800px; margin-left: auto; margin-right: auto;">
-                <h4>💡 Tips for Better Results:</h4>
-                <ul style="text-align: left;">
-                    <li>Use high-quality, well-lit images</li>
-                    <li>Ensure the object is clearly visible and not occluded</li>
-                    <li>Use masks to isolate specific objects in complex scenes</li>
-                    <li>Try different random seeds for variations</li>
-                    <li>Complex objects may take longer to process</li>
-                </ul>
-            </div>
         </div>
         """)
         
-        # System information
         gr.HTML(f"""
         <div style="margin-top: 30px; padding: 15px; background: #e8f5e8; border-radius: 10px; text-align: center;">
-            <p><strong>System Status:</strong></p>
-            <p style="font-size: 0.9em; color: #666;">
-                Inference Module: {"✓ Available" if INFERENCE_AVAILABLE else "✗ Demo Mode"} | 
-                SAM 3D Repository: {"✓ Cloned" if os.path.exists(REPO_DIR) else "✗ Not Available"}
-            </p>
+            <p><strong>System Status:</strong> ZeroGPU (NVIDIA H200) | SAM 3D Repository: {"✓ Cloned" if os.path.exists(REPO_DIR) else "✗ Not Available"}</p>
         </div>
         """)
     
     return demo
 
 if __name__ == "__main__":
-    # Create and launch the interface
     demo = create_interface()
-    
-    # Print available model paths for debugging
-    print("Checking for model checkpoints...")
-    if os.path.exists("checkpoints"):
-        for root, dirs, files in os.walk("checkpoints"):
-            print(f"Found in {root}: {files}")
-    else:
-        print("No checkpoints directory found")
-    
-    print("Available inference modules:", "✓" if INFERENCE_AVAILABLE else "✗")
-    
-    # Launch with proper configuration
     demo.launch(
         server_name="0.0.0.0",
         server_port=7860,
-        share=False,
         show_error=True,
-        debug=True,
-        inbrowser=True
     )