




import gc
import unittest

import numpy as np
import torch

from diffusers import AutoencoderKL, DDIMScheduler, DiTPipeline, DPMSolverMultistepScheduler, TransformeraDModel
from diffusers.utils import is_xformers_available, load_numpy, slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu

from ..pipeline_params import (
    CLASS_CONDITIONED_IMAGE_GENERATION_BATCH_PARAMS,
    CLASS_CONDITIONED_IMAGE_GENERATION_PARAMS,
)
from ..test_pipelines_common import PipelineTesterMixin


enable_full_determinism()
class     a__					(      _snake_case				,       unittest.TestCase   ):
						_a							:					str      =						DiTPipeline
						_a							:					Optional[int]      =						CLASS_CONDITIONED_IMAGE_GENERATION_PARAMS
						_a							:					Any      =						PipelineTesterMixin.required_optional_params - {
						    """latents""",
						    """num_images_per_prompt""",
						    """callback""",
						    """callback_steps""",
						}
						_a							:					Optional[int]      =						CLASS_CONDITIONED_IMAGE_GENERATION_BATCH_PARAMS
						_a							:					Optional[int]      =						False






						def        __SCREAMING_SNAKE_CASE(     self    ):




											"""simple docstring"""



											torch.manual_seed(0    )
											__lowerCAmelCase            = TransformeraDModel(
											    sample_size=1_6	,    num_layers=2	,    patch_size=4	,    attention_head_dim=8	,    num_attention_heads=2	,    in_channels=4	,    out_channels=8	,    attention_bias=snake_case_	,    activation_fn="gelu-approximate"	,    num_embeds_ada_norm=1_0_0_0	,    norm_type="ada_norm_zero"	,    norm_elementwise_affine=snake_case_	,    )
											__lowerCAmelCase            = AutoencoderKL()
											__lowerCAmelCase            = DDIMScheduler()
											__lowerCAmelCase            = {"transformer": transformer.eval(), "vae": vae.eval(), "scheduler": scheduler}
											return components






						def        __SCREAMING_SNAKE_CASE(     self	,    _A	,    _A=0    ):




											"""simple docstring"""



											if str(snake_case_    ).startswith("mps"    ):
																__lowerCAmelCase            = torch.manual_seed(snake_case_    )
											else:
																__lowerCAmelCase            = torch.Generator(device=snake_case_    ).manual_seed(snake_case_    )
											__lowerCAmelCase            = {
											    "class_labels": [1],
											    "generator": generator,
											    "num_inference_steps": 2,
											    "output_type": "numpy",
											}
											return inputs






						def        __SCREAMING_SNAKE_CASE(     self    ):




											"""simple docstring"""



											__lowerCAmelCase            = "cpu"

											__lowerCAmelCase            = self.get_dummy_components()
											__lowerCAmelCase            = self.pipeline_class(**snake_case_    )
											pipe.to(snake_case_    )
											pipe.set_progress_bar_config(disable=snake_case_    )

											__lowerCAmelCase            = self.get_dummy_inputs(snake_case_    )
											__lowerCAmelCase            = pipe(**snake_case_    ).images
											__lowerCAmelCase            = image[0, -3:, -3:, -1]

											self.assertEqual(image.shape	,    (1, 1_6, 1_6, 3)    )
											__lowerCAmelCase            = np.array([0.29_46, 0.66_01, 0.43_29, 0.32_96, 0.41_44, 0.53_19, 0.72_73, 0.50_13, 0.44_57]    )
											__lowerCAmelCase            = np.abs(image_slice.flatten() - expected_slice    ).max()
											self.assertLessEqual(snake_case_	,    1E-3    )






						def        __SCREAMING_SNAKE_CASE(     self    ):




											"""simple docstring"""



											self._test_inference_batch_single_identical(relax_max_difference=snake_case_	,    expected_max_diff=1E-3    )






						@unittest.skipIf(
						    torch_device != "cuda" or not is_xformers_available()	,    reason="XFormers attention is only available with CUDA and `xformers` installed"	,    )
						def        __SCREAMING_SNAKE_CASE(     self    ):




											"""simple docstring"""



											self._test_xformers_attention_forwardGenerator_pass(expected_max_diff=1E-3    )






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






						def        __SCREAMING_SNAKE_CASE(     self    ):




											"""simple docstring"""



											super().tearDown()
											gc.collect()
											torch.cuda.empty_cache()






						def        __SCREAMING_SNAKE_CASE(     self    ):




											"""simple docstring"""



											__lowerCAmelCase            = torch.manual_seed(0    )

											__lowerCAmelCase            = DiTPipeline.from_pretrained("facebook/DiT-XL-2-256"    )
											pipe.to("cuda"    )

											__lowerCAmelCase            = ["vase", "umbrella", "white shark", "white wolf"]
											__lowerCAmelCase            = pipe.get_label_ids(snake_case_    )

											__lowerCAmelCase            = pipe(snake_case_	,    generator=snake_case_	,    num_inference_steps=4_0	,    output_type="np"    ).images

											for word, image in zip(snake_case_	,    snake_case_    ):
																__lowerCAmelCase            = load_numpy(
																    f"""https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/dit/{word}.npy"""    )
																assert np.abs((expected_image - image).max()    ) < 1E-2






						def        __SCREAMING_SNAKE_CASE(     self    ):




											"""simple docstring"""



											__lowerCAmelCase            = DiTPipeline.from_pretrained("facebook/DiT-XL-2-512"    )
											__lowerCAmelCase            = DPMSolverMultistepScheduler.from_config(pipe.scheduler.config    )
											pipe.to("cuda"    )

											__lowerCAmelCase            = ["vase", "umbrella"]
											__lowerCAmelCase            = pipe.get_label_ids(snake_case_    )

											__lowerCAmelCase            = torch.manual_seed(0    )
											__lowerCAmelCase            = pipe(snake_case_	,    generator=snake_case_	,    num_inference_steps=2_5	,    output_type="np"    ).images

											for word, image in zip(snake_case_	,    snake_case_    ):
																__lowerCAmelCase            = load_numpy(
																    "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"
																    f"""/dit/{word}_512.npy"""    )

																assert np.abs((expected_image - image).max()    ) < 1E-1










