mps-carafe 0.2.1__tar.gz

mps_carafe-0.2.1/PKG-INFO

@@ -0,0 +1,153 @@
+Metadata-Version: 2.4
+Name: mps-carafe
+Version: 0.2.1
+Summary: CARAFE content-aware upsampling for Apple Silicon (MPS)
+Author: mpsops
+License-Expression: MIT
+Project-URL: Homepage, https://github.com/mpsops/mps-carafe
+Project-URL: Repository, https://github.com/mpsops/mps-carafe
+Project-URL: Issues, https://github.com/mpsops/mps-carafe/issues
+Keywords: pytorch,mps,apple-silicon,carafe,upsampling,segmentation
+Classifier: Development Status :: 4 - Beta
+Classifier: Intended Audience :: Developers
+Classifier: Intended Audience :: Science/Research
+Classifier: Operating System :: MacOS :: MacOS X
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
+Requires-Python: >=3.10
+Description-Content-Type: text/markdown
+Requires-Dist: torch>=2.0.0
+Dynamic: requires-python
+
+# MPS CARAFE
+
+CARAFE (Content-Aware ReAssembly of FEatures) for Apple Silicon (M1/M2/M3/M4).
+
+**Drop-in replacement** for mmcv's CARAFE op.
+
+## Why?
+
+CARAFE is a learnable upsampling operator used in:
+- **Mask R-CNN**: Instance segmentation
+- **FPN**: Feature Pyramid Networks
+- **YOLACT**: Real-time instance segmentation
+
+But mmcv's implementation is **CUDA-only**. On Mac you get:
+```
+NotImplementedError: carafe not implemented for MPS
+```
+
+This package provides a native Metal implementation.
+
+## Installation
+
+```bash
+pip install mps-carafe
+```
+
+Or from source:
+
+```bash
+git clone https://github.com/mpsops/mps-carafe
+cd mps-carafe
+pip install -e .
+```
+
+## Quick Start
+
+### Basic CARAFE Operation
+
+```python
+import torch
+from mps_carafe import carafe
+
+# Input features (N, C, H, W)
+features = torch.randn(1, 64, 32, 32, device='mps')
+
+# Reassembly masks (N, group_size * k^2, H*scale, W*scale)
+kernel_size = 5
+group_size = 1
+scale_factor = 2
+masks = torch.softmax(
+    torch.randn(1, group_size * kernel_size**2, 64, 64, device='mps'),
+    dim=1
+)
+
+# Upsample with CARAFE
+output = carafe(features, masks, kernel_size, group_size, scale_factor)
+# Output: (1, 64, 64, 64)
+```
+
+### CARAFE Module
+
+```python
+from mps_carafe import CARAFE
+
+carafe_layer = CARAFE(kernel_size=5, group_size=1, scale_factor=2)
+output = carafe_layer(features, masks)
+```
+
+### CARAFEPack (with mask predictor)
+
+```python
+from mps_carafe import CARAFEPack
+
+# Complete upsampling block with built-in mask predictor
+upsample = CARAFEPack(
+    channels=64,
+    kernel_size=5,
+    group_size=1,
+    scale_factor=2
+).to('mps')
+
+output = upsample(features)  # No need to provide masks
+```
+
+## API Reference
+
+### `carafe(features, masks, kernel_size, group_size, scale_factor)`
+
+| Parameter | Type | Description |
+|-----------|------|-------------|
+| `features` | Tensor | Input features (N, C, H, W) |
+| `masks` | Tensor | Reassembly kernels (N, group_size * k^2, H*scale, W*scale) |
+| `kernel_size` | int | Size of reassembly kernel (typically 5) |
+| `group_size` | int | Number of channel groups (typically 1) |
+| `scale_factor` | int | Upsampling factor (typically 2) |
+
+### `CARAFEPack`
+
+Complete CARAFE block with mask prediction convolutions.
+
+## How It Works
+
+CARAFE upsamples by:
+1. For each output pixel, identify the corresponding input neighborhood
+2. Use learned reassembly kernels (masks) to weight input pixels
+3. Sum the weighted inputs to produce the output
+
+Unlike bilinear upsampling which uses fixed weights, CARAFE learns content-aware weights that adapt to the image content.
+
+## Compatibility
+
+- **PyTorch**: 2.0+
+- **macOS**: 12.0+ (Monterey)
+- **Hardware**: Apple Silicon (M1/M2/M3/M4)
+
+## Features
+
+- Full forward and backward pass (training supported)
+- fp32 and fp16 supported
+- Compatible with mmcv CARAFE API
+
+## Credits
+
+- [CARAFE Paper](https://arxiv.org/abs/1905.02188) - Original research
+- [mmcv](https://github.com/open-mmlab/mmcv) - Reference implementation
+
+## License
+
+MIT

mps_carafe-0.2.1/README.md

@@ -0,0 +1,129 @@
+# MPS CARAFE
+
+CARAFE (Content-Aware ReAssembly of FEatures) for Apple Silicon (M1/M2/M3/M4).
+
+**Drop-in replacement** for mmcv's CARAFE op.
+
+## Why?
+
+CARAFE is a learnable upsampling operator used in:
+- **Mask R-CNN**: Instance segmentation
+- **FPN**: Feature Pyramid Networks
+- **YOLACT**: Real-time instance segmentation
+
+But mmcv's implementation is **CUDA-only**. On Mac you get:
+```
+NotImplementedError: carafe not implemented for MPS
+```
+
+This package provides a native Metal implementation.
+
+## Installation
+
+```bash
+pip install mps-carafe
+```
+
+Or from source:
+
+```bash
+git clone https://github.com/mpsops/mps-carafe
+cd mps-carafe
+pip install -e .
+```
+
+## Quick Start
+
+### Basic CARAFE Operation
+
+```python
+import torch
+from mps_carafe import carafe
+
+# Input features (N, C, H, W)
+features = torch.randn(1, 64, 32, 32, device='mps')
+
+# Reassembly masks (N, group_size * k^2, H*scale, W*scale)
+kernel_size = 5
+group_size = 1
+scale_factor = 2
+masks = torch.softmax(
+    torch.randn(1, group_size * kernel_size**2, 64, 64, device='mps'),
+    dim=1
+)
+
+# Upsample with CARAFE
+output = carafe(features, masks, kernel_size, group_size, scale_factor)
+# Output: (1, 64, 64, 64)
+```
+
+### CARAFE Module
+
+```python
+from mps_carafe import CARAFE
+
+carafe_layer = CARAFE(kernel_size=5, group_size=1, scale_factor=2)
+output = carafe_layer(features, masks)
+```
+
+### CARAFEPack (with mask predictor)
+
+```python
+from mps_carafe import CARAFEPack
+
+# Complete upsampling block with built-in mask predictor
+upsample = CARAFEPack(
+    channels=64,
+    kernel_size=5,
+    group_size=1,
+    scale_factor=2
+).to('mps')
+
+output = upsample(features)  # No need to provide masks
+```
+
+## API Reference
+
+### `carafe(features, masks, kernel_size, group_size, scale_factor)`
+
+| Parameter | Type | Description |
+|-----------|------|-------------|
+| `features` | Tensor | Input features (N, C, H, W) |
+| `masks` | Tensor | Reassembly kernels (N, group_size * k^2, H*scale, W*scale) |
+| `kernel_size` | int | Size of reassembly kernel (typically 5) |
+| `group_size` | int | Number of channel groups (typically 1) |
+| `scale_factor` | int | Upsampling factor (typically 2) |
+
+### `CARAFEPack`
+
+Complete CARAFE block with mask prediction convolutions.
+
+## How It Works
+
+CARAFE upsamples by:
+1. For each output pixel, identify the corresponding input neighborhood
+2. Use learned reassembly kernels (masks) to weight input pixels
+3. Sum the weighted inputs to produce the output
+
+Unlike bilinear upsampling which uses fixed weights, CARAFE learns content-aware weights that adapt to the image content.
+
+## Compatibility
+
+- **PyTorch**: 2.0+
+- **macOS**: 12.0+ (Monterey)
+- **Hardware**: Apple Silicon (M1/M2/M3/M4)
+
+## Features
+
+- Full forward and backward pass (training supported)
+- fp32 and fp16 supported
+- Compatible with mmcv CARAFE API
+
+## Credits
+
+- [CARAFE Paper](https://arxiv.org/abs/1905.02188) - Original research
+- [mmcv](https://github.com/open-mmlab/mmcv) - Reference implementation
+
+## License
+
+MIT

mps_carafe-0.2.1/mps_carafe/__init__.py

@@ -0,0 +1,262 @@
+"""
+MPS CARAFE - Content-Aware ReAssembly of FEatures for Apple Silicon
+
+Drop-in replacement for mmcv's CARAFE op.
+Used in Mask R-CNN, FPN, and other detection/segmentation networks.
+"""
+
+import torch
+from torch import nn, Tensor
+from torch.autograd import Function
+from typing import Tuple
+import threading
+
+__version__ = "0.2.0"
+
+# Thread-safe library loading
+_lib_lock = threading.Lock()
+
+
+def _load_library():
+    """Load the native extension."""
+    try:
+        from mps_carafe import _C as _lib
+        return _lib
+    except ImportError:
+        import os
+        from torch.utils.cpp_extension import load
+
+        src_dir = os.path.join(os.path.dirname(__file__), "csrc")
+        _lib = load(
+            name="mps_carafe",
+            sources=[os.path.join(src_dir, "carafe_mps.mm")],
+            extra_cflags=["-std=c++17"],
+            extra_ldflags=["-framework", "Metal", "-framework", "Foundation"],
+            verbose=False,
+        )
+        return _lib
+
+
+_lib_cache = None
+
+
+def _get_lib():
+    """Thread-safe library loading."""
+    global _lib_cache
+    if _lib_cache is None:
+        with _lib_lock:
+            # Double-check after acquiring lock
+            if _lib_cache is None:
+                _lib_cache = _load_library()
+    return _lib_cache
+
+
+def _validate_params(kernel_size: int, group_size: int, scale_factor: int, channels: int) -> None:
+    """Validate CARAFE parameters."""
+    if kernel_size <= 0:
+        raise ValueError(f"kernel_size must be positive, got {kernel_size}")
+    if group_size <= 0:
+        raise ValueError(f"group_size must be positive, got {group_size}")
+    if scale_factor <= 0:
+        raise ValueError(f"scale_factor must be positive, got {scale_factor}")
+    if channels % group_size != 0:
+        raise ValueError(f"channels ({channels}) must be divisible by group_size ({group_size})")
+
+
+def _check_tensor_device(tensor: Tensor, name: str, expected_device: torch.device) -> None:
+    """Check that tensor is on the expected device."""
+    if tensor.device != expected_device:
+        raise ValueError(
+            f"{name} must be on same device as features ({expected_device}), got {tensor.device}"
+        )
+
+
+class _CARAFEFunction(Function):
+    """Autograd function for CARAFE."""
+
+    @staticmethod
+    def forward(
+        ctx,
+        features: Tensor,
+        masks: Tensor,
+        kernel_size: int,
+        group_size: int,
+        scale_factor: int,
+    ) -> Tensor:
+        # Device validation
+        _check_tensor_device(masks, "masks", features.device)
+
+        ctx.save_for_backward(features, masks)
+        ctx.kernel_size = kernel_size
+        ctx.group_size = group_size
+        ctx.scale_factor = scale_factor
+
+        lib = _get_lib()
+        output = lib.carafe_forward(
+            features, masks,
+            kernel_size, group_size, scale_factor
+        )
+        return output
+
+    @staticmethod
+    def backward(ctx, grad_output: Tensor):
+        features, masks = ctx.saved_tensors
+
+        # Device validation for backward
+        if not grad_output.device.type == "mps":
+            raise ValueError(f"grad_output must be on MPS device, got {grad_output.device}")
+
+        lib = _get_lib()
+
+        grad_features, grad_masks = lib.carafe_backward(
+            grad_output.contiguous(),
+            features.contiguous(),
+            masks.contiguous(),
+            ctx.kernel_size,
+            ctx.group_size,
+            ctx.scale_factor
+        )
+
+        return grad_features, grad_masks, None, None, None
+
+
+def carafe(
+    features: Tensor,
+    masks: Tensor,
+    kernel_size: int,
+    group_size: int,
+    scale_factor: int,
+) -> Tensor:
+    """
+    CARAFE: Content-Aware ReAssembly of FEatures.
+
+    Performs content-aware upsampling using learned reassembly kernels.
+
+    Args:
+        features: Input feature map (N, C, H, W)
+        masks: Reassembly kernels (N, group_size * kernel_size^2, H*scale, W*scale)
+        kernel_size: Size of the reassembly kernel (typically 5, must be positive)
+        group_size: Number of groups for channel-wise reassembly (must be positive)
+        scale_factor: Upsampling factor (typically 2, must be positive)
+
+    Returns:
+        Upsampled feature map (N, C, H*scale_factor, W*scale_factor)
+
+    Raises:
+        ValueError: If tensors not on MPS or parameters invalid.
+    """
+    if features.device.type != "mps":
+        raise ValueError(f"Input must be on MPS device, got {features.device}")
+
+    channels = features.size(1)
+    _validate_params(kernel_size, group_size, scale_factor, channels)
+
+    return _CARAFEFunction.apply(
+        features, masks,
+        kernel_size, group_size, scale_factor
+    )
+
+
+class CARAFE(nn.Module):
+    """
+    CARAFE: Content-Aware ReAssembly of FEatures.
+
+    A learnable upsampling operator that uses content-aware kernels
+    for feature reassembly. Better than bilinear/nearest for dense prediction.
+
+    Args:
+        kernel_size: Size of reassembly kernel (default: 5)
+        group_size: Number of channel groups (default: 1)
+        scale_factor: Upsampling factor (default: 2)
+    """
+
+    def __init__(
+        self,
+        kernel_size: int = 5,
+        group_size: int = 1,
+        scale_factor: int = 2,
+    ):
+        super().__init__()
+        self.kernel_size = kernel_size
+        self.group_size = group_size
+        self.scale_factor = scale_factor
+
+    def forward(self, features: Tensor, masks: Tensor) -> Tensor:
+        return carafe(
+            features, masks,
+            self.kernel_size, self.group_size, self.scale_factor
+        )
+
+
+class CARAFEPack(nn.Module):
+    """
+    CARAFE with built-in mask predictor.
+
+    This module includes the convolutions to predict reassembly masks
+    from input features, making it a complete upsampling block.
+
+    Args:
+        channels: Number of input/output channels
+        kernel_size: Size of reassembly kernel (default: 5)
+        group_size: Number of channel groups (default: 1)
+        scale_factor: Upsampling factor (default: 2)
+        compressed_channels: Channels after compression (default: 64)
+    """
+
+    def __init__(
+        self,
+        channels: int,
+        kernel_size: int = 5,
+        group_size: int = 1,
+        scale_factor: int = 2,
+        compressed_channels: int = 64,
+    ):
+        super().__init__()
+        self.channels = channels
+        self.kernel_size = kernel_size
+        self.group_size = group_size
+        self.scale_factor = scale_factor
+
+        # Channel compressor
+        self.channel_compressor = nn.Conv2d(
+            channels, compressed_channels, kernel_size=1
+        )
+
+        # Content encoder - predicts reassembly masks
+        self.content_encoder = nn.Conv2d(
+            compressed_channels,
+            group_size * kernel_size * kernel_size * scale_factor * scale_factor,
+            kernel_size=3,
+            padding=1,
+        )
+
+        # CARAFE operator
+        self.carafe = CARAFE(kernel_size, group_size, scale_factor)
+
+    def forward(self, x: Tensor) -> Tensor:
+        # Compress channels
+        compressed = self.channel_compressor(x)
+
+        # Predict masks
+        masks = self.content_encoder(compressed)
+
+        # Upsample masks to output resolution
+        n, _, h, w = masks.shape
+        masks = masks.view(n, self.group_size, self.kernel_size * self.kernel_size,
+                          self.scale_factor, self.scale_factor, h, w)
+        masks = masks.permute(0, 1, 2, 5, 3, 6, 4).contiguous()
+        masks = masks.view(n, self.group_size * self.kernel_size * self.kernel_size,
+                          h * self.scale_factor, w * self.scale_factor)
+
+        # Softmax over kernel positions
+        masks = torch.softmax(masks.view(n, self.group_size, -1,
+                                         h * self.scale_factor, w * self.scale_factor), dim=2)
+        masks = masks.view(n, -1, h * self.scale_factor, w * self.scale_factor)
+
+        # Apply CARAFE
+        return self.carafe(x, masks)
+
+
+def is_available() -> bool:
+    """Check if MPS is available."""
+    return torch.backends.mps.is_available()