


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







   def 				lowerCamelCase	(			self		)			->	Tuple:

       '''simple docstring'''




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







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

       '''simple docstring'''




       snake_case					:   Union[str, Any]         =					1
       snake_case					:   Dict         =					3
       snake_case					:   Union[str, Any]         =					(32, 32)

       snake_case					:   Dict         =					floats_tensor((batch_size, num_channels) + sizes		,			rng=random.Random(0		)		).to(UpperCamelCase__		)
       return image







   @property
   def 				lowerCamelCase	(			self		)			->	Any:

       '''simple docstring'''




       torch.manual_seed(0		)
       snake_case					:   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 				lowerCamelCase	(			self		)			->	Union[str, Any]:

       '''simple docstring'''




       torch.manual_seed(0		)
       snake_case					:   List[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 				lowerCamelCase	(			self		)			->	Dict:

       '''simple docstring'''




       torch.manual_seed(0		)
       snake_case					:   str         =					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=5006		,			)
       return RobertaSeriesModelWithTransformation(UpperCamelCase__		)







   @property
   def 				lowerCamelCase	(			self		)			->	List[Any]:

       '''simple docstring'''




       def extract(*UpperCamelCase__		,			**UpperCamelCase__		):







           class _lowerCAmelCase       :







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

                  '''simple docstring'''




                  snake_case					:   str         =					torch.ones([0]		)







              def 				lowerCamelCase	(			self		,			UpperCamelCase__		)			->	Optional[Any]:

                  '''simple docstring'''




                  self.pixel_values.to(UpperCamelCase__		)
                  return self



           return Out()

       return extract







   def 				lowerCamelCase	(			self		)			->	int:

       '''simple docstring'''




       snake_case					:   List[str]         =					"cpu"  # ensure determinism for the device-dependent torch.Generator
       snake_case					:   List[Any]         =					self.dummy_cond_unet
       snake_case					:   Tuple         =					PNDMScheduler(skip_prk_steps=UpperCamelCase__		)
       snake_case					:   Optional[Any]         =					self.dummy_vae
       snake_case					:   List[str]         =					self.dummy_text_encoder
       snake_case					:   Dict         =					XLMRobertaTokenizer.from_pretrained("hf-internal-testing/tiny-xlm-roberta"		)
       snake_case					:   Tuple         =					77

       snake_case					:   Tuple         =					self.dummy_image.to(UpperCamelCase__		)
       snake_case					:   Optional[int]         =					init_image / 2 + 0.5

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

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

       snake_case					:   Any         =					output.images

       snake_case					:   Optional[int]         =					torch.Generator(device=UpperCamelCase__		).manual_seed(0		)
       snake_case					:   Any         =					alt_pipe(
           [prompt]		,			generator=UpperCamelCase__		,			guidance_scale=6.0		,			num_inference_steps=2		,			output_type="np"		,			image=UpperCamelCase__		,			return_dict=UpperCamelCase__		,			)[0]

       snake_case					:   Dict         =					image[0, -3:, -3:, -1]
       snake_case					:   Union[str, Any]         =					image_from_tuple[0, -3:, -3:, -1]

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

       '''simple docstring'''




       snake_case					:   Optional[Any]         =					self.dummy_cond_unet
       snake_case					:   List[str]         =					PNDMScheduler(skip_prk_steps=UpperCamelCase__		)
       snake_case					:   Any         =					self.dummy_vae
       snake_case					:   Tuple         =					self.dummy_text_encoder
       snake_case					:   Any         =					XLMRobertaTokenizer.from_pretrained("hf-internal-testing/tiny-xlm-roberta"		)
       snake_case					:   List[str]         =					77

       snake_case					:   Any         =					self.dummy_image.to(UpperCamelCase__		)

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

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

       snake_case					:   int         =					"A painting of a squirrel eating a burger"
       snake_case					:   Dict         =					torch.manual_seed(0		)
       snake_case					:   int         =					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 				lowerCamelCase	(			self		)			->	Tuple:

       '''simple docstring'''




       snake_case					:   Any         =					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
       snake_case					:   Optional[int]         =					init_image.resize((760, 504)		)

       snake_case					:   Optional[int]         =					"BAAI/AltDiffusion"
       snake_case					:   str         =					AltDiffusionImgaImgPipeline.from_pretrained(
           UpperCamelCase__		,			safety_checker=UpperCamelCase__		,			)
       pipe.to(UpperCamelCase__		)
       pipe.set_progress_bar_config(disable=UpperCamelCase__		)
       pipe.enable_attention_slicing()

       snake_case					:   Dict         =					"A fantasy landscape, trending on artstation"

       snake_case					:   int         =					torch.manual_seed(0		)
       snake_case					:   List[str]         =					pipe(
           prompt=UpperCamelCase__		,			image=UpperCamelCase__		,			strength=0.75		,			guidance_scale=7.5		,			generator=UpperCamelCase__		,			output_type="np"		,			)
       snake_case					:   int         =					output.images[0]

       snake_case					:   List[Any]         =					image[255:258, 383:386, -1]

       assert image.shape == (504, 760, 3)
       snake_case					:   Dict         =					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			):







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

       '''simple docstring'''




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







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

       '''simple docstring'''




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

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

       snake_case					:   List[Any]         =					"A fantasy landscape, trending on artstation"

       snake_case					:   Dict         =					torch.manual_seed(0		)
       snake_case					:   Union[str, Any]         =					pipe(
           prompt=UpperCamelCase__		,			image=UpperCamelCase__		,			strength=0.75		,			guidance_scale=7.5		,			generator=UpperCamelCase__		,			output_type="np"		,			)
       snake_case					:   str         =					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"""



from collections import OrderedDict
from typing import Mapping

from ...configuration_utils import PretrainedConfig
from ...onnx import OnnxConfig
from ...utils import logging


__snake_case         =					logging.get_logger(__name__)

__snake_case         =					{
    """facebook/xmod-base""": """https://huggingface.co/facebook/xmod-base/resolve/main/config.json""",
    """facebook/xmod-large-prenorm""": """https://huggingface.co/facebook/xmod-large-prenorm/resolve/main/config.json""",
    """facebook/xmod-base-13-125k""": """https://huggingface.co/facebook/xmod-base-13-125k/resolve/main/config.json""",
    """facebook/xmod-base-30-125k""": """https://huggingface.co/facebook/xmod-base-30-125k/resolve/main/config.json""",
    """facebook/xmod-base-30-195k""": """https://huggingface.co/facebook/xmod-base-30-195k/resolve/main/config.json""",
    """facebook/xmod-base-60-125k""": """https://huggingface.co/facebook/xmod-base-60-125k/resolve/main/config.json""",
    """facebook/xmod-base-60-265k""": """https://huggingface.co/facebook/xmod-base-60-265k/resolve/main/config.json""",
    """facebook/xmod-base-75-125k""": """https://huggingface.co/facebook/xmod-base-75-125k/resolve/main/config.json""",
    """facebook/xmod-base-75-269k""": """https://huggingface.co/facebook/xmod-base-75-269k/resolve/main/config.json""",
}







