diff --git a/RealCustom/inference/app.py b/RealCustom/inference/app.py new file mode 100644 index 0000000..f9cb4a6 --- /dev/null +++ b/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""" +
+ Build + Build +
+ """.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) \ No newline at end of file diff --git a/RealCustom/inference/inference_single_image.py b/RealCustom/inference/inference_single_image.py new file mode 100644 index 0000000..359ccc8 --- /dev/null +++ b/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") \ No newline at end of file diff --git a/RealCustom/inference/inference_utils.py b/RealCustom/inference/inference_utils.py new file mode 100644 index 0000000..0bcb60b --- /dev/null +++ b/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 \ No newline at end of file diff --git a/RealCustom/inference/mask_generation.py b/RealCustom/inference/mask_generation.py new file mode 100644 index 0000000..6fb5809 --- /dev/null +++ b/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 \ No newline at end of file diff --git a/RealCustom/inference/pipeline.py b/RealCustom/inference/pipeline.py new file mode 100644 index 0000000..a4b6a68 --- /dev/null +++ b/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 \ No newline at end of file