

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



 """simple docstring"""



 def  UpperCAmelCase__						(						self   :			Union[str, Any]     )    ->					Dict:


      """simple docstring"""




      super().tearDown()
      gc.collect()
      torch.cuda.empty_cache()
 @property
 def  UpperCAmelCase__						(						self   :			Optional[Any]     )    ->					Dict:


      """simple docstring"""




      __SCREAMING_SNAKE_CASE      = 1
      __SCREAMING_SNAKE_CASE      = 3
      __SCREAMING_SNAKE_CASE      = (32, 32)

      __SCREAMING_SNAKE_CASE      = floats_tensor((batch_size, num_channels) + sizes   ,				rng=random.Random(0     )     ).to(__SCREAMING_SNAKE_CASE     )
      return image
 @property
 def  UpperCAmelCase__						(						self   :			Union[str, Any]     )    ->					Dict:


      """simple docstring"""




      torch.manual_seed(0     )
      __SCREAMING_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   :			Any     )    ->					Tuple:


      """simple docstring"""




      torch.manual_seed(0     )
      __SCREAMING_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   :			Any     )    ->					Optional[int]:


      """simple docstring"""




      torch.manual_seed(0     )
      __SCREAMING_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=5_006   ,				)
      return RobertaSeriesModelWithTransformation(__SCREAMING_SNAKE_CASE     )
 @property
 def  UpperCAmelCase__						(						self   :			Union[str, Any]     )    ->					Optional[int]:


      """simple docstring"""




      def extract(*__SCREAMING_SNAKE_CASE   :			Tuple   ,				**__SCREAMING_SNAKE_CASE   :			Optional[Any]     ):







           class 		lowerCAmelCase__  :



            """simple docstring"""



            def __init__(						self   :			List[str]     )    ->					List[str]:


                 """simple docstring"""




                 __SCREAMING_SNAKE_CASE      = torch.ones([0]     )


            def  UpperCAmelCase__						(						self   :			List[str]   ,				__SCREAMING_SNAKE_CASE   :			int     )    ->					List[str]:


                 """simple docstring"""




                 self.pixel_values.to(__SCREAMING_SNAKE_CASE     )
                 return self


           return Out()

      return extract
 def  UpperCAmelCase__						(						self   :			Tuple     )    ->					List[str]:


      """simple docstring"""




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

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

      # make sure here that pndm scheduler skips prk
      __SCREAMING_SNAKE_CASE      = 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   ,				)
      __SCREAMING_SNAKE_CASE      = VaeImageProcessor(vae_scale_factor=alt_pipe.vae_scale_factor   ,				do_normalize=__SCREAMING_SNAKE_CASE     )
      __SCREAMING_SNAKE_CASE      = alt_pipe.to(__SCREAMING_SNAKE_CASE     )
      alt_pipe.set_progress_bar_config(disable=__SCREAMING_SNAKE_CASE     )

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

      __SCREAMING_SNAKE_CASE      = output.images

      __SCREAMING_SNAKE_CASE      = torch.Generator(device=__SCREAMING_SNAKE_CASE     ).manual_seed(0     )
      __SCREAMING_SNAKE_CASE      = 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]

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

      assert image.shape == (1, 32, 32, 3)
      __SCREAMING_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]     )    ->					Dict:


      """simple docstring"""




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

      __SCREAMING_SNAKE_CASE      = self.dummy_image.to(__SCREAMING_SNAKE_CASE     )

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

      # make sure here that pndm scheduler skips prk
      __SCREAMING_SNAKE_CASE      = 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   ,				)
      __SCREAMING_SNAKE_CASE      = VaeImageProcessor(vae_scale_factor=alt_pipe.vae_scale_factor   ,				do_normalize=__SCREAMING_SNAKE_CASE     )
      __SCREAMING_SNAKE_CASE      = alt_pipe.to(__SCREAMING_SNAKE_CASE     )
      alt_pipe.set_progress_bar_config(disable=__SCREAMING_SNAKE_CASE     )

      __SCREAMING_SNAKE_CASE      = """A painting of a squirrel eating a burger"""
      __SCREAMING_SNAKE_CASE      = torch.manual_seed(0     )
      __SCREAMING_SNAKE_CASE      = 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  UpperCAmelCase__						(						self   :			Any     )    ->					List[str]:


      """simple docstring"""




      __SCREAMING_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
      __SCREAMING_SNAKE_CASE      = init_image.resize((760, 504)     )

      __SCREAMING_SNAKE_CASE      = """BAAI/AltDiffusion"""
      __SCREAMING_SNAKE_CASE      = 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()

      __SCREAMING_SNAKE_CASE      = """A fantasy landscape, trending on artstation"""

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

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

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



 """simple docstring"""



 def  UpperCAmelCase__						(						self   :			Optional[Any]     )    ->					int:


      """simple docstring"""




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


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


      """simple docstring"""




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

      __SCREAMING_SNAKE_CASE      = """BAAI/AltDiffusion"""
      __SCREAMING_SNAKE_CASE      = 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()

      __SCREAMING_SNAKE_CASE      = """A fantasy landscape, trending on artstation"""

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









