





import gc
import unittest

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

from diffusers import (
    AutoencoderKL,
    DDIMScheduler,
    PNDMScheduler,
    StableDiffusionLDMaDPipeline,
    UNetaDConditionModel,
)
from diffusers.utils import nightly, slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu

from ..pipeline_params import TEXT_TO_IMAGE_BATCH_PARAMS, TEXT_TO_IMAGE_IMAGE_PARAMS, TEXT_TO_IMAGE_PARAMS


enable_full_determinism()






class 					__A    (    unittest.TestCase ):

			'''simple docstring'''







			lowerCAmelCase_							     =						StableDiffusionLDMaDPipeline
			lowerCAmelCase_							     =						TEXT_TO_IMAGE_PARAMS
			lowerCAmelCase_							     =						TEXT_TO_IMAGE_BATCH_PARAMS
			lowerCAmelCase_							     =						TEXT_TO_IMAGE_IMAGE_PARAMS





			def 		__lowerCamelCase				(    self							):



				'''simple docstring'''







				torch.manual_seed(0							)
				lowerCamelCase__								=		UNetaDConditionModel(
				    block_out_channels=(3_2, 6_4)       ,  layers_per_block=2       ,  sample_size=3_2       ,  in_channels=4       ,  out_channels=4       ,  down_block_types=('''DownBlock2D''', '''CrossAttnDownBlock2D''')       ,  up_block_types=('''CrossAttnUpBlock2D''', '''UpBlock2D''')       ,  cross_attention_dim=3_2       ,  )
				lowerCamelCase__								=		DDIMScheduler(
				    beta_start=0.0_0085       ,  beta_end=0.012       ,  beta_schedule='''scaled_linear'''       ,  clip_sample=__lowerCAmelCase       ,  set_alpha_to_one=__lowerCAmelCase       ,  )
				torch.manual_seed(0							)
				lowerCamelCase__								=		AutoencoderKL(
				    block_out_channels=[3_2, 6_4]       ,  in_channels=6       ,  out_channels=6       ,  down_block_types=['''DownEncoderBlock2D''', '''DownEncoderBlock2D''']       ,  up_block_types=['''UpDecoderBlock2D''', '''UpDecoderBlock2D''']       ,  latent_channels=4       ,  )
				torch.manual_seed(0							)
				lowerCamelCase__								=		CLIPTextConfig(
				    bos_token_id=0       ,  eos_token_id=2       ,  hidden_size=3_2       ,  intermediate_size=3_7       ,  layer_norm_eps=1E-05       ,  num_attention_heads=4       ,  num_hidden_layers=5       ,  pad_token_id=1       ,  vocab_size=1_0_0_0       ,  )
				lowerCamelCase__								=		CLIPTextModel(__lowerCAmelCase							)
				lowerCamelCase__								=		CLIPTokenizer.from_pretrained('''hf-internal-testing/tiny-random-clip'''							)

				lowerCamelCase__								=		{
				    '''unet''': unet,
				    '''scheduler''': scheduler,
				    '''vae''': vae,
				    '''text_encoder''': text_encoder,
				    '''tokenizer''': tokenizer,
				    '''safety_checker''': None,
				    '''feature_extractor''': None,
				}
				return components





			def 		__lowerCamelCase				(    self       ,  __lowerCAmelCase       ,  __lowerCAmelCase=0							):



				'''simple docstring'''







				if str(__lowerCAmelCase							).startswith('''mps'''							):
					lowerCamelCase__								=		torch.manual_seed(__lowerCAmelCase							)
				else:
					lowerCamelCase__								=		torch.Generator(device=__lowerCAmelCase							).manual_seed(__lowerCAmelCase							)
				lowerCamelCase__								=		{
				    '''prompt''': '''A painting of a squirrel eating a burger''',
				    '''generator''': generator,
				    '''num_inference_steps''': 2,
				    '''guidance_scale''': 6.0,
				    '''output_type''': '''numpy''',
				}
				return inputs





			def 		__lowerCamelCase				(    self							):



				'''simple docstring'''







				lowerCamelCase__								=		'''cpu'''  # ensure determinism for the device-dependent torch.Generator

				lowerCamelCase__								=		self.get_dummy_components()
				lowerCamelCase__								=		StableDiffusionLDMaDPipeline(**__lowerCAmelCase							)
				lowerCamelCase__								=		ldmad_pipe.to(__lowerCAmelCase							)
				ldmad_pipe.set_progress_bar_config(disable=__lowerCAmelCase							)

				lowerCamelCase__								=		self.get_dummy_inputs(__lowerCAmelCase							)
				lowerCamelCase__								=		ldmad_pipe(**__lowerCAmelCase							)
				lowerCamelCase__				,	lowerCamelCase__								=		output.rgb, output.depth

				lowerCamelCase__								=		rgb[0, -3:, -3:, -1]
				lowerCamelCase__								=		depth[0, -3:, -1]

				assert rgb.shape == (1, 6_4, 6_4, 3)
				assert depth.shape == (1, 6_4, 6_4)

				lowerCamelCase__								=		np.array(
				    [0.3733_8176, 0.7_0247, 0.7420_3193, 0.5164_3604, 0.5825_6793, 0.6093_2136, 0.418_1095, 0.4835_5877, 0.4653_5262]							)
				lowerCamelCase__								=		np.array([103.4_6727, 85.81_2004, 87.84_9236]							)

				assert np.abs(image_slice_rgb.flatten() - expected_slice_rgb							).max() < 1E-2
				assert np.abs(image_slice_depth.flatten() - expected_slice_depth							).max() < 1E-2





			def 		__lowerCamelCase				(    self							):



				'''simple docstring'''







				lowerCamelCase__								=		self.get_dummy_components()
				lowerCamelCase__								=		StableDiffusionLDMaDPipeline(**__lowerCAmelCase							)
				lowerCamelCase__								=		ldmad_pipe.to(__lowerCAmelCase							)
				ldmad_pipe.set_progress_bar_config(disable=__lowerCAmelCase							)

				lowerCamelCase__								=		self.get_dummy_inputs(__lowerCAmelCase							)
				lowerCamelCase__								=		3 * [inputs['''prompt''']]

				# forward
				lowerCamelCase__								=		ldmad_pipe(**__lowerCAmelCase							)
				lowerCamelCase__				,	lowerCamelCase__								=		output.rgb, output.depth
				lowerCamelCase__								=		rgb_slice_a[0, -3:, -3:, -1]
				lowerCamelCase__								=		depth_slice_a[0, -3:, -1]

				lowerCamelCase__								=		self.get_dummy_inputs(__lowerCAmelCase							)
				lowerCamelCase__								=		3 * [inputs.pop('''prompt'''							)]

				lowerCamelCase__								=		ldmad_pipe.tokenizer(
				    __lowerCAmelCase       ,  padding='''max_length'''       ,  max_length=ldmad_pipe.tokenizer.model_max_length       ,  truncation=__lowerCAmelCase       ,  return_tensors='''pt'''       ,  )
				lowerCamelCase__								=		text_inputs['''input_ids'''].to(__lowerCAmelCase							)

				lowerCamelCase__								=		ldmad_pipe.text_encoder(__lowerCAmelCase							)[0]

				lowerCamelCase__								=		prompt_embeds

				# forward
				lowerCamelCase__								=		ldmad_pipe(**__lowerCAmelCase							)
				lowerCamelCase__				,	lowerCamelCase__								=		output.rgb, output.depth
				lowerCamelCase__								=		rgb_slice_a[0, -3:, -3:, -1]
				lowerCamelCase__								=		depth_slice_a[0, -3:, -1]

				assert np.abs(rgb_slice_a.flatten() - rgb_slice_a.flatten()							).max() < 1E-4
				assert np.abs(depth_slice_a.flatten() - depth_slice_a.flatten()							).max() < 1E-4





			def 		__lowerCamelCase				(    self							):



				'''simple docstring'''







				lowerCamelCase__								=		'''cpu'''  # ensure determinism for the device-dependent torch.Generator
				lowerCamelCase__								=		self.get_dummy_components()
				lowerCamelCase__								=		PNDMScheduler(skip_prk_steps=__lowerCAmelCase							)
				lowerCamelCase__								=		StableDiffusionLDMaDPipeline(**__lowerCAmelCase							)
				lowerCamelCase__								=		ldmad_pipe.to(__lowerCAmelCase							)
				ldmad_pipe.set_progress_bar_config(disable=__lowerCAmelCase							)

				lowerCamelCase__								=		self.get_dummy_inputs(__lowerCAmelCase							)
				lowerCamelCase__								=		'''french fries'''
				lowerCamelCase__								=		ldmad_pipe(**__lowerCAmelCase       ,  negative_prompt=__lowerCAmelCase							)

				lowerCamelCase__				,	lowerCamelCase__								=		output.rgb, output.depth
				lowerCamelCase__								=		rgb[0, -3:, -3:, -1]
				lowerCamelCase__								=		depth[0, -3:, -1]

				assert rgb.shape == (1, 6_4, 6_4, 3)
				assert depth.shape == (1, 6_4, 6_4)

				lowerCamelCase__								=		np.array(
				    [0.3_7044, 0.7181_1503, 0.722_3251, 0.4860_3675, 0.563_8391, 0.636_4948, 0.4283_3704, 0.490_1315, 0.4792_6217]							)
				lowerCamelCase__								=		np.array([107.8_4738, 84.6_2802, 89.96_2135]							)
				assert np.abs(rgb_slice.flatten() - expected_slice_rgb							).max() < 1E-2
				assert np.abs(depth_slice.flatten() - expected_slice_depth							).max() < 1E-2






