app.py

import torch
import torch.nn as nn
import torch.nn.functional as F
import gradio as gr
from PIL import Image
import torchvision.transforms.functional as TF

# --- 1. MODEL ARCHITECTURE ---
class PureResBlock(nn.Module):
    def __init__(self, channels):
        super(PureResBlock, self).__init__()
        self.conv1 = nn.Conv2d(channels, channels, kernel_size=3, padding=1, padding_mode='replicate')
        self.act = nn.ReLU(inplace=True)
        self.conv2 = nn.Conv2d(channels, channels, kernel_size=3, padding=1, padding_mode='replicate')
        self.res_scale = 1.0 

    def forward(self, x):
        res = self.conv1(x)
        res = self.act(res)
        res = self.conv2(res)
        return x + (res * self.res_scale)

class FastEDSR(nn.Module):
    def __init__(self, scale_factor=2, num_blocks=8, channels=64):
        super(FastEDSR, self).__init__()
        self.scale_factor = scale_factor
        
        self.head = nn.Conv2d(3, channels, kernel_size=3, padding=1, padding_mode='replicate')
        self.body = nn.Sequential(*[PureResBlock(channels) for _ in range(num_blocks)])
        self.tail = nn.Conv2d(channels, channels, kernel_size=3, padding=1, padding_mode='replicate')
        self.sub_pixel = nn.Sequential(
            nn.Conv2d(channels, 3 * (scale_factor ** 2), kernel_size=3, padding=1, padding_mode='replicate'),
            nn.PixelShuffle(scale_factor)
        )

    def forward(self, x):
        base_upscaled = F.interpolate(x, scale_factor=self.scale_factor, mode='bicubic', align_corners=False)
        f0 = self.head(x)
        f_body = self.body(f0)
        f_body = self.tail(f_body)
        f_out = f0 + f_body
        details = self.sub_pixel(f_out)
        return base_upscaled + details

# --- 2. INITIALIZATION ---
device = torch.device('cpu') 
model = FastEDSR(scale_factor=2, num_blocks=8, channels=64)

# Load the weights
model_path = "FastEDSR_x2_31dB.pth" 
model.load_state_dict(torch.load(model_path, map_location=device))
model.eval()

def calc_psnr(pred, target):
    mse = torch.mean((pred - target) ** 2)
    if mse == 0:
        return 100.0
    return 10 * torch.log10(1.0 / mse).item()

# --- 3. INFERENCE FUNCTIONS ---

def standard_upscale(img):
    if img is None: return None, ""
        
    max_input_dim = 2048 
    w, h = img.size
    
    if w > max_input_dim or h > max_input_dim:
        gr.Warning(f"Input image exceeded the 2K ({max_input_dim}px) limit. It has been proportionally downscaled to ensure the 4K output fits in server memory and constraints.")
        scale = max_input_dim / max(w, h)
        w, h = int(w * scale), int(h * scale)
        img = img.resize((w, h), Image.BICUBIC)

    img = img.convert('RGB')
    input_tensor = TF.to_tensor(img).unsqueeze(0).to(device)

    with torch.no_grad():
        output_tensor = model(input_tensor)

    output_tensor = output_tensor.squeeze(0).clamp(0, 1)
    output_img = TF.to_pil_image(output_tensor)
    
    new_w, new_h = output_img.size
    
    details = (
        f"### Resolution Details\n"
        f"- **Before:** {w} x {h} ({w * h:,} pixels)\n\n"
        f"- **After:** {new_w} x {new_h} ({new_w * new_h:,} pixels)"
    )
    
    return output_img, details

