


"""simple docstring"""






import math
import warnings
from dataclasses import dataclass
from typing import List, Optional, Tuple, Union

import numpy as np
import torch
from scipy import integrate

from ..configuration_utils import ConfigMixin, register_to_config
from ..utils import BaseOutput
from .scheduling_utils import KarrasDiffusionSchedulers, SchedulerMixin


@dataclass
# Copied from diffusers.schedulers.scheduling_ddpm.DDPMSchedulerOutput with DDPM->LMSDiscrete
class 			SCREAMING_SNAKE_CASE_  (   _lowercase):




							'''simple docstring'''





							__magic_name__					:     torch.FloatTensor
							__magic_name__					:     Optional[torch.FloatTensor]				  =  None



def 			A__		(					_UpperCAmelCase    :      List[str]				,    _UpperCAmelCase    :      str=0.9_9_9				,    _UpperCAmelCase    :      List[Any]="cosine"				,    )							->	Tuple:

	'''simple docstring'''


	if alpha_transform_type == "cosine":

		def alpha_bar_fn(_UpperCAmelCase    :      List[str]	):
			return math.cos((t + 0.0_0_8) / 1.0_0_8 * math.pi / 2	) ** 2

	elif alpha_transform_type == "exp":

		def alpha_bar_fn(_UpperCAmelCase    :      Optional[int]	):
			return math.exp(t * -1_2.0	)

	else:
		raise ValueError(F"""Unsupported alpha_tranform_type: {alpha_transform_type}"""	)

	snake_case__  :						Any		=      []
	for i in range(_UpperCAmelCase	):
		snake_case__  :						Optional[Any]		=      i / num_diffusion_timesteps
		snake_case__  :						Dict		=      (i + 1) / num_diffusion_timesteps
		betas.append(min(1 - alpha_bar_fn(_UpperCAmelCase	) / alpha_bar_fn(_UpperCAmelCase	)				,    _UpperCAmelCase	)	)
	return torch.tensor(_UpperCAmelCase				,    dtype=torch.floataa	)



