


'''simple docstring'''






import unittest

import numpy as np
import torch

from diffusers import PNDMPipeline, PNDMScheduler, UNetaDModel
from diffusers.utils.testing_utils import enable_full_determinism, require_torch, slow, torch_device


enable_full_determinism()




class 	a__       (							unittest.TestCase						):





							'''simple docstring'''






							@property
							def 						__SCREAMING_SNAKE_CASE							(	self     ) ->							Optional[Any]:
														torch.manual_seed(0     )
														lowerCAmelCase__      =			UNetaDModel(
														    block_out_channels=(32, 64)      ,					layers_per_block=2      ,					sample_size=32      ,					in_channels=3      ,					out_channels=3      ,					down_block_types=('''DownBlock2D''', '''AttnDownBlock2D''')      ,					up_block_types=('''AttnUpBlock2D''', '''UpBlock2D''')      ,					)
														return model







							def 						__SCREAMING_SNAKE_CASE							(	self     ) ->							List[str]:
														lowerCAmelCase__      =			self.dummy_uncond_unet
														lowerCAmelCase__      =			PNDMScheduler()

														lowerCAmelCase__      =			PNDMPipeline(unet=lowerCamelCase_      ,					scheduler=lowerCamelCase_     )
														pndm.to(lowerCamelCase_     )
														pndm.set_progress_bar_config(disable=lowerCamelCase_     )

														lowerCAmelCase__      =			torch.manual_seed(0     )
														lowerCAmelCase__      =			pndm(generator=lowerCamelCase_      ,					num_inference_steps=20      ,					output_type='''numpy'''     ).images

														lowerCAmelCase__      =			torch.manual_seed(0     )
														lowerCAmelCase__      =			pndm(generator=lowerCamelCase_      ,					num_inference_steps=20      ,					output_type='''numpy'''      ,					return_dict=lowerCamelCase_     )[0]

														lowerCAmelCase__      =			image[0, -3:, -3:, -1]
														lowerCAmelCase__      =			image_from_tuple[0, -3:, -3:, -1]

														assert image.shape == (1, 32, 32, 3)
														lowerCAmelCase__      =			np.array([1.0, 1.0, 0.0, 1.0, 0.0, 1.0, 0.0, 0.0, 0.0]     )

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




@slow
@require_torch
class 	a__       (							unittest.TestCase						):





							'''simple docstring'''






							def 						__SCREAMING_SNAKE_CASE							(	self     ) ->							Tuple:
														lowerCAmelCase__      =			'''google/ddpm-cifar10-32'''

														lowerCAmelCase__      =			UNetaDModel.from_pretrained(lowerCamelCase_     )
														lowerCAmelCase__      =			PNDMScheduler()

														lowerCAmelCase__      =			PNDMPipeline(unet=lowerCamelCase_      ,					scheduler=lowerCamelCase_     )
														pndm.to(lowerCamelCase_     )
														pndm.set_progress_bar_config(disable=lowerCamelCase_     )
														lowerCAmelCase__      =			torch.manual_seed(0     )
														lowerCAmelCase__      =			pndm(generator=lowerCamelCase_      ,					output_type='''numpy'''     ).images

														lowerCAmelCase__      =			image[0, -3:, -3:, -1]

														assert image.shape == (1, 32, 32, 3)
														lowerCAmelCase__      =			np.array([0.1_564, 0.14_645, 0.1_406, 0.14_715, 0.12_425, 0.14_045, 0.13_115, 0.12_175, 0.125]     )

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



