




import argparse

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

from transformers import (
    CLIPTokenizer,
    CLIPTokenizerFast,
    VideoMAEImageProcessor,
    XCLIPConfig,
    XCLIPModel,
    XCLIPProcessor,
    XCLIPTextConfig,
    XCLIPVisionConfig,
)



def 					UpperCAmelCase__     (   lowerCamelCase,	lowerCamelCase      ):
  lowercase			:Union[str, Any]							=  XCLIPTextConfig()

  # derive patch size from model name
  lowercase			:int							=  model_name.find("patch"      )
  lowercase			:str							=  int(model_name[start_idx + len("patch"      ) : start_idx + len("patch"      ) + 2]      )
  lowercase			:List[Any]							=  XCLIPVisionConfig(patch_size=lowerCamelCase,	num_frames=lowerCamelCase      )

  if "large" in model_name:
    lowercase			:Optional[Any]							=  768
    lowercase			:str							=  3072
    lowercase			:Any							=  12

    lowercase			:Optional[Any]							=  1024
    lowercase			:Optional[Any]							=  4096
    lowercase			:Union[str, Any]							=  16
    lowercase			:Dict							=  24
    lowercase			:List[Any]							=  768
    lowercase			:Optional[int]							=  3072

  if model_name == "xclip-large-patch14-16-frames":
    lowercase			:Optional[Any]							=  336

  lowercase			:str							=  XCLIPConfig.from_text_vision_configs(lowerCamelCase,	lowerCamelCase      )

  if "large" in model_name:
    lowercase			:List[Any]							=  768

  return config



def 					UpperCAmelCase__     (   lowerCamelCase      ):
  # text encoder
  if name == "token_embedding.weight":
    lowercase			:Optional[int]							=  name.replace("token_embedding.weight",	"text_model.embeddings.token_embedding.weight"      )
  if name == "positional_embedding":
    lowercase			:Tuple							=  name.replace("positional_embedding",	"text_model.embeddings.position_embedding.weight"      )
  if "ln_1" in name:
    lowercase			:int							=  name.replace("ln_1",	"layer_norm1"      )
  if "ln_2" in name:
    lowercase			:Optional[Any]							=  name.replace("ln_2",	"layer_norm2"      )
  if "c_fc" in name:
    lowercase			:Optional[int]							=  name.replace("c_fc",	"fc1"      )
  if "c_proj" in name:
    lowercase			:Union[str, Any]							=  name.replace("c_proj",	"fc2"      )
  if name.startswith("transformer.resblocks"      ):
    lowercase			:Dict							=  name.replace("transformer.resblocks",	"text_model.encoder.layers"      )
  if "attn.out_proj" in name and "message" not in name:
    lowercase			:int							=  name.replace("attn.out_proj",	"self_attn.out_proj"      )
  if "ln_final" in name:
    lowercase			:Dict							=  name.replace("ln_final",	"text_model.final_layer_norm"      )
  # visual encoder
  if name == "visual.class_embedding":
    lowercase			:Any							=  name.replace("visual.class_embedding",	"vision_model.embeddings.class_embedding"      )
  if name == "visual.positional_embedding":
    lowercase			:int							=  name.replace("visual.positional_embedding",	"vision_model.embeddings.position_embedding.weight"      )
  if name.startswith("visual.transformer.resblocks"      ):
    lowercase			:str							=  name.replace("visual.transformer.resblocks",	"vision_model.encoder.layers"      )
  if "visual.conv1" in name:
    lowercase			:str							=  name.replace("visual.conv1",	"vision_model.embeddings.patch_embedding"      )
  if "visual.ln_pre" in name:
    lowercase			:Dict							=  name.replace("visual.ln_pre",	"vision_model.pre_layernorm"      )
  if "visual.ln_post" in name:
    lowercase			:List[Any]							=  name.replace("visual.ln_post",	"vision_model.post_layernorm"      )
  if "visual.proj" in name:
    lowercase			:Dict							=  name.replace("visual.proj",	"visual_projection.weight"      )
  if "text_projection" in name:
    lowercase			:int							=  name.replace("text_projection",	"text_projection.weight"      )
  # things on top
  if "prompts_visual_proj" in name:
    lowercase			:List[Any]							=  name.replace("prompts_visual_proj",	"prompts_visual_projection"      )
  if "prompts_visual_ln" in name:
    lowercase			:Union[str, Any]							=  name.replace("prompts_visual_ln",	"prompts_visual_layernorm"      )
  # mit
  if name == "mit.positional_embedding":
    lowercase			:List[Any]							=  name.replace("positional",	"position"      )
  if name.startswith("mit.resblocks"      ):
    lowercase			:str							=  name.replace("mit.resblocks",	"mit.encoder.layers"      )
  # prompts generator
  if name.startswith("prompts_generator.norm"      ):
    lowercase			:List[Any]							=  name.replace("prompts_generator.norm",	"prompts_generator.layernorm"      )

  return name



