





import gc
import random
import unittest

import numpy as np
import torch
from transformers import CLIPTextConfig, CLIPTextModel, CLIPTextModelWithProjection, CLIPTokenizer

from diffusers import (
    AutoencoderKL,
    DiffusionPipeline,
    EulerDiscreteScheduler,
    StableDiffusionXLImgaImgPipeline,
    UNetaDConditionModel,
)
from diffusers.utils import floats_tensor, slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu

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







 '''simple docstring'''
 lowercase_		=			StableDiffusionXLImgaImgPipeline
 lowercase_		=			TEXT_GUIDED_IMAGE_VARIATION_PARAMS - {"height", "width"}
 lowercase_		=			PipelineTesterMixin.required_optional_params - {"latents"}
 lowercase_		=			TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS
 lowercase_		=			IMAGE_TO_IMAGE_IMAGE_PARAMS
 lowercase_		=			IMAGE_TO_IMAGE_IMAGE_PARAMS

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



   '''simple docstring'''



   torch.manual_seed(0)
   lowerCamelCase__:      Tuple			  =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")       , attention_head_dim=(2, 4)       , use_linear_projection=UpperCAmelCase_       , addition_embed_type="text_time"       , addition_time_embed_dim=8       , transformer_layers_per_block=(1, 2)       , projection_class_embeddings_input_dim=80       , cross_attention_dim=64       , )
   lowerCamelCase__:      Tuple			  =EulerDiscreteScheduler(
       beta_start=0.0_0085       , beta_end=0.012       , steps_offset=1       , beta_schedule="scaled_linear"       , timestep_spacing="leading"       , )
   torch.manual_seed(0)
   lowerCamelCase__:      Tuple			  =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       , sample_size=128       , )
   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       , hidden_act="gelu"       , projection_dim=32       , )
   lowerCamelCase__:      Optional[Any]			  =CLIPTextModel(UpperCAmelCase_)
   lowerCamelCase__:      List[str]			  =CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip"       , local_files_only=UpperCAmelCase_)

   lowerCamelCase__:      Union[str, Any]			  =CLIPTextModelWithProjection(UpperCAmelCase_)
   lowerCamelCase__:      Optional[int]			  =CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip"       , local_files_only=UpperCAmelCase_)

   lowerCamelCase__:      int			  ={
       "unet": unet,
       "scheduler": scheduler,
       "vae": vae,
       "text_encoder": text_encoder,
       "tokenizer": tokenizer,
       "text_encoder_2": text_encoder_a,
       "tokenizer_2": tokenizer_a,
       # "safety_checker": None,
       # "feature_extractor": None,
   }
   return components

 def 							SCREAMING_SNAKE_CASE_					(self				:   Tuple       , UpperCAmelCase_				:   Union[str, Any]       , UpperCAmelCase_				:   Optional[int]=0)       ->List[Any]:



   '''simple docstring'''



   lowerCamelCase__:      Optional[Any]			  =floats_tensor((1, 3, 32, 32)       , rng=random.Random(UpperCAmelCase_)).to(UpperCAmelCase_)
   lowerCamelCase__:      Any			  =image / 2 + 0.5
   if str(UpperCAmelCase_).startswith("mps"):
     lowerCamelCase__:      str			  =torch.manual_seed(UpperCAmelCase_)
   else:
     lowerCamelCase__:      List[str]			  =torch.Generator(device=UpperCAmelCase_).manual_seed(UpperCAmelCase_)
   lowerCamelCase__:      Any			  ={
       "prompt": "A painting of a squirrel eating a burger",
       "image": image,
       "generator": generator,
       "num_inference_steps": 2,
       "guidance_scale": 5.0,
       "output_type": "numpy",
       "strength": 0.75,
   }
   return inputs

 def 							SCREAMING_SNAKE_CASE_					(self				:   List[str])       ->Dict:



   '''simple docstring'''



   lowerCamelCase__:      str			  ="cpu"  # ensure determinism for the device-dependent torch.Generator
   lowerCamelCase__:      List[str]			  =self.get_dummy_components()
   lowerCamelCase__:      Union[str, Any]			  =StableDiffusionXLImgaImgPipeline(**UpperCAmelCase_)
   lowerCamelCase__:      Dict			  =sd_pipe.to(UpperCAmelCase_)
   sd_pipe.set_progress_bar_config(disable=UpperCAmelCase_)

   lowerCamelCase__:      Dict			  =self.get_dummy_inputs(UpperCAmelCase_)
   lowerCamelCase__:      Union[str, Any]			  =sd_pipe(**UpperCAmelCase_).images
   lowerCamelCase__:      int			  =image[0, -3:, -3:, -1]

   assert image.shape == (1, 32, 32, 3)

   lowerCamelCase__:      List[Any]			  =np.array([0.4656, 0.4840, 0.4439, 0.6698, 0.5574, 0.4524, 0.5799, 0.5943, 0.5165])

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

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



   '''simple docstring'''



   super().test_attention_slicing_forward_pass(expected_max_diff=3E-3)

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



   '''simple docstring'''



   super().test_inference_batch_single_identical(expected_max_diff=3E-3)

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



   '''simple docstring'''



   pass







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



   '''simple docstring'''



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

   # forward without prompt embeds
   lowerCamelCase__:      int			  =self.get_dummy_inputs(UpperCAmelCase_)
   lowerCamelCase__:      List[Any]			  =3 * ["this is a negative prompt"]
   lowerCamelCase__:      Tuple			  =negative_prompt
   lowerCamelCase__:      int			  =3 * [inputs["prompt"]]

   lowerCamelCase__:      Tuple			  =sd_pipe(**UpperCAmelCase_)
   lowerCamelCase__:      Tuple			  =output.images[0, -3:, -3:, -1]

   # forward with prompt embeds
   lowerCamelCase__:      Union[str, Any]			  =self.get_dummy_inputs(UpperCAmelCase_)
   lowerCamelCase__:      Dict			  =3 * ["this is a negative prompt"]
   lowerCamelCase__:      Any			  =3 * [inputs.pop("prompt")]

   (
       (
       lowerCamelCase__
   )  ,					(
       lowerCamelCase__
   )  ,					(
       lowerCamelCase__
   )  ,					(
       lowerCamelCase__
   )  ,					
   ):      Tuple			  =sd_pipe.encode_prompt(UpperCAmelCase_       , negative_prompt=UpperCAmelCase_)

   lowerCamelCase__:      int			  =sd_pipe(
       **UpperCAmelCase_       , prompt_embeds=UpperCAmelCase_       , negative_prompt_embeds=UpperCAmelCase_       , pooled_prompt_embeds=UpperCAmelCase_       , negative_pooled_prompt_embeds=UpperCAmelCase_       , )
   lowerCamelCase__:      Optional[int]			  =output.images[0, -3:, -3:, -1]

   # make sure that it's equal
   assert np.abs(image_slice_a.flatten() - image_slice_a.flatten()).max() < 1E-4



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







 '''simple docstring'''

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



   '''simple docstring'''



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

 def 							SCREAMING_SNAKE_CASE_					(self				:   Optional[int]       , UpperCAmelCase_				:   List[Any]       , UpperCAmelCase_				:   Optional[Any]="cpu"       , UpperCAmelCase_				:   Optional[int]=torch.floataa       , UpperCAmelCase_				:   Any=0)       ->Tuple:



   '''simple docstring'''



   lowerCamelCase__:      Optional[int]			  =torch.Generator(device=UpperCAmelCase_).manual_seed(UpperCAmelCase_)
   lowerCamelCase__:      Union[str, Any]			  =np.random.RandomState(UpperCAmelCase_).standard_normal((1, 4, 64, 64))
   lowerCamelCase__:      Union[str, Any]			  =torch.from_numpy(UpperCAmelCase_).to(device=UpperCAmelCase_       , dtype=UpperCAmelCase_)
   lowerCamelCase__:      Any			  ={
       "prompt": "a photograph of an astronaut riding a horse",
       "latents": latents,
       "generator": generator,
       "num_inference_steps": 3,
       "guidance_scale": 7.5,
       "output_type": "numpy",
   }
   return inputs







 def 							SCREAMING_SNAKE_CASE_					(self				:   int)       ->List[str]:



   '''simple docstring'''



   lowerCamelCase__:      List[str]			  =DiffusionPipeline.from_pretrained("stabilityai/stable-diffusion-2-base")
   pipe.to(UpperCAmelCase_)
   pipe.set_progress_bar_config(disable=UpperCAmelCase_)

   lowerCamelCase__:      int			  =self.get_inputs(UpperCAmelCase_)
   lowerCamelCase__:      str			  =pipe(**UpperCAmelCase_).images
   lowerCamelCase__:      Tuple			  =image[0, -3:, -3:, -1].flatten()

   assert image.shape == (1, 512, 512, 3)
   lowerCamelCase__:      int			  =np.array([0.4_9493, 0.4_7896, 0.4_0798, 0.5_4214, 0.5_3212, 0.4_8202, 0.4_7656, 0.4_6329, 0.4_8506])
   assert np.abs(image_slice - expected_slice).max() < 7E-3



