





import os
import tempfile
import unittest

from transformers import FlaubertConfig, is_torch_available
from transformers.testing_utils import require_torch, require_torch_gpu, slow, torch_device

from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
							import torch

							from transformers import (
							    FlaubertForMultipleChoice,
							    FlaubertForQuestionAnswering,
							    FlaubertForQuestionAnsweringSimple,
							    FlaubertForSequenceClassification,
							    FlaubertForTokenClassification,
							    FlaubertModel,
							    FlaubertWithLMHeadModel,
							)
							from transformers.models.flaubert.modeling_flaubert import FLAUBERT_PRETRAINED_MODEL_ARCHIVE_LIST

class       __A     (						a    ):


				def __init__(		self   , UpperCAmelCase_   , UpperCAmelCase_=13   , UpperCAmelCase_=7   , UpperCAmelCase_=True   , UpperCAmelCase_=True   , UpperCAmelCase_=True   , UpperCAmelCase_=True   , UpperCAmelCase_=True   , UpperCAmelCase_=False   , UpperCAmelCase_=False   , UpperCAmelCase_=False   , UpperCAmelCase_=2   , UpperCAmelCase_=99   , UpperCAmelCase_=0   , UpperCAmelCase_=32   , UpperCAmelCase_=5   , UpperCAmelCase_=4   , UpperCAmelCase_=0.1   , UpperCAmelCase_=0.1   , UpperCAmelCase_=512   , UpperCAmelCase_=12   , UpperCAmelCase_=2   , UpperCAmelCase_=0.0_2   , UpperCAmelCase_=3   , UpperCAmelCase_=4   , UpperCAmelCase_="last"   , UpperCAmelCase_=None   , UpperCAmelCase_=None   , ):
								lowerCamelCase	 =parent
								lowerCamelCase	 =batch_size
								lowerCamelCase	 =seq_length
								lowerCamelCase	 =is_training
								lowerCamelCase	 =use_input_lengths
								lowerCamelCase	 =use_token_type_ids
								lowerCamelCase	 =use_labels
								lowerCamelCase	 =gelu_activation
								lowerCamelCase	 =sinusoidal_embeddings
								lowerCamelCase	 =causal
								lowerCamelCase	 =asm
								lowerCamelCase	 =n_langs
								lowerCamelCase	 =vocab_size
								lowerCamelCase	 =n_special
								lowerCamelCase	 =hidden_size
								lowerCamelCase	 =num_hidden_layers
								lowerCamelCase	 =num_attention_heads
								lowerCamelCase	 =hidden_dropout_prob
								lowerCamelCase	 =attention_probs_dropout_prob
								lowerCamelCase	 =max_position_embeddings
								lowerCamelCase	 =type_vocab_size
								lowerCamelCase	 =type_sequence_label_size
								lowerCamelCase	 =initializer_range
								lowerCamelCase	 =num_labels
								lowerCamelCase	 =num_choices
								lowerCamelCase	 =summary_type
								lowerCamelCase	 =use_proj
								lowerCamelCase	 =scope


				def _snake_case       (		self							):
								lowerCamelCase	 =ids_tensor([self.batch_size, self.seq_length]   , self.vocab_size							)
								lowerCamelCase	 =random_attention_mask([self.batch_size, self.seq_length]							)

								lowerCamelCase	 =None
								if self.use_input_lengths:
												lowerCamelCase	 =(
												    ids_tensor([self.batch_size]   , vocab_size=2							) + self.seq_length - 2
												)  # small variation of seq_length

								lowerCamelCase	 =None
								if self.use_token_type_ids:
												lowerCamelCase	 =ids_tensor([self.batch_size, self.seq_length]   , self.n_langs							)

								lowerCamelCase	 =None
								lowerCamelCase	 =None
								lowerCamelCase	 =None
								if self.use_labels:
												lowerCamelCase	 =ids_tensor([self.batch_size]   , self.type_sequence_label_size							)
												lowerCamelCase	 =ids_tensor([self.batch_size, self.seq_length]   , self.num_labels							)
												lowerCamelCase	 =ids_tensor([self.batch_size]   , 2							).float()
												lowerCamelCase	 =ids_tensor([self.batch_size]   , self.num_choices							)

								lowerCamelCase	 =self.get_config()

								return (
								    config,
								    input_ids,
								    token_type_ids,
								    input_lengths,
								    sequence_labels,
								    token_labels,
								    is_impossible_labels,
								    choice_labels,
								    input_mask,
								)


				def _snake_case       (		self							):
								return FlaubertConfig(
								    vocab_size=self.vocab_size   , n_special=self.n_special   , emb_dim=self.hidden_size   , n_layers=self.num_hidden_layers   , n_heads=self.num_attention_heads   , dropout=self.hidden_dropout_prob   , attention_dropout=self.attention_probs_dropout_prob   , gelu_activation=self.gelu_activation   , sinusoidal_embeddings=self.sinusoidal_embeddings   , asm=self.asm   , causal=self.causal   , n_langs=self.n_langs   , max_position_embeddings=self.max_position_embeddings   , initializer_range=self.initializer_range   , summary_type=self.summary_type   , use_proj=self.use_proj   , )


				def _snake_case       (		self   , UpperCAmelCase_   , UpperCAmelCase_   , UpperCAmelCase_   , UpperCAmelCase_   , UpperCAmelCase_   , UpperCAmelCase_   , UpperCAmelCase_   , UpperCAmelCase_   , UpperCAmelCase_   , ):
								lowerCamelCase	 =FlaubertModel(config=UpperCAmelCase_							)
								model.to(UpperCAmelCase_							)
								model.eval()
								lowerCamelCase	 =model(UpperCAmelCase_   , lengths=UpperCAmelCase_   , langs=UpperCAmelCase_							)
								lowerCamelCase	 =model(UpperCAmelCase_   , langs=UpperCAmelCase_							)
								lowerCamelCase	 =model(UpperCAmelCase_							)
								self.parent.assertEqual(result.last_hidden_state.shape   , (self.batch_size, self.seq_length, self.hidden_size)							)


				def _snake_case       (		self   , UpperCAmelCase_   , UpperCAmelCase_   , UpperCAmelCase_   , UpperCAmelCase_   , UpperCAmelCase_   , UpperCAmelCase_   , UpperCAmelCase_   , UpperCAmelCase_   , UpperCAmelCase_   , ):
								lowerCamelCase	 =FlaubertWithLMHeadModel(UpperCAmelCase_							)
								model.to(UpperCAmelCase_							)
								model.eval()

								lowerCamelCase	 =model(UpperCAmelCase_   , token_type_ids=UpperCAmelCase_   , labels=UpperCAmelCase_							)
								self.parent.assertEqual(result.loss.shape   , ()							)
								self.parent.assertEqual(result.logits.shape   , (self.batch_size, self.seq_length, self.vocab_size)							)


				def _snake_case       (		self   , UpperCAmelCase_   , UpperCAmelCase_   , UpperCAmelCase_   , UpperCAmelCase_   , UpperCAmelCase_   , UpperCAmelCase_   , UpperCAmelCase_   , UpperCAmelCase_   , UpperCAmelCase_   , ):
								lowerCamelCase	 =FlaubertForQuestionAnsweringSimple(UpperCAmelCase_							)
								model.to(UpperCAmelCase_							)
								model.eval()

								lowerCamelCase	 =model(UpperCAmelCase_							)

								lowerCamelCase	 =model(UpperCAmelCase_   , start_positions=UpperCAmelCase_   , end_positions=UpperCAmelCase_							)
								self.parent.assertEqual(result.start_logits.shape   , (self.batch_size, self.seq_length)							)
								self.parent.assertEqual(result.end_logits.shape   , (self.batch_size, self.seq_length)							)


				def _snake_case       (		self   , UpperCAmelCase_   , UpperCAmelCase_   , UpperCAmelCase_   , UpperCAmelCase_   , UpperCAmelCase_   , UpperCAmelCase_   , UpperCAmelCase_   , UpperCAmelCase_   , UpperCAmelCase_   , ):
								lowerCamelCase	 =FlaubertForQuestionAnswering(UpperCAmelCase_							)
								model.to(UpperCAmelCase_							)
								model.eval()

								lowerCamelCase	 =model(UpperCAmelCase_							)

								lowerCamelCase	 =model(
								    UpperCAmelCase_   , start_positions=UpperCAmelCase_   , end_positions=UpperCAmelCase_   , cls_index=UpperCAmelCase_   , is_impossible=UpperCAmelCase_   , p_mask=UpperCAmelCase_   , )

								lowerCamelCase	 =model(
								    UpperCAmelCase_   , start_positions=UpperCAmelCase_   , end_positions=UpperCAmelCase_   , cls_index=UpperCAmelCase_   , is_impossible=UpperCAmelCase_   , )

								((lowerCamelCase) ,  )	 =result_with_labels.to_tuple()

								lowerCamelCase	 =model(UpperCAmelCase_   , start_positions=UpperCAmelCase_   , end_positions=UpperCAmelCase_							)

								((lowerCamelCase) ,  )	 =result_with_labels.to_tuple()

								self.parent.assertEqual(result_with_labels.loss.shape   , ()							)
								self.parent.assertEqual(result.start_top_log_probs.shape   , (self.batch_size, model.config.start_n_top)							)
								self.parent.assertEqual(result.start_top_index.shape   , (self.batch_size, model.config.start_n_top)							)
								self.parent.assertEqual(
								    result.end_top_log_probs.shape   , (self.batch_size, model.config.start_n_top * model.config.end_n_top)							)
								self.parent.assertEqual(
								    result.end_top_index.shape   , (self.batch_size, model.config.start_n_top * model.config.end_n_top)							)
								self.parent.assertEqual(result.cls_logits.shape   , (self.batch_size,)							)


				def _snake_case       (		self   , UpperCAmelCase_   , UpperCAmelCase_   , UpperCAmelCase_   , UpperCAmelCase_   , UpperCAmelCase_   , UpperCAmelCase_   , UpperCAmelCase_   , UpperCAmelCase_   , UpperCAmelCase_   , ):
								lowerCamelCase	 =FlaubertForSequenceClassification(UpperCAmelCase_							)
								model.to(UpperCAmelCase_							)
								model.eval()

								lowerCamelCase	 =model(UpperCAmelCase_							)
								lowerCamelCase	 =model(UpperCAmelCase_   , labels=UpperCAmelCase_							)

								self.parent.assertEqual(result.loss.shape   , ()							)
								self.parent.assertEqual(result.logits.shape   , (self.batch_size, self.type_sequence_label_size)							)


				def _snake_case       (		self   , UpperCAmelCase_   , UpperCAmelCase_   , UpperCAmelCase_   , UpperCAmelCase_   , UpperCAmelCase_   , UpperCAmelCase_   , UpperCAmelCase_   , UpperCAmelCase_   , UpperCAmelCase_   , ):
								lowerCamelCase	 =self.num_labels
								lowerCamelCase	 =FlaubertForTokenClassification(UpperCAmelCase_							)
								model.to(UpperCAmelCase_							)
								model.eval()

								lowerCamelCase	 =model(UpperCAmelCase_   , attention_mask=UpperCAmelCase_   , labels=UpperCAmelCase_							)
								self.parent.assertEqual(result.logits.shape   , (self.batch_size, self.seq_length, self.num_labels)							)


				def _snake_case       (		self   , UpperCAmelCase_   , UpperCAmelCase_   , UpperCAmelCase_   , UpperCAmelCase_   , UpperCAmelCase_   , UpperCAmelCase_   , UpperCAmelCase_   , UpperCAmelCase_   , UpperCAmelCase_   , ):
								lowerCamelCase	 =self.num_choices
								lowerCamelCase	 =FlaubertForMultipleChoice(config=UpperCAmelCase_							)
								model.to(UpperCAmelCase_							)
								model.eval()
								lowerCamelCase	 =input_ids.unsqueeze(1							).expand(-1   , self.num_choices   , -1							).contiguous()
								lowerCamelCase	 =token_type_ids.unsqueeze(1							).expand(-1   , self.num_choices   , -1							).contiguous()
								lowerCamelCase	 =input_mask.unsqueeze(1							).expand(-1   , self.num_choices   , -1							).contiguous()
								lowerCamelCase	 =model(
								    UpperCAmelCase_   , attention_mask=UpperCAmelCase_   , token_type_ids=UpperCAmelCase_   , labels=UpperCAmelCase_   , )
								self.parent.assertEqual(result.logits.shape   , (self.batch_size, self.num_choices)							)


				def _snake_case       (		self							):
								lowerCamelCase	 =self.prepare_config_and_inputs()
								(
								    (
								    lowerCamelCase
								) ,  (
								    lowerCamelCase
								) ,  (
								    lowerCamelCase
								) ,  (
								    lowerCamelCase
								) ,  (
								    lowerCamelCase
								) ,  (
								    lowerCamelCase
								) ,  (
								    lowerCamelCase
								) ,  (
								    lowerCamelCase
								) ,  (
								    lowerCamelCase
								) ,  
								)	 =config_and_inputs
								lowerCamelCase	 ={
								    """input_ids""": input_ids,
								    """token_type_ids""": token_type_ids,
								    """lengths""": input_lengths,
								    """attention_mask""": input_mask,
								}
								return config, inputs_dict




