
'''simple docstring'''

from copy import deepcopy

import torch
import torch.nn.functional as F
from torch.optim import AdamW
from torch.optim.lr_scheduler import LambdaLR
from torch.utils.data import DataLoader

from accelerate.accelerator import Accelerator
from accelerate.state import GradientState
from accelerate.test_utils import RegressionDataset, RegressionModel
from accelerate.utils import DistributedType, is_torch_version, set_seed


def        UpperCAmelCase_ (  lowerCamelCase_      ,			lowerCamelCase_      ,			lowerCamelCase_      ,			lowerCamelCase_			):







 """simple docstring"""


 for param, grad_param in zip(model_a.parameters()      ,			model_b.parameters()			):
  if not param.requires_grad:
   continue
  if not did_step:
   # Grads should not be in sync
   assert (
       torch.allclose(param.grad      ,			grad_param.grad			) is False
   ), f'''Gradients in sync when they should not be at iteration {iteration}:\nmodel_a grad ({param.grad}) == model_b grad ({grad_param.grad})'''
  else:
   # Grads should be in sync
   assert (
       torch.allclose(param.grad      ,			grad_param.grad			) is True
   ), f'''Gradients not in sync when they should be at iteration {iteration}:\nmodel_a grad ({param.grad}) != model_b grad ({grad_param.grad})'''


def        UpperCAmelCase_ (  lowerCamelCase_      ,			lowerCamelCase_      ,			lowerCamelCase_      ,			lowerCamelCase_      ,			lowerCamelCase_=True			):







 """simple docstring"""


 model.train()
 lowerCAmelCase__     :   Dict     		=	model(lowerCamelCase_			)
 lowerCAmelCase__     :   Dict     		=	F.mse_loss(lowerCamelCase_      ,			target.to(output.device			)			)
 if not do_backward:
  loss /= accelerator.gradient_accumulation_steps
  loss.backward()
 else:
  accelerator.backward(lowerCamelCase_			)


def        UpperCAmelCase_ (  lowerCamelCase_      ,			lowerCamelCase_=False			):







 """simple docstring"""


 set_seed(4_2			)
 lowerCAmelCase__     :   Union[str, Any]     		=	RegressionModel()
 lowerCAmelCase__     :   List[str]     		=	deepcopy(lowerCamelCase_			)
 lowerCAmelCase__     :   Tuple     		=	RegressionDataset(length=8_0			)
 lowerCAmelCase__     :   Tuple     		=	DataLoader(lowerCamelCase_      ,			batch_size=1_6			)
 model.to(accelerator.device			)
 if sched:
  lowerCAmelCase__     :   Any     		=	AdamW(params=model.parameters()      ,			lr=1e-3			)
  lowerCAmelCase__     :   Union[str, Any]     		=	AdamW(params=ddp_model.parameters()      ,			lr=1e-3			)
  lowerCAmelCase__     :   Any     		=	LambdaLR(lowerCamelCase_      ,			lr_lambda=lambda lowerCamelCase_			: epoch**0.65			)
  lowerCAmelCase__     :   Tuple     		=	LambdaLR(lowerCamelCase_      ,			lr_lambda=lambda lowerCamelCase_			: epoch**0.65			)
 # Make a copy of `model`
 if sched:
  lowerCAmelCase__		,     lowerCAmelCase__		,     lowerCAmelCase__		,     lowerCAmelCase__     :   Tuple     		=	accelerator.prepare(lowerCamelCase_      ,			lowerCamelCase_      ,			lowerCamelCase_      ,			lowerCamelCase_			)
 else:
  lowerCAmelCase__		,     lowerCAmelCase__     :   int     		=	accelerator.prepare(lowerCamelCase_      ,			lowerCamelCase_			)
 if sched:
  return (model, opt, sched, dataloader, ddp_model, ddp_opt, ddp_sched)
 return model, ddp_model, dataloader