def 					UpperCAmelCase__     (   lowerCamelCase,	lowerCamelCase      ):
  for key in orig_state_dict.copy().keys():
    lowercase			:Any							=  orig_state_dict.pop(lowerCamelCase      )

    if "attn.in_proj" in key:
      lowercase			:Any							=  key.split("."      )
      if key.startswith("visual"      ):
        lowercase			:Optional[Any]							=  key_split[3]
        lowercase			:Tuple							=  config.vision_config.hidden_size
        if "message_attn" in key:
          if "weight" in key:
            lowercase			:Dict							=  val[
                :dim, :
            ]
            lowercase			:List[str]							=  val[
                dim : dim * 2, :
            ]
            lowercase			:List[Any]							=  val[
                -dim:, :
            ]
          else:
            lowercase			:Any							=  val[
                :dim
            ]
            lowercase			:Optional[Any]							=  val[
                dim : dim * 2
            ]
            lowercase			:List[Any]							=  val[
                -dim:
            ]
        else:
          if "weight" in key:
            lowercase			:List[Any]							=  val[
                :dim, :
            ]
            lowercase			:List[Any]							=  val[
                dim : dim * 2, :
            ]
            lowercase			:Any							=  val[
                -dim:, :
            ]
          else:
            lowercase			:Any							=  val[:dim]
            lowercase			:Optional[int]							=  val[
                dim : dim * 2
            ]
            lowercase			:int							=  val[-dim:]
      elif key.startswith("mit"      ):
        lowercase			:Any							=  key_split[2]
        lowercase			:str							=  config.vision_config.mit_hidden_size
        if "weight" in key:
          lowercase			:Any							=  val[:dim, :]
          lowercase			:Dict							=  val[dim : dim * 2, :]
          lowercase			:List[str]							=  val[-dim:, :]
        else:
          lowercase			:List[str]							=  val[:dim]
          lowercase			:int							=  val[dim : dim * 2]
          lowercase			:List[Any]							=  val[-dim:]
      else:
        lowercase			:int							=  key_split[2]
        lowercase			:int							=  config.text_config.hidden_size
        if "weight" in key:
          lowercase			:Any							=  val[:dim, :]
          lowercase			:Tuple							=  val[
              dim : dim * 2, :
          ]
          lowercase			:Optional[int]							=  val[-dim:, :]
        else:
          lowercase			:Optional[Any]							=  val[:dim]
          lowercase			:int							=  val[
              dim : dim * 2
          ]
          lowercase			:Optional[Any]							=  val[-dim:]
    else:
      lowercase			:str							=  rename_key(lowerCamelCase      )
      if new_key_name in ["visual_projection.weight", "text_projection.weight"]:
        lowercase			:Optional[int]							=  val.T
      lowercase			:Any							=  val

  return orig_state_dict



def 					UpperCAmelCase__     (   lowerCamelCase      ):
  if num_frames == 8:
    lowercase			:Optional[Any]							=  "eating_spaghetti_8_frames.npy"
  elif num_frames == 16:
    lowercase			:Optional[Any]							=  "eating_spaghetti.npy"
  elif num_frames == 32:
    lowercase			:List[Any]							=  "eating_spaghetti_32_frames.npy"
  lowercase			:Union[str, Any]							=  hf_hub_download(
      repo_id="hf-internal-testing/spaghetti-video",	filename=lowerCamelCase,	repo_type="dataset",	)
  lowercase			:List[Any]							=  np.load(lowerCamelCase      )
  return list(lowerCamelCase      )