'''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__       (							a__						):





							'''simple docstring'''






							def __init__(	self      ,					lowerCamelCase_      ,					lowerCamelCase_      ,					lowerCamelCase_      ,					lowerCamelCase_      ,					lowerCamelCase_      ,					lowerCamelCase_      ,					lowerCamelCase_      ,					) ->							Optional[Any]:
														super().__init__()
														self.register_modules(
														    vae=lowerCamelCase_      ,					text_encoder=lowerCamelCase_      ,					tokenizer=lowerCamelCase_      ,					unet=lowerCamelCase_      ,					scheduler=lowerCamelCase_      ,					safety_checker=lowerCamelCase_      ,					feature_extractor=lowerCamelCase_      ,					)







							def 						__SCREAMING_SNAKE_CASE							(	self      ,					lowerCamelCase_ = "auto"     ) ->							Any:
														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(lowerCamelCase_     )







							def 						__SCREAMING_SNAKE_CASE							(	self     ) ->							Optional[Any]:
														self.enable_attention_slicing(lowerCamelCase_     )







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

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

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

														# get prompt text embeddings
														lowerCAmelCase__      =			self.tokenizer(
														    lowerCamelCase_      ,					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__     ,      lowerCAmelCase__     ,      lowerCAmelCase__      =			text_embeddings.shape
														lowerCAmelCase__      =			text_embeddings.repeat(1      ,					lowerCamelCase_      ,					1     )
														lowerCAmelCase__      =			text_embeddings.view(bs_embed * num_images_per_prompt      ,					lowerCamelCase_      ,					-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(lowerCamelCase_     ) is not type(lowerCamelCase_     ):
																												raise TypeError(
																												    F"""`negative_prompt` should be the same type to `prompt`, but got {type(lowerCamelCase_     )} !="""
																												    F""" {type(lowerCamelCase_     )}."""     )
																					elif isinstance(lowerCamelCase_      ,					lowerCamelCase_     ):
																												lowerCAmelCase__      =			[negative_prompt]
																					elif batch_size != len(lowerCamelCase_     ):
																												raise ValueError(
																												    F"""`negative_prompt`: {negative_prompt} has batch size {len(lowerCamelCase_     )}, 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(
																					    lowerCamelCase_      ,					padding='''max_length'''      ,					max_length=lowerCamelCase_      ,					truncation=lowerCamelCase_      ,					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(lowerCamelCase_      ,					lowerCamelCase_      ,					1     )
																					lowerCAmelCase__      =			uncond_embeddings.view(batch_size * num_images_per_prompt      ,					lowerCamelCase_      ,					-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(
																												    lowerCamelCase_      ,					generator=lowerCamelCase_      ,					device='''cpu'''      ,					dtype=lowerCamelCase_     ).to(self.device     )
																												lowerCAmelCase__      =			torch.randn(lowerCamelCase_      ,					generator=lowerCamelCase_      ,					device='''cpu'''      ,					dtype=lowerCamelCase_     ).to(
																												    self.device     )
																					else:
																												lowerCAmelCase__      =			torch.randn(
																												    lowerCamelCase_      ,					generator=lowerCamelCase_      ,					device=self.device      ,					dtype=lowerCamelCase_     )
																												lowerCAmelCase__      =			torch.randn(lowerCamelCase_      ,					generator=lowerCamelCase_      ,					device=self.device      ,					dtype=lowerCamelCase_     )
														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(lowerCamelCase_     )

														# 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(lowerCamelCase_     )     ):
																					# 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(lowerCamelCase_      ,					lowerCamelCase_     )

																					# predict the noise residual
																					lowerCAmelCase__      =			self.unet(lowerCamelCase_      ,					lowerCamelCase_      ,					encoder_hidden_states=lowerCamelCase_     ).sample

																					# perform guidance
																					if do_classifier_free_guidance:
																												lowerCAmelCase__     ,      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(lowerCamelCase_      ,					lowerCamelCase_      ,					lowerCamelCase_      ,					**lowerCamelCase_     ).prev_sample

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

														lowerCAmelCase__      =			1 / 0.18_215 * latents
														lowerCAmelCase__      =			self.vae.decode(lowerCamelCase_     ).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(lowerCamelCase_     )      ,					return_tensors='''pt'''     ).to(
																					    self.device     )
																					lowerCAmelCase__     ,      lowerCAmelCase__      =			self.safety_checker(
																					    images=lowerCamelCase_      ,					clip_input=safety_checker_input.pixel_values.to(text_embeddings.dtype     )     )
														else:
																					lowerCAmelCase__      =			None

														if output_type == "pil":
																					lowerCAmelCase__      =			self.numpy_to_pil(lowerCamelCase_     )

														if not return_dict:
																					return (image, has_nsfw_concept)

														return StableDiffusionPipelineOutput(images=lowerCamelCase_      ,					nsfw_content_detected=lowerCamelCase_     )