File size: 4,228 Bytes
5761977 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 | """Minimal PiSA-SR inference on the diffusers repack, in default or adjustable mode."""
import argparse, os
import torch
from PIL import Image
from torchvision import transforms
import torchvision.transforms.functional as TF
from diffusers import AutoencoderKL, UNet2DConditionModel
from transformers import AutoTokenizer, CLIPTextModel
BASE = "stabilityai/stable-diffusion-2-1-base"
REPO = "ndtran0101/pisa-sr-diffusers"
def prepare(img, upscale=4):
"""Upsample a low-quality image to the diffusion input size, aligned to a multiple of 8.
The upsample happens before the UNet, so cost scales with the output size.
"""
img = img.resize((img.width * upscale, img.height * upscale))
return img.resize((img.width - img.width % 8, img.height - img.height % 8), Image.LANCZOS)
def adain(target, source):
"""Match the target's per-channel mean and std to the source.
The upstream pipeline applies this colour transfer after decoding; skipping it
shifts colour noticeably.
"""
t = TF.to_tensor(target).unsqueeze(0)
s = TF.to_tensor(source).unsqueeze(0)
t_mean, t_std = t.mean([2, 3], keepdim=True), t.std([2, 3], keepdim=True)
s_mean, s_std = s.mean([2, 3], keepdim=True), s.std([2, 3], keepdim=True)
out = ((t - t_mean) / (t_std + 1e-5)) * s_std + s_mean
return transforms.ToPILImage()(out.clamp(0, 1)[0])
@torch.no_grad()
def upscale(path, base=BASE, repo=REPO, lambda_pix=None, lambda_sem=None,
device="cuda", dtype=torch.float16, color_fix=True):
"""Upscale one image x4 and return it as a PIL image.
Runs one UNet pass at t=1 on an empty prompt and subtracts the predicted residual;
there is no scheduler and no noise. Passing both lambdas switches to the two-pass
adjustable mode, which additionally loads unet_pix.
path: low-quality input image
repo: HuggingFace id or local directory holding unet_full/ and unet_pix/
lambda_pix / lambda_sem: strength of the pixel and semantic branches
"""
adjustable = lambda_pix is not None and lambda_sem is not None
tok = AutoTokenizer.from_pretrained(base, subfolder="tokenizer")
te = CLIPTextModel.from_pretrained(base, subfolder="text_encoder").to(device, dtype).eval()
vae = AutoencoderKL.from_pretrained(base, subfolder="vae").to(device, dtype).eval()
unet = UNet2DConditionModel.from_pretrained(repo, subfolder="unet_full").to(device, dtype).eval()
src = Image.open(path).convert("RGB")
img = prepare(src)
x = TF.to_tensor(img).unsqueeze(0).to(device, dtype) * 2 - 1
ids = tok("", max_length=tok.model_max_length, padding="max_length",
truncation=True, return_tensors="pt").input_ids.to(device)
emb = te(ids)[0].to(dtype)
t = torch.tensor([1], device=device).long()
z = vae.encode(x).latent_dist.sample() * vae.config.scaling_factor
if adjustable:
unet_pix = UNet2DConditionModel.from_pretrained(
repo, subfolder="unet_pix").to(device, dtype).eval()
pred_sem = unet(z, t, encoder_hidden_states=emb).sample
pred_pix = unet_pix(z, t, encoder_hidden_states=emb).sample
pred = lambda_pix * pred_pix + lambda_sem * (pred_sem - pred_pix)
else:
pred = unet(z, t, encoder_hidden_states=emb).sample
out = vae.decode((z - pred) / vae.config.scaling_factor).sample.clamp(-1, 1)
pil = transforms.ToPILImage()((out * 0.5 + 0.5).clamp(0, 1)[0].float().cpu())
return adain(pil, img) if color_fix else pil
def main():
ap = argparse.ArgumentParser()
ap.add_argument("--input", required=True)
ap.add_argument("--output", default="sr.png")
ap.add_argument("--repo", default=REPO)
ap.add_argument("--base", default=BASE)
ap.add_argument("--lambda_pix", type=float, default=None)
ap.add_argument("--lambda_sem", type=float, default=None)
ap.add_argument("--no_color_fix", action="store_true")
a = ap.parse_args()
img = upscale(a.input, base=a.base, repo=a.repo,
lambda_pix=a.lambda_pix, lambda_sem=a.lambda_sem,
color_fix=not a.no_color_fix)
img.save(a.output)
print(f"{a.input} -> {a.output} {img.size[0]}x{img.size[1]}")
if __name__ == "__main__":
main()
|