import argparse
import pathlib
from pathlib import Path
from tempfile import TemporaryDirectory

import esm as esm_module
import torch
from esm.esmfold.va.misc import batch_encode_sequences as esmfold_encode_sequences
from esm.esmfold.va.pretrained import esmfold_va

from transformers.models.esm.configuration_esm import EsmConfig, EsmFoldConfig
from transformers.models.esm.modeling_esm import (
    EsmForMaskedLM,
    EsmForSequenceClassification,
    EsmIntermediate,
    EsmLayer,
    EsmOutput,
    EsmSelfAttention,
    EsmSelfOutput,
)
from transformers.models.esm.modeling_esmfold import EsmForProteinFolding
from transformers.models.esm.tokenization_esm import EsmTokenizer
from transformers.utils import logging


logging.set_verbosity_info()
UpperCAmelCase      		=					logging.get_logger(__name__)

UpperCAmelCase      		=					[
    (
        '''protein1''',
        '''MNGTEGPNFYVPFSNATGVVRSPFEYPQYYLAEPWQFSMLAAYMFLLIVLGFPINFLTLYVTVQHKKLRTPLNYILLNLAVADLFMVLGGFTSTLYTSLHGYFVFGPTGCNLEGFFATLGGEIALWSLVVLAIERYVVVCKPMSNFRFGENHAIMGVAFTWVMALACAAPPLAGWSRYIPEGLQCSCGIDYYTLKPEVNNESFVIYMFVVHFTIPMIIIFFCYGQLVFTVKEAAAQQQESATTQKAEKEVTRMVIIMVIAFLICWVPYASVAFYIFTHQGSNFGPIFMTIPAFFAKSAAIYNPVIYIMMNKQFRNCMLTTICCGKNPLGDDEASATVSKTETSQVAPA''',
    ),
    ('''protein2''', '''MKTVRQERLKSIVRILERSKEPVSGAQLAEELSVSRQVIVQDIAYLRSLGYNIVATPRGYVLA'''),
    ('''protein3''', '''MKTVRQERLKSI<mask>RILERSKEPVSGAQLAEELS<mask>SRQVIVQDIAYLRSLGYN<mask>VATPRGYVLAGG'''),
    ('''protein4''', '''MKTVRQERLKSI<mask>RILERSKEPVSGAQLAEELS<mask>SRQVIVQDIAYLRSLGYN<mask>VATPRGYVLA'''),
]

UpperCAmelCase      		=					{
    '''esm1b_t33_650M_UR50S''': esm_module.pretrained.esm1b_t33_650M_UR50S,
    '''esm1v_t33_650M_UR90S_1''': esm_module.pretrained.esm1v_t33_650M_UR90S_1,
    '''esm1v_t33_650M_UR90S_2''': esm_module.pretrained.esm1v_t33_650M_UR90S_2,
    '''esm1v_t33_650M_UR90S_3''': esm_module.pretrained.esm1v_t33_650M_UR90S_3,
    '''esm1v_t33_650M_UR90S_4''': esm_module.pretrained.esm1v_t33_650M_UR90S_4,
    '''esm1v_t33_650M_UR90S_5''': esm_module.pretrained.esm1v_t33_650M_UR90S_5,
    '''esm2_t48_15B_UR50D''': esm_module.pretrained.esm2_t48_15B_UR50D,
    '''esm2_t36_3B_UR50D''': esm_module.pretrained.esm2_t36_3B_UR50D,
    '''esm2_t33_650M_UR50D''': esm_module.pretrained.esm2_t33_650M_UR50D,
    '''esm2_t30_150M_UR50D''': esm_module.pretrained.esm2_t30_150M_UR50D,
    '''esm2_t12_35M_UR50D''': esm_module.pretrained.esm2_t12_35M_UR50D,
    '''esm2_t6_8M_UR50D''': esm_module.pretrained.esm2_t6_8M_UR50D,
    '''esmfold_v1''': esmfold_va,
}

UpperCAmelCase      		=					list('''ARNDCQEGHILKMFPSTWYV''')

UpperCAmelCase      		=					restypes + ['''X''']
UpperCAmelCase      		=					restypes_with_x + ['''<pad>''', '''<mask>''', '''<cls>''', '''<sep>''', '''<eos>''']


