


import gc
import unittest

import numpy as np
import torch
from transformers import CLIPTextConfig, CLIPTextModel, XLMRobertaTokenizer

from diffusers import AltDiffusionPipeline, AutoencoderKL, DDIMScheduler, PNDMScheduler, UNetaDConditionModel
from diffusers.pipelines.alt_diffusion.modeling_roberta_series import (
    RobertaSeriesConfig,
    RobertaSeriesModelWithTransformation,
)
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 PipelineKarrasSchedulerTesterMixin, PipelineLatentTesterMixin, PipelineTesterMixin


enable_full_determinism()

class       __magic_name__	(     __lowerCAmelCase					,				__lowerCAmelCase					,				__lowerCAmelCase					,				unittest.TestCase):
     A:		Optional[Any]			=      AltDiffusionPipeline
     A:		Tuple			=      TEXT_TO_IMAGE_PARAMS
     A:		Optional[int]			=      TEXT_TO_IMAGE_BATCH_PARAMS
     A:		Dict			=      TEXT_TO_IMAGE_IMAGE_PARAMS
     A:		Any			=      TEXT_TO_IMAGE_IMAGE_PARAMS


     def      UpperCAmelCase__	(      self							:     int			)  ->      Dict:





      '''simple docstring'''






      torch.manual_seed(0			)
      UpperCamelCase__					:    str          =   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							,	)
      UpperCamelCase__					:    Tuple          =   DDIMScheduler(
          beta_start=0.0_0085							,	beta_end=0.012							,	beta_schedule='''scaled_linear'''							,	clip_sample=lowerCamelCase__							,	set_alpha_to_one=lowerCamelCase__							,	)
      torch.manual_seed(0			)
      UpperCamelCase__					:    List[str]          =   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							,	)

      # TODO: address the non-deterministic text encoder (fails for save-load tests)
      # torch.manual_seed(0)
      # text_encoder_config = RobertaSeriesConfig(
      #     hidden_size=32,
      #     project_dim=32,
      #     intermediate_size=37,
      #     layer_norm_eps=1e-05,
      #     num_attention_heads=4,
      #     num_hidden_layers=5,
      #     vocab_size=5002,
      # )
      # text_encoder = RobertaSeriesModelWithTransformation(text_encoder_config)

      torch.manual_seed(0			)
      UpperCamelCase__					:    Union[str, Any]          =   CLIPTextConfig(
          bos_token_id=0							,	eos_token_id=2							,	hidden_size=32							,	projection_dim=32							,	intermediate_size=37							,	layer_norm_eps=1E-0_5							,	num_attention_heads=4							,	num_hidden_layers=5							,	pad_token_id=1							,	vocab_size=5002							,	)
      UpperCamelCase__					:    List[Any]          =   CLIPTextModel(lowerCamelCase__			)

      UpperCamelCase__					:    List[Any]          =   XLMRobertaTokenizer.from_pretrained('''hf-internal-testing/tiny-xlm-roberta'''			)
      UpperCamelCase__					:    Tuple          =   77

      UpperCamelCase__					:    Dict          =   {
          '''unet''': unet,
          '''scheduler''': scheduler,
          '''vae''': vae,
          '''text_encoder''': text_encoder,
          '''tokenizer''': tokenizer,
          '''safety_checker''': None,
          '''feature_extractor''': None,
      }
      return components


     def      UpperCAmelCase__	(      self							:     Tuple							,	lowerCamelCase__							:     Optional[int]							,	lowerCamelCase__							:     Tuple=0			)  ->      Union[str, Any]:





      '''simple docstring'''






      if str(lowerCamelCase__			).startswith('''mps'''			):
       UpperCamelCase__					:    Optional[int]          =   torch.manual_seed(lowerCamelCase__			)
      else:
       UpperCamelCase__					:    Dict          =   torch.Generator(device=lowerCamelCase__			).manual_seed(lowerCamelCase__			)
      UpperCamelCase__					:    Dict          =   {
          '''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							:     int			)  ->      Union[str, Any]:





      '''simple docstring'''






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


     def      UpperCAmelCase__	(      self							:     List[str]			)  ->      Dict:





      '''simple docstring'''






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


     def      UpperCAmelCase__	(      self							:     int			)  ->      Tuple:





      '''simple docstring'''






      UpperCamelCase__					:    Any          =   '''cpu'''  # ensure determinism for the device-dependent torch.Generator

      UpperCamelCase__					:    Optional[int]          =   self.get_dummy_components()
      torch.manual_seed(0			)
      UpperCamelCase__					:    str          =   RobertaSeriesConfig(
          hidden_size=32							,	project_dim=32							,	intermediate_size=37							,	layer_norm_eps=1E-0_5							,	num_attention_heads=4							,	num_hidden_layers=5							,	vocab_size=5002							,	)
      # TODO: remove after fixing the non-deterministic text encoder
      UpperCamelCase__					:    int          =   RobertaSeriesModelWithTransformation(lowerCamelCase__			)
      UpperCamelCase__					:    Any          =   text_encoder

      UpperCamelCase__					:    Dict          =   AltDiffusionPipeline(**lowerCamelCase__			)
      UpperCamelCase__					:    Any          =   alt_pipe.to(lowerCamelCase__			)
      alt_pipe.set_progress_bar_config(disable=lowerCamelCase__			)

      UpperCamelCase__					:    Dict          =   self.get_dummy_inputs(lowerCamelCase__			)
      UpperCamelCase__					:    Optional[int]          =   '''A photo of an astronaut'''
      UpperCamelCase__					:    int          =   alt_pipe(**lowerCamelCase__			)
      UpperCamelCase__					:    Tuple          =   output.images
      UpperCamelCase__					:    str          =   image[0, -3:, -3:, -1]

      assert image.shape == (1, 64, 64, 3)
      UpperCamelCase__					:    List[Any]          =   np.array(
          [0.574_8162, 0.6044_7145, 0.4882_1217, 0.5010_0636, 0.543_1185, 0.4576_3683, 0.4965_7696, 0.4813_2733, 0.4757_3093]			)

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


     def      UpperCAmelCase__	(      self							:     Optional[int]			)  ->      Dict:





      '''simple docstring'''






      UpperCamelCase__					:    List[str]          =   '''cpu'''  # ensure determinism for the device-dependent torch.Generator

      UpperCamelCase__					:    Optional[Any]          =   self.get_dummy_components()
      UpperCamelCase__					:    Optional[Any]          =   PNDMScheduler(skip_prk_steps=lowerCamelCase__			)
      torch.manual_seed(0			)
      UpperCamelCase__					:    Optional[int]          =   RobertaSeriesConfig(
          hidden_size=32							,	project_dim=32							,	intermediate_size=37							,	layer_norm_eps=1E-0_5							,	num_attention_heads=4							,	num_hidden_layers=5							,	vocab_size=5002							,	)
      # TODO: remove after fixing the non-deterministic text encoder
      UpperCamelCase__					:    int          =   RobertaSeriesModelWithTransformation(lowerCamelCase__			)
      UpperCamelCase__					:    str          =   text_encoder
      UpperCamelCase__					:    str          =   AltDiffusionPipeline(**lowerCamelCase__			)
      UpperCamelCase__					:    Optional[Any]          =   alt_pipe.to(lowerCamelCase__			)
      alt_pipe.set_progress_bar_config(disable=lowerCamelCase__			)

      UpperCamelCase__					:    List[Any]          =   self.get_dummy_inputs(lowerCamelCase__			)
      UpperCamelCase__					:    Tuple          =   alt_pipe(**lowerCamelCase__			)
      UpperCamelCase__					:    int          =   output.images
      UpperCamelCase__					:    Tuple          =   image[0, -3:, -3:, -1]

      assert image.shape == (1, 64, 64, 3)
      UpperCamelCase__					:    List[Any]          =   np.array(
          [0.5160_5093, 0.570_7241, 0.4736_5507, 0.5057_8886, 0.563_3877, 0.464_2503, 0.518_2081, 0.4876_3484, 0.4908_4237]			)

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



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


     def      UpperCAmelCase__	(      self							:     Dict			)  ->      List[str]:





      '''simple docstring'''






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


     def      UpperCAmelCase__	(      self							:     Any			)  ->      Optional[Any]:





      '''simple docstring'''






      UpperCamelCase__					:    int          =   AltDiffusionPipeline.from_pretrained('''BAAI/AltDiffusion'''							,	safety_checker=lowerCamelCase__			)
      UpperCamelCase__					:    Union[str, Any]          =   alt_pipe.to(lowerCamelCase__			)
      alt_pipe.set_progress_bar_config(disable=lowerCamelCase__			)

      UpperCamelCase__					:    Any          =   '''A painting of a squirrel eating a burger'''
      UpperCamelCase__					:    Optional[int]          =   torch.manual_seed(0			)
      UpperCamelCase__					:    Optional[Any]          =   alt_pipe([prompt]							,	generator=lowerCamelCase__							,	guidance_scale=6.0							,	num_inference_steps=20							,	output_type='''np'''			)

      UpperCamelCase__					:    Union[str, Any]          =   output.images

      UpperCamelCase__					:    Tuple          =   image[0, -3:, -3:, -1]

      assert image.shape == (1, 512, 512, 3)
      UpperCamelCase__					:    int          =   np.array([0.1010, 0.0800, 0.0794, 0.0885, 0.0843, 0.0762, 0.0769, 0.0729, 0.0586]			)

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


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





      '''simple docstring'''






      UpperCamelCase__					:    Any          =   DDIMScheduler.from_pretrained('''BAAI/AltDiffusion'''							,	subfolder='''scheduler'''			)

      UpperCamelCase__					:    Optional[int]          =   AltDiffusionPipeline.from_pretrained('''BAAI/AltDiffusion'''							,	scheduler=lowerCamelCase__							,	safety_checker=lowerCamelCase__			)
      UpperCamelCase__					:    Dict          =   alt_pipe.to(lowerCamelCase__			)
      alt_pipe.set_progress_bar_config(disable=lowerCamelCase__			)

      UpperCamelCase__					:    Optional[int]          =   '''A painting of a squirrel eating a burger'''
      UpperCamelCase__					:    Dict          =   torch.manual_seed(0			)

      UpperCamelCase__					:    List[str]          =   alt_pipe([prompt]							,	generator=lowerCamelCase__							,	num_inference_steps=2							,	output_type='''numpy'''			)
      UpperCamelCase__					:    Optional[int]          =   output.images

      UpperCamelCase__					:    Any          =   image[0, -3:, -3:, -1]

      assert image.shape == (1, 512, 512, 3)
      UpperCamelCase__					:    str          =   np.array([0.4019, 0.4052, 0.3810, 0.4119, 0.3916, 0.3982, 0.4651, 0.4195, 0.5323]			)

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






import contextlib
import importlib
import io
import unittest

import transformers

# Try to import everything from transformers to ensure every object can be loaded.
from transformers import *  # noqa F406
from transformers.testing_utils import DUMMY_UNKNOWN_IDENTIFIER, require_flax, require_tf, require_torch
from transformers.utils import ContextManagers, find_labels, is_flax_available, is_tf_available, is_torch_available


if is_torch_available():
       from transformers import BertForPreTraining, BertForQuestionAnswering, BertForSequenceClassification

if is_tf_available():
       from transformers import TFBertForPreTraining, TFBertForQuestionAnswering, TFBertForSequenceClassification

if is_flax_available():
       from transformers import FlaxBertForPreTraining, FlaxBertForQuestionAnswering, FlaxBertForSequenceClassification


__UpperCamelCase       :       Dict						       =				DUMMY_UNKNOWN_IDENTIFIER
# An actual model hosted on huggingface.co

__UpperCamelCase       :       Optional[int]						       =				"main"
# Default branch name
__UpperCamelCase       :       Optional[Any]						       =				"f2c752cfc5c0ab6f4bdec59acea69eefbee381c2"
# One particular commit (not the top of `main`)
__UpperCamelCase       :       str						       =				"aaaaaaa"
# This commit does not exist, so we should 404.

__UpperCamelCase       :       Optional[Any]						       =				"d9e9f15bc825e4b2c9249e9578f884bbcb5e3684"
# Sha-1 of config.json on the top of `main`, for checking purposes
__UpperCamelCase       :       Tuple						       =				"4b243c475af8d0a7754e87d7d096c92e5199ec2fe168a2ee7998e3b8e9bcb1d3"

@contextlib.contextmanager
def 			_a  (				):







 """simple docstring"""



 print('''Welcome!'''     )
 yield
 print('''Bye!'''     )

@contextlib.contextmanager
def 			_a  (				):







 """simple docstring"""



 print('''Bonjour!'''     )
 yield
 print('''Au revoir!'''     )

class       __magic_name__	(     unittest.TestCase):


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





      '''simple docstring'''






      assert transformers.__spec__ is not None
      assert importlib.util.find_spec('''transformers'''			) is not None



class       __magic_name__	(     unittest.TestCase):


     @unittest.mock.patch('''sys.stdout'''							,	new_callable=io.StringIO			)
     def      UpperCAmelCase__	(      self							:     Union[str, Any]							,	lowerCamelCase__							:     Union[str, Any]			)  ->      Optional[Any]:





      '''simple docstring'''






      with ContextManagers([]			):
       print('''Transformers are awesome!'''			)
      # The print statement adds a new line at the end of the output
      self.assertEqual(mock_stdout.getvalue()							,	'''Transformers are awesome!\n'''			)


     @unittest.mock.patch('''sys.stdout'''							,	new_callable=io.StringIO			)
     def      UpperCAmelCase__	(      self							:     Any							,	lowerCamelCase__							:     Union[str, Any]			)  ->      Dict:





      '''simple docstring'''






      with ContextManagers([context_en()]			):
       print('''Transformers are awesome!'''			)
      # The output should be wrapped with an English welcome and goodbye
      self.assertEqual(mock_stdout.getvalue()							,	'''Welcome!\nTransformers are awesome!\nBye!\n'''			)


     @unittest.mock.patch('''sys.stdout'''							,	new_callable=io.StringIO			)
     def      UpperCAmelCase__	(      self							:     List[Any]							,	lowerCamelCase__							:     List[str]			)  ->      Tuple:





      '''simple docstring'''






      with ContextManagers([context_fr(), context_en()]			):
       print('''Transformers are awesome!'''			)
      # The output should be wrapped with an English and French welcome and goodbye
      self.assertEqual(mock_stdout.getvalue()							,	'''Bonjour!\nWelcome!\nTransformers are awesome!\nBye!\nAu revoir!\n'''			)


     @require_torch
     def      UpperCAmelCase__	(      self							:     Tuple			)  ->      Optional[int]:





      '''simple docstring'''






      self.assertEqual(find_labels(lowerCamelCase__			)							,	['''labels''']			)
      self.assertEqual(find_labels(lowerCamelCase__			)							,	['''labels''', '''next_sentence_label''']			)
      self.assertEqual(find_labels(lowerCamelCase__			)							,	['''start_positions''', '''end_positions''']			)

      class       __magic_name__	(     __lowerCAmelCase):
           pass



      self.assertEqual(find_labels(lowerCamelCase__			)							,	['''labels''']			)


     @require_tf
     def      UpperCAmelCase__	(      self							:     Union[str, Any]			)  ->      Tuple:





      '''simple docstring'''






      self.assertEqual(find_labels(lowerCamelCase__			)							,	['''labels''']			)
      self.assertEqual(find_labels(lowerCamelCase__			)							,	['''labels''', '''next_sentence_label''']			)
      self.assertEqual(find_labels(lowerCamelCase__			)							,	['''start_positions''', '''end_positions''']			)

      class       __magic_name__	(     __lowerCAmelCase):
           pass



      self.assertEqual(find_labels(lowerCamelCase__			)							,	['''labels''']			)


     @require_flax
     def      UpperCAmelCase__	(      self							:     str			)  ->      Optional[Any]:





      '''simple docstring'''






      self.assertEqual(find_labels(lowerCamelCase__			)							,	[]			)
      self.assertEqual(find_labels(lowerCamelCase__			)							,	[]			)
      self.assertEqual(find_labels(lowerCamelCase__			)							,	[]			)

      class       __magic_name__	(     __lowerCAmelCase):
           pass



      self.assertEqual(find_labels(lowerCamelCase__			)							,	[]			)


