




import gc
import unittest

import torch
from transformers import CLIPTextConfig, CLIPTextModel, CLIPTextModelWithProjection, CLIPTokenizer

from diffusers import (
    AutoencoderKL,
    DDIMScheduler,
    DDPMScheduler,
    PriorTransformer,
    StableUnCLIPPipeline,
    UNetaDConditionModel,
)
from diffusers.pipelines.stable_diffusion.stable_unclip_image_normalizer import StableUnCLIPImageNormalizer
from diffusers.utils.testing_utils import enable_full_determinism, load_numpy, require_torch_gpu, slow, torch_device

from ..pipeline_params import TEXT_TO_IMAGE_BATCH_PARAMS, TEXT_TO_IMAGE_IMAGE_PARAMS, TEXT_TO_IMAGE_PARAMS
from ..test_pipelines_common import (
    PipelineKarrasSchedulerTesterMixin,
    PipelineLatentTesterMixin,
    PipelineTesterMixin,
    assert_mean_pixel_difference,
)


enable_full_determinism()





class        lowerCamelCase_							(      _SCREAMING_SNAKE_CASE       ,   _SCREAMING_SNAKE_CASE       ,   _SCREAMING_SNAKE_CASE       ,   unittest.TestCase					):


 '''simple docstring'''







 lowercase_				    =    StableUnCLIPPipeline
 lowercase_				    =    TEXT_TO_IMAGE_PARAMS
 lowercase_				    =    TEXT_TO_IMAGE_BATCH_PARAMS
 lowercase_				    =    TEXT_TO_IMAGE_IMAGE_PARAMS
 lowercase_				    =    TEXT_TO_IMAGE_IMAGE_PARAMS

 # TODO(will) Expected attn_bias.stride(1) == 0 to be true, but got false
 lowercase_				    =    False



 def       lowerCAmelCase_  (							self  :		str						):
       SCREAMING_SNAKE_CASE_							=			32
       SCREAMING_SNAKE_CASE_							=			embedder_hidden_size

       # prior components

       torch.manual_seed(0						)
       SCREAMING_SNAKE_CASE_							=			CLIPTokenizer.from_pretrained('hf-internal-testing/tiny-random-clip'						)

       torch.manual_seed(0						)
       SCREAMING_SNAKE_CASE_							=			CLIPTextModelWithProjection(
           CLIPTextConfig(
               bos_token_id=0				,				eos_token_id=2				,				hidden_size=_lowerCAmelCase				,				projection_dim=_lowerCAmelCase				,				intermediate_size=37				,				layer_norm_eps=1E-05				,				num_attention_heads=4				,				num_hidden_layers=5				,				pad_token_id=1				,				vocab_size=1_000				,				)						)

       torch.manual_seed(0						)
       SCREAMING_SNAKE_CASE_							=			PriorTransformer(
           num_attention_heads=2				,				attention_head_dim=12				,				embedding_dim=_lowerCAmelCase				,				num_layers=1				,				)

       torch.manual_seed(0						)
       SCREAMING_SNAKE_CASE_							=			DDPMScheduler(
           variance_type='fixed_small_log'				,				prediction_type='sample'				,				num_train_timesteps=1_000				,				clip_sample=_lowerCAmelCase				,				clip_sample_range=5.0				,				beta_schedule='squaredcos_cap_v2'				,				)

       # regular denoising components

       torch.manual_seed(0						)
       SCREAMING_SNAKE_CASE_							=			StableUnCLIPImageNormalizer(embedding_dim=_lowerCAmelCase						)
       SCREAMING_SNAKE_CASE_							=			DDPMScheduler(beta_schedule='squaredcos_cap_v2'						)

       torch.manual_seed(0						)
       SCREAMING_SNAKE_CASE_							=			CLIPTokenizer.from_pretrained('hf-internal-testing/tiny-random-clip'						)

       torch.manual_seed(0						)
       SCREAMING_SNAKE_CASE_							=			CLIPTextModel(
           CLIPTextConfig(
               bos_token_id=0				,				eos_token_id=2				,				hidden_size=_lowerCAmelCase				,				projection_dim=32				,				intermediate_size=37				,				layer_norm_eps=1E-05				,				num_attention_heads=4				,				num_hidden_layers=5				,				pad_token_id=1				,				vocab_size=1_000				,				)						)

       torch.manual_seed(0						)
       SCREAMING_SNAKE_CASE_							=			UNetaDConditionModel(
           sample_size=32				,				in_channels=4				,				out_channels=4				,				down_block_types=('CrossAttnDownBlock2D', 'DownBlock2D')				,				up_block_types=('UpBlock2D', 'CrossAttnUpBlock2D')				,				block_out_channels=(32, 64)				,				attention_head_dim=(2, 4)				,				class_embed_type='projection'				,				projection_class_embeddings_input_dim=embedder_projection_dim * 2				,				cross_attention_dim=_lowerCAmelCase				,				layers_per_block=1				,				upcast_attention=_lowerCAmelCase				,				use_linear_projection=_lowerCAmelCase				,				)

       torch.manual_seed(0						)
       SCREAMING_SNAKE_CASE_							=			DDIMScheduler(
           beta_schedule='scaled_linear'				,				beta_start=0.0_0085				,				beta_end=0.012				,				prediction_type='v_prediction'				,				set_alpha_to_one=_lowerCAmelCase				,				steps_offset=1				,				)

       torch.manual_seed(0						)
       SCREAMING_SNAKE_CASE_							=			AutoencoderKL()

       SCREAMING_SNAKE_CASE_							=			{
           # prior components
           'prior_tokenizer': prior_tokenizer,
           'prior_text_encoder': prior_text_encoder,
           'prior': prior,
           'prior_scheduler': prior_scheduler,
           # image noising components
           'image_normalizer': image_normalizer,
           'image_noising_scheduler': image_noising_scheduler,
           # regular denoising components
           'tokenizer': tokenizer,
           'text_encoder': text_encoder,
           'unet': unet,
           'scheduler': scheduler,
           'vae': vae,
       }

       return components



 def       lowerCAmelCase_  (							self  :		List[str]				,				_lowerCAmelCase  :		Dict				,				_lowerCAmelCase  :		Dict=0						):
       if str(_lowerCAmelCase						).startswith('mps'						):
             SCREAMING_SNAKE_CASE_							=			torch.manual_seed(_lowerCAmelCase						)
       else:
             SCREAMING_SNAKE_CASE_							=			torch.Generator(device=_lowerCAmelCase						).manual_seed(_lowerCAmelCase						)
       SCREAMING_SNAKE_CASE_							=			{
           'prompt': 'A painting of a squirrel eating a burger',
           'generator': generator,
           'num_inference_steps': 2,
           'prior_num_inference_steps': 2,
           'output_type': 'numpy',
       }
       return inputs



 def       lowerCAmelCase_  (							self  :		Any						):
       SCREAMING_SNAKE_CASE_							=			torch_device == 'cpu'

       self._test_attention_slicing_forward_pass(test_max_difference=_lowerCAmelCase						)




 def       lowerCAmelCase_  (							self  :		Any						):
       SCREAMING_SNAKE_CASE_							=			torch_device in ['cpu', 'mps']

       self._test_inference_batch_single_identical(test_max_difference=_lowerCAmelCase						)





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


 '''simple docstring'''







 def       lowerCAmelCase_  (							self  :		List[Any]						):
       # clean up the VRAM after each test
       super().tearDown()
       gc.collect()
       torch.cuda.empty_cache()



 def       lowerCAmelCase_  (							self  :		Optional[int]						):
       SCREAMING_SNAKE_CASE_							=			load_numpy(
           'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/stable_unclip/stable_unclip_2_1_l_anime_turtle_fp16.npy'						)

       SCREAMING_SNAKE_CASE_							=			StableUnCLIPPipeline.from_pretrained('fusing/stable-unclip-2-1-l'				,				torch_dtype=torch.floataa						)
       pipe.to(_lowerCAmelCase						)
       pipe.set_progress_bar_config(disable=_lowerCAmelCase						)
       # stable unclip will oom when integration tests are run on a V100,
       # so turn on memory savings
       pipe.enable_attention_slicing()
       pipe.enable_sequential_cpu_offload()

       SCREAMING_SNAKE_CASE_							=			torch.Generator(device='cpu'						).manual_seed(0						)
       SCREAMING_SNAKE_CASE_							=			pipe('anime turle'				,				generator=_lowerCAmelCase				,				output_type='np'						)

       SCREAMING_SNAKE_CASE_							=			output.images[0]

       assert image.shape == (768, 768, 3)

       assert_mean_pixel_difference(_lowerCAmelCase				,				_lowerCAmelCase						)




 def       lowerCAmelCase_  (							self  :		str						):
       torch.cuda.empty_cache()
       torch.cuda.reset_max_memory_allocated()
       torch.cuda.reset_peak_memory_stats()

       SCREAMING_SNAKE_CASE_							=			StableUnCLIPPipeline.from_pretrained('fusing/stable-unclip-2-1-l'				,				torch_dtype=torch.floataa						)
       SCREAMING_SNAKE_CASE_							=			pipe.to(_lowerCAmelCase						)
       pipe.set_progress_bar_config(disable=_lowerCAmelCase						)
       pipe.enable_attention_slicing()
       pipe.enable_sequential_cpu_offload()

       SCREAMING_SNAKE_CASE_							=			pipe(
           'anime turtle'				,				prior_num_inference_steps=2				,				num_inference_steps=2				,				output_type='np'				,				)

       SCREAMING_SNAKE_CASE_							=			torch.cuda.max_memory_allocated()
       # make sure that less than 7 GB is allocated
       assert mem_bytes < 7 * 10**9





