





import argparse
import glob
import logging
import os
import sys
import time
from collections import defaultdict
from pathlib import Path
from typing import Dict, List, Tuple

import numpy as np
import pytorch_lightning as pl
import torch
from callbacks import SeqaSeqLoggingCallback, get_checkpoint_callback, get_early_stopping_callback
from torch import nn
from torch.utils.data import DataLoader

from transformers import MBartTokenizer, TaForConditionalGeneration
from transformers.models.bart.modeling_bart import shift_tokens_right
from utils import (
    ROUGE_KEYS,
    LegacySeqaSeqDataset,
    SeqaSeqDataset,
    assert_all_frozen,
    calculate_bleu,
    calculate_rouge,
    check_output_dir,
    flatten_list,
    freeze_embeds,
    freeze_params,
    get_git_info,
    label_smoothed_nll_loss,
    lmap,
    pickle_save,
    save_git_info,
    save_json,
    use_task_specific_params,
)


# need the parent dir module
sys.path.insert(2, str(Path(__file__).resolve().parents[1]))
from lightning_base import BaseTransformer, add_generic_args, generic_train  # noqa


UpperCAmelCase__						:			List[str]									=logging.getLogger(__name__)

class       __A     (						a    ):
				__A      =				"""summarization"""
				__A      =				["""loss"""]
				__A      =				ROUGE_KEYS
				__A      =				"""rouge2"""


				def __init__(		self   , UpperCAmelCase_   , **UpperCAmelCase_							):
								if hparams.sortish_sampler and hparams.gpus > 1:
												lowerCamelCase	 =False
								elif hparams.max_tokens_per_batch is not None:
												if hparams.gpus > 1:
																raise NotImplementedError("""Dynamic Batch size does not work for multi-gpu training"""							)
												if hparams.sortish_sampler:
																raise ValueError("""--sortish_sampler and --max_tokens_per_batch may not be used simultaneously"""							)

								super().__init__(UpperCAmelCase_   , num_labels=UpperCAmelCase_   , mode=self.mode   , **UpperCAmelCase_							)
								use_task_specific_params(self.model   , """summarization"""							)
								save_git_info(self.hparams.output_dir							)
								lowerCamelCase	 =Path(self.output_dir							) / """metrics.json"""
								lowerCamelCase	 =Path(self.output_dir							) / """hparams.pkl"""
								pickle_save(self.hparams   , self.hparams_save_path							)
								lowerCamelCase	 =0
								lowerCamelCase	 =defaultdict(UpperCAmelCase_							)
								lowerCamelCase	 =self.config.model_type
								lowerCamelCase	 =self.config.tgt_vocab_size if self.model_type == """fsmt""" else self.config.vocab_size

								lowerCamelCase	 ={
								    "data_dir": self.hparams.data_dir,
								    "max_source_length": self.hparams.max_source_length,
								    "prefix": self.model.config.prefix or "",
								}
								lowerCamelCase	 ={
								    """train""": self.hparams.n_train,
								    """val""": self.hparams.n_val,
								    """test""": self.hparams.n_test,
								}
								lowerCamelCase	 ={k: v if v >= 0 else None for k, v in n_observations_per_split.items()}

								lowerCamelCase	 ={
								    """train""": self.hparams.max_target_length,
								    """val""": self.hparams.val_max_target_length,
								    """test""": self.hparams.test_max_target_length,
								}
								assert self.target_lens["train"] <= self.target_lens["val"], f"""target_lens: {self.target_lens}"""
								assert self.target_lens["train"] <= self.target_lens["test"], f"""target_lens: {self.target_lens}"""
								if self.hparams.freeze_embeds:
												freeze_embeds(self.model							)
								if self.hparams.freeze_encoder:
												freeze_params(self.model.get_encoder()							)
												assert_all_frozen(self.model.get_encoder()							)

								lowerCamelCase	 =get_git_info()["""repo_sha"""]
								lowerCamelCase	 =hparams.num_workers
								lowerCamelCase	 =None  # default to config
								if self.model.config.decoder_start_token_id is None and isinstance(self.tokenizer   , UpperCAmelCase_							):
												lowerCamelCase	 =self.tokenizer.lang_code_to_id[hparams.tgt_lang]
												lowerCamelCase	 =self.decoder_start_token_id
								lowerCamelCase	 =(
								    SeqaSeqDataset if hasattr(self.tokenizer   , """prepare_seq2seq_batch"""							) else LegacySeqaSeqDataset
								)
								lowerCamelCase	 =False
								lowerCamelCase	 =self.model.config.num_beams if self.hparams.eval_beams is None else self.hparams.eval_beams
								if self.hparams.eval_max_gen_length is not None:
												lowerCamelCase	 =self.hparams.eval_max_gen_length
								else:
												lowerCamelCase	 =self.model.config.max_length
								lowerCamelCase	 =self.default_val_metric if self.hparams.val_metric is None else self.hparams.val_metric


				def _snake_case       (		self   , UpperCAmelCase_							):
								lowerCamelCase	 ={
								    k: self.tokenizer.batch_decode(v.tolist()							) if """mask""" not in k else v.shape for k, v in batch.items()
								}
								save_json(UpperCAmelCase_   , Path(self.output_dir							) / """text_batch.json"""							)
								save_json({k: v.tolist() for k, v in batch.items()}   , Path(self.output_dir							) / """tok_batch.json"""							)

								lowerCamelCase	 =True
								return readable_batch


				def _snake_case       (		self   , UpperCAmelCase_   , **UpperCAmelCase_							):
								return self.model(UpperCAmelCase_   , **UpperCAmelCase_							)


				def _snake_case       (		self   , UpperCAmelCase_							):
								lowerCamelCase	 =self.tokenizer.batch_decode(
								    UpperCAmelCase_   , skip_special_tokens=UpperCAmelCase_   , clean_up_tokenization_spaces=UpperCAmelCase_							)
								return lmap(str.strip   , UpperCAmelCase_							)


				def _snake_case       (		self   , UpperCAmelCase_							):
								lowerCamelCase	 =self.tokenizer.pad_token_id
								lowerCamelCase ,  lowerCamelCase	 =batch["""input_ids"""], batch["""attention_mask"""]
								lowerCamelCase	 =batch["""labels"""]
								if isinstance(self.model   , UpperCAmelCase_							):
												lowerCamelCase	 =self.model._shift_right(UpperCAmelCase_							)
								else:
												lowerCamelCase	 =shift_tokens_right(UpperCAmelCase_   , UpperCAmelCase_							)
								if not self.already_saved_batch:  # This would be slightly better if it only happened on rank zero
												lowerCamelCase	 =decoder_input_ids
												self.save_readable_batch(UpperCAmelCase_							)

								lowerCamelCase	 =self(UpperCAmelCase_   , attention_mask=UpperCAmelCase_   , decoder_input_ids=UpperCAmelCase_   , use_cache=UpperCAmelCase_							)
								lowerCamelCase	 =outputs["""logits"""]
								if self.hparams.label_smoothing == 0:
												# Same behavior as modeling_bart.py, besides ignoring pad_token_id
												lowerCamelCase	 =nn.CrossEntropyLoss(ignore_index=UpperCAmelCase_							)

												assert lm_logits.shape[-1] == self.vocab_size
												lowerCamelCase	 =ce_loss_fct(lm_logits.view(-1   , lm_logits.shape[-1]							)   , tgt_ids.view(-1							)							)
								else:
												lowerCamelCase	 =nn.functional.log_softmax(UpperCAmelCase_   , dim=-1							)
												lowerCamelCase ,  lowerCamelCase	 =label_smoothed_nll_loss(
												    UpperCAmelCase_   , UpperCAmelCase_   , self.hparams.label_smoothing   , ignore_index=UpperCAmelCase_							)
								return (loss,)


				@property
				def _snake_case       (		self							):
								return self.tokenizer.pad_token_id


				def _snake_case       (		self   , UpperCAmelCase_   , UpperCAmelCase_							):
								lowerCamelCase	 =self._step(UpperCAmelCase_							)

								lowerCamelCase	 =dict(zip(self.loss_names   , UpperCAmelCase_							)							)
								# tokens per batch
								lowerCamelCase	 =batch["""input_ids"""].ne(self.pad							).sum() + batch["""labels"""].ne(self.pad							).sum()
								lowerCamelCase	 =batch["""input_ids"""].shape[0]
								lowerCamelCase	 =batch["""input_ids"""].eq(self.pad							).sum()
								lowerCamelCase	 =batch["""input_ids"""].eq(self.pad							).float().mean()
								# TODO(SS): make a wandb summary metric for this
								return {"loss": loss_tensors[0], "log": logs}


				def _snake_case       (		self   , UpperCAmelCase_   , UpperCAmelCase_							):
								return self._generative_step(UpperCAmelCase_							)


				def _snake_case       (		self   , UpperCAmelCase_   , UpperCAmelCase_="val"							):
								self.step_count += 1
								lowerCamelCase	 ={k: torch.stack([x[k] for x in outputs]							).mean() for k in self.loss_names}
								lowerCamelCase	 =losses["""loss"""]
								lowerCamelCase	 ={
								    k: np.array([x[k] for x in outputs]							).mean() for k in self.metric_names + ["""gen_time""", """gen_len"""]
								}
								lowerCamelCase	 =(
								    generative_metrics[self.val_metric] if self.val_metric in generative_metrics else losses[self.val_metric]
								)
								lowerCamelCase	 =torch.tensor(UpperCAmelCase_							).type_as(UpperCAmelCase_							)
								generative_metrics.update({k: v.item() for k, v in losses.items()}							)
								losses.update(UpperCAmelCase_							)
								lowerCamelCase	 ={f"""{prefix}_avg_{k}""": x for k, x in losses.items()}
								lowerCamelCase	 =self.step_count
								self.metrics[prefix].append(UpperCAmelCase_							)  # callback writes this to self.metrics_save_path
								lowerCamelCase	 =flatten_list([x["""preds"""] for x in outputs]							)
								return {
								    "log": all_metrics,
								    "preds": preds,
								    f"""{prefix}_loss""": loss,
								    f"""{prefix}_{self.val_metric}""": metric_tensor,
								}


				def _snake_case       (		self   , UpperCAmelCase_   , UpperCAmelCase_							):
								return calculate_rouge(UpperCAmelCase_   , UpperCAmelCase_							)


				def _snake_case       (		self   , UpperCAmelCase_							):
								lowerCamelCase	 =time.time()

								# parser.add_argument('--eval_max_gen_length', type=int, default=None, help='never generate more than n tokens')
								lowerCamelCase	 =self.model.generate(
								    batch["""input_ids"""]   , attention_mask=batch["""attention_mask"""]   , use_cache=UpperCAmelCase_   , decoder_start_token_id=self.decoder_start_token_id   , num_beams=self.eval_beams   , max_length=self.eval_max_length   , )
								lowerCamelCase	 =(time.time() - ta) / batch["""input_ids"""].shape[0]
								lowerCamelCase	 =self.ids_to_clean_text(UpperCAmelCase_							)
								lowerCamelCase	 =self.ids_to_clean_text(batch["""labels"""]							)
								lowerCamelCase	 =self._step(UpperCAmelCase_							)
								lowerCamelCase	 =dict(zip(self.loss_names   , UpperCAmelCase_							)							)
								lowerCamelCase	 =self.calc_generative_metrics(UpperCAmelCase_   , UpperCAmelCase_							)
								lowerCamelCase	 =np.mean(lmap(UpperCAmelCase_   , UpperCAmelCase_							)							)
								base_metrics.update(gen_time=UpperCAmelCase_   , gen_len=UpperCAmelCase_   , preds=UpperCAmelCase_   , target=UpperCAmelCase_   , **UpperCAmelCase_							)
								return base_metrics


				def _snake_case       (		self   , UpperCAmelCase_   , UpperCAmelCase_							):
								return self._generative_step(UpperCAmelCase_							)


				def _snake_case       (		self   , UpperCAmelCase_							):
								return self.validation_epoch_end(UpperCAmelCase_   , prefix="""test"""							)


				def _snake_case       (		self   , UpperCAmelCase_							):
								lowerCamelCase	 =self.n_obs[type_path]
								lowerCamelCase	 =self.target_lens[type_path]
								lowerCamelCase	 =self.dataset_class(
								    self.tokenizer   , type_path=UpperCAmelCase_   , n_obs=UpperCAmelCase_   , max_target_length=UpperCAmelCase_   , **self.dataset_kwargs   , )
								return dataset


				def _snake_case       (		self   , UpperCAmelCase_   , UpperCAmelCase_   , UpperCAmelCase_ = False							):
								lowerCamelCase	 =self.get_dataset(UpperCAmelCase_							)

								if self.hparams.sortish_sampler and type_path != "test" and type_path != "val":
												lowerCamelCase	 =dataset.make_sortish_sampler(UpperCAmelCase_   , distributed=self.hparams.gpus > 1							)
												return DataLoader(
												    UpperCAmelCase_   , batch_size=UpperCAmelCase_   , collate_fn=dataset.collate_fn   , shuffle=UpperCAmelCase_   , num_workers=self.num_workers   , sampler=UpperCAmelCase_   , )

								elif self.hparams.max_tokens_per_batch is not None and type_path != "test" and type_path != "val":
												lowerCamelCase	 =dataset.make_dynamic_sampler(
												    self.hparams.max_tokens_per_batch   , distributed=self.hparams.gpus > 1							)
												return DataLoader(
												    UpperCAmelCase_   , batch_sampler=UpperCAmelCase_   , collate_fn=dataset.collate_fn   , num_workers=self.num_workers   , )
								else:
												return DataLoader(
												    UpperCAmelCase_   , batch_size=UpperCAmelCase_   , collate_fn=dataset.collate_fn   , shuffle=UpperCAmelCase_   , num_workers=self.num_workers   , sampler=UpperCAmelCase_   , )


				def _snake_case       (		self							):
								lowerCamelCase	 =self.get_dataloader("""train"""   , batch_size=self.hparams.train_batch_size   , shuffle=UpperCAmelCase_							)
								return dataloader


				def _snake_case       (		self							):
								return self.get_dataloader("""val"""   , batch_size=self.hparams.eval_batch_size							)


				def _snake_case       (		self							):
								return self.get_dataloader("""test"""   , batch_size=self.hparams.eval_batch_size							)


				@staticmethod
				def _snake_case       (		UpperCAmelCase_   , UpperCAmelCase_							):
								BaseTransformer.add_model_specific_args(UpperCAmelCase_   , UpperCAmelCase_							)
								add_generic_args(UpperCAmelCase_   , UpperCAmelCase_							)
								parser.add_argument(
								    """--max_source_length"""   , default=1024   , type=UpperCAmelCase_   , help=(
								        """The maximum total input sequence length after tokenization. Sequences longer """
								        """than this will be truncated, sequences shorter will be padded."""
								    )   , )
								parser.add_argument(
								    """--max_target_length"""   , default=56   , type=UpperCAmelCase_   , help=(
								        """The maximum total input sequence length after tokenization. Sequences longer """
								        """than this will be truncated, sequences shorter will be padded."""
								    )   , )
								parser.add_argument(
								    """--val_max_target_length"""   , default=142   , type=UpperCAmelCase_   , help=(
								        """The maximum total input sequence length after tokenization. Sequences longer """
								        """than this will be truncated, sequences shorter will be padded."""
								    )   , )
								parser.add_argument(
								    """--test_max_target_length"""   , default=142   , type=UpperCAmelCase_   , help=(
								        """The maximum total input sequence length after tokenization. Sequences longer """
								        """than this will be truncated, sequences shorter will be padded."""
								    )   , )
								parser.add_argument("""--freeze_encoder"""   , action="""store_true"""							)
								parser.add_argument("""--freeze_embeds"""   , action="""store_true"""							)
								parser.add_argument("""--sortish_sampler"""   , action="""store_true"""   , default=UpperCAmelCase_							)
								parser.add_argument("""--overwrite_output_dir"""   , action="""store_true"""   , default=UpperCAmelCase_							)
								parser.add_argument("""--max_tokens_per_batch"""   , type=UpperCAmelCase_   , default=UpperCAmelCase_							)
								parser.add_argument("""--logger_name"""   , type=UpperCAmelCase_   , choices=["""default""", """wandb""", """wandb_shared"""]   , default="""default"""							)
								parser.add_argument("""--n_train"""   , type=UpperCAmelCase_   , default=-1   , required=UpperCAmelCase_   , help="""# examples. -1 means use all."""							)
								parser.add_argument("""--n_val"""   , type=UpperCAmelCase_   , default=500   , required=UpperCAmelCase_   , help="""# examples. -1 means use all."""							)
								parser.add_argument("""--n_test"""   , type=UpperCAmelCase_   , default=-1   , required=UpperCAmelCase_   , help="""# examples. -1 means use all."""							)
								parser.add_argument(
								    """--task"""   , type=UpperCAmelCase_   , default="""summarization"""   , required=UpperCAmelCase_   , help="""# examples. -1 means use all."""							)
								parser.add_argument("""--label_smoothing"""   , type=UpperCAmelCase_   , default=0.0   , required=UpperCAmelCase_							)
								parser.add_argument("""--src_lang"""   , type=UpperCAmelCase_   , default=""""""   , required=UpperCAmelCase_							)
								parser.add_argument("""--tgt_lang"""   , type=UpperCAmelCase_   , default=""""""   , required=UpperCAmelCase_							)
								parser.add_argument("""--eval_beams"""   , type=UpperCAmelCase_   , default=UpperCAmelCase_   , required=UpperCAmelCase_							)
								parser.add_argument(
								    """--val_metric"""   , type=UpperCAmelCase_   , default=UpperCAmelCase_   , required=UpperCAmelCase_   , choices=["""bleu""", """rouge2""", """loss""", None]							)
								parser.add_argument("""--eval_max_gen_length"""   , type=UpperCAmelCase_   , default=UpperCAmelCase_   , help="""never generate more than n tokens"""							)
								parser.add_argument("""--save_top_k"""   , type=UpperCAmelCase_   , default=1   , required=UpperCAmelCase_   , help="""How many checkpoints to save"""							)
								parser.add_argument(
								    """--early_stopping_patience"""   , type=UpperCAmelCase_   , default=-1   , required=UpperCAmelCase_   , help=(
								        """-1 means never early stop. early_stopping_patience is measured in validation checks, not epochs. So"""
								        """ val_check_interval will effect it."""
								    )   , )
								return parser




