


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 lowerCAmelCase_						(			_UpperCAmelCase	):

 '''simple docstring'''

 def __init__(     self		,				_UpperCAmelCase		,				_UpperCAmelCase		,				_UpperCAmelCase		,				_UpperCAmelCase		,				_UpperCAmelCase		,				_UpperCAmelCase		,				_UpperCAmelCase		,				):
       super().__init__()
       self.register_modules(
           vae=SCREAMING_SNAKE_CASE_		,				text_encoder=SCREAMING_SNAKE_CASE_		,				tokenizer=SCREAMING_SNAKE_CASE_		,				unet=SCREAMING_SNAKE_CASE_		,				scheduler=SCREAMING_SNAKE_CASE_		,				safety_checker=SCREAMING_SNAKE_CASE_		,				feature_extractor=SCREAMING_SNAKE_CASE_		,				)



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



 def 						UpperCamelCase__					(     self ):
       self.enable_attention_slicing(SCREAMING_SNAKE_CASE_ )





 @torch.no_grad()
 def __call__(     self		,				_UpperCAmelCase		,				_UpperCAmelCase = 5_12		,				_UpperCAmelCase = 5_12		,				_UpperCAmelCase = 50		,				_UpperCAmelCase = 7.5		,				_UpperCAmelCase = None		,				_UpperCAmelCase = 1		,				_UpperCAmelCase = 0.0		,				_UpperCAmelCase = None		,				_UpperCAmelCase = None		,				_UpperCAmelCase = "pil"		,				_UpperCAmelCase = True		,				_UpperCAmelCase = None		,				_UpperCAmelCase = 1		,				_UpperCAmelCase = None		,				**_UpperCAmelCase		,				):

       if isinstance(SCREAMING_SNAKE_CASE_		,				SCREAMING_SNAKE_CASE_ ):
             snake_case_	=				1
       elif isinstance(SCREAMING_SNAKE_CASE_		,				SCREAMING_SNAKE_CASE_ ):
             snake_case_	=				len(SCREAMING_SNAKE_CASE_ )
       else:
             raise ValueError(F'''`prompt` has to be of type `str` or `list` but is {type(SCREAMING_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(SCREAMING_SNAKE_CASE_		,				SCREAMING_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(SCREAMING_SNAKE_CASE_ )}.''' )

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

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

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

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

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

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

       # set timesteps
       self.scheduler.set_timesteps(SCREAMING_SNAKE_CASE_ )

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

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

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

             # predict the noise residual
             snake_case_	=				self.unet(SCREAMING_SNAKE_CASE_		,				SCREAMING_SNAKE_CASE_		,				encoder_hidden_states=SCREAMING_SNAKE_CASE_ ).sample

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

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

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

       snake_case_	=				1 / 0.18_215 * latents
       snake_case_	=				self.vae.decode(SCREAMING_SNAKE_CASE_ ).sample

       snake_case_	=				(image / 2 + 0.5).clamp(0		,				1 )

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

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

       if output_type == "pil":
             snake_case_	=				self.numpy_to_pil(SCREAMING_SNAKE_CASE_ )

       if not return_dict:
             return (image, has_nsfw_concept)

       return StableDiffusionPipelineOutput(images=SCREAMING_SNAKE_CASE_		,				nsfw_content_detected=SCREAMING_SNAKE_CASE_ )



import unittest

from transformers.utils.backbone_utils import (
    BackboneMixin,
    get_aligned_output_features_output_indices,
    verify_out_features_out_indices,
)







