


'''simple docstring'''




import gc
import unittest

import torch
from transformers import CLIPTextConfig, CLIPTextModel, CLIPTextModelWithProjection, CLIPTokenizer

from diffusers import (
    AutoencoderKL,
    DDIMScheduler,
    DDPMScheduler,
    PriorTransformer,
    StableUnCLIPPipeline,
    UNetaDConditionModel,
)
from diffusers.pipelines.stable_diffusion.stable_unclip_image_normalizer import StableUnCLIPImageNormalizer
from diffusers.utils.testing_utils import enable_full_determinism, load_numpy, require_torch_gpu, slow, torch_device

from ..pipeline_params import TEXT_TO_IMAGE_BATCH_PARAMS, TEXT_TO_IMAGE_IMAGE_PARAMS, TEXT_TO_IMAGE_PARAMS
from ..test_pipelines_common import (
    PipelineKarrasSchedulerTesterMixin,
    PipelineLatentTesterMixin,
    PipelineTesterMixin,
    assert_mean_pixel_difference,
)


enable_full_determinism()

class 		UpperCAmelCase_      (  __UpperCamelCase       ,	__UpperCamelCase       ,	__UpperCamelCase       ,	unittest.TestCase		):
     lowerCamelCase				:	Any		=     StableUnCLIPPipeline
     lowerCamelCase				:	str		=     TEXT_TO_IMAGE_PARAMS
     lowerCamelCase				:	Tuple		=     TEXT_TO_IMAGE_BATCH_PARAMS
     lowerCamelCase				:	str		=     TEXT_TO_IMAGE_IMAGE_PARAMS
     lowerCamelCase				:	Union[str, Any]		=     TEXT_TO_IMAGE_IMAGE_PARAMS

     # TODO(will) Expected attn_bias.stride(1) == 0 to be true, but got false
     lowerCamelCase				:	str		=     False





     def 							__UpperCAmelCase      ( self       :      Optional[int]    )				->  Optional[Any]:
        lowerCAmelCase     =			3_2
        lowerCAmelCase     =			embedder_hidden_size

        # prior components

        torch.manual_seed(0    )
        lowerCAmelCase     =			CLIPTokenizer.from_pretrained('hf-internal-testing/tiny-random-clip'    )

        torch.manual_seed(0    )
        lowerCAmelCase     =			CLIPTextModelWithProjection(
            CLIPTextConfig(
                bos_token_id=0    ,      eos_token_id=2    ,      hidden_size=a_    ,      projection_dim=a_    ,      intermediate_size=3_7    ,      layer_norm_eps=1E-05    ,      num_attention_heads=4    ,      num_hidden_layers=5    ,      pad_token_id=1    ,      vocab_size=1_0_0_0    ,      )    )

        torch.manual_seed(0    )
        lowerCAmelCase     =			PriorTransformer(
            num_attention_heads=2    ,      attention_head_dim=1_2    ,      embedding_dim=a_    ,      num_layers=1    ,      )

        torch.manual_seed(0    )
        lowerCAmelCase     =			DDPMScheduler(
            variance_type='fixed_small_log'    ,      prediction_type='sample'    ,      num_train_timesteps=1_0_0_0    ,      clip_sample=a_    ,      clip_sample_range=5.0    ,      beta_schedule='squaredcos_cap_v2'    ,      )

        # regular denoising components

        torch.manual_seed(0    )
        lowerCAmelCase     =			StableUnCLIPImageNormalizer(embedding_dim=a_    )
        lowerCAmelCase     =			DDPMScheduler(beta_schedule='squaredcos_cap_v2'    )

        torch.manual_seed(0    )
        lowerCAmelCase     =			CLIPTokenizer.from_pretrained('hf-internal-testing/tiny-random-clip'    )

        torch.manual_seed(0    )
        lowerCAmelCase     =			CLIPTextModel(
            CLIPTextConfig(
                bos_token_id=0    ,      eos_token_id=2    ,      hidden_size=a_    ,      projection_dim=3_2    ,      intermediate_size=3_7    ,      layer_norm_eps=1E-05    ,      num_attention_heads=4    ,      num_hidden_layers=5    ,      pad_token_id=1    ,      vocab_size=1_0_0_0    ,      )    )

        torch.manual_seed(0    )
        lowerCAmelCase     =			UNetaDConditionModel(
            sample_size=3_2    ,      in_channels=4    ,      out_channels=4    ,      down_block_types=('CrossAttnDownBlock2D', 'DownBlock2D')    ,      up_block_types=('UpBlock2D', 'CrossAttnUpBlock2D')    ,      block_out_channels=(3_2, 6_4)    ,      attention_head_dim=(2, 4)    ,      class_embed_type='projection'    ,      projection_class_embeddings_input_dim=embedder_projection_dim * 2    ,      cross_attention_dim=a_    ,      layers_per_block=1    ,      upcast_attention=a_    ,      use_linear_projection=a_    ,      )

        torch.manual_seed(0    )
        lowerCAmelCase     =			DDIMScheduler(
            beta_schedule='scaled_linear'    ,      beta_start=0.00_085    ,      beta_end=0.012    ,      prediction_type='v_prediction'    ,      set_alpha_to_one=a_    ,      steps_offset=1    ,      )

        torch.manual_seed(0    )
        lowerCAmelCase     =			AutoencoderKL()

        lowerCAmelCase     =			{
            # prior components
            '''prior_tokenizer''': prior_tokenizer,
            '''prior_text_encoder''': prior_text_encoder,
            '''prior''': prior,
            '''prior_scheduler''': prior_scheduler,
            # image noising components
            '''image_normalizer''': image_normalizer,
            '''image_noising_scheduler''': image_noising_scheduler,
            # regular denoising components
            '''tokenizer''': tokenizer,
            '''text_encoder''': text_encoder,
            '''unet''': unet,
            '''scheduler''': scheduler,
            '''vae''': vae,
        }

        return components





     def 							__UpperCAmelCase      ( self       :      Tuple    ,      UpperCAmelCase__       :      Tuple    ,      UpperCAmelCase__       :      Dict=0    )				->  Optional[int]:





        if str(a_    ).startswith('mps'    ):
           lowerCAmelCase     =			torch.manual_seed(a_    )
        else:
           lowerCAmelCase     =			torch.Generator(device=a_    ).manual_seed(a_    )
        lowerCAmelCase     =			{
            '''prompt''': '''A painting of a squirrel eating a burger''',
            '''generator''': generator,
            '''num_inference_steps''': 2,
            '''prior_num_inference_steps''': 2,
            '''output_type''': '''numpy''',
        }
        return inputs





     def 							__UpperCAmelCase      ( self       :      Any    )				->  int:
        lowerCAmelCase     =			torch_device == '''cpu'''

        self._test_attention_slicing_forward_pass(test_max_difference=a_    )





     def 							__UpperCAmelCase      ( self       :      str    )				->  Tuple:
        lowerCAmelCase     =			torch_device in ['''cpu''', '''mps''']

        self._test_inference_batch_single_identical(test_max_difference=a_    )




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





     def 							__UpperCAmelCase      ( self       :      Optional[int]    )				->  List[str]:
        super().tearDown()
        gc.collect()
        torch.cuda.empty_cache()





     def 							__UpperCAmelCase      ( self       :      Optional[Any]    )				->  Optional[int]:
        lowerCAmelCase     =			load_numpy(
            'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/stable_unclip/stable_unclip_2_1_l_anime_turtle_fp16.npy'    )

        lowerCAmelCase     =			StableUnCLIPPipeline.from_pretrained('fusing/stable-unclip-2-1-l'    ,      torch_dtype=torch.floataa    )
        pipe.to(a_    )
        pipe.set_progress_bar_config(disable=a_    )
        # stable unclip will oom when integration tests are run on a V100,
        # so turn on memory savings
        pipe.enable_attention_slicing()
        pipe.enable_sequential_cpu_offload()

        lowerCAmelCase     =			torch.Generator(device='cpu'    ).manual_seed(0    )
        lowerCAmelCase     =			pipe('anime turle'    ,      generator=a_    ,      output_type='np'    )

        lowerCAmelCase     =			output.images[0]

        assert image.shape == (7_6_8, 7_6_8, 3)

        assert_mean_pixel_difference(a_    ,      a_    )





     def 							__UpperCAmelCase      ( self       :      Dict    )				->  Optional[int]:
        torch.cuda.empty_cache()
        torch.cuda.reset_max_memory_allocated()
        torch.cuda.reset_peak_memory_stats()

        lowerCAmelCase     =			StableUnCLIPPipeline.from_pretrained('fusing/stable-unclip-2-1-l'    ,      torch_dtype=torch.floataa    )
        lowerCAmelCase     =			pipe.to(a_    )
        pipe.set_progress_bar_config(disable=a_    )
        pipe.enable_attention_slicing()
        pipe.enable_sequential_cpu_offload()

        lowerCAmelCase     =			pipe(
            'anime turtle'    ,      prior_num_inference_steps=2    ,      num_inference_steps=2    ,      output_type='np'    ,      )

        lowerCAmelCase     =			torch.cuda.max_memory_allocated()
        # make sure that less than 7 GB is allocated
        assert mem_bytes < 7 * 1_0**9










