File size: 2,294 Bytes
747091e |
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 |
from diffusers import LDMSuperResolutionPipeline
import numpy as np
import opensr_test
import torch
import pickle
from typing import Union
def create_stable_diffusion_model(
device: Union[str, torch.device] = "cuda"
) -> LDMSuperResolutionPipeline:
""" Create the stable diffusion model
Returns:
LDMSuperResolutionPipeline: The model to use for
super resolution.
"""
model_id = "CompVis/ldm-super-resolution-4x-openimages"
pipeline = LDMSuperResolutionPipeline.from_pretrained(model_id)
pipeline = pipeline.to(device)
return pipeline
def run_diffuser(
model: LDMSuperResolutionPipeline,
lr: torch.Tensor,
hr: torch.Tensor,
device: Union[str, torch.device] = "cuda"
) -> dict:
""" Run the model on the low resolution image
Args:
model (LDMSuperResolutionPipeline): The model to use
lr (torch.Tensor): The low resolution image
hr (torch.Tensor): The high resolution image
device (Union[str, torch.device], optional): The device
to use. Defaults to "cuda".
Returns:
dict: The results of the model
"""
# move the images to the device
lr = (torch.from_numpy(lr[[3, 2, 1]]) / 2000).to(device).clamp(0, 1)
if lr.shape[1] == 121:
# add padding
lr = torch.nn.functional.pad(
lr[None],
pad=(3, 4, 3, 4),
mode='reflect'
).squeeze()
# run the model
with torch.no_grad():
sr = model(lr[None], num_inference_steps=100, eta=1)
sr = torch.from_numpy(
np.array(sr.images[0])/255
).permute(2,0,1).float()
# remove padding
sr = sr[:, 3*4:-4*4, 3*4:-4*4]
lr = lr[:, 3:-4, 3:-4]
else:
# run the model
with torch.no_grad():
sr = model(lr[None], num_inference_steps=100, eta=1)
sr = torch.from_numpy(
np.array(sr.images[0])/255
).permute(2,0,1).float()
lr = (lr.cpu().numpy() * 2000).astype(np.uint16)
hr = ((hr[0:3] / 2000).clip(0, 1) * 2000).astype(np.uint16)
sr = (sr.cpu().numpy() * 2000).astype(np.uint16)
results = {
"lr": lr,
"hr": hr,
"sr": sr
}
return results |