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     a_			(						SCREAMING_SNAKE_CASE	,					SCREAMING_SNAKE_CASE	,					SCREAMING_SNAKE_CASE	,					unittest.TestCase						):






 '''simple docstring'''






 UpperCAmelCase_     			= AltDiffusionPipeline
 UpperCAmelCase_     			= TEXT_TO_IMAGE_PARAMS
 UpperCAmelCase_     			= TEXT_TO_IMAGE_BATCH_PARAMS
 UpperCAmelCase_     			= TEXT_TO_IMAGE_IMAGE_PARAMS
 UpperCAmelCase_     			= TEXT_TO_IMAGE_IMAGE_PARAMS




 def 		__snake_case						(      self :  Optional[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')      ,   cross_attention_dim=32      ,   )
        lowerCAmelCase__					     = DDIMScheduler(
            beta_start=0.00_085      ,   beta_end=0.012      ,   beta_schedule='scaled_linear'      ,   clip_sample=lowercase__      ,   set_alpha_to_one=lowercase__      ,   )
        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      ,   )

        # 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)
        lowerCAmelCase__					     = CLIPTextConfig(
            bos_token_id=0      ,   eos_token_id=2      ,   hidden_size=32      ,   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=5_002      ,   )
        lowerCAmelCase__					     = CLIPTextModel(lowercase__)

        lowerCAmelCase__					     = XLMRobertaTokenizer.from_pretrained('hf-internal-testing/tiny-xlm-roberta')
        lowerCAmelCase__					     = 77

        lowerCAmelCase__					     = {
            '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[str]      ,   lowercase__ :  Any      ,   lowercase__ :  Any=0):







        '''simple docstring'''



        if str(lowercase__).startswith('mps'):
               lowerCAmelCase__					     = torch.manual_seed(lowercase__)
        else:
               lowerCAmelCase__					     = torch.Generator(device=lowercase__).manual_seed(lowercase__)
        lowerCAmelCase__					     = {
            '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 :  Tuple):







        '''simple docstring'''



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




 def 		__snake_case						(      self :  Tuple):







        '''simple docstring'''



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




 def 		__snake_case						(      self :  Optional[int]):







        '''simple docstring'''



        lowerCAmelCase__					     = 'cpu'  # ensure determinism for the device-dependent torch.Generator

        lowerCAmelCase__					     = self.get_dummy_components()
        torch.manual_seed(0)
        lowerCAmelCase__					     = 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=5_002      ,   )
        # TODO: remove after fixing the non-deterministic text encoder
        lowerCAmelCase__					     = RobertaSeriesModelWithTransformation(lowercase__)
        lowerCAmelCase__					     = text_encoder

        lowerCAmelCase__					     = AltDiffusionPipeline(**lowercase__)
        lowerCAmelCase__					     = alt_pipe.to(lowercase__)
        alt_pipe.set_progress_bar_config(disable=lowercase__)

        lowerCAmelCase__					     = self.get_dummy_inputs(lowercase__)
        lowerCAmelCase__					     = 'A photo of an astronaut'
        lowerCAmelCase__					     = alt_pipe(**lowercase__)
        lowerCAmelCase__					     = output.images
        lowerCAmelCase__					     = image[0, -3:, -3:, -1]

        assert image.shape == (1, 64, 64, 3)
        lowerCAmelCase__					     = np.array(
            [0.5_748_162, 0.60_447_145, 0.48_821_217, 0.50_100_636, 0.5_431_185, 0.45_763_683, 0.49_657_696, 0.48_132_733, 0.47_573_093])

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




 def 		__snake_case						(      self :  int):







        '''simple docstring'''



        lowerCAmelCase__					     = 'cpu'  # ensure determinism for the device-dependent torch.Generator

        lowerCAmelCase__					     = self.get_dummy_components()
        lowerCAmelCase__					     = PNDMScheduler(skip_prk_steps=lowercase__)
        torch.manual_seed(0)
        lowerCAmelCase__					     = 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=5_002      ,   )
        # TODO: remove after fixing the non-deterministic text encoder
        lowerCAmelCase__					     = RobertaSeriesModelWithTransformation(lowercase__)
        lowerCAmelCase__					     = text_encoder
        lowerCAmelCase__					     = AltDiffusionPipeline(**lowercase__)
        lowerCAmelCase__					     = alt_pipe.to(lowercase__)
        alt_pipe.set_progress_bar_config(disable=lowercase__)

        lowerCAmelCase__					     = self.get_dummy_inputs(lowercase__)
        lowerCAmelCase__					     = alt_pipe(**lowercase__)
        lowerCAmelCase__					     = output.images
        lowerCAmelCase__					     = image[0, -3:, -3:, -1]

        assert image.shape == (1, 64, 64, 3)
        lowerCAmelCase__					     = np.array(
            [0.51_605_093, 0.5_707_241, 0.47_365_507, 0.50_578_886, 0.5_633_877, 0.4_642_503, 0.5_182_081, 0.48_763_484, 0.49_084_237])

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






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






 '''simple docstring'''






 def 		__snake_case						(      self :  Any):







        '''simple docstring'''



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




 def 		__snake_case						(      self :  Optional[Any]):







        '''simple docstring'''



        lowerCAmelCase__					     = AltDiffusionPipeline.from_pretrained('BAAI/AltDiffusion'      ,   safety_checker=lowercase__)
        lowerCAmelCase__					     = alt_pipe.to(lowercase__)
        alt_pipe.set_progress_bar_config(disable=lowercase__)

        lowerCAmelCase__					     = 'A painting of a squirrel eating a burger'
        lowerCAmelCase__					     = torch.manual_seed(0)
        lowerCAmelCase__					     = alt_pipe([prompt]      ,   generator=lowercase__      ,   guidance_scale=6.0      ,   num_inference_steps=20      ,   output_type='np')

        lowerCAmelCase__					     = output.images

        lowerCAmelCase__					     = image[0, -3:, -3:, -1]

        assert image.shape == (1, 512, 512, 3)
        lowerCAmelCase__					     = np.array([0.1_010, 0.0_800, 0.0_794, 0.0_885, 0.0_843, 0.0_762, 0.0_769, 0.0_729, 0.0_586])

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




 def 		__snake_case						(      self :  Tuple):







        '''simple docstring'''



        lowerCAmelCase__					     = DDIMScheduler.from_pretrained('BAAI/AltDiffusion'      ,   subfolder='scheduler')

        lowerCAmelCase__					     = AltDiffusionPipeline.from_pretrained('BAAI/AltDiffusion'      ,   scheduler=lowercase__      ,   safety_checker=lowercase__)
        lowerCAmelCase__					     = alt_pipe.to(lowercase__)
        alt_pipe.set_progress_bar_config(disable=lowercase__)

        lowerCAmelCase__					     = 'A painting of a squirrel eating a burger'
        lowerCAmelCase__					     = torch.manual_seed(0)

        lowerCAmelCase__					     = alt_pipe([prompt]      ,   generator=lowercase__      ,   num_inference_steps=2      ,   output_type='numpy')
        lowerCAmelCase__					     = output.images

        lowerCAmelCase__					     = image[0, -3:, -3:, -1]

        assert image.shape == (1, 512, 512, 3)
        lowerCAmelCase__					     = np.array([0.4_019, 0.4_052, 0.3_810, 0.4_119, 0.3_916, 0.3_982, 0.4_651, 0.4_195, 0.5_323])

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


import logging
import os

import quant_trainer
import torch
from torch.utils.data import DataLoader

from transformers import Trainer, is_torch_tpu_available
from transformers.trainer_utils import PredictionOutput


lowerCAmelCase__     						=			logging.getLogger(__name__)

if is_torch_tpu_available(check_device=False):
    import torch_xla.core.xla_model as xm
    import torch_xla.debug.metrics as met
class     a_			(						SCREAMING_SNAKE_CASE						):






 '''simple docstring'''






 def __init__(      self :  int      ,   *lowercase__ :  Dict      ,   lowercase__ :  List[Any]=None      ,   lowercase__ :  Optional[Any]=None      ,   lowercase__ :  Dict=None      ,   **lowercase__ :  Optional[int]):







        '''simple docstring'''



        super().__init__(*lowercase__      ,   **lowercase__)
        lowerCAmelCase__					     = eval_examples
        lowerCAmelCase__					     = post_process_function
        lowerCAmelCase__					     = quant_trainer_args
        lowerCAmelCase__					     = 128  # default number of calibration samples




 def 		__snake_case						(      self :  Tuple      ,   lowercase__ :  Any=None):







        '''simple docstring'''



        if calib_dataset is None and self.calib_dataset is None:
               raise ValueError('Trainer: calibration requires an calib_dataset.')
        lowerCAmelCase__					     = calib_dataset if calib_dataset is not None else self.calib_dataset

        lowerCAmelCase__					     = self._remove_unused_columns(lowercase__      ,   description='Calibration')

        return DataLoader(
            lowercase__      ,   batch_size=self.args.eval_batch_size      ,   collate_fn=self.data_collator      ,   drop_last=self.args.dataloader_drop_last      ,   num_workers=self.args.dataloader_num_workers      ,   pin_memory=self.args.dataloader_pin_memory      ,   shuffle=lowercase__      ,   )




 def 		__snake_case						(      self :  List[Any]      ,   lowercase__ :  Union[str, Any]=None):







        '''simple docstring'''



        lowerCAmelCase__					     = self.train_dataset if calib_dataset is None else calib_dataset
        lowerCAmelCase__					     = self.get_calib_dataloader(lowercase__)

        lowerCAmelCase__					     = self.model
        quant_trainer.configure_model(lowercase__      ,   self.quant_trainer_args      ,   calib=lowercase__)
        model.eval()
        quant_trainer.enable_calibration(lowercase__)

        logger.info('***** Running calibration *****')
        logger.info(F"""  Num examples = {self.calib_num}""")
        logger.info(F"""  Batch size = {calib_dataloader.batch_size}""")

        for step, inputs in enumerate(lowercase__):
               # Prediction step
               lowerCAmelCase__     , lowerCAmelCase__     , lowerCAmelCase__					     = self.prediction_step(lowercase__      ,   lowercase__      ,   prediction_loss_only=lowercase__)
               if (step + 1) * calib_dataloader.batch_size >= self.calib_num:
                      break

        quant_trainer.finish_calibration(lowercase__      ,   self.quant_trainer_args)
        lowerCAmelCase__					     = model




 def 		__snake_case						(      self :  Optional[Any]      ,   lowercase__ :  List[Any]=None      ,   lowercase__ :  Optional[Any]=None      ,   lowercase__ :  List[Any]=None      ,   lowercase__ :  str = "eval"):







        '''simple docstring'''



        lowerCAmelCase__					     = self.eval_dataset if eval_dataset is None else eval_dataset
        lowerCAmelCase__					     = self.get_eval_dataloader(lowercase__)
        lowerCAmelCase__					     = self.eval_examples if eval_examples is None else eval_examples

        # Temporarily disable metric computation, we will do it in the loop here.
        lowerCAmelCase__					     = self.compute_metrics
        lowerCAmelCase__					     = None
        lowerCAmelCase__					     = self.prediction_loop if self.args.use_legacy_prediction_loop else self.evaluation_loop
        try:
               lowerCAmelCase__					     = eval_loop(
                   lowercase__      ,   description='Evaluation'      ,   prediction_loss_only=True if compute_metrics is None else None      ,   ignore_keys=lowercase__      ,   )
        finally:
               lowerCAmelCase__					     = compute_metrics

        if self.post_process_function is not None and self.compute_metrics is not None:
               lowerCAmelCase__					     = self.post_process_function(lowercase__      ,   lowercase__      ,   output.predictions)
               lowerCAmelCase__					     = self.compute_metrics(lowercase__)

               # Prefix all keys with metric_key_prefix + '_'
               for key in list(metrics.keys()):
                      if not key.startswith(F"""{metric_key_prefix}_"""):
                             lowerCAmelCase__					     = metrics.pop(lowercase__)

               self.log(lowercase__)
        else:
               lowerCAmelCase__					     = {}

        if self.args.tpu_metrics_debug or self.args.debug:
               # tpu-comment: Logging debug metrics for PyTorch/XLA (compile, execute times, ops, etc.)
               xm.master_print(met.metrics_report())

        lowerCAmelCase__					     = self.callback_handler.on_evaluate(self.args      ,   self.state      ,   self.control      ,   lowercase__)
        return metrics




 def 		__snake_case						(      self :  Optional[int]      ,   lowercase__ :  str      ,   lowercase__ :  Any      ,   lowercase__ :  List[str]=None      ,   lowercase__ :  str = "test"):







        '''simple docstring'''



        lowerCAmelCase__					     = self.get_test_dataloader(lowercase__)

        # Temporarily disable metric computation, we will do it in the loop here.
        lowerCAmelCase__					     = self.compute_metrics
        lowerCAmelCase__					     = None
        lowerCAmelCase__					     = self.prediction_loop if self.args.use_legacy_prediction_loop else self.evaluation_loop
        try:
               lowerCAmelCase__					     = eval_loop(
                   lowercase__      ,   description='Prediction'      ,   prediction_loss_only=True if compute_metrics is None else None      ,   ignore_keys=lowercase__      ,   )
        finally:
               lowerCAmelCase__					     = compute_metrics

        if self.post_process_function is None or self.compute_metrics is None:
               return output

        lowerCAmelCase__					     = self.post_process_function(lowercase__      ,   lowercase__      ,   output.predictions      ,   'predict')
        lowerCAmelCase__					     = self.compute_metrics(lowercase__)

        # Prefix all keys with metric_key_prefix + '_'
        for key in list(metrics.keys()):
               if not key.startswith(F"""{metric_key_prefix}_"""):
                      lowerCAmelCase__					     = metrics.pop(lowercase__)

        return PredictionOutput(predictions=predictions.predictions      ,   label_ids=predictions.label_ids      ,   metrics=lowercase__)




 def 		__snake_case						(      self :  List[str]      ,   lowercase__ :  List[str]="./"):







        '''simple docstring'''



        lowerCAmelCase__					     = self.eval_dataset
        lowerCAmelCase__					     = self.get_eval_dataloader(lowercase__)

        lowerCAmelCase__					     = next(iter(lowercase__))

        # saving device - to make it consistent
        lowerCAmelCase__					     = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

        # convert to tuple
        lowerCAmelCase__					     = tuple(v.to(lowercase__) for k, v in batch.items())

        logger.info('Converting model to be onnx compatible')
        from pytorch_quantization.nn import TensorQuantizer

        lowerCAmelCase__					     = True

        lowerCAmelCase__					     = self.model.to(lowercase__)

        model.eval()
        model.float()

        lowerCAmelCase__					     = model.module if hasattr(lowercase__      ,   'module') else model
        quant_trainer.configure_model(lowercase__      ,   self.quant_trainer_args)

        lowerCAmelCase__					     = os.path.join(lowercase__      ,   'model.onnx')
        logger.info(F"""exporting model to {output_model_file}""")

        lowerCAmelCase__					     = {0: 'batch_size', 1: 'seq_len'}

        torch.onnx.export(
            lowercase__      ,   lowercase__      ,   lowercase__      ,   export_params=lowercase__      ,   opset_version=13      ,   do_constant_folding=lowercase__      ,   input_names=['input_ids', 'attention_mask', 'token_type_ids']      ,   output_names=['output_start_logits', 'output_end_logits']      ,   dynamic_axes={
                'input_ids': axes,
                'attention_mask': axes,
                'token_type_ids': axes,
                'output_start_logits': axes,
                'output_end_logits': axes,
            }      ,   verbose=lowercase__      ,   )
        logger.info('onnx export finished')
