

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  __lowerCAmelCase     (					SCREAMING_SNAKE_CASE			,			SCREAMING_SNAKE_CASE			,			SCREAMING_SNAKE_CASE			,			unittest.TestCase							):
  _a									=  AltDiffusionPipeline
  _a									=  TEXT_TO_IMAGE_PARAMS
  _a									=  TEXT_TO_IMAGE_BATCH_PARAMS
  _a									=  TEXT_TO_IMAGE_IMAGE_PARAMS
  _a									=  TEXT_TO_IMAGE_IMAGE_PARAMS




  def   A__     (				self					)  -> Union[str, Any]:
        '''simple docstring'''






        torch.manual_seed(0					)
        _lowercase	       =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							,   )
        _lowercase	       =DDIMScheduler(
            beta_start=0.00085							,   beta_end=0.012							,   beta_schedule='scaled_linear'							,   clip_sample=lowerCAmelCase							,   set_alpha_to_one=lowerCAmelCase							,   )
        torch.manual_seed(0					)
        _lowercase	       =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					)
        _lowercase	       =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							,   )
        _lowercase	       =CLIPTextModel(lowerCAmelCase					)

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

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




  def   A__     (				self							,   lowerCAmelCase							,   lowerCAmelCase=0					)  -> Optional[int]:
        '''simple docstring'''






        if str(lowerCAmelCase					).startswith('mps'					):
              _lowercase	       =torch.manual_seed(lowerCAmelCase					)
        else:
              _lowercase	       =torch.Generator(device=lowerCAmelCase					).manual_seed(lowerCAmelCase					)
        _lowercase	       ={
            '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   A__     (				self					)  -> Tuple:
        '''simple docstring'''






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




  def   A__     (				self					)  -> List[str]:
        '''simple docstring'''






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




  def   A__     (				self					)  -> Union[str, Any]:
        '''simple docstring'''






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

        _lowercase	       =self.get_dummy_components()
        torch.manual_seed(0					)
        _lowercase	       =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
        _lowercase	       =RobertaSeriesModelWithTransformation(lowerCAmelCase					)
        _lowercase	       =text_encoder

        _lowercase	       =AltDiffusionPipeline(**lowerCAmelCase					)
        _lowercase	       =alt_pipe.to(lowerCAmelCase					)
        alt_pipe.set_progress_bar_config(disable=lowerCAmelCase					)

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

        assert image.shape == (1, 64, 64, 3)
        _lowercase	       =np.array(
            [0.5748162, 0.60447145, 0.48821217, 0.50100636, 0.5431185, 0.45763683, 0.49657696, 0.48132733, 0.47573093]					)

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




  def   A__     (				self					)  -> List[str]:
        '''simple docstring'''






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

        _lowercase	       =self.get_dummy_components()
        _lowercase	       =PNDMScheduler(skip_prk_steps=lowerCAmelCase					)
        torch.manual_seed(0					)
        _lowercase	       =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
        _lowercase	       =RobertaSeriesModelWithTransformation(lowerCAmelCase					)
        _lowercase	       =text_encoder
        _lowercase	       =AltDiffusionPipeline(**lowerCAmelCase					)
        _lowercase	       =alt_pipe.to(lowerCAmelCase					)
        alt_pipe.set_progress_bar_config(disable=lowerCAmelCase					)

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

        assert image.shape == (1, 64, 64, 3)
        _lowercase	       =np.array(
            [0.51605093, 0.5707241, 0.47365507, 0.50578886, 0.5633877, 0.4642503, 0.5182081, 0.48763484, 0.49084237]					)

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





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




  def   A__     (				self					)  -> Optional[int]:
        '''simple docstring'''






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




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






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

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

        _lowercase	       =output.images

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

        assert image.shape == (1, 512, 512, 3)
        _lowercase	       =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   A__     (				self					)  -> Optional[int]:
        '''simple docstring'''






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

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

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

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

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

        assert image.shape == (1, 512, 512, 3)
        _lowercase	       =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





from __future__ import annotations
def      a      (			A__	:   list[int]     )    ->      int:





      """simple docstring"""



      if not nums:
            return 0
      _lowercase	       =nums[0]
      _lowercase	       =0
      for num in nums[1:]:
            _lowercase       ,       _lowercase	       =(
                max_excluding + num,
                max(A__      ,   A__     ),
            )
      return max(A__      ,   A__     )


if __name__ == "__main__":
  import doctest

  doctest.testmod()


