| 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)
|
| _, _, 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()
|
| )
|
|
|
