I

code 443

RealCustom/inference/app.py

@@ -0,0 +1,82 @@
+# Copyright (c) 2024 Bytedance Ltd. and/or its affiliates
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import gradio as gr
+
+from .pipeline import RealCustomInferencePipeline
+
+def create_demo():
+    pipeline = RealCustomInferencePipeline(
+        unet_config="configs/realcustom_sigdino_highres.json",
+        unet_checkpoint="ckpts/sdxl/unet/sdxl-unet.bin",
+        realcustom_checkpoint="ckpts/realcustom/RealCustom_highres.pth",
+        vae_config="ckpts/sdxl/vae/sdxl.json",
+        vae_checkpoint="ckpts/sdxl/vae/sdxl-vae.pth",
+    )
+
+    badges_text = r"""
+    <div style="text-align: center; display: flex; justify-content: left; gap: 5px;">
+    <a href="https://corleone-huang.github.io/RealCustom_plus_plus/"><img alt="Build" src="https://img.shields.io/badge/Project%20Page-RealCustom-yellow"></a> 
+    <a href="https://arxiv.org/pdf/2408.09744?"><img alt="Build" src="https://img.shields.io/badge/arXiv%20paper-RealCustom-b31b1b.svg"></a>
+    </div>
+    """.strip()
+
+    with gr.Blocks() as demo:
+        gr.Markdown(f"# RealCustom")
+        gr.Markdown(badges_text)
+        with gr.Row():
+            with gr.Column():
+                prompt = gr.Textbox(label="Prompt", value="")
+                target_phrase = gr.Textbox(label="Target Phrase", value="")
+                with gr.Row():
+                    image_prompt = gr.Image(label="Ref Img", visible=True, interactive=True, type="pil")
+
+                with gr.Row():
+                    with gr.Column():
+                        width = gr.Slider(512, 2048, 1024, step=16, label="Gneration Width")
+                        height = gr.Slider(512, 2048, 1024, step=16, label="Gneration Height")
+
+                with gr.Accordion("Advanced Options", open=False):
+                    with gr.Row():
+                        guidance = gr.Slider(1.0, 15, 3.5, step=0.5, label="Guidance Scale", interactive=True)
+                        mask_scope = gr.Slider(0.05, 1.0, 0.2, step=0.05, label="Mask Scope", interactive=True)
+                        seed = gr.Number(0, label="Seed (-1 for random)")
+                        num = gr.Number(4, label="Generation Number")
+                        new_unet_local_path = gr.Textbox(label="New Unet Local Path", value="")
+                        new_realcustom_local_path = gr.Textbox(label="New RealCustom Local Path", value="")
+
+                generate_btn = gr.Button("Generate")
+
+            with gr.Column():
+                output_image = gr.Image(label="Generated Image")
+                output_mask = gr.Image(label="Guidance Mask")
+
+            inputs = [
+                prompt, image_prompt, target_phrase, 
+                height, width, guidance, seed, num, 
+                mask_scope, 
+                new_unet_local_path, new_realcustom_local_path,
+            ]
+            generate_btn.click(
+                fn=pipeline.generation,
+                inputs=inputs,
+                outputs=[output_image, output_mask],
+            )
+
+    return demo
+
+
+if __name__ == "__main__":
+    demo = create_demo()
+    demo.launch(server_name='0.0.0.0', server_port=7860)

RealCustom/inference/inference_single_image.py

