



import gc
import unittest

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

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

from ..pipeline_params import TEXT_TO_IMAGE_BATCH_PARAMS, TEXT_TO_IMAGE_PARAMS
from ..test_pipelines_common import PipelineTesterMixin


enable_full_determinism()



class 	_lowercase							(   lowerCAmelCase,      unittest.TestCase	):





	"""simple docstring"""
	__A							=							LDMTextToImagePipeline
	__A							=							TEXT_TO_IMAGE_PARAMS - {
	    "negative_prompt",
	    "negative_prompt_embeds",
	    "cross_attention_kwargs",
	    "prompt_embeds",
	}
	__A							=							PipelineTesterMixin.required_optional_params - {
	    "num_images_per_prompt",
	    "callback",
	    "callback_steps",
	}
	__A							=							TEXT_TO_IMAGE_BATCH_PARAMS
	__A							=							False




	def      UpperCamelCase_       (self		):
						"""simple docstring"""

						torch.manual_seed(0		)
						a					     =   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		,       )
						a					     =   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		)
						a					     =   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		,       )
						torch.manual_seed(0		)
						a					     =   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		,       )
						a					     =   CLIPTextModel(lowerCamelCase_		)
						a					     =   CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip"		)

						a					     =   {
						    "unet": unet,
						    "scheduler": scheduler,
						    "vqvae": vae,
						    "bert": text_encoder,
						    "tokenizer": tokenizer,
						}
						return components




	def      UpperCamelCase_       (self		,       lowerCamelCase_		,       lowerCamelCase_=0		):
						"""simple docstring"""

						if str(lowerCamelCase_		).startswith("mps"		):
											a					     =   torch.manual_seed(lowerCamelCase_		)
						else:
											a					     =   torch.Generator(device=lowerCamelCase_		).manual_seed(lowerCamelCase_		)
						a					     =   {
						    "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      UpperCamelCase_       (self		):
						"""simple docstring"""

						a					     =   "cpu"  # ensure determinism for the device-dependent torch.Generator

						a					     =   self.get_dummy_components()
						a					     =   LDMTextToImagePipeline(**lowerCamelCase_		)
						pipe.to(lowerCamelCase_		)
						pipe.set_progress_bar_config(disable=lowerCamelCase_		)

						a					     =   self.get_dummy_inputs(lowerCamelCase_		)
						a					     =   pipe(**lowerCamelCase_		).images
						a					     =   image[0, -3:, -3:, -1]

						assert image.shape == (1, 16, 16, 3)
						a					     =   np.array([0.6101, 0.6156, 0.5622, 0.4895, 0.6661, 0.3804, 0.5748, 0.6136, 0.5014]		)

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




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





	"""simple docstring"""




	def      UpperCamelCase_       (self		):
						"""simple docstring"""

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




	def      UpperCamelCase_       (self		,       lowerCamelCase_		,       lowerCamelCase_=torch.floataa		,       lowerCamelCase_=0		):
						"""simple docstring"""

						a					     =   torch.manual_seed(lowerCamelCase_		)
						a					     =   np.random.RandomState(lowerCamelCase_		).standard_normal((1, 4, 32, 32)		)
						a					     =   torch.from_numpy(lowerCamelCase_		).to(device=lowerCamelCase_		,       dtype=lowerCamelCase_		)
						a					     =   {
						    "prompt": "A painting of a squirrel eating a burger",
						    "latents": latents,
						    "generator": generator,
						    "num_inference_steps": 3,
						    "guidance_scale": 6.0,
						    "output_type": "numpy",
						}
						return inputs




	def      UpperCamelCase_       (self		):
						"""simple docstring"""

						a					     =   LDMTextToImagePipeline.from_pretrained("CompVis/ldm-text2im-large-256"		).to(lowerCamelCase_		)
						pipe.set_progress_bar_config(disable=lowerCamelCase_		)

						a					     =   self.get_inputs(lowerCamelCase_		)
						a					     =   pipe(**lowerCamelCase_		).images
						a					     =   image[0, -3:, -3:, -1].flatten()

						assert image.shape == (1, 256, 256, 3)
						a					     =   np.array([0.5_1825, 0.5_2850, 0.5_2543, 0.5_4258, 0.5_2304, 0.5_2569, 0.5_4363, 0.5_5276, 0.5_6878]		)
						a					     =   np.abs(expected_slice - image_slice		).max()
						assert max_diff < 1E-3




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





	"""simple docstring"""




	def      UpperCamelCase_       (self		):
						"""simple docstring"""

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




	def      UpperCamelCase_       (self		,       lowerCamelCase_		,       lowerCamelCase_=torch.floataa		,       lowerCamelCase_=0		):
						"""simple docstring"""

						a					     =   torch.manual_seed(lowerCamelCase_		)
						a					     =   np.random.RandomState(lowerCamelCase_		).standard_normal((1, 4, 32, 32)		)
						a					     =   torch.from_numpy(lowerCamelCase_		).to(device=lowerCamelCase_		,       dtype=lowerCamelCase_		)
						a					     =   {
						    "prompt": "A painting of a squirrel eating a burger",
						    "latents": latents,
						    "generator": generator,
						    "num_inference_steps": 50,
						    "guidance_scale": 6.0,
						    "output_type": "numpy",
						}
						return inputs




	def      UpperCamelCase_       (self		):
						"""simple docstring"""

						a					     =   LDMTextToImagePipeline.from_pretrained("CompVis/ldm-text2im-large-256"		).to(lowerCamelCase_		)
						pipe.set_progress_bar_config(disable=lowerCamelCase_		)

						a					     =   self.get_inputs(lowerCamelCase_		)
						a					     =   pipe(**lowerCamelCase_		).images[0]

						a					     =   load_numpy(
						    "https://huggingface.co/datasets/diffusers/test-arrays/resolve/main/ldm_text2img/ldm_large_256_ddim.npy"		)
						a					     =   np.abs(expected_image - image		).max()
						assert max_diff < 1E-3





import torch

from diffusers import DDPMScheduler

from .test_schedulers import SchedulerCommonTest



class 	_lowercase							(   lowerCAmelCase	):





	"""simple docstring"""
	__A							=							(DDPMScheduler,)




	def      UpperCamelCase_       (self		,       **lowerCamelCase_		):
						"""simple docstring"""

						a					     =   {
						    "num_train_timesteps": 1000,
						    "beta_start": 0.0001,
						    "beta_end": 0.02,
						    "beta_schedule": "linear",
						    "variance_type": "fixed_small",
						    "clip_sample": True,
						}

						config.update(**lowerCamelCase_		)
						return config




	def      UpperCamelCase_       (self		):
						"""simple docstring"""

						for timesteps in [1, 5, 100, 1000]:
											self.check_over_configs(num_train_timesteps=lowerCamelCase_		)




	def      UpperCamelCase_       (self		):
						"""simple docstring"""

						for beta_start, beta_end in zip([0.0001, 0.001, 0.01, 0.1]		,       [0.002, 0.02, 0.2, 2]		):
											self.check_over_configs(beta_start=lowerCamelCase_		,       beta_end=lowerCamelCase_		)




	def      UpperCamelCase_       (self		):
						"""simple docstring"""

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




	def      UpperCamelCase_       (self		):
						"""simple docstring"""

						for variance in ["fixed_small", "fixed_large", "other"]:
											self.check_over_configs(variance_type=lowerCamelCase_		)




	def      UpperCamelCase_       (self		):
						"""simple docstring"""

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




	def      UpperCamelCase_       (self		):
						"""simple docstring"""

						self.check_over_configs(thresholding=lowerCamelCase_		)
						for threshold in [0.5, 1.0, 2.0]:
											for prediction_type in ["epsilon", "sample", "v_prediction"]:
																self.check_over_configs(
																    thresholding=lowerCamelCase_		,       prediction_type=lowerCamelCase_		,       sample_max_value=lowerCamelCase_		,       )




	def      UpperCamelCase_       (self		):
						"""simple docstring"""

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




	def      UpperCamelCase_       (self		):
						"""simple docstring"""

						for t in [0, 500, 999]:
											self.check_over_forward(time_step=lowerCamelCase_		)




	def      UpperCamelCase_       (self		):
						"""simple docstring"""

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

						assert torch.sum(torch.abs(scheduler._get_variance(0		) - 0.0		)		) < 1E-5
						assert torch.sum(torch.abs(scheduler._get_variance(487		) - 0.0_0979		)		) < 1E-5
						assert torch.sum(torch.abs(scheduler._get_variance(999		) - 0.02		)		) < 1E-5




	def      UpperCamelCase_       (self		):
						"""simple docstring"""

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

						a					     =   len(lowerCamelCase_		)

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

						for t in reversed(range(lowerCamelCase_		)		):
											# 1. predict noise residual
											a					     =   model(lowerCamelCase_		,       lowerCamelCase_		)

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

											# if t > 0:
											#     noise = self.dummy_sample_deter
											#     variance = scheduler.get_variance(t) ** (0.5) * noise
											#
											# sample = pred_prev_sample + variance
											a					     =   pred_prev_sample

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

						assert abs(result_sum.item() - 258.9606		) < 1E-2
						assert abs(result_mean.item() - 0.3372		) < 1E-3




	def      UpperCamelCase_       (self		):
						"""simple docstring"""

						a					     =   self.scheduler_classes[0]
						a					     =   self.get_scheduler_config(prediction_type="v_prediction"		)
						a					     =   scheduler_class(**lowerCamelCase_		)

						a					     =   len(lowerCamelCase_		)

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

						for t in reversed(range(lowerCamelCase_		)		):
											# 1. predict noise residual
											a					     =   model(lowerCamelCase_		,       lowerCamelCase_		)

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

											# if t > 0:
											#     noise = self.dummy_sample_deter
											#     variance = scheduler.get_variance(t) ** (0.5) * noise
											#
											# sample = pred_prev_sample + variance
											a					     =   pred_prev_sample

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

						assert abs(result_sum.item() - 202.0296		) < 1E-2
						assert abs(result_mean.item() - 0.2631		) < 1E-3




	def      UpperCamelCase_       (self		):
						"""simple docstring"""

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

						a					     =   [100, 87, 50, 1, 0]

						scheduler.set_timesteps(timesteps=lowerCamelCase_		)

						a					     =   scheduler.timesteps

						for i, timestep in enumerate(lowerCamelCase_		):
											if i == len(lowerCamelCase_		) - 1:
																a					     =   -1
											else:
																a					     =   timesteps[i + 1]

											a					     =   scheduler.previous_timestep(lowerCamelCase_		)
											a					     =   prev_t.item()

											self.assertEqual(lowerCamelCase_		,       lowerCamelCase_		)




	def      UpperCamelCase_       (self		):
						"""simple docstring"""

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

						a					     =   [100, 87, 50, 51, 0]

						with self.assertRaises(lowerCamelCase_		,       msg="`custom_timesteps` must be in descending order."		):
											scheduler.set_timesteps(timesteps=lowerCamelCase_		)




	def      UpperCamelCase_       (self		):
						"""simple docstring"""

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

						a					     =   [100, 87, 50, 1, 0]
						a					     =   len(lowerCamelCase_		)

						with self.assertRaises(lowerCamelCase_		,       msg="Can only pass one of `num_inference_steps` or `custom_timesteps`."		):
											scheduler.set_timesteps(num_inference_steps=lowerCamelCase_		,       timesteps=lowerCamelCase_		)




	def      UpperCamelCase_       (self		):
						"""simple docstring"""

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

						a					     =   [scheduler.config.num_train_timesteps]

						with self.assertRaises(
						    lowerCamelCase_		,       msg="`timesteps` must start before `self.config.train_timesteps`: {scheduler.config.num_train_timesteps}}"		,       ):
											scheduler.set_timesteps(timesteps=lowerCamelCase_		)
