748 lines
35 KiB
Python
748 lines
35 KiB
Python
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
|