


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 							lowerCamelCase__						(				UpperCAmelCase       ,						UpperCAmelCase       ,						UpperCAmelCase       ,						unittest.TestCase   ):






			"""simple docstring"""






			_UpperCamelCase		:       Tuple 			=		AltDiffusionPipeline
			_UpperCamelCase		:       Dict 			=		TEXT_TO_IMAGE_PARAMS
			_UpperCamelCase		:       List[str] 			=		TEXT_TO_IMAGE_BATCH_PARAMS
			_UpperCamelCase		:       Union[str, Any] 			=		TEXT_TO_IMAGE_IMAGE_PARAMS
			_UpperCamelCase		:       str 			=		TEXT_TO_IMAGE_IMAGE_PARAMS




			def 					snake_case__					(				self  ):

						'''simple docstring'''


						torch.manual_seed(0  )
						UpperCamelCase__       =							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   ,    )
						UpperCamelCase__       =							DDIMScheduler(
						    beta_start=0.00085   ,    beta_end=0.012   ,    beta_schedule="scaled_linear"   ,    clip_sample=snake_case   ,    set_alpha_to_one=snake_case   ,    )
						torch.manual_seed(0  )
						UpperCamelCase__       =							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  )
						UpperCamelCase__       =							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=5002   ,    )
						UpperCamelCase__       =							CLIPTextModel(snake_case  )

						UpperCamelCase__       =							XLMRobertaTokenizer.from_pretrained("hf-internal-testing/tiny-xlm-roberta"  )
						UpperCamelCase__       =							77

						UpperCamelCase__       =							{
						    "unet": unet,
						    "scheduler": scheduler,
						    "vae": vae,
						    "text_encoder": text_encoder,
						    "tokenizer": tokenizer,
						    "safety_checker": None,
						    "feature_extractor": None,
						}
						return components




			def 					snake_case__					(				self   ,    snake_case   ,    snake_case=0  ):

						'''simple docstring'''


						if str(snake_case  ).startswith("mps"  ):
									UpperCamelCase__       =							torch.manual_seed(snake_case  )
						else:
									UpperCamelCase__       =							torch.Generator(device=snake_case  ).manual_seed(snake_case  )
						UpperCamelCase__       =							{
						    "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  ):

						'''simple docstring'''


						super().test_attention_slicing_forward_pass(expected_max_diff=3E-3  )




			def 					snake_case__					(				self  ):

						'''simple docstring'''


						super().test_inference_batch_single_identical(expected_max_diff=3E-3  )




			def 					snake_case__					(				self  ):

						'''simple docstring'''


						UpperCamelCase__       =							"cpu"  # ensure determinism for the device-dependent torch.Generator

						UpperCamelCase__       =							self.get_dummy_components()
						torch.manual_seed(0  )
						UpperCamelCase__       =							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=5002   ,    )
						# TODO: remove after fixing the non-deterministic text encoder
						UpperCamelCase__       =							RobertaSeriesModelWithTransformation(snake_case  )
						UpperCamelCase__       =							text_encoder

						UpperCamelCase__       =							AltDiffusionPipeline(**snake_case  )
						UpperCamelCase__       =							alt_pipe.to(snake_case  )
						alt_pipe.set_progress_bar_config(disable=snake_case  )

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

						assert image.shape == (1, 64, 64, 3)
						UpperCamelCase__       =							np.array(
						    [0.5748162, 0.60447145, 0.48821217, 0.50100636, 0.5431185, 0.45763683, 0.49657696, 0.48132733, 0.47573093]  )

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







			def 					snake_case__					(				self  ):

						'''simple docstring'''


						UpperCamelCase__       =							"cpu"  # ensure determinism for the device-dependent torch.Generator

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

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

						assert image.shape == (1, 64, 64, 3)
						UpperCamelCase__       =							np.array(
						    [0.51605093, 0.5707241, 0.47365507, 0.50578886, 0.5633877, 0.4642503, 0.5182081, 0.48763484, 0.49084237]  )

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







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






			"""simple docstring"""






			def 					snake_case__					(				self  ):

						'''simple docstring'''


						super().tearDown()
						gc.collect()
						torch.cuda.empty_cache()




			def 					snake_case__					(				self  ):

						'''simple docstring'''


						UpperCamelCase__       =							AltDiffusionPipeline.from_pretrained("BAAI/AltDiffusion"   ,    safety_checker=snake_case  )
						UpperCamelCase__       =							alt_pipe.to(snake_case  )
						alt_pipe.set_progress_bar_config(disable=snake_case  )

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

						UpperCamelCase__       =							output.images

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

						assert image.shape == (1, 512, 512, 3)
						UpperCamelCase__       =							np.array([0.1010, 0.0800, 0.0794, 0.0885, 0.0843, 0.0762, 0.0769, 0.0729, 0.0586]  )

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







			def 					snake_case__					(				self  ):

						'''simple docstring'''


						UpperCamelCase__       =							DDIMScheduler.from_pretrained("BAAI/AltDiffusion"   ,    subfolder="scheduler"  )

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

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

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

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

						assert image.shape == (1, 512, 512, 3)
						UpperCamelCase__       =							np.array([0.4019, 0.4052, 0.3810, 0.4119, 0.3916, 0.3982, 0.4651, 0.4195, 0.5323]  )

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