from typing import TYPE_CHECKING

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


lowerCamelCase__ :   Union[str, Any]					=							{
    'configuration_layoutlmv2': ['LAYOUTLMV2_PRETRAINED_CONFIG_ARCHIVE_MAP', 'LayoutLMv2Config'],
    'processing_layoutlmv2': ['LayoutLMv2Processor'],
    'tokenization_layoutlmv2': ['LayoutLMv2Tokenizer'],
}

try:
   if not is_tokenizers_available():
      raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
   pass
else:
   lowerCamelCase__ :   Any					=							['LayoutLMv2TokenizerFast']

try:
   if not is_vision_available():
      raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
   pass
else:
   lowerCamelCase__ :   Dict					=							['LayoutLMv2FeatureExtractor']
   lowerCamelCase__ :   Dict					=							['LayoutLMv2ImageProcessor']

try:
   if not is_torch_available():
      raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
   pass
else:
   lowerCamelCase__ :   Union[str, Any]					=							[
       'LAYOUTLMV2_PRETRAINED_MODEL_ARCHIVE_LIST',
       'LayoutLMv2ForQuestionAnswering',
       'LayoutLMv2ForSequenceClassification',
       'LayoutLMv2ForTokenClassification',
       'LayoutLMv2Layer',
       'LayoutLMv2Model',
       'LayoutLMv2PreTrainedModel',
   ]

