

import gc
import unittest

import numpy as np
import torch
from transformers import CLIPTextConfig, CLIPTextModel, CLIPTokenizer

from diffusers import (
    AutoencoderKL,
    DDIMParallelScheduler,
    DDPMParallelScheduler,
    StableDiffusionParadigmsPipeline,
    UNetaDConditionModel,
)
from diffusers.utils import slow, torch_device
from diffusers.utils.testing_utils import (
    enable_full_determinism,
    require_torch_gpu,
)

from ..pipeline_params import TEXT_TO_IMAGE_BATCH_PARAMS, TEXT_TO_IMAGE_IMAGE_PARAMS, TEXT_TO_IMAGE_PARAMS
from ..test_pipelines_common import PipelineLatentTesterMixin, PipelineTesterMixin


enable_full_determinism()





class     snake_case  (	_snake_case ,  _snake_case ,  unittest.TestCase							):


   '''simple docstring'''







   UpperCamelCase__	:       Union[str, Any]     =      StableDiffusionParadigmsPipeline
   UpperCamelCase__	:       int     =      TEXT_TO_IMAGE_PARAMS
   UpperCamelCase__	:       List[str]     =      TEXT_TO_IMAGE_BATCH_PARAMS
   UpperCamelCase__	:       int     =      TEXT_TO_IMAGE_IMAGE_PARAMS
   UpperCamelCase__	:       Any     =      TEXT_TO_IMAGE_IMAGE_PARAMS




   def 	UpperCAmelCase       (					self     :	Optional[Any]			)					->Tuple:
     '''simple docstring'''







     torch.manual_seed(0			)
     UpperCAmelCase__    						=   UNetaDConditionModel(
         block_out_channels=(32, 64)						,    layers_per_block=2						,    sample_size=32						,    in_channels=4						,    out_channels=4						,    down_block_types=("""DownBlock2D""", """CrossAttnDownBlock2D""")						,    up_block_types=("""CrossAttnUpBlock2D""", """UpBlock2D""")						,    cross_attention_dim=32						,    attention_head_dim=(2, 4)						,    use_linear_projection=lowerCamelCase_						,    )
     UpperCAmelCase__    						=   DDIMParallelScheduler(
         beta_start=0.00085						,    beta_end=0.012						,    beta_schedule="""scaled_linear"""						,    clip_sample=lowerCamelCase_						,    set_alpha_to_one=lowerCamelCase_						,    )
     torch.manual_seed(0			)
     UpperCAmelCase__    						=   AutoencoderKL(
         block_out_channels=[32, 64]						,    in_channels=3						,    out_channels=3						,    down_block_types=["""DownEncoderBlock2D""", """DownEncoderBlock2D"""]						,    up_block_types=["""UpDecoderBlock2D""", """UpDecoderBlock2D"""]						,    latent_channels=4						,    sample_size=128						,    )
     torch.manual_seed(0			)
     UpperCAmelCase__    						=   CLIPTextConfig(
         bos_token_id=0						,    eos_token_id=2						,    hidden_size=32						,    intermediate_size=37						,    layer_norm_eps=1E-05						,    num_attention_heads=4						,    num_hidden_layers=5						,    pad_token_id=1						,    vocab_size=1000						,    hidden_act="""gelu"""						,    projection_dim=512						,    )
     UpperCAmelCase__    						=   CLIPTextModel(lowerCamelCase_			)
     UpperCAmelCase__    						=   CLIPTokenizer.from_pretrained("""hf-internal-testing/tiny-random-clip"""			)

     UpperCAmelCase__    						=   {
         """unet""": unet,
         """scheduler""": scheduler,
         """vae""": vae,
         """text_encoder""": text_encoder,
         """tokenizer""": tokenizer,
         """safety_checker""": None,
         """feature_extractor""": None,
     }
     return components




   def 	UpperCAmelCase       (					self     :	Tuple						,    lowerCamelCase_     :	Optional[int]						,    lowerCamelCase_     :	Union[str, Any]=0			)					->Any:
     '''simple docstring'''







     if str(lowerCamelCase_			).startswith("""mps"""			):
       UpperCAmelCase__    						=   torch.manual_seed(lowerCamelCase_			)
     else:
       UpperCAmelCase__    						=   torch.Generator(device=lowerCamelCase_			).manual_seed(lowerCamelCase_			)
     UpperCAmelCase__    						=   {
         """prompt""": """a photograph of an astronaut riding a horse""",
         """generator""": generator,
         """num_inference_steps""": 10,
         """guidance_scale""": 6.0,
         """output_type""": """numpy""",
         """parallel""": 3,
         """debug""": True,
     }
     return inputs




   def 	UpperCAmelCase       (					self     :	Optional[int]			)					->Dict:
     '''simple docstring'''







     UpperCAmelCase__    						=   """cpu"""  # ensure determinism for the device-dependent torch.Generator
     UpperCAmelCase__    						=   self.get_dummy_components()
     UpperCAmelCase__    						=   StableDiffusionParadigmsPipeline(**lowerCamelCase_			)
     UpperCAmelCase__    						=   sd_pipe.to(lowerCamelCase_			)
     sd_pipe.set_progress_bar_config(disable=lowerCamelCase_			)

     UpperCAmelCase__    						=   self.get_dummy_inputs(lowerCamelCase_			)
     UpperCAmelCase__    						=   sd_pipe(**lowerCamelCase_			).images
     UpperCAmelCase__    						=   image[0, -3:, -3:, -1]
     assert image.shape == (1, 64, 64, 3)

     UpperCAmelCase__    						=   np.array([0.4773, 0.5417, 0.4723, 0.4925, 0.5631, 0.4752, 0.5240, 0.4935, 0.5023]			)

     assert np.abs(image_slice.flatten() - expected_slice			).max() < 1E-2




   def 	UpperCAmelCase       (					self     :	List[str]			)					->List[Any]:
     '''simple docstring'''







     UpperCAmelCase__    						=   """cpu"""  # ensure determinism for the device-dependent torch.Generator
     UpperCAmelCase__    						=   self.get_dummy_components()
     torch.manual_seed(0			)
     UpperCAmelCase__    						=   DDPMParallelScheduler()
     torch.manual_seed(0			)
     UpperCAmelCase__    						=   StableDiffusionParadigmsPipeline(**lowerCamelCase_			)
     UpperCAmelCase__    						=   sd_pipe.to(lowerCamelCase_			)
     sd_pipe.set_progress_bar_config(disable=lowerCamelCase_			)

     UpperCAmelCase__    						=   self.get_dummy_inputs(lowerCamelCase_			)
     UpperCAmelCase__    						=   sd_pipe(**lowerCamelCase_			).images
     UpperCAmelCase__    						=   image[0, -3:, -3:, -1]
     assert image.shape == (1, 64, 64, 3)

     UpperCAmelCase__    						=   np.array([0.3573, 0.4420, 0.4960, 0.4799, 0.3796, 0.3879, 0.4819, 0.4365, 0.4468]			)

     assert np.abs(image_slice.flatten() - expected_slice			).max() < 1E-2




   def 	UpperCAmelCase       (					self     :	str			)					->str:
     '''simple docstring'''







     super().test_inference_batch_consistent(batch_sizes=[1, 2]			)




   def 	UpperCAmelCase       (					self     :	List[str]			)					->Optional[Any]:
     '''simple docstring'''







     super().test_inference_batch_single_identical(batch_size=2						,    expected_max_diff=3E-3			)







   def 	UpperCAmelCase       (					self     :	Union[str, Any]			)					->List[Any]:
     '''simple docstring'''







     UpperCAmelCase__    						=   """cpu"""  # ensure determinism for the device-dependent torch.Generator
     UpperCAmelCase__    						=   self.get_dummy_components()
     UpperCAmelCase__    						=   StableDiffusionParadigmsPipeline(**lowerCamelCase_			)
     UpperCAmelCase__    						=   sd_pipe.to(lowerCamelCase_			)
     sd_pipe.set_progress_bar_config(disable=lowerCamelCase_			)

     UpperCAmelCase__    						=   self.get_dummy_inputs(lowerCamelCase_			)
     UpperCAmelCase__    						=   """french fries"""
     UpperCAmelCase__    						=   sd_pipe(**lowerCamelCase_						,    negative_prompt=lowerCamelCase_			)
     UpperCAmelCase__    						=   output.images
     UpperCAmelCase__    						=   image[0, -3:, -3:, -1]

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

     UpperCAmelCase__    						=   np.array([0.4771, 0.5420, 0.4683, 0.4918, 0.5636, 0.4725, 0.5230, 0.4923, 0.5015]			)

     assert np.abs(image_slice.flatten() - expected_slice			).max() < 1E-2





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


   '''simple docstring'''







   def 	UpperCAmelCase       (					self     :	Tuple			)					->Any:
     '''simple docstring'''







     super().tearDown()
     gc.collect()
     torch.cuda.empty_cache()




   def 	UpperCAmelCase       (					self     :	List[str]						,    lowerCamelCase_     :	Dict=0			)					->int:
     '''simple docstring'''







     UpperCAmelCase__    						=   torch.Generator(device=lowerCamelCase_			).manual_seed(lowerCamelCase_			)
     UpperCAmelCase__    						=   {
         """prompt""": """a photograph of an astronaut riding a horse""",
         """generator""": generator,
         """num_inference_steps""": 10,
         """guidance_scale""": 7.5,
         """output_type""": """numpy""",
         """parallel""": 3,
         """debug""": True,
     }
     return inputs







   def 	UpperCAmelCase       (					self     :	List[Any]			)					->Any:
     '''simple docstring'''







     UpperCAmelCase__    						=   """stabilityai/stable-diffusion-2-base"""
     UpperCAmelCase__    						=   DDIMParallelScheduler.from_pretrained(lowerCamelCase_						,    subfolder="""scheduler"""			)
     UpperCAmelCase__    						=   StableDiffusionParadigmsPipeline.from_pretrained(lowerCamelCase_						,    scheduler=lowerCamelCase_						,    safety_checker=lowerCamelCase_			)
     pipe.to(lowerCamelCase_			)
     pipe.set_progress_bar_config(disable=lowerCamelCase_			)
     pipe.enable_attention_slicing()

     UpperCAmelCase__    						=   self.get_inputs()
     UpperCAmelCase__    						=   pipe(**lowerCamelCase_			).images
     UpperCAmelCase__    						=   image[0, -3:, -3:, -1].flatten()

     assert image.shape == (1, 512, 512, 3)

     UpperCAmelCase__    						=   np.array([0.9622, 0.9602, 0.9748, 0.9591, 0.9630, 0.9691, 0.9661, 0.9631, 0.9741]			)

     assert np.abs(expected_slice - image_slice			).max() < 1E-2