import argparse
import json

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

from transformers import (
    VideoMAEConfig,
    VideoMAEForPreTraining,
    VideoMAEForVideoClassification,
    VideoMAEImageProcessor,
)







def 			_a   (		SCREAMING_SNAKE_CASE_							:       Tuple  ):
					__lowerCAmelCase            = VideoMAEConfig()

					set_architecture_configs(SCREAMING_SNAKE_CASE_   ,   SCREAMING_SNAKE_CASE_  )

					if "finetuned" not in model_name:
										__lowerCAmelCase            = False

					if "finetuned" in model_name:
										__lowerCAmelCase            = "huggingface/label-files"
										if "kinetics" in model_name:
															__lowerCAmelCase            = 4_00
															__lowerCAmelCase            = "kinetics400-id2label.json"
										elif "ssv2" in model_name:
															__lowerCAmelCase            = 1_74
															__lowerCAmelCase            = "something-something-v2-id2label.json"
										else:
															raise ValueError("Model name should either contain 'kinetics' or 'ssv2' in case it's fine-tuned."  )
										__lowerCAmelCase            = json.load(open(hf_hub_download(SCREAMING_SNAKE_CASE_   ,   SCREAMING_SNAKE_CASE_   ,   repo_type="dataset"  )   ,   "r"  )  )
										__lowerCAmelCase            = {int(SCREAMING_SNAKE_CASE_  ): v for k, v in idalabel.items()}
										__lowerCAmelCase            = idalabel
										__lowerCAmelCase            = {v: k for k, v in idalabel.items()}

					return config







