





import random
import shutil
import tempfile
import unittest

import numpy as np

from transformers import TaTokenizer, TaTokenizerFast
from transformers.testing_utils import require_sentencepiece, require_torch
from transformers.utils.import_utils import is_speech_available, is_torch_available


if is_torch_available():
     pass

if is_speech_available():
     from transformers import EncodecFeatureExtractor, MusicgenProcessor


UpperCAmelCase			:			int	  = random.Random()
def   _A (	SCREAMING_SNAKE_CASE  :     Optional[int] , SCREAMING_SNAKE_CASE  :     Dict=1.0 , SCREAMING_SNAKE_CASE  :     Any=None , SCREAMING_SNAKE_CASE  :     Optional[Any]=None		):





  """simple docstring"""



  if rng is None:
    a__			:      Union[str, Any]            =global_rng

  a__			:      Dict            =[]
  for batch_idx in range(shape[0]		):
    values.append([]		)
    for _ in range(shape[1]		):
      values[-1].append(rng.random() * scale		)

  return values







@require_torch
@require_sentencepiece
class        __lowerCAmelCase (					unittest.TestCase):







      def      _lowercase						(					self							)     ->		Dict:



        '''simple docstring'''




        a__			:      Any            ="facebook/musicgen-small"
        a__			:      Any            =tempfile.mkdtemp()







      def      _lowercase						(					self	,   **lowerCAmelCase__							)     ->		Optional[int]:



        '''simple docstring'''




        return TaTokenizer.from_pretrained(self.checkpoint	,   **lowerCAmelCase__							)







      def      _lowercase						(					self	,   **lowerCAmelCase__							)     ->		str:



        '''simple docstring'''




        return EncodecFeatureExtractor.from_pretrained(self.checkpoint	,   **lowerCAmelCase__							)







      def      _lowercase						(					self							)     ->		List[str]:



        '''simple docstring'''




        shutil.rmtree(self.tmpdirname							)







      def      _lowercase						(					self							)     ->		List[str]:



        '''simple docstring'''




        a__			:      List[Any]            =self.get_tokenizer()
        a__			:      str            =self.get_feature_extractor()

        a__			:      Optional[Any]            =MusicgenProcessor(tokenizer=lowerCAmelCase__	,   feature_extractor=lowerCAmelCase__							)

        processor.save_pretrained(self.tmpdirname							)
        a__			:      List[str]            =MusicgenProcessor.from_pretrained(self.tmpdirname							)

        self.assertEqual(processor.tokenizer.get_vocab()	,   tokenizer.get_vocab()							)
        self.assertIsInstance(processor.tokenizer	,   lowerCAmelCase__							)

        self.assertEqual(processor.feature_extractor.to_json_string()	,   feature_extractor.to_json_string()							)
        self.assertIsInstance(processor.feature_extractor	,   lowerCAmelCase__							)







      def      _lowercase						(					self							)     ->		Optional[int]:



        '''simple docstring'''




        a__			:      List[Any]            =MusicgenProcessor(tokenizer=self.get_tokenizer()	,   feature_extractor=self.get_feature_extractor()							)
        processor.save_pretrained(self.tmpdirname							)

        a__			:      Optional[Any]            =self.get_tokenizer(bos_token="(BOS)"	,   eos_token="(EOS)"							)
        a__			:      Union[str, Any]            =self.get_feature_extractor(do_normalize=lowerCAmelCase__	,   padding_value=1.0							)

        a__			:      str            =MusicgenProcessor.from_pretrained(
            self.tmpdirname	,   bos_token="(BOS)"	,   eos_token="(EOS)"	,   do_normalize=lowerCAmelCase__	,   padding_value=1.0							)

        self.assertEqual(processor.tokenizer.get_vocab()	,   tokenizer_add_kwargs.get_vocab()							)
        self.assertIsInstance(processor.tokenizer	,   lowerCAmelCase__							)

        self.assertEqual(processor.feature_extractor.to_json_string()	,   feature_extractor_add_kwargs.to_json_string()							)
        self.assertIsInstance(processor.feature_extractor	,   lowerCAmelCase__							)







      def      _lowercase						(					self							)     ->		Dict:



        '''simple docstring'''




        a__			:      Tuple            =self.get_feature_extractor()
        a__			:      Optional[int]            =self.get_tokenizer()

        a__			:      str            =MusicgenProcessor(tokenizer=lowerCAmelCase__	,   feature_extractor=lowerCAmelCase__							)

        a__			:      int            =floats_list((3, 1_0_0_0)							)

        a__			:      List[Any]            =feature_extractor(lowerCAmelCase__	,   return_tensors="np"							)
        a__			:      Union[str, Any]            =processor(lowerCAmelCase__	,   return_tensors="np"							)

        for key in input_feat_extract.keys():
          self.assertAlmostEqual(input_feat_extract[key].sum()	,   input_processor[key].sum()	,   delta=1E-2							)







      def      _lowercase						(					self							)     ->		Dict:



        '''simple docstring'''




        a__			:      Any            =self.get_feature_extractor()
        a__			:      Dict            =self.get_tokenizer()

        a__			:      Optional[int]            =MusicgenProcessor(tokenizer=lowerCAmelCase__	,   feature_extractor=lowerCAmelCase__							)

        a__			:      Any            ="This is a test string"

        a__			:      Optional[int]            =processor(text=lowerCAmelCase__							)

        a__			:      int            =tokenizer(lowerCAmelCase__							)

        for key in encoded_tok.keys():
          self.assertListEqual(encoded_tok[key]	,   encoded_processor[key]							)







      def      _lowercase						(					self							)     ->		int:



        '''simple docstring'''




        a__			:      Optional[Any]            =self.get_feature_extractor()
        a__			:      Any            =self.get_tokenizer()

        a__			:      Tuple            =MusicgenProcessor(tokenizer=lowerCAmelCase__	,   feature_extractor=lowerCAmelCase__							)

        a__			:      List[Any]            =[[1, 4, 5, 8, 1, 0, 8], [3, 4, 3, 1, 1, 8, 9]]

        a__			:      Dict            =processor.batch_decode(sequences=lowerCAmelCase__							)
        a__			:      Any            =tokenizer.batch_decode(lowerCAmelCase__							)

        self.assertListEqual(lowerCAmelCase__	,   lowerCAmelCase__							)







      def      _lowercase						(					self							)     ->		Union[str, Any]:



        '''simple docstring'''




        a__			:      Any            =self.get_feature_extractor()
        a__			:      List[str]            =self.get_tokenizer()

        a__			:      List[str]            =MusicgenProcessor(tokenizer=lowerCAmelCase__	,   feature_extractor=lowerCAmelCase__							)

        self.assertListEqual(
            processor.model_input_names	,   feature_extractor.model_input_names	,   msg="`processor` and `feature_extractor` model input names do not match"	,   )







      def      _lowercase						(					self							)     ->		str:



        '''simple docstring'''




        a__			:      Dict            =self.get_feature_extractor(padding_side="left"							)
        a__			:      Any            =self.get_tokenizer()

        a__			:      List[str]            =MusicgenProcessor(tokenizer=lowerCAmelCase__	,   feature_extractor=lowerCAmelCase__							)

        a__			:      List[str]            =[floats_list((1, x)							)[0] for x in range(5	,   2_0	,   5							)]
        a__			:      List[Any]            =processor(lowerCAmelCase__							).padding_mask

        a__			:      Optional[int]            =np.asarray(floats_list((3, 2_0)							)							)[:, None, :]
        a__			:      Union[str, Any]            =processor.batch_decode(lowerCAmelCase__	,   padding_mask=lowerCAmelCase__							)

        self.assertIsInstance(lowerCAmelCase__	,   lowerCAmelCase__							)

        for audio in decoded_audios:
          self.assertIsInstance(lowerCAmelCase__	,   np.ndarray							)

        self.assertTrue(decoded_audios[0].shape == (1, 1_0)							)
        self.assertTrue(decoded_audios[1].shape == (1, 1_5)							)
        self.assertTrue(decoded_audios[2].shape == (1, 2_0)							)













