from typing import TYPE_CHECKING

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


__snake_case         =     {'''configuration_gptj''': ['''GPTJ_PRETRAINED_CONFIG_ARCHIVE_MAP''', '''GPTJConfig''', '''GPTJOnnxConfig''']}

try:
     if not is_torch_available():
          raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
     pass
else:
     __snake_case         =     [
         '''GPTJ_PRETRAINED_MODEL_ARCHIVE_LIST''',
         '''GPTJForCausalLM''',
         '''GPTJForQuestionAnswering''',
         '''GPTJForSequenceClassification''',
         '''GPTJModel''',
         '''GPTJPreTrainedModel''',
     ]

try:
     if not is_tf_available():
          raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
     pass
else:
     __snake_case         =     [
         '''TFGPTJForCausalLM''',
         '''TFGPTJForQuestionAnswering''',
         '''TFGPTJForSequenceClassification''',
         '''TFGPTJModel''',
         '''TFGPTJPreTrainedModel''',
     ]

try:
     if not is_flax_available():
          raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
     pass
else:
     __snake_case         =     [
         '''FlaxGPTJForCausalLM''',
         '''FlaxGPTJModel''',
         '''FlaxGPTJPreTrainedModel''',
     ]


if TYPE_CHECKING:
     from .configuration_gptj import GPTJ_PRETRAINED_CONFIG_ARCHIVE_MAP, GPTJConfig, GPTJOnnxConfig

     try:
          if not is_torch_available():
               raise OptionalDependencyNotAvailable()
     except OptionalDependencyNotAvailable:
          pass
     else:
          from .modeling_gptj import (
              GPTJ_PRETRAINED_MODEL_ARCHIVE_LIST,
              GPTJForCausalLM,
              GPTJForQuestionAnswering,
              GPTJForSequenceClassification,
              GPTJModel,
              GPTJPreTrainedModel,
          )

     try:
          if not is_tf_available():
               raise OptionalDependencyNotAvailable()
     except OptionalDependencyNotAvailable:
          pass
     else:
          from .modeling_tf_gptj import (
              TFGPTJForCausalLM,
              TFGPTJForQuestionAnswering,
              TFGPTJForSequenceClassification,
              TFGPTJModel,
              TFGPTJPreTrainedModel,
          )

     try:
          if not is_flax_available():
               raise OptionalDependencyNotAvailable()
     except OptionalDependencyNotAvailable:
          pass
     else:
          from .modeling_flax_gptj import FlaxGPTJForCausalLM, FlaxGPTJModel, FlaxGPTJPreTrainedModel

