




"""simple docstring"""




import argparse

import gdown
import numpy as np
import torch
from huggingface_hub import hf_hub_download

from transformers import (
    CLIPTokenizer,
    CLIPTokenizerFast,
    VideoMAEImageProcessor,
    XCLIPConfig,
    XCLIPModel,
    XCLIPProcessor,
    XCLIPTextConfig,
    XCLIPVisionConfig,
)

def lowercase_						(     _snake_case     ,_snake_case		):
							SCREAMING_SNAKE_CASE__		: Optional[Any]     =     XCLIPTextConfig()

							# derive patch size from model name
							SCREAMING_SNAKE_CASE__		: Any     =     model_name.find("""patch"""		)
							SCREAMING_SNAKE_CASE__		: List[str]     =     int(model_name[start_idx + len("""patch"""		) : start_idx + len("""patch"""		) + 2]		)
							SCREAMING_SNAKE_CASE__		: List[Any]     =     XCLIPVisionConfig(patch_size=lowercase_     ,num_frames=lowercase_		)

							if "large" in model_name:
														SCREAMING_SNAKE_CASE__		: str     =     768
														SCREAMING_SNAKE_CASE__		: str     =     3_072
														SCREAMING_SNAKE_CASE__		: Any     =     12

														SCREAMING_SNAKE_CASE__		: Tuple     =     1_024
														SCREAMING_SNAKE_CASE__		: str     =     4_096
														SCREAMING_SNAKE_CASE__		: Any     =     16
														SCREAMING_SNAKE_CASE__		: List[Any]     =     24
														SCREAMING_SNAKE_CASE__		: Optional[Any]     =     768
														SCREAMING_SNAKE_CASE__		: Any     =     3_072

							if model_name == "xclip-large-patch14-16-frames":
														SCREAMING_SNAKE_CASE__		: List[Any]     =     336

							SCREAMING_SNAKE_CASE__		: List[str]     =     XCLIPConfig.from_text_vision_configs(lowercase_     ,lowercase_		)

							if "large" in model_name:
														SCREAMING_SNAKE_CASE__		: Dict     =     768

							return config

