





from typing import List, Optional, Union

import torch

from ...models import UNetaDConditionModel, VQModel
from ...pipelines import DiffusionPipeline
from ...pipelines.pipeline_utils import ImagePipelineOutput
from ...schedulers import DDPMScheduler
from ...utils import (
    is_accelerate_available,
    is_accelerate_version,
    logging,
    randn_tensor,
    replace_example_docstring,
)


UpperCAmelCase__						:			List[str]									=logging.get_logger(__name__)  # pylint: disable=invalid-name

UpperCAmelCase__						:			Union[str, Any]									="\n    Examples:\n        ```py\n        >>> from diffusers import KandinskyV22Pipeline, KandinskyV22PriorPipeline\n        >>> import torch\n\n        >>> pipe_prior = KandinskyV22PriorPipeline.from_pretrained(\"kandinsky-community/kandinsky-2-2-prior\")\n        >>> pipe_prior.to(\"cuda\")\n        >>> prompt = \"red cat, 4k photo\"\n        >>> out = pipe_prior(prompt)\n        >>> image_emb = out.image_embeds\n        >>> zero_image_emb = out.negative_image_embeds\n        >>> pipe = KandinskyV22Pipeline.from_pretrained(\"kandinsky-community/kandinsky-2-2-decoder\")\n        >>> pipe.to(\"cuda\")\n        >>> image = pipe(\n        ...     image_embeds=image_emb,\n        ...     negative_image_embeds=zero_image_emb,\n        ...     height=768,\n        ...     width=768,\n        ...     num_inference_steps=50,\n        ... ).images\n        >>> image[0].save(\"cat.png\")\n        ```\n"





def      _lowercase					(						_UpperCAmelCase			,      _UpperCAmelCase			,      _UpperCAmelCase=8				)						->  Tuple:
				lowerCamelCase	 =height // scale_factor**2
				if height % scale_factor**2 != 0:
								new_height += 1
				lowerCamelCase	 =width // scale_factor**2
				if width % scale_factor**2 != 0:
								new_width += 1
				return new_height * scale_factor, new_width * scale_factor