def SCREAMING_SNAKE_CASE					(  ):
       with TemporaryDirectory() as tempdir:
              _UpperCamelCase  :							Union[str, Any]	    =	'\n'.join(lowercase_					)
              _UpperCamelCase  :							List[str]	    =	Path(lowercase_					) / 'vocab.txt'
              vocab_file.write_text(lowercase_					)
              _UpperCamelCase  :							List[Any]	    =	EsmTokenizer(vocab_file=str(lowercase_					)					)
       _UpperCamelCase  :							Any	    =	0  # Overlaps with 'A' but that seems to be what they want
       return hf_tokenizer


def SCREAMING_SNAKE_CASE					(  lowercase_:						Optional[Any]		,    lowercase_:						Dict					):
       _UpperCamelCase  :							str	    =	our_module.load_state_dict(original_module.state_dict()					)
       if status.missing_keys:
              raise ValueError(F"Missing keys: {status.missing_keys}"					)
       if status.unexpected_keys:
              raise ValueError(F"Unexpected keys: {status.unexpected_keys}"					)


def SCREAMING_SNAKE_CASE					(  lowercase_:						str		,    lowercase_:						str		,    lowercase_:						bool		,    lowercase_:						str		,    lowercase_:						str					):
       if model.startswith('esmfold'					):
              _UpperCamelCase  :							Dict	    =	MODEL_MAPPING[model]()
       else:
              _UpperCamelCase    ,   _UpperCamelCase  :							Any	    =	MODEL_MAPPING[model]()
       esm.eval()  # disable dropout

       if model.startswith('esmfold'					):
              _UpperCamelCase  :							Optional[int]	    =	esm.esm.embed_dim
              _UpperCamelCase  :							int	    =	esm.esm.num_layers
              _UpperCamelCase  :							List[str]	    =	esm.esm.attention_heads
              _UpperCamelCase  :							Any	    =	4 * embed_dim
              _UpperCamelCase  :							Union[str, Any]	    =	esm.esm.token_dropout
              _UpperCamelCase  :							List[str]	    =	False  # This code path does not exist in ESM-2
              _UpperCamelCase  :							Optional[Any]	    =	'rotary'
              _UpperCamelCase  :							Optional[int]	    =	True
              _UpperCamelCase  :							Optional[Any]	    =	EsmFoldConfig()
              for key, val in esm.cfg.items():
                     if hasattr(lowercase_		,    lowercase_					) and key != "trunk":
                            setattr(lowercase_		,    lowercase_		,    lowercase_					)
              for key, val in esm.cfg.trunk.items():
                     if hasattr(esmfold_config.trunk		,    lowercase_					) and key != "structure_module":
                            setattr(esmfold_config.trunk		,    lowercase_		,    lowercase_					)
              for key, val in esm.cfg.trunk.structure_module.items():
                     if hasattr(esmfold_config.trunk.structure_module		,    lowercase_					):
                            setattr(esmfold_config.trunk.structure_module		,    lowercase_		,    lowercase_					)
       elif hasattr(lowercase_		,    'args'					):
              # Indicates an ESM-1b or ESM-1v model
              _UpperCamelCase  :							List[str]	    =	esm.args.embed_dim
              _UpperCamelCase  :							Union[str, Any]	    =	esm.args.layers
              _UpperCamelCase  :							List[Any]	    =	esm.args.attention_heads
              _UpperCamelCase  :							str	    =	esm.args.ffn_embed_dim
              _UpperCamelCase  :							int	    =	esm.args.token_dropout
              _UpperCamelCase  :							int	    =	True if esm.emb_layer_norm_before else False
              _UpperCamelCase  :							Dict	    =	'absolute'
              _UpperCamelCase  :							List[str]	    =	False
              _UpperCamelCase  :							Optional[int]	    =	None
       else:
              # Indicates an ESM-2 model
              _UpperCamelCase  :							Optional[Any]	    =	esm.embed_dim
              _UpperCamelCase  :							str	    =	esm.num_layers
              _UpperCamelCase  :							List[Any]	    =	esm.attention_heads
              _UpperCamelCase  :							Optional[int]	    =	4 * embed_dim  # This is hardcoded in ESM-2
              _UpperCamelCase  :							Any	    =	esm.token_dropout
              _UpperCamelCase  :							Union[str, Any]	    =	False  # This code path does not exist in ESM-2
              _UpperCamelCase  :							List[str]	    =	'rotary'
              _UpperCamelCase  :							Tuple	    =	False
              _UpperCamelCase  :							Optional[int]	    =	None

       if is_folding_model:
              _UpperCamelCase  :							Optional[Any]	    =	esm.esm.alphabet
       _UpperCamelCase  :							int	    =	tuple(alphabet.all_toks					)
       _UpperCamelCase  :							int	    =	alphabet.mask_idx
       _UpperCamelCase  :							int	    =	alphabet.padding_idx

       if is_folding_model:
              _UpperCamelCase  :							Union[str, Any]	    =	esm.esm
       else:
              _UpperCamelCase  :							Optional[int]	    =	esm

       _UpperCamelCase  :							str	    =	EsmConfig(
           vocab_size=original_esm_model.embed_tokens.num_embeddings		,    mask_token_id=lowercase_		,    hidden_size=lowercase_		,    num_hidden_layers=lowercase_		,    num_attention_heads=lowercase_		,    intermediate_size=lowercase_		,    max_position_embeddings=1026		,    layer_norm_eps=1e-5		,    attention_probs_dropout_prob=0.0		,    hidden_dropout_prob=0.0		,    pad_token_id=lowercase_		,    emb_layer_norm_before=lowercase_		,    token_dropout=lowercase_		,    position_embedding_type=lowercase_		,    is_folding_model=lowercase_		,    esmfold_config=lowercase_		,    vocab_list=lowercase_		,    )
       if classification_head:
              _UpperCamelCase  :							Tuple	    =	esm.classification_heads['mnli'].out_proj.weight.shape[0]
       print('Our ESM config:'		,    lowercase_					)

       if model.startswith('esmfold'					):
              _UpperCamelCase  :							Union[str, Any]	    =	EsmForProteinFolding
       elif classification_head:
              _UpperCamelCase  :							Optional[Any]	    =	EsmForSequenceClassification
       else:
              _UpperCamelCase  :							int	    =	EsmForMaskedLM
       _UpperCamelCase  :							Dict	    =	model_class(lowercase_					)
       model.eval()

       # Now let's copy all the weights.
       # Embeddings
       _UpperCamelCase  :							Optional[Any]	    =	original_esm_model.embed_tokens.weight
       if position_embedding_type == "absolute":
              _UpperCamelCase  :							List[Any]	    =	original_esm_model.embed_positions.weight

       if config.emb_layer_norm_before:
              _UpperCamelCase  :							List[str]	    =	original_esm_model.emb_layer_norm_before.weight
              _UpperCamelCase  :							List[str]	    =	original_esm_model.emb_layer_norm_before.bias

       _UpperCamelCase  :							str	    =	original_esm_model.emb_layer_norm_after.weight
       _UpperCamelCase  :							List[str]	    =	original_esm_model.emb_layer_norm_after.bias

       for i in range(config.num_hidden_layers					):
              # Encoder: start of layer
              _UpperCamelCase  :							EsmLayer	    =	model.esm.encoder.layer[i]
              # esm_layer: TransformerSentenceEncoderLayer = original_esm_model.layers[i]
              _UpperCamelCase  :							List[Any]	    =	original_esm_model.layers[i]

              # self attention
              _UpperCamelCase  :							EsmSelfAttention	    =	layer.attention.self
              assert (
                  esm_layer.self_attn.k_proj.weight.data.shape
                  == esm_layer.self_attn.q_proj.weight.data.shape
                  == esm_layer.self_attn.v_proj.weight.data.shape
                  == torch.Size((config.hidden_size, config.hidden_size)					)
              )

              _UpperCamelCase  :							List[Any]	    =	esm_layer.self_attn.q_proj.weight
              _UpperCamelCase  :							List[Any]	    =	esm_layer.self_attn.q_proj.bias
              _UpperCamelCase  :							int	    =	esm_layer.self_attn.k_proj.weight
              _UpperCamelCase  :							Optional[int]	    =	esm_layer.self_attn.k_proj.bias
              _UpperCamelCase  :							Union[str, Any]	    =	esm_layer.self_attn.v_proj.weight
              _UpperCamelCase  :							Tuple	    =	esm_layer.self_attn.v_proj.bias

              if getattr(esm_layer.self_attn		,    'rot_emb'		,    lowercase_					) is not None:
                     # Matt: Although inv_freq is not a trainable weight, it is computed at model init and cached.
                     # During the training of ESM-2 the model was converted to float16 precision, which also converts
                     # the inv_freq tensor, and the loss of precision remains even if the model is loaded later as float32.
                     # If we recompute inv_freq without this loss of precision then we will get subtly different rotary
                     # embeddings, which are enough to cause significant discrepancies in model outputs. To avoid this,
                     # we make sure the new model copies the data from the old inv_freq.
                     _UpperCamelCase  :							Dict	    =	esm_layer.self_attn.rot_emb.inv_freq

              # LayerNorm changes for pre-activation
              _UpperCamelCase  :							Union[str, Any]	    =	esm_layer.self_attn_layer_norm.weight
              _UpperCamelCase  :							Optional[int]	    =	esm_layer.self_attn_layer_norm.bias
              _UpperCamelCase  :							Tuple	    =	esm_layer.final_layer_norm.weight
              _UpperCamelCase  :							Dict	    =	esm_layer.final_layer_norm.bias

              # self-attention output
              _UpperCamelCase  :							EsmSelfOutput	    =	layer.attention.output
              assert self_output.dense.weight.shape == esm_layer.self_attn.out_proj.weight.shape
              _UpperCamelCase  :							List[Any]	    =	esm_layer.self_attn.out_proj.weight
              _UpperCamelCase  :							Optional[Any]	    =	esm_layer.self_attn.out_proj.bias

              # intermediate
              _UpperCamelCase  :							EsmIntermediate	    =	layer.intermediate
              assert intermediate.dense.weight.shape == esm_layer.fca.weight.shape
              _UpperCamelCase  :							Optional[Any]	    =	esm_layer.fca.weight
              _UpperCamelCase  :							str	    =	esm_layer.fca.bias

              # output
              _UpperCamelCase  :							EsmOutput	    =	layer.output
              assert bert_output.dense.weight.shape == esm_layer.fca.weight.shape
              _UpperCamelCase  :							List[Any]	    =	esm_layer.fca.weight
              _UpperCamelCase  :							int	    =	esm_layer.fca.bias
              # end of layer

       if is_folding_model:
              _UpperCamelCase  :							List[Any]	    =	esm.esm_s_combine.data
              _UpperCamelCase  :							Union[str, Any]	    =	esm.afa_to_esm.data
              transfer_and_check_weights(esm.embedding		,    model.embedding					)
              transfer_and_check_weights(esm.esm_s_mlp		,    model.esm_s_mlp					)
              transfer_and_check_weights(esm.trunk		,    model.trunk					)
              transfer_and_check_weights(esm.distogram_head		,    model.distogram_head					)
              transfer_and_check_weights(esm.ptm_head		,    model.ptm_head					)
              transfer_and_check_weights(esm.lm_head		,    model.lm_head					)
              transfer_and_check_weights(esm.lddt_head		,    model.lddt_head					)

       elif classification_head:
              _UpperCamelCase  :							Optional[int]	    =	esm.esm.classification_heads['mnli'].dense.weight
              _UpperCamelCase  :							str	    =	esm.classification_heads['mnli'].dense.bias
              _UpperCamelCase  :							Optional[int]	    =	esm.classification_heads['mnli'].out_proj.weight
              _UpperCamelCase  :							Tuple	    =	esm.classification_heads['mnli'].out_proj.bias
       else:
              # LM Head
              _UpperCamelCase  :							Union[str, Any]	    =	esm.lm_head.dense.weight
              _UpperCamelCase  :							Any	    =	esm.lm_head.dense.bias
              _UpperCamelCase  :							Any	    =	esm.lm_head.layer_norm.weight
              _UpperCamelCase  :							Dict	    =	esm.lm_head.layer_norm.bias
              _UpperCamelCase  :							Any	    =	esm.lm_head.weight
              _UpperCamelCase  :							Optional[int]	    =	esm.lm_head.bias

       # Contact prediction head
       transfer_and_check_weights(esm.contact_head		,    model.esm.contact_head					)

       # Prepare data (first 2 sequences from ESMStructuralSplitDataset superfamily / 4)
       if is_folding_model:
              # Folding models aren't trained on masked inputs and don't like mask tokens.
              _UpperCamelCase  :							int	    =	SAMPLE_DATA[:2]
       else:
              _UpperCamelCase  :							List[Any]	    =	SAMPLE_DATA

       if is_folding_model:
              _UpperCamelCase  :							Optional[int]	    =	get_esmfold_tokenizer()
              _UpperCamelCase  :							Any	    =	hf_tokenizer(
                  [row[1] for row in sample_data]		,    return_tensors='pt'		,    padding=lowercase_		,    add_special_tokens=lowercase_					)
              _UpperCamelCase    ,   _UpperCamelCase    ,   _UpperCamelCase    ,   _UpperCamelCase    ,   _UpperCamelCase  :							Union[str, Any]	    =	esmfold_encode_sequences([row[1] for row in sample_data]					)
              _UpperCamelCase  :							Optional[Any]	    =	torch.all(hf_tokens['input_ids'] == esmfold_aas					) and torch.all(
                  hf_tokens['attention_mask'] == esmfold_mask					)
       else:
              # Let's check that we get the same results.
              _UpperCamelCase  :							List[Any]	    =	alphabet.get_batch_converter()
              _UpperCamelCase    ,   _UpperCamelCase    ,   _UpperCamelCase  :							Optional[int]	    =	batch_converter(lowercase_					)
              # Prepare tokenizer and make sure it matches
              with TemporaryDirectory() as tempdir:
                     _UpperCamelCase  :							List[str]	    =	'\n'.join(alphabet.all_toks					)
                     _UpperCamelCase  :							Optional[int]	    =	Path(lowercase_					) / 'vocab.txt'
                     vocab_file.write_text(lowercase_					)
                     _UpperCamelCase  :							Tuple	    =	EsmTokenizer(vocab_file=str(lowercase_					)					)

              _UpperCamelCase  :							Any	    =	hf_tokenizer([row[1] for row in sample_data]		,    return_tensors='pt'		,    padding=lowercase_					)
              _UpperCamelCase  :							List[str]	    =	torch.all(hf_tokens['input_ids'] == batch_tokens					)

       print('Do both models tokenizers output the same tokens?'		,    '🔥' if success else '💩'					)
       if not success:
              raise Exception('Tokenization does not match!'					)

       with torch.no_grad():
              if is_folding_model:
                     # Let's test the model in parts
                     # ESMFold always converts the ESM stem to float16, which requires float16 ops
                     # that don't exist on CPU. Therefore, to test it we need to run it on GPU. However,
                     # ESMFold is what we in the community call a "big boy" and so we desperately avoid putting both the
                     # original and the converted model on the GPU at the same time.
                     _UpperCamelCase  :							Tuple	    =	esm.cuda().infer([row[1] for row in sample_data]					)
                     _UpperCamelCase  :							List[Any]	    =	model.cuda()(
                         input_ids=hf_tokens['input_ids'].cuda()		,    attention_mask=hf_tokens['attention_mask'].cuda()					)
              else:
                     _UpperCamelCase  :							Tuple	    =	model(**lowercase_		,    output_hidden_states=lowercase_					)
                     _UpperCamelCase  :							str	    =	our_output['logits']
                     if classification_head:
                            _UpperCamelCase  :							Optional[int]	    =	esm.model.classification_heads['mnli'](esm.extract_features(lowercase_					)					)
                     else:
                            _UpperCamelCase  :							Any	    =	esm(hf_tokens['input_ids']		,    repr_layers=list(range(999					)					)					)
                            _UpperCamelCase  :							int	    =	their_output['logits']

              if is_folding_model:
                     _UpperCamelCase  :							Union[str, Any]	    =	torch.max(torch.abs(our_output['positions'] - their_output['positions']					)					).item()
                     _UpperCamelCase  :							List[Any]	    =	torch.allclose(our_output['positions']		,    their_output['positions']		,    atol=1e-5					)
              else:
                     _UpperCamelCase  :							Any	    =	torch.max(torch.abs(our_output - their_output					)					).item()
                     _UpperCamelCase  :							List[Any]	    =	torch.allclose(lowercase_		,    lowercase_		,    atol=1e-5					)

              print(F"max_absolute_diff = {max_absolute_diff}"					)  # ~ 1e-5
              print('Do both models output the same tensors?'		,    '🔥' if success else '💩'					)

              if not success:
                     raise Exception('Something went wRoNg'					)

              if not is_folding_model:
                     # Let's check contact prediction too
                     _UpperCamelCase  :							Union[str, Any]	    =	model.predict_contacts(hf_tokens['input_ids']		,    hf_tokens['attention_mask']					)
                     _UpperCamelCase  :							List[str]	    =	esm.predict_contacts(hf_tokens['input_ids']					)
                     _UpperCamelCase  :							Union[str, Any]	    =	torch.max(torch.abs(our_output - their_output					)					).item()
                     _UpperCamelCase  :							str	    =	torch.allclose(lowercase_		,    lowercase_		,    atol=1e-5					)

                     print('Contact prediction testing:'					)
                     print(F"max_absolute_diff = {max_absolute_diff}"					)  # ~ 1e-5
                     print('Do both models output the same tensors?'		,    '🔥' if success else '💩'					)

                     if not success:
                            raise Exception('Something went wRoNg'					)

              pathlib.Path(lowercase_					).mkdir(parents=lowercase_		,    exist_ok=lowercase_					)
              print(F"Saving model to {pytorch_dump_folder_path}"					)
              model.save_pretrained(lowercase_					)

              del esm  # Free up some memory before continuing

       print(F"Saving tokenizer to {pytorch_dump_folder_path}"					)
       hf_tokenizer.save_pretrained(lowercase_					)

       if push_to_repo:
              model.push_to_hub(repo_id=lowercase_		,    use_auth_token=lowercase_					)
              hf_tokenizer.push_to_hub(repo_id=lowercase_		,    use_auth_token=lowercase_					)