def 			_a   (		SCREAMING_SNAKE_CASE_							:       int   ,   SCREAMING_SNAKE_CASE_							:       List[str]  ):
					if "small" in model_name:
										__lowerCAmelCase            = 3_84
										__lowerCAmelCase            = 15_36
										__lowerCAmelCase            = 12
										__lowerCAmelCase            = 16
										__lowerCAmelCase            = 12
										__lowerCAmelCase            = 3
										__lowerCAmelCase            = 1_92
										__lowerCAmelCase            = 7_68
					elif "large" in model_name:
										__lowerCAmelCase            = 10_24
										__lowerCAmelCase            = 40_96
										__lowerCAmelCase            = 24
										__lowerCAmelCase            = 16
										__lowerCAmelCase            = 12
										__lowerCAmelCase            = 8
										__lowerCAmelCase            = 5_12
										__lowerCAmelCase            = 20_48
					elif "huge" in model_name:
										__lowerCAmelCase            = 12_80
										__lowerCAmelCase            = 51_20
										__lowerCAmelCase            = 32
										__lowerCAmelCase            = 16
										__lowerCAmelCase            = 12
										__lowerCAmelCase            = 8
										__lowerCAmelCase            = 6_40
										__lowerCAmelCase            = 25_60
					elif "base" not in model_name:
										raise ValueError("Model name should include either \"small\", \"base\", \"large\", or \"huge\""  )







def 			_a   (		SCREAMING_SNAKE_CASE_							:       Union[str, Any]  ):
					if "encoder." in name:
										__lowerCAmelCase            = name.replace("encoder."   ,   ""  )
					if "cls_token" in name:
										__lowerCAmelCase            = name.replace("cls_token"   ,   "videomae.embeddings.cls_token"  )
					if "decoder_pos_embed" in name:
										__lowerCAmelCase            = name.replace("decoder_pos_embed"   ,   "decoder.decoder_pos_embed"  )
					if "pos_embed" in name and "decoder" not in name:
										__lowerCAmelCase            = name.replace("pos_embed"   ,   "videomae.embeddings.position_embeddings"  )
					if "patch_embed.proj" in name:
										__lowerCAmelCase            = name.replace("patch_embed.proj"   ,   "videomae.embeddings.patch_embeddings.projection"  )
					if "patch_embed.norm" in name:
										__lowerCAmelCase            = name.replace("patch_embed.norm"   ,   "videomae.embeddings.norm"  )
					if "decoder.blocks" in name:
										__lowerCAmelCase            = name.replace("decoder.blocks"   ,   "decoder.decoder_layers"  )
					if "blocks" in name:
										__lowerCAmelCase            = name.replace("blocks"   ,   "videomae.encoder.layer"  )
					if "attn.proj" in name:
										__lowerCAmelCase            = name.replace("attn.proj"   ,   "attention.output.dense"  )
					if "attn" in name and "bias" not in name:
										__lowerCAmelCase            = name.replace("attn"   ,   "attention.self"  )
					if "attn" in name:
										__lowerCAmelCase            = name.replace("attn"   ,   "attention.attention"  )
					if "norm1" in name:
										__lowerCAmelCase            = name.replace("norm1"   ,   "layernorm_before"  )
					if "norm2" in name:
										__lowerCAmelCase            = name.replace("norm2"   ,   "layernorm_after"  )
					if "mlp.fc1" in name:
										__lowerCAmelCase            = name.replace("mlp.fc1"   ,   "intermediate.dense"  )
					if "mlp.fc2" in name:
										__lowerCAmelCase            = name.replace("mlp.fc2"   ,   "output.dense"  )
					if "decoder_embed" in name:
										__lowerCAmelCase            = name.replace("decoder_embed"   ,   "decoder.decoder_embed"  )
					if "decoder_norm" in name:
										__lowerCAmelCase            = name.replace("decoder_norm"   ,   "decoder.decoder_norm"  )
					if "decoder_pred" in name:
										__lowerCAmelCase            = name.replace("decoder_pred"   ,   "decoder.decoder_pred"  )
					if "norm.weight" in name and "decoder" not in name and "fc" not in name:
										__lowerCAmelCase            = name.replace("norm.weight"   ,   "videomae.layernorm.weight"  )
					if "norm.bias" in name and "decoder" not in name and "fc" not in name:
										__lowerCAmelCase            = name.replace("norm.bias"   ,   "videomae.layernorm.bias"  )
					if "head" in name and "decoder" not in name:
										__lowerCAmelCase            = name.replace("head"   ,   "classifier"  )

					return name







