



import gc
import random
import unittest

import numpy as np
import torch

from diffusers import (
    DDIMScheduler,
    KandinskyVaaControlnetPipeline,
    KandinskyVaaPriorPipeline,
    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 	lowerCamelCase__	(       _A							,  unittest.TestCase):



 """simple docstring"""

 a__					: Optional[int] =  KandinskyVaaControlnetPipeline
 a__					: Any =  ["image_embeds", "negative_image_embeds", "hint"]
 a__					: str =  ["image_embeds", "negative_image_embeds", "hint"]
 a__					: List[Any] =  [
     "generator",
     "height",
     "width",
     "latents",
     "guidance_scale",
     "num_inference_steps",
     "return_dict",
     "guidance_scale",
     "num_images_per_prompt",
     "output_type",
     "return_dict",
 ]
 a__					: Union[str, Any] =  False

 @property
 def     snake_case_      (    self					:					str      )   ->	List[Any]:
    return 32

 @property
 def     snake_case_      (    self					:					Optional[int]      )   ->	Optional[int]:
    return 32

 @property
 def     snake_case_      (    self					:					Union[str, Any]      )   ->	Tuple:
    return self.time_input_dim

 @property
 def     snake_case_      (    self					:					str      )   ->	Any:
    return self.time_input_dim * 4

 @property
 def     snake_case_      (    self					:					Optional[int]      )   ->	Tuple:
    return 1_00

 @property
 def     snake_case_      (    self					:					Tuple      )   ->	Union[str, Any]:
    torch.manual_seed(0      )

    _A      =       {
        '''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,
    }

    _A      =       UNetaDConditionModel(**__lowerCAmelCase      )
    return model

 @property
 def     snake_case_      (    self					:					Tuple      )   ->	List[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     snake_case_      (    self					:					int      )   ->	List[str]:
    torch.manual_seed(0      )
    _A      =       VQModel(**self.dummy_movq_kwargs      )
    return model

 def     snake_case_      (    self					:					List[str]      )   ->	Dict:
    _A      =       self.dummy_unet
    _A      =       self.dummy_movq

    _A      =       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			,						)

    _A      =       {
        '''unet''': unet,
        '''scheduler''': scheduler,
        '''movq''': movq,
    }
    return components

 def     snake_case_      (    self					:					Optional[Any]			,						__lowerCAmelCase					:					Union[str, Any]			,						__lowerCAmelCase					:					Tuple=0      )   ->	Optional[Any]:
    _A      =       floats_tensor((1, self.text_embedder_hidden_size)			,						rng=random.Random(__lowerCAmelCase      )      ).to(__lowerCAmelCase      )
    _A      =       floats_tensor((1, self.text_embedder_hidden_size)			,						rng=random.Random(seed + 1      )      ).to(
        __lowerCAmelCase      )

    # create hint
    _A      =       floats_tensor((1, 3, 64, 64)			,						rng=random.Random(__lowerCAmelCase      )      ).to(__lowerCAmelCase      )

    if str(__lowerCAmelCase      ).startswith('''mps'''      ):
       _A      =       torch.manual_seed(__lowerCAmelCase      )
    else:
       _A      =       torch.Generator(device=__lowerCAmelCase      ).manual_seed(__lowerCAmelCase      )
    _A      =       {
        '''image_embeds''': image_embeds,
        '''negative_image_embeds''': negative_image_embeds,
        '''hint''': hint,
        '''generator''': generator,
        '''height''': 64,
        '''width''': 64,
        '''guidance_scale''': 4.0,
        '''num_inference_steps''': 2,
        '''output_type''': '''np''',
    }
    return inputs

 def     snake_case_      (    self					:					Any      )   ->	Union[str, Any]:
    _A      =       '''cpu'''

    _A      =       self.get_dummy_components()

    _A      =       self.pipeline_class(**__lowerCAmelCase      )
    _A      =       pipe.to(__lowerCAmelCase      )

    pipe.set_progress_bar_config(disable=__lowerCAmelCase      )

    _A      =       pipe(**self.get_dummy_inputs(__lowerCAmelCase      )      )
    _A      =       output.images

    _A      =       pipe(
        **self.get_dummy_inputs(__lowerCAmelCase      )			,						return_dict=__lowerCAmelCase			,						)[0]

    _A      =       image[0, -3:, -3:, -1]
    _A      =       image_from_tuple[0, -3:, -3:, -1]

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

    _A      =       np.array(
        [0.695_9826, 0.86_8279, 0.755_8092, 0.6876_9467, 0.8580_5804, 0.6597_7496, 0.4488_5302, 0.595_9111, 0.425_1595]      )

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



 """simple docstring"""

 def     snake_case_      (    self					:					Any      )   ->	List[Any]:
    # clean up the VRAM after each test
    super().tearDown()
    gc.collect()
    torch.cuda.empty_cache()

 def     snake_case_      (    self					:					Tuple      )   ->	List[Any]:
    _A      =       load_numpy(
        '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'''
        '''/kandinskyv22/kandinskyv22_controlnet_robotcat_fp16.npy'''      )

    _A      =       load_image(
        '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'''
        '''/kandinskyv22/hint_image_cat.png'''      )
    _A      =       torch.from_numpy(np.array(__lowerCAmelCase      )      ).float() / 255.0
    _A      =       hint.permute(2			,						0			,						1      ).unsqueeze(0      )

    _A      =       KandinskyVaaPriorPipeline.from_pretrained(
        '''kandinsky-community/kandinsky-2-2-prior'''			,						torch_dtype=torch.floataa      )
    pipe_prior.to(__lowerCAmelCase      )

    _A      =       KandinskyVaaControlnetPipeline.from_pretrained(
        '''kandinsky-community/kandinsky-2-2-controlnet-depth'''			,						torch_dtype=torch.floataa      )
    _A      =       pipeline.to(__lowerCAmelCase      )
    pipeline.set_progress_bar_config(disable=__lowerCAmelCase      )

    _A      =       '''A robot, 4k photo'''

    _A      =       torch.Generator(device='''cuda'''      ).manual_seed(0      )
    _A  ,    _A      =       pipe_prior(
        __lowerCAmelCase			,						generator=__lowerCAmelCase			,						num_inference_steps=5			,						negative_prompt=''''''			,						).to_tuple()

    _A      =       torch.Generator(device='''cuda'''      ).manual_seed(0      )
    _A      =       pipeline(
        image_embeds=__lowerCAmelCase			,						negative_image_embeds=__lowerCAmelCase			,						hint=__lowerCAmelCase			,						generator=__lowerCAmelCase			,						num_inference_steps=1_00			,						output_type='''np'''			,						)

    _A      =       output.images[0]

    assert image.shape == (5_12, 5_12, 3)

    assert_mean_pixel_difference(__lowerCAmelCase			,						__lowerCAmelCase      )





import os
from dataclasses import dataclass, field
from io import BytesIO
from typing import TYPE_CHECKING, Any, ClassVar, Dict, Optional, Union

import numpy as np
import pyarrow as pa

from .. import config
from ..download.streaming_download_manager import xopen, xsplitext
from ..table import array_cast
from ..utils.py_utils import no_op_if_value_is_null, string_to_dict


if TYPE_CHECKING:
							from .features import FeatureType


lowerCamelCase_ ,							lowerCamelCase_ ,							lowerCamelCase_    :			Dict    =	False, False, False

@dataclass
class 							_UpperCamelCase					:






			'''simple docstring'''

			__UpperCamelCase   :    Optional[int]          =			None
			__UpperCamelCase   :    bool          =			True
			__UpperCamelCase   :    bool          =			True
			__UpperCamelCase   :    Optional[str]          =			None
			# Automatically constructed
			__UpperCamelCase   :    ClassVar[str]          =			"dict"
			__UpperCamelCase   :    ClassVar[Any]          =			pa.struct({"""bytes""": pa.binary(), """path""": pa.string()}	)
			__UpperCamelCase   :    str          =			field(default="""Audio"""  ,							init=_A  ,							repr=_A	)

			def __call__(							self			:							List[Any]   ):
							return self.pa_type

			def    lowerCAmelCase__						(							self			:							Optional[int]    ,						snake_case_			:							Union[str, bytes, dict]   ):
							try:
											import soundfile as sf  # soundfile is a dependency of librosa, needed to decode audio files.
							except ImportError as err:
											raise ImportError("""To support encoding audio data, please install 'soundfile'."""   ) from err
							if isinstance(snake_case_    ,						snake_case_   ):
											return {"bytes": None, "path": value}
							elif isinstance(snake_case_    ,						snake_case_   ):
											return {"bytes": value, "path": None}
							elif "array" in value:
											# convert the audio array to wav bytes
											UpperCamelCase_:				List[Any]						=    BytesIO()
											sf.write(snake_case_    ,						value["""array"""]    ,						value["""sampling_rate"""]    ,						format="""wav"""   )
											return {"bytes": buffer.getvalue(), "path": None}
							elif value.get("""path"""   ) is not None and os.path.isfile(value["""path"""]   ):
											# we set "bytes": None to not duplicate the data if they're already available locally
											if value["path"].endswith("""pcm"""   ):
															# "PCM" only has raw audio bytes
															if value.get("""sampling_rate"""   ) is None:
																			# At least, If you want to convert "PCM-byte" to "WAV-byte", you have to know sampling rate
																			raise KeyError("""To use PCM files, please specify a 'sampling_rate' in Audio object"""   )
															if value.get("""bytes"""   ):
																			# If we already had PCM-byte, we don`t have to make "read file, make bytes" (just use it!)
																			UpperCamelCase_:				str						=    np.frombuffer(value["""bytes"""]    ,						dtype=np.intaa   ).astype(np.floataa   ) / 3_2767
															else:
																			UpperCamelCase_:				Dict						=    np.memmap(value["""path"""]    ,						dtype="""h"""    ,						mode="""r"""   ).astype(np.floataa   ) / 3_2767

															UpperCamelCase_:				Dict						=    BytesIO(bytes()   )
															sf.write(snake_case_    ,						snake_case_    ,						value["""sampling_rate"""]    ,						format="""wav"""   )
															return {"bytes": buffer.getvalue(), "path": None}
											else:
															return {"bytes": None, "path": value.get("""path"""   )}
							elif value.get("""bytes"""   ) is not None or value.get("""path"""   ) is not None:
											# store the audio bytes, and path is used to infer the audio format using the file extension
											return {"bytes": value.get("""bytes"""   ), "path": value.get("""path"""   )}
							else:
											raise ValueError(
											    f'''An audio sample should have one of \'path\' or \'bytes\' but they are missing or None in {value}.'''   )

			def    lowerCAmelCase__						(							self			:							List[str]    ,						snake_case_			:							dict    ,						snake_case_			:							Optional[Dict[str, Union[str, bool, None]]] = None   ):
							if not self.decode:
											raise RuntimeError("""Decoding is disabled for this feature. Please use Audio(decode=True) instead."""   )

							UpperCamelCase_, UpperCamelCase_:				Any						=    (value["""path"""], BytesIO(value["""bytes"""]   )) if value["""bytes"""] is not None else (value["""path"""], None)
							if path is None and file is None:
											raise ValueError(f'''An audio sample should have one of \'path\' or \'bytes\' but both are None in {value}.'''   )

							try:
											import librosa
											import soundfile as sf
							except ImportError as err:
											raise ImportError("""To support decoding audio files, please install 'librosa' and 'soundfile'."""   ) from err

							UpperCamelCase_:				Optional[Any]						=    xsplitext(snake_case_   )[1][1:].lower() if path is not None else None
							if not config.IS_OPUS_SUPPORTED and audio_format == "opus":
											raise RuntimeError(
											    """Decoding 'opus' files requires system library 'libsndfile'>=1.0.31, """
											    """You can try to update `soundfile` python library: `pip install \"soundfile>=0.12.1\"`. """   )
							elif not config.IS_MP3_SUPPORTED and audio_format == "mp3":
											raise RuntimeError(
											    """Decoding 'mp3' files requires system library 'libsndfile'>=1.1.0, """
											    """You can try to update `soundfile` python library: `pip install \"soundfile>=0.12.1\"`. """   )

							if file is None:
											UpperCamelCase_:				Dict						=    token_per_repo_id or {}
											UpperCamelCase_:				str						=    path.split("""::"""   )[-1]
											try:
															UpperCamelCase_:				Any						=    string_to_dict(snake_case_    ,						config.HUB_DATASETS_URL   )["""repo_id"""]
															UpperCamelCase_:				List[str]						=    token_per_repo_id[repo_id]
											except (ValueError, KeyError):
															UpperCamelCase_:				Any						=    None

											with xopen(snake_case_    ,						"""rb"""    ,						use_auth_token=snake_case_   ) as f:
															UpperCamelCase_, UpperCamelCase_:				List[Any]						=    sf.read(snake_case_   )

							else:
											UpperCamelCase_, UpperCamelCase_:				Union[str, Any]						=    sf.read(snake_case_   )

							UpperCamelCase_:				List[str]						=    array.T
							if self.mono:
											UpperCamelCase_:				int						=    librosa.to_mono(snake_case_   )
							if self.sampling_rate and self.sampling_rate != sampling_rate:
											UpperCamelCase_:				Union[str, Any]						=    librosa.resample(snake_case_    ,						orig_sr=snake_case_    ,						target_sr=self.sampling_rate   )
											UpperCamelCase_:				Dict						=    self.sampling_rate

							return {"path": path, "array": array, "sampling_rate": sampling_rate}

			def    lowerCAmelCase__						(							self			:							Tuple   ):
							from .features import Value

							if self.decode:
											raise ValueError("""Cannot flatten a decoded Audio feature."""   )
							return {
							    "bytes": Value("""binary"""   ),
							    "path": Value("""string"""   ),
							}

			def    lowerCAmelCase__						(							self			:							Any    ,						snake_case_			:							Union[pa.StringArray, pa.StructArray]   ):
							if pa.types.is_string(storage.type   ):
											UpperCamelCase_:				Optional[Any]						=    pa.array([None] * len(snake_case_   )    ,						type=pa.binary()   )
											UpperCamelCase_:				List[str]						=    pa.StructArray.from_arrays([bytes_array, storage]    ,						["""bytes""", """path"""]    ,						mask=storage.is_null()   )
							elif pa.types.is_binary(storage.type   ):
											UpperCamelCase_:				str						=    pa.array([None] * len(snake_case_   )    ,						type=pa.string()   )
											UpperCamelCase_:				List[Any]						=    pa.StructArray.from_arrays([storage, path_array]    ,						["""bytes""", """path"""]    ,						mask=storage.is_null()   )
							elif pa.types.is_struct(storage.type   ) and storage.type.get_all_field_indices("""array"""   ):
											UpperCamelCase_:				Union[str, Any]						=    pa.array([Audio().encode_example(snake_case_   ) if x is not None else None for x in storage.to_pylist()]   )
							elif pa.types.is_struct(storage.type   ):
											if storage.type.get_field_index("""bytes"""   ) >= 0:
															UpperCamelCase_:				Optional[Any]						=    storage.field("""bytes"""   )
											else:
															UpperCamelCase_:				str						=    pa.array([None] * len(snake_case_   )    ,						type=pa.binary()   )
											if storage.type.get_field_index("""path"""   ) >= 0:
															UpperCamelCase_:				List[str]						=    storage.field("""path"""   )
											else:
															UpperCamelCase_:				int						=    pa.array([None] * len(snake_case_   )    ,						type=pa.string()   )
											UpperCamelCase_:				Dict						=    pa.StructArray.from_arrays([bytes_array, path_array]    ,						["""bytes""", """path"""]    ,						mask=storage.is_null()   )
							return array_cast(snake_case_    ,						self.pa_type   )




			def    lowerCAmelCase__						(							self			:							Optional[Any]    ,						snake_case_			:							pa.StructArray   ):

							@no_op_if_value_is_null
							def path_to_bytes(snake_case_			:							Optional[int]   ):
											with xopen(snake_case_    ,						"""rb"""   ) as f:
															UpperCamelCase_:				Optional[int]						=    f.read()
											return bytes_

							UpperCamelCase_:				int						=    pa.array(
							    [
							        (path_to_bytes(x["""path"""]   ) if x["""bytes"""] is None else x["""bytes"""]) if x is not None else None
							        for x in storage.to_pylist()
							    ]    ,						type=pa.binary()    ,						)
							UpperCamelCase_:				Union[str, Any]						=    pa.array(
							    [os.path.basename(snake_case_   ) if path is not None else None for path in storage.field("""path"""   ).to_pylist()]    ,						type=pa.string()    ,						)
							UpperCamelCase_:				Optional[int]						=    pa.StructArray.from_arrays([bytes_array, path_array]    ,						["""bytes""", """path"""]    ,						mask=bytes_array.is_null()   )
							return array_cast(snake_case_    ,						self.pa_type   )