def 					UpperCAmelCase__     (   lowerCamelCase,	lowerCamelCase=None,	lowerCamelCase=False      ):
  lowercase			:str							=  {
      # fully supervised kinetics-400 checkpoints
      "xclip-base-patch32": "https://github.com/nbl97/X-CLIP_Model_Zoo/releases/download/v1.0/k400_32_8.pth",
      "xclip-base-patch32-16-frames": (
          "https://github.com/nbl97/X-CLIP_Model_Zoo/releases/download/v1.0/k400_32_16.pth"
      ),
      "xclip-base-patch16": "https://github.com/nbl97/X-CLIP_Model_Zoo/releases/download/v1.0/k400_16_8.pth",
      "xclip-base-patch16-16-frames": (
          "https://github.com/nbl97/X-CLIP_Model_Zoo/releases/download/v1.0/k400_16_16.pth"
      ),
      "xclip-large-patch14": "https://drive.google.com/u/0/uc?id=1NUOImq0o5DlQTST17iIP3vG7DgmHQuCx&amp;export=download&amp;confirm=t&amp;uuid=b26caedc-88e2-473e-830a-9d158b653cdb",
      "xclip-large-patch14-16-frames": "https://drive.google.com/u/0/uc?id=1FOYgnJc097OJ4lGwtRCCydQyVPJEOH7d&amp;export=download&amp;confirm=t&amp;uuid=538fa810-e671-4050-b385-9a623f89804f",
      # fully supervised kinetics-600 checkpoints
      "xclip-base-patch16-kinetics-600": (
          "https://github.com/nbl97/X-CLIP_Model_Zoo/releases/download/v1.0/k600_16_8.pth"
      ),
      "xclip-base-patch16-kinetics-600-16-frames": (
          "https://github.com/nbl97/X-CLIP_Model_Zoo/releases/download/v1.0/k600_16_16.pth"
      ),
      "xclip-large-patch14-kinetics-600": "https://drive.google.com/u/0/uc?id=1FV8C1INuM91sLAN4ImjzePLIlpMSihwV&amp;export=download&amp;confirm=t&amp;uuid=141d4977-4a65-44ae-864f-4b0c19f838be",
      # few shot
      "xclip-base-patch16-hmdb-2-shot": (
          "https://github.com/nbl97/X-CLIP_Model_Zoo/releases/download/v1.0/few_hmdb_2.pth"
      ),
      "xclip-base-patch16-hmdb-4-shot": (
          "https://github.com/nbl97/X-CLIP_Model_Zoo/releases/download/v1.0/few_hmdb_4.pth"
      ),
      "xclip-base-patch16-hmdb-8-shot": (
          "https://github.com/nbl97/X-CLIP_Model_Zoo/releases/download/v1.0/few_hmdb_8.pth"
      ),
      "xclip-base-patch16-hmdb-16-shot": (
          "https://github.com/nbl97/X-CLIP_Model_Zoo/releases/download/v1.0/few_hmdb_16.pth"
      ),
      "xclip-base-patch16-ucf-2-shot": (
          "https://github.com/nbl97/X-CLIP_Model_Zoo/releases/download/v1.0/few_ucf_2.pth"
      ),
      "xclip-base-patch16-ucf-4-shot": (
          "https://github.com/nbl97/X-CLIP_Model_Zoo/releases/download/v1.0/few_ucf_4.pth"
      ),
      "xclip-base-patch16-ucf-8-shot": (
          "https://github.com/nbl97/X-CLIP_Model_Zoo/releases/download/v1.0/few_ucf_8.pth"
      ),
      "xclip-base-patch16-ucf-16-shot": (
          "https://github.com/nbl97/X-CLIP_Model_Zoo/releases/download/v1.0/few_ucf_16.pth"
      ),
      # zero shot
      "xclip-base-patch16-zero-shot": "https://github.com/nbl97/X-CLIP_Model_Zoo/releases/download/v1.0/zero.pth",
  }

  lowercase			:int							=  model_to_url[model_name]
  lowercase			:Optional[Any]							=  8
  if "16-frames" in model_name:
    lowercase			:str							=  16
  elif "shot" in model_name:
    lowercase			:Any							=  32

  lowercase			:Tuple							=  get_xclip_config(lowerCamelCase,	lowerCamelCase      )
  lowercase			:Any							=  XCLIPModel(lowerCamelCase      )
  model.eval()

  if "drive" in checkpoint_url:
    lowercase			:Optional[Any]							=  "pytorch_model.bin"
    gdown.cached_download(lowerCamelCase,	lowerCamelCase,	quiet=lowerCamelCase      )
    lowercase			:str							=  torch.load(lowerCamelCase,	map_location="cpu"      )["model"]
  else:
    lowercase			:List[str]							=  torch.hub.load_state_dict_from_url(lowerCamelCase      )["model"]

  lowercase			:List[str]							=  convert_state_dict(lowerCamelCase,	lowerCamelCase      )

  lowercase			:int							=  XCLIPModel(lowerCamelCase      )
  lowercase		,			lowercase			:int							=  model.load_state_dict(lowerCamelCase,	strict=lowerCamelCase      )
  assert missing_keys == ["text_model.embeddings.position_ids", "vision_model.embeddings.position_ids"]
  model.eval()

  lowercase			:Optional[Any]							=  336 if model_name == "xclip-large-patch14-16-frames" else 224
  lowercase			:Any							=  VideoMAEImageProcessor(size=lowerCamelCase      )
  lowercase			:Dict							=  CLIPTokenizer.from_pretrained("openai/clip-vit-base-patch32"      )
  lowercase			:List[str]							=  CLIPTokenizerFast.from_pretrained("openai/clip-vit-base-patch32"      )
  lowercase			:str							=  XCLIPProcessor(image_processor=lowerCamelCase,	tokenizer=lowerCamelCase      )

  lowercase			:Any							=  prepare_video(lowerCamelCase      )
  lowercase			:Tuple							=  processor(
      text=["playing sports", "eating spaghetti", "go shopping"],	videos=lowerCamelCase,	return_tensors="pt",	padding=lowerCamelCase      )

  print("Shape of pixel values:",	inputs.pixel_values.shape      )

  with torch.no_grad():
    lowercase			:int							=  model(**lowerCamelCase      )

  # Verify outputs
  lowercase			:List[str]							=  outputs.logits_per_video
  lowercase			:Tuple							=  logits_per_video.softmax(dim=1      )
  print("Probs:",	lowerCamelCase      )
  # kinetics-400
  if model_name == "xclip-base-patch32":
    lowercase			:Optional[int]							=  torch.tensor([[0.0_019, 0.9_951, 0.0_030]]      )
  elif model_name == "xclip-base-patch32-16-frames":
    lowercase			:int							=  torch.tensor([[7.0999e-04, 9.9883e-01, 4.5580e-04]]      )
  elif model_name == "xclip-base-patch16":
    lowercase			:List[str]							=  torch.tensor([[0.0_083, 0.9_681, 0.0_236]]      )
  elif model_name == "xclip-base-patch16-16-frames":
    lowercase			:Dict							=  torch.tensor([[7.6937e-04, 9.9728e-01, 1.9473e-03]]      )
  elif model_name == "xclip-large-patch14":
    lowercase			:List[Any]							=  torch.tensor([[0.0_062, 0.9_864, 0.0_075]]      )
  elif model_name == "xclip-large-patch14-16-frames":
    lowercase			:List[str]							=  torch.tensor([[3.3877e-04, 9.9937e-01, 2.8888e-04]]      )
  # kinetics-600
  elif model_name == "xclip-base-patch16-kinetics-600":
    lowercase			:List[str]							=  torch.tensor([[0.0_555, 0.8_914, 0.0_531]]      )
  elif model_name == "xclip-base-patch16-kinetics-600-16-frames":
    lowercase			:int							=  torch.tensor([[3.8554e-04, 9.9929e-01, 3.2754e-04]]      )
  elif model_name == "xclip-large-patch14-kinetics-600":
    lowercase			:Any							=  torch.tensor([[0.0_036, 0.9_920, 0.0_045]]      )
  # few shot
  elif model_name == "xclip-base-patch16-hmdb-2-shot":
    lowercase			:List[Any]							=  torch.tensor([[7.1890e-06, 9.9994e-01, 5.6559e-05]]      )
  elif model_name == "xclip-base-patch16-hmdb-4-shot":
    lowercase			:Dict							=  torch.tensor([[1.0320e-05, 9.9993e-01, 6.2435e-05]]      )
  elif model_name == "xclip-base-patch16-hmdb-8-shot":
    lowercase			:List[str]							=  torch.tensor([[4.1377e-06, 9.9990e-01, 9.8386e-05]]      )
  elif model_name == "xclip-base-patch16-hmdb-16-shot":
    lowercase			:List[str]							=  torch.tensor([[4.1347e-05, 9.9962e-01, 3.3411e-04]]      )
  elif model_name == "xclip-base-patch16-ucf-2-shot":
    lowercase			:Optional[int]							=  torch.tensor([[8.5857e-05, 9.9928e-01, 6.3291e-04]]      )
  elif model_name == "xclip-base-patch16-ucf-4-shot":
    lowercase			:List[Any]							=  torch.tensor([[8.5857e-05, 9.9928e-01, 6.3291e-04]]      )
  elif model_name == "xclip-base-patch16-ucf-8-shot":
    lowercase			:Union[str, Any]							=  torch.tensor([[0.0_027, 0.9_904, 0.0_070]]      )
  elif model_name == "xclip-base-patch16-ucf-16-shot":
    lowercase			:Optional[int]							=  torch.tensor([[9.8219e-04, 9.9593e-01, 3.0863e-03]]      )
  # zero shot
  elif model_name == "xclip-base-patch16-zero-shot":
    lowercase			:List[Any]							=  torch.tensor([[3.5082e-04, 9.9785e-01, 1.7966e-03]]      )
  else:
    raise ValueError(F"Model name {model_name} not supported"      )
  assert torch.allclose(lowerCamelCase,	lowerCamelCase,	atol=1e-3      )
  print("Looks ok!"      )

  if pytorch_dump_folder_path is not None:
    print(F"Saving model {model_name} to {pytorch_dump_folder_path}"      )
    model.save_pretrained(lowerCamelCase      )

  if push_to_hub:
    print("Pushing model, processor and slow tokenizer files to the hub..."      )
    model.push_to_hub(lowerCamelCase,	organization="nielsr"      )
    processor.push_to_hub(lowerCamelCase,	organization="nielsr"      )
    slow_tokenizer.push_to_hub(lowerCamelCase,	organization="nielsr"      )


