


'''simple docstring'''


import gc
import os
import tempfile
import unittest

import torch
import torch.nn as nn
import torch.nn.functional as F
from transformers import CLIPTextConfig, CLIPTextModel, CLIPTokenizer

from diffusers import AutoencoderKL, DDIMScheduler, StableDiffusionPipeline, UNetaDConditionModel
from diffusers.loaders import AttnProcsLayers, LoraLoaderMixin
from diffusers.models.attention_processor import (
    Attention,
    AttnProcessor,
    AttnProcessor2_0,
    LoRAAttnProcessor,
    LoRAAttnProcessor2_0,
    LoRAXFormersAttnProcessor,
    XFormersAttnProcessor,
)
from diffusers.utils import TEXT_ENCODER_ATTN_MODULE, floats_tensor, torch_device






def 			__lowerCAmelCase			(__lowerCAmelCase ):
		_UpperCAmelCase		:							Dict             =   {}
		for name in unet.attn_processors.keys():
				_UpperCAmelCase		:							int             =   None if name.endswith("attn1.processor" ) else unet.config.cross_attention_dim
				if name.startswith("mid_block" ):
						_UpperCAmelCase		:							Any             =   unet.config.block_out_channels[-1]
				elif name.startswith("up_blocks" ):
						_UpperCAmelCase		:							Union[str, Any]             =   int(name[len("up_blocks." )] )
						_UpperCAmelCase		:							Any             =   list(reversed(unet.config.block_out_channels ) )[block_id]
				elif name.startswith("down_blocks" ):
						_UpperCAmelCase		:							Dict             =   int(name[len("down_blocks." )] )
						_UpperCAmelCase		:							Optional[Any]             =   unet.config.block_out_channels[block_id]
				_UpperCAmelCase		:							Union[str, Any]             =   (
				    LoRAAttnProcessor2_0 if hasattr(__lowerCAmelCase   ,			"scaled_dot_product_attention" ) else LoRAAttnProcessor
				)
				_UpperCAmelCase		:							Optional[Any]             =   lora_attn_processor_class(
				    hidden_size=__lowerCAmelCase   ,			cross_attention_dim=__lowerCAmelCase )
		_UpperCAmelCase		:							Dict             =   AttnProcsLayers(__lowerCAmelCase )
		return lora_attn_procs, unet_lora_layers






def 			__lowerCAmelCase			(__lowerCAmelCase ):
		_UpperCAmelCase		:							List[str]             =   {}
		_UpperCAmelCase		:							Any             =   (
		    LoRAAttnProcessor2_0 if hasattr(__lowerCAmelCase   ,			"scaled_dot_product_attention" ) else LoRAAttnProcessor
		)
		for name, module in text_encoder.named_modules():
				if name.endswith(__lowerCAmelCase ):
						_UpperCAmelCase		:							Dict             =   lora_attn_processor_class(
						    hidden_size=module.out_proj.out_features   ,			cross_attention_dim=__lowerCAmelCase )
		return text_lora_attn_procs






def 			__lowerCAmelCase			(__lowerCAmelCase ):
		_UpperCAmelCase		:							int             =   create_text_encoder_lora_attn_procs(__lowerCAmelCase )
		_UpperCAmelCase		:							Optional[Any]             =   AttnProcsLayers(__lowerCAmelCase )
		return text_encoder_lora_layers






def 			__lowerCAmelCase			(__lowerCAmelCase   ,			__lowerCAmelCase=False ):
		for _, attn_proc in text_lora_attn_procs.items():
				# set up.weights
				for layer_name, layer_module in attn_proc.named_modules():
						if layer_name.endswith("_lora" ):
								_UpperCAmelCase		:							Optional[int]             =   (
								    torch.randn_like(layer_module.up.weight )
								    if randn_weight
								    else torch.zeros_like(layer_module.up.weight )
								)
								_UpperCAmelCase		:							Tuple             =   torch.nn.Parameter(__lowerCAmelCase )






