



'''simple docstring'''



import gc
import random
import unittest

import numpy as np
import torch
from PIL import Image
from transformers import XLMRobertaTokenizerFast

from diffusers import DDIMScheduler, KandinskyInpaintPipeline, KandinskyPriorPipeline, UNetaDConditionModel, VQModel
from diffusers.pipelines.kandinsky.text_encoder import MCLIPConfig, MultilingualCLIP
from diffusers.utils import floats_tensor, load_image, load_numpy, slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu

from ..test_pipelines_common import PipelineTesterMixin, assert_mean_pixel_difference


enable_full_determinism()
class 						lowerCamelCase_				(       snake_case_     ,					unittest.TestCase      ):
						_lowerCAmelCase  : str						  =					KandinskyInpaintPipeline
						_lowerCAmelCase  : List[Any]						  =					['prompt', 'image_embeds', 'negative_image_embeds', 'image', 'mask_image']
						_lowerCAmelCase  : int						  =					[
						    'prompt',
						    'negative_prompt',
						    'image_embeds',
						    'negative_image_embeds',
						    'image',
						    'mask_image',
						]
						_lowerCAmelCase  : List[str]						  =					[
						    'generator',
						    'height',
						    'width',
						    'latents',
						    'guidance_scale',
						    'negative_prompt',
						    'num_inference_steps',
						    'return_dict',
						    'guidance_scale',
						    'num_images_per_prompt',
						    'output_type',
						    'return_dict',
						]
						_lowerCAmelCase  : Union[str, Any]						  =					False







						@property
						def        __lowercase     (	self :   str					):

													"""simple docstring"""
													return 32







						@property
						def        __lowercase     (	self :   str					):

													"""simple docstring"""
													return 32







						@property
						def        __lowercase     (	self :   Optional[int]					):

													"""simple docstring"""
													return self.time_input_dim







						@property
						def        __lowercase     (	self :   int					):

													"""simple docstring"""
													return self.time_input_dim * 4







						@property
						def        __lowercase     (	self :   int					):

													"""simple docstring"""
													return 1_00







						@property
						def        __lowercase     (	self :   Tuple					):

													"""simple docstring"""
													SCREAMING_SNAKE_CASE       :   Optional[int]         =      XLMRobertaTokenizerFast.from_pretrained('''YiYiXu/tiny-random-mclip-base'''					)
													return tokenizer







						@property
						def        __lowercase     (	self :   List[Any]					):

													"""simple docstring"""
													torch.manual_seed(0					)
													SCREAMING_SNAKE_CASE       :   List[Any]         =      MCLIPConfig(
													    numDims=self.cross_attention_dim					,     transformerDimensions=self.text_embedder_hidden_size					,     hidden_size=self.text_embedder_hidden_size					,     intermediate_size=37					,     num_attention_heads=4					,     num_hidden_layers=5					,     vocab_size=10_05					,     )

													SCREAMING_SNAKE_CASE       :   Any         =      MultilingualCLIP(lowerCAmelCase__					)
													SCREAMING_SNAKE_CASE       :   Union[str, Any]         =      text_encoder.eval()

													return text_encoder







						@property
						def        __lowercase     (	self :   List[Any]					):

													"""simple docstring"""
													torch.manual_seed(0					)

													SCREAMING_SNAKE_CASE       :   Tuple         =      {
													    '''in_channels''': 9,
													    # Out channels is double in channels because predicts mean and variance
													    '''out_channels''': 8,
													    '''addition_embed_type''': '''text_image''',
													    '''down_block_types''': ('''ResnetDownsampleBlock2D''', '''SimpleCrossAttnDownBlock2D'''),
													    '''up_block_types''': ('''SimpleCrossAttnUpBlock2D''', '''ResnetUpsampleBlock2D'''),
													    '''mid_block_type''': '''UNetMidBlock2DSimpleCrossAttn''',
													    '''block_out_channels''': (self.block_out_channels_a, self.block_out_channels_a * 2),
													    '''layers_per_block''': 1,
													    '''encoder_hid_dim''': self.text_embedder_hidden_size,
													    '''encoder_hid_dim_type''': '''text_image_proj''',
													    '''cross_attention_dim''': self.cross_attention_dim,
													    '''attention_head_dim''': 4,
													    '''resnet_time_scale_shift''': '''scale_shift''',
													    '''class_embed_type''': None,
													}

													SCREAMING_SNAKE_CASE       :   Union[str, Any]         =      UNetaDConditionModel(**lowerCAmelCase__					)
													return model







						@property
						def        __lowercase     (	self :   Optional[Any]					):

													"""simple docstring"""
													return {
													    "block_out_channels": [32, 64],
													    "down_block_types": ["DownEncoderBlock2D", "AttnDownEncoderBlock2D"],
													    "in_channels": 3,
													    "latent_channels": 4,
													    "layers_per_block": 1,
													    "norm_num_groups": 8,
													    "norm_type": "spatial",
													    "num_vq_embeddings": 12,
													    "out_channels": 3,
													    "up_block_types": [
													        "AttnUpDecoderBlock2D",
													        "UpDecoderBlock2D",
													    ],
													    "vq_embed_dim": 4,
													}







						@property
						def        __lowercase     (	self :   Union[str, Any]					):

													"""simple docstring"""
													torch.manual_seed(0					)
													SCREAMING_SNAKE_CASE       :   Optional[Any]         =      VQModel(**self.dummy_movq_kwargs					)
													return model







						def        __lowercase     (	self :   Optional[int]					):

													"""simple docstring"""
													SCREAMING_SNAKE_CASE       :   Any         =      self.dummy_text_encoder
													SCREAMING_SNAKE_CASE       :   Tuple         =      self.dummy_tokenizer
													SCREAMING_SNAKE_CASE       :   int         =      self.dummy_unet
													SCREAMING_SNAKE_CASE       :   str         =      self.dummy_movq

													SCREAMING_SNAKE_CASE       :   Optional[int]         =      DDIMScheduler(
													    num_train_timesteps=10_00					,     beta_schedule='''linear'''					,     beta_start=0.0_0085					,     beta_end=0.012					,     clip_sample=lowerCAmelCase__					,     set_alpha_to_one=lowerCAmelCase__					,     steps_offset=1					,     prediction_type='''epsilon'''					,     thresholding=lowerCAmelCase__					,     )

													SCREAMING_SNAKE_CASE       :   Optional[int]         =      {
													    '''text_encoder''': text_encoder,
													    '''tokenizer''': tokenizer,
													    '''unet''': unet,
													    '''scheduler''': scheduler,
													    '''movq''': movq,
													}

													return components







						def        __lowercase     (	self :   Dict					,     lowerCAmelCase__ :   Optional[Any]					,     lowerCAmelCase__ :   List[Any]=0					):

													"""simple docstring"""
													SCREAMING_SNAKE_CASE       :   Union[str, Any]         =      floats_tensor((1, self.cross_attention_dim)					,     rng=random.Random(lowerCAmelCase__					)					).to(lowerCAmelCase__					)
													SCREAMING_SNAKE_CASE       :   List[Any]         =      floats_tensor((1, self.cross_attention_dim)					,     rng=random.Random(seed + 1					)					).to(lowerCAmelCase__					)
													# create init_image
													SCREAMING_SNAKE_CASE       :   Optional[Any]         =      floats_tensor((1, 3, 64, 64)					,     rng=random.Random(lowerCAmelCase__					)					).to(lowerCAmelCase__					)
													SCREAMING_SNAKE_CASE       :   Tuple         =      image.cpu().permute(0					,     2					,     3					,     1					)[0]
													SCREAMING_SNAKE_CASE       :   Union[str, Any]         =      Image.fromarray(np.uinta(lowerCAmelCase__					)					).convert('''RGB'''					).resize((2_56, 2_56)					)
													# create mask
													SCREAMING_SNAKE_CASE       :   Union[str, Any]         =      np.ones((64, 64)					,     dtype=np.floataa					)
													SCREAMING_SNAKE_CASE       :   str         =      0

													if str(lowerCAmelCase__					).startswith('''mps'''					):
																				SCREAMING_SNAKE_CASE       :   Tuple         =      torch.manual_seed(lowerCAmelCase__					)
													else:
																				SCREAMING_SNAKE_CASE       :   Union[str, Any]         =      torch.Generator(device=lowerCAmelCase__					).manual_seed(lowerCAmelCase__					)
													SCREAMING_SNAKE_CASE       :   Tuple         =      {
													    '''prompt''': '''horse''',
													    '''image''': init_image,
													    '''mask_image''': mask,
													    '''image_embeds''': image_embeds,
													    '''negative_image_embeds''': negative_image_embeds,
													    '''generator''': generator,
													    '''height''': 64,
													    '''width''': 64,
													    '''num_inference_steps''': 2,
													    '''guidance_scale''': 4.0,
													    '''output_type''': '''np''',
													}
													return inputs







						def        __lowercase     (	self :   List[Any]					):

													"""simple docstring"""
													SCREAMING_SNAKE_CASE       :   Any         =      '''cpu'''

													SCREAMING_SNAKE_CASE       :   Dict         =      self.get_dummy_components()

													SCREAMING_SNAKE_CASE       :   Optional[int]         =      self.pipeline_class(**lowerCAmelCase__					)
													SCREAMING_SNAKE_CASE       :   int         =      pipe.to(lowerCAmelCase__					)

													pipe.set_progress_bar_config(disable=lowerCAmelCase__					)

													SCREAMING_SNAKE_CASE       :   str         =      pipe(**self.get_dummy_inputs(lowerCAmelCase__					)					)
													SCREAMING_SNAKE_CASE       :   int         =      output.images

													SCREAMING_SNAKE_CASE       :   int         =      pipe(
													    **self.get_dummy_inputs(lowerCAmelCase__					)					,     return_dict=lowerCAmelCase__					,     )[0]

													SCREAMING_SNAKE_CASE       :   Optional[Any]         =      image[0, -3:, -3:, -1]
													SCREAMING_SNAKE_CASE       :   Dict         =      image_from_tuple[0, -3:, -3:, -1]

													print(F"""image.shape {image.shape}"""					)

													assert image.shape == (1, 64, 64, 3)

													SCREAMING_SNAKE_CASE       :   List[str]         =      np.array(
													    [0.832_6919, 0.7379_0467, 0.2091_8581, 0.930_9612, 0.551_1791, 0.4371_3328, 0.551_3321, 0.4992_2934, 0.5949_7786]					)

													assert (
													    np.abs(image_slice.flatten() - expected_slice					).max() < 1e-2
													), F""" expected_slice {expected_slice}, but got {image_slice.flatten()}"""
													assert (
													    np.abs(image_from_tuple_slice.flatten() - expected_slice					).max() < 1e-2
													), F""" expected_slice {expected_slice}, but got {image_from_tuple_slice.flatten()}"""







						def        __lowercase     (	self :   str					):

													"""simple docstring"""
													super().test_inference_batch_single_identical(expected_max_diff=3e-3					)


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







						def        __lowercase     (	self :   List[str]					):

													"""simple docstring"""
													# clean up the VRAM after each test
													super().tearDown()
													gc.collect()
													torch.cuda.empty_cache()







						def        __lowercase     (	self :   Dict					):

													"""simple docstring"""
													SCREAMING_SNAKE_CASE       :   List[Any]         =      load_numpy(
													    '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'''
													    '''/kandinsky/kandinsky_inpaint_cat_with_hat_fp16.npy'''					)

													SCREAMING_SNAKE_CASE       :   Optional[Any]         =      load_image(
													    '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main''' '''/kandinsky/cat.png'''					)
													SCREAMING_SNAKE_CASE       :   Union[str, Any]         =      np.ones((7_68, 7_68)					,     dtype=np.floataa					)
													SCREAMING_SNAKE_CASE       :   List[str]         =      0

													SCREAMING_SNAKE_CASE       :   Dict         =      '''a hat'''

													SCREAMING_SNAKE_CASE       :   str         =      KandinskyPriorPipeline.from_pretrained(
													    '''kandinsky-community/kandinsky-2-1-prior'''					,     torch_dtype=torch.floataa					)
													pipe_prior.to(lowerCAmelCase__					)

													SCREAMING_SNAKE_CASE       :   Union[str, Any]         =      KandinskyInpaintPipeline.from_pretrained(
													    '''kandinsky-community/kandinsky-2-1-inpaint'''					,     torch_dtype=torch.floataa					)
													SCREAMING_SNAKE_CASE       :   Dict         =      pipeline.to(lowerCAmelCase__					)
													pipeline.set_progress_bar_config(disable=lowerCAmelCase__					)

													SCREAMING_SNAKE_CASE       :   Optional[int]         =      torch.Generator(device='''cpu'''					).manual_seed(0					)
													SCREAMING_SNAKE_CASE       ,SCREAMING_SNAKE_CASE       :   Any         =      pipe_prior(
													    lowerCAmelCase__					,     generator=lowerCAmelCase__					,     num_inference_steps=5					,     negative_prompt=''''''					,     ).to_tuple()

													SCREAMING_SNAKE_CASE       :   Tuple         =      pipeline(
													    lowerCAmelCase__					,     image=lowerCAmelCase__					,     mask_image=lowerCAmelCase__					,     image_embeds=lowerCAmelCase__					,     negative_image_embeds=lowerCAmelCase__					,     generator=lowerCAmelCase__					,     num_inference_steps=1_00					,     height=7_68					,     width=7_68					,     output_type='''np'''					,     )

													SCREAMING_SNAKE_CASE       :   Tuple         =      output.images[0]

													assert image.shape == (7_68, 7_68, 3)

													assert_mean_pixel_difference(lowerCAmelCase__					,     lowerCAmelCase__					)









