




'''simple docstring'''




from dataclasses import dataclass
from typing import Optional

import numpy as np
import torch
import torch.nn as nn

from ..utils import BaseOutput, is_torch_version, randn_tensor
from .attention_processor import SpatialNorm
from .unet_ad_blocks import UNetMidBlockaD, get_down_block, get_up_block






@dataclass
class      UpperCAmelCase_     (  lowerCamelCase__     ):


	"""simple docstring"""


	snake_case									=		4_2






class      UpperCAmelCase_     (  nn.Module     ):


	"""simple docstring"""


	def __init__(						self						, UpperCAmelCase_=3						, UpperCAmelCase_=3						, UpperCAmelCase_=("DownEncoderBlock2D",)						, UpperCAmelCase_=(64,)						, UpperCAmelCase_=2						, UpperCAmelCase_=32						, UpperCAmelCase_="silu"						, UpperCAmelCase_=True						, ):
						super().__init__()
						snake_case_	=      layers_per_block

						snake_case_	=      torch.nn.Convad(
						    __snake_case						, block_out_channels[0]						, kernel_size=3						, stride=1						, padding=1						, )

						snake_case_	=      None
						snake_case_	=      nn.ModuleList([]				)

						# down
						snake_case_	=      block_out_channels[0]
						for i, down_block_type in enumerate(__snake_case				):
											snake_case_	=      output_channel
											snake_case_	=      block_out_channels[i]
											snake_case_	=      i == len(__snake_case				) - 1

											snake_case_	=      get_down_block(
											    __snake_case						, num_layers=self.layers_per_block						, in_channels=__snake_case						, out_channels=__snake_case						, add_downsample=not is_final_block						, resnet_eps=1e-6						, downsample_padding=0						, resnet_act_fn=__snake_case						, resnet_groups=__snake_case						, attention_head_dim=__snake_case						, temb_channels=__snake_case						, )
											self.down_blocks.append(__snake_case				)

						# mid
						snake_case_	=      UNetMidBlockaD(
						    in_channels=block_out_channels[-1]						, resnet_eps=1e-6						, resnet_act_fn=__snake_case						, output_scale_factor=1						, resnet_time_scale_shift="default"						, attention_head_dim=block_out_channels[-1]						, resnet_groups=__snake_case						, temb_channels=__snake_case						, )

						# out
						snake_case_	=      nn.GroupNorm(num_channels=block_out_channels[-1]						, num_groups=__snake_case						, eps=1e-6				)
						snake_case_	=      nn.SiLU()

						snake_case_	=      2 * out_channels if double_z else out_channels
						snake_case_	=      nn.Convad(block_out_channels[-1]						, __snake_case						, 3						, padding=1				)

						snake_case_	=      False







	def  _lowercase    (						self						, UpperCAmelCase_				):
						snake_case_	=      x
						snake_case_	=      self.conv_in(__snake_case				)

						if self.training and self.gradient_checkpointing:

											def create_custom_forward(UpperCAmelCase_				):
																def custom_forward(*UpperCAmelCase_				):
																					return module(*__snake_case				)

																return custom_forward

											# down
											if is_torch_version(">="						, "1.11.0"				):
																for down_block in self.down_blocks:
																					snake_case_	=      torch.utils.checkpoint.checkpoint(
																					    create_custom_forward(__snake_case				)						, __snake_case						, use_reentrant=__snake_case				)
																# middle
																snake_case_	=      torch.utils.checkpoint.checkpoint(
																    create_custom_forward(self.mid_block				)						, __snake_case						, use_reentrant=__snake_case				)
											else:
																for down_block in self.down_blocks:
																					snake_case_	=      torch.utils.checkpoint.checkpoint(create_custom_forward(__snake_case				)						, __snake_case				)
																# middle
																snake_case_	=      torch.utils.checkpoint.checkpoint(create_custom_forward(self.mid_block				)						, __snake_case				)

						else:
											# down
											for down_block in self.down_blocks:
																snake_case_	=      down_block(__snake_case				)

											# middle
											snake_case_	=      self.mid_block(__snake_case				)

						# post-process
						snake_case_	=      self.conv_norm_out(__snake_case				)
						snake_case_	=      self.conv_act(__snake_case				)
						snake_case_	=      self.conv_out(__snake_case				)

						return sample






