






'''simple docstring'''



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 							__SCREAMING_SNAKE_CASE       (							_UpperCamelCase      ,      _UpperCamelCase      ,      unittest.TestCase     ):
      __a  					=CycleDiffusionPipeline
      __a  					=TEXT_GUIDED_IMAGE_VARIATION_PARAMS - {
          "negative_prompt",
          "height",
          "width",
          "negative_prompt_embeds",
      }
      __a  					=PipelineTesterMixin.required_optional_params - {"latents"}
      __a  					=TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS.union({"source_prompt"}     )
      __a  					=IMAGE_TO_IMAGE_IMAGE_PARAMS
      __a  					=IMAGE_TO_IMAGE_IMAGE_PARAMS




      def  __UpperCamelCase							(  self      )							->Union[str, Any]:



         '''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.00085 ,   beta_end=0.012 ,   beta_schedule='scaled_linear' ,   num_train_timesteps=1000 ,   clip_sample=_UpperCAmelCase ,   set_alpha_to_one=_UpperCAmelCase ,   )
         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=1000 ,   )
         __a				     =    CLIPTextModel(_UpperCAmelCase      )
         __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  __UpperCamelCase							(  self ,   lowerCamelCase ,   lowerCamelCase=0      )							->Any:



         '''simple docstring'''
         __a				     =    floats_tensor((1, 3, 32, 32) ,   rng=random.Random(_UpperCAmelCase      )      ).to(_UpperCAmelCase      )
         __a				     =    image / 2 + 0.5
         if str(_UpperCAmelCase      ).startswith('mps'      ):
            __a				     =    torch.manual_seed(_UpperCAmelCase      )
         else:
            __a				     =    torch.Generator(device=_UpperCAmelCase      ).manual_seed(_UpperCAmelCase      )
         __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  __UpperCamelCase							(  self      )							->List[Any]:



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

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

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

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

         assert images.shape == (1, 32, 32, 3)
         __a				     =    np.array([0.4459, 0.4943, 0.4544, 0.6643, 0.5474, 0.4327, 0.5701, 0.5959, 0.5179]      )

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




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



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

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

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

         assert images.shape == (1, 32, 32, 3)
         __a				     =    np.array([0.3506, 0.4543, 0.446, 0.4575, 0.5195, 0.4155, 0.5273, 0.518, 0.4116]      )

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




      @skip_mps
      def  __UpperCamelCase							(  self      )							->Any:



         '''simple docstring'''
         return super().test_save_load_local()




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



         '''simple docstring'''
         return super().test_inference_batch_single_identical()




      @skip_mps
      def  __UpperCamelCase							(  self      )							->Optional[int]:



         '''simple docstring'''
         return super().test_dict_tuple_outputs_equivalent()




      @skip_mps
      def  __UpperCamelCase							(  self      )							->Dict:



         '''simple docstring'''
         return super().test_save_load_optional_components()






      @skip_mps
      def  __UpperCamelCase							(  self      )							->Tuple:



         '''simple docstring'''
         return super().test_attention_slicing_forward_pass()







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




      def  __UpperCamelCase							(  self      )							->List[Any]:



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




      def  __UpperCamelCase							(  self      )							->Dict:



         '''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(_UpperCAmelCase ,   subfolder='scheduler'      )
         __a				     =    CycleDiffusionPipeline.from_pretrained(
             _UpperCAmelCase ,   scheduler=_UpperCAmelCase ,   safety_checker=_UpperCAmelCase ,   torch_dtype=torch.floataa ,   revision='fp16'      )

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

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

         __a				     =    torch.manual_seed(0      )
         __a				     =    pipe(
             prompt=_UpperCAmelCase ,   source_prompt=_UpperCAmelCase ,   image=_UpperCAmelCase ,   num_inference_steps=100 ,   eta=0.1 ,   strength=0.85 ,   guidance_scale=3 ,   source_guidance_scale=1 ,   generator=_UpperCAmelCase ,   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  __UpperCamelCase							(  self      )							->int:



         '''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(_UpperCAmelCase ,   subfolder='scheduler'      )
         __a				     =    CycleDiffusionPipeline.from_pretrained(_UpperCAmelCase ,   scheduler=_UpperCAmelCase ,   safety_checker=_UpperCAmelCase      )

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

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

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

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







'''simple docstring'''



import gc
import random
import unittest

import torch

from diffusers import (
    IFImgaImgPipeline,
    IFImgaImgSuperResolutionPipeline,
    IFInpaintingPipeline,
    IFInpaintingSuperResolutionPipeline,
    IFPipeline,
    IFSuperResolutionPipeline,
)
from diffusers.models.attention_processor import AttnAddedKVProcessor
from diffusers.utils.import_utils import is_xformers_available
from diffusers.utils.testing_utils import floats_tensor, load_numpy, require_torch_gpu, skip_mps, slow, torch_device

from ..pipeline_params import TEXT_TO_IMAGE_BATCH_PARAMS, TEXT_TO_IMAGE_PARAMS
from ..test_pipelines_common import PipelineTesterMixin, assert_mean_pixel_difference
from . import IFPipelineTesterMixin
@skip_mps
class 							__SCREAMING_SNAKE_CASE       (							_lowerCAmelCase      ,      _lowerCAmelCase      ,      unittest.TestCase     ):
      __a  					=IFPipeline
      __a  					=TEXT_TO_IMAGE_PARAMS - {"width", "height", "latents"}
      __a  					=TEXT_TO_IMAGE_BATCH_PARAMS
      __a  					=PipelineTesterMixin.required_optional_params - {"latents"}




      def  __UpperCamelCase							(  self      )							->int:



         '''simple docstring'''
         return self._get_dummy_components()




      def  __UpperCamelCase							(  self ,   lowerCamelCase ,   lowerCamelCase=0      )							->Tuple:



         '''simple docstring'''
         if str(lowerCamelCase      ).startswith('mps'      ):
            __a				     =    torch.manual_seed(lowerCamelCase      )
         else:
            __a				     =    torch.Generator(device=lowerCamelCase      ).manual_seed(lowerCamelCase      )

         __a				     =    {
             'prompt': 'A painting of a squirrel eating a burger',
             'generator': generator,
             'num_inference_steps': 2,
             'output_type': 'numpy',
         }

         return inputs




      def  __UpperCamelCase							(  self      )							->Optional[Any]:



         '''simple docstring'''
         self._test_save_load_optional_components()




      @unittest.skipIf(torch_device != 'cuda' ,   reason='float16 requires CUDA'      )
      def  __UpperCamelCase							(  self      )							->Optional[Any]:



         '''simple docstring'''
         # Due to non-determinism in save load of the hf-internal-testing/tiny-random-t5 text encoder
         super().test_save_load_floataa(expected_max_diff=1e-1      )




      def  __UpperCamelCase							(  self      )							->Optional[int]:



         '''simple docstring'''
         self._test_attention_slicing_forward_pass(expected_max_diff=1e-2      )




      def  __UpperCamelCase							(  self      )							->Any:



         '''simple docstring'''
         self._test_save_load_local()




      def  __UpperCamelCase							(  self      )							->Tuple:



         '''simple docstring'''
         self._test_inference_batch_single_identical(
             expected_max_diff=1e-2 ,   )






      @unittest.skipIf(
          torch_device != 'cuda' or not is_xformers_available() ,   reason='XFormers attention is only available with CUDA and `xformers` installed' ,   )
      def  __UpperCamelCase							(  self      )							->List[str]:



         '''simple docstring'''
         self._test_xformers_attention_forwardGenerator_pass(expected_max_diff=1e-3      )







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




      def  __UpperCamelCase							(  self      )							->Optional[Any]:



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




      def  __UpperCamelCase							(  self      )							->Any:



         '''simple docstring'''
         # if

         __a				     =    IFPipeline.from_pretrained('DeepFloyd/IF-I-XL-v1.0' ,   variant='fp16' ,   torch_dtype=torch.floataa      )

         __a				     =    IFSuperResolutionPipeline.from_pretrained(
             'DeepFloyd/IF-II-L-v1.0' ,   variant='fp16' ,   torch_dtype=torch.floataa ,   text_encoder=lowerCamelCase ,   tokenizer=lowerCamelCase      )

         # pre compute text embeddings and remove T5 to save memory

         pipe_a.text_encoder.to('cuda'      )

         __a     ,			__a				     =    pipe_a.encode_prompt('anime turtle' ,   device='cuda'      )

         del pipe_a.tokenizer
         del pipe_a.text_encoder
         gc.collect()

         __a				     =    None
         __a				     =    None

         pipe_a.enable_model_cpu_offload()
         pipe_a.enable_model_cpu_offload()

         pipe_a.unet.set_attn_processor(AttnAddedKVProcessor()      )
         pipe_a.unet.set_attn_processor(AttnAddedKVProcessor()      )

         self._test_if(lowerCamelCase ,   lowerCamelCase ,   lowerCamelCase ,   lowerCamelCase      )

         pipe_a.remove_all_hooks()
         pipe_a.remove_all_hooks()

         # img2img

         __a				     =    IFImgaImgPipeline(**pipe_a.components      )
         __a				     =    IFImgaImgSuperResolutionPipeline(**pipe_a.components      )

         pipe_a.enable_model_cpu_offload()
         pipe_a.enable_model_cpu_offload()

         pipe_a.unet.set_attn_processor(AttnAddedKVProcessor()      )
         pipe_a.unet.set_attn_processor(AttnAddedKVProcessor()      )

         self._test_if_imgaimg(lowerCamelCase ,   lowerCamelCase ,   lowerCamelCase ,   lowerCamelCase      )

         pipe_a.remove_all_hooks()
         pipe_a.remove_all_hooks()

         # inpainting

         __a				     =    IFInpaintingPipeline(**pipe_a.components      )
         __a				     =    IFInpaintingSuperResolutionPipeline(**pipe_a.components      )

         pipe_a.enable_model_cpu_offload()
         pipe_a.enable_model_cpu_offload()

         pipe_a.unet.set_attn_processor(AttnAddedKVProcessor()      )
         pipe_a.unet.set_attn_processor(AttnAddedKVProcessor()      )

         self._test_if_inpainting(lowerCamelCase ,   lowerCamelCase ,   lowerCamelCase ,   lowerCamelCase      )




      def  __UpperCamelCase							(  self ,   lowerCamelCase ,   lowerCamelCase ,   lowerCamelCase ,   lowerCamelCase      )							->List[Any]:



         '''simple docstring'''
         # pipeline 1

         _start_torch_memory_measurement()

         __a				     =    torch.Generator(device='cpu'      ).manual_seed(0      )
         __a				     =    pipe_a(
             prompt_embeds=lowerCamelCase ,   negative_prompt_embeds=lowerCamelCase ,   num_inference_steps=2 ,   generator=lowerCamelCase ,   output_type='np' ,   )

         __a				     =    output.images[0]

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

         __a				     =    torch.cuda.max_memory_allocated()
         assert mem_bytes < 13 * 10**9

         __a				     =    load_numpy(
             'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/if/test_if.npy'      )
         assert_mean_pixel_difference(lowerCamelCase ,   lowerCamelCase      )

         # pipeline 2

         _start_torch_memory_measurement()

         __a				     =    torch.Generator(device='cpu'      ).manual_seed(0      )

         __a				     =    floats_tensor((1, 3, 64, 64) ,   rng=random.Random(0      )      ).to(lowerCamelCase      )

         __a				     =    pipe_a(
             prompt_embeds=lowerCamelCase ,   negative_prompt_embeds=lowerCamelCase ,   image=lowerCamelCase ,   generator=lowerCamelCase ,   num_inference_steps=2 ,   output_type='np' ,   )

         __a				     =    output.images[0]

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

         __a				     =    torch.cuda.max_memory_allocated()
         assert mem_bytes < 4 * 10**9

         __a				     =    load_numpy(
             'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/if/test_if_superresolution_stage_II.npy'      )
         assert_mean_pixel_difference(lowerCamelCase ,   lowerCamelCase      )




      def  __UpperCamelCase							(  self ,   lowerCamelCase ,   lowerCamelCase ,   lowerCamelCase ,   lowerCamelCase      )							->Any:



         '''simple docstring'''
         # pipeline 1

         _start_torch_memory_measurement()

         __a				     =    floats_tensor((1, 3, 64, 64) ,   rng=random.Random(0      )      ).to(lowerCamelCase      )

         __a				     =    torch.Generator(device='cpu'      ).manual_seed(0      )

         __a				     =    pipe_a(
             prompt_embeds=lowerCamelCase ,   negative_prompt_embeds=lowerCamelCase ,   image=lowerCamelCase ,   num_inference_steps=2 ,   generator=lowerCamelCase ,   output_type='np' ,   )

         __a				     =    output.images[0]

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

         __a				     =    torch.cuda.max_memory_allocated()
         assert mem_bytes < 10 * 10**9

         __a				     =    load_numpy(
             'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/if/test_if_img2img.npy'      )
         assert_mean_pixel_difference(lowerCamelCase ,   lowerCamelCase      )

         # pipeline 2

         _start_torch_memory_measurement()

         __a				     =    torch.Generator(device='cpu'      ).manual_seed(0      )

         __a				     =    floats_tensor((1, 3, 256, 256) ,   rng=random.Random(0      )      ).to(lowerCamelCase      )
         __a				     =    floats_tensor((1, 3, 64, 64) ,   rng=random.Random(0      )      ).to(lowerCamelCase      )

         __a				     =    pipe_a(
             prompt_embeds=lowerCamelCase ,   negative_prompt_embeds=lowerCamelCase ,   image=lowerCamelCase ,   original_image=lowerCamelCase ,   generator=lowerCamelCase ,   num_inference_steps=2 ,   output_type='np' ,   )

         __a				     =    output.images[0]

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

         __a				     =    torch.cuda.max_memory_allocated()
         assert mem_bytes < 4 * 10**9

         __a				     =    load_numpy(
             'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/if/test_if_img2img_superresolution_stage_II.npy'      )
         assert_mean_pixel_difference(lowerCamelCase ,   lowerCamelCase      )






      def  __UpperCamelCase							(  self ,   lowerCamelCase ,   lowerCamelCase ,   lowerCamelCase ,   lowerCamelCase      )							->Optional[int]:



         '''simple docstring'''
         # pipeline 1

         _start_torch_memory_measurement()

         __a				     =    floats_tensor((1, 3, 64, 64) ,   rng=random.Random(0      )      ).to(lowerCamelCase      )
         __a				     =    floats_tensor((1, 3, 64, 64) ,   rng=random.Random(1      )      ).to(lowerCamelCase      )

         __a				     =    torch.Generator(device='cpu'      ).manual_seed(0      )
         __a				     =    pipe_a(
             prompt_embeds=lowerCamelCase ,   negative_prompt_embeds=lowerCamelCase ,   image=lowerCamelCase ,   mask_image=lowerCamelCase ,   num_inference_steps=2 ,   generator=lowerCamelCase ,   output_type='np' ,   )

         __a				     =    output.images[0]

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

         __a				     =    torch.cuda.max_memory_allocated()
         assert mem_bytes < 10 * 10**9

         __a				     =    load_numpy(
             'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/if/test_if_inpainting.npy'      )
         assert_mean_pixel_difference(lowerCamelCase ,   lowerCamelCase      )

         # pipeline 2

         _start_torch_memory_measurement()

         __a				     =    torch.Generator(device='cpu'      ).manual_seed(0      )

         __a				     =    floats_tensor((1, 3, 64, 64) ,   rng=random.Random(0      )      ).to(lowerCamelCase      )
         __a				     =    floats_tensor((1, 3, 256, 256) ,   rng=random.Random(0      )      ).to(lowerCamelCase      )
         __a				     =    floats_tensor((1, 3, 256, 256) ,   rng=random.Random(1      )      ).to(lowerCamelCase      )

         __a				     =    pipe_a(
             prompt_embeds=lowerCamelCase ,   negative_prompt_embeds=lowerCamelCase ,   image=lowerCamelCase ,   mask_image=lowerCamelCase ,   original_image=lowerCamelCase ,   generator=lowerCamelCase ,   num_inference_steps=2 ,   output_type='np' ,   )

         __a				     =    output.images[0]

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

         __a				     =    torch.cuda.max_memory_allocated()
         assert mem_bytes < 4 * 10**9

         __a				     =    load_numpy(
             'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/if/test_if_inpainting_superresolution_stage_II.npy'      )
         assert_mean_pixel_difference(lowerCamelCase ,   lowerCamelCase      )







def    __UpperCAmelCase	(				)  ->			Dict:

   """simple docstring"""
   torch.cuda.empty_cache()
   torch.cuda.reset_max_memory_allocated()
   torch.cuda.reset_peak_memory_stats()