
"""simple docstring"""
import gc
import unittest

import numpy as np
import torch
from transformers import CLIPTextConfig, CLIPTextModel, CLIPTokenizer

from diffusers import (
    AutoencoderKL,
    DDIMScheduler,
    EulerAncestralDiscreteScheduler,
    PNDMScheduler,
    StableDiffusionModelEditingPipeline,
    UNetaDConditionModel,
)
from diffusers.utils import slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu, skip_mps

from ..pipeline_params import TEXT_TO_IMAGE_BATCH_PARAMS, TEXT_TO_IMAGE_IMAGE_PARAMS, TEXT_TO_IMAGE_PARAMS
from ..test_pipelines_common import PipelineKarrasSchedulerTesterMixin, PipelineLatentTesterMixin, PipelineTesterMixin


enable_full_determinism()




@skip_mps
class    _snake_case  (						_lowercase ,				_lowercase ,				_lowercase ,				unittest.TestCase       ):
							SCREAMING_SNAKE_CASE		     =						StableDiffusionModelEditingPipeline
							SCREAMING_SNAKE_CASE		     =						TEXT_TO_IMAGE_PARAMS
							SCREAMING_SNAKE_CASE		     =						TEXT_TO_IMAGE_BATCH_PARAMS
							SCREAMING_SNAKE_CASE		     =						TEXT_TO_IMAGE_IMAGE_PARAMS
							SCREAMING_SNAKE_CASE		     =						TEXT_TO_IMAGE_IMAGE_PARAMS






							def 							__lowercase  (					self						:			List[Any]    ):






												"""simple docstring"""



												torch.manual_seed(0    )
												__a:						Optional[int]		    =					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						, )
												__a:						str		    =					DDIMScheduler()
												torch.manual_seed(0    )
												__a:						Any		    =					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						, )
												torch.manual_seed(0    )
												__a:						Dict		    =					CLIPTextConfig(
												    bos_token_id=0						, eos_token_id=2						, hidden_size=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=1_0_0_0						, )
												__a:						Dict		    =					CLIPTextModel(snake_case__    )
												__a:						Dict		    =					CLIPTokenizer.from_pretrained("""hf-internal-testing/tiny-random-clip"""    )

												__a:						str		    =					{
												    """unet""": unet,
												    """scheduler""": scheduler,
												    """vae""": vae,
												    """text_encoder""": text_encoder,
												    """tokenizer""": tokenizer,
												    """safety_checker""": None,
												    """feature_extractor""": None,
												}
												return components






							def 							__lowercase  (					self						:			List[str]						, snake_case__						:			Dict						, snake_case__						:			Tuple=0    ):






												"""simple docstring"""



												__a:						str		    =					torch.manual_seed(snake_case__    )
												__a:						List[Any]		    =					{
												    """prompt""": """A field of roses""",
												    """generator""": generator,
												    # Setting height and width to None to prevent OOMs on CPU.
												    """height""": None,
												    """width""": None,
												    """num_inference_steps""": 2,
												    """guidance_scale""": 6.0,
												    """output_type""": """numpy""",
												}
												return inputs






							def 							__lowercase  (					self						:			Dict    ):






												"""simple docstring"""



												__a:						List[Any]		    =					"""cpu"""  # ensure determinism for the device-dependent torch.Generator
												__a:						List[str]		    =					self.get_dummy_components()
												__a:						Any		    =					StableDiffusionModelEditingPipeline(**snake_case__    )
												__a:						Dict		    =					sd_pipe.to(snake_case__    )
												sd_pipe.set_progress_bar_config(disable=snake_case__    )

												__a:						str		    =					self.get_dummy_inputs(snake_case__    )
												__a:						int		    =					sd_pipe(**snake_case__    ).images
												__a:						Union[str, Any]		    =					image[0, -3:, -3:, -1]
												assert image.shape == (1, 6_4, 6_4, 3)

												__a:						int		    =					np.array([0.4_755, 0.5_132, 0.4_976, 0.3_904, 0.3_554, 0.4_765, 0.5_139, 0.5_158, 0.4_889]    )

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






							def 							__lowercase  (					self						:			Tuple    ):






												"""simple docstring"""



												__a:						Union[str, Any]		    =					"""cpu"""  # ensure determinism for the device-dependent torch.Generator
												__a:						Optional[int]		    =					self.get_dummy_components()
												__a:						str		    =					StableDiffusionModelEditingPipeline(**snake_case__    )
												__a:						Union[str, Any]		    =					sd_pipe.to(snake_case__    )
												sd_pipe.set_progress_bar_config(disable=snake_case__    )

												__a:						Union[str, Any]		    =					self.get_dummy_inputs(snake_case__    )
												__a:						List[Any]		    =					"""french fries"""
												__a:						str		    =					sd_pipe(**snake_case__						, negative_prompt=snake_case__    )
												__a:						str		    =					output.images
												__a:						str		    =					image[0, -3:, -3:, -1]

												assert image.shape == (1, 6_4, 6_4, 3)

												__a:						Dict		    =					np.array([0.4_992, 0.5_101, 0.5_004, 0.3_949, 0.3_604, 0.4_735, 0.5_216, 0.5_204, 0.4_913]    )

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






							def 							__lowercase  (					self						:			Union[str, Any]    ):






												"""simple docstring"""



												__a:						Dict		    =					"""cpu"""  # ensure determinism for the device-dependent torch.Generator
												__a:						Dict		    =					self.get_dummy_components()
												__a:						Optional[Any]		    =					EulerAncestralDiscreteScheduler(
												    beta_start=0.00_085						, beta_end=0.012						, beta_schedule="""scaled_linear"""    )
												__a:						int		    =					StableDiffusionModelEditingPipeline(**snake_case__    )
												__a:						Optional[int]		    =					sd_pipe.to(snake_case__    )
												sd_pipe.set_progress_bar_config(disable=snake_case__    )

												__a:						Optional[Any]		    =					self.get_dummy_inputs(snake_case__    )
												__a:						Union[str, Any]		    =					sd_pipe(**snake_case__    ).images
												__a:						str		    =					image[0, -3:, -3:, -1]

												assert image.shape == (1, 6_4, 6_4, 3)

												__a:						List[str]		    =					np.array([0.4_747, 0.5_372, 0.4_779, 0.4_982, 0.5_543, 0.4_816, 0.5_238, 0.4_904, 0.5_027]    )

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






							def 							__lowercase  (					self						:			int    ):






												"""simple docstring"""



												__a:						Any		    =					"""cpu"""  # ensure determinism for the device-dependent torch.Generator
												__a:						List[Any]		    =					self.get_dummy_components()
												__a:						str		    =					PNDMScheduler()
												__a:						Any		    =					StableDiffusionModelEditingPipeline(**snake_case__    )
												__a:						Tuple		    =					sd_pipe.to(snake_case__    )
												sd_pipe.set_progress_bar_config(disable=snake_case__    )

												__a:						Any		    =					self.get_dummy_inputs(snake_case__    )
												# the pipeline does not expect pndm so test if it raises error.
												with self.assertRaises(snake_case__    ):
																	__a:						Optional[int]		    =					sd_pipe(**snake_case__    ).images






							def 							__lowercase  (					self						:			Optional[int]    ):






												"""simple docstring"""



												super().test_inference_batch_single_identical(expected_max_diff=5e-3    )






							def 							__lowercase  (					self						:			Tuple    ):






												"""simple docstring"""



												super().test_attention_slicing_forward_pass(expected_max_diff=5e-3    )






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






							def 							__lowercase  (					self						:			Optional[int]    ):






												"""simple docstring"""



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






							def 							__lowercase  (					self						:			Union[str, Any]						, snake_case__						:			Optional[int]=0    ):






												"""simple docstring"""



												__a:						Optional[Any]		    =					torch.manual_seed(snake_case__    )
												__a:						Optional[int]		    =					{
												    """prompt""": """A field of roses""",
												    """generator""": generator,
												    """num_inference_steps""": 3,
												    """guidance_scale""": 7.5,
												    """output_type""": """numpy""",
												}
												return inputs






							def 							__lowercase  (					self						:			List[str]    ):






												"""simple docstring"""



												__a:						List[str]		    =					"""CompVis/stable-diffusion-v1-4"""
												__a:						Dict		    =					StableDiffusionModelEditingPipeline.from_pretrained(snake_case__						, safety_checker=snake_case__    )
												pipe.to(snake_case__    )
												pipe.set_progress_bar_config(disable=snake_case__    )
												pipe.enable_attention_slicing()

												__a:						str		    =					self.get_inputs()
												__a:						str		    =					pipe(**snake_case__    ).images
												__a:						Dict		    =					image[0, -3:, -3:, -1].flatten()

												assert image.shape == (1, 5_1_2, 5_1_2, 3)

												__a:						Tuple		    =					np.array(
												    [0.6_749_496, 0.6_386_453, 0.51_443_267, 0.66_094_905, 0.61_921_215, 0.5_491_332, 0.5_744_417, 0.58_075_106, 0.5_174_658]    )

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

												# make sure image changes after editing
												pipe.edit_model("""A pack of roses"""						, """A pack of blue roses"""    )

												__a:						Any		    =					pipe(**snake_case__    ).images
												__a:						Optional[int]		    =					image[0, -3:, -3:, -1].flatten()

												assert image.shape == (1, 5_1_2, 5_1_2, 3)

												assert np.abs(expected_slice - image_slice    ).max() > 1e-1






							def 							__lowercase  (					self						:			Optional[int]    ):






												"""simple docstring"""



												torch.cuda.empty_cache()
												torch.cuda.reset_max_memory_allocated()
												torch.cuda.reset_peak_memory_stats()

												__a:						List[Any]		    =					"""CompVis/stable-diffusion-v1-4"""
												__a:						Dict		    =					DDIMScheduler.from_pretrained(snake_case__						, subfolder="""scheduler"""    )
												__a:						str		    =					StableDiffusionModelEditingPipeline.from_pretrained(
												    snake_case__						, scheduler=snake_case__						, safety_checker=snake_case__    )
												__a:						str		    =					pipe.to(snake_case__    )
												pipe.set_progress_bar_config(disable=snake_case__    )
												pipe.enable_attention_slicing(1    )
												pipe.enable_sequential_cpu_offload()

												__a:						int		    =					self.get_inputs()
												__a:						Optional[Any]		    =					pipe(**snake_case__    )

												__a:						int		    =					torch.cuda.max_memory_allocated()
												# make sure that less than 4.4 GB is allocated
												assert mem_bytes < 4.4 * 1_0**9


