


'''simple docstring'''






import gc
import random
import unittest

import numpy as np
import torch
from PIL import Image

from diffusers import (
    DDIMScheduler,
    KandinskyVaaImgaImgPipeline,
    KandinskyVaaPriorPipeline,
    UNetaDConditionModel,
    VQModel,
)
from diffusers.utils import floats_tensor, load_image, load_numpy, slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu

from ..test_pipelines_common import PipelineTesterMixin, assert_mean_pixel_difference


enable_full_determinism()




class 						a__							(				a__							,							unittest.TestCase    ):





							'''simple docstring'''






							lowercase__     :  List[str]      		= KandinskyVaaImgaImgPipeline
							lowercase__     :  Any      		= ["image_embeds", "negative_image_embeds", "image"]
							lowercase__     :  Union[str, Any]      		= [
							    "image_embeds",
							    "negative_image_embeds",
							    "image",
							]
							lowercase__     :  List[str]      		= [
							    "generator",
							    "height",
							    "width",
							    "strength",
							    "guidance_scale",
							    "num_inference_steps",
							    "return_dict",
							    "guidance_scale",
							    "num_images_per_prompt",
							    "output_type",
							    "return_dict",
							]
							lowercase__     :  Union[str, Any]      		= False







							@property
							def     __SCREAMING_SNAKE_CASE				(						self  )			->				int:
														return 32







							@property
							def     __SCREAMING_SNAKE_CASE				(						self  )			->				Tuple:
														return 32







							@property
							def     __SCREAMING_SNAKE_CASE				(						self  )			->				List[Any]:
														return self.time_input_dim







							@property
							def     __SCREAMING_SNAKE_CASE				(						self  )			->				int:
														return self.time_input_dim * 4







							@property
							def     __SCREAMING_SNAKE_CASE				(						self  )			->				List[Any]:
														return 1_00







							@property
							def     __SCREAMING_SNAKE_CASE				(						self  )			->				Dict:
														torch.manual_seed(0  )

														lowerCAmelCase__			  =      {
														    '''in_channels''': 4,
														    # Out channels is double in channels because predicts mean and variance
														    '''out_channels''': 8,
														    '''addition_embed_type''': '''image''',
														    '''down_block_types''': ('''ResnetDownsampleBlock2D''', '''SimpleCrossAttnDownBlock2D'''),
														    '''up_block_types''': ('''SimpleCrossAttnUpBlock2D''', '''ResnetUpsampleBlock2D'''),
														    '''mid_block_type''': '''UNetMidBlock2DSimpleCrossAttn''',
														    '''block_out_channels''': (self.block_out_channels_a, self.block_out_channels_a * 2),
														    '''layers_per_block''': 1,
														    '''encoder_hid_dim''': self.text_embedder_hidden_size,
														    '''encoder_hid_dim_type''': '''image_proj''',
														    '''cross_attention_dim''': self.cross_attention_dim,
														    '''attention_head_dim''': 4,
														    '''resnet_time_scale_shift''': '''scale_shift''',
														    '''class_embed_type''': None,
														}

														lowerCAmelCase__			  =      UNetaDConditionModel(**lowerCamelCase_  )
														return model







							@property
							def     __SCREAMING_SNAKE_CASE				(						self  )			->				int:
														return {
														    "block_out_channels": [32, 64],
														    "down_block_types": ["DownEncoderBlock2D", "AttnDownEncoderBlock2D"],
														    "in_channels": 3,
														    "latent_channels": 4,
														    "layers_per_block": 1,
														    "norm_num_groups": 8,
														    "norm_type": "spatial",
														    "num_vq_embeddings": 12,
														    "out_channels": 3,
														    "up_block_types": [
														        "AttnUpDecoderBlock2D",
														        "UpDecoderBlock2D",
														    ],
														    "vq_embed_dim": 4,
														}







							@property
							def     __SCREAMING_SNAKE_CASE				(						self  )			->				Tuple:
														torch.manual_seed(0  )
														lowerCAmelCase__			  =      VQModel(**self.dummy_movq_kwargs  )
														return model







							def     __SCREAMING_SNAKE_CASE				(						self  )			->				Any:
														lowerCAmelCase__			  =      self.dummy_unet
														lowerCAmelCase__			  =      self.dummy_movq

														lowerCAmelCase__			  =      {
														    '''num_train_timesteps''': 10_00,
														    '''beta_schedule''': '''linear''',
														    '''beta_start''': 0.00_085,
														    '''beta_end''': 0.012,
														    '''clip_sample''': False,
														    '''set_alpha_to_one''': False,
														    '''steps_offset''': 0,
														    '''prediction_type''': '''epsilon''',
														    '''thresholding''': False,
														}

														lowerCAmelCase__			  =      DDIMScheduler(**lowerCamelCase_  )

														lowerCAmelCase__			  =      {
														    '''unet''': unet,
														    '''scheduler''': scheduler,
														    '''movq''': movq,
														}

														return components







							def     __SCREAMING_SNAKE_CASE				(						self   ,       lowerCamelCase_   ,       lowerCamelCase_=0  )			->				Dict:
														lowerCAmelCase__			  =      floats_tensor((1, self.text_embedder_hidden_size)   ,       rng=random.Random(lowerCamelCase_  )  ).to(lowerCamelCase_  )
														lowerCAmelCase__			  =      floats_tensor((1, self.text_embedder_hidden_size)   ,       rng=random.Random(seed + 1  )  ).to(
														    lowerCamelCase_  )
														# create init_image
														lowerCAmelCase__			  =      floats_tensor((1, 3, 64, 64)   ,       rng=random.Random(lowerCamelCase_  )  ).to(lowerCamelCase_  )
														lowerCAmelCase__			  =      image.cpu().permute(0   ,       2   ,       3   ,       1  )[0]
														lowerCAmelCase__			  =      Image.fromarray(np.uinta(lowerCamelCase_  )  ).convert('''RGB'''  ).resize((2_56, 2_56)  )

														if str(lowerCamelCase_  ).startswith('''mps'''  ):
																					lowerCAmelCase__			  =      torch.manual_seed(lowerCamelCase_  )
														else:
																					lowerCAmelCase__			  =      torch.Generator(device=lowerCamelCase_  ).manual_seed(lowerCamelCase_  )
														lowerCAmelCase__			  =      {
														    '''image''': init_image,
														    '''image_embeds''': image_embeds,
														    '''negative_image_embeds''': negative_image_embeds,
														    '''generator''': generator,
														    '''height''': 64,
														    '''width''': 64,
														    '''num_inference_steps''': 10,
														    '''guidance_scale''': 7.0,
														    '''strength''': 0.2,
														    '''output_type''': '''np''',
														}
														return inputs







							def     __SCREAMING_SNAKE_CASE				(						self  )			->				Any:
														lowerCAmelCase__			  =      '''cpu'''

														lowerCAmelCase__			  =      self.get_dummy_components()

														lowerCAmelCase__			  =      self.pipeline_class(**lowerCamelCase_  )
														lowerCAmelCase__			  =      pipe.to(lowerCamelCase_  )

														pipe.set_progress_bar_config(disable=lowerCamelCase_  )

														lowerCAmelCase__			  =      pipe(**self.get_dummy_inputs(lowerCamelCase_  )  )
														lowerCAmelCase__			  =      output.images

														lowerCAmelCase__			  =      pipe(
														    **self.get_dummy_inputs(lowerCamelCase_  )   ,       return_dict=lowerCamelCase_   ,       )[0]

														lowerCAmelCase__			  =      image[0, -3:, -3:, -1]
														lowerCAmelCase__			  =      image_from_tuple[0, -3:, -3:, -1]

														assert image.shape == (1, 64, 64, 3)

														lowerCAmelCase__			  =      np.array(
														    [0.6_199_778, 0.63_984_406, 0.46_145_785, 0.62_944_984, 0.5_622_215, 0.47_306_132, 0.47_441_456, 0.4_607_606, 0.48_719_263]  )
														assert (
														    np.abs(image_slice.flatten() - expected_slice  ).max() < 1e-2
														), F""" expected_slice {expected_slice}, but got {image_slice.flatten()}"""
														assert (
														    np.abs(image_from_tuple_slice.flatten() - expected_slice  ).max() < 1e-2
														), F""" expected_slice {expected_slice}, but got {image_from_tuple_slice.flatten()}"""




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





							'''simple docstring'''






							def     __SCREAMING_SNAKE_CASE				(						self  )			->				str:
														# clean up the VRAM after each test
														super().tearDown()
														gc.collect()
														torch.cuda.empty_cache()







							def     __SCREAMING_SNAKE_CASE				(						self  )			->				Optional[int]:
														lowerCAmelCase__			  =      load_numpy(
														    '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'''
														    '''/kandinskyv22/kandinskyv22_img2img_frog.npy'''  )

														lowerCAmelCase__			  =      load_image(
														    '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main''' '''/kandinsky/cat.png'''  )
														lowerCAmelCase__			  =      '''A red cartoon frog, 4k'''

														lowerCAmelCase__			  =      KandinskyVaaPriorPipeline.from_pretrained(
														    '''kandinsky-community/kandinsky-2-2-prior'''   ,       torch_dtype=torch.floataa  )
														pipe_prior.to(lowerCamelCase_  )

														lowerCAmelCase__			  =      KandinskyVaaImgaImgPipeline.from_pretrained(
														    '''kandinsky-community/kandinsky-2-2-decoder'''   ,       torch_dtype=torch.floataa  )
														lowerCAmelCase__			  =      pipeline.to(lowerCamelCase_  )

														pipeline.set_progress_bar_config(disable=lowerCamelCase_  )

														lowerCAmelCase__			  =      torch.Generator(device='''cpu'''  ).manual_seed(0  )
														lowerCAmelCase__  ,   lowerCAmelCase__			  =      pipe_prior(
														    lowerCamelCase_   ,       generator=lowerCamelCase_   ,       num_inference_steps=5   ,       negative_prompt=''''''   ,       ).to_tuple()

														lowerCAmelCase__			  =      pipeline(
														    image=lowerCamelCase_   ,       image_embeds=lowerCamelCase_   ,       negative_image_embeds=lowerCamelCase_   ,       generator=lowerCamelCase_   ,       num_inference_steps=1_00   ,       height=7_68   ,       width=7_68   ,       strength=0.2   ,       output_type='''np'''   ,       )

														lowerCAmelCase__			  =      output.images[0]

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

														assert_mean_pixel_difference(lowerCamelCase_   ,       lowerCamelCase_  )



