spandrel_autodetect_model_architecture_and_image_inference.py

1import torch
2from spandrel import ModelLoader
3from PIL import Image
4import torchvision.transforms.functional as F
5
6# 1. Load the model
7# Spandrel will automatically detect the architecture and configuration
8model_path = "path/to/your/model.pth"
9loader = ModelLoader()
10model = loader.load_from_file(model_path)
11
12# 2. Prepare the model for inference
13model.eval()
14if torch.cuda.is_available():
15    model.cuda()
16
17# 3. Load and process an image
18img = Image.open("path/to/image.png").convert("RGB")
19# Convert to tensor and add batch dimension (C, H, W) -> (1, C, H, W)
20input_tensor = F.to_tensor(img).unsqueeze(0)
21
22# Move to same device as model
23if torch.cuda.is_available():
24    input_tensor = input_tensor.cuda()
25
26# 4. Run the model
27with torch.no_grad():
28    output_tensor = model(input_tensor)
29
30# 5. Convert back to image
31# Remove batch dimension and convert to PIL
32output_img = F.to_pil_image(output_tensor.squeeze(0).clamp(0, 1))
33output_img.save("output.png")