



import gc
import random
import unittest

import numpy as np
import torch
from transformers import CLIPTextConfig, CLIPTextModel, CLIPTokenizer

from diffusers import AutoencoderKL, CycleDiffusionPipeline, DDIMScheduler, UNetaDConditionModel
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, skip_mps

from ..pipeline_params import (
    IMAGE_TO_IMAGE_IMAGE_PARAMS,
    TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS,
    TEXT_GUIDED_IMAGE_VARIATION_PARAMS,
)
from ..test_pipelines_common import PipelineLatentTesterMixin, PipelineTesterMixin


enable_full_determinism()

class 					_A (			__UpperCAmelCase   ,__UpperCAmelCase   ,unittest.TestCase					):
       UpperCamelCase__						:					str		      =		CycleDiffusionPipeline
       UpperCamelCase__						:					Tuple		      =		TEXT_GUIDED_IMAGE_VARIATION_PARAMS - {
           '''negative_prompt''',
           '''height''',
           '''width''',
           '''negative_prompt_embeds''',
       }
       UpperCamelCase__						:					Dict		      =		PipelineTesterMixin.required_optional_params - {'''latents'''}
       UpperCamelCase__						:					Union[str, Any]		      =		TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS.union({'''source_prompt'''}					)
       UpperCamelCase__						:					int		      =		IMAGE_TO_IMAGE_IMAGE_PARAMS
       UpperCamelCase__						:					Optional[Any]		      =		IMAGE_TO_IMAGE_IMAGE_PARAMS



       def 	_lowerCamelCase    (	self							:  Tuple):


           '''simple docstring'''

           torch.manual_seed(0)
           __a		 =       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  ,	)
           __a		 =       DDIMScheduler(
               beta_start=0.0_00_85  ,	beta_end=0.0_12  ,	beta_schedule='''scaled_linear'''  ,	num_train_timesteps=1_000  ,	clip_sample=__SCREAMING_SNAKE_CASE  ,	set_alpha_to_one=__SCREAMING_SNAKE_CASE  ,	)
           torch.manual_seed(0)
           __a		 =       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  ,	)
           torch.manual_seed(0)
           __a		 =       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=1_000  ,	)
           __a		 =       CLIPTextModel(__SCREAMING_SNAKE_CASE)
           __a		 =       CLIPTokenizer.from_pretrained('''hf-internal-testing/tiny-random-clip''')

           __a		 =       {
               '''unet''': unet,
               '''scheduler''': scheduler,
               '''vae''': vae,
               '''text_encoder''': text_encoder,
               '''tokenizer''': tokenizer,
               '''safety_checker''': None,
               '''feature_extractor''': None,
           }
           return components



       def 	_lowerCamelCase    (	self							:  List[str]  ,	__SCREAMING_SNAKE_CASE							:  Union[str, Any]  ,	__SCREAMING_SNAKE_CASE							:  List[str]=0):


           '''simple docstring'''

           __a		 =       floats_tensor((1, 3, 32, 32)  ,	rng=random.Random(__SCREAMING_SNAKE_CASE)).to(__SCREAMING_SNAKE_CASE)
           __a		 =       image / 2 + 0.5
           if str(__SCREAMING_SNAKE_CASE).startswith('''mps'''):
               __a		 =       torch.manual_seed(__SCREAMING_SNAKE_CASE)
           else:
               __a		 =       torch.Generator(device=__SCREAMING_SNAKE_CASE).manual_seed(__SCREAMING_SNAKE_CASE)
           __a		 =       {
               '''prompt''': '''An astronaut riding an elephant''',
               '''source_prompt''': '''An astronaut riding a horse''',
               '''image''': image,
               '''generator''': generator,
               '''num_inference_steps''': 2,
               '''eta''': 0.1,
               '''strength''': 0.8,
               '''guidance_scale''': 3,
               '''source_guidance_scale''': 1,
               '''output_type''': '''numpy''',
           }
           return inputs



       def 	_lowerCamelCase    (	self							:  List[Any]):


           '''simple docstring'''

           __a		 =       '''cpu'''  # ensure determinism for the device-dependent torch.Generator

           __a		 =       self.get_dummy_components()
           __a		 =       CycleDiffusionPipeline(**__SCREAMING_SNAKE_CASE)
           __a		 =       pipe.to(__SCREAMING_SNAKE_CASE)
           pipe.set_progress_bar_config(disable=__SCREAMING_SNAKE_CASE)

           __a		 =       self.get_dummy_inputs(__SCREAMING_SNAKE_CASE)
           __a		 =       pipe(**__SCREAMING_SNAKE_CASE)
           __a		 =       output.images

           __a		 =       images[0, -3:, -3:, -1]

           assert images.shape == (1, 32, 32, 3)
           __a		 =       np.array([0.44_59, 0.49_43, 0.45_44, 0.66_43, 0.54_74, 0.43_27, 0.57_01, 0.59_59, 0.51_79])

           assert np.abs(image_slice.flatten() - expected_slice).max() < 1E-2



       @unittest.skipIf(torch_device != '''cuda'''  ,	'''This test requires a GPU''')
       def 	_lowerCamelCase    (	self							:  Optional[Any]):


           '''simple docstring'''

           __a		 =       self.get_dummy_components()
           for name, module in components.items():
               if hasattr(__SCREAMING_SNAKE_CASE  ,	'''half'''):
                   __a		 =       module.half()
           __a		 =       CycleDiffusionPipeline(**__SCREAMING_SNAKE_CASE)
           __a		 =       pipe.to(__SCREAMING_SNAKE_CASE)
           pipe.set_progress_bar_config(disable=__SCREAMING_SNAKE_CASE)

           __a		 =       self.get_dummy_inputs(__SCREAMING_SNAKE_CASE)
           __a		 =       pipe(**__SCREAMING_SNAKE_CASE)
           __a		 =       output.images

           __a		 =       images[0, -3:, -3:, -1]

           assert images.shape == (1, 32, 32, 3)
           __a		 =       np.array([0.35_06, 0.45_43, 0.4_46, 0.45_75, 0.51_95, 0.41_55, 0.52_73, 0.5_18, 0.41_16])

           assert np.abs(image_slice.flatten() - expected_slice).max() < 1E-2



       @skip_mps
       def 	_lowerCamelCase    (	self							:  Tuple):


           '''simple docstring'''

           return super().test_save_load_local()



       @unittest.skip('''non-deterministic pipeline''')
       def 	_lowerCamelCase    (	self							:  Dict):


           '''simple docstring'''

           return super().test_inference_batch_single_identical()



       @skip_mps
       def 	_lowerCamelCase    (	self							:  Optional[int]):


           '''simple docstring'''

           return super().test_dict_tuple_outputs_equivalent()



       @skip_mps
       def 	_lowerCamelCase    (	self							:  int):


           '''simple docstring'''

           return super().test_save_load_optional_components()





       @skip_mps
       def 	_lowerCamelCase    (	self							:  Tuple):


           '''simple docstring'''

           return super().test_attention_slicing_forward_pass()


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



       def 	_lowerCamelCase    (	self							:  Any):


           '''simple docstring'''

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



       def 	_lowerCamelCase    (	self							:  str):


           '''simple docstring'''

           __a		 =       load_image(
               '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'''
               '''/cycle-diffusion/black_colored_car.png''')
           __a		 =       load_numpy(
               '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/cycle-diffusion/blue_colored_car_fp16.npy''')
           __a		 =       init_image.resize((512, 512))

           __a		 =       '''CompVis/stable-diffusion-v1-4'''
           __a		 =       DDIMScheduler.from_pretrained(__SCREAMING_SNAKE_CASE  ,	subfolder='''scheduler''')
           __a		 =       CycleDiffusionPipeline.from_pretrained(
               __SCREAMING_SNAKE_CASE  ,	scheduler=__SCREAMING_SNAKE_CASE  ,	safety_checker=__SCREAMING_SNAKE_CASE  ,	torch_dtype=torch.floataa  ,	revision='''fp16''')

           pipe.to(__SCREAMING_SNAKE_CASE)
           pipe.set_progress_bar_config(disable=__SCREAMING_SNAKE_CASE)
           pipe.enable_attention_slicing()

           __a		 =       '''A black colored car'''
           __a		 =       '''A blue colored car'''

           __a		 =       torch.manual_seed(0)
           __a		 =       pipe(
               prompt=__SCREAMING_SNAKE_CASE  ,	source_prompt=__SCREAMING_SNAKE_CASE  ,	image=__SCREAMING_SNAKE_CASE  ,	num_inference_steps=100  ,	eta=0.1  ,	strength=0.85  ,	guidance_scale=3  ,	source_guidance_scale=1  ,	generator=__SCREAMING_SNAKE_CASE  ,	output_type='''np'''  ,	)
           __a		 =       output.images

           # the values aren't exactly equal, but the images look the same visually
           assert np.abs(image - expected_image).max() < 5E-1





       def 	_lowerCamelCase    (	self							:  str):


           '''simple docstring'''

           __a		 =       load_image(
               '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'''
               '''/cycle-diffusion/black_colored_car.png''')
           __a		 =       load_numpy(
               '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/cycle-diffusion/blue_colored_car.npy''')
           __a		 =       init_image.resize((512, 512))

           __a		 =       '''CompVis/stable-diffusion-v1-4'''
           __a		 =       DDIMScheduler.from_pretrained(__SCREAMING_SNAKE_CASE  ,	subfolder='''scheduler''')
           __a		 =       CycleDiffusionPipeline.from_pretrained(__SCREAMING_SNAKE_CASE  ,	scheduler=__SCREAMING_SNAKE_CASE  ,	safety_checker=__SCREAMING_SNAKE_CASE)

           pipe.to(__SCREAMING_SNAKE_CASE)
           pipe.set_progress_bar_config(disable=__SCREAMING_SNAKE_CASE)
           pipe.enable_attention_slicing()

           __a		 =       '''A black colored car'''
           __a		 =       '''A blue colored car'''

           __a		 =       torch.manual_seed(0)
           __a		 =       pipe(
               prompt=__SCREAMING_SNAKE_CASE  ,	source_prompt=__SCREAMING_SNAKE_CASE  ,	image=__SCREAMING_SNAKE_CASE  ,	num_inference_steps=100  ,	eta=0.1  ,	strength=0.85  ,	guidance_scale=3  ,	source_guidance_scale=1  ,	generator=__SCREAMING_SNAKE_CASE  ,	output_type='''np'''  ,	)
           __a		 =       output.images

           assert np.abs(image - expected_image).max() < 2E-2










