





'''simple docstring'''


import gc
import random
import unittest

import numpy as np
import torch
from transformers import CLIPImageProcessor, CLIPVisionConfig, CLIPVisionModel

from diffusers import HeunDiscreteScheduler, PriorTransformer, ShapEImgaImgPipeline
from diffusers.pipelines.shap_e import ShapERenderer
from diffusers.utils import floats_tensor, load_image, load_numpy, slow
from diffusers.utils.testing_utils import require_torch_gpu, torch_device

from ..test_pipelines_common import PipelineTesterMixin, assert_mean_pixel_difference






class       UpperCamelCase__  (		UpperCamelCase_	,unittest.TestCase     ):
	lowerCAmelCase__					:			Union[str, Any]								=						ShapEImgaImgPipeline
	lowerCAmelCase__					:			Tuple								=						["""image"""]
	lowerCAmelCase__					:			Any								=						["""image"""]
	lowerCAmelCase__					:			int								=						[
	    """num_images_per_prompt""",
	    """num_inference_steps""",
	    """generator""",
	    """latents""",
	    """guidance_scale""",
	    """frame_size""",
	    """output_type""",
	    """return_dict""",
	]
	lowerCAmelCase__					:			Tuple								=						False







	@property
	def 						__a   (				self			:   Tuple		):




				'''simple docstring'''
				return 3_2







	@property
	def 						__a   (				self			:   Any		):




				'''simple docstring'''
				return 3_2







	@property
	def 						__a   (				self			:   Optional[Any]		):




				'''simple docstring'''
				return self.time_input_dim * 4







	@property
	def 						__a   (				self			:   Dict		):




				'''simple docstring'''
				return 8







	@property
	def 						__a   (				self			:   Optional[Any]		):




				'''simple docstring'''
				torch.manual_seed(0		)
				a__     = CLIPVisionConfig(
				    hidden_size=self.text_embedder_hidden_size , image_size=6_4 , projection_dim=self.text_embedder_hidden_size , intermediate_size=3_7 , num_attention_heads=4 , num_channels=3 , num_hidden_layers=5 , patch_size=1 , )

				a__     = CLIPVisionModel(_a		)
				return model







	@property
	def 						__a   (				self			:   int		):




				'''simple docstring'''
				a__     = CLIPImageProcessor(
				    crop_size=2_2_4 , do_center_crop=_a , do_normalize=_a , do_resize=_a , image_mean=[0.48145466, 0.4578275, 0.40821073] , image_std=[0.26862954, 0.26130258, 0.27577711] , resample=3 , size=2_2_4 , )

				return image_processor







	@property
	def 						__a   (				self			:   Tuple		):




				'''simple docstring'''
				torch.manual_seed(0		)

				a__     = {
				    """num_attention_heads""": 2,
				    """attention_head_dim""": 1_6,
				    """embedding_dim""": self.time_input_dim,
				    """num_embeddings""": 3_2,
				    """embedding_proj_dim""": self.text_embedder_hidden_size,
				    """time_embed_dim""": self.time_embed_dim,
				    """num_layers""": 1,
				    """clip_embed_dim""": self.time_input_dim * 2,
				    """additional_embeddings""": 0,
				    """time_embed_act_fn""": """gelu""",
				    """norm_in_type""": """layer""",
				    """embedding_proj_norm_type""": """layer""",
				    """encoder_hid_proj_type""": None,
				    """added_emb_type""": None,
				}

				a__     = PriorTransformer(**_a		)
				return model







	@property
	def 						__a   (				self			:   List[Any]		):




				'''simple docstring'''
				torch.manual_seed(0		)

				a__     = {
				    """param_shapes""": (
				        (self.renderer_dim, 9_3),
				        (self.renderer_dim, 8),
				        (self.renderer_dim, 8),
				        (self.renderer_dim, 8),
				    ),
				    """d_latent""": self.time_input_dim,
				    """d_hidden""": self.renderer_dim,
				    """n_output""": 1_2,
				    """background""": (
				        0.1,
				        0.1,
				        0.1,
				    ),
				}
				a__     = ShapERenderer(**_a		)
				return model







	def 						__a   (				self			:   List[Any]		):




				'''simple docstring'''
				a__     = self.dummy_prior
				a__     = self.dummy_image_encoder
				a__     = self.dummy_image_processor
				a__     = self.dummy_renderer

				a__     = HeunDiscreteScheduler(
				    beta_schedule="exp" , num_train_timesteps=1_0_2_4 , prediction_type="sample" , use_karras_sigmas=_a , clip_sample=_a , clip_sample_range=1.0 , )
				a__     = {
				    """prior""": prior,
				    """image_encoder""": image_encoder,
				    """image_processor""": image_processor,
				    """renderer""": renderer,
				    """scheduler""": scheduler,
				}

				return components







	def 						__a   (				self			:   str , lowerCamelCase			:   Optional[int] , lowerCamelCase			:   Tuple=0		):




				'''simple docstring'''
				a__     = floats_tensor((1, 3, 6_4, 6_4) , rng=random.Random(_a		)		).to(_a		)

				if str(_a		).startswith("mps"		):
							a__     = torch.manual_seed(_a		)
				else:
							a__     = torch.Generator(device=_a		).manual_seed(_a		)
				a__     = {
				    """image""": input_image,
				    """generator""": generator,
				    """num_inference_steps""": 1,
				    """frame_size""": 3_2,
				    """output_type""": """np""",
				}
				return inputs







	def 						__a   (				self			:   Union[str, Any]		):




				'''simple docstring'''
				a__     = """cpu"""

				a__     = self.get_dummy_components()

				a__     = self.pipeline_class(**_a		)
				a__     = pipe.to(_a		)

				pipe.set_progress_bar_config(disable=_a		)

				a__     = pipe(**self.get_dummy_inputs(_a		)		)
				a__     = output.images[0]
				a__     = image[0, -3:, -3:, -1]

				assert image.shape == (2_0, 3_2, 3_2, 3)

				a__     = np.array(
				    [
				        0.00039216,
				        0.00039216,
				        0.00039216,
				        0.00039216,
				        0.00039216,
				        0.00039216,
				        0.00039216,
				        0.00039216,
				        0.00039216,
				    ]		)

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







	def 						__a   (				self			:   Union[str, Any]		):




				'''simple docstring'''
				self._test_inference_batch_consistent(batch_sizes=[1, 2]		)







	def 						__a   (				self			:   Optional[int]		):




				'''simple docstring'''
				a__     = torch_device == """cpu"""
				a__     = True
				self._test_inference_batch_single_identical(
				    batch_size=2 , test_max_difference=_a , relax_max_difference=_a , )







	def 						__a   (				self			:   int		):




				'''simple docstring'''
				a__     = self.get_dummy_components()
				a__     = self.pipeline_class(**_a		)
				a__     = pipe.to(_a		)
				pipe.set_progress_bar_config(disable=_a		)

				a__     = 1
				a__     = 2

				a__     = self.get_dummy_inputs(_a		)

				for key in inputs.keys():
							if key in self.batch_params:
										a__     = batch_size * [inputs[key]]

				a__     = pipe(**_a , num_images_per_prompt=_a		)[0]

				assert images.shape[0] == batch_size * num_images_per_prompt






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







	def 						__a   (				self			:   Any		):




				'''simple docstring'''
				super().tearDown()
				gc.collect()
				torch.cuda.empty_cache()







	def 						__a   (				self			:   Optional[Any]		):




				'''simple docstring'''
				a__     = load_image(
				    "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main" "/shap_e/corgi.png"		)
				a__     = load_numpy(
				    "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"
				    "/shap_e/test_shap_e_img2img_out.npy"		)
				a__     = ShapEImgaImgPipeline.from_pretrained("openai/shap-e-img2img"		)
				a__     = pipe.to(_a		)
				pipe.set_progress_bar_config(disable=_a		)

				a__     = torch.Generator(device=_a		).manual_seed(0		)

				a__     = pipe(
				    _a , generator=_a , guidance_scale=3.0 , num_inference_steps=6_4 , frame_size=6_4 , output_type="np" , ).images[0]

				assert images.shape == (2_0, 6_4, 6_4, 3)

				assert_mean_pixel_difference(_a , _a		)










