




'''simple docstring'''




import enum
import inspect
from typing import Iterable, List, Optional, Tuple, Union

class 						lowerCAmelCase_   ( enum.Enum   ):







  '''simple docstring'''







  lowerCAmelCase_				:					str				     =						"""timm"""
  lowerCAmelCase_				:					Any				     =						"""transformers"""

def    _UpperCamelCase      (				SCREAMING_SNAKE_CASE__			:      Optional[Iterable[str]]       ,						SCREAMING_SNAKE_CASE__			:      Optional[Iterable[int]]       ,						SCREAMING_SNAKE_CASE__			:      Optional[Iterable[str]]       ):
 '''simple docstring'''





 if stage_names is None:
  raise ValueError("""Stage_names must be set for transformers backbones"""       )

 if out_features is not None:
  if not isinstance(SCREAMING_SNAKE_CASE__       ,						(list,)       ):
   raise ValueError(F'''out_features must be a list {type(SCREAMING_SNAKE_CASE__       )}'''       )
  if any(feat not in stage_names for feat in out_features       ):
   raise ValueError(F'''out_features must be a subset of stage_names: {stage_names} got {out_features}'''       )

 if out_indices is not None:
  if not isinstance(SCREAMING_SNAKE_CASE__       ,						(list, tuple)       ):
   raise ValueError(F'''out_indices must be a list or tuple, got {type(SCREAMING_SNAKE_CASE__       )}'''       )
  if any(idx >= len(SCREAMING_SNAKE_CASE__       ) for idx in out_indices       ):
   raise ValueError("""out_indices must be valid indices for stage_names {stage_names}, got {out_indices}"""       )

 if out_features is not None and out_indices is not None:
  if len(SCREAMING_SNAKE_CASE__       ) != len(SCREAMING_SNAKE_CASE__       ):
   raise ValueError("""out_features and out_indices should have the same length if both are set"""       )
  if out_features != [stage_names[idx] for idx in out_indices]:
   raise ValueError("""out_features and out_indices should correspond to the same stages if both are set"""       )

def    _UpperCamelCase      (				SCREAMING_SNAKE_CASE__			:      Optional[List[str]]       ,						SCREAMING_SNAKE_CASE__			:      Optional[Union[List[int], Tuple[int]]]       ,						SCREAMING_SNAKE_CASE__			:      List[str]       ,						):
 '''simple docstring'''





 if out_indices is None and out_features is None:
  UpperCAmelCase__										=    [len(SCREAMING_SNAKE_CASE__       ) - 1]
  UpperCAmelCase__										=    [stage_names[-1]]
 elif out_indices is None and out_features is not None:
  UpperCAmelCase__										=    [stage_names.index(SCREAMING_SNAKE_CASE__       ) for layer in out_features]
 elif out_features is None and out_indices is not None:
  UpperCAmelCase__										=    [stage_names[idx] for idx in out_indices]
 return out_features, out_indices

def    _UpperCamelCase      (				SCREAMING_SNAKE_CASE__			:      Optional[List[str]]       ,						SCREAMING_SNAKE_CASE__			:      Optional[Union[List[int], Tuple[int]]]       ,						SCREAMING_SNAKE_CASE__			:      List[str]       ,						):
 '''simple docstring'''





 verify_out_features_out_indices(out_features=SCREAMING_SNAKE_CASE__       ,						out_indices=SCREAMING_SNAKE_CASE__       ,						stage_names=SCREAMING_SNAKE_CASE__       )
 UpperCAmelCase__							,			UpperCAmelCase__										=    _align_output_features_output_indices(
     out_features=SCREAMING_SNAKE_CASE__       ,						out_indices=SCREAMING_SNAKE_CASE__       ,						stage_names=SCREAMING_SNAKE_CASE__       )
 # Verify that the aligned out_features and out_indices are valid
 verify_out_features_out_indices(out_features=SCREAMING_SNAKE_CASE__       ,						out_indices=SCREAMING_SNAKE_CASE__       ,						stage_names=SCREAMING_SNAKE_CASE__       )
 return output_features, output_indices

