






'''simple docstring'''







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					(      lowercase_       :  Optional[int]   ,       lowercase_       :  Union[str, Any]   ):
       """simple docstring"""




       a_						     =			XCLIPTextConfig()

       # derive patch size from model name
       a_						     =			model_name.find('patch'   )
       a_						     =			int(model_name[start_idx + len('patch'   ) : start_idx + len('patch'   ) + 2]   )
       a_						     =			XCLIPVisionConfig(patch_size=lowercase_   ,       num_frames=lowercase_   )

       if "large" in model_name:
              a_						     =			768
              a_						     =			3_072
              a_						     =			12

              a_						     =			1_024
              a_						     =			4_096
              a_						     =			16
              a_						     =			24
              a_						     =			768
              a_						     =			3_072

       if model_name == "xclip-large-patch14-16-frames":
              a_						     =			336

       a_						     =			XCLIPConfig.from_text_vision_configs(lowercase_   ,       lowercase_   )

       if "large" in model_name:
              a_						     =			768

       return config


def       __UpperCamelCase					(      lowercase_       :  int   ):
       """simple docstring"""




       if name == "token_embedding.weight":
              a_						     =			name.replace('token_embedding.weight'   ,       'text_model.embeddings.token_embedding.weight'   )
       if name == "positional_embedding":
              a_						     =			name.replace('positional_embedding'   ,       'text_model.embeddings.position_embedding.weight'   )
       if "ln_1" in name:
              a_						     =			name.replace('ln_1'   ,       'layer_norm1'   )
       if "ln_2" in name:
              a_						     =			name.replace('ln_2'   ,       'layer_norm2'   )
       if "c_fc" in name:
              a_						     =			name.replace('c_fc'   ,       'fc1'   )
       if "c_proj" in name:
              a_						     =			name.replace('c_proj'   ,       'fc2'   )
       if name.startswith('transformer.resblocks'   ):
              a_						     =			name.replace('transformer.resblocks'   ,       'text_model.encoder.layers'   )
       if "attn.out_proj" in name and "message" not in name:
              a_						     =			name.replace('attn.out_proj'   ,       'self_attn.out_proj'   )
       if "ln_final" in name:
              a_						     =			name.replace('ln_final'   ,       'text_model.final_layer_norm'   )
       # visual encoder
       if name == "visual.class_embedding":
              a_						     =			name.replace('visual.class_embedding'   ,       'vision_model.embeddings.class_embedding'   )
       if name == "visual.positional_embedding":
              a_						     =			name.replace('visual.positional_embedding'   ,       'vision_model.embeddings.position_embedding.weight'   )
       if name.startswith('visual.transformer.resblocks'   ):
              a_						     =			name.replace('visual.transformer.resblocks'   ,       'vision_model.encoder.layers'   )
       if "visual.conv1" in name:
              a_						     =			name.replace('visual.conv1'   ,       'vision_model.embeddings.patch_embedding'   )
       if "visual.ln_pre" in name:
              a_						     =			name.replace('visual.ln_pre'   ,       'vision_model.pre_layernorm'   )
       if "visual.ln_post" in name:
              a_						     =			name.replace('visual.ln_post'   ,       'vision_model.post_layernorm'   )
       if "visual.proj" in name:
              a_						     =			name.replace('visual.proj'   ,       'visual_projection.weight'   )
       if "text_projection" in name:
              a_						     =			name.replace('text_projection'   ,       'text_projection.weight'   )
       # things on top
       if "prompts_visual_proj" in name:
              a_						     =			name.replace('prompts_visual_proj'   ,       'prompts_visual_projection'   )
       if "prompts_visual_ln" in name:
              a_						     =			name.replace('prompts_visual_ln'   ,       'prompts_visual_layernorm'   )
       # mit
       if name == "mit.positional_embedding":
              a_						     =			name.replace('positional'   ,       'position'   )
       if name.startswith('mit.resblocks'   ):
              a_						     =			name.replace('mit.resblocks'   ,       'mit.encoder.layers'   )
       # prompts generator
       if name.startswith('prompts_generator.norm'   ):
              a_						     =			name.replace('prompts_generator.norm'   ,       'prompts_generator.layernorm'   )

       return name