'''simple docstring'''



import unittest

import numpy as np
import torch

from diffusers import KarrasVePipeline, KarrasVeScheduler, UNetaDModel
from diffusers.utils.testing_utils import enable_full_determinism, require_torch, slow, torch_device


enable_full_determinism()







class 		lowerCAmelCase__  ( unittest.TestCase   ):



 """simple docstring"""



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


      """simple docstring"""




      torch.manual_seed(0     )
      __SCREAMING_SNAKE_CASE      = UNetaDModel(
          block_out_channels=(32, 64)   ,				layers_per_block=2   ,				sample_size=32   ,				in_channels=3   ,				out_channels=3   ,				down_block_types=("""DownBlock2D""", """AttnDownBlock2D""")   ,				up_block_types=("""AttnUpBlock2D""", """UpBlock2D""")   ,				)
      return model


 def  UpperCAmelCase__						(						self   :			Optional[Any]     )    ->					Tuple:


      """simple docstring"""




      __SCREAMING_SNAKE_CASE      = self.dummy_uncond_unet
      __SCREAMING_SNAKE_CASE      = KarrasVeScheduler()

      __SCREAMING_SNAKE_CASE      = KarrasVePipeline(unet=__SCREAMING_SNAKE_CASE   ,				scheduler=__SCREAMING_SNAKE_CASE     )
      pipe.to(__SCREAMING_SNAKE_CASE     )
      pipe.set_progress_bar_config(disable=__SCREAMING_SNAKE_CASE     )

      __SCREAMING_SNAKE_CASE      = torch.manual_seed(0     )
      __SCREAMING_SNAKE_CASE      = pipe(num_inference_steps=2   ,				generator=__SCREAMING_SNAKE_CASE   ,				output_type="""numpy"""     ).images

      __SCREAMING_SNAKE_CASE      = torch.manual_seed(0     )
      __SCREAMING_SNAKE_CASE      = pipe(num_inference_steps=2   ,				generator=__SCREAMING_SNAKE_CASE   ,				output_type="""numpy"""   ,				return_dict=__SCREAMING_SNAKE_CASE     )[0]

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

      assert image.shape == (1, 32, 32, 3)
      __SCREAMING_SNAKE_CASE      = np.array([0.0, 1.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0]     )

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







@slow
@require_torch
class 		lowerCAmelCase__  ( unittest.TestCase   ):



 """simple docstring"""



 def  UpperCAmelCase__						(						self   :			List[Any]     )    ->					Union[str, Any]:


      """simple docstring"""




      __SCREAMING_SNAKE_CASE      = """google/ncsnpp-celebahq-256"""
      __SCREAMING_SNAKE_CASE      = UNetaDModel.from_pretrained(__SCREAMING_SNAKE_CASE     )
      __SCREAMING_SNAKE_CASE      = KarrasVeScheduler()

      __SCREAMING_SNAKE_CASE      = KarrasVePipeline(unet=__SCREAMING_SNAKE_CASE   ,				scheduler=__SCREAMING_SNAKE_CASE     )
      pipe.to(__SCREAMING_SNAKE_CASE     )
      pipe.set_progress_bar_config(disable=__SCREAMING_SNAKE_CASE     )

      __SCREAMING_SNAKE_CASE      = torch.manual_seed(0     )
      __SCREAMING_SNAKE_CASE      = pipe(num_inference_steps=20   ,				generator=__SCREAMING_SNAKE_CASE   ,				output_type="""numpy"""     ).images

      __SCREAMING_SNAKE_CASE      = image[0, -3:, -3:, -1]
      assert image.shape == (1, 256, 256, 3)
      __SCREAMING_SNAKE_CASE      = np.array([0.578, 0.5811, 0.5924, 0.5809, 0.587, 0.5886, 0.5861, 0.5802, 0.586]     )

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






