






"""simple docstring"""
import argparse
import hashlib
import os
import urllib
import warnings

import torch
from torch import nn
from tqdm import tqdm

from transformers import WhisperConfig, WhisperForConditionalGeneration


lowerCAmelCase__    						=			{
    'tiny.en': 'https://openaipublic.azureedge.net/main/whisper/models/d3dd57d32accea0b295c96e26691aa14d8822fac7d9d27d5dc00b4ca2826dd03/tiny.en.pt',
    'tiny': 'https://openaipublic.azureedge.net/main/whisper/models/65147644a518d12f04e32d6f3b26facc3f8dd46e5390956a9424a650c0ce22b9/tiny.pt',
    'base.en': 'https://openaipublic.azureedge.net/main/whisper/models/25a8566e1d0c1e2231d1c762132cd20e0f96a85d16145c3a00adf5d1ac670ead/base.en.pt',
    'base': 'https://openaipublic.azureedge.net/main/whisper/models/ed3a0b6b1c0edf879ad9b11b1af5a0e6ab5db9205f891f668f8b0e6c6326e34e/base.pt',
    'small.en': 'https://openaipublic.azureedge.net/main/whisper/models/f953ad0fd29cacd07d5a9eda5624af0f6bcf2258be67c92b79389873d91e0872/small.en.pt',
    'small': 'https://openaipublic.azureedge.net/main/whisper/models/9ecf779972d90ba49c06d968637d720dd632c55bbf19d441fb42bf17a411e794/small.pt',
    'medium.en': 'https://openaipublic.azureedge.net/main/whisper/models/d7440d1dc186f76616474e0ff0b3b6b879abc9d1a4926b7adfa41db2d497ab4f/medium.en.pt',
    'medium': 'https://openaipublic.azureedge.net/main/whisper/models/345ae4da62f9b3d59415adc60127b97c714f32e89e936602e85993674d08dcb1/medium.pt',
    'large': 'https://openaipublic.azureedge.net/main/whisper/models/e4b87e7e0bf463eb8e6956e646f1e277e901512310def2c24bf0e11bd3c28e9a/large.pt',
    'large-v2': 'https://openaipublic.azureedge.net/main/whisper/models/81f7c96c852ee8fc832187b0132e569d6c3065a3252ed18e56effd0b6a73e524/large-v2.pt',
}





def 				lowercase__		(					lowerCamelCase       ):
				_SCREAMING_SNAKE_CASE   :		str       		=    ['layers', 'blocks']
				for k in ignore_keys:
								state_dict.pop(lowerCamelCase,						lowerCamelCase       )


lowerCAmelCase__    						=			{
    'blocks': 'layers',
    'mlp.0': 'fc1',
    'mlp.2': 'fc2',
    'mlp_ln': 'final_layer_norm',
    '.attn.query': '.self_attn.q_proj',
    '.attn.key': '.self_attn.k_proj',
    '.attn.value': '.self_attn.v_proj',
    '.attn_ln': '.self_attn_layer_norm',
    '.attn.out': '.self_attn.out_proj',
    '.cross_attn.query': '.encoder_attn.q_proj',
    '.cross_attn.key': '.encoder_attn.k_proj',
    '.cross_attn.value': '.encoder_attn.v_proj',
    '.cross_attn_ln': '.encoder_attn_layer_norm',
    '.cross_attn.out': '.encoder_attn.out_proj',
    'decoder.ln.': 'decoder.layer_norm.',
    'encoder.ln.': 'encoder.layer_norm.',
    'token_embedding': 'embed_tokens',
    'encoder.positional_embedding': 'encoder.embed_positions.weight',
    'decoder.positional_embedding': 'decoder.embed_positions.weight',
    'ln_post': 'layer_norm',
}





