317 lines
14 KiB
Python
317 lines
14 KiB
Python
# 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") |