metalgrow 0.1.0__py3-none-any.whl

metalgrow/__init__.py

@@ -0,0 +1,5 @@
+from metalgrow.device import get_device
+from metalgrow.upscaler import Upscaler
+
+__all__ = ["Upscaler", "get_device"]
+__version__ = "0.0.1"

metalgrow/backbones/__init__.py

@@ -0,0 +1,8 @@
+# Register built-in backbones by importing their modules for side effects.
+from metalgrow.backbones import bicubic as _bicubic  # noqa: F401
+from metalgrow.backbones import realesrgan as _realesrgan  # noqa: F401
+from metalgrow.backbones import swinir as _swinir  # noqa: F401
+from metalgrow.backbones.base import Backbone
+from metalgrow.backbones.registry import get_backbone, list_backbones, register
+
+__all__ = ["Backbone", "get_backbone", "list_backbones", "register"]

metalgrow/backbones/base.py

@@ -0,0 +1,47 @@
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from typing import ClassVar
+
+import torch
+
+
+class Backbone(ABC):
+    """Pluggable super-resolution backbone.
+
+    Implementations accept a float tensor in [0, 1] of shape [N, C, H, W] on
+    ``self.device`` and return an upscaled tensor of the same dtype/device.
+    Preprocessing, model inference, and postprocessing all live inside
+    :meth:`upscale` so that the shared :class:`~metalgrow.upscaler.Upscaler`
+    stays agnostic to backbone internals.
+    """
+
+    name: ClassVar[str]
+    # None means any scale > 1.0 is accepted (e.g. analytical resamplers).
+    supported_scales: ClassVar[tuple[float, ...] | None] = None
+    # None means any channel count is supported. 3 means RGB-only — the
+    # Upscaler will handle alpha separately for such backbones.
+    input_channels: ClassVar[int | None] = None
+    # Default tile size for the shared Upscaler-level tiler. None disables
+    # tiling by default (analytical backbones have no VRAM ceiling); learned
+    # backbones should set a size that keeps a single tile's activations in
+    # VRAM on the target device.
+    default_tile: ClassVar[int | None] = None
+    # Pixels of input context padded around each tile. Needs to cover the
+    # backbone's receptive field so the blended seam is free of artefacts.
+    default_tile_pad: ClassVar[int] = 16
+
+    def __init__(self, device: torch.device, dtype: torch.dtype = torch.float32):
+        self.device = device
+        self.dtype = dtype
+
+    def validate_scale(self, scale: float) -> None:
+        if scale <= 1.0:
+            raise ValueError("scale must be > 1.0")
+        if self.supported_scales is not None and scale not in self.supported_scales:
+            allowed = ", ".join(f"{s:g}" for s in self.supported_scales)
+            raise ValueError(f"backbone {self.name!r} supports scales [{allowed}], got {scale:g}")
+
+    @abstractmethod
+    def upscale(self, tensor: torch.Tensor, scale: float) -> torch.Tensor:
+        """Return an upscaled copy of ``tensor``."""

metalgrow/backbones/bicubic.py

@@ -0,0 +1,27 @@
+from __future__ import annotations
+
+import torch
+import torch.nn.functional as F
+
+from metalgrow.backbones.base import Backbone
+from metalgrow.backbones.registry import register
+
+
+class BicubicBackbone(Backbone):
+    """Analytical bicubic resampler. Always available; no weights required."""
+
+    name = "bicubic"
+    supported_scales = None  # any scale > 1.0
+
+    def upscale(self, tensor: torch.Tensor, scale: float) -> torch.Tensor:
+        self.validate_scale(scale)
+        _, _, h, w = tensor.shape
+        return F.interpolate(
+            tensor,
+            size=(int(h * scale), int(w * scale)),
+            mode="bicubic",
+            align_corners=False,
+        ).clamp(0.0, 1.0)
+
+
+register("bicubic", lambda device, dtype: BicubicBackbone(device, dtype))

metalgrow/backbones/realesrgan.py