class       __A     (						a    ):
				__A      =				"""translation"""
				__A      =				["""loss"""]
				__A      =				["""bleu"""]
				__A      =				"""bleu"""


				def __init__(		self   , UpperCAmelCase_   , **UpperCAmelCase_							):
								super().__init__(UpperCAmelCase_   , **UpperCAmelCase_							)
								lowerCamelCase	 =hparams.src_lang
								lowerCamelCase	 =hparams.tgt_lang


				def _snake_case       (		self   , UpperCAmelCase_   , UpperCAmelCase_							):
								return calculate_bleu(UpperCAmelCase_   , UpperCAmelCase_							)





def      _lowercase					(						_UpperCAmelCase			,      _UpperCAmelCase=None				)						->  SummarizationModule:
				Path(args.output_dir				).mkdir(exist_ok=_UpperCAmelCase				)
				check_output_dir(_UpperCAmelCase			,      expected_items=3				)

				if model is None:
								if "summarization" in args.task:
												lowerCamelCase	 =SummarizationModule(_UpperCAmelCase				)
								else:
												lowerCamelCase	 =TranslationModule(_UpperCAmelCase				)
				lowerCamelCase	 =Path(args.data_dir				).name
				if (
				    args.logger_name == "default"
				    or args.fast_dev_run
				    or str(args.output_dir				).startswith("""/tmp"""				)
				    or str(args.output_dir				).startswith("""/var"""				)
				):
								lowerCamelCase	 =True  # don't pollute wandb logs unnecessarily
				elif args.logger_name == "wandb":
								from pytorch_lightning.loggers import WandbLogger

								lowerCamelCase	 =os.environ.get("""WANDB_PROJECT"""			,      _UpperCAmelCase				)
								lowerCamelCase	 =WandbLogger(name=model.output_dir.name			,      project=_UpperCAmelCase				)

				elif args.logger_name == "wandb_shared":
								from pytorch_lightning.loggers import WandbLogger

								lowerCamelCase	 =WandbLogger(name=model.output_dir.name			,      project=F"""hf_{dataset}"""				)

				if args.early_stopping_patience >= 0:
								lowerCamelCase	 =get_early_stopping_callback(model.val_metric			,      args.early_stopping_patience				)
				else:
								lowerCamelCase	 =False

				lowerCamelCase	 =args.val_metric == """loss"""
				lowerCamelCase	 =generic_train(
				    _UpperCAmelCase			,      _UpperCAmelCase			,      logging_callback=SeqaSeqLoggingCallback()			,      checkpoint_callback=get_checkpoint_callback(
				        args.output_dir			,      model.val_metric			,      args.save_top_k			,      _UpperCAmelCase				)			,      early_stopping_callback=_UpperCAmelCase			,      logger=_UpperCAmelCase			,      )
				pickle_save(model.hparams			,      model.output_dir / """hparams.pkl"""				)
				if not args.do_predict:
								return model

				lowerCamelCase	 =""""""
				lowerCamelCase	 =sorted(glob.glob(os.path.join(args.output_dir			,      """*.ckpt"""				)			,      recursive=_UpperCAmelCase				)				)
				if checkpoints:
								lowerCamelCase	 =checkpoints[-1]
								lowerCamelCase	 =checkpoints[-1]
				trainer.logger.log_hyperparams(model.hparams				)

				# test() without a model tests using the best checkpoint automatically
				trainer.test()
				return model