class lowerCAmelCase_						(			unittest.TestCase	):

 '''simple docstring'''

 def 						UpperCamelCase__					(     self ):
       snake_case_	=				['''a''', '''b''', '''c''']

       # Defaults to last layer if both are None
       snake_case_	,   snake_case_	=				get_aligned_output_features_output_indices(_UpperCAmelCase		,				_UpperCAmelCase		,				_UpperCAmelCase )
       self.assertEqual(_UpperCAmelCase		,				['''c'''] )
       self.assertEqual(_UpperCAmelCase		,				[2] )

       # Out indices set to match out features
       snake_case_	,   snake_case_	=				get_aligned_output_features_output_indices(['''a''', '''c''']		,				_UpperCAmelCase		,				_UpperCAmelCase )
       self.assertEqual(_UpperCAmelCase		,				['''a''', '''c'''] )
       self.assertEqual(_UpperCAmelCase		,				[0, 2] )

       # Out features set to match out indices
       snake_case_	,   snake_case_	=				get_aligned_output_features_output_indices(_UpperCAmelCase		,				[0, 2]		,				_UpperCAmelCase )
       self.assertEqual(_UpperCAmelCase		,				['''a''', '''c'''] )
       self.assertEqual(_UpperCAmelCase		,				[0, 2] )

       # Out features selected from negative indices
       snake_case_	,   snake_case_	=				get_aligned_output_features_output_indices(_UpperCAmelCase		,				[-3, -1]		,				_UpperCAmelCase )
       self.assertEqual(_UpperCAmelCase		,				['''a''', '''c'''] )
       self.assertEqual(_UpperCAmelCase		,				[-3, -1] )



 def 						UpperCamelCase__					(     self ):
       # Stage names must be set
       with self.assertRaises(_UpperCAmelCase ):
             verify_out_features_out_indices(['''a''', '''b''']		,				(0, 1)		,				_UpperCAmelCase )

       # Out features must be a list
       with self.assertRaises(_UpperCAmelCase ):
             verify_out_features_out_indices(('''a''', '''b''')		,				(0, 1)		,				['''a''', '''b'''] )

       # Out features must be a subset of stage names
       with self.assertRaises(_UpperCAmelCase ):
             verify_out_features_out_indices(['''a''', '''b''']		,				(0, 1)		,				['''a'''] )

       # Out indices must be a list or tuple
       with self.assertRaises(_UpperCAmelCase ):
             verify_out_features_out_indices(_UpperCAmelCase		,				0		,				['''a''', '''b'''] )

       # Out indices must be a subset of stage names
       with self.assertRaises(_UpperCAmelCase ):
             verify_out_features_out_indices(_UpperCAmelCase		,				(0, 1)		,				['''a'''] )

       # Out features and out indices must be the same length
       with self.assertRaises(_UpperCAmelCase ):
             verify_out_features_out_indices(['''a''', '''b''']		,				(0,)		,				['''a''', '''b''', '''c'''] )

       # Out features should match out indices
       with self.assertRaises(_UpperCAmelCase ):
             verify_out_features_out_indices(['''a''', '''b''']		,				(0, 2)		,				['''a''', '''b''', '''c'''] )

       # Out features and out indices should be in order
       with self.assertRaises(_UpperCAmelCase ):
             verify_out_features_out_indices(['''b''', '''a''']		,				(0, 1)		,				['''a''', '''b'''] )

       # Check passes with valid inputs
       verify_out_features_out_indices(['''a''', '''b''', '''d''']		,				(0, 1, -1)		,				['''a''', '''b''', '''c''', '''d'''] )





 def 						UpperCamelCase__					(     self ):
       snake_case_	=				BackboneMixin()

       snake_case_	=				['''a''', '''b''', '''c''']
       snake_case_	=				['''a''', '''c''']
       snake_case_	=				[0, 2]

       # Check that the output features and indices are set correctly
       self.assertEqual(backbone.out_features		,				['''a''', '''c'''] )
       self.assertEqual(backbone.out_indices		,				[0, 2] )

       # Check out features and indices are updated correctly
       snake_case_	=				['''a''', '''b''']
       self.assertEqual(backbone.out_features		,				['''a''', '''b'''] )
       self.assertEqual(backbone.out_indices		,				[0, 1] )

       snake_case_	=				[-3, -1]
       self.assertEqual(backbone.out_features		,				['''a''', '''c'''] )
       self.assertEqual(backbone.out_indices		,				[-3, -1] )