@slow
@require_torch_gpu
class 					__A    (    unittest.TestCase ):

			'''simple docstring'''







			def 		__lowerCamelCase				(    self							):



				'''simple docstring'''







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





			def 		__lowerCamelCase				(    self       ,  __lowerCAmelCase       ,  __lowerCAmelCase="cpu"       ,  __lowerCAmelCase=torch.floataa       ,  __lowerCAmelCase=0							):



				'''simple docstring'''







				lowerCamelCase__								=		torch.Generator(device=__lowerCAmelCase							).manual_seed(__lowerCAmelCase							)
				lowerCamelCase__								=		np.random.RandomState(__lowerCAmelCase							).standard_normal((1, 4, 6_4, 6_4)							)
				lowerCamelCase__								=		torch.from_numpy(__lowerCAmelCase							).to(device=__lowerCAmelCase       ,  dtype=__lowerCAmelCase							)
				lowerCamelCase__								=		{
				    '''prompt''': '''a photograph of an astronaut riding a horse''',
				    '''latents''': latents,
				    '''generator''': generator,
				    '''num_inference_steps''': 3,
				    '''guidance_scale''': 7.5,
				    '''output_type''': '''numpy''',
				}
				return inputs





			def 		__lowerCamelCase				(    self							):



				'''simple docstring'''







				lowerCamelCase__								=		StableDiffusionLDMaDPipeline.from_pretrained('''Intel/ldm3d'''							)
				lowerCamelCase__								=		ldmad_pipe.to(__lowerCAmelCase							)
				ldmad_pipe.set_progress_bar_config(disable=__lowerCAmelCase							)

				lowerCamelCase__								=		self.get_inputs(__lowerCAmelCase							)
				lowerCamelCase__								=		ldmad_pipe(**__lowerCAmelCase							)
				lowerCamelCase__				,	lowerCamelCase__								=		output.rgb, output.depth
				lowerCamelCase__								=		rgb[0, -3:, -3:, -1].flatten()
				lowerCamelCase__								=		rgb[0, -3:, -1].flatten()

				assert rgb.shape == (1, 5_1_2, 5_1_2, 3)
				assert depth.shape == (1, 5_1_2, 5_1_2)

				lowerCamelCase__								=		np.array(
				    [0.5380_5465, 0.5670_7305, 0.548_6515, 0.5701_2236, 0.581_4511, 0.5625_3487, 0.5484_3014, 0.5509_2263, 0.645_9706]							)
				lowerCamelCase__								=		np.array(
				    [0.926_3781, 0.667_8672, 0.548_6515, 0.9220_2145, 0.6783_1135, 0.5625_3487, 0.924_1694, 0.755_1478, 0.645_9706]							)
				assert np.abs(rgb_slice - expected_slice_rgb							).max() < 3E-3
				assert np.abs(depth_slice - expected_slice_depth							).max() < 3E-3






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

			'''simple docstring'''







			def 		__lowerCamelCase				(    self							):



				'''simple docstring'''







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





			def 		__lowerCamelCase				(    self       ,  __lowerCAmelCase       ,  __lowerCAmelCase="cpu"       ,  __lowerCAmelCase=torch.floataa       ,  __lowerCAmelCase=0							):



				'''simple docstring'''







				lowerCamelCase__								=		torch.Generator(device=__lowerCAmelCase							).manual_seed(__lowerCAmelCase							)
				lowerCamelCase__								=		np.random.RandomState(__lowerCAmelCase							).standard_normal((1, 4, 6_4, 6_4)							)
				lowerCamelCase__								=		torch.from_numpy(__lowerCAmelCase							).to(device=__lowerCAmelCase       ,  dtype=__lowerCAmelCase							)
				lowerCamelCase__								=		{
				    '''prompt''': '''a photograph of an astronaut riding a horse''',
				    '''latents''': latents,
				    '''generator''': generator,
				    '''num_inference_steps''': 5_0,
				    '''guidance_scale''': 7.5,
				    '''output_type''': '''numpy''',
				}
				return inputs





			def 		__lowerCamelCase				(    self							):



				'''simple docstring'''







				lowerCamelCase__								=		StableDiffusionLDMaDPipeline.from_pretrained('''Intel/ldm3d'''							).to(__lowerCAmelCase							)
				ldmad_pipe.set_progress_bar_config(disable=__lowerCAmelCase							)

				lowerCamelCase__								=		self.get_inputs(__lowerCAmelCase							)
				lowerCamelCase__								=		ldmad_pipe(**__lowerCAmelCase							)
				lowerCamelCase__				,	lowerCamelCase__								=		output.rgb, output.depth

				lowerCamelCase__								=		0.49_5586
				lowerCamelCase__								=		0.3379_5515
				lowerCamelCase__								=		112.4_8518
				lowerCamelCase__								=		98.48_9746
				assert np.abs(expected_rgb_mean - rgb.mean()							) < 1E-3
				assert np.abs(expected_rgb_std - rgb.std()							) < 1E-3
				assert np.abs(expected_depth_mean - depth.mean()							) < 1E-3
				assert np.abs(expected_depth_std - depth.std()							) < 1E-3





			def 		__lowerCamelCase				(    self							):



				'''simple docstring'''







				lowerCamelCase__								=		StableDiffusionLDMaDPipeline.from_pretrained('''Intel/ldm3d-4c'''							).to(__lowerCAmelCase							)
				ldmad_pipe.set_progress_bar_config(disable=__lowerCAmelCase							)

				lowerCamelCase__								=		self.get_inputs(__lowerCAmelCase							)
				lowerCamelCase__								=		ldmad_pipe(**__lowerCAmelCase							)
				lowerCamelCase__				,	lowerCamelCase__								=		output.rgb, output.depth

				lowerCamelCase__								=		0.419_4127
				lowerCamelCase__								=		0.3537_5586
				lowerCamelCase__								=		0.563_8502
				lowerCamelCase__								=		0.3468_6103

				assert rgb.shape == (1, 5_1_2, 5_1_2, 3)
				assert depth.shape == (1, 5_1_2, 5_1_2, 1)
				assert np.abs(expected_rgb_mean - rgb.mean()							) < 1E-3
				assert np.abs(expected_rgb_std - rgb.std()							) < 1E-3
				assert np.abs(expected_depth_mean - depth.mean()							) < 1E-3
				assert np.abs(expected_depth_std - depth.std()							) < 1E-3










