import torch import torch.nn as nn import torch.nn.functional as F import torchvision.transforms.functional as TF from PIL import Image import math import gradio as gr import os class PatchEmbedding(nn.Module): def __init__(self, in_channels = 3, embed_dim = 64, patch_size = 2, img_size = 64): super().__init__() self.in_channels = in_channels self.embed_dim = embed_dim num_patches = (img_size // patch_size)**2 self.proj = nn.Conv2d(in_channels, embed_dim, kernel_size=patch_size, stride=patch_size) self.pos_embed = nn.Parameter(torch.zeros(1, num_patches, embed_dim)) def forward(self, x): x = self.proj(x) x = x.flatten(2) x = x.transpose(1,2) x = x + self.pos_embed return x class MultiHeadSelfAttention(nn.Module): def __init__(self, embed_dim = 64, num_heads = 4): super().__init__() assert embed_dim % num_heads == 0, f"embed_dim {embed_dim} must be divisible by number of heads {num_heads}" self.embed_dim = embed_dim self.num_heads = num_heads self.head_dim = embed_dim // num_heads self.q_proj = nn.Linear(embed_dim, embed_dim) self.k_proj = nn.Linear(embed_dim, embed_dim) self.v_proj = nn.Linear(embed_dim, embed_dim) def forward(self, x): q = self.q_proj(x) k = self.k_proj(x) v = self.v_proj(x) B,N,D = q.shape q = q.reshape(B, N, self.num_heads, self.head_dim) k = k.reshape(B, N, self.num_heads, self.head_dim) v = v.reshape(B, N, self.num_heads, self.head_dim) q = q.transpose(1,2) k = k.transpose(1,2) v = v.transpose(1,2) scores = q @ k.transpose(2,3) scores = scores / (math.sqrt(self.head_dim)) weights = F.softmax(scores, dim = -1) output = weights @ v output = output.transpose(1,2) output = output.contiguous() output = output.reshape(B,N, self.embed_dim) return output class TransformerBlock(nn.Module): def __init__(self, embed_dim = 64,num_heads = 4): super().__init__() self.embed_dim = embed_dim hidden_dim = 4 * embed_dim self.num_heads = num_heads self.norm1 = nn.LayerNorm(embed_dim) self.norm2 = nn.LayerNorm(embed_dim) self.attn = MultiHeadSelfAttention(embed_dim, num_heads) self.mlp = nn.Sequential( nn.Linear(embed_dim, hidden_dim), nn.GELU(), nn.Linear(hidden_dim, embed_dim), ) def forward(self, x): x = x + self.attn(self.norm1(x)) x = x + self.mlp(self.norm2(x)) return x class ImageSRTransformer(nn.Module): def __init__(self, embed_dim = 64, num_heads = 4, depth = 6, patch_size = 2, img_size = 64): super().__init__() self.embed_dim = embed_dim self.num_heads = num_heads self.depth = depth self.grid_size = img_size // patch_size self.patch_size = patch_size self.img_size = img_size self.patch_embed = PatchEmbedding(embed_dim = embed_dim , patch_size = patch_size, img_size = img_size) self.blocks = nn.ModuleList([ TransformerBlock(embed_dim, num_heads) for _ in range(depth) ]) self.head = nn.Sequential( nn.Conv2d(embed_dim, embed_dim * 4, 3, padding=1), nn.PixelShuffle(2), nn.Conv2d(embed_dim, embed_dim * 4, 3, padding=1), nn.PixelShuffle(2), nn.Conv2d(embed_dim, embed_dim * 4, 3, padding=1), nn.PixelShuffle(2), nn.Conv2d(embed_dim, 3, 3, padding=1) ) def forward(self, x): x = self.patch_embed(x) for block in self.blocks: x = block(x) B, N, D = x.shape x = x.transpose(1,2) x = x.contiguous() x = x.reshape(B,D, self.grid_size, self.grid_size) x = self.head(x) return x device = "cuda" if torch.cuda.is_available() else "cpu" model = ImageSRTransformer().to(device) checkpoint = torch.load("sr_best.pt", map_location=device) model.load_state_dict(checkpoint["model_state_dict"]) model.eval() def upscale(pil_img, tile=64, scale=4): if pil_img is None: return None if pil_img.mode != "RGB": pil_img = pil_img.convert("RGB") x = TF.to_tensor(pil_img).unsqueeze(0).to(device) # [1,3,H,W] _, _, H, W = x.shape pad_h = (tile - H % tile) % tile pad_w = (tile - W % tile) % tile x = F.pad(x, (0, pad_w, 0, pad_h), mode="reflect") _, _, Hp, Wp = x.shape out = torch.zeros(1, 3, Hp * scale, Wp * scale, device=device) with torch.no_grad(): for i in range(0, Hp, tile): for j in range(0, Wp, tile): tile_in = x[:, :, i:i+tile, j:j+tile] if device == "cuda": with torch.autocast(device_type="cuda", dtype=torch.float16): tile_out = model(tile_in) else: tile_out = model(tile_in) out[:, :, i*scale:(i+tile)*scale, j*scale:(j+tile)*scale] = tile_out.float() out = out[:, :, :H*scale, :W*scale] out = out.clamp(0, 1).squeeze(0).cpu() return TF.to_pil_image(out) with gr.Blocks(title="ViT-SR") as demo: gr.Markdown(""" # ViT-SR: Vision Transformer for ×4 Super-Resolution **Built from scratch in PyTorch** — no pretrained weights, no existing repos. Trained on LSDIR (76,716 images) · ~786K params · Test PSNR: 23.30 dB """) with gr.Row(): with gr.Column(): input_img = gr.Image(type="pil", label="Low-Resolution Input") submit_btn = gr.Button("Super-Resolve ×4", variant="primary") with gr.Column(): output_img = gr.Image(type="pil", label="Super-Resolved Output") gr.Markdown(""" ### How it works The model uses a Vision Transformer architecture: images are split into 2×2 patches, embedded into 64-dim tokens, processed through 6 transformer blocks with 4-head attention, then reconstructed at 4× resolution via PixelShuffle upsampling. ### Notes - Upload a **low-resolution image** (ideally a bicubic-downscaled version of a sharp photo) - The model tiles large inputs into 64×64 patches and stitches the output - This is a from-scratch baseline; a larger model improves PSNR """) gr.Examples( examples=[ ["1.png"], ["2.png"], ["3.png"], ["4.png"], ["5.png"], ], inputs=input_img, outputs=output_img, fn=upscale, cache_examples=False, label="Example inputs (click to try)" ) submit_btn.click(fn=upscale, inputs=input_img, outputs=output_img) if __name__ == "__main__": demo.launch( allowed_paths=["/app"], theme=gr.themes.Soft() )