def       __UpperCamelCase					(      lowercase_       :  Optional[Any]   ,       lowercase_       :  Union[str, Any]   ):
       """simple docstring"""




       for key in orig_state_dict.copy().keys():
              a_						     =			orig_state_dict.pop(lowercase_   )

              if "attn.in_proj" in key:
                     a_						     =			key.split('.'   )
                     if key.startswith('visual'   ):
                            a_						     =			key_split[3]
                            a_						     =			config.vision_config.hidden_size
                            if "message_attn" in key:
                                   if "weight" in key:
                                          a_						     =			val[
                                              :dim, :
                                          ]
                                          a_						     =			val[
                                              dim : dim * 2, :
                                          ]
                                          a_						     =			val[
                                              -dim:, :
                                          ]
                                   else:
                                          a_						     =			val[
                                              :dim
                                          ]
                                          a_						     =			val[
                                              dim : dim * 2
                                          ]
                                          a_						     =			val[
                                              -dim:
                                          ]
                            else:
                                   if "weight" in key:
                                          a_						     =			val[
                                              :dim, :
                                          ]
                                          a_						     =			val[
                                              dim : dim * 2, :
                                          ]
                                          a_						     =			val[
                                              -dim:, :
                                          ]
                                   else:
                                          a_						     =			val[:dim]
                                          a_						     =			val[
                                              dim : dim * 2
                                          ]
                                          a_						     =			val[-dim:]
                     elif key.startswith('mit'   ):
                            a_						     =			key_split[2]
                            a_						     =			config.vision_config.mit_hidden_size
                            if "weight" in key:
                                   a_						     =			val[:dim, :]
                                   a_						     =			val[dim : dim * 2, :]
                                   a_						     =			val[-dim:, :]
                            else:
                                   a_						     =			val[:dim]
                                   a_						     =			val[dim : dim * 2]
                                   a_						     =			val[-dim:]
                     else:
                            a_						     =			key_split[2]
                            a_						     =			config.text_config.hidden_size
                            if "weight" in key:
                                   a_						     =			val[:dim, :]
                                   a_						     =			val[
                                       dim : dim * 2, :
                                   ]
                                   a_						     =			val[-dim:, :]
                            else:
                                   a_						     =			val[:dim]
                                   a_						     =			val[
                                       dim : dim * 2
                                   ]
                                   a_						     =			val[-dim:]
              else:
                     a_						     =			rename_key(lowercase_   )
                     if new_key_name in ["visual_projection.weight", "text_projection.weight"]:
                            a_						     =			val.T
                     a_						     =			val

       return orig_state_dict


def       __UpperCamelCase					(      lowercase_       :  Optional[int]   ):
       """simple docstring"""




       if num_frames == 8:
              a_						     =			'eating_spaghetti_8_frames.npy'
       elif num_frames == 16:
              a_						     =			'eating_spaghetti.npy'
       elif num_frames == 32:
              a_						     =			'eating_spaghetti_32_frames.npy'
       a_						     =			hf_hub_download(
           repo_id='hf-internal-testing/spaghetti-video'   ,       filename=lowercase_   ,       repo_type='dataset'   ,       )
       a_						     =			np.load(lowercase_   )
       return list(lowercase_   )


