### Editable Installs

Source: https://github.com/chainner-org/spandrel/blob/main/CONTRIBUTING.md

Install the Spandrel library and its extra architectures in editable mode. This also installs their dependencies.

```bash
pip install -e libs/spandrel -e libs/spandrel_extra_arches
```

--------------------------------

### Install Spandrel and Extra Architectures

Source: https://github.com/chainner-org/spandrel/blob/main/libs/spandrel_extra_arches/README.md

Install the necessary libraries using pip. This command installs both the core spandrel library and the extra architectures.

```shell
pip install spandrel spandrel_extra_arches
```

--------------------------------

### Install Pre-commit Hooks

Source: https://github.com/chainner-org/spandrel/blob/main/CONTRIBUTING.md

Install pre-commit to automatically run Ruff checks before each commit. Requires pre-commit to be installed.

```bash
pre-commit install
```

--------------------------------

### Install Development Dependencies

Source: https://github.com/chainner-org/spandrel/blob/main/CLAUDE.md

Install all necessary development dependencies, including the main library and extra architectures, in editable mode.

```bash
pip install -r requirements-dev.txt
```

```bash
pip install -e libs/spandrel -e libs/spandrel_extra_arches
```

--------------------------------

### Complete Inference Pipeline Example

Source: https://context7.com/chainner-org/spandrel/llms.txt

A full example demonstrating image loading, model inference, and saving results using OpenCV and PyTorch. Ensure you have the necessary model file and input image.

```python
from spandrel import ImageModelDescriptor, ModelLoader
import torch
import cv2
import numpy as np

def upscale_image(model_path: str, input_image_path: str, output_path: str):
    # Load model
    model = ModelLoader(device="cuda").load_from_file(model_path)
    assert isinstance(model, ImageModelDescriptor), "Expected image model"
    model.eval()

    # Load and preprocess image
    image = cv2.imread(input_image_path, cv2.IMREAD_COLOR)
    image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)

    # Convert to tensor: HWC -> CHW, [0,255] -> [0,1]
    tensor = torch.from_numpy(image).permute(2, 0, 1).float() / 255.0
    tensor = tensor.unsqueeze(0).to("cuda")  # Add batch dimension

    # Run inference
    with torch.no_grad():
        output = model(tensor)

    # Convert back to image: CHW -> HWC, [0,1] -> [0,255]
    output = output.squeeze(0).permute(1, 2, 0).cpu().numpy()
    output = (output * 255).clip(0, 255).astype(np.uint8)
    output = cv2.cvtColor(output, cv2.COLOR_RGB2BGR)

    cv2.imwrite(output_path, output)
    print(f"Upscaled {image.shape[:2]} -> {output.shape[:2]} ({model.scale}x)")

# Usage
upscale_image("RealESRGAN_x4plus.pth", "input.jpg", "output.png")
```

--------------------------------

### Install Dev Dependencies

Source: https://github.com/chainner-org/spandrel/blob/main/CONTRIBUTING.md

Install development dependencies using pip. Ensure you have the requirements-dev.txt file.

```bash
pip install -r requirements-dev.txt
```

--------------------------------

### Install and List Extra Architectures

Source: https://context7.com/chainner-org/spandrel/llms.txt

Install additional architectures from the `spandrel_extra_arches` package and list them. Note that these may have more restrictive licenses.

```python
import spandrel_extra_arches

spandrel_extra_arches.install(ignore_duplicates=True)
print("\nExtra architectures:")
for arch in spandrel_extra_arches.EXTRA_REGISTRY.architectures():
    print(f"  - {arch.architecture.id}")
```

--------------------------------

### Install Spandrel via Pip

Source: https://github.com/chainner-org/spandrel/blob/main/README.md

Use this command to install the Spandrel library. It provides basic model loading capabilities.

```shell
pip install spandrel
```

--------------------------------

### DITN Architecture `load` Method Example

Source: https://github.com/chainner-org/spandrel/blob/main/CONTRIBUTING.md

An example of a filled-in `load` function for the DITN architecture, demonstrating how to set hyperparameters, purpose based on scale, and other descriptor properties.

```python
from ...__helpers.model_descriptor import ImageModelDescriptor, StateDict
from .__arch.DITN_Real import DITN_Real as DITN


def load(state_dict: StateDict) -> ImageModelDescriptor[DITN]:
    # default values
    inp_channels = 3
    dim = 60
    ITL_blocks = 4
    SAL_blocks = 4
    UFONE_blocks = 1
    ffn_expansion_factor = 2
    bias = False
    LayerNorm_type = "WithBias"
    patch_size = 8
    upscale = 4

    model = DITN(
        inp_channels=inp_channels,
        dim=dim,
        ITL_blocks=ITL_blocks,
        SAL_blocks=SAL_blocks,
        UFONE_blocks=UFONE_blocks,
        ffn_expansion_factor=ffn_expansion_factor,
        bias=bias,
        LayerNorm_type=LayerNorm_type,
        patch_size=patch_size,
        upscale=upscale,
    )

    return ImageModelDescriptor(
        model,
        state_dict,
        architecture="DITN",
        purpose="Restoration" if upscale == 1 else "SR",
        tags=[],
        supports_half=True,
        supports_bfloat16=True,
        scale=upscale,
        input_channels=inp_channels,
        output_channels=inp_channels,
    )
```

--------------------------------

### Install and Use Extra Architectures