'''simple docstring'''






import gc
import unittest

import numpy as np
import torch
from transformers import CLIPTextConfig, CLIPTextModel, XLMRobertaTokenizer

from diffusers import AltDiffusionPipeline, AutoencoderKL, DDIMScheduler, PNDMScheduler, UNetaDConditionModel
from diffusers.pipelines.alt_diffusion.modeling_roberta_series import (
    RobertaSeriesConfig,
    RobertaSeriesModelWithTransformation,
)
from diffusers.utils import slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu

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


enable_full_determinism()




class 						a__							(				a__							,							a__							,							a__							,							unittest.TestCase    ):





							'''simple docstring'''






							lowercase__     :  Optional[Any]      		= AltDiffusionPipeline
							lowercase__     :  Dict      		= TEXT_TO_IMAGE_PARAMS
							lowercase__     :  str      		= TEXT_TO_IMAGE_BATCH_PARAMS
							lowercase__     :  Any      		= TEXT_TO_IMAGE_IMAGE_PARAMS
							lowercase__     :  List[str]      		= TEXT_TO_IMAGE_IMAGE_PARAMS







							def     __SCREAMING_SNAKE_CASE				(						self  )			->				str:
														torch.manual_seed(0  )
														lowerCAmelCase__			  =      UNetaDConditionModel(
														    block_out_channels=(32, 64)   ,       layers_per_block=2   ,       sample_size=32   ,       in_channels=4   ,       out_channels=4   ,       down_block_types=('''DownBlock2D''', '''CrossAttnDownBlock2D''')   ,       up_block_types=('''CrossAttnUpBlock2D''', '''UpBlock2D''')   ,       cross_attention_dim=32   ,       )
														lowerCAmelCase__			  =      DDIMScheduler(
														    beta_start=0.00_085   ,       beta_end=0.012   ,       beta_schedule='''scaled_linear'''   ,       clip_sample=lowerCamelCase_   ,       set_alpha_to_one=lowerCamelCase_   ,       )
														torch.manual_seed(0  )
														lowerCAmelCase__			  =      AutoencoderKL(
														    block_out_channels=[32, 64]   ,       in_channels=3   ,       out_channels=3   ,       down_block_types=['''DownEncoderBlock2D''', '''DownEncoderBlock2D''']   ,       up_block_types=['''UpDecoderBlock2D''', '''UpDecoderBlock2D''']   ,       latent_channels=4   ,       )

														# TODO: address the non-deterministic text encoder (fails for save-load tests)
														# torch.manual_seed(0)
														# text_encoder_config = RobertaSeriesConfig(
														#     hidden_size=32,
														#     project_dim=32,
														#     intermediate_size=37,
														#     layer_norm_eps=1e-05,
														#     num_attention_heads=4,
														#     num_hidden_layers=5,
														#     vocab_size=5002,
														# )
														# text_encoder = RobertaSeriesModelWithTransformation(text_encoder_config)

														torch.manual_seed(0  )
														lowerCAmelCase__			  =      CLIPTextConfig(
														    bos_token_id=0   ,       eos_token_id=2   ,       hidden_size=32   ,       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=50_02   ,       )
														lowerCAmelCase__			  =      CLIPTextModel(lowerCamelCase_  )

														lowerCAmelCase__			  =      XLMRobertaTokenizer.from_pretrained('''hf-internal-testing/tiny-xlm-roberta'''  )
														lowerCAmelCase__			  =      77

														lowerCAmelCase__			  =      {
														    '''unet''': unet,
														    '''scheduler''': scheduler,
														    '''vae''': vae,
														    '''text_encoder''': text_encoder,
														    '''tokenizer''': tokenizer,
														    '''safety_checker''': None,
														    '''feature_extractor''': None,
														}
														return components







							def     __SCREAMING_SNAKE_CASE				(						self   ,       lowerCamelCase_   ,       lowerCamelCase_=0  )			->				List[str]:
														if str(lowerCamelCase_  ).startswith('''mps'''  ):
																					lowerCAmelCase__			  =      torch.manual_seed(lowerCamelCase_  )
														else:
																					lowerCAmelCase__			  =      torch.Generator(device=lowerCamelCase_  ).manual_seed(lowerCamelCase_  )
														lowerCAmelCase__			  =      {
														    '''prompt''': '''A painting of a squirrel eating a burger''',
														    '''generator''': generator,
														    '''num_inference_steps''': 2,
														    '''guidance_scale''': 6.0,
														    '''output_type''': '''numpy''',
														}
														return inputs







							def     __SCREAMING_SNAKE_CASE				(						self  )			->				Union[str, Any]:
														super().test_attention_slicing_forward_pass(expected_max_diff=3e-3  )







							def     __SCREAMING_SNAKE_CASE				(						self  )			->				Optional[int]:
														super().test_inference_batch_single_identical(expected_max_diff=3e-3  )







							def     __SCREAMING_SNAKE_CASE				(						self  )			->				List[str]:
														lowerCAmelCase__			  =      '''cpu'''  # ensure determinism for the device-dependent torch.Generator

														lowerCAmelCase__			  =      self.get_dummy_components()
														torch.manual_seed(0  )
														lowerCAmelCase__			  =      RobertaSeriesConfig(
														    hidden_size=32   ,       project_dim=32   ,       intermediate_size=37   ,       layer_norm_eps=1e-05   ,       num_attention_heads=4   ,       num_hidden_layers=5   ,       vocab_size=50_02   ,       )
														# TODO: remove after fixing the non-deterministic text encoder
														lowerCAmelCase__			  =      RobertaSeriesModelWithTransformation(lowerCamelCase_  )
														lowerCAmelCase__			  =      text_encoder

														lowerCAmelCase__			  =      AltDiffusionPipeline(**lowerCamelCase_  )
														lowerCAmelCase__			  =      alt_pipe.to(lowerCamelCase_  )
														alt_pipe.set_progress_bar_config(disable=lowerCamelCase_  )

														lowerCAmelCase__			  =      self.get_dummy_inputs(lowerCamelCase_  )
														lowerCAmelCase__			  =      '''A photo of an astronaut'''
														lowerCAmelCase__			  =      alt_pipe(**lowerCamelCase_  )
														lowerCAmelCase__			  =      output.images
														lowerCAmelCase__			  =      image[0, -3:, -3:, -1]

														assert image.shape == (1, 64, 64, 3)
														lowerCAmelCase__			  =      np.array(
														    [0.5_748_162, 0.60_447_145, 0.48_821_217, 0.50_100_636, 0.5_431_185, 0.45_763_683, 0.49_657_696, 0.48_132_733, 0.47_573_093]  )

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







							def     __SCREAMING_SNAKE_CASE				(						self  )			->				List[Any]:
														lowerCAmelCase__			  =      '''cpu'''  # ensure determinism for the device-dependent torch.Generator

														lowerCAmelCase__			  =      self.get_dummy_components()
														lowerCAmelCase__			  =      PNDMScheduler(skip_prk_steps=lowerCamelCase_  )
														torch.manual_seed(0  )
														lowerCAmelCase__			  =      RobertaSeriesConfig(
														    hidden_size=32   ,       project_dim=32   ,       intermediate_size=37   ,       layer_norm_eps=1e-05   ,       num_attention_heads=4   ,       num_hidden_layers=5   ,       vocab_size=50_02   ,       )
														# TODO: remove after fixing the non-deterministic text encoder
														lowerCAmelCase__			  =      RobertaSeriesModelWithTransformation(lowerCamelCase_  )
														lowerCAmelCase__			  =      text_encoder
														lowerCAmelCase__			  =      AltDiffusionPipeline(**lowerCamelCase_  )
														lowerCAmelCase__			  =      alt_pipe.to(lowerCamelCase_  )
														alt_pipe.set_progress_bar_config(disable=lowerCamelCase_  )

														lowerCAmelCase__			  =      self.get_dummy_inputs(lowerCamelCase_  )
														lowerCAmelCase__			  =      alt_pipe(**lowerCamelCase_  )
														lowerCAmelCase__			  =      output.images
														lowerCAmelCase__			  =      image[0, -3:, -3:, -1]

														assert image.shape == (1, 64, 64, 3)
														lowerCAmelCase__			  =      np.array(
														    [0.51_605_093, 0.5_707_241, 0.47_365_507, 0.50_578_886, 0.5_633_877, 0.4_642_503, 0.5_182_081, 0.48_763_484, 0.49_084_237]  )

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




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





							'''simple docstring'''






							def     __SCREAMING_SNAKE_CASE				(						self  )			->				Tuple:
														# clean up the VRAM after each test
														super().tearDown()
														gc.collect()
														torch.cuda.empty_cache()







							def     __SCREAMING_SNAKE_CASE				(						self  )			->				List[Any]:
														# make sure here that pndm scheduler skips prk
														lowerCAmelCase__			  =      AltDiffusionPipeline.from_pretrained('''BAAI/AltDiffusion'''   ,       safety_checker=lowerCamelCase_  )
														lowerCAmelCase__			  =      alt_pipe.to(lowerCamelCase_  )
														alt_pipe.set_progress_bar_config(disable=lowerCamelCase_  )

														lowerCAmelCase__			  =      '''A painting of a squirrel eating a burger'''
														lowerCAmelCase__			  =      torch.manual_seed(0  )
														lowerCAmelCase__			  =      alt_pipe([prompt]   ,       generator=lowerCamelCase_   ,       guidance_scale=6.0   ,       num_inference_steps=20   ,       output_type='''np'''  )

														lowerCAmelCase__			  =      output.images

														lowerCAmelCase__			  =      image[0, -3:, -3:, -1]

														assert image.shape == (1, 5_12, 5_12, 3)
														lowerCAmelCase__			  =      np.array([0.1_010, 0.0_800, 0.0_794, 0.0_885, 0.0_843, 0.0_762, 0.0_769, 0.0_729, 0.0_586]  )

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







							def     __SCREAMING_SNAKE_CASE				(						self  )			->				Any:
														lowerCAmelCase__			  =      DDIMScheduler.from_pretrained('''BAAI/AltDiffusion'''   ,       subfolder='''scheduler'''  )

														lowerCAmelCase__			  =      AltDiffusionPipeline.from_pretrained('''BAAI/AltDiffusion'''   ,       scheduler=lowerCamelCase_   ,       safety_checker=lowerCamelCase_  )
														lowerCAmelCase__			  =      alt_pipe.to(lowerCamelCase_  )
														alt_pipe.set_progress_bar_config(disable=lowerCamelCase_  )

														lowerCAmelCase__			  =      '''A painting of a squirrel eating a burger'''
														lowerCAmelCase__			  =      torch.manual_seed(0  )

														lowerCAmelCase__			  =      alt_pipe([prompt]   ,       generator=lowerCamelCase_   ,       num_inference_steps=2   ,       output_type='''numpy'''  )
														lowerCAmelCase__			  =      output.images

														lowerCAmelCase__			  =      image[0, -3:, -3:, -1]

														assert image.shape == (1, 5_12, 5_12, 3)
														lowerCAmelCase__			  =      np.array([0.4_019, 0.4_052, 0.3_810, 0.4_119, 0.3_916, 0.3_982, 0.4_651, 0.4_195, 0.5_323]  )

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