@@ -0,0 +1,54 @@
+from __future__ import annotations
+
+import torch
+
+from metalgrow.backbones.base import Backbone
+from metalgrow.backbones.registry import register
+from metalgrow.weights import ensure_weight
+
+
+class RealESRGANBackbone(Backbone):
+    """Real-ESRGAN RRDBNet backbone loaded via :mod:`spandrel`."""
+
+    name = "realesrgan"
+    supported_scales = (2.0, 4.0)
+    input_channels = 3
+    # Real-ESRGAN is convolutional; the Upscaler-level tiler wraps the forward
+    # call with overlap padding large enough to cover the receptive field.
+    default_tile = 256
+    default_tile_pad = 16
+
+    def __init__(self, device: torch.device, dtype: torch.dtype, registry_key: str):
+        super().__init__(device, dtype)
+        self._registry_key = registry_key
+        self._model: torch.nn.Module | None = None
+
+    def _load(self) -> torch.nn.Module:
+        if self._model is not None:
+            return self._model
+        from spandrel import ModelLoader
+
+        weights_path = ensure_weight(self._registry_key)
+        descriptor = ModelLoader().load_from_file(str(weights_path))
+        model = descriptor.model.to(self.device, self.dtype)
+        model.eval()
+        self._model = model
+        return model
+
+    def upscale(self, tensor: torch.Tensor, scale: float) -> torch.Tensor:
+        self.validate_scale(scale)
+        model = self._load()
+        tensor = tensor.to(self.device, self.dtype)
+        with torch.no_grad():
+            return model(tensor).clamp(0.0, 1.0)
+
+
+def _factory(registry_key: str):
+    def make(device: torch.device, dtype: torch.dtype) -> RealESRGANBackbone:
+        return RealESRGANBackbone(device, dtype, registry_key=registry_key)
+
+    return make
+
+
+register("realesrgan-x2", _factory("realesrgan-x2"))
+register("realesrgan-x4", _factory("realesrgan-x4"))

metalgrow/backbones/registry.py

@@ -0,0 +1,30 @@
+from __future__ import annotations
+
+from collections.abc import Callable
+
+import torch
+
+from metalgrow.backbones.base import Backbone
+
+BackboneFactory = Callable[[torch.device, torch.dtype], Backbone]
+
+_REGISTRY: dict[str, BackboneFactory] = {}
+
+
+def register(name: str, factory: BackboneFactory) -> None:
+    if name in _REGISTRY:
+        raise ValueError(f"backbone {name!r} already registered")
+    _REGISTRY[name] = factory
+
+
+def get_backbone(name: str, device: torch.device, dtype: torch.dtype) -> Backbone:
+    try:
+        factory = _REGISTRY[name]
+    except KeyError:
+        available = ", ".join(sorted(_REGISTRY)) or "<none>"
+        raise ValueError(f"unknown backbone {name!r}; available: {available}") from None
+    return factory(device, dtype)
+
+
+def list_backbones() -> list[str]:
+    return sorted(_REGISTRY)

metalgrow/backbones/swinir.py

@@ -0,0 +1,58 @@
+from __future__ import annotations
+
+import torch
+
+from metalgrow.backbones.base import Backbone
+from metalgrow.backbones.registry import register
+from metalgrow.weights import ensure_weight
+
+
+class SwinIRBackbone(Backbone):
+    """SwinIR (Swin Transformer) backbone loaded via :mod:`spandrel`.
+
+    SwinIR trades more memory and latency for generally sharper outputs than
+    Real-ESRGAN on clean, photographic content. Attention-heavy, so we tile
+    at a smaller default than Real-ESRGAN to keep per-tile activations in
+    VRAM on the target device.
+    """
+
+    name = "swinir"
+    supported_scales = (2.0, 4.0)
+    input_channels = 3
+    default_tile = 128
+    default_tile_pad = 16
+
+    def __init__(self, device: torch.device, dtype: torch.dtype, registry_key: str):
+        super().__init__(device, dtype)
+        self._registry_key = registry_key
+        self._model: torch.nn.Module | None = None
+
+    def _load(self) -> torch.nn.Module:
+        if self._model is not None:
+            return self._model
+        from spandrel import ModelLoader
+
+        weights_path = ensure_weight(self._registry_key)
+        descriptor = ModelLoader().load_from_file(str(weights_path))
+        model = descriptor.model.to(self.device, self.dtype)
+        model.eval()
+        self._model = model
+        return model
+
+    def upscale(self, tensor: torch.Tensor, scale: float) -> torch.Tensor:
+        self.validate_scale(scale)
+        model = self._load()
+        tensor = tensor.to(self.device, self.dtype)
+        with torch.no_grad():
+            return model(tensor).clamp(0.0, 1.0)
+
+
+def _factory(registry_key: str):
+    def make(device: torch.device, dtype: torch.dtype) -> SwinIRBackbone:
+        return SwinIRBackbone(device, dtype, registry_key=registry_key)
+
+    return make
+
+
+register("swinir-x2", _factory("swinir-x2"))
+register("swinir-x4", _factory("swinir-x4"))

metalgrow/batch.py