class      UpperCAmelCase_     (  nn.Module     ):


	"""simple docstring"""


	def __init__(						self						, UpperCAmelCase_=3						, UpperCAmelCase_=3						, UpperCAmelCase_=("UpDecoderBlock2D",)						, UpperCAmelCase_=(64,)						, UpperCAmelCase_=2						, UpperCAmelCase_=32						, UpperCAmelCase_="silu"						, UpperCAmelCase_="group"						, ):
						super().__init__()
						snake_case_	=      layers_per_block

						snake_case_	=      nn.Convad(
						    __snake_case						, block_out_channels[-1]						, kernel_size=3						, stride=1						, padding=1						, )

						snake_case_	=      None
						snake_case_	=      nn.ModuleList([]				)

						snake_case_	=      in_channels if norm_type == 'spatial' else None

						# mid
						snake_case_	=      UNetMidBlockaD(
						    in_channels=block_out_channels[-1]						, resnet_eps=1e-6						, resnet_act_fn=__snake_case						, output_scale_factor=1						, resnet_time_scale_shift="default" if norm_type == "group" else norm_type						, attention_head_dim=block_out_channels[-1]						, resnet_groups=__snake_case						, temb_channels=__snake_case						, )

						# up
						snake_case_	=      list(reversed(__snake_case				)				)
						snake_case_	=      reversed_block_out_channels[0]
						for i, up_block_type in enumerate(__snake_case				):
											snake_case_	=      output_channel
											snake_case_	=      reversed_block_out_channels[i]

											snake_case_	=      i == len(__snake_case				) - 1

											snake_case_	=      get_up_block(
											    __snake_case						, num_layers=self.layers_per_block + 1						, in_channels=__snake_case						, out_channels=__snake_case						, prev_output_channel=__snake_case						, add_upsample=not is_final_block						, resnet_eps=1e-6						, resnet_act_fn=__snake_case						, resnet_groups=__snake_case						, attention_head_dim=__snake_case						, temb_channels=__snake_case						, resnet_time_scale_shift=__snake_case						, )
											self.up_blocks.append(__snake_case				)
											snake_case_	=      output_channel

						# out
						if norm_type == "spatial":
											snake_case_	=      SpatialNorm(block_out_channels[0]						, __snake_case				)
						else:
											snake_case_	=      nn.GroupNorm(num_channels=block_out_channels[0]						, num_groups=__snake_case						, eps=1e-6				)
						snake_case_	=      nn.SiLU()
						snake_case_	=      nn.Convad(block_out_channels[0]						, __snake_case						, 3						, padding=1				)

						snake_case_	=      False







	def  _lowercase    (						self						, UpperCAmelCase_						, UpperCAmelCase_=None				):
						snake_case_	=      z
						snake_case_	=      self.conv_in(__snake_case				)

						snake_case_	=      next(iter(self.up_blocks.parameters()				)				).dtype
						if self.training and self.gradient_checkpointing:

											def create_custom_forward(UpperCAmelCase_				):
																def custom_forward(*UpperCAmelCase_				):
																					return module(*__snake_case				)

																return custom_forward

											if is_torch_version(">="						, "1.11.0"				):
																# middle
																snake_case_	=      torch.utils.checkpoint.checkpoint(
																    create_custom_forward(self.mid_block				)						, __snake_case						, __snake_case						, use_reentrant=__snake_case				)
																snake_case_	=      sample.to(__snake_case				)

																# up
																for up_block in self.up_blocks:
																					snake_case_	=      torch.utils.checkpoint.checkpoint(
																					    create_custom_forward(__snake_case				)						, __snake_case						, __snake_case						, use_reentrant=__snake_case				)
											else:
																# middle
																snake_case_	=      torch.utils.checkpoint.checkpoint(
																    create_custom_forward(self.mid_block				)						, __snake_case						, __snake_case				)
																snake_case_	=      sample.to(__snake_case				)

																# up
																for up_block in self.up_blocks:
																					snake_case_	=      torch.utils.checkpoint.checkpoint(create_custom_forward(__snake_case				)						, __snake_case						, __snake_case				)
						else:
											# middle
											snake_case_	=      self.mid_block(__snake_case						, __snake_case				)
											snake_case_	=      sample.to(__snake_case				)

											# up
											for up_block in self.up_blocks:
																snake_case_	=      up_block(__snake_case						, __snake_case				)

											# post-process
						if latent_embeds is None:
											snake_case_	=      self.conv_norm_out(__snake_case				)
						else:
											snake_case_	=      self.conv_norm_out(__snake_case						, __snake_case				)
						snake_case_	=      self.conv_act(__snake_case				)
						snake_case_	=      self.conv_out(__snake_case				)

						return sample






