| import os
|
| import logging
|
| from spandrel import ModelLoader, ImageModelDescriptor
|
| from comfy import model_management
|
| import torch
|
| import comfy.utils
|
| import folder_paths
|
|
|
| try:
|
| from spandrel_extra_arches import EXTRA_REGISTRY
|
| from spandrel import MAIN_REGISTRY
|
| MAIN_REGISTRY.add(*EXTRA_REGISTRY)
|
| logging.info("Successfully imported spandrel_extra_arches: support for non commercial upscale models.")
|
| except:
|
| pass
|
|
|
| class UpscaleModelLoader:
|
| @classmethod
|
| def INPUT_TYPES(s):
|
| return {"required": { "model_name": (folder_paths.get_filename_list("upscale_models"), ),
|
| }}
|
| RETURN_TYPES = ("UPSCALE_MODEL",)
|
| FUNCTION = "load_model"
|
|
|
| CATEGORY = "loaders"
|
|
|
| def load_model(self, model_name):
|
| model_path = folder_paths.get_full_path("upscale_models", model_name)
|
| sd = comfy.utils.load_torch_file(model_path, safe_load=True)
|
| if "module.layers.0.residual_group.blocks.0.norm1.weight" in sd:
|
| sd = comfy.utils.state_dict_prefix_replace(sd, {"module.":""})
|
| out = ModelLoader().load_from_state_dict(sd).eval()
|
|
|
| if not isinstance(out, ImageModelDescriptor):
|
| raise Exception("Upscale model must be a single-image model.")
|
|
|
| return (out, )
|
|
|
|
|
| class ImageUpscaleWithModel:
|
| @classmethod
|
| def INPUT_TYPES(s):
|
| return {"required": { "upscale_model": ("UPSCALE_MODEL",),
|
| "image": ("IMAGE",),
|
| }}
|
| RETURN_TYPES = ("IMAGE",)
|
| FUNCTION = "upscale"
|
|
|
| CATEGORY = "image/upscaling"
|
|
|
| def upscale(self, upscale_model, image):
|
| device = model_management.get_torch_device()
|
|
|
| memory_required = model_management.module_size(upscale_model.model)
|
| memory_required += (512 * 512 * 3) * image.element_size() * max(upscale_model.scale, 1.0) * 384.0
|
| memory_required += image.nelement() * image.element_size()
|
| model_management.free_memory(memory_required, device)
|
|
|
| upscale_model.to(device)
|
| in_img = image.movedim(-1,-3).to(device)
|
|
|
| tile = 512
|
| overlap = 32
|
|
|
| oom = True
|
| while oom:
|
| try:
|
| steps = in_img.shape[0] * comfy.utils.get_tiled_scale_steps(in_img.shape[3], in_img.shape[2], tile_x=tile, tile_y=tile, overlap=overlap)
|
| pbar = comfy.utils.ProgressBar(steps)
|
| s = comfy.utils.tiled_scale(in_img, lambda a: upscale_model(a), tile_x=tile, tile_y=tile, overlap=overlap, upscale_amount=upscale_model.scale, pbar=pbar)
|
| oom = False
|
| except model_management.OOM_EXCEPTION as e:
|
| tile //= 2
|
| if tile < 128:
|
| raise e
|
|
|
| upscale_model.to("cpu")
|
| s = torch.clamp(s.movedim(-3,-1), min=0, max=1.0)
|
| return (s,)
|
|
|
| NODE_CLASS_MAPPINGS = {
|
| "UpscaleModelLoader": UpscaleModelLoader,
|
| "ImageUpscaleWithModel": ImageUpscaleWithModel
|
| }
|
|
|