@@ -0,0 +1,163 @@
+"""Batch / directory upscaling.
+
+Reads and writes happen on a thread pool so I/O can overlap with model
+inference, but the inference call itself is serialized on the main thread:
+running multiple GPU forwards concurrently from Python threads doesn't speed
+anything up and risks contention on MPS, so we prefetch reads and queue
+writes around a single in-flight forward.
+"""
+
+from __future__ import annotations
+
+import glob as glob_module
+from collections import deque
+from collections.abc import Iterable, Iterator
+from concurrent.futures import Future, ThreadPoolExecutor
+from dataclasses import dataclass
+from pathlib import Path
+
+from PIL import Image
+from rich.progress import (
+    BarColumn,
+    Progress,
+    SpinnerColumn,
+    TextColumn,
+    TimeElapsedColumn,
+)
+
+from metalgrow.upscaler import Upscaler
+
+IMAGE_EXTS = frozenset({".png", ".jpg", ".jpeg", ".webp", ".bmp", ".tif", ".tiff"})
+
+
+@dataclass(frozen=True)
+class BatchResult:
+    processed: int
+    skipped: int
+    failed: int
+
+    @property
+    def total(self) -> int:
+        return self.processed + self.skipped + self.failed
+
+
+def discover_inputs(src: str) -> list[Path]:
+    """Resolve ``src`` to a sorted list of input image paths.
+
+    Accepts a single file path, a directory (non-recursive walk), or a glob
+    pattern containing wildcards. Directories filter to known image
+    extensions; globs trust the user.
+    """
+    if any(ch in src for ch in "*?["):
+        matches = sorted(Path(p) for p in glob_module.glob(src))
+        return [p for p in matches if p.is_file()]
+    path = Path(src)
+    if path.is_file():
+        return [path]
+    if path.is_dir():
+        return sorted(p for p in path.iterdir() if p.is_file() and p.suffix.lower() in IMAGE_EXTS)
+    raise FileNotFoundError(src)
+
+
+def plan_outputs(inputs: Iterable[Path], dst_dir: Path) -> list[tuple[Path, Path]]:
+    """Map each input to ``dst_dir / input.name``."""
+    return [(src, dst_dir / src.name) for src in inputs]
+
+
+def _load(src: Path) -> tuple[Image.Image, str]:
+    image = Image.open(src)
+    has_alpha = image.mode in ("RGBA", "LA") or "transparency" in image.info
+    mode = "RGBA" if has_alpha else "RGB"
+    return image.convert(mode), mode
+
+
+def _save(image: Image.Image, dst: Path) -> None:
+    dst.parent.mkdir(parents=True, exist_ok=True)
+    image.save(dst)
+
+
+def _prefetch(
+    executor: ThreadPoolExecutor,
+    items: Iterable[tuple[Path, Path]],
+    depth: int,
+) -> Iterator[tuple[Path, Path, Future]]:
+    """Stream items keeping at most ``depth`` reads in flight at once."""
+    pending: deque[tuple[Path, Path, Future]] = deque()
+    iterator = iter(items)
+
+    def submit_next() -> bool:
+        try:
+            src, dst = next(iterator)
+        except StopIteration:
+            return False
+        pending.append((src, dst, executor.submit(_load, src)))
+        return True
+
+    for _ in range(max(1, depth)):
+        if not submit_next():
+            break
+
+    while pending:
+        item = pending.popleft()
+        yield item
+        submit_next()
+
+
+def run_batch(
+    upscaler: Upscaler,
+    items: list[tuple[Path, Path]],
+    *,
+    scale: float,
+    tile: int | None,
+    tile_pad: int | None,
+    workers: int = 4,
+    skip_existing: bool = False,
+    progress: bool = True,
+) -> BatchResult:
+    """Run ``upscaler`` over each (src, dst) pair, returning counts."""
+    pending = [(s, d) for s, d in items if not (skip_existing and d.exists())]
+    skipped = len(items) - len(pending)
+    processed = 0
+    failed = 0
+    write_futs: list[Future] = []
+
+    columns = (
+        SpinnerColumn(),
+        TextColumn("[progress.description]{task.description}"),
+        BarColumn(),
+        TextColumn("{task.completed}/{task.total}"),
+        TimeElapsedColumn(),
+    )
+
+    with ThreadPoolExecutor(max_workers=max(1, workers)) as pool:
+        with Progress(*columns, disable=not progress) as bar:
+            task = bar.add_task("upscaling", total=len(pending))
+            for src, dst, read_fut in _prefetch(pool, pending, depth=workers):
+                try:
+                    image, _mode = read_fut.result()
+                except Exception as exc:  # noqa: BLE001 — surface but keep batch going
+                    failed += 1
+                    bar.console.log(f"[red]read failed[/red] {src}: {exc}")
+                    bar.advance(task)
+                    continue
+                try:
+                    out = upscaler.upscale(image, scale=scale, tile=tile, tile_pad=tile_pad)
+                except Exception as exc:  # noqa: BLE001
+                    failed += 1
+                    bar.console.log(f"[red]upscale failed[/red] {src}: {exc}")
+                    bar.advance(task)
+                    continue
+                write_futs.append(pool.submit(_save, out, dst))
+                processed += 1
+                bar.advance(task)
+
+        for fut in write_futs:
+            try:
+                fut.result()
+            except Exception as exc:  # noqa: BLE001
+                failed += 1
+                processed -= 1
+                # Best-effort surfacing; a dedicated logger would be nicer here.
+                print(f"write failed: {exc}")
+
+    return BatchResult(processed=processed, skipped=skipped, failed=failed)