Source: https://context7.com/chainner-org/spandrel/llms.txt

Install additional model architectures from the spandrel_extra_arches package. Call install() before using ModelLoader to make these architectures available. Duplicates can be ignored.

```python
import spandrel
import spandrel_extra_arches

# Install extra architectures into the main registry
# Call this BEFORE using ModelLoader
installed = spandrel_extra_arches.install()
print(f"Installed {len(installed)} extra architectures")

# Now ModelLoader can detect extra architectures like CodeFormer, Restormer, etc.
model = spandrel.ModelLoader().load_from_file("codeformer.pth")

# Alternatively, ignore duplicates if already installed
spandrel_extra_arches.install(ignore_duplicates=True)

# Access the extra registry directly if needed
from spandrel_extra_arches import EXTRA_REGISTRY
print(f"Extra architectures: {[a.architecture.id for a in EXTRA_REGISTRY]}")
```

--------------------------------

### Diff Example: ffn_expansion_factor Impact

Source: https://github.com/chainner-org/spandrel/blob/main/CONTRIBUTING.md

This diff shows how changing the `ffn_expansion_factor` parameter affects the dimensions of weight tensors in the FeedForward network layers. Observe the scaling of values like 240 to 360 and 120 to 180.

```diff
       ffn
-        UFONE.0.ITLs.0.ffn.project_in.weight:  Tensor float32 Size([240, 60, 1, 1])
-        UFONE.0.ITLs.0.ffn.dwconv.weight:      Tensor float32 Size([240, 1, 3, 3])
-        UFONE.0.ITLs.0.ffn.project_out.weight: Tensor float32 Size([60, 120, 1, 1])
+        UFONE.0.ITLs.0.ffn.project_in.weight:  Tensor float32 Size([360, 60, 1, 1])
+        UFONE.0.ITLs.0.ffn.dwconv.weight:      Tensor float32 Size([360, 1, 3, 3])
+        UFONE.0.ITLs.0.ffn.project_out.weight: Tensor float32 Size([60, 180, 1, 1])
 ...
       ffn
-        UFONE.0.ITLs.1.ffn.project_in.weight:  Tensor float32 Size([240, 60, 1, 1])
-        UFONE.0.ITLs.1.ffn.dwconv.weight:      Tensor float32 Size([240, 1, 3, 3])
-        UFONE.0.ITLs.1.ffn.project_out.weight: Tensor float32 Size([60, 120, 1, 1])
+        UFONE.0.ITLs.1.ffn.project_in.weight:  Tensor float32 Size([360, 60, 1, 1])
+        UFONE.0.ITLs.1.ffn.dwconv.weight:      Tensor float32 Size([360, 1, 3, 3])
+        UFONE.0.ITLs.1.ffn.project_out.weight: Tensor float32 Size([60, 180, 1, 1])
 ...
       ffn
-        UFONE.0.ITLs.2.ffn.project_in.weight:  Tensor float32 Size([240, 60, 1, 1])
-        UFONE.0.ITLs.2.ffn.dwconv.weight:      Tensor float32 Size([240, 1, 3, 3])
-        UFONE.0.ITLs.2.ffn.project_out.weight: Tensor float32 Size([60, 120, 1, 1])
+        UFONE.0.ITLs.2.ffn.project_in.weight:  Tensor float32 Size([360, 60, 1, 1])
+        UFONE.0.ITLs.2.ffn.dwconv.weight:      Tensor float32 Size([360, 1, 3, 3])
+        UFONE.0.ITLs.2.ffn.project_out.weight: Tensor float32 Size([60, 180, 1, 1])
 ...
```

--------------------------------

### Run All Tests

Source: https://github.com/chainner-org/spandrel/blob/main/CLAUDE.md

Execute all tests in the project. Use this to ensure the entire library functions as expected.

```bash
pytest tests
```

--------------------------------

### Run Tests and Update Snapshots

Source: https://github.com/chainner-org/spandrel/blob/main/CLAUDE.md

Run tests for a specific architecture and update the snapshot files. Use when changes to expected outputs are intentional.

```bash
pytest tests/test_ESRGAN.py --snapshot-update
```

--------------------------------

### Load and Use an Image Model with Spandrel

Source: https://github.com/chainner-org/spandrel/blob/main/README.md

Demonstrates loading a super-resolution model from a file, asserting it's an image-to-image model, moving it to the GPU, and performing inference. Note that `ImageModelDescriptor` does not handle image-to-tensor conversion.

```python
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)
```

--------------------------------

### Run Pre-commit on All Files

Source: https://github.com/chainner-org/spandrel/blob/main/CONTRIBUTING.md

Manually run all pre-commit hooks on all files in the repository. Useful for ensuring compliance before committing.

```bash
pre-commit run --all-files
```

--------------------------------

### Run All Tests and Update Snapshots

Source: https://github.com/chainner-org/spandrel/blob/main/CONTRIBUTING.md

Execute all tests and update any snapshots that have changed. Use this when you have intentionally modified expected outputs.

```bash
pytest tests --snapshot-update
```

--------------------------------

### Write Test for `load` Function

Source: https://github.com/chainner-org/spandrel/blob/main/CONTRIBUTING.md

Create a test file (`tests/test_ARCH.py`) to verify the correctness of your `load` function using the `assert_loads_correctly` utility. This function saves, loads, and compares models.

