


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()
SCREAMING_SNAKE_CASE       :		int         =				logging.get_logger(__name__)


SCREAMING_SNAKE_CASE       :		Any         =				["model.decoder.embed_positions.weights"]




def     UpperCamelCase    (      _a		)					->							Tuple:
     '''simple docstring'''





     if "emb" in name:
          lowercase_       :Any     =       name.replace('''emb'''       ,		'''model.decoder.embed_tokens'''		)
     if "transformer" in name:
          lowercase_       :Optional[int]     =       name.replace('''transformer'''       ,		'''model.decoder'''		)
     if "cross_attention" in name:
          lowercase_       :Any     =       name.replace('''cross_attention'''       ,		'''encoder_attn'''		)
     if "linear1" in name:
          lowercase_       :Union[str, Any]     =       name.replace('''linear1'''       ,		'''fc1'''		)
     if "linear2" in name:
          lowercase_       :Optional[int]     =       name.replace('''linear2'''       ,		'''fc2'''		)
     if "norm1" in name:
          lowercase_       :Union[str, Any]     =       name.replace('''norm1'''       ,		'''self_attn_layer_norm'''		)
     if "norm_cross" in name:
          lowercase_       :List[Any]     =       name.replace('''norm_cross'''       ,		'''encoder_attn_layer_norm'''		)
     if "norm2" in name:
          lowercase_       :List[str]     =       name.replace('''norm2'''       ,		'''final_layer_norm'''		)
     if "out_norm" in name:
          lowercase_       :Any     =       name.replace('''out_norm'''       ,		'''model.decoder.layer_norm'''		)
     if "linears" in name:
          lowercase_       :List[str]     =       name.replace('''linears'''       ,		'''lm_heads'''		)
     if "condition_provider.conditioners.description.output_proj" in name:
          lowercase_       :Tuple     =       name.replace('''condition_provider.conditioners.description.output_proj'''       ,		'''enc_to_dec_proj'''		)
     return name




def     UpperCamelCase    (      _a       ,		_a		)					->							int:
     '''simple docstring'''





     lowercase_       :Union[str, Any]     =       list(state_dict.keys()		)
     lowercase_       :int     =       {}
     for key in keys:
          lowercase_       :str     =       state_dict.pop(lowerCamelCase_		)
          lowercase_       :Union[str, Any]     =       rename_keys(lowerCamelCase_		)
          if "in_proj_weight" in key:
               # split fused qkv proj
               lowercase_       :List[Any]     =       val[:hidden_size, :]
               lowercase_       :Union[str, Any]     =       val[hidden_size : 2 * hidden_size, :]
               lowercase_       :Dict     =       val[-hidden_size:, :]
          elif "enc_to_dec_proj" in key:
               lowercase_       :str     =       val
          else:
               lowercase_       :Union[str, Any]     =       val
     return state_dict, enc_dec_proj_state_dict




def     UpperCamelCase    (      _a		)					->							Tuple:
     '''simple docstring'''





     if checkpoint == "small":
          # default config values
          lowercase_       :Optional[int]     =       1_0_2_4
          lowercase_       :Tuple     =       2_4
          lowercase_       :Dict     =       1_6
     elif checkpoint == "medium":
          lowercase_       :int     =       1_5_3_6
          lowercase_       :int     =       4_8
          lowercase_       :int     =       2_4
     elif checkpoint == "large":
          lowercase_       :Optional[Any]     =       2_0_4_8
          lowercase_       :Any     =       4_8
          lowercase_       :Any     =       3_2
     else:
          raise ValueError(f"Checkpoint should be one of `[\'small\', \'medium\', \'large\']`, got {checkpoint}."		)
     lowercase_       :Dict     =       MusicgenDecoderConfig(
         hidden_size=lowerCamelCase_       ,		ffn_dim=hidden_size * 4       ,		num_hidden_layers=lowerCamelCase_       ,		num_attention_heads=lowerCamelCase_       ,		)
     return config




