feat: enable pipeline to output fused result

app.py

@@ -10,9 +10,8 @@ import argparse
 import numpy as np
 import torchvision.transforms.functional as TF
 from scipy.ndimage import convolve, zoom
-import cv2
-import time
 import spaces
+from utils import resize_by_short_side
 
 from tools.interact_tools import SamControler
 from tools.misc import get_device
@@ -33,106 +32,6 @@ def parse_augment():
 
     return args 
 
-
-def pad_to_multiple(image: np.ndarray, multiple: int = 8):
-    h, w = image.shape[:2]
-    pad_h = (multiple - h % multiple) % multiple
-    pad_w = (multiple - w % multiple) % multiple
-    if image.ndim == 3:
-        padded = np.pad(image, ((0, pad_h), (0, pad_w), (0,0)), mode='reflect')
-    else:
-        padded = np.pad(image, ((0, pad_h), (0, pad_w)), mode='reflect')
-    return padded, h, w
-
-def crop_to_original(image: np.ndarray, h: int, w: int):
-    return image[:h, :w]
-
-def wavelet_blur_np(image: np.ndarray, radius: int):
-    kernel = np.array([
-        [0.0625, 0.125, 0.0625],
-        [0.125,  0.25,  0.125],
-        [0.0625, 0.125, 0.0625]
-    ], dtype=np.float32)
-
-    blurred = np.empty_like(image)
-    for c in range(image.shape[0]):
-        blurred_c = convolve(image[c], kernel, mode='nearest')
-        if radius > 1:
-            blurred_c = zoom(zoom(blurred_c, 1 / radius, order=1), radius, order=1)
-        blurred[c] = blurred_c
-    return blurred
-
-def wavelet_decomposition_np(image: np.ndarray, levels=5):
-    high_freq = np.zeros_like(image)
-    for i in range(levels):
-        radius = 2 ** i
-        low_freq = wavelet_blur_np(image, radius)
-        high_freq += (image - low_freq)
-        image = low_freq
-    return high_freq, low_freq
-
-def wavelet_reconstruction_np(content_feat: np.ndarray, style_feat: np.ndarray):
-    content_high, _ = wavelet_decomposition_np(content_feat)
-    _, style_low = wavelet_decomposition_np(style_feat)
-    return content_high + style_low
-
-def wavelet_color_fix_np(fused: np.ndarray, mask: np.ndarray) -> np.ndarray:
-    fused_np = fused.astype(np.float32) / 255.0
-    mask_np = mask.astype(np.float32) / 255.0
-
-    fused_np = fused_np.transpose(2, 0, 1)
-    mask_np = mask_np.transpose(2, 0, 1)
-
-    result_np = wavelet_reconstruction_np(fused_np, mask_np)
-
-    result_np = result_np.transpose(1, 2, 0)
-    result_np = np.clip(result_np * 255.0, 0, 255).astype(np.uint8)
-
-    return result_np
-
-def fuse_with_wavelet(ori: np.ndarray, removed: np.ndarray, attn_map: np.ndarray, multiple: int = 8):
-    H, W = ori.shape[:2]
-    attn_map = attn_map.astype(np.float32)
-    _, attn_map = cv2.threshold(attn_map, 128, 255, cv2.THRESH_BINARY)
-    am = attn_map.astype(np.float32)
-    am = am/255.0
-    am_up = cv2.resize(am, (W, H), interpolation=cv2.INTER_NEAREST)
-
-    kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (21,21))
-    am_d = cv2.dilate(am_up, kernel, iterations=1)
-    am_d = cv2.GaussianBlur(am_d.astype(np.float32), (9,9), sigmaX=2)
-
-    am_merged = np.maximum(am_up, am_d)
-    am_merged = np.clip(am_merged, 0, 1)
-
-    attn_up_3c = np.stack([am_merged]*3, axis=-1)
-    attn_up_ori_3c = np.stack([am_up]*3, axis=-1)
-
-    ori_out = ori * (1 - attn_up_ori_3c)
-    rem_out = removed * (1 - attn_up_ori_3c)
-
-    ori_pad, h0, w0 = pad_to_multiple(ori_out, multiple)
-    rem_pad, _, _   = pad_to_multiple(rem_out, multiple)
-
-    wave_rgb = wavelet_color_fix_np(ori_pad, rem_pad)
-    wave = crop_to_original(wave_rgb, h0, w0)
-    # fusion
-    fused = (wave * (1 - attn_up_3c) + removed * attn_up_3c).astype(np.uint8)
-    return fused
-
-
-def resize_by_short_side(image, target_short=512, resample=Image.BICUBIC):
-    w, h = image.size
-    if w < h:
-        new_w = target_short
-        new_h = int(h * target_short / w)
-        new_h = (new_h + 15) // 16 * 16 
-    else:
-        new_h = target_short
-        new_w = int(w * target_short / h)
-        new_w = (new_w + 15) // 16 * 16
-    return image.resize((new_w, new_h), resample=resample)
-
 # convert points input to prompt state
 def get_prompt(click_state, click_input):
     inputs = json.loads(click_input)
