



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       _SCREAMING_SNAKE_CASE	(UpperCamelCase			,    UpperCamelCase			,    UpperCamelCase			,    unittest.TestCase ):

     '''simple docstring'''





     lowerCAmelCase					=			AltDiffusionPipeline
     lowerCAmelCase					=			TEXT_TO_IMAGE_PARAMS
     lowerCAmelCase					=			TEXT_TO_IMAGE_BATCH_PARAMS
     lowerCAmelCase					=			TEXT_TO_IMAGE_IMAGE_PARAMS
     lowerCAmelCase					=			TEXT_TO_IMAGE_IMAGE_PARAMS





     def 	__snake_case    (       self :							int				)->Any:
      torch.manual_seed(0				)
      __SCREAMING_SNAKE_CASE       :    List[Any]       =	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     ,  )
      __SCREAMING_SNAKE_CASE       :    str       =	DDIMScheduler(
          beta_start=0.0_0_0_8_5     ,  beta_end=0.0_1_2     ,  beta_schedule="scaled_linear"     ,  clip_sample=UpperCamelCase     ,  set_alpha_to_one=UpperCamelCase     ,  )
      torch.manual_seed(0				)
      __SCREAMING_SNAKE_CASE       :    Tuple       =	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     ,  )

      # 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				)
      __SCREAMING_SNAKE_CASE       :    int       =	CLIPTextConfig(
          bos_token_id=0     ,  eos_token_id=2     ,  hidden_size=3_2     ,  projection_dim=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=5_0_0_2     ,  )
      __SCREAMING_SNAKE_CASE       :    Any       =	CLIPTextModel(UpperCamelCase				)

      __SCREAMING_SNAKE_CASE       :    Union[str, Any]       =	XLMRobertaTokenizer.from_pretrained("hf-internal-testing/tiny-xlm-roberta"				)
      __SCREAMING_SNAKE_CASE       :    List[Any]       =	7_7

      __SCREAMING_SNAKE_CASE       :    int       =	{
          "unet": unet,
          "scheduler": scheduler,
          "vae": vae,
          "text_encoder": text_encoder,
          "tokenizer": tokenizer,
          "safety_checker": None,
          "feature_extractor": None,
      }
      return components





     def 	__snake_case    (       self :							List[Any]     ,  UpperCamelCase :							Dict     ,  UpperCamelCase :							Tuple=0				)->Any:
      if str(UpperCamelCase				).startswith("mps"				):
       __SCREAMING_SNAKE_CASE       :    Optional[Any]       =	torch.manual_seed(UpperCamelCase				)
      else:
       __SCREAMING_SNAKE_CASE       :    Dict       =	torch.Generator(device=UpperCamelCase				).manual_seed(UpperCamelCase				)
      __SCREAMING_SNAKE_CASE       :    Optional[int]       =	{
          "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				)->List[str]:
      super().test_attention_slicing_forward_pass(expected_max_diff=3E-3				)





     def 	__snake_case    (       self :							Optional[int]				)->List[str]:
      super().test_inference_batch_single_identical(expected_max_diff=3E-3				)





     def 	__snake_case    (       self :							int				)->Union[str, Any]:
      __SCREAMING_SNAKE_CASE       :    Tuple       =	"cpu"  # ensure determinism for the device-dependent torch.Generator

      __SCREAMING_SNAKE_CASE       :    List[Any]       =	self.get_dummy_components()
      torch.manual_seed(0				)
      __SCREAMING_SNAKE_CASE       :    str       =	RobertaSeriesConfig(
          hidden_size=3_2     ,  project_dim=3_2     ,  intermediate_size=3_7     ,  layer_norm_eps=1E-05     ,  num_attention_heads=4     ,  num_hidden_layers=5     ,  vocab_size=5_0_0_2     ,  )
      # TODO: remove after fixing the non-deterministic text encoder
      __SCREAMING_SNAKE_CASE       :    Tuple       =	RobertaSeriesModelWithTransformation(UpperCamelCase				)
      __SCREAMING_SNAKE_CASE       :    Optional[int]       =	text_encoder

      __SCREAMING_SNAKE_CASE       :    Tuple       =	AltDiffusionPipeline(**UpperCamelCase				)
      __SCREAMING_SNAKE_CASE       :    Optional[int]       =	alt_pipe.to(UpperCamelCase				)
      alt_pipe.set_progress_bar_config(disable=UpperCamelCase				)

      __SCREAMING_SNAKE_CASE       :    Any       =	self.get_dummy_inputs(UpperCamelCase				)
      __SCREAMING_SNAKE_CASE       :    List[Any]       =	"A photo of an astronaut"
      __SCREAMING_SNAKE_CASE       :    str       =	alt_pipe(**UpperCamelCase				)
      __SCREAMING_SNAKE_CASE       :    List[str]       =	output.images
      __SCREAMING_SNAKE_CASE       :    int       =	image[0, -3:, -3:, -1]

      assert image.shape == (1, 6_4, 6_4, 3)
      __SCREAMING_SNAKE_CASE       :    List[Any]       =	np.array(
          [0.5_7_4_8_1_6_2, 0.6_0_4_4_7_1_4_5, 0.4_8_8_2_1_2_1_7, 0.5_0_1_0_0_6_3_6, 0.5_4_3_1_1_8_5, 0.4_5_7_6_3_6_8_3, 0.4_9_6_5_7_6_9_6, 0.4_8_1_3_2_7_3_3, 0.4_7_5_7_3_0_9_3]				)

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






     def 	__snake_case    (       self :							Dict				)->Any:
      __SCREAMING_SNAKE_CASE       :    str       =	"cpu"  # ensure determinism for the device-dependent torch.Generator

      __SCREAMING_SNAKE_CASE       :    List[str]       =	self.get_dummy_components()
      __SCREAMING_SNAKE_CASE       :    Dict       =	PNDMScheduler(skip_prk_steps=UpperCamelCase				)
      torch.manual_seed(0				)
      __SCREAMING_SNAKE_CASE       :    List[str]       =	RobertaSeriesConfig(
          hidden_size=3_2     ,  project_dim=3_2     ,  intermediate_size=3_7     ,  layer_norm_eps=1E-05     ,  num_attention_heads=4     ,  num_hidden_layers=5     ,  vocab_size=5_0_0_2     ,  )
      # TODO: remove after fixing the non-deterministic text encoder
      __SCREAMING_SNAKE_CASE       :    Tuple       =	RobertaSeriesModelWithTransformation(UpperCamelCase				)
      __SCREAMING_SNAKE_CASE       :    List[str]       =	text_encoder
      __SCREAMING_SNAKE_CASE       :    Any       =	AltDiffusionPipeline(**UpperCamelCase				)
      __SCREAMING_SNAKE_CASE       :    List[str]       =	alt_pipe.to(UpperCamelCase				)
      alt_pipe.set_progress_bar_config(disable=UpperCamelCase				)

      __SCREAMING_SNAKE_CASE       :    Tuple       =	self.get_dummy_inputs(UpperCamelCase				)
      __SCREAMING_SNAKE_CASE       :    List[str]       =	alt_pipe(**UpperCamelCase				)
      __SCREAMING_SNAKE_CASE       :    str       =	output.images
      __SCREAMING_SNAKE_CASE       :    Optional[int]       =	image[0, -3:, -3:, -1]

      assert image.shape == (1, 6_4, 6_4, 3)
      __SCREAMING_SNAKE_CASE       :    Union[str, Any]       =	np.array(
          [0.5_1_6_0_5_0_9_3, 0.5_7_0_7_2_4_1, 0.4_7_3_6_5_5_0_7, 0.5_0_5_7_8_8_8_6, 0.5_6_3_3_8_7_7, 0.4_6_4_2_5_0_3, 0.5_1_8_2_0_8_1, 0.4_8_7_6_3_4_8_4, 0.4_9_0_8_4_2_3_7]				)

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






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

     '''simple docstring'''





     def 	__snake_case    (       self :							Union[str, Any]				)->List[Any]:
      # clean up the VRAM after each test
      super().tearDown()
      gc.collect()
      torch.cuda.empty_cache()





     def 	__snake_case    (       self :							Optional[Any]				)->Optional[int]:
      # make sure here that pndm scheduler skips prk
      __SCREAMING_SNAKE_CASE       :    List[str]       =	AltDiffusionPipeline.from_pretrained("BAAI/AltDiffusion"     ,  safety_checker=UpperCamelCase				)
      __SCREAMING_SNAKE_CASE       :    Optional[int]       =	alt_pipe.to(UpperCamelCase				)
      alt_pipe.set_progress_bar_config(disable=UpperCamelCase				)

      __SCREAMING_SNAKE_CASE       :    int       =	"A painting of a squirrel eating a burger"
      __SCREAMING_SNAKE_CASE       :    Any       =	torch.manual_seed(0				)
      __SCREAMING_SNAKE_CASE       :    Union[str, Any]       =	alt_pipe([prompt]     ,  generator=UpperCamelCase     ,  guidance_scale=6.0     ,  num_inference_steps=2_0     ,  output_type="np"				)

      __SCREAMING_SNAKE_CASE       :    str       =	output.images

      __SCREAMING_SNAKE_CASE       :    Dict       =	image[0, -3:, -3:, -1]

      assert image.shape == (1, 5_1_2, 5_1_2, 3)
      __SCREAMING_SNAKE_CASE       :    Union[str, Any]       =	np.array([0.1_0_1_0, 0.0_8_0_0, 0.0_7_9_4, 0.0_8_8_5, 0.0_8_4_3, 0.0_7_6_2, 0.0_7_6_9, 0.0_7_2_9, 0.0_5_8_6]				)

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






     def 	__snake_case    (       self :							Tuple				)->Any:
      __SCREAMING_SNAKE_CASE       :    List[str]       =	DDIMScheduler.from_pretrained("BAAI/AltDiffusion"     ,  subfolder="scheduler"				)

      __SCREAMING_SNAKE_CASE       :    List[Any]       =	AltDiffusionPipeline.from_pretrained("BAAI/AltDiffusion"     ,  scheduler=UpperCamelCase     ,  safety_checker=UpperCamelCase				)
      __SCREAMING_SNAKE_CASE       :    Tuple       =	alt_pipe.to(UpperCamelCase				)
      alt_pipe.set_progress_bar_config(disable=UpperCamelCase				)

      __SCREAMING_SNAKE_CASE       :    Dict       =	"A painting of a squirrel eating a burger"
      __SCREAMING_SNAKE_CASE       :    List[Any]       =	torch.manual_seed(0				)

      __SCREAMING_SNAKE_CASE       :    List[str]       =	alt_pipe([prompt]     ,  generator=UpperCamelCase     ,  num_inference_steps=2     ,  output_type="numpy"				)
      __SCREAMING_SNAKE_CASE       :    str       =	output.images

      __SCREAMING_SNAKE_CASE       :    Union[str, Any]       =	image[0, -3:, -3:, -1]

      assert image.shape == (1, 5_1_2, 5_1_2, 3)
      __SCREAMING_SNAKE_CASE       :    List[Any]       =	np.array([0.4_0_1_9, 0.4_0_5_2, 0.3_8_1_0, 0.4_1_1_9, 0.3_9_1_6, 0.3_9_8_2, 0.4_6_5_1, 0.4_1_9_5, 0.5_3_2_3]				)

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











import logging
import os
from dataclasses import dataclass
from typing import List, Optional, Union

import tqdm
from filelock import FileLock

from transformers import (
    BartTokenizer,
    BartTokenizerFast,
    DataProcessor,
    PreTrainedTokenizer,
    RobertaTokenizer,
    RobertaTokenizerFast,
    XLMRobertaTokenizer,
    is_tf_available,
    is_torch_available,
)


_lowerCamelCase						    =  logging.getLogger(__name__)

@dataclass(frozen=UpperCamelCase )
class       _SCREAMING_SNAKE_CASE	:
     lowerCAmelCase					=			42
     lowerCAmelCase					=			42
     lowerCAmelCase					=			None
     lowerCAmelCase					=			None
     lowerCAmelCase					=			None






@dataclass(frozen=UpperCamelCase )
class       _SCREAMING_SNAKE_CASE	:
     lowerCAmelCase					=			42
     lowerCAmelCase					=			None
     lowerCAmelCase					=			None
     lowerCAmelCase					=			None
     lowerCAmelCase					=			None


if is_torch_available():
      import torch
      from torch.utils.data import Dataset

      class       _SCREAMING_SNAKE_CASE	(UpperCamelCase ):
           lowerCAmelCase					=			42





           def __init__(       self :							str     ,  UpperCamelCase :							str     ,  UpperCamelCase :							PreTrainedTokenizer     ,  UpperCamelCase :							str     ,  UpperCamelCase :							Optional[int] = None     ,  UpperCamelCase :							Union[str, Any]=False     ,  UpperCamelCase :							bool = False     ,  )->List[str]:
            __SCREAMING_SNAKE_CASE       :    int       =	hans_processors[task]()

            __SCREAMING_SNAKE_CASE       :    Optional[Any]       =	os.path.join(
                UpperCamelCase     ,  "cached_{}_{}_{}_{}".format(
                    "dev" if evaluate else "train"     ,  tokenizer.__class__.__name__     ,  str(UpperCamelCase				)     ,  UpperCamelCase     ,  )     ,  )
            __SCREAMING_SNAKE_CASE       :    Optional[Any]       =	processor.get_labels()
            if tokenizer.__class__ in (
                RobertaTokenizer,
                RobertaTokenizerFast,
                XLMRobertaTokenizer,
                BartTokenizer,
                BartTokenizerFast,
            ):
             # HACK(label indices are swapped in RoBERTa pretrained model)
             __SCREAMING_SNAKE_CASE,		__SCREAMING_SNAKE_CASE       :    Union[str, Any]       =	label_list[2], label_list[1]
            __SCREAMING_SNAKE_CASE       :    Any       =	label_list

            # Make sure only the first process in distributed training processes the dataset,
            # and the others will use the cache.
            __SCREAMING_SNAKE_CASE       :    str       =	cached_features_file + ".lock"
            with FileLock(UpperCamelCase				):
             if os.path.exists(UpperCamelCase				) and not overwrite_cache:
              logger.info(F"""Loading features from cached file {cached_features_file}"""				)
              __SCREAMING_SNAKE_CASE       :    Optional[Any]       =	torch.load(UpperCamelCase				)
             else:
              logger.info(F"""Creating features from dataset file at {data_dir}"""				)

              __SCREAMING_SNAKE_CASE       :    Union[str, Any]       =	(
                  processor.get_dev_examples(UpperCamelCase				) if evaluate else processor.get_train_examples(UpperCamelCase				)
              )

              logger.info("Training examples: %s"     ,  len(UpperCamelCase				)				)
              __SCREAMING_SNAKE_CASE       :    Any       =	hans_convert_examples_to_features(UpperCamelCase     ,  UpperCamelCase     ,  UpperCamelCase     ,  UpperCamelCase				)
              logger.info("Saving features into cached file %s"     ,  UpperCamelCase				)
              torch.save(self.features     ,  UpperCamelCase				)





           def __len__(       self :							Tuple				)->Any:
            return len(self.features				)





           def __getitem__(       self :							Union[str, Any]     ,  UpperCamelCase :							int				)->InputFeatures:
            return self.features[i]






           def 	__snake_case    (       self :							int				)->Tuple:
            return self.label_list


if is_tf_available():
      import tensorflow as tf

      class       _SCREAMING_SNAKE_CASE	:
           lowerCAmelCase					=			42





           def __init__(       self :							Dict     ,  UpperCamelCase :							str     ,  UpperCamelCase :							PreTrainedTokenizer     ,  UpperCamelCase :							str     ,  UpperCamelCase :							Optional[int] = 1_2_8     ,  UpperCamelCase :							Union[str, Any]=False     ,  UpperCamelCase :							bool = False     ,  )->str:
            __SCREAMING_SNAKE_CASE       :    str       =	hans_processors[task]()
            __SCREAMING_SNAKE_CASE       :    Union[str, Any]       =	processor.get_labels()
            if tokenizer.__class__ in (
                RobertaTokenizer,
                RobertaTokenizerFast,
                XLMRobertaTokenizer,
                BartTokenizer,
                BartTokenizerFast,
            ):
             # HACK(label indices are swapped in RoBERTa pretrained model)
             __SCREAMING_SNAKE_CASE,		__SCREAMING_SNAKE_CASE       :    List[str]       =	label_list[2], label_list[1]
            __SCREAMING_SNAKE_CASE       :    Union[str, Any]       =	label_list

            __SCREAMING_SNAKE_CASE       :    str       =	processor.get_dev_examples(UpperCamelCase				) if evaluate else processor.get_train_examples(UpperCamelCase				)
            __SCREAMING_SNAKE_CASE       :    str       =	hans_convert_examples_to_features(UpperCamelCase     ,  UpperCamelCase     ,  UpperCamelCase     ,  UpperCamelCase				)

            def gen():
             for ex_index, ex in tqdm.tqdm(enumerate(self.features				)     ,  desc="convert examples to features"				):
              if ex_index % 1_0_0_0_0 == 0:
               logger.info("Writing example %d of %d" % (ex_index, len(UpperCamelCase				))				)

              yield (
                  {
                      "example_id": 0,
                      "input_ids": ex.input_ids,
                      "attention_mask": ex.attention_mask,
                      "token_type_ids": ex.token_type_ids,
                  },
                  ex.label,
              )

            __SCREAMING_SNAKE_CASE       :    List[str]       =	tf.data.Dataset.from_generator(
                UpperCamelCase     ,  (
                    {
                        "example_id": tf.intaa,
                        "input_ids": tf.intaa,
                        "attention_mask": tf.intaa,
                        "token_type_ids": tf.intaa,
                    },
                    tf.intaa,
                )     ,  (
                    {
                        "example_id": tf.TensorShape([]				),
                        "input_ids": tf.TensorShape([None, None]				),
                        "attention_mask": tf.TensorShape([None, None]				),
                        "token_type_ids": tf.TensorShape([None, None]				),
                    },
                    tf.TensorShape([]				),
                )     ,  )





           def 	__snake_case    (       self :							List[Any]				)->str:
            return self.dataset





           def __len__(       self :							Tuple				)->List[str]:
            return len(self.features				)





           def __getitem__(       self :							Optional[Any]     ,  UpperCamelCase :							Tuple				)->InputFeatures:
            return self.features[i]






           def 	__snake_case    (       self :							List[Any]				)->Optional[int]:
            return self.label_list

class       _SCREAMING_SNAKE_CASE	(UpperCamelCase ):





     def 	__snake_case    (       self :							List[Any]     ,  UpperCamelCase :							Union[str, Any]				)->Tuple:
      return self._create_examples(self._read_tsv(os.path.join(UpperCamelCase     ,  "heuristics_train_set.txt"				)				)     ,  "train"				)





     def 	__snake_case    (       self :							List[str]     ,  UpperCamelCase :							Optional[int]				)->Any:
      return self._create_examples(self._read_tsv(os.path.join(UpperCamelCase     ,  "heuristics_evaluation_set.txt"				)				)     ,  "dev"				)





     def 	__snake_case    (       self :							Optional[int]				)->Tuple:
      return ["contradiction", "entailment", "neutral"]






     def 	__snake_case    (       self :							List[str]     ,  UpperCamelCase :							str     ,  UpperCamelCase :							Any				)->Tuple:
      __SCREAMING_SNAKE_CASE       :    Optional[int]       =	[]
      for i, line in enumerate(UpperCamelCase				):
       if i == 0:
        continue
       __SCREAMING_SNAKE_CASE       :    str       =	"%s-%s" % (set_type, line[0])
       __SCREAMING_SNAKE_CASE       :    List[str]       =	line[5]
       __SCREAMING_SNAKE_CASE       :    List[str]       =	line[6]
       __SCREAMING_SNAKE_CASE       :    Optional[Any]       =	line[7][2:] if line[7].startswith("ex"				) else line[7]
       __SCREAMING_SNAKE_CASE       :    Optional[Any]       =	line[0]
       examples.append(InputExample(guid=UpperCamelCase     ,  text_a=UpperCamelCase     ,  text_b=UpperCamelCase     ,  label=UpperCamelCase     ,  pairID=UpperCamelCase				)				)
      return examples






def 	_lowerCAmelCase						(				__lowerCamelCase   :     List[InputExample]		,	__lowerCamelCase   :     List[str]		,	__lowerCamelCase   :     int		,	__lowerCamelCase   :     PreTrainedTokenizer		,	):



 """simple docstring"""





 __SCREAMING_SNAKE_CASE       :    str       =	{label: i for i, label in enumerate(__lowerCamelCase     )}

 __SCREAMING_SNAKE_CASE       :    Any       =	[]
 for ex_index, example in tqdm.tqdm(enumerate(__lowerCamelCase     )		,	desc="convert examples to features"     ):
  if ex_index % 10000 == 0:
   logger.info("Writing example %d" % (ex_index)     )

  __SCREAMING_SNAKE_CASE       :    Union[str, Any]       =	tokenizer(
      example.text_a		,	example.text_b		,	add_special_tokens=__lowerCamelCase		,	max_length=__lowerCamelCase		,	padding="max_length"		,	truncation=__lowerCamelCase		,	return_overflowing_tokens=__lowerCamelCase		,	)

  __SCREAMING_SNAKE_CASE       :    List[str]       =	label_map[example.label] if example.label in label_map else 0

  __SCREAMING_SNAKE_CASE       :    Union[str, Any]       =	int(example.pairID     )

  features.append(InputFeatures(**__lowerCamelCase		,	label=__lowerCamelCase		,	pairID=__lowerCamelCase     )     )

 for i, example in enumerate(examples[:5]     ):
  logger.info("*** Example ***"     )
  logger.info(F"""guid: {example}"""     )
  logger.info(F"""features: {features[i]}"""     )

 return features


_lowerCamelCase						    =  {
    """hans""": 3,
}






_lowerCamelCase						    =  {
    """hans""": HansProcessor,
}






