|
|
|
|
|
|
|
import os |
|
import shutil |
|
from omegaconf import OmegaConf |
|
from cog import BasePredictor, Input, Path |
|
|
|
from sampler import ResShiftSampler |
|
|
|
class Predictor(BasePredictor): |
|
def setup(self) -> None: |
|
"""Load the model into memory to make running multiple predictions efficient""" |
|
self.configs = { |
|
"realsr": OmegaConf.load('./configs/realsr_swinunet_realesrgan256_journal.yaml'), |
|
"bicsr": configs = OmegaConf.load('./configs/bicx4_swinunet_lpips.yaml'), |
|
} |
|
|
|
def predict( |
|
self, |
|
image: Path = Input(description="Grayscale input image"), |
|
scale: int = Input(description="Factor to scale image by.", default=4), |
|
chop_size: int = Input( |
|
choices=[512, 256], description="Chopping forward.", default=512 |
|
), |
|
task: str = Input( |
|
choices=["realsr", "bicsr"], |
|
description="Choose a task", |
|
default="realsr", |
|
), |
|
seed: int = Input( |
|
description="Random seed. Leave blank to randomize the seed.", default=12345 |
|
), |
|
) -> Path: |
|
"""Run a single prediction on the model""" |
|
if seed is None: |
|
seed = int.from_bytes(os.urandom(2), "big") |
|
print(f"Using seed: {seed}") |
|
|
|
configs = self.configs[task] |
|
|
|
if task == 'realsr': |
|
ckpt_path = f"weights/resshift_realsrx4_s4_v3.pth" |
|
configs.model.ckpt_path = ckpt_path |
|
else: |
|
ckpt_path = f"weights/resshift_bicsrx4_s4.pth" |
|
configs.model.ckpt_path = ckpt_path |
|
configs.diffusion.params.steps = 4 |
|
configs.diffusion.params.sf = scale |
|
configs.autoencoder.ckpt_path = f"weights/autoencoder_vq_f4.pth" |
|
|
|
chop_stride = 448 if chop_size == 512 else 224 |
|
|
|
resshift_sampler = ResShiftSampler( |
|
configs, |
|
sf=scale, |
|
chop_size=chop_size, |
|
chop_stride=chop_stride, |
|
chop_bs=1, |
|
use_amp=True, |
|
seed=seed, |
|
padding_offset=configs.model.params.get('lq_size', 64), |
|
) |
|
|
|
out_path = "out_dir" |
|
if os.path.exists(out_path): |
|
shutil.rmtree(out_path) |
|
resshift_sampler.inference( |
|
str(image), |
|
out_path, |
|
mask_path=None, |
|
bs=1, |
|
noise_repeat=False |
|
) |
|
|
|
out = "/tmp/out.png" |
|
shutil.copy(os.path.join(out_path, os.listdir(out_path)[0]), out) |
|
|
|
return Path(out) |
|
|
