extensions/multidiffusion-upscaler-for-automatic1111/tile_methods/abstractdiffusion.py

from tile_utils.utils import *


class AbstractDiffusion:

    def __init__(self, p: Processing, sampler: Sampler):
        self.method = self.__class__.__name__
        self.p: Processing = p
        self.pbar = None

        # sampler
        self.sampler_name = p.sampler_name
        self.sampler_raw = sampler
        self.sampler = sampler

        # fix. Kdiff 'AND' support and image editing model support
        if self.is_kdiff and not hasattr(self, 'is_edit_model'):
            self.is_edit_model = (shared.sd_model.cond_stage_key == "edit"      # "txt"
                and self.sampler.model_wrap_cfg.image_cfg_scale is not None 
                and self.sampler.model_wrap_cfg.image_cfg_scale != 1.0)

        # cache. final result of current sampling step, [B, C=4, H//8, W//8]
        # avoiding overhead of creating new tensors and weight summing
        self.x_buffer: Tensor = None
        self.w: int = int(self.p.width  // opt_f)       # latent size
        self.h: int = int(self.p.height // opt_f)
        # weights for background & grid bboxes
        self.weights: Tensor = torch.zeros((1, 1, self.h, self.w), device=devices.device, dtype=torch.float32)

        # FIXME: I'm trying to count the step correctly but it's not working
        self.step_count = 0         
        self.inner_loop_count = 0  
        self.kdiff_step = -1

        # ext. Grid tiling painting (grid bbox)
        self.enable_grid_bbox: bool = False
        self.tile_w: int = None
        self.tile_h: int = None
        self.tile_bs: int = None
        self.num_tiles: int = None
        self.num_batches: int = None
        self.batched_bboxes: List[List[BBox]] = []

        # ext. Region Prompt Control (custom bbox)
        self.enable_custom_bbox: bool = False
        self.custom_bboxes: List[CustomBBox] = []
        self.cond_basis: Cond = None
        self.uncond_basis: Uncond = None
        self.draw_background: bool = True       # by default we draw major prompts in grid tiles
        self.causal_layers: bool = None

        # ext. Noise Inversion (noise inversion)
        self.noise_inverse_enabled: bool = False
        self.noise_inverse_steps: int = 0
        self.noise_inverse_retouch: float = None
        self.noise_inverse_renoise_strength: float = None
        self.noise_inverse_renoise_kernel: int = None
        self.noise_inverse_get_cache = None
        self.noise_inverse_set_cache = None
        self.sample_img2img_original = None

        # ext. ControlNet
        self.enable_controlnet: bool = False
        self.controlnet_script: ModuleType = None
        self.control_tensor_batch: List[List[Tensor]] = []
        self.control_params: Dict[str, Tensor] = {}
        self.control_tensor_cpu: bool = None
        self.control_tensor_custom: List[List[Tensor]] = []

        # ext. StableSR
        self.enable_stablesr: bool = False
        self.stablesr_script: ModuleType = None
        self.stablesr_tensor: Tensor = None
        self.stablesr_tensor_batch: List[Tensor] = []
        self.stablesr_tensor_custom: List[Tensor] = []

    @property
    def is_kdiff(self):
        return isinstance(self.sampler_raw, KDiffusionSampler)

    @property
    def is_ddim(self):
        return isinstance(self.sampler_raw, CompVisSampler)

    def update_pbar(self):
        if self.pbar.n >= self.pbar.total:
            self.pbar.close()
        else:
            if self.step_count == state.sampling_step:
                self.inner_loop_count += 1
                if self.inner_loop_count < self.total_bboxes:
                    self.pbar.update()
            else:
                self.step_count = state.sampling_step
                self.inner_loop_count = 0

    def reset_buffer(self, x_in:Tensor):
        # Judge if the shape of x_in is the same as the shape of x_buffer
        if self.x_buffer is None or self.x_buffer.shape != x_in.shape:
            self.x_buffer = torch.zeros_like(x_in, device=x_in.device, dtype=x_in.dtype)
        else:
            self.x_buffer.zero_()

    def init_done(self):
        '''
          Call this after all `init_*`, settings are done, now perform:
            - settings sanity check
            - pre-computations, cache init
            - anything thing needed before denoising starts
        '''

        self.total_bboxes = 0
        if self.enable_grid_bbox:   self.total_bboxes += self.num_batches
        if self.enable_custom_bbox: self.total_bboxes += len(self.custom_bboxes)
        assert self.total_bboxes > 0, "Nothing to paint! No background to draw and no custom bboxes were provided."

        self.pbar = tqdm(total=(self.total_bboxes) * state.sampling_steps, desc=f"{self.method} Sampling: ")

    ''' ↓↓↓ cond_dict utils ↓↓↓ '''

    def _tcond_key(self, cond_dict:CondDict) -> str:
        return 'crossattn' if 'crossattn' in cond_dict else 'c_crossattn'

    def get_tcond(self, cond_dict:CondDict) -> Tensor:
        tcond = cond_dict[self._tcond_key(cond_dict)]
        if isinstance(tcond, list): tcond = tcond[0]
        return tcond

    def set_tcond(self, cond_dict:CondDict, tcond:Tensor):
        key = self._tcond_key(cond_dict)
        if isinstance(cond_dict[key], list): tcond = [tcond]
        cond_dict[key] = tcond

    def _icond_key(self, cond_dict:CondDict) -> str:
        return 'c_adm' if shared.sd_model.model.conditioning_key in ['crossattn-adm', 'adm'] else 'c_concat'

    def get_icond(self, cond_dict:CondDict) -> Tensor:
        ''' icond differs for different models (inpaint/unclip model) '''
        key = self._icond_key(cond_dict)
        icond = cond_dict[key]
        if isinstance(icond, list): icond = icond[0]
        return icond

    def set_icond(self, cond_dict:CondDict, icond:Tensor):
        key = self._icond_key(cond_dict)
        if isinstance(cond_dict[key], list): icond = [icond]
        cond_dict[key] = icond

    def _vcond_key(self, cond_dict:CondDict) -> Optional[str]:
        return 'vector' if 'vector' in cond_dict else None

    def get_vcond(self, cond_dict:CondDict) -> Optional[Tensor]:
        ''' vector for SDXL '''
        key = self._vcond_key(cond_dict)
        return cond_dict.get(key)

    def set_vcond(self, cond_dict:CondDict, vcond:Optional[Tensor]):
        key = self._vcond_key(cond_dict)
        if key is not None:
            cond_dict[key] = vcond

    def make_cond_dict(self, cond_in:CondDict, tcond:Tensor, icond:Tensor, vcond:Tensor=None) -> CondDict:
        ''' copy & replace the content, returns a new object '''
        cond_out = cond_in.copy()
        self.set_tcond(cond_out, tcond)
        self.set_icond(cond_out, icond)
        self.set_vcond(cond_out, vcond)
        return cond_out

    ''' ↓↓↓ extensive functionality ↓↓↓ '''

    @grid_bbox
    def init_grid_bbox(self, tile_w:int, tile_h:int, overlap:int, tile_bs:int):
        self.enable_grid_bbox = True

        self.tile_w = min(tile_w, self.w)
        self.tile_h = min(tile_h, self.h)
        overlap = max(0, min(overlap, min(tile_w, tile_h) - 4))
        # split the latent into overlapped tiles, then batching
        # weights basically indicate how many times a pixel is painted
        bboxes, weights = split_bboxes(self.w, self.h, self.tile_w, self.tile_h, overlap, self.get_tile_weights())
        self.weights += weights
        self.num_tiles = len(bboxes)
        self.num_batches = math.ceil(self.num_tiles / tile_bs)
        self.tile_bs = math.ceil(len(bboxes) / self.num_batches)          # optimal_batch_size
        self.batched_bboxes = [bboxes[i*self.tile_bs:(i+1)*self.tile_bs] for i in range(self.num_batches)]

    @grid_bbox
    def get_tile_weights(self) -> Union[Tensor, float]:
        return 1.0


    @custom_bbox
    def init_custom_bbox(self, bbox_settings:Dict[int,BBoxSettings], draw_background:bool, causal_layers:bool):
        self.enable_custom_bbox = True

        self.causal_layers = causal_layers
        self.draw_background = draw_background
        if not draw_background:
            self.enable_grid_bbox = False
            self.weights.zero_()

        self.custom_bboxes: List[CustomBBox] = []
        for bbox_setting in bbox_settings.values():
            e, x, y, w, h, p, n, blend_mode, feather_ratio, seed = bbox_setting
            if not e or x > 1.0 or y > 1.0 or w <= 0.0 or h <= 0.0: continue
            x = int(x * self.w)
            y = int(y * self.h)
            w = math.ceil(w * self.w)
            h = math.ceil(h * self.h)
            x = max(0, x)
            y = max(0, y)
            w = min(self.w - x, w)
            h = min(self.h - y, h)
            self.custom_bboxes.append(CustomBBox(x, y, w, h, p, n, blend_mode, feather_ratio, seed))

        if len(self.custom_bboxes) == 0:
            self.enable_custom_bbox = False
            return

        # prepare cond
        p = self.p
        prompts = p.all_prompts[:p.batch_size]
        neg_prompts = p.all_negative_prompts[:p.batch_size]
        for bbox in self.custom_bboxes:
            bbox.cond, bbox.extra_network_data = Condition.get_custom_cond(prompts, bbox.prompt, p.steps, p.styles)
            bbox.uncond = Condition.get_uncond(Prompt.append_prompt(neg_prompts, bbox.neg_prompt), p.steps, p.styles)
        self.cond_basis = Condition.get_cond(prompts, p.steps)
        self.uncond_basis = Condition.get_uncond(neg_prompts, p.steps)

    @custom_bbox
    def reconstruct_custom_cond(self, org_cond:CondDict, custom_cond:Cond, custom_uncond:Uncond, bbox:CustomBBox) -> Tuple[List, Tensor, Uncond, Tensor]:
        image_conditioning = None
        if isinstance(org_cond, dict):
            icond = self.get_icond(org_cond)
            if icond.shape[2:] == (self.h, self.w):    # img2img
                icond = icond[bbox.slicer]
            image_conditioning = icond

        sampler_step = self.sampler.model_wrap_cfg.step
        tensor = Condition.reconstruct_cond(custom_cond, sampler_step)
        custom_uncond = Condition.reconstruct_uncond(custom_uncond, sampler_step)
        return tensor, custom_uncond, image_conditioning

    @custom_bbox
    def kdiff_custom_forward(self, x_tile:Tensor, sigma_in:Tensor, original_cond:CondDict, bbox_id:int, bbox:CustomBBox, forward_func:Callable) -> Tensor:
        '''
        The inner kdiff noise prediction is usually batched.
        We need to unwrap the inside loop to simulate the batched behavior.
        This can be extremely tricky.
        '''

        sampler_step = self.sampler.model_wrap_cfg.step
        if self.kdiff_step != sampler_step:
            self.kdiff_step = sampler_step
            self.kdiff_step_bbox = [-1 for _ in range(len(self.custom_bboxes))]
            self.tensor = {}        # {int: Tensor[cond]}
            self.uncond = {}        # {int: Tensor[cond]}
            self.image_cond_in = {}
            # Initialize global prompts just for estimate the behavior of kdiff
            self.real_tensor = Condition.reconstruct_cond(self.cond_basis, sampler_step)
            self.real_uncond = Condition.reconstruct_uncond(self.uncond_basis, sampler_step)
            # reset the progress for all bboxes
            self.a = [0 for _ in range(len(self.custom_bboxes))]

        if self.kdiff_step_bbox[bbox_id] != sampler_step:
            # When a new step starts for a bbox, we need to judge whether the tensor is batched.
            self.kdiff_step_bbox[bbox_id] = sampler_step

            tensor, uncond, image_cond_in = self.reconstruct_custom_cond(original_cond, bbox.cond, bbox.uncond, bbox)

            if self.real_tensor.shape[1] == self.real_uncond.shape[1]:
                if shared.batch_cond_uncond:
                    # when the real tensor is with equal length, all information is contained in x_tile.
                    # we simulate the batched behavior and compute all the tensors in one go.
                    if tensor.shape[1] == uncond.shape[1]:
                        # When our prompt tensor is with equal length, we can directly their code.
                        if not self.is_edit_model:
                            cond = torch.cat([tensor, uncond])
                        else:
                            cond = torch.cat([tensor, uncond, uncond])
                        self.set_custom_controlnet_tensors(bbox_id, x_tile.shape[0])
                        self.set_custom_stablesr_tensors(bbox_id)
                        return forward_func(
                            x_tile, 
                            sigma_in, 
                            cond=self.make_cond_dict(original_cond, cond, image_cond_in),
                        )
                    else:
                        # When not, we need to pass the tensor to UNet separately.
                        x_out = torch.zeros_like(x_tile)
                        cond_size = tensor.shape[0]
                        self.set_custom_controlnet_tensors(bbox_id, cond_size)
                        self.set_custom_stablesr_tensors(bbox_id)
                        cond_out = forward_func(
                            x_tile  [:cond_size], 
                            sigma_in[:cond_size], 
                            cond=self.make_cond_dict(original_cond, tensor, image_cond_in[:cond_size]),
                        )
                        uncond_size = uncond.shape[0]
                        self.set_custom_controlnet_tensors(bbox_id, uncond_size)
                        self.set_custom_stablesr_tensors(bbox_id)
                        uncond_out = forward_func(
                            x_tile  [cond_size:cond_size+uncond_size], 
                            sigma_in[cond_size:cond_size+uncond_size], 
                            cond=self.make_cond_dict(original_cond, uncond, image_cond_in[cond_size:cond_size+uncond_size]),
                        )
                        x_out[:cond_size] = cond_out
                        x_out[cond_size:cond_size+uncond_size] = uncond_out
                        if self.is_edit_model:
                            x_out[cond_size+uncond_size:] = uncond_out
                        return x_out
                
            # otherwise, the x_tile is only a partial batch. 
            # We have to denoise in different runs.
            # We store the prompt and neg_prompt tensors for current bbox
            self.tensor[bbox_id] = tensor
            self.uncond[bbox_id] = uncond
            self.image_cond_in[bbox_id] = image_cond_in

        # Now we get current batch of prompt and neg_prompt tensors
        tensor: Tensor = self.tensor[bbox_id]
        uncond: Tensor = self.uncond[bbox_id]
        batch_size = x_tile.shape[0]
        # get the start and end index of the current batch
        a = self.a[bbox_id]
        b = a + batch_size
        self.a[bbox_id] += batch_size

        if self.real_tensor.shape[1] == self.real_uncond.shape[1]:
            # When use --lowvram or --medvram, kdiff will slice the cond and uncond with [a:b]
            # So we need to slice our tensor and uncond with the same index as original kdiff.
            
            # --- original code in kdiff ---
            # if not self.is_edit_model:
            #     cond = torch.cat([tensor, uncond])
            # else:
            #     cond = torch.cat([tensor, uncond, uncond])
            # cond = cond[a:b]
            # ------------------------------
            
            # The original kdiff code is to concat and then slice, but this cannot apply to
            # our custom prompt tensor when tensor.shape[1] != uncond.shape[1]. So we adapt it.
            cond_in, uncond_in = None, None
            # Slice the [prompt, neg prompt, (possibly) neg prompt] with [a:b]
            if not self.is_edit_model:
                if   b <= tensor.shape[0]: cond_in = tensor[a:b]
                elif a >= tensor.shape[0]: cond_in = uncond[a-tensor.shape[0]:b-tensor.shape[0]]
                else:
                    cond_in = tensor[a:]
                    uncond_in = uncond[:b-tensor.shape[0]]
            else:
                if b <= tensor.shape[0]:
                    cond_in = tensor[a:b]
                elif b > tensor.shape[0] and b <= tensor.shape[0] + uncond.shape[0]:
                    if a>= tensor.shape[0]:
                        cond_in = uncond[a-tensor.shape[0]:b-tensor.shape[0]]
                    else:
                        cond_in = tensor[a:]
                        uncond_in = uncond[:b-tensor.shape[0]]
                else:
                    if a >= tensor.shape[0] + uncond.shape[0]:
                        cond_in = uncond[a-tensor.shape[0]-uncond.shape[0]:b-tensor.shape[0]-uncond.shape[0]]
                    elif a >= tensor.shape[0]:
                        cond_in = torch.cat([uncond[a-tensor.shape[0]:], uncond[:b-tensor.shape[0]-uncond.shape[0]]])

            if uncond_in is None or tensor.shape[1] == uncond.shape[1]:
                # If the tensor can be passed to UNet in one go, do it.
                if uncond_in is not None:
                    cond_in = torch.cat([cond_in, uncond_in])
                self.set_custom_controlnet_tensors(bbox_id, x_tile.shape[0])
                self.set_custom_stablesr_tensors(bbox_id)
                return forward_func(
                    x_tile, 
                    sigma_in, 
                    cond=self.make_cond_dict(original_cond, cond_in, self.image_cond_in[bbox_id]),
                )
            else:
                # If not, we need to pass the tensor to UNet separately.
                x_out = torch.zeros_like(x_tile)
                cond_size = cond_in.shape[0]
                self.set_custom_controlnet_tensors(bbox_id, cond_size)
                self.set_custom_stablesr_tensors(bbox_id)
                cond_out = forward_func(
                    x_tile  [:cond_size], 
                    sigma_in[:cond_size], 
                    cond=self.make_cond_dict(original_cond, cond_in, self.image_cond_in[bbox_id])
                )
                self.set_custom_controlnet_tensors(bbox_id, uncond_in.shape[0])
                self.set_custom_stablesr_tensors(bbox_id)
                uncond_out = forward_func(
                    x_tile  [cond_size:], 
                    sigma_in[cond_size:], 
                    cond=self.make_cond_dict(original_cond, uncond_in, self.image_cond_in[bbox_id])
                )
                x_out[:cond_size] = cond_out
                x_out[cond_size:] = uncond_out
                return x_out

        # If the original prompt is with different length, 
        # kdiff will deal with the cond and uncond separately.
        # Hence we also deal with the tensor and uncond separately.
        # get the start and end index of the current batch

        if a < tensor.shape[0]:
            # Deal with custom prompt tensor
            if not self.is_edit_model:
                c_crossattn = tensor[a:b]
            else:
                c_crossattn = torch.cat([tensor[a:b]], uncond)
            self.set_custom_controlnet_tensors(bbox_id, x_tile.shape[0])
            self.set_custom_stablesr_tensors(bbox_id)
            # complete this batch.
            return forward_func(
                x_tile, 
                sigma_in, 
                cond=self.make_cond_dict(original_cond, c_crossattn, self.image_cond_in[bbox_id])
            )
        else:
            # if the cond is finished, we need to process the uncond.
            self.set_custom_controlnet_tensors(bbox_id, uncond.shape[0])
            self.set_custom_stablesr_tensors(bbox_id)
            return forward_func(
                x_tile, 
                sigma_in, 
                cond=self.make_cond_dict(original_cond, uncond, self.image_cond_in[bbox_id])
            )

    @custom_bbox
    def ddim_custom_forward(self, x:Tensor, cond_in:CondDict, bbox:CustomBBox, ts:Tensor, forward_func:Callable, *args, **kwargs) -> Tensor:
        ''' draw custom bbox '''

        tensor, uncond, image_conditioning = self.reconstruct_custom_cond(cond_in, bbox.cond, bbox.uncond, bbox)

        cond = tensor
        # for DDIM, shapes definitely match. So we dont need to do the same thing as in the KDIFF sampler.
        if uncond.shape[1] < cond.shape[1]:
            last_vector = uncond[:, -1:]
            last_vector_repeated = last_vector.repeat([1, cond.shape[1] - uncond.shape[1], 1])
            uncond = torch.hstack([uncond, last_vector_repeated])
        elif uncond.shape[1] > cond.shape[1]:
            uncond = uncond[:, :cond.shape[1]]

        # Wrap the image conditioning back up since the DDIM code can accept the dict directly.
        # Note that they need to be lists because it just concatenates them later.
        if image_conditioning is not None:
            cond   = self.make_cond_dict(cond_in, cond,   image_conditioning)
            uncond = self.make_cond_dict(cond_in, uncond, image_conditioning)
        
        # We cannot determine the batch size here for different methods, so delay it to the forward_func.
        return forward_func(x, cond, ts, unconditional_conditioning=uncond, *args, **kwargs)


    @controlnet
    def init_controlnet(self, controlnet_script:ModuleType, control_tensor_cpu:bool):
        self.enable_controlnet = True

        self.controlnet_script = controlnet_script
        self.control_tensor_cpu = control_tensor_cpu
        self.control_tensor_batch = None
        self.control_params = None
        self.control_tensor_custom = []
        
        self.prepare_controlnet_tensors()

    @controlnet
    def reset_controlnet_tensors(self):
        if not self.enable_controlnet: return
        if self.control_tensor_batch is None: return

        for param_id in range(len(self.control_params)):
            self.control_params[param_id].hint_cond = self.org_control_tensor_batch[param_id]

    @controlnet
    def prepare_controlnet_tensors(self, refresh:bool=False):
        ''' Crop the control tensor into tiles and cache them '''

        if not refresh:
            if self.control_tensor_batch is not None or self.control_params is not None: return

        if not self.enable_controlnet or self.controlnet_script is None: return

        latest_network = self.controlnet_script.latest_network
        if latest_network is None or not hasattr(latest_network, 'control_params'): return

        self.control_params = latest_network.control_params
        tensors = [param.hint_cond for param in latest_network.control_params]
        self.org_control_tensor_batch = tensors

        if len(tensors) == 0: return

        self.control_tensor_batch = []
        for i in range(len(tensors)):
            control_tile_list = []
            control_tensor = tensors[i]
            for bboxes in self.batched_bboxes:
                single_batch_tensors = []
                for bbox in bboxes:
                    if len(control_tensor.shape) == 3:
                        control_tensor.unsqueeze_(0)
                    control_tile = control_tensor[:, :, bbox[1]*opt_f:bbox[3]*opt_f, bbox[0]*opt_f:bbox[2]*opt_f]
                    single_batch_tensors.append(control_tile)
                control_tile = torch.cat(single_batch_tensors, dim=0)
                if self.control_tensor_cpu:
                    control_tile = control_tile.cpu()
                control_tile_list.append(control_tile)
            self.control_tensor_batch.append(control_tile_list)

            if len(self.custom_bboxes) > 0:
                custom_control_tile_list = []
                for bbox in self.custom_bboxes:
                    if len(control_tensor.shape) == 3:
                        control_tensor.unsqueeze_(0)
                    control_tile = control_tensor[:, :, bbox[1]*opt_f:bbox[3]*opt_f, bbox[0]*opt_f:bbox[2]*opt_f]
                    if self.control_tensor_cpu:
                        control_tile = control_tile.cpu()
                    custom_control_tile_list.append(control_tile)
                self.control_tensor_custom.append(custom_control_tile_list)

    @controlnet
    def switch_controlnet_tensors(self, batch_id:int, x_batch_size:int, tile_batch_size:int, is_denoise=False):
        if not self.enable_controlnet: return
        if self.control_tensor_batch is None: return

        for param_id in range(len(self.control_params)):
            control_tile = self.control_tensor_batch[param_id][batch_id]
            if self.is_kdiff:
                all_control_tile = []
                for i in range(tile_batch_size):
                    this_control_tile = [control_tile[i].unsqueeze(0)] * x_batch_size
                    all_control_tile.append(torch.cat(this_control_tile, dim=0))
                control_tile = torch.cat(all_control_tile, dim=0)                                           
            else:
                control_tile = control_tile.repeat([x_batch_size if is_denoise else x_batch_size * 2, 1, 1, 1])
            self.control_params[param_id].hint_cond = control_tile.to(devices.device)

    @controlnet
    def set_custom_controlnet_tensors(self, bbox_id:int, repeat_size:int):
        if not self.enable_controlnet: return
        if not len(self.control_tensor_custom): return
        
        for param_id in range(len(self.control_params)):
            control_tensor = self.control_tensor_custom[param_id][bbox_id].to(devices.device)
            self.control_params[param_id].hint_cond = control_tensor.repeat((repeat_size, 1, 1, 1))


    @stablesr
    def init_stablesr(self, stablesr_script:ModuleType):
        if stablesr_script.stablesr_model is None: return
        self.stablesr_script = stablesr_script
        def set_image_hook(latent_image):
            self.enable_stablesr = True
            self.stablesr_tensor = latent_image
            self.stablesr_tensor_batch = []
            for bboxes in self.batched_bboxes:
                single_batch_tensors = []
                for bbox in bboxes:
                    stablesr_tile = self.stablesr_tensor[:, :, bbox[1]:bbox[3], bbox[0]:bbox[2]]
                    single_batch_tensors.append(stablesr_tile)
                stablesr_tile = torch.cat(single_batch_tensors, dim=0)
                self.stablesr_tensor_batch.append(stablesr_tile)
            if len(self.custom_bboxes) > 0:
                self.stablesr_tensor_custom = []
                for bbox in self.custom_bboxes:
                    stablesr_tile = self.stablesr_tensor[:, :, bbox[1]:bbox[3], bbox[0]:bbox[2]]
                    self.stablesr_tensor_custom.append(stablesr_tile)

        stablesr_script.stablesr_model.set_image_hooks['TiledDiffusion'] = set_image_hook

    @stablesr
    def reset_stablesr_tensors(self):
        if not self.enable_stablesr: return
        if self.stablesr_script.stablesr_model is None: return
        self.stablesr_script.stablesr_model.latent_image = self.stablesr_tensor

    @stablesr
    def switch_stablesr_tensors(self, batch_id:int):
        if not self.enable_stablesr: return
        if self.stablesr_script.stablesr_model is None: return
        if self.stablesr_tensor_batch is None: return
        self.stablesr_script.stablesr_model.latent_image = self.stablesr_tensor_batch[batch_id]

    @stablesr
    def set_custom_stablesr_tensors(self, bbox_id:int):
        if not self.enable_stablesr: return
        if self.stablesr_script.stablesr_model is None: return
        if not len(self.stablesr_tensor_custom): return
        self.stablesr_script.stablesr_model.latent_image = self.stablesr_tensor_custom[bbox_id]


    @noise_inverse
    def init_noise_inverse(self, steps:int, retouch:float, get_cache_callback, set_cache_callback, renoise_strength:float, renoise_kernel:int):
        self.noise_inverse_enabled = True
        self.noise_inverse_steps = steps
        self.noise_inverse_retouch = float(retouch)
        self.noise_inverse_renoise_strength = float(renoise_strength)
        self.noise_inverse_renoise_kernel = int(renoise_kernel)
        if self.sample_img2img_original is None:
            self.sample_img2img_original = self.sampler_raw.sample_img2img
        self.sampler_raw.sample_img2img = MethodType(self.sample_img2img, self.sampler_raw)
        self.noise_inverse_set_cache = set_cache_callback
        self.noise_inverse_get_cache = get_cache_callback

    @noise_inverse
    @keep_signature
    def sample_img2img(self, sampler: KDiffusionSampler, p:ProcessingImg2Img, 
                       x:Tensor, noise:Tensor, conditioning, unconditional_conditioning,
                       steps=None, image_conditioning=None):
        # noise inverse sampling - renoise mask
        import torch.nn.functional as F
        renoise_mask = None
        if self.noise_inverse_renoise_strength > 0:
            image = p.init_images[0]
            # convert to grayscale with PIL
            image = image.convert('L')
            np_mask = get_retouch_mask(np.asarray(image), self.noise_inverse_renoise_kernel)
            renoise_mask = torch.from_numpy(np_mask).to(noise.device)
            # resize retouch mask to match noise size
            renoise_mask = 1 - F.interpolate(renoise_mask.unsqueeze(0).unsqueeze(0), size=noise.shape[-2:], mode='bilinear').squeeze(0).squeeze(0)
            renoise_mask *= self.noise_inverse_renoise_strength
            renoise_mask = torch.clamp(renoise_mask, 0, 1)

        prompts = p.all_prompts[:p.batch_size]
        
        latent = None
        # try to use cached latent to save huge amount of time.
        cached_latent: NoiseInverseCache = self.noise_inverse_get_cache()
        if cached_latent is not None and \
            cached_latent.model_hash == p.sd_model.sd_model_hash and \
            cached_latent.noise_inversion_steps == self.noise_inverse_steps and \
            len(cached_latent.prompts) == len(prompts) and \
            all([cached_latent.prompts[i] == prompts[i] for i in range(len(prompts))]) and \
            abs(cached_latent.retouch - self.noise_inverse_retouch) < 0.01 and \
            cached_latent.x0.shape == p.init_latent.shape and \
            torch.abs(cached_latent.x0.to(p.init_latent.device) - p.init_latent).sum() < 100: # the 100 is an arbitrary threshold copy-pasted from the img2img alt code
                # use cached noise
                print('[Tiled Diffusion] Your checkpoint, image, prompts, inverse steps, and retouch params are all unchanged.')
                print('[Tiled Diffusion] Noise Inversion will use the cached noise from the previous run. To clear the cache, click the Free GPU button.')
                latent = cached_latent.xt.to(noise.device)
        if latent is None:
            # run noise inversion
            shared.state.job_count += 1
            latent = self.find_noise_for_image_sigma_adjustment(sampler.model_wrap, self.noise_inverse_steps, prompts)
            shared.state.nextjob()
            self.noise_inverse_set_cache(p.init_latent.clone().cpu(), latent.clone().cpu(), prompts)
            # The cache is only 1 latent image and is very small (16 MB for 8192 * 8192 image), so we don't need to worry about memory leakage.

        # calculate sampling steps
        adjusted_steps, _ = sd_samplers_common.setup_img2img_steps(p, steps)
        sigmas = sampler.get_sigmas(p, adjusted_steps)
        inverse_noise = latent - (p.init_latent / sigmas[0])

        # inject noise to high-frequency area so that the details won't lose too much
        if renoise_mask is not None:
            # If the background is not drawn, we need to filter out the un-drawn pixels and reweight foreground with feather mask
            # This is to enable the renoise mask in regional inpainting
            if not self.enable_grid_bbox:
                background_count  = torch.zeros((1, 1, noise.shape[2], noise.shape[3]), device=noise.device)
                foreground_noise  = torch.zeros_like(noise)
                foreground_weight = torch.zeros((1, 1, noise.shape[2], noise.shape[3]), device=noise.device)
                foreground_count  = torch.zeros((1, 1, noise.shape[2], noise.shape[3]), device=noise.device)
                for bbox in self.custom_bboxes:
                    if bbox.blend_mode == BlendMode.BACKGROUND:
                        background_count[bbox.slicer] += 1
                    elif bbox.blend_mode == BlendMode.FOREGROUND:
                        foreground_noise [bbox.slicer] += noise[bbox.slicer]
                        foreground_weight[bbox.slicer] += bbox.feather_mask
                        foreground_count [bbox.slicer] += 1
                background_noise  = torch.where(background_count > 0, noise, 0)
                foreground_noise  = torch.where(foreground_count > 0, foreground_noise  / foreground_count, 0)
                foreground_weight = torch.where(foreground_count > 0, foreground_weight / foreground_count, 0)
                noise = background_noise * (1 - foreground_weight) + foreground_noise * foreground_weight
                del background_noise, foreground_noise, foreground_weight, background_count, foreground_count
            combined_noise = ((1 - renoise_mask) * inverse_noise + renoise_mask * noise) / ((renoise_mask**2 + (1 - renoise_mask)**2) ** 0.5)
        else:
            combined_noise = inverse_noise

        # use the estimated noise for the original img2img sampling
        return self.sample_img2img_original(p, x, combined_noise, conditioning, unconditional_conditioning, steps, image_conditioning)

    @noise_inverse
    @torch.no_grad()
    def find_noise_for_image_sigma_adjustment(self, dnw, steps, prompts:List[str]) -> Tensor:
        '''
        Migrate from the built-in script img2imgalt.py
        Tiled noise inverse for better image upscaling
        '''
        import k_diffusion as K
        assert self.p.sampler_name == 'Euler'

        x = self.p.init_latent
        s_in = x.new_ones([x.shape[0]])
        skip = 1 if shared.sd_model.parameterization == "v" else 0
        sigmas = dnw.get_sigmas(steps).flip(0)
    
        cond = self.p.sd_model.get_learned_conditioning(prompts)
        if isinstance(cond, Tensor):    # SD1/SD2
            cond_dict_dummy = {
                'c_crossattn': [],      # List[Tensor]
                'c_concat': [],         # List[Tensor]
            }
            cond_in = self.make_cond_dict(cond_dict_dummy, cond, self.p.image_conditioning)
        else:                           # SDXL
            cond_dict_dummy = {
                'crossattn': None,      # Tensor
                'vector': None,         # Tensor
                'c_concat': [],         # List[Tensor]
            }
            cond_in = self.make_cond_dict(cond_dict_dummy, cond['crossattn'], self.p.image_conditioning, cond['vector'])

        state.sampling_steps = steps
        pbar = tqdm(total=steps, desc='Noise Inversion')
        for i in range(1, len(sigmas)):
            if state.interrupted: return x

            state.sampling_step += 1

            x_in = x
            sigma_in = torch.cat([sigmas[i] * s_in])
            c_out, c_in = [K.utils.append_dims(k, x_in.ndim) for k in dnw.get_scalings(sigma_in)[skip:]]

            t = dnw.sigma_to_t(sigma_in)
            t = t / self.noise_inverse_retouch

            eps = self.get_noise(x_in * c_in, t, cond_in, steps - i)
            denoised = x_in + eps * c_out

            # Euler method:
            d = (x - denoised) / sigmas[i]
            dt = sigmas[i] - sigmas[i - 1]
            x = x + d * dt

            sd_samplers_common.store_latent(x)

            # This is neccessary to save memory before the next iteration
            del x_in, sigma_in, c_out, c_in, t,
            del eps, denoised, d, dt

            pbar.update(1)
        pbar.close()

        return x / sigmas[-1]
    
    @noise_inverse
    @torch.no_grad()
    def get_noise(self, x_in: Tensor, sigma_in:Tensor, cond_in:Dict[str, Tensor], step:int) -> Tensor:
        raise NotImplementedError