```python
from spandrel.architectures.ARCH_NAME import ARCH_NAME, load

from .util import assert_loads_correctly


def test_load():
    assert_loads_correctly(
        load,
        lambda: ARCH_NAME(),
        condition=lambda a, b: (
            a.param1 == b.param1
            and a.param2 == b.param2
        ),
    )
```

--------------------------------

### Load Model with Detected Input Channels in Python

Source: https://github.com/chainner-org/spandrel/blob/main/CONTRIBUTING.md

Loads a model by dynamically detecting the input channels from the state dictionary's SFT weight shape.

```python
def load(state_dict: StateDict) -> ImageModelDescriptor[DITN]:
    # default values
    inp_channels = 3
    dim = 60
    # ...

    inp_channels = state_dict["sft.weight"].shape[1]

    model = DITN(
        inp_channels=inp_channels,
        dim=dim,
        # ...
    )
```

--------------------------------

### Initialize UFONE Blocks in Python

Source: https://github.com/chainner-org/spandrel/blob/main/CONTRIBUTING.md

Initializes a sequence of UFONE blocks for a model. The number of blocks is determined by UFONE_blocks.

```python
UFONE_body = [
    UFONE(
        dim,
        ffn_expansion_factor,
        bias,
        LayerNorm_type,
        ITL_blocks,
        SAL_blocks,
        patch_size,
    )
    for _ in range(UFONE_blocks)
]
self.UFONE = nn.Sequential(*UFONE_body)
```

--------------------------------

### Initialize Spandrel with Extra Architectures

Source: https://github.com/chainner-org/spandrel/blob/main/libs/spandrel_extra_arches/README.md

Before using ModelLoader, ensure extra architectures are registered by calling spandrel_extra_arches.install(). This makes additional model architectures available for loading.

```python
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 model
```

--------------------------------

### Inspect Model State Dictionary with dump_dummy.py

Source: https://github.com/chainner-org/spandrel/blob/main/CONTRIBUTING.md

This script generates a YAML file representing a model's state dictionary, useful for understanding parameter structure and shapes. Ensure the `create_dummy` function in the script returns a model of your target architecture.

```yaml
# DITN.DITN()
sft
  sft.weight: Tensor float32 Size([60, 3, 3, 3])
  sft.bias:   Tensor float32 Size([60])
UFONE.0
  ITLs
    0
      attn
        temperature
          UFONE.0.ITLs.0.attn.temperature: Tensor float32 Size([1, 1, 1])
        qkv
          UFONE.0.ITLs.0.attn.qkv.weight: Tensor float32 Size([180, 60])
          UFONE.0.ITLs.0.attn.qkv.bias:   Tensor float32 Size([180])
        project_out
          UFONE.0.ITLs.0.attn.project_out.weight: Tensor float32 Size([60, 60, 1, 1])
      conv1
        UFONE.0.ITLs.0.conv1.weight: Tensor float32 Size([60, 60, 1, 1])
        UFONE.0.ITLs.0.conv1.bias:   Tensor float32 Size([60])
      conv2
        UFONE.0.ITLs.0.conv2.weight: Tensor float32 Size([60, 60, 1, 1])
        UFONE.0.ITLs.0.conv2.bias:   Tensor float32 Size([60])
      ffn
        UFONE.0.ITLs.0.ffn.project_in.weight:  Tensor float32 Size([240, 60, 1, 1])
        UFONE.0.ITLs.0.ffn.dwconv.weight:      Tensor float32 Size([240, 1, 3, 3])
        UFONE.0.ITLs.0.ffn.project_out.weight: Tensor float32 Size([60, 120, 1, 1])
    1
      ...
    2
      ...
    3
      ...
  SALs
    ...
conv_after_body
  conv_after_body.weight: Tensor float32 Size([60, 60, 3, 3])
  conv_after_body.bias:   Tensor float32 Size([60])
upsample.0
  upsample.0.weight: Tensor float32 Size([48, 60, 3, 3])
  upsample.0.bias:   Tensor float32 Size([48])

```

--------------------------------

### Register DITN Architecture Support

Source: https://github.com/chainner-org/spandrel/blob/main/CONTRIBUTING.md

This snippet shows how to register a new architecture, 'DITN', with spandrel. It defines how to detect the architecture from a state dict using a list of keys and specifies the loading function.

```python
ArchSupport(
    id="DITN",
    detect=_has_keys(
        "sft.weight",
        "UFONE.0.ITLs.0.attn.temperature",
        "UFONE.0.ITLs.0.ffn.project_in.weight",
        "UFONE.0.ITLs.0.ffn.dwconv.weight",
        "UFONE.0.ITLs.0.ffn.project_out.weight",
        "conv_after_body.weight",
        "upsample.0.weight",
    ),
    load=DITN.load,
)
```

--------------------------------

### List Main Architectures

Source: https://context7.com/chainner-org/spandrel/llms.txt

List the main architectures available in the Spandrel library, which typically have permissive licenses (MIT/Apache/public domain).

```python
from spandrel import MAIN_REGISTRY

print("Main architectures:")
for arch in MAIN_REGISTRY.architectures():
    print(f"  - {arch.architecture.id}")
```

--------------------------------

### Dump Actual Model State Dict

Source: https://github.com/chainner-org/spandrel/blob/main/CLAUDE.md

Generate a state dictionary from a specified model file to `dump.yml`. Use this to inspect the structure of actual pretrained models.

