






import gc
import unittest

import torch
from transformers import CLIPTextConfig, CLIPTextModel, CLIPTextModelWithProjection, CLIPTokenizer

from diffusers import (
    AutoencoderKL,
    DDIMScheduler,
    DDPMScheduler,
    PriorTransformer,
    StableUnCLIPPipeline,
    UNetaDConditionModel,
)
from diffusers.pipelines.stable_diffusion.stable_unclip_image_normalizer import StableUnCLIPImageNormalizer
from diffusers.utils.testing_utils import enable_full_determinism, load_numpy, require_torch_gpu, slow, torch_device

from ..pipeline_params import TEXT_TO_IMAGE_BATCH_PARAMS, TEXT_TO_IMAGE_IMAGE_PARAMS, TEXT_TO_IMAGE_PARAMS
from ..test_pipelines_common import (
    PipelineKarrasSchedulerTesterMixin,
    PipelineLatentTesterMixin,
    PipelineTesterMixin,
    assert_mean_pixel_difference,
)


enable_full_determinism()
class 						lowerCamelCase_				(      __lowerCAmelCase			,				__lowerCAmelCase			,				__lowerCAmelCase			,				unittest.TestCase ):
				a__					       =  StableUnCLIPPipeline
				a__					       =  TEXT_TO_IMAGE_PARAMS
				a__					       =  TEXT_TO_IMAGE_BATCH_PARAMS
				a__					       =  TEXT_TO_IMAGE_IMAGE_PARAMS
				a__					       =  TEXT_TO_IMAGE_IMAGE_PARAMS

				# TODO(will) Expected attn_bias.stride(1) == 0 to be true, but got false
				a__					       =  False







				def 			A	(						self     ):
								"""simple docstring"""
								__magic_name__		:Dict        =      3_2
								__magic_name__		:Any        =      embedder_hidden_size

								# prior components

								torch.manual_seed(0     )
								__magic_name__		:Optional[int]        =      CLIPTokenizer.from_pretrained('''hf-internal-testing/tiny-random-clip'''     )

								torch.manual_seed(0     )
								__magic_name__		:int        =      CLIPTextModelWithProjection(
								    CLIPTextConfig(
								        bos_token_id=0       ,					eos_token_id=2       ,					hidden_size=lowerCamelCase__       ,					projection_dim=lowerCamelCase__       ,					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       ,					)     )

								torch.manual_seed(0     )
								__magic_name__		:List[str]        =      PriorTransformer(
								    num_attention_heads=2       ,					attention_head_dim=1_2       ,					embedding_dim=lowerCamelCase__       ,					num_layers=1       ,					)

								torch.manual_seed(0     )
								__magic_name__		:Any        =      DDPMScheduler(
								    variance_type='''fixed_small_log'''       ,					prediction_type='''sample'''       ,					num_train_timesteps=1_0_0_0       ,					clip_sample=lowerCamelCase__       ,					clip_sample_range=5.0       ,					beta_schedule='''squaredcos_cap_v2'''       ,					)

								# regular denoising components

								torch.manual_seed(0     )
								__magic_name__		:List[Any]        =      StableUnCLIPImageNormalizer(embedding_dim=lowerCamelCase__     )
								__magic_name__		:List[str]        =      DDPMScheduler(beta_schedule='''squaredcos_cap_v2'''     )

								torch.manual_seed(0     )
								__magic_name__		:Any        =      CLIPTokenizer.from_pretrained('''hf-internal-testing/tiny-random-clip'''     )

								torch.manual_seed(0     )
								__magic_name__		:Dict        =      CLIPTextModel(
								    CLIPTextConfig(
								        bos_token_id=0       ,					eos_token_id=2       ,					hidden_size=lowerCamelCase__       ,					projection_dim=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       ,					)     )

								torch.manual_seed(0     )
								__magic_name__		:List[Any]        =      UNetaDConditionModel(
								    sample_size=3_2       ,					in_channels=4       ,					out_channels=4       ,					down_block_types=('''CrossAttnDownBlock2D''', '''DownBlock2D''')       ,					up_block_types=('''UpBlock2D''', '''CrossAttnUpBlock2D''')       ,					block_out_channels=(3_2, 6_4)       ,					attention_head_dim=(2, 4)       ,					class_embed_type='''projection'''       ,					projection_class_embeddings_input_dim=embedder_projection_dim * 2       ,					cross_attention_dim=lowerCamelCase__       ,					layers_per_block=1       ,					upcast_attention=lowerCamelCase__       ,					use_linear_projection=lowerCamelCase__       ,					)

								torch.manual_seed(0     )
								__magic_name__		:str        =      DDIMScheduler(
								    beta_schedule='''scaled_linear'''       ,					beta_start=0.00085       ,					beta_end=0.012       ,					prediction_type='''v_prediction'''       ,					set_alpha_to_one=lowerCamelCase__       ,					steps_offset=1       ,					)

								torch.manual_seed(0     )
								__magic_name__		:List[str]        =      AutoencoderKL()

								__magic_name__		:List[str]        =      {
								    # prior components
								    '''prior_tokenizer''': prior_tokenizer,
								    '''prior_text_encoder''': prior_text_encoder,
								    '''prior''': prior,
								    '''prior_scheduler''': prior_scheduler,
								    # image noising components
								    '''image_normalizer''': image_normalizer,
								    '''image_noising_scheduler''': image_noising_scheduler,
								    # regular denoising components
								    '''tokenizer''': tokenizer,
								    '''text_encoder''': text_encoder,
								    '''unet''': unet,
								    '''scheduler''': scheduler,
								    '''vae''': vae,
								}

								return components







				def 			A	(						self       ,					__lowerCAmelCase       ,					__lowerCAmelCase=0     ):
								"""simple docstring"""






								if str(lowerCamelCase__     ).startswith('''mps'''     ):
												__magic_name__		:Union[str, Any]        =      torch.manual_seed(lowerCamelCase__     )
								else:
												__magic_name__		:Tuple        =      torch.Generator(device=lowerCamelCase__     ).manual_seed(lowerCamelCase__     )
								__magic_name__		:List[str]        =      {
								    '''prompt''': '''A painting of a squirrel eating a burger''',
								    '''generator''': generator,
								    '''num_inference_steps''': 2,
								    '''prior_num_inference_steps''': 2,
								    '''output_type''': '''numpy''',
								}
								return inputs







				def 			A	(						self     ):
								"""simple docstring"""
								__magic_name__		:List[str]        =      torch_device == '''cpu'''

								self._test_attention_slicing_forward_pass(test_max_difference=lowerCamelCase__     )







				def 			A	(						self     ):
								"""simple docstring"""
								__magic_name__		:int        =      torch_device in ['''cpu''', '''mps''']

								self._test_inference_batch_single_identical(test_max_difference=lowerCamelCase__     )







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







				def 			A	(						self     ):
								"""simple docstring"""
								super().tearDown()
								gc.collect()
								torch.cuda.empty_cache()







				def 			A	(						self     ):
								"""simple docstring"""
								__magic_name__		:Union[str, Any]        =      load_numpy(
								    '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/stable_unclip/stable_unclip_2_1_l_anime_turtle_fp16.npy'''     )

								__magic_name__		:Any        =      StableUnCLIPPipeline.from_pretrained('''fusing/stable-unclip-2-1-l'''       ,					torch_dtype=torch.floataa     )
								pipe.to(lowerCamelCase__     )
								pipe.set_progress_bar_config(disable=lowerCamelCase__     )
								# stable unclip will oom when integration tests are run on a V100,
								# so turn on memory savings
								pipe.enable_attention_slicing()
								pipe.enable_sequential_cpu_offload()

								__magic_name__		:Dict        =      torch.Generator(device='''cpu'''     ).manual_seed(0     )
								__magic_name__		:Tuple        =      pipe('''anime turle'''       ,					generator=lowerCamelCase__       ,					output_type='''np'''     )

								__magic_name__		:Union[str, Any]        =      output.images[0]

								assert image.shape == (7_6_8, 7_6_8, 3)

								assert_mean_pixel_difference(lowerCamelCase__       ,					lowerCamelCase__     )







				def 			A	(						self     ):
								"""simple docstring"""
								torch.cuda.empty_cache()
								torch.cuda.reset_max_memory_allocated()
								torch.cuda.reset_peak_memory_stats()

								__magic_name__		:Dict        =      StableUnCLIPPipeline.from_pretrained('''fusing/stable-unclip-2-1-l'''       ,					torch_dtype=torch.floataa     )
								__magic_name__		:List[str]        =      pipe.to(lowerCamelCase__     )
								pipe.set_progress_bar_config(disable=lowerCamelCase__     )
								pipe.enable_attention_slicing()
								pipe.enable_sequential_cpu_offload()

								__magic_name__		:Optional[Any]        =      pipe(
								    '''anime turtle'''       ,					prior_num_inference_steps=2       ,					num_inference_steps=2       ,					output_type='''np'''       ,					)

								__magic_name__		:int        =      torch.cuda.max_memory_allocated()
								# make sure that less than 7 GB is allocated
								assert mem_bytes < 7 * 1_0**9