class 						lowerCAmelCase__      (      unittest.TestCase						):




			def 				lowerCAmelCase__     (		self	:  Union[str, Any]			) ->Dict:



					'''simple docstring'''







					torch.manual_seed(0			)
					_UpperCAmelCase		:							Tuple             =   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		:							Union[str, Any]             =   DDIMScheduler(
					    beta_start=0.0_0_0_8_5						,	beta_end=0.0_1_2						,	beta_schedule="scaled_linear"						,	clip_sample=lowerCamelCase__						,	set_alpha_to_one=lowerCamelCase__						,	steps_offset=1						,	)
					torch.manual_seed(0			)
					_UpperCAmelCase		:							Optional[int]             =   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						,	)
					_UpperCAmelCase		:							int             =   CLIPTextConfig(
					    bos_token_id=0						,	eos_token_id=2						,	hidden_size=32						,	intermediate_size=37						,	layer_norm_eps=1E-05						,	num_attention_heads=4						,	num_hidden_layers=5						,	pad_token_id=1						,	vocab_size=10_00						,	)
					_UpperCAmelCase		:							Dict             =   CLIPTextModel(lowerCamelCase__			)
					_UpperCAmelCase		:							List[str]             =   CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip"			)

					_UpperCAmelCase       ,     _UpperCAmelCase		:							str             =   create_unet_lora_layers(lowerCamelCase__			)
					_UpperCAmelCase		:							str             =   create_text_encoder_lora_layers(lowerCamelCase__			)

					_UpperCAmelCase		:							str             =   {
					    "unet": unet,
					    "scheduler": scheduler,
					    "vae": vae,
					    "text_encoder": text_encoder,
					    "tokenizer": tokenizer,
					    "safety_checker": None,
					    "feature_extractor": None,
					}
					_UpperCAmelCase		:							Any             =   {
					    "unet_lora_layers": unet_lora_layers,
					    "text_encoder_lora_layers": text_encoder_lora_layers,
					    "unet_lora_attn_procs": unet_lora_attn_procs,
					}
					return pipeline_components, lora_components




			def 				lowerCAmelCase__     (		self	:  str			) ->List[str]:



					'''simple docstring'''







					_UpperCAmelCase		:							List[str]             =   1
					_UpperCAmelCase		:							Any             =   10
					_UpperCAmelCase		:							Union[str, Any]             =   4
					_UpperCAmelCase		:							List[str]             =   (32, 32)

					_UpperCAmelCase		:							Optional[Any]             =   torch.manual_seed(0			)
					_UpperCAmelCase		:							List[Any]             =   floats_tensor((batch_size, num_channels) + sizes			)
					_UpperCAmelCase		:							str             =   torch.randint(1						,	lowerCamelCase__						,	size=(batch_size, sequence_length)						,	generator=lowerCamelCase__			)

					_UpperCAmelCase		:							Optional[int]             =   {
					    "prompt": "A painting of a squirrel eating a burger",
					    "generator": generator,
					    "num_inference_steps": 2,
					    "guidance_scale": 6.0,
					    "output_type": "numpy",
					}

					return noise, input_ids, pipeline_inputs

					# copied from: https://colab.research.google.com/gist/sayakpaul/df2ef6e1ae6d8c10a49d859883b10860/scratchpad.ipynb




			def 				lowerCAmelCase__     (		self	:  List[str]			) ->str:



					'''simple docstring'''







					_UpperCAmelCase		:							Any             =   77

					_UpperCAmelCase		:							str             =   torch.randint(2						,	56						,	size=(1, max_seq_length)						,	generator=torch.manual_seed(0			)			)

					_UpperCAmelCase		:							Dict             =   {}
					_UpperCAmelCase		:							Union[str, Any]             =   inputs
					return prepared_inputs




			def 				lowerCAmelCase__     (		self	:  Any						,	lowerCamelCase__	:  Tuple			) ->Optional[int]:



					'''simple docstring'''







					_UpperCAmelCase       ,     _UpperCAmelCase		:							Any             =   self.get_dummy_components()
					LoraLoaderMixin.save_lora_weights(
					    save_directory=lowerCamelCase__						,	unet_lora_layers=lora_components["unet_lora_layers"]						,	text_encoder_lora_layers=lora_components["text_encoder_lora_layers"]						,	)
					self.assertTrue(os.path.isfile(os.path.join(lowerCamelCase__						,	"pytorch_lora_weights.bin"			)			)			)




			def 				lowerCAmelCase__     (		self	:  Optional[int]			) ->int:



					'''simple docstring'''







					_UpperCAmelCase       ,     _UpperCAmelCase		:							Union[str, Any]             =   self.get_dummy_components()
					_UpperCAmelCase		:							Union[str, Any]             =   StableDiffusionPipeline(**lowerCamelCase__			)
					_UpperCAmelCase		:							Optional[int]             =   sd_pipe.to(lowerCamelCase__			)
					sd_pipe.set_progress_bar_config(disable=lowerCamelCase__			)

					_UpperCAmelCase       ,     _UpperCAmelCase       ,     _UpperCAmelCase		:							Dict             =   self.get_dummy_inputs()

					_UpperCAmelCase		:							List[str]             =   sd_pipe(**lowerCamelCase__			).images
					_UpperCAmelCase		:							List[str]             =   original_images[0, -3:, -3:, -1]

					with tempfile.TemporaryDirectory() as tmpdirname:
							LoraLoaderMixin.save_lora_weights(
							    save_directory=lowerCamelCase__						,	unet_lora_layers=lora_components["unet_lora_layers"]						,	text_encoder_lora_layers=lora_components["text_encoder_lora_layers"]						,	)
							self.assertTrue(os.path.isfile(os.path.join(lowerCamelCase__						,	"pytorch_lora_weights.bin"			)			)			)
							sd_pipe.load_lora_weights(lowerCamelCase__			)

					_UpperCAmelCase		:							Optional[Any]             =   sd_pipe(**lowerCamelCase__			).images
					_UpperCAmelCase		:							Optional[Any]             =   lora_images[0, -3:, -3:, -1]

					# Outputs shouldn't match.
					self.assertFalse(torch.allclose(torch.from_numpy(lowerCamelCase__			)						,	torch.from_numpy(lowerCamelCase__			)			)			)




			def 				lowerCAmelCase__     (		self	:  Any			) ->List[str]:



					'''simple docstring'''







					_UpperCAmelCase       ,     _UpperCAmelCase		:							str             =   self.get_dummy_components()
					_UpperCAmelCase		:							List[str]             =   StableDiffusionPipeline(**lowerCamelCase__			)
					_UpperCAmelCase		:							Any             =   sd_pipe.to(lowerCamelCase__			)
					sd_pipe.set_progress_bar_config(disable=lowerCamelCase__			)

					_UpperCAmelCase       ,     _UpperCAmelCase       ,     _UpperCAmelCase		:							List[Any]             =   self.get_dummy_inputs()

					_UpperCAmelCase		:							int             =   sd_pipe(**lowerCamelCase__			).images
					_UpperCAmelCase		:							Union[str, Any]             =   original_images[0, -3:, -3:, -1]

					with tempfile.TemporaryDirectory() as tmpdirname:
							LoraLoaderMixin.save_lora_weights(
							    save_directory=lowerCamelCase__						,	unet_lora_layers=lora_components["unet_lora_layers"]						,	text_encoder_lora_layers=lora_components["text_encoder_lora_layers"]						,	safe_serialization=lowerCamelCase__						,	)
							self.assertTrue(os.path.isfile(os.path.join(lowerCamelCase__						,	"pytorch_lora_weights.safetensors"			)			)			)
							sd_pipe.load_lora_weights(lowerCamelCase__			)

					_UpperCAmelCase		:							Optional[Any]             =   sd_pipe(**lowerCamelCase__			).images
					_UpperCAmelCase		:							Optional[Any]             =   lora_images[0, -3:, -3:, -1]

					# Outputs shouldn't match.
					self.assertFalse(torch.allclose(torch.from_numpy(lowerCamelCase__			)						,	torch.from_numpy(lowerCamelCase__			)			)			)




			def 				lowerCAmelCase__     (		self	:  Optional[Any]			) ->Union[str, Any]:



					'''simple docstring'''







					_UpperCAmelCase       ,     _UpperCAmelCase		:							Tuple             =   self.get_dummy_components()
					_UpperCAmelCase		:							Tuple             =   lora_components["unet_lora_attn_procs"]
					_UpperCAmelCase		:							Dict             =   StableDiffusionPipeline(**lowerCamelCase__			)
					_UpperCAmelCase		:							List[Any]             =   sd_pipe.to(lowerCamelCase__			)
					sd_pipe.set_progress_bar_config(disable=lowerCamelCase__			)

					_UpperCAmelCase       ,     _UpperCAmelCase       ,     _UpperCAmelCase		:							Dict             =   self.get_dummy_inputs()

					_UpperCAmelCase		:							Any             =   sd_pipe(**lowerCamelCase__			).images
					_UpperCAmelCase		:							Tuple             =   original_images[0, -3:, -3:, -1]

					with tempfile.TemporaryDirectory() as tmpdirname:
							_UpperCAmelCase		:							Dict             =   sd_pipe.unet
							unet.set_attn_processor(lowerCamelCase__			)
							unet.save_attn_procs(lowerCamelCase__			)
							self.assertTrue(os.path.isfile(os.path.join(lowerCamelCase__						,	"pytorch_lora_weights.bin"			)			)			)
							sd_pipe.load_lora_weights(lowerCamelCase__			)

					_UpperCAmelCase		:							Dict             =   sd_pipe(**lowerCamelCase__			).images
					_UpperCAmelCase		:							Optional[int]             =   lora_images[0, -3:, -3:, -1]

					# Outputs shouldn't match.
					self.assertFalse(torch.allclose(torch.from_numpy(lowerCamelCase__			)						,	torch.from_numpy(lowerCamelCase__			)			)			)




			def 				lowerCAmelCase__     (		self	:  Dict			) ->Dict:



					'''simple docstring'''







					_UpperCAmelCase       ,     _UpperCAmelCase		:							Any             =   self.get_dummy_components()
					_UpperCAmelCase		:							int             =   StableDiffusionPipeline(**lowerCamelCase__			)

					_UpperCAmelCase		:							List[Any]             =   self.get_dummy_tokens()

					# inference without lora
					_UpperCAmelCase		:							Any             =   pipe.text_encoder(**lowerCamelCase__			)[0]
					assert outputs_without_lora.shape == (1, 77, 32)

					# create lora_attn_procs with zeroed out up.weights
					_UpperCAmelCase		:							Dict             =   create_text_encoder_lora_attn_procs(pipe.text_encoder			)
					set_lora_up_weights(lowerCamelCase__						,	randn_weight=lowerCamelCase__			)

					# monkey patch
					pipe._modify_text_encoder(lowerCamelCase__			)

					# verify that it's okay to release the text_attn_procs which holds the LoRAAttnProcessor.
					del text_attn_procs
					gc.collect()

					# inference with lora
					_UpperCAmelCase		:							Union[str, Any]             =   pipe.text_encoder(**lowerCamelCase__			)[0]
					assert outputs_with_lora.shape == (1, 77, 32)

					assert torch.allclose(
					    lowerCamelCase__						,	lowerCamelCase__			), "lora_up_weight are all zero, so the lora outputs should be the same to without lora outputs"

					# create lora_attn_procs with randn up.weights
					_UpperCAmelCase		:							Any             =   create_text_encoder_lora_attn_procs(pipe.text_encoder			)
					set_lora_up_weights(lowerCamelCase__						,	randn_weight=lowerCamelCase__			)

					# monkey patch
					pipe._modify_text_encoder(lowerCamelCase__			)

					# verify that it's okay to release the text_attn_procs which holds the LoRAAttnProcessor.
					del text_attn_procs
					gc.collect()

					# inference with lora
					_UpperCAmelCase		:							Dict             =   pipe.text_encoder(**lowerCamelCase__			)[0]
					assert outputs_with_lora.shape == (1, 77, 32)

					assert not torch.allclose(
					    lowerCamelCase__						,	lowerCamelCase__			), "lora_up_weight are not zero, so the lora outputs should be different to without lora outputs"




			def 				lowerCAmelCase__     (		self	:  str			) ->str:



					'''simple docstring'''







					_UpperCAmelCase       ,     _UpperCAmelCase		:							Tuple             =   self.get_dummy_components()
					_UpperCAmelCase		:							Dict             =   StableDiffusionPipeline(**lowerCamelCase__			)

					_UpperCAmelCase		:							str             =   self.get_dummy_tokens()

					# inference without lora
					_UpperCAmelCase		:							List[str]             =   pipe.text_encoder(**lowerCamelCase__			)[0]
					assert outputs_without_lora.shape == (1, 77, 32)

					# create lora_attn_procs with randn up.weights
					_UpperCAmelCase		:							Optional[Any]             =   create_text_encoder_lora_attn_procs(pipe.text_encoder			)
					set_lora_up_weights(lowerCamelCase__						,	randn_weight=lowerCamelCase__			)

					# monkey patch
					pipe._modify_text_encoder(lowerCamelCase__			)

					# verify that it's okay to release the text_attn_procs which holds the LoRAAttnProcessor.
					del text_attn_procs
					gc.collect()

					# inference with lora
					_UpperCAmelCase		:							Dict             =   pipe.text_encoder(**lowerCamelCase__			)[0]
					assert outputs_with_lora.shape == (1, 77, 32)

					assert not torch.allclose(
					    lowerCamelCase__						,	lowerCamelCase__			), "lora outputs should be different to without lora outputs"

					# remove monkey patch
					pipe._remove_text_encoder_monkey_patch()

					# inference with removed lora
					_UpperCAmelCase		:							int             =   pipe.text_encoder(**lowerCamelCase__			)[0]
					assert outputs_without_lora_removed.shape == (1, 77, 32)

					assert torch.allclose(
					    lowerCamelCase__						,	lowerCamelCase__			), "remove lora monkey patch should restore the original outputs"




			def 				lowerCAmelCase__     (		self	:  Tuple			) ->Any:



					'''simple docstring'''







					_UpperCAmelCase       ,     _UpperCAmelCase		:							Any             =   self.get_dummy_components()
					_UpperCAmelCase		:							List[str]             =   StableDiffusionPipeline(**lowerCamelCase__			)
					_UpperCAmelCase		:							Dict             =   sd_pipe.to(lowerCamelCase__			)
					sd_pipe.set_progress_bar_config(disable=lowerCamelCase__			)

					_UpperCAmelCase       ,     _UpperCAmelCase       ,     _UpperCAmelCase		:							Optional[int]             =   self.get_dummy_inputs()

					with tempfile.TemporaryDirectory() as tmpdirname:
							LoraLoaderMixin.save_lora_weights(
							    save_directory=lowerCamelCase__						,	unet_lora_layers=lora_components["unet_lora_layers"]						,	text_encoder_lora_layers=lora_components["text_encoder_lora_layers"]						,	)
							self.assertTrue(os.path.isfile(os.path.join(lowerCamelCase__						,	"pytorch_lora_weights.bin"			)			)			)
							sd_pipe.load_lora_weights(lowerCamelCase__			)

					_UpperCAmelCase		:							Any             =   sd_pipe(**lowerCamelCase__			).images
					_UpperCAmelCase		:							List[Any]             =   lora_images[0, -3:, -3:, -1]

					_UpperCAmelCase		:							Optional[int]             =   sd_pipe(**lowerCamelCase__						,	cross_attention_kwargs={"scale": 0.5}			).images
					_UpperCAmelCase		:							Optional[int]             =   lora_images_with_scale[0, -3:, -3:, -1]

					# Outputs shouldn't match.
					self.assertFalse(
					    torch.allclose(torch.from_numpy(lowerCamelCase__			)						,	torch.from_numpy(lowerCamelCase__			)			)			)




			def 				lowerCAmelCase__     (		self	:  Dict			) ->Union[str, Any]:



					'''simple docstring'''







					with tempfile.TemporaryDirectory() as tmpdirname:
							self.create_lora_weight_file(lowerCamelCase__			)

							_UpperCAmelCase       ,     _UpperCAmelCase		:							Optional[int]             =   self.get_dummy_components()
							_UpperCAmelCase		:							List[str]             =   StableDiffusionPipeline(**lowerCamelCase__			)
							_UpperCAmelCase		:							List[str]             =   sd_pipe.to(lowerCamelCase__			)
							sd_pipe.set_progress_bar_config(disable=lowerCamelCase__			)

							# check if vanilla attention processors are used
							for _, module in sd_pipe.unet.named_modules():
									if isinstance(lowerCamelCase__						,	lowerCamelCase__			):
											self.assertIsInstance(module.processor						,	(AttnProcessor, AttnProcessor2_0)			)

            # load LoRA weight file
							sd_pipe.load_lora_weights(lowerCamelCase__			)

							# check if lora attention processors are used
							for _, module in sd_pipe.unet.named_modules():
									if isinstance(lowerCamelCase__						,	lowerCamelCase__			):
											_UpperCAmelCase		:							Any             =   (
											    LoRAAttnProcessor2_0 if hasattr(lowerCamelCase__						,	"scaled_dot_product_attention"			) else LoRAAttnProcessor
											)
											self.assertIsInstance(module.processor						,	lowerCamelCase__			)




			@unittest.skipIf(torch_device != "cuda"						,	"This test is supposed to run on GPU"			)
			def 				lowerCAmelCase__     (		self	:  Optional[Any]			) ->int:



					'''simple docstring'''







					with tempfile.TemporaryDirectory() as tmpdirname:
							self.create_lora_weight_file(lowerCamelCase__			)

							_UpperCAmelCase       ,     _UpperCAmelCase		:							str             =   self.get_dummy_components()
							_UpperCAmelCase		:							Tuple             =   StableDiffusionPipeline(**lowerCamelCase__			)
							_UpperCAmelCase		:							Union[str, Any]             =   sd_pipe.to(lowerCamelCase__			)
							sd_pipe.set_progress_bar_config(disable=lowerCamelCase__			)

							# enable XFormers
							sd_pipe.enable_xformers_memory_efficient_attention()

							# check if xFormers attention processors are used
							for _, module in sd_pipe.unet.named_modules():
									if isinstance(lowerCamelCase__						,	lowerCamelCase__			):
											self.assertIsInstance(module.processor						,	lowerCamelCase__			)

            # load LoRA weight file
							sd_pipe.load_lora_weights(lowerCamelCase__			)

							# check if lora attention processors are used
							for _, module in sd_pipe.unet.named_modules():
									if isinstance(lowerCamelCase__						,	lowerCamelCase__			):
											self.assertIsInstance(module.processor						,	lowerCamelCase__			)






			@unittest.skipIf(torch_device != "cuda"						,	"This test is supposed to run on GPU"			)
			def 				lowerCAmelCase__     (		self	:  Optional[Any]			) ->Dict:



					'''simple docstring'''







					_UpperCAmelCase       ,     _UpperCAmelCase		:							List[str]             =   self.get_dummy_components()
					_UpperCAmelCase		:							int             =   StableDiffusionPipeline(**lowerCamelCase__			)
					_UpperCAmelCase		:							Optional[Any]             =   sd_pipe.to(lowerCamelCase__			)
					sd_pipe.set_progress_bar_config(disable=lowerCamelCase__			)

					_UpperCAmelCase       ,     _UpperCAmelCase       ,     _UpperCAmelCase		:							List[str]             =   self.get_dummy_inputs()

					# enable XFormers
					sd_pipe.enable_xformers_memory_efficient_attention()

					_UpperCAmelCase		:							Union[str, Any]             =   sd_pipe(**lowerCamelCase__			).images
					_UpperCAmelCase		:							str             =   original_images[0, -3:, -3:, -1]

					with tempfile.TemporaryDirectory() as tmpdirname:
							LoraLoaderMixin.save_lora_weights(
							    save_directory=lowerCamelCase__						,	unet_lora_layers=lora_components["unet_lora_layers"]						,	text_encoder_lora_layers=lora_components["text_encoder_lora_layers"]						,	)
							self.assertTrue(os.path.isfile(os.path.join(lowerCamelCase__						,	"pytorch_lora_weights.bin"			)			)			)
							sd_pipe.load_lora_weights(lowerCamelCase__			)

					_UpperCAmelCase		:							Dict             =   sd_pipe(**lowerCamelCase__			).images
					_UpperCAmelCase		:							Optional[int]             =   lora_images[0, -3:, -3:, -1]

					# Outputs shouldn't match.
					self.assertFalse(torch.allclose(torch.from_numpy(lowerCamelCase__			)						,	torch.from_numpy(lowerCamelCase__			)			)			)







