
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        _SCREAMING_SNAKE_CASE	( snake_case_     ,			snake_case_     ,			snake_case_     ,			unittest.TestCase		):
	lowerCAmelCase__   =					AltDiffusionPipeline
	lowerCAmelCase__   =					TEXT_TO_IMAGE_PARAMS
	lowerCAmelCase__   =					TEXT_TO_IMAGE_BATCH_PARAMS
	lowerCAmelCase__   =					TEXT_TO_IMAGE_IMAGE_PARAMS
	lowerCAmelCase__   =					TEXT_TO_IMAGE_IMAGE_PARAMS



	def        SCREAMING_SNAKE_CASE_(       self					)							->			int:
					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=lowercase		,    set_alpha_to_one=lowercase		,    )
					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=5002		,    )
					lowerCamelCase_         =     CLIPTextModel(lowercase					)

					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		,    lowercase		,    lowercase=0					)							->			Union[str, Any]:
					if str(lowercase					).startswith("mps"					):
									lowerCamelCase_         =     torch.manual_seed(lowercase					)
					else:
									lowerCamelCase_         =     torch.Generator(device=lowercase					).manual_seed(lowercase					)
					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					)							->			List[str]:
					super().test_attention_slicing_forward_pass(expected_max_diff=3e-3					)



	def        SCREAMING_SNAKE_CASE_(       self					)							->			Dict:
					super().test_inference_batch_single_identical(expected_max_diff=3e-3					)



	def        SCREAMING_SNAKE_CASE_(       self					)							->			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=5002		,    )
					# TODO: remove after fixing the non-deterministic text encoder
					lowerCamelCase_         =     RobertaSeriesModelWithTransformation(lowercase					)
					lowerCamelCase_         =     text_encoder

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

					lowerCamelCase_         =     self.get_dummy_inputs(lowercase					)
					lowerCamelCase_         =     "A photo of an astronaut"
					lowerCamelCase_         =     alt_pipe(**lowercase					)
					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        SCREAMING_SNAKE_CASE_(       self					)							->			Optional[Any]:
					lowerCamelCase_         =     "cpu"  # ensure determinism for the device-dependent torch.Generator

					lowerCamelCase_         =     self.get_dummy_components()
					lowerCamelCase_         =     PNDMScheduler(skip_prk_steps=lowercase					)
					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=5002		,    )
					# TODO: remove after fixing the non-deterministic text encoder
					lowerCamelCase_         =     RobertaSeriesModelWithTransformation(lowercase					)
					lowerCamelCase_         =     text_encoder
					lowerCamelCase_         =     AltDiffusionPipeline(**lowercase					)
					lowerCamelCase_         =     alt_pipe.to(lowercase					)
					alt_pipe.set_progress_bar_config(disable=lowercase					)

					lowerCamelCase_         =     self.get_dummy_inputs(lowercase					)
					lowerCamelCase_         =     alt_pipe(**lowercase					)
					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        _SCREAMING_SNAKE_CASE	( unittest.TestCase		):



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



	def        SCREAMING_SNAKE_CASE_(       self					)							->			str:
					# make sure here that pndm scheduler skips prk
					lowerCamelCase_         =     AltDiffusionPipeline.from_pretrained("BAAI/AltDiffusion"		,    safety_checker=lowercase					)
					lowerCamelCase_         =     alt_pipe.to(lowercase					)
					alt_pipe.set_progress_bar_config(disable=lowercase					)

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

					lowerCamelCase_         =     output.images

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

					assert image.shape == (1, 512, 512, 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        SCREAMING_SNAKE_CASE_(       self					)							->			Union[str, Any]:
					lowerCamelCase_         =     DDIMScheduler.from_pretrained("BAAI/AltDiffusion"		,    subfolder="scheduler"					)

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

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

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

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

					assert image.shape == (1, 512, 512, 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




import unittest

from transformers import AlbertTokenizer, AlbertTokenizerFast
from transformers.testing_utils import get_tests_dir, require_sentencepiece, require_tokenizers, slow

from ...test_tokenization_common import TokenizerTesterMixin


__A							=get_tests_dir('''fixtures/spiece.model''')




@require_sentencepiece
@require_tokenizers
class        _SCREAMING_SNAKE_CASE	( snake_case_     ,			unittest.TestCase		):
	lowerCAmelCase__   =					AlbertTokenizer
	lowerCAmelCase__   =					AlbertTokenizerFast
	lowerCAmelCase__   =					True
	lowerCAmelCase__   =					True
	lowerCAmelCase__   =					True



	def        SCREAMING_SNAKE_CASE_(       self					)							->			Any:
					super().setUp()

					# We have a SentencePiece fixture for testing
					lowerCamelCase_         =     AlbertTokenizer(lowercase					)
					tokenizer.save_pretrained(self.tmpdirname					)



	def        SCREAMING_SNAKE_CASE_(       self		,    lowercase					)							->			str:
					lowerCamelCase_         =     "this is a test"
					lowerCamelCase_         =     "this is a test"
					return input_text, output_text



	def        SCREAMING_SNAKE_CASE_(       self					)							->			Dict:
					lowerCamelCase_         =     "<pad>"
					lowerCamelCase_         =     0

					self.assertEqual(self.get_tokenizer()._convert_token_to_id(lowercase					)		,    lowercase					)
					self.assertEqual(self.get_tokenizer()._convert_id_to_token(lowercase					)		,    lowercase					)



	def        SCREAMING_SNAKE_CASE_(       self					)							->			Any:
					lowerCamelCase_         =     list(self.get_tokenizer().get_vocab().keys()					)

					self.assertEqual(vocab_keys[0]		,    "<pad>"					)
					self.assertEqual(vocab_keys[1]		,    "<unk>"					)
					self.assertEqual(vocab_keys[-1]		,    "▁eloquent"					)
					self.assertEqual(len(lowercase					)		,    30000					)



	def        SCREAMING_SNAKE_CASE_(       self					)							->			Union[str, Any]:
					self.assertEqual(self.get_tokenizer().vocab_size		,    30000					)



	def        SCREAMING_SNAKE_CASE_(       self					)							->			Union[str, Any]:
					if not self.test_rust_tokenizer:
									return

					lowerCamelCase_         =     self.get_tokenizer()
					lowerCamelCase_         =     self.get_rust_tokenizer()

					lowerCamelCase_         =     "I was born in 92000, and this is falsé."

					lowerCamelCase_         =     tokenizer.tokenize(lowercase					)
					lowerCamelCase_         =     rust_tokenizer.tokenize(lowercase					)
					self.assertListEqual(lowercase		,    lowercase					)

					lowerCamelCase_         =     tokenizer.encode(lowercase		,    add_special_tokens=lowercase					)
					lowerCamelCase_         =     rust_tokenizer.encode(lowercase		,    add_special_tokens=lowercase					)
					self.assertListEqual(lowercase		,    lowercase					)

					lowerCamelCase_         =     self.get_rust_tokenizer()
					lowerCamelCase_         =     tokenizer.encode(lowercase					)
					lowerCamelCase_         =     rust_tokenizer.encode(lowercase					)
					self.assertListEqual(lowercase		,    lowercase					)



	def        SCREAMING_SNAKE_CASE_(       self					)							->			Optional[Any]:
					lowerCamelCase_         =     AlbertTokenizer(lowercase		,    keep_accents=lowercase					)

					lowerCamelCase_         =     tokenizer.tokenize("This is a test"					)
					self.assertListEqual(lowercase		,    ["▁this", "▁is", "▁a", "▁test"]					)

					self.assertListEqual(tokenizer.convert_tokens_to_ids(lowercase					)		,    [48, 25, 21, 1289]					)

					lowerCamelCase_         =     tokenizer.tokenize("I was born in 92000, and this is falsé."					)
					self.assertListEqual(
					    lowercase		,    ["▁i", "▁was", "▁born", "▁in", "▁9", "2000", ",", "▁and", "▁this", "▁is", "▁fal", "s", "é", "."]					)
					lowerCamelCase_         =     tokenizer.convert_tokens_to_ids(lowercase					)
					self.assertListEqual(lowercase		,    [31, 23, 386, 19, 561, 3050, 15, 17, 48, 25, 8256, 18, 1, 9]					)

					lowerCamelCase_         =     tokenizer.convert_ids_to_tokens(lowercase					)
					self.assertListEqual(
					    lowercase		,    ["▁i", "▁was", "▁born", "▁in", "▁9", "2000", ",", "▁and", "▁this", "▁is", "▁fal", "s", "<unk>", "."]		,    )



	def        SCREAMING_SNAKE_CASE_(       self					)							->			Any:
					lowerCamelCase_         =     AlbertTokenizer(lowercase					)

					lowerCamelCase_         =     tokenizer.encode("sequence builders"					)
					lowerCamelCase_         =     tokenizer.encode("multi-sequence build"					)

					lowerCamelCase_         =     tokenizer.build_inputs_with_special_tokens(lowercase					)
					lowerCamelCase_         =     tokenizer.build_inputs_with_special_tokens(lowercase		,    lowercase					)

					assert encoded_sentence == [tokenizer.cls_token_id] + text + [tokenizer.sep_token_id]
					assert encoded_pair == [tokenizer.cls_token_id] + text + [tokenizer.sep_token_id] + text_a + [
					    tokenizer.sep_token_id
					]



	@slow
	def        SCREAMING_SNAKE_CASE_(       self					)							->			Optional[Any]:
					# fmt: off
					lowerCamelCase_         =     {"attention_mask": [[1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1], [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]], "input_ids": [[2, 21970, 13, 5, 6092, 167, 28, 7103, 2153, 673, 8, 7028, 12051, 18, 17, 7103, 2153, 673, 8, 3515, 18684, 8, 4461, 6, 1927, 297, 8, 12060, 2607, 18, 13, 5, 4461, 15, 10538, 38, 8, 135, 15, 822, 58, 15, 993, 10363, 15, 1460, 8005, 4461, 15, 993, 255, 2328, 9, 9, 9, 6, 26, 1112, 816, 3260, 13, 5, 103, 2377, 6, 17, 1112, 816, 2782, 13, 5, 103, 10641, 6, 29, 84, 2512, 2430, 782, 18684, 2761, 19, 808, 2430, 2556, 17, 855, 1480, 9477, 4091, 128, 11712, 15, 7103, 2153, 673, 17, 24883, 9990, 9, 3], [2, 11502, 25, 1006, 20, 782, 8, 11809, 855, 1732, 19393, 18667, 37, 367, 21018, 69, 1854, 34, 11860, 19124, 27, 156, 225, 17, 193, 4141, 19, 65, 9124, 9, 3, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], [2, 14, 2231, 886, 2385, 17659, 84, 14, 16792, 1952, 9, 3, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]], "token_type_ids": [[0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]]}  # noqa: E501
					# fmt: on

					self.tokenizer_integration_test_util(
					    expected_encoding=lowercase		,    model_name="albert-base-v2"		,    revision="6b6560eaf5ff2e250b00c50f380c5389a9c2d82e"		,    )


