import spaces
import gradio as gr
import torch
from PIL import Image
import numpy as np
from transformers import Swin2SRForImageSuperResolution, Swin2SRImageProcessor

MODEL_ID = "caidas/swin2SR-classical-sr-x4-64"

print("Loading Swin2SR x4 model...")
processor = Swin2SRImageProcessor.from_pretrained(MODEL_ID)
model = Swin2SRForImageSuperResolution.from_pretrained(MODEL_ID)
model.to("cuda")
model.eval()
print("Model ready.")


@spaces.GPU(duration=30)
def upscale_image(image: Image.Image):
    if image is None:
        raise gr.Error("Please upload a photo first.")

    try:
        image = image.convert("RGB")

        inputs = processor(image, return_tensors="pt").to("cuda")

        with torch.no_grad():
            outputs = model(**inputs)

        output = outputs.reconstruction.data.squeeze().float().cpu().clamp_(0, 1).numpy()
        output = np.moveaxis(output, source=0, destination=-1)
        output = (output * 255.0).round().astype(np.uint8)
        result = Image.fromarray(output)

        return result
    except Exception as e:
        import traceback
        traceback.print_exc()
        raise gr.Error(f"Upscaling failed: {e}")


css = """
#header {
    text-align: center;
    padding: 24px 0 8px;
}
#header h1 {
    font-size: 32px;
    font-weight: 700;
    background: linear-gradient(135deg, #6366f1, #06b6d4);
    -webkit-background-clip: text;
    -webkit-text-fill-color: transparent;
    margin-bottom: 4px;
}
#header p {
    color: #888;
    font-size: 14px;
}
#run-btn {
    background: linear-gradient(135deg, #6366f1, #06b6d4) !important;
    color: white !important;
    font-weight: 600 !important;
    border: none !important;
}
"""

with gr.Blocks(title="Peace Network Upscaler", css=css) as demo:
    gr.HTML(
        """
        <div id="header">
            <h1>🔍 Peace Network Upscaler</h1>
            <p>Upload a low-res photo — get a sharp 4x upscaled version, powered by Swin2SR.</p>
        </div>
        """
    )

    with gr.Row():
        inp = gr.Image(type="pil", label="Upload photo")
        out = gr.Image(type="pil", label="Upscaled result (4x)", format="png")

    btn = gr.Button("✨ Upscale Image", variant="primary", elem_id="run-btn")
    btn.click(upscale_image, inputs=inp, outputs=out)

demo.queue().launch()