def benchmark_upscale(hr_img):
    if hr_img is None: return "", None, None
    
    hr_img = hr_img.convert('RGB')
    w, h = hr_img.size
    
    w = w - (w % 2)
    h = h - (h % 2)
    hr_img = hr_img.crop((0, 0, w, h))
    
    max_hr_dim = 4096 
    if w > max_hr_dim or h > max_hr_dim:
        gr.Warning(f"Ground truth image exceeded the 4K ({max_hr_dim}px) limit. It has been proportionally downscaled.")
        scale = max_hr_dim / max(w, h)
        w, h = int(w * scale), int(h * scale)
        w = w - (w % 2)
        h = h - (h % 2)
        hr_img = hr_img.resize((w, h), Image.BICUBIC)

    lr_w, lr_h = w // 2, h // 2
    lr_img = hr_img.resize((lr_w, lr_h), Image.BICUBIC)
    
    lr_tensor = TF.to_tensor(lr_img).unsqueeze(0).to(device)
    hr_tensor = TF.to_tensor(hr_img).unsqueeze(0).to(device)

    with torch.no_grad():
        pred_tensor = model(lr_tensor).clamp(0, 1)

    psnr = calc_psnr(pred_tensor, hr_tensor)
    pred_img = TF.to_pil_image(pred_tensor.squeeze(0))
    
    lr_slider_img = lr_img.resize((w, h), Image.NEAREST)
    
    details = (
        f"### Benchmark Results\n"
        f"- **PSNR:** {psnr:.2f} dB\n\n"
        f"- **Low-Res Input:** {lr_w} x {lr_h} ({lr_w * lr_h:,} pixels)\n\n"
        f"- **Model Output & Ground Truth:** {w} x {h} ({w * h:,} pixels)"
    )
    
    # Gradio's native ImageSlider expects a tuple of (image1, image2)
    return details, (lr_slider_img, pred_img), (hr_img, pred_img)

# --- 4. GRADIO UI ---
with gr.Blocks() as app:
    gr.Markdown(
        """
        # ⚡ FastEDSR 2x Image Upscaler
        Upload an image to enhance and upscale it by 2x. Supports up to 4K resolution output.

        For more information on the model, training, and use, and for a local demo, visit our [website](https://infinitode.netlify.app).
        
        This model was trained on DIV2K:
        ```
        @inproceedings{Agustsson2017,
          title={NTIRE 2017 Challenge on Single Image Super-Resolution: Dataset and Study},
          author={Agustsson, Eirikur and Timofte, Radu},
          booktitle={IEEE Conference on Computer Vision and Pattern Recognition Workshops},
          year={2017}
        }
        ```
        """
    )
    
    with gr.Tabs():
        # TAB 1: STANDARD
        with gr.TabItem("⚡ Standard Upscaling"):
            gr.Markdown("Directly upscale any low-resolution image. Inputs over 2K (2048px) will be scaled down to prevent memory limits.")
            with gr.Row():
                with gr.Column():
                    std_input = gr.Image(type="pil", label="Low Resolution Input")
                    std_btn = gr.Button("Upscale Image", variant="primary")
                with gr.Column():
                    std_output = gr.Image(type="pil", label="2x High Resolution Output")
                    std_details = gr.Markdown()
                    
            std_btn.click(fn=standard_upscale, inputs=std_input, outputs=[std_output, std_details])

        # TAB 2: BENCHMARK
        with gr.TabItem("📊 Benchmark Mode"):
            gr.Markdown("Upload a high-quality image (up to 4K). The app will compress it to 2x lower resolution, upscale it using FastEDSR, and measure the PSNR quality against the original. It will also generate side-by-side comparisons for you to view.")
            with gr.Row():
                with gr.Column():
                    bm_input = gr.Image(type="pil", label="Ground Truth (High Res) Image")
                    bm_btn = gr.Button("Run Benchmark", variant="primary")
                    bm_details = gr.Markdown()
                with gr.Column():
                    gr.Markdown("### Low-Res vs. Model Prediction")
                    slider_lr_pred = gr.ImageSlider(label="Left: Pixelated Low-Res | Right: FastEDSR")
                    
                    gr.Markdown("### Ground Truth vs. Model Prediction")
                    slider_hr_pred = gr.ImageSlider(label="Left: Original HR | Right: FastEDSR")
                    
            bm_btn.click(fn=benchmark_upscale, inputs=bm_input, outputs=[bm_details, slider_lr_pred, slider_hr_pred])

if __name__ == "__main__":
    app.launch()