


import gc
import random
import tempfile
import unittest

import numpy as np
import torch
from transformers import CLIPTextConfig, CLIPTextModel, CLIPTokenizer

from diffusers import AutoencoderKL, DDIMScheduler, LMSDiscreteScheduler, PNDMScheduler, UNetaDConditionModel
from diffusers.pipelines.stable_diffusion_safe import StableDiffusionPipelineSafe as StableDiffusionPipeline
from diffusers.utils import floats_tensor, nightly, torch_device
from diffusers.utils.testing_utils import require_torch_gpu







class 							lowerCamelCase__						(				unittest.TestCase   ):






			"""simple docstring"""






			def 					snake_case__					(				self  ):

						'''simple docstring'''


						super().tearDown()
						gc.collect()
						torch.cuda.empty_cache()




			@property
			def 					snake_case__					(				self  ):

						'''simple docstring'''


						UpperCamelCase__       =							1
						UpperCamelCase__       =							3
						UpperCamelCase__       =							(32, 32)

						UpperCamelCase__       =							floats_tensor((batch_size, num_channels) + sizes   ,    rng=random.Random(0  )  ).to(snake_case  )
						return image




			@property
			def 					snake_case__					(				self  ):

						'''simple docstring'''


						torch.manual_seed(0  )
						UpperCamelCase__       =							UNetaDConditionModel(
						    block_out_channels=(32, 64)   ,    layers_per_block=2   ,    sample_size=32   ,    in_channels=4   ,    out_channels=4   ,    down_block_types=("DownBlock2D", "CrossAttnDownBlock2D")   ,    up_block_types=("CrossAttnUpBlock2D", "UpBlock2D")   ,    cross_attention_dim=32   ,    )
						return model




			@property
			def 					snake_case__					(				self  ):

						'''simple docstring'''


						torch.manual_seed(0  )
						UpperCamelCase__       =							AutoencoderKL(
						    block_out_channels=[32, 64]   ,    in_channels=3   ,    out_channels=3   ,    down_block_types=["DownEncoderBlock2D", "DownEncoderBlock2D"]   ,    up_block_types=["UpDecoderBlock2D", "UpDecoderBlock2D"]   ,    latent_channels=4   ,    )
						return model




			@property
			def 					snake_case__					(				self  ):

						'''simple docstring'''


						torch.manual_seed(0  )
						UpperCamelCase__       =							CLIPTextConfig(
						    bos_token_id=0   ,    eos_token_id=2   ,    hidden_size=32   ,    intermediate_size=37   ,    layer_norm_eps=1E-0_5   ,    num_attention_heads=4   ,    num_hidden_layers=5   ,    pad_token_id=1   ,    vocab_size=1000   ,    )
						return CLIPTextModel(snake_case  )




			@property
			def 					snake_case__					(				self  ):

						'''simple docstring'''


						def extract(*snake_case   ,    **snake_case  ):







									class 							lowerCamelCase__						:






												"""simple docstring"""






												def __init__(				self  ):

															'''simple docstring'''


															UpperCamelCase__       =							torch.ones([0]  )







												def 					snake_case__					(				self   ,    snake_case  ):

															'''simple docstring'''


															self.pixel_values.to(snake_case  )
															return self



									return Out()

						return extract




			def 					snake_case__					(				self  ):

						'''simple docstring'''


						UpperCamelCase__       =							"cpu"  # ensure determinism for the device-dependent torch.Generator
						UpperCamelCase__       =							self.dummy_cond_unet
						UpperCamelCase__       =							DDIMScheduler(
						    beta_start=0.00085   ,    beta_end=0.012   ,    beta_schedule="scaled_linear"   ,    clip_sample=snake_case   ,    set_alpha_to_one=snake_case   ,    )

						UpperCamelCase__       =							self.dummy_vae
						UpperCamelCase__       =							self.dummy_text_encoder
						UpperCamelCase__       =							CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip"  )

						# make sure here that pndm scheduler skips prk
						UpperCamelCase__       =							StableDiffusionPipeline(
						    unet=snake_case   ,    scheduler=snake_case   ,    vae=snake_case   ,    text_encoder=snake_case   ,    tokenizer=snake_case   ,    safety_checker=snake_case   ,    feature_extractor=self.dummy_extractor   ,    )
						UpperCamelCase__       =							sd_pipe.to(snake_case  )
						sd_pipe.set_progress_bar_config(disable=snake_case  )

						UpperCamelCase__       =							"A painting of a squirrel eating a burger"

						UpperCamelCase__       =							torch.Generator(device=snake_case  ).manual_seed(0  )
						UpperCamelCase__       =							sd_pipe([prompt]   ,    generator=snake_case   ,    guidance_scale=6.0   ,    num_inference_steps=2   ,    output_type="np"  )
						UpperCamelCase__       =							output.images

						UpperCamelCase__       =							torch.Generator(device=snake_case  ).manual_seed(0  )
						UpperCamelCase__       =							sd_pipe(
						    [prompt]   ,    generator=snake_case   ,    guidance_scale=6.0   ,    num_inference_steps=2   ,    output_type="np"   ,    return_dict=snake_case   ,    )[0]

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

						assert image.shape == (1, 64, 64, 3)
						UpperCamelCase__       =							np.array([0.5756, 0.6118, 0.5005, 0.5041, 0.5471, 0.4726, 0.4976, 0.4865, 0.4864]  )

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




			def 					snake_case__					(				self  ):

						'''simple docstring'''


						UpperCamelCase__       =							"cpu"  # ensure determinism for the device-dependent torch.Generator
						UpperCamelCase__       =							self.dummy_cond_unet
						UpperCamelCase__       =							PNDMScheduler(skip_prk_steps=snake_case  )
						UpperCamelCase__       =							self.dummy_vae
						UpperCamelCase__       =							self.dummy_text_encoder
						UpperCamelCase__       =							CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip"  )

						# make sure here that pndm scheduler skips prk
						UpperCamelCase__       =							StableDiffusionPipeline(
						    unet=snake_case   ,    scheduler=snake_case   ,    vae=snake_case   ,    text_encoder=snake_case   ,    tokenizer=snake_case   ,    safety_checker=snake_case   ,    feature_extractor=self.dummy_extractor   ,    )
						UpperCamelCase__       =							sd_pipe.to(snake_case  )
						sd_pipe.set_progress_bar_config(disable=snake_case  )

						UpperCamelCase__       =							"A painting of a squirrel eating a burger"
						UpperCamelCase__       =							torch.Generator(device=snake_case  ).manual_seed(0  )
						UpperCamelCase__       =							sd_pipe([prompt]   ,    generator=snake_case   ,    guidance_scale=6.0   ,    num_inference_steps=2   ,    output_type="np"  )

						UpperCamelCase__       =							output.images

						UpperCamelCase__       =							torch.Generator(device=snake_case  ).manual_seed(0  )
						UpperCamelCase__       =							sd_pipe(
						    [prompt]   ,    generator=snake_case   ,    guidance_scale=6.0   ,    num_inference_steps=2   ,    output_type="np"   ,    return_dict=snake_case   ,    )[0]

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

						assert image.shape == (1, 64, 64, 3)
						UpperCamelCase__       =							np.array([0.5125, 0.5716, 0.4828, 0.5060, 0.5650, 0.4768, 0.5185, 0.4895, 0.4993]  )

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




			def 					snake_case__					(				self  ):

						'''simple docstring'''


						UpperCamelCase__       =							StableDiffusionPipeline.from_pretrained(
						    "hf-internal-testing/tiny-stable-diffusion-lms-pipe"   ,    safety_checker=snake_case  )
						assert isinstance(snake_case   ,    snake_case  )
						assert isinstance(pipe.scheduler   ,    snake_case  )
						assert pipe.safety_checker is None

						UpperCamelCase__       =							pipe("example prompt"   ,    num_inference_steps=2  ).images[0]
						assert image is not None

						# check that there's no error when saving a pipeline with one of the models being None
						with tempfile.TemporaryDirectory() as tmpdirname:
									pipe.save_pretrained(snake_case  )
									UpperCamelCase__       =							StableDiffusionPipeline.from_pretrained(snake_case  )

						# sanity check that the pipeline still works
						assert pipe.safety_checker is None
						UpperCamelCase__       =							pipe("example prompt"   ,    num_inference_steps=2  ).images[0]
						assert image is not None







			@unittest.skipIf(torch_device != "cuda"   ,    "This test requires a GPU"  )
			def 					snake_case__					(				self  ):

						'''simple docstring'''


						UpperCamelCase__       =							self.dummy_cond_unet
						UpperCamelCase__       =							PNDMScheduler(skip_prk_steps=snake_case  )
						UpperCamelCase__       =							self.dummy_vae
						UpperCamelCase__       =							self.dummy_text_encoder
						UpperCamelCase__       =							CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip"  )

						# put models in fp16
						UpperCamelCase__       =							unet.half()
						UpperCamelCase__       =							vae.half()
						UpperCamelCase__       =							bert.half()

						# make sure here that pndm scheduler skips prk
						UpperCamelCase__       =							StableDiffusionPipeline(
						    unet=snake_case   ,    scheduler=snake_case   ,    vae=snake_case   ,    text_encoder=snake_case   ,    tokenizer=snake_case   ,    safety_checker=snake_case   ,    feature_extractor=self.dummy_extractor   ,    )
						UpperCamelCase__       =							sd_pipe.to(snake_case  )
						sd_pipe.set_progress_bar_config(disable=snake_case  )

						UpperCamelCase__       =							"A painting of a squirrel eating a burger"
						UpperCamelCase__       =							sd_pipe([prompt]   ,    num_inference_steps=2   ,    output_type="np"  ).images

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







