nodes/kijai-ComfyUI-SUPIR
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Add tiled_vae and progress bar

Browse Source
This commit is contained in:
kijai
2024-02-29 02:34:08 +02:00
parent 06c1b3b6c9
commit fef3a9c11a
2 changed files with 12 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files
nodes.py
+8 -4
View File
@@ -41,6 +41,9 @@ class SUPIR_Upscale:
"default": 'Wavelet'
}),
"keep_model_loaded": ("BOOLEAN", {"default": True}),
"use_tiled_vae": ("BOOLEAN", {"default": True}),
"encoder_tile_size_pixels": ("INT", {"default": 512, "min": 64, "max": 8192, "step": 64}),
"decoder_tile_size_latent": ("INT", {"default": 64, "min": 64, "max": 8192, "step": 64}),
},
"optional": {
"captions": ("STRING", {"forceInput": True, "multiline": False, "default": "",}),
@@ -55,9 +58,9 @@ class SUPIR_Upscale:
CATEGORY = "SUPIR"
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=""):
def process(self, steps, image, color_fix_type, seed, scale_by, cfg_scale, resize_method, s_churn, s_noise, encoder_tile_size_pixels, decoder_tile_size_latent,
control_scale, cfg_scale_start, control_scale_start, restoration_scale, keep_model_loaded,
a_prompt, n_prompt, sdxl_model, supir_model, use_tiled_vae, captions=""):
device = comfy.model_management.get_torch_device()
@@ -77,6 +80,8 @@ class SUPIR_Upscale:
self.model.load_state_dict(supir_state_dict, strict=False)
self.model.load_state_dict(sdxl_state_dict, strict=False)
self.model.to(device).to(dtype)
if use_tiled_vae:
self.model.init_tile_vae(encoder_tile_size=encoder_tile_size_pixels, decoder_tile_size=decoder_tile_size_latent)
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():
@@ -102,7 +107,6 @@ class SUPIR_Upscale:
# save
if not keep_model_loaded:
self.model = None
print(samples.shape)
samples = samples.permute(0, 2, 3, 1).cpu()
return(samples,)
sgm/modules/diffusionmodules/sampling.py
+4 -2
View File
@@ -366,7 +366,8 @@ class DPMPP2MSampler(BaseDiffusionSampler):
return x
import comfy.utils
class RestoreEDMSampler(SingleStepDiffusionSampler):
def __init__(
self, s_churn=0.0, s_tmin=0.0, s_tmax=float("inf"), s_noise=1.0, restore_cfg=4.0,
@@ -417,7 +418,7 @@ class RestoreEDMSampler(SingleStepDiffusionSampler):
x, s_in, sigmas, num_sigmas, cond, uc = self.prepare_sampling_loop(
x, cond, uc, num_steps
)
pbar_comfy = comfy.utils.ProgressBar(num_sigmas)
for _idx, i in enumerate(self.get_sigma_gen(num_sigmas)):
gamma = (
min(self.s_churn / (num_sigmas - 1), 2**0.5 - 1)
@@ -437,6 +438,7 @@ class RestoreEDMSampler(SingleStepDiffusionSampler):
use_linear_control_scale=use_linear_control_scale,
control_scale_start=control_scale_start,
)
pbar_comfy.update(1)
return x
def to_d_center(denoised, x_center, x):