def 			_a   (		SCREAMING_SNAKE_CASE_							:       Any   ,   SCREAMING_SNAKE_CASE_							:       Tuple  ):
					for key in orig_state_dict.copy().keys():
										__lowerCAmelCase            = orig_state_dict.pop(SCREAMING_SNAKE_CASE_  )

										if key.startswith("encoder."  ):
															__lowerCAmelCase            = key.replace("encoder."   ,   ""  )

										if "qkv" in key:
															__lowerCAmelCase            = key.split("."  )
															if key.startswith("decoder.blocks"  ):
																				__lowerCAmelCase            = config.decoder_hidden_size
																				__lowerCAmelCase            = int(key_split[2]  )
																				__lowerCAmelCase            = "decoder.decoder_layers."
																				if "weight" in key:
																									__lowerCAmelCase            = val[:dim, :]
																									__lowerCAmelCase            = val[dim : dim * 2, :]
																									__lowerCAmelCase            = val[-dim:, :]
															else:
																				__lowerCAmelCase            = config.hidden_size
																				__lowerCAmelCase            = int(key_split[1]  )
																				__lowerCAmelCase            = "videomae.encoder.layer."
																				if "weight" in key:
																									__lowerCAmelCase            = val[:dim, :]
																									__lowerCAmelCase            = val[dim : dim * 2, :]
																									__lowerCAmelCase            = val[-dim:, :]
										else:
															__lowerCAmelCase            = val

					return orig_state_dict







def 			_a   (		):
					__lowerCAmelCase            = hf_hub_download(
					    repo_id="hf-internal-testing/spaghetti-video"   ,   filename="eating_spaghetti.npy"   ,   repo_type="dataset"  )
					__lowerCAmelCase            = np.load(SCREAMING_SNAKE_CASE_  )
					return list(SCREAMING_SNAKE_CASE_  )







