repository·main·Indexed 43 days ago
https://github.com/chainner-org/spandrelA library for loading and running pre-trained PyTorch models (v0.4.2). It automatically detects model architecture and hyperparameters from files (.pth, .pt, .ckpt, .safetensors) and provides a unified interface for inference. It supports a wide variety of image processing tasks, including Single Image Super Resolution (ESRGAN, SwinIR, HAT), Face Restoration (GFPGAN), Inpainting (LaMa), and Denoising (SCUNet). Additional architectures are available via the spandrel_extra_arches extension.
The Spandrel repository itself is licensed under the MIT license.
Important: The specific code for implemented architectures (located within __arch__/ directories) is governed by their own original respective licenses. Users should check the individual architecture directories for their specific licensing terms.
While the spandrel_extra_arches library itself is licensed under the MIT license, it contains the source code for specific architectures that carry their own restrictive licenses (such as non-commercial or copyleft licenses).
Install the core Spandrel library using pip to enable automatic detection of model architectures and hyperparameters from model files.
pip install spandrelTo run an image-to-image model (such as ESRGAN, SwinIR, or HAT), use ModelLoader to load the file, verify it is an ImageModelDescriptor, and move it to the desired device.
Note: ImageModelDescriptor does not convert images to tensors; you must provide a torch.Tensor for the forward pass. It is a wrapper around the PyTorch model that provides a unified interface and metadata.
from spandrel import ImageModelDescriptor, ModelLoader
import torch
# load a model from disk
model = ModelLoader().load_from_file(r"path/to/model.pth")
# make sure it's an image to image model
assert isinstance(model, ImageModelDescriptor)
# send it to the GPU and put it in inference mode
model.cuda().eval()
# use the model
def process(image: torch.Tensor) -> torch.Tensor:
with torch.no_grad():
return model(image)To use the extra architectures, you must install both the core spandrel library and the spandrel_extra_arches extension via pip.
pip install spandrel spandrel_extra_archesThe spandrel package only includes architectures with permissive licenses (MIT, Apache 2.0, public domain). For architectures with restrictive or non-commercial licenses, install and use the spandrel_extra_arches package.
You must call spandrel_extra_arches.install() before initializing ModelLoader to ensure the extra architectures are registered.
import spandrel
import spandrel_extra_arches
# add extra architectures before `ModelLoader` is used
spandrel_extra_arches.install()
# load a model from disk
model = spandrel.ModelLoader().load_from_file(r"path/to/model.pth")
... # use modelSpandrel uses descriptors to define the properties, requirements, and purpose of a model architecture. These descriptors allow the library to understand what a model does and what it needs to run.
Key descriptor types include:
ModelDescriptor: The base class for all descriptors.ImageModelDescriptor: For models that perform image-to-image or text-to-image tasks.MaskedImageModelDescriptor: For models that operate on masked images.Architecture: Represents the specific architecture definition.ArchId: A unique identifier for an architecture.Purpose: Describes the intended use of the model.SizeRequirements: Defines the input size constraints of the model.ModelTiling: Information regarding whether the model supports tiling.StateDict: The dictionary containing the model weights.When choosing model files, it is highly recommended to use .safetensors files for guaranteed security.
Loading .pth files poses a security risk because they use Python's pickle module, which is inherently unsafe and vulnerable to arbitrary code execution (ACE). While Spandrel attempts to mitigate this by only allowing the deserialization of certain data types, it does not fully eliminate the risk of ACE. Always prioritize .safetensors whenever available.
To enable the additional architectures in your project, you must call spandrel_extra_arches.install() before initializing or using spandrel.ModelLoader(). This registers the extra architectures so that the loader can recognize and instantiate them when loading model files.
import spandrel
import spandrel_extra_arches
# add extra architectures before `ModelLoader` is used
spandrel_extra_arches.install()
# load a model from disk
model = spandrel.ModelLoader().load_from_file(r"path/to/model.pth")
... # use modelFor torch.nn.MultiheadAttention with embed_dim=D and num_heads=H (where D % H == 0):
If kdim=D and vdim=D (default behavior):
p.in_proj_weight: [3*D, D]p.in_proj_bias: [3*D]p.out_proj.weight: [D, D]p.out_proj.bias: [D]If kdim=K, vdim=V, and add_bias_kv=True:
p.q_proj_weight: [D, D]p.k_proj_weight: [D, K]p.v_proj_weight: [D, V]p.in_proj_bias: [3*D]p.bias_k: [1, 1, D]p.bias_v: [1, 1, D]p.out_proj.weight: [D, D]p.out_proj.bias: [D]If add_bias_kv=False:
p.q_proj_weight: [D, D]p.k_proj_weight: [D, K]p.v_proj_weight: [D, V]p.in_proj_bias: [3*D]p.out_proj.weight: [D, D]p.out_proj.bias: [D]# Example: Standard MultiheadAttention
assert D % H == 0
p = nn.MultiheadAttention(embed_dim=D, num_heads=H, bias=True)
# p.in_proj_weight: Tensor Size([3*D, D])
# p.in_proj_bias: Tensor Size([3*D])
# p.out_proj.weight: Tensor Size([D, D])
# p.out_proj.bias: Tensor Size([D])Spandrel supports a wide variety of neural network architectures for image processing tasks. Architectures are categorized by their primary function. Note that some architectures are included in the core spandrel package, while others require the spandrel_extra_arches package (marked with a + in the documentation).
Includes ESRGAN (and variants like Real-ESRGAN, BSRGAN), SwinIR, Swin2SR, HAT, Omni-SR, SRFormer (+), DAT, FeMaSR (+), Real-CUGAN, and many others.
Includes GFPGAN, RestoreFormer, and CodeFormer (+).
Includes LaMa and MAT (+).
Includes SCUNet, Uformer, KBNet, NAFNet, Restormer (+), MPRNet (+), MIRNet2 (+), and others.
The parameters stored in the state dict for torch.nn.BatchNorm2d depend on the affine and track_running_stats settings. Given num_features=N:
affine=True and track_running_stats=True:p.weight: [N]p.bias: [N]p.running_mean: [N]p.running_var: [N]p.num_batches_tracked: [] (scalar)affine=True and track_running_stats=False:p.weight: [N]p.bias: [N]affine=False:# Example: affine=True, track_running_stats=True
p = nn.BatchNorm2d(num_features=N, affine=True, track_running_stats=True)
# p.weight: Tensor Size([N])
# p.bias: Tensor Size([N])
# p.running_mean: Tensor Size([N])
# p.running_var: Tensor Size([N])
# p.num_batches_tracked: Tensor Size([])