

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   lowercase_   (  unittest.TestCase					):



      def 						UpperCamelCase_		(      self				:    Any		)   ->					Dict:
        # clean up the VRAM after each test
        super().tearDown()
        gc.collect()
        torch.cuda.empty_cache()



      @property
      def 						UpperCamelCase_		(      self				:    Tuple		)   ->					Optional[int]:
        _snake_case						=							1
        _snake_case						=							3
        _snake_case						=							(32, 32)

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



      @property
      def 						UpperCamelCase_		(      self				:    Any		)   ->					List[str]:
        torch.manual_seed(0		)
        _snake_case						=							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 						UpperCamelCase_		(      self				:    Optional[Any]		)   ->					Dict:
        torch.manual_seed(0		)
        _snake_case						=							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 						UpperCamelCase_		(      self				:    Optional[int]		)   ->					Optional[int]:
        torch.manual_seed(0		)
        _snake_case						=							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=5006      , )
        return RobertaSeriesModelWithTransformation(A__		)



      @property
      def 						UpperCamelCase_		(      self				:    Tuple		)   ->					Tuple:
        def extract(*A__				:    int      , **A__				:    List[Any]		):



          class   lowercase_   :



                def __init__(      self				:    List[str]		)   ->					Union[str, Any]:
                  _snake_case						=							torch.ones([0]		)







                def 						UpperCamelCase_		(      self				:    List[Any]      , A__				:    Optional[int]		)   ->					int:
                  self.pixel_values.to(A__		)
                  return self

          return Out()

        return extract



      def 						UpperCamelCase_		(      self				:    Union[str, Any]		)   ->					Tuple:
        _snake_case						=							'''cpu'''  # ensure determinism for the device-dependent torch.Generator
        _snake_case						=							self.dummy_cond_unet
        _snake_case						=							PNDMScheduler(skip_prk_steps=A__		)
        _snake_case						=							self.dummy_vae
        _snake_case						=							self.dummy_text_encoder
        _snake_case						=							XLMRobertaTokenizer.from_pretrained('''hf-internal-testing/tiny-xlm-roberta'''		)
        _snake_case						=							77

        _snake_case						=							self.dummy_image.to(A__		)
        _snake_case						=							init_image / 2 + 0.5

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

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

        _snake_case						=							output.images

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

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

        assert image.shape == (1, 32, 32, 3)
        _snake_case						=							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 						UpperCamelCase_		(      self				:    Optional[Any]		)   ->					List[str]:
        _snake_case						=							self.dummy_cond_unet
        _snake_case						=							PNDMScheduler(skip_prk_steps=A__		)
        _snake_case						=							self.dummy_vae
        _snake_case						=							self.dummy_text_encoder
        _snake_case						=							XLMRobertaTokenizer.from_pretrained('''hf-internal-testing/tiny-xlm-roberta'''		)
        _snake_case						=							77

        _snake_case						=							self.dummy_image.to(A__		)

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

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

        _snake_case						=							'''A painting of a squirrel eating a burger'''
        _snake_case						=							torch.manual_seed(0		)
        _snake_case						=							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 						UpperCamelCase_		(      self				:    str		)   ->					Dict:
        _snake_case						=							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
        _snake_case						=							init_image.resize((760, 504)		)

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

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

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

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

        assert image.shape == (504, 760, 3)
        _snake_case						=							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   lowercase_   (  unittest.TestCase					):



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







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

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

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

        _snake_case						=							torch.manual_seed(0		)
        _snake_case						=							pipe(
            prompt=A__      , image=A__      , strength=0.75      , guidance_scale=7.5      , generator=A__      , output_type='''np'''      , )
        _snake_case						=							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




__A      		=     '''ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/'''



def 				snake_case_(_UpperCamelCase					)							->     bytes:



  """simple docstring"""






  if not isinstance(_UpperCamelCase			,	_UpperCamelCase					):
    _snake_case						=							F"""a bytes-like object is required, not '{data.__class__.__name__}'"""
    raise TypeError(_UpperCamelCase					)

  _snake_case						=							''''''.join(bin(_UpperCamelCase					)[2:].zfill(8					) for byte in data					)

  _snake_case						=							len(_UpperCamelCase					) % 6 != 0

  if padding_needed:
    # The padding that will be added later
    _snake_case						=							b'''=''' * ((6 - len(_UpperCamelCase					) % 6) // 2)

    # Append binary_stream with arbitrary binary digits (0's by default) to make its
    # length a multiple of 6.
    binary_stream += "0" * (6 - len(_UpperCamelCase					) % 6)
  else:
    _snake_case						=							b''''''

  # Encode every 6 binary digits to their corresponding Base64 character
  return (
      "".join(
          B64_CHARSET[int(binary_stream[index : index + 6]			,	2					)]
          for index in range(0			,	len(_UpperCamelCase					)			,	6					)					).encode()
      + padding
  )



def 				snake_case_(_UpperCamelCase					)							->     bytes:



  """simple docstring"""






  if not isinstance(_UpperCamelCase			,	_UpperCamelCase					) and not isinstance(_UpperCamelCase			,	_UpperCamelCase					):
    _snake_case						=							(
        '''argument should be a bytes-like object or ASCII string, '''
        F"""not '{encoded_data.__class__.__name__}'"""
    )
    raise TypeError(_UpperCamelCase					)

  # In case encoded_data is a bytes-like object, make sure it contains only
  # ASCII characters so we convert it to a string object
  if isinstance(_UpperCamelCase			,	_UpperCamelCase					):
    try:
      _snake_case						=							encoded_data.decode('''utf-8'''					)
    except UnicodeDecodeError:
      raise ValueError('''base64 encoded data should only contain ASCII characters'''					)

  _snake_case						=							encoded_data.count('''='''					)

  # Check if the encoded string contains non base64 characters
  if padding:
    assert all(
        char in B64_CHARSET for char in encoded_data[:-padding]					), "Invalid base64 character(s) found."
  else:
    assert all(
        char in B64_CHARSET for char in encoded_data					), "Invalid base64 character(s) found."

  # Check the padding
  assert len(_UpperCamelCase					) % 4 == 0 and padding < 3, "Incorrect padding"

  if padding:
    # Remove padding if there is one
    _snake_case						=							encoded_data[:-padding]

    _snake_case						=							''''''.join(
        bin(B64_CHARSET.index(_UpperCamelCase					)					)[2:].zfill(6					) for char in encoded_data					)[: -padding * 2]
  else:
    _snake_case						=							''''''.join(
        bin(B64_CHARSET.index(_UpperCamelCase					)					)[2:].zfill(6					) for char in encoded_data					)

  _snake_case						=							[
      int(binary_stream[index : index + 8]			,	2					)
      for index in range(0			,	len(_UpperCamelCase					)			,	8					)
  ]

  return bytes(_UpperCamelCase					)


if __name__ == "__main__":
 import doctest

 doctest.testmod()

