

"""simple docstring"""

import collections
import types

import numpy as np

from ..utils import (
    add_end_docstrings,
    is_tensorflow_probability_available,
    is_tf_available,
    is_torch_available,
    requires_backends,
)
from .base import PIPELINE_INIT_ARGS, ArgumentHandler, Dataset, Pipeline, PipelineException


if is_torch_available():
							import torch

							from ..models.auto.modeling_auto import (
							    MODEL_FOR_SEQ_TO_SEQ_CAUSAL_LM_MAPPING,
							    MODEL_FOR_TABLE_QUESTION_ANSWERING_MAPPING,
							)

if is_tf_available() and is_tensorflow_probability_available():
							import tensorflow as tf
							import tensorflow_probability as tfp

							from ..models.auto.modeling_tf_auto import (
							    TF_MODEL_FOR_SEQ_TO_SEQ_CAUSAL_LM_MAPPING,
							    TF_MODEL_FOR_TABLE_QUESTION_ANSWERING_MAPPING,
							)


class 			_snake_case     (						UpperCamelCase       ):







						def __call__(self     ,							UpperCAmelCase_=None     ,							UpperCAmelCase_=None     ,							**UpperCAmelCase_							)			->						Optional[int]:
											# Returns tqa_pipeline_inputs of shape:
											# [
											#   {"table": pd.DataFrame, "query": List[str]},
											#   ...,
											#   {"table": pd.DataFrame, "query" : List[str]}
											# ]
											requires_backends(self     ,							'''pandas'''							)
											import pandas as pd

											if table is None:
																raise ValueError('''Keyword argument `table` cannot be None.'''							)
											elif query is None:
																if isinstance(UpperCAmelCase_     ,							UpperCAmelCase_							) and table.get('''query'''							) is not None and table.get('''table'''							) is not None:
																					__lowercase												=			[table]
																elif isinstance(UpperCAmelCase_     ,							UpperCAmelCase_							) and len(UpperCAmelCase_							) > 0:
																					if not all(isinstance(UpperCAmelCase_     ,							UpperCAmelCase_							) for d in table							):
																										raise ValueError(
																										    f'''Keyword argument `table` should be a list of dict, but is {(type(UpperCAmelCase_							) for d in table)}'''							)

																					if table[0].get('''query'''							) is not None and table[0].get('''table'''							) is not None:
																										__lowercase												=			table
																					else:
																										raise ValueError(
																										    '''If keyword argument `table` is a list of dictionaries, each dictionary should have a `table`'''
																										    f''' and `query` key, but only dictionary has keys {table[0].keys()} `table` and `query` keys.'''							)
																elif Dataset is not None and isinstance(UpperCAmelCase_     ,							UpperCAmelCase_							) or isinstance(UpperCAmelCase_     ,							types.GeneratorType							):
																					return table
																else:
																					raise ValueError(
																					    '''Invalid input. Keyword argument `table` should be either of type `dict` or `list`, but '''
																					    f'''is {type(UpperCAmelCase_							)})'''							)
											else:
																__lowercase												=			[{'''table''': table, '''query''': query}]

											for tqa_pipeline_input in tqa_pipeline_inputs:
																if not isinstance(tqa_pipeline_input['''table''']     ,							pd.DataFrame							):
																					if tqa_pipeline_input["table"] is None:
																										raise ValueError('''Table cannot be None.'''							)

																					__lowercase												=			pd.DataFrame(tqa_pipeline_input['''table''']							)

											return tqa_pipeline_inputs