@torch.no_grad()
def     UpperCamelCase    (      _a       ,		_a=None       ,		_a=None       ,		_a="cpu"		)					->							Optional[int]:
     '''simple docstring'''





     lowercase_       :Union[str, Any]     =       MusicGen.get_pretrained(lowerCamelCase_       ,		device=lowerCamelCase_		)
     lowercase_       :List[Any]     =       decoder_config_from_checkpoint(lowerCamelCase_		)

     lowercase_       :Any     =       fairseq_model.lm.state_dict()
     lowercase_       :Dict     =       rename_state_dict(
         lowerCamelCase_       ,		hidden_size=decoder_config.hidden_size		)

     lowercase_       :int     =       TaEncoderModel.from_pretrained('''t5-base'''		)
     lowercase_       :Any     =       EncodecModel.from_pretrained('''facebook/encodec_32khz'''		)
     lowercase_       :Dict     =       MusicgenForCausalLM(lowerCamelCase_		).eval()

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

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

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

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

     # init the composite model
     lowercase_       :Optional[Any]     =       MusicgenForConditionalGeneration(text_encoder=lowerCamelCase_       ,		audio_encoder=lowerCamelCase_       ,		decoder=lowerCamelCase_		)

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

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

     with torch.no_grad():
          lowercase_       :List[Any]     =       model(input_ids=lowerCamelCase_       ,		decoder_input_ids=lowerCamelCase_		).logits

     if logits.shape != (8, 1, 2_0_4_8):
          raise ValueError('''Incorrect shape for logits'''		)

     # now construct the processor
     lowercase_       :Optional[Any]     =       AutoTokenizer.from_pretrained('''t5-base'''		)
     lowercase_       :List[str]     =       AutoFeatureExtractor.from_pretrained('''facebook/encodec_32khz'''       ,		padding_side='''left'''		)

     lowercase_       :str     =       MusicgenProcessor(feature_extractor=lowerCamelCase_       ,		tokenizer=lowerCamelCase_		)

     # set the appropriate bos/pad token ids
     lowercase_       :int     =       2_0_4_8
     lowercase_       :Optional[int]     =       2_0_4_8

     # set other default generation config params
     lowercase_       :Optional[int]     =       int(3_0 * audio_encoder.config.frame_rate		)
     lowercase_       :List[str]     =       True
     lowercase_       :str     =       3.0

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

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


if __name__ == "__main__":
   SCREAMING_SNAKE_CASE       :		Union[str, 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."
   )

   SCREAMING_SNAKE_CASE       :		str         =				parser.parse_args()
   convert_musicgen_checkpoint(args.checkpoint, args.pytorch_dump_folder, args.push_to_hub)








import gc
import random
import unittest

import numpy as np
import torch
from transformers import XLMRobertaTokenizer

from diffusers import (
    AltDiffusionImgaImgPipeline,
    AutoencoderKL,
    PNDMScheduler,
    UNetaDConditionModel,
)
from diffusers.image_processor import VaeImageProcessor
from diffusers.pipelines.alt_diffusion.modeling_roberta_series import (
    RobertaSeriesConfig,
    RobertaSeriesModelWithTransformation,
)
from diffusers.utils import floats_tensor, load_image, load_numpy, slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu


enable_full_determinism()