def        UpperCAmelCase_ (  lowerCamelCase_			):







 """simple docstring"""


 lowerCAmelCase__		,     lowerCAmelCase__		,     lowerCAmelCase__     :   Optional[Any]     		=	get_training_setup(lowerCamelCase_			)
 # Use a single batch
 lowerCAmelCase__		,     lowerCAmelCase__     :   int     		=	next(iter(lowerCamelCase_			)			).values()
 for iteration in range(3			):
  # Gather the distributed inputs and targs for the base model
  lowerCAmelCase__		,     lowerCAmelCase__     :   Optional[int]     		=	accelerator.gather((ddp_input, ddp_target)			)
  lowerCAmelCase__		,     lowerCAmelCase__     :   List[str]     		=	input.to(accelerator.device			), target.to(accelerator.device			)
  # Perform our initial ground truth step in non "DDP"
  step_model(lowerCamelCase_      ,			lowerCamelCase_      ,			lowerCamelCase_      ,			lowerCamelCase_			)
  # Do "gradient accumulation" (noop)
  if iteration % 2 == 0:
   # Accumulate grads locally
   with accelerator.no_sync(lowerCamelCase_			):
    step_model(lowerCamelCase_      ,			lowerCamelCase_      ,			lowerCamelCase_      ,			lowerCamelCase_			)
  else:
   # Sync grads
   step_model(lowerCamelCase_      ,			lowerCamelCase_      ,			lowerCamelCase_      ,			lowerCamelCase_			)

  # Since `no_sync` is a noop, `ddp_model` and `model` grads should always be in sync
  check_model_parameters(lowerCamelCase_      ,			lowerCamelCase_      ,			lowerCamelCase_      ,			lowerCamelCase_			)
  for param, ddp_param in zip(model.parameters()      ,			ddp_model.parameters()			):
   if not param.requires_grad:
    continue
   assert torch.allclose(
       param.grad      ,			ddp_param.grad			), f'''Gradients not in sync when they should be:\nModel grad ({param.grad}) != DDP grad ({ddp_param.grad})'''

  # Shuffle ddp_input on each iteration
  torch.manual_seed(1_3_3_7 + iteration			)
  lowerCAmelCase__     :   List[str]     		=	ddp_input[torch.randperm(len(lowerCamelCase_			)			)]


def        UpperCAmelCase_ (  lowerCamelCase_			):







 """simple docstring"""


 lowerCAmelCase__		,     lowerCAmelCase__		,     lowerCAmelCase__     :   List[str]     		=	get_training_setup(lowerCamelCase_			)
 # Use a single batch
 lowerCAmelCase__		,     lowerCAmelCase__     :   int     		=	next(iter(lowerCamelCase_			)			).values()
 for iteration in range(3			):
  # Gather the distributed inputs and targs for the base model
  lowerCAmelCase__		,     lowerCAmelCase__     :   Tuple     		=	accelerator.gather((ddp_input, ddp_target)			)
  lowerCAmelCase__		,     lowerCAmelCase__     :   List[str]     		=	input.to(accelerator.device			), target.to(accelerator.device			)
  # Perform our initial ground truth step in non "DDP"
  step_model(lowerCamelCase_      ,			lowerCamelCase_      ,			lowerCamelCase_      ,			lowerCamelCase_			)
  # Do "gradient accumulation" (noop)
  if iteration % 2 == 0:
   # Accumulate grads locally
   with accelerator.no_sync(lowerCamelCase_			):
    step_model(lowerCamelCase_      ,			lowerCamelCase_      ,			lowerCamelCase_      ,			lowerCamelCase_			)
  else:
   # Sync grads
   step_model(lowerCamelCase_      ,			lowerCamelCase_      ,			lowerCamelCase_      ,			lowerCamelCase_			)

  # DDP model and model should only be in sync when not (iteration % 2 == 0)
  for param, ddp_param in zip(model.parameters()      ,			ddp_model.parameters()			):
   if not param.requires_grad:
    continue
   if iteration % 2 == 0:
    # Grads should not be in sync
    assert (
        torch.allclose(param.grad      ,			ddp_param.grad			) is False
    ), f'''Gradients in sync when they should not be:\nModel grad ({param.grad}) == DDP grad ({ddp_param.grad})'''
   else:
    # Grads should be in sync
    assert (
        torch.allclose(param.grad      ,			ddp_param.grad			) is True
    ), f'''Gradients not in sync when they should be:\nModel grad ({param.grad}) != DDP grad ({ddp_param.grad})'''

        # Shuffle ddp_input on each iteration
  torch.manual_seed(1_3_3_7 + iteration			)
  lowerCAmelCase__     :   Dict     		=	ddp_input[torch.randperm(len(lowerCamelCase_			)			)]