def 				lowercase__		(					lowerCamelCase       ):
				_SCREAMING_SNAKE_CASE   :		Union[str, Any]       		=    list(s_dict.keys()       )
				for key in keys:
								_SCREAMING_SNAKE_CASE   :		Any       		=    key
								for k, v in WHISPER_MAPPING.items():
												if k in key:
																_SCREAMING_SNAKE_CASE   :		List[str]       		=    new_key.replace(lowerCamelCase,						lowerCamelCase       )

								print(f"""{key} -> {new_key}"""       )

								_SCREAMING_SNAKE_CASE   :		str       		=    s_dict.pop(lowerCamelCase       )
				return s_dict





def 				lowercase__		(					lowerCamelCase       ):
				_SCREAMING_SNAKE_CASE		,						_SCREAMING_SNAKE_CASE   :		Optional[int]       		=    emb.weight.shape
				_SCREAMING_SNAKE_CASE   :		Dict       		=    nn.Linear(lowerCamelCase,						lowerCamelCase,						bias=lowerCamelCase       )
				_SCREAMING_SNAKE_CASE   :		List[str]       		=    emb.weight.data
				return lin_layer





def 				lowercase__		(					lowerCamelCase,						lowerCamelCase       ):
				os.makedirs(lowerCamelCase,						exist_ok=lowerCamelCase       )
				_SCREAMING_SNAKE_CASE   :		List[str]       		=    os.path.basename(lowerCamelCase       )

				_SCREAMING_SNAKE_CASE   :		Tuple       		=    url.split('/'       )[-2]
				_SCREAMING_SNAKE_CASE   :		Dict       		=    os.path.join(lowerCamelCase,						lowerCamelCase       )

				if os.path.exists(lowerCamelCase       ) and not os.path.isfile(lowerCamelCase       ):
								raise RuntimeError(f"""{download_target} exists and is not a regular file"""       )

				if os.path.isfile(lowerCamelCase       ):
								_SCREAMING_SNAKE_CASE   :		Optional[int]       		=    open(lowerCamelCase,						'rb'       ).read()
								if hashlib.shaaaa(lowerCamelCase       ).hexdigest() == expected_shaaaa:
												return model_bytes
								else:
												warnings.warn(f"""{download_target} exists, but the SHA256 checksum does not match; re-downloading the file"""       )

				with urllib.request.urlopen(lowerCamelCase       ) as source, open(lowerCamelCase,						'wb'       ) as output:
								with tqdm(
								    total=int(source.info().get('Content-Length'       )       ),						ncols=80,						unit='iB',						unit_scale=lowerCamelCase,						unit_divisor=1_024       ) as loop:
												while True:
																_SCREAMING_SNAKE_CASE   :		List[str]       		=    source.read(8_192       )
																if not buffer:
																				break

																output.write(lowerCamelCase       )
																loop.update(len(lowerCamelCase       )       )

				_SCREAMING_SNAKE_CASE   :		List[str]       		=    open(lowerCamelCase,						'rb'       ).read()
				if hashlib.shaaaa(lowerCamelCase       ).hexdigest() != expected_shaaaa:
								raise RuntimeError(
								    'Model has been downloaded but the SHA256 checksum does not not match. Please retry loading the model.'       )

				return model_bytes