# Copyright 2022 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.

from pathlib import Path

from .config_args import default_config_file, load_config_from_file
from .config_utils import SubcommandHelpFormatter


snake_case__ :						str     =			'''Update an existing config file with the latest defaults while maintaining the old configuration.'''


def 						lowercase		(     _lowerCAmelCase      ):
  UpperCAmelCase__    						=   args.config_file
  if config_file is None and Path(_lowerCAmelCase      ).exists():
    UpperCAmelCase__    						=   default_config_file
  elif not Path(_lowerCAmelCase      ).exists():
    raise ValueError(F'''The passed config file located at {config_file} doesn\'t exist.'''      )
  UpperCAmelCase__    						=   load_config_from_file(_lowerCAmelCase      )

  if config_file.endswith(""".json"""      ):
    config.to_json_file(_lowerCAmelCase      )
  else:
    config.to_yaml_file(_lowerCAmelCase      )
  return config_file


def 						lowercase		(     _lowerCAmelCase   ,     _lowerCAmelCase      ):
  UpperCAmelCase__    						=   parser.add_parser("""update"""   ,     parents=_lowerCAmelCase   ,     help=_lowerCAmelCase   ,     formatter_class=_lowerCAmelCase      )
  parser.add_argument(
      """--config_file"""   ,     default=_lowerCAmelCase   ,     help=(
          """The path to the config file to update. Will default to a file named default_config.yaml in the cache """
          """location, which is the content of the environment `HF_HOME` suffixed with 'accelerate', or if you don't have """
          """such an environment variable, your cache directory ('~/.cache' or the content of `XDG_CACHE_HOME`) suffixed """
          """with 'huggingface'."""
      )   ,     )

  parser.set_defaults(func=_lowerCAmelCase      )
  return parser





def 						lowercase		(     _lowerCAmelCase      ):
  UpperCAmelCase__    						=   update_config(_lowerCAmelCase      )
  print(F'''Sucessfully updated the configuration file at {config_file}.'''      )




