| | 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) |
| |
|
