mirror of
https://github.com/vladmandic/automatic
synced 2026-09-19 01:04:32 +02:00
refactor img2img processing
This commit is contained in:
Vladimir Mandic
+39
-59
@@ -49,19 +49,20 @@ debug('Trace: PROCESS')
|
||||
|
||||
|
||||
def setup_color_correction(image):
|
||||
shared.log.debug("Calibrating color correction.")
|
||||
debug("Calibrating color correction")
|
||||
correction_target = cv2.cvtColor(np.asarray(image.copy()), cv2.COLOR_RGB2LAB)
|
||||
return correction_target
|
||||
|
||||
|
||||
def apply_color_correction(correction, original_image):
|
||||
shared.log.debug("Applying color correction.")
|
||||
shared.log.debug(f"Applying color correction: correction={correction} image={original_image}")
|
||||
image = Image.fromarray(cv2.cvtColor(exposure.match_histograms(cv2.cvtColor(np.asarray(original_image), cv2.COLOR_RGB2LAB), correction, channel_axis=2), cv2.COLOR_LAB2RGB).astype("uint8"))
|
||||
image = blendLayers(image, original_image, BlendType.LUMINOSITY)
|
||||
return image
|
||||
|
||||
|
||||
def apply_overlay(image: Image, paste_loc, index, overlays):
|
||||
debug(f'Apply overlay: image={image} loc={paste_loc} index={index} overlays={overlays}')
|
||||
if overlays is None or index >= len(overlays):
|
||||
return image
|
||||
overlay = overlays[index]
|
||||
@@ -1240,8 +1241,8 @@ class StableDiffusionProcessingImg2Img(StableDiffusionProcessing):
|
||||
self.image_mask = mask
|
||||
self.latent_mask = None
|
||||
self.mask_for_overlay = None
|
||||
self.mask_blur_x: int = 4
|
||||
self.mask_blur_y: int = 4
|
||||
self.mask_blur_x = mask_blur # a1111 compatibility item
|
||||
self.mask_blur_y = mask_blur # a1111 compatibility item
|
||||
self.mask_blur = mask_blur
|
||||
self.inpainting_fill = inpainting_fill
|
||||
self.inpaint_full_res = inpaint_full_res
|
||||
@@ -1262,16 +1263,6 @@ class StableDiffusionProcessingImg2Img(StableDiffusionProcessing):
|
||||
self.scripts = None
|
||||
self.script_args = []
|
||||
|
||||
@property
|
||||
def mask_blur(self):
|
||||
mask_blur = max(self.mask_blur_x, self.mask_blur_y)
|
||||
return mask_blur
|
||||
|
||||
@mask_blur.setter
|
||||
def mask_blur(self, value):
|
||||
self.mask_blur_x = value
|
||||
self.mask_blur_y = value
|
||||
|
||||
def init(self, all_prompts, all_seeds, all_subseeds):
|
||||
if shared.backend == shared.Backend.DIFFUSERS and self.image_mask is not None and not self.is_control:
|
||||
shared.sd_model = modules.sd_models.set_diffuser_pipe(self.sd_model, modules.sd_models.DiffusersTaskType.INPAINTING)
|
||||
@@ -1290,46 +1281,40 @@ class StableDiffusionProcessingImg2Img(StableDiffusionProcessing):
|
||||
self.ops.append('img2img')
|
||||
crop_region = None
|
||||
|
||||
image_mask = self.image_mask
|
||||
if image_mask is not None:
|
||||
if type(image_mask) == list:
|
||||
image_mask = image_mask[0]
|
||||
image_mask = create_binary_mask(image_mask)
|
||||
if self.image_mask is not None:
|
||||
if type(self.image_mask) == list:
|
||||
self.image_mask = self.image_mask[0]
|
||||
self.image_mask = create_binary_mask(self.image_mask)
|
||||
if self.inpainting_mask_invert:
|
||||
image_mask = ImageOps.invert(image_mask)
|
||||
if self.mask_blur_x > 0:
|
||||
np_mask = np.array(image_mask)
|
||||
kernel_size = 2 * int(2.5 * self.mask_blur_x + 0.5) + 1
|
||||
np_mask = cv2.GaussianBlur(np_mask, (kernel_size, 1), self.mask_blur_x)
|
||||
image_mask = Image.fromarray(np_mask)
|
||||
if self.mask_blur_y > 0:
|
||||
np_mask = np.array(image_mask)
|
||||
kernel_size = 2 * int(2.5 * self.mask_blur_y + 0.5) + 1
|
||||
np_mask = cv2.GaussianBlur(np_mask, (1, kernel_size), self.mask_blur_y)
|
||||
image_mask = Image.fromarray(np_mask)
|
||||
self.image_mask = ImageOps.invert(self.image_mask)
|
||||
if self.mask_blur > 0:
|
||||
np_mask = np.array(self.image_mask)
|
||||
kernel_size = 2 * int(2.5 * self.mask_blur + 0.5) + 1
|
||||
np_mask = cv2.GaussianBlur(np_mask, (kernel_size, 1), self.mask_blur)
|
||||
np_mask = cv2.GaussianBlur(np_mask, (1, kernel_size), self.mask_blur)
|
||||
self.image_mask = Image.fromarray(np_mask)
|
||||
if self.inpaint_full_res:
|
||||
self.mask_for_overlay = image_mask
|
||||
mask = image_mask.convert('L')
|
||||
self.mask_for_overlay = self.image_mask
|
||||
mask = self.image_mask.convert('L')
|
||||
crop_region = modules.masking.get_crop_region(np.array(mask), self.inpaint_full_res_padding)
|
||||
crop_region = modules.masking.expand_crop_region(crop_region, self.width, self.height, mask.width, mask.height)
|
||||
x1, y1, x2, y2 = crop_region
|
||||
mask = mask.crop(crop_region)
|
||||
image_mask = images.resize_image(2, mask, self.width, self.height)
|
||||
self.image_mask = images.resize_image(2, mask, self.width, self.height)
|
||||
self.paste_to = (x1, y1, x2-x1, y2-y1)
|
||||
else:
|
||||
image_mask = images.resize_image(self.resize_mode, image_mask, self.width, self.height)
|
||||
np_mask = np.array(image_mask)
|
||||
self.image_mask = images.resize_image(self.resize_mode, self.image_mask, self.width, self.height)
|
||||
np_mask = np.array(self.image_mask)
|
||||
np_mask = np.clip((np_mask.astype(np.float32)) * 2, 0, 255).astype(np.uint8)
|
||||
self.mask_for_overlay = Image.fromarray(np_mask)
|
||||
self.overlay_images = []
|
||||
|
||||
latent_mask = self.latent_mask if self.latent_mask is not None else image_mask
|
||||
latent_mask = self.latent_mask if self.latent_mask is not None else self.image_mask
|
||||
|
||||
add_color_corrections = shared.opts.img2img_color_correction and self.color_corrections is None
|
||||
if add_color_corrections:
|
||||
self.color_corrections = []
|
||||
imgs = []
|
||||
unprocessed = []
|
||||
processed = []
|
||||
if getattr(self, 'init_images', None) is None:
|
||||
return
|
||||
if not isinstance(self.init_images, list):
|
||||
@@ -1349,7 +1334,7 @@ class StableDiffusionProcessingImg2Img(StableDiffusionProcessing):
|
||||
image = images.resize_image(self.resize_mode, image, self.width, self.height, self.resize_name)
|
||||
self.width = image.width
|
||||
self.height = image.height
|
||||
if image_mask is not None:
|
||||
if self.image_mask is not None:
|
||||
try:
|
||||
image_masked = Image.new('RGBa', (image.width, image.height))
|
||||
image_to_paste = image.convert("RGBA").convert("RGBa")
|
||||
@@ -1363,37 +1348,32 @@ class StableDiffusionProcessingImg2Img(StableDiffusionProcessing):
|
||||
image = image.crop(crop_region)
|
||||
if image.width != self.width or image.height != self.height:
|
||||
image = images.resize_image(3, image, self.width, self.height, self.resize_name)
|
||||
if image_mask is not None and self.inpainting_fill != 1:
|
||||
if self.image_mask is not None and self.inpainting_fill != 1:
|
||||
image = modules.masking.fill(image, latent_mask)
|
||||
if add_color_corrections:
|
||||
self.color_corrections.append(setup_color_correction(image))
|
||||
if shared.backend == shared.Backend.DIFFUSERS:
|
||||
unprocessed.append(image) # assign early for diffusers
|
||||
image = np.array(image).astype(np.float32) / 255.0
|
||||
image = np.moveaxis(image, 2, 0)
|
||||
imgs.append(image)
|
||||
self.init_images = unprocessed if shared.backend == shared.Backend.DIFFUSERS else imgs
|
||||
if len(imgs) == 1:
|
||||
batch_images = np.expand_dims(imgs[0], axis=0).repeat(self.batch_size, axis=0)
|
||||
if self.overlay_images is not None:
|
||||
self.overlay_images = self.overlay_images * self.batch_size
|
||||
if self.color_corrections is not None and len(self.color_corrections) == 1:
|
||||
self.color_corrections = self.color_corrections * self.batch_size
|
||||
elif len(imgs) <= self.batch_size:
|
||||
self.batch_size = len(imgs)
|
||||
batch_images = np.array(imgs)
|
||||
else:
|
||||
raise RuntimeError(f"Incorrect number of of images={len(imgs)} expected={self.batch_size} or less")
|
||||
processed.append(image)
|
||||
self.init_images = processed
|
||||
self.batch_size = len(self.init_images)
|
||||
if self.overlay_images is not None:
|
||||
self.overlay_images = self.overlay_images * self.batch_size
|
||||
if self.color_corrections is not None and len(self.color_corrections) == 1:
|
||||
self.color_corrections = self.color_corrections * self.batch_size
|
||||
if shared.backend == shared.Backend.DIFFUSERS:
|
||||
return # we've already set self.init_images and self.mask and we dont need any more processing
|
||||
|
||||
self.init_images = [np.moveaxis((np.array(image).astype(np.float32) / 255.0), 2, 0) for image in self.init_images]
|
||||
if len(self.init_images) == 1:
|
||||
batch_images = np.expand_dims(self.init_images[0], axis=0).repeat(self.batch_size, axis=0)
|
||||
elif len(self.init_images) <= self.batch_size:
|
||||
batch_images = np.array(self.init_images)
|
||||
image = torch.from_numpy(batch_images)
|
||||
image = 2. * image - 1.
|
||||
image = image.to(device=shared.device, dtype=devices.dtype_vae)
|
||||
self.init_latent = self.sd_model.get_first_stage_encoding(self.sd_model.encode_first_stage(image))
|
||||
if self.resize_mode == 4:
|
||||
self.init_latent = torch.nn.functional.interpolate(self.init_latent, size=(self.height // 8, self.width // 8), mode="bilinear")
|
||||
if image_mask is not None:
|
||||
if self.image_mask is not None:
|
||||
init_mask = latent_mask
|
||||
latmask = init_mask.convert('RGB').resize((self.init_latent.shape[3], self.init_latent.shape[2]))
|
||||
latmask = np.moveaxis(np.array(latmask, dtype=np.float32), 2, 0) / 255
|
||||
@@ -1406,7 +1386,7 @@ class StableDiffusionProcessingImg2Img(StableDiffusionProcessing):
|
||||
self.init_latent = self.init_latent * self.mask + create_random_tensors(self.init_latent.shape[1:], all_seeds[0:self.init_latent.shape[0]]) * self.nmask
|
||||
elif self.inpainting_fill == 3:
|
||||
self.init_latent = self.init_latent * self.mask
|
||||
self.image_conditioning = self.img2img_image_conditioning(image, self.init_latent, image_mask)
|
||||
self.image_conditioning = self.img2img_image_conditioning(image, self.init_latent, self.image_mask)
|
||||
|
||||
def sample(self, conditioning, unconditional_conditioning, seeds, subseeds, subseed_strength, prompts):
|
||||
hypertile_set(self)
|
||||
|
||||
Reference in New Issue
Block a user