@require_torch
class       __A     (						a    ,					a    ,					unittest.TestCase    ):
				__A      =				(
				    (
				        FlaubertModel,
				        FlaubertWithLMHeadModel,
				        FlaubertForQuestionAnswering,
				        FlaubertForQuestionAnsweringSimple,
				        FlaubertForSequenceClassification,
				        FlaubertForTokenClassification,
				        FlaubertForMultipleChoice,
				    )
				    if is_torch_available()
				    else ()
				)
				__A      =				(
				    {
				        """feature-extraction""": FlaubertModel,
				        """fill-mask""": FlaubertWithLMHeadModel,
				        """question-answering""": FlaubertForQuestionAnsweringSimple,
				        """text-classification""": FlaubertForSequenceClassification,
				        """token-classification""": FlaubertForTokenClassification,
				        """zero-shot""": FlaubertForSequenceClassification,
				    }
				    if is_torch_available()
				    else {}
				)


				def _snake_case       (		self   , UpperCAmelCase_   , UpperCAmelCase_   , UpperCAmelCase_   , UpperCAmelCase_   , UpperCAmelCase_							):
								if (
								    pipeline_test_casse_name == "QAPipelineTests"
								    and tokenizer_name is not None
								    and not tokenizer_name.endswith("""Fast"""							)
								):
												# `QAPipelineTests` fails for a few models when the slower tokenizer are used.
												# (The slower tokenizers were never used for pipeline tests before the pipeline testing rework)
												# TODO: check (and possibly fix) the `QAPipelineTests` with slower tokenizer
												return True

								return False


				def _snake_case       (		self   , UpperCAmelCase_   , UpperCAmelCase_   , UpperCAmelCase_=False							):
								lowerCamelCase	 =super()._prepare_for_class(UpperCAmelCase_   , UpperCAmelCase_   , return_labels=UpperCAmelCase_							)

								if return_labels:
												if model_class.__name__ == "FlaubertForQuestionAnswering":
																lowerCamelCase	 =torch.zeros(
																    self.model_tester.batch_size   , dtype=torch.long   , device=UpperCAmelCase_							)
																lowerCamelCase	 =torch.zeros(
																    self.model_tester.batch_size   , dtype=torch.long   , device=UpperCAmelCase_							)

								return inputs_dict


				def _snake_case       (		self							):
								lowerCamelCase	 =FlaubertModelTester(self							)
								lowerCamelCase	 =ConfigTester(self   , config_class=UpperCAmelCase_   , emb_dim=37							)


				def _snake_case       (		self							):
								self.config_tester.run_common_tests()


				def _snake_case       (		self							):
								lowerCamelCase	 =self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_flaubert_model(*UpperCAmelCase_							)


				def _snake_case       (		self							):
								lowerCamelCase	 =self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_flaubert_lm_head(*UpperCAmelCase_							)


				def _snake_case       (		self							):
								lowerCamelCase	 =self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_flaubert_simple_qa(*UpperCAmelCase_							)


				def _snake_case       (		self							):
								lowerCamelCase	 =self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_flaubert_qa(*UpperCAmelCase_							)


				def _snake_case       (		self							):
								lowerCamelCase	 =self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_flaubert_sequence_classif(*UpperCAmelCase_							)


				def _snake_case       (		self							):
								lowerCamelCase	 =self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_flaubert_token_classif(*UpperCAmelCase_							)


				def _snake_case       (		self							):
								lowerCamelCase	 =self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_flaubert_multiple_choice(*UpperCAmelCase_							)


				@slow
				def _snake_case       (		self							):
								for model_name in FLAUBERT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
												lowerCamelCase	 =FlaubertModel.from_pretrained(UpperCAmelCase_							)
												self.assertIsNotNone(UpperCAmelCase_							)


				@slow
				@require_torch_gpu
				def _snake_case       (		self							):
								lowerCamelCase ,  lowerCamelCase	 =self.model_tester.prepare_config_and_inputs_for_common()
								for model_class in self.all_model_classes:
												# FlauBertForMultipleChoice behaves incorrectly in JIT environments.
												if model_class == FlaubertForMultipleChoice:
																return

												lowerCamelCase	 =True
												lowerCamelCase	 =model_class(config=UpperCAmelCase_							)

												lowerCamelCase	 =self._prepare_for_class(UpperCAmelCase_   , UpperCAmelCase_							)
												lowerCamelCase	 =torch.jit.trace(
												    UpperCAmelCase_   , (inputs_dict["""input_ids"""].to("""cpu"""							), inputs_dict["""attention_mask"""].to("""cpu"""							))							)

												with tempfile.TemporaryDirectory() as tmp:
																torch.jit.save(UpperCAmelCase_   , os.path.join(UpperCAmelCase_   , """traced_model.pt"""							)							)
																lowerCamelCase	 =torch.jit.load(os.path.join(UpperCAmelCase_   , """traced_model.pt"""							)   , map_location=UpperCAmelCase_							)
																loaded(inputs_dict["""input_ids"""].to(UpperCAmelCase_							)   , inputs_dict["""attention_mask"""].to(UpperCAmelCase_							)							)




@require_torch
class       __A     (						unittest.TestCase    ):


				@slow
				def _snake_case       (		self							):
								lowerCamelCase	 =FlaubertModel.from_pretrained("""flaubert/flaubert_base_cased"""							)
								lowerCamelCase	 =torch.tensor([[0, 345, 232, 328, 740, 140, 1695, 69, 6078, 1588, 2]]							)
								with torch.no_grad():
												lowerCamelCase	 =model(UpperCAmelCase_							)[0]
								lowerCamelCase	 =torch.Size((1, 11, 768)							)
								self.assertEqual(output.shape   , UpperCAmelCase_							)
								lowerCamelCase	 =torch.tensor(
								    [[[-2.6_2_5_1, -1.4_2_9_8, -0.0_2_2_7], [-2.8_5_1_0, -1.6_3_8_7, 0.2_2_5_8], [-2.8_1_1_4, -1.1_8_3_2, -0.3_0_6_6]]]							)

								self.assertTrue(torch.allclose(output[:, :3, :3]   , UpperCAmelCase_   , atol=1E-4							)							)












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