metalgrow/cli.py

@@ -0,0 +1,137 @@
+from pathlib import Path
+
+import torch
+import typer
+
+from metalgrow.backbones import list_backbones
+from metalgrow.batch import discover_inputs, plan_outputs, run_batch
+from metalgrow.upscaler import Upscaler
+from metalgrow.weights import (
+    REGISTRY,
+    cached_path,
+    ensure_weight,
+    remove_cached,
+)
+
+_DTYPES = {"fp32": torch.float32, "fp16": torch.float16}
+
+app = typer.Typer(help="metalgrow — AI image upscaler on Apple Metal.")
+
+
+@app.command()
+def upscale(
+    src: str = typer.Argument(..., help="Image file, directory, or glob (e.g. 'in/*.png')"),
+    dst: Path = typer.Argument(..., help="Output file (single src) or directory (batch)"),
+    scale: float = typer.Option(2.0, "--scale", "-s", min=1.01, max=8.0),
+    device: str = typer.Option("auto", "--device", "-d", help="auto | mps | cuda | cpu"),
+    backbone: str = typer.Option(
+        "bicubic",
+        "--backbone",
+        "-b",
+        help=f"SR backbone: {', '.join(list_backbones())}",
+    ),
+    dtype: str = typer.Option(
+        "fp32", "--dtype", help="Inference dtype: fp32 | fp16 (fp16 MPS-only, noisier)"
+    ),
+    tile: int | None = typer.Option(
+        None, "--tile", help="Tile size in input px (0 disables; omit = backbone default)"
+    ),
+    tile_pad: int | None = typer.Option(
+        None, "--tile-pad", help="Context padding per tile edge (omit = backbone default)"
+    ),
+    skip_existing: bool = typer.Option(
+        False, "--skip-existing", help="Skip outputs that already exist (batch mode)"
+    ),
+    workers: int = typer.Option(
+        4, "--workers", "-j", min=1, help="Parallel I/O workers (inference stays serial)"
+    ),
+):
+    if dtype not in _DTYPES:
+        raise typer.BadParameter(f"dtype must be one of {list(_DTYPES)}")
+    try:
+        inputs = discover_inputs(src)
+    except FileNotFoundError as exc:
+        raise typer.BadParameter(f"source not found: {exc}") from None
+    if not inputs:
+        raise typer.BadParameter(f"no images found at {src!r}")
+
+    upscaler = Upscaler(backbone=backbone, device=device, dtype=_DTYPES[dtype])
+    typer.echo(f"device: {upscaler.device}")
+    typer.echo(f"backbone: {backbone}")
+    typer.echo(f"dtype: {dtype}")
+
+    batch_mode = len(inputs) > 1 or Path(src).is_dir() or any(ch in src for ch in "*?[")
+
+    if not batch_mode:
+        out = upscaler.upscale_file(inputs[0], dst, scale=scale, tile=tile, tile_pad=tile_pad)
+        typer.echo(f"wrote: {out}")
+        return
+
+    if dst.exists() and not dst.is_dir():
+        raise typer.BadParameter(f"batch destination must be a directory: {dst}")
+    dst.mkdir(parents=True, exist_ok=True)
+
+    items = plan_outputs(inputs, dst)
+    result = run_batch(
+        upscaler,
+        items,
+        scale=scale,
+        tile=tile,
+        tile_pad=tile_pad,
+        workers=workers,
+        skip_existing=skip_existing,
+    )
+    typer.echo(
+        f"done: {result.processed} processed, {result.skipped} skipped, "
+        f"{result.failed} failed (of {result.total})"
+    )
+
+
+@app.command()
+def info():
+    import torch
+
+    typer.echo(f"torch: {torch.__version__}")
+    typer.echo(f"mps available: {torch.backends.mps.is_available()}")
+    typer.echo(f"cuda available: {torch.cuda.is_available()}")
+
+
+models_app = typer.Typer(help="Manage cached model weights.")
+app.add_typer(models_app, name="models")
+
+
+@models_app.command("list")
+def models_list():
+    """List registered models with cache status."""
+    header = f"{'NAME':<20} {'SIZE':>10}  {'SHA256':<16}  STATUS"
+    typer.echo(header)
+    for name, spec in sorted(REGISTRY.items()):
+        path = cached_path(name)
+        if path.exists():
+            size = f"{path.stat().st_size / 1e6:.1f} MB"
+            status = "cached"
+        else:
+            size = "-"
+            status = "missing"
+        typer.echo(f"{name:<20} {size:>10}  {spec.sha256[:16]}  {status}")
+
+
+@models_app.command("download")
+def models_download(name: str = typer.Argument(...)):
+    """Fetch weights for NAME and verify the sha256."""
+    if name not in REGISTRY:
+        raise typer.BadParameter(f"unknown model {name!r}; see `metalgrow models list`")
+    path = ensure_weight(name)
+    typer.echo(f"ok: {path}")
+
+
+@models_app.command("rm")
+def models_rm(name: str = typer.Argument(...)):
+    """Remove the cached weight for NAME."""
+    if name not in REGISTRY:
+        raise typer.BadParameter(f"unknown model {name!r}; see `metalgrow models list`")
+    typer.echo("removed" if remove_cached(name) else "not cached")
+
+
+if __name__ == "__main__":
+    app()

