Spaces:
Running
Running
| from impact.utils import * | |
| from impact import impact_sampling | |
| from comfy import model_management | |
| from comfy.cli_args import args | |
| import nodes | |
| try: | |
| from comfy_extras import nodes_differential_diffusion | |
| except Exception: | |
| print(f"[Impact Pack] ComfyUI is an outdated version. The DifferentialDiffusion feature will be disabled.") | |
| # Implementation based on `https://github.com/lingondricka2/Upscaler-Detailer` | |
| # code from comfyroll ---> | |
| # https://github.com/Suzie1/ComfyUI_Comfyroll_CustomNodes/blob/main/nodes/functions_upscale.py | |
| 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 | |
| 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 | |
| def upscaler(image, upscale_model, rescale_factor, resampling_method, supersample, rounding_modulus): | |
| if upscale_model is not None: | |
| up_image = upscale_with_model(upscale_model, image) | |
| else: | |
| up_image = image | |
| pil_img = tensor2pil(image) | |
| original_width, original_height = pil_img.size | |
| scaled_image = pil2tensor(apply_resize_image(tensor2pil(up_image), original_width, original_height, rounding_modulus, 'rescale', | |
| supersample, rescale_factor, 1024, resampling_method)) | |
| return scaled_image | |
| # <--- | |
| def img2img_segs(image, model, clip, vae, seed, steps, cfg, sampler_name, scheduler, | |
| positive, negative, denoise, noise_mask, control_net_wrapper=None, | |
| inpaint_model=False, noise_mask_feather=0, scheduler_func_opt=None): | |
| original_image_size = image.shape[1:3] | |
| # Match to original image size | |
| if original_image_size[0] % 8 > 0 or original_image_size[1] % 8 > 0: | |
| scale = 8/min(original_image_size[0], original_image_size[1]) + 1 | |
| w = int(original_image_size[1] * scale) | |
| h = int(original_image_size[0] * scale) | |
| image = tensor_resize(image, w, h) | |
| if noise_mask is not None: | |
| noise_mask = tensor_gaussian_blur_mask(noise_mask, noise_mask_feather) | |
| noise_mask = noise_mask.squeeze(3) | |
| if noise_mask_feather > 0 and 'denoise_mask_function' not in model.model_options: | |
| model = nodes_differential_diffusion.DifferentialDiffusion().apply(model)[0] | |
| if control_net_wrapper is not None: | |
| positive, negative, _ = control_net_wrapper.apply(positive, negative, image, noise_mask) | |
| # prepare mask | |
| if noise_mask is not None and inpaint_model: | |
| positive, negative, latent_image = nodes.InpaintModelConditioning().encode(positive, negative, image, vae, noise_mask) | |
| else: | |
| latent_image = to_latent_image(image, vae) | |
| if noise_mask is not None: | |
| latent_image['noise_mask'] = noise_mask | |
| refined_latent = latent_image | |
| # ksampler | |
| refined_latent = impact_sampling.ksampler_wrapper(model, seed, steps, cfg, sampler_name, scheduler, positive, negative, refined_latent, denoise, scheduler_func=scheduler_func_opt) | |
| # non-latent downscale - latent downscale cause bad quality | |
| refined_image = vae.decode(refined_latent['samples']) | |
| # prevent mixing of device | |
| refined_image = refined_image.cpu() | |
| # Match to original image size | |
| if refined_image.shape[1:3] != original_image_size: | |
| refined_image = tensor_resize(refined_image, original_image_size[1], original_image_size[0]) | |
| # don't convert to latent - latent break image | |
| # preserving pil is much better | |
| return refined_image | |