

'''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 							lowercase_						(		__A			,							unittest.TestCase			):

  '''simple docstring'''





  __lowerCAmelCase       :			Any     	= LDMTextToImagePipeline
  __lowerCAmelCase       :			Union[str, Any]     	= TEXT_TO_IMAGE_PARAMS - {
      "negative_prompt",
      "negative_prompt_embeds",
      "cross_attention_kwargs",
      "prompt_embeds",
  }
  __lowerCAmelCase       :			Optional[Any]     	= PipelineTesterMixin.required_optional_params - {
      "num_images_per_prompt",
      "callback",
      "callback_steps",
  }
  __lowerCAmelCase       :			str     	= TEXT_TO_IMAGE_BATCH_PARAMS
  __lowerCAmelCase       :			Tuple     	= False
  def     snake_case_      (  self  )				->					List[str]:


    """simple docstring"""
    torch.manual_seed(0  )
    UpperCAmelCase    				=						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 ,			)
    UpperCAmelCase    				=						DDIMScheduler(
        beta_start=0.0_0085 ,			beta_end=0.012 ,			beta_schedule='scaled_linear' ,			clip_sample=__A ,			set_alpha_to_one=__A ,			)
    torch.manual_seed(0  )
    UpperCAmelCase    				=						AutoencoderKL(
        block_out_channels=(3_2, 6_4) ,			in_channels=3 ,			out_channels=3 ,			down_block_types=('DownEncoderBlock2D', 'DownEncoderBlock2D') ,			up_block_types=('UpDecoderBlock2D', 'UpDecoderBlock2D') ,			latent_channels=4 ,			)
    torch.manual_seed(0  )
    UpperCAmelCase    				=						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 ,			)
    UpperCAmelCase    				=						CLIPTextModel(__A  )
    UpperCAmelCase    				=						CLIPTokenizer.from_pretrained('hf-internal-testing/tiny-random-clip'  )

    UpperCAmelCase    				=						{
        'unet': unet,
        'scheduler': scheduler,
        'vqvae': vae,
        'bert': text_encoder,
        'tokenizer': tokenizer,
    }
    return components
  def     snake_case_      (  self ,			a_ ,			a_=0  )				->					Optional[int]:


    """simple docstring"""
    if str(__A  ).startswith('mps'  ):
      UpperCAmelCase    				=						torch.manual_seed(__A  )
    else:
      UpperCAmelCase    				=						torch.Generator(device=__A  ).manual_seed(__A  )
    UpperCAmelCase    				=						{
        '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     snake_case_      (  self  )				->					Dict:


    """simple docstring"""
    UpperCAmelCase    				=						'cpu'  # ensure determinism for the device-dependent torch.Generator

    UpperCAmelCase    				=						self.get_dummy_components()
    UpperCAmelCase    				=						LDMTextToImagePipeline(**__A  )
    pipe.to(__A  )
    pipe.set_progress_bar_config(disable=__A  )

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

    assert image.shape == (1, 1_6, 1_6, 3)
    UpperCAmelCase    				=						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     snake_case_      (  self  )				->					Optional[int]:


    """simple docstring"""
    super().tearDown()
    gc.collect()
    torch.cuda.empty_cache()
  def     snake_case_      (  self ,			a_ ,			a_=torch.floataa ,			a_=0  )				->					Dict:


    """simple docstring"""
    UpperCAmelCase    				=						torch.manual_seed(__A  )
    UpperCAmelCase    				=						np.random.RandomState(__A  ).standard_normal((1, 4, 3_2, 3_2)  )
    UpperCAmelCase    				=						torch.from_numpy(__A  ).to(device=__A ,			dtype=__A  )
    UpperCAmelCase    				=						{
        '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     snake_case_      (  self  )				->					Tuple:


    """simple docstring"""
    UpperCAmelCase    				=						LDMTextToImagePipeline.from_pretrained('CompVis/ldm-text2im-large-256'  ).to(__A  )
    pipe.set_progress_bar_config(disable=__A  )

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

    assert image.shape == (1, 2_5_6, 2_5_6, 3)
    UpperCAmelCase    				=						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]  )
    UpperCAmelCase    				=						np.abs(expected_slice - image_slice  ).max()
    assert max_diff < 1E-3




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

  '''simple docstring'''





  def     snake_case_      (  self  )				->					List[Any]:


    """simple docstring"""
    super().tearDown()
    gc.collect()
    torch.cuda.empty_cache()
  def     snake_case_      (  self ,			a_ ,			a_=torch.floataa ,			a_=0  )				->					List[str]:


    """simple docstring"""
    UpperCAmelCase    				=						torch.manual_seed(__A  )
    UpperCAmelCase    				=						np.random.RandomState(__A  ).standard_normal((1, 4, 3_2, 3_2)  )
    UpperCAmelCase    				=						torch.from_numpy(__A  ).to(device=__A ,			dtype=__A  )
    UpperCAmelCase    				=						{
        'prompt': 'A painting of a squirrel eating a burger',
        'latents': latents,
        'generator': generator,
        'num_inference_steps': 5_0,
        'guidance_scale': 6.0,
        'output_type': 'numpy',
    }
    return inputs





  def     snake_case_      (  self  )				->					Union[str, Any]:


    """simple docstring"""
    UpperCAmelCase    				=						LDMTextToImagePipeline.from_pretrained('CompVis/ldm-text2im-large-256'  ).to(__A  )
    pipe.set_progress_bar_config(disable=__A  )

    UpperCAmelCase    				=						self.get_inputs(__A  )
    UpperCAmelCase    				=						pipe(**__A  ).images[0]

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






from typing import TYPE_CHECKING

# rely on isort to merge the imports
from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_torch_available


SCREAMING_SNAKE_CASE							    =							{
    'configuration_autoformer': [
        'AUTOFORMER_PRETRAINED_CONFIG_ARCHIVE_MAP',
        'AutoformerConfig',
    ],
}

try:
  if not is_torch_available():
    raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
  pass
else:
  SCREAMING_SNAKE_CASE							    =							[
      'AUTOFORMER_PRETRAINED_MODEL_ARCHIVE_LIST',
      'AutoformerForPrediction',
      'AutoformerModel',
      'AutoformerPreTrainedModel',
  ]


if TYPE_CHECKING:
  from .configuration_autoformer import (
      AUTOFORMER_PRETRAINED_CONFIG_ARCHIVE_MAP,
      AutoformerConfig,
  )

  try:
    if not is_torch_available():
      raise OptionalDependencyNotAvailable()
  except OptionalDependencyNotAvailable:
    pass
  else:
    from .modeling_autoformer import (
        AUTOFORMER_PRETRAINED_MODEL_ARCHIVE_LIST,
        AutoformerForPrediction,
        AutoformerModel,
        AutoformerPreTrainedModel,
    )

else:
  import sys

  SCREAMING_SNAKE_CASE							    =							_LazyModule(__name__, globals()['__file__'], _import_structure, module_spec=__spec__)



