#!/usr/bin/env python
# -*- coding:utf-8 -*-
# Power by Zongsheng Yue 2024-12-11 17:17:41
import spaces
import warnings
warnings.filterwarnings("ignore")
import argparse
import numpy as np
import gradio as gr
from pathlib import Path
from omegaconf import OmegaConf
from sampler_invsr import InvSamplerSR
from utils import util_common
from utils import util_image
from basicsr.utils.download_util import load_file_from_url
def get_configs(num_steps=1, chopping_size=128, seed=12345):
configs = OmegaConf.load("./configs/sample-sd-turbo.yaml")
if num_steps == 1:
configs.timesteps = [200,]
elif num_steps == 2:
configs.timesteps = [200, 100]
elif num_steps == 3:
configs.timesteps = [200, 100, 50]
elif num_steps == 4:
configs.timesteps = [200, 150, 100, 50]
elif num_steps == 5:
configs.timesteps = [250, 200, 150, 100, 50]
else:
assert num_steps <= 250
configs.timesteps = np.linspace(
start=250, stop=0, num=num_steps, endpoint=False, dtype=np.int64()
).tolist()
print(f'Setting timesteps for inference: {configs.timesteps}')
# path to save noise predictor
started_ckpt_name = "noise_predictor_sd_turbo_v5.pth"
started_ckpt_dir = "./weights"
util_common.mkdir(started_ckpt_dir, delete=False, parents=True)
started_ckpt_path = Path(started_ckpt_dir) / started_ckpt_name
if not started_ckpt_path.exists():
load_file_from_url(
url="https://huggingface.co/OAOA/InvSR/resolve/main/noise_predictor_sd_turbo_v5.pth",
model_dir=started_ckpt_dir,
progress=True,
file_name=started_ckpt_name,
)
configs.model_start.ckpt_path = str(started_ckpt_path)
configs.bs = 1
configs.seed = 12345
configs.basesr.chopping.pch_size = chopping_size
if chopping_size == 128:
configs.basesr.chopping.extra_bs = 8
elif chopping_size == 256:
configs.basesr.chopping.extra_bs = 4
else:
configs.basesr.chopping.extra_bs = 1
return configs
@spaces.GPU
def predict(in_path, num_steps=1, chopping_size=128, seed=12345):
configs = get_configs(num_steps=num_steps, chopping_size=chopping_size, seed=12345)
sampler = InvSamplerSR(configs)
out_dir = Path('invsr_output')
if not out_dir.exists():
out_dir.mkdir()
sampler.inference(in_path, out_path=out_dir, bs=1)
out_path = out_dir / f"{Path(in_path).stem}.png"
assert out_path.exists(), 'Super-resolution failed!'
im_sr = util_image.imread(out_path, chn="rgb", dtype="uint8")
return im_sr, str(out_path)
title = "Arbitrary-steps Image Super-resolution via Diffusion Inversion"
description = r"""
Official Gradio demo for Arbitrary-steps Image Super-resolution via Diffuion Inversion.
🔥 InvSR is an image super-resolution method via Diffusion Inversion, supporting arbitrary sampling steps.
"""
article = r"""
If you've found InvSR useful for your research or projects, please show your support by ⭐ the Github Repo. Thanks!
[![GitHub Stars](https://img.shields.io/github/stars/zsyOAOA/InvSR?affiliations=OWNER&color=green&style=social)](https://github.com/zsyOAOA/InvSR)
---
If our work is useful for your research, please consider citing:
```bibtex
@article{yue2024InvSR,
title={Arbitrary-steps Image Super-resolution via Diffusion Inversion},
author={Yue, Zongsheng and Kang, Liao and Loy, Chen Change},
journal = {arXiv preprint arXiv:2412.09013},
year={2024},
}
```
📋 **License**
This project is licensed under S-Lab License 1.0.
Redistribution and use for non-commercial purposes should follow this license.
📧 **Contact**
If you have any questions, please feel free to contact me via zsyzam@gmail.com.
![visitors](https://visitor-badge.laobi.icu/badge?page_id=zsyOAOA/InvSR)
"""
demo = gr.Interface(
fn=predict,
inputs=[
gr.Image(type="filepath", label="Input: Low Quality Image"),
gr.Dropdown(
choices=[1,2,3,4,5],
value=1,
label="Number of steps",
),
gr.Dropdown(
choices=[128, 256, 512],
value=128,
label="Chopping size",
),
gr.Number(value=12345, precision=0, label="Ranom seed")
],
outputs=[
gr.Image(type="numpy", label="Output: High Quality Image"),
gr.File(label="Download the output")
],
title=title,
description=description,
article=article,
examples=[
['./testdata/RealSet80/29.jpg', 3, 128, 12345],
['./testdata/RealSet80/32.jpg', 1, 128, 12345],
['./testdata/RealSet80/0030.jpg', 1, 128, 12345],
['./testdata/RealSet80/2684538-PH.jpg', 1, 128, 12345],
['./testdata/RealSet80/oldphoto6.png', 1, 128, 12345],
]
)
demo.queue(max_size=5)
demo.launch(share=True)