"""simple docstring"""
import torch
import torch.nn as nn
from transformers.modeling_utils import ModuleUtilsMixin
from transformers.models.ta.modeling_ta import (
    TaBlock,
    TaConfig,
    TaLayerNorm,
)

from ...configuration_utils import ConfigMixin, register_to_config
from ...models import ModelMixin




class    _snake_case  (						_lowercase ,				_lowercase ,				_lowercase       ):






							@register_to_config
							def __init__(					self						:			str						, snake_case__						:			int						, snake_case__						:			int						, snake_case__						:			int						, snake_case__						:			float						, snake_case__						:			int						, snake_case__						:			int						, snake_case__						:			int						, snake_case__						:			int						, snake_case__						:			str						, snake_case__						:			bool = False						, ):






												"""simple docstring"""



												super().__init__()

												__a:						Optional[Any]		    =					nn.Linear(snake_case__						, snake_case__						, bias=snake_case__    )

												__a:						Optional[Any]		    =					nn.Embedding(snake_case__						, snake_case__    )
												__a:						Any		    =					False

												__a:						Any		    =					nn.Dropout(p=snake_case__    )

												__a:						List[str]		    =					TaConfig(
												    d_model=snake_case__						, num_heads=snake_case__						, d_kv=snake_case__						, d_ff=snake_case__						, feed_forward_proj=snake_case__						, dropout_rate=snake_case__						, is_decoder=snake_case__						, is_encoder_decoder=snake_case__						, )
												__a:						Optional[int]		    =					nn.ModuleList()
												for lyr_num in range(snake_case__    ):
																	__a:						Union[str, Any]		    =					TaBlock(snake_case__    )
																	self.encoders.append(snake_case__    )

												__a:						List[str]		    =					TaLayerNorm(snake_case__    )
												__a:						Optional[int]		    =					nn.Dropout(p=snake_case__    )






							def 							__lowercase  (					self						:			str						, snake_case__						:			Optional[Any]						, snake_case__						:			int    ):






												"""simple docstring"""



												__a:						str		    =					self.input_proj(snake_case__    )

												# terminal relative positional encodings
												__a:						Optional[int]		    =					encoder_inputs.shape[1]
												__a:						List[Any]		    =					torch.arange(snake_case__						, device=encoder_inputs.device    )

												__a:						Union[str, Any]		    =					encoder_inputs_mask.sum(-1    )
												__a:						str		    =					torch.roll(input_positions.unsqueeze(0    )						, tuple(seq_lens.tolist()    )						, dims=0    )
												x += self.position_encoding(snake_case__    )

												__a:						Any		    =					self.dropout_pre(snake_case__    )

												# inverted the attention mask
												__a:						Any		    =					encoder_inputs.size()
												__a:						str		    =					self.get_extended_attention_mask(snake_case__						, snake_case__    )

												for lyr in self.encoders:
																	__a:						Any		    =					lyr(snake_case__						, snake_case__    )[0]
												__a:						int		    =					self.layer_norm(snake_case__    )

												return self.dropout_post(snake_case__    ), encoder_inputs_mask