def        UpperCAmelCase_ (  lowerCamelCase_=False      ,			lowerCamelCase_=False			):







 """simple docstring"""


 lowerCAmelCase__     :   List[str]     		=	Accelerator(
     split_batches=lowerCamelCase_      ,			dispatch_batches=lowerCamelCase_      ,			gradient_accumulation_steps=2			)
 # Test that context manager behaves properly
 lowerCAmelCase__		,     lowerCAmelCase__		,     lowerCAmelCase__     :   Any     		=	get_training_setup(lowerCamelCase_			)
 for iteration, batch in enumerate(lowerCamelCase_			):
  lowerCAmelCase__		,     lowerCAmelCase__     :   Tuple     		=	batch.values()
  # Gather the distributed inputs and targs for the base model
  lowerCAmelCase__		,     lowerCAmelCase__     :   Optional[int]     		=	accelerator.gather((ddp_input, ddp_target)			)
  lowerCAmelCase__		,     lowerCAmelCase__     :   Optional[Any]     		=	input.to(accelerator.device			), target.to(accelerator.device			)
  # Perform our initial ground truth step in non "DDP"
  step_model(lowerCamelCase_      ,			lowerCamelCase_      ,			lowerCamelCase_      ,			lowerCamelCase_      ,			lowerCamelCase_			)
  # Do "gradient accumulation" (noop)
  with accelerator.accumulate(lowerCamelCase_			):
   step_model(lowerCamelCase_      ,			lowerCamelCase_      ,			lowerCamelCase_      ,			lowerCamelCase_			)

  # DDP model and model should only be in sync when not (iteration % 2 == 0)
  for param, ddp_param in zip(model.parameters()      ,			ddp_model.parameters()			):
   if not param.requires_grad:
    continue
   if ((iteration + 1) % 2 == 0) or (iteration == len(lowerCamelCase_			) - 1):
    # Grads should be in sync
    assert (
        torch.allclose(param.grad      ,			ddp_param.grad			) is True
    ), f'''Gradients not in sync when they should be at iteration {iteration}:\nModel grad ({param.grad}) != DDP grad ({ddp_param.grad})'''
   else:
    # Grads should not be in sync
    assert (
        torch.allclose(param.grad      ,			ddp_param.grad			) is False
    ), f'''Gradients in sync when they should not be at iteration {iteration}:\nModel grad ({param.grad}) == DDP grad ({ddp_param.grad})'''

        # Shuffle ddp_input on each iteration
  torch.manual_seed(1_3_3_7 + iteration			)
  lowerCAmelCase__     :   Any     		=	ddp_input[torch.randperm(len(lowerCamelCase_			)			)]
 GradientState._reset_state()


