




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 ):






  def    __a     (    self						:      Optional[Any] )						->       Dict:






       """simple docstring"""







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






  @property
  def    __a     (    self						:      str )						->       str:






       """simple docstring"""







       lowercase   :  List[Any]        =					1
       lowercase   :  Optional[int]        =					3
       lowercase   :  List[Any]        =					(32, 32)

       lowercase   :  List[Any]        =					floats_tensor((batch_size, num_channels) + sizes ,		rng=random.Random(0 ) ).to(_UpperCAmelCase )
       return image






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






       """simple docstring"""







       torch.manual_seed(0 )
       lowercase   :  Tuple        =					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						:      Tuple )						->       Dict:






       """simple docstring"""







       torch.manual_seed(0 )
       lowercase   :  Optional[Any]        =					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						:      List[str] )						->       List[str]:






       """simple docstring"""







       torch.manual_seed(0 )
       lowercase   :  Tuple        =					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(_UpperCAmelCase )






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






       """simple docstring"""







       def extract(*_A						:      Dict ,		**_A						:      int ):
            class 		_A   :






              def __init__(    self						:      Any )						->       Dict:






                   """simple docstring"""







                   lowercase   :  Optional[int]        =					torch.ones([0] )






              def    __a     (    self						:      Union[str, Any] ,		_A						:      int )						->       int:






                   """simple docstring"""







                   self.pixel_values.to(_UpperCAmelCase )
                   return self


            return Out()

       return extract






  def    __a     (    self						:      Optional[Any] )						->       Optional[Any]:






       """simple docstring"""







       lowercase   :  Any        =					'cpu'  # ensure determinism for the device-dependent torch.Generator
       lowercase   :  int        =					self.dummy_cond_unet
       lowercase   :  List[str]        =					PNDMScheduler(skip_prk_steps=_UpperCAmelCase )
       lowercase   :  str        =					self.dummy_vae
       lowercase   :  str        =					self.dummy_text_encoder
       lowercase   :  int        =					XLMRobertaTokenizer.from_pretrained('''hf-internal-testing/tiny-xlm-roberta''' )
       lowercase   :  Union[str, Any]        =					77

       lowercase   :  Union[str, Any]        =					self.dummy_image.to(_UpperCAmelCase )
       lowercase   :  Tuple        =					init_image / 2 + 0.5

       # make sure here that pndm scheduler skips prk
       lowercase   :  Dict        =					AltDiffusionImgaImgPipeline(
           unet=_UpperCAmelCase ,		scheduler=_UpperCAmelCase ,		vae=_UpperCAmelCase ,		text_encoder=_UpperCAmelCase ,		tokenizer=_UpperCAmelCase ,		safety_checker=_UpperCAmelCase ,		feature_extractor=self.dummy_extractor ,		)
       lowercase   :  Union[str, Any]        =					VaeImageProcessor(vae_scale_factor=alt_pipe.vae_scale_factor ,		do_normalize=_UpperCAmelCase )
       lowercase   :  int        =					alt_pipe.to(_UpperCAmelCase )
       alt_pipe.set_progress_bar_config(disable=_UpperCAmelCase )

       lowercase   :  Dict        =					'A painting of a squirrel eating a burger'
       lowercase   :  Optional[Any]        =					torch.Generator(device=_UpperCAmelCase ).manual_seed(0 )
       lowercase   :  Tuple        =					alt_pipe(
           [prompt] ,		generator=_UpperCAmelCase ,		guidance_scale=6.0 ,		num_inference_steps=2 ,		output_type='''np''' ,		image=_UpperCAmelCase ,		)

       lowercase   :  Union[str, Any]        =					output.images

       lowercase   :  List[Any]        =					torch.Generator(device=_UpperCAmelCase ).manual_seed(0 )
       lowercase   :  Dict        =					alt_pipe(
           [prompt] ,		generator=_UpperCAmelCase ,		guidance_scale=6.0 ,		num_inference_steps=2 ,		output_type='''np''' ,		image=_UpperCAmelCase ,		return_dict=_UpperCAmelCase ,		)[0]

       lowercase   :  Optional[int]        =					image[0, -3:, -3:, -1]
       lowercase   :  Optional[Any]        =					image_from_tuple[0, -3:, -3:, -1]

       assert image.shape == (1, 32, 32, 3)
       lowercase   :  int        =					np.array([0.4_427, 0.3_731, 0.4_249, 0.4_941, 0.4_546, 0.4_148, 0.4_193, 0.4_666, 0.4_499] )

       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						:      str )						->       str:






       """simple docstring"""







       lowercase   :  Dict        =					self.dummy_cond_unet
       lowercase   :  Dict        =					PNDMScheduler(skip_prk_steps=_UpperCAmelCase )
       lowercase   :  Optional[int]        =					self.dummy_vae
       lowercase   :  Union[str, Any]        =					self.dummy_text_encoder
       lowercase   :  Tuple        =					XLMRobertaTokenizer.from_pretrained('''hf-internal-testing/tiny-xlm-roberta''' )
       lowercase   :  Optional[Any]        =					77

       lowercase   :  Any        =					self.dummy_image.to(_UpperCAmelCase )

       # put models in fp16
       lowercase   :  Optional[Any]        =					unet.half()
       lowercase   :  Optional[int]        =					vae.half()
       lowercase   :  str        =					bert.half()

       # make sure here that pndm scheduler skips prk
       lowercase   :  Any        =					AltDiffusionImgaImgPipeline(
           unet=_UpperCAmelCase ,		scheduler=_UpperCAmelCase ,		vae=_UpperCAmelCase ,		text_encoder=_UpperCAmelCase ,		tokenizer=_UpperCAmelCase ,		safety_checker=_UpperCAmelCase ,		feature_extractor=self.dummy_extractor ,		)
       lowercase   :  Dict        =					VaeImageProcessor(vae_scale_factor=alt_pipe.vae_scale_factor ,		do_normalize=_UpperCAmelCase )
       lowercase   :  List[Any]        =					alt_pipe.to(_UpperCAmelCase )
       alt_pipe.set_progress_bar_config(disable=_UpperCAmelCase )

       lowercase   :  Dict        =					'A painting of a squirrel eating a burger'
       lowercase   :  Tuple        =					torch.manual_seed(0 )
       lowercase   :  List[str]        =					alt_pipe(
           [prompt] ,		generator=_UpperCAmelCase ,		num_inference_steps=2 ,		output_type='''np''' ,		image=_UpperCAmelCase ,		).images

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






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






       """simple docstring"""







       lowercase   :  int        =					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
       lowercase   :  int        =					init_image.resize((760, 504) )

       lowercase   :  List[str]        =					'BAAI/AltDiffusion'
       lowercase   :  Dict        =					AltDiffusionImgaImgPipeline.from_pretrained(
           _UpperCAmelCase ,		safety_checker=_UpperCAmelCase ,		)
       pipe.to(_UpperCAmelCase )
       pipe.set_progress_bar_config(disable=_UpperCAmelCase )
       pipe.enable_attention_slicing()

       lowercase   :  Any        =					'A fantasy landscape, trending on artstation'

       lowercase   :  List[str]        =					torch.manual_seed(0 )
       lowercase   :  Any        =					pipe(
           prompt=_UpperCAmelCase ,		image=_UpperCAmelCase ,		strength=0.75 ,		guidance_scale=7.5 ,		generator=_UpperCAmelCase ,		output_type='''np''' ,		)
       lowercase   :  Tuple        =					output.images[0]

       lowercase   :  List[str]        =					image[255:258, 383:386, -1]

       assert image.shape == (504, 760, 3)
       lowercase   :  List[str]        =					np.array([0.9_358, 0.9_397, 0.9_599, 0.9_901, 1.0_000, 1.0_000, 0.9_882, 1.0_000, 1.0_000] )

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


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






  def    __a     (    self						:      List[str] )						->       Tuple:






       """simple docstring"""







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






  def    __a     (    self						:      Optional[Any] )						->       List[Any]:






       """simple docstring"""







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

       lowercase   :  Tuple        =					'BAAI/AltDiffusion'
       lowercase   :  Optional[Any]        =					AltDiffusionImgaImgPipeline.from_pretrained(
           _UpperCAmelCase ,		safety_checker=_UpperCAmelCase ,		)
       pipe.to(_UpperCAmelCase )
       pipe.set_progress_bar_config(disable=_UpperCAmelCase )
       pipe.enable_attention_slicing()

       lowercase   :  Optional[int]        =					'A fantasy landscape, trending on artstation'

       lowercase   :  List[str]        =					torch.manual_seed(0 )
       lowercase   :  Optional[Any]        =					pipe(
           prompt=_UpperCAmelCase ,		image=_UpperCAmelCase ,		strength=0.75 ,		guidance_scale=7.5 ,		generator=_UpperCAmelCase ,		output_type='''np''' ,		)
       lowercase   :  Any        =					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'''




import mpmath  # for roots of unity
import numpy as np


class 					SCREAMING_SNAKE_CASE       :


     def __init__(		self		,      _UpperCAmelCase=None		,      _UpperCAmelCase=None):

          '''simple docstring'''






          __A		:					List[Any]						=			list(poly_a or [0])[:]
          __A		:					Optional[int]						=			list(poly_b or [0])[:]

          # Remove leading zero coefficients
          while self.polyA[-1] == 0:
               self.polyA.pop()
          __A		:					Union[str, Any]						=			len(self.polyA)

          while self.polyB[-1] == 0:
               self.polyB.pop()
          __A		:					Optional[int]						=			len(self.polyB)

          # Add 0 to make lengths equal a power of 2
          __A		:					Optional[Any]						=			int(
              2 ** np.ceil(np.loga(len(self.polyA) + len(self.polyB) - 1)))

          while len(self.polyA) < self.c_max_length:
               self.polyA.append(0)
          while len(self.polyB) < self.c_max_length:
               self.polyB.append(0)
          # A complex root used for the fourier transform
          __A		:					str						=			complex(mpmath.root(x=1		,      n=self.c_max_length		,      k=1))

          # The product
          __A		:					Tuple						=			self.__multiply()


     def 				SCREAMING_SNAKE_CASE    (		self		,      _UpperCAmelCase):

          '''simple docstring'''






          __A		:					Optional[int]						=			[[x] for x in self.polyA] if which == 'A' else [[x] for x in self.polyB]
          # Corner case
          if len(_UpperCAmelCase) <= 1:
               return dft[0]
          #
          __A		:					Dict						=			self.c_max_length // 2
          while next_ncol > 0:
               __A		:					Optional[Any]						=			[[] for i in range(_UpperCAmelCase)]
               __A		:					Tuple						=			self.root**next_ncol

               # First half of next step
               __A		:					Optional[Any]						=			1
               for j in range(self.c_max_length // (next_ncol * 2)):
                    for i in range(_UpperCAmelCase):
                         new_dft[i].append(dft[i][j] + current_root * dft[i + next_ncol][j])
                    current_root *= root
               # Second half of next step
               __A		:					List[str]						=			1
               for j in range(self.c_max_length // (next_ncol * 2)):
                    for i in range(_UpperCAmelCase):
                         new_dft[i].append(dft[i][j] - current_root * dft[i + next_ncol][j])
                    current_root *= root
               # Update
               __A		:					Optional[int]						=			new_dft
               __A		:					Tuple						=			next_ncol // 2
          return dft[0]


     def 				SCREAMING_SNAKE_CASE    (		self):

          '''simple docstring'''






          __A		:					int						=			self.__dft('A')
          __A		:					Optional[Any]						=			self.__dft('B')
          __A		:					str						=			[[dft_a[i] * dft_b[i] for i in range(self.c_max_length)]]
          del dft_a
          del dft_b

          # Corner Case
          if len(inverce_c[0]) <= 1:
               return inverce_c[0]
          # Inverse DFT
          __A		:					Dict						=			2
          while next_ncol <= self.c_max_length:
               __A		:					Optional[int]						=			[[] for i in range(_UpperCAmelCase)]
               __A		:					Any						=			self.root ** (next_ncol // 2)
               __A		:					Tuple						=			1
               # First half of next step
               for j in range(self.c_max_length // next_ncol):
                    for i in range(next_ncol // 2):
                         # Even positions
                         new_inverse_c[i].append(
                             (
                                 inverce_c[i][j]
                                 + inverce_c[i][j + self.c_max_length // next_ncol]
                             )
                             / 2)
                         # Odd positions
                         new_inverse_c[i + next_ncol // 2].append(
                             (
                                 inverce_c[i][j]
                                 - inverce_c[i][j + self.c_max_length // next_ncol]
                             )
                             / (2 * current_root))
                    current_root *= root
               # Update
               __A		:					int						=			new_inverse_c
               next_ncol *= 2
          # Unpack
          __A		:					Optional[int]						=			[round(x[0].real		,      8) + round(x[0].imag		,      8) * 1j for x in inverce_c]

          # Remove leading 0's
          while inverce_c[-1] == 0:
               inverce_c.pop()
          return inverce_c


     def __str__(		self):

          '''simple docstring'''






          __A		:					int						=			'A = ' + ' + '.join(
              F'{coef}*x^{i}' for coef, i in enumerate(self.polyA[: self.len_A]))
          __A		:					Optional[Any]						=			'B = ' + ' + '.join(
              F'{coef}*x^{i}' for coef, i in enumerate(self.polyB[: self.len_B]))
          __A		:					str						=			'A*B = ' + ' + '.join(
              F'{coef}*x^{i}' for coef, i in enumerate(self.product))

          return F'{a}\n{b}\n{c}'


# Unit tests
if __name__ == "__main__":
  import doctest

  doctest.testmod()