@nightly
@require_torch_gpu
class 							lowerCamelCase__						(				unittest.TestCase   ):






			"""simple docstring"""






			def 					snake_case__					(				self  ):

						'''simple docstring'''


						super().tearDown()
						gc.collect()
						torch.cuda.empty_cache()




			def 					snake_case__					(				self  ):

						'''simple docstring'''


						UpperCamelCase__       =							StableDiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5"   ,    safety_checker=snake_case  )
						UpperCamelCase__       =							LMSDiscreteScheduler.from_config(sd_pipe.scheduler.config  )
						UpperCamelCase__       =							sd_pipe.to(snake_case  )
						sd_pipe.set_progress_bar_config(disable=snake_case  )

						UpperCamelCase__       =							(
						    "portrait of girl with smokey eyes makeup in abandoned hotel, grange clothes, redshift, wide high angle"
						    " coloured polaroid photograph with flash, kodak film, hyper real, stunning moody cinematography, with"
						    " anamorphic lenses, by maripol, fallen angels by wong kar - wai, style of suspiria and neon demon and"
						    " children from bahnhof zoo, detailed "
						)
						UpperCamelCase__       =							4003660346
						UpperCamelCase__       =							7

						# without safety guidance (sld_guidance_scale = 0)
						UpperCamelCase__       =							torch.manual_seed(snake_case  )
						UpperCamelCase__       =							sd_pipe(
						    [prompt]   ,    generator=snake_case   ,    guidance_scale=snake_case   ,    num_inference_steps=50   ,    output_type="np"   ,    width=512   ,    height=512   ,    sld_guidance_scale=0   ,    )

						UpperCamelCase__       =							output.images
						UpperCamelCase__       =							image[0, -3:, -3:, -1]
						UpperCamelCase__       =							[0.2278, 0.2231, 0.2249, 0.2333, 0.2303, 0.1885, 0.2273, 0.2144, 0.2176]

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

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

						# without safety guidance (strong configuration)
						UpperCamelCase__       =							torch.manual_seed(snake_case  )
						UpperCamelCase__       =							sd_pipe(
						    [prompt]   ,    generator=snake_case   ,    guidance_scale=snake_case   ,    num_inference_steps=50   ,    output_type="np"   ,    width=512   ,    height=512   ,    sld_guidance_scale=2000   ,    sld_warmup_steps=7   ,    sld_threshold=0.025   ,    sld_momentum_scale=0.5   ,    sld_mom_beta=0.7   ,    )

						UpperCamelCase__       =							output.images
						UpperCamelCase__       =							image[0, -3:, -3:, -1]
						UpperCamelCase__       =							[0.2383, 0.2276, 0.236, 0.2192, 0.2186, 0.2053, 0.1971, 0.1901, 0.1719]

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

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




			def 					snake_case__					(				self  ):

						'''simple docstring'''


						UpperCamelCase__       =							StableDiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5"   ,    safety_checker=snake_case  )
						UpperCamelCase__       =							LMSDiscreteScheduler.from_config(sd_pipe.scheduler.config  )
						UpperCamelCase__       =							sd_pipe.to(snake_case  )
						sd_pipe.set_progress_bar_config(disable=snake_case  )

						UpperCamelCase__       =							"padme amidala taking a bath artwork, safe for work, no nudity"
						UpperCamelCase__       =							2734971755
						UpperCamelCase__       =							7

						UpperCamelCase__       =							torch.manual_seed(snake_case  )
						UpperCamelCase__       =							sd_pipe(
						    [prompt]   ,    generator=snake_case   ,    guidance_scale=snake_case   ,    num_inference_steps=50   ,    output_type="np"   ,    width=512   ,    height=512   ,    sld_guidance_scale=0   ,    )

						UpperCamelCase__       =							output.images
						UpperCamelCase__       =							image[0, -3:, -3:, -1]
						UpperCamelCase__       =							[0.3502, 0.3622, 0.3396, 0.3642, 0.3478, 0.3318, 0.35, 0.3348, 0.3297]

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

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

						UpperCamelCase__       =							torch.manual_seed(snake_case  )
						UpperCamelCase__       =							sd_pipe(
						    [prompt]   ,    generator=snake_case   ,    guidance_scale=snake_case   ,    num_inference_steps=50   ,    output_type="np"   ,    width=512   ,    height=512   ,    sld_guidance_scale=2000   ,    sld_warmup_steps=7   ,    sld_threshold=0.025   ,    sld_momentum_scale=0.5   ,    sld_mom_beta=0.7   ,    )

						UpperCamelCase__       =							output.images
						UpperCamelCase__       =							image[0, -3:, -3:, -1]
						UpperCamelCase__       =							[0.5531, 0.5206, 0.4895, 0.5156, 0.5182, 0.4751, 0.4802, 0.4803, 0.4443]

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

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







			def 					snake_case__					(				self  ):

						'''simple docstring'''


						UpperCamelCase__       =							StableDiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5"  )
						UpperCamelCase__       =							sd_pipe.to(snake_case  )
						sd_pipe.set_progress_bar_config(disable=snake_case  )

						UpperCamelCase__       =							(
						    "the four horsewomen of the apocalypse, painting by tom of finland, gaston bussiere, craig mullins, j. c."
						    " leyendecker"
						)
						UpperCamelCase__       =							1044355234
						UpperCamelCase__       =							12

						UpperCamelCase__       =							torch.manual_seed(snake_case  )
						UpperCamelCase__       =							sd_pipe(
						    [prompt]   ,    generator=snake_case   ,    guidance_scale=snake_case   ,    num_inference_steps=50   ,    output_type="np"   ,    width=512   ,    height=512   ,    sld_guidance_scale=0   ,    )

						UpperCamelCase__       =							output.images
						UpperCamelCase__       =							image[0, -3:, -3:, -1]
						UpperCamelCase__       =							np.array([0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]  )

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

						assert np.abs(image_slice.flatten() - expected_slice  ).max() < 1E-7

						UpperCamelCase__       =							torch.manual_seed(snake_case  )
						UpperCamelCase__       =							sd_pipe(
						    [prompt]   ,    generator=snake_case   ,    guidance_scale=snake_case   ,    num_inference_steps=50   ,    output_type="np"   ,    width=512   ,    height=512   ,    sld_guidance_scale=2000   ,    sld_warmup_steps=7   ,    sld_threshold=0.025   ,    sld_momentum_scale=0.5   ,    sld_mom_beta=0.7   ,    )

						UpperCamelCase__       =							output.images
						UpperCamelCase__       =							image[0, -3:, -3:, -1]
						UpperCamelCase__       =							np.array([0.5818, 0.6285, 0.6835, 0.6019, 0.625, 0.6754, 0.6096, 0.6334, 0.6561]  )
						assert image.shape == (1, 512, 512, 3)

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