def       __UpperCamelCase					(      lowercase_       :  Optional[Any]   ,       lowercase_       :  Any=None   ,       lowercase_       :  List[Any]=False   ):
       """simple docstring"""




       a_						     =			{
           # 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',
       }

       a_						     =			model_to_url[model_name]
       a_						     =			8
       if "16-frames" in model_name:
              a_						     =			16
       elif "shot" in model_name:
              a_						     =			32

       a_						     =			get_xclip_config(lowercase_   ,       lowercase_   )
       a_						     =			XCLIPModel(lowercase_   )
       model.eval()

       if "drive" in checkpoint_url:
              a_						     =			'pytorch_model.bin'
              gdown.cached_download(lowercase_   ,       lowercase_   ,       quiet=lowercase_   )
              a_						     =			torch.load(lowercase_   ,       map_location='cpu'   )['model']
       else:
              a_						     =			torch.hub.load_state_dict_from_url(lowercase_   )['model']

       a_						     =			convert_state_dict(lowercase_   ,       lowercase_   )

       a_						     =			XCLIPModel(lowercase_   )
       a_     ,	a_						     =			model.load_state_dict(lowercase_   ,       strict=lowercase_   )
       assert missing_keys == ["text_model.embeddings.position_ids", "vision_model.embeddings.position_ids"]
       model.eval()

       a_						     =			336 if model_name == 'xclip-large-patch14-16-frames' else 224
       a_						     =			VideoMAEImageProcessor(size=lowercase_   )
       a_						     =			CLIPTokenizer.from_pretrained('openai/clip-vit-base-patch32'   )
       a_						     =			CLIPTokenizerFast.from_pretrained('openai/clip-vit-base-patch32'   )
       a_						     =			XCLIPProcessor(image_processor=lowercase_   ,       tokenizer=lowercase_   )

       a_						     =			prepare_video(lowercase_   )
       a_						     =			processor(
           text=['playing sports', 'eating spaghetti', 'go shopping']   ,       videos=lowercase_   ,       return_tensors='pt'   ,       padding=lowercase_   )

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

       with torch.no_grad():
              a_						     =			model(**lowercase_   )

       # Verify outputs
       a_						     =			outputs.logits_per_video
       a_						     =			logits_per_video.softmax(dim=1   )
       print('Probs:'   ,       lowercase_   )
       # kinetics-400
       if model_name == "xclip-base-patch32":
              a_						     =			torch.tensor([[0.0019, 0.9951, 0.0030]]   )
       elif model_name == "xclip-base-patch32-16-frames":
              a_						     =			torch.tensor([[7.0999E-04, 9.9883E-01, 4.5580E-04]]   )
       elif model_name == "xclip-base-patch16":
              a_						     =			torch.tensor([[0.0083, 0.9681, 0.0236]]   )
       elif model_name == "xclip-base-patch16-16-frames":
              a_						     =			torch.tensor([[7.6937E-04, 9.9728E-01, 1.9473E-03]]   )
       elif model_name == "xclip-large-patch14":
              a_						     =			torch.tensor([[0.0062, 0.9864, 0.0075]]   )
       elif model_name == "xclip-large-patch14-16-frames":
              a_						     =			torch.tensor([[3.3877E-04, 9.9937E-01, 2.8888E-04]]   )
       # kinetics-600
       elif model_name == "xclip-base-patch16-kinetics-600":
              a_						     =			torch.tensor([[0.0555, 0.8914, 0.0531]]   )
       elif model_name == "xclip-base-patch16-kinetics-600-16-frames":
              a_						     =			torch.tensor([[3.8554E-04, 9.9929E-01, 3.2754E-04]]   )
       elif model_name == "xclip-large-patch14-kinetics-600":
              a_						     =			torch.tensor([[0.0036, 0.9920, 0.0045]]   )
       # few shot
       elif model_name == "xclip-base-patch16-hmdb-2-shot":
              a_						     =			torch.tensor([[7.1890E-06, 9.9994E-01, 5.6559E-05]]   )
       elif model_name == "xclip-base-patch16-hmdb-4-shot":
              a_						     =			torch.tensor([[1.0320E-05, 9.9993E-01, 6.2435E-05]]   )
       elif model_name == "xclip-base-patch16-hmdb-8-shot":
              a_						     =			torch.tensor([[4.1377E-06, 9.9990E-01, 9.8386E-05]]   )
       elif model_name == "xclip-base-patch16-hmdb-16-shot":
              a_						     =			torch.tensor([[4.1347E-05, 9.9962E-01, 3.3411E-04]]   )
       elif model_name == "xclip-base-patch16-ucf-2-shot":
              a_						     =			torch.tensor([[8.5857E-05, 9.9928E-01, 6.3291E-04]]   )
       elif model_name == "xclip-base-patch16-ucf-4-shot":
              a_						     =			torch.tensor([[8.5857E-05, 9.9928E-01, 6.3291E-04]]   )
       elif model_name == "xclip-base-patch16-ucf-8-shot":
              a_						     =			torch.tensor([[0.0027, 0.9904, 0.0070]]   )
       elif model_name == "xclip-base-patch16-ucf-16-shot":
              a_						     =			torch.tensor([[9.8219E-04, 9.9593E-01, 3.0863E-03]]   )
       # zero shot
       elif model_name == "xclip-base-patch16-zero-shot":
              a_						     =			torch.tensor([[3.5082E-04, 9.9785E-01, 1.7966E-03]]   )
       else:
              raise ValueError(F'Model name {model_name} not supported'   )
       assert torch.allclose(lowercase_   ,       lowercase_   ,       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(lowercase_   )

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


if __name__ == "__main__":
  __lowerCAmelCase        =							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."
  )

  __lowerCAmelCase        =							parser.parse_args()
  convert_xclip_checkpoint(args.model_name, args.pytorch_dump_folder_path, args.push_to_hub)