if __name__ == "__main__":
 UpperCAmelCase      		=					argparse.ArgumentParser()
 # Required parameters
 parser.add_argument(
     '''--pytorch_dump_folder_path''', type=str, required=True, help='''Path to the output PyTorch model.'''
 )
 parser.add_argument(
     '''--classification_head''', action='''store_true''', help='''Whether to convert a final classification head.'''
 )
 parser.add_argument('''--model''', default=None, type=str, required=True, help='''Name of model to convert.''')
 parser.add_argument('''--push_to_repo''', type=str, help='''Repo to upload to (including username!).''')
 parser.add_argument('''--auth_token''', type=str, help='''HuggingFace auth token.''')
 UpperCAmelCase      		=					parser.parse_args()
 convert_esm_checkpoint_to_pytorch(
     args.model, args.pytorch_dump_folder_path, args.classification_head, args.push_to_repo, args.auth_token
 )









import gc
import random
import unittest

import numpy as np
import torch
from transformers import XLMRobertaTokenizerFast

from diffusers import DDIMScheduler, KandinskyPipeline, KandinskyPriorPipeline, UNetaDConditionModel, VQModel
from diffusers.pipelines.kandinsky.text_encoder import MCLIPConfig, MultilingualCLIP
from diffusers.utils import floats_tensor, load_numpy, slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu

from ..test_pipelines_common import PipelineTesterMixin, assert_mean_pixel_difference


enable_full_determinism()

class   UpperCAmelCase    (      A ,       unittest.TestCase ):
   __SCREAMING_SNAKE_CASE		:				int      					=       KandinskyPipeline
   __SCREAMING_SNAKE_CASE		:				Union[str, Any]      					=       [
       "prompt",
       "image_embeds",
       "negative_image_embeds",
   ]
   __SCREAMING_SNAKE_CASE		:				Tuple      					=       ["prompt", "negative_prompt", "image_embeds", "negative_image_embeds"]
   __SCREAMING_SNAKE_CASE		:				Tuple      					=       [
       "generator",
       "height",
       "width",
       "latents",
       "guidance_scale",
       "negative_prompt",
       "num_inference_steps",
       "return_dict",
       "guidance_scale",
       "num_images_per_prompt",
       "output_type",
       "return_dict",
   ]
   __SCREAMING_SNAKE_CASE		:				Dict      					=       False






   @property
   def 			_a    (						self :							str  )->      Dict:
     return 32






   @property
   def 			_a    (						self :							Tuple  )->      Any:
     return 32






   @property
   def 			_a    (						self :							Union[str, Any]  )->      List[str]:
     return self.time_input_dim






   @property
   def 			_a    (						self :							Optional[Any]  )->      Tuple:
     return self.time_input_dim * 4






   @property
   def 			_a    (						self :							int  )->      Dict:
     return 1_00






   @property
   def 			_a    (						self :							Tuple  )->      Tuple:
     _lowercase:			List[Any]      =     XLMRobertaTokenizerFast.from_pretrained('YiYiXu/tiny-random-mclip-base'  )
     return tokenizer






   @property
   def 			_a    (						self :							List[str]  )->      Optional[int]:
     torch.manual_seed(0  )
     _lowercase:			int      =     MCLIPConfig(
         numDims=self.cross_attention_dim    ,       transformerDimensions=self.text_embedder_hidden_size    ,       hidden_size=self.text_embedder_hidden_size    ,       intermediate_size=37    ,       num_attention_heads=4    ,       num_hidden_layers=5    ,       vocab_size=10_05    ,       )

     _lowercase:			Optional[Any]      =     MultilingualCLIP(UpperCamelCase_  )
     _lowercase:			int      =     text_encoder.eval()

     return text_encoder






   @property
   def 			_a    (						self :							Optional[int]  )->      Optional[Any]:
     torch.manual_seed(0  )

     _lowercase:			Optional[Any]      =     {
         'in_channels': 4,
         # Out channels is double in channels because predicts mean and variance
         'out_channels': 8,
         'addition_embed_type': 'text_image',
         'down_block_types': ('ResnetDownsampleBlock2D', 'SimpleCrossAttnDownBlock2D'),
         'up_block_types': ('SimpleCrossAttnUpBlock2D', 'ResnetUpsampleBlock2D'),
         'mid_block_type': 'UNetMidBlock2DSimpleCrossAttn',
         'block_out_channels': (self.block_out_channels_a, self.block_out_channels_a * 2),
         'layers_per_block': 1,
         'encoder_hid_dim': self.text_embedder_hidden_size,
         'encoder_hid_dim_type': 'text_image_proj',
         'cross_attention_dim': self.cross_attention_dim,
         'attention_head_dim': 4,
         'resnet_time_scale_shift': 'scale_shift',
         'class_embed_type': None,
     }

     _lowercase:			str      =     UNetaDConditionModel(**UpperCamelCase_  )
     return model






   @property
   def 			_a    (						self :							Dict  )->      Optional[Any]:
     return {
         "block_out_channels": [32, 64],
         "down_block_types": ["DownEncoderBlock2D", "AttnDownEncoderBlock2D"],
         "in_channels": 3,
         "latent_channels": 4,
         "layers_per_block": 1,
         "norm_num_groups": 8,
         "norm_type": "spatial",
         "num_vq_embeddings": 12,
         "out_channels": 3,
         "up_block_types": [
             "AttnUpDecoderBlock2D",
             "UpDecoderBlock2D",
         ],
         "vq_embed_dim": 4,
     }






   @property
   def 			_a    (						self :							List[str]  )->      int:
     torch.manual_seed(0  )
     _lowercase:			List[str]      =     VQModel(**self.dummy_movq_kwargs  )
     return model






   def 			_a    (						self :							int  )->      List[Any]:
     _lowercase:			List[str]      =     self.dummy_text_encoder
     _lowercase:			List[Any]      =     self.dummy_tokenizer
     _lowercase:			Optional[Any]      =     self.dummy_unet
     _lowercase:			List[str]      =     self.dummy_movq

     _lowercase:			Optional[Any]      =     DDIMScheduler(
         num_train_timesteps=10_00    ,       beta_schedule='linear'    ,       beta_start=0.00_085    ,       beta_end=0.012    ,       clip_sample=UpperCamelCase_    ,       set_alpha_to_one=UpperCamelCase_    ,       steps_offset=1    ,       prediction_type='epsilon'    ,       thresholding=UpperCamelCase_    ,       )

     _lowercase:			List[Any]      =     {
         'text_encoder': text_encoder,
         'tokenizer': tokenizer,
         'unet': unet,
         'scheduler': scheduler,
         'movq': movq,
     }
     return components






   def 			_a    (						self :							str    ,       UpperCamelCase_ :							int    ,       UpperCamelCase_ :							Optional[int]=0  )->      Union[str, Any]:
     _lowercase:			str      =     floats_tensor((1, self.cross_attention_dim)    ,       rng=random.Random(UpperCamelCase_  )  ).to(UpperCamelCase_  )
     _lowercase:			Optional[int]      =     floats_tensor((1, self.cross_attention_dim)    ,       rng=random.Random(seed + 1  )  ).to(UpperCamelCase_  )
     if str(UpperCamelCase_  ).startswith('mps'  ):
       _lowercase:			Any      =     torch.manual_seed(UpperCamelCase_  )
     else:
       _lowercase:			Tuple      =     torch.Generator(device=UpperCamelCase_  ).manual_seed(UpperCamelCase_  )
     _lowercase:			List[str]      =     {
         'prompt': 'horse',
         'image_embeds': image_embeds,
         'negative_image_embeds': negative_image_embeds,
         'generator': generator,
         'height': 64,
         'width': 64,
         'guidance_scale': 4.0,
         'num_inference_steps': 2,
         'output_type': 'np',
     }
     return inputs






   def 			_a    (						self :							List[str]  )->      List[str]:
     _lowercase:			str      =     'cpu'

     _lowercase:			Optional[Any]      =     self.get_dummy_components()

     _lowercase:			int      =     self.pipeline_class(**UpperCamelCase_  )
     _lowercase:			List[str]      =     pipe.to(UpperCamelCase_  )

     pipe.set_progress_bar_config(disable=UpperCamelCase_  )

     _lowercase:			Tuple      =     pipe(**self.get_dummy_inputs(UpperCamelCase_  )  )
     _lowercase:			Union[str, Any]      =     output.images

     _lowercase:			Union[str, Any]      =     pipe(
         **self.get_dummy_inputs(UpperCamelCase_  )    ,       return_dict=UpperCamelCase_    ,       )[0]

     _lowercase:			int      =     image[0, -3:, -3:, -1]
     _lowercase:			List[str]      =     image_from_tuple[0, -3:, -3:, -1]

     assert image.shape == (1, 64, 64, 3)

     _lowercase:			Union[str, Any]      =     np.array(
         [0.328_663, 1.0, 0.23_216_873, 1.0, 0.92_717_564, 0.4_639_046, 0.96_894_777, 0.31_713_378, 0.6_293_953]  )

     assert (
         np.abs(image_slice.flatten() - expected_slice  ).max() < 1E-2
     ), f' expected_slice {expected_slice}, but got {image_slice.flatten()}'
     assert (
         np.abs(image_from_tuple_slice.flatten() - expected_slice  ).max() < 1E-2
     ), f' expected_slice {expected_slice}, but got {image_from_tuple_slice.flatten()}'