if __name__ == "__main__":
							UpperCAmelCase__						:			str									=argparse.ArgumentParser()
							UpperCAmelCase__						:			Any									=pl.Trainer.add_argparse_args(parser)
							UpperCAmelCase__						:			Any									=SummarizationModule.add_model_specific_args(parser, os.getcwd())

							UpperCAmelCase__						:			List[str]									=parser.parse_args()

							main(args)












from dataclasses import dataclass
from typing import Optional

import numpy as np
import torch
import torch.nn as nn

from ..utils import BaseOutput, is_torch_version, randn_tensor
from .attention_processor import SpatialNorm
from .unet_ad_blocks import UNetMidBlockaD, get_down_block, get_up_block

@dataclass
class       __A     (						a    ):
				__A      =				42




class       __A     (						nn.Module    ):


				def __init__(		self   , UpperCAmelCase_=3   , UpperCAmelCase_=3   , UpperCAmelCase_=("DownEncoderBlock2D",)   , UpperCAmelCase_=(64,)   , UpperCAmelCase_=2   , UpperCAmelCase_=32   , UpperCAmelCase_="silu"   , UpperCAmelCase_=True   , ):
								super().__init__()
								lowerCamelCase	 =layers_per_block

								lowerCamelCase	 =torch.nn.Convad(
								    UpperCAmelCase_   , block_out_channels[0]   , kernel_size=3   , stride=1   , padding=1   , )

								lowerCamelCase	 =None
								lowerCamelCase	 =nn.ModuleList([]							)

								# down
								lowerCamelCase	 =block_out_channels[0]
								for i, down_block_type in enumerate(UpperCAmelCase_							):
												lowerCamelCase	 =output_channel
												lowerCamelCase	 =block_out_channels[i]
												lowerCamelCase	 =i == len(UpperCAmelCase_							) - 1

												lowerCamelCase	 =get_down_block(
												    UpperCAmelCase_   , num_layers=self.layers_per_block   , in_channels=UpperCAmelCase_   , out_channels=UpperCAmelCase_   , add_downsample=not is_final_block   , resnet_eps=1E-6   , downsample_padding=0   , resnet_act_fn=UpperCAmelCase_   , resnet_groups=UpperCAmelCase_   , attention_head_dim=UpperCAmelCase_   , temb_channels=UpperCAmelCase_   , )
												self.down_blocks.append(UpperCAmelCase_							)

								# mid
								lowerCamelCase	 =UNetMidBlockaD(
								    in_channels=block_out_channels[-1]   , resnet_eps=1E-6   , resnet_act_fn=UpperCAmelCase_   , output_scale_factor=1   , resnet_time_scale_shift="""default"""   , attention_head_dim=block_out_channels[-1]   , resnet_groups=UpperCAmelCase_   , temb_channels=UpperCAmelCase_   , )

								# out
								lowerCamelCase	 =nn.GroupNorm(num_channels=block_out_channels[-1]   , num_groups=UpperCAmelCase_   , eps=1E-6							)
								lowerCamelCase	 =nn.SiLU()

								lowerCamelCase	 =2 * out_channels if double_z else out_channels
								lowerCamelCase	 =nn.Convad(block_out_channels[-1]   , UpperCAmelCase_   , 3   , padding=1							)

								lowerCamelCase	 =False


				def _snake_case       (		self   , UpperCAmelCase_							):
								lowerCamelCase	 =x
								lowerCamelCase	 =self.conv_in(UpperCAmelCase_							)

								if self.training and self.gradient_checkpointing:

												def create_custom_forward(UpperCAmelCase_							):
																def custom_forward(*UpperCAmelCase_							):
																				return module(*UpperCAmelCase_							)

																return custom_forward

												# down
												if is_torch_version(""">="""   , """1.11.0"""							):
																for down_block in self.down_blocks:
																				lowerCamelCase	 =torch.utils.checkpoint.checkpoint(
																				    create_custom_forward(UpperCAmelCase_							)   , UpperCAmelCase_   , use_reentrant=UpperCAmelCase_							)
																# middle
																lowerCamelCase	 =torch.utils.checkpoint.checkpoint(
																    create_custom_forward(self.mid_block							)   , UpperCAmelCase_   , use_reentrant=UpperCAmelCase_							)
												else:
																for down_block in self.down_blocks:
																				lowerCamelCase	 =torch.utils.checkpoint.checkpoint(create_custom_forward(UpperCAmelCase_							)   , UpperCAmelCase_							)
																# middle
																lowerCamelCase	 =torch.utils.checkpoint.checkpoint(create_custom_forward(self.mid_block							)   , UpperCAmelCase_							)

								else:
												# down
												for down_block in self.down_blocks:
																lowerCamelCase	 =down_block(UpperCAmelCase_							)

												# middle
												lowerCamelCase	 =self.mid_block(UpperCAmelCase_							)

								# post-process
								lowerCamelCase	 =self.conv_norm_out(UpperCAmelCase_							)
								lowerCamelCase	 =self.conv_act(UpperCAmelCase_							)
								lowerCamelCase	 =self.conv_out(UpperCAmelCase_							)

								return sample