if __name__ == "__main__":
     _UpperCAmelCase		:  Tuple	=   argparse.ArgumentParser()
     # Required parameters
     parser.add_argument(
         "--model_name",
         default="xclip-base-patch32",
         type=str,
         help="Name of the model.",
     )
     parser.add_argument(
         "--pytorch_dump_folder_path", default=None, type=str, help="Path to the output PyTorch model directory."
     )
     parser.add_argument(
         "--push_to_hub", action="store_true", help="Whether or not to push the converted model to the 🤗 hub."
     )

     _UpperCAmelCase		:  Dict	=   parser.parse_args()
     convert_xclip_checkpoint(args.model_name, args.pytorch_dump_folder_path, args.push_to_hub)









import gc
import unittest

import numpy as np
import torch

from diffusers import (
    AudioDiffusionPipeline,
    AutoencoderKL,
    DDIMScheduler,
    DDPMScheduler,
    DiffusionPipeline,
    Mel,
    UNetaDConditionModel,
    UNetaDModel,
)
from diffusers.utils import slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu


enable_full_determinism()




class       __lowerCAmelCase				(							unittest.TestCase):






       def     SCREAMING_SNAKE_CASE			(      self:    Union[str, Any] ):
         # clean up the VRAM after each test
         super().tearDown()
         gc.collect()
         torch.cuda.empty_cache()






       @property
       def     SCREAMING_SNAKE_CASE			(      self:    int ):
         torch.manual_seed(0 )
         lowercase			:str							=  UNetaDModel(
             sample_size=(32, 64)				,    in_channels=1				,    out_channels=1				,    layers_per_block=2				,    block_out_channels=(1_28, 1_28)				,    down_block_types=("AttnDownBlock2D", "DownBlock2D")				,    up_block_types=("UpBlock2D", "AttnUpBlock2D")				,    )
         return model






       @property
       def     SCREAMING_SNAKE_CASE			(      self:    Any ):
         torch.manual_seed(0 )
         lowercase			:Dict							=  UNetaDConditionModel(
             sample_size=(64, 32)				,    in_channels=1				,    out_channels=1				,    layers_per_block=2				,    block_out_channels=(1_28, 1_28)				,    down_block_types=("CrossAttnDownBlock2D", "DownBlock2D")				,    up_block_types=("UpBlock2D", "CrossAttnUpBlock2D")				,    cross_attention_dim=10				,    )
         return model






       @property
       def     SCREAMING_SNAKE_CASE			(      self:    Union[str, Any] ):
         torch.manual_seed(0 )
         lowercase			:List[str]							=  AutoencoderKL(
             sample_size=(1_28, 64)				,    in_channels=1				,    out_channels=1				,    latent_channels=1				,    layers_per_block=2				,    block_out_channels=(1_28, 1_28)				,    down_block_types=("DownEncoderBlock2D", "DownEncoderBlock2D")				,    up_block_types=("UpDecoderBlock2D", "UpDecoderBlock2D")				,    )
         lowercase			:List[str]							=  UNetaDModel(
             sample_size=(64, 32)				,    in_channels=1				,    out_channels=1				,    layers_per_block=2				,    block_out_channels=(1_28, 1_28)				,    down_block_types=("AttnDownBlock2D", "DownBlock2D")				,    up_block_types=("UpBlock2D", "AttnUpBlock2D")				,    )
         return vqvae, unet






       @slow
       def     SCREAMING_SNAKE_CASE			(      self:    str ):
         lowercase			:Optional[Any]							=  "cpu"  # ensure determinism for the device-dependent torch.Generator
         lowercase			:Optional[int]							=  Mel(
             x_res=self.dummy_unet.config.sample_size[1]				,    y_res=self.dummy_unet.config.sample_size[0]				,    )

         lowercase			:List[str]							=  DDPMScheduler()
         lowercase			:Tuple							=  AudioDiffusionPipeline(vqvae=_lowerCAmelCase				,    unet=self.dummy_unet				,    mel=_lowerCAmelCase				,    scheduler=_lowerCAmelCase )
         lowercase			:Dict							=  pipe.to(_lowerCAmelCase )
         pipe.set_progress_bar_config(disable=_lowerCAmelCase )

         lowercase			:Tuple							=  torch.Generator(device=_lowerCAmelCase ).manual_seed(42 )
         lowercase			:Optional[int]							=  pipe(generator=_lowerCAmelCase				,    steps=4 )
         lowercase			:List[str]							=  output.audios[0]
         lowercase			:List[str]							=  output.images[0]

         lowercase			:List[str]							=  torch.Generator(device=_lowerCAmelCase ).manual_seed(42 )
         lowercase			:Optional[int]							=  pipe(generator=_lowerCAmelCase				,    steps=4				,    return_dict=_lowerCAmelCase )
         lowercase			:int							=  output[0][0]

         assert audio.shape == (1, (self.dummy_unet.config.sample_size[1] - 1) * mel.hop_length)
         assert (
             image.height == self.dummy_unet.config.sample_size[0]
             and image.width == self.dummy_unet.config.sample_size[1]
         )
         lowercase			:Any							=  np.frombuffer(image.tobytes()				,    dtype="uint8" )[:10]
         lowercase			:Dict							=  np.frombuffer(image_from_tuple.tobytes()				,    dtype="uint8" )[:10]
         lowercase			:List[Any]							=  np.array([69, 2_55, 2_55, 2_55, 0, 0, 77, 1_81, 12, 1_27] )

         assert np.abs(image_slice.flatten() - expected_slice ).max() == 0
         assert np.abs(image_from_tuple_slice.flatten() - expected_slice ).max() == 0

         lowercase			:Optional[Any]							=  Mel(
             x_res=self.dummy_vqvae_and_unet[0].config.sample_size[1]				,    y_res=self.dummy_vqvae_and_unet[0].config.sample_size[0]				,    )

         lowercase			:List[str]							=  DDIMScheduler()
         lowercase			:Tuple							=  self.dummy_vqvae_and_unet
         lowercase			:str							=  AudioDiffusionPipeline(
             vqvae=self.dummy_vqvae_and_unet[0]				,    unet=dummy_vqvae_and_unet[1]				,    mel=_lowerCAmelCase				,    scheduler=_lowerCAmelCase )
         lowercase			:Union[str, Any]							=  pipe.to(_lowerCAmelCase )
         pipe.set_progress_bar_config(disable=_lowerCAmelCase )

         np.random.seed(0 )
         lowercase			:Dict							=  np.random.uniform(-1				,    1				,    ((dummy_vqvae_and_unet[0].config.sample_size[1] - 1) * mel.hop_length,) )
         lowercase			:List[Any]							=  torch.Generator(device=_lowerCAmelCase ).manual_seed(42 )
         lowercase			:Dict							=  pipe(raw_audio=_lowerCAmelCase				,    generator=_lowerCAmelCase				,    start_step=5				,    steps=10 )
         lowercase			:str							=  output.images[0]

         assert (
             image.height == self.dummy_vqvae_and_unet[0].config.sample_size[0]
             and image.width == self.dummy_vqvae_and_unet[0].config.sample_size[1]
         )
         lowercase			:Optional[Any]							=  np.frombuffer(image.tobytes()				,    dtype="uint8" )[:10]
         lowercase			:List[str]							=  np.array([1_20, 1_17, 1_10, 1_09, 1_38, 1_67, 1_38, 1_48, 1_32, 1_21] )

         assert np.abs(image_slice.flatten() - expected_slice ).max() == 0

         lowercase			:Optional[Any]							=  self.dummy_unet_condition
         lowercase			:Optional[int]							=  AudioDiffusionPipeline(
             vqvae=self.dummy_vqvae_and_unet[0]				,    unet=_lowerCAmelCase				,    mel=_lowerCAmelCase				,    scheduler=_lowerCAmelCase )
         lowercase			:int							=  pipe.to(_lowerCAmelCase )
         pipe.set_progress_bar_config(disable=_lowerCAmelCase )

         np.random.seed(0 )
         lowercase			:List[str]							=  torch.rand((1, 1, 10) )
         lowercase			:Union[str, Any]							=  pipe(generator=_lowerCAmelCase				,    encoding=_lowerCAmelCase )
         lowercase			:Tuple							=  output.images[0]
         lowercase			:Optional[Any]							=  np.frombuffer(image.tobytes()				,    dtype="uint8" )[:10]
         lowercase			:Any							=  np.array([1_07, 1_03, 1_20, 1_27, 1_42, 1_22, 1_13, 1_22, 97, 1_11] )

         assert np.abs(image_slice.flatten() - expected_slice ).max() == 0







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






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






       def     SCREAMING_SNAKE_CASE			(      self:    str ):
         lowercase			:Tuple							=  torch_device

         lowercase			:List[Any]							=  DiffusionPipeline.from_pretrained("teticio/audio-diffusion-ddim-256" )
         lowercase			:Dict							=  pipe.to(_lowerCAmelCase )
         pipe.set_progress_bar_config(disable=_lowerCAmelCase )

         lowercase			:Any							=  torch.Generator(device=_lowerCAmelCase ).manual_seed(42 )
         lowercase			:Optional[int]							=  pipe(generator=_lowerCAmelCase )
         lowercase			:List[str]							=  output.audios[0]
         lowercase			:Any							=  output.images[0]

         assert audio.shape == (1, (pipe.unet.config.sample_size[1] - 1) * pipe.mel.hop_length)
         assert image.height == pipe.unet.config.sample_size[0] and image.width == pipe.unet.config.sample_size[1]
         lowercase			:Dict							=  np.frombuffer(image.tobytes()				,    dtype="uint8" )[:10]
         lowercase			:List[str]							=  np.array([1_51, 1_67, 1_54, 1_44, 1_22, 1_34, 1_21, 1_05, 70, 26] )

         assert np.abs(image_slice.flatten() - expected_slice ).max() == 0