def        UpperCAmelCase_ (  lowerCamelCase_=False      ,			lowerCamelCase_=False			):







 """simple docstring"""


 lowerCAmelCase__     :   Tuple     		=	Accelerator(
     split_batches=lowerCamelCase_      ,			dispatch_batches=lowerCamelCase_      ,			gradient_accumulation_steps=2			)
 # Test that context manager behaves properly
 lowerCAmelCase__		,     lowerCAmelCase__		,     lowerCAmelCase__		,     lowerCAmelCase__		,     lowerCAmelCase__		,     lowerCAmelCase__		,     lowerCAmelCase__     :   Dict     		=	get_training_setup(lowerCamelCase_      ,			lowerCamelCase_			)
 for iteration, batch in enumerate(lowerCamelCase_			):
  lowerCAmelCase__		,     lowerCAmelCase__     :   Any     		=	batch.values()
  # Gather the distributed inputs and targs for the base model
  lowerCAmelCase__		,     lowerCAmelCase__     :   Tuple     		=	accelerator.gather((ddp_input, ddp_target)			)
  lowerCAmelCase__		,     lowerCAmelCase__     :   int     		=	input.to(accelerator.device			), target.to(accelerator.device			)
  # Perform our initial ground truth step in non "DDP"
  model.train()
  ddp_model.train()
  step_model(lowerCamelCase_      ,			lowerCamelCase_      ,			lowerCamelCase_      ,			lowerCamelCase_      ,			lowerCamelCase_			)
  opt.step()

  if ((iteration + 1) % 2 == 0) or ((iteration + 1) == len(lowerCamelCase_			)):
   if split_batches:
    sched.step()
   else:
    for _ in range(accelerator.num_processes			):
     sched.step()
  opt.zero_grad()
  # Perform gradient accumulation under wrapper
  with accelerator.accumulate(lowerCamelCase_			):
   step_model(lowerCamelCase_      ,			lowerCamelCase_      ,			lowerCamelCase_      ,			lowerCamelCase_			)
   ddp_opt.step()
   ddp_sched.step()
   ddp_opt.zero_grad()

  # Learning rates should be the same
  assert (
      opt.param_groups[0]["lr"] == ddp_opt.param_groups[0]["lr"]
  ), f'''Learning rates found in each optimizer did not align\nopt: {opt.param_groups[0]['lr']}\nDDP opt: {ddp_opt.param_groups[0]['lr']}\n'''
  lowerCAmelCase__     :   int     		=	(((iteration + 1) % 2) == 0) or ((iteration + 1) == len(lowerCamelCase_			))
  if accelerator.num_processes > 1:
   check_model_parameters(lowerCamelCase_      ,			lowerCamelCase_      ,			lowerCamelCase_      ,			lowerCamelCase_			)
  # Shuffle ddp_input on each iteration
  torch.manual_seed(1_3_3_7 + iteration			)
 GradientState._reset_state()


def        UpperCAmelCase_ (  ):







 """simple docstring"""


 lowerCAmelCase__     :   Optional[Any]     		=	Accelerator()

 lowerCAmelCase__     :   Dict     		=	RegressionDataset(length=8_0			)
 lowerCAmelCase__     :   List[Any]     		=	DataLoader(lowerCamelCase_      ,			batch_size=1_6			)
 lowerCAmelCase__     :   Tuple     		=	RegressionDataset(length=9_6			)
 lowerCAmelCase__     :   Tuple     		=	DataLoader(lowerCamelCase_      ,			batch_size=1_6			)
 lowerCAmelCase__		,     lowerCAmelCase__     :   Dict     		=	accelerator.prepare(lowerCamelCase_      ,			lowerCamelCase_			)
 assert accelerator.gradient_state.active_dataloader is None
 for iteration, _ in enumerate(lowerCamelCase_			):
  assert id(accelerator.gradient_state.active_dataloader			) == id(lowerCamelCase_			)
  if iteration < len(lowerCamelCase_			) - 1:
   assert not accelerator.gradient_state.end_of_dataloader
   if iteration == 1:
    for batch_num, _ in enumerate(lowerCamelCase_			):
     assert id(accelerator.gradient_state.active_dataloader			) == id(lowerCamelCase_			)
     if batch_num < len(lowerCamelCase_			) - 1:
      assert not accelerator.gradient_state.end_of_dataloader
     else:
      assert accelerator.gradient_state.end_of_dataloader
  else:
   assert accelerator.gradient_state.end_of_dataloader
 assert accelerator.gradient_state.active_dataloader is None