```bash
python scripts/dump_state_dict.py /path/to.pth
```

--------------------------------

### Dump Dummy Model State Dict

Source: https://github.com/chainner-org/spandrel/blob/main/CLAUDE.md

Generate a dummy model state dictionary to `dump.yml`. Useful for understanding the expected structure without needing a real model.

```bash
python scripts/dump_dummy.py
```

--------------------------------

### Format Code

Source: https://github.com/chainner-org/spandrel/blob/main/CLAUDE.md

Apply code formatting rules using Ruff to maintain a consistent style across the project.

```bash
ruff format libs tests
```

--------------------------------

### Implement `load` Method for Architecture

Source: https://github.com/chainner-org/spandrel/blob/main/CONTRIBUTING.md

Define the `load` function within your architecture's `__init__.py` file. This function initializes a model with default or detected hyperparameters and wraps it in an `ImageModelDescriptor`.

```python
from ...__helpers.model_descriptor import ImageModelDescriptor, StateDict
from .__arch.ARCH_NAME import ARCH_NAME


def load(state_dict: StateDict) -> ImageModelDescriptor[ARCH_NAME]:
    # default values
    param1 = 1
    param2 = 2

    model = ARCH_NAME(
        param1=param1,
        param2=param2,
    )

    return ImageModelDescriptor(
        model,
        state_dict,
        architecture="ARCH_NAME",
        purpose="SR",
        tags=[],
        supports_half=True,
        supports_bfloat16=True,
        scale=1,  # TODO: fix me
        input_channels=3,  # TODO: fix me
        output_channels=3,  # TODO: fix me
    )
```

--------------------------------

### Run Super Resolution with ImageModelDescriptor

Source: https://context7.com/chainner-org/spandrel/llms.txt

Wrap image-to-image models using `ImageModelDescriptor` for automatic input padding and output clamping. Ensure the model is moved to the desired device and set to evaluation mode before inference. Input tensors should be in the range [0, 1].

```python
from spandrel import ImageModelDescriptor, ModelLoader
import torch

# Load a super resolution model
model = ModelLoader().load_from_file("RealESRGAN_x4plus.pth")
assert isinstance(model, ImageModelDescriptor)

# Move to GPU and set to evaluation mode
model.cuda().eval()
# Or chain: model.to("cuda").eval()
# Or with dtype: model.to("cuda", torch.float16)

# Prepare input tensor: (batch=1, channels, height, width), range [0, 1]
input_tensor = torch.rand(1, 3, 64, 64, device="cuda")

# Run inference - handles padding automatically
with torch.no_grad():
    output = model(input_tensor)  # Returns tensor in range [0, 1]

# Output shape depends on scale
print(f"Input: {input_tensor.shape}")   # torch.Size([1, 3, 64, 64])
print(f"Output: {output.shape}")        # torch.Size([1, 3, 256, 256]) for 4x scale

# Check size requirements if doing manual tiling
req = model.size_requirements
print(f"Minimum size: {req.minimum}")
print(f"Multiple of: {req.multiple_of}")
print(f"Must be square: {req.square}")
print(f"Input satisfies requirements: {req.check(64, 64)}")
padding = req.get_padding(64, 64)  # Returns (pad_w, pad_h) needed
```

--------------------------------

### BatchNorm2d with track_running_stats=True

Source: https://github.com/chainner-org/spandrel/blob/main/TENSOR_SHAPE_REFERENCE.md

Defines BatchNorm2d with tracking enabled. Parameters like weight, bias, running_mean, running_var, and num_batches_tracked are stored in the state dict.

```python
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([])
```

--------------------------------

### Run Inference Tests with MPS Device

Source: https://github.com/chainner-org/spandrel/blob/main/CONTRIBUTING.md

Set environment variables to run inference tests using the MPS (Metal Performance Shaders) device on Apple Silicon. This helps in comparing outputs across different devices.

```bash
env PYTORCH_ENABLE_MPS_FALLBACK=1 SPANDREL_TEST_DEVICE=mps SPANDREL_TEST_OUTPUTS_DIR=outputs-mps pytest --snapshot-update
```

--------------------------------

### Test Model Loading Correctness in Python

Source: https://github.com/chainner-org/spandrel/blob/main/CONTRIBUTING.md

Asserts that the model loading function correctly initializes models with various input channel configurations.

```python
assert_loads_correctly(
    load,
    lambda: DITN(),
    lambda: DITN(inp_channels=4),
    lambda: DITN(inp_channels=1),
    condition=lambda a, b: (...),
)
```

--------------------------------

### Update Snapshots for Specific Architecture

Source: https://github.com/chainner-org/spandrel/blob/main/CONTRIBUTING.md

Run tests for a specific architecture and update snapshots. Use this command when tests fail due to snapshot changes.

```bash
pytest tests/test_<arch>.py --snapshot-update
```

--------------------------------

### Manage Device and Precision for Models

Source: https://context7.com/chainner-org/spandrel/llms.txt

Move models between devices (CPU/GPU) and change their data type precision (fp32, fp16, bfloat16) using convenient methods. Check for support before attempting unsupported conversions.

