RobinsAIWorld/ml-sharp

1"""SHARP MCP Server for programmatic access to 3D Gaussian prediction.2 3Run standalone:4    uv run python mcp_server.py5 6Or integrate with MCP clients via stdio transport.7"""8 9from __future__ import annotations10 11import json12import os13from pathlib import Path14from typing import Literal15 16import torch17from mcp.server.fastmcp import FastMCP18 19from model_utils import (20    DEFAULT_OUTPUTS_DIR,21    ModelWrapper,22    TrajectoryType,23    get_global_model,24)25 26MCP_PORT: int = int(os.getenv("SHARP_MCP_PORT", "49201"))27 28mcp = FastMCP(29    "sharp",30    description="SHARP: Single-image 3D Gaussian scene prediction",31)32 33# -----------------------------------------------------------------------------34# Tools35# -----------------------------------------------------------------------------36 37 38@mcp.tool()39def sharp_predict(40    image_path: str,41    render_video: bool = True,42    trajectory_type: TrajectoryType = "rotate_forward",43    num_frames: int = 60,44    fps: int = 30,45    output_long_side: int | None = None,46) -> dict:47    """Predict 3D Gaussians from a single image.48 49    Args:50        image_path: Absolute path to input image (jpg/png/webp).51        render_video: Whether to render a camera trajectory video (requires CUDA).52        trajectory_type: Camera trajectory type (swipe/shake/rotate/rotate_forward).53        num_frames: Number of frames for video rendering.54        fps: Frames per second for video.55        output_long_side: Output resolution (longest side). None = match input.56 57    Returns:58        dict with keys:59            - ply_path: Path to exported PLY file60            - video_path: Path to rendered MP4 (or null if not rendered)61            - cuda_available: Whether CUDA was available62    """63    image_path_obj = Path(image_path)64    if not image_path_obj.exists():65        raise FileNotFoundError(f"Image not found: {image_path}")66 67    model = get_global_model()68    video_path, ply_path = model.predict_and_maybe_render(69        image_path_obj,70        trajectory_type=trajectory_type,71        num_frames=num_frames,72        fps=fps,73        output_long_side=output_long_side,74        render_video=render_video,75    )76 77    return {78        "ply_path": str(ply_path),79        "video_path": str(video_path) if video_path else None,80        "cuda_available": torch.cuda.is_available(),81    }82 83 84@mcp.tool()85def sharp_render(86    ply_path: str,87    trajectory_type: TrajectoryType = "rotate_forward",88    num_frames: int = 60,89    fps: int = 30,90    output_long_side: int | None = None,91) -> dict:92    """Render a video from an existing PLY file.93 94    Note: This requires re-predicting from the original image since Gaussians95    are not stored in standard PLY format. For now, returns an error.96    Future versions may support loading Gaussians from PLY.97 98    Args:99        ply_path: Path to PLY file (from previous prediction).100        trajectory_type: Camera trajectory type.101        num_frames: Number of frames.102        fps: Frames per second.103        output_long_side: Output resolution.104 105    Returns:106        dict with error message (feature not yet implemented).107    """108    return {109        "error": "Rendering from PLY not yet implemented. Use sharp_predict with render_video=True.",110        "hint": "PLY files store only point data, not the full Gaussian parameters needed for rendering.",111    }112 113 114@mcp.tool()115def list_outputs() -> dict:116    """List all generated output files (PLY and MP4).117 118    Returns:119        dict with keys:120            - outputs_dir: Path to outputs directory121            - ply_files: List of PLY file paths122            - video_files: List of MP4 file paths123    """124    outputs_dir = DEFAULT_OUTPUTS_DIR125    ply_files = sorted(outputs_dir.glob("*.ply"))126    video_files = sorted(outputs_dir.glob("*.mp4"))127 128    return {129        "outputs_dir": str(outputs_dir),130        "ply_files": [str(f) for f in ply_files],131        "video_files": [str(f) for f in video_files],132    }133 134 135# -----------------------------------------------------------------------------136# Resources137# -----------------------------------------------------------------------------138 139 140@mcp.resource("sharp://info")141def get_info() -> str:142    """Get SHARP server info including GPU status and configuration."""143    cuda_available = torch.cuda.is_available()144    gpu_info = []145 146    if cuda_available:147        for i in range(torch.cuda.device_count()):148            props = torch.cuda.get_device_properties(i)149            gpu_info.append({150                "index": i,151                "name": props.name,152                "total_memory_gb": round(props.total_memory / (1024**3), 2),153                "compute_capability": f"{props.major}.{props.minor}",154            })155 156    info = {157        "model": "SHARP (Apple ml-sharp)",158        "description": "Single-image 3D Gaussian scene prediction",159        "cuda_available": cuda_available,160        "cuda_device_count": torch.cuda.device_count() if cuda_available else 0,161        "gpus": gpu_info,162        "outputs_dir": str(DEFAULT_OUTPUTS_DIR),163        "checkpoint_sources": [164            "SHARP_CHECKPOINT_PATH env var",165            "HuggingFace Hub (apple/Sharp)",166            "Upstream CDN (torch.hub)",167        ],168        "env_vars": {169            "SHARP_CHECKPOINT_PATH": os.getenv("SHARP_CHECKPOINT_PATH", "(not set)"),170            "SHARP_KEEP_MODEL_ON_DEVICE": os.getenv("SHARP_KEEP_MODEL_ON_DEVICE", "1"),171            "CUDA_VISIBLE_DEVICES": os.getenv("CUDA_VISIBLE_DEVICES", "(not set)"),172        },173    }174 175    return json.dumps(info, indent=2)176 177 178@mcp.resource("sharp://help")179def get_help() -> str:180    """Get usage help for the SHARP MCP server."""181    help_text = """182# SHARP MCP Server183 184## Tools185 186### sharp_predict187Predict 3D Gaussians from a single image.188 189Parameters:190- image_path (required): Absolute path to input image191- render_video: Whether to render MP4 (default: true, requires CUDA)192- trajectory_type: swipe | shake | rotate | rotate_forward (default: rotate_forward)193- num_frames: Number of video frames (default: 60)194- fps: Video frame rate (default: 30)195- output_long_side: Output resolution, null = match input196 197### list_outputs198List all generated PLY and MP4 files.199 200## Resources201 202### sharp://info203Server info, GPU status, configuration.204 205### sharp://help206This help text.207 208## Environment Variables209 210- SHARP_MCP_PORT: MCP server port (default: 49201)211- SHARP_CHECKPOINT_PATH: Local checkpoint path override212- SHARP_KEEP_MODEL_ON_DEVICE: Keep model on GPU (default: 1)213- CUDA_VISIBLE_DEVICES: GPU selection (e.g., "0" or "0,1")214"""215    return help_text.strip()216 217 218# -----------------------------------------------------------------------------219# Main220# -----------------------------------------------------------------------------221 222if __name__ == "__main__":223    # Run as stdio transport for MCP clients224    mcp.run()225