'''simple docstring'''


import os
from argparse import ArgumentParser
from pathlib import Path
from shutil import copyfile
from typing import List

from datasets import config
from datasets.builder import DatasetBuilder
from datasets.commands import BaseDatasetsCLICommand
from datasets.download.download_config import DownloadConfig
from datasets.download.download_manager import DownloadMode
from datasets.load import dataset_module_factory, import_main_class
from datasets.utils.info_utils import VerificationMode






def 			__lowerCAmelCase			(__lowerCAmelCase   ,			**__lowerCAmelCase ):
		return RunBeamCommand(
		    args.dataset   ,			args.name   ,			args.cache_dir   ,			args.beam_pipeline_options   ,			args.data_dir   ,			args.all_configs   ,			args.save_info or args.save_infos   ,			args.ignore_verifications   ,			args.force_redownload   ,			**__lowerCAmelCase   ,			)






class 						lowerCAmelCase__      (      UpperCAmelCase__						):




			@staticmethod
			def 				lowerCAmelCase__     (		lowerCamelCase__	:  ArgumentParser			) ->Optional[Any]:



					'''simple docstring'''







					_UpperCAmelCase		:							Union[str, Any]             =   parser.add_parser("run_beam"						,	help="Run a Beam dataset processing pipeline"			)
					run_beam_parser.add_argument("dataset"						,	type=lowerCamelCase__						,	help="Name of the dataset to download"			)
					run_beam_parser.add_argument("--name"						,	type=lowerCamelCase__						,	default=lowerCamelCase__						,	help="Dataset config name"			)
					run_beam_parser.add_argument(
					    "--cache_dir"						,	type=lowerCamelCase__						,	default=lowerCamelCase__						,	help="Cache directory where the datasets are stored"						,	)
					run_beam_parser.add_argument(
					    "--beam_pipeline_options"						,	type=lowerCamelCase__						,	default=""						,	help="Beam pipeline options, separated by commas. Example:: `--beam_pipeline_options=job_name=my-job,project=my-project`"						,	)
					run_beam_parser.add_argument(
					    "--data_dir"						,	type=lowerCamelCase__						,	default=lowerCamelCase__						,	help="Can be used to specify a manual directory to get the files from"						,	)
					run_beam_parser.add_argument("--all_configs"						,	action="store_true"						,	help="Test all dataset configurations"			)
					run_beam_parser.add_argument("--save_info"						,	action="store_true"						,	help="Save the dataset infos file"			)
					run_beam_parser.add_argument(
					    "--ignore_verifications"						,	action="store_true"						,	help="Run the test without checksums and splits checks"			)
					run_beam_parser.add_argument("--force_redownload"						,	action="store_true"						,	help="Force dataset redownload"			)
					# aliases
					run_beam_parser.add_argument("--save_infos"						,	action="store_true"						,	help="alias for save_info"			)
					run_beam_parser.set_defaults(func=lowerCamelCase__			)




			def __init__(		self	:  str						,	lowerCamelCase__	:  str						,	lowerCamelCase__	:  str						,	lowerCamelCase__	:  str						,	lowerCamelCase__	:  str						,	lowerCamelCase__	:  str						,	lowerCamelCase__	:  bool						,	lowerCamelCase__	:  bool						,	lowerCamelCase__	:  bool						,	lowerCamelCase__	:  bool						,	**lowerCamelCase__	:  Optional[Any]						,	) ->Tuple:



					'''simple docstring'''







					_UpperCAmelCase		:							Any             =   dataset
					_UpperCAmelCase		:							str             =   name
					_UpperCAmelCase		:							int             =   cache_dir
					_UpperCAmelCase		:							Optional[int]             =   beam_pipeline_options
					_UpperCAmelCase		:							int             =   data_dir
					_UpperCAmelCase		:							Optional[int]             =   all_configs
					_UpperCAmelCase		:							Any             =   save_infos
					_UpperCAmelCase		:							Optional[Any]             =   ignore_verifications
					_UpperCAmelCase		:							Any             =   force_redownload
					_UpperCAmelCase		:							Tuple             =   config_kwargs






			def 				lowerCAmelCase__     (		self	:  Dict			) ->Union[str, Any]:



					'''simple docstring'''







					import apache_beam as beam

					if self._name is not None and self._all_configs:
							print("Both parameters `name` and `all_configs` can't be used at once."			)
							exit(1			)
					_UpperCAmelCase       ,     _UpperCAmelCase		:							Union[str, Any]             =   self._dataset, self._name
					_UpperCAmelCase		:							Union[str, Any]             =   dataset_module_factory(lowerCamelCase__			)
					_UpperCAmelCase		:							Any             =   import_main_class(dataset_module.module_path			)
					_UpperCAmelCase		:							List[DatasetBuilder]             =   []
					if self._beam_pipeline_options:
							_UpperCAmelCase		:							Optional[int]             =   beam.options.pipeline_options.PipelineOptions(
							    flags=[F"""--{opt.strip()}""" for opt in self._beam_pipeline_options.split(","			) if opt]			)
					else:
							_UpperCAmelCase		:							str             =   None
					if self._all_configs and len(builder_cls.BUILDER_CONFIGS			) > 0:
							for builder_config in builder_cls.BUILDER_CONFIGS:
									builders.append(
									    builder_cls(
									        config_name=builder_config.name						,	data_dir=self._data_dir						,	hash=dataset_module.hash						,	beam_options=lowerCamelCase__						,	cache_dir=self._cache_dir						,	base_path=dataset_module.builder_kwargs.get("base_path"			)						,	)			)
					else:
							builders.append(
							    builder_cls(
							        config_name=lowerCamelCase__						,	data_dir=self._data_dir						,	beam_options=lowerCamelCase__						,	cache_dir=self._cache_dir						,	base_path=dataset_module.builder_kwargs.get("base_path"			)						,	**self._config_kwargs						,	)			)

					for builder in builders:
							builder.download_and_prepare(
							    download_mode=DownloadMode.REUSE_CACHE_IF_EXISTS
							    if not self._force_redownload
							    else DownloadMode.FORCE_REDOWNLOAD						,	download_config=DownloadConfig(cache_dir=config.DOWNLOADED_DATASETS_PATH			)						,	verification_mode=VerificationMode.NO_CHECKS
							    if self._ignore_verifications
							    else VerificationMode.ALL_CHECKS						,	try_from_hf_gcs=lowerCamelCase__						,	)
							if self._save_infos:
									builder._save_infos()

					print("Apache beam run successful."			)

					# If save_infos=True, the dataset infos file is created next to the loaded module file.
					# Let's move it to the original directory of the dataset script, to allow the user to
					# upload them on S3 at the same time afterwards.
					if self._save_infos:
							_UpperCAmelCase		:							Optional[Any]             =   os.path.join(builder_cls.get_imported_module_dir()						,	config.DATASETDICT_INFOS_FILENAME			)

							_UpperCAmelCase		:							List[str]             =   Path(lowerCamelCase__			).name + ".py"

							_UpperCAmelCase		:							int             =   os.path.join(lowerCamelCase__						,	lowerCamelCase__			)
							if os.path.isfile(lowerCamelCase__			):
									_UpperCAmelCase		:							Any             =   os.path.dirname(lowerCamelCase__			)
							elif os.path.isfile(lowerCamelCase__			):
									_UpperCAmelCase		:							Tuple             =   path
							else:  # in case of a remote dataset
									print(F"""Dataset Infos file saved at {dataset_infos_path}"""			)
									exit(1			)

							# Move datasetinfo back to the user
							_UpperCAmelCase		:							Union[str, Any]             =   os.path.join(lowerCamelCase__						,	config.DATASETDICT_INFOS_FILENAME			)
							copyfile(lowerCamelCase__						,	lowerCamelCase__			)
							print(F"""Dataset Infos file saved at {user_dataset_infos_path}"""			)