# limitations under the License.


from typing import Optional, Tuple, Union

import torch

from diffusers import DiffusionPipeline, ImagePipelineOutput







class 							lowerCamelCase__						(				UpperCAmelCase   ):






			"""simple docstring"""






			def __init__(				self   ,    snake_case   ,    snake_case  ):

						'''simple docstring'''


						super().__init__()
						self.register_modules(unet=snake_case   ,    scheduler=snake_case  )







			@torch.no_grad()
			def __call__(				self   ,    snake_case = 1   ,    snake_case = None   ,    snake_case = 50   ,    snake_case = "pil"   ,    snake_case = True   ,    **snake_case   ,    ):

						'''simple docstring'''


						UpperCamelCase__       =							torch.randn(
						    (batch_size, self.unet.config.in_channels, self.unet.config.sample_size, self.unet.config.sample_size)   ,    generator=snake_case   ,    )
						UpperCamelCase__       =							image.to(self.device  )

						# set step values
						self.scheduler.set_timesteps(snake_case  )

						for t in self.progress_bar(self.scheduler.timesteps  ):
									# 1. predict noise model_output
									UpperCamelCase__       =							self.unet(snake_case   ,    snake_case  ).sample

									# 2. predict previous mean of image x_t-1 and add variance depending on eta
									# eta corresponds to η in paper and should be between [0, 1]
									# do x_t -> x_t-1
									UpperCamelCase__       =							self.scheduler.step(snake_case   ,    snake_case   ,    snake_case  ).prev_sample

						UpperCamelCase__       =							(image / 2 + 0.5).clamp(0   ,    1  )
						UpperCamelCase__       =							image.cpu().permute(0   ,    2   ,    3   ,    1  ).numpy()
						if output_type == "pil":
									UpperCamelCase__       =							self.numpy_to_pil(snake_case  )

						if not return_dict:
									return (image,), "This is a local test"

						return ImagePipelineOutput(images=snake_case  ), "This is a local test"


