






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

     '''simple docstring'''






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







     def   _SCREAMING_SNAKE_CASE     ( self      :    Optional[int] ):






           """simple docstring"""




           torch.manual_seed(0 )
           UpperCamelCase					       =  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   ,						)
           UpperCamelCase					       =  DDIMScheduler(
               beta_start=0.0_0_0_8_5   ,						beta_end=0.0_1_2   ,						beta_schedule='scaled_linear'   ,						num_train_timesteps=1000   ,						clip_sample=lowerCamelCase__   ,						set_alpha_to_one=lowerCamelCase__   ,						)
           torch.manual_seed(0 )
           UpperCamelCase					       =  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 )
           UpperCamelCase					       =  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=1000   ,						)
           UpperCamelCase					       =  CLIPTextModel(lowerCamelCase__ )
           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   _SCREAMING_SNAKE_CASE     ( self      :    List[str]   ,						_SCREAMING_SNAKE_CASE      :    Optional[Any]   ,						_SCREAMING_SNAKE_CASE      :    Dict=0 ):






           """simple docstring"""




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






           """simple docstring"""




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

           UpperCamelCase					       =  self.get_dummy_components()
           UpperCamelCase					       =  CycleDiffusionPipeline(**lowerCamelCase__ )
           UpperCamelCase					       =  pipe.to(lowerCamelCase__ )
           pipe.set_progress_bar_config(disable=lowerCamelCase__ )

           UpperCamelCase					       =  self.get_dummy_inputs(lowerCamelCase__ )
           UpperCamelCase					       =  pipe(**lowerCamelCase__ )
           UpperCamelCase					       =  output.images

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

           assert images.shape == (1, 32, 32, 3)
           UpperCamelCase					       =  np.array([0.4_4_5_9, 0.4_9_4_3, 0.4_5_4_4, 0.6_6_4_3, 0.5_4_7_4, 0.4_3_2_7, 0.5_7_0_1, 0.5_9_5_9, 0.5_1_7_9] )

           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      :    Dict ):






           """simple docstring"""




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

           UpperCamelCase					       =  self.get_dummy_inputs(lowerCamelCase__ )
           UpperCamelCase					       =  pipe(**lowerCamelCase__ )
           UpperCamelCase					       =  output.images

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

           assert images.shape == (1, 32, 32, 3)
           UpperCamelCase					       =  np.array([0.3_5_0_6, 0.4_5_4_3, 0.4_4_6, 0.4_5_7_5, 0.5_1_9_5, 0.4_1_5_5, 0.5_2_7_3, 0.5_1_8, 0.4_1_1_6] )

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







     @skip_mps
     def   _SCREAMING_SNAKE_CASE     ( self      :    Any ):






           """simple docstring"""




           return super().test_save_load_local()







     @unittest.skip('non-deterministic pipeline' )
     def   _SCREAMING_SNAKE_CASE     ( self      :    Union[str, Any] ):






           """simple docstring"""




           return super().test_inference_batch_single_identical()







     @skip_mps
     def   _SCREAMING_SNAKE_CASE     ( self      :    Any ):






           """simple docstring"""




           return super().test_dict_tuple_outputs_equivalent()







     @skip_mps
     def   _SCREAMING_SNAKE_CASE     ( self      :    Any ):






           """simple docstring"""




           return super().test_save_load_optional_components()







     @skip_mps
     def   _SCREAMING_SNAKE_CASE     ( self      :    Any ):






           """simple docstring"""




           return super().test_attention_slicing_forward_pass()





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

     '''simple docstring'''






     def   _SCREAMING_SNAKE_CASE     ( self      :    Optional[int] ):






           """simple docstring"""




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







     def   _SCREAMING_SNAKE_CASE     ( self      :    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_fp16.npy' )
           UpperCamelCase					       =  init_image.resize((512, 512) )

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

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

           UpperCamelCase					       =  "A black colored car"
           UpperCamelCase					       =  "A blue colored car"

           UpperCamelCase					       =  torch.manual_seed(0 )
           UpperCamelCase					       =  pipe(
               prompt=lowerCamelCase__   ,						source_prompt=lowerCamelCase__   ,						image=lowerCamelCase__   ,						num_inference_steps=100   ,						eta=0.1   ,						strength=0.8_5   ,						guidance_scale=3   ,						source_guidance_scale=1   ,						generator=lowerCamelCase__   ,						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   _SCREAMING_SNAKE_CASE     ( self      :    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.npy' )
           UpperCamelCase					       =  init_image.resize((512, 512) )

           UpperCamelCase					       =  "CompVis/stable-diffusion-v1-4"
           UpperCamelCase					       =  DDIMScheduler.from_pretrained(lowerCamelCase__   ,						subfolder='scheduler' )
           UpperCamelCase					       =  CycleDiffusionPipeline.from_pretrained(lowerCamelCase__   ,						scheduler=lowerCamelCase__   ,						safety_checker=lowerCamelCase__ )

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

           UpperCamelCase					       =  "A black colored car"
           UpperCamelCase					       =  "A blue colored car"

           UpperCamelCase					       =  torch.manual_seed(0 )
           UpperCamelCase					       =  pipe(
               prompt=lowerCamelCase__   ,						source_prompt=lowerCamelCase__   ,						image=lowerCamelCase__   ,						num_inference_steps=100   ,						eta=0.1   ,						strength=0.8_5   ,						guidance_scale=3   ,						source_guidance_scale=1   ,						generator=lowerCamelCase__   ,						output_type='np'   ,						)
           UpperCamelCase					       =  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 __UpperCamelCase     (				_a		):






					'''simple docstring'''



					def __init__(				self				,      lowerCamelCase__				,      lowerCamelCase__=None				,      lowerCamelCase__=None				,      lowerCamelCase__=0						):
						UpperCAmelCase__:	Any									= 1.0 if scale is None else scale
						UpperCAmelCase__:	int									= 0.0 if loc is None else loc

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





					@property
					def    _UpperCAmelCase    (				self						):
						return self.base_dist.mean * self.scale + self.loc





					@property
					def    _UpperCAmelCase    (				self						):
						return self.base_dist.variance * self.scale**2





					@property
					def    _UpperCAmelCase    (				self						):
						return self.variance.sqrt()

class __UpperCamelCase     (				nn.Module		):






					'''simple docstring'''



					def __init__(				self				,      lowerCamelCase__				,      lowerCamelCase__				,      lowerCamelCase__				,      **lowerCamelCase__						):
						super().__init__(**lowerCamelCase__						)
						UpperCAmelCase__:	Any									= args_dim
						UpperCAmelCase__:	List[str]									= nn.ModuleList([nn.Linear(lowerCamelCase__				,      lowerCamelCase__						) for dim in args_dim.values()]						)
						UpperCAmelCase__:	Any									= domain_map





					def    _UpperCAmelCase    (				self				,      lowerCamelCase__						):
						UpperCAmelCase__:	str									= [proj(lowerCamelCase__						) for proj in self.proj]

						return self.domain_map(*lowerCamelCase__						)

class __UpperCamelCase     (				nn.Module		):






					'''simple docstring'''



					def __init__(				self				,      lowerCamelCase__						):
						super().__init__()
						UpperCAmelCase__:	str									= function





					def    _UpperCAmelCase    (				self				,      lowerCamelCase__				,      *lowerCamelCase__						):
						return self.function(lowerCamelCase__				,      *lowerCamelCase__						)

class __UpperCamelCase     :






					'''simple docstring'''



					__magic_name__        =   42
					__magic_name__        =   42
					__magic_name__        =   42





					def __init__(				self				,      lowerCamelCase__ = 1						):
						UpperCAmelCase__:	str									= dim
						UpperCAmelCase__:	Any									= {k: dim * self.args_dim[k] for k in self.args_dim}





					def    _UpperCAmelCase    (				self				,      lowerCamelCase__						):
						if self.dim == 1:
							return self.distribution_class(*lowerCamelCase__						)
						else:
							return Independent(self.distribution_class(*lowerCamelCase__						)				,      1						)





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





					@property
					def    _UpperCAmelCase    (				self						):
						return () if self.dim == 1 else (self.dim,)





					@property
					def    _UpperCAmelCase    (				self						):
						return len(self.event_shape						)





					@property
					def    _UpperCAmelCase    (				self						):
						return 0.0





					def    _UpperCAmelCase    (				self				,      lowerCamelCase__						):
						return ParameterProjection(
						    in_features=lowerCamelCase__				,      args_dim=self.args_dim				,      domain_map=LambdaLayer(self.domain_map						)				,      )





					def    _UpperCAmelCase    (				self				,      *lowerCamelCase__						):
						raise NotImplementedError()





					@staticmethod
					def    _UpperCAmelCase    (				lowerCamelCase__						):
						return (x + torch.sqrt(torch.square(lowerCamelCase__						) + 4.0						)) / 2.0

class __UpperCamelCase     (				_a		):






					'''simple docstring'''



					__magic_name__        =   {"df": 1, "loc": 1, "scale": 1}
					__magic_name__        =   StudentT





					@classmethod
					def    _UpperCAmelCase    (				cls				,      lowerCamelCase__				,      lowerCamelCase__				,      lowerCamelCase__						):
						UpperCAmelCase__:	Dict									= cls.squareplus(lowerCamelCase__						).clamp_min(torch.finfo(scale.dtype						).eps						)
						UpperCAmelCase__:	List[Any]									= 2.0 + cls.squareplus(lowerCamelCase__						)
						return df.squeeze(-1						), loc.squeeze(-1						), scale.squeeze(-1						)

class __UpperCamelCase     (				_a		):






					'''simple docstring'''



					__magic_name__        =   {"loc": 1, "scale": 1}
					__magic_name__        =   Normal





					@classmethod
					def    _UpperCAmelCase    (				cls				,      lowerCamelCase__				,      lowerCamelCase__						):
						UpperCAmelCase__:	str									= cls.squareplus(lowerCamelCase__						).clamp_min(torch.finfo(scale.dtype						).eps						)
						return loc.squeeze(-1						), scale.squeeze(-1						)

class __UpperCamelCase     (				_a		):






					'''simple docstring'''



					__magic_name__        =   {"total_count": 1, "logits": 1}
					__magic_name__        =   NegativeBinomial





					@classmethod
					def    _UpperCAmelCase    (				cls				,      lowerCamelCase__				,      lowerCamelCase__						):
						UpperCAmelCase__:	Dict									= cls.squareplus(lowerCamelCase__						)
						return total_count.squeeze(-1						), logits.squeeze(-1						)





					def    _UpperCAmelCase    (				self				,      lowerCamelCase__						):
						UpperCAmelCase__							,						UpperCAmelCase__:	str									= distr_args
						if self.dim == 1:
							return self.distribution_class(total_count=lowerCamelCase__				,      logits=lowerCamelCase__						)
						else:
							return Independent(self.distribution_class(total_count=lowerCamelCase__				,      logits=lowerCamelCase__						)				,      1						)





					def    _UpperCAmelCase    (				self				,      lowerCamelCase__				,      lowerCamelCase__ = None				,      lowerCamelCase__ = None						):
						UpperCAmelCase__							,						UpperCAmelCase__:	List[str]									= distr_args

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

						return self._base_distribution((total_count, logits)						)