class 						lowerCAmelCase_   :







  '''simple docstring'''







  lowerCAmelCase_				:					Optional[BackboneType]				     =						None


  def 							SCREAMING_SNAKE_CASE__			(						self    :	Optional[int]						,		_UpperCAmelCase    :	int ):






   """simple docstring"""






   if getattr(self						,		"""_backbone"""						,		_UpperCAmelCase ) is None:
    raise ValueError("""self._backbone must be set before calling _init_timm_backbone""" )

   # These will diagree with the defaults for the transformers models e.g. for resnet50
   # the transformer model has out_features = ['stem', 'stage1', 'stage2', 'stage3', 'stage4']
   # the timm model has out_features = ['act', 'layer1', 'layer2', 'layer3', 'layer4']
   UpperCAmelCase__										=    [stage["""module"""] for stage in self._backbone.feature_info.info]
   UpperCAmelCase__										=    [stage["""num_chs"""] for stage in self._backbone.feature_info.info]
   UpperCAmelCase__										=    self._backbone.feature_info.out_indices
   UpperCAmelCase__										=    self._backbone.feature_info.module_name()

   # We verify the out indices and out features are valid
   verify_out_features_out_indices(
       out_features=_UpperCAmelCase						,		out_indices=_UpperCAmelCase						,		stage_names=self.stage_names )
   UpperCAmelCase__							,			UpperCAmelCase__										=    out_features, out_indices


  def 							SCREAMING_SNAKE_CASE__			(						self    :	Tuple						,		_UpperCAmelCase    :	Dict ):






   """simple docstring"""






   UpperCAmelCase__										=    getattr(_UpperCAmelCase						,		"""stage_names""" )
   UpperCAmelCase__										=    getattr(_UpperCAmelCase						,		"""out_features"""						,		_UpperCAmelCase )
   UpperCAmelCase__										=    getattr(_UpperCAmelCase						,		"""out_indices"""						,		_UpperCAmelCase )

   UpperCAmelCase__										=    stage_names
   UpperCAmelCase__							,			UpperCAmelCase__										=    get_aligned_output_features_output_indices(
       out_features=_UpperCAmelCase						,		out_indices=_UpperCAmelCase						,		stage_names=_UpperCAmelCase )
   # Number of channels for each stage. This is set in the transformer backbone model init
   UpperCAmelCase__										=    None


  def 							SCREAMING_SNAKE_CASE__			(						self    :	List[Any]						,		_UpperCAmelCase    :	Tuple ):






   """simple docstring"""






   UpperCAmelCase__										=    config

   UpperCAmelCase__										=    getattr(_UpperCAmelCase						,		"""use_timm_backbone"""						,		_UpperCAmelCase )
   UpperCAmelCase__										=    BackboneType.TIMM if self.use_timm_backbone else BackboneType.TRANSFORMERS

   if self.backbone_type == BackboneType.TIMM:
    self._init_timm_backbone(_UpperCAmelCase )
   elif self.backbone_type == BackboneType.TRANSFORMERS:
    self._init_transformers_backbone(_UpperCAmelCase )
   else:
    raise ValueError(f'''backbone_type {self.backbone_type} not supported.''' )


  @property
  def 							SCREAMING_SNAKE_CASE__			(						self    :	Optional[int] ):






   """simple docstring"""






   return self._out_features


  @out_features.setter
  def 							SCREAMING_SNAKE_CASE__			(						self    :	Union[str, Any]						,		_UpperCAmelCase    :	List[str] ):






   """simple docstring"""






   UpperCAmelCase__							,			UpperCAmelCase__										=    get_aligned_output_features_output_indices(
       out_features=_UpperCAmelCase						,		out_indices=_UpperCAmelCase						,		stage_names=self.stage_names )


  @property
  def 							SCREAMING_SNAKE_CASE__			(						self    :	int ):






   """simple docstring"""






   return self._out_indices


  @out_indices.setter
  def 							SCREAMING_SNAKE_CASE__			(						self    :	List[Any]						,		_UpperCAmelCase    :	Union[Tuple[int], List[int]] ):






   """simple docstring"""






   UpperCAmelCase__							,			UpperCAmelCase__										=    get_aligned_output_features_output_indices(
       out_features=_UpperCAmelCase						,		out_indices=_UpperCAmelCase						,		stage_names=self.stage_names )


  @property
  def 							SCREAMING_SNAKE_CASE__			(						self    :	Dict ):






   """simple docstring"""






   return {stage: self.num_features[i] for i, stage in enumerate(self.stage_names )}


  @property
  def 							SCREAMING_SNAKE_CASE__			(						self    :	Optional[int] ):






   """simple docstring"""






   return [self.out_feature_channels[name] for name in self.out_features]


  def 							SCREAMING_SNAKE_CASE__			(						self    :	Optional[Any]						,		*_UpperCAmelCase    :	Optional[int]						,		**_UpperCAmelCase    :	int ):






   """simple docstring"""






   UpperCAmelCase__										=    dict(inspect.signature(self.forward ).parameters )
   UpperCAmelCase__										=    {k: v for k, v in kwargs.items() if k in signature}
   return self(*_UpperCAmelCase						,		**_UpperCAmelCase )


  def 							SCREAMING_SNAKE_CASE__			(						self    :	Optional[int]						,		_UpperCAmelCase    :	str						,		_UpperCAmelCase    :	Optional[bool] = None						,		_UpperCAmelCase    :	Optional[bool] = None						,		_UpperCAmelCase    :	Optional[bool] = None						,		):






   """simple docstring"""






   raise NotImplementedError("""This method should be implemented by the derived class.""" )



  def 							SCREAMING_SNAKE_CASE__			(						self    :	Any ):






   """simple docstring"""






   UpperCAmelCase__										=    super().to_dict()
   UpperCAmelCase__										=    output.pop("""_out_features""" )
   UpperCAmelCase__										=    output.pop("""_out_indices""" )
   return output

