





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

  """simple docstring"""


  def     lowercase		(    self			:		int					)							->						Optional[Any]:
     # clean up the VRAM after each test
     super().tearDown()
     gc.collect()
     torch.cuda.empty_cache()






  @property
  def     lowercase		(    self			:		int					)							->						List[Any]:
     __lowerCAmelCase							     =     1
     __lowerCAmelCase							     =     3
     __lowerCAmelCase							     =     (3_2, 3_2)

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






  @property
  def     lowercase		(    self			:		List[str]					)							->						Dict:
     torch.manual_seed(0					)
     __lowerCAmelCase							     =     UNetaDConditionModel(
         block_out_channels=(3_2, 6_4) ,  layers_per_block=2 ,  sample_size=3_2 ,  in_channels=4 ,  out_channels=4 ,  down_block_types=('DownBlock2D', 'CrossAttnDownBlock2D') ,  up_block_types=('CrossAttnUpBlock2D', 'UpBlock2D') ,  cross_attention_dim=3_2 ,  )
     return model






  @property
  def     lowercase		(    self			:		Union[str, Any]					)							->						int:
     torch.manual_seed(0					)
     __lowerCAmelCase							     =     AutoencoderKL(
         block_out_channels=[3_2, 6_4] ,  in_channels=3 ,  out_channels=3 ,  down_block_types=['DownEncoderBlock2D', 'DownEncoderBlock2D'] ,  up_block_types=['UpDecoderBlock2D', 'UpDecoderBlock2D'] ,  latent_channels=4 ,  )
     return model






  @property
  def     lowercase		(    self			:		str					)							->						Dict:
     torch.manual_seed(0					)
     __lowerCAmelCase							     =     RobertaSeriesConfig(
         hidden_size=3_2 ,  project_dim=3_2 ,  intermediate_size=3_7 ,  layer_norm_eps=1e-05 ,  num_attention_heads=4 ,  num_hidden_layers=5 ,  pad_token_id=1 ,  vocab_size=5_0_0_6 ,  )
     return RobertaSeriesModelWithTransformation(lowerCAmelCase_					)






  @property
  def     lowercase		(    self			:		List[str]					)							->						Optional[int]:
     def extract(*lowerCAmelCase_			:		Dict ,  **lowerCAmelCase_			:		Dict					):



        class     _UpperCAmelCase   :

          """simple docstring"""


          def __init__(    self			:		int					)							->						Dict:
             __lowerCAmelCase							     =     torch.ones([0]					)






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





        return Out()

     return extract






  def     lowercase		(    self			:		Optional[Any]					)							->						Optional[Any]:
     __lowerCAmelCase							     =     'cpu'  # ensure determinism for the device-dependent torch.Generator
     __lowerCAmelCase							     =     self.dummy_cond_unet
     __lowerCAmelCase							     =     PNDMScheduler(skip_prk_steps=lowerCAmelCase_					)
     __lowerCAmelCase							     =     self.dummy_vae
     __lowerCAmelCase							     =     self.dummy_text_encoder
     __lowerCAmelCase							     =     XLMRobertaTokenizer.from_pretrained('hf-internal-testing/tiny-xlm-roberta'					)
     __lowerCAmelCase							     =     7_7

     __lowerCAmelCase							     =     self.dummy_image.to(lowerCAmelCase_					)
     __lowerCAmelCase							     =     init_image / 2 + 0.5

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

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

     __lowerCAmelCase							     =     output.images

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

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

     assert image.shape == (1, 3_2, 3_2, 3)
     __lowerCAmelCase							     =     np.array([0.44_27, 0.37_31, 0.42_49, 0.49_41, 0.45_46, 0.41_48, 0.41_93, 0.46_66, 0.44_99]					)

     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     lowercase		(    self			:		List[str]					)							->						Union[str, Any]:
     __lowerCAmelCase							     =     self.dummy_cond_unet
     __lowerCAmelCase							     =     PNDMScheduler(skip_prk_steps=lowerCAmelCase_					)
     __lowerCAmelCase							     =     self.dummy_vae
     __lowerCAmelCase							     =     self.dummy_text_encoder
     __lowerCAmelCase							     =     XLMRobertaTokenizer.from_pretrained('hf-internal-testing/tiny-xlm-roberta'					)
     __lowerCAmelCase							     =     7_7

     __lowerCAmelCase							     =     self.dummy_image.to(lowerCAmelCase_					)

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

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

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

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






  @unittest.skipIf(torch_device != 'cuda' ,  'This test requires a GPU'					)
  def     lowercase		(    self			:		Any					)							->						List[Any]:
     __lowerCAmelCase							     =     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
     __lowerCAmelCase							     =     init_image.resize((7_6_0, 5_0_4)					)

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

     __lowerCAmelCase							     =     'A fantasy landscape, trending on artstation'

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

     __lowerCAmelCase							     =     image[2_5_5:2_5_8, 3_8_3:3_8_6, -1]

     assert image.shape == (5_0_4, 7_6_0, 3)
     __lowerCAmelCase							     =     np.array([0.93_58, 0.93_97, 0.95_99, 0.99_01, 1.00_00, 1.00_00, 0.98_82, 1.00_00, 1.00_00]					)

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





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

  """simple docstring"""


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






  def     lowercase		(    self			:		List[Any]					)							->						Any:
     __lowerCAmelCase							     =     load_image(
         'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'
         '/img2img/sketch-mountains-input.jpg'					)
     __lowerCAmelCase							     =     init_image.resize((7_6_8, 5_1_2)					)
     __lowerCAmelCase							     =     load_numpy(
         'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/img2img/fantasy_landscape_alt.npy'					)

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

     __lowerCAmelCase							     =     'A fantasy landscape, trending on artstation'

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

     assert image.shape == (5_1_2, 7_6_8, 3)
     # img2img is flaky across GPUs even in fp32, so using MAE here
     assert np.abs(expected_image - image					).max() < 1e-2