def lowercase_						(     _snake_case		):
							# text encoder
							if name == "token_embedding.weight":
														SCREAMING_SNAKE_CASE__		: Dict     =     name.replace("""token_embedding.weight"""     ,"""text_model.embeddings.token_embedding.weight"""		)
							if name == "positional_embedding":
														SCREAMING_SNAKE_CASE__		: Union[str, Any]     =     name.replace("""positional_embedding"""     ,"""text_model.embeddings.position_embedding.weight"""		)
							if "ln_1" in name:
														SCREAMING_SNAKE_CASE__		: Optional[Any]     =     name.replace("""ln_1"""     ,"""layer_norm1"""		)
							if "ln_2" in name:
														SCREAMING_SNAKE_CASE__		: Optional[int]     =     name.replace("""ln_2"""     ,"""layer_norm2"""		)
							if "c_fc" in name:
														SCREAMING_SNAKE_CASE__		: Optional[Any]     =     name.replace("""c_fc"""     ,"""fc1"""		)
							if "c_proj" in name:
														SCREAMING_SNAKE_CASE__		: str     =     name.replace("""c_proj"""     ,"""fc2"""		)
							if name.startswith("""transformer.resblocks"""		):
														SCREAMING_SNAKE_CASE__		: Dict     =     name.replace("""transformer.resblocks"""     ,"""text_model.encoder.layers"""		)
							if "attn.out_proj" in name and "message" not in name:
														SCREAMING_SNAKE_CASE__		: Optional[Any]     =     name.replace("""attn.out_proj"""     ,"""self_attn.out_proj"""		)
							if "ln_final" in name:
														SCREAMING_SNAKE_CASE__		: List[Any]     =     name.replace("""ln_final"""     ,"""text_model.final_layer_norm"""		)
							# visual encoder
							if name == "visual.class_embedding":
														SCREAMING_SNAKE_CASE__		: List[str]     =     name.replace("""visual.class_embedding"""     ,"""vision_model.embeddings.class_embedding"""		)
							if name == "visual.positional_embedding":
														SCREAMING_SNAKE_CASE__		: Union[str, Any]     =     name.replace("""visual.positional_embedding"""     ,"""vision_model.embeddings.position_embedding.weight"""		)
							if name.startswith("""visual.transformer.resblocks"""		):
														SCREAMING_SNAKE_CASE__		: Any     =     name.replace("""visual.transformer.resblocks"""     ,"""vision_model.encoder.layers"""		)
							if "visual.conv1" in name:
														SCREAMING_SNAKE_CASE__		: str     =     name.replace("""visual.conv1"""     ,"""vision_model.embeddings.patch_embedding"""		)
							if "visual.ln_pre" in name:
														SCREAMING_SNAKE_CASE__		: List[str]     =     name.replace("""visual.ln_pre"""     ,"""vision_model.pre_layernorm"""		)
							if "visual.ln_post" in name:
														SCREAMING_SNAKE_CASE__		: str     =     name.replace("""visual.ln_post"""     ,"""vision_model.post_layernorm"""		)
							if "visual.proj" in name:
														SCREAMING_SNAKE_CASE__		: Dict     =     name.replace("""visual.proj"""     ,"""visual_projection.weight"""		)
							if "text_projection" in name:
														SCREAMING_SNAKE_CASE__		: Any     =     name.replace("""text_projection"""     ,"""text_projection.weight"""		)
							# things on top
							if "prompts_visual_proj" in name:
														SCREAMING_SNAKE_CASE__		: str     =     name.replace("""prompts_visual_proj"""     ,"""prompts_visual_projection"""		)
							if "prompts_visual_ln" in name:
														SCREAMING_SNAKE_CASE__		: List[Any]     =     name.replace("""prompts_visual_ln"""     ,"""prompts_visual_layernorm"""		)
							# mit
							if name == "mit.positional_embedding":
														SCREAMING_SNAKE_CASE__		: List[str]     =     name.replace("""positional"""     ,"""position"""		)
							if name.startswith("""mit.resblocks"""		):
														SCREAMING_SNAKE_CASE__		: Union[str, Any]     =     name.replace("""mit.resblocks"""     ,"""mit.encoder.layers"""		)
							# prompts generator
							if name.startswith("""prompts_generator.norm"""		):
														SCREAMING_SNAKE_CASE__		: List[Any]     =     name.replace("""prompts_generator.norm"""     ,"""prompts_generator.layernorm"""		)

							return name

