






'''simple docstring'''







import gc
import random
import unittest

import numpy as np
import torch
from PIL import Image

from diffusers import (
    DDIMScheduler,
    KandinskyVaaControlnetImgaImgPipeline,
    KandinskyVaaPriorEmbaEmbPipeline,
    UNetaDConditionModel,
    VQModel,
)
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 _lowercase			(       _lowercase       ,   unittest.TestCase	):
							a          = KandinskyVaaControlnetImgaImgPipeline
							a          = ["""image_embeds""", """negative_image_embeds""", """image""", """hint"""]
							a          = ["""image_embeds""", """negative_image_embeds""", """image""", """hint"""]
							a          = [
							    """generator""",
							    """height""",
							    """width""",
							    """strength""",
							    """guidance_scale""",
							    """num_inference_steps""",
							    """return_dict""",
							    """guidance_scale""",
							    """num_images_per_prompt""",
							    """output_type""",
							    """return_dict""",
							]
							a          = False




							@property
							def 	lowerCamelCase_       (							self:       Union[str, Any]  ):
														return 32




							@property
							def 	lowerCamelCase_       (							self:       int  ):
														return 32




							@property
							def 	lowerCamelCase_       (							self:       Dict  ):
														return self.time_input_dim




							@property
							def 	lowerCamelCase_       (							self:       int  ):
														return self.time_input_dim * 4




							@property
							def 	lowerCamelCase_       (							self:       Optional[Any]  ):
														return 100




							@property
							def 	lowerCamelCase_       (							self:       str  ):
														torch.manual_seed(0  )

														lowerCamelCase__      :			Optional[Any] 						= {
														    """in_channels""": 8,
														    # Out channels is double in channels because predicts mean and variance
														    """out_channels""": 8,
														    """addition_embed_type""": """image_hint""",
														    """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""": """image_proj""",
														    """cross_attention_dim""": self.cross_attention_dim,
														    """attention_head_dim""": 4,
														    """resnet_time_scale_shift""": """scale_shift""",
														    """class_embed_type""": None,
														}

														lowerCamelCase__      :			Optional[Any] 						= UNetaDConditionModel(**UpperCamelCase__  )
														return model




							@property
							def 	lowerCamelCase_       (							self:       Union[str, Any]  ):
														return {
														    "block_out_channels": [32, 32, 64, 64],
														    "down_block_types": [
														        "DownEncoderBlock2D",
														        "DownEncoderBlock2D",
														        "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", "UpDecoderBlock2D", "UpDecoderBlock2D"],
														    "vq_embed_dim": 4,
														}




							@property
							def 	lowerCamelCase_       (							self:       Any  ):
														torch.manual_seed(0  )
														lowerCamelCase__      :			Optional[int] 						= VQModel(**self.dummy_movq_kwargs  )
														return model




							def 	lowerCamelCase_       (							self:       List[str]  ):
														lowerCamelCase__      :			Tuple 						= self.dummy_unet
														lowerCamelCase__      :			Any 						= self.dummy_movq

														lowerCamelCase__      :			str 						= {
														    """num_train_timesteps""": 1_000,
														    """beta_schedule""": """linear""",
														    """beta_start""": 0.00_085,
														    """beta_end""": 0.012,
														    """clip_sample""": False,
														    """set_alpha_to_one""": False,
														    """steps_offset""": 0,
														    """prediction_type""": """epsilon""",
														    """thresholding""": False,
														}

														lowerCamelCase__      :			int 						= DDIMScheduler(**UpperCamelCase__  )

														lowerCamelCase__      :			List[str] 						= {
														    """unet""": unet,
														    """scheduler""": scheduler,
														    """movq""": movq,
														}

														return components




							def 	lowerCamelCase_       (							self:       int    ,							UpperCamelCase__:       List[Any]    ,							UpperCamelCase__:       List[str]=0  ):
														lowerCamelCase__      :			Tuple 						= floats_tensor((1, self.text_embedder_hidden_size)    ,							rng=random.Random(UpperCamelCase__  )  ).to(UpperCamelCase__  )
														lowerCamelCase__      :			Optional[Any] 						= floats_tensor((1, self.text_embedder_hidden_size)    ,							rng=random.Random(seed + 1  )  ).to(
														    UpperCamelCase__  )
														# create init_image
														lowerCamelCase__      :			int 						= floats_tensor((1, 3, 64, 64)    ,							rng=random.Random(UpperCamelCase__  )  ).to(UpperCamelCase__  )
														lowerCamelCase__      :			int 						= image.cpu().permute(0    ,							2    ,							3    ,							1  )[0]
														lowerCamelCase__      :			Tuple 						= Image.fromarray(np.uinta(UpperCamelCase__  )  ).convert("""RGB"""  ).resize((256, 256)  )
														# create hint
														lowerCamelCase__      :			Any 						= floats_tensor((1, 3, 64, 64)    ,							rng=random.Random(UpperCamelCase__  )  ).to(UpperCamelCase__  )

														if str(UpperCamelCase__  ).startswith("""mps"""  ):
																					lowerCamelCase__      :			List[str] 						= torch.manual_seed(UpperCamelCase__  )
														else:
																					lowerCamelCase__      :			Optional[int] 						= torch.Generator(device=UpperCamelCase__  ).manual_seed(UpperCamelCase__  )
														lowerCamelCase__      :			Optional[int] 						= {
														    """image""": init_image,
														    """image_embeds""": image_embeds,
														    """negative_image_embeds""": negative_image_embeds,
														    """hint""": hint,
														    """generator""": generator,
														    """height""": 64,
														    """width""": 64,
														    """num_inference_steps""": 10,
														    """guidance_scale""": 7.0,
														    """strength""": 0.2,
														    """output_type""": """np""",
														}
														return inputs




							def 	lowerCamelCase_       (							self:       Optional[int]  ):
														lowerCamelCase__      :			Optional[Any] 						= """cpu"""

														lowerCamelCase__      :			Any 						= self.get_dummy_components()

														lowerCamelCase__      :			Optional[int] 						= self.pipeline_class(**UpperCamelCase__  )
														lowerCamelCase__      :			int 						= pipe.to(UpperCamelCase__  )

														pipe.set_progress_bar_config(disable=UpperCamelCase__  )

														lowerCamelCase__      :			List[Any] 						= pipe(**self.get_dummy_inputs(UpperCamelCase__  )  )
														lowerCamelCase__      :			Dict 						= output.images

														lowerCamelCase__      :			List[Any] 						= pipe(
														    **self.get_dummy_inputs(UpperCamelCase__  )    ,							return_dict=UpperCamelCase__    ,							)[0]

														lowerCamelCase__      :			int 						= image[0, -3:, -3:, -1]
														lowerCamelCase__      :			List[str] 						= image_from_tuple[0, -3:, -3:, -1]

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

														lowerCamelCase__      :			int 						= np.array(
														    [0.54_985_034, 0.55_509_365, 0.52_561_504, 0.5_570_494, 0.5_593_818, 0.5_263_979, 0.50_285_643, 0.5_069_846, 0.51_196_736]  )
														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()}'''



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




							def 	lowerCamelCase_       (							self:       str  ):
														# clean up the VRAM after each test
														super().tearDown()
														gc.collect()
														torch.cuda.empty_cache()




							def 	lowerCamelCase_       (							self:       str  ):
														lowerCamelCase__      :			str 						= load_numpy(
														    """https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"""
														    """/kandinskyv22/kandinskyv22_controlnet_img2img_robotcat_fp16.npy"""  )

														lowerCamelCase__      :			List[str] 						= load_image(
														    """https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main""" """/kandinsky/cat.png"""  )
														lowerCamelCase__      :			Optional[int] 						= init_image.resize((512, 512)  )

														lowerCamelCase__      :			int 						= load_image(
														    """https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"""
														    """/kandinskyv22/hint_image_cat.png"""  )
														lowerCamelCase__      :			Tuple 						= torch.from_numpy(np.array(UpperCamelCase__  )  ).float() / 255.0
														lowerCamelCase__      :			Optional[Any] 						= hint.permute(2    ,							0    ,							1  ).unsqueeze(0  )

														lowerCamelCase__      :			Dict 						= """A robot, 4k photo"""

														lowerCamelCase__      :			List[str] 						= KandinskyVaaPriorEmbaEmbPipeline.from_pretrained(
														    """kandinsky-community/kandinsky-2-2-prior"""    ,							torch_dtype=torch.floataa  )
														pipe_prior.to(UpperCamelCase__  )

														lowerCamelCase__      :			Union[str, Any] 						= KandinskyVaaControlnetImgaImgPipeline.from_pretrained(
														    """kandinsky-community/kandinsky-2-2-controlnet-depth"""    ,							torch_dtype=torch.floataa  )
														lowerCamelCase__      :			Union[str, Any] 						= pipeline.to(UpperCamelCase__  )

														pipeline.set_progress_bar_config(disable=UpperCamelCase__  )

														lowerCamelCase__      :			Any 						= torch.Generator(device="""cpu"""  ).manual_seed(0  )

														lowerCamelCase__						,	lowerCamelCase__      :			Optional[int] 						= pipe_prior(
														    UpperCamelCase__    ,							image=UpperCamelCase__    ,							strength=0.85    ,							generator=UpperCamelCase__    ,							negative_prompt=""""""    ,							).to_tuple()

														lowerCamelCase__      :			int 						= pipeline(
														    image=UpperCamelCase__    ,							image_embeds=UpperCamelCase__    ,							negative_image_embeds=UpperCamelCase__    ,							hint=UpperCamelCase__    ,							generator=UpperCamelCase__    ,							num_inference_steps=100    ,							height=512    ,							width=512    ,							strength=0.5    ,							output_type="""np"""    ,							)

														lowerCamelCase__      :			Union[str, Any] 						= output.images[0]

														assert image.shape == (512, 512, 3)

														assert_mean_pixel_difference(UpperCamelCase__    ,							UpperCamelCase__  )









'''simple docstring'''







import os
from typing import Optional

import fsspec
from fsspec.archive import AbstractArchiveFileSystem
from fsspec.utils import DEFAULT_BLOCK_SIZE

class _lowercase			(       _lowercase	):
							a          = """"""
							a          = (
							    None  # protocol passed in prefix to the url. ex: "gzip", for gzip://file.txt::http://foo.bar/file.txt.gz
							)
							a          = None  # compression type in fsspec. ex: "gzip"
							a          = None  # extension of the filename to strip. ex: "".gz" to get file.txt from file.txt.gz




							def __init__(							self:       str    ,							UpperCamelCase__:       str = ""    ,							UpperCamelCase__:       Optional[str] = None    ,							UpperCamelCase__:       Optional[dict] = None    ,							**UpperCamelCase__:       List[Any]  ):
														super().__init__(self    ,							**UpperCamelCase__  )
														# always open as "rb" since fsspec can then use the TextIOWrapper to make it work for "r" mode
														lowerCamelCase__      :			List[Any] 						= fsspec.open(
														    UpperCamelCase__    ,							mode="""rb"""    ,							protocol=UpperCamelCase__    ,							compression=self.compression    ,							client_kwargs={
														        """requote_redirect_url""": False,  # see https://github.com/huggingface/datasets/pull/5459
														        """trust_env""": True,  # Enable reading proxy env variables.
														        **(target_options or {}).pop("""client_kwargs"""    ,							{}  ),  # To avoid issues if it was already passed.
														    }    ,							**(target_options or {})    ,							)
														lowerCamelCase__      :			str 						= os.path.basename(self.file.path.split("""::"""  )[0]  )
														lowerCamelCase__      :			Union[str, Any] 						= (
														    self.compressed_name[: self.compressed_name.rindex("""."""  )]
														    if """.""" in self.compressed_name
														    else self.compressed_name
														)
														lowerCamelCase__      :			Tuple 						= None




							@classmethod
							def 	lowerCamelCase_       (							cls:       Optional[int]    ,							UpperCamelCase__:       Optional[int]  ):
														# compressed file paths are always relative to the archive root
														return super()._strip_protocol(UpperCamelCase__  ).lstrip("""/"""  )




							def 	lowerCamelCase_       (							self:       Tuple  ):
														if self.dir_cache is None:
																					lowerCamelCase__      :			Dict 						= {**self.file.fs.info(self.file.path  ), """name""": self.uncompressed_name}
																					lowerCamelCase__      :			int 						= {f["""name"""]: f}




							def 	lowerCamelCase_       (							self:       List[str]    ,							UpperCamelCase__:       str  ):
														return self.file.open().read()




							def 	lowerCamelCase_       (							self:       Optional[int]    ,							UpperCamelCase__:       str    ,							UpperCamelCase__:       str = "rb"    ,							UpperCamelCase__:       Optional[int]=None    ,							UpperCamelCase__:       Tuple=True    ,							UpperCamelCase__:       Tuple=None    ,							**UpperCamelCase__:       Optional[Any]    ,							):
														lowerCamelCase__      :			Union[str, Any] 						= self._strip_protocol(UpperCamelCase__  )
														if mode != "rb":
																					raise ValueError(F'''Tried to read with mode {mode} on file {self.file.path} opened with mode \'rb\''''  )
														return self.file.open()