class       __A     (						nn.Module    ):


				def __init__(		self   , UpperCAmelCase_=3   , UpperCAmelCase_=3   , UpperCAmelCase_=("UpDecoderBlock2D",)   , UpperCAmelCase_=(64,)   , UpperCAmelCase_=2   , UpperCAmelCase_=32   , UpperCAmelCase_="silu"   , UpperCAmelCase_="group"   , ):
								super().__init__()
								lowerCamelCase	 =layers_per_block

								lowerCamelCase	 =nn.Convad(
								    UpperCAmelCase_   , block_out_channels[-1]   , kernel_size=3   , stride=1   , padding=1   , )

								lowerCamelCase	 =None
								lowerCamelCase	 =nn.ModuleList([]							)

								lowerCamelCase	 =in_channels if norm_type == """spatial""" else None

								# mid
								lowerCamelCase	 =UNetMidBlockaD(
								    in_channels=block_out_channels[-1]   , resnet_eps=1E-6   , resnet_act_fn=UpperCAmelCase_   , output_scale_factor=1   , resnet_time_scale_shift="""default""" if norm_type == """group""" else norm_type   , attention_head_dim=block_out_channels[-1]   , resnet_groups=UpperCAmelCase_   , temb_channels=UpperCAmelCase_   , )

								# up
								lowerCamelCase	 =list(reversed(UpperCAmelCase_							)							)
								lowerCamelCase	 =reversed_block_out_channels[0]
								for i, up_block_type in enumerate(UpperCAmelCase_							):
												lowerCamelCase	 =output_channel
												lowerCamelCase	 =reversed_block_out_channels[i]

												lowerCamelCase	 =i == len(UpperCAmelCase_							) - 1

												lowerCamelCase	 =get_up_block(
												    UpperCAmelCase_   , num_layers=self.layers_per_block + 1   , in_channels=UpperCAmelCase_   , out_channels=UpperCAmelCase_   , prev_output_channel=UpperCAmelCase_   , add_upsample=not is_final_block   , resnet_eps=1E-6   , resnet_act_fn=UpperCAmelCase_   , resnet_groups=UpperCAmelCase_   , attention_head_dim=UpperCAmelCase_   , temb_channels=UpperCAmelCase_   , resnet_time_scale_shift=UpperCAmelCase_   , )
												self.up_blocks.append(UpperCAmelCase_							)
												lowerCamelCase	 =output_channel

								# out
								if norm_type == "spatial":
												lowerCamelCase	 =SpatialNorm(block_out_channels[0]   , UpperCAmelCase_							)
								else:
												lowerCamelCase	 =nn.GroupNorm(num_channels=block_out_channels[0]   , num_groups=UpperCAmelCase_   , eps=1E-6							)
								lowerCamelCase	 =nn.SiLU()
								lowerCamelCase	 =nn.Convad(block_out_channels[0]   , UpperCAmelCase_   , 3   , padding=1							)

								lowerCamelCase	 =False


				def _snake_case       (		self   , UpperCAmelCase_   , UpperCAmelCase_=None							):
								lowerCamelCase	 =z
								lowerCamelCase	 =self.conv_in(UpperCAmelCase_							)

								lowerCamelCase	 =next(iter(self.up_blocks.parameters()							)							).dtype
								if self.training and self.gradient_checkpointing:

												def create_custom_forward(UpperCAmelCase_							):
																def custom_forward(*UpperCAmelCase_							):
																				return module(*UpperCAmelCase_							)

																return custom_forward

												if is_torch_version(""">="""   , """1.11.0"""							):
																# middle
																lowerCamelCase	 =torch.utils.checkpoint.checkpoint(
																    create_custom_forward(self.mid_block							)   , UpperCAmelCase_   , UpperCAmelCase_   , use_reentrant=UpperCAmelCase_							)
																lowerCamelCase	 =sample.to(UpperCAmelCase_							)

																# up
																for up_block in self.up_blocks:
																				lowerCamelCase	 =torch.utils.checkpoint.checkpoint(
																				    create_custom_forward(UpperCAmelCase_							)   , UpperCAmelCase_   , UpperCAmelCase_   , use_reentrant=UpperCAmelCase_							)
												else:
																# middle
																lowerCamelCase	 =torch.utils.checkpoint.checkpoint(
																    create_custom_forward(self.mid_block							)   , UpperCAmelCase_   , UpperCAmelCase_							)
																lowerCamelCase	 =sample.to(UpperCAmelCase_							)

																# up
																for up_block in self.up_blocks:
																				lowerCamelCase	 =torch.utils.checkpoint.checkpoint(create_custom_forward(UpperCAmelCase_							)   , UpperCAmelCase_   , UpperCAmelCase_							)
								else:
												# middle
												lowerCamelCase	 =self.mid_block(UpperCAmelCase_   , UpperCAmelCase_							)
												lowerCamelCase	 =sample.to(UpperCAmelCase_							)

												# up
												for up_block in self.up_blocks:
																lowerCamelCase	 =up_block(UpperCAmelCase_   , UpperCAmelCase_							)

        # post-process
								if latent_embeds is None:
												lowerCamelCase	 =self.conv_norm_out(UpperCAmelCase_							)
								else:
												lowerCamelCase	 =self.conv_norm_out(UpperCAmelCase_   , UpperCAmelCase_							)
								lowerCamelCase	 =self.conv_act(UpperCAmelCase_							)
								lowerCamelCase	 =self.conv_out(UpperCAmelCase_							)

								return sample




