



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







 '''simple docstring'''
 lowercase_		=      CycleDiffusionPipeline
 lowercase_		=      TEXT_GUIDED_IMAGE_VARIATION_PARAMS - {
     "negative_prompt",
     "height",
     "width",
     "negative_prompt_embeds",
 }
 lowercase_		=      PipelineTesterMixin.required_optional_params - {"latents"}
 lowercase_		=      TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS.union({"source_prompt"}   )
 lowercase_		=      IMAGE_TO_IMAGE_IMAGE_PARAMS
 lowercase_		=      IMAGE_TO_IMAGE_IMAGE_PARAMS

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



   '''simple docstring'''



   torch.manual_seed(0)
   lowerCamelCase__:				Union[str, Any]        =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	,			)
   lowerCamelCase__:				List[str]        =DDIMScheduler(
       beta_start=0.0_0085	,			beta_end=0.012	,			beta_schedule="scaled_linear"	,			num_train_timesteps=1_000	,			clip_sample=UpperCAmelCase_	,			set_alpha_to_one=UpperCAmelCase_	,			)
   torch.manual_seed(0)
   lowerCamelCase__:				Dict        =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)
   lowerCamelCase__:				Union[str, Any]        =CLIPTextConfig(
       bos_token_id=0	,			eos_token_id=2	,			hidden_size=32	,			intermediate_size=37	,			layer_norm_eps=1E-0_5	,			num_attention_heads=4	,			num_hidden_layers=5	,			pad_token_id=1	,			vocab_size=1_000	,			)
   lowerCamelCase__:				str        =CLIPTextModel(UpperCAmelCase_)
   lowerCamelCase__:				Union[str, Any]        =CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip")

   lowerCamelCase__:				List[str]        ={
       "unet": unet,
       "scheduler": scheduler,
       "vae": vae,
       "text_encoder": text_encoder,
       "tokenizer": tokenizer,
       "safety_checker": None,
       "feature_extractor": None,
   }
   return components

 def 							SCREAMING_SNAKE_CASE_ (self    :       Dict	,			UpperCAmelCase_    :       List[Any]	,			UpperCAmelCase_    :       Any=0)	->List[str]:



   '''simple docstring'''



   lowerCamelCase__:				List[str]        =floats_tensor((1, 3, 32, 32)	,			rng=random.Random(UpperCAmelCase_)).to(UpperCAmelCase_)
   lowerCamelCase__:				Dict        =image / 2 + 0.5
   if str(UpperCAmelCase_).startswith("mps"):
     lowerCamelCase__:				Optional[Any]        =torch.manual_seed(UpperCAmelCase_)
   else:
     lowerCamelCase__:				Union[str, Any]        =torch.Generator(device=UpperCAmelCase_).manual_seed(UpperCAmelCase_)
   lowerCamelCase__:				Dict        ={
       "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 							SCREAMING_SNAKE_CASE_ (self    :       str)	->List[Any]:



   '''simple docstring'''



   lowerCamelCase__:				List[str]        ="cpu"  # ensure determinism for the device-dependent torch.Generator

   lowerCamelCase__:				Optional[int]        =self.get_dummy_components()
   lowerCamelCase__:				List[Any]        =CycleDiffusionPipeline(**UpperCAmelCase_)
   lowerCamelCase__:				List[str]        =pipe.to(UpperCAmelCase_)
   pipe.set_progress_bar_config(disable=UpperCAmelCase_)

   lowerCamelCase__:				Optional[int]        =self.get_dummy_inputs(UpperCAmelCase_)
   lowerCamelCase__:				List[Any]        =pipe(**UpperCAmelCase_)
   lowerCamelCase__:				Optional[Any]        =output.images

   lowerCamelCase__:				Any        =images[0, -3:, -3:, -1]

   assert images.shape == (1, 32, 32, 3)
   lowerCamelCase__:				List[str]        =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 							SCREAMING_SNAKE_CASE_ (self    :       int)	->Optional[Any]:



   '''simple docstring'''



   lowerCamelCase__:				Optional[int]        =self.get_dummy_components()
   for name, module in components.items():
     if hasattr(UpperCAmelCase_	,			"half"):
       lowerCamelCase__:				int        =module.half()
   lowerCamelCase__:				List[Any]        =CycleDiffusionPipeline(**UpperCAmelCase_)
   lowerCamelCase__:				int        =pipe.to(UpperCAmelCase_)
   pipe.set_progress_bar_config(disable=UpperCAmelCase_)

   lowerCamelCase__:				Union[str, Any]        =self.get_dummy_inputs(UpperCAmelCase_)
   lowerCamelCase__:				Optional[Any]        =pipe(**UpperCAmelCase_)
   lowerCamelCase__:				List[Any]        =output.images

   lowerCamelCase__:				Optional[int]        =images[0, -3:, -3:, -1]

   assert images.shape == (1, 32, 32, 3)
   lowerCamelCase__:				Any        =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 							SCREAMING_SNAKE_CASE_ (self    :       List[Any])	->Union[str, Any]:



   '''simple docstring'''



   return super().test_save_load_local()

 @unittest.skip("non-deterministic pipeline")
 def 							SCREAMING_SNAKE_CASE_ (self    :       List[Any])	->int:



   '''simple docstring'''



   return super().test_inference_batch_single_identical()

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



   '''simple docstring'''



   return super().test_dict_tuple_outputs_equivalent()

 @skip_mps
 def 							SCREAMING_SNAKE_CASE_ (self    :       List[str])	->Tuple:



   '''simple docstring'''



   return super().test_save_load_optional_components()







 @skip_mps
 def 							SCREAMING_SNAKE_CASE_ (self    :       int)	->List[Any]:



   '''simple docstring'''



   return super().test_attention_slicing_forward_pass()



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







 '''simple docstring'''

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



   '''simple docstring'''



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

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



   '''simple docstring'''



   lowerCamelCase__:				Any        =load_image(
       "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"
       "/cycle-diffusion/black_colored_car.png")
   lowerCamelCase__:				Tuple        =load_numpy(
       "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/cycle-diffusion/blue_colored_car_fp16.npy")
   lowerCamelCase__:				List[str]        =init_image.resize((512, 512))

   lowerCamelCase__:				List[str]        ="CompVis/stable-diffusion-v1-4"
   lowerCamelCase__:				str        =DDIMScheduler.from_pretrained(UpperCAmelCase_	,			subfolder="scheduler")
   lowerCamelCase__:				List[Any]        =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()

   lowerCamelCase__:				Dict        ="A black colored car"
   lowerCamelCase__:				int        ="A blue colored car"

   lowerCamelCase__:				Union[str, Any]        =torch.manual_seed(0)
   lowerCamelCase__:				Optional[Any]        =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"	,			)
   lowerCamelCase__:				Tuple        =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 							SCREAMING_SNAKE_CASE_ (self    :       Dict)	->Optional[Any]:



   '''simple docstring'''



   lowerCamelCase__:				str        =load_image(
       "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"
       "/cycle-diffusion/black_colored_car.png")
   lowerCamelCase__:				List[Any]        =load_numpy(
       "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/cycle-diffusion/blue_colored_car.npy")
   lowerCamelCase__:				Optional[int]        =init_image.resize((512, 512))

   lowerCamelCase__:				Union[str, Any]        ="CompVis/stable-diffusion-v1-4"
   lowerCamelCase__:				List[str]        =DDIMScheduler.from_pretrained(UpperCAmelCase_	,			subfolder="scheduler")
   lowerCamelCase__:				Tuple        =CycleDiffusionPipeline.from_pretrained(UpperCAmelCase_	,			scheduler=UpperCAmelCase_	,			safety_checker=UpperCAmelCase_)

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

   lowerCamelCase__:				Union[str, Any]        ="A black colored car"
   lowerCamelCase__:				Optional[Any]        ="A blue colored car"

   lowerCamelCase__:				Optional[int]        =torch.manual_seed(0)
   lowerCamelCase__:				Optional[Any]        =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"	,			)
   lowerCamelCase__:				Union[str, Any]        =output.images

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








from typing import Callable, Dict, Optional, Tuple

import torch
from torch import nn
from torch.distributions import (
    AffineTransform,
    Distribution,
    Independent,
    NegativeBinomial,
    Normal,
    StudentT,
    TransformedDistribution,
)



class        _SCREAMING_SNAKE_CASE						(     __SCREAMING_SNAKE_CASE   ):







 '''simple docstring'''

 def __init__(self    :       int	,			UpperCAmelCase_    :       Distribution	,			UpperCAmelCase_    :       Dict=None	,			UpperCAmelCase_    :       Tuple=None	,			UpperCAmelCase_    :       Tuple=0)	->Optional[int]:



   '''simple docstring'''



   lowerCamelCase__:				Optional[Any]        =1.0 if scale is None else scale
   lowerCamelCase__:				List[Any]        =0.0 if loc is None else loc

   super().__init__(UpperCAmelCase_	,			[AffineTransform(loc=self.loc	,			scale=self.scale	,			event_dim=UpperCAmelCase_)])

 @property
 def 							SCREAMING_SNAKE_CASE_ (self    :       Union[str, Any])	->int:



   '''simple docstring'''



   return self.base_dist.mean * self.scale + self.loc

 @property
 def 							SCREAMING_SNAKE_CASE_ (self    :       int)	->Dict:



   '''simple docstring'''



   return self.base_dist.variance * self.scale**2







 @property
 def 							SCREAMING_SNAKE_CASE_ (self    :       int)	->Optional[Any]:



   '''simple docstring'''



   return self.variance.sqrt()



class        _SCREAMING_SNAKE_CASE						(     nn.Module   ):







 '''simple docstring'''

 def __init__(self    :       Tuple	,			UpperCAmelCase_    :       int	,			UpperCAmelCase_    :       Dict[str, int]	,			UpperCAmelCase_    :       Callable[..., Tuple[torch.Tensor]]	,			**UpperCAmelCase_    :       Dict)	->None:



   '''simple docstring'''



   super().__init__(**UpperCAmelCase_)
   lowerCamelCase__:				Tuple        =args_dim
   lowerCamelCase__:				Any        =nn.ModuleList([nn.Linear(UpperCAmelCase_	,			UpperCAmelCase_) for dim in args_dim.values()])
   lowerCamelCase__:				Any        =domain_map







 def 							SCREAMING_SNAKE_CASE_ (self    :       Optional[Any]	,			UpperCAmelCase_    :       torch.Tensor)	->Tuple[torch.Tensor]:



   '''simple docstring'''



   lowerCamelCase__:				Any        =[proj(UpperCAmelCase_) for proj in self.proj]

   return self.domain_map(*UpperCAmelCase_)



class        _SCREAMING_SNAKE_CASE						(     nn.Module   ):







 '''simple docstring'''

 def __init__(self    :       str	,			UpperCAmelCase_    :       Optional[Any])	->Optional[int]:



   '''simple docstring'''



   super().__init__()
   lowerCamelCase__:				Any        =function







 def 							SCREAMING_SNAKE_CASE_ (self    :       Dict	,			UpperCAmelCase_    :       Optional[Any]	,			*UpperCAmelCase_    :       Union[str, Any])	->Union[str, Any]:



   '''simple docstring'''



   return self.function(UpperCAmelCase_	,			*UpperCAmelCase_)



class        _SCREAMING_SNAKE_CASE						:







 '''simple docstring'''
 lowercase_		=      42
 lowercase_		=      42
 lowercase_		=      42

 def __init__(self    :       str	,			UpperCAmelCase_    :       int = 1)	->None:



   '''simple docstring'''



   lowerCamelCase__:				List[str]        =dim
   lowerCamelCase__:				int        ={k: dim * self.args_dim[k] for k in self.args_dim}

 def 							SCREAMING_SNAKE_CASE_ (self    :       int	,			UpperCAmelCase_    :       str)	->Optional[Any]:



   '''simple docstring'''



   if self.dim == 1:
     return self.distribution_class(*UpperCAmelCase_)
   else:
     return Independent(self.distribution_class(*UpperCAmelCase_)	,			1)

 def 							SCREAMING_SNAKE_CASE_ (self    :       Any	,			UpperCAmelCase_    :       List[Any]	,			UpperCAmelCase_    :       Optional[torch.Tensor] = None	,			UpperCAmelCase_    :       Optional[torch.Tensor] = None	,			)	->Distribution:



   '''simple docstring'''



   lowerCamelCase__:				Optional[Any]        =self._base_distribution(UpperCAmelCase_)
   if loc is None and scale is None:
     return distr
   else:
     return AffineTransformed(UpperCAmelCase_	,			loc=UpperCAmelCase_	,			scale=UpperCAmelCase_	,			event_dim=self.event_dim)

 @property
 def 							SCREAMING_SNAKE_CASE_ (self    :       int)	->Tuple:



   '''simple docstring'''



   return () if self.dim == 1 else (self.dim,)

 @property
 def 							SCREAMING_SNAKE_CASE_ (self    :       Optional[Any])	->int:



   '''simple docstring'''



   return len(self.event_shape)

 @property
 def 							SCREAMING_SNAKE_CASE_ (self    :       Tuple)	->float:



   '''simple docstring'''



   return 0.0

 def 							SCREAMING_SNAKE_CASE_ (self    :       str	,			UpperCAmelCase_    :       int)	->nn.Module:



   '''simple docstring'''



   return ParameterProjection(
       in_features=UpperCAmelCase_	,			args_dim=self.args_dim	,			domain_map=LambdaLayer(self.domain_map)	,			)

 def 							SCREAMING_SNAKE_CASE_ (self    :       Any	,			*UpperCAmelCase_    :       torch.Tensor)	->Any:



   '''simple docstring'''



   raise NotImplementedError()







 @staticmethod
 def 							SCREAMING_SNAKE_CASE_ (UpperCAmelCase_    :       torch.Tensor)	->torch.Tensor:



   '''simple docstring'''



   return (x + torch.sqrt(torch.square(UpperCAmelCase_) + 4.0)) / 2.0



class        _SCREAMING_SNAKE_CASE						(     __SCREAMING_SNAKE_CASE   ):







 '''simple docstring'''
 lowercase_		=      {"df": 1, "loc": 1, "scale": 1}
 lowercase_		=      StudentT

 @classmethod
 def 							SCREAMING_SNAKE_CASE_ (cls    :       Optional[int]	,			UpperCAmelCase_    :       torch.Tensor	,			UpperCAmelCase_    :       torch.Tensor	,			UpperCAmelCase_    :       torch.Tensor)	->Dict:



   '''simple docstring'''



   lowerCamelCase__:				str        =cls.squareplus(UpperCAmelCase_).clamp_min(torch.finfo(scale.dtype).eps)
   lowerCamelCase__:				Any        =2.0 + cls.squareplus(UpperCAmelCase_)
   return df.squeeze(-1), loc.squeeze(-1), scale.squeeze(-1)



class        _SCREAMING_SNAKE_CASE						(     __SCREAMING_SNAKE_CASE   ):







 '''simple docstring'''
 lowercase_		=      {"loc": 1, "scale": 1}
 lowercase_		=      Normal

 @classmethod
 def 							SCREAMING_SNAKE_CASE_ (cls    :       Optional[Any]	,			UpperCAmelCase_    :       torch.Tensor	,			UpperCAmelCase_    :       torch.Tensor)	->Optional[Any]:



   '''simple docstring'''



   lowerCamelCase__:				List[Any]        =cls.squareplus(UpperCAmelCase_).clamp_min(torch.finfo(scale.dtype).eps)
   return loc.squeeze(-1), scale.squeeze(-1)



class        _SCREAMING_SNAKE_CASE						(     __SCREAMING_SNAKE_CASE   ):







 '''simple docstring'''
 lowercase_		=      {"total_count": 1, "logits": 1}
 lowercase_		=      NegativeBinomial

 @classmethod
 def 							SCREAMING_SNAKE_CASE_ (cls    :       Optional[int]	,			UpperCAmelCase_    :       torch.Tensor	,			UpperCAmelCase_    :       torch.Tensor)	->List[str]:



   '''simple docstring'''



   lowerCamelCase__:				Tuple        =cls.squareplus(UpperCAmelCase_)
   return total_count.squeeze(-1), logits.squeeze(-1)

 def 							SCREAMING_SNAKE_CASE_ (self    :       Tuple	,			UpperCAmelCase_    :       Tuple)	->Distribution:



   '''simple docstring'''



   lowerCamelCase__  , lowerCamelCase__:				Any        =distr_args
   if self.dim == 1:
     return self.distribution_class(total_count=UpperCAmelCase_	,			logits=UpperCAmelCase_)
   else:
     return Independent(self.distribution_class(total_count=UpperCAmelCase_	,			logits=UpperCAmelCase_)	,			1)







 def 							SCREAMING_SNAKE_CASE_ (self    :       Union[str, Any]	,			UpperCAmelCase_    :       Optional[Any]	,			UpperCAmelCase_    :       Optional[torch.Tensor] = None	,			UpperCAmelCase_    :       Optional[torch.Tensor] = None)	->Distribution:



   '''simple docstring'''



   lowerCamelCase__  , lowerCamelCase__:				Union[str, Any]        =distr_args

   if scale is not None:
     # See scaling property of Gamma.
     logits += scale.log()

   return self._base_distribution((total_count, logits))

