



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__  :    List[str]											=						CycleDiffusionPipeline
       UpperCamelCase__  :    str											=						TEXT_GUIDED_IMAGE_VARIATION_PARAMS - {
           '''negative_prompt''',
           '''height''',
           '''width''',
           '''negative_prompt_embeds''',
       }
       UpperCamelCase__  :    Any											=						PipelineTesterMixin.required_optional_params - {'''latents'''}
       UpperCamelCase__  :    Tuple											=						TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS.union({'''source_prompt'''}    )
       UpperCamelCase__  :    Union[str, Any]											=						IMAGE_TO_IMAGE_IMAGE_PARAMS
       UpperCamelCase__  :    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							:     Optional[int]     ,		__SCREAMING_SNAKE_CASE							:     int     ,		__SCREAMING_SNAKE_CASE							:     Dict=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							:     Dict):


           '''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							:     List[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							:     Union[str, Any]):


           '''simple docstring'''

           return super().test_save_load_local()



       @unittest.skip('''non-deterministic pipeline''')
       def 		_lowerCamelCase			(		self							:     Optional[Any]):


           '''simple docstring'''

           return super().test_inference_batch_single_identical()



       @skip_mps
       def 		_lowerCamelCase			(		self							:     Tuple):


           '''simple docstring'''

           return super().test_dict_tuple_outputs_equivalent()



       @skip_mps
       def 		_lowerCamelCase			(		self							:     List[Any]):


           '''simple docstring'''

           return super().test_save_load_optional_components()





       @skip_mps
       def 		_lowerCamelCase			(		self							:     int):


           '''simple docstring'''

           return super().test_attention_slicing_forward_pass()


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



       def 		_lowerCamelCase			(		self							:     List[str]):


           '''simple docstring'''

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



       def 		_lowerCamelCase			(		self							:     List[Any]):


           '''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							:     Union[str, Any]):


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










from typing import Optional

from .. import Features, NamedSplit
from ..packaged_modules.text.text import Text
from ..utils.typing import NestedDataStructureLike, PathLike
from .abc import AbstractDatasetReader

class     _A       (	__UpperCAmelCase    ):



       def __init__(		self							:     Optional[int]     ,		__SCREAMING_SNAKE_CASE							:     NestedDataStructureLike[PathLike]     ,		__SCREAMING_SNAKE_CASE							:     Optional[NamedSplit] = None     ,		__SCREAMING_SNAKE_CASE							:     Optional[Features] = None     ,		__SCREAMING_SNAKE_CASE							:     str = None     ,		__SCREAMING_SNAKE_CASE							:     bool = False     ,		__SCREAMING_SNAKE_CASE							:     bool = False     ,		__SCREAMING_SNAKE_CASE							:     Optional[int] = None     ,		**__SCREAMING_SNAKE_CASE							:     List[str]     ,		):


           '''simple docstring'''

           super().__init__(
               __SCREAMING_SNAKE_CASE     ,		split=__SCREAMING_SNAKE_CASE     ,		features=__SCREAMING_SNAKE_CASE     ,		cache_dir=__SCREAMING_SNAKE_CASE     ,		keep_in_memory=__SCREAMING_SNAKE_CASE     ,		streaming=__SCREAMING_SNAKE_CASE     ,		num_proc=__SCREAMING_SNAKE_CASE     ,		**__SCREAMING_SNAKE_CASE     ,		)
           __a						       =				path_or_paths if isinstance(__SCREAMING_SNAKE_CASE     ,		__SCREAMING_SNAKE_CASE) else {self.split: path_or_paths}
           __a						       =				Text(
               cache_dir=__SCREAMING_SNAKE_CASE     ,		data_files=__SCREAMING_SNAKE_CASE     ,		features=__SCREAMING_SNAKE_CASE     ,		**__SCREAMING_SNAKE_CASE     ,		)





       def 		_lowerCamelCase			(		self							:     List[str]):


           '''simple docstring'''

           if self.streaming:
               __a						       =				self.builder.as_streaming_dataset(split=self.split)
           # Build regular (map-style) dataset
           else:
               __a						       =				None
               __a						       =				None
               __a						       =				None
               __a						       =				None

               self.builder.download_and_prepare(
                   download_config=__SCREAMING_SNAKE_CASE     ,		download_mode=__SCREAMING_SNAKE_CASE     ,		verification_mode=__SCREAMING_SNAKE_CASE     ,		base_path=__SCREAMING_SNAKE_CASE     ,		num_proc=self.num_proc     ,		)
               __a						       =				self.builder.as_dataset(
                   split=self.split     ,		verification_mode=__SCREAMING_SNAKE_CASE     ,		in_memory=self.keep_in_memory)
           return dataset





