




'''simple docstring'''







import os
import re
import urllib.parse
from pathlib import Path
from typing import Callable, List, Optional, Union
from zipfile import ZipFile

from ..utils.file_utils import cached_path, hf_github_url
from ..utils.logging import get_logger
from ..utils.version import Version


A									=get_logger(__name__)


class     _a						:
       __a       :  Tuple            =      """dummy_data"""
       __a       :  List[str]            =      """datasets"""
       __a       :  Tuple            =      False
       def __init__(			self		:			int     ,       lowercase		:			str     ,       lowercase		:			str     ,       lowercase		:			Union[Version, str]     ,       lowercase		:			Optional[str] = None     ,       lowercase		:			bool = False     ,       lowercase		:			bool = True     ,       lowercase		:			Optional[List[Callable]] = None     ,       ):



           '''simple docstring'''

           UpperCAmelCase     =	0
           UpperCAmelCase     =	dataset_name
           UpperCAmelCase     =	cache_dir
           UpperCAmelCase     =	use_local_dummy_data
           UpperCAmelCase     =	config
           # download_callbacks take a single url as input
           UpperCAmelCase     =	download_callbacks or []
           # if False, it doesn't load existing files and it returns the paths of the dummy files relative
           # to the dummy_data zip file root
           UpperCAmelCase     =	load_existing_dummy_data

           # TODO(PVP, QL) might need to make this more general
           UpperCAmelCase     =	str(lowercase       )
           # to be downloaded
           UpperCAmelCase     =	None
           UpperCAmelCase     =	None
       @property
       def 					A   (			self		:			Dict       ):



           '''simple docstring'''

           if self._dummy_file is None:
               UpperCAmelCase     =	self.download_dummy_data()
           return self._dummy_file
       @property
       def 					A   (			self		:			Union[str, Any]       ):



           '''simple docstring'''

           if self.config is not None:
               # structure is dummy / config_name / version_name
               return os.path.join('''dummy'''     ,       self.config.name     ,       self.version_name       )
           # structure is dummy / version_name
           return os.path.join('''dummy'''     ,       self.version_name       )
       @property
       def 					A   (			self		:			int       ):



           '''simple docstring'''

           return os.path.join(self.dummy_data_folder     ,       '''dummy_data.zip'''       )
       def 					A   (			self		:			Any       ):



           '''simple docstring'''

           UpperCAmelCase     =	(
               self.local_path_to_dummy_data if self.use_local_dummy_data is True else self.github_path_to_dummy_data
           )

           UpperCAmelCase     =	cached_path(
               lowercase     ,       cache_dir=self.cache_dir     ,       extract_compressed_file=lowercase     ,       force_extract=lowercase       )

           return os.path.join(lowercase     ,       self.dummy_file_name       )
       @property
       def 					A   (			self		:			Dict       ):



           '''simple docstring'''

           return os.path.join(self.datasets_scripts_dir     ,       self.dataset_name     ,       self.dummy_zip_file       )
       @property
       def 					A   (			self		:			Optional[Any]       ):



           '''simple docstring'''

           if self._bucket_url is None:
               UpperCAmelCase     =	hf_github_url(self.dataset_name     ,       self.dummy_zip_file.replace(os.sep     ,       '''/'''       )       )
           return self._bucket_url
       @property
       def 					A   (			self		:			str       ):



           '''simple docstring'''

           if os.path.isdir(self.dummy_file       ):
               return self.dummy_file
           # else cut off path to file -> example `xsum`.
           return "/".join(self.dummy_file.replace(os.sep     ,       '''/'''       ).split('''/'''       )[:-1]       )
       def 					A   (			self		:			Tuple     ,       lowercase		:			List[str]     ,       *lowercase		:			Optional[int]       ):



           '''simple docstring'''

           if self.load_existing_dummy_data:
               # dummy data is downloaded and tested
               UpperCAmelCase     =	self.dummy_file
           else:
               # dummy data cannot be downloaded and only the path to dummy file is returned
               UpperCAmelCase     =	self.dummy_file_name

           # special case when data_url is a dict
           if isinstance(lowercase     ,       lowercase       ):
               return self.create_dummy_data_dict(lowercase     ,       lowercase       )
           elif isinstance(lowercase     ,       (list, tuple)       ):
               return self.create_dummy_data_list(lowercase     ,       lowercase       )
           else:
               return self.create_dummy_data_single(lowercase     ,       lowercase       )
       def 					A   (			self		:			Optional[int]     ,       lowercase		:			List[str]     ,       *lowercase		:			List[Any]       ):



           '''simple docstring'''

           return self.download_and_extract(lowercase       )
       def 					A   (			self		:			Union[str, Any]     ,       lowercase		:			List[str]     ,       lowercase		:			Any       ):



           '''simple docstring'''

           return self.download_and_extract(lowercase       )
       def 					A   (			self		:			Optional[Any]     ,       lowercase		:			str     ,       *lowercase		:			Optional[int]     ,       **lowercase		:			str       ):



           '''simple docstring'''

           return path
       def 					A   (			self		:			Dict       ):



           '''simple docstring'''

           return {}
       def 					A   (			self		:			str     ,       lowercase		:			Optional[Any]     ,       lowercase		:			Optional[Any]       ):



           '''simple docstring'''

           UpperCAmelCase     =	{}
           for key, single_urls in data_url.items():
               for download_callback in self.download_callbacks:
                   if isinstance(lowercase     ,       lowercase       ):
                       for single_url in single_urls:
                           download_callback(lowercase       )
                   else:
                       UpperCAmelCase     =	single_urls
                       download_callback(lowercase       )
            # we force the name of each key to be the last file / folder name of the url path
            # if the url has arguments, we need to encode them with urllib.parse.quote_plus
               if isinstance(lowercase     ,       lowercase       ):
                   UpperCAmelCase     =	[os.path.join(lowercase     ,       urllib.parse.quote_plus(Path(lowercase       ).name       )       ) for x in single_urls]
               else:
                   UpperCAmelCase     =	single_urls
                   UpperCAmelCase     =	os.path.join(lowercase     ,       urllib.parse.quote_plus(Path(lowercase       ).name       )       )
               UpperCAmelCase     =	value

           # make sure that values are unique
           if all(isinstance(lowercase     ,       lowercase       ) for i in dummy_data_dict.values()       ) and len(set(dummy_data_dict.values()       )       ) < len(
               dummy_data_dict.values()       ):
               # append key to value to make its name unique
               UpperCAmelCase     =	{key: value + key for key, value in dummy_data_dict.items()}

           return dummy_data_dict
       def 					A   (			self		:			Union[str, Any]     ,       lowercase		:			List[Any]     ,       lowercase		:			Optional[int]       ):



           '''simple docstring'''

           UpperCAmelCase     =	[]
           # trick: if there are many shards named like `data.txt-000001-of-00300`, only use the first one
           UpperCAmelCase     =	all(bool(re.findall('''[0-9]{3,}-of-[0-9]{3,}'''     ,       lowercase       )       ) for url in data_url       )
           UpperCAmelCase     =	all(
               url.startswith('''https://ftp.ncbi.nlm.nih.gov/pubmed/baseline/pubmed'''       ) for url in data_url       )
           if data_url and (is_tf_records or is_pubmed_records):
               UpperCAmelCase     =	[data_url[0]] * len(lowercase       )
           for single_url in data_url:
               for download_callback in self.download_callbacks:
                   download_callback(lowercase       )
               # we force the name of each key to be the last file / folder name of the url path
               # if the url has arguments, we need to encode them with urllib.parse.quote_plus
               UpperCAmelCase     =	os.path.join(lowercase     ,       urllib.parse.quote_plus(single_url.split('''/'''       )[-1]       )       )
               dummy_data_list.append(lowercase       )
           return dummy_data_list
       def 					A   (			self		:			List[Any]     ,       lowercase		:			Tuple     ,       lowercase		:			Any       ):



           '''simple docstring'''

           for download_callback in self.download_callbacks:
               download_callback(lowercase       )
           # we force the name of each key to be the last file / folder name of the url path
           # if the url has arguments, we need to encode them with urllib.parse.quote_plus
           UpperCAmelCase     =	os.path.join(lowercase     ,       urllib.parse.quote_plus(data_url.split('''/'''       )[-1]       )       )
           if os.path.exists(lowercase       ) or not self.load_existing_dummy_data:
               return value
           else:
               # Backward compatibility, maybe deprecate at one point.
               # For many datasets with single url calls to dl_manager.download_and_extract,
               # the dummy_data.zip file is actually the zipped downloaded file
               # while now we expected the dummy_data.zip file to be a directory containing
               # the downloaded file.
               return path_to_dummy_data
       def 					A   (			self		:			Any       ):



           '''simple docstring'''

           pass
       def 					A   (			self		:			Union[str, Any]       ):



           '''simple docstring'''

           pass
       def 					A   (			self		:			Dict     ,       lowercase		:			List[Any]       ):



           '''simple docstring'''

           def _iter_archive_members(lowercase		:			Union[str, Any]       ):
               # this preserves the order of the members inside the ZIP archive
               UpperCAmelCase     =	Path(self.dummy_file       ).parent
               UpperCAmelCase     =	path.relative_to(lowercase       )
               with ZipFile(self.local_path_to_dummy_data       ) as zip_file:
                   UpperCAmelCase     =	zip_file.namelist()
               for member in members:
                   if member.startswith(relative_path.as_posix()       ):
                       yield dummy_parent_path.joinpath(lowercase       )

           UpperCAmelCase     =	Path(lowercase       )
           UpperCAmelCase     =	_iter_archive_members(lowercase       ) if self.use_local_dummy_data else path.rglob('''*'''       )
           for file_path in file_paths:
               if file_path.is_file() and not file_path.name.startswith(('''.''', '''__''')       ):
                   yield file_path.relative_to(lowercase       ).as_posix(), file_path.open('''rb'''       )







       def 					A   (			self		:			str     ,       lowercase		:			Any       ):



           '''simple docstring'''

           if not isinstance(lowercase     ,       lowercase       ):
               UpperCAmelCase     =	[paths]
           for path in paths:
               if os.path.isfile(lowercase       ):
                   if os.path.basename(lowercase       ).startswith(('''.''', '''__''')       ):
                       return
                   yield path
               else:
                   for dirpath, dirnames, filenames in os.walk(lowercase       ):
                       if os.path.basename(lowercase       ).startswith(('''.''', '''__''')       ):
                           continue
                       dirnames.sort()
                       for filename in sorted(lowercase       ):
                           if filename.startswith(('''.''', '''__''')       ):
                               continue
                           yield os.path.join(lowercase     ,       lowercase       )