import torch

from diffusers import UnCLIPScheduler

from .test_schedulers import SchedulerCommonTest






class 					__A    (    lowerCAmelCase ):

			'''simple docstring'''







			lowerCAmelCase_							     =						(UnCLIPScheduler,)





			def 		__lowerCamelCase				(    self       ,  **__lowerCAmelCase							):



				'''simple docstring'''







				lowerCamelCase__								=		{
				    '''num_train_timesteps''': 1_0_0_0,
				    '''variance_type''': '''fixed_small_log''',
				    '''clip_sample''': True,
				    '''clip_sample_range''': 1.0,
				    '''prediction_type''': '''epsilon''',
				}

				config.update(**__lowerCAmelCase							)
				return config





			def 		__lowerCamelCase				(    self							):



				'''simple docstring'''







				for timesteps in [1, 5, 1_0_0, 1_0_0_0]:
					self.check_over_configs(num_train_timesteps=__lowerCAmelCase							)





			def 		__lowerCamelCase				(    self							):



				'''simple docstring'''







				for variance in ["fixed_small_log", "learned_range"]:
					self.check_over_configs(variance_type=__lowerCAmelCase							)





			def 		__lowerCamelCase				(    self							):



				'''simple docstring'''







				for clip_sample in [True, False]:
					self.check_over_configs(clip_sample=__lowerCAmelCase							)





			def 		__lowerCamelCase				(    self							):



				'''simple docstring'''







				for clip_sample_range in [1, 5, 1_0, 2_0]:
					self.check_over_configs(clip_sample_range=__lowerCAmelCase							)





			def 		__lowerCamelCase				(    self							):



				'''simple docstring'''







				for prediction_type in ["epsilon", "sample"]:
					self.check_over_configs(prediction_type=__lowerCAmelCase							)





			def 		__lowerCamelCase				(    self							):



				'''simple docstring'''







				for time_step in [0, 5_0_0, 9_9_9]:
					for prev_timestep in [None, 5, 1_0_0, 2_5_0, 5_0_0, 7_5_0]:
						if prev_timestep is not None and prev_timestep >= time_step:
							continue

						self.check_over_forward(time_step=__lowerCAmelCase       ,  prev_timestep=__lowerCAmelCase							)





			def 		__lowerCamelCase				(    self							):



				'''simple docstring'''







				lowerCamelCase__								=		self.scheduler_classes[0]
				lowerCamelCase__								=		self.get_scheduler_config(variance_type='''fixed_small_log'''							)
				lowerCamelCase__								=		scheduler_class(**__lowerCAmelCase							)

				assert torch.sum(torch.abs(scheduler._get_variance(0							) - 1.0000E-10							)							) < 1E-5
				assert torch.sum(torch.abs(scheduler._get_variance(4_8_7							) - 0.054_9625							)							) < 1E-5
				assert torch.sum(torch.abs(scheduler._get_variance(9_9_9							) - 0.999_4987							)							) < 1E-5





			def 		__lowerCamelCase				(    self							):



				'''simple docstring'''







				lowerCamelCase__								=		self.scheduler_classes[0]
				lowerCamelCase__								=		self.get_scheduler_config(variance_type='''learned_range'''							)
				lowerCamelCase__								=		scheduler_class(**__lowerCAmelCase							)

				lowerCamelCase__								=		0.5

				assert scheduler._get_variance(1       ,  predicted_variance=__lowerCAmelCase							) - -10.171_2790 < 1E-5
				assert scheduler._get_variance(4_8_7       ,  predicted_variance=__lowerCAmelCase							) - -5.799_8052 < 1E-5
				assert scheduler._get_variance(9_9_9       ,  predicted_variance=__lowerCAmelCase							) - -0.001_0011 < 1E-5





			def 		__lowerCamelCase				(    self							):



				'''simple docstring'''







				lowerCamelCase__								=		self.scheduler_classes[0]
				lowerCamelCase__								=		self.get_scheduler_config()
				lowerCamelCase__								=		scheduler_class(**__lowerCAmelCase							)

				lowerCamelCase__								=		scheduler.timesteps

				lowerCamelCase__								=		self.dummy_model()
				lowerCamelCase__								=		self.dummy_sample_deter
				lowerCamelCase__								=		torch.manual_seed(0							)

				for i, t in enumerate(__lowerCAmelCase							):
					# 1. predict noise residual
					lowerCamelCase__								=		model(__lowerCAmelCase       ,  __lowerCAmelCase							)

					# 2. predict previous mean of sample x_t-1
					lowerCamelCase__								=		scheduler.step(__lowerCAmelCase       ,  __lowerCAmelCase       ,  __lowerCAmelCase       ,  generator=__lowerCAmelCase							).prev_sample

					lowerCamelCase__								=		pred_prev_sample

				lowerCamelCase__								=		torch.sum(torch.abs(__lowerCAmelCase							)							)
				lowerCamelCase__								=		torch.mean(torch.abs(__lowerCAmelCase							)							)

				assert abs(result_sum.item() - 252.268_2495							) < 1E-2
				assert abs(result_mean.item() - 0.328_4743							) < 1E-3





			def 		__lowerCamelCase				(    self							):



				'''simple docstring'''







				lowerCamelCase__								=		self.scheduler_classes[0]
				lowerCamelCase__								=		self.get_scheduler_config()
				lowerCamelCase__								=		scheduler_class(**__lowerCAmelCase							)

				scheduler.set_timesteps(2_5							)

				lowerCamelCase__								=		scheduler.timesteps

				lowerCamelCase__								=		self.dummy_model()
				lowerCamelCase__								=		self.dummy_sample_deter
				lowerCamelCase__								=		torch.manual_seed(0							)

				for i, t in enumerate(__lowerCAmelCase							):
					# 1. predict noise residual
					lowerCamelCase__								=		model(__lowerCAmelCase       ,  __lowerCAmelCase							)

					if i + 1 == timesteps.shape[0]:
						lowerCamelCase__								=		None
					else:
						lowerCamelCase__								=		timesteps[i + 1]

					# 2. predict previous mean of sample x_t-1
					lowerCamelCase__								=		scheduler.step(
					    __lowerCAmelCase       ,  __lowerCAmelCase       ,  __lowerCAmelCase       ,  prev_timestep=__lowerCAmelCase       ,  generator=__lowerCAmelCase							).prev_sample

					lowerCamelCase__								=		pred_prev_sample

				lowerCamelCase__								=		torch.sum(torch.abs(__lowerCAmelCase							)							)
				lowerCamelCase__								=		torch.mean(torch.abs(__lowerCAmelCase							)							)

				assert abs(result_sum.item() - 258.204_4983							) < 1E-2
				assert abs(result_mean.item() - 0.336_2038							) < 1E-3





			def 		__lowerCamelCase				(    self							):



				'''simple docstring'''







				pass





			def 		__lowerCamelCase				(    self							):



				'''simple docstring'''







				pass