class 						lowerCAmelCase_   :







  '''simple docstring'''







  @property
  def 							SCREAMING_SNAKE_CASE__			(						self    :	List[str] ):






   """simple docstring"""






   return self._out_features


  @out_features.setter
  def 							SCREAMING_SNAKE_CASE__			(						self    :	Dict						,		_UpperCAmelCase    :	List[str] ):






   """simple docstring"""






   UpperCAmelCase__							,			UpperCAmelCase__										=    get_aligned_output_features_output_indices(
       out_features=_UpperCAmelCase						,		out_indices=_UpperCAmelCase						,		stage_names=self.stage_names )


  @property
  def 							SCREAMING_SNAKE_CASE__			(						self    :	Optional[int] ):






   """simple docstring"""






   return self._out_indices


  @out_indices.setter
  def 							SCREAMING_SNAKE_CASE__			(						self    :	Optional[Any]						,		_UpperCAmelCase    :	Union[Tuple[int], List[int]] ):






   """simple docstring"""






   UpperCAmelCase__							,			UpperCAmelCase__										=    get_aligned_output_features_output_indices(
       out_features=_UpperCAmelCase						,		out_indices=_UpperCAmelCase						,		stage_names=self.stage_names )



  def 							SCREAMING_SNAKE_CASE__			(						self    :	List[Any] ):






   """simple docstring"""






   UpperCAmelCase__										=    super().to_dict()
   UpperCAmelCase__										=    output.pop("""_out_features""" )
   UpperCAmelCase__										=    output.pop("""_out_indices""" )
   return output






'''simple docstring'''




import random
import unittest

import torch

from diffusers import IFImgaImgPipeline
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_VARIATION_BATCH_PARAMS,
    TEXT_GUIDED_IMAGE_VARIATION_PARAMS,
)
from ..test_pipelines_common import PipelineTesterMixin
from . import IFPipelineTesterMixin

@skip_mps
class 						lowerCAmelCase_   ( lowerCamelCase_			,  lowerCamelCase_			,  unittest.TestCase   ):







  '''simple docstring'''







  lowerCAmelCase_				:					Any				     =						IFImgaImgPipeline
  lowerCAmelCase_				:					Dict				     =						TEXT_GUIDED_IMAGE_VARIATION_PARAMS - {"""width""", """height"""}
  lowerCAmelCase_				:					Union[str, Any]				     =						TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS
  lowerCAmelCase_				:					Dict				     =						PipelineTesterMixin.required_optional_params - {"""latents"""}


  def 							SCREAMING_SNAKE_CASE__			(						self    :	Optional[int] ):






   """simple docstring"""






   return self._get_dummy_components()


  def 							SCREAMING_SNAKE_CASE__			(						self    :	Tuple						,		_UpperCAmelCase    :	Any						,		_UpperCAmelCase    :	Any=0 ):






   """simple docstring"""






   if str(_UpperCAmelCase ).startswith("""mps""" ):
    UpperCAmelCase__										=    torch.manual_seed(_UpperCAmelCase )
   else:
    UpperCAmelCase__										=    torch.Generator(device=_UpperCAmelCase ).manual_seed(_UpperCAmelCase )

   UpperCAmelCase__										=    floats_tensor((1, 3, 32, 32)						,		rng=random.Random(_UpperCAmelCase ) ).to(_UpperCAmelCase )

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

   return inputs


  def 							SCREAMING_SNAKE_CASE__			(						self    :	List[Any] ):






   """simple docstring"""






   self._test_save_load_optional_components()


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






   """simple docstring"""






   self._test_xformers_attention_forwardGenerator_pass(expected_max_diff=1E-3 )


  @unittest.skipIf(torch_device != """cuda"""						,		reason="""float16 requires CUDA""" )
  def 							SCREAMING_SNAKE_CASE__			(						self    :	int ):






   """simple docstring"""






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


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






   """simple docstring"""






   super().test_floataa_inference(expected_max_diff=1E-1 )


  def 							SCREAMING_SNAKE_CASE__			(						self    :	List[str] ):






   """simple docstring"""






   self._test_attention_slicing_forward_pass(expected_max_diff=1E-2 )


  def 							SCREAMING_SNAKE_CASE__			(						self    :	int ):






   """simple docstring"""






   self._test_save_load_local()



  def 							SCREAMING_SNAKE_CASE__			(						self    :	Any ):






   """simple docstring"""






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