


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    UpperCAmelCase__	:



    """simple docstring"""






    def _a (      self   )  ->	Any:
         torch.manual_seed(0   )
         __UpperCamelCase   							=TaEncoderModel.from_pretrained('hf-internal-testing/tiny-random-t5'   )

         torch.manual_seed(0   )
         __UpperCamelCase   							=AutoTokenizer.from_pretrained('hf-internal-testing/tiny-random-t5'   )

         torch.manual_seed(0   )
         __UpperCamelCase   							=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   )
         __UpperCamelCase   							=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   )
         __UpperCamelCase   							=IFWatermarker()

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






    def _a (      self   )  ->	Union[str, Any]:
         torch.manual_seed(0   )
         __UpperCamelCase   							=TaEncoderModel.from_pretrained('hf-internal-testing/tiny-random-t5'   )

         torch.manual_seed(0   )
         __UpperCamelCase   							=AutoTokenizer.from_pretrained('hf-internal-testing/tiny-random-t5'   )

         torch.manual_seed(0   )
         __UpperCamelCase   							=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   )
         __UpperCamelCase   							=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   )
         __UpperCamelCase   							=DDPMScheduler(
             num_train_timesteps=1000 ,       beta_schedule='squaredcos_cap_v2' ,       beta_start=0.0001 ,       beta_end=0.02 ,       )

         torch.manual_seed(0   )
         __UpperCamelCase   							=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   )  ->	Tuple:
         __UpperCamelCase   							=self.get_dummy_components()
         __UpperCamelCase   							=self.pipeline_class(**A_   )
         pipe.to(A_   )
         pipe.set_progress_bar_config(disable=A_   )

         __UpperCamelCase   							=self.get_dummy_inputs(A_   )

         __UpperCamelCase   							=inputs['prompt']
         __UpperCamelCase   							=inputs['generator']
         __UpperCamelCase   							=inputs['num_inference_steps']
         __UpperCamelCase   							=inputs['output_type']

         if "image" in inputs:
              __UpperCamelCase   							=inputs['image']
         else:
              __UpperCamelCase   							=None

         if "mask_image" in inputs:
              __UpperCamelCase   							=inputs['mask_image']
         else:
              __UpperCamelCase   							=None

         if "original_image" in inputs:
              __UpperCamelCase   							=inputs['original_image']
         else:
              __UpperCamelCase   							=None

         __UpperCamelCase							,				__UpperCamelCase   							=pipe.encode_prompt(A_   )

         # inputs with prompt converted to embeddings
         __UpperCamelCase   							={
             '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:
              __UpperCamelCase   							=image

         if mask_image is not None:
              __UpperCamelCase   							=mask_image

         if original_image is not None:
              __UpperCamelCase   							=original_image

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

         __UpperCamelCase   							=pipe(**A_   )[0]

         with tempfile.TemporaryDirectory() as tmpdir:
              pipe.save_pretrained(A_   )
              __UpperCamelCase   							=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.' ,       )

         __UpperCamelCase   							=self.get_dummy_inputs(A_   )

         __UpperCamelCase   							=inputs['generator']
         __UpperCamelCase   							=inputs['num_inference_steps']
         __UpperCamelCase   							=inputs['output_type']

         # inputs with prompt converted to embeddings
         __UpperCamelCase   							={
             '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:
              __UpperCamelCase   							=image

         if mask_image is not None:
              __UpperCamelCase   							=mask_image

         if original_image is not None:
              __UpperCamelCase   							=original_image

         __UpperCamelCase   							=pipe_loaded(**A_   )[0]

         __UpperCamelCase   							=np.abs(to_np(A_   ) - to_np(A_   )   ).max()
         self.assertLess(A_ ,       1E-4   )






    def _a (      self   )  ->	int:
         __UpperCamelCase   							=self.get_dummy_components()
         __UpperCamelCase   							=self.pipeline_class(**A_   )
         pipe.to(A_   )
         pipe.set_progress_bar_config(disable=A_   )

         __UpperCamelCase   							=self.get_dummy_inputs(A_   )
         __UpperCamelCase   							=pipe(**A_   )[0]

         with tempfile.TemporaryDirectory() as tmpdir:
              pipe.save_pretrained(A_   )
              __UpperCamelCase   							=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

         __UpperCamelCase   							=self.get_dummy_inputs(A_   )
         __UpperCamelCase   							=pipe_loaded(**A_   )[0]

         __UpperCamelCase   							=np.abs(to_np(A_   ) - to_np(A_   )   ).max()
         self.assertLess(A_ ,       1E-4   )









# Copyright 2021 The HuggingFace Team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import argparse

from ...utils.dataclasses import (
    ComputeEnvironment,
    DistributedType,
    DynamoBackend,
    PrecisionType,
    SageMakerDistributedType,
)
from ..menu import BulletMenu


_A			=						[
    'EAGER',
    'AOT_EAGER',
    'INDUCTOR',
    'NVFUSER',
    'AOT_NVFUSER',
    'AOT_CUDAGRAPHS',
    'OFI',
    'FX2TRT',
    'ONNXRT',
    'IPEX',
]

def 						_UpperCAmelCase    (      SCREAMING_SNAKE_CASE__      : Optional[int]      ,					SCREAMING_SNAKE_CASE__      : Optional[Any]=None      ,					SCREAMING_SNAKE_CASE__      : Optional[Any]=None      ,					SCREAMING_SNAKE_CASE__      : str=None    ):
     __UpperCamelCase   							=True
     while ask_again:
          __UpperCamelCase   							=input(SCREAMING_SNAKE_CASE__    )
          try:
               if default is not None and len(SCREAMING_SNAKE_CASE__    ) == 0:
                    return default
               return convert_value(SCREAMING_SNAKE_CASE__    ) if convert_value is not None else result
          except Exception:
               if error_message is not None:
                    print(SCREAMING_SNAKE_CASE__    )

def 						_UpperCAmelCase    (      SCREAMING_SNAKE_CASE__      : List[Any]      ,					SCREAMING_SNAKE_CASE__      : List[str]=[]      ,					SCREAMING_SNAKE_CASE__      : Union[str, Any]=None      ,					SCREAMING_SNAKE_CASE__      : str=0    ):
     __UpperCamelCase   							=BulletMenu(SCREAMING_SNAKE_CASE__      ,					SCREAMING_SNAKE_CASE__    )
     __UpperCamelCase   							=menu.run(default_choice=SCREAMING_SNAKE_CASE__    )
     return convert_value(SCREAMING_SNAKE_CASE__    ) if convert_value is not None else result

def 						_UpperCAmelCase    (      SCREAMING_SNAKE_CASE__      : Union[str, Any]    ):
     __UpperCamelCase   							=int(SCREAMING_SNAKE_CASE__    )
     return ComputeEnvironment(['LOCAL_MACHINE', 'AMAZON_SAGEMAKER'][value]    )

def 						_UpperCAmelCase    (      SCREAMING_SNAKE_CASE__      : List[str]    ):
     __UpperCamelCase   							=int(SCREAMING_SNAKE_CASE__    )
     return DistributedType(['NO', 'MULTI_CPU', 'MULTI_XPU', 'MULTI_GPU', 'MULTI_NPU', 'TPU'][value]    )

def 						_UpperCAmelCase    (      SCREAMING_SNAKE_CASE__      : Optional[int]    ):
     __UpperCamelCase   							=int(SCREAMING_SNAKE_CASE__    )
     return DynamoBackend(DYNAMO_BACKENDS[value]    ).value

def 						_UpperCAmelCase    (      SCREAMING_SNAKE_CASE__      : Dict    ):
     __UpperCamelCase   							=int(SCREAMING_SNAKE_CASE__    )
     return PrecisionType(['no', 'fp16', 'bf16', 'fp8'][value]    )

def 						_UpperCAmelCase    (      SCREAMING_SNAKE_CASE__      : Optional[Any]    ):
     __UpperCamelCase   							=int(SCREAMING_SNAKE_CASE__    )
     return SageMakerDistributedType(['NO', 'DATA_PARALLEL', 'MODEL_PARALLEL'][value]    )

def 						_UpperCAmelCase    (      SCREAMING_SNAKE_CASE__      : Any    ):
     return {"yes": True, "no": False}[value.lower()]






class    UpperCAmelCase__	( argparse.RawDescriptionHelpFormatter		):



    """simple docstring"""






    def _a (      self ,       A_ ,       A_ ,       A_ ,       A_   )  ->	Optional[int]:
         __UpperCamelCase   							=super()._format_usage(A_ ,       A_ ,       A_ ,       A_   )
         __UpperCamelCase   							=usage.replace('<command> [<args>] ' ,       ''   )
         return usage





