






"""simple docstring"""






from typing import Optional

import numpy as np
import torch
from torch import nn
from transformers import GPTaConfig, GPTaLMHeadModel
from transformers.modeling_utils import ModuleUtilsMixin

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





class       snake_case_       (    _lowerCamelCase    ,							_lowerCamelCase    ,							_lowerCamelCase       ):



	"""simple docstring"""

	SCREAMING_SNAKE_CASE_:							int   			=					[R"""h\.\d+\.attn\.bias""", R"""h\.\d+\.attn\.masked_bias"""]


	@register_to_config
	def __init__(      self  ,					__a  ,					__a  ,					__a = None  ,					__a = 5_0257  ,					__a = 1024  ,					__a = 768  ,					__a = 12  ,					__a = 12  ,					__a = None  ,					__a = "gelu_new"  ,					__a = 0.1  ,					__a = 0.1  ,					__a = 0.1  ,					__a = 1E-5  ,					__a = 0.02  ,					__a = True  ,					__a = True  ,					__a = False  ,					__a = False  ,					):


					"""simple docstring"""

					super().__init__()

					A__							 =       prefix_length

					if prefix_inner_dim != n_embd and prefix_hidden_dim is None:
									raise ValueError(
									    f'''`prefix_hidden_dim` cannot be `None` when `prefix_inner_dim`: {prefix_hidden_dim} and'''
									    f''' `n_embd`: {n_embd} are not equal.'''						)

					A__							 =       prefix_inner_dim
					A__							 =       prefix_hidden_dim

					A__							 =       (
					    nn.Linear(self.prefix_inner_dim  ,					self.prefix_hidden_dim						)
					    if self.prefix_hidden_dim is not None
					    else nn.Identity()
					)
					A__							 =       (
					    nn.Linear(self.prefix_hidden_dim  ,					UpperCamelCase_						) if self.prefix_hidden_dim is not None else nn.Identity()
					)

					A__							 =       GPTaConfig(
					    vocab_size=UpperCamelCase_  ,					n_positions=UpperCamelCase_  ,					n_embd=UpperCamelCase_  ,					n_layer=UpperCamelCase_  ,					n_head=UpperCamelCase_  ,					n_inner=UpperCamelCase_  ,					activation_function=UpperCamelCase_  ,					resid_pdrop=UpperCamelCase_  ,					embd_pdrop=UpperCamelCase_  ,					attn_pdrop=UpperCamelCase_  ,					layer_norm_epsilon=UpperCamelCase_  ,					initializer_range=UpperCamelCase_  ,					scale_attn_weights=UpperCamelCase_  ,					use_cache=UpperCamelCase_  ,					scale_attn_by_inverse_layer_idx=UpperCamelCase_  ,					reorder_and_upcast_attn=UpperCamelCase_  ,					)
					A__							 =       GPTaLMHeadModel(UpperCamelCase_						)


	def       _UpperCAmelCase						(      self  ,					__a  ,					__a  ,					__a = None  ,					__a = None  ,					):


					"""simple docstring"""

					A__							 =       self.transformer.transformer.wte(UpperCamelCase_						)
					A__							 =       self.encode_prefix(UpperCamelCase_						)
					A__							 =       self.decode_prefix(UpperCamelCase_						)
					A__							 =       torch.cat((prefix_embeds, embedding_text)  ,					dim=1						)

					if labels is not None:
									A__							 =       self.get_dummy_token(input_ids.shape[0]  ,					input_ids.device						)
									A__							 =       torch.cat((dummy_token, input_ids)  ,					dim=1						)
					A__							 =       self.transformer(inputs_embeds=UpperCamelCase_  ,					labels=UpperCamelCase_  ,					attention_mask=UpperCamelCase_						)
					if self.prefix_hidden_dim is not None:
									return out, hidden
					else:
									return out


	def       _UpperCAmelCase						(      self  ,					__a  ,					__a						):


					"""simple docstring"""

					return torch.zeros(UpperCamelCase_  ,					self.prefix_length  ,					dtype=torch.intaa  ,					device=UpperCamelCase_						)


	def       _UpperCAmelCase						(      self  ,					__a						):


					"""simple docstring"""

					return self.encode_prefix(UpperCamelCase_						)


	@torch.no_grad()
	def       _UpperCAmelCase						(      self  ,					__a  ,					__a  ,					__a						):


					"""simple docstring"""

					A__							 =       torch.split(UpperCamelCase_  ,					1  ,					dim=0						)
					A__							 =       []
					A__							 =       []
					for feature in features:
									A__							 =       self.decode_prefix(feature.to(UpperCamelCase_						)						)  # back to the clip feature
									# Only support beam search for now
									A__							 =       self.generate_beam(
									    input_embeds=UpperCamelCase_  ,					device=UpperCamelCase_  ,					eos_token_id=UpperCamelCase_						)
									generated_tokens.append(output_tokens[0]						)
									generated_seq_lengths.append(seq_lengths[0]						)
					A__							 =       torch.stack(UpperCamelCase_						)
					A__							 =       torch.stack(UpperCamelCase_						)
					return generated_tokens, generated_seq_lengths






	@torch.no_grad()
	def       _UpperCAmelCase						(      self  ,					__a=None  ,					__a=None  ,					__a=None  ,					__a = 5  ,					__a = 67  ,					__a = 1.0  ,					__a = None  ,					):


					"""simple docstring"""

					A__							 =       eos_token_id
					A__							 =       None
					A__							 =       None
					A__							 =       torch.ones(UpperCamelCase_  ,					device=UpperCamelCase_  ,					dtype=torch.int						)
					A__							 =       torch.zeros(UpperCamelCase_  ,					device=UpperCamelCase_  ,					dtype=torch.bool						)

					if input_embeds is not None:
									A__							 =       input_embeds
					else:
									A__							 =       self.transformer.transformer.wte(UpperCamelCase_						)

					for i in range(UpperCamelCase_						):
									A__							 =       self.transformer(inputs_embeds=UpperCamelCase_						)
									A__							 =       outputs.logits
									A__							 =       logits[:, -1, :] / (temperature if temperature > 0 else 1.0)
									A__							 =       logits.softmax(-1						).log()

									if scores is None:
													A__							 =       logits.topk(UpperCamelCase_  ,					-1						)
													A__							 =       generated.expand(UpperCamelCase_  ,					*generated.shape[1:]						)
													A__							 =       next_tokens.permute(1  ,					0						), scores.squeeze(0						)
													if tokens is None:
																	A__							 =       next_tokens
													else:
																	A__							 =       tokens.expand(UpperCamelCase_  ,					*tokens.shape[1:]						)
																	A__							 =       torch.cat((tokens, next_tokens)  ,					dim=1						)
									else:
													A__							 =       -float(np.inf						)
													A__							 =       0
													A__							 =       scores[:, None] + logits
													seq_lengths[~is_stopped] += 1
													A__							 =       scores_sum / seq_lengths[:, None]
													A__							 =       scores_sum_average.view(-1						).topk(UpperCamelCase_  ,					-1						)
													A__							 =       next_tokens // scores_sum.shape[1]
													A__							 =       seq_lengths[next_tokens_source]
													A__							 =       next_tokens % scores_sum.shape[1]
													A__							 =       next_tokens.unsqueeze(1						)
													A__							 =       tokens[next_tokens_source]
													A__							 =       torch.cat((tokens, next_tokens)  ,					dim=1						)
													A__							 =       generated[next_tokens_source]
													A__							 =       scores_sum_average * seq_lengths
													A__							 =       is_stopped[next_tokens_source]

									A__							 =       self.transformer.transformer.wte(next_tokens.squeeze()						).view(generated.shape[0]  ,					1  ,					-1						)
									A__							 =       torch.cat((generated, next_token_embed)  ,					dim=1						)
									A__							 =       is_stopped + next_tokens.eq(UpperCamelCase_						).squeeze()
									if is_stopped.all():
													break

					A__							 =       scores / seq_lengths
					A__							 =       scores.argsort(descending=UpperCamelCase_						)
					# tokens tensors are already padded to max_seq_length
					A__							 =       [tokens[i] for i in order]
					A__							 =       torch.stack(UpperCamelCase_  ,					dim=0						)
					A__							 =       torch.tensor([seq_lengths[i] for i in order]  ,					dtype=seq_lengths.dtype						)
					return output_texts, seq_lengths