else:
     import sys

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

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       __snake_case (  lowerCamelCase__       , lowerCamelCase__       , unittest.TestCase ):
 __lowerCamelCase     : int 	=   StableDiffusionXLPipeline
 __lowerCamelCase     : Any 	=   TEXT_TO_IMAGE_PARAMS
 __lowerCamelCase     : List[str] 	=   TEXT_TO_IMAGE_BATCH_PARAMS
 __lowerCamelCase     : Optional[int] 	=   TEXT_TO_IMAGE_IMAGE_PARAMS
 __lowerCamelCase     : Union[str, Any] 	=   TEXT_TO_IMAGE_IMAGE_PARAMS
 def 						UpperCAmelCase__							(	self    )						->				Optional[Any]:
       '''simple docstring'''





       torch.manual_seed(0    )
       UpperCAmelCase     :			Union[str, Any]     	=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=snake_case__		, 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		, )
       UpperCAmelCase     :			Optional[int]     	=EulerDiscreteScheduler(
           beta_start=0.0_0085		, beta_end=0.012		, steps_offset=1		, beta_schedule='''scaled_linear'''		, timestep_spacing='''leading'''		, )
       torch.manual_seed(0    )
       UpperCAmelCase     :			int     	=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=128		, )
       torch.manual_seed(0    )
       UpperCAmelCase     :			Optional[Any]     	=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		, hidden_act='''gelu'''		, projection_dim=32		, )
       UpperCAmelCase     :			List[str]     	=CLIPTextModel(snake_case__    )
       UpperCAmelCase     :			int     	=CLIPTokenizer.from_pretrained('''hf-internal-testing/tiny-random-clip'''		, local_files_only=snake_case__    )

       UpperCAmelCase     :			Tuple     	=CLIPTextModelWithProjection(snake_case__    )
       UpperCAmelCase     :			int     	=CLIPTokenizer.from_pretrained('''hf-internal-testing/tiny-random-clip'''		, local_files_only=snake_case__    )

       UpperCAmelCase     :			Optional[int]     	={
           '''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 						UpperCAmelCase__							(	self		, snake_case__		, snake_case__=0    )						->				Optional[Any]:
       '''simple docstring'''





       if str(snake_case__    ).startswith('''mps'''    ):
             UpperCAmelCase     :			Union[str, Any]     	=torch.manual_seed(snake_case__    )
       else:
             UpperCAmelCase     :			Optional[int]     	=torch.Generator(device=snake_case__    ).manual_seed(snake_case__    )
       UpperCAmelCase     :			Any     	={
           '''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 						UpperCAmelCase__							(	self    )						->				List[Any]:
       '''simple docstring'''





       UpperCAmelCase     :			Optional[Any]     	='''cpu'''  # ensure determinism for the device-dependent torch.Generator
       UpperCAmelCase     :			List[Any]     	=self.get_dummy_components()
       UpperCAmelCase     :			Tuple     	=StableDiffusionXLPipeline(**snake_case__    )
       UpperCAmelCase     :			Optional[Any]     	=sd_pipe.to(snake_case__    )
       sd_pipe.set_progress_bar_config(disable=snake_case__    )

       UpperCAmelCase     :			Optional[Any]     	=self.get_dummy_inputs(snake_case__    )
       UpperCAmelCase     :			Optional[int]     	=sd_pipe(**snake_case__    ).images
       UpperCAmelCase     :			str     	=image[0, -3:, -3:, -1]

       assert image.shape == (1, 64, 64, 3)
       UpperCAmelCase     :			Dict     	=np.array([0.5873, 0.6128, 0.4797, 0.5122, 0.5674, 0.4639, 0.5227, 0.5149, 0.4747]    )

       assert np.abs(image_slice.flatten() - expected_slice    ).max() < 1e-2
 def 						UpperCAmelCase__							(	self    )						->				List[str]:
       '''simple docstring'''





       UpperCAmelCase     :			Optional[Any]     	=self.get_dummy_components()
       UpperCAmelCase     :			List[Any]     	=StableDiffusionXLPipeline(**snake_case__    )
       UpperCAmelCase     :			List[str]     	=sd_pipe.to(snake_case__    )
       UpperCAmelCase     :			Optional[int]     	=sd_pipe.to(snake_case__    )
       sd_pipe.set_progress_bar_config(disable=snake_case__    )

       # forward without prompt embeds
       UpperCAmelCase     :			int     	=self.get_dummy_inputs(snake_case__    )
       UpperCAmelCase     :			int     	=3 * ['''this is a negative prompt''']
       UpperCAmelCase     :			Union[str, Any]     	=negative_prompt
       UpperCAmelCase     :			Union[str, Any]     	=3 * [inputs['''prompt''']]

       UpperCAmelCase     :			int     	=sd_pipe(**snake_case__    )
       UpperCAmelCase     :			str     	=output.images[0, -3:, -3:, -1]

       # forward with prompt embeds
       UpperCAmelCase     :			Any     	=self.get_dummy_inputs(snake_case__    )
       UpperCAmelCase     :			List[str]     	=3 * ['''this is a negative prompt''']
       UpperCAmelCase     :			int     	=3 * [inputs.pop('''prompt'''    )]

       (
           (
           UpperCAmelCase
       )	,    (
           UpperCAmelCase
       )	,    (
           UpperCAmelCase
       )	,    (
           UpperCAmelCase
       )	,    
       )     :			Union[str, Any]     	=sd_pipe.encode_prompt(snake_case__		, negative_prompt=snake_case__    )

       UpperCAmelCase     :			Union[str, Any]     	=sd_pipe(
           **snake_case__		, prompt_embeds=snake_case__		, negative_prompt_embeds=snake_case__		, pooled_prompt_embeds=snake_case__		, negative_pooled_prompt_embeds=snake_case__		, )
       UpperCAmelCase     :			Union[str, Any]     	=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 						UpperCAmelCase__							(	self    )						->				List[Any]:
       '''simple docstring'''





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






 def 						UpperCAmelCase__							(	self    )						->				Any:
       '''simple docstring'''





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







@slow
@require_torch_gpu
class       __snake_case (  unittest.TestCase ):
 def 						UpperCAmelCase__							(	self    )						->				List[str]:
       '''simple docstring'''





       super().tearDown()
       gc.collect()
       torch.cuda.empty_cache()
 def 						UpperCAmelCase__							(	self		, snake_case__		, snake_case__="cpu"		, snake_case__=torch.floataa		, snake_case__=0    )						->				Any:
       '''simple docstring'''





       UpperCAmelCase     :			str     	=torch.Generator(device=snake_case__    ).manual_seed(snake_case__    )
       UpperCAmelCase     :			Optional[Any]     	=np.random.RandomState(snake_case__    ).standard_normal((1, 4, 64, 64)    )
       UpperCAmelCase     :			List[Any]     	=torch.from_numpy(snake_case__    ).to(device=snake_case__		, dtype=snake_case__    )
       UpperCAmelCase     :			Optional[Any]     	={
           '''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 						UpperCAmelCase__							(	self    )						->				str:
       '''simple docstring'''





       UpperCAmelCase     :			List[Any]     	=DiffusionPipeline.from_pretrained('''stabilityai/stable-diffusion-2-base'''    )
       pipe.to(snake_case__    )
       pipe.set_progress_bar_config(disable=snake_case__    )

       UpperCAmelCase     :			List[Any]     	=self.get_inputs(snake_case__    )
       UpperCAmelCase     :			str     	=pipe(**snake_case__    ).images
       UpperCAmelCase     :			int     	=image[0, -3:, -3:, -1].flatten()

       assert image.shape == (1, 512, 512, 3)
       UpperCAmelCase     :			Any     	=np.array([0.4_9493, 0.4_7896, 0.4_0798, 0.5_4214, 0.5_3212, 0.4_8202, 0.4_7656, 0.4_6329, 0.4_8506]    )
       assert np.abs(image_slice - expected_slice    ).max() < 7e-3
