


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


     def        __snake_case(  self    ):
       super().tearDown()
       gc.collect()
       torch.cuda.empty_cache()


     @property
     def        __snake_case(  self    ):
       _UpperCAmelCase		:							str							=     1
       _UpperCAmelCase		:							Any							=     3
       _UpperCAmelCase		:							Optional[int]							=     (32, 32)

       _UpperCAmelCase		:							str							=     floats_tensor((batch_size, num_channels) + sizes							,					rng=random.Random(0    )    ).to(_SCREAMING_SNAKE_CASE    )
       return image


     @property
     def        __snake_case(  self    ):
       torch.manual_seed(0    )
       _UpperCAmelCase		:							List[Any]							=     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        __snake_case(  self    ):
       torch.manual_seed(0    )
       _UpperCAmelCase		:							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        __snake_case(  self    ):
       torch.manual_seed(0    )
       _UpperCAmelCase		:							Optional[Any]							=     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=50_06							,					)
       return RobertaSeriesModelWithTransformation(_SCREAMING_SNAKE_CASE    )


     @property
     def        __snake_case(  self    ):




       def extract(*A_							,					**A_    ):


         class 		_SCREAMING_SNAKE_CASE     :


              def __init__(  self    ):
                _UpperCAmelCase		:							Optional[int]							=     torch.ones([0]    )




              def        __snake_case(  self							,					A_    ):
                self.pixel_values.to(_SCREAMING_SNAKE_CASE    )
                return self


         return Out()

       return extract


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

       _UpperCAmelCase		:							int							=     self.dummy_image.to(_SCREAMING_SNAKE_CASE    )
       _UpperCAmelCase		:							int							=     init_image / 2 + 0.5

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

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

       _UpperCAmelCase		:							List[str]							=     output.images

       _UpperCAmelCase		:							Any							=     torch.Generator(device=_SCREAMING_SNAKE_CASE    ).manual_seed(0    )
       _UpperCAmelCase		:							str							=     alt_pipe(
           [prompt]							,					generator=_SCREAMING_SNAKE_CASE							,					guidance_scale=6.0							,					num_inference_steps=2							,					output_type="""np"""							,					image=_SCREAMING_SNAKE_CASE							,					return_dict=_SCREAMING_SNAKE_CASE							,					)[0]

       _UpperCAmelCase		:							int							=     image[0, -3:, -3:, -1]
       _UpperCAmelCase		:							Union[str, Any]							=     image_from_tuple[0, -3:, -3:, -1]

       assert image.shape == (1, 32, 32, 3)
       _UpperCAmelCase		:							List[str]							=     np.array([0.4_4_2_7, 0.3_7_3_1, 0.4_2_4_9, 0.4_9_4_1, 0.4_5_4_6, 0.4_1_4_8, 0.4_1_9_3, 0.4_6_6_6, 0.4_4_9_9]    )

       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        __snake_case(  self    ):
       _UpperCAmelCase		:							List[str]							=     self.dummy_cond_unet
       _UpperCAmelCase		:							List[str]							=     PNDMScheduler(skip_prk_steps=_SCREAMING_SNAKE_CASE    )
       _UpperCAmelCase		:							Tuple							=     self.dummy_vae
       _UpperCAmelCase		:							int							=     self.dummy_text_encoder
       _UpperCAmelCase		:							int							=     XLMRobertaTokenizer.from_pretrained("""hf-internal-testing/tiny-xlm-roberta"""    )
       _UpperCAmelCase		:							str							=     77

       _UpperCAmelCase		:							List[str]							=     self.dummy_image.to(_SCREAMING_SNAKE_CASE    )

       # put models in fp16
       _UpperCAmelCase		:							Dict							=     unet.half()
       _UpperCAmelCase		:							Any							=     vae.half()
       _UpperCAmelCase		:							Union[str, Any]							=     bert.half()

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

       _UpperCAmelCase		:							Any							=     """A painting of a squirrel eating a burger"""
       _UpperCAmelCase		:							List[str]							=     torch.manual_seed(0    )
       _UpperCAmelCase		:							Optional[Any]							=     alt_pipe(
           [prompt]							,					generator=_SCREAMING_SNAKE_CASE							,					num_inference_steps=2							,					output_type="""np"""							,					image=_SCREAMING_SNAKE_CASE							,					).images

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




     @unittest.skipIf(torch_device != """cuda"""							,					"""This test requires a GPU"""    )
     def        __snake_case(  self    ):
       _UpperCAmelCase		:							str							=     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
       _UpperCAmelCase		:							str							=     init_image.resize((7_60, 5_04)    )

       _UpperCAmelCase		:							Union[str, Any]							=     """BAAI/AltDiffusion"""
       _UpperCAmelCase		:							List[str]							=     AltDiffusionImgaImgPipeline.from_pretrained(
           _SCREAMING_SNAKE_CASE							,					safety_checker=_SCREAMING_SNAKE_CASE							,					)
       pipe.to(_SCREAMING_SNAKE_CASE    )
       pipe.set_progress_bar_config(disable=_SCREAMING_SNAKE_CASE    )
       pipe.enable_attention_slicing()

       _UpperCAmelCase		:							str							=     """A fantasy landscape, trending on artstation"""

       _UpperCAmelCase		:							List[str]							=     torch.manual_seed(0    )
       _UpperCAmelCase		:							Any							=     pipe(
           prompt=_SCREAMING_SNAKE_CASE							,					image=_SCREAMING_SNAKE_CASE							,					strength=0.7_5							,					guidance_scale=7.5							,					generator=_SCREAMING_SNAKE_CASE							,					output_type="""np"""							,					)
       _UpperCAmelCase		:							Union[str, Any]							=     output.images[0]

       _UpperCAmelCase		:							Dict							=     image[2_55:2_58, 3_83:3_86, -1]

       assert image.shape == (5_04, 7_60, 3)
       _UpperCAmelCase		:							List[Any]							=     np.array([0.9_3_5_8, 0.9_3_9_7, 0.9_5_9_9, 0.9_9_0_1, 1.0_0_0_0, 1.0_0_0_0, 0.9_8_8_2, 1.0_0_0_0, 1.0_0_0_0]    )

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


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


     def        __snake_case(  self    ):
       super().tearDown()
       gc.collect()
       torch.cuda.empty_cache()




     def        __snake_case(  self    ):
       _UpperCAmelCase		:							str							=     load_image(
           """https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"""
           """/img2img/sketch-mountains-input.jpg"""    )
       _UpperCAmelCase		:							Dict							=     init_image.resize((7_68, 5_12)    )
       _UpperCAmelCase		:							Any							=     load_numpy(
           """https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/img2img/fantasy_landscape_alt.npy"""    )

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

       _UpperCAmelCase		:							Union[str, Any]							=     """A fantasy landscape, trending on artstation"""

       _UpperCAmelCase		:							Optional[Any]							=     torch.manual_seed(0    )
       _UpperCAmelCase		:							str							=     pipe(
           prompt=_SCREAMING_SNAKE_CASE							,					image=_SCREAMING_SNAKE_CASE							,					strength=0.7_5							,					guidance_scale=7.5							,					generator=_SCREAMING_SNAKE_CASE							,					output_type="""np"""							,					)
       _UpperCAmelCase		:							Optional[Any]							=     output.images[0]

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













def 		lowercase__  (				_UpperCamelCase       ,							_UpperCamelCase) ->			bool:
      """simple docstring"""

      return numa ^ numa < 0


if __name__ == "__main__":
       import doctest

       doctest.testmod()



