



'''simple docstring'''





import gc
import random
import unittest

import numpy as np
import torch
from transformers import XLMRobertaTokenizer

from diffusers import (
    AltDiffusionImgaImgPipeline,
    AutoencoderKL,
    PNDMScheduler,
    UNetaDConditionModel,
)
from diffusers.image_processor import VaeImageProcessor
from diffusers.pipelines.alt_diffusion.modeling_roberta_series import (
    RobertaSeriesConfig,
    RobertaSeriesModelWithTransformation,
)
from diffusers.utils import floats_tensor, load_image, load_numpy, slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu


enable_full_determinism()






class __A			(  unittest.TestCase						):



 '''simple docstring'''







 def 					a__  (self       )    ->     str:


      """simple docstring"""



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



 @property
 def 					a__  (self       )    ->     Union[str, Any]:


      """simple docstring"""



      _a      =  1
      _a      =  3
      _a      =  (32, 32)

      _a      =  floats_tensor((batch_size, num_channels) + sizes ,   rng=random.Random(0       )       ).to(A       )
      return image



 @property
 def 					a__  (self       )    ->     Dict:


      """simple docstring"""



      torch.manual_seed(0       )
      _a      =  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 ,   )
      return model



 @property
 def 					a__  (self       )    ->     Optional[int]:


      """simple docstring"""



      torch.manual_seed(0       )
      _a      =  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 ,   )
      return model



 @property
 def 					a__  (self       )    ->     Optional[Any]:


      """simple docstring"""



      torch.manual_seed(0       )
      _a      =  RobertaSeriesConfig(
          hidden_size=32 ,   project_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_006 ,   )
      return RobertaSeriesModelWithTransformation(A       )



 @property
 def 					a__  (self       )    ->     List[str]:


      """simple docstring"""



      def extract(*A ,   **A       ):






           class __A			:



            '''simple docstring'''







            def __init__(self       )    ->     Any:


                 """simple docstring"""



                 _a      =  torch.ones([0]       )



            def 					a__  (self ,   A       )    ->     Tuple:


                 """simple docstring"""



                 self.pixel_values.to(A       )
                 return self





           return Out()

      return extract



 def 					a__  (self       )    ->     str:


      """simple docstring"""



      _a      =  '''cpu'''  # ensure determinism for the device-dependent torch.Generator
      _a      =  self.dummy_cond_unet
      _a      =  PNDMScheduler(skip_prk_steps=A       )
      _a      =  self.dummy_vae
      _a      =  self.dummy_text_encoder
      _a      =  XLMRobertaTokenizer.from_pretrained('''hf-internal-testing/tiny-xlm-roberta'''       )
      _a      =  77

      _a      =  self.dummy_image.to(A       )
      _a      =  init_image / 2 + 0.5

      # make sure here that pndm scheduler skips prk
      _a      =  AltDiffusionImgaImgPipeline(
          unet=A ,   scheduler=A ,   vae=A ,   text_encoder=A ,   tokenizer=A ,   safety_checker=A ,   feature_extractor=self.dummy_extractor ,   )
      _a      =  VaeImageProcessor(vae_scale_factor=alt_pipe.vae_scale_factor ,   do_normalize=A       )
      _a      =  alt_pipe.to(A       )
      alt_pipe.set_progress_bar_config(disable=A       )

      _a      =  '''A painting of a squirrel eating a burger'''
      _a      =  torch.Generator(device=A       ).manual_seed(0       )
      _a      =  alt_pipe(
          [prompt] ,   generator=A ,   guidance_scale=6.0 ,   num_inference_steps=2 ,   output_type='''np''' ,   image=A ,   )

      _a      =  output.images

      _a      =  torch.Generator(device=A       ).manual_seed(0       )
      _a      =  alt_pipe(
          [prompt] ,   generator=A ,   guidance_scale=6.0 ,   num_inference_steps=2 ,   output_type='''np''' ,   image=A ,   return_dict=A ,   )[0]

      _a      =  image[0, -3:, -3:, -1]
      _a      =  image_from_tuple[0, -3:, -3:, -1]

      assert image.shape == (1, 32, 32, 3)
      _a      =  np.array([0.4427, 0.3731, 0.4249, 0.4941, 0.4546, 0.4148, 0.4193, 0.4666, 0.4499]       )

      assert np.abs(image_slice.flatten() - expected_slice       ).max() < 5E-3
      assert np.abs(image_from_tuple_slice.flatten() - expected_slice       ).max() < 5E-3



 @unittest.skipIf(torch_device != '''cuda''' ,   '''This test requires a GPU'''       )
 def 					a__  (self       )    ->     List[Any]:


      """simple docstring"""



      _a      =  self.dummy_cond_unet
      _a      =  PNDMScheduler(skip_prk_steps=A       )
      _a      =  self.dummy_vae
      _a      =  self.dummy_text_encoder
      _a      =  XLMRobertaTokenizer.from_pretrained('''hf-internal-testing/tiny-xlm-roberta'''       )
      _a      =  77

      _a      =  self.dummy_image.to(A       )

      # put models in fp16
      _a      =  unet.half()
      _a      =  vae.half()
      _a      =  bert.half()

      # make sure here that pndm scheduler skips prk
      _a      =  AltDiffusionImgaImgPipeline(
          unet=A ,   scheduler=A ,   vae=A ,   text_encoder=A ,   tokenizer=A ,   safety_checker=A ,   feature_extractor=self.dummy_extractor ,   )
      _a      =  VaeImageProcessor(vae_scale_factor=alt_pipe.vae_scale_factor ,   do_normalize=A       )
      _a      =  alt_pipe.to(A       )
      alt_pipe.set_progress_bar_config(disable=A       )

      _a      =  '''A painting of a squirrel eating a burger'''
      _a      =  torch.manual_seed(0       )
      _a      =  alt_pipe(
          [prompt] ,   generator=A ,   num_inference_steps=2 ,   output_type='''np''' ,   image=A ,   ).images

      assert image.shape == (1, 32, 32, 3)



 @unittest.skipIf(torch_device != '''cuda''' ,   '''This test requires a GPU'''       )
 def 					a__  (self       )    ->     int:


      """simple docstring"""



      _a      =  load_image(
          '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'''
          '''/img2img/sketch-mountains-input.jpg'''       )
      # resize to resolution that is divisible by 8 but not 16 or 32
      _a      =  init_image.resize((760, 504)       )

      _a      =  '''BAAI/AltDiffusion'''
      _a      =  AltDiffusionImgaImgPipeline.from_pretrained(
          A ,   safety_checker=A ,   )
      pipe.to(A       )
      pipe.set_progress_bar_config(disable=A       )
      pipe.enable_attention_slicing()

      _a      =  '''A fantasy landscape, trending on artstation'''

      _a      =  torch.manual_seed(0       )
      _a      =  pipe(
          prompt=A ,   image=A ,   strength=0.75 ,   guidance_scale=7.5 ,   generator=A ,   output_type='''np''' ,   )
      _a      =  output.images[0]

      _a      =  image[255:258, 383:386, -1]

      assert image.shape == (504, 760, 3)
      _a      =  np.array([0.9358, 0.9397, 0.9599, 0.9901, 1.0000, 1.0000, 0.9882, 1.0000, 1.0000]       )

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






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



 '''simple docstring'''







 def 					a__  (self       )    ->     str:


      """simple docstring"""



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



 def 					a__  (self       )    ->     Dict:


      """simple docstring"""



      _a      =  load_image(
          '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'''
          '''/img2img/sketch-mountains-input.jpg'''       )
      _a      =  init_image.resize((768, 512)       )
      _a      =  load_numpy(
          '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/img2img/fantasy_landscape_alt.npy'''       )

      _a      =  '''BAAI/AltDiffusion'''
      _a      =  AltDiffusionImgaImgPipeline.from_pretrained(
          A ,   safety_checker=A ,   )
      pipe.to(A       )
      pipe.set_progress_bar_config(disable=A       )
      pipe.enable_attention_slicing()

      _a      =  '''A fantasy landscape, trending on artstation'''

      _a      =  torch.manual_seed(0       )
      _a      =  pipe(
          prompt=A ,   image=A ,   strength=0.75 ,   guidance_scale=7.5 ,   generator=A ,   output_type='''np''' ,   )
      _a      =  output.images[0]

      assert image.shape == (512, 768, 3)
      # img2img is flaky across GPUs even in fp32, so using MAE here
      assert np.abs(expected_image - image       ).max() < 1E-2








