





'''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 (					__lowercase						,					__lowercase						,					__lowercase						,					unittest.TestCase		):
		lowercase__:   Optional[int]				     =   AltDiffusionPipeline
		lowercase__:   Tuple				     =   TEXT_TO_IMAGE_PARAMS
		lowercase__:   List[str]				     =   TEXT_TO_IMAGE_BATCH_PARAMS
		lowercase__:   List[Any]				     =   TEXT_TO_IMAGE_IMAGE_PARAMS
		lowercase__:   List[str]				     =   TEXT_TO_IMAGE_IMAGE_PARAMS

		def 	lowercase__     (   self					:   Dict			)					->    List[str]:


					"""simple docstring"""

					torch.manual_seed(0			)
					__snake_case					:      Dict							      =     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  ,      )
					__snake_case					:      Any							      =     DDIMScheduler(
					    beta_start=0.00085  ,      beta_end=0.012  ,      beta_schedule="""scaled_linear"""  ,      clip_sample=__magic_name__  ,      set_alpha_to_one=__magic_name__  ,      )
					torch.manual_seed(0			)
					__snake_case					:      Any							      =     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			)
					__snake_case					:      int							      =     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  ,      )
					__snake_case					:      Dict							      =     CLIPTextModel(__magic_name__			)

					__snake_case					:      Optional[int]							      =     XLMRobertaTokenizer.from_pretrained("""hf-internal-testing/tiny-xlm-roberta"""			)
					__snake_case					:      str							      =     77

					__snake_case					:      List[Any]							      =     {
					    """unet""": unet,
					    """scheduler""": scheduler,
					    """vae""": vae,
					    """text_encoder""": text_encoder,
					    """tokenizer""": tokenizer,
					    """safety_checker""": None,
					    """feature_extractor""": None,
					}
					return components

		def 	lowercase__     (   self					:   Any  ,      __magic_name__					:   Union[str, Any]  ,      __magic_name__					:   Any=0			)					->    Optional[int]:


					"""simple docstring"""

					if str(__magic_name__			).startswith("""mps"""			):
								__snake_case					:      int							      =     torch.manual_seed(__magic_name__			)
					else:
								__snake_case					:      Optional[Any]							      =     torch.Generator(device=__magic_name__			).manual_seed(__magic_name__			)
					__snake_case					:      List[Any]							      =     {
					    """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 	lowercase__     (   self					:   Union[str, Any]			)					->    List[str]:


					"""simple docstring"""

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

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


					"""simple docstring"""

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

		def 	lowercase__     (   self					:   Dict			)					->    List[Any]:


					"""simple docstring"""

					__snake_case					:      Dict							      =     """cpu"""  # ensure determinism for the device-dependent torch.Generator

					__snake_case					:      Optional[Any]							      =     self.get_dummy_components()
					torch.manual_seed(0			)
					__snake_case					:      str							      =     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
					__snake_case					:      Optional[int]							      =     RobertaSeriesModelWithTransformation(__magic_name__			)
					__snake_case					:      Optional[Any]							      =     text_encoder

					__snake_case					:      Any							      =     AltDiffusionPipeline(**__magic_name__			)
					__snake_case					:      Any							      =     alt_pipe.to(__magic_name__			)
					alt_pipe.set_progress_bar_config(disable=__magic_name__			)

					__snake_case					:      int							      =     self.get_dummy_inputs(__magic_name__			)
					__snake_case					:      Union[str, Any]							      =     """A photo of an astronaut"""
					__snake_case					:      Any							      =     alt_pipe(**__magic_name__			)
					__snake_case					:      List[Any]							      =     output.images
					__snake_case					:      Optional[Any]							      =     image[0, -3:, -3:, -1]

					assert image.shape == (1, 64, 64, 3)
					__snake_case					:      int							      =     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 	lowercase__     (   self					:   Optional[int]			)					->    Optional[int]:


					"""simple docstring"""

					__snake_case					:      Any							      =     """cpu"""  # ensure determinism for the device-dependent torch.Generator

					__snake_case					:      Any							      =     self.get_dummy_components()
					__snake_case					:      List[Any]							      =     PNDMScheduler(skip_prk_steps=__magic_name__			)
					torch.manual_seed(0			)
					__snake_case					:      int							      =     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
					__snake_case					:      Dict							      =     RobertaSeriesModelWithTransformation(__magic_name__			)
					__snake_case					:      Tuple							      =     text_encoder
					__snake_case					:      Optional[Any]							      =     AltDiffusionPipeline(**__magic_name__			)
					__snake_case					:      Optional[int]							      =     alt_pipe.to(__magic_name__			)
					alt_pipe.set_progress_bar_config(disable=__magic_name__			)

					__snake_case					:      Dict							      =     self.get_dummy_inputs(__magic_name__			)
					__snake_case					:      List[Any]							      =     alt_pipe(**__magic_name__			)
					__snake_case					:      Dict							      =     output.images
					__snake_case					:      Any							      =     image[0, -3:, -3:, -1]

					assert image.shape == (1, 64, 64, 3)
					__snake_case					:      str							      =     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  _A (					unittest.TestCase		):

		def 	lowercase__     (   self					:   Tuple			)					->    Union[str, Any]:


					"""simple docstring"""

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

		def 	lowercase__     (   self					:   List[Any]			)					->    Optional[Any]:


					"""simple docstring"""

					__snake_case					:      List[str]							      =     AltDiffusionPipeline.from_pretrained("""BAAI/AltDiffusion"""  ,      safety_checker=__magic_name__			)
					__snake_case					:      List[str]							      =     alt_pipe.to(__magic_name__			)
					alt_pipe.set_progress_bar_config(disable=__magic_name__			)

					__snake_case					:      int							      =     """A painting of a squirrel eating a burger"""
					__snake_case					:      Tuple							      =     torch.manual_seed(0			)
					__snake_case					:      Dict							      =     alt_pipe([prompt]  ,      generator=__magic_name__  ,      guidance_scale=6.0  ,      num_inference_steps=20  ,      output_type="""np"""			)

					__snake_case					:      Optional[Any]							      =     output.images

					__snake_case					:      Optional[Any]							      =     image[0, -3:, -3:, -1]

					assert image.shape == (1, 5_12, 5_12, 3)
					__snake_case					:      str							      =     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 	lowercase__     (   self					:   Tuple			)					->    Optional[Any]:


					"""simple docstring"""

					__snake_case					:      Union[str, Any]							      =     DDIMScheduler.from_pretrained("""BAAI/AltDiffusion"""  ,      subfolder="""scheduler"""			)

					__snake_case					:      Dict							      =     AltDiffusionPipeline.from_pretrained("""BAAI/AltDiffusion"""  ,      scheduler=__magic_name__  ,      safety_checker=__magic_name__			)
					__snake_case					:      List[str]							      =     alt_pipe.to(__magic_name__			)
					alt_pipe.set_progress_bar_config(disable=__magic_name__			)

					__snake_case					:      Dict							      =     """A painting of a squirrel eating a burger"""
					__snake_case					:      Dict							      =     torch.manual_seed(0			)

					__snake_case					:      Any							      =     alt_pipe([prompt]  ,      generator=__magic_name__  ,      num_inference_steps=2  ,      output_type="""numpy"""			)
					__snake_case					:      List[Any]							      =     output.images

					__snake_case					:      Tuple							      =     image[0, -3:, -3:, -1]

					assert image.shape == (1, 5_12, 5_12, 3)
					__snake_case					:      Union[str, Any]							      =     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









'''simple docstring'''






import json
import os
from pathlib import Path

import pytest

from datasets.download.download_config import DownloadConfig
from datasets.download.download_manager import DownloadManager
from datasets.utils.file_utils import hash_url_to_filename


__UpperCamelCase					    =	"http://www.mocksite.com/file1.txt"
__UpperCamelCase					    =	"\"text\": [\"foo\", \"foo\"]"
__UpperCamelCase					    =	"6d8ce9aa78a471c7477201efbeabd3bb01ac2e7d100a6dc024ba1608361f90a8"




class  _A :
		lowercase__:   str				     =   200
		lowercase__:   List[str]				     =   {'''Content-Length''': '''100'''}
		lowercase__:   Union[str, Any]				     =   {}

		def 	lowercase__     (   self					:   Any  ,      **__magic_name__					:   List[Any]			)					->    Dict:


					"""simple docstring"""

					return [bytes(__magic_name__  ,      """utf-8"""			)]







def 					_a						(       *_lowerCamelCase    ,	**_lowerCamelCase ) ->				List[str]:



			"""simple docstring"""


			return MockResponse()







@pytest.mark.parametrize("""urls_type"""    ,	[str, list, dict] )
def 					_a						(       _lowerCamelCase    ,	_lowerCamelCase    ,	_lowerCamelCase ) ->				List[str]:



			"""simple docstring"""


			import requests

			monkeypatch.setattr(_lowerCamelCase    ,	"""request"""    ,	_lowerCamelCase )

			__snake_case					:      Union[str, Any]							      =     URL
			if issubclass(_lowerCamelCase    ,	_lowerCamelCase ):
						__snake_case					:      str							      =     url
			elif issubclass(_lowerCamelCase    ,	_lowerCamelCase ):
						__snake_case					:      Dict							      =     [url]
			elif issubclass(_lowerCamelCase    ,	_lowerCamelCase ):
						__snake_case					:      Union[str, Any]							      =     {"""train""": url}
			__snake_case					:      Dict							      =     """dummy"""
			__snake_case					:      List[str]							      =     """downloads"""
			__snake_case					:      List[Any]							      =     tmp_path
			__snake_case					:      List[Any]							      =     DownloadConfig(
			    cache_dir=os.path.join(_lowerCamelCase    ,	_lowerCamelCase )    ,	use_etag=_lowerCamelCase    ,	)
			__snake_case					:      List[str]							      =     DownloadManager(dataset_name=_lowerCamelCase    ,	download_config=_lowerCamelCase )
			__snake_case					:      int							      =     dl_manager.download(_lowerCamelCase )
			__snake_case					:      Tuple							      =     urls
			for downloaded_paths in [downloaded_paths]:
						if isinstance(_lowerCamelCase    ,	_lowerCamelCase ):
									__snake_case					:      Any							      =     [downloaded_paths]
									__snake_case					:      List[Any]							      =     [urls]
						elif isinstance(_lowerCamelCase    ,	_lowerCamelCase ):
									assert "train" in downloaded_paths.keys()
									__snake_case					:      Tuple							      =     downloaded_paths.values()
									__snake_case					:      Optional[int]							      =     urls.values()
						assert downloaded_paths
						for downloaded_path, input_url in zip(_lowerCamelCase    ,	_lowerCamelCase ):
									assert downloaded_path == dl_manager.downloaded_paths[input_url]
									__snake_case					:      List[str]							      =     Path(_lowerCamelCase )
									__snake_case					:      Any							      =     downloaded_path.parts
									assert parts[-1] == HASH
									assert parts[-2] == cache_subdir
									assert downloaded_path.exists()
									__snake_case					:      Union[str, Any]							      =     downloaded_path.read_text()
									assert content == CONTENT
									__snake_case					:      List[str]							      =     downloaded_path.with_suffix(""".json""" )
									assert metadata_downloaded_path.exists()
									__snake_case					:      Union[str, Any]							      =     json.loads(metadata_downloaded_path.read_text() )
									assert metadata_content == {"url": URL, "etag": None}







@pytest.mark.parametrize("""paths_type"""    ,	[str, list, dict] )
def 					_a						(       _lowerCamelCase    ,	_lowerCamelCase    ,	_lowerCamelCase ) ->				List[Any]:



			"""simple docstring"""


			__snake_case					:      Any							      =     str(_lowerCamelCase )
			if issubclass(_lowerCamelCase    ,	_lowerCamelCase ):
						__snake_case					:      Optional[int]							      =     filename
			elif issubclass(_lowerCamelCase    ,	_lowerCamelCase ):
						__snake_case					:      Tuple							      =     [filename]
			elif issubclass(_lowerCamelCase    ,	_lowerCamelCase ):
						__snake_case					:      Dict							      =     {"""train""": filename}
			__snake_case					:      Optional[Any]							      =     """dummy"""
			__snake_case					:      List[Any]							      =     xz_file.parent
			__snake_case					:      int							      =     """extracted"""
			__snake_case					:      Dict							      =     DownloadConfig(
			    cache_dir=_lowerCamelCase    ,	use_etag=_lowerCamelCase    ,	)
			__snake_case					:      List[str]							      =     DownloadManager(dataset_name=_lowerCamelCase    ,	download_config=_lowerCamelCase )
			__snake_case					:      Optional[Any]							      =     dl_manager.extract(_lowerCamelCase )
			__snake_case					:      Union[str, Any]							      =     paths
			for extracted_paths in [extracted_paths]:
						if isinstance(_lowerCamelCase    ,	_lowerCamelCase ):
									__snake_case					:      Dict							      =     [extracted_paths]
									__snake_case					:      int							      =     [paths]
						elif isinstance(_lowerCamelCase    ,	_lowerCamelCase ):
									assert "train" in extracted_paths.keys()
									__snake_case					:      int							      =     extracted_paths.values()
									__snake_case					:      int							      =     paths.values()
						assert extracted_paths
						for extracted_path, input_path in zip(_lowerCamelCase    ,	_lowerCamelCase ):
									assert extracted_path == dl_manager.extracted_paths[input_path]
									__snake_case					:      Any							      =     Path(_lowerCamelCase )
									__snake_case					:      str							      =     extracted_path.parts
									assert parts[-1] == hash_url_to_filename(_lowerCamelCase    ,	etag=_lowerCamelCase )
									assert parts[-2] == extracted_subdir
									assert extracted_path.exists()
									__snake_case					:      Optional[int]							      =     extracted_path.read_text()
									__snake_case					:      str							      =     text_file.read_text()
									assert extracted_file_content == expected_file_content







def 					_a						(       _lowerCamelCase    ,	_lowerCamelCase ) ->				Optional[int]:



			"""simple docstring"""


			assert path.endswith(""".jsonl""" )
			for num_items, line in enumerate(_lowerCamelCase    ,	start=1 ):
						__snake_case					:      Tuple							      =     json.loads(line.decode("""utf-8""" ) )
						assert item.keys() == {"col_1", "col_2", "col_3"}
			assert num_items == 4







@pytest.mark.parametrize("""archive_jsonl"""    ,	["""tar_jsonl_path""", """zip_jsonl_path"""] )
def 					_a						(       _lowerCamelCase    ,	_lowerCamelCase ) ->				Optional[int]:



			"""simple docstring"""


			__snake_case					:      Any							      =     request.getfixturevalue(_lowerCamelCase )
			__snake_case					:      str							      =     DownloadManager()
			for num_jsonl, (path, file) in enumerate(dl_manager.iter_archive(_lowerCamelCase )    ,	start=1 ):
						_test_jsonl(_lowerCamelCase    ,	_lowerCamelCase )
			assert num_jsonl == 2







@pytest.mark.parametrize("""archive_nested_jsonl"""    ,	["""tar_nested_jsonl_path""", """zip_nested_jsonl_path"""] )
def 					_a						(       _lowerCamelCase    ,	_lowerCamelCase ) ->				List[str]:



			"""simple docstring"""


			__snake_case					:      int							      =     request.getfixturevalue(_lowerCamelCase )
			__snake_case					:      List[str]							      =     DownloadManager()
			for num_tar, (path, file) in enumerate(dl_manager.iter_archive(_lowerCamelCase )    ,	start=1 ):
						for num_jsonl, (subpath, subfile) in enumerate(dl_manager.iter_archive(_lowerCamelCase )    ,	start=1 ):
									_test_jsonl(_lowerCamelCase    ,	_lowerCamelCase )
			assert num_tar == 1
			assert num_jsonl == 2







def 					_a						(       _lowerCamelCase ) ->				Any:



			"""simple docstring"""


			__snake_case					:      List[str]							      =     DownloadManager()
			for num_file, file in enumerate(dl_manager.iter_files(_lowerCamelCase )    ,	start=1 ):
						assert os.path.basename(_lowerCamelCase ) == ("test.txt" if num_file == 1 else "train.txt")
			assert num_file == 2