metalgrow/device.py

@@ -0,0 +1,30 @@
+import os
+
+import torch
+
+
+def get_device(prefer: str = "auto") -> torch.device:
+    if prefer == "cpu":
+        return torch.device("cpu")
+    if prefer in ("mps", "auto") and torch.backends.mps.is_available():
+        _enable_mps_fallback()
+        return torch.device("mps")
+    if prefer in ("cuda", "auto") and torch.cuda.is_available():
+        return torch.device("cuda")
+    return torch.device("cpu")
+
+
+def _enable_mps_fallback() -> None:
+    """Set PYTORCH_ENABLE_MPS_FALLBACK=1 so ops missing on MPS fall back to CPU.
+
+    Several torch ops (notably some upsampling and reflection-pad variants)
+    still have gaps on the MPS backend. Without the fallback, a single
+    unsupported op aborts the forward pass; with it, torch silently executes
+    the op on CPU and returns to MPS for the rest.
+
+    Opt out by setting ``METALGROW_DISABLE_MPS_FALLBACK=1`` — useful when
+    debugging which op is falling back.
+    """
+    if os.environ.get("METALGROW_DISABLE_MPS_FALLBACK") == "1":
+        return
+    os.environ.setdefault("PYTORCH_ENABLE_MPS_FALLBACK", "1")

metalgrow/upscaler.py