'''simple docstring'''


import os
from shutil import copyfile
from typing import Any, Dict, List, Optional, Tuple

import sentencepiece as spm

from ...tokenization_utils import PreTrainedTokenizer
from ...utils import logging


lowerCAmelCase_   :   str       =					logging.get_logger(__name__)

lowerCAmelCase_   :   Optional[Any]       =					"▁"

lowerCAmelCase_   :   Union[str, Any]       =					{"vocab_file": "sentencepiece.bpe.model"}

lowerCAmelCase_   :   str       =					{
    "vocab_file": {
        "facebook/xglm-564M": "https://huggingface.co/facebook/xglm-564M/resolve/main/sentencepiece.bpe.model",
    }
}

lowerCAmelCase_   :   Optional[int]       =					{
    "facebook/xglm-564M": 2048,
}






class       UpperCamelCase__  (		__lowerCAmelCase     ):
	lowerCAmelCase__					:			int								=						VOCAB_FILES_NAMES
	lowerCAmelCase__					:			Optional[int]								=						PRETRAINED_VOCAB_FILES_MAP
	lowerCAmelCase__					:			Tuple								=						PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
	lowerCAmelCase__					:			Optional[int]								=						["input_ids", "attention_mask"]







	def __init__(				self			:   Optional[Any] , lowerCamelCase			:   Union[str, Any] , lowerCamelCase			:   int="<s>" , lowerCamelCase			:   List[Any]="</s>" , lowerCamelCase			:   Any="</s>" , lowerCamelCase			:   Tuple="<s>" , lowerCamelCase			:   Optional[Any]="<unk>" , lowerCamelCase			:   Tuple="<pad>" , lowerCamelCase			:   Optional[Dict[str, Any]] = None , **lowerCamelCase			:   Union[str, Any] , ):




				'''simple docstring'''
				a__     = {} if sp_model_kwargs is None else sp_model_kwargs

				# Compatibility with the original tokenizer
				a__     = 7
				a__     = [F'''<madeupword{i}>''' for i in range(self.num_madeup_words		)]

				a__     = kwargs.get("additional_special_tokens" , []		)
				kwargs["additional_special_tokens"] += [
				    word for word in madeup_words if word not in kwargs["additional_special_tokens"]
				]

				super().__init__(
				    bos_token=lowerCamelCase , eos_token=lowerCamelCase , unk_token=lowerCamelCase , sep_token=lowerCamelCase , cls_token=lowerCamelCase , pad_token=lowerCamelCase , sp_model_kwargs=self.sp_model_kwargs , **lowerCamelCase , )

				a__     = spm.SentencePieceProcessor(**self.sp_model_kwargs		)
				self.sp_model.Load(str(lowerCamelCase		)		)
				a__     = vocab_file

				# Original fairseq vocab and spm vocab must be "aligned":
				# Vocab    |    0    |    1    |   2    |    3    |  4  |  5  |  6  |   7   |   8   |  9
				# -------- | ------- | ------- | ------ | ------- | --- | --- | --- | ----- | ----- | ----
				# fairseq  | '<s>'   | '<pad>' | '</s>' | '<unk>' | ',' | '.' | '▁' | 's'   | '▁de' | '-'
				# spm      | '<unk>' | '<s>'   | '</s>' | ','     | '.' | '▁' | 's' | '▁de' | '-'   | '▁a'

				# The first "real" token "," has position 4 in the original fairseq vocab and position 3 in the spm vocab
				a__     = 1

				# Mimic fairseq token-to-id alignment for the first 4 token
				a__     = {"<s>": 0, "<pad>": 1, "</s>": 2, "<unk>": 3}

				a__     = len(self.sp_model		)
				a__     = {F'''<madeupword{i}>''': sp_size + i + self.fairseq_offset for i in range(self.num_madeup_words		)}
				self.fairseq_tokens_to_ids.update(lowerCamelCase		)

				a__     = {v: k for k, v in self.fairseq_tokens_to_ids.items()}







	def __getstate__(				self			:   Optional[Any]		):




				'''simple docstring'''
				a__     = self.__dict__.copy()
				a__     = None
				a__     = self.sp_model.serialized_model_proto()
				return state







	def __setstate__(				self			:   Dict , lowerCamelCase			:   List[Any]		):




				'''simple docstring'''
				a__     = d

				# for backward compatibility
				if not hasattr(self , "sp_model_kwargs"		):
							a__     = {}

				a__     = spm.SentencePieceProcessor(**self.sp_model_kwargs		)
				self.sp_model.LoadFromSerializedProto(self.sp_model_proto		)







	def 						__a   (				self			:   Union[str, Any] , lowerCamelCase			:   List[int] , lowerCamelCase			:   Optional[List[int]] = None		):




				'''simple docstring'''

				if token_ids_a is None:
							return [self.sep_token_id] + token_ids_a
				a__     = [self.sep_token_id]
				return sep + token_ids_a + sep + sep + token_ids_a







	def 						__a   (				self			:   int , lowerCamelCase			:   List[int] , lowerCamelCase			:   Optional[List[int]] = None , lowerCamelCase			:   bool = False		):




				'''simple docstring'''

				if already_has_special_tokens:
							return super().get_special_tokens_mask(
							    token_ids_a=lowerCamelCase , token_ids_a=lowerCamelCase , already_has_special_tokens=lowerCamelCase		)

				if token_ids_a is None:
							return [1] + ([0] * len(lowerCamelCase		))
				return [1] + ([0] * len(lowerCamelCase		)) + [1, 1] + ([0] * len(lowerCamelCase		))







	def 						__a   (				self			:   Dict , lowerCamelCase			:   List[int] , lowerCamelCase			:   Optional[List[int]] = None		):




				'''simple docstring'''
				a__     = [self.sep_token_id]

				if token_ids_a is None:
							return len(sep + token_ids_a		) * [0]
				return len(sep + token_ids_a + sep + sep + token_ids_a		) * [0]







	@property
	def 						__a   (				self			:   List[str]		):




				'''simple docstring'''
				return len(self.sp_model		) + self.fairseq_offset + self.num_madeup_words







	def 						__a   (				self			:   Dict		):




				'''simple docstring'''
				a__     = {self.convert_ids_to_tokens(lowerCamelCase		): i for i in range(self.vocab_size		)}
				vocab.update(self.added_tokens_encoder		)
				return vocab







	def 						__a   (				self			:   int , lowerCamelCase			:   str		):




				'''simple docstring'''
				return self.sp_model.encode(lowerCamelCase , out_type=lowerCamelCase		)







	def 						__a   (				self			:   List[str] , lowerCamelCase			:   Union[str, Any]		):




				'''simple docstring'''
				if token in self.fairseq_tokens_to_ids:
							return self.fairseq_tokens_to_ids[token]
				a__     = self.sp_model.PieceToId(lowerCamelCase		)

				# Need to return unknown token if the SP model returned 0
				return spm_id + self.fairseq_offset if spm_id else self.unk_token_id







	def 						__a   (				self			:   Any , lowerCamelCase			:   List[str]		):




				'''simple docstring'''
				if index in self.fairseq_ids_to_tokens:
							return self.fairseq_ids_to_tokens[index]
				return self.sp_model.IdToPiece(index - self.fairseq_offset		)







	def 						__a   (				self			:   int , lowerCamelCase			:   Optional[Any]		):




				'''simple docstring'''
				a__     = "".join(lowerCamelCase		).replace(lowerCamelCase , " "		).strip()
				return out_string







	def 						__a   (				self			:   int , lowerCamelCase			:   str , lowerCamelCase			:   Optional[str] = None		):




				'''simple docstring'''
				if not os.path.isdir(lowerCamelCase		):
							logger.error(F'''Vocabulary path ({save_directory}) should be a directory'''		)
							return
				a__     = os.path.join(
				    lowerCamelCase , (filename_prefix + "-" if filename_prefix else "") + VOCAB_FILES_NAMES["vocab_file"]		)

				if os.path.abspath(self.vocab_file		) != os.path.abspath(lowerCamelCase		) and os.path.isfile(self.vocab_file		):
							copyfile(self.vocab_file , lowerCamelCase		)
				elif not os.path.isfile(self.vocab_file		):
							with open(lowerCamelCase , "wb"		) as fi:
										a__     = self.sp_model.serialized_model_proto()
										fi.write(lowerCamelCase		)

				return (out_vocab_file,)