def lowercase_						(     _snake_case     ,_snake_case		):
							for key in orig_state_dict.copy().keys():
														SCREAMING_SNAKE_CASE__		: Dict     =     orig_state_dict.pop(lowercase_		)

														if "attn.in_proj" in key:
																					SCREAMING_SNAKE_CASE__		: Union[str, Any]     =     key.split("""."""		)
																					if key.startswith("""visual"""		):
																												SCREAMING_SNAKE_CASE__		: Any     =     key_split[3]
																												SCREAMING_SNAKE_CASE__		: Optional[int]     =     config.vision_config.hidden_size
																												if "message_attn" in key:
																																			if "weight" in key:
																																										SCREAMING_SNAKE_CASE__		: List[Any]     =     val[
																																										    :dim, :
																																										]
																																										SCREAMING_SNAKE_CASE__		: Optional[Any]     =     val[
																																										    dim : dim * 2, :
																																										]
																																										SCREAMING_SNAKE_CASE__		: Tuple     =     val[
																																										    -dim:, :
																																										]
																																			else:
																																										SCREAMING_SNAKE_CASE__		: List[str]     =     val[
																																										    :dim
																																										]
																																										SCREAMING_SNAKE_CASE__		: Union[str, Any]     =     val[
																																										    dim : dim * 2
																																										]
																																										SCREAMING_SNAKE_CASE__		: List[Any]     =     val[
																																										    -dim:
																																										]
																												else:
																																			if "weight" in key:
																																										SCREAMING_SNAKE_CASE__		: List[Any]     =     val[
																																										    :dim, :
																																										]
																																										SCREAMING_SNAKE_CASE__		: Union[str, Any]     =     val[
																																										    dim : dim * 2, :
																																										]
																																										SCREAMING_SNAKE_CASE__		: List[str]     =     val[
																																										    -dim:, :
																																										]
																																			else:
																																										SCREAMING_SNAKE_CASE__		: Optional[Any]     =     val[:dim]
																																										SCREAMING_SNAKE_CASE__		: Dict     =     val[
																																										    dim : dim * 2
																																										]
																																										SCREAMING_SNAKE_CASE__		: str     =     val[-dim:]
																					elif key.startswith("""mit"""		):
																												SCREAMING_SNAKE_CASE__		: Optional[Any]     =     key_split[2]
																												SCREAMING_SNAKE_CASE__		: str     =     config.vision_config.mit_hidden_size
																												if "weight" in key:
																																			SCREAMING_SNAKE_CASE__		: Dict     =     val[:dim, :]
																																			SCREAMING_SNAKE_CASE__		: int     =     val[dim : dim * 2, :]
																																			SCREAMING_SNAKE_CASE__		: List[str]     =     val[-dim:, :]
																												else:
																																			SCREAMING_SNAKE_CASE__		: List[str]     =     val[:dim]
																																			SCREAMING_SNAKE_CASE__		: List[Any]     =     val[dim : dim * 2]
																																			SCREAMING_SNAKE_CASE__		: Optional[Any]     =     val[-dim:]
																					else:
																												SCREAMING_SNAKE_CASE__		: Optional[int]     =     key_split[2]
																												SCREAMING_SNAKE_CASE__		: Dict     =     config.text_config.hidden_size
																												if "weight" in key:
																																			SCREAMING_SNAKE_CASE__		: int     =     val[:dim, :]
																																			SCREAMING_SNAKE_CASE__		: Any     =     val[
																																			    dim : dim * 2, :
																																			]
																																			SCREAMING_SNAKE_CASE__		: List[str]     =     val[-dim:, :]
																												else:
																																			SCREAMING_SNAKE_CASE__		: Any     =     val[:dim]
																																			SCREAMING_SNAKE_CASE__		: Union[str, Any]     =     val[
																																			    dim : dim * 2
																																			]
																																			SCREAMING_SNAKE_CASE__		: str     =     val[-dim:]
														else:
																					SCREAMING_SNAKE_CASE__		: List[Any]     =     rename_key(lowercase_		)
																					if new_key_name in ["visual_projection.weight", "text_projection.weight"]:
																												SCREAMING_SNAKE_CASE__		: List[Any]     =     val.T
																					SCREAMING_SNAKE_CASE__		: List[Any]     =     val

							return orig_state_dict

def lowercase_						(     _snake_case		):
							if num_frames == 8:
														SCREAMING_SNAKE_CASE__		: Tuple     =     """eating_spaghetti_8_frames.npy"""
							elif num_frames == 16:
														SCREAMING_SNAKE_CASE__		: Tuple     =     """eating_spaghetti.npy"""
							elif num_frames == 32:
														SCREAMING_SNAKE_CASE__		: Optional[int]     =     """eating_spaghetti_32_frames.npy"""
							SCREAMING_SNAKE_CASE__		: Any     =     hf_hub_download(
							    repo_id="""hf-internal-testing/spaghetti-video"""     ,filename=lowercase_     ,repo_type="""dataset"""     ,)
							SCREAMING_SNAKE_CASE__		: str     =     np.load(lowercase_		)
							return list(lowercase_		)