class 			SCREAMING_SNAKE_CASE_  (   _lowercase	,      _lowercase):




							'''simple docstring'''





							__magic_name__					:     Any				  =  [e.name for e in KarrasDiffusionSchedulers]
							__magic_name__					:     Optional[Any]				  =  1




							@register_to_config
							def __init__(						self	,					lowerCamelCase__ = 1_000	,					lowerCamelCase__ = 0.00_01	,					lowerCamelCase__ = 0.02	,					lowerCamelCase__ = "linear"	,					lowerCamelCase__ = None	,					lowerCamelCase__ = False	,					lowerCamelCase__ = "epsilon"	,					lowerCamelCase__ = "linspace"	,					lowerCamelCase__ = 0	,					)   ->	Any:







								'''simple docstring'''




								if trained_betas is not None:
									snake_case__  :						List[Any]		=      torch.tensor(lowerCamelCase__	,					dtype=torch.floataa)
								elif beta_schedule == "linear":
									snake_case__  :						Union[str, Any]		=      torch.linspace(lowerCamelCase__	,					lowerCamelCase__	,					lowerCamelCase__	,					dtype=torch.floataa)
								elif beta_schedule == "scaled_linear":
									# this schedule is very specific to the latent diffusion model.
									snake_case__  :						Optional[int]		=      (
									    torch.linspace(beta_start**0.5	,					beta_end**0.5	,					lowerCamelCase__	,					dtype=torch.floataa) ** 2
									)
								elif beta_schedule == "squaredcos_cap_v2":
									# Glide cosine schedule
									snake_case__  :						Tuple		=      betas_for_alpha_bar(lowerCamelCase__)
								else:
									raise NotImplementedError(f"""{beta_schedule} does is not implemented for {self.__class__}""")

								snake_case__  :						Union[str, Any]		=      1.0 - self.betas
								snake_case__  :						Optional[int]		=      torch.cumprod(self.alphas	,					dim=0)

								snake_case__  :						Any		=      np.array(((1 - self.alphas_cumprod) / self.alphas_cumprod) ** 0.5)
								snake_case__  :						Union[str, Any]		=      np.concatenate([sigmas[::-1], [0.0]]).astype(np.floataa)
								snake_case__  :						str		=      torch.from_numpy(lowerCamelCase__)

								# setable values
								snake_case__  :						str		=      None
								snake_case__  :						Dict		=      use_karras_sigmas
								self.set_timesteps(lowerCamelCase__	,					lowerCamelCase__)
								snake_case__  :						Dict		=      []
								snake_case__  :						int		=      False




							@property
							def        UpperCAmelCase (						self)   ->	int:







								'''simple docstring'''




								if self.config.timestep_spacing in ["linspace", "trailing"]:
									return self.sigmas.max()

								return (self.sigmas.max() ** 2 + 1) ** 0.5




							def        UpperCAmelCase (						self	,					lowerCamelCase__	,					lowerCamelCase__)   ->	torch.FloatTensor:







								'''simple docstring'''




								if isinstance(lowerCamelCase__	,					torch.Tensor):
									snake_case__  :						int		=      timestep.to(self.timesteps.device)
								snake_case__  :						Dict		=      (self.timesteps == timestep).nonzero().item()
								snake_case__  :						int		=      self.sigmas[step_index]
								snake_case__  :						Any		=      sample / ((sigma**2 + 1) ** 0.5)
								snake_case__  :						int		=      True
								return sample




							def        UpperCAmelCase (						self	,					lowerCamelCase__	,					lowerCamelCase__	,					lowerCamelCase__)   ->	Optional[int]:







								'''simple docstring'''




								def lms_derivative(lowerCamelCase__):
									snake_case__  :						Any		=      1.0
									for k in range(lowerCamelCase__):
										if current_order == k:
											continue
										prod *= (tau - self.sigmas[t - k]) / (self.sigmas[t - current_order] - self.sigmas[t - k])
									return prod

								snake_case__  :						List[str]		=      integrate.quad(lowerCamelCase__	,					self.sigmas[t]	,					self.sigmas[t + 1]	,					epsrel=1E-4)[0]

								return integrated_coeff




							def        UpperCAmelCase (						self	,					lowerCamelCase__	,					lowerCamelCase__ = None)   ->	Union[str, Any]:







								'''simple docstring'''




								snake_case__  :						Union[str, Any]		=      num_inference_steps

								# "linspace", "leading", "trailing" corresponds to annotation of Table 2. of https://arxiv.org/abs/2305.08891
								if self.config.timestep_spacing == "linspace":
									snake_case__  :						int		=      np.linspace(0	,					self.config.num_train_timesteps - 1	,					lowerCamelCase__	,					dtype=lowerCamelCase__)[
									    ::-1
									].copy()
								elif self.config.timestep_spacing == "leading":
									snake_case__  :						Tuple		=      self.config.num_train_timesteps // self.num_inference_steps
									# creates integer timesteps by multiplying by ratio
									# casting to int to avoid issues when num_inference_step is power of 3
									snake_case__  :						Tuple		=      (np.arange(0	,					lowerCamelCase__) * step_ratio).round()[::-1].copy().astype(lowerCamelCase__)
									timesteps += self.config.steps_offset
								elif self.config.timestep_spacing == "trailing":
									snake_case__  :						Union[str, Any]		=      self.config.num_train_timesteps / self.num_inference_steps
									# creates integer timesteps by multiplying by ratio
									# casting to int to avoid issues when num_inference_step is power of 3
									snake_case__  :						Optional[int]		=      (np.arange(self.config.num_train_timesteps	,					0	,					-step_ratio)).round().copy().astype(lowerCamelCase__)
									timesteps -= 1
								else:
									raise ValueError(
									    f"""{self.config.timestep_spacing} is not supported. Please make sure to choose one of 'linspace', 'leading' or 'trailing'.""")

								snake_case__  :						Optional[int]		=      np.array(((1 - self.alphas_cumprod) / self.alphas_cumprod) ** 0.5)
								snake_case__  :						List[str]		=      np.log(lowerCamelCase__)
								snake_case__  :						int		=      np.interp(lowerCamelCase__	,					np.arange(0	,					len(lowerCamelCase__))	,					lowerCamelCase__)

								if self.use_karras_sigmas:
									snake_case__  :						Optional[Any]		=      self._convert_to_karras(in_sigmas=lowerCamelCase__)
									snake_case__  :						int		=      np.array([self._sigma_to_t(lowerCamelCase__	,					lowerCamelCase__) for sigma in sigmas])

								snake_case__  :						int		=      np.concatenate([sigmas, [0.0]]).astype(np.floataa)

								snake_case__  :						Any		=      torch.from_numpy(lowerCamelCase__).to(device=lowerCamelCase__)
								if str(lowerCamelCase__).startswith("mps"):
									# mps does not support float64
									snake_case__  :						int		=      torch.from_numpy(lowerCamelCase__).to(lowerCamelCase__	,					dtype=torch.floataa)
								else:
									snake_case__  :						Optional[int]		=      torch.from_numpy(lowerCamelCase__).to(device=lowerCamelCase__)

								snake_case__  :						List[Any]		=      []




							def        UpperCAmelCase (						self	,					lowerCamelCase__	,					lowerCamelCase__)   ->	Optional[Any]:







								'''simple docstring'''




								snake_case__  :						List[str]		=      np.log(lowerCamelCase__)

								# get distribution
								snake_case__  :						int		=      log_sigma - log_sigmas[:, np.newaxis]

								# get sigmas range
								snake_case__  :						Any		=      np.cumsum((dists >= 0)	,					axis=0).argmax(axis=0).clip(max=log_sigmas.shape[0] - 2)
								snake_case__  :						Union[str, Any]		=      low_idx + 1

								snake_case__  :						List[str]		=      log_sigmas[low_idx]
								snake_case__  :						int		=      log_sigmas[high_idx]

								# interpolate sigmas
								snake_case__  :						List[str]		=      (low - log_sigma) / (low - high)
								snake_case__  :						int		=      np.clip(lowerCamelCase__	,					0	,					1)

								# transform interpolation to time range
								snake_case__  :						Tuple		=      (1 - w) * low_idx + w * high_idx
								snake_case__  :						str		=      t.reshape(sigma.shape)
								return t




							def        UpperCAmelCase (						self	,					lowerCamelCase__)   ->	torch.FloatTensor:







								'''simple docstring'''




								snake_case__  :						float		=      in_sigmas[-1].item()
								snake_case__  :						float		=      in_sigmas[0].item()

								snake_case__  :						Optional[Any]		=      7.0  # 7.0 is the value used in the paper
								snake_case__  :						Tuple		=      np.linspace(0	,					1	,					self.num_inference_steps)
								snake_case__  :						Union[str, Any]		=      sigma_min ** (1 / rho)
								snake_case__  :						Dict		=      sigma_max ** (1 / rho)
								snake_case__  :						int		=      (max_inv_rho + ramp * (min_inv_rho - max_inv_rho)) ** rho
								return sigmas




							def        UpperCAmelCase (						self	,					lowerCamelCase__	,					lowerCamelCase__	,					lowerCamelCase__	,					lowerCamelCase__ = 4	,					lowerCamelCase__ = True	,					)   ->	Union[LMSDiscreteSchedulerOutput, Tuple]:







								'''simple docstring'''




								if not self.is_scale_input_called:
									warnings.warn(
									    "The `scale_model_input` function should be called before `step` to ensure correct denoising. "
									    "See `StableDiffusionPipeline` for a usage example.")

								if isinstance(lowerCamelCase__	,					torch.Tensor):
									snake_case__  :						List[Any]		=      timestep.to(self.timesteps.device)
								snake_case__  :						Dict		=      (self.timesteps == timestep).nonzero().item()
								snake_case__  :						str		=      self.sigmas[step_index]

								# 1. compute predicted original sample (x_0) from sigma-scaled predicted noise
								if self.config.prediction_type == "epsilon":
									snake_case__  :						Optional[int]		=      sample - sigma * model_output
								elif self.config.prediction_type == "v_prediction":
									# * c_out + input * c_skip
									snake_case__  :						Optional[Any]		=      model_output * (-sigma / (sigma**2 + 1) ** 0.5) + (sample / (sigma**2 + 1))
								elif self.config.prediction_type == "sample":
									snake_case__  :						Tuple		=      model_output
								else:
									raise ValueError(
									    f"""prediction_type given as {self.config.prediction_type} must be one of `epsilon`, or `v_prediction`""")

								# 2. Convert to an ODE derivative
								snake_case__  :						List[Any]		=      (sample - pred_original_sample) / sigma
								self.derivatives.append(lowerCamelCase__)
								if len(self.derivatives) > order:
									self.derivatives.pop(0)

								# 3. Compute linear multistep coefficients
								snake_case__  :						Any		=      min(step_index + 1	,					lowerCamelCase__)
								snake_case__  :						Union[str, Any]		=      [self.get_lms_coefficient(lowerCamelCase__	,					lowerCamelCase__	,					lowerCamelCase__) for curr_order in range(lowerCamelCase__)]

								# 4. Compute previous sample based on the derivatives path
								snake_case__  :						List[Any]		=      sample + sum(
								    coeff * derivative for coeff, derivative in zip(lowerCamelCase__	,					reversed(self.derivatives)))

								if not return_dict:
									return (prev_sample,)

								return LMSDiscreteSchedulerOutput(prev_sample=lowerCamelCase__	,					pred_original_sample=lowerCamelCase__)




							def        UpperCAmelCase (						self	,					lowerCamelCase__	,					lowerCamelCase__	,					lowerCamelCase__	,					)   ->	torch.FloatTensor:







								'''simple docstring'''




								snake_case__  :						Union[str, Any]		=      self.sigmas.to(device=original_samples.device	,					dtype=original_samples.dtype)
								if original_samples.device.type == "mps" and torch.is_floating_point(lowerCamelCase__):
									# mps does not support float64
									snake_case__  :						str		=      self.timesteps.to(original_samples.device	,					dtype=torch.floataa)
									snake_case__  :						int		=      timesteps.to(original_samples.device	,					dtype=torch.floataa)
								else:
									snake_case__  :						List[str]		=      self.timesteps.to(original_samples.device)
									snake_case__  :						Any		=      timesteps.to(original_samples.device)

								snake_case__  :						Dict		=      [(schedule_timesteps == t).nonzero().item() for t in timesteps]

								snake_case__  :						Optional[int]		=      sigmas[step_indices].flatten()
								while len(sigma.shape) < len(original_samples.shape):
									snake_case__  :						Tuple		=      sigma.unsqueeze(-1)

								snake_case__  :						Tuple		=      original_samples + noise * sigma
								return noisy_samples







							def __len__(						self)   ->	Any:







								'''simple docstring'''




								return self.config.num_train_timesteps













