






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      :    Union[str, Any]   		=   False



class 							SCREAMING_SNAKE_CASE__		(       unittest.TestCase		):





					"""simple docstring"""






					def  _lowercase  (      self  :						Dict ):
						super().tearDown()
						gc.collect()
						torch.cuda.empty_cache()



					@property
					def  _lowercase  (      self  :						Tuple ):
						return 1_2



					@property
					def  _lowercase  (      self  :						Optional[Any] ):
						return 1_2



					@property
					def  _lowercase  (      self  :						List[Any] ):
						return 3_2



					@property
					def  _lowercase  (      self  :						Any ):
						torch.manual_seed(0 )
						snake_case__				:   Tuple     =						VQModel(
						    block_out_channels=[3_2, 6_4]							,	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  _lowercase  (      self  :						List[Any] ):
						snake_case__				:   Optional[int]     =						CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip" )
						return tokenizer



					@property
					def  _lowercase  (      self  :						List[str] ):
						torch.manual_seed(0 )
						snake_case__				:   int     =						CLIPTextConfig(
						    bos_token_id=0							,	eos_token_id=2							,	hidden_size=self.text_embedder_hidden_size							,	intermediate_size=3_7							,	layer_norm_eps=1e-0_5							,	num_attention_heads=4							,	num_hidden_layers=5							,	pad_token_id=1							,	vocab_size=1_0_0_0							,	)
						return CLIPTextModel(UpperCAmelCase__ )



					@property
					def  _lowercase  (      self  :						Dict ):
						torch.manual_seed(0 )

						snake_case__				:   Tuple     =						1_2
						snake_case__				:   Optional[Any]     =						1_2

						snake_case__				:   Tuple     =						{
						    '''attention_bias''': True,
						    '''cross_attention_dim''': 3_2,
						    '''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''': 3_2,
						    '''sample_size''': width,
						    '''activation_fn''': '''geglu-approximate''',
						}

						snake_case__				:   List[Any]     =						TransformeraDModel(**UpperCAmelCase__ )
						return model



					def  _lowercase  (      self  :						Optional[Any] ):
						snake_case__				:   int     =						'''cpu'''

						snake_case__				:   List[str]     =						self.dummy_vqvae
						snake_case__				:   Optional[Any]     =						self.dummy_text_encoder
						snake_case__				:   List[str]     =						self.dummy_tokenizer
						snake_case__				:   List[Any]     =						self.dummy_transformer
						snake_case__				:   str     =						VQDiffusionScheduler(self.num_embed )
						snake_case__				:   Optional[int]     =						LearnedClassifierFreeSamplingEmbeddings(learnable=UpperCAmelCase__ )

						snake_case__				:   str     =						VQDiffusionPipeline(
						    vqvae=UpperCAmelCase__							,	text_encoder=UpperCAmelCase__							,	tokenizer=UpperCAmelCase__							,	transformer=UpperCAmelCase__							,	scheduler=UpperCAmelCase__							,	learned_classifier_free_sampling_embeddings=UpperCAmelCase__							,	)
						snake_case__				:   Dict     =						pipe.to(UpperCAmelCase__ )
						pipe.set_progress_bar_config(disable=UpperCAmelCase__ )

						snake_case__				:   Tuple     =						'''teddy bear playing in the pool'''

						snake_case__				:   Any     =						torch.Generator(device=UpperCAmelCase__ ).manual_seed(0 )
						snake_case__				:   Dict     =						pipe([prompt]							,	generator=UpperCAmelCase__							,	num_inference_steps=2							,	output_type="np" )
						snake_case__				:   List[str]     =						output.images

						snake_case__				:   int     =						torch.Generator(device=UpperCAmelCase__ ).manual_seed(0 )
						snake_case__				:   List[Any]     =						pipe(
						    [prompt]							,	generator=UpperCAmelCase__							,	output_type="np"							,	return_dict=UpperCAmelCase__							,	num_inference_steps=2 )[0]

						snake_case__				:   Optional[Any]     =						image[0, -3:, -3:, -1]
						snake_case__				:   int     =						image_from_tuple[0, -3:, -3:, -1]

						assert image.shape == (1, 2_4, 2_4, 3)

						snake_case__				:   str     =						np.array([0.6_5_5_1, 0.6_1_6_8, 0.5_0_0_8, 0.5_6_7_6, 0.5_6_5_9, 0.4_2_9_5, 0.6_0_7_3, 0.5_5_9_9, 0.4_9_9_2] )

						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  _lowercase  (      self  :						str ):
						snake_case__				:   Dict     =						'''cpu'''

						snake_case__				:   Any     =						self.dummy_vqvae
						snake_case__				:   List[str]     =						self.dummy_text_encoder
						snake_case__				:   int     =						self.dummy_tokenizer
						snake_case__				:   str     =						self.dummy_transformer
						snake_case__				:   Dict     =						VQDiffusionScheduler(self.num_embed )
						snake_case__				:   Dict     =						LearnedClassifierFreeSamplingEmbeddings(
						    learnable=UpperCAmelCase__							,	hidden_size=self.text_embedder_hidden_size							,	length=tokenizer.model_max_length )

						snake_case__				:   Dict     =						VQDiffusionPipeline(
						    vqvae=UpperCAmelCase__							,	text_encoder=UpperCAmelCase__							,	tokenizer=UpperCAmelCase__							,	transformer=UpperCAmelCase__							,	scheduler=UpperCAmelCase__							,	learned_classifier_free_sampling_embeddings=UpperCAmelCase__							,	)
						snake_case__				:   int     =						pipe.to(UpperCAmelCase__ )
						pipe.set_progress_bar_config(disable=UpperCAmelCase__ )

						snake_case__				:   Optional[Any]     =						'''teddy bear playing in the pool'''

						snake_case__				:   Optional[int]     =						torch.Generator(device=UpperCAmelCase__ ).manual_seed(0 )
						snake_case__				:   Optional[int]     =						pipe([prompt]							,	generator=UpperCAmelCase__							,	num_inference_steps=2							,	output_type="np" )
						snake_case__				:   Any     =						output.images

						snake_case__				:   Optional[int]     =						torch.Generator(device=UpperCAmelCase__ ).manual_seed(0 )
						snake_case__				:   Dict     =						pipe(
						    [prompt]							,	generator=UpperCAmelCase__							,	output_type="np"							,	return_dict=UpperCAmelCase__							,	num_inference_steps=2 )[0]

						snake_case__				:   Optional[Any]     =						image[0, -3:, -3:, -1]
						snake_case__				:   List[Any]     =						image_from_tuple[0, -3:, -3:, -1]

						assert image.shape == (1, 2_4, 2_4, 3)

						snake_case__				:   List[Any]     =						np.array([0.6_6_9_3, 0.6_0_7_5, 0.4_9_5_9, 0.5_7_0_1, 0.5_5_8_3, 0.4_3_3_3, 0.6_1_7_1, 0.5_6_8_4, 0.4_9_8_8] )

						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  _lowercase  (      self  :						str ):
						super().tearDown()
						gc.collect()
						torch.cuda.empty_cache()






					def  _lowercase  (      self  :						Any ):
						snake_case__				:   str     =						load_numpy(
						    "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"
						    "/vq_diffusion/teddy_bear_pool_classifier_free_sampling.npy" )

						snake_case__				:   Optional[Any]     =						VQDiffusionPipeline.from_pretrained("microsoft/vq-diffusion-ithq" )
						snake_case__				:   int     =						pipeline.to(UpperCAmelCase__ )
						pipeline.set_progress_bar_config(disable=UpperCAmelCase__ )

						# requires GPU generator for gumbel softmax
						# don't use GPU generator in tests though
						snake_case__				:   int     =						torch.Generator(device=UpperCAmelCase__ ).manual_seed(0 )
						snake_case__				:   Optional[int]     =						pipeline(
						    "teddy bear playing in the pool"							,	num_images_per_prompt=1							,	generator=UpperCAmelCase__							,	output_type="np"							,	)

						snake_case__				:   Optional[Any]     =						output.images[0]

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











# Lint as: python3
# pylint: enable=line-too-long
# pylint: disable=g-import-not-at-top,g-bad-import-order,wrong-import-position

__lowerCamelCase      :    Union[str, Any]   		=   """2.13.1"""

import platform

import pyarrow
from packaging import version


if version.parse(platform.python_version()) < version.parse("""3.7"""):
	raise ImportWarning(
	    """To use `datasets`, Python>=3.7 is required, and the current version of Python doesn't match this condition."""
	)

if version.parse(pyarrow.__version__).major < 8:
	raise ImportWarning(
	    """To use `datasets`, the module `pyarrow>=8.0.0` is required, and the current version of `pyarrow` doesn't match this condition.\n"""
	    """If you are running this in a Google Colab, you should probably just restart the runtime to use the right version of `pyarrow`."""
	)

del platform
del pyarrow
del version

from .arrow_dataset import Dataset
from .arrow_reader import ReadInstruction
from .builder import ArrowBasedBuilder, BeamBasedBuilder, BuilderConfig, DatasetBuilder, GeneratorBasedBuilder
from .combine import concatenate_datasets, interleave_datasets
from .dataset_dict import DatasetDict, IterableDatasetDict
from .download import *
from .features import *
from .fingerprint import disable_caching, enable_caching, is_caching_enabled, set_caching_enabled
from .info import DatasetInfo, MetricInfo
from .inspect import (
    get_dataset_config_info,
    get_dataset_config_names,
    get_dataset_infos,
    get_dataset_split_names,
    inspect_dataset,
    inspect_metric,
    list_datasets,
    list_metrics,
)
from .iterable_dataset import IterableDataset
from .load import load_dataset, load_dataset_builder, load_from_disk, load_metric
from .metric import Metric
from .splits import (
    NamedSplit,
    NamedSplitAll,
    Split,
    SplitBase,
    SplitDict,
    SplitGenerator,
    SplitInfo,
    SubSplitInfo,
    percent,
)
from .tasks import *
from .utils import *
from .utils import logging


# deprecated modules
from datasets import arrow_dataset as _arrow_dataset  # isort:skip
from datasets import utils as _utils  # isort:skip
from datasets.utils import download_manager as _deprecated_download_manager  # isort:skip

__lowerCamelCase      :    List[Any]   		=   concatenate_datasets
__lowerCamelCase      :    List[str]   		=   DownloadConfig
__lowerCamelCase      :    Union[str, Any]   		=   DownloadManager
__lowerCamelCase      :    str   		=   DownloadMode
__lowerCamelCase      :    Union[str, Any]   		=   DownloadConfig
__lowerCamelCase      :    List[str]   		=   DownloadMode
__lowerCamelCase      :    Dict   		=   DownloadManager

del _arrow_dataset, _utils, _deprecated_download_manager