class      UpperCAmelCase_     (  nn.Module     ):


	"""simple docstring"""


	def __init__(						self						, UpperCAmelCase_						, UpperCAmelCase_						, UpperCAmelCase_						, UpperCAmelCase_=None						, UpperCAmelCase_="random"						, UpperCAmelCase_=False						, UpperCAmelCase_=True				):
						super().__init__()
						snake_case_	=      n_e
						snake_case_	=      vq_embed_dim
						snake_case_	=      beta
						snake_case_	=      legacy

						snake_case_	=      nn.Embedding(self.n_e						, self.vq_embed_dim				)
						self.embedding.weight.data.uniform_(-1.0 / self.n_e						, 1.0 / self.n_e				)

						snake_case_	=      remap
						if self.remap is not None:
											self.register_buffer("used"						, torch.tensor(np.load(self.remap				)				)				)
											snake_case_	=      self.used.shape[0]
											snake_case_	=      unknown_index  # "random" or "extra" or integer
											if self.unknown_index == "extra":
																snake_case_	=      self.re_embed
																snake_case_	=      self.re_embed + 1
											print(
											    f'''Remapping {self.n_e} indices to {self.re_embed} indices. '''
											    f'''Using {self.unknown_index} for unknown indices.'''				)
						else:
											snake_case_	=      n_e

						snake_case_	=      sane_index_shape







	def  _lowercase    (						self						, UpperCAmelCase_				):
						snake_case_	=      inds.shape
						assert len(__snake_case				) > 1
						snake_case_	=      inds.reshape(ishape[0]						, -1				)
						snake_case_	=      self.used.to(__snake_case				)
						snake_case_	=      (inds[:, :, None] == used[None, None, ...]).long()
						snake_case_	=      match.argmax(-1				)
						snake_case_	=      match.sum(2				) < 1
						if self.unknown_index == "random":
											snake_case_	=      torch.randint(0						, self.re_embed						, size=new[unknown].shape				).to(device=new.device				)
						else:
											snake_case_	=      self.unknown_index
						return new.reshape(__snake_case				)







	def  _lowercase    (						self						, UpperCAmelCase_				):
						snake_case_	=      inds.shape
						assert len(__snake_case				) > 1
						snake_case_	=      inds.reshape(ishape[0]						, -1				)
						snake_case_	=      self.used.to(__snake_case				)
						if self.re_embed > self.used.shape[0]:  # extra token
											snake_case_	=      0  # simply set to zero
						snake_case_	=      torch.gather(used[None, :][inds.shape[0] * [0], :]						, 1						, __snake_case				)
						return back.reshape(__snake_case				)







	def  _lowercase    (						self						, UpperCAmelCase_				):
						# reshape z -> (batch, height, width, channel) and flatten
						snake_case_	=      z.permute(0						, 2						, 3						, 1				).contiguous()
						snake_case_	=      z.view(-1						, self.vq_embed_dim				)

						# distances from z to embeddings e_j (z - e)^2 = z^2 + e^2 - 2 e * z
						snake_case_	=      torch.argmin(torch.cdist(__snake_case						, self.embedding.weight				)						, dim=1				)

						snake_case_	=      self.embedding(__snake_case				).view(z.shape				)
						snake_case_	=      None
						snake_case_	=      None

						# compute loss for embedding
						if not self.legacy:
											snake_case_	=      self.beta * torch.mean((z_q.detach() - z) ** 2				) + torch.mean((z_q - z.detach()) ** 2				)
						else:
											snake_case_	=      torch.mean((z_q.detach() - z) ** 2				) + self.beta * torch.mean((z_q - z.detach()) ** 2				)

						# preserve gradients
						snake_case_	=      z + (z_q - z).detach()

						# reshape back to match original input shape
						snake_case_	=      z_q.permute(0						, 3						, 1						, 2				).contiguous()

						if self.remap is not None:
											snake_case_	=      min_encoding_indices.reshape(z.shape[0]						, -1				)  # add batch axis
											snake_case_	=      self.remap_to_used(__snake_case				)
											snake_case_	=      min_encoding_indices.reshape(-1						, 1				)  # flatten

						if self.sane_index_shape:
											snake_case_	=      min_encoding_indices.reshape(z_q.shape[0]						, z_q.shape[2]						, z_q.shape[3]				)

						return z_q, loss, (perplexity, min_encodings, min_encoding_indices)







	def  _lowercase    (						self						, UpperCAmelCase_						, UpperCAmelCase_				):
						# shape specifying (batch, height, width, channel)
						if self.remap is not None:
											snake_case_	=      indices.reshape(shape[0]						, -1				)  # add batch axis
											snake_case_	=      self.unmap_to_all(__snake_case				)
											snake_case_	=      indices.reshape(-1				)  # flatten again

						# get quantized latent vectors
						snake_case_	=      self.embedding(__snake_case				)

						if shape is not None:
											snake_case_	=      z_q.view(__snake_case				)
											# reshape back to match original input shape
											snake_case_	=      z_q.permute(0						, 3						, 1						, 2				).contiguous()

						return z_q