def lowercase_						(     _snake_case     ,_snake_case=None     ,_snake_case=False		):
							SCREAMING_SNAKE_CASE__		: int     =     {
							    # fully supervised kinetics-400 checkpoints
							    """xclip-base-patch32""": """https://github.com/nbl97/X-CLIP_Model_Zoo/releases/download/v1.0/k400_32_8.pth""",
							    """xclip-base-patch32-16-frames""": (
							        """https://github.com/nbl97/X-CLIP_Model_Zoo/releases/download/v1.0/k400_32_16.pth"""
							    ),
							    """xclip-base-patch16""": """https://github.com/nbl97/X-CLIP_Model_Zoo/releases/download/v1.0/k400_16_8.pth""",
							    """xclip-base-patch16-16-frames""": (
							        """https://github.com/nbl97/X-CLIP_Model_Zoo/releases/download/v1.0/k400_16_16.pth"""
							    ),
							    """xclip-large-patch14""": """https://drive.google.com/u/0/uc?id=1NUOImq0o5DlQTST17iIP3vG7DgmHQuCx&amp;export=download&amp;confirm=t&amp;uuid=b26caedc-88e2-473e-830a-9d158b653cdb""",
							    """xclip-large-patch14-16-frames""": """https://drive.google.com/u/0/uc?id=1FOYgnJc097OJ4lGwtRCCydQyVPJEOH7d&amp;export=download&amp;confirm=t&amp;uuid=538fa810-e671-4050-b385-9a623f89804f""",
							    # fully supervised kinetics-600 checkpoints
							    """xclip-base-patch16-kinetics-600""": (
							        """https://github.com/nbl97/X-CLIP_Model_Zoo/releases/download/v1.0/k600_16_8.pth"""
							    ),
							    """xclip-base-patch16-kinetics-600-16-frames""": (
							        """https://github.com/nbl97/X-CLIP_Model_Zoo/releases/download/v1.0/k600_16_16.pth"""
							    ),
							    """xclip-large-patch14-kinetics-600""": """https://drive.google.com/u/0/uc?id=1FV8C1INuM91sLAN4ImjzePLIlpMSihwV&amp;export=download&amp;confirm=t&amp;uuid=141d4977-4a65-44ae-864f-4b0c19f838be""",
							    # few shot
							    """xclip-base-patch16-hmdb-2-shot""": (
							        """https://github.com/nbl97/X-CLIP_Model_Zoo/releases/download/v1.0/few_hmdb_2.pth"""
							    ),
							    """xclip-base-patch16-hmdb-4-shot""": (
							        """https://github.com/nbl97/X-CLIP_Model_Zoo/releases/download/v1.0/few_hmdb_4.pth"""
							    ),
							    """xclip-base-patch16-hmdb-8-shot""": (
							        """https://github.com/nbl97/X-CLIP_Model_Zoo/releases/download/v1.0/few_hmdb_8.pth"""
							    ),
							    """xclip-base-patch16-hmdb-16-shot""": (
							        """https://github.com/nbl97/X-CLIP_Model_Zoo/releases/download/v1.0/few_hmdb_16.pth"""
							    ),
							    """xclip-base-patch16-ucf-2-shot""": (
							        """https://github.com/nbl97/X-CLIP_Model_Zoo/releases/download/v1.0/few_ucf_2.pth"""
							    ),
							    """xclip-base-patch16-ucf-4-shot""": (
							        """https://github.com/nbl97/X-CLIP_Model_Zoo/releases/download/v1.0/few_ucf_4.pth"""
							    ),
							    """xclip-base-patch16-ucf-8-shot""": (
							        """https://github.com/nbl97/X-CLIP_Model_Zoo/releases/download/v1.0/few_ucf_8.pth"""
							    ),
							    """xclip-base-patch16-ucf-16-shot""": (
							        """https://github.com/nbl97/X-CLIP_Model_Zoo/releases/download/v1.0/few_ucf_16.pth"""
							    ),
							    # zero shot
							    """xclip-base-patch16-zero-shot""": """https://github.com/nbl97/X-CLIP_Model_Zoo/releases/download/v1.0/zero.pth""",
							}

							SCREAMING_SNAKE_CASE__		: List[Any]     =     model_to_url[model_name]
							SCREAMING_SNAKE_CASE__		: List[str]     =     8
							if "16-frames" in model_name:
														SCREAMING_SNAKE_CASE__		: Any     =     16
							elif "shot" in model_name:
														SCREAMING_SNAKE_CASE__		: Dict     =     32

							SCREAMING_SNAKE_CASE__		: Any     =     get_xclip_config(lowercase_     ,lowercase_		)
							SCREAMING_SNAKE_CASE__		: str     =     XCLIPModel(lowercase_		)
							model.eval()

							if "drive" in checkpoint_url:
														SCREAMING_SNAKE_CASE__		: Any     =     """pytorch_model.bin"""
														gdown.cached_download(lowercase_     ,lowercase_     ,quiet=lowercase_		)
														SCREAMING_SNAKE_CASE__		: int     =     torch.load(lowercase_     ,map_location="""cpu"""		)["""model"""]
							else:
														SCREAMING_SNAKE_CASE__		: Dict     =     torch.hub.load_state_dict_from_url(lowercase_		)["""model"""]

							SCREAMING_SNAKE_CASE__		: List[Any]     =     convert_state_dict(lowercase_     ,lowercase_		)

							SCREAMING_SNAKE_CASE__		: Tuple     =     XCLIPModel(lowercase_		)
							SCREAMING_SNAKE_CASE__		: List[Any]     =     model.load_state_dict(lowercase_     ,strict=lowercase_		)
							assert missing_keys == ["text_model.embeddings.position_ids", "vision_model.embeddings.position_ids"]
							model.eval()

							SCREAMING_SNAKE_CASE__		: List[Any]     =     336 if model_name == """xclip-large-patch14-16-frames""" else 224
							SCREAMING_SNAKE_CASE__		: Optional[int]     =     VideoMAEImageProcessor(size=lowercase_		)
							SCREAMING_SNAKE_CASE__		: Union[str, Any]     =     CLIPTokenizer.from_pretrained("""openai/clip-vit-base-patch32"""		)
							SCREAMING_SNAKE_CASE__		: Tuple     =     CLIPTokenizerFast.from_pretrained("""openai/clip-vit-base-patch32"""		)
							SCREAMING_SNAKE_CASE__		: Dict     =     XCLIPProcessor(image_processor=lowercase_     ,tokenizer=lowercase_		)

							SCREAMING_SNAKE_CASE__		: str     =     prepare_video(lowercase_		)
							SCREAMING_SNAKE_CASE__		: Tuple     =     processor(
							    text=["""playing sports""", """eating spaghetti""", """go shopping"""]     ,videos=lowercase_     ,return_tensors="""pt"""     ,padding=lowercase_		)

							print("""Shape of pixel values:"""     ,inputs.pixel_values.shape		)

							with torch.no_grad():
														SCREAMING_SNAKE_CASE__		: List[str]     =     model(**lowercase_		)

							# Verify outputs
							SCREAMING_SNAKE_CASE__		: int     =     outputs.logits_per_video
							SCREAMING_SNAKE_CASE__		: str     =     logits_per_video.softmax(dim=1		)
							print("""Probs:"""     ,lowercase_		)
							# kinetics-400
							if model_name == "xclip-base-patch32":
														SCREAMING_SNAKE_CASE__		: Optional[Any]     =     torch.tensor([[0.0019, 0.9951, 0.0030]]		)
							elif model_name == "xclip-base-patch32-16-frames":
														SCREAMING_SNAKE_CASE__		: Union[str, Any]     =     torch.tensor([[7.09_99E-04, 9.98_83E-01, 4.55_80E-04]]		)
							elif model_name == "xclip-base-patch16":
														SCREAMING_SNAKE_CASE__		: Tuple     =     torch.tensor([[0.0083, 0.9681, 0.0236]]		)
							elif model_name == "xclip-base-patch16-16-frames":
														SCREAMING_SNAKE_CASE__		: Union[str, Any]     =     torch.tensor([[7.69_37E-04, 9.97_28E-01, 1.94_73E-03]]		)
							elif model_name == "xclip-large-patch14":
														SCREAMING_SNAKE_CASE__		: Any     =     torch.tensor([[0.0062, 0.9864, 0.0075]]		)
							elif model_name == "xclip-large-patch14-16-frames":
														SCREAMING_SNAKE_CASE__		: List[Any]     =     torch.tensor([[3.38_77E-04, 9.99_37E-01, 2.88_88E-04]]		)
							# kinetics-600
							elif model_name == "xclip-base-patch16-kinetics-600":
														SCREAMING_SNAKE_CASE__		: Tuple     =     torch.tensor([[0.0555, 0.8914, 0.0531]]		)
							elif model_name == "xclip-base-patch16-kinetics-600-16-frames":
														SCREAMING_SNAKE_CASE__		: Tuple     =     torch.tensor([[3.85_54E-04, 9.99_29E-01, 3.27_54E-04]]		)
							elif model_name == "xclip-large-patch14-kinetics-600":
														SCREAMING_SNAKE_CASE__		: List[Any]     =     torch.tensor([[0.0036, 0.9920, 0.0045]]		)
							# few shot
							elif model_name == "xclip-base-patch16-hmdb-2-shot":
														SCREAMING_SNAKE_CASE__		: List[Any]     =     torch.tensor([[7.18_90E-06, 9.99_94E-01, 5.65_59E-05]]		)
							elif model_name == "xclip-base-patch16-hmdb-4-shot":
														SCREAMING_SNAKE_CASE__		: Optional[int]     =     torch.tensor([[1.03_20E-05, 9.99_93E-01, 6.24_35E-05]]		)
							elif model_name == "xclip-base-patch16-hmdb-8-shot":
														SCREAMING_SNAKE_CASE__		: Union[str, Any]     =     torch.tensor([[4.13_77E-06, 9.99_90E-01, 9.83_86E-05]]		)
							elif model_name == "xclip-base-patch16-hmdb-16-shot":
														SCREAMING_SNAKE_CASE__		: Any     =     torch.tensor([[4.13_47E-05, 9.99_62E-01, 3.34_11E-04]]		)
							elif model_name == "xclip-base-patch16-ucf-2-shot":
														SCREAMING_SNAKE_CASE__		: Any     =     torch.tensor([[8.58_57E-05, 9.99_28E-01, 6.32_91E-04]]		)
							elif model_name == "xclip-base-patch16-ucf-4-shot":
														SCREAMING_SNAKE_CASE__		: Optional[int]     =     torch.tensor([[8.58_57E-05, 9.99_28E-01, 6.32_91E-04]]		)
							elif model_name == "xclip-base-patch16-ucf-8-shot":
														SCREAMING_SNAKE_CASE__		: Union[str, Any]     =     torch.tensor([[0.0027, 0.9904, 0.0070]]		)
							elif model_name == "xclip-base-patch16-ucf-16-shot":
														SCREAMING_SNAKE_CASE__		: List[Any]     =     torch.tensor([[9.82_19E-04, 9.95_93E-01, 3.08_63E-03]]		)
							# zero shot
							elif model_name == "xclip-base-patch16-zero-shot":
														SCREAMING_SNAKE_CASE__		: Union[str, Any]     =     torch.tensor([[3.50_82E-04, 9.97_85E-01, 1.79_66E-03]]		)
							else:
														raise ValueError(f'''Model name {model_name} not supported'''		)
							assert torch.allclose(lowercase_     ,lowercase_     ,atol=1E-3		)
							print("""Looks ok!"""		)

							if pytorch_dump_folder_path is not None:
														print(f'''Saving model {model_name} to {pytorch_dump_folder_path}'''		)
														model.save_pretrained(lowercase_		)

							if push_to_hub:
														print("""Pushing model, processor and slow tokenizer files to the hub..."""		)
														model.push_to_hub(lowercase_     ,organization="""nielsr"""		)
														processor.push_to_hub(lowercase_     ,organization="""nielsr"""		)
														slow_tokenizer.push_to_hub(lowercase_     ,organization="""nielsr"""		)


