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 		_lowerCamelCase		( UpperCamelCase_     ,  UpperCamelCase_     ,  UpperCamelCase_     ,  unittest.TestCase     ):
    __a						=   StableUnCLIPPipeline
    __a						=   TEXT_TO_IMAGE_PARAMS
    __a						=   TEXT_TO_IMAGE_BATCH_PARAMS
    __a						=   TEXT_TO_IMAGE_IMAGE_PARAMS
    __a						=   TEXT_TO_IMAGE_IMAGE_PARAMS

    # TODO(will) Expected attn_bias.stride(1) == 0 to be true, but got false
    __a						=   False
    def     UpperCamelCase_					(	self )							->       Dict:
     SCREAMING_SNAKE_CASE__: List[str]=   32
     SCREAMING_SNAKE_CASE__: Tuple=   embedder_hidden_size

     # prior components

     torch.manual_seed(0 )
     SCREAMING_SNAKE_CASE__: Union[str, Any]=   CLIPTokenizer.from_pretrained('''hf-internal-testing/tiny-random-clip''' )

     torch.manual_seed(0 )
     SCREAMING_SNAKE_CASE__: Tuple=   CLIPTextModelWithProjection(
         CLIPTextConfig(
             bos_token_id=0			,			eos_token_id=2			,			hidden_size=lowerCAmelCase			,			projection_dim=lowerCAmelCase			,			intermediate_size=37			,			layer_norm_eps=1e-05			,			num_attention_heads=4			,			num_hidden_layers=5			,			pad_token_id=1			,			vocab_size=1000			,			) )

     torch.manual_seed(0 )
     SCREAMING_SNAKE_CASE__: List[Any]=   PriorTransformer(
         num_attention_heads=2			,			attention_head_dim=12			,			embedding_dim=lowerCAmelCase			,			num_layers=1			,			)

     torch.manual_seed(0 )
     SCREAMING_SNAKE_CASE__: Optional[Any]=   DDPMScheduler(
         variance_type='''fixed_small_log'''			,			prediction_type='''sample'''			,			num_train_timesteps=1000			,			clip_sample=lowerCAmelCase			,			clip_sample_range=5.0			,			beta_schedule='''squaredcos_cap_v2'''			,			)

     # regular denoising components

     torch.manual_seed(0 )
     SCREAMING_SNAKE_CASE__: Tuple=   StableUnCLIPImageNormalizer(embedding_dim=lowerCAmelCase )
     SCREAMING_SNAKE_CASE__: Dict=   DDPMScheduler(beta_schedule='''squaredcos_cap_v2''' )

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

     torch.manual_seed(0 )
     SCREAMING_SNAKE_CASE__: Any=   CLIPTextModel(
         CLIPTextConfig(
             bos_token_id=0			,			eos_token_id=2			,			hidden_size=lowerCAmelCase			,			projection_dim=32			,			intermediate_size=37			,			layer_norm_eps=1e-05			,			num_attention_heads=4			,			num_hidden_layers=5			,			pad_token_id=1			,			vocab_size=1000			,			) )

     torch.manual_seed(0 )
     SCREAMING_SNAKE_CASE__: Union[str, Any]=   UNetaDConditionModel(
         sample_size=32			,			in_channels=4			,			out_channels=4			,			down_block_types=('''CrossAttnDownBlock2D''', '''DownBlock2D''')			,			up_block_types=('''UpBlock2D''', '''CrossAttnUpBlock2D''')			,			block_out_channels=(32, 64)			,			attention_head_dim=(2, 4)			,			class_embed_type='''projection'''			,			projection_class_embeddings_input_dim=embedder_projection_dim * 2			,			cross_attention_dim=lowerCAmelCase			,			layers_per_block=1			,			upcast_attention=lowerCAmelCase			,			use_linear_projection=lowerCAmelCase			,			)

     torch.manual_seed(0 )
     SCREAMING_SNAKE_CASE__: Any=   DDIMScheduler(
         beta_schedule='''scaled_linear'''			,			beta_start=0.00085			,			beta_end=0.012			,			prediction_type='''v_prediction'''			,			set_alpha_to_one=lowerCAmelCase			,			steps_offset=1			,			)

     torch.manual_seed(0 )
     SCREAMING_SNAKE_CASE__: Optional[int]=   AutoencoderKL()

     SCREAMING_SNAKE_CASE__: str=   {
         # 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			,			lowerCAmelCase			,			lowerCAmelCase=0 )							->       Dict:
     if str(lowerCAmelCase ).startswith('''mps''' ):
      SCREAMING_SNAKE_CASE__: int=   torch.manual_seed(lowerCAmelCase )
     else:
      SCREAMING_SNAKE_CASE__: List[str]=   torch.Generator(device=lowerCAmelCase ).manual_seed(lowerCAmelCase )
     SCREAMING_SNAKE_CASE__: Dict=   {
         '''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 )							->       Optional[int]:
     SCREAMING_SNAKE_CASE__: List[str]=   torch_device == '''cpu'''

     self._test_attention_slicing_forward_pass(test_max_difference=lowerCAmelCase )



    def     UpperCamelCase_					(	self )							->       Optional[Any]:
     SCREAMING_SNAKE_CASE__: str=   torch_device in ['''cpu''', '''mps''']

     self._test_inference_batch_single_identical(test_max_difference=lowerCAmelCase )






@slow
@require_torch_gpu
class 		_lowerCamelCase		( unittest.TestCase     ):
    def     UpperCamelCase_					(	self )							->       Union[str, Any]:
     # clean up the VRAM after each test
     super().tearDown()
     gc.collect()
     torch.cuda.empty_cache()
    def     UpperCamelCase_					(	self )							->       int:
     SCREAMING_SNAKE_CASE__: List[Any]=   load_numpy(
         '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/stable_unclip/stable_unclip_2_1_l_anime_turtle_fp16.npy''' )

     SCREAMING_SNAKE_CASE__: Union[str, Any]=   StableUnCLIPPipeline.from_pretrained('''fusing/stable-unclip-2-1-l'''			,			torch_dtype=torch.floataa )
     pipe.to(lowerCAmelCase )
     pipe.set_progress_bar_config(disable=lowerCAmelCase )
     # 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()

     SCREAMING_SNAKE_CASE__: str=   torch.Generator(device='''cpu''' ).manual_seed(0 )
     SCREAMING_SNAKE_CASE__: Any=   pipe('''anime turle'''			,			generator=lowerCAmelCase			,			output_type='''np''' )

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

     assert image.shape == (768, 768, 3)

     assert_mean_pixel_difference(lowerCAmelCase			,			lowerCAmelCase )



    def     UpperCamelCase_					(	self )							->       List[Any]:
     torch.cuda.empty_cache()
     torch.cuda.reset_max_memory_allocated()
     torch.cuda.reset_peak_memory_stats()

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

     SCREAMING_SNAKE_CASE__: Optional[Any]=   pipe(
         '''anime turtle'''			,			prior_num_inference_steps=2			,			num_inference_steps=2			,			output_type='''np'''			,			)

     SCREAMING_SNAKE_CASE__: str=   torch.cuda.max_memory_allocated()
     # make sure that less than 7 GB is allocated
     assert mem_bytes < 7 * 10**9





'''simple docstring'''


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 (
    center_crop,
    convert_to_rgb,
    get_resize_output_image_size,
    normalize,
    rescale,
    resize,
    to_channel_dimension_format,
)
from ...image_utils import (
    OPENAI_CLIP_MEAN,
    OPENAI_CLIP_STD,
    ChannelDimension,
    ImageInput,
    PILImageResampling,
    make_list_of_images,
    to_numpy_array,
    valid_images,
)
from ...utils import TensorType, is_vision_available, logging


UpperCAmelCase_						:  int				  =						logging.get_logger(__name__)


if is_vision_available():
				import PIL





class        lowercase__    (		_snake_case   ):





							'''simple docstring'''


							A_				: int       =    ["""pixel_values"""]






							def __init__(       self       ,   __snake_case = True       ,   __snake_case = None       ,   __snake_case = PILImageResampling.BICUBIC       ,   __snake_case = True       ,   __snake_case = None       ,   __snake_case = True       ,   __snake_case = 1 / 255       ,   __snake_case = True       ,   __snake_case = None       ,   __snake_case = None       ,   __snake_case = True       ,   **__snake_case       ,   ):
											super().__init__(**__snake_case    )
											_SCREAMING_SNAKE_CASE					:			int							= size if size is not None else {"""shortest_edge""": 224}
											_SCREAMING_SNAKE_CASE					:			Optional[Any]							= get_size_dict(__snake_case       ,   default_to_square=__snake_case    )
											_SCREAMING_SNAKE_CASE					:			Optional[Any]							= crop_size if crop_size is not None else {"""height""": 224, """width""": 224}
											_SCREAMING_SNAKE_CASE					:			Any							= get_size_dict(__snake_case       ,   default_to_square=__snake_case       ,   param_name="""crop_size"""    )

											_SCREAMING_SNAKE_CASE					:			str							= do_resize
											_SCREAMING_SNAKE_CASE					:			Union[str, Any]							= size
											_SCREAMING_SNAKE_CASE					:			List[str]							= resample
											_SCREAMING_SNAKE_CASE					:			str							= do_center_crop
											_SCREAMING_SNAKE_CASE					:			Optional[Any]							= crop_size
											_SCREAMING_SNAKE_CASE					:			Tuple							= do_rescale
											_SCREAMING_SNAKE_CASE					:			Dict							= rescale_factor
											_SCREAMING_SNAKE_CASE					:			Optional[Any]							= do_normalize
											_SCREAMING_SNAKE_CASE					:			List[Any]							= image_mean if image_mean is not None else OPENAI_CLIP_MEAN
											_SCREAMING_SNAKE_CASE					:			Optional[Any]							= image_std if image_std is not None else OPENAI_CLIP_STD
											_SCREAMING_SNAKE_CASE					:			Optional[int]							= do_convert_rgb






							def UpperCAmelCase_      (       self       ,   __snake_case       ,   __snake_case       ,   __snake_case = PILImageResampling.BICUBIC       ,   __snake_case = None       ,   **__snake_case       ,   ):
											_SCREAMING_SNAKE_CASE					:			int							= get_size_dict(__snake_case       ,   default_to_square=__snake_case    )
											if "shortest_edge" not in size:
															raise ValueError(f"""The `size` parameter must contain the key `shortest_edge`. Got {size.keys()}"""    )
											_SCREAMING_SNAKE_CASE					:			Any							= get_resize_output_image_size(__snake_case       ,   size=size["""shortest_edge"""]       ,   default_to_square=__snake_case    )
											return resize(__snake_case       ,   size=__snake_case       ,   resample=__snake_case       ,   data_format=__snake_case       ,   **__snake_case    )






							def UpperCAmelCase_      (       self       ,   __snake_case       ,   __snake_case       ,   __snake_case = None       ,   **__snake_case       ,   ):
											_SCREAMING_SNAKE_CASE					:			List[str]							= get_size_dict(__snake_case    )
											if "height" not in size or "width" not in size:
															raise ValueError(f"""The `size` parameter must contain the keys (height, width). Got {size.keys()}"""    )
											return center_crop(__snake_case       ,   size=(size["""height"""], size["""width"""])       ,   data_format=__snake_case       ,   **__snake_case    )






							def UpperCAmelCase_      (       self       ,   __snake_case       ,   __snake_case       ,   __snake_case = None       ,   **__snake_case       ,   ):
											return rescale(__snake_case       ,   scale=__snake_case       ,   data_format=__snake_case       ,   **__snake_case    )






							def UpperCAmelCase_      (       self       ,   __snake_case       ,   __snake_case       ,   __snake_case       ,   __snake_case = None       ,   **__snake_case       ,   ):
											return normalize(__snake_case       ,   mean=__snake_case       ,   std=__snake_case       ,   data_format=__snake_case       ,   **__snake_case    )






							def UpperCAmelCase_      (       self       ,   __snake_case       ,   __snake_case = None       ,   __snake_case = None       ,   __snake_case = None       ,   __snake_case = None       ,   __snake_case = None       ,   __snake_case = None       ,   __snake_case = None       ,   __snake_case = None       ,   __snake_case = None       ,   __snake_case = None       ,   __snake_case = None       ,   __snake_case = None       ,   __snake_case = ChannelDimension.FIRST       ,   **__snake_case       ,   ):
											_SCREAMING_SNAKE_CASE					:			Optional[int]							= do_resize if do_resize is not None else self.do_resize
											_SCREAMING_SNAKE_CASE					:			Any							= size if size is not None else self.size
											_SCREAMING_SNAKE_CASE					:			List[str]							= get_size_dict(__snake_case       ,   param_name="""size"""       ,   default_to_square=__snake_case    )
											_SCREAMING_SNAKE_CASE					:			Any							= resample if resample is not None else self.resample
											_SCREAMING_SNAKE_CASE					:			Optional[Any]							= do_center_crop if do_center_crop is not None else self.do_center_crop
											_SCREAMING_SNAKE_CASE					:			str							= crop_size if crop_size is not None else self.crop_size
											_SCREAMING_SNAKE_CASE					:			Optional[int]							= get_size_dict(__snake_case       ,   param_name="""crop_size"""       ,   default_to_square=__snake_case    )
											_SCREAMING_SNAKE_CASE					:			Optional[int]							= do_rescale if do_rescale is not None else self.do_rescale
											_SCREAMING_SNAKE_CASE					:			int							= rescale_factor if rescale_factor is not None else self.rescale_factor
											_SCREAMING_SNAKE_CASE					:			Tuple							= do_normalize if do_normalize is not None else self.do_normalize
											_SCREAMING_SNAKE_CASE					:			List[str]							= image_mean if image_mean is not None else self.image_mean
											_SCREAMING_SNAKE_CASE					:			List[Any]							= image_std if image_std is not None else self.image_std
											_SCREAMING_SNAKE_CASE					:			Any							= do_convert_rgb if do_convert_rgb is not None else self.do_convert_rgb

											_SCREAMING_SNAKE_CASE					:			Optional[Any]							= make_list_of_images(__snake_case    )

											if not valid_images(__snake_case    ):
															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:
															raise ValueError("""Size must be specified if do_resize is True."""    )

											if do_center_crop and crop_size is None:
															raise ValueError("""Crop size must be specified if do_center_crop is True."""    )

											if do_rescale and rescale_factor is None:
															raise ValueError("""Rescale factor must be specified if do_rescale is True."""    )

											if do_normalize and (image_mean is None or image_std is None):
															raise ValueError("""Image mean and std must be specified if do_normalize is True."""    )

											# PIL RGBA images are converted to RGB
											if do_convert_rgb:
															_SCREAMING_SNAKE_CASE					:			Optional[int]							= [convert_to_rgb(__snake_case    ) for image in images]

											# All transformations expect numpy arrays.
											_SCREAMING_SNAKE_CASE					:			str							= [to_numpy_array(__snake_case    ) for image in images]

											if do_resize:
															_SCREAMING_SNAKE_CASE					:			Optional[Any]							= [self.resize(image=__snake_case       ,   size=__snake_case       ,   resample=__snake_case    ) for image in images]

											if do_center_crop:
															_SCREAMING_SNAKE_CASE					:			Tuple							= [self.center_crop(image=__snake_case       ,   size=__snake_case    ) for image in images]

											if do_rescale:
															_SCREAMING_SNAKE_CASE					:			List[Any]							= [self.rescale(image=__snake_case       ,   scale=__snake_case    ) for image in images]

											if do_normalize:
															_SCREAMING_SNAKE_CASE					:			Dict							= [self.normalize(image=__snake_case       ,   mean=__snake_case       ,   std=__snake_case    ) for image in images]

											_SCREAMING_SNAKE_CASE					:			str							= [to_channel_dimension_format(__snake_case       ,   __snake_case    ) for image in images]

											_SCREAMING_SNAKE_CASE					:			int							= {"""pixel_values""": images}
											return BatchFeature(data=__snake_case       ,   tensor_type=__snake_case    )