@@ -0,0 +1,177 @@
+from collections.abc import Callable
+from pathlib import Path
+
+import torch
+import torch.nn.functional as F
+from PIL import Image
+from torchvision.transforms.functional import pil_to_tensor, to_pil_image
+
+from metalgrow.backbones import get_backbone
+from metalgrow.device import get_device
+
+
+class Upscaler:
+    """Upscales images using a pluggable backbone on the selected torch device.
+
+    The ``bicubic`` backbone is always available and requires no weights, so the
+    pipeline is runnable end-to-end in CI. Learned backbones (Real-ESRGAN,
+    SwinIR, …) are selected by name via the registry in
+    :mod:`metalgrow.backbones`.
+    """
+
+    def __init__(
+        self,
+        backbone: str = "bicubic",
+        device: str = "auto",
+        dtype: torch.dtype = torch.float32,
+    ):
+        self.device = get_device(device)
+        self.dtype = dtype
+        self.backbone = get_backbone(backbone, self.device, dtype)
+
+    def upscale(
+        self,
+        image: Image.Image,
+        scale: float = 2.0,
+        tile: int | None = None,
+        tile_pad: int | None = None,
+    ) -> Image.Image:
+        tensor = pil_to_tensor(image).to(self.device, self.dtype).unsqueeze(0) / 255.0
+        _, c, _, _ = tensor.shape
+
+        if c == 4 and self.backbone.input_channels == 3:
+            # Alpha is out-of-distribution for an RGB-trained SR model. Run
+            # the backbone on the RGB plane and bicubic-upscale the alpha
+            # channel, then recombine. Cheap and preserves transparency.
+            rgb = self._run_backbone(tensor[:, :3], scale, tile, tile_pad).clamp(0.0, 1.0)
+            alpha = F.interpolate(
+                tensor[:, 3:4], size=rgb.shape[-2:], mode="bicubic", align_corners=False
+            ).clamp(0.0, 1.0)
+            out = torch.cat([rgb, alpha], dim=1)
+        else:
+            out = self._run_backbone(tensor, scale, tile, tile_pad).clamp(0.0, 1.0)
+
+        mode = {1: "L", 3: "RGB", 4: "RGBA"}.get(out.shape[1])
+        return to_pil_image(out.squeeze(0).float().cpu(), mode=mode)
+
+    def upscale_file(
+        self,
+        src: Path,
+        dst: Path,
+        scale: float = 2.0,
+        tile: int | None = None,
+        tile_pad: int | None = None,
+    ) -> Path:
+        image = Image.open(src)
+        has_alpha = image.mode in ("RGBA", "LA") or "transparency" in image.info
+        image = image.convert("RGBA" if has_alpha else "RGB")
+        result = self.upscale(image, scale=scale, tile=tile, tile_pad=tile_pad)
+        dst.parent.mkdir(parents=True, exist_ok=True)
+        result.save(dst)
+        return dst
+
+    def _run_backbone(
+        self,
+        tensor: torch.Tensor,
+        scale: float,
+        tile: int | None,
+        tile_pad: int | None,
+    ) -> torch.Tensor:
+        if tile is None:
+            tile = self.backbone.default_tile
+        if tile_pad is None:
+            tile_pad = self.backbone.default_tile_pad
+
+        _, _, h, w = tensor.shape
+        if not tile or tile <= 0 or (h <= tile and w <= tile):
+            return self.backbone.upscale(tensor, scale)
+
+        return _tiled_forward(
+            lambda x: self.backbone.upscale(x, scale),
+            tensor,
+            scale=scale,
+            tile=int(tile),
+            pad=int(tile_pad),
+        )
+
+
+def _tiled_forward(
+    fn: Callable[[torch.Tensor], torch.Tensor],
+    x: torch.Tensor,
+    scale: float,
+    tile: int,
+    pad: int,
+) -> torch.Tensor:
+    """Run ``fn`` over tiled crops of ``x`` and blend with a feathered overlap.
+
+    Each tile is padded by ``pad`` pixels of real image context on every edge
+    that isn't at the image boundary, giving the backbone enough receptive
+    field to produce seam-free interior pixels. Overlapping regions are
+    combined with a linear feather weight so any residual edge mismatch fades
+    instead of appearing as a hard seam.
+    """
+    if tile <= 0:
+        raise ValueError("tile must be positive")
+
+    n, c, h, w = x.shape
+    out_h = int(h * scale)
+    out_w = int(w * scale)
+    out_dtype = x.dtype
+    out_device = x.device
+
+    accum: torch.Tensor | None = None
+    norm: torch.Tensor | None = None
+
+    for y0 in range(0, h, tile):
+        for x0 in range(0, w, tile):
+            y1 = min(y0 + tile, h)
+            x1 = min(x0 + tile, w)
+            py0, py1 = max(0, y0 - pad), min(h, y1 + pad)
+            px0, px1 = max(0, x0 - pad), min(w, x1 + pad)
+
+            patch_out = fn(x[:, :, py0:py1, px0:px1])
+
+            if accum is None:
+                accum = torch.zeros(
+                    n, patch_out.shape[1], out_h, out_w, dtype=out_dtype, device=out_device
+                )
+                norm = torch.zeros(1, 1, out_h, out_w, dtype=out_dtype, device=out_device)
+
+            sy = patch_out.shape[-2] / (py1 - py0)
+            sx = patch_out.shape[-1] / (px1 - px0)
+            oy0, oy1 = round(py0 * sy), round(py1 * sy)
+            ox0, ox1 = round(px0 * sx), round(px1 * sx)
+            ph = oy1 - oy0
+            pw = ox1 - ox0
+            if patch_out.shape[-2:] != (ph, pw):
+                # Fall back to the patch's actual size if the backbone didn't
+                # scale integrally (e.g. non-integer scale factors).
+                ph, pw = patch_out.shape[-2:]
+                oy1 = oy0 + ph
+                ox1 = ox0 + pw
+
+            left_pad = round((y0 - py0) * sy)
+            right_pad = round((py1 - y1) * sy)
+            top_pad = round((x0 - px0) * sx)
+            bot_pad = round((px1 - x1) * sx)
+
+            wy = _feather(ph, left_pad, right_pad, out_dtype, out_device)
+            wx = _feather(pw, top_pad, bot_pad, out_dtype, out_device)
+            weight = wy.view(-1, 1) * wx.view(1, -1)
+
+            accum[:, :, oy0:oy1, ox0:ox1] += patch_out * weight
+            norm[:, :, oy0:oy1, ox0:ox1] += weight
+
+    assert accum is not None and norm is not None
+    return accum / norm.clamp_min(1e-8)
+
+
+def _feather(length: int, left: int, right: int, dtype: torch.dtype, device: torch.device):
+    w = torch.ones(length, dtype=dtype, device=device)
+    if left > 0:
+        ramp = torch.linspace(0.0, 1.0, left + 2, dtype=dtype, device=device)[1:-1]
+        w[:left] = ramp
+    if right > 0:
+        ramp = torch.linspace(1.0, 0.0, right + 2, dtype=dtype, device=device)[1:-1]
+        w[length - right :] = ramp
+    return w

metalgrow/weights.py