"""simple docstring"""
import unittest

import numpy as np
import torch
from torch import nn
from transformers import (
    CLIPImageProcessor,
    CLIPTextConfig,
    CLIPTextModelWithProjection,
    CLIPTokenizer,
    CLIPVisionConfig,
    CLIPVisionModelWithProjection,
)

from diffusers import KandinskyPriorPipeline, PriorTransformer, UnCLIPScheduler
from diffusers.utils import torch_device
from diffusers.utils.testing_utils import enable_full_determinism, skip_mps

from ..test_pipelines_common import PipelineTesterMixin


enable_full_determinism()



class 						a_ (      snake_case_    ,		unittest.TestCase					):







      '''simple docstring'''

      lowerCamelCase__			:							Optional[Any]					=      KandinskyPriorPipeline
      lowerCamelCase__			:							Any					=      ['prompt']
      lowerCamelCase__			:							Any					=      ['prompt', 'negative_prompt']
      lowerCamelCase__			:							str					=      [
          'num_images_per_prompt',
          'generator',
          'num_inference_steps',
          'latents',
          'negative_prompt',
          'guidance_scale',
          'output_type',
          'return_dict',
      ]
      lowerCamelCase__			:							Dict					=      False




      @property
      def      a__				(self   ):

       '''simple docstring'''



       return 3_2




      @property
      def      a__				(self   ):

       '''simple docstring'''



       return 3_2




      @property
      def      a__				(self   ):

       '''simple docstring'''



       return self.time_input_dim




      @property
      def      a__				(self   ):

       '''simple docstring'''



       return self.time_input_dim * 4




      @property
      def      a__				(self   ):

       '''simple docstring'''



       return 1_0_0




      @property
      def      a__				(self   ):

       '''simple docstring'''



       lowerCamelCase__	:     Any								=		CLIPTokenizer.from_pretrained('hf-internal-testing/tiny-random-clip'   )
       return tokenizer




      @property
      def      a__				(self   ):

       '''simple docstring'''



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




      @property
      def      a__				(self   ):

       '''simple docstring'''



       torch.manual_seed(0   )

       lowerCamelCase__	:     Optional[int]								=		{
           'num_attention_heads': 2,
           'attention_head_dim': 1_2,
           'embedding_dim': self.text_embedder_hidden_size,
           'num_layers': 1,
       }

       lowerCamelCase__	:     Optional[int]								=		PriorTransformer(**lowerCamelCase_   )
       # clip_std and clip_mean is initialized to be 0 so PriorTransformer.post_process_latents will always return 0 - set clip_std to be 1 so it won't return 0
       lowerCamelCase__	:     Dict								=		nn.Parameter(torch.ones(model.clip_std.shape   )   )
       return model




      @property
      def      a__				(self   ):

       '''simple docstring'''



       torch.manual_seed(0   )
       lowerCamelCase__	:     List[str]								=		CLIPVisionConfig(
           hidden_size=self.text_embedder_hidden_size,	image_size=2_2_4,	projection_dim=self.text_embedder_hidden_size,	intermediate_size=3_7,	num_attention_heads=4,	num_channels=3,	num_hidden_layers=5,	patch_size=1_4,	)

       lowerCamelCase__	:     int								=		CLIPVisionModelWithProjection(lowerCamelCase_   )
       return model




      @property
      def      a__				(self   ):

       '''simple docstring'''



       lowerCamelCase__	:     Union[str, Any]								=		CLIPImageProcessor(
           crop_size=2_2_4,	do_center_crop=lowerCamelCase_,	do_normalize=lowerCamelCase_,	do_resize=lowerCamelCase_,	image_mean=[0.48_145_466, 0.4_578_275, 0.40_821_073],	image_std=[0.26_862_954, 0.26_130_258, 0.27_577_711],	resample=3,	size=2_2_4,	)

       return image_processor




      def      a__				(self   ):

       '''simple docstring'''



       lowerCamelCase__	:     str								=		self.dummy_prior
       lowerCamelCase__	:     Optional[Any]								=		self.dummy_image_encoder
       lowerCamelCase__	:     Optional[Any]								=		self.dummy_text_encoder
       lowerCamelCase__	:     Union[str, Any]								=		self.dummy_tokenizer
       lowerCamelCase__	:     List[Any]								=		self.dummy_image_processor

       lowerCamelCase__	:     List[Any]								=		UnCLIPScheduler(
           variance_type='fixed_small_log',	prediction_type='sample',	num_train_timesteps=1_0_0_0,	clip_sample=lowerCamelCase_,	clip_sample_range=10.0,	)

       lowerCamelCase__	:     Union[str, Any]								=		{
           'prior': prior,
           'image_encoder': image_encoder,
           'text_encoder': text_encoder,
           'tokenizer': tokenizer,
           'scheduler': scheduler,
           'image_processor': image_processor,
       }

       return components




      def      a__				(self,	lowerCamelCase_,	lowerCamelCase_=0   ):

       '''simple docstring'''



       if str(lowerCamelCase_   ).startswith('mps'   ):
        lowerCamelCase__	:     List[Any]								=		torch.manual_seed(lowerCamelCase_   )
       else:
        lowerCamelCase__	:     Optional[int]								=		torch.Generator(device=lowerCamelCase_   ).manual_seed(lowerCamelCase_   )
       lowerCamelCase__	:     List[Any]								=		{
           'prompt': 'horse',
           'generator': generator,
           'guidance_scale': 4.0,
           'num_inference_steps': 2,
           'output_type': 'np',
       }
       return inputs




      def      a__				(self   ):

       '''simple docstring'''



       lowerCamelCase__	:     List[str]								=		'cpu'

       lowerCamelCase__	:     str								=		self.get_dummy_components()

       lowerCamelCase__	:     Any								=		self.pipeline_class(**lowerCamelCase_   )
       lowerCamelCase__	:     Dict								=		pipe.to(lowerCamelCase_   )

       pipe.set_progress_bar_config(disable=lowerCamelCase_   )

       lowerCamelCase__	:     Tuple								=		pipe(**self.get_dummy_inputs(lowerCamelCase_   )   )
       lowerCamelCase__	:     Dict								=		output.image_embeds

       lowerCamelCase__	:     Optional[Any]								=		pipe(
           **self.get_dummy_inputs(lowerCamelCase_   ),	return_dict=lowerCamelCase_,	)[0]

       lowerCamelCase__	:     Dict								=		image[0, -1_0:]
       lowerCamelCase__	:     Tuple								=		image_from_tuple[0, -1_0:]

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

       lowerCamelCase__	:     Dict								=		np.array(
           [-0.0_532, 1.7_120, 0.3_656, -1.0_852, -0.8_946, -1.1_756, 0.4_348, 0.2_482, 0.5_146, -0.1_156]   )

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




      @skip_mps
      def      a__				(self   ):

       '''simple docstring'''



       lowerCamelCase__	:     Optional[int]								=		torch_device == 'cpu'
       lowerCamelCase__	:     Any								=		True
       lowerCamelCase__	:     List[Any]								=		False

       self._test_inference_batch_single_identical(
           test_max_difference=lowerCamelCase_,	relax_max_difference=lowerCamelCase_,	test_mean_pixel_difference=lowerCamelCase_,	)





      @skip_mps
      def      a__				(self   ):

       '''simple docstring'''



       lowerCamelCase__	:     str								=		torch_device == 'cpu'
       lowerCamelCase__	:     int								=		False

       self._test_attention_slicing_forward_pass(
           test_max_difference=lowerCamelCase_,	test_mean_pixel_difference=lowerCamelCase_,	)