import json
import logging
import os
import sys
from pathlib import Path

import finetune_rag

from transformers.file_utils import is_apex_available
from transformers.testing_utils import (
    TestCasePlus,
    execute_subprocess_async,
    require_ray,
    require_torch_gpu,
    require_torch_multi_gpu,
)


logging.basicConfig(level=logging.DEBUG)
__UpperCamelCase     		=							logging.getLogger()

__UpperCamelCase     		=							logging.StreamHandler(sys.stdout)
logger.addHandler(stream_handler)







class 							lowerCamelCase__						(				UpperCAmelCase   ):






			"""simple docstring"""






			def 					snake_case__					(				self   ,    snake_case  ):

						'''simple docstring'''


						os.makedirs(snake_case   ,    exist_ok=snake_case  )
						UpperCamelCase__       =							{"source": "What is love ?", "target": "life"}
						UpperCamelCase__       =							{"train": 12, "val": 2, "test": 2}
						for split in ["train", "test", "val"]:
									for field in ["source", "target"]:
												UpperCamelCase__       =							"\n".join([contents[field]] * n_lines[split]  )
												with open(os.path.join(snake_case   ,    F'''{split}.{field}'''  )   ,    "w"  ) as f:
															f.write(snake_case  )




			def 					snake_case__					(				self   ,    snake_case   ,    snake_case = "pytorch"  ):

						'''simple docstring'''


						UpperCamelCase__       =							self.get_auto_remove_tmp_dir()
						UpperCamelCase__       =							os.path.join(snake_case   ,    "output"  )
						UpperCamelCase__       =							os.path.join(snake_case   ,    "data"  )
						self._create_dummy_data(data_dir=snake_case  )

						UpperCamelCase__       =							F'''
                --data_dir {data_dir} \
                --output_dir {output_dir} \
                --model_name_or_path facebook/rag-sequence-base \
                --model_type rag_sequence \
                --do_train \
                --do_predict \
                --n_val -1 \
                --val_check_interval 1.0 \
                --train_batch_size 2 \
                --eval_batch_size 1 \
                --max_source_length 25 \
                --max_target_length 25 \
                --val_max_target_length 25 \
                --test_max_target_length 25 \
                --label_smoothing 0.1 \
                --dropout 0.1 \
                --attention_dropout 0.1 \
                --weight_decay 0.001 \
                --adam_epsilon 1e-08 \
                --max_grad_norm 0.1 \
                --lr_scheduler polynomial \
                --learning_rate 3e-04 \
                --num_train_epochs 1 \
                --warmup_steps 4 \
                --gradient_accumulation_steps 1 \
                --distributed-port 8787 \
                --use_dummy_dataset 1 \
                --distributed_retriever {distributed_retriever} \
            '''.split()

						if gpus > 0:
									testargs.append(F'''--gpus={gpus}'''  )
									if is_apex_available():
												testargs.append("--fp16"  )
						else:
									testargs.append("--gpus=0"  )
									testargs.append("--distributed_backend=ddp_cpu"  )
									testargs.append("--num_processes=2"  )

						UpperCamelCase__       =							[sys.executable, str(Path(finetune_rag.__file__  ).resolve()  )] + testargs
						execute_subprocess_async(snake_case   ,    env=self.get_env()  )

						UpperCamelCase__       =							os.path.join(snake_case   ,    "metrics.json"  )
						with open(snake_case  ) as f:
									UpperCamelCase__       =							json.load(snake_case  )
						return result




			@require_torch_gpu
			def 					snake_case__					(				self  ):

						'''simple docstring'''


						UpperCamelCase__       =							self._run_finetune(gpus=1  )
						self.assertGreaterEqual(result["test"][0]["test_avg_em"]   ,    0.2  )




			@require_torch_multi_gpu
			def 					snake_case__					(				self  ):

						'''simple docstring'''


						UpperCamelCase__       =							self._run_finetune(gpus=2  )
						self.assertGreaterEqual(result["test"][0]["test_avg_em"]   ,    0.2  )




			@require_torch_gpu
			@require_ray
			def 					snake_case__					(				self  ):

						'''simple docstring'''


						UpperCamelCase__       =							self._run_finetune(gpus=1   ,    distributed_retriever="ray"  )
						self.assertGreaterEqual(result["test"][0]["test_avg_em"]   ,    0.2  )







			@require_torch_multi_gpu
			@require_ray
			def 					snake_case__					(				self  ):

						'''simple docstring'''


						UpperCamelCase__       =							self._run_finetune(gpus=1   ,    distributed_retriever="ray"  )
						self.assertGreaterEqual(result["test"][0]["test_avg_em"]   ,    0.2  )