@@ -0,0 +1,317 @@
+# Copyright (c) 2024 Bytedance Ltd. and/or its affiliates
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import argparse
+import torch
+import json
+import os
+import torchvision
+from torchvision.utils import make_grid
+from torchvision.transforms.functional import to_pil_image
+from tqdm import tqdm
+from PIL import Image
+
+from ..models.text import TextModel
+from ..models.vae import AutoencoderKL
+
+from ..models.unet_2d_condition_custom import UNet2DConditionModel as UNet2DConditionModelDiffusers
+from ..schedulers.ddim import DDIMScheduler
+from ..schedulers.dpm_s import DPMSolverSingleStepScheduler
+from ..schedulers.utils import get_betas
+
+from .inference_utils import find_phrase_positions_in_text, classifier_free_guidance_image_prompt_cascade
+from .mask_generation import mask_generation
+from ..utils import instantiate_from_config
+
+# Argument parser
+parser = argparse.ArgumentParser()
+parser.add_argument("--width", type=int, default=512)
+parser.add_argument("--height", type=int, default=512)
+
+parser.add_argument("--samples_per_prompt", type=int, required=True)
+parser.add_argument("--nrow", type=int, default=4)
+parser.add_argument("--sample_steps", type=int, required=True)
+parser.add_argument("--schedule_type", type=str, default="squared_linear")                                            # default, `squared_linear
+parser.add_argument("--scheduler_type", type=str, default="dpm", choices=["ddim", "dpm"])                             # default, "dpm"
+parser.add_argument("--schedule_shift_snr", type=float, default=1)                                                    # default, 1
+
+parser.add_argument("--text_encoder_variant", type=str, nargs="+")
+parser.add_argument("--vae_config", type=str, default="configs/vae.json")                                             # default
+parser.add_argument("--vae_checkpoint", type=str, required=True)
+parser.add_argument("--unet_config", type=str, required=True)
+parser.add_argument("--unet_checkpoint", type=str, required=True)
+parser.add_argument("--unet_checkpoint_base_model", type=str, default="")
+parser.add_argument("--unet_prediction", type=str, choices=DDIMScheduler.prediction_types, default="epsilon")          # default, "epsilon"
+
+parser.add_argument("--negative_prompt", type=str, default="prompts/validation_negative.txt")                          # default
+
+parser.add_argument("--compile", action="store_true", default=False)
+parser.add_argument("--output_dir", type=str, required=True)
+
+parser.add_argument("--guidance_weight", type=float, default=7.5)
+parser.add_argument("--seed", type=int, default=666)
+parser.add_argument("--device", type=str, default="cuda")
+
+parser.add_argument("--text_prompt", type=str, required=True)
+parser.add_argument("--image_prompt_path", type=str, required=True)
+parser.add_argument("--target_phrase", type=str, required=True)
+parser.add_argument("--mask_scope", type=float, default=0.20)
+parser.add_argument("--mask_strategy",  type=str, nargs="+", default=["max_norm"])
+parser.add_argument("--mask_reused_step", type=int, default=12)
+
+args = parser.parse_args()
+
+# Initialize unet model
+with open(args.unet_config) as unet_config_file:
+    unet_config = json.load(unet_config_file)
+
+    # Settings for image encoder
+    vision_model_config = unet_config.pop("vision_model_config", None)
+    args.vision_model_config = vision_model_config.pop("vision_model_config", None)
+
+    unet_type = unet_config.pop("type", None)
+    unet_model = UNet2DConditionModelDiffusers(**unet_config)
+
+unet_model.eval().to(args.device)
+unet_model.load_state_dict(torch.load(args.unet_checkpoint, map_location=args.device), strict=False)
+print("loading unet model finished.")
+
+if args.unet_checkpoint_base_model != "":
+    if "safetensors" in args.unet_checkpoint_base_model:
+        from safetensors import safe_open
+        tensors = {}
+        with safe_open(args.unet_checkpoint_base_model, framework="pt", device='cpu') as f:
+            for k in f.keys():
+                new_k = k.replace("model.diffusion_model.", "")
+                tensors[k] = f.get_tensor(k)
+        unet_model.load_state_dict(tensors, strict=False)
+    else:
+        unet_model.load_state_dict(torch.load(args.unet_checkpoint_base_model, map_location=args.device), strict=False)
+unet_model = torch.compile(unet_model, disable=not args.compile)
+print("loading unet base model finished.")
+
+# Initialize vae model
+with open(args.vae_config) as vae_config_file:
+    vae_config = json.load(vae_config_file)
+vae_downsample_factor = 2 ** (len(vae_config["block_out_channels"]) - 1) # 2 ** 3 = 8
+vae_model = AutoencoderKL(**vae_config)
+vae_model.eval().to(args.device)
+vae_model.load_state_dict(torch.load(args.vae_checkpoint, map_location=args.device))
+vae_decoder = torch.compile(lambda x: vae_model.decode(x / vae_model.scaling_factor).sample.clip(-1, 1), disable=not args.compile)
+vae_encoder = torch.compile(lambda x: vae_model.encode(x).latent_dist.mode().mul_(vae_model.scaling_factor), disable=not args.compile)
+print("loading vae finished.")
+
+# Initialize ddim scheduler
+ddim_train_steps = 1000
+ddim_betas = get_betas(name=args.schedule_type, num_steps=ddim_train_steps, shift_snr=args.schedule_shift_snr, terminal_pure_noise=False)
+scheduler_class = DPMSolverSingleStepScheduler if args.scheduler_type == 'dpm' else DDIMScheduler
+scheduler = scheduler_class(betas=ddim_betas, num_train_timesteps=ddim_train_steps, num_inference_timesteps=args.sample_steps, device=args.device)
+infer_timesteps = scheduler.timesteps
+
+# Initialize text model
+text_model = TextModel(args.text_encoder_variant, ["penultimate_nonorm"])
+text_model.eval().to(args.device)
+print("loading text model finished.")
+
+# Initialize image model.
+vision_model = instantiate_from_config(args.vision_model_config)
+vision_model = vision_model.eval().to(args.device)
+print("loading image model finished.")
+
+negative_prompt = ""
+if args.negative_prompt:
+    with open(args.negative_prompt) as f:
+        negative_prompt = f.read().strip()
+
+image_metadata_validate = torch.tensor(
+    data=[
+        args.width,     # original_height
+        args.height,    # original_width
+        0,              # coordinate top
+        0,              # coordinate left
+        args.width,     # target_height
+        args.height,    # target_width
+    ],
+    device=args.device,
+    dtype=torch.float32
+).view(1, -1).repeat(args.samples_per_prompt, 1)
+
+# Create output directory
+os.makedirs(args.output_dir, exist_ok=True)
+args.output_image_grid_dir = os.path.join(args.output_dir, "images_grid")
+args.output_image_dir = os.path.join(args.output_dir, "images")
+args.output_mask_grid_dir = os.path.join(args.output_dir, "masks_grid")
+args.output_mask_dir = os.path.join(args.output_dir, "masks")
+os.makedirs(args.output_image_grid_dir, exist_ok=True)
+os.makedirs(args.output_image_dir, exist_ok=True)
+os.makedirs(args.output_mask_grid_dir, exist_ok=True)
+os.makedirs(args.output_mask_dir, exist_ok=True)
+
+with torch.no_grad():
+    # Prepare negative prompt.
+    if args.guidance_weight != 1:
+        text_negative_output = text_model(negative_prompt)
+
+    positive_prompt = args.text_prompt
+    positive_promt_image_path = args.image_prompt_path
+    target_phrase = args.target_phrase
+
+    # Compute target phrases
+    target_token = torch.zeros(1, 77).to(args.device)
+    positions = find_phrase_positions_in_text(positive_prompt, target_phrase)
+    for position in positions:
+        prompt_before = positive_prompt[:position] # NOTE We do not need -1 here because the SDXL text encoder does not encode the trailing space.
+        prompt_include = positive_prompt[:position+len(target_phrase)]
+        print("prompt before: ", prompt_before, ", prompt_include: ", prompt_include)
+        prompt_before_length = text_model.get_vaild_token_length(prompt_before) + 1
+        prompt_include_length = text_model.get_vaild_token_length(prompt_include) + 1
+        print("prompt_before_length: ", prompt_before_length, ", prompt_include_length: ", prompt_include_length)
+        target_token[:, prompt_before_length:prompt_include_length] = 1
+
+    # Text used for progress bar
+    pbar_text = positive_prompt[:40]
+
+    # Compute text embeddings
+    text_positive_output = text_model(positive_prompt)
+    text_positive_embeddings = text_positive_output.embeddings.repeat_interleave(args.samples_per_prompt, dim=0)
+    text_positive_pooled = text_positive_output.pooled[-1].repeat_interleave(args.samples_per_prompt, dim=0)
+    if args.guidance_weight != 1:
+        text_negative_embeddings = text_negative_output.embeddings.repeat_interleave(args.samples_per_prompt, dim=0)
+        text_negative_pooled = text_negative_output.pooled[-1].repeat_interleave(args.samples_per_prompt, dim=0)
+    
+    # Compute image embeddings
+    positive_image = Image.open(positive_promt_image_path).convert("RGB")
+    positive_image = torchvision.transforms.ToTensor()(positive_image)
+
+    positive_image = positive_image.unsqueeze(0).repeat_interleave(args.samples_per_prompt, dim=0)
+    positive_image = torch.nn.functional.interpolate(
+        positive_image, 
+        size=(768, 768), 
+        mode="bilinear", 
+        align_corners=False
+    )
+    negative_image = torch.zeros_like(positive_image)
+    print(positive_image.size(), negative_image.size())
+    positive_image = positive_image.to(args.device)
+    negative_image = negative_image.to(args.device)
+
+    positive_image_dict = {"image_ref": positive_image}
+    positive_image_output = vision_model(positive_image_dict, device=args.device)
+
+    negative_image_dict = {"image_ref": negative_image}
+    negative_image_output = vision_model(negative_image_dict, device=args.device)
+
+    # Initialize latent with input latent + noise (i2i) / pure noise (t2i)
+    latent = torch.randn(
+        size=[
+            args.samples_per_prompt,
+            vae_config["latent_channels"],
+            args.height // vae_downsample_factor,
+            args.width // vae_downsample_factor
+        ],
+        device=args.device,
+        generator=torch.Generator(args.device).manual_seed(args.seed))
+    target_h = (args.height // vae_downsample_factor) // 2
+    target_w = (args.width // vae_downsample_factor) // 2
+
+    # Real Reverse diffusion process.
+    text2image_crossmap_2d_all_timesteps_list = []
+    current_step = 0
+    for timestep in tqdm(iterable=infer_timesteps, desc=f"[{pbar_text}]", dynamic_ncols=True):
+        if current_step < args.mask_reused_step:
+            pred_cond, pred_cond_dict = unet_model(
+                sample=latent,
+                timestep=timestep,
+                encoder_hidden_states=text_positive_embeddings,
+                encoder_attention_mask=None,
+                added_cond_kwargs=dict(
+                    text_embeds=text_positive_pooled,
+                    time_ids=image_metadata_validate
+                ),
+                vision_input_dict=None,
+                vision_guided_mask=None,
+                return_as_origin=False,
+                return_text2image_mask=True,
+            )
+            crossmap_2d_avg = mask_generation(
+                crossmap_2d_list=pred_cond_dict["text2image_crossmap_2d"], selfmap_2d_list=pred_cond_dict.get("self_attention_map", []), 
+                target_token=target_token, mask_scope=args.mask_scope,
+                mask_target_h=target_h, mask_target_w=target_w, mask_mode=args.mask_strategy,
+            )
+        else:
+            # using previous step's mask
+            crossmap_2d_avg = text2image_crossmap_2d_all_timesteps_list[-1].squeeze(1)
+        if crossmap_2d_avg.dim() == 5: # Means that each layer uses a separate mask weight.
+            text2image_crossmap_2d_all_timesteps_list.append(crossmap_2d_avg.mean(dim=2).unsqueeze(1))
+        else:
+            text2image_crossmap_2d_all_timesteps_list.append(crossmap_2d_avg.unsqueeze(1))
+
+        pred_cond, pred_cond_dict = unet_model(
+            sample=latent,
+            timestep=timestep,
+            encoder_hidden_states=text_positive_embeddings,
+            encoder_attention_mask=None,
+            added_cond_kwargs=dict(
+                text_embeds=text_positive_pooled,
+                time_ids=image_metadata_validate
+            ),
+            vision_input_dict=positive_image_output,
+            vision_guided_mask=crossmap_2d_avg,
+            return_as_origin=False,
+            return_text2image_mask=True,
+            multiple_reference_image=False
+        )
+
+        crossmap_2d_avg_neg = crossmap_2d_avg.mean(dim=1, keepdim=True)
+        pred_negative, pred_negative_dict = unet_model(
+            sample=latent,
+            timestep=timestep,
+            encoder_hidden_states=text_negative_embeddings,
+            encoder_attention_mask=None,
+            added_cond_kwargs=dict(
+                text_embeds=text_negative_pooled,
+                time_ids=image_metadata_validate
+            ),
+            vision_input_dict=negative_image_output,
+            vision_guided_mask=crossmap_2d_avg,
+            return_as_origin=False,
+            return_text2image_mask=True,
+            multiple_reference_image=False
+        )
+
+        pred = classifier_free_guidance_image_prompt_cascade(
+            pred_t_cond=None, pred_ti_cond=pred_cond, pred_uncond=pred_negative, 
+            guidance_weight_t=args.guidance_weight, guidance_weight_i=args.guidance_weight, 
+            guidance_stdev_rescale_factor=0, cfg_rescale_mode="naive_global_direct"
+        )
+        step = scheduler.step(
+            model_output=pred,
+            model_output_type=args.unet_prediction,
+            timestep=timestep,
+            sample=latent)
+
+        latent = step.prev_sample
+
+        current_step += 1
+
+    sample = vae_decoder(step.pred_original_sample)
+
+    # save each image 
+    for sample_i in range(sample.size(0)):
+        sample_i_image = torch.clamp(sample[sample_i] * 0.5 + 0.5, min=0, max=1).float()
+        to_pil_image(sample_i_image).save(args.output_image_dir + "/output_{}.jpg".format(sample_i))
+
+    # save grid images
+    sample = make_grid(sample, normalize=True, value_range=(-1, 1), nrow=args.nrow).float()
+    to_pil_image(sample).save(args.output_image_grid_dir + "/grid_image.jpg")

RealCustom/inference/inference_utils.py

@@ -0,0 +1,76 @@
+# Copyright (c) 2024 Bytedance Ltd. and/or its affiliates
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+def find_phrase_positions_in_text(text, phrase):
+    """
+    Return the position of the first character of the phrase in the text.
+    """
+    
+    position = -1
+    positions = []
+    while True:
+        position = text.find(phrase, position + 1)
+        if position == -1:
+            break
+        positions.append(position)
+    return positions
+
+def classifier_free_guidance_image_prompt_cascade(
+    pred_t_cond, pred_ti_cond, pred_uncond, guidance_weight_t=7.5, guidance_weight_i=7.5, 
+    guidance_stdev_rescale_factor=0.7, cfg_rescale_mode="none", super_cross_mask=None
+    ):
+
+    if cfg_rescale_mode == "none":
+        pred = pred_uncond + guidance_weight_t * (pred_t_cond - pred_uncond) + guidance_weight_i * (pred_ti_cond - pred_t_cond)
+    elif cfg_rescale_mode == "none_direct":
+        pred = pred_uncond + guidance_weight_i * (pred_ti_cond - pred_uncond)
+    elif cfg_rescale_mode == "naive":
+        assert super_cross_mask is not None
+        pred_std_t_before = pred_t_cond.std([1, 2, 3], keepdim=True)
+        pred_std_ti_before = pred_ti_cond.std([1, 2, 3], keepdim=True)
+
+        pred = pred_uncond + guidance_weight_t * (pred_t_cond - pred_uncond) + guidance_weight_i * (pred_ti_cond - pred_t_cond)
+
+        pred_std_after = pred.std([1, 2, 3], keepdim=True)
+
+        pred_rescale_t_factor = guidance_stdev_rescale_factor * (pred_std_t_before / pred_std_after) + (1 - guidance_stdev_rescale_factor)
+        pred_rescale_ti_factor = guidance_stdev_rescale_factor * (pred_std_ti_before / pred_std_after) + (1 - guidance_stdev_rescale_factor)
+
+        pred_ti = pred * super_cross_mask
+        pred_t = pred * (1 - super_cross_mask)
+        pred = pred_ti * pred_rescale_ti_factor + pred_t * pred_rescale_t_factor
+    elif cfg_rescale_mode == "naive_global":
+        pred_std_ti_before = pred_ti_cond.std([1, 2, 3], keepdim=True)
+
+        pred = pred_uncond + guidance_weight_t * (pred_t_cond - pred_uncond) + guidance_weight_i * (pred_ti_cond - pred_t_cond)
+
+        pred_std_after = pred.std([1, 2, 3], keepdim=True)
+
+        pred_rescale_ti_factor = guidance_stdev_rescale_factor * (pred_std_ti_before / pred_std_after) + (1 - guidance_stdev_rescale_factor)
+
+        pred = pred * pred_rescale_ti_factor
+    elif cfg_rescale_mode == "naive_global_direct":
+        pred_std_ti_before = pred_ti_cond.std([1, 2, 3], keepdim=True)
+
+        pred = pred_uncond + guidance_weight_i * (pred_ti_cond - pred_uncond)
+
+        pred_std_after = pred.std([1, 2, 3], keepdim=True)
+
+        pred_rescale_ti_factor = guidance_stdev_rescale_factor * (pred_std_ti_before / pred_std_after) + (1 - guidance_stdev_rescale_factor)
+
+        pred = pred * pred_rescale_ti_factor
+    else:
+        raise NotImplementedError()
+
+    return pred

RealCustom/inference/mask_generation.py

@@ -0,0 +1,114 @@
+# Copyright (c) 2024 Bytedance Ltd. and/or its affiliates
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import torch
+import torch.nn.functional as F
+from einops import rearrange
+
+def mask_generation(
+    crossmap_2d_list, selfmap_2d_list=None, 
+    target_token=None, mask_scope=None,
+    mask_target_h=64, mask_target_w=64,
+    mask_mode=["binary"],
+):  
+    if len(selfmap_2d_list) > 0:
+        target_hw_selfmap = mask_target_h * mask_target_w
+        selfmap_2ds = []
+        for i in range(len(selfmap_2d_list)):
+            selfmap_ = selfmap_2d_list[i]
+            selfmap_ = F.interpolate(selfmap_, size=(target_hw_selfmap, target_hw_selfmap), mode='bilinear')
+            selfmap_2ds.append(selfmap_ )
+        selfmap_2ds = torch.cat(selfmap_2ds, dim=1)
+        if "selfmap_min_max_per_channel" in mask_mode:
+            selfmap_1ds = rearrange(selfmap_2ds, "b c h w -> b c (h w)")
+            channel_max_self = torch.max(selfmap_1ds, dim=-1, keepdim=True)[0].unsqueeze(-1)
+            channel_min_self = torch.min(selfmap_1ds, dim=-1, keepdim=True)[0].unsqueeze(-1)
+            selfmap_2ds = (selfmap_2ds - channel_min_self) / (channel_max_self - channel_min_self + 1e-6)
+        elif "selfmap_max_norm" in mask_mode:
+            selfmap_1ds = rearrange(selfmap_2ds, "b c h w -> b c (h w)")
+            b = selfmap_1ds.size(0)
+            batch_max = torch.max(selfmap_1ds.view(b, -1), dim=-1, keepdim=True)[0].unsqueeze(-1).unsqueeze(-1)
+            selfmap_2ds  = selfmap_2ds / (batch_max + 1e-10)
+
+        selfmap_2d = selfmap_2ds.mean(dim=1, keepdim=True)
+    else:
+        selfmap_2d = None
+    
+    crossmap_2ds = []
+    for i in range(len(crossmap_2d_list)):
+        crossmap = crossmap_2d_list[i]
+        crossmap = crossmap.mean(dim=1)  # average on head dim
+        crossmap = crossmap * target_token.unsqueeze(-1).unsqueeze(-1) # target token valid
+        crossmap = crossmap.sum(dim=1, keepdim=True)
+
+        crossmap = F.interpolate(crossmap, size=(mask_target_h, mask_target_w), mode='bilinear')
+        crossmap_2ds.append(crossmap)
+    crossmap_2ds = torch.cat(crossmap_2ds, dim=1)
+    crossmap_1ds = rearrange(crossmap_2ds, "b c h w -> b c (h w)")
+
+    if "max_norm" in mask_mode:
+        crossmap_1d_avg = torch.mean(crossmap_1ds, dim=1, keepdim=True)  # [b, 1, (h w)]
+        if selfmap_2d is not None:
+            crossmap_1d_avg = torch.matmul(selfmap_2d, crossmap_1d_avg.unsqueeze(-1)).squeeze(-1)
+        b, c, n = crossmap_1ds.shape
+        batch_max = torch.max(crossmap_1d_avg.view(b, -1), dim=-1, keepdim=True)[0].unsqueeze(1)
+        crossmap_1d_avg = crossmap_1d_avg / (batch_max + 1e-6)
+    elif "min_max_norm" in mask_mode:
+        crossmap_1d_avg = torch.mean(crossmap_1ds, dim=1, keepdim=True)  # [b, 1, (h w)]
+        if selfmap_2d is not None:
+            crossmap_1d_avg = torch.matmul(selfmap_2d, crossmap_1d_avg.unsqueeze(-1)).squeeze(-1)
+        b, c, n = crossmap_1ds.shape
+        batch_max = torch.max(crossmap_1d_avg.view(b, -1), dim=-1, keepdim=True)[0].unsqueeze(1) # NOTE unsqueeze
+        batch_min = torch.min(crossmap_1d_avg.view(b, -1), dim=-1, keepdim=True)[0].unsqueeze(1) # NOTE unsqueeze
+        crossmap_1d_avg = (crossmap_1d_avg - batch_min) / (batch_max - batch_min + 1e-6)
+    elif "min_max_per_channel" in mask_mode:
+        channel_max = torch.max(crossmap_1ds, dim=-1, keepdim=True)[0]
+        channel_min = torch.min(crossmap_1ds, dim=-1, keepdim=True)[0]
+        crossmap_1ds = (crossmap_1ds - channel_min) / (channel_max - channel_min + 1e-6)
+        crossmap_1d_avg = torch.mean(crossmap_1ds, dim=1, keepdim=True)  # [b, 1, (h w)]
+        if selfmap_2d is not None:
+            crossmap_1d_avg = torch.matmul(selfmap_2d, crossmap_1d_avg.unsqueeze(-1)).squeeze(-1)
+
+        # renormalize to 0-1
+        b, c, n = crossmap_1d_avg.shape
+        batch_max = torch.max(crossmap_1d_avg.view(b, -1), dim=-1, keepdim=True)[0].unsqueeze(1)
+        batch_min = torch.min(crossmap_1d_avg.view(b, -1), dim=-1, keepdim=True)[0].unsqueeze(1)
+        crossmap_1d_avg = (crossmap_1d_avg - batch_min) / (batch_max - batch_min + 1e-6)
+    else:
+        crossmap_1d_avg = torch.mean(crossmap_1ds, dim=1, keepdim=True)  # [b, 1, (h w)]
+        
+        
+    if "threshold" in mask_mode:
+        threshold = 1 - mask_scope
+        crossmap_1d_avg[crossmap_1d_avg < threshold] = 0.0
+        if "binary" in mask_mode:
+            crossmap_1d_avg[crossmap_1d_avg > threshold] = 1.0
+    else:
+        # topk
+        topk_num = int(crossmap_1d_avg.size(-1) * mask_scope)
+        sort_score, sort_order = crossmap_1d_avg.sort(descending=True, dim=-1)
+        sort_topk = sort_order[:, :, :topk_num]
+        sort_topk_remain = sort_order[:, :, topk_num:]
+        crossmap_1d_avg = crossmap_1d_avg.scatter(2, sort_topk_remain, 0.)
+        if "binary" in mask_mode:
+            crossmap_1d_avg = crossmap_1d_avg.scatter(2, sort_topk, 1.0)
+
+    crossmap_2d_avg = rearrange(crossmap_1d_avg, "b c (h w) -> b c h w", h=mask_target_h, w=mask_target_w)
+    crossmap_2d_avg = crossmap_2d_avg
+
+    output = crossmap_2d_avg.unsqueeze(1)  # torch.Size([4, 1, 60, 64, 64]), The second dimension is the dimension of the number of reference images.
+    if output.size(2) == 1: # The dimension of the layer.  
+        output = output.squeeze(2)  # If there is only a single dimension, then all layers will share the same mask.
+
+    return output

RealCustom/inference/pipeline.py

@@ -0,0 +1,360 @@
+# Copyright (c) 2024 Bytedance Ltd. and/or its affiliates
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import os
+import json
+import torch
+import torchvision
+from torchvision.utils import make_grid
+from torchvision.transforms.functional import to_pil_image
+from PIL import Image
+
+from ..models.text import TextModel
+from ..models.vae import AutoencoderKL
+from ..models.unet_2d_condition_custom import UNet2DConditionModel as UNet2DConditionModelDiffusers
+
+from ..schedulers.ddim import DDIMScheduler
+from ..schedulers.dpm_s import DPMSolverSingleStepScheduler
+from ..schedulers.utils import get_betas
+
+from .inference_utils import find_phrase_positions_in_text, classifier_free_guidance_image_prompt_cascade
+from .mask_generation import mask_generation
+from ..utils import instantiate_from_config
+
+from tqdm import tqdm
+from einops import rearrange
+
+class RealCustomInferencePipeline:
+    def __init__(
+        self,
+        unet_config,
+        unet_checkpoint,
+        realcustom_checkpoint,
+        vae_config="ckpts/sdxl/vae/sdxl.json",
+        vae_checkpoint="ckpts/sdxl/vae/sdxl-vae.pth",
+        model_type="bf16", 
+        device="cuda",
+    ):
+        if model_type == "bf16":
+            self.torch_dtype = torch.bfloat16
+        else:
+            self.torch_dtype = torch.float32
+
+        if not os.path.exists("ckpts/"):
+            from huggingface_hub import snapshot_download
+            print("Downloading RealCustom ...")
+            snapshot_download(
+                repo_id="bytedance-research/RealCustom",
+                repo_type="model",
+                local_dir="ckpts",  # 指定本地目录
+                allow_patterns="ckpts/**",  # 只下载 ckpts 文件夹内容
+                local_dir_use_symlinks=False,  # 直接存储文件而非符号链接
+            )
+        
+        self.device = device
+        self.unet_checkpoint = unet_checkpoint
+        self.realcustom_checkpoint = realcustom_checkpoint
+        self._load_unet_checkpoint(unet_config, unet_checkpoint, realcustom_checkpoint)
+        self._load_vae_checkpoint(vae_config, vae_checkpoint)
+        self._load_encoder_checkpoint()
+        self._init_scheduler()
+        self._load_negative_prompt()
+
+    
+    def _load_unet_checkpoint(self, unet_config, unet_checkpoint, realcustom_checkpoint):
+        # Initialize unet model
+        with open(unet_config) as unet_config_file:
+            unet_config = json.load(unet_config_file)
+        self.unet_prediction = "epsilon"
+
+        # Settings for image encoder
+        vision_model_config = unet_config.pop("vision_model_config", None)
+        self.vision_model_config = vision_model_config.pop("vision_model_config", None)
+
+        self.unet_model = UNet2DConditionModelDiffusers(**unet_config)
+
+        self.unet_model.eval().to(self.device).to(self.torch_dtype)
+        self.unet_model.load_state_dict(torch.load(unet_checkpoint, map_location=self.device), strict=False)
+        self.unet_model.load_state_dict(torch.load(realcustom_checkpoint, map_location=self.device), strict=False)
+        print("loading unet model finished.")
+
+    def _reload_unet_checkpoint(self, unet_checkpoint, realcustom_checkpoint):
+        self.unet_model.load_state_dict(torch.load(unet_checkpoint, map_location=self.device), strict=False)
+        self.unet_model.load_state_dict(torch.load(realcustom_checkpoint, map_location=self.device), strict=False)
+        print("reloading unet model finished.")
+    
+    def _load_vae_checkpoint(self, vae_config, vae_checkpoint):
+        # Initialize vae model
+        with open(vae_config) as vae_config_file:
+            vae_config = json.load(vae_config_file)
+        self.latent_channels = vae_config["latent_channels"]
+        self.vae_downsample_factor = 2 ** (len(vae_config["block_out_channels"]) - 1) # 2 ** 3 = 8
+
+        vae_model = AutoencoderKL(**vae_config)
+        vae_model.eval().to(self.device).to(self.torch_dtype)
+        vae_model.load_state_dict(torch.load(vae_checkpoint, map_location=self.device))
+        self.vae_decoder = torch.compile(lambda x: vae_model.decode(x / vae_model.scaling_factor).sample.clip(-1, 1), disable=True)
+        self.vae_encoder = torch.compile(lambda x: vae_model.encode(x).latent_dist.mode().mul_(vae_model.scaling_factor), disable=True)
+
+        print("loading vae finished.")
+    
+    def _load_encoder_checkpoint(self, ):
+        # Initialize text encoder
+        text_encoder_variant = ["ckpts/sdxl/clip-sdxl-1", "ckpts/sdxl/clip-sdxl-2"]
+        text_encoder_mode = ["penultimate_nonorm"]
+        self.text_model = TextModel(text_encoder_variant, text_encoder_mode)
+        self.text_model.eval().to(self.device).to(self.torch_dtype)
+        print("loading text model finished.")
+
+        # Initialize image encoder
+        self.vision_model = instantiate_from_config(self.vision_model_config)
+        self.vision_model.eval().to(self.device).to(self.torch_dtype)
+        print("loading image model finished.")
+    
+    def _init_scheduler(self, ):
+        # Initialize ddim scheduler
+        ddim_train_steps = 1000
+        schedule_type = "squared_linear"
+        scheduler_type = "dpm"
+        schedule_shift_snr = 1
+        self.sample_steps = 25
+        ddim_betas = get_betas(name=schedule_type, num_steps=ddim_train_steps, shift_snr=schedule_shift_snr, terminal_pure_noise=False)
+        scheduler_class = DPMSolverSingleStepScheduler if scheduler_type == 'dpm' else DDIMScheduler
+
+        self.scheduler = scheduler_class(betas=ddim_betas, num_train_timesteps=ddim_train_steps, num_inference_timesteps=self.sample_steps, device=self.device)
+        self.infer_timesteps = self.scheduler.timesteps
+    
+    def _load_negative_prompt(self, ):
+        with open("prompts/validation_negative.txt") as f:
+            self.negative_prompt = f.read().strip()
+        self.text_negative_output = self.text_model(self.negative_prompt)
+    
+    def generation(
+        self, 
+        text, 
+        image_pil, 
+        target_phrase,
+
+        height=1024,
+        width=1024,
+        guidance_scale=3.5,
+        seed=1234,
+        samples_per_prompt=4,
+
+        mask_scope=0.25,
+
+        new_unet_checkpoint="",          # in case you want to change
+        new_realcustom_checkpoint="",    # in case you want to change
+        mask_strategy=["min_max_per_channel"],
+        mask_reused_step=12,
+        return_each_image=False,
+    ):
+        
+        if new_unet_checkpoint != "" and new_unet_checkpoint != self.unet_checkpoint:
+            self.unet_checkpoint = new_unet_checkpoint
+            self.unet_model.load_state_dict(torch.load(new_unet_checkpoint, map_location=self.device), strict=False)
+            print("Reloading Unet {} finised.".format(new_unet_checkpoint))
+        if new_realcustom_checkpoint != "" and new_realcustom_checkpoint != self.realcustom_checkpoint:
+            self.realcustom_checkpoint = new_realcustom_checkpoint
+            self.unet_model.load_state_dict(torch.load(new_realcustom_checkpoint, map_location=self.device), strict=False)
+            print("Reloading RealCustom {} finised.".format(new_realcustom_checkpoint))
+
+        samples_per_prompt = int(samples_per_prompt)
+        image_metadata_validate = self._get_metadata(height, width, samples_per_prompt)
+        if seed == -1:
+            seed = torch.randint(0, 1000000, (1,)).item()
+        seed = int(seed)
+
+        with torch.no_grad(), torch.autocast(self.device, self.torch_dtype):
+            target_token = self._find_phrase_positions_in_text(text, target_phrase)
+
+            # Compute text embeddings
+            text_positive_output = self.text_model(text)
+            text_positive_embeddings = text_positive_output.embeddings.repeat_interleave(samples_per_prompt, dim=0)
+            text_positive_pooled = text_positive_output.pooled[-1].repeat_interleave(samples_per_prompt, dim=0)
+            if guidance_scale != 1:
+                text_negative_embeddings = self.text_negative_output.embeddings.repeat_interleave(samples_per_prompt, dim=0)
+                text_negative_pooled = self.text_negative_output.pooled[-1].repeat_interleave(samples_per_prompt, dim=0)
+
+            # Compute image embeddings
+            # positive_image = Image.open(image_path).convert("RGB")
+            positive_image = image_pil
+            positive_image = torchvision.transforms.ToTensor()(positive_image)
+
+            positive_image = positive_image.unsqueeze(0).repeat_interleave(samples_per_prompt, dim=0)
+            positive_image = torch.nn.functional.interpolate(
+                positive_image, 
+                size=(768, 768), 
+                mode="bilinear", 
+                align_corners=False
+            )
+            negative_image = torch.zeros_like(positive_image)
+            positive_image = positive_image.to(self.device).to(self.torch_dtype)
+            negative_image = negative_image.to(self.device).to(self.torch_dtype)
+
+            positive_image_dict = {"image_ref": positive_image}
+            positive_image_output = self.vision_model(positive_image_dict, device=self.device)
+
+            negative_image_dict = {"image_ref": negative_image}
+            negative_image_output = self.vision_model(negative_image_dict, device=self.device)
+
+            # Initialize latent with input latent
+            latent = torch.randn(
+                size=[
+                    samples_per_prompt,
+                    self.latent_channels,
+                    height // self.vae_downsample_factor,
+                    width // self.vae_downsample_factor
+                ],
+                device=self.device,
+                generator=torch.Generator(self.device).manual_seed(seed)).to(self.torch_dtype)
+            target_h = (height // self.vae_downsample_factor) // 2
+            target_w = (width // self.vae_downsample_factor) // 2
+
+            text2image_crossmap_2d_all_timesteps_list = []
+            current_step = 0
+            pbar_text = text[:40]
+            for timestep in tqdm(iterable=self.infer_timesteps, desc=f"[{pbar_text}]", dynamic_ncols=True):
+                if current_step < mask_reused_step:
+                    pred_cond, pred_cond_dict = self.unet_model(
+                        sample=latent,
+                        timestep=timestep,
+                        encoder_hidden_states=text_positive_embeddings,
+                        encoder_attention_mask=None,
+                        added_cond_kwargs=dict(
+                            text_embeds=text_positive_pooled,
+                            time_ids=image_metadata_validate
+                        ),
+                        vision_input_dict=None,
+                        vision_guided_mask=None,
+                        return_as_origin=False,
+                        return_text2image_mask=True,
+                    )
+                   
+                    crossmap_2d_avg = mask_generation(
+                        crossmap_2d_list=pred_cond_dict.get("text2image_crossmap_2d",[]), selfmap_2d_list=pred_cond_dict.get("self_attention_map", []), 
+                        target_token=target_token, mask_scope=mask_scope,
+                        mask_target_h=target_h, mask_target_w=target_w, mask_mode=mask_strategy,
+                    )
+                else:
+                    # using previous step's mask
+                    crossmap_2d_avg = text2image_crossmap_2d_all_timesteps_list[-1].squeeze(1)
+                if crossmap_2d_avg.dim() == 5: # Means that each layer uses a separate mask weight.
+                    text2image_crossmap_2d_all_timesteps_list.append(crossmap_2d_avg.mean(dim=2).unsqueeze(1))
+                else:
+                    text2image_crossmap_2d_all_timesteps_list.append(crossmap_2d_avg.unsqueeze(1))
+
+                pred_cond, pred_cond_dict = self.unet_model(
+                    sample=latent,
+                    timestep=timestep,
+                    encoder_hidden_states=text_positive_embeddings,
+                    encoder_attention_mask=None,
+                    added_cond_kwargs=dict(
+                        text_embeds=text_positive_pooled,
+                        time_ids=image_metadata_validate
+                    ),
+                    vision_input_dict=positive_image_output,
+                    vision_guided_mask=crossmap_2d_avg,
+                    return_as_origin=False,
+                    return_text2image_mask=True,
+                    multiple_reference_image=False
+                )
+
+                pred_negative, pred_negative_dict = self.unet_model(
+                    sample=latent,
+                    timestep=timestep,
+                    encoder_hidden_states=text_negative_embeddings,
+                    encoder_attention_mask=None,
+                    added_cond_kwargs=dict(
+                        text_embeds=text_negative_pooled,
+                        time_ids=image_metadata_validate
+                    ),
+                    vision_input_dict=negative_image_output,
+                    vision_guided_mask=crossmap_2d_avg,
+                    return_as_origin=False,
+                    return_text2image_mask=True,
+                    multiple_reference_image=False
+                )
+
+                pred = classifier_free_guidance_image_prompt_cascade(
+                    pred_t_cond=None, pred_ti_cond=pred_cond, pred_uncond=pred_negative, 
+                    guidance_weight_t=guidance_scale, guidance_weight_i=guidance_scale, 
+                    guidance_stdev_rescale_factor=0, cfg_rescale_mode="naive_global_direct"
+                )
+                step = self.scheduler.step(
+                    model_output=pred,
+                    model_output_type=self.unet_prediction,
+                    timestep=timestep,
+                    sample=latent)
+
+                latent = step.prev_sample
+
+                current_step += 1
+            sample = self.vae_decoder(step.pred_original_sample)
+        
+        # save each image
+        images_pil_list = []
+        for sample_i in range(sample.size(0)):
+            sample_i_image = torch.clamp(sample[sample_i] * 0.5 + 0.5, min=0, max=1).float()
+
+            images_pil_list.append(to_pil_image(sample_i_image))
+            # to_pil_image(sample_i_image).save("./test_{}.jpg".format(sample_i))
+
+        # save grid images
+        sample = make_grid(sample, normalize=True, value_range=(-1, 1), nrow=int(samples_per_prompt ** 0.5)).float()
+        # to_pil_image(sample).save("./output_grid_image.jpg")
+
+        # save all masks
+        text2image_crossmap_2d_all_timesteps = torch.cat(text2image_crossmap_2d_all_timesteps_list, dim=1)
+        text2image_crossmap_2d_all_timesteps = rearrange(text2image_crossmap_2d_all_timesteps, "b t c h w -> (b t) c h w")
+        c = text2image_crossmap_2d_all_timesteps.size(1)
+        text2image_crossmap_2d_all_timesteps = rearrange(text2image_crossmap_2d_all_timesteps, "B (c 1) h w -> (B c) 1 h w")
+        sample_mask = make_grid(text2image_crossmap_2d_all_timesteps, normalize=False, value_range=(-1, 1), nrow=int(self.sample_steps * c))
+        # to_pil_image(sample_mask).save("./output_grid_mask.jpg")
+
+        if return_each_image:
+            return images_pil_list, to_pil_image(sample), to_pil_image(sample_mask)
+        else:
+            return to_pil_image(sample), to_pil_image(sample_mask)
+
+    def _get_metadata(self, height, width, samples_per_prompt):
+        image_metadata_validate = torch.tensor(
+            data=[
+                width,     # original_height
+                height,    # original_width
+                0,         # coordinate top
+                0,         # coordinate left
+                width,     # target_height
+                height,    # target_width
+            ],
+            device=self.device,
+            dtype=self.torch_dtype
+        ).view(1, -1).repeat(samples_per_prompt, 1)
+
+        return image_metadata_validate
+    
+    def _find_phrase_positions_in_text(self, text, target_phrase):
+        # Compute target phrases
+        target_token = torch.zeros(1, 77).to(self.device)
+        positions = find_phrase_positions_in_text(text, target_phrase)
+        for position in positions:
+            prompt_before = text[:position] # NOTE We do not need -1 here because the SDXL text encoder does not encode the trailing space.
+            prompt_include = text[:position+len(target_phrase)]
+            print("prompt before: ", prompt_before, ", prompt_include: ", prompt_include)
+            prompt_before_length = self.text_model.get_vaild_token_length(prompt_before) + 1
+            prompt_include_length = self.text_model.get_vaild_token_length(prompt_include) + 1
+            print("prompt_before_length: ", prompt_before_length, ", prompt_include_length: ", prompt_include_length)
+            target_token[:, prompt_before_length:prompt_include_length] = 1
+        
+        return target_token