@@ -0,0 +1,152 @@
+"""Weight file registry + download-on-first-use cache for learned backbones."""
+
+from __future__ import annotations
+
+import hashlib
+import os
+import shutil
+import sys
+from dataclasses import dataclass
+from pathlib import Path
+from urllib.request import urlopen
+
+
+@dataclass(frozen=True)
+class WeightSpec:
+    name: str
+    url: str
+    sha256: str
+
+
+REGISTRY: dict[str, WeightSpec] = {
+    "realesrgan-x2": WeightSpec(
+        name="RealESRGAN_x2plus.pth",
+        url=(
+            "https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.1/RealESRGAN_x2plus.pth"
+        ),
+        sha256="49fafd45f8fd7aa8d31ab2a22d14d91b536c34494a5cfe31eb5d89c2fa266abb",
+    ),
+    "realesrgan-x4": WeightSpec(
+        name="RealESRGAN_x4plus.pth",
+        url=(
+            "https://github.com/xinntao/Real-ESRGAN/releases/download/v0.1.0/RealESRGAN_x4plus.pth"
+        ),
+        sha256="4fa0d38905f75ac06eb49a7951b426670021be3018265fd191d2125df9d682f1",
+    ),
+    "swinir-x2": WeightSpec(
+        name="001_classicalSR_DIV2K_s48w8_SwinIR-M_x2.pth",
+        url=(
+            "https://github.com/JingyunLiang/SwinIR/releases/download/v0.0/"
+            "001_classicalSR_DIV2K_s48w8_SwinIR-M_x2.pth"
+        ),
+        sha256="ac6d4895f73aa1b3d563937e29dcedcfdff4d401280e248f23db3c6f68e8343b",
+    ),
+    "swinir-x4": WeightSpec(
+        name="001_classicalSR_DIV2K_s48w8_SwinIR-M_x4.pth",
+        url=(
+            "https://github.com/JingyunLiang/SwinIR/releases/download/v0.0/"
+            "001_classicalSR_DIV2K_s48w8_SwinIR-M_x4.pth"
+        ),
+        sha256="129dc773ba2d4c07f3eb0bb116fbe692011b7cc072d9ca12797cd3748198610a",
+    ),
+}
+
+
+def cache_dir() -> Path:
+    """Return the weight cache directory, honoring METALGROW_CACHE_DIR."""
+    override = os.environ.get("METALGROW_CACHE_DIR")
+    root = Path(override) if override else Path.home() / ".cache" / "metalgrow"
+    root.mkdir(parents=True, exist_ok=True)
+    return root
+
+
+def _sha256_of(path: Path, chunk: int = 1 << 20) -> str:
+    h = hashlib.sha256()
+    with path.open("rb") as f:
+        while True:
+            buf = f.read(chunk)
+            if not buf:
+                break
+            h.update(buf)
+    return h.hexdigest()
+
+
+def _download(url: str, dst: Path) -> None:
+    """Stream ``url`` to ``dst`` atomically, printing coarse progress to stderr."""
+    tmp = dst.with_suffix(dst.suffix + ".part")
+    # `\r` is only meaningful on a TTY. Under captured / piped stderr (CI logs,
+    # Typer's CliRunner, `| cat`) every update would concatenate onto one line,
+    # so we fall back to periodic line-per-decile output there.
+    tty = sys.stderr.isatty()
+    with urlopen(url) as resp:  # noqa: S310 — URLs come from a checksum-verified registry
+        total = int(resp.headers.get("Content-Length", 0) or 0)
+        read = 0
+        last_decile = -1
+        with tmp.open("wb") as out:
+            while True:
+                chunk = resp.read(1 << 20)
+                if not chunk:
+                    break
+                out.write(chunk)
+                read += len(chunk)
+                if not total:
+                    continue
+                pct = read * 100 // total
+                if tty:
+                    print(f"\rdownloading {dst.name}: {pct}%", end="", file=sys.stderr)
+                else:
+                    decile = pct // 10
+                    if decile != last_decile:
+                        print(f"downloading {dst.name}: {decile * 10}%", file=sys.stderr)
+                        last_decile = decile
+        if total and tty:
+            print("", file=sys.stderr)
+    shutil.move(tmp, dst)
+
+
+def cached_path(backbone: str) -> Path:
+    """Return the on-disk cache path for ``backbone`` (may not exist)."""
+    try:
+        spec = REGISTRY[backbone]
+    except KeyError:
+        raise KeyError(f"no weights registered for backbone {backbone!r}") from None
+    return cache_dir() / spec.name
+
+
+def is_cached(backbone: str) -> bool:
+    return cached_path(backbone).exists()
+
+
+def remove_cached(backbone: str) -> bool:
+    """Delete the cached weight for ``backbone``; return True if a file was removed."""
+    path = cached_path(backbone)
+    if path.exists():
+        path.unlink()
+        return True
+    return False
+
+
+def ensure_weight(backbone: str) -> Path:
+    """Return a local path to the verified weight file for ``backbone``.
+
+    Downloads on first use into :func:`cache_dir`; re-verifies the checksum on
+    every call. Raises ``RuntimeError`` if the downloaded file fails the hash
+    check (and removes the bad file so a retry can re-download).
+    """
+    try:
+        spec = REGISTRY[backbone]
+    except KeyError:
+        raise KeyError(f"no weights registered for backbone {backbone!r}") from None
+
+    dst = cache_dir() / spec.name
+    if not dst.exists():
+        _download(spec.url, dst)
+
+    actual = _sha256_of(dst)
+    if actual != spec.sha256:
+        dst.unlink(missing_ok=True)
+        raise RuntimeError(
+            f"checksum mismatch for {spec.name}: expected {spec.sha256}, got {actual}. "
+            "Removed the corrupt file; retry to re-download."
+        )
+    return dst