from typing import Dict, List, Optional, Union

import numpy as np

from ...image_processing_utils import BaseImageProcessor, BatchFeature, get_size_dict
from ...image_transforms import rescale, resize, to_channel_dimension_format
from ...image_utils import (
    ChannelDimension,
    ImageInput,
    PILImageResampling,
    make_list_of_images,
    to_numpy_array,
    valid_images,
)
from ...utils import TensorType, is_vision_available, logging


if is_vision_available():
				import PIL


__A   			=logging.get_logger(__name__)

def        a						(      _UpperCAmelCase							:					Optional[Any] ,   _UpperCAmelCase							:					Dict						):







		'''simple docstring'''





		__UpperCAmelCase    :     List[str]													=	b.T
		__UpperCAmelCase    :     Any													=	np.sum(np.square(_UpperCAmelCase						) ,   axis=1						)
		__UpperCAmelCase    :     int													=	np.sum(np.square(_UpperCAmelCase						) ,   axis=0						)
		__UpperCAmelCase    :     Optional[int]													=	np.matmul(_UpperCAmelCase ,   _UpperCAmelCase						)
		__UpperCAmelCase    :     List[str]													=	aa[:, None] - 2 * ab + ba[None, :]
		return d

def        a						(      _UpperCAmelCase							:					Optional[int] ,   _UpperCAmelCase							:					Optional[Any]						):







		'''simple docstring'''





		__UpperCAmelCase    :     str													=	x.reshape(-1 ,   3						)
		__UpperCAmelCase    :     Optional[int]													=	squared_euclidean_distance(_UpperCAmelCase ,   _UpperCAmelCase						)
		return np.argmin(_UpperCAmelCase ,   axis=1						)

