Spaces:
Sleeping
Sleeping
| #!/usr/bin/env python | |
| # -*- coding:utf-8 -*- | |
| import os | |
| os.environ["CUDA_VISIBLE_DEVICES"] = "" # 🔥 Force CPU | |
| 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 | |
| # Optional: enforce CPU in torch if used internally | |
| try: | |
| import torch | |
| torch.set_default_device("cpu") | |
| except: | |
| pass | |
| 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}') | |
| 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 = seed | |
| 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 | |
| # ❌ Removed @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=seed) | |
| sampler = InvSamplerSR(configs) | |
| out_dir = Path('invsr_output') | |
| out_dir.mkdir(exist_ok=True) | |
| 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 = """ | |
| <b>CPU version</b> of InvSR demo.<br> | |
| ⚠️ Note: Running on CPU will be significantly slower than GPU. | |
| """ | |
| demo = gr.Interface( | |
| fn=predict, | |
| inputs=[ | |
| gr.Image(type="filepath", label="Input Image"), | |
| gr.Dropdown([1,2,3,4,5], value=1, label="Steps"), | |
| gr.Dropdown([128, 256, 512], value=128, label="Chopping size"), | |
| gr.Number(value=12345, precision=0, label="Seed") | |
| ], | |
| outputs=[ | |
| gr.Image(type="numpy", label="Output Image"), | |
| gr.File(label="Download") | |
| ], | |
| title=title, | |
| description=description | |
| ) | |
| demo.queue(max_size=5) | |
| demo.launch() |