'''simple docstring'''







# DISCLAIMER: This file is strongly influenced by https://github.com/ermongroup/ddim

from dataclasses import dataclass
from typing import Optional, Tuple, Union

import flax
import jax
import jax.numpy as jnp

from ..configuration_utils import ConfigMixin, register_to_config
from .scheduling_utils_flax import (
    CommonSchedulerState,
    FlaxKarrasDiffusionSchedulers,
    FlaxSchedulerMixin,
    FlaxSchedulerOutput,
    add_noise_common,
    get_velocity_common,
)






@flax.struct.dataclass
class  lowerCAmelCase__   :
					lowerCAmelCase_        =     42

					# setable values
					lowerCAmelCase_        =     42
					lowerCAmelCase_        =     42
					lowerCAmelCase_        =     None






					@classmethod
					def 	_snake_case    (   cls				,							__SCREAMING_SNAKE_CASE				,							__SCREAMING_SNAKE_CASE				,							__SCREAMING_SNAKE_CASE		):


											"""simple docstring"""





											return cls(common=__SCREAMING_SNAKE_CASE				,							init_noise_sigma=__SCREAMING_SNAKE_CASE				,							timesteps=__SCREAMING_SNAKE_CASE		)







@dataclass
class  lowerCAmelCase__   (  lowerCamelCase_     ):
					lowerCAmelCase_        =     42