class      UpperCamelCase					(		unittest.TestCase     ):


   '''simple docstring'''







   def 		UpperCamelCase  (     self		):
        # clean up the VRAM after each test
        super().tearDown()
        gc.collect()
        torch.cuda.empty_cache()
   @property
   def 		UpperCamelCase  (     self		):
        lowercase_       :Dict     =       1
        lowercase_       :Optional[Any]     =       3
        lowercase_       :Optional[int]     =       (32, 32)

        lowercase_       :Dict     =       floats_tensor((batch_size, num_channels) + sizes   ,							rng=random.Random(0		)		).to(UpperCamelCase_		)
        return image
   @property
   def 		UpperCamelCase  (     self		):
        torch.manual_seed(0		)
        lowercase_       :Optional[int]     =       UNetaDConditionModel(
            block_out_channels=(32, 64)   ,							layers_per_block=2   ,							sample_size=32   ,							in_channels=4   ,							out_channels=4   ,							down_block_types=('''DownBlock2D''', '''CrossAttnDownBlock2D''')   ,							up_block_types=('''CrossAttnUpBlock2D''', '''UpBlock2D''')   ,							cross_attention_dim=32   ,							)
        return model
   @property
   def 		UpperCamelCase  (     self		):
        torch.manual_seed(0		)
        lowercase_       :Dict     =       AutoencoderKL(
            block_out_channels=[32, 64]   ,							in_channels=3   ,							out_channels=3   ,							down_block_types=['''DownEncoderBlock2D''', '''DownEncoderBlock2D''']   ,							up_block_types=['''UpDecoderBlock2D''', '''UpDecoderBlock2D''']   ,							latent_channels=4   ,							)
        return model
   @property
   def 		UpperCamelCase  (     self		):
        torch.manual_seed(0		)
        lowercase_       :str     =       RobertaSeriesConfig(
            hidden_size=32   ,							project_dim=32   ,							intermediate_size=37   ,							layer_norm_eps=1E-05   ,							num_attention_heads=4   ,							num_hidden_layers=5   ,							pad_token_id=1   ,							vocab_size=5006   ,							)
        return RobertaSeriesModelWithTransformation(UpperCamelCase_		)
   @property
   def 		UpperCamelCase  (     self		):
        def extract(*UpperCamelCase_   ,							**UpperCamelCase_		):






             class      UpperCamelCase					:


                '''simple docstring'''







                def __init__(     self		):
                     lowercase_       :List[str]     =       torch.ones([0]		)




                def 		UpperCamelCase  (     self   ,							UpperCamelCase_		):
                     self.pixel_values.to(UpperCamelCase_		)
                     return self
             return Out()

        return extract
   def 		UpperCamelCase  (     self		):
        lowercase_       :List[str]     =       '''cpu'''  # ensure determinism for the device-dependent torch.Generator
        lowercase_       :Dict     =       self.dummy_cond_unet
        lowercase_       :Optional[Any]     =       PNDMScheduler(skip_prk_steps=UpperCamelCase_		)
        lowercase_       :Union[str, Any]     =       self.dummy_vae
        lowercase_       :Any     =       self.dummy_text_encoder
        lowercase_       :List[Any]     =       XLMRobertaTokenizer.from_pretrained('''hf-internal-testing/tiny-xlm-roberta'''		)
        lowercase_       :List[str]     =       77

        lowercase_       :int     =       self.dummy_image.to(UpperCamelCase_		)
        lowercase_       :Optional[Any]     =       init_image / 2 + 0.5

        # make sure here that pndm scheduler skips prk
        lowercase_       :str     =       AltDiffusionImgaImgPipeline(
            unet=UpperCamelCase_   ,							scheduler=UpperCamelCase_   ,							vae=UpperCamelCase_   ,							text_encoder=UpperCamelCase_   ,							tokenizer=UpperCamelCase_   ,							safety_checker=UpperCamelCase_   ,							feature_extractor=self.dummy_extractor   ,							)
        lowercase_       :List[str]     =       VaeImageProcessor(vae_scale_factor=alt_pipe.vae_scale_factor   ,							do_normalize=UpperCamelCase_		)
        lowercase_       :int     =       alt_pipe.to(UpperCamelCase_		)
        alt_pipe.set_progress_bar_config(disable=UpperCamelCase_		)

        lowercase_       :Any     =       '''A painting of a squirrel eating a burger'''
        lowercase_       :Tuple     =       torch.Generator(device=UpperCamelCase_		).manual_seed(0		)
        lowercase_       :Optional[int]     =       alt_pipe(
            [prompt]   ,							generator=UpperCamelCase_   ,							guidance_scale=6.0   ,							num_inference_steps=2   ,							output_type='''np'''   ,							image=UpperCamelCase_   ,							)

        lowercase_       :Dict     =       output.images

        lowercase_       :Tuple     =       torch.Generator(device=UpperCamelCase_		).manual_seed(0		)
        lowercase_       :Optional[int]     =       alt_pipe(
            [prompt]   ,							generator=UpperCamelCase_   ,							guidance_scale=6.0   ,							num_inference_steps=2   ,							output_type='''np'''   ,							image=UpperCamelCase_   ,							return_dict=UpperCamelCase_   ,							)[0]

        lowercase_       :Optional[int]     =       image[0, -3:, -3:, -1]
        lowercase_       :Dict     =       image_from_tuple[0, -3:, -3:, -1]

        assert image.shape == (1, 32, 32, 3)
        lowercase_       :int     =       np.array([0.4427, 0.3731, 0.4249, 0.4941, 0.4546, 0.4148, 0.4193, 0.4666, 0.4499]		)

        assert np.abs(image_slice.flatten() - expected_slice		).max() < 5E-3
        assert np.abs(image_from_tuple_slice.flatten() - expected_slice		).max() < 5E-3
   @unittest.skipIf(torch_device != '''cuda'''   ,							'''This test requires a GPU'''		)
   def 		UpperCamelCase  (     self		):
        lowercase_       :List[str]     =       self.dummy_cond_unet
        lowercase_       :Dict     =       PNDMScheduler(skip_prk_steps=UpperCamelCase_		)
        lowercase_       :Tuple     =       self.dummy_vae
        lowercase_       :Dict     =       self.dummy_text_encoder
        lowercase_       :Tuple     =       XLMRobertaTokenizer.from_pretrained('''hf-internal-testing/tiny-xlm-roberta'''		)
        lowercase_       :str     =       77

        lowercase_       :str     =       self.dummy_image.to(UpperCamelCase_		)

        # put models in fp16
        lowercase_       :Union[str, Any]     =       unet.half()
        lowercase_       :Union[str, Any]     =       vae.half()
        lowercase_       :List[str]     =       bert.half()

        # make sure here that pndm scheduler skips prk
        lowercase_       :List[Any]     =       AltDiffusionImgaImgPipeline(
            unet=UpperCamelCase_   ,							scheduler=UpperCamelCase_   ,							vae=UpperCamelCase_   ,							text_encoder=UpperCamelCase_   ,							tokenizer=UpperCamelCase_   ,							safety_checker=UpperCamelCase_   ,							feature_extractor=self.dummy_extractor   ,							)
        lowercase_       :List[Any]     =       VaeImageProcessor(vae_scale_factor=alt_pipe.vae_scale_factor   ,							do_normalize=UpperCamelCase_		)
        lowercase_       :List[str]     =       alt_pipe.to(UpperCamelCase_		)
        alt_pipe.set_progress_bar_config(disable=UpperCamelCase_		)

        lowercase_       :str     =       '''A painting of a squirrel eating a burger'''
        lowercase_       :Union[str, Any]     =       torch.manual_seed(0		)
        lowercase_       :Any     =       alt_pipe(
            [prompt]   ,							generator=UpperCamelCase_   ,							num_inference_steps=2   ,							output_type='''np'''   ,							image=UpperCamelCase_   ,							).images

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




   @unittest.skipIf(torch_device != '''cuda'''   ,							'''This test requires a GPU'''		)
   def 		UpperCamelCase  (     self		):
        lowercase_       :Optional[int]     =       load_image(
            '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'''
            '''/img2img/sketch-mountains-input.jpg'''		)
        # resize to resolution that is divisible by 8 but not 16 or 32
        lowercase_       :Optional[Any]     =       init_image.resize((760, 504)		)

        lowercase_       :List[str]     =       '''BAAI/AltDiffusion'''
        lowercase_       :Optional[Any]     =       AltDiffusionImgaImgPipeline.from_pretrained(
            UpperCamelCase_   ,							safety_checker=UpperCamelCase_   ,							)
        pipe.to(UpperCamelCase_		)
        pipe.set_progress_bar_config(disable=UpperCamelCase_		)
        pipe.enable_attention_slicing()

        lowercase_       :Optional[Any]     =       '''A fantasy landscape, trending on artstation'''

        lowercase_       :Optional[Any]     =       torch.manual_seed(0		)
        lowercase_       :str     =       pipe(
            prompt=UpperCamelCase_   ,							image=UpperCamelCase_   ,							strength=0.75   ,							guidance_scale=7.5   ,							generator=UpperCamelCase_   ,							output_type='''np'''   ,							)
        lowercase_       :Optional[Any]     =       output.images[0]

        lowercase_       :Optional[int]     =       image[255:258, 383:386, -1]

        assert image.shape == (504, 760, 3)
        lowercase_       :Any     =       np.array([0.9358, 0.9397, 0.9599, 0.9901, 1.0000, 1.0000, 0.9882, 1.0000, 1.0000]		)

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






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


   '''simple docstring'''







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




   def 		UpperCamelCase  (     self		):
        lowercase_       :Union[str, Any]     =       load_image(
            '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'''
            '''/img2img/sketch-mountains-input.jpg'''		)
        lowercase_       :Any     =       init_image.resize((768, 512)		)
        lowercase_       :List[str]     =       load_numpy(
            '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/img2img/fantasy_landscape_alt.npy'''		)

        lowercase_       :List[Any]     =       '''BAAI/AltDiffusion'''
        lowercase_       :Tuple     =       AltDiffusionImgaImgPipeline.from_pretrained(
            UpperCamelCase_   ,							safety_checker=UpperCamelCase_   ,							)
        pipe.to(UpperCamelCase_		)
        pipe.set_progress_bar_config(disable=UpperCamelCase_		)
        pipe.enable_attention_slicing()

        lowercase_       :List[str]     =       '''A fantasy landscape, trending on artstation'''

        lowercase_       :Optional[int]     =       torch.manual_seed(0		)
        lowercase_       :Tuple     =       pipe(
            prompt=UpperCamelCase_   ,							image=UpperCamelCase_   ,							strength=0.75   ,							guidance_scale=7.5   ,							generator=UpperCamelCase_   ,							output_type='''np'''   ,							)
        lowercase_       :int     =       output.images[0]

        assert image.shape == (512, 768, 3)
        # img2img is flaky across GPUs even in fp32, so using MAE here
        assert np.abs(expected_image - image		).max() < 1E-2