if __name__ == "__main__":
					UpperCAmelCase__   :					List[str]												=				argparse.ArgumentParser()
					# Required parameters
					parser.add_argument(
					    '--model_name',
					    default='xclip-base-patch32',
					    type=str,
					    help='Name of the model.',
					)
					parser.add_argument(
					    '--pytorch_dump_folder_path', default=None, type=str, help='Path to the output PyTorch model directory.'
					)
					parser.add_argument(
					    '--push_to_hub', action='store_true', help='Whether or not to push the converted model to the 🤗 hub.'
					)

					UpperCAmelCase__   :					str												=				parser.parse_args()
					convert_xclip_checkpoint(args.model_name, args.pytorch_dump_folder_path, args.push_to_hub)






"""simple docstring"""




from typing import Optional

from torch import nn

from .transformer_ad import TransformeraDModel, TransformeraDModelOutput



class 			lowerCAmelCase_		(nn.Module						):


				"""simple docstring"""



				def __init__(self     , SCREAMING_SNAKE_CASE__ = 16     , SCREAMING_SNAKE_CASE__ = 88     , SCREAMING_SNAKE_CASE__ = None     , SCREAMING_SNAKE_CASE__ = 1     , SCREAMING_SNAKE_CASE__ = 0.0     , SCREAMING_SNAKE_CASE__ = 32     , SCREAMING_SNAKE_CASE__ = None     , SCREAMING_SNAKE_CASE__ = False     , SCREAMING_SNAKE_CASE__ = None     , SCREAMING_SNAKE_CASE__ = None     , SCREAMING_SNAKE_CASE__ = "geglu"     , SCREAMING_SNAKE_CASE__ = None     , )      ->	Optional[Any]:



											"""simple docstring"""




											super().__init__()
											SCREAMING_SNAKE_CASE__		: Union[str, Any]     =     nn.ModuleList(
											    [
											        TransformeraDModel(
											            num_attention_heads=SCREAMING_SNAKE_CASE__     , attention_head_dim=SCREAMING_SNAKE_CASE__     , in_channels=SCREAMING_SNAKE_CASE__     , num_layers=SCREAMING_SNAKE_CASE__     , dropout=SCREAMING_SNAKE_CASE__     , norm_num_groups=SCREAMING_SNAKE_CASE__     , cross_attention_dim=SCREAMING_SNAKE_CASE__     , attention_bias=SCREAMING_SNAKE_CASE__     , sample_size=SCREAMING_SNAKE_CASE__     , num_vector_embeds=SCREAMING_SNAKE_CASE__     , activation_fn=SCREAMING_SNAKE_CASE__     , num_embeds_ada_norm=SCREAMING_SNAKE_CASE__     , )
											        for _ in range(2					)
											    ]					)

											# Variables that can be set by a pipeline:

											# The ratio of transformer1 to transformer2's output states to be combined during inference
											SCREAMING_SNAKE_CASE__		: Tuple     =     0.5

											# The shape of `encoder_hidden_states` is expected to be
											# `(batch_size, condition_lengths[0]+condition_lengths[1], num_features)`
											SCREAMING_SNAKE_CASE__		: Union[str, Any]     =     [77, 2_57]

											# Which transformer to use to encode which condition.
											# E.g. `(1, 0)` means that we'll use `transformers[1](conditions[0])` and `transformers[0](conditions[1])`
											SCREAMING_SNAKE_CASE__		: int     =     [1, 0]



				def 			__magic_name__					(self     , SCREAMING_SNAKE_CASE__     , SCREAMING_SNAKE_CASE__     , SCREAMING_SNAKE_CASE__=None     , SCREAMING_SNAKE_CASE__=None     , SCREAMING_SNAKE_CASE__=None     , SCREAMING_SNAKE_CASE__ = True     , )      ->	Optional[int]:



											"""simple docstring"""




											SCREAMING_SNAKE_CASE__		: Any     =     hidden_states

											SCREAMING_SNAKE_CASE__		: Tuple     =     []
											SCREAMING_SNAKE_CASE__		: List[str]     =     0
											# attention_mask is not used yet
											for i in range(2					):
																		# for each of the two transformers, pass the corresponding condition tokens
																		SCREAMING_SNAKE_CASE__		: str     =     encoder_hidden_states[:, tokens_start : tokens_start + self.condition_lengths[i]]
																		SCREAMING_SNAKE_CASE__		: Any     =     self.transformer_index_for_condition[i]
																		SCREAMING_SNAKE_CASE__		: Union[str, Any]     =     self.transformers[transformer_index](
																		    SCREAMING_SNAKE_CASE__     , encoder_hidden_states=SCREAMING_SNAKE_CASE__     , timestep=SCREAMING_SNAKE_CASE__     , cross_attention_kwargs=SCREAMING_SNAKE_CASE__     , return_dict=SCREAMING_SNAKE_CASE__     , )[0]
																		encoded_states.append(encoded_state - input_states					)
																		tokens_start += self.condition_lengths[i]

											SCREAMING_SNAKE_CASE__		: Tuple     =     encoded_states[0] * self.mix_ratio + encoded_states[1] * (1 - self.mix_ratio)
											SCREAMING_SNAKE_CASE__		: Union[str, Any]     =     output_states + input_states

											if not return_dict:
																		return (output_states,)

											return TransformeraDModelOutput(sample=SCREAMING_SNAKE_CASE__					)