```python
from spandrel import ModelLoader, UnsupportedDtypeError
import torch

model = ModelLoader().load_from_file("model.pth")

# Move to GPU
model.cuda()        # Move to default CUDA device
model.cuda(0)       # Move to specific GPU
model.cpu()         # Move back to CPU

# Change precision
model.float()       # fp32 (default)
model.half()        # fp16 - throws UnsupportedDtypeError if not supported
model.bfloat16()    # bfloat16 - throws UnsupportedDtypeError if not supported

# Check support before conversion
if model.supports_half:
    model.half()
if model.supports_bfloat16:
    model.bfloat16()

# Combined device and dtype
model.to("cuda", torch.float16)
model.to(device="cuda:0", dtype=torch.bfloat16)

```

--------------------------------

### DITN Model Initialization Snippet

Source: https://github.com/chainner-org/spandrel/blob/main/CONTRIBUTING.md

This Python code snippet shows a section of the `DITN` class constructor, illustrating how convolutional layers are initialized with specific input/output channels, kernel size, stride, and padding.

```python
self.sft = nn.Conv2d(inp_channels, dim, 3, 1, 1)

```

--------------------------------

### Calculating ffn_expansion_factor from State Dict

Source: https://github.com/chainner-org/spandrel/blob/main/CONTRIBUTING.md

This Python code demonstrates how to calculate the `ffn_expansion_factor` by using the shapes of weight tensors from the state dictionary and the known `dim` value. This method is useful for detecting parameters when their direct values are not explicitly stored.

```python
hidden_features = state_dict["UFONE.0.ITLs.0.ffn.project_in.weight"].shape[0] / 2
ffn_expansion_factor = hidden_features / dim
```

--------------------------------

### Access Model Weight Shape in Python

Source: https://github.com/chainner-org/spandrel/blob/main/CONTRIBUTING.md

Demonstrates direct access to a model's field, specifically the weight shape of the SFT layer.

```python
model = DITN()
print(model.sft.weight.shape)
```

--------------------------------

### Run Tests for Specific Architecture

Source: https://github.com/chainner-org/spandrel/blob/main/CLAUDE.md

Execute tests for a single architecture, useful for targeted debugging or development.

```bash
pytest tests/test_ESRGAN.py
```

--------------------------------

### Type Checking

Source: https://github.com/chainner-org/spandrel/blob/main/CLAUDE.md

Perform static type checking on the library and test code using Pyright to catch type-related errors.

```bash
pyright libs tests
```

--------------------------------

### Run Specific Architecture Tests

Source: https://github.com/chainner-org/spandrel/blob/main/CONTRIBUTING.md

Execute tests for a particular architecture using pytest. Replace <arch> with the architecture name (e.g., ESRGAN).

```bash
pytest tests/test_<arch>.py
```

--------------------------------

### Manage Architectures with ArchRegistry

Source: https://context7.com/chainner-org/spandrel/llms.txt

Use ArchRegistry to manage available architectures. You can view, check, and create custom registries for ModelLoader. Registries can be copied and modified.

```python
from spandrel import MAIN_REGISTRY, ArchRegistry, ArchSupport, ModelLoader

# View all registered architectures
for arch in MAIN_REGISTRY:
    print(f"{arch.architecture.id}: {arch.architecture.name}")

# Check if an architecture is registered
if "ESRGAN" in MAIN_REGISTRY:
    esrgan_support = MAIN_REGISTRY["ESRGAN"]
    print(f"ESRGAN registered: {esrgan_support.architecture.name}")

# Create a custom registry with specific architectures
custom_registry = ArchRegistry()
custom_registry.add(
    MAIN_REGISTRY["ESRGAN"],
    MAIN_REGISTRY["SwinIR"],
)

# Use custom registry with ModelLoader
loader = ModelLoader(registry=custom_registry)
model = loader.load_from_file("model.pth")  # Only detects ESRGAN or SwinIR

# Copy and modify registry
my_registry = MAIN_REGISTRY.copy()
# Add architectures with ordering (detected before specified arch)
# my_registry.add(ArchSupport.from_architecture(MyArch(), before=(ArchId("ESRGAN"),)))

```

--------------------------------

### Test Model Loading with Assertions

Source: https://github.com/chainner-org/spandrel/blob/main/CONTRIBUTING.md

Use this function to test if a model's parameters are loaded correctly by comparing fields between the original and loaded models. It requires a load function, a lambda to create the original model, and an optional condition function for field comparison.

```python
from spandrel.architectures.DITN import DITN, load

from .util import assert_loads_correctly


def test_load():
    assert_loads_correctly(
        load,
        lambda: DITN(),
        condition=lambda a, b: (
            a.patch_size == b.patch_size
            and a.dim == b.dim
            and a.scale == b.scale
            and a.SAL_blocks == b.SAL_blocks
            and a.ITL_blocks == b.ITL_blocks
        ),
    )
```

--------------------------------

### Load PyTorch Models with ModelLoader

Source: https://context7.com/chainner-org/spandrel/llms.txt

Use `ModelLoader` to automatically detect and load PyTorch models from various file formats. Specify the device (e.g., 'cuda' or 'cpu') during initialization. You can also load just the state dictionary for manual processing.