'''simple docstring'''





from collections import OrderedDict
from typing import Any, List, Mapping, Optional

from ... import PreTrainedTokenizer, TensorType, is_torch_available
from ...configuration_utils import PretrainedConfig
from ...onnx import OnnxConfigWithPast, PatchingSpec
from ...utils import logging


lowercase_		       =  logging.get_logger(__name__)

lowercase_		       =  {
    "EleutherAI/gpt-j-6B": "https://huggingface.co/EleutherAI/gpt-j-6B/resolve/main/config.json",
    # See all GPT-J models at https://huggingface.co/models?filter=gpt_j
}






class __A			(  A						):



 '''simple docstring'''







 __lowerCamelCase    :       Any          =       'gptj'
 __lowerCamelCase    :       List[str]          =       {
     'max_position_embeddings': 'n_positions',
     'hidden_size': 'n_embd',
     'num_attention_heads': 'n_head',
     'num_hidden_layers': 'n_layer',
 }



 def __init__(self ,   A=50_400 ,   A=2_048 ,   A=4_096 ,   A=28 ,   A=16 ,   A=64 ,   A=None ,   A="gelu_new" ,   A=0.0 ,   A=0.0 ,   A=0.0 ,   A=1E-5 ,   A=0.02 ,   A=True ,   A=50_256 ,   A=50_256 ,   A=False ,   **A ,   )    ->     Tuple:


      """simple docstring"""



      _a      =  vocab_size
      _a      =  n_positions
      _a      =  n_embd
      _a      =  n_layer
      _a      =  n_head
      _a      =  n_inner
      _a      =  rotary_dim
      _a      =  activation_function
      _a      =  resid_pdrop
      _a      =  embd_pdrop
      _a      =  attn_pdrop
      _a      =  layer_norm_epsilon
      _a      =  initializer_range
      _a      =  use_cache

      _a      =  bos_token_id
      _a      =  eos_token_id

      super().__init__(
          bos_token_id=A ,   eos_token_id=A ,   tie_word_embeddings=A ,   **A       )






