diff --git a/sgm/modules/diffusionmodules/sampling.py b/sgm/modules/diffusionmodules/sampling.py index 40915e4..9a4f9f8 100644 --- a/sgm/modules/diffusionmodules/sampling.py +++ b/sgm/modules/diffusionmodules/sampling.py @@ -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, )