# 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_				:
				a__					       =  42

				# setable values
				a__					       =  42
				a__					       =  42
				a__					       =  None







				@classmethod
				def 			A	(						cls       ,					__lowerCAmelCase       ,					__lowerCAmelCase       ,					__lowerCAmelCase     ):
								"""simple docstring"""
								return cls(common=__lowerCAmelCase       ,					init_noise_sigma=__lowerCAmelCase       ,					timesteps=__lowerCAmelCase     )







@dataclass
class 						lowerCamelCase_				(      lowerCamelCase ):
				a__					       =  42







class 						lowerCamelCase_				(      lowerCamelCase			,				lowerCamelCase ):
				a__					       =  [e.name for e in FlaxKarrasDiffusionSchedulers]

				a__					       =  42







				@property
				def 			A	(						self     ):
								"""simple docstring"""
								return True







				@register_to_config
				def __init__(						self       ,					__lowerCAmelCase = 1_0_0_0       ,					__lowerCAmelCase = 0.0001       ,					__lowerCAmelCase = 0.02       ,					__lowerCAmelCase = "linear"       ,					__lowerCAmelCase = None       ,					__lowerCAmelCase = "fixed_small"       ,					__lowerCAmelCase = True       ,					__lowerCAmelCase = "epsilon"       ,					__lowerCAmelCase = jnp.floataa       ,					):
								"""simple docstring"""
								__magic_name__		:Optional[int]        =      dtype







				def 			A	(						self       ,					__lowerCAmelCase = None     ):
								"""simple docstring"""
								if common is None:
												__magic_name__		:Dict        =      CommonSchedulerState.create(self     )

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

								__magic_name__		:str        =      jnp.arange(0       ,					self.config.num_train_timesteps     ).round()[::-1]

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







				def 			A	(						self       ,					__lowerCAmelCase       ,					__lowerCAmelCase       ,					__lowerCAmelCase = None     ):
								"""simple docstring"""
								return sample







				def 			A	(						self       ,					__lowerCAmelCase       ,					__lowerCAmelCase       ,					__lowerCAmelCase = ()     ):
								"""simple docstring"""
								__magic_name__		:int        =      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
								__magic_name__		:List[Any]        =      (jnp.arange(0       ,					__lowerCAmelCase     ) * step_ratio).round()[::-1]

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







				def 			A	(						self       ,					__lowerCAmelCase       ,					__lowerCAmelCase       ,					__lowerCAmelCase=None       ,					__lowerCAmelCase=None     ):
								"""simple docstring"""
								__magic_name__		:Optional[Any]        =      state.common.alphas_cumprod[t]
								__magic_name__		:Optional[Any]        =      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
								__magic_name__		:Tuple        =      (1 - alpha_prod_t_prev) / (1 - alpha_prod_t) * state.common.betas[t]

								if variance_type is None:
												__magic_name__		:Optional[Any]        =      self.config.variance_type

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

								return variance







				def 			A	(						self       ,					__lowerCAmelCase       ,					__lowerCAmelCase       ,					__lowerCAmelCase       ,					__lowerCAmelCase       ,					__lowerCAmelCase = None       ,					__lowerCAmelCase = True       ,					):
								"""simple docstring"""
								__magic_name__		:List[str]        =      timestep

								if key is None:
												__magic_name__		:Union[str, Any]        =      jax.random.PRNGKey(0     )

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

								# 1. compute alphas, betas
								__magic_name__		:Any        =      state.common.alphas_cumprod[t]
								__magic_name__		:int        =      jnp.where(t > 0       ,					state.common.alphas_cumprod[t - 1]       ,					jnp.array(1.0       ,					dtype=self.dtype     )     )
								__magic_name__		:Optional[int]        =      1 - alpha_prod_t
								__magic_name__		:Union[str, Any]        =      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":
												__magic_name__		:List[Any]        =      (sample - beta_prod_t ** 0.5 * model_output) / alpha_prod_t ** 0.5
								elif self.config.prediction_type == "sample":
												__magic_name__		:Tuple        =      model_output
								elif self.config.prediction_type == "v_prediction":
												__magic_name__		:Dict        =      (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:
												__magic_name__		:Union[str, Any]        =      jnp.clip(__lowerCAmelCase       ,					-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
								__magic_name__		:Union[str, Any]        =      (alpha_prod_t_prev ** 0.5 * state.common.betas[t]) / beta_prod_t
								__magic_name__		:Union[str, 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
								__magic_name__		:Union[str, Any]        =      pred_original_sample_coeff * pred_original_sample + current_sample_coeff * sample

								# 6. Add noise
								def random_variance():
												__magic_name__		:Tuple        =      jax.random.split(__lowerCAmelCase       ,					num=1     )
												__magic_name__		:Dict        =      jax.random.normal(__lowerCAmelCase       ,					shape=model_output.shape       ,					dtype=self.dtype     )
												return (self._get_variance(__lowerCAmelCase       ,					__lowerCAmelCase       ,					predicted_variance=__lowerCAmelCase     ) ** 0.5) * noise

								__magic_name__		:List[str]        =      jnp.where(t > 0       ,					random_variance()       ,					jnp.zeros(model_output.shape       ,					dtype=self.dtype     )     )

								__magic_name__		:int        =      pred_prev_sample + variance

								if not return_dict:
												return (pred_prev_sample, state)

								return FlaxDDPMSchedulerOutput(prev_sample=__lowerCAmelCase       ,					state=__lowerCAmelCase     )







				def 			A	(						self       ,					__lowerCAmelCase       ,					__lowerCAmelCase       ,					__lowerCAmelCase       ,					__lowerCAmelCase       ,					):
								"""simple docstring"""
								return add_noise_common(state.common       ,					__lowerCAmelCase       ,					__lowerCAmelCase       ,					__lowerCAmelCase     )







				def 			A	(						self       ,					__lowerCAmelCase       ,					__lowerCAmelCase       ,					__lowerCAmelCase       ,					__lowerCAmelCase       ,					):
								"""simple docstring"""
								return get_velocity_common(state.common       ,					__lowerCAmelCase       ,					__lowerCAmelCase       ,					__lowerCAmelCase     )







				def __len__(						self     ):
								"""simple docstring"""
								return self.config.num_train_timesteps