if TYPE_CHECKING:
   from .configuration_layoutlmva import LAYOUTLMV2_PRETRAINED_CONFIG_ARCHIVE_MAP, LayoutLMvaConfig
   from .processing_layoutlmva import LayoutLMvaProcessor
   from .tokenization_layoutlmva import LayoutLMvaTokenizer

   try:
      if not is_tokenizers_available():
         raise OptionalDependencyNotAvailable()
   except OptionalDependencyNotAvailable:
      pass
   else:
      from .tokenization_layoutlmva_fast import LayoutLMvaTokenizerFast

   try:
      if not is_vision_available():
         raise OptionalDependencyNotAvailable()
   except OptionalDependencyNotAvailable:
      pass
   else:
      from .feature_extraction_layoutlmva import LayoutLMvaFeatureExtractor, LayoutLMvaImageProcessor

   try:
      if not is_torch_available():
         raise OptionalDependencyNotAvailable()
   except OptionalDependencyNotAvailable:
      pass
   else:
      from .modeling_layoutlmva import (
          LAYOUTLMV2_PRETRAINED_MODEL_ARCHIVE_LIST,
          LayoutLMvaForQuestionAnswering,
          LayoutLMvaForSequenceClassification,
          LayoutLMvaForTokenClassification,
          LayoutLMvaLayer,
          LayoutLMvaModel,
          LayoutLMvaPreTrainedModel,
      )
else:
   import sys

   lowerCamelCase__ :   Any					=							_LazyModule(__name__, globals()['__file__'], _import_structure, module_spec=__spec__)