From 420cf09a5e68ebcaab0d6c24f58b8141927fc59f Mon Sep 17 00:00:00 2001 From: kijai <40791699+kijai@users.noreply.github.com> Date: Wed, 28 Feb 2024 22:52:09 +0200 Subject: [PATCH] Add rest of the parameters --- SUPIR/utils/tilevae.py | 3 ++ nodes.py | 62 ++++++++++++++++++++-------------------- sgm/modules/attention.py | 20 ++++++------- 3 files changed, 44 insertions(+), 41 deletions(-) diff --git a/SUPIR/utils/tilevae.py b/SUPIR/utils/tilevae.py index 6e9a3a5..dc1e31b 100644 --- a/SUPIR/utils/tilevae.py +++ b/SUPIR/utils/tilevae.py @@ -892,6 +892,8 @@ class VAEHook: # Task queue execution pbar = tqdm(total=num_tiles * len(task_queues[0]), desc=f"[Tiled VAE]: Executing {'Decoder' if is_decoder else 'Encoder'} Task Queue: ") + import comfy.utils + pbar_comfy = comfy.utils.ProgressBar(num_tiles) # execute the task back and forth when switch tiles so that we always # keep one tile on the GPU to reduce unnecessary data transfer @@ -934,6 +936,7 @@ class VAEHook: tile = task[1](tile) #print(tiles[i].shape, tile.shape, task) pbar.update(1) + pbar_comfy.update(1) if interrupted: break diff --git a/nodes.py b/nodes.py index c02b806..0d4b87b 100644 --- a/nodes.py +++ b/nodes.py @@ -7,10 +7,7 @@ from omegaconf import OmegaConf import comfy.model_management import folder_paths from nodes import ImageScaleBy -from nodes import ImageScale import torch.cuda -from .SUPIR.models.SUPIR_model import SUPIRModel -from PIL import Image from .sgm.util import instantiate_from_config script_directory = os.path.dirname(os.path.abspath(__file__)) @@ -22,15 +19,19 @@ class SUPIR_Upscale: "supir_model": (folder_paths.get_filename_list("checkpoints"), ), "sdxl_model": (folder_paths.get_filename_list("checkpoints"), ), "image": ("IMAGE", ), + "seed": ("INT", {"default": 123,"min": 0, "max": 0xffffffffffffffff, "step": 1}), "resize_method": (s.upscale_methods, {"default": "lanczos"}), "scale_by": ("FLOAT", {"default": 1.0, "min": 0.01, "max": 20.0, "step": 0.01}), "steps": ("INT", {"default": 45, "min": 3, "max": 4096, "step": 1}), + "restoration_scale": ("FLOAT", {"default": -1.0, "min": -1.0, "max": 6.0, "step": 1.0}), "cfg_scale": ("FLOAT", {"default": 7.5,"min": 0, "max": 20, "step": 0.01}), - "a_prompt": ("STRING", {"multiline": True, "default": "high quality",}), - "n_prompt": ("STRING", {"multiline": True, "default": "illustration",}), - - "min_size": ("INT", {"default": 1024, "min": 1, "max": 4096, "step": 1}), - + "a_prompt": ("STRING", {"multiline": True, "default": "high quality, detailed",}), + "n_prompt": ("STRING", {"multiline": True, "default": "bad quality, blurry, messy",}), + "s_churn": ("INT", {"default": 5,"min": 0, "max": 40, "step": 1}), + "s_noise": ("FLOAT", {"default": 1.003,"min": 1.0, "max": 1.1, "step": 0.001}), + "control_scale": ("FLOAT", {"default": 1.0, "min": 0, "max": 1, "step": 0.05}), + "cfg_scale_start": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 9.0, "step": 0.05}), + "control_scale_start": ("FLOAT", {"default": 0.0, "min": 0, "max": 1.0, "step": 0.05}), "color_fix_type": ( [ 'None', @@ -39,9 +40,11 @@ class SUPIR_Upscale: ], { "default": 'Wavelet' }), - "keep_model_loaded": ("BOOLEAN", {"default": False}), - "seed": ("INT", {"default": 123,"min": 0, "max": 0xffffffffffffffff, "step": 1}), + "keep_model_loaded": ("BOOLEAN", {"default": True}), }, + "optional": { + "captions": ("STRING", {"forceInput": True, "multiline": False, "default": "",}), + } } @@ -52,11 +55,11 @@ class SUPIR_Upscale: CATEGORY = "SUPIR" - def process(self, steps, image, color_fix_type, seed, scale_by, min_size, cfg_scale, resize_method, - a_prompt, n_prompt, sdxl_model, supir_model, keep_model_loaded): - torch.manual_seed(seed) - torch.cuda.manual_seed_all(seed) - comfy.model_management.unload_all_models() + def process(self, steps, image, color_fix_type, seed, scale_by, cfg_scale, resize_method, s_churn, s_noise, + control_scale, cfg_scale_start, control_scale_start, restoration_scale, keep_model_loaded, + a_prompt, n_prompt, sdxl_model, supir_model, captions=""): + + device = comfy.model_management.get_torch_device() image = image.to(device) SUPIR_MODEL_PATH = folder_paths.get_full_path("checkpoints", supir_model) @@ -78,30 +81,27 @@ class SUPIR_Upscale: autocast_condition = dtype == torch.float16 or torch.bfloat16 and not comfy.model_management.is_device_mps(device) with torch.autocast(comfy.model_management.get_autocast_device(device), dtype=dtype) if autocast_condition else nullcontext(): image, = ImageScaleBy.upscale(self, image, resize_method, scale_by) - - # Assuming 'image' is a PyTorch tensor with shape [B, H, W, C] and you want to resize it. B, H, W, C = image.shape - - # Calculate the new height and width, rounding down to the nearest multiple of 64. new_height = H // 64 * 64 new_width = W // 64 * 64 - - # Reorder to [B, C, H, W] before using interpolate. image = image.permute(0, 3, 1, 2).contiguous() - - # Resize the image tensor. - resized_image = F.interpolate(image, size=(new_height, new_width), mode='bicubic', align_corners=False) + resized_image = F.interpolate(image, size=(new_height, new_width), mode=resize_method, align_corners=False) - captions = [''] - print(captions) - + captions_list = [] + captions_list.append(captions) + print(captions_list) + + use_linear_CFG = cfg_scale_start > 0 + use_linear_control_scale = control_scale_start > 0 # # step 3: Diffusion Process - samples = self.model.batchify_sample(resized_image, captions, num_steps=steps, restoration_scale= -1, s_churn=5, - s_noise=1.003, cfg_scale=cfg_scale, control_scale= 1, seed=seed, + samples = self.model.batchify_sample(resized_image, captions_list, num_steps=steps, restoration_scale= restoration_scale, s_churn=s_churn, + s_noise=s_noise, cfg_scale=cfg_scale, control_scale=control_scale, seed=seed, num_samples=1, p_p=a_prompt, n_p=n_prompt, color_fix_type=color_fix_type, - use_linear_CFG=False, use_linear_control_scale=False, - cfg_scale_start=1.0, control_scale_start=0) + use_linear_CFG=use_linear_CFG, use_linear_control_scale=use_linear_control_scale, + cfg_scale_start=cfg_scale_start, control_scale_start=control_scale_start) # save + if not keep_model_loaded: + self.model = None print(samples.shape) samples = samples.permute(0, 2, 3, 1).cpu() diff --git a/sgm/modules/attention.py b/sgm/modules/attention.py index 8d1f64b..7cca67f 100644 --- a/sgm/modules/attention.py +++ b/sgm/modules/attention.py @@ -291,10 +291,10 @@ class MemoryEfficientCrossAttention(nn.Module): self, query_dim, context_dim=None, heads=8, dim_head=64, dropout=0.0, **kwargs ): super().__init__() - print( - f"Setting up {self.__class__.__name__}. Query dim is {query_dim}, context_dim is {context_dim} and using " - f"{heads} heads with a dimension of {dim_head}." - ) + #print( + # f"Setting up {self.__class__.__name__}. Query dim is {query_dim}, context_dim is {context_dim} and using " + # f"{heads} heads with a dimension of {dim_head}." + #) inner_dim = dim_head * heads context_dim = default(context_dim, query_dim) @@ -438,8 +438,8 @@ class BasicTransformerBlock(nn.Module): self.norm2 = nn.LayerNorm(dim) self.norm3 = nn.LayerNorm(dim) self.checkpoint = checkpoint - if self.checkpoint: - print(f"{self.__class__.__name__} is using checkpointing") + #if self.checkpoint: + #print(f"{self.__class__.__name__} is using checkpointing") def forward( self, x, context=None, additional_tokens=None, n_times_crossframe_attn_in_self=0 @@ -565,10 +565,10 @@ class SpatialTransformer(nn.Module): context_dim = [context_dim] if exists(context_dim) and isinstance(context_dim, list): if depth != len(context_dim): - print( - f"WARNING: {self.__class__.__name__}: Found context dims {context_dim} of depth {len(context_dim)}, " - f"which does not match the specified 'depth' of {depth}. Setting context_dim to {depth * [context_dim[0]]} now." - ) + #print( + # f"WARNING: {self.__class__.__name__}: Found context dims {context_dim} of depth {len(context_dim)}, " + # f"which does not match the specified 'depth' of {depth}. Setting context_dim to {depth * [context_dim[0]]} now." + # ) # depth does not match context dims. assert all( map(lambda x: x == context_dim[0], context_dim)