@add_end_docstrings(UpperCamelCase       )
class 			_snake_case     (						UpperCamelCase       ):
						__snake_case 			=     'table,query'







						def __init__(self     ,							UpperCAmelCase_=TableQuestionAnsweringArgumentHandler()     ,							*UpperCAmelCase_     ,							**UpperCAmelCase_							)			->						Tuple:
											super().__init__(*UpperCAmelCase_     ,							**UpperCAmelCase_							)
											__lowercase												=			args_parser

											self.check_model_type(
											    dict(
											        TF_MODEL_FOR_TABLE_QUESTION_ANSWERING_MAPPING.items()
											        + TF_MODEL_FOR_SEQ_TO_SEQ_CAUSAL_LM_MAPPING.items()							)
											    if self.framework == '''tf'''
											    else dict(
											        MODEL_FOR_TABLE_QUESTION_ANSWERING_MAPPING.items() + MODEL_FOR_SEQ_TO_SEQ_CAUSAL_LM_MAPPING.items()							)							)

											__lowercase												=			bool(getattr(self.model.config     ,							'''aggregation_labels'''     ,							UpperCAmelCase_							)							) and bool(
											    getattr(self.model.config     ,							'''num_aggregation_labels'''     ,							UpperCAmelCase_							)							)
											__lowercase												=			'''tapas''' if hasattr(self.model.config     ,							'''aggregation_labels'''							) else None







						def        _UpperCamelCase    (self     ,							**UpperCAmelCase_							)			->						Optional[Any]:
											return self.model(**UpperCAmelCase_							)







						def        _UpperCamelCase    (self     ,							**UpperCAmelCase_							)			->						Any:
											if self.framework == "pt":
																__lowercase												=			[]
																__lowercase												=			[]
																__lowercase												=			None
																__lowercase												=			inputs['''input_ids'''].shape[0]

																__lowercase												=			inputs['''input_ids'''].to(self.device							)
																__lowercase												=			inputs['''attention_mask'''].to(self.device							)
																__lowercase												=			inputs['''token_type_ids'''].to(self.device							)
																__lowercase												=			None

																for index in range(UpperCAmelCase_							):
																					# If sequences have already been processed, the token type IDs will be created according to the previous
																					# answer.
																					if prev_answers is not None:
																										__lowercase												=			token_type_ids_example[:, 3]  # shape (seq_len,)
																										__lowercase												=			np.zeros_like(prev_labels_example.cpu().numpy()							)  # shape (seq_len,)

																										__lowercase												=			token_type_ids[index]  # shape (seq_len, 7)
																										for i in range(model_labels.shape[0]							):
																															__lowercase												=			token_type_ids_example[:, 0].tolist()[i]
																															__lowercase												=			token_type_ids_example[:, 1].tolist()[i] - 1
																															__lowercase												=			token_type_ids_example[:, 2].tolist()[i] - 1

																															if row_id >= 0 and col_id >= 0 and segment_id == 1:
																																				__lowercase												=			int(prev_answers[(col_id, row_id)]							)

																										__lowercase												=			torch.from_numpy(UpperCAmelCase_							).type(torch.long							).to(self.device							)

																					__lowercase												=			input_ids[index]
																					__lowercase												=			attention_mask[index]  # shape (seq_len,)
																					__lowercase												=			token_type_ids[index]  # shape (seq_len, 7)
																					__lowercase												=			self.model(
																					    input_ids=input_ids_example.unsqueeze(0							)     ,							attention_mask=attention_mask_example.unsqueeze(0							)     ,							token_type_ids=token_type_ids_example.unsqueeze(0							)     ,							)
																					__lowercase												=			outputs.logits

																					if self.aggregate:
																										all_aggregations.append(outputs.logits_aggregation							)

																					all_logits.append(UpperCAmelCase_							)

																					__lowercase												=			torch.distributions.Bernoulli(logits=UpperCAmelCase_							)
																					__lowercase												=			dist_per_token.probs * attention_mask_example.type(torch.floataa							).to(
																					    dist_per_token.probs.device							)

																					__lowercase												=			collections.defaultdict(UpperCAmelCase_							)
																					for i, p in enumerate(probabilities.squeeze().tolist()							):
																										__lowercase												=			token_type_ids_example[:, 0].tolist()[i]
																										__lowercase												=			token_type_ids_example[:, 1].tolist()[i] - 1
																										__lowercase												=			token_type_ids_example[:, 2].tolist()[i] - 1
																										if col >= 0 and row >= 0 and segment_id == 1:
																															coords_to_probs[(col, row)].append(UpperCAmelCase_							)

																					__lowercase												=			{key: np.array(coords_to_probs[key]							).mean() > 0.5 for key in coords_to_probs}

																__lowercase												=			torch.cat(tuple(UpperCAmelCase_							)     ,							0							)

																return (logits_batch,) if not self.aggregate else (logits_batch, torch.cat(tuple(UpperCAmelCase_							)     ,							0							))
											else:
																__lowercase												=			[]
																__lowercase												=			[]
																__lowercase												=			None
																__lowercase												=			inputs['''input_ids'''].shape[0]

																__lowercase												=			inputs['''input_ids''']
																__lowercase												=			inputs['''attention_mask''']
																__lowercase												=			inputs['''token_type_ids'''].numpy()
																__lowercase												=			None

																for index in range(UpperCAmelCase_							):
																					# If sequences have already been processed, the token type IDs will be created according to the previous
																					# answer.
																					if prev_answers is not None:
																										__lowercase												=			token_type_ids_example[:, 3]  # shape (seq_len,)
																										__lowercase												=			np.zeros_like(UpperCAmelCase_     ,							dtype=np.intaa							)  # shape (seq_len,)

																										__lowercase												=			token_type_ids[index]  # shape (seq_len, 7)
																										for i in range(model_labels.shape[0]							):
																															__lowercase												=			token_type_ids_example[:, 0].tolist()[i]
																															__lowercase												=			token_type_ids_example[:, 1].tolist()[i] - 1
																															__lowercase												=			token_type_ids_example[:, 2].tolist()[i] - 1

																															if row_id >= 0 and col_id >= 0 and segment_id == 1:
																																				__lowercase												=			int(prev_answers[(col_id, row_id)]							)

																										__lowercase												=			model_labels

																					__lowercase												=			input_ids[index]
																					__lowercase												=			attention_mask[index]  # shape (seq_len,)
																					__lowercase												=			token_type_ids[index]  # shape (seq_len, 7)
																					__lowercase												=			self.model(
																					    input_ids=np.expand_dims(UpperCAmelCase_     ,							axis=0							)     ,							attention_mask=np.expand_dims(UpperCAmelCase_     ,							axis=0							)     ,							token_type_ids=np.expand_dims(UpperCAmelCase_     ,							axis=0							)     ,							)
																					__lowercase												=			outputs.logits

																					if self.aggregate:
																										all_aggregations.append(outputs.logits_aggregation							)

																					all_logits.append(UpperCAmelCase_							)

																					__lowercase												=			tfp.distributions.Bernoulli(logits=UpperCAmelCase_							)
																					__lowercase												=			dist_per_token.probs_parameter() * tf.cast(UpperCAmelCase_     ,							tf.floataa							)

																					__lowercase												=			collections.defaultdict(UpperCAmelCase_							)
																					__lowercase												=			token_type_ids_example
																					for i, p in enumerate(tf.squeeze(UpperCAmelCase_							).numpy().tolist()							):
																										__lowercase												=			token_type_ids_example[:, 0].tolist()[i]
																										__lowercase												=			token_type_ids_example[:, 1].tolist()[i] - 1
																										__lowercase												=			token_type_ids_example[:, 2].tolist()[i] - 1
																										if col >= 0 and row >= 0 and segment_id == 1:
																															coords_to_probs[(col, row)].append(UpperCAmelCase_							)

																					__lowercase												=			{key: np.array(coords_to_probs[key]							).mean() > 0.5 for key in coords_to_probs}

																__lowercase												=			tf.concat(tuple(UpperCAmelCase_							)     ,							0							)

																return (logits_batch,) if not self.aggregate else (logits_batch, tf.concat(tuple(UpperCAmelCase_							)     ,							0							))







						def __call__(self     ,							*UpperCAmelCase_     ,							**UpperCAmelCase_							)			->						Tuple:
											__lowercase												=			self._args_parser(*UpperCAmelCase_     ,							**UpperCAmelCase_							)

											__lowercase												=			super().__call__(UpperCAmelCase_     ,							**UpperCAmelCase_							)
											if len(UpperCAmelCase_							) == 1:
																return results[0]
											return results







						def        _UpperCamelCase    (self     ,							UpperCAmelCase_=None     ,							UpperCAmelCase_=None     ,							UpperCAmelCase_=None     ,							**UpperCAmelCase_							)			->						List[Any]:
											__lowercase												=			{}
											if padding is not None:
																__lowercase												=			padding
											if truncation is not None:
																__lowercase												=			truncation

											__lowercase												=			{}
											if sequential is not None:
																__lowercase												=			sequential
											return preprocess_params, forward_params, {}







						def        _UpperCamelCase    (self     ,							UpperCAmelCase_     ,							UpperCAmelCase_=None     ,							UpperCAmelCase_=True     ,							UpperCAmelCase_=None							)			->						Any:
											if truncation is None:
																if self.type == "tapas":
																					__lowercase												=			'''drop_rows_to_fit'''
																else:
																					__lowercase												=			'''do_not_truncate'''

											__lowercase						, __lowercase												=			pipeline_input['''table'''], pipeline_input['''query''']
											if table.empty:
																raise ValueError('''table is empty'''							)
											if query is None or query == "":
																raise ValueError('''query is empty'''							)
											__lowercase												=			self.tokenizer(UpperCAmelCase_     ,							UpperCAmelCase_     ,							return_tensors=self.framework     ,							truncation=UpperCAmelCase_     ,							padding=UpperCAmelCase_							)
											__lowercase												=			table
											return inputs







						def        _UpperCamelCase    (self     ,							UpperCAmelCase_     ,							UpperCAmelCase_=False							)			->						int:
											__lowercase												=			model_inputs.pop('''table'''							)

											if self.type == "tapas":
																if sequential:
																					__lowercase												=			self.sequential_inference(**UpperCAmelCase_							)
																else:
																					__lowercase												=			self.batch_inference(**UpperCAmelCase_							)
											else:
																__lowercase												=			self.model.generate(**UpperCAmelCase_							)
											__lowercase												=			{'''model_inputs''': model_inputs, '''table''': table, '''outputs''': outputs}
											return model_outputs







						def        _UpperCamelCase    (self     ,							UpperCAmelCase_							)			->						List[str]:
											__lowercase												=			model_outputs['''model_inputs''']
											__lowercase												=			model_outputs['''table''']
											__lowercase												=			model_outputs['''outputs''']
											if self.type == "tapas":
																if self.aggregate:
																					__lowercase						, __lowercase												=			outputs[:2]
																					__lowercase												=			self.tokenizer.convert_logits_to_predictions(UpperCAmelCase_     ,							UpperCAmelCase_     ,							UpperCAmelCase_							)
																					__lowercase						, __lowercase												=			predictions
																					__lowercase												=			{i: self.model.config.aggregation_labels[pred] for i, pred in enumerate(UpperCAmelCase_							)}

																					__lowercase												=			self.model.config.no_aggregation_label_index
																					__lowercase												=			{
																					    i: aggregators[i] + ''' > ''' for i, pred in enumerate(UpperCAmelCase_							) if pred != no_agg_label_index
																					}
																else:
																					__lowercase												=			outputs[0]
																					__lowercase												=			self.tokenizer.convert_logits_to_predictions(UpperCAmelCase_     ,							UpperCAmelCase_							)
																					__lowercase												=			predictions[0]
																					__lowercase												=			{}
																					__lowercase												=			{}
																__lowercase												=			[]
																for index, coordinates in enumerate(UpperCAmelCase_							):
																					__lowercase												=			[table.iat[coordinate] for coordinate in coordinates]
																					__lowercase												=			aggregators.get(UpperCAmelCase_     ,							''''''							)
																					__lowercase												=			aggregators_prefix.get(UpperCAmelCase_     ,							''''''							)
																					__lowercase												=			{
																					    '''answer''': aggregator_prefix + ''', '''.join(UpperCAmelCase_							),
																					    '''coordinates''': coordinates,
																					    '''cells''': [table.iat[coordinate] for coordinate in coordinates],
																					}
																					if aggregator:
																										__lowercase												=			aggregator

																					answers.append(UpperCAmelCase_							)
																if len(UpperCAmelCase_							) == 0:
																					raise PipelineException('''Empty answer'''							)
											else:
																__lowercase												=			[{'''answer''': answer} for answer in self.tokenizer.batch_decode(UpperCAmelCase_     ,							skip_special_tokens=UpperCAmelCase_							)]

											return answers if len(UpperCAmelCase_							) > 1 else answers[0]





