rastermoves 0.3.1__py3-none-any.whl

rastermoves/__init__.py

@@ -0,0 +1,14 @@
+"""RasterMoves: modular image upscaling and enhancement for DanceFlow."""
+__version__ = "0.3.1"
+
+# Keep importing the package lightweight; neural frameworks load only on use.
+def __getattr__(name):
+    if name == "Upscaler":
+        from .pipeline import Upscaler
+        return Upscaler
+    if name == "Registry":
+        from .registry import Registry
+        return Registry
+    raise AttributeError(name)
+
+__all__ = ["Upscaler", "Registry", "__version__"]

rastermoves/__main__.py

@@ -0,0 +1,2 @@
+from .cli import main
+raise SystemExit(main())

rastermoves/backends/base.py

@@ -0,0 +1,17 @@
+from __future__ import annotations
+from typing import Protocol
+import numpy as np
+
+
+class LoadedModel(Protocol):
+    """Backend contract: HWC float32 RGB/gray arrays in [0,1], no batch axis."""
+    scale: int
+    input_channels: int
+    output_channels: int
+    tiling: str  # supported | discouraged | internal
+    device: str
+    precision: str
+
+    def predict(self, image: np.ndarray) -> np.ndarray: ...
+    def clear_cache(self) -> None: ...
+    def close(self) -> None: ...

rastermoves/backends/onnx_backend.py

