'''
Part of the implementation is borrowed and modified from TextDiffuser, publicly available at https://github.com/microsoft/unilm/blob/master/textdiffuser/inference.py
'''
import os
import json
import random
from tqdm import tqdm
import argparse
import numpy as np
from packaging import version
from termcolor import colored
from PIL import Image
from datasets import disable_caching
import torch
import torch.utils.checkpoint
from torchvision import transforms
from accelerate.utils import set_seed
from diffusers import AutoencoderKL, DDPMScheduler, UNet2DConditionModel
from diffusers.utils import check_min_version
from diffusers.utils.import_utils import is_xformers_available
from transformers import CLIPTextModel, CLIPTokenizer
from util import filter_segmentation_mask
from model.text_segmenter.unet import UNet

disable_caching()
check_min_version("0.15.0.dev0")
PLACE_HOLDER = '*'


def parse_args():
 parser = argparse.ArgumentParser(description="Simple example of a training script.")
 parser.add_argument(
 "--pretrained_model_name_or_path",
 type=str,
 default='runwayml/stable-diffusion-v1-5', # no need to modify this
 help="Path to pretrained model or model identifier from huggingface.co/models. Please do not modify this.",
 )
 parser.add_argument(
 "--revision",
 type=str,
 default=None,
 required=False,
 help="Revision of pretrained model identifier from huggingface.co/models.",
 )
 parser.add_argument(
 "--mode",
 type=str,
 default="text-to-image-with-template",
 choices=["text-to-image-with-template"],
 help="Three modes can be used.",
 )
 parser.add_argument(
 "--output_dir",
 type=str,
 default="./textdiffuser_laion_generated/",
 help="output path",
 )
 parser.add_argument(
 "--seed",
 type=int,
 default=100,
 help="A seed for reproducible training."
 )
 parser.add_argument(
 "--classifier_free_scale",
 type=float,
 default=9.0, # following stable diffusion (https://github.com/CompVis/stable-diffusion)
 help="Classifier free scale following https://arxiv.org/abs/2207.12598.",
 )
 parser.add_argument(
 "--model_path",
 type=str,
 default='textdiffuser-ckpt/diffusion_backbone',
 help='path of model'
 )
 parser.add_argument(
 "--character_segmenter_path",
 type=str,
 default='textdiffuser-ckpt/text_segmenter.pth',
 help="checkpoint of character-level segmenter"
 )
 parser.add_argument(
 "--enable_xformers_memory_efficient_attention",
 action="store_true",
 default=True,
 help="Whether or not to use xformers."
 )
 parser.add_argument(
 "--sample_steps",
 type=int,
 default=20,
 help="Diffusion steps for sampling."
 )
 parser.add_argument(
 "--vis_num",
 type=int,
 default=4,
 help="Number of images to be sample. Please decrease it when encountering out of memory error."
 )
 parser.add_argument(
 "--a_prompt",
 type=str,
 default='best quality, extremely detailed',
 help="additional prompt"
 )
 parser.add_argument(
 "--n_prompt",
 type=str,
 default='longbody, lowres, bad anatomy, bad hands, missing fingers, extra digit, fewer digits, cropped, worst quality, low quality, watermark',
 help="negative prompt"
 )
 parser.add_argument(
 "--glyph_dir",
 type=str,
 default='/data/vdb/yuxiang.tyx/AIGC/data/laion_word/glyph_laion',
 help="path of glyph images from anytext evaluation dataset"
 )
 parser.add_argument(
 "--json_path",
 type=str,
 default='/data/vdb/yuxiang.tyx/AIGC/data/laion_word/test1k.json',
 help="json path for evaluation dataset"
 )

 args = parser.parse_args()
 print(f'{colored("[√]", "green")} Arguments are loaded.')
 return args


def load_json(file_path: str):
 with open(file_path, 'r', encoding='utf8') as f:
 content = json.load(f)
 return content