"""simple docstring"""






import gc
import random
import tempfile
import unittest

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

from diffusers import (
    AutoencoderKL,
    DDIMInverseScheduler,
    DDIMScheduler,
    DPMSolverMultistepInverseScheduler,
    DPMSolverMultistepScheduler,
    StableDiffusionDiffEditPipeline,
    UNetaDConditionModel,
)
from diffusers.utils import load_image, slow
from diffusers.utils.testing_utils import enable_full_determinism, floats_tensor, require_torch_gpu, torch_device

from ..pipeline_params import TEXT_GUIDED_IMAGE_INPAINTING_BATCH_PARAMS, TEXT_GUIDED_IMAGE_INPAINTING_PARAMS
from ..test_pipelines_common import PipelineLatentTesterMixin, PipelineTesterMixin


enable_full_determinism()





class       snake_case_       (    _lowerCamelCase    ,							_lowerCamelCase    ,							unittest.TestCase       ):



	"""simple docstring"""

	SCREAMING_SNAKE_CASE_:							Union[str, Any]   			=					StableDiffusionDiffEditPipeline
	SCREAMING_SNAKE_CASE_:							Union[str, Any]   			=					TEXT_GUIDED_IMAGE_INPAINTING_PARAMS - {"""height""", """width""", """image"""} | {"""image_latents"""}
	SCREAMING_SNAKE_CASE_:							int   			=					TEXT_GUIDED_IMAGE_INPAINTING_BATCH_PARAMS - {"""image"""} | {"""image_latents"""}
	SCREAMING_SNAKE_CASE_:							str   			=					frozenset(
	    []       )  # TO-DO: update image_params once pipeline is refactored with VaeImageProcessor.preprocess
	SCREAMING_SNAKE_CASE_:							int   			=					frozenset([]       )


	def       _UpperCAmelCase						(      self						):


					"""simple docstring"""

					torch.manual_seed(0						)
					A__							 =       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  ,					attention_head_dim=(2, 4)  ,					use_linear_projection=__a  ,					)
					A__							 =       DDIMScheduler(
					    beta_start=0.0_0085  ,					beta_end=0.012  ,					beta_schedule='scaled_linear'  ,					clip_sample=__a  ,					set_alpha_to_one=__a  ,					)
					A__							 =       DDIMInverseScheduler(
					    beta_start=0.0_0085  ,					beta_end=0.012  ,					beta_schedule='scaled_linear'  ,					clip_sample=__a  ,					set_alpha_to_zero=__a  ,					)
					torch.manual_seed(0						)
					A__							 =       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  ,					sample_size=128  ,					)
					torch.manual_seed(0						)
					A__							 =       CLIPTextConfig(
					    bos_token_id=0  ,					eos_token_id=2  ,					hidden_size=32  ,					intermediate_size=37  ,					layer_norm_eps=1E-05  ,					num_attention_heads=4  ,					num_hidden_layers=5  ,					pad_token_id=1  ,					vocab_size=1000  ,					hidden_act='gelu'  ,					projection_dim=512  ,					)
					A__							 =       CLIPTextModel(__a						)
					A__							 =       CLIPTokenizer.from_pretrained('hf-internal-testing/tiny-random-clip'						)

					A__							 =       {
					    'unet': unet,
					    'scheduler': scheduler,
					    'inverse_scheduler': inverse_scheduler,
					    'vae': vae,
					    'text_encoder': text_encoder,
					    'tokenizer': tokenizer,
					    'safety_checker': None,
					    'feature_extractor': None,
					}

					return components


	def       _UpperCAmelCase						(      self  ,					__a  ,					__a=0						):


					"""simple docstring"""

					A__							 =       floats_tensor((1, 16, 16)  ,					rng=random.Random(__a						)						).to(__a						)
					A__							 =       floats_tensor((1, 2, 4, 16, 16)  ,					rng=random.Random(__a						)						).to(__a						)
					if str(__a						).startswith('mps'						):
									A__							 =       torch.manual_seed(__a						)
					else:
									A__							 =       torch.Generator(device=__a						).manual_seed(__a						)
					A__							 =       {
					    'prompt': 'a dog and a newt',
					    'mask_image': mask,
					    'image_latents': latents,
					    'generator': generator,
					    'num_inference_steps': 2,
					    'inpaint_strength': 1.0,
					    'guidance_scale': 6.0,
					    'output_type': 'numpy',
					}

					return inputs


	def       _UpperCAmelCase						(      self  ,					__a  ,					__a=0						):


					"""simple docstring"""

					A__							 =       floats_tensor((1, 3, 32, 32)  ,					rng=random.Random(__a						)						).to(__a						)
					A__							 =       image.cpu().permute(0  ,					2  ,					3  ,					1						)[0]
					A__							 =       Image.fromarray(np.uinta(__a						)						).convert('RGB'						)
					if str(__a						).startswith('mps'						):
									A__							 =       torch.manual_seed(__a						)
					else:
									A__							 =       torch.Generator(device=__a						).manual_seed(__a						)
					A__							 =       {
					    'image': image,
					    'source_prompt': 'a cat and a frog',
					    'target_prompt': 'a dog and a newt',
					    'generator': generator,
					    'num_inference_steps': 2,
					    'num_maps_per_mask': 2,
					    'mask_encode_strength': 1.0,
					    'guidance_scale': 6.0,
					    'output_type': 'numpy',
					}

					return inputs


	def       _UpperCAmelCase						(      self  ,					__a  ,					__a=0						):


					"""simple docstring"""

					A__							 =       floats_tensor((1, 3, 32, 32)  ,					rng=random.Random(__a						)						).to(__a						)
					A__							 =       image.cpu().permute(0  ,					2  ,					3  ,					1						)[0]
					A__							 =       Image.fromarray(np.uinta(__a						)						).convert('RGB'						)
					if str(__a						).startswith('mps'						):
									A__							 =       torch.manual_seed(__a						)
					else:
									A__							 =       torch.Generator(device=__a						).manual_seed(__a						)
					A__							 =       {
					    'image': image,
					    'prompt': 'a cat and a frog',
					    'generator': generator,
					    'num_inference_steps': 2,
					    'inpaint_strength': 1.0,
					    'guidance_scale': 6.0,
					    'decode_latents': True,
					    'output_type': 'numpy',
					}
					return inputs


	def       _UpperCAmelCase						(      self						):


					"""simple docstring"""

					if not hasattr(self.pipeline_class  ,					'_optional_components'						):
									return

					A__							 =       self.get_dummy_components()
					A__							 =       self.pipeline_class(**__a						)
					pipe.to(__a						)
					pipe.set_progress_bar_config(disable=__a						)

					# set all optional components to None and update pipeline config accordingly
					for optional_component in pipe._optional_components:
									setattr(__a  ,					__a  ,					__a						)
					pipe.register_modules(**{optional_component: None for optional_component in pipe._optional_components}						)

					A__							 =       self.get_dummy_inputs(__a						)
					A__							 =       pipe(**__a						)[0]

					with tempfile.TemporaryDirectory() as tmpdir:
									pipe.save_pretrained(__a						)
									A__							 =       self.pipeline_class.from_pretrained(__a						)
									pipe_loaded.to(__a						)
									pipe_loaded.set_progress_bar_config(disable=__a						)

					for optional_component in pipe._optional_components:
									self.assertTrue(
									    getattr(__a  ,					__a						) is None  ,					f'''`{optional_component}` did not stay set to None after loading.'''  ,					)

					A__							 =       self.get_dummy_inputs(__a						)
					A__							 =       pipe_loaded(**__a						)[0]

					A__							 =       np.abs(output - output_loaded						).max()
					self.assertLess(__a  ,					1E-4						)


	def       _UpperCAmelCase						(      self						):


					"""simple docstring"""

					A__							 =       'cpu'

					A__							 =       self.get_dummy_components()
					A__							 =       self.pipeline_class(**__a						)
					pipe.to(__a						)
					pipe.set_progress_bar_config(disable=__a						)

					A__							 =       self.get_dummy_mask_inputs(__a						)
					A__							 =       pipe.generate_mask(**__a						)
					A__							 =       mask[0, -3:, -3:]

					self.assertEqual(mask.shape  ,					(1, 16, 16)						)
					A__							 =       np.array([0] * 9						)
					A__							 =       np.abs(mask_slice.flatten() - expected_slice						).max()
					self.assertLessEqual(__a  ,					1E-3						)
					self.assertEqual(mask[0, -3, -4]  ,					0						)


	def       _UpperCAmelCase						(      self						):


					"""simple docstring"""

					A__							 =       'cpu'

					A__							 =       self.get_dummy_components()
					A__							 =       self.pipeline_class(**__a						)
					pipe.to(__a						)
					pipe.set_progress_bar_config(disable=__a						)

					A__							 =       self.get_dummy_inversion_inputs(__a						)
					A__							 =       pipe.invert(**__a						).images
					A__							 =       image[0, -1, -3:, -3:]

					self.assertEqual(image.shape  ,					(2, 32, 32, 3)						)
					A__							 =       np.array(
					    [0.5150, 0.5134, 0.5043, 0.5376, 0.4694, 0.5_1050, 0.5015, 0.4407, 0.4799]  ,					)
					A__							 =       np.abs(image_slice.flatten() - expected_slice						).max()
					self.assertLessEqual(__a  ,					1E-3						)


	def       _UpperCAmelCase						(      self						):


					"""simple docstring"""

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






	def       _UpperCAmelCase						(      self						):


					"""simple docstring"""

					A__							 =       'cpu'

					A__							 =       self.get_dummy_components()

					A__							 =       {'beta_start': 0.0_0085, 'beta_end': 0.012, 'beta_schedule': 'scaled_linear'}
					A__							 =       DPMSolverMultistepScheduler(**__a						)
					A__							 =       DPMSolverMultistepInverseScheduler(**__a						)

					A__							 =       self.pipeline_class(**__a						)
					pipe.to(__a						)
					pipe.set_progress_bar_config(disable=__a						)

					A__							 =       self.get_dummy_inversion_inputs(__a						)
					A__							 =       pipe.invert(**__a						).images
					A__							 =       image[0, -1, -3:, -3:]

					self.assertEqual(image.shape  ,					(2, 32, 32, 3)						)
					A__							 =       np.array(
					    [0.5150, 0.5134, 0.5043, 0.5376, 0.4694, 0.5_1050, 0.5015, 0.4407, 0.4799]  ,					)
					A__							 =       np.abs(image_slice.flatten() - expected_slice						).max()
					self.assertLessEqual(__a  ,					1E-3						)





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



	"""simple docstring"""

	def       _UpperCAmelCase						(      self						):


					"""simple docstring"""

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


	@classmethod
	def       _UpperCAmelCase						(      cls						):


					"""simple docstring"""

					A__							 =       load_image(
					    'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/diffedit/fruit.png'						)

					A__							 =       raw_image.convert('RGB'						).resize((768, 768)						)

					A__							 =       raw_image


	def       _UpperCAmelCase						(      self						):


					"""simple docstring"""

					A__							 =       torch.manual_seed(0						)

					A__							 =       StableDiffusionDiffEditPipeline.from_pretrained(
					    'stabilityai/stable-diffusion-2-1'  ,					safety_checker=__a  ,					torch_dtype=torch.floataa						)
					A__							 =       DDIMScheduler.from_config(pipe.scheduler.config						)
					A__							 =       DDIMInverseScheduler.from_config(pipe.scheduler.config						)
					pipe.enable_model_cpu_offload()
					pipe.set_progress_bar_config(disable=__a						)

					A__							 =       'a bowl of fruit'
					A__							 =       'a bowl of pears'

					A__							 =       pipe.generate_mask(
					    image=self.raw_image  ,					source_prompt=__a  ,					target_prompt=__a  ,					generator=__a  ,					)

					A__							 =       pipe.invert(
					    prompt=__a  ,					image=self.raw_image  ,					inpaint_strength=0.7  ,					generator=__a						).latents

					A__							 =       pipe(
					    prompt=__a  ,					mask_image=__a  ,					image_latents=__a  ,					generator=__a  ,					negative_prompt=__a  ,					inpaint_strength=0.7  ,					output_type='numpy'  ,					).images[0]

					A__							 =       (
					    np.array(
					        load_image(
					            'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'
					            '/diffedit/pears.png'						).resize((768, 768)						)						)
					    / 255
					)
					assert np.abs((expected_image - image).max()						) < 5E-1






	def       _UpperCAmelCase						(      self						):


					"""simple docstring"""

					A__							 =       torch.manual_seed(0						)

					A__							 =       StableDiffusionDiffEditPipeline.from_pretrained(
					    'stabilityai/stable-diffusion-2-1'  ,					safety_checker=__a  ,					torch_dtype=torch.floataa						)
					A__							 =       DPMSolverMultistepScheduler.from_config(pipe.scheduler.config						)
					A__							 =       DPMSolverMultistepInverseScheduler.from_config(pipe.scheduler.config						)
					pipe.enable_model_cpu_offload()
					pipe.set_progress_bar_config(disable=__a						)

					A__							 =       'a bowl of fruit'
					A__							 =       'a bowl of pears'

					A__							 =       pipe.generate_mask(
					    image=self.raw_image  ,					source_prompt=__a  ,					target_prompt=__a  ,					generator=__a  ,					)

					A__							 =       pipe.invert(
					    prompt=__a  ,					image=self.raw_image  ,					inpaint_strength=0.7  ,					generator=__a  ,					num_inference_steps=25  ,					).latents

					A__							 =       pipe(
					    prompt=__a  ,					mask_image=__a  ,					image_latents=__a  ,					generator=__a  ,					negative_prompt=__a  ,					inpaint_strength=0.7  ,					num_inference_steps=25  ,					output_type='numpy'  ,					).images[0]

					A__							 =       (
					    np.array(
					        load_image(
					            'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'
					            '/diffedit/pears.png'						).resize((768, 768)						)						)
					    / 255
					)
					assert np.abs((expected_image - image).max()						) < 5E-1


