



'''simple docstring'''



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 					UpperCAmelCase__	(		UpperCamelCase__							,       unittest.TestCase ):
					a      :     Tuple        =							LDMTextToImagePipeline
					a      :     Optional[Any]        =							TEXT_TO_IMAGE_PARAMS - {
					    """negative_prompt""",
					    """negative_prompt_embeds""",
					    """cross_attention_kwargs""",
					    """prompt_embeds""",
					}
					a      :     Optional[Any]        =							PipelineTesterMixin.required_optional_params - {
					    """num_images_per_prompt""",
					    """callback""",
					    """callback_steps""",
					}
					a      :     str        =							TEXT_TO_IMAGE_BATCH_PARAMS
					a      :     List[Any]        =							False



					def        UpperCAmelCase_ (    self     )     ->			Optional[int]:
						torch.manual_seed(0     )
						__lowerCAmelCase												=  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			,				)
						__lowerCAmelCase												=  DDIMScheduler(
						    beta_start=0.0_00_85			,				beta_end=0.0_12			,				beta_schedule="scaled_linear"			,				clip_sample=UpperCamelCase			,				set_alpha_to_one=UpperCamelCase			,				)
						torch.manual_seed(0     )
						__lowerCAmelCase												=  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     )
						__lowerCAmelCase												=  CLIPTextConfig(
						    bos_token_id=0			,				eos_token_id=2			,				hidden_size=32			,				intermediate_size=37			,				layer_norm_eps=1E-05			,				num_attention_heads=4			,				num_hidden_layers=5			,				pad_token_id=1			,				vocab_size=1000			,				)
						__lowerCAmelCase												=  CLIPTextModel(UpperCamelCase     )
						__lowerCAmelCase												=  CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip"     )

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



					def        UpperCAmelCase_ (    self			,				UpperCamelCase			,				UpperCamelCase=0     )     ->			Dict:
						if str(UpperCamelCase     ).startswith("mps"     ):
							__lowerCAmelCase												=  torch.manual_seed(UpperCamelCase     )
						else:
							__lowerCAmelCase												=  torch.Generator(device=UpperCamelCase     ).manual_seed(UpperCamelCase     )
						__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        UpperCAmelCase_ (    self     )     ->			Dict:
						__lowerCAmelCase												=  "cpu"  # ensure determinism for the device-dependent torch.Generator

						__lowerCAmelCase												=  self.get_dummy_components()
						__lowerCAmelCase												=  LDMTextToImagePipeline(**UpperCamelCase     )
						pipe.to(UpperCamelCase     )
						pipe.set_progress_bar_config(disable=UpperCamelCase     )

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

						assert image.shape == (1, 16, 16, 3)
						__lowerCAmelCase												=  np.array([0.61_01, 0.61_56, 0.56_22, 0.48_95, 0.66_61, 0.38_04, 0.57_48, 0.61_36, 0.50_14]     )

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






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



					def        UpperCAmelCase_ (    self     )     ->			Dict:
						super().tearDown()
						gc.collect()
						torch.cuda.empty_cache()



					def        UpperCAmelCase_ (    self			,				UpperCamelCase			,				UpperCamelCase=torch.floataa			,				UpperCamelCase=0     )     ->			Union[str, Any]:
						__lowerCAmelCase												=  torch.manual_seed(UpperCamelCase     )
						__lowerCAmelCase												=  np.random.RandomState(UpperCamelCase     ).standard_normal((1, 4, 32, 32)     )
						__lowerCAmelCase												=  torch.from_numpy(UpperCamelCase     ).to(device=UpperCamelCase			,				dtype=UpperCamelCase     )
						__lowerCAmelCase												=  {
						    "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     )     ->			Any:
						__lowerCAmelCase												=  LDMTextToImagePipeline.from_pretrained("CompVis/ldm-text2im-large-256"     ).to(UpperCamelCase     )
						pipe.set_progress_bar_config(disable=UpperCamelCase     )

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

						assert image.shape == (1, 256, 256, 3)
						__lowerCAmelCase												=  np.array([0.5_18_25, 0.5_28_50, 0.5_25_43, 0.5_42_58, 0.5_23_04, 0.5_25_69, 0.5_43_63, 0.5_52_76, 0.5_68_78]     )
						__lowerCAmelCase												=  np.abs(expected_slice - image_slice     ).max()
						assert max_diff < 1E-3






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



					def        UpperCAmelCase_ (    self     )     ->			Union[str, Any]:
						super().tearDown()
						gc.collect()
						torch.cuda.empty_cache()



					def        UpperCAmelCase_ (    self			,				UpperCamelCase			,				UpperCamelCase=torch.floataa			,				UpperCamelCase=0     )     ->			Tuple:
						__lowerCAmelCase												=  torch.manual_seed(UpperCamelCase     )
						__lowerCAmelCase												=  np.random.RandomState(UpperCamelCase     ).standard_normal((1, 4, 32, 32)     )
						__lowerCAmelCase												=  torch.from_numpy(UpperCamelCase     ).to(device=UpperCamelCase			,				dtype=UpperCamelCase     )
						__lowerCAmelCase												=  {
						    "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     )     ->			Optional[Any]:
						__lowerCAmelCase												=  LDMTextToImagePipeline.from_pretrained("CompVis/ldm-text2im-large-256"     ).to(UpperCamelCase     )
						pipe.set_progress_bar_config(disable=UpperCamelCase     )

						__lowerCAmelCase												=  self.get_inputs(UpperCamelCase     )
						__lowerCAmelCase												=  pipe(**UpperCamelCase     ).images[0]

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




'''simple docstring'''



import importlib
import os
from dataclasses import dataclass
from enum import Enum
from typing import Any, Dict, Optional, Union

import torch

from ..utils import BaseOutput


lowerCAmelCase	:		Optional[Any]				=   '''scheduler_config.json'''






class 					UpperCAmelCase__	(		UpperCamelCase__ ):
					a      :     str        =							1
					a      :     Optional[int]        =							2
					a      :     int        =							3
					a      :     Union[str, Any]        =							4
					a      :     int        =							5
					a      :     Optional[int]        =							6
					a      :     str        =							7
					a      :     List[Any]        =							8
					a      :     List[str]        =							9
					a      :     List[str]        =							1_0
					a      :     int        =							1_1
					a      :     Any        =							1_2
					a      :     Any        =							1_3
					a      :     Tuple        =							1_4






@dataclass
class 					UpperCAmelCase__	(		UpperCamelCase__ ):
					a      :     torch.FloatTensor






class 					UpperCAmelCase__	:
					a      :     Tuple        =							SCHEDULER_CONFIG_NAME
					a      :     Union[str, Any]        =							[]
					a      :     str        =							True



					@classmethod
					def        UpperCAmelCase_ (    cls			,				UpperCamelCase = None			,				UpperCamelCase = None			,				UpperCamelCase=False			,				**UpperCamelCase			,				)     ->			int:
						__lowerCAmelCase							,       __lowerCAmelCase							,       __lowerCAmelCase												=  cls.load_config(
						    pretrained_model_name_or_path=UpperCamelCase			,				subfolder=UpperCamelCase			,				return_unused_kwargs=UpperCamelCase			,				return_commit_hash=UpperCamelCase			,				**UpperCamelCase			,				)
						return cls.from_config(UpperCamelCase			,				return_unused_kwargs=UpperCamelCase			,				**UpperCamelCase     )



					def        UpperCAmelCase_ (    self			,				UpperCamelCase			,				UpperCamelCase = False			,				**UpperCamelCase     )     ->			Dict:
						self.save_config(save_directory=UpperCamelCase			,				push_to_hub=UpperCamelCase			,				**UpperCamelCase     )



					@property
					def        UpperCAmelCase_ (    self     )     ->			str:
						return self._get_compatibles()



					@classmethod
					def        UpperCAmelCase_ (    cls     )     ->			Tuple:
						__lowerCAmelCase												=  list(set([cls.__name__] + cls._compatibles     )     )
						__lowerCAmelCase												=  importlib.import_module(__name__.split("."     )[0]     )
						__lowerCAmelCase												=  [
						    getattr(UpperCamelCase			,				UpperCamelCase     ) for c in compatible_classes_str if hasattr(UpperCamelCase			,				UpperCamelCase     )
						]
						return compatible_classes