def 			_a   (		SCREAMING_SNAKE_CASE_							:       Tuple   ,   SCREAMING_SNAKE_CASE_							:       Optional[int]   ,   SCREAMING_SNAKE_CASE_							:       Union[str, Any]   ,   SCREAMING_SNAKE_CASE_							:       List[str]  ):
					__lowerCAmelCase            = get_videomae_config(SCREAMING_SNAKE_CASE_  )

					if "finetuned" in model_name:
										__lowerCAmelCase            = VideoMAEForVideoClassification(SCREAMING_SNAKE_CASE_  )
					else:
										__lowerCAmelCase            = VideoMAEForPreTraining(SCREAMING_SNAKE_CASE_  )

					# download original checkpoint, hosted on Google Drive
					__lowerCAmelCase            = "pytorch_model.bin"
					gdown.cached_download(SCREAMING_SNAKE_CASE_   ,   SCREAMING_SNAKE_CASE_   ,   quiet=SCREAMING_SNAKE_CASE_  )
					__lowerCAmelCase            = torch.load(SCREAMING_SNAKE_CASE_   ,   map_location="cpu"  )
					if "model" in files:
										__lowerCAmelCase            = files["model"]
					else:
										__lowerCAmelCase            = files["module"]
					__lowerCAmelCase            = convert_state_dict(SCREAMING_SNAKE_CASE_   ,   SCREAMING_SNAKE_CASE_  )

					model.load_state_dict(SCREAMING_SNAKE_CASE_  )
					model.eval()

					# verify model on basic input
					__lowerCAmelCase            = VideoMAEImageProcessor(image_mean=[0.5, 0.5, 0.5]   ,   image_std=[0.5, 0.5, 0.5]  )
					__lowerCAmelCase            = prepare_video()
					__lowerCAmelCase            = image_processor(SCREAMING_SNAKE_CASE_   ,   return_tensors="pt"  )

					if "finetuned" not in model_name:
										__lowerCAmelCase            = hf_hub_download(repo_id="hf-internal-testing/bool-masked-pos"   ,   filename="bool_masked_pos.pt"  )
										__lowerCAmelCase            = torch.load(SCREAMING_SNAKE_CASE_  )

					__lowerCAmelCase            = model(**SCREAMING_SNAKE_CASE_  )
					__lowerCAmelCase            = outputs.logits

					__lowerCAmelCase            = [
					    "videomae-small-finetuned-kinetics",
					    "videomae-small-finetuned-ssv2",
					    # Kinetics-400 checkpoints (short = pretrained only for 800 epochs instead of 1600)
					    "videomae-base-short",
					    "videomae-base-short-finetuned-kinetics",
					    "videomae-base",
					    "videomae-base-finetuned-kinetics",
					    "videomae-large",
					    "videomae-large-finetuned-kinetics",
					    "videomae-huge-finetuned-kinetics",
					    # Something-Something-v2 checkpoints (short = pretrained only for 800 epochs instead of 2400)
					    "videomae-base-short-ssv2",
					    "videomae-base-short-finetuned-ssv2",
					    "videomae-base-ssv2",
					    "videomae-base-finetuned-ssv2",
					]

					# NOTE: logits were tested with image_mean and image_std equal to [0.5, 0.5, 0.5] and [0.5, 0.5, 0.5]
					if model_name == "videomae-small-finetuned-kinetics":
										__lowerCAmelCase            = torch.Size([1, 4_00]  )
										__lowerCAmelCase            = torch.tensor([-0.92_91, -0.40_61, -0.93_07]  )
					elif model_name == "videomae-small-finetuned-ssv2":
										__lowerCAmelCase            = torch.Size([1, 1_74]  )
										__lowerCAmelCase            = torch.tensor([0.26_71, -0.46_89, -0.82_35]  )
					elif model_name == "videomae-base":
										__lowerCAmelCase            = torch.Size([1, 14_08, 15_36]  )
										__lowerCAmelCase            = torch.tensor([[0.77_39, 0.79_68, 0.70_89], [0.67_01, 0.74_87, 0.62_09], [0.42_87, 0.51_58, 0.47_73]]  )
					elif model_name == "videomae-base-short":
										__lowerCAmelCase            = torch.Size([1, 14_08, 15_36]  )
										__lowerCAmelCase            = torch.tensor([[0.79_94, 0.96_12, 0.85_08], [0.74_01, 0.89_58, 0.83_02], [0.58_62, 0.74_68, 0.73_25]]  )
										# we verified the loss both for normalized and unnormalized targets for this one
										__lowerCAmelCase            = torch.tensor([0.51_42]  ) if config.norm_pix_loss else torch.tensor([0.64_69]  )
					elif model_name == "videomae-large":
										__lowerCAmelCase            = torch.Size([1, 14_08, 15_36]  )
										__lowerCAmelCase            = torch.tensor([[0.71_49, 0.79_97, 0.69_66], [0.67_68, 0.78_69, 0.69_48], [0.51_39, 0.62_21, 0.56_05]]  )
					elif model_name == "videomae-large-finetuned-kinetics":
										__lowerCAmelCase            = torch.Size([1, 4_00]  )
										__lowerCAmelCase            = torch.tensor([0.07_71, 0.00_11, -0.36_25]  )
					elif model_name == "videomae-huge-finetuned-kinetics":
										__lowerCAmelCase            = torch.Size([1, 4_00]  )
										__lowerCAmelCase            = torch.tensor([0.24_33, 0.16_32, -0.48_94]  )
					elif model_name == "videomae-base-short-finetuned-kinetics":
										__lowerCAmelCase            = torch.Size([1, 4_00]  )
										__lowerCAmelCase            = torch.tensor([0.65_88, 0.09_90, -0.24_93]  )
					elif model_name == "videomae-base-finetuned-kinetics":
										__lowerCAmelCase            = torch.Size([1, 4_00]  )
										__lowerCAmelCase            = torch.tensor([0.36_69, -0.06_88, -0.24_21]  )
					elif model_name == "videomae-base-short-ssv2":
										__lowerCAmelCase            = torch.Size([1, 14_08, 15_36]  )
										__lowerCAmelCase            = torch.tensor([[0.47_12, 0.52_96, 0.57_86], [0.22_78, 0.27_29, 0.40_26], [0.03_52, 0.07_30, 0.25_06]]  )
					elif model_name == "videomae-base-short-finetuned-ssv2":
										__lowerCAmelCase            = torch.Size([1, 1_74]  )
										__lowerCAmelCase            = torch.tensor([-0.05_37, -0.15_39, -0.32_66]  )
					elif model_name == "videomae-base-ssv2":
										__lowerCAmelCase            = torch.Size([1, 14_08, 15_36]  )
										__lowerCAmelCase            = torch.tensor([[0.81_31, 0.87_27, 0.85_46], [0.73_66, 0.93_77, 0.88_70], [0.59_35, 0.88_74, 0.85_64]]  )
					elif model_name == "videomae-base-finetuned-ssv2":
										__lowerCAmelCase            = torch.Size([1, 1_74]  )
										__lowerCAmelCase            = torch.tensor([0.19_61, -0.83_37, -0.63_89]  )
					else:
										raise ValueError(F"""Model name not supported. Should be one of {model_names}"""  )

					# verify logits
					assert logits.shape == expected_shape
					if "finetuned" in model_name:
										assert torch.allclose(logits[0, :3]   ,   SCREAMING_SNAKE_CASE_   ,   atol=1E-4  )
					else:
										print("Logits:"   ,   logits[0, :3, :3]  )
										assert torch.allclose(logits[0, :3, :3]   ,   SCREAMING_SNAKE_CASE_   ,   atol=1E-4  )
					print("Logits ok!"  )

					# verify loss, if applicable
					if model_name == "videomae-base-short":
										__lowerCAmelCase            = outputs.loss
										assert torch.allclose(SCREAMING_SNAKE_CASE_   ,   SCREAMING_SNAKE_CASE_   ,   atol=1E-4  )
										print("Loss ok!"  )

					if pytorch_dump_folder_path is not None:
										print(F"""Saving model and image processor to {pytorch_dump_folder_path}"""  )
										image_processor.save_pretrained(SCREAMING_SNAKE_CASE_  )
										model.save_pretrained(SCREAMING_SNAKE_CASE_  )

					if push_to_hub:
										print("Pushing to the hub..."  )
										model.push_to_hub(SCREAMING_SNAKE_CASE_   ,   organization="nielsr"  )


