


import inspect
from typing import Callable, List, Optional, Union

import torch
from transformers import CLIPImageProcessor, CLIPTextModel, CLIPTokenizer

from diffusers import DiffusionPipeline
from diffusers.models import AutoencoderKL, UNetaDConditionModel
from diffusers.pipelines.stable_diffusion import StableDiffusionPipelineOutput
from diffusers.pipelines.stable_diffusion.safety_checker import StableDiffusionSafetyChecker
from diffusers.schedulers import DDIMScheduler, LMSDiscreteScheduler, PNDMScheduler
from diffusers.utils import logging


lowerCamelCase		   =			logging.get_logger(__name__)  # pylint: disable=invalid-name




class A    (  UpperCamelCase_	):






      def __init__(   self		:	int     ,  lowercase_		:	List[str]     ,  lowercase_		:	Dict     ,  lowercase_		:	Optional[int]     ,  lowercase_		:	List[str]     ,  lowercase_		:	Dict     ,  lowercase_		:	List[str]     ,  lowercase_		:	Union[str, Any]     ,  )   ->		Tuple:
             """simple docstring"""




             super().__init__()
             self.register_modules(
                 vae=__A     ,  text_encoder=__A     ,  tokenizer=__A     ,  unet=__A     ,  scheduler=__A     ,  safety_checker=__A     ,  feature_extractor=__A     ,  )






      def 			lowerCamelCase			(   self		:	Tuple     ,  lowercase_		:	List[Any] = "auto"					)   ->		Tuple:
             """simple docstring"""




             if slice_size == "auto":
                    # half the attention head size is usually a good trade-off between
                    # speed and memory
                    _lowerCamelCase			:						Optional[int]            =self.unet.config.attention_head_dim // 2
             self.unet.set_attention_slice(__A					)






      def 			lowerCamelCase			(   self		:	Dict					)   ->		str:
             """simple docstring"""




             self.enable_attention_slicing(__A					)






      @torch.no_grad()
      def __call__(   self		:	Dict     ,  lowercase_		:	List[Any]     ,  lowercase_		:	List[Any] = 512     ,  lowercase_		:	str = 512     ,  lowercase_		:	Optional[int] = 50     ,  lowercase_		:	Dict = 7.5     ,  lowercase_		:	Any = None     ,  lowercase_		:	Optional[int] = 1     ,  lowercase_		:	List[Any] = 0.0     ,  lowercase_		:	int = None     ,  lowercase_		:	List[Any] = None     ,  lowercase_		:	Tuple = "pil"     ,  lowercase_		:	Optional[int] = True     ,  lowercase_		:	Optional[int] = None     ,  lowercase_		:	int = 1     ,  lowercase_		:	List[str] = None     ,  **lowercase_		:	str     ,  )   ->		Tuple:
             """simple docstring"""




             if isinstance(__A     ,  __A					):
                    _lowerCamelCase			:						List[str]            =1
             elif isinstance(__A     ,  __A					):
                    _lowerCamelCase			:						Optional[int]            =len(__A					)
             else:
                    raise ValueError(F'''`prompt` has to be of type `str` or `list` but is {type(__A					)}'''					)

             if height % 8 != 0 or width % 8 != 0:
                    raise ValueError(F'''`height` and `width` have to be divisible by 8 but are {height} and {width}.'''					)

             if (callback_steps is None) or (
                 callback_steps is not None and (not isinstance(__A     ,  __A					) or callback_steps <= 0)
             ):
                    raise ValueError(
                        F'''`callback_steps` has to be a positive integer but is {callback_steps} of type'''
                        F''' {type(__A					)}.'''					)

             # get prompt text embeddings
             _lowerCamelCase			:						Dict            =self.tokenizer(
                 __A     ,  padding='max_length'     ,  max_length=self.tokenizer.model_max_length     ,  return_tensors='pt'     ,  )
             _lowerCamelCase			:						List[Any]            =text_inputs.input_ids

             if text_input_ids.shape[-1] > self.tokenizer.model_max_length:
                    _lowerCamelCase			:						List[str]            =self.tokenizer.batch_decode(text_input_ids[:, self.tokenizer.model_max_length :]					)
                    logger.warning(
                        'The following part of your input was truncated because CLIP can only handle sequences up to'
                        F''' {self.tokenizer.model_max_length} tokens: {removed_text}'''					)
                    _lowerCamelCase			:						Optional[Any]            =text_input_ids[:, : self.tokenizer.model_max_length]

             if text_embeddings is None:
                    _lowerCamelCase			:						Optional[int]            =self.text_encoder(text_input_ids.to(self.device					)					)[0]

             # duplicate text embeddings for each generation per prompt, using mps friendly method
             _lowerCamelCase			:						str            =text_embeddings.shape
             _lowerCamelCase			:						Dict            =text_embeddings.repeat(1     ,  __A     ,  1					)
             _lowerCamelCase			:						Union[str, Any]            =text_embeddings.view(bs_embed * num_images_per_prompt     ,  __A     ,  -1					)

             # here `guidance_scale` is defined analog to the guidance weight `w` of equation (2)
             # of the Imagen paper: https://arxiv.org/pdf/2205.11487.pdf . `guidance_scale = 1`
             # corresponds to doing no classifier free guidance.
             _lowerCamelCase			:						List[Any]            =guidance_scale > 1.0
             # get unconditional embeddings for classifier free guidance
             if do_classifier_free_guidance:
                    _lowerCamelCase			:						List[str]
                    if negative_prompt is None:
                           _lowerCamelCase			:						int            =['''''']
                    elif type(__A					) is not type(__A					):
                           raise TypeError(
                               F'''`negative_prompt` should be the same type to `prompt`, but got {type(__A					)} !='''
                               F''' {type(__A					)}.'''					)
                    elif isinstance(__A     ,  __A					):
                           _lowerCamelCase			:						List[Any]            =[negative_prompt]
                    elif batch_size != len(__A					):
                           raise ValueError(
                               F'''`negative_prompt`: {negative_prompt} has batch size {len(__A					)}, but `prompt`:'''
                               F''' {prompt} has batch size {batch_size}. Please make sure that passed `negative_prompt` matches'''
                               ' the batch size of `prompt`.'					)
                    else:
                           _lowerCamelCase			:						List[Any]            =negative_prompt

                    _lowerCamelCase			:						Tuple            =text_input_ids.shape[-1]
                    _lowerCamelCase			:						List[Any]            =self.tokenizer(
                        __A     ,  padding='max_length'     ,  max_length=__A     ,  truncation=__A     ,  return_tensors='pt'     ,  )
                    _lowerCamelCase			:						Any            =self.text_encoder(uncond_input.input_ids.to(self.device					)					)[0]

                    # duplicate unconditional embeddings for each generation per prompt, using mps friendly method
                    _lowerCamelCase			:						List[Any]            =uncond_embeddings.shape[1]
                    _lowerCamelCase			:						Optional[int]            =uncond_embeddings.repeat(__A     ,  __A     ,  1					)
                    _lowerCamelCase			:						int            =uncond_embeddings.view(batch_size * num_images_per_prompt     ,  __A     ,  -1					)

                    # For classifier free guidance, we need to do two forward passes.
                    # Here we concatenate the unconditional and text embeddings into a single batch
                    # to avoid doing two forward passes
                    _lowerCamelCase			:						Union[str, Any]            =torch.cat([uncond_embeddings, text_embeddings]					)

             # get the initial random noise unless the user supplied it

             # Unlike in other pipelines, latents need to be generated in the target device
             # for 1-to-1 results reproducibility with the CompVis implementation.
             # However this currently doesn't work in `mps`.
             _lowerCamelCase			:						Tuple            =(batch_size * num_images_per_prompt, self.unet.config.in_channels, height // 8, width // 8)
             _lowerCamelCase			:						Optional[int]            =(batch_size * num_images_per_prompt, self.unet.config.in_channels, 64, 64)
             _lowerCamelCase			:						Optional[int]            =text_embeddings.dtype
             if latents is None:
                    if self.device.type == "mps":
                           # randn does not exist on mps
                           _lowerCamelCase			:						Optional[int]            =torch.randn(
                               __A     ,  generator=__A     ,  device='cpu'     ,  dtype=__A					).to(self.device					)
                           _lowerCamelCase			:						Any            =torch.randn(__A     ,  generator=__A     ,  device='cpu'     ,  dtype=__A					).to(
                               self.device					)
                    else:
                           _lowerCamelCase			:						Optional[Any]            =torch.randn(
                               __A     ,  generator=__A     ,  device=self.device     ,  dtype=__A					)
                           _lowerCamelCase			:						List[str]            =torch.randn(__A     ,  generator=__A     ,  device=self.device     ,  dtype=__A					)
             else:
                    if latents_reference.shape != latents_shape:
                           raise ValueError(F'''Unexpected latents shape, got {latents.shape}, expected {latents_shape}'''					)
                    _lowerCamelCase			:						Tuple            =latents_reference.to(self.device					)
                    _lowerCamelCase			:						str            =latents.to(self.device					)

             # This is the key part of the pipeline where we
             # try to ensure that the generated images w/ the same seed
             # but different sizes actually result in similar images
             _lowerCamelCase			:						Optional[Any]            =(latents_shape[3] - latents_shape_reference[3]) // 2
             _lowerCamelCase			:						Optional[Any]            =(latents_shape[2] - latents_shape_reference[2]) // 2
             _lowerCamelCase			:						Optional[Any]            =latents_shape_reference[3] if dx >= 0 else latents_shape_reference[3] + 2 * dx
             _lowerCamelCase			:						List[str]            =latents_shape_reference[2] if dy >= 0 else latents_shape_reference[2] + 2 * dy
             _lowerCamelCase			:						int            =0 if dx < 0 else dx
             _lowerCamelCase			:						Dict            =0 if dy < 0 else dy
             _lowerCamelCase			:						Dict            =max(-dx     ,  0					)
             _lowerCamelCase			:						List[Any]            =max(-dy     ,  0					)
             # import pdb
             # pdb.set_trace()
             _lowerCamelCase			:						Dict            =latents_reference[:, :, dy : dy + h, dx : dx + w]

             # set timesteps
             self.scheduler.set_timesteps(__A					)

             # Some schedulers like PNDM have timesteps as arrays
             # It's more optimized to move all timesteps to correct device beforehand
             _lowerCamelCase			:						List[str]            =self.scheduler.timesteps.to(self.device					)

             # scale the initial noise by the standard deviation required by the scheduler
             _lowerCamelCase			:						Any            =latents * self.scheduler.init_noise_sigma

             # prepare extra kwargs for the scheduler step, since not all schedulers have the same signature
             # eta (η) is only used with the DDIMScheduler, it will be ignored for other schedulers.
             # eta corresponds to η in DDIM paper: https://arxiv.org/abs/2010.02502
             # and should be between [0, 1]
             _lowerCamelCase			:						Union[str, Any]            ='''eta''' in set(inspect.signature(self.scheduler.step					).parameters.keys()					)
             _lowerCamelCase			:						Optional[Any]            ={}
             if accepts_eta:
                    _lowerCamelCase			:						Any            =eta

             for i, t in enumerate(self.progress_bar(__A					)					):
                    # expand the latents if we are doing classifier free guidance
                    _lowerCamelCase			:						Tuple            =torch.cat([latents] * 2					) if do_classifier_free_guidance else latents
                    _lowerCamelCase			:						int            =self.scheduler.scale_model_input(__A     ,  __A					)

                    # predict the noise residual
                    _lowerCamelCase			:						int            =self.unet(__A     ,  __A     ,  encoder_hidden_states=__A					).sample

                    # perform guidance
                    if do_classifier_free_guidance:
                           _lowerCamelCase			:						Optional[int]            =noise_pred.chunk(2					)
                           _lowerCamelCase			:						List[Any]            =noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)

                    # compute the previous noisy sample x_t -> x_t-1
                    _lowerCamelCase			:						str            =self.scheduler.step(__A     ,  __A     ,  __A     ,  **__A					).prev_sample

                    # call the callback, if provided
                    if callback is not None and i % callback_steps == 0:
                           callback(__A     ,  __A     ,  __A					)

             _lowerCamelCase			:						List[Any]            =1 / 0.18215 * latents
             _lowerCamelCase			:						int            =self.vae.decode(__A					).sample

             _lowerCamelCase			:						Tuple            =(image / 2 + 0.5).clamp(0     ,  1					)

             # we always cast to float32 as this does not cause significant overhead and is compatible with bfloat16
             _lowerCamelCase			:						int            =image.cpu().permute(0     ,  2     ,  3     ,  1					).float().numpy()

             if self.safety_checker is not None:
                    _lowerCamelCase			:						Optional[Any]            =self.feature_extractor(self.numpy_to_pil(__A					)     ,  return_tensors='pt'					).to(
                        self.device					)
                    _lowerCamelCase			:						int            =self.safety_checker(
                        images=__A     ,  clip_input=safety_checker_input.pixel_values.to(text_embeddings.dtype					)					)
             else:
                    _lowerCamelCase			:						List[Any]            =None

             if output_type == "pil":
                    _lowerCamelCase			:						Optional[int]            =self.numpy_to_pil(__A					)

             if not return_dict:
                    return (image, has_nsfw_concept)

             return StableDiffusionPipelineOutput(images=__A     ,  nsfw_content_detected=__A					)