@slow
@require_torch_gpu
class   UpperCAmelCase    (      unittest.TestCase ):






   def 			_a    (						self :							Optional[int]  )->      Any:
     # clean up the VRAM after each test
     super().tearDown()
     gc.collect()
     torch.cuda.empty_cache()






   def 			_a    (						self :							List[Any]  )->      Optional[int]:
     _lowercase:			Optional[int]      =     load_numpy(
         'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'
         '/kandinsky/kandinsky_text2img_cat_fp16.npy'  )

     _lowercase:			List[Any]      =     KandinskyPriorPipeline.from_pretrained(
         'kandinsky-community/kandinsky-2-1-prior'    ,       torch_dtype=torch.floataa  )
     pipe_prior.to(UpperCamelCase_  )

     _lowercase:			str      =     KandinskyPipeline.from_pretrained('kandinsky-community/kandinsky-2-1'    ,       torch_dtype=torch.floataa  )
     _lowercase:			List[str]      =     pipeline.to(UpperCamelCase_  )
     pipeline.set_progress_bar_config(disable=UpperCamelCase_  )

     _lowercase:			Dict      =     'red cat, 4k photo'

     _lowercase:			List[Any]      =     torch.Generator(device='cuda'  ).manual_seed(0  )
     _lowercase  ,       _lowercase:			Dict      =     pipe_prior(
         UpperCamelCase_    ,       generator=UpperCamelCase_    ,       num_inference_steps=5    ,       negative_prompt=''    ,       ).to_tuple()

     _lowercase:			str      =     torch.Generator(device='cuda'  ).manual_seed(0  )
     _lowercase:			Optional[int]      =     pipeline(
         UpperCamelCase_    ,       image_embeds=UpperCamelCase_    ,       negative_image_embeds=UpperCamelCase_    ,       generator=UpperCamelCase_    ,       num_inference_steps=1_00    ,       output_type='np'    ,       )

     _lowercase:			Optional[int]      =     output.images[0]

     assert image.shape == (5_12, 5_12, 3)

     assert_mean_pixel_difference(UpperCamelCase_    ,       UpperCamelCase_  )
