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

Tiled sampler consequential runs workaround

Browse Source 
Copy the conds so they are not permanently modified by the Tiled samplers
This commit is contained in:
kijai
2024-03-26 10:34:17 +02:00
parent 3f340a6b96
commit 2ce3ced581
1 changed files with 23 additions and 18 deletions
Download Patch File Download Diff File Expand all files Collapse all files
sgm/modules/diffusionmodules/sampling.py
+23 -18
View File
@@ -18,6 +18,7 @@ from ...modules.diffusionmodules.sampling_utils import (
to_sigma,
)
from ...util import append_dims, default, instantiate_from_config
import copy
DEFAULT_GUIDER = {"target": ".sgm.modules.diffusionmodules.guiders.IdentityGuider"}
@@ -462,18 +463,20 @@ class TiledRestoreEDMSampler(RestoreEDMSampler):
def __call__(self, denoiser, x, cond, uc=None, num_steps=None, x_center=None, control_scale=1.0,
use_linear_control_scale=False, control_scale_start=0.0):
use_local_prompt = isinstance(cond, list)
cond_copy = copy.deepcopy(cond)
uc_copy = copy.deepcopy(uc)
use_local_prompt = isinstance(cond_copy, list)
b, _, h, w = x.shape
latent_tiles_iterator = _sliding_windows(h, w, self.tile_size, self.tile_stride)
tile_weights = self.tile_weights.repeat(b, 1, 1, 1)
if not use_local_prompt:
LQ_latent = cond['control']
LQ_latent = cond_copy['control']
else:
assert len(cond) == len(latent_tiles_iterator), "Number of local prompts should be equal to number of tiles"
LQ_latent = cond[0]['control']
assert len(cond_copy) == len(latent_tiles_iterator), "Number of local prompts should be equal to number of tiles"
LQ_latent = cond_copy[0]['control']
clean_LQ_latent = x_center
x, s_in, sigmas, num_sigmas, cond, uc = self.prepare_sampling_loop(
x, cond, uc, num_steps
x, s_in, sigmas, num_sigmas, cond_copy, uc_copy = self.prepare_sampling_loop(
x, cond_copy, uc_copy, num_steps
)
pbar_comfy = comfy.utils.ProgressBar(num_sigmas)
for _idx, i in enumerate(self.get_sigma_gen(num_sigmas)):
@@ -490,18 +493,18 @@ class TiledRestoreEDMSampler(RestoreEDMSampler):
_eps_noise = eps_noise[:, :, hi:hi_end, wi:wi_end]
x_center_tile = clean_LQ_latent[:, :, hi:hi_end, wi:wi_end]
if use_local_prompt:
_cond = cond[j]
_cond = cond_copy[j]
else:
_cond = cond
_cond = cond_copy
_cond['control'] = LQ_latent[:, :, hi:hi_end, wi:wi_end]
uc['control'] = LQ_latent[:, :, hi:hi_end, wi:wi_end]
uc_copy['control'] = LQ_latent[:, :, hi:hi_end, wi:wi_end]
_x = self.sampler_step(
s_in * sigmas[i],
s_in * sigmas[i + 1],
denoiser,
x_tile,
_cond,
uc,
uc_copy,
gamma,
x_center_tile,
eps_noise=_eps_noise,
@@ -684,15 +687,17 @@ class TiledRestoreDPMPP2MSampler(RestoreDPMPP2MSampler):
latent_tiles_iterator = _sliding_windows(h, w, self.tile_size, self.tile_stride)
print(f"Image divided into {len(latent_tiles_iterator)} tiles")
print("Conds received: ", len(cond))
cond_copy = copy.deepcopy(cond)
uc_copy = copy.deepcopy(uc)
tile_weights = self.tile_weights.repeat(b, 1, 1, 1)
if not use_local_prompt:
LQ_latent = cond['control']
else:
assert len(cond) == len(latent_tiles_iterator), "Number of local prompts should be equal to number of tiles"
LQ_latent = cond[0]['control']
assert len(cond_copy) == len(latent_tiles_iterator), "Number of local prompts should be equal to number of tiles"
LQ_latent = cond_copy[0]['control']
print("LQ_latent shape: ",LQ_latent.shape)
x, s_in, sigmas, num_sigmas, cond, uc = self.prepare_sampling_loop(
x, cond, uc, num_steps
x, s_in, sigmas, num_sigmas, cond_copy, uc_copy = self.prepare_sampling_loop(
x, cond_copy, uc_copy, num_steps
)
sigmas_min, sigmas_max = sigmas[-2].cpu(), sigmas[0].cpu()
sigmas_new = get_sigmas_karras(self.num_steps, sigmas_min, sigmas_max, device=x.device)
@@ -718,11 +723,11 @@ class TiledRestoreDPMPP2MSampler(RestoreDPMPP2MSampler):
else:
old_denoised_tile = None
if use_local_prompt:
_cond = cond[j]
_cond = cond_copy[j]
else:
_cond = cond
_cond = cond_copy
_cond['control'] = LQ_latent[:, :, hi:hi_end, wi:wi_end]
uc['control'] = LQ_latent[:, :, hi:hi_end, wi:wi_end]
uc_copy['control'] = LQ_latent[:, :, hi:hi_end, wi:wi_end]
_x, _old_denoised = self.sampler_step(
old_denoised_tile,
None if i == 0 else s_in * sigmas[i - 1],
@@ -731,7 +736,7 @@ class TiledRestoreDPMPP2MSampler(RestoreDPMPP2MSampler):
denoiser,
x_tile,
_cond,
uc=uc,
uc=uc_copy,
eps_noise=_eps_noise,
control_scale=control_scale,
)