import torch

from diffusers import DDIMParallelScheduler

from .test_schedulers import SchedulerCommonTest


class 			lowercase_		(      A       ):
       __lowerCamelCase				       =						(DDIMParallelScheduler,)
       __lowerCamelCase				       =						(("eta", 0.0), ("num_inference_steps", 5_0))


       def 	_snake_case				(    self						,       **__A				)						->    List[str]:
              SCREAMING_SNAKE_CASE_							:				Optional[Any]		={
                  '''num_train_timesteps''': 1_000,
                  '''beta_start''': 0.0_001,
                  '''beta_end''': 0.02,
                  '''beta_schedule''': '''linear''',
                  '''clip_sample''': True,
              }

              config.update(**__A				)
              return config


       def 	_snake_case				(    self						,       **__A				)						->    List[Any]:
              SCREAMING_SNAKE_CASE_							:				int		=self.scheduler_classes[0]
              SCREAMING_SNAKE_CASE_							:				str		=self.get_scheduler_config(**__A				)
              SCREAMING_SNAKE_CASE_							:				List[Any]		=scheduler_class(**__A				)

              SCREAMING_SNAKE_CASE_	,     SCREAMING_SNAKE_CASE_							:				Union[str, Any]		=10, 0.0

              SCREAMING_SNAKE_CASE_							:				List[str]		=self.dummy_model()
              SCREAMING_SNAKE_CASE_							:				Union[str, Any]		=self.dummy_sample_deter

              scheduler.set_timesteps(__A				)

              for t in scheduler.timesteps:
                     SCREAMING_SNAKE_CASE_							:				Optional[int]		=model(__A						,       __A				)
                     SCREAMING_SNAKE_CASE_							:				List[Any]		=scheduler.step(__A						,       __A						,       __A						,       __A				).prev_sample

              return sample


       def 	_snake_case				(    self				)						->    Optional[Any]:
              for timesteps in [100, 500, 1_000]:
                     self.check_over_configs(num_train_timesteps=__A				)


       def 	_snake_case				(    self				)						->    Optional[int]:
              for steps_offset in [0, 1]:
                     self.check_over_configs(steps_offset=__A				)

              SCREAMING_SNAKE_CASE_							:				Dict		=self.scheduler_classes[0]
              SCREAMING_SNAKE_CASE_							:				Optional[Any]		=self.get_scheduler_config(steps_offset=1				)
              SCREAMING_SNAKE_CASE_							:				Dict		=scheduler_class(**__A				)
              scheduler.set_timesteps(5				)
              assert torch.equal(scheduler.timesteps						,       torch.LongTensor([801, 601, 401, 201, 1]				)				)


       def 	_snake_case				(    self				)						->    Dict:
              for beta_start, beta_end in zip([0.0_001, 0.001, 0.01, 0.1]						,       [0.002, 0.02, 0.2, 2]				):
                     self.check_over_configs(beta_start=__A						,       beta_end=__A				)


       def 	_snake_case				(    self				)						->    str:
              for schedule in ["linear", "squaredcos_cap_v2"]:
                     self.check_over_configs(beta_schedule=__A				)


       def 	_snake_case				(    self				)						->    List[str]:
              for prediction_type in ["epsilon", "v_prediction"]:
                     self.check_over_configs(prediction_type=__A				)


       def 	_snake_case				(    self				)						->    List[Any]:
              for clip_sample in [True, False]:
                     self.check_over_configs(clip_sample=__A				)


       def 	_snake_case				(    self				)						->    int:
              for timestep_spacing in ["trailing", "leading"]:
                     self.check_over_configs(timestep_spacing=__A				)


       def 	_snake_case				(    self				)						->    Tuple:
              for rescale_betas_zero_snr in [True, False]:
                     self.check_over_configs(rescale_betas_zero_snr=__A				)


       def 	_snake_case				(    self				)						->    List[Any]:
              self.check_over_configs(thresholding=__A				)
              for threshold in [0.5, 1.0, 2.0]:
                     for prediction_type in ["epsilon", "v_prediction"]:
                            self.check_over_configs(
                                thresholding=__A						,       prediction_type=__A						,       sample_max_value=__A						,       )


       def 	_snake_case				(    self				)						->    Dict:
              for t in [1, 10, 49]:
                     self.check_over_forward(time_step=__A				)


       def 	_snake_case				(    self				)						->    Dict:
              for t, num_inference_steps in zip([1, 10, 50]						,       [10, 50, 500]				):
                     self.check_over_forward(time_step=__A						,       num_inference_steps=__A				)


       def 	_snake_case				(    self				)						->    int:
              for t, eta in zip([1, 10, 49]						,       [0.0, 0.5, 1.0]				):
                     self.check_over_forward(time_step=__A						,       eta=__A				)


       def 	_snake_case				(    self				)						->    Optional[int]:
              SCREAMING_SNAKE_CASE_							:				List[Any]		=self.scheduler_classes[0]
              SCREAMING_SNAKE_CASE_							:				List[str]		=self.get_scheduler_config()
              SCREAMING_SNAKE_CASE_							:				List[str]		=scheduler_class(**__A				)

              assert torch.sum(torch.abs(scheduler._get_variance(0						,       0				) - 0.0				)				) < 1e-5
              assert torch.sum(torch.abs(scheduler._get_variance(420						,       400				) - 0.14_771				)				) < 1e-5
              assert torch.sum(torch.abs(scheduler._get_variance(980						,       960				) - 0.32_460				)				) < 1e-5
              assert torch.sum(torch.abs(scheduler._get_variance(0						,       0				) - 0.0				)				) < 1e-5
              assert torch.sum(torch.abs(scheduler._get_variance(487						,       486				) - 0.00_979				)				) < 1e-5
              assert torch.sum(torch.abs(scheduler._get_variance(999						,       998				) - 0.02				)				) < 1e-5


       def 	_snake_case				(    self				)						->    Union[str, Any]:
              SCREAMING_SNAKE_CASE_							:				List[str]		=self.scheduler_classes[0]
              SCREAMING_SNAKE_CASE_							:				List[str]		=self.get_scheduler_config()
              SCREAMING_SNAKE_CASE_							:				str		=scheduler_class(**__A				)

              SCREAMING_SNAKE_CASE_	,     SCREAMING_SNAKE_CASE_							:				List[str]		=10, 0.0
              scheduler.set_timesteps(__A				)

              SCREAMING_SNAKE_CASE_							:				List[str]		=self.dummy_model()
              SCREAMING_SNAKE_CASE_							:				str		=self.dummy_sample_deter
              SCREAMING_SNAKE_CASE_							:				int		=self.dummy_sample_deter + 0.1
              SCREAMING_SNAKE_CASE_							:				Tuple		=self.dummy_sample_deter - 0.1

              SCREAMING_SNAKE_CASE_							:				Union[str, Any]		=samplea.shape[0]
              SCREAMING_SNAKE_CASE_							:				List[str]		=torch.stack([samplea, samplea, samplea]						,       dim=0				)
              SCREAMING_SNAKE_CASE_							:				Any		=torch.arange(__A				)[0:3, None].repeat(1						,       __A				)

              SCREAMING_SNAKE_CASE_							:				Dict		=model(samples.flatten(0						,       1				)						,       timesteps.flatten(0						,       1				)				)
              SCREAMING_SNAKE_CASE_							:				Optional[int]		=scheduler.batch_step_no_noise(__A						,       timesteps.flatten(0						,       1				)						,       samples.flatten(0						,       1				)						,       __A				)

              SCREAMING_SNAKE_CASE_							:				str		=torch.sum(torch.abs(__A				)				)
              SCREAMING_SNAKE_CASE_							:				int		=torch.mean(torch.abs(__A				)				)

              assert abs(result_sum.item() - 1_147.7_904				) < 1e-2
              assert abs(result_mean.item() - 0.4_982				) < 1e-3


       def 	_snake_case				(    self				)						->    Any:
              SCREAMING_SNAKE_CASE_							:				List[str]		=self.full_loop()

              SCREAMING_SNAKE_CASE_							:				List[Any]		=torch.sum(torch.abs(__A				)				)
              SCREAMING_SNAKE_CASE_							:				Optional[Any]		=torch.mean(torch.abs(__A				)				)

              assert abs(result_sum.item() - 172.0_067				) < 1e-2
              assert abs(result_mean.item() - 0.223_967				) < 1e-3


       def 	_snake_case				(    self				)						->    Optional[Any]:
              SCREAMING_SNAKE_CASE_							:				str		=self.full_loop(prediction_type='''v_prediction'''				)

              SCREAMING_SNAKE_CASE_							:				List[Any]		=torch.sum(torch.abs(__A				)				)
              SCREAMING_SNAKE_CASE_							:				str		=torch.mean(torch.abs(__A				)				)

              assert abs(result_sum.item() - 52.5_302				) < 1e-2
              assert abs(result_mean.item() - 0.0_684				) < 1e-3


       def 	_snake_case				(    self				)						->    Dict:
              # We specify different beta, so that the first alpha is 0.99
              SCREAMING_SNAKE_CASE_							:				Tuple		=self.full_loop(set_alpha_to_one=__A						,       beta_start=0.01				)
              SCREAMING_SNAKE_CASE_							:				Tuple		=torch.sum(torch.abs(__A				)				)
              SCREAMING_SNAKE_CASE_							:				int		=torch.mean(torch.abs(__A				)				)

              assert abs(result_sum.item() - 149.8_295				) < 1e-2
              assert abs(result_mean.item() - 0.1_951				) < 1e-3





       def 	_snake_case				(    self				)						->    Optional[Any]:
              # We specify different beta, so that the first alpha is 0.99
              SCREAMING_SNAKE_CASE_							:				List[str]		=self.full_loop(set_alpha_to_one=__A						,       beta_start=0.01				)
              SCREAMING_SNAKE_CASE_							:				Optional[int]		=torch.sum(torch.abs(__A				)				)
              SCREAMING_SNAKE_CASE_							:				Tuple		=torch.mean(torch.abs(__A				)				)

              assert abs(result_sum.item() - 149.0_784				) < 1e-2
              assert abs(result_mean.item() - 0.1_941				) < 1e-3