def        UpperCAmelCase_ (  ):







 """simple docstring"""


 lowerCAmelCase__     :   List[str]     		=	Accelerator()
 lowerCAmelCase__     :   int     		=	accelerator.state
 if state.local_process_index == 0:
  print("**Test `accumulate` gradient accumulation with dataloader break**"			)
 test_dataloader_break()
 if state.distributed_type == DistributedType.NO:
  if state.local_process_index == 0:
   print("**Test NOOP `no_sync` context manager**"			)
  test_noop_sync(lowerCamelCase_			)
 if state.distributed_type in (DistributedType.MULTI_GPU, DistributedType.MULTI_CPU):
  if state.local_process_index == 0:
   print("**Test Distributed `no_sync` context manager**"			)
  test_distributed_sync(lowerCamelCase_			)
 if state.distributed_type == DistributedType.MULTI_GPU:
  for split_batch in [True, False]:
   for dispatch_batches in [True, False]:
    if state.local_process_index == 0:
     print(
         "**Test `accumulate` gradient accumulation, "      ,			f'''`split_batches={split_batch}` and `dispatch_batches={dispatch_batches}`**'''      ,			)
    test_gradient_accumulation(lowerCamelCase_      ,			lowerCamelCase_			)

    # Currently will break on torch 2.0 +, need to investigate why
 if is_torch_version("<"      ,			"2.0"			) or state.distributed_type == DistributedType.NO:
  if state.local_process_index == 0:
   print(
       "**Test `accumulate` gradient accumulation with optimizer and scheduler, "      ,			"`split_batches=False`, `dispatch_batches=False`**"      ,			)
  test_gradient_accumulation_with_opt_and_scheduler()
  if state.distributed_type == DistributedType.MULTI_GPU:
   for split_batch in [True, False]:
    for dispatch_batches in [True, False]:
     if not split_batch and not dispatch_batches:
      continue
     if state.local_process_index == 0:
      print(
          "**Test `accumulate` gradient accumulation with optimizer and scheduler, "      ,			f'''`split_batches={split_batch}` and `dispatch_batches={dispatch_batches}`**'''      ,			)
     test_gradient_accumulation_with_opt_and_scheduler(lowerCamelCase_      ,			lowerCamelCase_			)


def        UpperCAmelCase_ (  lowerCamelCase_			):







 """simple docstring"""


 main()


if __name__ == "__main__":
     main()







'''simple docstring'''

import tempfile

import numpy as np
import torch
from transformers import AutoTokenizer, TaEncoderModel

from diffusers import DDPMScheduler, UNetaDConditionModel
from diffusers.models.attention_processor import AttnAddedKVProcessor
from diffusers.pipelines.deepfloyd_if import IFWatermarker
from diffusers.utils.testing_utils import torch_device

from ..test_pipelines_common import to_np