import mpmath  # for roots of unity
import numpy as np



class     _UpperCAmelCase   :

  """simple docstring"""


  def __init__(    self			:		List[Any] ,  lowerCAmelCase_			:		Dict=None ,  lowerCAmelCase_			:		str=None					)							->						List[Any]:
     # Input as list
     __lowerCAmelCase							     =     list(poly_a or [0]					)[:]
     __lowerCAmelCase							     =     list(poly_b or [0]					)[:]

     # Remove leading zero coefficients
     while self.polyA[-1] == 0:
        self.polyA.pop()
     __lowerCAmelCase							     =     len(self.polyA					)

     while self.polyB[-1] == 0:
        self.polyB.pop()
     __lowerCAmelCase							     =     len(self.polyB					)

     # Add 0 to make lengths equal a power of 2
     __lowerCAmelCase							     =     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
     __lowerCAmelCase							     =     complex(mpmath.root(x=1 ,  n=self.c_max_length ,  k=1					)					)

     # The product
     __lowerCAmelCase							     =     self.__multiply()






  def     lowercase		(    self			:		Optional[int] ,  lowerCAmelCase_			:		str					)							->						Optional[int]:
     __lowerCAmelCase							     =     [[x] for x in self.polyA] if which == 'A' else [[x] for x in self.polyB]
     # Corner case
     if len(lowerCAmelCase_					) <= 1:
        return dft[0]
     #
     __lowerCAmelCase							     =     self.c_max_length // 2
     while next_ncol > 0:
        __lowerCAmelCase							     =     [[] for i in range(lowerCAmelCase_					)]
        __lowerCAmelCase							     =     self.root**next_ncol

        # First half of next step
        __lowerCAmelCase							     =     1
        for j in range(self.c_max_length // (next_ncol * 2)					):
           for i in range(lowerCAmelCase_					):
              new_dft[i].append(dft[i][j] + current_root * dft[i + next_ncol][j]					)
           current_root *= root
        # Second half of next step
        __lowerCAmelCase							     =     1
        for j in range(self.c_max_length // (next_ncol * 2)					):
           for i in range(lowerCAmelCase_					):
              new_dft[i].append(dft[i][j] - current_root * dft[i + next_ncol][j]					)
           current_root *= root
        # Update
        __lowerCAmelCase							     =     new_dft
        __lowerCAmelCase							     =     next_ncol // 2
     return dft[0]






  def     lowercase		(    self			:		Optional[int]					)							->						Any:
     __lowerCAmelCase							     =     self.__dft('A'					)
     __lowerCAmelCase							     =     self.__dft('B'					)
     __lowerCAmelCase							     =     [[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
     __lowerCAmelCase							     =     2
     while next_ncol <= self.c_max_length:
        __lowerCAmelCase							     =     [[] for i in range(lowerCAmelCase_					)]
        __lowerCAmelCase							     =     self.root ** (next_ncol // 2)
        __lowerCAmelCase							     =     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
        __lowerCAmelCase							     =     new_inverse_c
        next_ncol *= 2
     # Unpack
     __lowerCAmelCase							     =     [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			:		Dict					)							->						int:
     __lowerCAmelCase							     =     'A = ' + ' + '.join(
         f"""{coef}*x^{i}""" for coef, i in enumerate(self.polyA[: self.len_A]					)					)
     __lowerCAmelCase							     =     'B = ' + ' + '.join(
         f"""{coef}*x^{i}""" for coef, i in enumerate(self.polyB[: self.len_B]					)					)
     __lowerCAmelCase							     =     '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()