class       __A     (						__UpperCAmelCase    ):


				def __init__(		self   , UpperCAmelCase_   , UpperCAmelCase_   , UpperCAmelCase_   , ):
								super().__init__()

								self.register_modules(
								    unet=lowerCAmelCase_   , scheduler=lowerCAmelCase_   , movq=lowerCAmelCase_   , )
								lowerCamelCase	 =2 ** (len(self.movq.config.block_out_channels							) - 1)


				def _snake_case       (		self   , UpperCAmelCase_   , UpperCAmelCase_   , UpperCAmelCase_   , UpperCAmelCase_   , UpperCAmelCase_   , UpperCAmelCase_							):







								if latents is None:
												lowerCamelCase	 =randn_tensor(lowerCAmelCase_   , generator=lowerCAmelCase_   , device=lowerCAmelCase_   , dtype=lowerCAmelCase_							)
								else:
												if latents.shape != shape:
																raise ValueError(f"""Unexpected latents shape, got {latents.shape}, expected {shape}"""							)
												lowerCamelCase	 =latents.to(lowerCAmelCase_							)

								lowerCamelCase	 =latents * scheduler.init_noise_sigma
								return latents


				def _snake_case       (		self   , UpperCAmelCase_=0							):







								if is_accelerate_available():
												from accelerate import cpu_offload
								else:
												raise ImportError("""Please install accelerate via `pip install accelerate`"""							)

								lowerCamelCase	 =torch.device(f"""cuda:{gpu_id}"""							)

								lowerCamelCase	 =[
								    self.unet,
								    self.movq,
								]
								for cpu_offloaded_model in models:
												if cpu_offloaded_model is not None:
																cpu_offload(lowerCAmelCase_   , lowerCAmelCase_							)


				def _snake_case       (		self   , UpperCAmelCase_=0							):







								if is_accelerate_available() and is_accelerate_version(""">="""   , """0.17.0.dev0"""							):
												from accelerate import cpu_offload_with_hook
								else:
												raise ImportError("""`enable_model_cpu_offload` requires `accelerate v0.17.0` or higher."""							)

								lowerCamelCase	 =torch.device(f"""cuda:{gpu_id}"""							)

								if self.device.type != "cpu":
												self.to("""cpu"""   , silence_dtype_warnings=lowerCAmelCase_							)
												torch.cuda.empty_cache()  # otherwise we don't see the memory savings (but they probably exist)

								lowerCamelCase	 =None
								for cpu_offloaded_model in [self.unet, self.movq]:
												lowerCamelCase ,  lowerCamelCase	 =cpu_offload_with_hook(lowerCAmelCase_   , lowerCAmelCase_   , prev_module_hook=lowerCAmelCase_							)

								# We'll offload the last model manually.
								lowerCamelCase	 =hook


				@property
				# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline._execution_device
				def _snake_case       (		self							):







								if not hasattr(self.unet   , """_hf_hook"""							):
												return self.device
								for module in self.unet.modules():
												if (
												    hasattr(lowerCAmelCase_   , """_hf_hook"""							)
												    and hasattr(module._hf_hook   , """execution_device"""							)
												    and module._hf_hook.execution_device is not None
												):
																return torch.device(module._hf_hook.execution_device							)
								return self.device


				@torch.no_grad()
				@replace_example_docstring(lowerCAmelCase_							)
				def __call__(		self   , UpperCAmelCase_   , UpperCAmelCase_   , UpperCAmelCase_ = 512   , UpperCAmelCase_ = 512   , UpperCAmelCase_ = 100   , UpperCAmelCase_ = 4.0   , UpperCAmelCase_ = 1   , UpperCAmelCase_ = None   , UpperCAmelCase_ = None   , UpperCAmelCase_ = "pil"   , UpperCAmelCase_ = True   , ):
								lowerCamelCase	 =self._execution_device

								lowerCamelCase	 =guidance_scale > 1.0

								if isinstance(lowerCAmelCase_   , lowerCAmelCase_							):
												lowerCamelCase	 =torch.cat(lowerCAmelCase_   , dim=0							)
								lowerCamelCase	 =image_embeds.shape[0] * num_images_per_prompt
								if isinstance(lowerCAmelCase_   , lowerCAmelCase_							):
												lowerCamelCase	 =torch.cat(lowerCAmelCase_   , dim=0							)

								if do_classifier_free_guidance:
												lowerCamelCase	 =image_embeds.repeat_interleave(lowerCAmelCase_   , dim=0							)
												lowerCamelCase	 =negative_image_embeds.repeat_interleave(lowerCAmelCase_   , dim=0							)

								lowerCamelCase	 =torch.cat([negative_image_embeds, image_embeds]   , dim=0							).to(dtype=self.unet.dtype   , device=lowerCAmelCase_							)

								self.scheduler.set_timesteps(lowerCAmelCase_   , device=lowerCAmelCase_							)
								lowerCamelCase	 =self.scheduler.timesteps

								lowerCamelCase	 =self.unet.config.in_channels

								lowerCamelCase ,  lowerCamelCase	 =downscale_height_and_width(lowerCAmelCase_   , lowerCAmelCase_   , self.movq_scale_factor							)

								# create initial latent
								lowerCamelCase	 =self.prepare_latents(
								    (batch_size, num_channels_latents, height, width)   , image_embeds.dtype   , lowerCAmelCase_   , lowerCAmelCase_   , lowerCAmelCase_   , self.scheduler   , )

								for i, t in enumerate(self.progress_bar(lowerCAmelCase_							)							):
												# expand the latents if we are doing classifier free guidance
												lowerCamelCase	 =torch.cat([latents] * 2							) if do_classifier_free_guidance else latents

												lowerCamelCase	 ={"""image_embeds""": image_embeds}
												lowerCamelCase	 =self.unet(
												    sample=lowerCAmelCase_   , timestep=lowerCAmelCase_   , encoder_hidden_states=lowerCAmelCase_   , added_cond_kwargs=lowerCAmelCase_   , return_dict=lowerCAmelCase_   , )[0]

												if do_classifier_free_guidance:
																lowerCamelCase ,  lowerCamelCase	 =noise_pred.split(latents.shape[1]   , dim=1							)
																lowerCamelCase ,  lowerCamelCase	 =noise_pred.chunk(2							)
																lowerCamelCase ,  lowerCamelCase	 =variance_pred.chunk(2							)
																lowerCamelCase	 =noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)
																lowerCamelCase	 =torch.cat([noise_pred, variance_pred_text]   , dim=1							)

												if not (
												    hasattr(self.scheduler.config   , """variance_type"""							)
												    and self.scheduler.config.variance_type in ["learned", "learned_range"]
												):
																lowerCamelCase ,  lowerCamelCase	 =noise_pred.split(latents.shape[1]   , dim=1							)

												# compute the previous noisy sample x_t -> x_t-1
												lowerCamelCase	 =self.scheduler.step(
												    lowerCAmelCase_   , lowerCAmelCase_   , lowerCAmelCase_   , generator=lowerCAmelCase_   , )[0]
								# post-processing
								lowerCamelCase	 =self.movq.decode(lowerCAmelCase_   , force_not_quantize=lowerCAmelCase_							)["""sample"""]

								if output_type not in ["pt", "np", "pil"]:
												raise ValueError(f"""Only the output types `pt`, `pil` and `np` are supported not output_type={output_type}"""							)

								if output_type in ["np", "pil"]:
												lowerCamelCase	 =image * 0.5 + 0.5
												lowerCamelCase	 =image.clamp(0   , 1							)
												lowerCamelCase	 =image.cpu().permute(0   , 2   , 3   , 1							).float().numpy()

								if output_type == "pil":
												lowerCamelCase	 =self.numpy_to_pil(lowerCAmelCase_							)

								if not return_dict:
												return (image,)

								return ImagePipelineOutput(images=lowerCAmelCase_							)













