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






       def 	_SCREAMING_SNAKE_CASE       (  self :							List[Any]			)->  Tuple:







              """simple docstring"""
              # clean up the VRAM after each test
              super().tearDown()
              gc.collect()
              torch.cuda.empty_cache()






       @property
       def 	_SCREAMING_SNAKE_CASE       (  self :							int			)->  str:







              """simple docstring"""
              UpperCamelCase					 =		1
              UpperCamelCase					 =		3
              UpperCamelCase					 =		(32, 32)

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






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







              """simple docstring"""
              torch.manual_seed(0			)
              UpperCamelCase					 =		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 	_SCREAMING_SNAKE_CASE       (  self :							Optional[Any]			)->  Any:







              """simple docstring"""
              torch.manual_seed(0			)
              UpperCamelCase					 =		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 	_SCREAMING_SNAKE_CASE       (  self :							Tuple			)->  Union[str, Any]:







              """simple docstring"""
              torch.manual_seed(0			)
              UpperCamelCase					 =		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 	_SCREAMING_SNAKE_CASE       (  self :							str			)->  List[Any]:







              """simple docstring"""
              def extract(*UpperCAmelCase_ :							Optional[int]       ,      **UpperCAmelCase_ :							str			):



                     class   __a  :






                            def __init__(  self :							Any			)->  int:







                                   """simple docstring"""
                                   UpperCamelCase					 =		torch.ones([0]			)






                            def 	_SCREAMING_SNAKE_CASE       (  self :							Any       ,      UpperCAmelCase_ :							Optional[Any]			)->  List[str]:







                                   """simple docstring"""
                                   self.pixel_values.to(UpperCAmelCase_			)
                                   return self





                     return Out()

              return extract






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







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

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

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

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

              UpperCamelCase					 =		output.images

              UpperCamelCase					 =		torch.Generator(device=UpperCAmelCase_			).manual_seed(0			)
              UpperCamelCase					 =		alt_pipe(
                  [prompt]       ,      generator=UpperCAmelCase_       ,      guidance_scale=6.0       ,      num_inference_steps=2       ,      output_type="np"       ,      image=UpperCAmelCase_       ,      return_dict=UpperCAmelCase_       ,      )[0]

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

              assert image.shape == (1, 32, 32, 3)
              UpperCamelCase					 =		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 	_SCREAMING_SNAKE_CASE       (  self :							Union[str, Any]			)->  Optional[int]:







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

              UpperCamelCase					 =		self.dummy_image.to(UpperCAmelCase_			)

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

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

              UpperCamelCase					 =		"A painting of a squirrel eating a burger"
              UpperCamelCase					 =		torch.manual_seed(0			)
              UpperCamelCase					 =		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 	_SCREAMING_SNAKE_CASE       (  self :							Optional[int]			)->  Tuple:







              """simple docstring"""
              UpperCamelCase					 =		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					 =		init_image.resize((760, 504)			)

              UpperCamelCase					 =		"BAAI/AltDiffusion"
              UpperCamelCase					 =		AltDiffusionImgaImgPipeline.from_pretrained(
                  UpperCAmelCase_       ,      safety_checker=UpperCAmelCase_       ,      )
              pipe.to(UpperCAmelCase_			)
              pipe.set_progress_bar_config(disable=UpperCAmelCase_			)
              pipe.enable_attention_slicing()

              UpperCamelCase					 =		"A fantasy landscape, trending on artstation"

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

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

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






       def 	_SCREAMING_SNAKE_CASE       (  self :							Tuple			)->  List[Any]:







              """simple docstring"""
              # clean up the VRAM after each test
              super().tearDown()
              gc.collect()
              torch.cuda.empty_cache()






       def 	_SCREAMING_SNAKE_CASE       (  self :							Dict			)->  Any:







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

              UpperCamelCase					 =		"BAAI/AltDiffusion"
              UpperCamelCase					 =		AltDiffusionImgaImgPipeline.from_pretrained(
                  UpperCAmelCase_       ,      safety_checker=UpperCAmelCase_       ,      )
              pipe.to(UpperCAmelCase_			)
              pipe.set_progress_bar_config(disable=UpperCAmelCase_			)
              pipe.enable_attention_slicing()

              UpperCamelCase					 =		"A fantasy landscape, trending on artstation"

              UpperCamelCase					 =		torch.manual_seed(0			)
              UpperCamelCase					 =		pipe(
                  prompt=UpperCAmelCase_       ,      image=UpperCAmelCase_       ,      strength=0.75       ,      guidance_scale=7.5       ,      generator=UpperCAmelCase_       ,      output_type="np"       ,      )
              UpperCamelCase					 =		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 inspect
import os
import unittest

import torch

import accelerate
from accelerate import debug_launcher
from accelerate.test_utils import (
    execute_subprocess_async,
    require_cpu,
    require_huggingface_suite,
    require_multi_gpu,
    require_single_gpu,
)
from accelerate.utils import patch_environment



@require_huggingface_suite
class   __a  (      unittest.TestCase	):






       def 	_SCREAMING_SNAKE_CASE       (  self :							Optional[int]			)->  Dict:







              """simple docstring"""
              UpperCamelCase					 =		inspect.getfile(accelerate.test_utils			)
              UpperCamelCase					 =		os.path.sep.join(
                  mod_file.split(os.path.sep			)[:-1] + ["scripts", "external_deps", "test_metrics.py"]			)

              from accelerate.test_utils.scripts.external_deps import test_metrics  # noqa: F401

              UpperCamelCase					 =		test_metrics






       @require_cpu
       def 	_SCREAMING_SNAKE_CASE       (  self :							Dict			)->  Union[str, Any]:







              """simple docstring"""
              debug_launcher(self.test_metrics.main       ,      num_processes=1			)






       @require_cpu
       def 	_SCREAMING_SNAKE_CASE       (  self :							Tuple			)->  List[str]:







              """simple docstring"""
              debug_launcher(self.test_metrics.main			)






       @require_single_gpu
       def 	_SCREAMING_SNAKE_CASE       (  self :							List[Any]			)->  str:







              """simple docstring"""
              self.test_metrics.main()






       @require_multi_gpu
       def 	_SCREAMING_SNAKE_CASE       (  self :							Optional[Any]			)->  Optional[Any]:







              """simple docstring"""
              print(f"Found {torch.cuda.device_count()} devices."			)
              UpperCamelCase					 =		["torchrun", f"--nproc_per_node={torch.cuda.device_count()}", self.test_file_path]
              with patch_environment(omp_num_threads=1			):
                     execute_subprocess_async(UpperCAmelCase_       ,      env=os.environ.copy()			)





