|
|
import comfy.utils |
|
|
import torch |
|
|
import gc |
|
|
import logging |
|
|
import comfy.model_management as model_management |
|
|
|
|
|
|
|
|
def clear_gpu_and_ram_cache(): |
|
|
gc.collect() |
|
|
if torch.cuda.is_available(): |
|
|
torch.cuda.empty_cache() |
|
|
torch.cuda.ipc_collect() |
|
|
|
|
|
|
|
|
def _smart_decode(vae, latent, tile_size=512): |
|
|
try: |
|
|
images = vae.decode(latent["samples"]) |
|
|
except model_management.OOM_EXCEPTION: |
|
|
logging.warning("VAE decode OOM, using tiled decode") |
|
|
compression = vae.spacial_compression_decode() |
|
|
images = vae.decode_tiled( |
|
|
latent["samples"], |
|
|
tile_x=tile_size // compression, |
|
|
tile_y=tile_size // compression, |
|
|
overlap=(tile_size // 4) // compression, |
|
|
) |
|
|
if len(images.shape) == 5: |
|
|
images = images.reshape(-1, images.shape[-3], images.shape[-2], images.shape[-1]) |
|
|
return images |
|
|
|
|
|
|
|
|
class MagicUpscaleModule: |
|
|
"""Moved into mod/ as mg_upscale_module keeping class/key name.""" |
|
|
upscale_methods = ["nearest-exact", "bilinear", "area", "bicubic", "lanczos"] |
|
|
|
|
|
@classmethod |
|
|
def INPUT_TYPES(cls): |
|
|
return { |
|
|
"required": { |
|
|
"samples": ("LATENT", {}), |
|
|
"vae": ("VAE", {}), |
|
|
"upscale_method": (cls.upscale_methods, {"default": "bilinear"}), |
|
|
"scale_by": ("FLOAT", {"default": 1.2, "min": 0.01, "max": 8.0, "step": 0.01}), |
|
|
} |
|
|
} |
|
|
|
|
|
RETURN_TYPES = ("LATENT", "IMAGE") |
|
|
RETURN_NAMES = ("LATENT", "Upscaled Image") |
|
|
FUNCTION = "process_upscale" |
|
|
CATEGORY = "MagicNodes" |
|
|
|
|
|
def process_upscale(self, samples, vae, upscale_method, scale_by): |
|
|
clear_gpu_and_ram_cache() |
|
|
images = _smart_decode(vae, samples) |
|
|
samples_t = images.movedim(-1, 1) |
|
|
width = round(samples_t.shape[3] * scale_by) |
|
|
height = round(samples_t.shape[2] * scale_by) |
|
|
|
|
|
try: |
|
|
stride = int(vae.spacial_compression_decode()) |
|
|
except Exception: |
|
|
stride = 8 |
|
|
if stride <= 0: |
|
|
stride = 8 |
|
|
def _align_up(x, s): |
|
|
return int(((x + s - 1) // s) * s) |
|
|
width_al = _align_up(width, stride) |
|
|
height_al = _align_up(height, stride) |
|
|
up = comfy.utils.common_upscale(samples_t, width_al, height_al, upscale_method, "disabled") |
|
|
up = up.movedim(1, -1) |
|
|
encoded = vae.encode(up[:, :, :, :3]) |
|
|
return ({"samples": encoded}, up) |
|
|
|