import argparse
from pathlib import Path
from typing import Dict, OrderedDict, Tuple

import torch
from audiocraft.models import MusicGen

from transformers import (
    AutoFeatureExtractor,
    AutoTokenizer,
    EncodecModel,
    MusicgenDecoderConfig,
    MusicgenForConditionalGeneration,
    MusicgenProcessor,
    TaEncoderModel,
)
from transformers.models.musicgen.modeling_musicgen import MusicgenForCausalLM
from transformers.utils import logging


logging.set_verbosity_info()
UpperCAmelCase__						:			Tuple									=logging.get_logger(__name__)


UpperCAmelCase__						:			str									=['''model.decoder.embed_positions.weights''']





def      _lowercase					(						_UpperCAmelCase				)						->  List[Any]:
				if "emb" in name:
								lowerCamelCase	 =name.replace("""emb"""			,      """model.decoder.embed_tokens"""				)
				if "transformer" in name:
								lowerCamelCase	 =name.replace("""transformer"""			,      """model.decoder"""				)
				if "cross_attention" in name:
								lowerCamelCase	 =name.replace("""cross_attention"""			,      """encoder_attn"""				)
				if "linear1" in name:
								lowerCamelCase	 =name.replace("""linear1"""			,      """fc1"""				)
				if "linear2" in name:
								lowerCamelCase	 =name.replace("""linear2"""			,      """fc2"""				)
				if "norm1" in name:
								lowerCamelCase	 =name.replace("""norm1"""			,      """self_attn_layer_norm"""				)
				if "norm_cross" in name:
								lowerCamelCase	 =name.replace("""norm_cross"""			,      """encoder_attn_layer_norm"""				)
				if "norm2" in name:
								lowerCamelCase	 =name.replace("""norm2"""			,      """final_layer_norm"""				)
				if "out_norm" in name:
								lowerCamelCase	 =name.replace("""out_norm"""			,      """model.decoder.layer_norm"""				)
				if "linears" in name:
								lowerCamelCase	 =name.replace("""linears"""			,      """lm_heads"""				)
				if "condition_provider.conditioners.description.output_proj" in name:
								lowerCamelCase	 =name.replace("""condition_provider.conditioners.description.output_proj"""			,      """enc_to_dec_proj"""				)
				return name





def      _lowercase					(						_UpperCAmelCase			,      _UpperCAmelCase				)						->  Tuple[Dict, Dict]:
				lowerCamelCase	 =list(state_dict.keys()				)
				lowerCamelCase	 ={}
				for key in keys:
								lowerCamelCase	 =state_dict.pop(_UpperCAmelCase				)
								lowerCamelCase	 =rename_keys(_UpperCAmelCase				)
								if "in_proj_weight" in key:
												# split fused qkv proj
												lowerCamelCase	 =val[:hidden_size, :]
												lowerCamelCase	 =val[hidden_size : 2 * hidden_size, :]
												lowerCamelCase	 =val[-hidden_size:, :]
								elif "enc_to_dec_proj" in key:
												lowerCamelCase	 =val
								else:
												lowerCamelCase	 =val
				return state_dict, enc_dec_proj_state_dict





def      _lowercase					(						_UpperCAmelCase				)						->  MusicgenDecoderConfig:
				if checkpoint == "small":
								# default config values
								lowerCamelCase	 =10_24
								lowerCamelCase	 =24
								lowerCamelCase	 =16
				elif checkpoint == "medium":
								lowerCamelCase	 =15_36
								lowerCamelCase	 =48
								lowerCamelCase	 =24
				elif checkpoint == "large":
								lowerCamelCase	 =20_48
								lowerCamelCase	 =48
								lowerCamelCase	 =32
				else:
								raise ValueError(F"""Checkpoint should be one of `['small', 'medium', 'large']`, got {checkpoint}."""				)
				lowerCamelCase	 =MusicgenDecoderConfig(
				    hidden_size=_UpperCAmelCase			,      ffn_dim=hidden_size * 4			,      num_hidden_layers=_UpperCAmelCase			,      num_attention_heads=_UpperCAmelCase			,      )
				return config