metalgrow-0.1.0.dist-info/METADATA

@@ -0,0 +1,16 @@
+Metadata-Version: 2.4
+Name: metalgrow
+Version: 0.1.0
+Summary: AI-powered image upscaler accelerated on Apple Metal (MPS).
+License-File: LICENSE
+Requires-Python: >=3.11
+Requires-Dist: numpy>=1.26
+Requires-Dist: pillow>=10.0
+Requires-Dist: spandrel>=0.4
+Requires-Dist: torch>=2.3
+Requires-Dist: torchvision>=0.18
+Requires-Dist: typer>=0.12
+Provides-Extra: dev
+Requires-Dist: pre-commit>=3.7; extra == 'dev'
+Requires-Dist: pytest>=8.0; extra == 'dev'
+Requires-Dist: ruff>=0.6; extra == 'dev'

metalgrow-0.1.0.dist-info/RECORD

@@ -0,0 +1,17 @@
+metalgrow/__init__.py,sha256=EqvSxKYXIx6SmXlRoVkzPxziRFHMrYYq0GgcXk5u9h4,140
+metalgrow/batch.py,sha256=eHcxj9sMWsw4LKj-ftqWiE6ZTbGQa2I94h-Ig6Ri1uE,5239
+metalgrow/cli.py,sha256=tgee9XVCJw3WrazY_dojLuEIfGifOCsn2QjUow0I2SQ,4521
+metalgrow/device.py,sha256=vpCkJoeennXgOdRx7XmyeB5rh4z5bSRxdqgtS65-NP0,1070
+metalgrow/upscaler.py,sha256=XOTe52GF4a2slvESzrEItT_Enb-iRZ0-TCBWT9egBUY,6286
+metalgrow/weights.py,sha256=9HakZjxaWEJVHbqEnsl6yiPrQJud0p07DZtbgtq1naE,5119
+metalgrow/backbones/__init__.py,sha256=wurbYBxxl_m842-qix-feiDZ9nYLzGuV7D8ezWRJc64,473
+metalgrow/backbones/base.py,sha256=bG_MWuHhRpCjKdxT2Ax8sRVyLvfol-BXUte_H91wDAI,2055
+metalgrow/backbones/bicubic.py,sha256=QWPTan6Tf--tAH7dW4DAslUWUFMLu394vtc-DsGm87g,770
+metalgrow/backbones/realesrgan.py,sha256=_DDaIb76qXOk5tBU3gHN0o1bjpgCAZFJZ8VHb8pwn-Q,1791
+metalgrow/backbones/registry.py,sha256=mcxHyj-iJ0kCM2dc0ML4VXYGmlSG5xf0mOwQsmt0XwA,827
+metalgrow/backbones/swinir.py,sha256=Pt5PZ0WkAmyAQeT-yPR1lRSsEacdSdL6y__4fELMYaQ,1875
+metalgrow-0.1.0.dist-info/METADATA,sha256=JA3FxqY7QVrAUVCEZpiTKPaWaJjGz4hfmfPpHNRhpdE,489
+metalgrow-0.1.0.dist-info/WHEEL,sha256=QccIxa26bgl1E6uMy58deGWi-0aeIkkangHcxk2kWfw,87
+metalgrow-0.1.0.dist-info/entry_points.txt,sha256=gMCvW17BMAsgCD1uzRJ_Cle3iRGiIHHpFKp1KY2AFmY,48
+metalgrow-0.1.0.dist-info/licenses/LICENSE,sha256=WiZDw4NgLukfYYsGRSR4UG0ppehO0W8U_PiD-m69SDA,1068
+metalgrow-0.1.0.dist-info/RECORD,,

metalgrow-0.1.0.dist-info/WHEEL

@@ -0,0 +1,4 @@
+Wheel-Version: 1.0
+Generator: hatchling 1.29.0
+Root-Is-Purelib: true
+Tag: py3-none-any

metalgrow-0.1.0.dist-info/entry_points.txt

@@ -0,0 +1,2 @@
+[console_scripts]
+metalgrow = metalgrow.cli:app

metalgrow-0.1.0.dist-info/licenses/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2026 joaodotwork
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.