class __A			(  A						):



 '''simple docstring'''







 def __init__(self ,   A ,   A = "default" ,   A = None ,   A = False ,   )    ->     List[str]:


      """simple docstring"""



      super().__init__(A ,   task=A ,   patching_specs=A ,   use_past=A       )
      if not getattr(self._config ,   '''pad_token_id''' ,   A       ):
           # TODO: how to do that better?
           _a      =  0



 @property
 def 					a__  (self       )    ->     Mapping[str, Mapping[int, str]]:


      """simple docstring"""



      _a      =  OrderedDict({'''input_ids''': {0: '''batch''', 1: '''sequence'''}}       )
      if self.use_past:
           self.fill_with_past_key_values_(A ,   direction='''inputs'''       )
           _a      =  {0: '''batch''', 1: '''past_sequence + sequence'''}
      else:
           _a      =  {0: '''batch''', 1: '''sequence'''}

      return common_inputs



 @property
 def 					a__  (self       )    ->     int:


      """simple docstring"""



      return self._config.n_layer



 @property
 def 					a__  (self       )    ->     int:


      """simple docstring"""



      return self._config.n_head



 def 					a__  (self ,   A ,   A = -1 ,   A = -1 ,   A = False ,   A = None ,   )    ->     Mapping[str, Any]:


      """simple docstring"""



      _a      =  super(A ,   self       ).generate_dummy_inputs(
          A ,   batch_size=A ,   seq_length=A ,   is_pair=A ,   framework=A       )

      # We need to order the input in the way they appears in the forward()
      _a      =  OrderedDict({'''input_ids''': common_inputs['''input_ids''']}       )

      # Need to add the past_keys
      if self.use_past:
           if not is_torch_available():
                raise ValueError('''Cannot generate dummy past_keys inputs without PyTorch installed.'''       )
           else:
                import torch

                _a  ,    _a      =  common_inputs['''input_ids'''].shape
                # Not using the same length for past_key_values
                _a      =  seqlen + 2
                _a      =  (
                    batch,
                    self.num_attention_heads,
                    past_key_values_length,
                    self._config.hidden_size // self.num_attention_heads,
                )
                _a      =  [
                    (torch.zeros(A       ), torch.zeros(A       )) for _ in range(self.num_layers       )
                ]

      _a      =  common_inputs['''attention_mask''']
      if self.use_past:
           _a      =  ordered_inputs['''attention_mask'''].dtype
           _a      =  torch.cat(
               [ordered_inputs['''attention_mask'''], torch.ones(A ,   A ,   dtype=A       )] ,   dim=1       )

      return ordered_inputs



 @property
 def 					a__  (self       )    ->     int:


      """simple docstring"""



      return 13