def 				lowercase__		(					lowerCamelCase,						lowerCamelCase       ):
				if ".pt" not in checkpoint_path:
								_SCREAMING_SNAKE_CASE   :		Union[str, Any]       		=    _download(_MODELS[checkpoint_path]       )
				else:
								_SCREAMING_SNAKE_CASE   :		str       		=    torch.load(lowerCamelCase,						map_location='cpu'       )
				_SCREAMING_SNAKE_CASE   :		Tuple       		=    original_checkpoint['dims']
				_SCREAMING_SNAKE_CASE   :		str       		=    original_checkpoint['model_state_dict']
				_SCREAMING_SNAKE_CASE   :		List[Any]       		=    state_dict['decoder.token_embedding.weight']
				remove_ignore_keys_(lowerCamelCase       )
				rename_keys(lowerCamelCase       )
				_SCREAMING_SNAKE_CASE   :		Dict       		=    True
				_SCREAMING_SNAKE_CASE   :		Dict       		=    state_dict['decoder.layers.0.fc1.weight'].shape[0]

				_SCREAMING_SNAKE_CASE   :		int       		=    WhisperConfig(
				    vocab_size=dimensions['n_vocab'],						encoder_ffn_dim=lowerCamelCase,						decoder_ffn_dim=lowerCamelCase,						num_mel_bins=dimensions['n_mels'],						d_model=dimensions['n_audio_state'],						max_target_positions=dimensions['n_text_ctx'],						encoder_layers=dimensions['n_audio_layer'],						encoder_attention_heads=dimensions['n_audio_head'],						decoder_layers=dimensions['n_text_layer'],						decoder_attention_heads=dimensions['n_text_state'],						max_source_positions=dimensions['n_audio_ctx'],						)

				_SCREAMING_SNAKE_CASE   :		Union[str, Any]       		=    WhisperForConditionalGeneration(lowerCamelCase       )
				_SCREAMING_SNAKE_CASE		,						_SCREAMING_SNAKE_CASE   :		Optional[Any]       		=    model.model.load_state_dict(lowerCamelCase,						strict=lowerCamelCase       )
				if len(lowerCamelCase       ) > 0 and not set(lowerCamelCase       ) <= {
				    "encoder.embed_positions.weights",
				    "decoder.embed_positions.weights",
				}:
								raise ValueError(
								    'Only `encoder.embed_positions.weights` and `decoder.embed_positions.weights`  are allowed to be missing,'
								    f""" but all the following weights are missing {missing}"""       )

				if tie_embeds:
								_SCREAMING_SNAKE_CASE   :		int       		=    make_linear_from_emb(model.model.decoder.embed_tokens       )
				else:
								_SCREAMING_SNAKE_CASE   :		List[Any]       		=    proj_out_weights

				model.save_pretrained(lowerCamelCase       )


if __name__ == "__main__":
			lowerCAmelCase__    						=			argparse.ArgumentParser()
			# # Required parameters
			parser.add_argument('--checkpoint_path', type=str, help='Patht to the downloaded checkpoints')
			parser.add_argument('--pytorch_dump_folder_path', default=None, type=str, help='Path to the output PyTorch model.')
			lowerCAmelCase__    						=			parser.parse_args()

			convert_openai_whisper_to_tfms(args.checkpoint_path, args.pytorch_dump_folder_path)









"""simple docstring"""
import gc
import random
import unittest

import numpy as np
import torch
from PIL import Image
from transformers import CLIPTextConfig, CLIPTextModel, CLIPTokenizer

from diffusers import AutoencoderKL, PNDMScheduler, StableDiffusionInpaintPipeline, UNetaDConditionModel
from diffusers.utils import floats_tensor, load_image, load_numpy, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu, slow

from ..pipeline_params import TEXT_GUIDED_IMAGE_INPAINTING_BATCH_PARAMS, TEXT_GUIDED_IMAGE_INPAINTING_PARAMS
from ..test_pipelines_common import PipelineKarrasSchedulerTesterMixin, PipelineLatentTesterMixin, PipelineTesterMixin


enable_full_determinism()