class _lowerCAmelCase       (   snake_case_			):
   __UpperCAmelCase				: Union[str, Any]													=						'''xmod'''







   def __init__(			self		,			UpperCamelCase__=3_0522		,			UpperCamelCase__=768		,			UpperCamelCase__=12		,			UpperCamelCase__=12		,			UpperCamelCase__=3072		,			UpperCamelCase__="gelu"		,			UpperCamelCase__=0.1		,			UpperCamelCase__=0.1		,			UpperCamelCase__=512		,			UpperCamelCase__=2		,			UpperCamelCase__=0.02		,			UpperCamelCase__=1e-12		,			UpperCamelCase__=1		,			UpperCamelCase__=0		,			UpperCamelCase__=2		,			UpperCamelCase__="absolute"		,			UpperCamelCase__=True		,			UpperCamelCase__=None		,			UpperCamelCase__=False		,			UpperCamelCase__=2		,			UpperCamelCase__=False		,			UpperCamelCase__=True		,			UpperCamelCase__=True		,			UpperCamelCase__=("en_XX",)		,			UpperCamelCase__=None		,			**UpperCamelCase__		,			)			->	List[str]:

       '''simple docstring'''




       super().__init__(pad_token_id=UpperCamelCase__		,			bos_token_id=UpperCamelCase__		,			eos_token_id=UpperCamelCase__		,			**UpperCamelCase__		)

       snake_case					:   Optional[int]         =					vocab_size
       snake_case					:   Optional[Any]         =					hidden_size
       snake_case					:   Union[str, Any]         =					num_hidden_layers
       snake_case					:   int         =					num_attention_heads
       snake_case					:   Dict         =					hidden_act
       snake_case					:   Union[str, Any]         =					intermediate_size
       snake_case					:   str         =					hidden_dropout_prob
       snake_case					:   List[str]         =					attention_probs_dropout_prob
       snake_case					:   Dict         =					max_position_embeddings
       snake_case					:   Union[str, Any]         =					type_vocab_size
       snake_case					:   Union[str, Any]         =					initializer_range
       snake_case					:   str         =					layer_norm_eps
       snake_case					:   Dict         =					position_embedding_type
       snake_case					:   Any         =					use_cache
       snake_case					:   str         =					classifier_dropout
       snake_case					:   Any         =					pre_norm
       snake_case					:   Union[str, Any]         =					adapter_reduction_factor
       snake_case					:   Optional[int]         =					adapter_layer_norm
       snake_case					:   int         =					adapter_reuse_layer_norm
       snake_case					:   List[Any]         =					ln_before_adapter
       snake_case					:   Tuple         =					list(UpperCamelCase__		)
       snake_case					:   Tuple         =					default_language







class _lowerCAmelCase       (   snake_case_			):







   @property
   def 				lowerCamelCase	(			self		)			->	Mapping[str, Mapping[int, str]]:

       '''simple docstring'''




       if self.task == "multiple-choice":
           snake_case					:   str         =					{0: "batch", 1: "choice", 2: "sequence"}
       else:
           snake_case					:   Tuple         =					{0: "batch", 1: "sequence"}
       return OrderedDict(
           [
               ("input_ids", dynamic_axis),
               ("attention_mask", dynamic_axis),
           ]		)