@@ -281,7 +180,7 @@ pipe = ObjectClearPipeline.from_pretrained_with_custom_modules(
     "jixin0101/ObjectClear",
     torch_dtype=torch.float16,
     variant='fp16',
-    save_cross_attn=True
+    apply_attention_guided_fusion=True
 )
 
 pipe.to(device)
@@ -325,13 +224,7 @@ def process(image_state, interactive_state, mask_dropdown, guidance_scale, seed,
         height=h,
         width=w,
     )
-    
-    inpainted_img = result[0].images[0]
-    attn_map = result[1]
-    attn_np = attn_map.mean(dim=1)[0].cpu().numpy() * 255.
-
-    fused_img = fuse_with_wavelet(np.array(image), np.array(inpainted_img), attn_np)
-    fused_img_pil = Image.fromarray(fused_img.astype(np.uint8))
+    fused_img_pil = result.images[0]
 
     return fused_img_pil.resize((image_or.size[:2])), (image.resize((image_or.size[:2])), fused_img_pil.resize((image_or.size[:2])))
 

pipeline_objectclear.py

@@ -58,6 +58,7 @@ from diffusers.pipelines.pipeline_utils import DiffusionPipeline, StableDiffusio
 from diffusers.pipelines.stable_diffusion_xl.pipeline_output import StableDiffusionXLPipelineOutput
 
 from model import CLIPImageEncoder, PostfuseModule
+from utils import attention_guided_fusion
 import gc
 import torch.nn.functional as F
 
