'''simple docstring'''




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


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




class    a__				(       SCREAMING_SNAKE_CASE_			):






     def __init__(			self		:		Optional[int]  ,    a		:		List[str]  ,    a		:		Optional[Any]  ,    a		:		List[Any]  ,    a		:		List[str]  ,    a		:		Any  ,    a		:		Optional[int]  ,    a		:		List[Any]  ,    ):





         """simple docstring"""



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






     def        SCREAMING_SNAKE_CASE__       (			self		:		Optional[int]  ,    a		:		List[Any] = "auto"   ):





         """simple docstring"""



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






     def        SCREAMING_SNAKE_CASE__       (			self		:		Dict   ):





         """simple docstring"""



         self.enable_attention_slicing(snake_case__   )






     @torch.no_grad()
     def __call__(			self		:		Dict  ,    a		:		Any  ,    a		:		Union[str, Any] = 5_12  ,    a		:		int = 5_12  ,    a		:		Any = 50  ,    a		:		Optional[int] = 7.5  ,    a		:		int = None  ,    a		:		str = 1  ,    a		:		Optional[Any] = 0.0  ,    a		:		Any = None  ,    a		:		Any = None  ,    a		:		Optional[int] = "pil"  ,    a		:		Optional[Any] = True  ,    a		:		str = None  ,    a		:		Tuple = 1  ,    a		:		List[Any] = None  ,    **a		:		Optional[int]  ,    ):





         """simple docstring"""



         if isinstance(snake_case__  ,    snake_case__   ):
             __lowerCamelCase          =  1
         elif isinstance(snake_case__  ,    snake_case__   ):
             __lowerCamelCase          =  len(snake_case__   )
         else:
             raise ValueError(f"""`prompt` has to be of type `str` or `list` but is {type(snake_case__   )}"""   )

         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(snake_case__  ,    snake_case__   ) or callback_steps <= 0)
         ):
             raise ValueError(
                 f"""`callback_steps` has to be a positive integer but is {callback_steps} of type"""
                 f""" {type(snake_case__   )}."""   )

         # get prompt text embeddings
         __lowerCamelCase          =  self.tokenizer(
             snake_case__  ,    padding='''max_length'''  ,    max_length=self.tokenizer.model_max_length  ,    return_tensors='''pt'''  ,    )
         __lowerCamelCase          =  text_inputs.input_ids

         if text_input_ids.shape[-1] > self.tokenizer.model_max_length:
             __lowerCamelCase          =  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          =  text_input_ids[:, : self.tokenizer.model_max_length]

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

         # duplicate text embeddings for each generation per prompt, using mps friendly method
         __lowerCamelCase          =  text_embeddings.shape
         __lowerCamelCase          =  text_embeddings.repeat(1  ,    snake_case__  ,    1   )
         __lowerCamelCase          =  text_embeddings.view(bs_embed * num_images_per_prompt  ,    snake_case__  ,    -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          =  guidance_scale > 1.0
         # get unconditional embeddings for classifier free guidance
         if do_classifier_free_guidance:
             __lowerCamelCase          =  42
             if negative_prompt is None:
                 __lowerCamelCase          =  ['']
             elif type(snake_case__   ) is not type(snake_case__   ):
                 raise TypeError(
                     f"""`negative_prompt` should be the same type to `prompt`, but got {type(snake_case__   )} !="""
                     f""" {type(snake_case__   )}."""   )
             elif isinstance(snake_case__  ,    snake_case__   ):
                 __lowerCamelCase          =  [negative_prompt]
             elif batch_size != len(snake_case__   ):
                 raise ValueError(
                     f"""`negative_prompt`: {negative_prompt} has batch size {len(snake_case__   )}, but `prompt`:"""
                     f""" {prompt} has batch size {batch_size}. Please make sure that passed `negative_prompt` matches"""
                     ''' the batch size of `prompt`.'''   )
             else:
                 __lowerCamelCase          =  negative_prompt

             __lowerCamelCase          =  text_input_ids.shape[-1]
             __lowerCamelCase          =  self.tokenizer(
                 snake_case__  ,    padding='''max_length'''  ,    max_length=snake_case__  ,    truncation=snake_case__  ,    return_tensors='''pt'''  ,    )
             __lowerCamelCase          =  self.text_encoder(uncond_input.input_ids.to(self.device   )   )[0]

             # duplicate unconditional embeddings for each generation per prompt, using mps friendly method
             __lowerCamelCase          =  uncond_embeddings.shape[1]
             __lowerCamelCase          =  uncond_embeddings.repeat(snake_case__  ,    snake_case__  ,    1   )
             __lowerCamelCase          =  uncond_embeddings.view(batch_size * num_images_per_prompt  ,    snake_case__  ,    -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          =  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          =  (batch_size * num_images_per_prompt, self.unet.config.in_channels, height // 8, width // 8)
         __lowerCamelCase          =  (batch_size * num_images_per_prompt, self.unet.config.in_channels, 64, 64)
         __lowerCamelCase          =  text_embeddings.dtype
         if latents is None:
             if self.device.type == "mps":
                 # randn does not exist on mps
                 __lowerCamelCase          =  torch.randn(
                     snake_case__  ,    generator=snake_case__  ,    device='''cpu'''  ,    dtype=snake_case__   ).to(self.device   )
                 __lowerCamelCase          =  torch.randn(snake_case__  ,    generator=snake_case__  ,    device='''cpu'''  ,    dtype=snake_case__   ).to(
                     self.device   )
             else:
                 __lowerCamelCase          =  torch.randn(
                     snake_case__  ,    generator=snake_case__  ,    device=self.device  ,    dtype=snake_case__   )
                 __lowerCamelCase          =  torch.randn(snake_case__  ,    generator=snake_case__  ,    device=self.device  ,    dtype=snake_case__   )
         else:
             if latents_reference.shape != latents_shape:
                 raise ValueError(f"""Unexpected latents shape, got {latents.shape}, expected {latents_shape}"""   )
             __lowerCamelCase          =  latents_reference.to(self.device   )
             __lowerCamelCase          =  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          =  (latents_shape[3] - latents_shape_reference[3]) // 2
         __lowerCamelCase          =  (latents_shape[2] - latents_shape_reference[2]) // 2
         __lowerCamelCase          =  latents_shape_reference[3] if dx >= 0 else latents_shape_reference[3] + 2 * dx
         __lowerCamelCase          =  latents_shape_reference[2] if dy >= 0 else latents_shape_reference[2] + 2 * dy
         __lowerCamelCase          =  0 if dx < 0 else dx
         __lowerCamelCase          =  0 if dy < 0 else dy
         __lowerCamelCase          =  max(-dx  ,    0   )
         __lowerCamelCase          =  max(-dy  ,    0   )
         # import pdb
         # pdb.set_trace()
         __lowerCamelCase          =  latents_reference[:, :, dy : dy + h, dx : dx + w]

         # set timesteps
         self.scheduler.set_timesteps(snake_case__   )

         # Some schedulers like PNDM have timesteps as arrays
         # It's more optimized to move all timesteps to correct device beforehand
         __lowerCamelCase          =  self.scheduler.timesteps.to(self.device   )

         # scale the initial noise by the standard deviation required by the scheduler
         __lowerCamelCase          =  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          =  'eta' in set(inspect.signature(self.scheduler.step   ).parameters.keys()   )
         __lowerCamelCase          =  {}
         if accepts_eta:
             __lowerCamelCase          =  eta

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

             # predict the noise residual
             __lowerCamelCase          =  self.unet(snake_case__  ,    snake_case__  ,    encoder_hidden_states=snake_case__   ).sample

             # perform guidance
             if do_classifier_free_guidance:
                 __lowerCamelCase          =  noise_pred.chunk(2   )
                 __lowerCamelCase          =  noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)

             # compute the previous noisy sample x_t -> x_t-1
             __lowerCamelCase          =  self.scheduler.step(snake_case__  ,    snake_case__  ,    snake_case__  ,    **snake_case__   ).prev_sample

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

         __lowerCamelCase          =  1 / 0.1_82_15 * latents
         __lowerCamelCase          =  self.vae.decode(snake_case__   ).sample

         __lowerCamelCase          =  (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          =  image.cpu().permute(0  ,    2  ,    3  ,    1   ).float().numpy()

         if self.safety_checker is not None:
             __lowerCamelCase          =  self.feature_extractor(self.numpy_to_pil(snake_case__   )  ,    return_tensors='''pt'''   ).to(
                 self.device   )
             __lowerCamelCase          =  self.safety_checker(
                 images=snake_case__  ,    clip_input=safety_checker_input.pixel_values.to(text_embeddings.dtype   )   )
         else:
             __lowerCamelCase          =  None

         if output_type == "pil":
             __lowerCamelCase          =  self.numpy_to_pil(snake_case__   )

         if not return_dict:
             return (image, has_nsfw_concept)

         return StableDiffusionPipelineOutput(images=snake_case__  ,    nsfw_content_detected=snake_case__   )

'''simple docstring'''




import torch

from diffusers import DPMSolverSDEScheduler
from diffusers.utils import torch_device
from diffusers.utils.testing_utils import require_torchsde

from .test_schedulers import SchedulerCommonTest




@require_torchsde
class    a__				(       UpperCAmelCase__			):
     lowerCamelCase  :					Dict							=(DPMSolverSDEScheduler,)
     lowerCamelCase  :					List[str]							=1_0






     def        SCREAMING_SNAKE_CASE__       (			self		:		Optional[Any]  ,    **a		:		Optional[int]   ):





         """simple docstring"""



         __lowerCamelCase          =  {
             '''num_train_timesteps''': 11_00,
             '''beta_start''': 0.00_01,
             '''beta_end''': 0.02,
             '''beta_schedule''': '''linear''',
             '''noise_sampler_seed''': 0,
         }

         config.update(**a   )
         return config






     def        SCREAMING_SNAKE_CASE__       (			self		:		Optional[Any]   ):





         """simple docstring"""



         for timesteps in [10, 50, 1_00, 10_00]:
             self.check_over_configs(num_train_timesteps=a   )






     def        SCREAMING_SNAKE_CASE__       (			self		:		str   ):





         """simple docstring"""



         for beta_start, beta_end in zip([0.0_00_01, 0.00_01, 0.0_01]  ,    [0.00_02, 0.0_02, 0.02]   ):
             self.check_over_configs(beta_start=a  ,    beta_end=a   )






     def        SCREAMING_SNAKE_CASE__       (			self		:		List[Any]   ):





         """simple docstring"""



         for schedule in ["linear", "scaled_linear"]:
             self.check_over_configs(beta_schedule=a   )






     def        SCREAMING_SNAKE_CASE__       (			self		:		str   ):





         """simple docstring"""



         for prediction_type in ["epsilon", "v_prediction"]:
             self.check_over_configs(prediction_type=a   )






     def        SCREAMING_SNAKE_CASE__       (			self		:		Tuple   ):





         """simple docstring"""



         __lowerCamelCase          =  self.scheduler_classes[0]
         __lowerCamelCase          =  self.get_scheduler_config()
         __lowerCamelCase          =  scheduler_class(**a   )

         scheduler.set_timesteps(self.num_inference_steps   )

         __lowerCamelCase          =  self.dummy_model()
         __lowerCamelCase          =  self.dummy_sample_deter * scheduler.init_noise_sigma
         __lowerCamelCase          =  sample.to(a   )

         for i, t in enumerate(scheduler.timesteps   ):
             __lowerCamelCase          =  scheduler.scale_model_input(a  ,    a   )

             __lowerCamelCase          =  model(a  ,    a   )

             __lowerCamelCase          =  scheduler.step(a  ,    a  ,    a   )
             __lowerCamelCase          =  output.prev_sample

         __lowerCamelCase          =  torch.sum(torch.abs(a   )   )
         __lowerCamelCase          =  torch.mean(torch.abs(a   )   )

         if torch_device in ["mps"]:
             assert abs(result_sum.item() - 1_67.47_82_10_44_92_18_75   ) < 1e-2
             assert abs(result_mean.item() - 0.21_78_70_59_64_56_52_77   ) < 1e-3
         elif torch_device in ["cuda"]:
             assert abs(result_sum.item() - 1_71.59_35_21_11_81_64_06   ) < 1e-2
             assert abs(result_mean.item() - 0.2_23_42_90_68_92_29_96_52   ) < 1e-3
         else:
             assert abs(result_sum.item() - 1_62.52_38_34_22_85_15_62   ) < 1e-2
             assert abs(result_mean.item() - 0.2_11_61_95_70_85_13_26   ) < 1e-3






     def        SCREAMING_SNAKE_CASE__       (			self		:		List[str]   ):





         """simple docstring"""



         __lowerCamelCase          =  self.scheduler_classes[0]
         __lowerCamelCase          =  self.get_scheduler_config(prediction_type='''v_prediction'''   )
         __lowerCamelCase          =  scheduler_class(**a   )

         scheduler.set_timesteps(self.num_inference_steps   )

         __lowerCamelCase          =  self.dummy_model()
         __lowerCamelCase          =  self.dummy_sample_deter * scheduler.init_noise_sigma
         __lowerCamelCase          =  sample.to(a   )

         for i, t in enumerate(scheduler.timesteps   ):
             __lowerCamelCase          =  scheduler.scale_model_input(a  ,    a   )

             __lowerCamelCase          =  model(a  ,    a   )

             __lowerCamelCase          =  scheduler.step(a  ,    a  ,    a   )
             __lowerCamelCase          =  output.prev_sample

         __lowerCamelCase          =  torch.sum(torch.abs(a   )   )
         __lowerCamelCase          =  torch.mean(torch.abs(a   )   )

         if torch_device in ["mps"]:
             assert abs(result_sum.item() - 1_24.77_14_92_00_43_94_53   ) < 1e-2
             assert abs(result_mean.item() - 0.1_62_26_28_90_14_81_62_84   ) < 1e-3
         elif torch_device in ["cuda"]:
             assert abs(result_sum.item() - 1_28.1_66_33_60_59_57_03   ) < 1e-2
             assert abs(result_mean.item() - 0.1_66_88_32_60_01_16_72_97   ) < 1e-3
         else:
             assert abs(result_sum.item() - 1_19.8_48_75_48_82_81_25   ) < 1e-2
             assert abs(result_mean.item() - 0.15_60_53_06_62_53_66_21   ) < 1e-3






     def        SCREAMING_SNAKE_CASE__       (			self		:		Dict   ):





         """simple docstring"""



         __lowerCamelCase          =  self.scheduler_classes[0]
         __lowerCamelCase          =  self.get_scheduler_config()
         __lowerCamelCase          =  scheduler_class(**a   )

         scheduler.set_timesteps(self.num_inference_steps  ,    device=a   )

         __lowerCamelCase          =  self.dummy_model()
         __lowerCamelCase          =  self.dummy_sample_deter.to(a   ) * scheduler.init_noise_sigma

         for t in scheduler.timesteps:
             __lowerCamelCase          =  scheduler.scale_model_input(a  ,    a   )

             __lowerCamelCase          =  model(a  ,    a   )

             __lowerCamelCase          =  scheduler.step(a  ,    a  ,    a   )
             __lowerCamelCase          =  output.prev_sample

         __lowerCamelCase          =  torch.sum(torch.abs(a   )   )
         __lowerCamelCase          =  torch.mean(torch.abs(a   )   )

         if torch_device in ["mps"]:
             assert abs(result_sum.item() - 1_67.46_95_73_97_46_09_38   ) < 1e-2
             assert abs(result_mean.item() - 0.2_18_05_93_46_07_98_26_35   ) < 1e-3
         elif torch_device in ["cuda"]:
             assert abs(result_sum.item() - 1_71.59_35_36_37_69_53_12   ) < 1e-2
             assert abs(result_mean.item() - 0.2_23_42_90_83_82_41_57_71   ) < 1e-3
         else:
             assert abs(result_sum.item() - 1_62.52_38_34_22_85_15_62   ) < 1e-2
             assert abs(result_mean.item() - 0.2_11_61_95_70_85_13_26   ) < 1e-3






     def        SCREAMING_SNAKE_CASE__       (			self		:		Optional[int]   ):





         """simple docstring"""



         __lowerCamelCase          =  self.scheduler_classes[0]
         __lowerCamelCase          =  self.get_scheduler_config()
         __lowerCamelCase          =  scheduler_class(**a  ,    use_karras_sigmas=a   )

         scheduler.set_timesteps(self.num_inference_steps  ,    device=a   )

         __lowerCamelCase          =  self.dummy_model()
         __lowerCamelCase          =  self.dummy_sample_deter.to(a   ) * scheduler.init_noise_sigma
         __lowerCamelCase          =  sample.to(a   )

         for t in scheduler.timesteps:
             __lowerCamelCase          =  scheduler.scale_model_input(a  ,    a   )

             __lowerCamelCase          =  model(a  ,    a   )

             __lowerCamelCase          =  scheduler.step(a  ,    a  ,    a   )
             __lowerCamelCase          =  output.prev_sample

         __lowerCamelCase          =  torch.sum(torch.abs(a   )   )
         __lowerCamelCase          =  torch.mean(torch.abs(a   )   )

         if torch_device in ["mps"]:
             assert abs(result_sum.item() - 1_76.66_97_41_35_74_21_88   ) < 1e-2
             assert abs(result_mean.item() - 0.2_30_03_87_27_30_98_18_11   ) < 1e-2
         elif torch_device in ["cuda"]:
             assert abs(result_sum.item() - 1_77.63_65_35_64_45_31_25   ) < 1e-2
             assert abs(result_mean.item() - 0.2_30_03_87_27_30_98_18_11   ) < 1e-2
         else:
             assert abs(result_sum.item() - 1_70.3_13_52_23_38_86_72   ) < 1e-2
             assert abs(result_mean.item() - 0.2_30_03_87_27_30_98_18_11   ) < 1e-2