'''simple docstring'''







import gc
import unittest

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

from diffusers import HeunDiscreteScheduler, PriorTransformer, ShapEPipeline
from diffusers.pipelines.shap_e import ShapERenderer
from diffusers.utils import 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     _a						(			__a							,							unittest.TestCase				):
       __a       :  Tuple            =      ShapEPipeline
       __a       :  Tuple            =      ["""prompt"""]
       __a       :  Tuple            =      ["""prompt"""]
       __a       :  str            =      [
           """num_images_per_prompt""",
           """num_inference_steps""",
           """generator""",
           """latents""",
           """guidance_scale""",
           """frame_size""",
           """output_type""",
           """return_dict""",
       ]
       __a       :  Union[str, Any]            =      False
       @property
       def 					A   (			self		:			Optional[Any]       ):



           '''simple docstring'''

           return 32
       @property
       def 					A   (			self		:			List[Any]       ):



           '''simple docstring'''

           return 32
       @property
       def 					A   (			self		:			Union[str, 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'''

           UpperCAmelCase     =	CLIPTokenizer.from_pretrained('''hf-internal-testing/tiny-random-clip'''       )
           return tokenizer
       @property
       def 					A   (			self		:			Union[str, Any]       ):



           '''simple docstring'''

           torch.manual_seed(0       )
           UpperCAmelCase     =	CLIPTextConfig(
               bos_token_id=0     ,       eos_token_id=2     ,       hidden_size=self.text_embedder_hidden_size     ,       projection_dim=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 CLIPTextModelWithProjection(lowercase       )
       @property
       def 					A   (			self		:			Union[str, Any]       ):



           '''simple docstring'''

           torch.manual_seed(0       )

           UpperCAmelCase     =	{
               '''num_attention_heads''': 2,
               '''attention_head_dim''': 16,
               '''embedding_dim''': self.time_input_dim,
               '''num_embeddings''': 32,
               '''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''',
               '''encoder_hid_proj_type''': None,
               '''added_emb_type''': None,
           }

           UpperCAmelCase     =	PriorTransformer(**lowercase       )
           return model
       @property
       def 					A   (			self		:			List[Any]       ):



           '''simple docstring'''

           torch.manual_seed(0       )

           UpperCAmelCase     =	{
               '''param_shapes''': (
                   (self.renderer_dim, 93),
                   (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''': 12,
               '''background''': (
                   0.1,
                   0.1,
                   0.1,
               ),
           }
           UpperCAmelCase     =	ShapERenderer(**lowercase       )
           return model
       def 					A   (			self		:			Tuple       ):



           '''simple docstring'''

           UpperCAmelCase     =	self.dummy_prior
           UpperCAmelCase     =	self.dummy_text_encoder
           UpperCAmelCase     =	self.dummy_tokenizer
           UpperCAmelCase     =	self.dummy_renderer

           UpperCAmelCase     =	HeunDiscreteScheduler(
               beta_schedule='''exp'''     ,       num_train_timesteps=1_024     ,       prediction_type='''sample'''     ,       use_karras_sigmas=lowercase     ,       clip_sample=lowercase     ,       clip_sample_range=1.0     ,       )
           UpperCAmelCase     =	{
               '''prior''': prior,
               '''text_encoder''': text_encoder,
               '''tokenizer''': tokenizer,
               '''renderer''': renderer,
               '''scheduler''': scheduler,
           }

           return components
       def 					A   (			self		:			Dict     ,       lowercase		:			Optional[int]     ,       lowercase		:			Optional[int]=0       ):



           '''simple docstring'''

           if str(lowercase       ).startswith('''mps'''       ):
               UpperCAmelCase     =	torch.manual_seed(lowercase       )
           else:
               UpperCAmelCase     =	torch.Generator(device=lowercase       ).manual_seed(lowercase       )
           UpperCAmelCase     =	{
               '''prompt''': '''horse''',
               '''generator''': generator,
               '''num_inference_steps''': 1,
               '''frame_size''': 32,
               '''output_type''': '''np''',
           }
           return inputs
       def 					A   (			self		:			str       ):



           '''simple docstring'''

           UpperCAmelCase     =	'''cpu'''

           UpperCAmelCase     =	self.get_dummy_components()

           UpperCAmelCase     =	self.pipeline_class(**lowercase       )
           UpperCAmelCase     =	pipe.to(lowercase       )

           pipe.set_progress_bar_config(disable=lowercase       )

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

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

           UpperCAmelCase     =	np.array(
               [
                   0.0003_9216,
                   0.0003_9216,
                   0.0003_9216,
                   0.0003_9216,
                   0.0003_9216,
                   0.0003_9216,
                   0.0003_9216,
                   0.0003_9216,
                   0.0003_9216,
               ]       )

           assert np.abs(image_slice.flatten() - expected_slice       ).max() < 1E-2
       def 					A   (			self		:			int       ):



           '''simple docstring'''

           self._test_inference_batch_consistent(batch_sizes=[1, 2]       )
       def 					A   (			self		:			Dict       ):



           '''simple docstring'''

           UpperCAmelCase     =	torch_device == '''cpu'''
           UpperCAmelCase     =	True

           self._test_inference_batch_single_identical(
               batch_size=2     ,       test_max_difference=lowercase     ,       relax_max_difference=lowercase     ,       )







       def 					A   (			self		:			Any       ):



           '''simple docstring'''

           UpperCAmelCase     =	self.get_dummy_components()
           UpperCAmelCase     =	self.pipeline_class(**lowercase       )
           UpperCAmelCase     =	pipe.to(lowercase       )
           pipe.set_progress_bar_config(disable=lowercase       )

           UpperCAmelCase     =	1
           UpperCAmelCase     =	2

           UpperCAmelCase     =	self.get_dummy_inputs(lowercase       )

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

           UpperCAmelCase     =	pipe(**lowercase     ,       num_images_per_prompt=lowercase       )[0]

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






@slow
@require_torch_gpu
class     _a						(			unittest.TestCase				):
       def 					A   (			self		:			str       ):



           '''simple docstring'''

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







       def 					A   (			self		:			int       ):



           '''simple docstring'''

           UpperCAmelCase     =	load_numpy(
               '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'''
               '''/shap_e/test_shap_e_np_out.npy'''       )
           UpperCAmelCase     =	ShapEPipeline.from_pretrained('''openai/shap-e'''       )
           UpperCAmelCase     =	pipe.to(lowercase       )
           pipe.set_progress_bar_config(disable=lowercase       )

           UpperCAmelCase     =	torch.Generator(device=lowercase       ).manual_seed(0       )

           UpperCAmelCase     =	pipe(
               '''a shark'''     ,       generator=lowercase     ,       guidance_scale=15.0     ,       num_inference_steps=64     ,       frame_size=64     ,       output_type='''np'''     ,       ).images[0]

           assert images.shape == (20, 64, 64, 3)

           assert_mean_pixel_difference(lowercase     ,       lowercase       )