"""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     (						UpperCamelCase						,			UpperCamelCase						,			UpperCamelCase						,			unittest.TestCase       ):
						__snake_case 			=     StableDiffusionModelEditingPipeline
						__snake_case 			=     TEXT_TO_IMAGE_PARAMS
						__snake_case 			=     TEXT_TO_IMAGE_BATCH_PARAMS
						__snake_case 			=     TEXT_TO_IMAGE_IMAGE_PARAMS
						__snake_case 			=     TEXT_TO_IMAGE_IMAGE_PARAMS







						def        _UpperCamelCase    (self							)			->						int:
											torch.manual_seed(0							)
											__lowercase												=			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     ,							)
											__lowercase												=			DDIMScheduler()
											torch.manual_seed(0							)
											__lowercase												=			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							)
											__lowercase												=			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     ,							)
											__lowercase												=			CLIPTextModel(UpperCAmelCase_							)
											__lowercase												=			CLIPTokenizer.from_pretrained('''hf-internal-testing/tiny-random-clip'''							)

											__lowercase												=			{
											    '''unet''': unet,
											    '''scheduler''': scheduler,
											    '''vae''': vae,
											    '''text_encoder''': text_encoder,
											    '''tokenizer''': tokenizer,
											    '''safety_checker''': None,
											    '''feature_extractor''': None,
											}
											return components







						def        _UpperCamelCase    (self     ,							UpperCAmelCase_     ,							UpperCAmelCase_=0							)			->						List[str]:
											__lowercase												=			torch.manual_seed(UpperCAmelCase_							)
											__lowercase												=			{
											    '''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        _UpperCamelCase    (self							)			->						Optional[int]:
											__lowercase												=			'''cpu'''  # ensure determinism for the device-dependent torch.Generator
											__lowercase												=			self.get_dummy_components()
											__lowercase												=			StableDiffusionModelEditingPipeline(**UpperCAmelCase_							)
											__lowercase												=			sd_pipe.to(UpperCAmelCase_							)
											sd_pipe.set_progress_bar_config(disable=UpperCAmelCase_							)

											__lowercase												=			self.get_dummy_inputs(UpperCAmelCase_							)
											__lowercase												=			sd_pipe(**UpperCAmelCase_							).images
											__lowercase												=			image[0, -3:, -3:, -1]
											assert image.shape == (1, 6_4, 6_4, 3)

											__lowercase												=			np.array([0.4_7_5_5, 0.5_1_3_2, 0.4_9_7_6, 0.3_9_0_4, 0.3_5_5_4, 0.4_7_6_5, 0.5_1_3_9, 0.5_1_5_8, 0.4_8_8_9]							)

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







						def        _UpperCamelCase    (self							)			->						Tuple:
											__lowercase												=			'''cpu'''  # ensure determinism for the device-dependent torch.Generator
											__lowercase												=			self.get_dummy_components()
											__lowercase												=			StableDiffusionModelEditingPipeline(**UpperCAmelCase_							)
											__lowercase												=			sd_pipe.to(UpperCAmelCase_							)
											sd_pipe.set_progress_bar_config(disable=UpperCAmelCase_							)

											__lowercase												=			self.get_dummy_inputs(UpperCAmelCase_							)
											__lowercase												=			'''french fries'''
											__lowercase												=			sd_pipe(**UpperCAmelCase_     ,							negative_prompt=UpperCAmelCase_							)
											__lowercase												=			output.images
											__lowercase												=			image[0, -3:, -3:, -1]

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

											__lowercase												=			np.array([0.4_9_9_2, 0.5_1_0_1, 0.5_0_0_4, 0.3_9_4_9, 0.3_6_0_4, 0.4_7_3_5, 0.5_2_1_6, 0.5_2_0_4, 0.4_9_1_3]							)

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







						def        _UpperCamelCase    (self							)			->						Dict:
											__lowercase												=			'''cpu'''  # ensure determinism for the device-dependent torch.Generator
											__lowercase												=			self.get_dummy_components()
											__lowercase												=			EulerAncestralDiscreteScheduler(
											    beta_start=0.0_0_0_8_5     ,							beta_end=0.0_1_2     ,							beta_schedule='''scaled_linear'''							)
											__lowercase												=			StableDiffusionModelEditingPipeline(**UpperCAmelCase_							)
											__lowercase												=			sd_pipe.to(UpperCAmelCase_							)
											sd_pipe.set_progress_bar_config(disable=UpperCAmelCase_							)

											__lowercase												=			self.get_dummy_inputs(UpperCAmelCase_							)
											__lowercase												=			sd_pipe(**UpperCAmelCase_							).images
											__lowercase												=			image[0, -3:, -3:, -1]

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

											__lowercase												=			np.array([0.4_7_4_7, 0.5_3_7_2, 0.4_7_7_9, 0.4_9_8_2, 0.5_5_4_3, 0.4_8_1_6, 0.5_2_3_8, 0.4_9_0_4, 0.5_0_2_7]							)

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







						def        _UpperCamelCase    (self							)			->						str:
											__lowercase												=			'''cpu'''  # ensure determinism for the device-dependent torch.Generator
											__lowercase												=			self.get_dummy_components()
											__lowercase												=			PNDMScheduler()
											__lowercase												=			StableDiffusionModelEditingPipeline(**UpperCAmelCase_							)
											__lowercase												=			sd_pipe.to(UpperCAmelCase_							)
											sd_pipe.set_progress_bar_config(disable=UpperCAmelCase_							)

											__lowercase												=			self.get_dummy_inputs(UpperCAmelCase_							)
											# the pipeline does not expect pndm so test if it raises error.
											with self.assertRaises(UpperCAmelCase_							):
																__lowercase												=			sd_pipe(**UpperCAmelCase_							).images







						def        _UpperCamelCase    (self							)			->						Tuple:
											super().test_inference_batch_single_identical(expected_max_diff=5e-3							)







						def        _UpperCamelCase    (self							)			->						str:
											super().test_attention_slicing_forward_pass(expected_max_diff=5e-3							)



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







						def        _UpperCamelCase    (self							)			->						int:
											super().tearDown()
											gc.collect()
											torch.cuda.empty_cache()







						def        _UpperCamelCase    (self     ,							UpperCAmelCase_=0							)			->						int:
											__lowercase												=			torch.manual_seed(UpperCAmelCase_							)
											__lowercase												=			{
											    '''prompt''': '''A field of roses''',
											    '''generator''': generator,
											    '''num_inference_steps''': 3,
											    '''guidance_scale''': 7.5,
											    '''output_type''': '''numpy''',
											}
											return inputs







						def        _UpperCamelCase    (self							)			->						Optional[int]:
											__lowercase												=			'''CompVis/stable-diffusion-v1-4'''
											__lowercase												=			StableDiffusionModelEditingPipeline.from_pretrained(UpperCAmelCase_     ,							safety_checker=UpperCAmelCase_							)
											pipe.to(UpperCAmelCase_							)
											pipe.set_progress_bar_config(disable=UpperCAmelCase_							)
											pipe.enable_attention_slicing()

											__lowercase												=			self.get_inputs()
											__lowercase												=			pipe(**UpperCAmelCase_							).images
											__lowercase												=			image[0, -3:, -3:, -1].flatten()

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

											__lowercase												=			np.array(
											    [0.6_7_4_9_4_9_6, 0.6_3_8_6_4_5_3, 0.5_1_4_4_3_2_6_7, 0.6_6_0_9_4_9_0_5, 0.6_1_9_2_1_2_1_5, 0.5_4_9_1_3_3_2, 0.5_7_4_4_4_1_7, 0.5_8_0_7_5_1_0_6, 0.5_1_7_4_6_5_8]							)

											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'''							)

											__lowercase												=			pipe(**UpperCAmelCase_							).images
											__lowercase												=			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        _UpperCamelCase    (self							)			->						Optional[int]:
											torch.cuda.empty_cache()
											torch.cuda.reset_max_memory_allocated()
											torch.cuda.reset_peak_memory_stats()

											__lowercase												=			'''CompVis/stable-diffusion-v1-4'''
											__lowercase												=			DDIMScheduler.from_pretrained(UpperCAmelCase_     ,							subfolder='''scheduler'''							)
											__lowercase												=			StableDiffusionModelEditingPipeline.from_pretrained(
											    UpperCAmelCase_     ,							scheduler=UpperCAmelCase_     ,							safety_checker=UpperCAmelCase_							)
											__lowercase												=			pipe.to(UpperCAmelCase_							)
											pipe.set_progress_bar_config(disable=UpperCAmelCase_							)
											pipe.enable_attention_slicing(1							)
											pipe.enable_sequential_cpu_offload()

											__lowercase												=			self.get_inputs()
											__lowercase												=			pipe(**UpperCAmelCase_							)

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