class       __A     (						nn.Module    ):


				def __init__(		self   , UpperCAmelCase_   , UpperCAmelCase_   , UpperCAmelCase_   , UpperCAmelCase_=None   , UpperCAmelCase_="random"   , UpperCAmelCase_=False   , UpperCAmelCase_=True							):
								super().__init__()
								lowerCamelCase	 =n_e
								lowerCamelCase	 =vq_embed_dim
								lowerCamelCase	 =beta
								lowerCamelCase	 =legacy

								lowerCamelCase	 =nn.Embedding(self.n_e   , self.vq_embed_dim							)
								self.embedding.weight.data.uniform_(-1.0 / self.n_e   , 1.0 / self.n_e							)

								lowerCamelCase	 =remap
								if self.remap is not None:
												self.register_buffer("""used"""   , torch.tensor(np.load(self.remap							)							)							)
												lowerCamelCase	 =self.used.shape[0]
												lowerCamelCase	 =unknown_index  # "random" or "extra" or integer
												if self.unknown_index == "extra":
																lowerCamelCase	 =self.re_embed
																lowerCamelCase	 =self.re_embed + 1
												print(
												    f"""Remapping {self.n_e} indices to {self.re_embed} indices. """
												    f"""Using {self.unknown_index} for unknown indices."""							)
								else:
												lowerCamelCase	 =n_e

								lowerCamelCase	 =sane_index_shape


				def _snake_case       (		self   , UpperCAmelCase_							):
								lowerCamelCase	 =inds.shape
								assert len(UpperCAmelCase_							) > 1
								lowerCamelCase	 =inds.reshape(ishape[0]   , -1							)
								lowerCamelCase	 =self.used.to(UpperCAmelCase_							)
								lowerCamelCase	 =(inds[:, :, None] == used[None, None, ...]).long()
								lowerCamelCase	 =match.argmax(-1							)
								lowerCamelCase	 =match.sum(2							) < 1
								if self.unknown_index == "random":
												lowerCamelCase	 =torch.randint(0   , self.re_embed   , size=new[unknown].shape							).to(device=new.device							)
								else:
												lowerCamelCase	 =self.unknown_index
								return new.reshape(UpperCAmelCase_							)


				def _snake_case       (		self   , UpperCAmelCase_							):
								lowerCamelCase	 =inds.shape
								assert len(UpperCAmelCase_							) > 1
								lowerCamelCase	 =inds.reshape(ishape[0]   , -1							)
								lowerCamelCase	 =self.used.to(UpperCAmelCase_							)
								if self.re_embed > self.used.shape[0]:  # extra token
												lowerCamelCase	 =0  # simply set to zero
								lowerCamelCase	 =torch.gather(used[None, :][inds.shape[0] * [0], :]   , 1   , UpperCAmelCase_							)
								return back.reshape(UpperCAmelCase_							)


				def _snake_case       (		self   , UpperCAmelCase_							):
								# reshape z -> (batch, height, width, channel) and flatten
								lowerCamelCase	 =z.permute(0   , 2   , 3   , 1							).contiguous()
								lowerCamelCase	 =z.view(-1   , self.vq_embed_dim							)

								# distances from z to embeddings e_j (z - e)^2 = z^2 + e^2 - 2 e * z
								lowerCamelCase	 =torch.argmin(torch.cdist(UpperCAmelCase_   , self.embedding.weight							)   , dim=1							)

								lowerCamelCase	 =self.embedding(UpperCAmelCase_							).view(z.shape							)
								lowerCamelCase	 =None
								lowerCamelCase	 =None

								# compute loss for embedding
								if not self.legacy:
												lowerCamelCase	 =self.beta * torch.mean((z_q.detach() - z) ** 2							) + torch.mean((z_q - z.detach()) ** 2							)
								else:
												lowerCamelCase	 =torch.mean((z_q.detach() - z) ** 2							) + self.beta * torch.mean((z_q - z.detach()) ** 2							)

								# preserve gradients
								lowerCamelCase	 =z + (z_q - z).detach()

								# reshape back to match original input shape
								lowerCamelCase	 =z_q.permute(0   , 3   , 1   , 2							).contiguous()

								if self.remap is not None:
												lowerCamelCase	 =min_encoding_indices.reshape(z.shape[0]   , -1							)  # add batch axis
												lowerCamelCase	 =self.remap_to_used(UpperCAmelCase_							)
												lowerCamelCase	 =min_encoding_indices.reshape(-1   , 1							)  # flatten

								if self.sane_index_shape:
												lowerCamelCase	 =min_encoding_indices.reshape(z_q.shape[0]   , z_q.shape[2]   , z_q.shape[3]							)

								return z_q, loss, (perplexity, min_encodings, min_encoding_indices)


				def _snake_case       (		self   , UpperCAmelCase_   , UpperCAmelCase_							):
								# shape specifying (batch, height, width, channel)
								if self.remap is not None:
												lowerCamelCase	 =indices.reshape(shape[0]   , -1							)  # add batch axis
												lowerCamelCase	 =self.unmap_to_all(UpperCAmelCase_							)
												lowerCamelCase	 =indices.reshape(-1							)  # flatten again

								# get quantized latent vectors
								lowerCamelCase	 =self.embedding(UpperCAmelCase_							)

								if shape is not None:
												lowerCamelCase	 =z_q.view(UpperCAmelCase_							)
												# reshape back to match original input shape
												lowerCamelCase	 =z_q.permute(0   , 3   , 1   , 2							).contiguous()

								return z_q