```python
from spandrel import ModelLoader, ImageModelDescriptor
import torch

# Create a loader (optionally specify device)
loader = ModelLoader(device="cuda")  # or "cpu", torch.device("cuda:0")

# Load a model from file - architecture is auto-detected
model = loader.load_from_file("path/to/model.pth")

# Check the model type and access metadata
print(f"Architecture: {model.architecture.id}")  # e.g., "ESRGAN", "SwinIR"
print(f"Scale: {model.scale}")                   # e.g., 4 for 4x upscaling
print(f"Input channels: {model.input_channels}") # e.g., 3 for RGB
print(f"Output channels: {model.output_channels}")
print(f"Purpose: {model.purpose}")               # "SR", "FaceSR", "Restoration", "Inpainting"
print(f"Tags: {model.tags}")                     # e.g., ["64nf", "23nb"]
print(f"Supports half: {model.supports_half}")   # fp16 support
print(f"Supports bfloat16: {model.supports_bfloat16}")

# Alternative: Load just the state dict for manual processing
state_dict = loader.load_state_dict_from_file("path/to/model.pth")
model = loader.load_from_state_dict(state_dict)
```

--------------------------------

### Pyright Type Checking (Specific Directory)

Source: https://github.com/chainner-org/spandrel/blob/main/CONTRIBUTING.md

Check for type errors specifically within the 'libs' directory using Pyright.

```bash
pyright libs
```

--------------------------------

### Handle Model Input Constraints with SizeRequirements

Source: https://context7.com/chainner-org/spandrel/llms.txt

Use SizeRequirements to define and check input size constraints for models. It handles minimum size, divisibility, and square requirements. Padding can be calculated manually.

```python
from spandrel import SizeRequirements, ModelLoader

# Load model and access its size requirements
model = ModelLoader().load_from_file("model.pth")
req = model.size_requirements

# Check requirements
print(f"Minimum: {req.minimum}")      # Minimum input size (0 = no minimum)
print(f"Multiple of: {req.multiple_of}")  # Size must be divisible by this
print(f"Square: {req.square}")        # True if input must be square

# Check if a size satisfies requirements
if req.check(width=256, height=256):
    print("256x256 is valid")

# Get padding needed for a given size
pad_w, pad_h = req.get_padding(width=100, height=100)
print(f"Need to pad by: {pad_w}x{pad_h}")

# Create custom requirements for manual use
custom_req = SizeRequirements(
    minimum=64,      # At least 64x64
    multiple_of=16,  # Divisible by 16
    square=False     # Non-square allowed
)
```

--------------------------------

### Dump State Dict Contents

Source: https://github.com/chainner-org/spandrel/blob/main/CONTRIBUTING.md

Use this script to dump the contents of a PyTorch state dictionary to a YAML file. This is useful for inspecting model weights and potentially adding new architectures.

```python
python scripts/dump_state_dict.py /path/to/model.pth
```

--------------------------------

### BatchNorm2d with track_running_stats=False

Source: https://github.com/chainner-org/spandrel/blob/main/TENSOR_SHAPE_REFERENCE.md

Defines BatchNorm2d without tracking running statistics. Only weight and bias are stored in the state dict.

```python
p = nn.BatchNorm2d(num_features=N, affine=True, track_running_stats=False)

# p.weight: Tensor Size([N])
# p.bias:   Tensor Size([N])
```

--------------------------------

### Test Architecture Model Inference

Source: https://github.com/chainner-org/spandrel/blob/main/CONTRIBUTING.md

This test function verifies model loading and image inference for a specific architecture using official models. It requires a model file URL and asserts correct model instantiation and inference results.

```python
from spandrel.architectures.ARCH import ARCH, load

from .util import (
    ModelFile,
    TestImage,
    assert_image_inference,
    assert_loads_correctly,
    disallowed_props,
    skip_if_unchanged,
)

skip_if_unchanged(__file__)


def test_load():
    ...


def test_ARCH_model_name(snapshot):
    file = ModelFile.from_url(
        "https://example.com/path/to/model_name.pth"
    )
    model = file.load_model()
    assert model == snapshot(exclude=disallowed_props)
    assert isinstance(model.model, ARCH)
    assert_image_inference(
        file,
        model,
        [TestImage.SR_16, TestImage.SR_32, TestImage.SR_64],
    )
```

--------------------------------

### BatchNorm2d with affine=False

Source: https://github.com/chainner-org/spandrel/blob/main/TENSOR_SHAPE_REFERENCE.md

Defines BatchNorm2d without affine parameters. No tensors are stored in the state dict.

```python
p = nn.BatchNorm2d(num_features=N, affine=False, track_running_stats=False)

# nothing is stored in state dict
```

--------------------------------

### MultiheadAttention with bias=True and add_bias_kv=True

Source: https://github.com/chainner-org/spandrel/blob/main/TENSOR_SHAPE_REFERENCE.md

Defines a MultiheadAttention layer with bias and bias for key/value projections. embed_dim must be divisible by num_heads. Stores multiple weight and bias tensors.

```python
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])
```

--------------------------------

### Lint and Auto-fix Code

Source: https://github.com/chainner-org/spandrel/blob/main/CLAUDE.md

Check code for style and potential errors using Ruff and automatically fix violations where possible. Ensures code consistency.

```bash
ruff check libs tests --fix --unsafe-fixes
```

--------------------------------

### MultiheadAttention with custom kdim/vdim, bias=True, add_bias_kv=True

Source: https://github.com/chainner-org/spandrel/blob/main/TENSOR_SHAPE_REFERENCE.md

Defines MultiheadAttention with custom key/value dimensions, bias, and bias for key/value projections. embed_dim must be divisible by num_heads. Stores multiple weight and bias tensors.