@@ -328,6 +329,10 @@ def retrieve_timesteps(
     return timesteps, num_inference_steps
 
 
+@dataclass
+class ObjectClearPipelineOutput(StableDiffusionXLPipelineOutput):
+    attns: Optional[List[PIL.Image.Image]] = None
+
 class ObjectClearPipeline(
     DiffusionPipeline,
     StableDiffusionMixin,
@@ -422,7 +427,7 @@ class ObjectClearPipeline(
         requires_aesthetics_score: bool = False,
         force_zeros_for_empty_prompt: bool = True,
         add_watermarker: Optional[bool] = None,
-        save_cross_attn: bool = False,
+        apply_attention_guided_fusion: bool = False,
     ):
         super().__init__()
 
@@ -441,7 +446,7 @@ class ObjectClearPipeline(
         )
         self.register_to_config(force_zeros_for_empty_prompt=force_zeros_for_empty_prompt)
         self.register_to_config(requires_aesthetics_score=requires_aesthetics_score)
-        self.register_to_config(save_cross_attn=save_cross_attn)
+        self.register_to_config(apply_attention_guided_fusion=apply_attention_guided_fusion)
         self.vae_scale_factor = 2 ** (len(self.vae.config.block_out_channels) - 1)
         self.image_processor = VaeImageProcessor(vae_scale_factor=self.vae_scale_factor)
         self.mask_processor = VaeImageProcessor(
@@ -455,7 +460,7 @@ class ObjectClearPipeline(
         else:
             self.watermark = None
         
-        if self.config.save_cross_attn:
+        if self.config.apply_attention_guided_fusion:
             self.cross_attention_scores = {}
             self.unet = self.unet_store_cross_attention_scores(
                 self.unet, self.cross_attention_scores
@@ -1367,6 +1372,7 @@ class ObjectClearPipeline(
         ip_adapter_image_embeds: Optional[List[torch.Tensor]] = None,
         output_type: Optional[str] = "pil",
         return_dict: bool = True,
+        return_attn_map: bool = False,
         cross_attention_kwargs: Optional[Dict[str, Any]] = None,
         guidance_rescale: float = 0.0,
         original_size: Tuple[int, int] = None,
@@ -1859,7 +1865,7 @@ class ObjectClearPipeline(
             ).to(device=device, dtype=latents.dtype)
 
         self._num_timesteps = len(timesteps)
-        self.attn_map = None
+        attn_map = None
         with self.progress_bar(total=num_inference_steps) as progress_bar:
             for i, t in enumerate(timesteps):
                 if self.interrupt:
@@ -1906,16 +1912,16 @@ class ObjectClearPipeline(
 
                 # progressive attention mask blending
                 fuse_index = 5
-                if self.config.save_cross_attn:
+                if self.config.apply_attention_guided_fusion:
                     if i == len(timesteps) - 1:
                         attn_key, attn_map = next(iter(self.cross_attention_scores.items()))
-                        self.attn_map = self.resize_attn_map_divide2(attn_map, mask, fuse_index)
+                        attn_map = self.resize_attn_map_divide2(attn_map, mask, fuse_index)
                         init_latents_proper = image_latents
                         if self.do_classifier_free_guidance:
-                            _, init_mask = self.attn_map.chunk(2)
+                            _, init_mask = attn_map.chunk(2)
                         else:
-                            init_mask = self.attn_map
-                        self.attn_map = init_mask
+                            init_mask = attn_map
+                        attn_map = init_mask
                     self.clear_cross_attention_scores(self.cross_attention_scores)
                 
                 if num_channels_unet == 4:
@@ -1994,7 +2000,7 @@ class ObjectClearPipeline(
             if needs_upcasting:
                 self.vae.to(dtype=torch.float16)
         else:
-            return StableDiffusionXLPipelineOutput(images=latents)
+            return ObjectClearPipelineOutput(images=latents)
 
         # apply watermark if available
         if self.watermark is not None:
@@ -2004,11 +2010,39 @@ class ObjectClearPipeline(
 
         if padding_mask_crop is not None:
             image = [self.image_processor.apply_overlay(mask_image, original_image, i, crops_coords) for i in image]
+            
+        attn_pils = []
+        if output_type == "pil" and attn_map is not None:
+            for i in range(len(attn_map)):
+                attn_np = attn_map[i].mean(dim=0).cpu().numpy() * 255.
+                attn_pil = PIL.Image.fromarray(attn_np.astype(np.uint8)).convert("L")
+                attn_pils.append(attn_pil)
+            
+            original_pils = self.image_processor.postprocess(init_image, output_type="pil")
 
-        # Offload all models
-        self.maybe_free_model_hooks()
+            generated_pils = image
+
+            fused_images = []
+            for i in range(len(generated_pils)):
+                ori_pil = original_pils[i]
+                gen_pil = generated_pils[i]
+                attn_pil = attn_pils[i]
 
-        if not return_dict:
-            return (image,)
+                fused_np = attention_guided_fusion(np.array(ori_pil), np.array(gen_pil), np.array(attn_pil))
+                fused_pil = PIL.Image.fromarray(fused_np.astype(np.uint8)).resize(ori_pil.size)
 
-        return StableDiffusionXLPipelineOutput(images=image), self.attn_map
+                fused_images.append(fused_pil)
+
+            image = fused_images
+
+        # Offload all models
+        self.maybe_free_model_hooks()
+    
+        if return_attn_map and len(attn_pils) > 0:
+            if not return_dict:
+                return (image, attn_pils)
+            return ObjectClearPipelineOutput(images=image, attns=attn_pils)
+        else:
+            if not return_dict:
+                return (image,)
+            return ObjectClearPipelineOutput(images=image)

utils.py

@@ -0,0 +1,105 @@
+import numpy as np
+import cv2
+from scipy.ndimage import convolve, zoom
+from PIL import Image
+
+
+
+def pad_to_multiple(image: np.ndarray, multiple: int = 8):
+    h, w = image.shape[:2]
+    pad_h = (multiple - h % multiple) % multiple
+    pad_w = (multiple - w % multiple) % multiple
+    if image.ndim == 3:
+        padded = np.pad(image, ((0, pad_h), (0, pad_w), (0,0)), mode='reflect')
+    else:
+        padded = np.pad(image, ((0, pad_h), (0, pad_w)), mode='reflect')
+    return padded, h, w
+
+def crop_to_original(image: np.ndarray, h: int, w: int):
+    return image[:h, :w]
+
+def wavelet_blur_np(image: np.ndarray, radius: int):
+    kernel = np.array([
+        [0.0625, 0.125, 0.0625],
+        [0.125,  0.25,  0.125],
+        [0.0625, 0.125, 0.0625]
+    ], dtype=np.float32)
+
+    blurred = np.empty_like(image)
+    for c in range(image.shape[0]):
+        blurred_c = convolve(image[c], kernel, mode='nearest')
+        if radius > 1:
+            blurred_c = zoom(zoom(blurred_c, 1 / radius, order=1), radius, order=1)
+        blurred[c] = blurred_c
+    return blurred
+
+def wavelet_decomposition_np(image: np.ndarray, levels=5):
+    high_freq = np.zeros_like(image)
+    for i in range(levels):
+        radius = 2 ** i
+        low_freq = wavelet_blur_np(image, radius)
+        high_freq += (image - low_freq)
+        image = low_freq
+    return high_freq, low_freq
+
+def wavelet_reconstruction_np(content_feat: np.ndarray, style_feat: np.ndarray):
+    content_high, _ = wavelet_decomposition_np(content_feat)
+    _, style_low = wavelet_decomposition_np(style_feat)
+    return content_high + style_low
+
+def wavelet_color_fix_np(fused: np.ndarray, mask: np.ndarray) -> np.ndarray:
+    fused_np = fused.astype(np.float32) / 255.0
+    mask_np = mask.astype(np.float32) / 255.0
+
+    fused_np = fused_np.transpose(2, 0, 1)
+    mask_np = mask_np.transpose(2, 0, 1)
+
+    result_np = wavelet_reconstruction_np(fused_np, mask_np)
+
+    result_np = result_np.transpose(1, 2, 0)
+    result_np = np.clip(result_np * 255.0, 0, 255).astype(np.uint8)
+
+    return result_np
+
+def attention_guided_fusion(ori: np.ndarray, removed: np.ndarray, attn_map: np.ndarray, multiple: int = 8):
+    H, W = ori.shape[:2]
+    attn_map = attn_map.astype(np.float32)
+    _, attn_map = cv2.threshold(attn_map, 128, 255, cv2.THRESH_BINARY)
+    am = attn_map.astype(np.float32)
+    am = am/255.0
+    am_up = cv2.resize(am, (W, H), interpolation=cv2.INTER_NEAREST)
+
+    kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (21,21))
+    am_d = cv2.dilate(am_up, kernel, iterations=1)
+    am_d = cv2.GaussianBlur(am_d.astype(np.float32), (9,9), sigmaX=2)
+
+    am_merged = np.maximum(am_up, am_d)
+    am_merged = np.clip(am_merged, 0, 1)
+
+    attn_up_3c = np.stack([am_merged]*3, axis=-1)
+    attn_up_ori_3c = np.stack([am_up]*3, axis=-1)
+
+    ori_out = ori * (1 - attn_up_ori_3c)
+    rem_out = removed * (1 - attn_up_ori_3c)
+
+    ori_pad, h0, w0 = pad_to_multiple(ori_out, multiple)
+    rem_pad, _, _   = pad_to_multiple(rem_out, multiple)
+
+    wave_rgb = wavelet_color_fix_np(ori_pad, rem_pad)
+    wave = crop_to_original(wave_rgb, h0, w0)
+    # fusion
+    fused = (wave * (1 - attn_up_3c) + removed * attn_up_3c).astype(np.uint8)
+    return fused
+
+
+def resize_by_short_side(image, target_short=512, resample=Image.BICUBIC):
+    w, h = image.size
+    if w < h:
+        new_w = target_short
+        new_h = int(h * target_short / w)
+        new_h = (new_h + 15) // 16 * 16 
+    else:
+        new_h = target_short
+        new_w = int(w * target_short / h)
+        new_w = (new_w + 15) // 16 * 16
+    return image.resize((new_w, new_h), resample=resample)