@@ -0,0 +1,110 @@
+from __future__ import annotations
+
+import gc
+from pathlib import Path
+import numpy as np
+
+from ..errors import BackendOOM, UnsupportedModelError, UpscaleError
+
+
+class OnnxModel:
+    """Single-file, single-image ONNX models; normalized float NCHW/NHWC tensors."""
+    def __init__(self, path: Path, spec, *, device="auto", precision="fp32", **_):
+        try:
+            import onnxruntime as ort
+        except ImportError as e:
+            raise UpscaleError("Install the onnx extra (CPU), or onnx-gpu for CUDA. Do not install both runtime packages.") from e
+        available = ort.get_available_providers()
+        if device == "auto":
+            device = "cuda" if "CUDAExecutionProvider" in available else "cpu"
+        if device == "cpu":
+            providers = ["CPUExecutionProvider"]
+        elif device == "cuda" or device.startswith("cuda:"):
+            if "CUDAExecutionProvider" not in available:
+                raise UpscaleError("The installed ONNX Runtime has no CUDA provider. Install onnxruntime-gpu or use cpu.")
+            index = int(device.split(":")[1]) if ":" in device else 0
+            providers = [("CUDAExecutionProvider", {"device_id": index}), "CPUExecutionProvider"]
+        else:
+            raise UpscaleError("The ONNX backend currently accepts auto, cpu, cuda, or cuda:N.")
+        try:
+            # Bytes disable external-data path resolution; multi-file ONNX needs a custom backend.
+            self.session = ort.InferenceSession(path.read_bytes(), providers=providers)
+        except Exception as e:
+            raise UnsupportedModelError(f"ONNX load failed ({type(e).__name__}). Only self-contained ONNX files are supported.") from e
+        active = self.session.get_providers()
+        if device.startswith("cuda") and "CUDAExecutionProvider" not in active:
+            raise UpscaleError("ONNX CUDA initialization failed and fell back to CPU; check CUDA/cuDNN libraries.")
+        inputs, outputs = self.session.get_inputs(), self.session.get_outputs()
+        if len(inputs) != 1 or len(outputs) != 1 or len(inputs[0].shape) != 4 or len(outputs[0].shape) != 4:
+            raise UnsupportedModelError("ONNX model must expose one rank-four image input and one rank-four image output.")
+        self.input, self.output = inputs[0], outputs[0]
+        types = {"tensor(float)": np.float32, "tensor(float16)": np.float16}
+        if self.input.type not in types:
+            raise UnsupportedModelError("Only normalized float32/float16 ONNX image inputs are supported.")
+        self.dtype = types[self.input.type]
+        actual = "fp16" if self.dtype == np.float16 else "fp32"
+        if precision != actual:
+            raise UpscaleError(f"ONNX precision is fixed by the export ({actual}); select --precision {actual}.")
+        self.layout = spec.options.get("layout", "NCHW")
+        if self.layout not in {"NCHW", "NHWC"}:
+            raise UnsupportedModelError("ONNX layout must be NCHW or NHWC.")
+        self.input_channels, self.output_channels, self.scale = spec.input_channels, spec.output_channels, spec.scale
+        if self.input_channels not in (1, 3) or self.output_channels not in (1, 3):
+            raise UnsupportedModelError("Only RGB/grayscale ONNX models are supported.")
+        channel_axis = 1 if self.layout == "NCHW" else 3
+        for info, expected in ((self.input, self.input_channels), (self.output, self.output_channels)):
+            if isinstance(info.shape[0], int) and info.shape[0] != 1:
+                raise UnsupportedModelError("ONNX batch size must be one or dynamic.")
+            if isinstance(info.shape[channel_axis], int) and info.shape[channel_axis] != expected:
+                raise UnsupportedModelError("ONNX channel shape disagrees with the model plugin.")
+        dims = self.input.shape[2:4] if self.layout == "NCHW" else self.input.shape[1:3]
+        self.fixed = tuple(v if isinstance(v, int) and v > 0 else None for v in dims)
+        self.multiple = int(spec.options.get("multiple_of", 1))
+        self.minimum = int(spec.options.get("minimum_size", 1))
+        self.square = bool(spec.options.get("square", False))
+        self.bgr = spec.options.get("channel_order", "RGB") == "BGR"
+        if self.multiple < 1 or self.minimum < 1:
+            raise UpscaleError("ONNX padding requirements must be positive.")
+        self.tiling = spec.options.get("tiling", "supported")
+        self.device, self.precision = device, actual
+
+    def predict(self, image: np.ndarray) -> np.ndarray:
+        h, w, _ = image.shape
+        sizes = [max(self.minimum, ((v + self.multiple - 1) // self.multiple) * self.multiple) for v in (h, w)]
+        if self.square:
+            sizes = [max(sizes)] * 2
+        sizes = [fixed if fixed is not None else v for fixed, v in zip(self.fixed, sizes)]
+        if sizes[0] < h or sizes[1] < w:
+            raise UpscaleError(f"ONNX input exceeds fixed shape {self.fixed}; lower --tile and --tile-pad, or use a dynamic export.")
+        x = np.pad(image, ((0, sizes[0] - h), (0, sizes[1] - w), (0, 0)), mode="edge")
+        if self.bgr and self.input_channels == 3:
+            x = x[..., ::-1]
+        if self.layout == "NCHW":
+            x = x.transpose(2, 0, 1)
+        try:
+            out = self.session.run([self.output.name], {
+                self.input.name: np.ascontiguousarray(x[None], dtype=self.dtype)
+            })[0]
+        except Exception as e:
+            if any(k in str(e).lower() for k in ("out of memory", "cuda_error_out_of_memory")):
+                raise BackendOOM("ONNX device ran out of memory.") from None
+            raise UpscaleError(f"ONNX inference failed ({type(e).__name__}); check shape/padding options in the model plugin.") from e
+        if not isinstance(out, np.ndarray) or out.ndim != 4 or out.shape[0] != 1:
+            raise UnsupportedModelError("Unexpected ONNX image output.")
+        out = out[0]
+        if self.layout == "NCHW":
+            out = out.transpose(1, 2, 0)
+        expected = (sizes[0] * self.scale, sizes[1] * self.scale, self.output_channels)
+        if out.shape != expected:
+            raise UnsupportedModelError(f"ONNX output shape {out.shape} does not match {expected}; check native scale/layout.")
+        out = out[:h * self.scale, :w * self.scale]
+        if self.bgr and self.output_channels == 3:
+            out = out[..., ::-1]
+        return np.ascontiguousarray(out, dtype=np.float32)
+
+    def clear_cache(self):
+        gc.collect()
+
+    def close(self):
+        self.session = None
+        self.clear_cache()

rastermoves/backends/spandrel_backend.py

@@ -0,0 +1,158 @@
+from __future__ import annotations
+
+from collections.abc import Mapping
+import gc
+from pathlib import Path
+import zipfile
+
+import numpy as np
+
+from ..errors import BackendOOM, UnsupportedModelError, UpscaleError
+
+
+def safe_state_dict(path: Path):
+    import torch
+    major, minor = map(int, torch.__version__.split("+")[0].split(".")[:2])
+    if (major, minor) < (2, 6):
+        raise UpscaleError("PyTorch >=2.6 is required for restricted checkpoint loading. Upgrade torch/torchvision.")
+    if path.suffix.lower() == ".safetensors":
+        from safetensors.torch import load_file
+        state = load_file(str(path), device="cpu")
+    else:
+        if zipfile.is_zipfile(path):
+            with zipfile.ZipFile(path) as z:
+                if any("/code/" in n or n.endswith("/constants.pkl") for n in z.namelist()):
+                    raise UnsupportedModelError("TorchScript archives are not accepted; use state-dict weights or safetensors.")
+        # Never fall back to weights_only=False or torch.hub/trust_remote_code.
+        state = torch.load(path, map_location="cpu", weights_only=True)
+    if not isinstance(state, Mapping):
+        raise UnsupportedModelError("Expected a tensor state dictionary, not a serialized Python model.")
+    return state
+
+
+def normalize_state_dict(state):
+    """Unwrap common training checkpoints without invoking private Spandrel APIs.
+
+    Spandrel's public load_from_state_dict expects the actual parameter mapping,
+    unlike its load_from_file convenience method. Only passive mappings are read.
+    """
+    wrappers = ("model_state_dict", "state_dict", "params_ema", "params-ema", "params", "model", "net")
+    for _ in range(8):
+        child = next((state[k] for k in wrappers if k in state and isinstance(state[k], Mapping)), None)
+        if child is None and len(state) == 1:
+            value = next(iter(state.values()))
+            child = value if isinstance(value, Mapping) else None
+        if child is None:
+            break
+        state = child
+    if not state or not all(isinstance(k, str) for k in state):
+        raise UnsupportedModelError("Expected a nonempty state dictionary with string parameter names.")
+    for _ in range(8):
+        prefix = next((p for p in ("module.", "netG.") if all(k.startswith(p) for k in state)), None)
+        if prefix is None:
+            break
+        state = {k[len(prefix):]: v for k, v in state.items()}
+    return state
+
+
+_EXTRA_ARCHES_INSTALLED = False
+
+
+def resolve_device(requested: str):
+    import torch
+    if requested == "auto":
+        requested = "cuda" if torch.cuda.is_available() else (
+            "mps" if torch.backends.mps.is_available() else "cpu")
+    if requested not in {"cpu", "mps", "cuda"} and not requested.startswith("cuda:"):
+        raise UpscaleError("PyTorch device must be auto, cpu, mps, cuda, or cuda:N.")
+    try:
+        device = torch.device(requested)
+    except (RuntimeError, ValueError) as e:
+        raise UpscaleError("Invalid PyTorch device; use cpu, mps, cuda, or cuda:N.") from e
+    if device.type == "cuda":
+        if not torch.cuda.is_available():
+            raise UpscaleError("CUDA is unavailable. Install the matching PyTorch GPU build or use --device cpu.")
+        if device.index is not None and device.index >= torch.cuda.device_count():
+            raise UpscaleError("Requested CUDA device index does not exist.")
+    if device.type == "mps" and not torch.backends.mps.is_available():
+        raise UpscaleError("Apple MPS is unavailable on this machine/PyTorch build.")
+    return device
+
+
+class SpandrelModel:
+    def __init__(self, path: Path, spec, *, device="auto", precision="fp32", extra_arches=False,
+                 validate_metadata=True):
+        try:
+            import torch
+            from spandrel import ImageModelDescriptor, ModelLoader
+        except (ImportError, RuntimeError, OSError) as e:
+            raise UpscaleError("PyTorch backend unavailable. Install the torch extra and matching torch/torchvision builds.") from e
+        global _EXTRA_ARCHES_INSTALLED
+        if extra_arches and not _EXTRA_ARCHES_INSTALLED:
+            try:
+                import spandrel_extra_arches
+                spandrel_extra_arches.install()
+                _EXTRA_ARCHES_INSTALLED = True
+            except ImportError as e:
+                raise UpscaleError("Install the extra-arches extra from this project and review the architecture licences.") from e
+        target = resolve_device(device)
+        try:
+            state = normalize_state_dict(safe_state_dict(path))
+            descriptor = ModelLoader(device="cpu").load_from_state_dict(state)
+            if not isinstance(descriptor, ImageModelDescriptor):
+                raise UnsupportedModelError("This is not an image-to-image model; masked/inpainting models need another plugin.")
+            if descriptor.purpose == "FaceSR":
+                raise UnsupportedModelError("Face restoration needs face detection/alignment; this is not a general image upscaler.")
+            if descriptor.input_channels not in (1, 3) or descriptor.output_channels not in (1, 3):
+                raise UnsupportedModelError("This image pipeline supports only RGB and grayscale neural models.")
+            if validate_metadata and (descriptor.scale != spec.scale or
+                    descriptor.input_channels != spec.input_channels or
+                    descriptor.output_channels != spec.output_channels):
+                raise UnsupportedModelError("Detected scale/channels disagree with the manifest. Fix the plugin or use a standalone local model.")
+            if precision not in {"fp32", "fp16"}:
+                raise UpscaleError("Precision must be fp32 or fp16.")
+            if precision == "fp16" and (target.type != "cuda" or not descriptor.supports_half):
+                raise UpscaleError("FP16 requires CUDA and a model advertising half-precision support; use fp32.")
+            descriptor.to(device=target, dtype=torch.float16 if precision == "fp16" else torch.float32)
+            descriptor.eval()
+        except UpscaleError:
+            raise
+        except Exception as e:
+            raise UnsupportedModelError(f"Unable to safely load this checkpoint ({type(e).__name__}). "
+                                        "Check the architecture, dependency versions and weight format; unsafe pickle loading is disabled.") from e
+        self.descriptor = descriptor
+        self.scale = int(descriptor.scale)
+        self.input_channels = descriptor.input_channels
+        self.output_channels = descriptor.output_channels
+        self.tiling = descriptor.tiling.name.lower()
+        self.device, self.precision = str(target), precision
+        self.architecture = str(descriptor.architecture.id)
+
+    def predict(self, image: np.ndarray) -> np.ndarray:
+        import torch
+        tensor = output = None
+        try:
+            tensor = torch.from_numpy(np.ascontiguousarray(image.transpose(2, 0, 1))).unsqueeze(0)
+            tensor = tensor.to(device=self.descriptor.device, dtype=self.descriptor.dtype)
+            with torch.inference_mode():
+                # The descriptor pads to its size/window requirements and crops output back.
+                output = self.descriptor(tensor)
+            return output.detach().float().cpu().squeeze(0).permute(1, 2, 0).numpy()
+        except RuntimeError as e:
+            if isinstance(e, torch.OutOfMemoryError) or "out of memory" in str(e).lower():
+                raise BackendOOM("Device ran out of memory during inference.") from None
+            raise UpscaleError(f"PyTorch inference failed: {e}") from e
+        finally:
+            del tensor, output
+
+    def clear_cache(self):
+        import torch
+        gc.collect()
+        if self.device.startswith("cuda"):
+            torch.cuda.empty_cache()
+        elif self.device == "mps":
+            torch.mps.empty_cache()
+
+    def close(self):
+        self.descriptor = None
+        self.clear_cache()

rastermoves/catalog.py

@@ -0,0 +1,127 @@
+"""Import OpenModelDB's exported JSON, individual entries, or a checked-out data directory."""
+from __future__ import annotations
+
+from datetime import datetime, timezone
+import json
+from pathlib import Path, PurePosixPath
+import tarfile
+import tempfile
+
+from filelock import FileLock
+
+from .errors import DownloadError, UpscaleError
+from .network import atomic_json, stream_download
+from .specs import ModelSpec, Resource, validate_id
+
+API_URL = "https://openmodeldb.info/api/v1/models.json"
+ARCHIVE_URL = "https://codeload.github.com/OpenModelDB/open-model-database/tar.gz/refs/heads/main"
+RAW_BASE = "https://raw.githubusercontent.com/OpenModelDB/open-model-database/main/data/models/"
+
+
+def from_openmodeldb(model_id: str, data: dict) -> ModelSpec:
+    validate_id(model_id)
+    resources = []
+    for r in data.get("resources", []):
+        fmt = str(r.get("type", "unknown")).lower().lstrip(".")
+        platform = str(r.get("platform", "unknown")).lower()
+        backend = {"pytorch": "spandrel", "onnx": "onnx"}.get(platform, platform)
+        resources.append(Resource(format=fmt, urls=tuple(r.get("urls", [])),
+                                  sha256=r.get("sha256"), size=r.get("size"), backend=backend))
+    author = data.get("author", "unknown")
+    if isinstance(author, list):
+        author = ", ".join(str(x) for x in author)
+    license_value = data.get("license", "unknown")
+    if not isinstance(license_value, str):
+        license_value = json.dumps(license_value, ensure_ascii=False)
+    return ModelSpec(id=model_id, name=data["name"], scale=data["scale"],
+                     architecture=data.get("architecture", "unknown"), license=license_value,
+                     author=str(author), description=data.get("description", ""),
+                     tags=tuple(data.get("tags", [])), resources=tuple(resources),
+                     source_page=f"https://openmodeldb.info/models/{model_id}",
+                     input_channels=data.get("inputChannels", 3),
+                     output_channels=data.get("outputChannels", 3))
+
+
+def convert_catalog(raw) -> tuple[list[ModelSpec], list[str]]:
+    if not isinstance(raw, dict):
+        raise UpscaleError("Expected an OpenModelDB JSON object keyed by model ID.")
+    if "models" in raw and isinstance(raw["models"], dict):
+        raw = raw["models"]
+    converted, skipped = [], []
+    for key, value in raw.items():
+        try:
+            if not isinstance(value, dict):
+                raise UpscaleError("model entry is not an object")
+            converted.append(from_openmodeldb(key, value))
+        except (KeyError, TypeError, ValueError, UpscaleError) as e:
+            skipped.append(f"{key}: {e}")
+    if not converted:
+        raise UpscaleError("No valid models found; the previous catalogue was not changed.")
+    return converted, skipped
+
+
+def _read_archive(path: Path) -> dict:
+    data, expanded = {}, 0
+    with tarfile.open(path, "r:gz") as archive:
+        # Read model JSON entries only. Do not extract any paths or execute repository code.
+        for member in archive:
+            parts = PurePosixPath(member.name).parts
+            if (len(parts) != 4 or parts[1:3] != ("data", "models")
+                    or not parts[-1].endswith(".json") or not member.isfile()):
+                continue
+            if member.size > 1024 * 1024:
+                raise UpscaleError("Unexpectedly large model metadata entry.")
+            expanded += member.size
+            if expanded > 64 * 1024**2 or len(data) > 20000:
+                raise UpscaleError("Catalogue archive exceeds metadata limits.")
+            f = archive.extractfile(member)
+            if f is not None:
+                with f:
+                    data[parts[-1][:-5]] = json.load(f)
+    return data
+
+
+def sync_catalog(cache: Path, source: str | None = None, *, offline=False) -> dict:
+    cache = Path(cache)
+    cache.mkdir(parents=True, exist_ok=True)
+    origin = source or API_URL
+    with FileLock(str(cache / "catalog.lock"), timeout=600), tempfile.TemporaryDirectory() as tmp:
+        path = Path(origin).expanduser() if not origin.startswith("https://") else None
+        if path is not None and path.is_dir():
+            data_dir = path / "data" / "models" if (path / "data" / "models").is_dir() else path
+            raw = {p.stem: json.loads(p.read_text(encoding="utf-8")) for p in sorted(data_dir.glob("*.json"))}
+        elif path is not None and path.is_file():
+            raw = json.loads(path.read_text(encoding="utf-8"))
+        elif path is not None:
+            raise UpscaleError(f"Catalogue source does not exist: {path}")
+        else:
+            if offline:
+                raise DownloadError("Offline mode forbids catalogue network requests; use a local JSON file.")
+            dest = Path(tmp) / "catalog.json"
+            try:
+                stream_download(origin, dest, max_bytes=64 * 1024**2)
+                raw = json.loads(dest.read_text(encoding="utf-8"))
+            except (DownloadError, ValueError):
+                if source is not None:
+                    raise
+                origin = ARCHIVE_URL
+                dest = Path(tmp) / "catalog.tar.gz"
+                stream_download(origin, dest, max_bytes=128 * 1024**2)
+                raw = _read_archive(dest)
+        specs, skipped = convert_catalog(raw)
+        payload = {"schema_version": 1, "synced_at": datetime.now(timezone.utc).isoformat(),
+                   "source": origin, "models": [s.to_dict() for s in specs], "skipped": skipped}
+        atomic_json(cache / "catalog.json", payload)
+        return {"imported": len(specs), "skipped": skipped, "source": origin}
+
+
+def fetch_one(model_id: str, cache: Path, *, offline=False) -> ModelSpec:
+    validate_id(model_id)
+    if offline:
+        raise DownloadError(f"{model_id} is not in the local catalogue; run sync online first.")
+    with tempfile.TemporaryDirectory() as tmp:
+        path = Path(tmp) / "model.json"
+        stream_download(RAW_BASE + model_id + ".json", path, max_bytes=1024**2)
+        spec = from_openmodeldb(model_id, json.loads(path.read_text(encoding="utf-8")))
+        atomic_json(Path(cache) / "imported-models" / f"{model_id}.json", spec.to_dict())
+        return spec