class      UpperCAmelCase_     (  lowerCamelCase__     ):


	"""simple docstring"""


	def __init__(						self						, UpperCAmelCase_						, UpperCAmelCase_=False				):
						snake_case_	=      parameters
						snake_case_	=      torch.chunk(__snake_case						, 2						, dim=1				)
						snake_case_	=      torch.clamp(self.logvar						, -30.0						, 20.0				)
						snake_case_	=      deterministic
						snake_case_	=      torch.exp(0.5 * self.logvar				)
						snake_case_	=      torch.exp(self.logvar				)
						if self.deterministic:
											snake_case_	=      torch.zeros_like(
											    self.mean						, device=self.parameters.device						, dtype=self.parameters.dtype				)







	def  _lowercase    (						self						, UpperCAmelCase_ = None				):
						# make sure sample is on the same device as the parameters and has same dtype
						snake_case_	=      randn_tensor(
						    self.mean.shape						, generator=__snake_case						, device=self.parameters.device						, dtype=self.parameters.dtype				)
						snake_case_	=      self.mean + self.std * sample
						return x







	def  _lowercase    (						self						, UpperCAmelCase_=None				):
						if self.deterministic:
											return torch.Tensor([0.0]				)
						else:
											if other is None:
																return 0.5 * torch.sum(torch.pow(self.mean						, 2				) + self.var - 1.0 - self.logvar						, dim=[1, 2, 3]				)
											else:
																return 0.5 * torch.sum(
																    torch.pow(self.mean - other.mean						, 2				) / other.var
																    + self.var / other.var
																    - 1.0
																    - self.logvar
																    + other.logvar						, dim=[1, 2, 3]						, )







	def  _lowercase    (						self						, UpperCAmelCase_						, UpperCAmelCase_=[1, 2, 3]				):
						if self.deterministic:
											return torch.Tensor([0.0]				)
						snake_case_	=      np.log(2.0 * np.pi				)
						return 0.5 * torch.sum(logtwopi + self.logvar + torch.pow(sample - self.mean						, 2				) / self.var						, dim=__snake_case				)







	def  _lowercase    (						self				):
						return self.mean










'''simple docstring'''






import gc
import random
import unittest

import torch

from diffusers import (
    IFImgaImgPipeline,
    IFImgaImgSuperResolutionPipeline,
    IFInpaintingPipeline,
    IFInpaintingSuperResolutionPipeline,
    IFPipeline,
    IFSuperResolutionPipeline,
)
from diffusers.models.attention_processor import AttnAddedKVProcessor
from diffusers.utils.import_utils import is_xformers_available
from diffusers.utils.testing_utils import floats_tensor, load_numpy, require_torch_gpu, skip_mps, slow, torch_device