def load_data(input_path):
 content = load_json(input_path)
 d = []
 count = 0
 for gt in content['data_list']:
 info = {}
 info['img_name'] = gt['img_name']
 info['caption'] = gt['caption']
 if PLACE_HOLDER in info['caption']:
 count += 1
 info['caption'] = info['caption'].replace(PLACE_HOLDER, " ")
 if 'annotations' in gt:
 polygons = []
 texts = []
 pos = []
 for annotation in gt['annotations']:
 if len(annotation['polygon']) == 0:
 continue
 if annotation['valid'] is False:
 continue
 polygons.append(annotation['polygon'])
 texts.append(annotation['text'])
 pos.append(annotation['pos'])
 info['polygons'] = [np.array(i) for i in polygons]
 info['texts'] = texts
 info['pos'] = pos
 d.append(info)
 print(f'{input_path} loaded, imgs={len(d)}')
 if count > 0:
 print(f"Found {count} image's caption contain placeholder: {PLACE_HOLDER}, change to ' '...")
 return d


def get_item(data_list, item):
 item_dict = {}
 cur_item = data_list[item]
 item_dict['img_name'] = cur_item['img_name']
 item_dict['caption'] = cur_item['caption']
 return item_dict


def main():
 args = parse_args()
 seed = args.seed if args.seed is not None else random.randint(0, 1000000)
 set_seed(seed)
 if os.path.exists(args.output_dir) is not True:
 os.makedirs(args.output_dir)
 # Load scheduler, tokenizer and models.
 tokenizer = CLIPTokenizer.from_pretrained(
 args.pretrained_model_name_or_path, subfolder="tokenizer", revision=args.revision
 )
 text_encoder = CLIPTextModel.from_pretrained(
 args.pretrained_model_name_or_path, subfolder="text_encoder", revision=args.revision
 ).cuda()
 vae = AutoencoderKL.from_pretrained(args.pretrained_model_name_or_path, subfolder="vae", revision=args.revision).cuda()
 unet = UNet2DConditionModel.from_pretrained(
 args.model_path, subfolder="unet", revision=None
 ).cuda()
 # load character-level segmenter
 segmenter = UNet(3, 96, True).cuda()
 segmenter = torch.nn.DataParallel(segmenter)
 segmenter.load_state_dict(torch.load(args.character_segmenter_path))
 segmenter.eval()
 # Freeze vae and text_encoder
 vae.requires_grad_(False)
 text_encoder.requires_grad_(False)

 if args.enable_xformers_memory_efficient_attention:
 if is_xformers_available():
 import xformers
 xformers_version = version.parse(xformers.__version__)
 if xformers_version == version.parse("0.0.16"):
 print(
 "xFormers 0.0.16 cannot be used for training in some GPUs. If you observe problems during training, please update xFormers to at least 0.0.17. See https://huggingface.co/docs/diffusers/main/en/optimization/xformers for more details."
 )
 unet.enable_xformers_memory_efficient_attention()
 else:
 raise ValueError("xformers is not available. Make sure it is installed correctly")
 # setup schedulers
 scheduler = DDPMScheduler.from_pretrained(args.pretrained_model_name_or_path, subfolder="scheduler")
 scheduler.set_timesteps(args.sample_steps)
 sample_num = args.vis_num
 # inference loop
 data_list = load_data(args.json_path)
 for i in tqdm(range(len(data_list)), desc='generator'):
 item_dict = get_item(data_list, i)
 p = item_dict['img_name']
 img_name = item_dict['img_name'].split('.')[0] + '_3.jpg'
 if os.path.exists(os.path.join(args.output_dir, img_name)):
 continue
 input_image_path = os.path.join(args.glyph_dir, p)
 prompt = item_dict['caption']
 template_image = Image.open(input_image_path).resize((256, 256)).convert('RGB')

 noise = torch.randn((sample_num, 4, 64, 64)).to("cuda") # (b, 4, 64, 64)
 input = noise # (b, 4, 64, 64)
 captions = [prompt + ', ' + args.a_prompt] * sample_num
 captions_nocond = [args.n_prompt] * sample_num
 # encode text prompts
 inputs = tokenizer(
 captions, max_length=tokenizer.model_max_length, padding="max_length", truncation=True, return_tensors="pt"
 ).input_ids.cuda() # (b, 77)
 encoder_hidden_states = text_encoder(inputs)[0].cuda() # (b, 77, 768)

 inputs_nocond = tokenizer(
 captions_nocond, max_length=tokenizer.model_max_length, padding="max_length", truncation=True, return_tensors="pt"
 ).input_ids.cuda() # (b, 77)
 encoder_hidden_states_nocond = text_encoder(inputs_nocond)[0].cuda() # (b, 77, 768)
 if args.mode == 'text-to-image-with-template':
 to_tensor = transforms.ToTensor()
 image_tensor = to_tensor(template_image).unsqueeze(0).cuda().sub_(0.5).div_(0.5) # (b, 3, 256, 256)
 with torch.no_grad():
 segmentation_mask = segmenter(image_tensor) # (b, 96, 256, 256)
 segmentation_mask = segmentation_mask.max(1)[1].squeeze(0) # (256, 256)
 segmentation_mask = filter_segmentation_mask(segmentation_mask) # (256, 256)
 segmentation_mask = torch.nn.functional.interpolate(segmentation_mask.unsqueeze(0).unsqueeze(0).float(), size=(256, 256), mode='nearest') # (b, 1, 256, 256)
 segmentation_mask = segmentation_mask.squeeze(1).repeat(sample_num, 1, 1).long().to('cuda') # (b, 1, 256, 256)

 feature_mask = torch.ones(sample_num, 1, 64, 64).to('cuda') # (b, 1, 64, 64)
 masked_image = torch.zeros(sample_num, 3, 512, 512).to('cuda') # (b, 3, 512, 512)
 masked_feature = vae.encode(masked_image).latent_dist.sample() # (b, 4, 64, 64)
 masked_feature = masked_feature * vae.config.scaling_factor # (b, 4, 64, 64)

 # diffusion process
 intermediate_images = []
 for t in tqdm(scheduler.timesteps):
 with torch.no_grad():
 noise_pred_cond = unet(sample=input, timestep=t, encoder_hidden_states=encoder_hidden_states, segmentation_mask=segmentation_mask, feature_mask=feature_mask, masked_feature=masked_feature).sample # b, 4, 64, 64
 noise_pred_uncond = unet(sample=input, timestep=t, encoder_hidden_states=encoder_hidden_states_nocond, segmentation_mask=segmentation_mask, feature_mask=feature_mask, masked_feature=masked_feature).sample # b, 4, 64, 64
 noisy_residual = noise_pred_uncond + args.classifier_free_scale * (noise_pred_cond - noise_pred_uncond) # b, 4, 64, 64
 prev_noisy_sample = scheduler.step(noisy_residual, t, input).prev_sample
 input = prev_noisy_sample
 intermediate_images.append(prev_noisy_sample)

 # decode and visualization
 input = 1 / vae.config.scaling_factor * input
 sample_images = vae.decode(input.float(), return_dict=False)[0] # (b, 3, 512, 512)

 # save pred_img
 pred_image_list = []
 for image in sample_images.float():
 image = (image / 2 + 0.5).clamp(0, 1).unsqueeze(0)
 image = image.cpu().permute(0, 2, 3, 1).numpy()[0]
 image = Image.fromarray((image * 255).round().astype("uint8")).convert('RGB')
 pred_image_list.append(image)

 for idx, img in enumerate(pred_image_list):
 img_name = item_dict['img_name'].split('.')[0]+f'_{idx}' + '.jpg'
 img.save(os.path.join(args.output_dir, img_name))


if __name__ == "__main__":
 main()