class    lowerCAmelCase       :


 def    _A	(    self  : Union[str, Any]    ):
  '''simple docstring'''

  torch.manual_seed(0    )
  lowerCAmelCase__     :   List[str]     		=	TaEncoderModel.from_pretrained("hf-internal-testing/tiny-random-t5"    )

  torch.manual_seed(0    )
  lowerCAmelCase__     :   Optional[int]     		=	AutoTokenizer.from_pretrained("hf-internal-testing/tiny-random-t5"    )

  torch.manual_seed(0    )
  lowerCAmelCase__     :   Any     		=	UNetaDConditionModel(
      sample_size=32			,							layers_per_block=1			,							block_out_channels=[32, 64]			,							down_block_types=[
          "ResnetDownsampleBlock2D",
          "SimpleCrossAttnDownBlock2D",
      ]			,							mid_block_type="UNetMidBlock2DSimpleCrossAttn"			,							up_block_types=["SimpleCrossAttnUpBlock2D", "ResnetUpsampleBlock2D"]			,							in_channels=3			,							out_channels=6			,							cross_attention_dim=32			,							encoder_hid_dim=32			,							attention_head_dim=8			,							addition_embed_type="text"			,							addition_embed_type_num_heads=2			,							cross_attention_norm="group_norm"			,							resnet_time_scale_shift="scale_shift"			,							act_fn="gelu"			,							)
  unet.set_attn_processor(AttnAddedKVProcessor()    )  # For reproducibility tests

  torch.manual_seed(0    )
  lowerCAmelCase__     :   Union[str, Any]     		=	DDPMScheduler(
      num_train_timesteps=1000			,							beta_schedule="squaredcos_cap_v2"			,							beta_start=0.0001			,							beta_end=0.02			,							thresholding=a__			,							dynamic_thresholding_ratio=0.95			,							sample_max_value=1.0			,							prediction_type="epsilon"			,							variance_type="learned_range"			,							)

  torch.manual_seed(0    )
  lowerCAmelCase__     :   Tuple     		=	IFWatermarker()

  return {
      "text_encoder": text_encoder,
      "tokenizer": tokenizer,
      "unet": unet,
      "scheduler": scheduler,
      "watermarker": watermarker,
      "safety_checker": None,
      "feature_extractor": None,
  }


 def    _A	(    self  : Dict    ):
  '''simple docstring'''

  torch.manual_seed(0    )
  lowerCAmelCase__     :   Any     		=	TaEncoderModel.from_pretrained("hf-internal-testing/tiny-random-t5"    )

  torch.manual_seed(0    )
  lowerCAmelCase__     :   List[Any]     		=	AutoTokenizer.from_pretrained("hf-internal-testing/tiny-random-t5"    )

  torch.manual_seed(0    )
  lowerCAmelCase__     :   int     		=	UNetaDConditionModel(
      sample_size=32			,							layers_per_block=[1, 2]			,							block_out_channels=[32, 64]			,							down_block_types=[
          "ResnetDownsampleBlock2D",
          "SimpleCrossAttnDownBlock2D",
      ]			,							mid_block_type="UNetMidBlock2DSimpleCrossAttn"			,							up_block_types=["SimpleCrossAttnUpBlock2D", "ResnetUpsampleBlock2D"]			,							in_channels=6			,							out_channels=6			,							cross_attention_dim=32			,							encoder_hid_dim=32			,							attention_head_dim=8			,							addition_embed_type="text"			,							addition_embed_type_num_heads=2			,							cross_attention_norm="group_norm"			,							resnet_time_scale_shift="scale_shift"			,							act_fn="gelu"			,							class_embed_type="timestep"			,							mid_block_scale_factor=1.414			,							time_embedding_act_fn="gelu"			,							time_embedding_dim=32			,							)
  unet.set_attn_processor(AttnAddedKVProcessor()    )  # For reproducibility tests

  torch.manual_seed(0    )
  lowerCAmelCase__     :   Tuple     		=	DDPMScheduler(
      num_train_timesteps=1000			,							beta_schedule="squaredcos_cap_v2"			,							beta_start=0.0001			,							beta_end=0.02			,							thresholding=a__			,							dynamic_thresholding_ratio=0.95			,							sample_max_value=1.0			,							prediction_type="epsilon"			,							variance_type="learned_range"			,							)

  torch.manual_seed(0    )
  lowerCAmelCase__     :   Tuple     		=	DDPMScheduler(
      num_train_timesteps=1000			,							beta_schedule="squaredcos_cap_v2"			,							beta_start=0.0001			,							beta_end=0.02			,							)

  torch.manual_seed(0    )
  lowerCAmelCase__     :   Dict     		=	IFWatermarker()

  return {
      "text_encoder": text_encoder,
      "tokenizer": tokenizer,
      "unet": unet,
      "scheduler": scheduler,
      "image_noising_scheduler": image_noising_scheduler,
      "watermarker": watermarker,
      "safety_checker": None,
      "feature_extractor": None,
  }


 def    _A	(    self  : Optional[Any]    ):
  '''simple docstring'''

  lowerCAmelCase__     :   Union[str, Any]     		=	self.get_dummy_components()
  lowerCAmelCase__     :   Dict     		=	self.pipeline_class(**a__    )
  pipe.to(a__    )
  pipe.set_progress_bar_config(disable=a__    )

  lowerCAmelCase__     :   Tuple     		=	self.get_dummy_inputs(a__    )

  lowerCAmelCase__     :   str     		=	inputs["prompt"]
  lowerCAmelCase__     :   Optional[Any]     		=	inputs["generator"]
  lowerCAmelCase__     :   str     		=	inputs["num_inference_steps"]
  lowerCAmelCase__     :   int     		=	inputs["output_type"]

  if "image" in inputs:
   lowerCAmelCase__     :   Optional[int]     		=	inputs["image"]
  else:
   lowerCAmelCase__     :   str     		=	None

  if "mask_image" in inputs:
   lowerCAmelCase__     :   List[str]     		=	inputs["mask_image"]
  else:
   lowerCAmelCase__     :   str     		=	None

  if "original_image" in inputs:
   lowerCAmelCase__     :   Optional[int]     		=	inputs["original_image"]
  else:
   lowerCAmelCase__     :   Union[str, Any]     		=	None

  lowerCAmelCase__		,     lowerCAmelCase__     :   Dict     		=	pipe.encode_prompt(a__    )

  # inputs with prompt converted to embeddings
  lowerCAmelCase__     :   Dict     		=	{
      "prompt_embeds": prompt_embeds,
      "negative_prompt_embeds": negative_prompt_embeds,
      "generator": generator,
      "num_inference_steps": num_inference_steps,
      "output_type": output_type,
  }

  if image is not None:
   lowerCAmelCase__     :   Optional[Any]     		=	image

  if mask_image is not None:
   lowerCAmelCase__     :   Optional[int]     		=	mask_image

  if original_image is not None:
   lowerCAmelCase__     :   Optional[Any]     		=	original_image

  # set all optional components to None
  for optional_component in pipe._optional_components:
   setattr(a__			,							a__			,							a__    )

  lowerCAmelCase__     :   Optional[int]     		=	pipe(**a__    )[0]

  with tempfile.TemporaryDirectory() as tmpdir:
   pipe.save_pretrained(a__    )
   lowerCAmelCase__     :   Tuple     		=	self.pipeline_class.from_pretrained(a__    )
   pipe_loaded.to(a__    )
   pipe_loaded.set_progress_bar_config(disable=a__    )

  pipe_loaded.unet.set_attn_processor(AttnAddedKVProcessor()    )  # For reproducibility tests

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

  lowerCAmelCase__     :   Union[str, Any]     		=	self.get_dummy_inputs(a__    )

  lowerCAmelCase__     :   Any     		=	inputs["generator"]
  lowerCAmelCase__     :   Union[str, Any]     		=	inputs["num_inference_steps"]
  lowerCAmelCase__     :   str     		=	inputs["output_type"]

  # inputs with prompt converted to embeddings
  lowerCAmelCase__     :   int     		=	{
      "prompt_embeds": prompt_embeds,
      "negative_prompt_embeds": negative_prompt_embeds,
      "generator": generator,
      "num_inference_steps": num_inference_steps,
      "output_type": output_type,
  }

  if image is not None:
   lowerCAmelCase__     :   Tuple     		=	image

  if mask_image is not None:
   lowerCAmelCase__     :   int     		=	mask_image

  if original_image is not None:
   lowerCAmelCase__     :   List[Any]     		=	original_image

  lowerCAmelCase__     :   Dict     		=	pipe_loaded(**a__    )[0]

  lowerCAmelCase__     :   int     		=	np.abs(to_np(a__    ) - to_np(a__    )    ).max()
  self.assertLess(a__			,							1e-4    )


 def    _A	(    self  : Any    ):
  '''simple docstring'''

  lowerCAmelCase__     :   str     		=	self.get_dummy_components()
  lowerCAmelCase__     :   int     		=	self.pipeline_class(**a__    )
  pipe.to(a__    )
  pipe.set_progress_bar_config(disable=a__    )

  lowerCAmelCase__     :   Any     		=	self.get_dummy_inputs(a__    )
  lowerCAmelCase__     :   Union[str, Any]     		=	pipe(**a__    )[0]

  with tempfile.TemporaryDirectory() as tmpdir:
   pipe.save_pretrained(a__    )
   lowerCAmelCase__     :   str     		=	self.pipeline_class.from_pretrained(a__    )
   pipe_loaded.to(a__    )
   pipe_loaded.set_progress_bar_config(disable=a__    )

  pipe_loaded.unet.set_attn_processor(AttnAddedKVProcessor()    )  # For reproducibility tests

  lowerCAmelCase__     :   str     		=	self.get_dummy_inputs(a__    )
  lowerCAmelCase__     :   Union[str, Any]     		=	pipe_loaded(**a__    )[0]

  lowerCAmelCase__     :   Union[str, Any]     		=	np.abs(to_np(a__    ) - to_np(a__    )    ).max()
  self.assertLess(a__			,							1e-4    )