from ..pipeline_params import TEXT_TO_IMAGE_BATCH_PARAMS, TEXT_TO_IMAGE_PARAMS
from ..test_pipelines_common import PipelineTesterMixin, assert_mean_pixel_difference
from . import IFPipelineTesterMixin

@skip_mps
class 				a__(   lowerCamelCase__						,					lowerCamelCase__						,					unittest.TestCase      ):
    lowercase__					=       IFPipeline
    lowercase__					=       TEXT_TO_IMAGE_PARAMS - {"""width""", """height""", """latents"""}
    lowercase__					=       TEXT_TO_IMAGE_BATCH_PARAMS
    lowercase__					=       PipelineTesterMixin.required_optional_params - {"""latents"""}







    def      lowercase_			( self					:    List[Any]      ):
      return self._get_dummy_components()







    def      lowercase_			( self					:    int		,     __snake_case					:    List[str]		,     __snake_case					:    List[Any]=0      ):
      if str(__snake_case      ).startswith('mps'      ):
        a		:   List[Any]	  =	torch.manual_seed(__snake_case      )
      else:
        a		:   str	  =	torch.Generator(device=__snake_case      ).manual_seed(__snake_case      )

      a		:   Dict	  =	{
          'prompt': 'A painting of a squirrel eating a burger',
          'generator': generator,
          'num_inference_steps': 2,
          'output_type': 'numpy',
      }

      return inputs







    def      lowercase_			( self					:    Optional[Any]      ):
      self._test_save_load_optional_components()







    @unittest.skipIf(torch_device != 'cuda'		,     reason='float16 requires CUDA'      )
    def      lowercase_			( self					:    Union[str, Any]      ):
      # Due to non-determinism in save load of the hf-internal-testing/tiny-random-t5 text encoder
      super().test_save_load_floataa(expected_max_diff=1e-1      )







    def      lowercase_			( self					:    Tuple      ):
      self._test_attention_slicing_forward_pass(expected_max_diff=1e-2      )







    def      lowercase_			( self					:    str      ):
      self._test_save_load_local()







    def      lowercase_			( self					:    Any      ):
      self._test_inference_batch_single_identical(
          expected_max_diff=1e-2		,     )







    @unittest.skipIf(
        torch_device != 'cuda' or not is_xformers_available()		,     reason='XFormers attention is only available with CUDA and `xformers` installed'		,     )
    def      lowercase_			( self					:    List[str]      ):
      self._test_xformers_attention_forwardGenerator_pass(expected_max_diff=1e-3      )

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







    def      lowercase_			( self					:    Optional[Any]      ):
      # clean up the VRAM after each test
      super().tearDown()
      gc.collect()
      torch.cuda.empty_cache()







    def      lowercase_			( self					:    Tuple      ):
      # if

      a		:   List[Any]	  =	IFPipeline.from_pretrained('DeepFloyd/IF-I-XL-v1.0'		,     variant='fp16'		,     torch_dtype=torch.floataa      )

      a		:   Optional[int]	  =	IFSuperResolutionPipeline.from_pretrained(
          'DeepFloyd/IF-II-L-v1.0'		,     variant='fp16'		,     torch_dtype=torch.floataa		,     text_encoder=__snake_case		,     tokenizer=__snake_case      )

      # pre compute text embeddings and remove T5 to save memory

      pipe_a.text_encoder.to('cuda'      )

      a  ,      a		:   Optional[Any]	  =	pipe_a.encode_prompt('anime turtle'		,     device='cuda'      )

      del pipe_a.tokenizer
      del pipe_a.text_encoder
      gc.collect()

      a		:   List[Any]	  =	None
      a		:   Dict	  =	None

      pipe_a.enable_model_cpu_offload()
      pipe_a.enable_model_cpu_offload()

      pipe_a.unet.set_attn_processor(AttnAddedKVProcessor()      )
      pipe_a.unet.set_attn_processor(AttnAddedKVProcessor()      )

      self._test_if(__snake_case		,     __snake_case		,     __snake_case		,     __snake_case      )

      pipe_a.remove_all_hooks()
      pipe_a.remove_all_hooks()

      # img2img

      a		:   List[str]	  =	IFImgaImgPipeline(**pipe_a.components      )
      a		:   int	  =	IFImgaImgSuperResolutionPipeline(**pipe_a.components      )

      pipe_a.enable_model_cpu_offload()
      pipe_a.enable_model_cpu_offload()

      pipe_a.unet.set_attn_processor(AttnAddedKVProcessor()      )
      pipe_a.unet.set_attn_processor(AttnAddedKVProcessor()      )

      self._test_if_imgaimg(__snake_case		,     __snake_case		,     __snake_case		,     __snake_case      )

      pipe_a.remove_all_hooks()
      pipe_a.remove_all_hooks()

      # inpainting

      a		:   Tuple	  =	IFInpaintingPipeline(**pipe_a.components      )
      a		:   int	  =	IFInpaintingSuperResolutionPipeline(**pipe_a.components      )

      pipe_a.enable_model_cpu_offload()
      pipe_a.enable_model_cpu_offload()

      pipe_a.unet.set_attn_processor(AttnAddedKVProcessor()      )
      pipe_a.unet.set_attn_processor(AttnAddedKVProcessor()      )

      self._test_if_inpainting(__snake_case		,     __snake_case		,     __snake_case		,     __snake_case      )







    def      lowercase_			( self					:    Optional[Any]		,     __snake_case					:    Optional[int]		,     __snake_case					:    Union[str, Any]		,     __snake_case					:    Tuple		,     __snake_case					:    List[str]      ):
      # pipeline 1

      _start_torch_memory_measurement()

      a		:   int	  =	torch.Generator(device='cpu'      ).manual_seed(0      )
      a		:   int	  =	pipe_a(
          prompt_embeds=__snake_case		,     negative_prompt_embeds=__snake_case		,     num_inference_steps=2		,     generator=__snake_case		,     output_type='np'		,     )

      a		:   Dict	  =	output.images[0]

      assert image.shape == (64, 64, 3)

      a		:   Any	  =	torch.cuda.max_memory_allocated()
      assert mem_bytes < 13 * 10**9

      a		:   Optional[int]	  =	load_numpy(
          'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/if/test_if.npy'      )
      assert_mean_pixel_difference(__snake_case		,     __snake_case      )

      # pipeline 2

      _start_torch_memory_measurement()

      a		:   List[str]	  =	torch.Generator(device='cpu'      ).manual_seed(0      )

      a		:   List[Any]	  =	floats_tensor((1, 3, 64, 64)		,     rng=random.Random(0      )      ).to(__snake_case      )

      a		:   Any	  =	pipe_a(
          prompt_embeds=__snake_case		,     negative_prompt_embeds=__snake_case		,     image=__snake_case		,     generator=__snake_case		,     num_inference_steps=2		,     output_type='np'		,     )

      a		:   Tuple	  =	output.images[0]

      assert image.shape == (2_56, 2_56, 3)

      a		:   Optional[Any]	  =	torch.cuda.max_memory_allocated()
      assert mem_bytes < 4 * 10**9

      a		:   Tuple	  =	load_numpy(
          'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/if/test_if_superresolution_stage_II.npy'      )
      assert_mean_pixel_difference(__snake_case		,     __snake_case      )







    def      lowercase_			( self					:    List[str]		,     __snake_case					:    str		,     __snake_case					:    int		,     __snake_case					:    Any		,     __snake_case					:    Any      ):
      # pipeline 1

      _start_torch_memory_measurement()

      a		:   Optional[Any]	  =	floats_tensor((1, 3, 64, 64)		,     rng=random.Random(0      )      ).to(__snake_case      )

      a		:   Optional[int]	  =	torch.Generator(device='cpu'      ).manual_seed(0      )

      a		:   Any	  =	pipe_a(
          prompt_embeds=__snake_case		,     negative_prompt_embeds=__snake_case		,     image=__snake_case		,     num_inference_steps=2		,     generator=__snake_case		,     output_type='np'		,     )

      a		:   List[Any]	  =	output.images[0]

      assert image.shape == (64, 64, 3)

      a		:   Optional[Any]	  =	torch.cuda.max_memory_allocated()
      assert mem_bytes < 10 * 10**9

      a		:   Tuple	  =	load_numpy(
          'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/if/test_if_img2img.npy'      )
      assert_mean_pixel_difference(__snake_case		,     __snake_case      )

      # pipeline 2

      _start_torch_memory_measurement()

      a		:   Optional[int]	  =	torch.Generator(device='cpu'      ).manual_seed(0      )

      a		:   str	  =	floats_tensor((1, 3, 2_56, 2_56)		,     rng=random.Random(0      )      ).to(__snake_case      )
      a		:   Any	  =	floats_tensor((1, 3, 64, 64)		,     rng=random.Random(0      )      ).to(__snake_case      )

      a		:   Dict	  =	pipe_a(
          prompt_embeds=__snake_case		,     negative_prompt_embeds=__snake_case		,     image=__snake_case		,     original_image=__snake_case		,     generator=__snake_case		,     num_inference_steps=2		,     output_type='np'		,     )

      a		:   Optional[Any]	  =	output.images[0]

      assert image.shape == (2_56, 2_56, 3)

      a		:   Any	  =	torch.cuda.max_memory_allocated()
      assert mem_bytes < 4 * 10**9

      a		:   Dict	  =	load_numpy(
          'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/if/test_if_img2img_superresolution_stage_II.npy'      )
      assert_mean_pixel_difference(__snake_case		,     __snake_case      )







    def      lowercase_			( self					:    List[Any]		,     __snake_case					:    Tuple		,     __snake_case					:    Any		,     __snake_case					:    str		,     __snake_case					:    int      ):
      # pipeline 1

      _start_torch_memory_measurement()

      a		:   Tuple	  =	floats_tensor((1, 3, 64, 64)		,     rng=random.Random(0      )      ).to(__snake_case      )
      a		:   Union[str, Any]	  =	floats_tensor((1, 3, 64, 64)		,     rng=random.Random(1      )      ).to(__snake_case      )

      a		:   Union[str, Any]	  =	torch.Generator(device='cpu'      ).manual_seed(0      )
      a		:   str	  =	pipe_a(
          prompt_embeds=__snake_case		,     negative_prompt_embeds=__snake_case		,     image=__snake_case		,     mask_image=__snake_case		,     num_inference_steps=2		,     generator=__snake_case		,     output_type='np'		,     )

      a		:   List[Any]	  =	output.images[0]

      assert image.shape == (64, 64, 3)

      a		:   List[str]	  =	torch.cuda.max_memory_allocated()
      assert mem_bytes < 10 * 10**9

      a		:   Optional[int]	  =	load_numpy(
          'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/if/test_if_inpainting.npy'      )
      assert_mean_pixel_difference(__snake_case		,     __snake_case      )

      # pipeline 2

      _start_torch_memory_measurement()

      a		:   Union[str, Any]	  =	torch.Generator(device='cpu'      ).manual_seed(0      )

      a		:   Optional[int]	  =	floats_tensor((1, 3, 64, 64)		,     rng=random.Random(0      )      ).to(__snake_case      )
      a		:   Tuple	  =	floats_tensor((1, 3, 2_56, 2_56)		,     rng=random.Random(0      )      ).to(__snake_case      )
      a		:   Optional[int]	  =	floats_tensor((1, 3, 2_56, 2_56)		,     rng=random.Random(1      )      ).to(__snake_case      )

      a		:   int	  =	pipe_a(
          prompt_embeds=__snake_case		,     negative_prompt_embeds=__snake_case		,     image=__snake_case		,     mask_image=__snake_case		,     original_image=__snake_case		,     generator=__snake_case		,     num_inference_steps=2		,     output_type='np'		,     )

      a		:   List[str]	  =	output.images[0]

      assert image.shape == (2_56, 2_56, 3)

      a		:   Union[str, Any]	  =	torch.cuda.max_memory_allocated()
      assert mem_bytes < 4 * 10**9

      a		:   Any	  =	load_numpy(
          'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/if/test_if_inpainting_superresolution_stage_II.npy'      )
      assert_mean_pixel_difference(__snake_case		,     __snake_case      )
def  lowerCamelCase__			(		):
  torch.cuda.empty_cache()
  torch.cuda.reset_max_memory_allocated()
  torch.cuda.reset_peak_memory_stats()