




"""simple docstring"""



import random
import unittest

import torch

from diffusers import IFInpaintingSuperResolutionPipeline
from diffusers.utils import floats_tensor
from diffusers.utils.import_utils import is_xformers_available
from diffusers.utils.testing_utils import skip_mps, torch_device

from ..pipeline_params import (
    TEXT_GUIDED_IMAGE_INPAINTING_BATCH_PARAMS,
    TEXT_GUIDED_IMAGE_INPAINTING_PARAMS,
)
from ..test_pipelines_common import PipelineTesterMixin
from . import IFPipelineTesterMixin




@skip_mps
class      UpperCAmelCase_  (						A_,       A_,       unittest.TestCase							):
    lowercase__						 =				IFInpaintingSuperResolutionPipeline
    lowercase__						 =				TEXT_GUIDED_IMAGE_INPAINTING_PARAMS - {'''width''', '''height'''}
    lowercase__						 =				TEXT_GUIDED_IMAGE_INPAINTING_BATCH_PARAMS.union({'''original_image'''}							)
    lowercase__						 =				PipelineTesterMixin.required_optional_params - {'''latents'''}




    def 							__magic_name__							(   self						:	Union[str, Any]				)				->							Union[str, Any]:




     '''simple docstring'''

     return self._get_superresolution_dummy_components()




    def 							__magic_name__							(   self						:	Tuple     ,			snake_case_						:	List[str]     ,			snake_case_						:	Any=0				)				->							Optional[int]:




     '''simple docstring'''

     if str(snake_case_				).startswith("mps"				):
      A__	=    torch.manual_seed(snake_case_				)
     else:
      A__	=    torch.Generator(device=snake_case_				).manual_seed(snake_case_				)

     A__	=    floats_tensor((1, 3, 16, 16)     ,			rng=random.Random(snake_case_				)				).to(snake_case_				)
     A__	=    floats_tensor((1, 3, 32, 32)     ,			rng=random.Random(snake_case_				)				).to(snake_case_				)
     A__	=    floats_tensor((1, 3, 32, 32)     ,			rng=random.Random(snake_case_				)				).to(snake_case_				)

     A__	=    {
         "prompt": "A painting of a squirrel eating a burger",
         "image": image,
         "original_image": original_image,
         "mask_image": mask_image,
         "generator": generator,
         "num_inference_steps": 2,
         "output_type": "numpy",
     }

     return inputs




    @unittest.skipIf(
        torch_device != "cuda" or not is_xformers_available()     ,			reason="XFormers attention is only available with CUDA and `xformers` installed"     ,			)
    def 							__magic_name__							(   self						:	Tuple				)				->							Dict:




     '''simple docstring'''

     self._test_xformers_attention_forwardGenerator_pass(expected_max_diff=1e-3				)




    def 							__magic_name__							(   self						:	str				)				->							List[Any]:




     '''simple docstring'''

     self._test_save_load_optional_components()




    @unittest.skipIf(torch_device != "cuda"     ,			reason="float16 requires CUDA"				)
    def 							__magic_name__							(   self						:	Optional[int]				)				->							Optional[Any]:




     '''simple docstring'''

     super().test_save_load_floataa(expected_max_diff=1e-1				)




    def 							__magic_name__							(   self						:	str				)				->							str:




     '''simple docstring'''

     self._test_attention_slicing_forward_pass(expected_max_diff=1e-2				)




    def 							__magic_name__							(   self						:	str				)				->							Dict:




     '''simple docstring'''

     self._test_save_load_local()




    def 							__magic_name__							(   self						:	str				)				->							List[str]:




     '''simple docstring'''

     self._test_inference_batch_single_identical(
         expected_max_diff=1e-2     ,			)












"""simple docstring"""



import re

from flax.core.frozen_dict import freeze
from flax.traverse_util import flatten_dict, unflatten_dict
from jax.experimental import PartitionSpec as P


# Sentinels
SCREAMING_SNAKE_CASE    =	object()

# For specifying empty leaf dict `{}`
SCREAMING_SNAKE_CASE    =	object()







def   _SCREAMING_SNAKE_CASE       (							lowercase_      ,			lowercase_						)						->		List[Any]:
 A__	=    tuple((re.compile(x + "$"						) for x in qs)						)
 for i in range(len(lowercase_						) - len(lowercase_						) + 1						):
  A__	=    [x.match(lowercase_						) for x, y in zip(lowercase_      ,			ks[i:]						)]
  if matches and all(lowercase_						):
   return True
 return False







def   _SCREAMING_SNAKE_CASE       (							lowercase_						)						->		Optional[int]:
 def replace(lowercase_      ,			lowercase_						):
  for rule, replacement in rules:
   if _match(lowercase_      ,			lowercase_						):
    return replacement
  return val

 return replace







def   _SCREAMING_SNAKE_CASE       (							)						->		List[str]:
 return [
     # embeddings
     (("transformer", "wpe", "embedding"), P("mp"      ,			lowercase_						)),
     (("transformer", "wte", "embedding"), P("mp"      ,			lowercase_						)),
     # atention
     (("attention", "(q_proj|k_proj|v_proj)", "kernel"), P(lowercase_      ,			"mp"						)),
     (("attention", "out_proj", "kernel"), P("mp"      ,			lowercase_						)),
     (("attention", "out_proj", "bias"), None),
     # mlp
     (("mlp", "c_fc", "kernel"), P(lowercase_      ,			"mp"						)),
     (("mlp", "c_fc", "bias"), P("mp"						)),
     (("mlp", "c_proj", "kernel"), P("mp"      ,			lowercase_						)),
     (("mlp", "c_proj", "bias"), None),
     # layer norms
     ((r"ln_\d+", "bias"), None),
     ((r"\d+", r"ln_\d+", "scale"), None),
     (("ln_f", "bias"), None),
     (("ln_f", "scale"), None),
 ]







def   _SCREAMING_SNAKE_CASE       (							lowercase_						)						->		int:
 A__	=    _get_partition_rules()
 A__	=    _replacement_rules(lowercase_						)
 A__	=    {k: _unmatched for k in flatten_dict(lowercase_						)}
 A__	=    {k: replace(lowercase_      ,			lowercase_						) for k, v in initd.items()}
 assert _unmatched not in result.values(), "Incomplete partition spec."
 return freeze(unflatten_dict(lowercase_						)						)