@torch.no_grad()
def      _lowercase					(						_UpperCAmelCase			,      _UpperCAmelCase=None			,      _UpperCAmelCase=None			,      _UpperCAmelCase="cpu"				)						->  Dict:
				lowerCamelCase	 =MusicGen.get_pretrained(_UpperCAmelCase			,      device=_UpperCAmelCase				)
				lowerCamelCase	 =decoder_config_from_checkpoint(_UpperCAmelCase				)

				lowerCamelCase	 =fairseq_model.lm.state_dict()
				lowerCamelCase ,  lowerCamelCase	 =rename_state_dict(
				    _UpperCAmelCase			,      hidden_size=decoder_config.hidden_size				)

				lowerCamelCase	 =TaEncoderModel.from_pretrained("""t5-base"""				)
				lowerCamelCase	 =EncodecModel.from_pretrained("""facebook/encodec_32khz"""				)
				lowerCamelCase	 =MusicgenForCausalLM(_UpperCAmelCase				).eval()

				# load all decoder weights - expect that we'll be missing embeddings and enc-dec projection
				lowerCamelCase ,  lowerCamelCase	 =decoder.load_state_dict(_UpperCAmelCase			,      strict=_UpperCAmelCase				)

				for key in missing_keys.copy():
								if key.startswith(("""text_encoder""", """audio_encoder""")				) or key in EXPECTED_MISSING_KEYS:
												missing_keys.remove(_UpperCAmelCase				)

				if len(_UpperCAmelCase				) > 0:
								raise ValueError(F"""Missing key(s) in state_dict: {missing_keys}"""				)

				if len(_UpperCAmelCase				) > 0:
								raise ValueError(F"""Unexpected key(s) in state_dict: {unexpected_keys}"""				)

				# init the composite model
				lowerCamelCase	 =MusicgenForConditionalGeneration(text_encoder=_UpperCAmelCase			,      audio_encoder=_UpperCAmelCase			,      decoder=_UpperCAmelCase				)

				# load the pre-trained enc-dec projection (from the decoder state dict)
				model.enc_to_dec_proj.load_state_dict(_UpperCAmelCase				)

				# check we can do a forward pass
				lowerCamelCase	 =torch.arange(0			,      8			,      dtype=torch.long				).reshape(2			,      -1				)
				lowerCamelCase	 =input_ids.reshape(2 * 4			,      -1				)

				with torch.no_grad():
								lowerCamelCase	 =model(input_ids=_UpperCAmelCase			,      decoder_input_ids=_UpperCAmelCase				).logits

				if logits.shape != (8, 1, 20_48):
								raise ValueError("""Incorrect shape for logits"""				)

				# now construct the processor
				lowerCamelCase	 =AutoTokenizer.from_pretrained("""t5-base"""				)
				lowerCamelCase	 =AutoFeatureExtractor.from_pretrained("""facebook/encodec_32khz"""			,      padding_side="""left"""				)

				lowerCamelCase	 =MusicgenProcessor(feature_extractor=_UpperCAmelCase			,      tokenizer=_UpperCAmelCase				)

				# set the appropriate bos/pad token ids
				lowerCamelCase	 =20_48
				lowerCamelCase	 =20_48

				# set other default generation config params
				lowerCamelCase	 =int(30 * audio_encoder.config.frame_rate				)
				lowerCamelCase	 =True
				lowerCamelCase	 =3.0

				if pytorch_dump_folder is not None:
								Path(_UpperCAmelCase				).mkdir(exist_ok=_UpperCAmelCase				)
								logger.info(F"""Saving model {checkpoint} to {pytorch_dump_folder}"""				)
								model.save_pretrained(_UpperCAmelCase				)
								processor.save_pretrained(_UpperCAmelCase				)

				if repo_id:
								logger.info(F"""Pushing model {checkpoint} to {repo_id}"""				)
								model.push_to_hub(_UpperCAmelCase				)
								processor.push_to_hub(_UpperCAmelCase				)


if __name__ == "__main__":
							UpperCAmelCase__						:			List[Any]									=argparse.ArgumentParser()
							# Required parameters
							parser.add_argument(
							    '''--checkpoint''',
							    default='''small''',
							    type=str,
							    help='''Checkpoint size of the MusicGen model you\'d like to convert. Can be one of: `[\'small\', \'medium\', \'large\']`.''',
							)
							parser.add_argument(
							    '''--pytorch_dump_folder''',
							    required=True,
							    default=None,
							    type=str,
							    help='''Path to the output PyTorch model directory.''',
							)
							parser.add_argument(
							    '''--push_to_hub''', default=None, type=str, help='''Where to upload the converted model on the 🤗 hub.'''
							)
							parser.add_argument(
							    '''--device''', default='''cpu''', type=str, help='''Torch device to run the conversion, either cpu or cuda.'''
							)

							UpperCAmelCase__						:			Optional[Any]									=parser.parse_args()
							convert_musicgen_checkpoint(args.checkpoint, args.pytorch_dump_folder, args.push_to_hub)