'''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


lowerCAmelCase_  :							Any             =    'http://www.mocksite.com/file1.txt'
lowerCAmelCase_  :							List[str]             =    '"text": ["foo", "foo"]'
lowerCAmelCase_  :							Any             =    '6d8ce9aa78a471c7477201efbeabd3bb01ac2e7d100a6dc024ba1608361f90a8'
class 						lowerCamelCase_				:
						_lowerCAmelCase  : Dict						  =					2_0_0
						_lowerCAmelCase  : int						  =					{'Content-Length': '100'}
						_lowerCAmelCase  : Optional[Any]						  =					{}







						def        __lowercase     (	self :   str					,     **lowerCAmelCase__ :   Any					):

													"""simple docstring"""
													return [bytes(lowerCAmelCase__					,     '''utf-8'''					)]


def        UpperCAmelCase (    *A	:					int       ,		**A	:					Optional[Any]					):
							return MockResponse()
@pytest.mark.parametrize('''urls_type'''       ,		[str, list, dict]					)
def        UpperCAmelCase (    A	:					int       ,		A	:					Union[str, Any]       ,		A	:					Union[str, Any]					):
							import requests

							monkeypatch.setattr(A       ,		'''request'''       ,		A					)

							SCREAMING_SNAKE_CASE       :   Optional[Any]         =      URL
							if issubclass(A       ,		A					):
														SCREAMING_SNAKE_CASE       :   List[Any]         =      url
							elif issubclass(A       ,		A					):
														SCREAMING_SNAKE_CASE       :   List[Any]         =      [url]
							elif issubclass(A       ,		A					):
														SCREAMING_SNAKE_CASE       :   str         =      {'''train''': url}
							SCREAMING_SNAKE_CASE       :   Any         =      '''dummy'''
							SCREAMING_SNAKE_CASE       :   Optional[Any]         =      '''downloads'''
							SCREAMING_SNAKE_CASE       :   Optional[Any]         =      tmp_path
							SCREAMING_SNAKE_CASE       :   List[str]         =      DownloadConfig(
							    cache_dir=os.path.join(A       ,		A					)       ,		use_etag=A       ,		)
							SCREAMING_SNAKE_CASE       :   Union[str, Any]         =      DownloadManager(dataset_name=A       ,		download_config=A					)
							SCREAMING_SNAKE_CASE       :   List[str]         =      dl_manager.download(A					)
							SCREAMING_SNAKE_CASE       :   Dict         =      urls
							for downloaded_paths in [downloaded_paths]:
														if isinstance(A       ,		A					):
																					SCREAMING_SNAKE_CASE       :   List[str]         =      [downloaded_paths]
																					SCREAMING_SNAKE_CASE       :   int         =      [urls]
														elif isinstance(A       ,		A					):
																					assert "train" in downloaded_paths.keys()
																					SCREAMING_SNAKE_CASE       :   Optional[int]         =      downloaded_paths.values()
																					SCREAMING_SNAKE_CASE       :   Union[str, Any]         =      urls.values()
														assert downloaded_paths
														for downloaded_path, input_url in zip(A       ,		A					):
																					assert downloaded_path == dl_manager.downloaded_paths[input_url]
																					SCREAMING_SNAKE_CASE       :   Dict         =      Path(A					)
																					SCREAMING_SNAKE_CASE       :   Optional[Any]         =      downloaded_path.parts
																					assert parts[-1] == HASH
																					assert parts[-2] == cache_subdir
																					assert downloaded_path.exists()
																					SCREAMING_SNAKE_CASE       :   str         =      downloaded_path.read_text()
																					assert content == CONTENT
																					SCREAMING_SNAKE_CASE       :   str         =      downloaded_path.with_suffix('''.json'''					)
																					assert metadata_downloaded_path.exists()
																					SCREAMING_SNAKE_CASE       :   Optional[int]         =      json.loads(metadata_downloaded_path.read_text()					)
																					assert metadata_content == {"url": URL, "etag": None}
@pytest.mark.parametrize('''paths_type'''       ,		[str, list, dict]					)
def        UpperCAmelCase (    A	:					Tuple       ,		A	:					Union[str, Any]       ,		A	:					Dict					):
							SCREAMING_SNAKE_CASE       :   List[str]         =      str(A					)
							if issubclass(A       ,		A					):
														SCREAMING_SNAKE_CASE       :   Optional[int]         =      filename
							elif issubclass(A       ,		A					):
														SCREAMING_SNAKE_CASE       :   Optional[int]         =      [filename]
							elif issubclass(A       ,		A					):
														SCREAMING_SNAKE_CASE       :   List[Any]         =      {'''train''': filename}
							SCREAMING_SNAKE_CASE       :   Dict         =      '''dummy'''
							SCREAMING_SNAKE_CASE       :   Optional[Any]         =      xz_file.parent
							SCREAMING_SNAKE_CASE       :   List[Any]         =      '''extracted'''
							SCREAMING_SNAKE_CASE       :   Optional[Any]         =      DownloadConfig(
							    cache_dir=A       ,		use_etag=A       ,		)
							SCREAMING_SNAKE_CASE       :   Optional[int]         =      DownloadManager(dataset_name=A       ,		download_config=A					)
							SCREAMING_SNAKE_CASE       :   Optional[int]         =      dl_manager.extract(A					)
							SCREAMING_SNAKE_CASE       :   Tuple         =      paths
							for extracted_paths in [extracted_paths]:
														if isinstance(A       ,		A					):
																					SCREAMING_SNAKE_CASE       :   List[Any]         =      [extracted_paths]
																					SCREAMING_SNAKE_CASE       :   str         =      [paths]
														elif isinstance(A       ,		A					):
																					assert "train" in extracted_paths.keys()
																					SCREAMING_SNAKE_CASE       :   Optional[int]         =      extracted_paths.values()
																					SCREAMING_SNAKE_CASE       :   Union[str, Any]         =      paths.values()
														assert extracted_paths
														for extracted_path, input_path in zip(A       ,		A					):
																					assert extracted_path == dl_manager.extracted_paths[input_path]
																					SCREAMING_SNAKE_CASE       :   Any         =      Path(A					)
																					SCREAMING_SNAKE_CASE       :   Optional[Any]         =      extracted_path.parts
																					assert parts[-1] == hash_url_to_filename(A       ,		etag=A					)
																					assert parts[-2] == extracted_subdir
																					assert extracted_path.exists()
																					SCREAMING_SNAKE_CASE       :   int         =      extracted_path.read_text()
																					SCREAMING_SNAKE_CASE       :   Optional[int]         =      text_file.read_text()
																					assert extracted_file_content == expected_file_content
def        UpperCAmelCase (    A	:					Optional[int]       ,		A	:					int					):
							assert path.endswith('''.jsonl'''					)
							for num_items, line in enumerate(A       ,		start=1					):
														SCREAMING_SNAKE_CASE       :   List[Any]         =      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        UpperCAmelCase (    A	:					Any       ,		A	:					str					):
							SCREAMING_SNAKE_CASE       :   Union[str, Any]         =      request.getfixturevalue(A					)
							SCREAMING_SNAKE_CASE       :   List[Any]         =      DownloadManager()
							for num_jsonl, (path, file) in enumerate(dl_manager.iter_archive(A					)       ,		start=1					):
														_test_jsonl(A       ,		A					)
							assert num_jsonl == 2
@pytest.mark.parametrize('''archive_nested_jsonl'''       ,		['''tar_nested_jsonl_path''', '''zip_nested_jsonl_path''']					)
def        UpperCAmelCase (    A	:					Dict       ,		A	:					Optional[Any]					):
							SCREAMING_SNAKE_CASE       :   Union[str, Any]         =      request.getfixturevalue(A					)
							SCREAMING_SNAKE_CASE       :   Optional[int]         =      DownloadManager()
							for num_tar, (path, file) in enumerate(dl_manager.iter_archive(A					)       ,		start=1					):
														for num_jsonl, (subpath, subfile) in enumerate(dl_manager.iter_archive(A					)       ,		start=1					):
																					_test_jsonl(A       ,		A					)
							assert num_tar == 1
							assert num_jsonl == 2






def        UpperCAmelCase (    A	:					List[Any]					):
							SCREAMING_SNAKE_CASE       :   Optional[Any]         =      DownloadManager()
							for num_file, file in enumerate(dl_manager.iter_files(A					)       ,		start=1					):
														assert os.path.basename(A					) == ("test.txt" if num_file == 1 else "train.txt")
							assert num_file == 2