class _lowercase			(       _lowercase	):
							a          = """bz2"""
							a          = """bz2"""
							a          = """.bz2"""



class _lowercase			(       _lowercase	):
							a          = """gzip"""
							a          = """gzip"""
							a          = """.gz"""



class _lowercase			(       _lowercase	):
							a          = """lz4"""
							a          = """lz4"""
							a          = """.lz4"""



class _lowercase			(       _lowercase	):
							a          = """xz"""
							a          = """xz"""
							a          = """.xz"""



class _lowercase			(       _lowercase	):
							a          = """zstd"""
							a          = """zstd"""
							a          = """.zst"""




							def __init__(							self:       int    ,							UpperCamelCase__:       str    ,							UpperCamelCase__:       str = "rb"    ,							UpperCamelCase__:       Optional[str] = None    ,							UpperCamelCase__:       Optional[dict] = None    ,							UpperCamelCase__:       int = DEFAULT_BLOCK_SIZE    ,							**UpperCamelCase__:       Dict    ,							):
														super().__init__(
														    fo=UpperCamelCase__    ,							mode=UpperCamelCase__    ,							target_protocol=UpperCamelCase__    ,							target_options=UpperCamelCase__    ,							block_size=UpperCamelCase__    ,							**UpperCamelCase__    ,							)
														# We need to wrap the zstd decompressor to avoid this error in fsspec==2021.7.0 and zstandard==0.15.2:
														#
														# File "/Users/user/.virtualenvs/hf-datasets/lib/python3.7/site-packages/fsspec/core.py", line 145, in open
														#     out.close = close
														# AttributeError: 'zstd.ZstdDecompressionReader' object attribute 'close' is read-only
														#
														# see https://github.com/intake/filesystem_spec/issues/725
														lowerCamelCase__      :			Tuple 						= self.file.__enter__

														class _lowercase			:




																					def __init__(							self:       Optional[int]    ,							UpperCamelCase__:       Any  ):
																												lowerCamelCase__      :			Optional[int] 						= file_




																					def __enter__(							self:       List[Any]  ):
																												self._file.__enter__()
																												return self




																					def __exit__(							self:       Any    ,							*UpperCamelCase__:       str    ,							**UpperCamelCase__:       Any  ):
																												self._file.__exit__(*UpperCamelCase__    ,							**UpperCamelCase__  )




																					def __iter__(							self:       Any  ):
																												return iter(self._file  )




																					def 	lowerCamelCase_       (							self:       List[Any]  ):
																												return next(self._file  )




																					def __getattr__(							self:       List[str]    ,							UpperCamelCase__:       Dict  ):
																												return getattr(self._file    ,							UpperCamelCase__  )

														def fixed_enter(*UpperCamelCase__:       Union[str, Any]    ,							**UpperCamelCase__:       List[str]  ):
																					return WrappedFile(_enter(*UpperCamelCase__    ,							**UpperCamelCase__  )  )

														lowerCamelCase__      :			Optional[Any] 						= fixed_enter

