AIupscale_run.py

import torch
import torch.nn as nn
import torch.optim as optim
import coremltools as ct

from torch.utils.data import Dataset, DataLoader
from torchvision import transforms
import os
from PIL import Image
import torchvision.transforms.functional as TF

device = torch.device("mps")
class UPSC(nn.Module):
    def __init__(self):
        super(UPSC,self).__init__()
        self.model = nn.Sequential(
            nn.Conv2d(in_channels=3, out_channels=64, kernel_size=5, padding=2),
            nn.ReLU(),
            nn.Conv2d(in_channels=64, out_channels=32, kernel_size=3, padding=1),
            nn.ReLU(),
            # This convolution outputs channels that are scale_factor^2 * number_of_channels.
            nn.Conv2d(in_channels=32, out_channels=3 * 3 * 3, kernel_size=3, padding=1),
            # PixelShuffle rearranges channels into spatial dimensions.
            nn.PixelShuffle(3)
        )
    def forward(self, x):
        return self.model(x)

model = UPSC().to(device)
model.load_state_dict(torch.load("upscaling.pth", weights_only=True))
model.eval()

img = Image.open("test.png").convert("RGB")

# Resize it to match what the model expects (e.g. 256x256)
transform = transforms.Compose([
    transforms.Resize((256, 256)),  # match training input size
    transforms.ToTensor()
])

lr_tensor = transform(img).unsqueeze(0).to(device)

with torch.no_grad():
    sr_tensor = model(lr_tensor)
    traced_model = torch.jit.trace(model, lr_tensor)


# Remove batch dimension and convert to PIL
sr_image = TF.to_pil_image(sr_tensor.squeeze(0).clamp(0, 1))
sr_image.save("upscaled_output_5.jpg")

mlmodel = ct.convert(
    traced_model,
    inputs=[ct.ImageType(name="input", shape=lr_tensor.shape)],
    compute_units=ct.ComputeUnit.ALL  # Use ANE if available
)

mlmodel.save("upscaling.mlmodel")