class  lowerCAmelCase__   (  lowerCamelCase_    ,						lowerCamelCase_     ):
					lowerCAmelCase_        =     [e.name for e in FlaxKarrasDiffusionSchedulers]

					lowerCAmelCase_        =     42






					@property
					def 	_snake_case    (   self		):


											"""simple docstring"""





											return True






					@register_to_config
					def __init__(   self				,							__SCREAMING_SNAKE_CASE = 10_00				,							__SCREAMING_SNAKE_CASE = 0.0_001				,							__SCREAMING_SNAKE_CASE = 0.02				,							__SCREAMING_SNAKE_CASE = "linear"				,							__SCREAMING_SNAKE_CASE = None				,							__SCREAMING_SNAKE_CASE = "fixed_small"				,							__SCREAMING_SNAKE_CASE = True				,							__SCREAMING_SNAKE_CASE = "epsilon"				,							__SCREAMING_SNAKE_CASE = jnp.floataa				,							):


											"""simple docstring"""





											lowercase_	:     Dict													=       dtype






					def 	_snake_case    (   self				,							__SCREAMING_SNAKE_CASE = None		):


											"""simple docstring"""





											if common is None:
																	lowercase_	:     Tuple													=       CommonSchedulerState.create(self		)

											# standard deviation of the initial noise distribution
											lowercase_	:     Union[str, Any]													=       jnp.array(1.0				,							dtype=self.dtype		)

											lowercase_	:     List[Any]													=       jnp.arange(0				,							self.config.num_train_timesteps		).round()[::-1]

											return DDPMSchedulerState.create(
											    common=__SCREAMING_SNAKE_CASE				,							init_noise_sigma=__SCREAMING_SNAKE_CASE				,							timesteps=__SCREAMING_SNAKE_CASE				,							)






					def 	_snake_case    (   self				,							__SCREAMING_SNAKE_CASE				,							__SCREAMING_SNAKE_CASE				,							__SCREAMING_SNAKE_CASE = None		):


											"""simple docstring"""





											return sample






					def 	_snake_case    (   self				,							__SCREAMING_SNAKE_CASE				,							__SCREAMING_SNAKE_CASE				,							__SCREAMING_SNAKE_CASE = ()		):


											"""simple docstring"""





											lowercase_	:     Optional[Any]													=       self.config.num_train_timesteps // num_inference_steps
											# creates integer timesteps by multiplying by ratio
											# rounding to avoid issues when num_inference_step is power of 3
											lowercase_	:     int													=       (jnp.arange(0				,							__SCREAMING_SNAKE_CASE		) * step_ratio).round()[::-1]

											return state.replace(
											    num_inference_steps=__SCREAMING_SNAKE_CASE				,							timesteps=__SCREAMING_SNAKE_CASE				,							)






					def 	_snake_case    (   self				,							__SCREAMING_SNAKE_CASE				,							__SCREAMING_SNAKE_CASE				,							__SCREAMING_SNAKE_CASE=None				,							__SCREAMING_SNAKE_CASE=None		):


											"""simple docstring"""





											lowercase_	:     List[Any]													=       state.common.alphas_cumprod[t]
											lowercase_	:     str													=       jnp.where(t > 0				,							state.common.alphas_cumprod[t - 1]				,							jnp.array(1.0				,							dtype=self.dtype		)		)

											# For t > 0, compute predicted variance βt (see formula (6) and (7) from https://arxiv.org/pdf/2006.11239.pdf)
											# and sample from it to get previous sample
											# x_{t-1} ~ N(pred_prev_sample, variance) == add variance to pred_sample
											lowercase_	:     int													=       (1 - alpha_prod_t_prev) / (1 - alpha_prod_t) * state.common.betas[t]

											if variance_type is None:
																	lowercase_	:     str													=       self.config.variance_type

											# hacks - were probably added for training stability
											if variance_type == "fixed_small":
																	lowercase_	:     int													=       jnp.clip(__SCREAMING_SNAKE_CASE				,							a_min=1E-2_0		)
											# for rl-diffuser https://arxiv.org/abs/2205.09991
											elif variance_type == "fixed_small_log":
																	lowercase_	:     List[str]													=       jnp.log(jnp.clip(__SCREAMING_SNAKE_CASE				,							a_min=1E-2_0		)		)
											elif variance_type == "fixed_large":
																	lowercase_	:     List[Any]													=       state.common.betas[t]
											elif variance_type == "fixed_large_log":
																	# Glide max_log
																	lowercase_	:     List[Any]													=       jnp.log(state.common.betas[t]		)
											elif variance_type == "learned":
																	return predicted_variance
											elif variance_type == "learned_range":
																	lowercase_	:     Optional[Any]													=       variance
																	lowercase_	:     Union[str, Any]													=       state.common.betas[t]
																	lowercase_	:     Union[str, Any]													=       (predicted_variance + 1) / 2
																	lowercase_	:     Any													=       frac * max_log + (1 - frac) * min_log

											return variance






					def 	_snake_case    (   self				,							__SCREAMING_SNAKE_CASE				,							__SCREAMING_SNAKE_CASE				,							__SCREAMING_SNAKE_CASE				,							__SCREAMING_SNAKE_CASE				,							__SCREAMING_SNAKE_CASE = None				,							__SCREAMING_SNAKE_CASE = True				,							):


											"""simple docstring"""





											lowercase_	:     Optional[int]													=       timestep

											if key is None:
																	lowercase_	:     int													=       jax.random.PRNGKey(0		)

											if model_output.shape[1] == sample.shape[1] * 2 and self.config.variance_type in ["learned", "learned_range"]:
																	lowercase_						, lowercase_	:     Optional[Any]													=       jnp.split(__SCREAMING_SNAKE_CASE				,							sample.shape[1]				,							axis=1		)
											else:
																	lowercase_	:     int													=       None

											# 1. compute alphas, betas
											lowercase_	:     Any													=       state.common.alphas_cumprod[t]
											lowercase_	:     Optional[int]													=       jnp.where(t > 0				,							state.common.alphas_cumprod[t - 1]				,							jnp.array(1.0				,							dtype=self.dtype		)		)
											lowercase_	:     int													=       1 - alpha_prod_t
											lowercase_	:     str													=       1 - alpha_prod_t_prev

											# 2. compute predicted original sample from predicted noise also called
											# "predicted x_0" of formula (15) from https://arxiv.org/pdf/2006.11239.pdf
											if self.config.prediction_type == "epsilon":
																	lowercase_	:     Tuple													=       (sample - beta_prod_t ** 0.5 * model_output) / alpha_prod_t ** 0.5
											elif self.config.prediction_type == "sample":
																	lowercase_	:     Any													=       model_output
											elif self.config.prediction_type == "v_prediction":
																	lowercase_	:     List[Any]													=       (alpha_prod_t**0.5) * sample - (beta_prod_t**0.5) * model_output
											else:
																	raise ValueError(
																	    F'''prediction_type given as {self.config.prediction_type} must be one of `epsilon`, `sample` '''
																	    ''' for the FlaxDDPMScheduler.'''		)

											# 3. Clip "predicted x_0"
											if self.config.clip_sample:
																	lowercase_	:     Optional[Any]													=       jnp.clip(__SCREAMING_SNAKE_CASE				,							-1				,							1		)

											# 4. Compute coefficients for pred_original_sample x_0 and current sample x_t
											# See formula (7) from https://arxiv.org/pdf/2006.11239.pdf
											lowercase_	:     List[Any]													=       (alpha_prod_t_prev ** 0.5 * state.common.betas[t]) / beta_prod_t
											lowercase_	:     Optional[Any]													=       state.common.alphas[t] ** 0.5 * beta_prod_t_prev / beta_prod_t

											# 5. Compute predicted previous sample µ_t
											# See formula (7) from https://arxiv.org/pdf/2006.11239.pdf
											lowercase_	:     Optional[Any]													=       pred_original_sample_coeff * pred_original_sample + current_sample_coeff * sample

											# 6. Add noise
											def random_variance():
																	lowercase_	:     str													=       jax.random.split(__SCREAMING_SNAKE_CASE				,							num=1		)
																	lowercase_	:     List[Any]													=       jax.random.normal(__SCREAMING_SNAKE_CASE				,							shape=model_output.shape				,							dtype=self.dtype		)
																	return (self._get_variance(__SCREAMING_SNAKE_CASE				,							__SCREAMING_SNAKE_CASE				,							predicted_variance=__SCREAMING_SNAKE_CASE		) ** 0.5) * noise

											lowercase_	:     Optional[Any]													=       jnp.where(t > 0				,							random_variance()				,							jnp.zeros(model_output.shape				,							dtype=self.dtype		)		)

											lowercase_	:     Any													=       pred_prev_sample + variance

											if not return_dict:
																	return (pred_prev_sample, state)

											return FlaxDDPMSchedulerOutput(prev_sample=__SCREAMING_SNAKE_CASE				,							state=__SCREAMING_SNAKE_CASE		)






					def 	_snake_case    (   self				,							__SCREAMING_SNAKE_CASE				,							__SCREAMING_SNAKE_CASE				,							__SCREAMING_SNAKE_CASE				,							__SCREAMING_SNAKE_CASE				,							):


											"""simple docstring"""





											return add_noise_common(state.common				,							__SCREAMING_SNAKE_CASE				,							__SCREAMING_SNAKE_CASE				,							__SCREAMING_SNAKE_CASE		)






					def 	_snake_case    (   self				,							__SCREAMING_SNAKE_CASE				,							__SCREAMING_SNAKE_CASE				,							__SCREAMING_SNAKE_CASE				,							__SCREAMING_SNAKE_CASE				,							):


											"""simple docstring"""





											return get_velocity_common(state.common				,							__SCREAMING_SNAKE_CASE				,							__SCREAMING_SNAKE_CASE				,							__SCREAMING_SNAKE_CASE		)






					def __len__(   self		):


											"""simple docstring"""





											return self.config.num_train_timesteps