if __name__ == "__main__":
					UpperCamelCase__					    =    argparse.ArgumentParser()
					# Required parameters
					parser.add_argument(
					    """--checkpoint_url""",
					    default="""https://drive.google.com/u/1/uc?id=1tEhLyskjb755TJ65ptsrafUG2llSwQE1&amp;export=download&amp;confirm=t&amp;uuid=aa3276eb-fb7e-482a-adec-dc7171df14c4""",
					    type=str,
					    help=(
					        """URL of the original PyTorch checkpoint (on Google Drive) you'd like to convert. Should be a direct"""
					        """ download link."""
					    ),
					)
					parser.add_argument(
					    """--pytorch_dump_folder_path""",
					    default="""/Users/nielsrogge/Documents/VideoMAE/Test""",
					    type=str,
					    help="""Path to the output PyTorch model directory.""",
					)
					parser.add_argument("""--model_name""", default="""videomae-base""", type=str, help="""Name of the model.""")
					parser.add_argument(
					    """--push_to_hub""", action="""store_true""", help="""Whether or not to push the converted model to the 🤗 hub."""
					)

					UpperCamelCase__					    =    parser.parse_args()
					convert_videomae_checkpoint(args.checkpoint_url, args.pytorch_dump_folder_path, args.model_name, args.push_to_hub)




