


"""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   _UpperCamelCase  (   lowerCAmelCase__       ,lowerCAmelCase__       ,lowerCAmelCase__       ,unittest.TestCase		):
							'''simple docstring'''







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

							def 			snake_case				(     self	):
										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.0_0_0_8_5    ,   beta_end=0.0_1_2    ,   beta_schedule="scaled_linear"    ,   clip_sample=__a    ,   set_alpha_to_one=__a    ,   )
										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-0_5    ,   num_attention_heads=4    ,   num_hidden_layers=5    ,   pad_token_id=1    ,   vocab_size=50_02    ,   )
										__lowerCAmelCase				    =			CLIPTextModel(__a	)

										__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 			snake_case				(     self    ,   __a    ,   __a=0	):
										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,
										    "guidance_scale": 6.0,
										    "output_type": "numpy",
										}
										return inputs

							def 			snake_case				(     self	):
										super().test_attention_slicing_forward_pass(expected_max_diff=3e-3	)

							def 			snake_case				(     self	):
										super().test_inference_batch_single_identical(expected_max_diff=3e-3	)

							def 			snake_case				(     self	):
										__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-0_5    ,   num_attention_heads=4    ,   num_hidden_layers=5    ,   vocab_size=50_02    ,   )
										# TODO: remove after fixing the non-deterministic text encoder
										__lowerCAmelCase				    =			RobertaSeriesModelWithTransformation(__a	)
										__lowerCAmelCase				    =			text_encoder

										__lowerCAmelCase				    =			AltDiffusionPipeline(**__a	)
										__lowerCAmelCase				    =			alt_pipe.to(__a	)
										alt_pipe.set_progress_bar_config(disable=__a	)

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

										assert image.shape == (1, 64, 64, 3)
										__lowerCAmelCase				    =			np.array(
										    [0.5_7_4_8_1_6_2, 0.6_0_4_4_7_1_4_5, 0.4_8_8_2_1_2_1_7, 0.5_0_1_0_0_6_3_6, 0.5_4_3_1_1_8_5, 0.4_5_7_6_3_6_8_3, 0.4_9_6_5_7_6_9_6, 0.4_8_1_3_2_7_3_3, 0.4_7_5_7_3_0_9_3]	)

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



							def 			snake_case				(     self	):
										__lowerCAmelCase				    =			"cpu"  # ensure determinism for the device-dependent torch.Generator

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

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

										assert image.shape == (1, 64, 64, 3)
										__lowerCAmelCase				    =			np.array(
										    [0.5_1_6_0_5_0_9_3, 0.5_7_0_7_2_4_1, 0.4_7_3_6_5_5_0_7, 0.5_0_5_7_8_8_8_6, 0.5_6_3_3_8_7_7, 0.4_6_4_2_5_0_3, 0.5_1_8_2_0_8_1, 0.4_8_7_6_3_4_8_4, 0.4_9_0_8_4_2_3_7]	)

										assert np.abs(image_slice.flatten() - expected_slice	).max() < 1e-2
@slow
@require_torch_gpu
class   _UpperCamelCase  (   unittest.TestCase		):
							'''simple docstring'''







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

							def 			snake_case				(     self	):
										# make sure here that pndm scheduler skips prk
										__lowerCAmelCase				    =			AltDiffusionPipeline.from_pretrained("BAAI/AltDiffusion"    ,   safety_checker=__a	)
										__lowerCAmelCase				    =			alt_pipe.to(__a	)
										alt_pipe.set_progress_bar_config(disable=__a	)

										__lowerCAmelCase				    =			"A painting of a squirrel eating a burger"
										__lowerCAmelCase				    =			torch.manual_seed(0	)
										__lowerCAmelCase				    =			alt_pipe([prompt]    ,   generator=__a    ,   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_0_1_0, 0.0_8_0_0, 0.0_7_9_4, 0.0_8_8_5, 0.0_8_4_3, 0.0_7_6_2, 0.0_7_6_9, 0.0_7_2_9, 0.0_5_8_6]	)

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



							def 			snake_case				(     self	):
										__lowerCAmelCase				    =			DDIMScheduler.from_pretrained("BAAI/AltDiffusion"    ,   subfolder="scheduler"	)

										__lowerCAmelCase				    =			AltDiffusionPipeline.from_pretrained("BAAI/AltDiffusion"    ,   scheduler=__a    ,   safety_checker=__a	)
										__lowerCAmelCase				    =			alt_pipe.to(__a	)
										alt_pipe.set_progress_bar_config(disable=__a	)

										__lowerCAmelCase				    =			"A painting of a squirrel eating a burger"
										__lowerCAmelCase				    =			torch.manual_seed(0	)

										__lowerCAmelCase				    =			alt_pipe([prompt]    ,   generator=__a    ,   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_0_1_9, 0.4_0_5_2, 0.3_8_1_0, 0.4_1_1_9, 0.3_9_1_6, 0.3_9_8_2, 0.4_6_5_1, 0.4_1_9_5, 0.5_3_2_3]	)

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








"""simple docstring"""


import unittest

from transformers import (
    MODEL_FOR_SEQ_TO_SEQ_CAUSAL_LM_MAPPING,
    TF_MODEL_FOR_SEQ_TO_SEQ_CAUSAL_LM_MAPPING,
    TextaTextGenerationPipeline,
    pipeline,
)
from transformers.testing_utils import is_pipeline_test, require_tf, require_torch
from transformers.utils import is_torch_available

from .test_pipelines_common import ANY


if is_torch_available():
	import torch
@is_pipeline_test
class   _UpperCamelCase  (   unittest.TestCase		):
							'''simple docstring'''







							__UpperCAmelCase						:    Union[str, Any]				      =MODEL_FOR_SEQ_TO_SEQ_CAUSAL_LM_MAPPING
							__UpperCAmelCase						:    Union[str, Any]				      =TF_MODEL_FOR_SEQ_TO_SEQ_CAUSAL_LM_MAPPING

							def 			snake_case				(     self    ,   __a    ,   __a    ,   __a	):
										__lowerCAmelCase				    =			TextaTextGenerationPipeline(model=__a    ,   tokenizer=__a	)
										return generator, ["Something to write", "Something else"]

							def 			snake_case				(     self    ,   __a    ,   __a	):
										__lowerCAmelCase				    =			generator("Something there"	)
										self.assertEqual(__a    ,   [{"generated_text": ANY(__a	)}]	)
										# These are encoder decoder, they don't just append to incoming string
										self.assertFalse(outputs[0]["generated_text"].startswith("Something there"	)	)

										__lowerCAmelCase				    =			generator(["This is great !", "Something else"]    ,   num_return_sequences=2    ,   do_sample=__a	)
										self.assertEqual(
										    __a    ,   [
										        [{"generated_text": ANY(__a	)}, {"generated_text": ANY(__a	)}],
										        [{"generated_text": ANY(__a	)}, {"generated_text": ANY(__a	)}],
										    ]    ,   )

										__lowerCAmelCase				    =			generator(
										    ["This is great !", "Something else"]    ,   num_return_sequences=2    ,   batch_size=2    ,   do_sample=__a	)
										self.assertEqual(
										    __a    ,   [
										        [{"generated_text": ANY(__a	)}, {"generated_text": ANY(__a	)}],
										        [{"generated_text": ANY(__a	)}, {"generated_text": ANY(__a	)}],
										    ]    ,   )

										with self.assertRaises(__a	):
													generator(4	)

							@require_torch
							def 			snake_case				(     self	):
										__lowerCAmelCase				    =			pipeline("text2text-generation"    ,   model="patrickvonplaten/t5-tiny-random"    ,   framework="pt"	)
										# do_sample=False necessary for reproducibility
										__lowerCAmelCase				    =			generator("Something there"    ,   do_sample=__a	)
										self.assertEqual(__a    ,   [{"generated_text": ""}]	)

										__lowerCAmelCase				    =			3
										__lowerCAmelCase				    =			generator(
										    "Something there"    ,   num_return_sequences=__a    ,   num_beams=__a    ,   )
										__lowerCAmelCase				    =			[
										    {"generated_text": "Beide Beide Beide Beide Beide Beide Beide Beide Beide"},
										    {"generated_text": "Beide Beide Beide Beide Beide Beide Beide Beide"},
										    {"generated_text": ""},
										]
										self.assertEqual(__a    ,   __a	)

										__lowerCAmelCase				    =			generator("This is a test"    ,   do_sample=__a    ,   num_return_sequences=2    ,   return_tensors=__a	)
										self.assertEqual(
										    __a    ,   [
										        {"generated_token_ids": ANY(torch.Tensor	)},
										        {"generated_token_ids": ANY(torch.Tensor	)},
										    ]    ,   )
										__lowerCAmelCase				    =			generator.model.config.eos_token_id
										__lowerCAmelCase				    =			"<pad>"
										__lowerCAmelCase				    =			generator(
										    ["This is a test", "This is a second test"]    ,   do_sample=__a    ,   num_return_sequences=2    ,   batch_size=2    ,   return_tensors=__a    ,   )
										self.assertEqual(
										    __a    ,   [
										        [
										            {"generated_token_ids": ANY(torch.Tensor	)},
										            {"generated_token_ids": ANY(torch.Tensor	)},
										        ],
										        [
										            {"generated_token_ids": ANY(torch.Tensor	)},
										            {"generated_token_ids": ANY(torch.Tensor	)},
										        ],
										    ]    ,   )



							@require_tf
							def 			snake_case				(     self	):
										__lowerCAmelCase				    =			pipeline("text2text-generation"    ,   model="patrickvonplaten/t5-tiny-random"    ,   framework="tf"	)
										# do_sample=False necessary for reproducibility
										__lowerCAmelCase				    =			generator("Something there"    ,   do_sample=__a	)
										self.assertEqual(__a    ,   [{"generated_text": ""}]	)