```python
assert D % H == 0
p = nn.MultiheadAttention(embed_dim=D, num_heads=H, bias=True, add_bias_kv=True, kdim=K, vdim=V)

# p.q_proj_weight:   Tensor Size([D, D])
# p.k_proj_weight:   Tensor Size([D, K])
# p.v_proj_weight:   Tensor Size([D, V])
# p.in_proj_bias:    Tensor Size([3*D])
# p.bias_k:          Tensor Size([1, 1, D])
# p.bias_v:          Tensor Size([1, 1, D])
# p.out_proj.weight: Tensor Size([D, D])
# p.out_proj.bias:   Tensor Size([D])
```

--------------------------------

### Determine Sequence Length from State Dict in Python

Source: https://github.com/chainner-org/spandrel/blob/main/CONTRIBUTING.md

Illustrates how to use a helper function 'get_seq_len' to determine the length of a sequence within a model's state dictionary, often indicated by numbered groupings.

```python
def load(state_dict: StateDict) -> ImageModelDescriptor[DITN]:
    # ...

    # To get the length of UFONE_blocks, we use get_seq_len
    UFONE_blocks = get_seq_len(state_dict, "UFONE")

    # ... rest of the load function
```

--------------------------------

### MultiheadAttention with custom kdim/vdim, bias=False, add_bias_kv=False

Source: https://github.com/chainner-org/spandrel/blob/main/TENSOR_SHAPE_REFERENCE.md

Defines MultiheadAttention with custom key/value dimensions, without bias, and without bias for key/value projections. embed_dim must be divisible by num_heads. Stores only weight tensors.

```python
assert D % H == 0
p = nn.MultiheadAttention(embed_dim=D, num_heads=H, bias=False, add_bias_kv=False, kdim=K, vdim=V)

# p.q_proj_weight:   Tensor Size([D, D])
# p.k_proj_weight:   Tensor Size([D, K])
# p.v_proj_weight:   Tensor Size([D, V])
# p.out_proj.weight: Tensor Size([D, D])
```

--------------------------------

### MultiheadAttention with custom kdim/vdim, bias=True, add_bias_kv=False

Source: https://github.com/chainner-org/spandrel/blob/main/TENSOR_SHAPE_REFERENCE.md

Defines MultiheadAttention with custom key/value dimensions and bias, but without bias for key/value projections. embed_dim must be divisible by num_heads. Stores weight and bias tensors.

```python
assert D % H == 0
p = nn.MultiheadAttention(embed_dim=D, num_heads=H, bias=True, add_bias_kv=False, kdim=K, vdim=V)

# p.q_proj_weight:   Tensor Size([D, D])
# p.k_proj_weight:   Tensor Size([D, K])
# p.v_proj_weight:   Tensor Size([D, V])
# p.in_proj_bias:    Tensor Size([3*D])
# p.out_proj.weight: Tensor Size([D, D])
# p.out_proj.bias:   Tensor Size([D])
```

--------------------------------

### Set Model to Evaluation Mode

Source: https://context7.com/chainner-org/spandrel/llms.txt

Set the model to evaluation mode, which is recommended for inference. This ensures layers like dropout and batch normalization behave correctly during inference.

```python
model.eval()
```

--------------------------------

### Run Inpainting with MaskedImageModelDescriptor

Source: https://context7.com/chainner-org/spandrel/llms.txt

Utilize `MaskedImageModelDescriptor` for inpainting models that require both an image and a mask. The mask indicates regions to be inpainted (1 for inpaint, 0 for keep). Input and output tensors are in the range [0, 1].

```python
from spandrel import MaskedImageModelDescriptor, ModelLoader
import torch

# Load an inpainting model (e.g., LaMa, MAT)
model = ModelLoader().load_from_file("big-lama.pt")
assert isinstance(model, MaskedImageModelDescriptor)

model.cuda().eval()

# Prepare inputs
height, width = 512, 512
image = torch.rand(1, 3, height, width, device="cuda")  # RGB image [0, 1]
mask = torch.zeros(1, 1, height, width, device="cuda")  # Mask: 0=keep, 1=inpaint

# Create a mask region to inpaint (e.g., center square)
mask[:, :, 100:200, 100:200] = 1.0

# Run inpainting
with torch.no_grad():
    output = model(image, mask)  # Returns inpainted image [0, 1]

print(f"Output shape: {output.shape}")  # Same as input for inpainting (scale=1)
```

--------------------------------

### Access Underlying PyTorch Module

Source: https://context7.com/chainner-org/spandrel/llms.txt

Access the underlying PyTorch module directly from the Spandrel model object. This allows for direct interaction with PyTorch functionalities if needed.

```python
pytorch_model = model.model  # torch.nn.Module
print(f"Model type: {type(pytorch_model)}")
print(f"Device: {model.device}")
print(f"Dtype: {model.dtype}")
```

--------------------------------

### MultiheadAttention with bias=False

Source: https://github.com/chainner-org/spandrel/blob/main/TENSOR_SHAPE_REFERENCE.md

Defines a MultiheadAttention layer without bias. embed_dim must be divisible by num_heads. Stores only weight tensors.

```python
assert D % H == 0
p = nn.MultiheadAttention(embed_dim=D, num_heads=H, bias=False)

# p.in_proj_weight:  Tensor Size([3*D, D])
# p.out_proj.weight: Tensor Size([D, D])
```

--------------------------------

### Normalize State Dictionaries

Source: https://context7.com/chainner-org/spandrel/llms.txt

