spandrel_extra_arches_model_loader_image_upscale_quickstart.py

1import torch
2from spandrel import ModelLoader
3import spandrel_extra_arches # Must be imported or installed to register extra architectures
4
5def main():
6    # Path to a supported model file (e.g., a .pth or .ckpt file)
7    # The 'spandrel-extra-arches' package allows ModelLoader to recognize
8    # architectures not included in the base spandrel library.
9    model_path = "path/to/your/model.pth"
10
11    # Load the model
12    # ModelLoader automatically detects the architecture and loads weights
13    model = ModelLoader().load_from_file(model_path)
14
15    # Move model to GPU if available and set to evaluation mode
16    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
17    model.to(device)
18    model.eval()
19
20    # Create a dummy input tensor (B, C, H, W)
21    # Most spandrel models expect 3-channel RGB images
22    dummy_input = torch.randn(1, 3, 64, 64).to(device)
23
24    # Perform inference
25    with torch.no_grad():
26        output = model(dummy_input)
27
28    print(f"Model Name: {model.architecture}")
29    print(f"Output Shape: {output.shape}")
30
31if __name__ == "__main__":
32    main()