class 		UpperCAmelCase__			(			__UpperCamelCase   ):




			'''simple docstring'''






			UpperCamelCase									=       ["""pixel_values"""]




			def __init__(   self							: str       ,       a_							: Optional[Union[List[List[int]], np.ndarray]] = None       ,       a_							: bool = True       ,       a_							: Dict[str, int] = None       ,       a_							: PILImageResampling = PILImageResampling.BILINEAR       ,       a_							: bool = True       ,       a_							: bool = True       ,       **a_							: List[str]       ,       ):
					'''simple docstring'''





					super().__init__(**a_ )
					__UpperCAmelCase    :     Optional[int]													=	size if size is not None else {'''height''': 2_56, '''width''': 2_56}
					__UpperCAmelCase    :     List[str]													=	get_size_dict(a_ )
					__UpperCAmelCase    :     str													=	np.array(a_ ) if clusters is not None else None
					__UpperCAmelCase    :     Dict													=	do_resize
					__UpperCAmelCase    :     Tuple													=	size
					__UpperCAmelCase    :     Union[str, Any]													=	resample
					__UpperCAmelCase    :     Tuple													=	do_normalize
					__UpperCAmelCase    :     Optional[int]													=	do_color_quantize




			def  snake_case__     (   self							: Optional[Any]       ,       a_							: np.ndarray       ,       a_							: Dict[str, int]       ,       a_							: PILImageResampling = PILImageResampling.BILINEAR       ,       a_							: Optional[Union[str, ChannelDimension]] = None       ,       **a_							: Dict       ,       ):
					'''simple docstring'''





					__UpperCAmelCase    :     Tuple													=	get_size_dict(a_ )
					if "height" not in size or "width" not in size:
							raise ValueError(F'Size dictionary must contain both height and width keys. Got {size.keys()}' )
					return resize(
					    a_       ,       size=(size['''height'''], size['''width'''])       ,       resample=a_       ,       data_format=a_       ,       **a_ )




			def  snake_case__     (   self							: Tuple       ,       a_							: np.ndarray       ,       a_							: Optional[Union[str, ChannelDimension]] = None       ,       ):
					'''simple docstring'''





					__UpperCAmelCase    :     Dict													=	rescale(image=a_       ,       scale=1 / 1_2_7.5       ,       data_format=a_ )
					__UpperCAmelCase    :     Union[str, Any]													=	image - 1
					return image




			def  snake_case__     (   self							: int       ,       a_							: ImageInput       ,       a_							: bool = None       ,       a_							: Dict[str, int] = None       ,       a_							: PILImageResampling = None       ,       a_							: bool = None       ,       a_							: Optional[bool] = None       ,       a_							: Optional[Union[List[List[int]], np.ndarray]] = None       ,       a_							: Optional[Union[str, TensorType]] = None       ,       a_							: Optional[Union[str, ChannelDimension]] = ChannelDimension.FIRST       ,       **a_							: Any       ,       ):
					'''simple docstring'''





					__UpperCAmelCase    :     Any													=	do_resize if do_resize is not None else self.do_resize
					__UpperCAmelCase    :     List[str]													=	size if size is not None else self.size
					__UpperCAmelCase    :     Any													=	get_size_dict(a_ )
					__UpperCAmelCase    :     Optional[int]													=	resample if resample is not None else self.resample
					__UpperCAmelCase    :     List[str]													=	do_normalize if do_normalize is not None else self.do_normalize
					__UpperCAmelCase    :     int													=	do_color_quantize if do_color_quantize is not None else self.do_color_quantize
					__UpperCAmelCase    :     Optional[int]													=	clusters if clusters is not None else self.clusters
					__UpperCAmelCase    :     Any													=	np.array(a_ )

					__UpperCAmelCase    :     Optional[int]													=	make_list_of_images(a_ )

					if not valid_images(a_ ):
							raise ValueError(
							    '''Invalid image type. Must be of type PIL.Image.Image, numpy.ndarray, '''
							    '''torch.Tensor, tf.Tensor or jax.ndarray.''' )

					if do_resize and size is None or resample is None:
							raise ValueError('''Size and resample must be specified if do_resize is True.''' )

					if do_color_quantize and clusters is None:
							raise ValueError('''Clusters must be specified if do_color_quantize is True.''' )

					# All transformations expect numpy arrays.
					__UpperCAmelCase    :     List[Any]													=	[to_numpy_array(a_ ) for image in images]

					if do_resize:
							__UpperCAmelCase    :     List[str]													=	[self.resize(image=a_       ,       size=a_       ,       resample=a_ ) for image in images]

					if do_normalize:
							__UpperCAmelCase    :     Dict													=	[self.normalize(image=a_ ) for image in images]

					if do_color_quantize:
							__UpperCAmelCase    :     int													=	[to_channel_dimension_format(a_       ,       ChannelDimension.LAST ) for image in images]
							# color quantize from (batch_size, height, width, 3) to (batch_size, height, width)
							__UpperCAmelCase    :     List[str]													=	np.array(a_ )
							__UpperCAmelCase    :     Dict													=	color_quantize(a_       ,       a_ ).reshape(images.shape[:-1] )

							# flatten to (batch_size, height*width)
							__UpperCAmelCase    :     Any													=	images.shape[0]
							__UpperCAmelCase    :     Any													=	images.reshape(a_       ,       -1 )

							# We need to convert back to a list of images to keep consistent behaviour across processors.
							__UpperCAmelCase    :     List[Any]													=	list(a_ )
					else:
							__UpperCAmelCase    :     int													=	[to_channel_dimension_format(a_       ,       a_ ) for image in images]

					__UpperCAmelCase    :     int													=	{'''input_ids''': images}
					return BatchFeature(data=a_       ,       tensor_type=a_ )