Normalize state dictionaries using canonicalize_state_dict to handle nested structures and common prefixes like 'module.' or 'netG.'. This prepares the dictionary for direct use with an architecture's load method.

```python
from spandrel import canonicalize_state_dict
import torch

# State dicts may come wrapped in various formats
raw_state_dict = torch.load("model.pth", map_location="cpu")

# Common wrapper formats that are automatically handled:
# {"model_state_dict": {...}}, {"state_dict": {...}}, {"params_ema": {...}},
# {"params": {...}}, {"model": {...}}, {"net": {...}}

# Canonicalize removes wrappers and common prefixes like "module.", "netG."
clean_state_dict = canonicalize_state_dict(raw_state_dict)

# Now use with architecture's load method directly
from spandrel.architectures.ESRGAN import ESRGANArch
arch = ESRGANArch()
if arch.detect(clean_state_dict):
    model = arch.load(clean_state_dict)
```

--------------------------------

### Linear layer with bias=True

Source: https://github.com/chainner-org/spandrel/blob/main/TENSOR_SHAPE_REFERENCE.md

Defines a linear (fully connected) layer with bias. Stores weight and bias tensors.

```python
p = nn.Linear(in_features=I, out_features=O, bias=True)

# p.weight: Tensor Size([O, I])
# p.bias:   Tensor Size([O])
```

--------------------------------

### Conv2d with bias=True

Source: https://github.com/chainner-org/spandrel/blob/main/TENSOR_SHAPE_REFERENCE.md

Defines a 2D convolution layer with bias. Stores weight and bias tensors in the state dict.

```python
p = nn.Conv2d(in_channels=I, out_channels=O, kernel_size=K, bias=True)

# p.weight: Tensor Size([O, I, K, K])
# p.bias:   Tensor Size([O])
```

--------------------------------

### Conv2d with non-square kernel and bias=True

Source: https://github.com/chainner-org/spandrel/blob/main/TENSOR_SHAPE_REFERENCE.md

Defines a 2D convolution layer with a non-square kernel and bias. Stores weight and bias tensors.

```python
p = nn.Conv2d(in_channels=I, out_channels=O, kernel_size=(K1, K2), bias=True)

# p.weight: Tensor Size([O, I, K1, K2])
# p.bias:   Tensor Size([O])
```

--------------------------------

### Linear layer with bias=False

Source: https://github.com/chainner-org/spandrel/blob/main/TENSOR_SHAPE_REFERENCE.md

Defines a linear layer without bias. Only the weight tensor is stored.

```python
p = p = nn.Linear(in_features=I, out_features=O, bias=False)

# p.weight: Tensor Size([O, I])
```

--------------------------------

### FeedForward Network Class Definition

Source: https://github.com/chainner-org/spandrel/blob/main/CONTRIBUTING.md

This Python class defines a FeedForward network, where `ffn_expansion_factor` directly influences the `hidden_features` calculation, impacting the dimensions of the `project_in`, `dwconv`, and `project_out` layers.

```python
class FeedForward(nn.Module):
    def __init__(self, dim, ffn_expansion_factor, bias):
        super().__init__()
        hidden_features = int(dim * ffn_expansion_factor)
        self.project_in = nn.Conv2d(dim, hidden_features * 2, kernel_size=1, bias=bias)

        self.dwconv = nn.Conv2d(
            hidden_features * 2,
            hidden_features * 2,
            kernel_size=3,
            stride=1,
            padding=1,
            groups=hidden_features * 2,
            bias=bias,
        )

        self.project_out = nn.Conv2d(hidden_features, dim, kernel_size=1, bias=bias)

    def forward(self, x):
        ... 
```

--------------------------------

### Define SFT Layer in Python

Source: https://github.com/chainner-org/spandrel/blob/main/CONTRIBUTING.md

Defines the SFT (Spatial Feature Transformation) layer, typically a 2D convolution, used in image models.

```python
self.sft = nn.Conv2d(inp_channels, dim, 3, 1, 1)
```

--------------------------------

### Conv2d with bias=False

Source: https://github.com/chainner-org/spandrel/blob/main/TENSOR_SHAPE_REFERENCE.md

Defines a 2D convolution layer without bias. Only the weight tensor is stored in the state dict.

```python
p = nn.Conv2d(in_channels=I, out_channels=O, kernel_size=K, bias=False)

# p.weight: Tensor Size([O, I, K, K])
```

--------------------------------

### Conv2d with groups and bias=True

Source: https://github.com/chainner-org/spandrel/blob/main/TENSOR_SHAPE_REFERENCE.md

Defines a grouped 2D convolution layer. The input and output channels must be divisible by the group count. Stores weight and bias tensors.

```python
assert I % G == 0 and O % G == 0
p = nn.Conv2d(in_channels=I, out_channels=O, kernel_size=(K1, K2), group=G, bias=True)

# p.weight: Tensor Size([O, I/G, K1, K2])
# p.bias:   Tensor Size([O])
```

--------------------------------

### Python Function for Sequence Length

Source: https://github.com/chainner-org/spandrel/blob/main/CONTRIBUTING.md

This function retrieves the sequence length from a state dictionary, likely used for processing sequential data within the model.

```python
from ..__arch_helpers.state import get_seq_len

def load(state_dict: StateDict) -> ImageModelDescriptor[DITN]:
    # ...
    UFONE_blocks = get_seq_len(state_dict, "UFONE")
```