Create app.py

app.py

@@ -0,0 +1,79 @@
+import os
+import torch
+import gradio as gr
+import spaces
+from PIL import Image
+import numpy as np
+from omegaconf import OmegaConf
+import requests
+from tqdm import tqdm
+
+def download_file(url, filename):
+    response = requests.get(url, stream=True)
+    total_size = int(response.headers.get('content-length', 0))
+    block_size = 1024
+    with open(filename, 'wb') as file, tqdm(
+        desc=filename,
+        total=total_size,
+        unit='iB',
+        unit_scale=True,
+        unit_divisor=1024,
+    ) as progress_bar:
+        for data in response.iter_content(block_size):
+            size = file.write(data)
+            progress_bar.update(size)
+
+def setup_environment():
+    os.makedirs("weights", exist_ok=True)
+    if not os.path.exists("weights/real-world_ccsr.ckpt"):
+        print("Downloading model checkpoint...")
+        download_file(
+            "https://huggingface.co/camenduru/CCSR/resolve/main/real-world_ccsr.ckpt",
+            "weights/real-world_ccsr.ckpt"
+        )
+    else:
+        print("Model checkpoint already exists. Skipping download.")
+
+setup_environment()
+
+from ccsr.models.ccsr import CCSR
+from ccsr.utils.util import instantiate_from_config
+
+config = OmegaConf.load("configs/model/ccsr_stage2.yaml")
+model = instantiate_from_config(config.model)
+ckpt = torch.load("weights/real-world_ccsr.ckpt", map_location="cpu")
+model.load_state_dict(ckpt["state_dict"], strict=False)
+model.cuda().eval()
+
+@spaces.GPU
+@torch.inference_mode()
+def infer(image, sr_scale, t_max, t_min, color_fix_type):
+    image = Image.open(image).convert("RGB").resize((256, 256), Image.LANCZOS)
+    image = torch.from_numpy(np.array(image)).float().cuda() / 127.5 - 1
+    image = image.permute(2, 0, 1).unsqueeze(0)
+
+    output = model.super_resolution(
+        image,
+        sr_scale=sr_scale,
+        t_max=t_max,
+        t_min=t_min,
+        color_fix_type=color_fix_type
+    )
+
+    output = ((output.squeeze().permute(1, 2, 0).cpu().numpy() + 1) * 127.5).clip(0, 255).astype(np.uint8)
+    return Image.fromarray(output)
+
+interface = gr.Interface(
+    fn=infer,
+    inputs=[
+        gr.Image(type="filepath"),
+        gr.Slider(minimum=1, maximum=8, step=1, value=4),
+        gr.Slider(minimum=0, maximum=1, step=0.0001, value=0.6667),
+        gr.Slider(minimum=0, maximum=1, step=0.0001, value=0.3333),
+        gr.Dropdown(choices=["adain", "wavelet", "none"], value="adain"),
+    ],
+    outputs=gr.Image(type="pil"),
+    title="CCSR: Continuous Contrastive Super-Resolution",
+)
+
+interface.launch()