class       __A     (						a    ):


				def __init__(		self   , UpperCAmelCase_   , UpperCAmelCase_=False							):
								lowerCamelCase	 =parameters
								lowerCamelCase ,  lowerCamelCase	 =torch.chunk(UpperCAmelCase_   , 2   , dim=1							)
								lowerCamelCase	 =torch.clamp(self.logvar   , -3_0.0   , 2_0.0							)
								lowerCamelCase	 =deterministic
								lowerCamelCase	 =torch.exp(0.5 * self.logvar							)
								lowerCamelCase	 =torch.exp(self.logvar							)
								if self.deterministic:
												lowerCamelCase	 =lowerCamelCase	 =torch.zeros_like(
												    self.mean   , device=self.parameters.device   , dtype=self.parameters.dtype							)


				def _snake_case       (		self   , UpperCAmelCase_ = None							):
								# make sure sample is on the same device as the parameters and has same dtype
								lowerCamelCase	 =randn_tensor(
								    self.mean.shape   , generator=UpperCAmelCase_   , device=self.parameters.device   , dtype=self.parameters.dtype							)
								lowerCamelCase	 =self.mean + self.std * sample
								return x


				def _snake_case       (		self   , UpperCAmelCase_=None							):
								if self.deterministic:
												return torch.Tensor([0.0]							)
								else:
												if other is None:
																return 0.5 * torch.sum(torch.pow(self.mean   , 2							) + self.var - 1.0 - self.logvar   , dim=[1, 2, 3]							)
												else:
																return 0.5 * torch.sum(
																    torch.pow(self.mean - other.mean   , 2							) / other.var
																    + self.var / other.var
																    - 1.0
																    - self.logvar
																    + other.logvar   , dim=[1, 2, 3]   , )


				def _snake_case       (		self   , UpperCAmelCase_   , UpperCAmelCase_=[1, 2, 3]							):
								if self.deterministic:
												return torch.Tensor([0.0]							)
								lowerCamelCase	 =np.log(2.0 * np.pi							)
								return 0.5 * torch.sum(logtwopi + self.logvar + torch.pow(sample - self.mean   , 2							) / self.var   , dim=UpperCAmelCase_							)


				def _snake_case       (		self							):
								return self.mean