'''simple docstring'''







import gc
import unittest

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

from diffusers import TransformeraDModel, VQDiffusionPipeline, VQDiffusionScheduler, VQModel
from diffusers.pipelines.vq_diffusion.pipeline_vq_diffusion import LearnedClassifierFreeSamplingEmbeddings
from diffusers.utils import load_numpy, slow, torch_device
from diffusers.utils.testing_utils import require_torch_gpu


__lowerCAmelCase        =							False




class   __SCREAMING_SNAKE_CASE       (unittest.TestCase       ):


  """simple docstring"""


  def    _a       (			self	):






         """simple docstring"""


         super().tearDown()
         gc.collect()
         torch.cuda.empty_cache()







  @property
  def    _a       (			self	):






         """simple docstring"""


         return 12







  @property
  def    _a       (			self	):






         """simple docstring"""


         return 12







  @property
  def    _a       (			self	):






         """simple docstring"""


         return 32







  @property
  def    _a       (			self	):






         """simple docstring"""


         torch.manual_seed(0	)
         a_						     =			VQModel(
             block_out_channels=[32, 64]   ,	in_channels=3   ,	out_channels=3   ,	down_block_types=['DownEncoderBlock2D', 'DownEncoderBlock2D']   ,	up_block_types=['UpDecoderBlock2D', 'UpDecoderBlock2D']   ,	latent_channels=3   ,	num_vq_embeddings=self.num_embed   ,	vq_embed_dim=3   ,	)
         return model







  @property
  def    _a       (			self	):






         """simple docstring"""


         a_						     =			CLIPTokenizer.from_pretrained('hf-internal-testing/tiny-random-clip'	)
         return tokenizer







  @property
  def    _a       (			self	):






         """simple docstring"""


         torch.manual_seed(0	)
         a_						     =			CLIPTextConfig(
             bos_token_id=0   ,	eos_token_id=2   ,	hidden_size=self.text_embedder_hidden_size   ,	intermediate_size=37   ,	layer_norm_eps=1e-05   ,	num_attention_heads=4   ,	num_hidden_layers=5   ,	pad_token_id=1   ,	vocab_size=1_000   ,	)
         return CLIPTextModel(UpperCamelCase__	)







  @property
  def    _a       (			self	):






         """simple docstring"""


         torch.manual_seed(0	)

         a_						     =			12
         a_						     =			12

         a_						     =			{
             'attention_bias': True,
             'cross_attention_dim': 32,
             'attention_head_dim': height * width,
             'num_attention_heads': 1,
             'num_vector_embeds': self.num_embed,
             'num_embeds_ada_norm': self.num_embeds_ada_norm,
             'norm_num_groups': 32,
             'sample_size': width,
             'activation_fn': 'geglu-approximate',
         }

         a_						     =			TransformeraDModel(**UpperCamelCase__	)
         return model







  def    _a       (			self	):






         """simple docstring"""


         a_						     =			'cpu'

         a_						     =			self.dummy_vqvae
         a_						     =			self.dummy_text_encoder
         a_						     =			self.dummy_tokenizer
         a_						     =			self.dummy_transformer
         a_						     =			VQDiffusionScheduler(self.num_embed	)
         a_						     =			LearnedClassifierFreeSamplingEmbeddings(learnable=UpperCamelCase__	)

         a_						     =			VQDiffusionPipeline(
             vqvae=UpperCamelCase__   ,	text_encoder=UpperCamelCase__   ,	tokenizer=UpperCamelCase__   ,	transformer=UpperCamelCase__   ,	scheduler=UpperCamelCase__   ,	learned_classifier_free_sampling_embeddings=UpperCamelCase__   ,	)
         a_						     =			pipe.to(UpperCamelCase__	)
         pipe.set_progress_bar_config(disable=UpperCamelCase__	)

         a_						     =			'teddy bear playing in the pool'

         a_						     =			torch.Generator(device=UpperCamelCase__	).manual_seed(0	)
         a_						     =			pipe([prompt]   ,	generator=UpperCamelCase__   ,	num_inference_steps=2   ,	output_type='np'	)
         a_						     =			output.images

         a_						     =			torch.Generator(device=UpperCamelCase__	).manual_seed(0	)
         a_						     =			pipe(
             [prompt]   ,	generator=UpperCamelCase__   ,	output_type='np'   ,	return_dict=UpperCamelCase__   ,	num_inference_steps=2	)[0]

         a_						     =			image[0, -3:, -3:, -1]
         a_						     =			image_from_tuple[0, -3:, -3:, -1]

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

         a_						     =			np.array([0.6_551, 0.6_168, 0.5_008, 0.5_676, 0.5_659, 0.4_295, 0.6_073, 0.5_599, 0.4_992]	)

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







  def    _a       (			self	):






         """simple docstring"""


         a_						     =			'cpu'

         a_						     =			self.dummy_vqvae
         a_						     =			self.dummy_text_encoder
         a_						     =			self.dummy_tokenizer
         a_						     =			self.dummy_transformer
         a_						     =			VQDiffusionScheduler(self.num_embed	)
         a_						     =			LearnedClassifierFreeSamplingEmbeddings(
             learnable=UpperCamelCase__   ,	hidden_size=self.text_embedder_hidden_size   ,	length=tokenizer.model_max_length	)

         a_						     =			VQDiffusionPipeline(
             vqvae=UpperCamelCase__   ,	text_encoder=UpperCamelCase__   ,	tokenizer=UpperCamelCase__   ,	transformer=UpperCamelCase__   ,	scheduler=UpperCamelCase__   ,	learned_classifier_free_sampling_embeddings=UpperCamelCase__   ,	)
         a_						     =			pipe.to(UpperCamelCase__	)
         pipe.set_progress_bar_config(disable=UpperCamelCase__	)

         a_						     =			'teddy bear playing in the pool'

         a_						     =			torch.Generator(device=UpperCamelCase__	).manual_seed(0	)
         a_						     =			pipe([prompt]   ,	generator=UpperCamelCase__   ,	num_inference_steps=2   ,	output_type='np'	)
         a_						     =			output.images

         a_						     =			torch.Generator(device=UpperCamelCase__	).manual_seed(0	)
         a_						     =			pipe(
             [prompt]   ,	generator=UpperCamelCase__   ,	output_type='np'   ,	return_dict=UpperCamelCase__   ,	num_inference_steps=2	)[0]

         a_						     =			image[0, -3:, -3:, -1]
         a_						     =			image_from_tuple[0, -3:, -3:, -1]

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

         a_						     =			np.array([0.6_693, 0.6_075, 0.4_959, 0.5_701, 0.5_583, 0.4_333, 0.6_171, 0.5_684, 0.4_988]	)

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





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


  """simple docstring"""


  def    _a       (			self	):






         """simple docstring"""


         super().tearDown()
         gc.collect()
         torch.cuda.empty_cache()







  def    _a       (			self	):






         """simple docstring"""


         a_						     =			load_numpy(
             'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'
             '/vq_diffusion/teddy_bear_pool_classifier_free_sampling.npy'	)

         a_						     =			VQDiffusionPipeline.from_pretrained('microsoft/vq-diffusion-ithq'	)
         a_						     =			pipeline.to(UpperCamelCase__	)
         pipeline.set_progress_bar_config(disable=UpperCamelCase__	)

         # requires GPU generator for gumbel softmax
         # don't use GPU generator in tests though
         a_						     =			torch.Generator(device=UpperCamelCase__	).manual_seed(0	)
         a_						     =			pipeline(
             'teddy bear playing in the pool'   ,	num_images_per_prompt=1   ,	generator=UpperCamelCase__   ,	output_type='np'   ,	)

         a_						     =			output.images[0]

         assert image.shape == (256, 256, 3)
         assert np.abs(expected_image - image	).max() < 2.0