import tempfile
import unittest

import torch

from diffusers import UNetaDConditionModel
from diffusers.training_utils import EMAModel
from diffusers.utils.testing_utils import enable_full_determinism, skip_mps, torch_device


enable_full_determinism()







class        __lowerCAmelCase (					unittest.TestCase):
      _lowercase       :      str        =    """hf-internal-testing/tiny-stable-diffusion-pipe"""
      _lowercase       :      Optional[int]        =    1
      _lowercase       :      List[str]        =    77
      _lowercase       :      int        =    32
      _lowercase       :      Dict        =    4
      _lowercase       :      str        =    64
      _lowercase       :      Optional[Any]        =    torch.manual_seed(0)







      def      _lowercase						(					self	,   lowerCAmelCase__=0.99_99							)     ->		Tuple:



        '''simple docstring'''




        a__			:      int            =UNetaDConditionModel.from_pretrained(self.model_id	,   subfolder="unet"							)
        a__			:      Any            =unet.to(lowerCAmelCase__							)
        a__			:      int            =EMAModel(unet.parameters()	,   decay=lowerCAmelCase__	,   model_cls=lowerCAmelCase__	,   model_config=unet.config							)
        return unet, ema_unet







      def      _lowercase						(					self							)     ->		Dict:



        '''simple docstring'''




        a__			:      str            =torch.randn(
            self.batch_size	,   self.num_in_channels	,   self.latent_height	,   self.latent_width	,   generator=self.generator							).to(lowerCAmelCase__							)
        a__			:      Any            =torch.randint(0	,   1_0_0_0	,   size=(self.batch_size,)	,   generator=self.generator							).to(lowerCAmelCase__							)
        a__			:      Tuple            =torch.randn(
            self.batch_size	,   self.prompt_length	,   self.text_encoder_hidden_dim	,   generator=self.generator							).to(lowerCAmelCase__							)
        return noisy_latents, timesteps, encoder_hidden_states







      def      _lowercase						(					self	,   lowerCAmelCase__							)     ->		Any:



        '''simple docstring'''




        a__			:      Any            ={}
        for k, param in unet.state_dict().items():
          a__			:      Tuple            =torch.randn_like(lowerCAmelCase__							) + (param * torch.randn_like(lowerCAmelCase__							))
          updated_state_dict.update({k: updated_param}							)
        unet.load_state_dict(lowerCAmelCase__							)
        return unet







      def      _lowercase						(					self							)     ->		Tuple:



        '''simple docstring'''




        a__      ,							a__			:      List[Any]            =self.get_models()
        # Take the first (hypothetical) EMA step.
        ema_unet.step(unet.parameters()							)
        assert ema_unet.optimization_step == 1

        # Take two more.
        for _ in range(2							):
          ema_unet.step(unet.parameters()							)
        assert ema_unet.optimization_step == 3







      def      _lowercase						(					self							)     ->		List[str]:



        '''simple docstring'''




        a__      ,							a__			:      Any            =self.get_models()
        # Since the `unet` is not being updated (i.e., backprop'd)
        # there won't be any difference between the `params` of `unet`
        # and `ema_unet` even if we call `ema_unet.step(unet.parameters())`.
        ema_unet.step(unet.parameters()							)
        a__			:      str            =list(unet.parameters()							)
        for s_param, param in zip(ema_unet.shadow_params	,   lowerCAmelCase__							):
          assert torch.allclose(lowerCAmelCase__	,   lowerCAmelCase__							)

        # The above holds true even if we call `ema.step()` multiple times since
        # `unet` params are still not being updated.
        for _ in range(4							):
          ema_unet.step(unet.parameters()							)
        for s_param, param in zip(ema_unet.shadow_params	,   lowerCAmelCase__							):
          assert torch.allclose(lowerCAmelCase__	,   lowerCAmelCase__							)







      def      _lowercase						(					self							)     ->		List[Any]:



        '''simple docstring'''




        a__      ,							a__			:      Optional[Any]            =self.get_models()
        # Here we simulate the parameter updates for `unet`. Since there might
        # be some parameters which are initialized to zero we take extra care to
        # initialize their values to something non-zero before the multiplication.
        a__			:      List[str]            =self.simulate_backprop(lowerCAmelCase__							)

        # Take the EMA step.
        ema_unet.step(unet_pseudo_updated_step_one.parameters()							)

        # Now the EMA'd parameters won't be equal to the original model parameters.
        a__			:      Optional[int]            =list(unet_pseudo_updated_step_one.parameters()							)
        for s_param, param in zip(ema_unet.shadow_params	,   lowerCAmelCase__							):
          assert ~torch.allclose(lowerCAmelCase__	,   lowerCAmelCase__							)

        # Ensure this is the case when we take multiple EMA steps.
        for _ in range(4							):
          ema_unet.step(unet.parameters()							)
        for s_param, param in zip(ema_unet.shadow_params	,   lowerCAmelCase__							):
          assert ~torch.allclose(lowerCAmelCase__	,   lowerCAmelCase__							)







      def      _lowercase						(					self							)     ->		Union[str, Any]:



        '''simple docstring'''




        a__      ,							a__			:      Dict            =self.get_models()

        # First backprop + EMA
        a__			:      Tuple            =self.simulate_backprop(lowerCAmelCase__							)
        ema_unet.step(unet_step_one.parameters()							)
        a__			:      Dict            =ema_unet.shadow_params

        # Second backprop + EMA
        a__			:      str            =self.simulate_backprop(lowerCAmelCase__							)
        ema_unet.step(unet_step_two.parameters()							)
        a__			:      List[str]            =ema_unet.shadow_params

        for step_one, step_two in zip(lowerCAmelCase__	,   lowerCAmelCase__							):
          assert ~torch.allclose(lowerCAmelCase__	,   lowerCAmelCase__							)







      def      _lowercase						(					self							)     ->		List[str]:



        '''simple docstring'''




        a__      ,							a__			:      int            =self.get_models(decay=0.0							)
        a__			:      Optional[int]            =self.simulate_backprop(lowerCAmelCase__							)
        ema_unet.step(unet_step_one.parameters()							)
        a__			:      List[Any]            =ema_unet.shadow_params

        a__			:      Tuple            =self.simulate_backprop(lowerCAmelCase__							)
        ema_unet.step(unet_step_two.parameters()							)
        a__			:      List[Any]            =ema_unet.shadow_params

        for step_one, step_two in zip(lowerCAmelCase__	,   lowerCAmelCase__							):
          assert torch.allclose(lowerCAmelCase__	,   lowerCAmelCase__							)







      @skip_mps
      def      _lowercase						(					self							)     ->		Optional[Any]:



        '''simple docstring'''




        a__      ,							a__			:      Union[str, Any]            =self.get_models()
        a__      ,							a__      ,							a__			:      Tuple            =self.get_dummy_inputs()

        with tempfile.TemporaryDirectory() as tmpdir:
          ema_unet.save_pretrained(lowerCAmelCase__							)
          a__			:      List[str]            =UNetaDConditionModel.from_pretrained(lowerCAmelCase__	,   model_cls=lowerCAmelCase__							)
          a__			:      Tuple            =loaded_unet.to(unet.device							)

        # Since no EMA step has been performed the outputs should match.
        a__			:      Optional[Any]            =unet(lowerCAmelCase__	,   lowerCAmelCase__	,   lowerCAmelCase__							).sample
        a__			:      Optional[int]            =loaded_unet(lowerCAmelCase__	,   lowerCAmelCase__	,   lowerCAmelCase__							).sample

        assert torch.allclose(lowerCAmelCase__	,   lowerCAmelCase__	,   atol=1E-4							)






