'''simple docstring'''




import gc
import unittest

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

from diffusers import (
    AutoencoderKL,
    DiffusionPipeline,
    EulerDiscreteScheduler,
    StableDiffusionXLPipeline,
    UNetaDConditionModel,
)
from diffusers.utils import 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
from ..test_pipelines_common import PipelineLatentTesterMixin, PipelineTesterMixin


enable_full_determinism()




class 					a__   (     UpperCAmelCase__							,     UpperCAmelCase__							,     unittest.TestCase				):
     lowerCamelCase						:  Any							=StableDiffusionXLPipeline
     lowerCamelCase						:  Union[str, Any]							=TEXT_TO_IMAGE_PARAMS
     lowerCamelCase						:  Optional[Any]							=TEXT_TO_IMAGE_BATCH_PARAMS
     lowerCamelCase						:  int							=TEXT_TO_IMAGE_IMAGE_PARAMS
     lowerCamelCase						:  List[Any]							=TEXT_TO_IMAGE_IMAGE_PARAMS






     def      SCREAMING_SNAKE_CASE__     (				self	:	Any					):





         """simple docstring"""



         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''')						,       attention_head_dim=(2, 4)						,       use_linear_projection=a						,       addition_embed_type='''text_time'''						,       addition_time_embed_dim=8						,       transformer_layers_per_block=(1, 2)						,       projection_class_embeddings_input_dim=80						,       cross_attention_dim=64						,       )
         __lowerCamelCase						=						EulerDiscreteScheduler(
             beta_start=0.0_00_85						,       beta_end=0.0_12						,       steps_offset=1						,       beta_schedule='''scaled_linear'''						,       timestep_spacing='''leading'''						,       )
         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						,       sample_size=1_28						,       )
         torch.manual_seed(0					)
         __lowerCamelCase						=						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=10_00						,       hidden_act='''gelu'''						,       projection_dim=32						,       )
         __lowerCamelCase						=						CLIPTextModel(a					)
         __lowerCamelCase						=						CLIPTokenizer.from_pretrained('''hf-internal-testing/tiny-random-clip'''						,       local_files_only=a					)

         __lowerCamelCase						=						CLIPTextModelWithProjection(a					)
         __lowerCamelCase						=						CLIPTokenizer.from_pretrained('''hf-internal-testing/tiny-random-clip'''						,       local_files_only=a					)

         __lowerCamelCase						=						{
             '''unet''': unet,
             '''scheduler''': scheduler,
             '''vae''': vae,
             '''text_encoder''': text_encoder,
             '''tokenizer''': tokenizer,
             '''text_encoder_2''': text_encoder_a,
             '''tokenizer_2''': tokenizer_a,
             # "safety_checker": None,
             # "feature_extractor": None,
         }
         return components






     def      SCREAMING_SNAKE_CASE__     (				self	:	List[Any]						,       a	:	Optional[Any]						,       a	:	Any=0					):





         """simple docstring"""



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






     def      SCREAMING_SNAKE_CASE__     (				self	:	Any					):





         """simple docstring"""



         __lowerCamelCase						=						'''cpu'''  # ensure determinism for the device-dependent torch.Generator
         __lowerCamelCase						=						self.get_dummy_components()
         __lowerCamelCase						=						StableDiffusionXLPipeline(**a					)
         __lowerCamelCase						=						sd_pipe.to(a					)
         sd_pipe.set_progress_bar_config(disable=a					)

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

         assert image.shape == (1, 64, 64, 3)
         __lowerCamelCase						=						np.array([0.58_73, 0.61_28, 0.47_97, 0.51_22, 0.56_74, 0.46_39, 0.52_27, 0.51_49, 0.47_47]					)

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






     def      SCREAMING_SNAKE_CASE__     (				self	:	Optional[Any]					):





         """simple docstring"""



         __lowerCamelCase						=						self.get_dummy_components()
         __lowerCamelCase						=						StableDiffusionXLPipeline(**a					)
         __lowerCamelCase						=						sd_pipe.to(a					)
         __lowerCamelCase						=						sd_pipe.to(a					)
         sd_pipe.set_progress_bar_config(disable=a					)

         # forward without prompt embeds
         __lowerCamelCase						=						self.get_dummy_inputs(a					)
         __lowerCamelCase						=						3 * ['''this is a negative prompt''']
         __lowerCamelCase						=						negative_prompt
         __lowerCamelCase						=						3 * [inputs['''prompt''']]

         __lowerCamelCase						=						sd_pipe(**a					)
         __lowerCamelCase						=						output.images[0, -3:, -3:, -1]

         # forward with prompt embeds
         __lowerCamelCase						=						self.get_dummy_inputs(a					)
         __lowerCamelCase						=						3 * ['''this is a negative prompt''']
         __lowerCamelCase						=						3 * [inputs.pop('''prompt'''					)]

         (
             (
             __lowerCamelCase
         )			,						(
             __lowerCamelCase
         )			,						(
             __lowerCamelCase
         )			,						(
             __lowerCamelCase
         )			,						
         )						=						sd_pipe.encode_prompt(a						,       negative_prompt=a					)

         __lowerCamelCase						=						sd_pipe(
             **a						,       prompt_embeds=a						,       negative_prompt_embeds=a						,       pooled_prompt_embeds=a						,       negative_pooled_prompt_embeds=a						,       )
         __lowerCamelCase						=						output.images[0, -3:, -3:, -1]

         # make sure that it's equal
         assert np.abs(image_slice_a.flatten() - image_slice_a.flatten()					).max() < 1e-4






     def      SCREAMING_SNAKE_CASE__     (				self	:	List[str]					):





         """simple docstring"""



         super().test_attention_slicing_forward_pass(expected_max_diff=3e-3					)






     def      SCREAMING_SNAKE_CASE__     (				self	:	Optional[int]					):





         """simple docstring"""



         super().test_inference_batch_single_identical(expected_max_diff=3e-3					)




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






     def      SCREAMING_SNAKE_CASE__     (				self	:	Optional[int]					):





         """simple docstring"""



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






     def      SCREAMING_SNAKE_CASE__     (				self	:	Optional[int]						,       a	:	Optional[int]						,       a	:	Any="cpu"						,       a	:	Optional[int]=torch.floataa						,       a	:	Union[str, Any]=0					):





         """simple docstring"""



         __lowerCamelCase						=						torch.Generator(device=a					).manual_seed(a					)
         __lowerCamelCase						=						np.random.RandomState(a					).standard_normal((1, 4, 64, 64)					)
         __lowerCamelCase						=						torch.from_numpy(a					).to(device=a						,       dtype=a					)
         __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      SCREAMING_SNAKE_CASE__     (				self	:	Dict					):





         """simple docstring"""



         __lowerCamelCase						=						DiffusionPipeline.from_pretrained('''stabilityai/stable-diffusion-2-base'''					)
         pipe.to(a					)
         pipe.set_progress_bar_config(disable=a					)

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

         assert image.shape == (1, 5_12, 5_12, 3)
         __lowerCamelCase						=						np.array([0.4_94_93, 0.4_78_96, 0.4_07_98, 0.5_42_14, 0.5_32_12, 0.4_82_02, 0.4_76_56, 0.4_63_29, 0.4_85_06]					)
         assert np.abs(image_slice - expected_slice					).max() < 7e-3

'''simple docstring'''




from typing import TYPE_CHECKING

from ...utils import (
    OptionalDependencyNotAvailable,
    _LazyModule,
    is_torch_available,
)


__UpperCAmelCase  							={
    "configuration_rwkv": ["RWKV_PRETRAINED_CONFIG_ARCHIVE_MAP", "RwkvConfig", "RwkvOnnxConfig"],
}

try:
  if not is_torch_available():
    raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
  pass
else:
  __UpperCAmelCase  							=[
      "RWKV_PRETRAINED_MODEL_ARCHIVE_LIST",
      "RwkvForCausalLM",
      "RwkvModel",
      "RwkvPreTrainedModel",
  ]


if TYPE_CHECKING:
  from .configuration_rwkv import RWKV_PRETRAINED_CONFIG_ARCHIVE_MAP, RwkvConfig, RwkvOnnxConfig

  try:
    if not is_torch_available():
      raise OptionalDependencyNotAvailable()
  except OptionalDependencyNotAvailable:
    pass
  else:
    from .modeling_rwkv import (
        RWKV_PRETRAINED_MODEL_ARCHIVE_LIST,
        RwkvForCausalLM,
        RwkvModel,
        RwkvPreTrainedModel,
    )
else:
  import sys

  __UpperCAmelCase  							=_LazyModule(__name__, globals()["__file__"], _import_structure, module_spec=__spec__)