# Lint as: python3
import sys
from collections.abc import Mapping
from typing import TYPE_CHECKING, Dict, Optional

import numpy as np
import pyarrow as pa

from .. import config
from ..utils.logging import get_logger
from ..utils.py_utils import map_nested
from .formatting import TensorFormatter


if TYPE_CHECKING:
 import jax
 import jaxlib

__snake_case		:str   	=  get_logger()

__snake_case		:Optional[dict]   	=  None

class 					_A (			TensorFormatter[Mapping, '''jax.Array''', Mapping]					):



       def __init__(	self							:  str  ,	__SCREAMING_SNAKE_CASE							:  int=None  ,	__SCREAMING_SNAKE_CASE							:  List[Any]=None  ,	**__SCREAMING_SNAKE_CASE							:  Any):


           '''simple docstring'''

           super().__init__(features=__SCREAMING_SNAKE_CASE)
           import jax
           from jaxlib.xla_client import Device

           if isinstance(__SCREAMING_SNAKE_CASE  ,	__SCREAMING_SNAKE_CASE):
               raise ValueError(
                   F'Expected {device} to be a `str` not {type(__SCREAMING_SNAKE_CASE)}, as `jaxlib.xla_extension.Device` '
                   '''is not serializable neither with `pickle` nor with `dill`. Instead you can surround '''
                   '''the device with `str()` to get its string identifier that will be internally mapped '''
                   '''to the actual `jaxlib.xla_extension.Device`.''')
           __a		 =       device if isinstance(__SCREAMING_SNAKE_CASE  ,	__SCREAMING_SNAKE_CASE) else str(jax.devices()[0])
           # using global variable since `jaxlib.xla_extension.Device` is not serializable neither
           # with `pickle` nor with `dill`, so we need to use a global variable instead
           global DEVICE_MAPPING
           if DEVICE_MAPPING is None:
               __a		 =       self._map_devices_to_str()
           if self.device not in list(DEVICE_MAPPING.keys()):
               logger.warning(
                   F'Device with string identifier {self.device} not listed among the available '
                   F'devices: {list(DEVICE_MAPPING.keys())}, so falling back to the default '
                   F'device: {str(jax.devices()[0])}.')
               __a		 =       str(jax.devices()[0])
           __a		 =       jnp_array_kwargs



       @staticmethod
       def 	_lowerCamelCase    (	):


           '''simple docstring'''

           import jax

           return {str(__SCREAMING_SNAKE_CASE): device for device in jax.devices()}



       def 	_lowerCamelCase    (	self							:  Union[str, Any]  ,	__SCREAMING_SNAKE_CASE							:  Tuple):


           '''simple docstring'''

           import jax
           import jax.numpy as jnp

           if isinstance(__SCREAMING_SNAKE_CASE  ,	__SCREAMING_SNAKE_CASE) and column:
               if all(
                   isinstance(__SCREAMING_SNAKE_CASE  ,	jax.Array) and x.shape == column[0].shape and x.dtype == column[0].dtype for x in column):
                   return jnp.stack(__SCREAMING_SNAKE_CASE  ,	axis=0)
           return column



       def 	_lowerCamelCase    (	self							:  str  ,	__SCREAMING_SNAKE_CASE							:  str):


           '''simple docstring'''

           import jax
           import jax.numpy as jnp

           if isinstance(__SCREAMING_SNAKE_CASE  ,	(str, bytes, type(__SCREAMING_SNAKE_CASE))):
               return value
           elif isinstance(__SCREAMING_SNAKE_CASE  ,	(np.character, np.ndarray)) and np.issubdtype(value.dtype  ,	np.character):
               return value.tolist()

           __a		 =       {}

           if isinstance(__SCREAMING_SNAKE_CASE  ,	(np.number, np.ndarray)) and np.issubdtype(value.dtype  ,	np.integer):
               # the default int precision depends on the jax config
               # see https://jax.readthedocs.io/en/latest/notebooks/Common_Gotchas_in_JAX.html#double-64bit-precision
               if jax.config.jax_enable_xaa:
                   __a		 =       {'''dtype''': jnp.intaa}
               else:
                   __a		 =       {'''dtype''': jnp.intaa}
           elif isinstance(__SCREAMING_SNAKE_CASE  ,	(np.number, np.ndarray)) and np.issubdtype(value.dtype  ,	np.floating):
               __a		 =       {'''dtype''': jnp.floataa}
           elif config.PIL_AVAILABLE and "PIL" in sys.modules:
               import PIL.Image

               if isinstance(__SCREAMING_SNAKE_CASE  ,	PIL.Image.Image):
                   __a		 =       np.asarray(__SCREAMING_SNAKE_CASE)

        # using global variable since `jaxlib.xla_extension.Device` is not serializable neither
        # with `pickle` nor with `dill`, so we need to use a global variable instead
           global DEVICE_MAPPING
           if DEVICE_MAPPING is None:
               __a		 =       self._map_devices_to_str()

           with jax.default_device(DEVICE_MAPPING[self.device]):
               # calling jnp.array on a np.ndarray does copy the data
               # see https://github.com/google/jax/issues/4486
               return jnp.array(__SCREAMING_SNAKE_CASE  ,	**{**default_dtype, **self.jnp_array_kwargs})



       def 	_lowerCamelCase    (	self							:  Tuple  ,	__SCREAMING_SNAKE_CASE							:  str):


           '''simple docstring'''

           import jax

           # support for torch, tf, jax etc.
           if config.TORCH_AVAILABLE and "torch" in sys.modules:
               import torch

               if isinstance(__SCREAMING_SNAKE_CASE  ,	torch.Tensor):
                   return self._tensorize(data_struct.detach().cpu().numpy()[()])
           if hasattr(__SCREAMING_SNAKE_CASE  ,	'''__array__''') and not isinstance(__SCREAMING_SNAKE_CASE  ,	jax.Array):
               __a		 =       data_struct.__array__()
           # support for nested types like struct of list of struct
           if isinstance(__SCREAMING_SNAKE_CASE  ,	np.ndarray):
               if data_struct.dtype == object:  # jax arrays cannot be instantied from an array of objects
                   return self._consolidate([self.recursive_tensorize(__SCREAMING_SNAKE_CASE) for substruct in data_struct])
           elif isinstance(__SCREAMING_SNAKE_CASE  ,	(list, tuple)):
               return self._consolidate([self.recursive_tensorize(__SCREAMING_SNAKE_CASE) for substruct in data_struct])
           return self._tensorize(__SCREAMING_SNAKE_CASE)



       def 	_lowerCamelCase    (	self							:  Tuple  ,	__SCREAMING_SNAKE_CASE							:  dict):


           '''simple docstring'''

           return map_nested(self._recursive_tensorize  ,	__SCREAMING_SNAKE_CASE  ,	map_list=__SCREAMING_SNAKE_CASE)



       def 	_lowerCamelCase    (	self							:  str  ,	__SCREAMING_SNAKE_CASE							:  pa.Table):


           '''simple docstring'''

           __a		 =       self.numpy_arrow_extractor().extract_row(__SCREAMING_SNAKE_CASE)
           __a		 =       self.python_features_decoder.decode_row(__SCREAMING_SNAKE_CASE)
           return self.recursive_tensorize(__SCREAMING_SNAKE_CASE)



       def 	_lowerCamelCase    (	self							:  Tuple  ,	__SCREAMING_SNAKE_CASE							:  pa.Table):


           '''simple docstring'''

           __a		 =       self.numpy_arrow_extractor().extract_column(__SCREAMING_SNAKE_CASE)
           __a		 =       self.python_features_decoder.decode_column(__SCREAMING_SNAKE_CASE  ,	pa_table.column_names[0])
           __a		 =       self.recursive_tensorize(__SCREAMING_SNAKE_CASE)
           __a		 =       self._consolidate(__SCREAMING_SNAKE_CASE)
           return column





       def 	_lowerCamelCase    (	self							:  Union[str, Any]  ,	__SCREAMING_SNAKE_CASE							:  pa.Table):


           '''simple docstring'''

           __a		 =       self.numpy_arrow_extractor().extract_batch(__SCREAMING_SNAKE_CASE)
           __a		 =       self.python_features_decoder.decode_batch(__SCREAMING_SNAKE_CASE)
           __a		 =       self.recursive_tensorize(__SCREAMING_SNAKE_CASE)
           for column_name in batch:
               __a		 =       self._consolidate(batch[column_name])
           return batch





