Spaces:
Sleeping
Sleeping
File size: 3,891 Bytes
4450790 |
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 82 83 84 85 86 87 |
#---------------------------------------------------------------------------------------------------------------------#
# Comfyroll Studio custom nodes by RockOfFire and Akatsuzi https://github.com/Suzie1/ComfyUI_Comfyroll_CustomNodes
# for ComfyUI https://github.com/comfyanonymous/ComfyUI
#---------------------------------------------------------------------------------------------------------------------#
#---------------------------------------------------------------------------------------------------------------------#
# UPSCALE FUNCTIONS
#---------------------------------------------------------------------------------------------------------------------#
# These functions are based on WAS nodes Image Resize and the Comfy Extras upscale with model nodes
import torch
#import os
from comfy_extras.chainner_models import model_loading
from comfy import model_management
import numpy as np
import comfy.utils
import folder_paths
from PIL import Image
# PIL to Tensor
def pil2tensor(image):
return torch.from_numpy(np.array(image).astype(np.float32) / 255.0).unsqueeze(0)
# Tensor to PIL
def tensor2pil(image):
return Image.fromarray(np.clip(255. * image.cpu().numpy().squeeze(), 0, 255).astype(np.uint8))
def load_model(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 = model_loading.load_state_dict(sd).eval()
return out
def upscale_with_model(upscale_model, image):
device = model_management.get_torch_device()
upscale_model.to(device)
in_img = image.movedim(-1,-3).to(device)
free_memory = model_management.get_free_memory(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.cpu()
s = torch.clamp(s.movedim(-3,-1), min=0, max=1.0)
return s
def apply_resize_image(image: Image.Image, original_width, original_height, rounding_modulus, mode='scale', supersample='true', factor: int = 2, width: int = 1024, height: int = 1024, resample='bicubic'):
# Calculate the new width and height based on the given mode and parameters
if mode == 'rescale':
new_width, new_height = int(original_width * factor), int(original_height * factor)
else:
m = rounding_modulus
original_ratio = original_height / original_width
height = int(width * original_ratio)
new_width = width if width % m == 0 else width + (m - width % m)
new_height = height if height % m == 0 else height + (m - height % m)
# Define a dictionary of resampling filters
resample_filters = {'nearest': 0, 'bilinear': 2, 'bicubic': 3, 'lanczos': 1}
# Apply supersample
if supersample == 'true':
image = image.resize((new_width * 8, new_height * 8), resample=Image.Resampling(resample_filters[resample]))
# Resize the image using the given resampling filter
resized_image = image.resize((new_width, new_height), resample=Image.Resampling(resample_filters[resample]))
return resized_image
|