Spaces:
Sleeping
Sleeping
| 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() | |
| ) | |