import torch
from model import Swin2SR

model_weights = "model-70.pt"
model_params = {
    "upscale": 2,
    "in_chans": 4,
    "img_size": 64,
    "window_size": 16,
    "img_range": 1.,
    "depths": [6, 6, 6, 6],
    "embed_dim": 90,
    "num_heads": [6, 6, 6, 6],
    "mlp_ratio": 2,
    "upsampler": "pixelshuffledirect",
    "resi_connection": "1conv"
}

sr_model = Swin2SR(**model_params)
sr_model.load_state_dict(torch.load(model_weights))