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


   def 			lowercase_      (self     :						str		)						->				List[Any]:
         """simple docstring"""



         torch.manual_seed(0		)
         UpperCAmelCase__								=  UNetaDConditionModel(
             block_out_channels=(3_2, 6_4)       ,  layers_per_block=2       ,  sample_size=3_2       ,  in_channels=4       ,  out_channels=4       ,  down_block_types=("DownBlock2D", "CrossAttnDownBlock2D")       ,  up_block_types=("CrossAttnUpBlock2D", "UpBlock2D")       ,  cross_attention_dim=3_2       ,  )
         UpperCAmelCase__								=  DDIMScheduler(
             beta_start=0.00085       ,  beta_end=0.012       ,  beta_schedule="scaled_linear"       ,  num_train_timesteps=1_0_0_0       ,  clip_sample=__UpperCAmelCase       ,  set_alpha_to_one=__UpperCAmelCase       ,  )
         torch.manual_seed(0		)
         UpperCAmelCase__								=  AutoencoderKL(
             block_out_channels=[3_2, 6_4]       ,  in_channels=3       ,  out_channels=3       ,  down_block_types=["DownEncoderBlock2D", "DownEncoderBlock2D"]       ,  up_block_types=["UpDecoderBlock2D", "UpDecoderBlock2D"]       ,  latent_channels=4       ,  )
         torch.manual_seed(0		)
         UpperCAmelCase__								=  CLIPTextConfig(
             bos_token_id=0       ,  eos_token_id=2       ,  hidden_size=3_2       ,  intermediate_size=3_7       ,  layer_norm_eps=1E-05       ,  num_attention_heads=4       ,  num_hidden_layers=5       ,  pad_token_id=1       ,  vocab_size=1_0_0_0       ,  )
         UpperCAmelCase__								=  CLIPTextModel(__UpperCAmelCase		)
         UpperCAmelCase__								=  CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip"		)

         UpperCAmelCase__								=  {
             "unet": unet,
             "scheduler": scheduler,
             "vae": vae,
             "text_encoder": text_encoder,
             "tokenizer": tokenizer,
             "safety_checker": None,
             "feature_extractor": None,
         }
         return components


   def 			lowercase_      (self     :						str       ,  __UpperCAmelCase     :						Union[str, Any]       ,  __UpperCAmelCase     :						List[Any]=0		)						->				int:
         """simple docstring"""



         UpperCAmelCase__								=  floats_tensor((1, 3, 3_2, 3_2)       ,  rng=random.Random(__UpperCAmelCase		)		).to(__UpperCAmelCase		)
         UpperCAmelCase__								=  image / 2 + 0.5
         if str(__UpperCAmelCase		).startswith("mps"		):
               UpperCAmelCase__								=  torch.manual_seed(__UpperCAmelCase		)
         else:
               UpperCAmelCase__								=  torch.Generator(device=__UpperCAmelCase		).manual_seed(__UpperCAmelCase		)
         UpperCAmelCase__								=  {
             "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 			lowercase_      (self     :						List[Any]		)						->				Union[str, Any]:
         """simple docstring"""



         UpperCAmelCase__								=  "cpu"  # ensure determinism for the device-dependent torch.Generator

         UpperCAmelCase__								=  self.get_dummy_components()
         UpperCAmelCase__								=  CycleDiffusionPipeline(**__UpperCAmelCase		)
         UpperCAmelCase__								=  pipe.to(__UpperCAmelCase		)
         pipe.set_progress_bar_config(disable=__UpperCAmelCase		)

         UpperCAmelCase__								=  self.get_dummy_inputs(__UpperCAmelCase		)
         UpperCAmelCase__								=  pipe(**__UpperCAmelCase		)
         UpperCAmelCase__								=  output.images

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

         assert images.shape == (1, 3_2, 3_2, 3)
         UpperCAmelCase__								=  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 			lowercase_      (self     :						int		)						->				int:
         """simple docstring"""



         UpperCAmelCase__								=  self.get_dummy_components()
         for name, module in components.items():
               if hasattr(__UpperCAmelCase       ,  "half"		):
                     UpperCAmelCase__								=  module.half()
         UpperCAmelCase__								=  CycleDiffusionPipeline(**__UpperCAmelCase		)
         UpperCAmelCase__								=  pipe.to(__UpperCAmelCase		)
         pipe.set_progress_bar_config(disable=__UpperCAmelCase		)

         UpperCAmelCase__								=  self.get_dummy_inputs(__UpperCAmelCase		)
         UpperCAmelCase__								=  pipe(**__UpperCAmelCase		)
         UpperCAmelCase__								=  output.images

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

         assert images.shape == (1, 3_2, 3_2, 3)
         UpperCAmelCase__								=  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 			lowercase_      (self     :						Optional[int]		)						->				Tuple:
         """simple docstring"""



         return super().test_save_load_local()


   @unittest.skip("non-deterministic pipeline"		)
   def 			lowercase_      (self     :						Any		)						->				str:
         """simple docstring"""



         return super().test_inference_batch_single_identical()


   @skip_mps
   def 			lowercase_      (self     :						List[str]		)						->				Any:
         """simple docstring"""



         return super().test_dict_tuple_outputs_equivalent()


   @skip_mps
   def 			lowercase_      (self     :						Optional[int]		)						->				Any:
         """simple docstring"""



         return super().test_save_load_optional_components()





   @skip_mps
   def 			lowercase_      (self     :						Union[str, Any]		)						->				Union[str, Any]:
         """simple docstring"""



         return super().test_attention_slicing_forward_pass()





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


   def 			lowercase_      (self     :						Dict		)						->				Optional[Any]:
         """simple docstring"""



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


   def 			lowercase_      (self     :						Union[str, Any]		)						->				Optional[int]:
         """simple docstring"""



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

         UpperCAmelCase__								=  "CompVis/stable-diffusion-v1-4"
         UpperCAmelCase__								=  DDIMScheduler.from_pretrained(__UpperCAmelCase       ,  subfolder="scheduler"		)
         UpperCAmelCase__								=  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()

         UpperCAmelCase__								=  "A black colored car"
         UpperCAmelCase__								=  "A blue colored car"

         UpperCAmelCase__								=  torch.manual_seed(0		)
         UpperCAmelCase__								=  pipe(
             prompt=__UpperCAmelCase       ,  source_prompt=__UpperCAmelCase       ,  image=__UpperCAmelCase       ,  num_inference_steps=1_0_0       ,  eta=0.1       ,  strength=0.85       ,  guidance_scale=3       ,  source_guidance_scale=1       ,  generator=__UpperCAmelCase       ,  output_type="np"       ,  )
         UpperCAmelCase__								=  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 			lowercase_      (self     :						Any		)						->				List[Any]:
         """simple docstring"""



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

         UpperCAmelCase__								=  "CompVis/stable-diffusion-v1-4"
         UpperCAmelCase__								=  DDIMScheduler.from_pretrained(__UpperCAmelCase       ,  subfolder="scheduler"		)
         UpperCAmelCase__								=  CycleDiffusionPipeline.from_pretrained(__UpperCAmelCase       ,  scheduler=__UpperCAmelCase       ,  safety_checker=__UpperCAmelCase		)

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

         UpperCAmelCase__								=  "A black colored car"
         UpperCAmelCase__								=  "A blue colored car"

         UpperCAmelCase__								=  torch.manual_seed(0		)
         UpperCAmelCase__								=  pipe(
             prompt=__UpperCAmelCase       ,  source_prompt=__UpperCAmelCase       ,  image=__UpperCAmelCase       ,  num_inference_steps=1_0_0       ,  eta=0.1       ,  strength=0.85       ,  guidance_scale=3       ,  source_guidance_scale=1       ,  generator=__UpperCAmelCase       ,  output_type="np"       ,  )
         UpperCAmelCase__								=  output.images

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




from manim import *





class      A							(				UpperCAmelCase_     ):


   def 			lowercase_      (self     :						Union[str, Any]		)						->				List[str]:
         """simple docstring"""



         UpperCAmelCase__								=  Rectangle(height=0.5       ,  width=0.5		)
         UpperCAmelCase__								=  Rectangle(height=0.46       ,  width=0.46		).set_stroke(width=0		)

         UpperCAmelCase__								=  Rectangle(height=0.25       ,  width=0.25		)

         UpperCAmelCase__								=  [mem.copy() for i in range(6		)]
         UpperCAmelCase__								=  [mem.copy() for i in range(6		)]
         UpperCAmelCase__								=  VGroup(*__UpperCAmelCase		).arrange(__UpperCAmelCase       ,  buff=0		)
         UpperCAmelCase__								=  VGroup(*__UpperCAmelCase		).arrange(__UpperCAmelCase       ,  buff=0		)
         UpperCAmelCase__								=  VGroup(__UpperCAmelCase       ,  __UpperCAmelCase		).arrange(__UpperCAmelCase       ,  buff=0		)
         UpperCAmelCase__								=  Text("CPU"       ,  font_size=2_4		)
         UpperCAmelCase__								=  Group(__UpperCAmelCase       ,  __UpperCAmelCase		).arrange(__UpperCAmelCase       ,  buff=0.5       ,  aligned_edge=__UpperCAmelCase		)
         cpu.move_to([-2.5, -0.5, 0]		)
         self.add(__UpperCAmelCase		)

         UpperCAmelCase__								=  [mem.copy() for i in range(4		)]
         UpperCAmelCase__								=  VGroup(*__UpperCAmelCase		).arrange(__UpperCAmelCase       ,  buff=0		)
         UpperCAmelCase__								=  Text("GPU"       ,  font_size=2_4		)
         UpperCAmelCase__								=  Group(__UpperCAmelCase       ,  __UpperCAmelCase		).arrange(__UpperCAmelCase       ,  buff=0.5       ,  aligned_edge=__UpperCAmelCase		)
         gpu.move_to([-1, -1, 0]		)
         self.add(__UpperCAmelCase		)

         UpperCAmelCase__								=  [mem.copy() for i in range(6		)]
         UpperCAmelCase__								=  VGroup(*__UpperCAmelCase		).arrange(__UpperCAmelCase       ,  buff=0		)

         UpperCAmelCase__								=  Text("Model"       ,  font_size=2_4		)
         UpperCAmelCase__								=  Group(__UpperCAmelCase       ,  __UpperCAmelCase		).arrange(__UpperCAmelCase       ,  buff=0.5       ,  aligned_edge=__UpperCAmelCase		)
         model.move_to([3, -1.0, 0]		)
         self.add(__UpperCAmelCase		)

         UpperCAmelCase__								=  []
         UpperCAmelCase__								=  []

         for i, rect in enumerate(__UpperCAmelCase		):
               UpperCAmelCase__								=  fill.copy().set_fill(__UpperCAmelCase       ,  opacity=0.8		)
               target.move_to(__UpperCAmelCase		)
               model_arr.append(__UpperCAmelCase		)

               UpperCAmelCase__								=  Rectangle(height=0.46       ,  width=0.46		).set_stroke(width=0.0		).set_fill(__UpperCAmelCase       ,  opacity=0.8		)
               cpu_target.move_to(cpu_left_col_base[i]		)
               model_cpu_arr.append(__UpperCAmelCase		)

         self.add(*__UpperCAmelCase       ,  *__UpperCAmelCase		)

         UpperCAmelCase__								=  [meta_mem.copy() for i in range(6		)]
         UpperCAmelCase__								=  [meta_mem.copy() for i in range(6		)]
         UpperCAmelCase__								=  VGroup(*__UpperCAmelCase		).arrange(__UpperCAmelCase       ,  buff=0		)
         UpperCAmelCase__								=  VGroup(*__UpperCAmelCase		).arrange(__UpperCAmelCase       ,  buff=0		)
         UpperCAmelCase__								=  VGroup(__UpperCAmelCase       ,  __UpperCAmelCase		).arrange(__UpperCAmelCase       ,  buff=0		)
         UpperCAmelCase__								=  Text("Disk"       ,  font_size=2_4		)
         UpperCAmelCase__								=  Group(__UpperCAmelCase       ,  __UpperCAmelCase		).arrange(__UpperCAmelCase       ,  buff=0.5       ,  aligned_edge=__UpperCAmelCase		)
         disk.move_to([-4, -1.25, 0]		)
         self.add(__UpperCAmelCase       ,  __UpperCAmelCase		)

         UpperCAmelCase__								=  Square(side_length=2.2		)
         key.move_to([-5, 2, 0]		)

         UpperCAmelCase__								=  MarkupText(
             f"""<b>Key:</b>\n\n<span fgcolor='{YELLOW}'>●</span> Empty Model"""       ,  font_size=1_8       ,  )

         key_text.move_to([-5, 2.4, 0]		)

         self.add(__UpperCAmelCase       ,  __UpperCAmelCase		)

         UpperCAmelCase__								=  MarkupText(
             f"""<span fgcolor='{BLUE}'>●</span> Checkpoint"""       ,  font_size=1_8       ,  )

         blue_text.next_to(__UpperCAmelCase       ,  DOWN * 2.4       ,  aligned_edge=key_text.get_left()		)
         self.add(__UpperCAmelCase		)

         UpperCAmelCase__								=  MarkupText(
             f"""Now watch as an input is passed through the model\nand how the memory is utilized and handled."""       ,  font_size=2_4       ,  )
         step_a.move_to([2, 2, 0]		)

         self.play(Write(__UpperCAmelCase		)		)

         UpperCAmelCase__								=  Square(0.3		)
         input.set_fill(__UpperCAmelCase       ,  opacity=1.0		)
         input.set_stroke(width=0.0		)
         input.next_to(model_base[0]       ,  __UpperCAmelCase       ,  buff=0.5		)

         self.play(Write(__UpperCAmelCase		)		)

         input.generate_target()
         input.target.next_to(model_arr[0]       ,  direction=__UpperCAmelCase       ,  buff=0.02		)
         self.play(MoveToTarget(__UpperCAmelCase		)		)

         self.play(FadeOut(__UpperCAmelCase		)		)

         UpperCAmelCase__								=  Arrow(start=__UpperCAmelCase       ,  end=__UpperCAmelCase       ,  color=__UpperCAmelCase       ,  buff=0.5		)
         a.next_to(model_arr[0].get_left()       ,  __UpperCAmelCase       ,  buff=0.2		)

         model_cpu_arr[0].generate_target()
         model_cpu_arr[0].target.move_to(gpu_rect[0]		)

         UpperCAmelCase__								=  MarkupText(
             f"""As the input reaches a layer, the hook triggers\nand weights are moved from the CPU\nto the GPU and back."""       ,  font_size=2_4       ,  )
         step_a.move_to([2, 2, 0]		)

         self.play(Write(__UpperCAmelCase       ,  run_time=3		)		)

         UpperCAmelCase__								=  {"run_time": 1, "fade_in": True, "fade_out": True, "buff": 0.02}

         self.play(
             Write(__UpperCAmelCase		)       ,  Circumscribe(model_arr[0]       ,  color=__UpperCAmelCase       ,  **__UpperCAmelCase		)       ,  Circumscribe(model_cpu_arr[0]       ,  color=__UpperCAmelCase       ,  **__UpperCAmelCase		)       ,  Circumscribe(gpu_rect[0]       ,  color=__UpperCAmelCase       ,  **__UpperCAmelCase		)       ,  )
         self.play(MoveToTarget(model_cpu_arr[0]		)		)

         UpperCAmelCase__								=  a.copy()
         for i in range(6		):
               a_c.next_to(model_arr[i].get_right() + 0.02       ,  __UpperCAmelCase       ,  buff=0.2		)

               input.generate_target()
               input.target.move_to(model_arr[i].get_right() + 0.02		)

               UpperCAmelCase__								=  AnimationGroup(
                   FadeOut(__UpperCAmelCase       ,  run_time=0.5		)       ,  MoveToTarget(__UpperCAmelCase       ,  run_time=0.5		)       ,  FadeIn(__UpperCAmelCase       ,  run_time=0.5		)       ,  lag_ratio=0.2		)

               self.play(__UpperCAmelCase		)

               model_cpu_arr[i].generate_target()
               model_cpu_arr[i].target.move_to(cpu_left_col_base[i]		)

               if i < 5:
                     model_cpu_arr[i + 1].generate_target()
                     model_cpu_arr[i + 1].target.move_to(gpu_rect[0]		)
                     if i >= 1:
                           UpperCAmelCase__								=  0.7

                     self.play(
                         Circumscribe(model_arr[i]       ,  **__UpperCAmelCase		)       ,  Circumscribe(cpu_left_col_base[i]       ,  **__UpperCAmelCase		)       ,  Circumscribe(cpu_left_col_base[i + 1]       ,  color=__UpperCAmelCase       ,  **__UpperCAmelCase		)       ,  Circumscribe(gpu_rect[0]       ,  color=__UpperCAmelCase       ,  **__UpperCAmelCase		)       ,  Circumscribe(model_arr[i + 1]       ,  color=__UpperCAmelCase       ,  **__UpperCAmelCase		)       ,  )
                     if i < 1:
                           self.play(
                               MoveToTarget(model_cpu_arr[i]		)       ,  MoveToTarget(model_cpu_arr[i + 1]		)       ,  )
                     else:
                           self.play(
                               MoveToTarget(model_cpu_arr[i]       ,  run_time=0.7		)       ,  MoveToTarget(model_cpu_arr[i + 1]       ,  run_time=0.7		)       ,  )
               else:
                     model_cpu_arr[i].generate_target()
                     model_cpu_arr[i].target.move_to(cpu_left_col_base[-1]		)
                     input.generate_target()
                     input.target.next_to(model_arr[-1].get_right()       ,  RIGHT + 0.02       ,  buff=0.2		)

                     self.play(
                         Circumscribe(model_arr[-1]       ,  color=__UpperCAmelCase       ,  **__UpperCAmelCase		)       ,  Circumscribe(cpu_left_col_base[-1]       ,  color=__UpperCAmelCase       ,  **__UpperCAmelCase		)       ,  Circumscribe(gpu_rect[0]       ,  color=__UpperCAmelCase       ,  **__UpperCAmelCase		)       ,  )

                     self.play(MoveToTarget(model_cpu_arr[i]		)		)

               UpperCAmelCase__								=  a_c
               UpperCAmelCase__								=  a_c.copy()

         input.generate_target()
         input.target.next_to(model_base[-1]       ,  RIGHT + 0.02       ,  buff=0.5		)
         self.play(
             FadeOut(__UpperCAmelCase		)       ,  FadeOut(__UpperCAmelCase       ,  run_time=0.5		)       ,  )

         UpperCAmelCase__								=  MarkupText(f"""Inference on a model too large for GPU memory\nis successfully completed."""       ,  font_size=2_4		)
         step_a.move_to([2, 2, 0]		)

         self.play(Write(__UpperCAmelCase       ,  run_time=3		)       ,  MoveToTarget(__UpperCAmelCase		)		)

         self.wait()



