70 lines
2.5 KiB
Python
70 lines
2.5 KiB
Python
from comfy_extras.nodes_upscale_model import ImageUpscaleWithModel
|
|
from PIL import Image
|
|
|
|
from .tensor_image_ops import (
|
|
pilimgbatch_to_torch,
|
|
torch_to_pilimgbatch,
|
|
)
|
|
from .utils import scale_dim
|
|
|
|
|
|
class Upscale:
|
|
def __init__(
|
|
self,
|
|
*,
|
|
resample_mode="bicubic",
|
|
rescale_increment=64,
|
|
upscale_model=None,
|
|
force_upscale_model=False,
|
|
):
|
|
self.resample_mode = resample_mode
|
|
self.rescale_increment = scale_dim(max(8, rescale_increment), increment=8)
|
|
self.upscale_model = upscale_model
|
|
self.force_upscale_model = upscale_model is not None and force_upscale_model
|
|
|
|
def __call__(self, imgbatch, scale_factor, *, pbar=None, use_upscale_model=True):
|
|
if scale_factor == 1.0 and not self.force_upscale_model:
|
|
return imgbatch
|
|
_batch, height, width, _channels = imgbatch.shape
|
|
target_height = scale_dim(
|
|
height,
|
|
scale_factor,
|
|
increment=self.rescale_increment,
|
|
)
|
|
target_width = scale_dim(
|
|
width,
|
|
scale_factor,
|
|
increment=self.rescale_increment,
|
|
)
|
|
# tqdm.write(f">> UPSCALE: {width}x{height} -> {target_width}x{target_height}")
|
|
if (target_height, target_width) == (
|
|
height,
|
|
width,
|
|
) and not self.force_upscale_model:
|
|
return imgbatch
|
|
if use_upscale_model and self.upscale_model is not None:
|
|
if pbar is not None:
|
|
pbar.set_description(
|
|
f"upscale with model: {width}x{height} -> {target_width}x{target_height}",
|
|
)
|
|
# tqdm.write("** Upscaling with model")
|
|
imgbatch = ImageUpscaleWithModel().upscale(self.upscale_model, imgbatch)[0]
|
|
if imgbatch.shape[1:3] == (target_height, target_width):
|
|
return imgbatch
|
|
# tqdm.write(
|
|
# f"** PIL upscale {imgbatch.shape[2]}x{imgbatch.shape[1]} -> {target_width}x{target_height}",
|
|
# )
|
|
if pbar is not None and self.upscale_model is None:
|
|
pbar.set_description(
|
|
f"upscale (simple): {imgbatch.shape[2]}x{imgbatch.shape[1]} -> {target_width}x{target_height}",
|
|
)
|
|
return pilimgbatch_to_torch(
|
|
tuple(
|
|
i.resize(
|
|
(target_width, target_height),
|
|
resample=getattr(Image.Resampling, self.resample_mode.upper()),
|
|
)
|
|
for i in torch_to_pilimgbatch(imgbatch)
|
|
),
|
|
)
|