class _lowerCAmelCase						(    __UpperCAmelCase   ,  __UpperCAmelCase   ,  __UpperCAmelCase   ,  unittest.TestCase       ):
						SCREAMING_SNAKE_CASE_:					Optional[Any] 					=				StableDiffusionInpaintPipeline
						SCREAMING_SNAKE_CASE_:					List[str] 					=				TEXT_GUIDED_IMAGE_INPAINTING_PARAMS
						SCREAMING_SNAKE_CASE_:					Dict 					=				TEXT_GUIDED_IMAGE_INPAINTING_BATCH_PARAMS
						SCREAMING_SNAKE_CASE_:					Any 					=				frozenset(
						    []       )  # TO-DO: update image_params once pipeline is refactored with VaeImageProcessor.preprocess
						SCREAMING_SNAKE_CASE_:					Any 					=				frozenset([]       )







						def 			A      (    self							)		->							Optional[Any]:
										torch.manual_seed(0							)
										_SCREAMING_SNAKE_CASE   :		Tuple       		=    UNetaDConditionModel(
										    block_out_channels=(3_2, 6_4)   ,				layers_per_block=2   ,				sample_size=3_2   ,				in_channels=9   ,				out_channels=4   ,				down_block_types=('DownBlock2D', 'CrossAttnDownBlock2D')   ,				up_block_types=('CrossAttnUpBlock2D', 'UpBlock2D')   ,				cross_attention_dim=3_2   ,				attention_head_dim=(2, 4)   ,				use_linear_projection=lowerCAmelCase_   ,				)
										_SCREAMING_SNAKE_CASE   :		List[str]       		=    PNDMScheduler(skip_prk_steps=lowerCAmelCase_							)
										torch.manual_seed(0							)
										_SCREAMING_SNAKE_CASE   :		Dict       		=    AutoencoderKL(
										    block_out_channels=[3_2, 6_4]   ,				in_channels=3   ,				out_channels=3   ,				down_block_types=['DownEncoderBlock2D', 'DownEncoderBlock2D']   ,				up_block_types=['UpDecoderBlock2D', 'UpDecoderBlock2D']   ,				latent_channels=4   ,				sample_size=1_2_8   ,				)
										torch.manual_seed(0							)
										_SCREAMING_SNAKE_CASE   :		Optional[Any]       		=    CLIPTextConfig(
										    bos_token_id=0   ,				eos_token_id=2   ,				hidden_size=3_2   ,				intermediate_size=3_7   ,				layer_norm_eps=1e-05   ,				num_attention_heads=4   ,				num_hidden_layers=5   ,				pad_token_id=1   ,				vocab_size=1_0_0_0   ,				hidden_act='gelu'   ,				projection_dim=5_1_2   ,				)
										_SCREAMING_SNAKE_CASE   :		Dict       		=    CLIPTextModel(lowerCAmelCase_							)
										_SCREAMING_SNAKE_CASE   :		Tuple       		=    CLIPTokenizer.from_pretrained('hf-internal-testing/tiny-random-clip'							)

										_SCREAMING_SNAKE_CASE   :		Tuple       		=    {
										    'unet': unet,
										    'scheduler': scheduler,
										    'vae': vae,
										    'text_encoder': text_encoder,
										    'tokenizer': tokenizer,
										    'safety_checker': None,
										    'feature_extractor': None,
										}
										return components







						def 			A      (    self   ,				lowerCAmelCase_   ,				lowerCAmelCase_=0							)		->							Optional[int]:
										# TODO: use tensor inputs instead of PIL, this is here just to leave the old expected_slices untouched
										_SCREAMING_SNAKE_CASE   :		Union[str, Any]       		=    floats_tensor((1, 3, 3_2, 3_2)   ,				rng=random.Random(lowerCAmelCase_							)							).to(lowerCAmelCase_							)
										_SCREAMING_SNAKE_CASE   :		Tuple       		=    image.cpu().permute(0   ,				2   ,				3   ,				1							)[0]
										_SCREAMING_SNAKE_CASE   :		Any       		=    Image.fromarray(np.uinta(lowerCAmelCase_							)							).convert('RGB'							).resize((6_4, 6_4)							)
										_SCREAMING_SNAKE_CASE   :		Optional[Any]       		=    Image.fromarray(np.uinta(image + 4							)							).convert('RGB'							).resize((6_4, 6_4)							)
										if str(lowerCAmelCase_							).startswith('mps'							):
														_SCREAMING_SNAKE_CASE   :		List[str]       		=    torch.manual_seed(lowerCAmelCase_							)
										else:
														_SCREAMING_SNAKE_CASE   :		Tuple       		=    torch.Generator(device=lowerCAmelCase_							).manual_seed(lowerCAmelCase_							)
										_SCREAMING_SNAKE_CASE   :		Optional[Any]       		=    {
										    'prompt': 'A painting of a squirrel eating a burger',
										    'image': init_image,
										    'mask_image': mask_image,
										    'generator': generator,
										    'num_inference_steps': 2,
										    'guidance_scale': 6.0,
										    'output_type': 'numpy',
										}
										return inputs







						def 			A      (    self							)		->							str:
										_SCREAMING_SNAKE_CASE   :		List[str]       		=    'cpu'  # ensure determinism for the device-dependent torch.Generator
										_SCREAMING_SNAKE_CASE   :		str       		=    self.get_dummy_components()
										_SCREAMING_SNAKE_CASE   :		Any       		=    StableDiffusionInpaintPipeline(**lowerCAmelCase_							)
										_SCREAMING_SNAKE_CASE   :		List[str]       		=    sd_pipe.to(lowerCAmelCase_							)
										sd_pipe.set_progress_bar_config(disable=lowerCAmelCase_							)

										_SCREAMING_SNAKE_CASE   :		Optional[int]       		=    self.get_dummy_inputs(lowerCAmelCase_							)
										_SCREAMING_SNAKE_CASE   :		Optional[Any]       		=    sd_pipe(**lowerCAmelCase_							).images
										_SCREAMING_SNAKE_CASE   :		Optional[int]       		=    image[0, -3:, -3:, -1]

										assert image.shape == (1, 6_4, 6_4, 3)
										_SCREAMING_SNAKE_CASE   :		List[str]       		=    np.array([0.4_727, 0.5_735, 0.3_941, 0.5_446, 0.5_926, 0.4_394, 0.5_062, 0.4_654, 0.4_476]							)

										assert np.abs(image_slice.flatten() - expected_slice							).max() < 1e-2







						def 			A      (    self							)		->							str:
										super().test_inference_batch_single_identical(expected_max_diff=3e-3							)







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







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







						def 			A      (    self							)		->							Optional[int]:
										_SCREAMING_SNAKE_CASE   :		List[str]       		=    load_image(
										    'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'
										    '/sd2-inpaint/init_image.png'							)
										_SCREAMING_SNAKE_CASE   :		Any       		=    load_image(
										    'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/sd2-inpaint/mask.png'							)
										_SCREAMING_SNAKE_CASE   :		Dict       		=    load_numpy(
										    'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/sd2-inpaint'
										    '/yellow_cat_sitting_on_a_park_bench.npy'							)

										_SCREAMING_SNAKE_CASE   :		Optional[Any]       		=    'stabilityai/stable-diffusion-2-inpainting'
										_SCREAMING_SNAKE_CASE   :		Optional[int]       		=    StableDiffusionInpaintPipeline.from_pretrained(lowerCAmelCase_   ,				safety_checker=lowerCAmelCase_							)
										pipe.to(lowerCAmelCase_							)
										pipe.set_progress_bar_config(disable=lowerCAmelCase_							)
										pipe.enable_attention_slicing()

										_SCREAMING_SNAKE_CASE   :		List[str]       		=    'Face of a yellow cat, high resolution, sitting on a park bench'

										_SCREAMING_SNAKE_CASE   :		Optional[Any]       		=    torch.manual_seed(0							)
										_SCREAMING_SNAKE_CASE   :		List[Any]       		=    pipe(
										    prompt=lowerCAmelCase_   ,				image=lowerCAmelCase_   ,				mask_image=lowerCAmelCase_   ,				generator=lowerCAmelCase_   ,				output_type='np'   ,				)
										_SCREAMING_SNAKE_CASE   :		int       		=    output.images[0]

										assert image.shape == (5_1_2, 5_1_2, 3)
										assert np.abs(expected_image - image							).max() < 9e-3







						def 			A      (    self							)		->							int:
										_SCREAMING_SNAKE_CASE   :		Dict       		=    load_image(
										    'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'
										    '/sd2-inpaint/init_image.png'							)
										_SCREAMING_SNAKE_CASE   :		str       		=    load_image(
										    'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/sd2-inpaint/mask.png'							)
										_SCREAMING_SNAKE_CASE   :		int       		=    load_numpy(
										    'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/sd2-inpaint'
										    '/yellow_cat_sitting_on_a_park_bench_fp16.npy'							)

										_SCREAMING_SNAKE_CASE   :		Optional[Any]       		=    'stabilityai/stable-diffusion-2-inpainting'
										_SCREAMING_SNAKE_CASE   :		str       		=    StableDiffusionInpaintPipeline.from_pretrained(
										    lowerCAmelCase_   ,				torch_dtype=torch.floataa   ,				safety_checker=lowerCAmelCase_   ,				)
										pipe.to(lowerCAmelCase_							)
										pipe.set_progress_bar_config(disable=lowerCAmelCase_							)
										pipe.enable_attention_slicing()

										_SCREAMING_SNAKE_CASE   :		str       		=    'Face of a yellow cat, high resolution, sitting on a park bench'

										_SCREAMING_SNAKE_CASE   :		Tuple       		=    torch.manual_seed(0							)
										_SCREAMING_SNAKE_CASE   :		List[Any]       		=    pipe(
										    prompt=lowerCAmelCase_   ,				image=lowerCAmelCase_   ,				mask_image=lowerCAmelCase_   ,				generator=lowerCAmelCase_   ,				output_type='np'   ,				)
										_SCREAMING_SNAKE_CASE   :		Optional[int]       		=    output.images[0]

										assert image.shape == (5_1_2, 5_1_2, 3)
										assert np.abs(expected_image - image							).max() < 5e-1







						def 			A      (    self							)		->							List[str]:
										torch.cuda.empty_cache()
										torch.cuda.reset_max_memory_allocated()
										torch.cuda.reset_peak_memory_stats()

										_SCREAMING_SNAKE_CASE   :		str       		=    load_image(
										    'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'
										    '/sd2-inpaint/init_image.png'							)
										_SCREAMING_SNAKE_CASE   :		Optional[Any]       		=    load_image(
										    'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/sd2-inpaint/mask.png'							)

										_SCREAMING_SNAKE_CASE   :		str       		=    'stabilityai/stable-diffusion-2-inpainting'
										_SCREAMING_SNAKE_CASE   :		List[Any]       		=    PNDMScheduler.from_pretrained(lowerCAmelCase_   ,				subfolder='scheduler'							)
										_SCREAMING_SNAKE_CASE   :		Optional[Any]       		=    StableDiffusionInpaintPipeline.from_pretrained(
										    lowerCAmelCase_   ,				safety_checker=lowerCAmelCase_   ,				scheduler=lowerCAmelCase_   ,				torch_dtype=torch.floataa   ,				)
										pipe.to(lowerCAmelCase_							)
										pipe.set_progress_bar_config(disable=lowerCAmelCase_							)
										pipe.enable_attention_slicing(1							)
										pipe.enable_sequential_cpu_offload()

										_SCREAMING_SNAKE_CASE   :		str       		=    'Face of a yellow cat, high resolution, sitting on a park bench'

										_SCREAMING_SNAKE_CASE   :		Union[str, Any]       		=    torch.manual_seed(0							)
										_SCREAMING_SNAKE_CASE   :		Dict       		=    pipe(
										    prompt=lowerCAmelCase_   ,				image=lowerCAmelCase_   ,				mask_image=lowerCAmelCase_   ,				generator=lowerCAmelCase_   ,				num_inference_steps=2   ,				output_type='np'   ,				)

										_SCREAMING_SNAKE_CASE   :		Union[str, Any]       		=    torch.cuda.max_memory_allocated()
										# make sure that less than 2.65 GB is allocated
										assert mem_bytes < 2.65 * 1_0**9

