spandrel

package documentation
(source)

Spandrel is a library for loading and running pre-trained PyTorch models. It automatically detects the model architecture and hyper parameters from model files, and provides a unified interface for running models.

Package architectures The package containing the implementations of all supported architectures. Not necessary for most user code.
Package util A module containing commonly-used functionality to implement architectures.

From __init__.py:

Class Architecture The abstract base class for all architectures.
Class ArchRegistry A registry of architectures.
Class ArchSupport An entry in an ArchRegistry that describes how to detect and load a model architecture.
Class ImageModelDescriptor A model that takes an image as input and returns an image. See __call__ for more information.
Class MaskedImageModelDescriptor A model that takes an image and a mask for that image as input and returns an image. See __call__ for more information.
Class ModelBase The base class of all model descriptors.
Class ModelLoader Class for automatically loading a pth file into any architecture
Class ModelTiling Describes whether and how a model supports tiling.
Class SizeRequirements A set of requirements for the size of an input image.
Exception DuplicateArchitectureError An error that will be thrown by ArchRegistry if the same architecture is added twice.
Exception UnsupportedDtypeError An error that will be thrown by .to if the model does not support the given dtype.
Exception UnsupportedModelError An error that will be thrown by ArchRegistry and ModelLoader if a model architecture is not supported.
Constant MAIN_REGISTRY The main architecture registry of spandrel.
Type Alias ModelDescriptor A model descriptor is a loaded model with metadata. Metadata includes the architecture, purpose, tags, and other information about the model.
Type Alias Purpose A short string describing the purpose of the model.
Type Alias StateDict Spandrel's type alias for PyTorch state dicts.
Variable ArchId A unique identifier for an Architecture.
Function canonicalize_state_dict Canonicalize a state dict.
def canonicalize_state_dict(state_dict: StateDict) -> StateDict: (source)

Canonicalize a state dict.

This function is used to canonicalize a state dict, so that it can be used for architecture detection and loading.

This function is not intended to be used in production code.

MAIN_REGISTRY = (source)

The main architecture registry of spandrel.

Modifying this registry will affect all ModelLoader instances without a custom registry.

Value
ArchRegistry()
ArchId = (source)

A unique identifier for an Architecture.

ModelDescriptor = (source)

A model descriptor is a loaded model with metadata. Metadata includes the architecture, purpose, tags, and other information about the model.

The API of a model is described by the type of the model descriptor. E.g. a SISR model will have a descriptor of type ImageModelDescriptor.

Value
Union[ImageModelDescriptor[torch.nn.Module],
      MaskedImageModelDescriptor[torch.nn.Module]]
Purpose = (source)

A short string describing the purpose of the model.

  • SR: Super resolution
  • FaceSR: Face super resolution
  • Inpainting: Image inpainting
  • Restoration: Image restoration (denoising, deblurring, JPEG, etc.)
Value
Literal['SR', 'FaceSR', 'Inpainting', 'Restoration']
StateDict = (source)

Spandrel's type alias for PyTorch state dicts.

See https://pytorch.org/tutorials/recipes/recipes/what_is_state_dict.html

Value
Dict[str, Any]