






"""simple docstring"""


import gc
import random
import unittest

import numpy as np
import torch
from PIL import Image
from transformers import CLIPImageProcessor, CLIPVisionConfig, CLIPVisionModelWithProjection

from diffusers import (
    AutoencoderKL,
    DPMSolverMultistepScheduler,
    PNDMScheduler,
    StableDiffusionImageVariationPipeline,
    UNetaDConditionModel,
)
from diffusers.utils import floats_tensor, load_image, load_numpy, nightly, slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu

from ..pipeline_params import IMAGE_VARIATION_BATCH_PARAMS, IMAGE_VARIATION_PARAMS
from ..test_pipelines_common import PipelineKarrasSchedulerTesterMixin, PipelineLatentTesterMixin, PipelineTesterMixin


enable_full_determinism()






class 		snake_case    (    UpperCAmelCase  ,					UpperCAmelCase  ,					UpperCAmelCase  ,					unittest.TestCase			):
    __magic_name__       	=	StableDiffusionImageVariationPipeline
    __magic_name__       	=	IMAGE_VARIATION_PARAMS
    __magic_name__       	=	IMAGE_VARIATION_BATCH_PARAMS
    __magic_name__       	=	frozenset([]			)
    # TO-DO: update image_params once pipeline is refactored with VaeImageProcessor.preprocess
    __magic_name__       	=	frozenset([]			)





    def 							lowerCamelCase__ (				self      :   Optional[Any]     ):


      '''simple docstring'''




      torch.manual_seed(0     )
      a	:     List[str]							=  UNetaDConditionModel(
          block_out_channels=(3_2, 6_4)							,						layers_per_block=2							,						sample_size=3_2							,						in_channels=4							,						out_channels=4							,						down_block_types=('DownBlock2D', 'CrossAttnDownBlock2D')							,						up_block_types=('CrossAttnUpBlock2D', 'UpBlock2D')							,						cross_attention_dim=3_2							,						)
      a	:     Union[str, Any]							=  PNDMScheduler(skip_prk_steps=A     )
      torch.manual_seed(0     )
      a	:     int							=  AutoencoderKL(
          block_out_channels=[3_2, 6_4]							,						in_channels=3							,						out_channels=3							,						down_block_types=['DownEncoderBlock2D', 'DownEncoderBlock2D']							,						up_block_types=['UpDecoderBlock2D', 'UpDecoderBlock2D']							,						latent_channels=4							,						)
      torch.manual_seed(0     )
      a	:     Tuple							=  CLIPVisionConfig(
          hidden_size=3_2							,						projection_dim=3_2							,						intermediate_size=3_7							,						layer_norm_eps=1E-05							,						num_attention_heads=4							,						num_hidden_layers=5							,						image_size=3_2							,						patch_size=4							,						)
      a	:     Any							=  CLIPVisionModelWithProjection(A     )
      a	:     Union[str, Any]							=  CLIPImageProcessor(crop_size=3_2							,						size=3_2     )

      a	:     Dict							=  {
          'unet': unet,
          'scheduler': scheduler,
          'vae': vae,
          'image_encoder': image_encoder,
          'feature_extractor': feature_extractor,
          'safety_checker': None,
      }
      return components





    def 							lowerCamelCase__ (				self      :   int							,						A      :   int							,						A      :   Union[str, Any]=0     ):


      '''simple docstring'''




      a	:     str							=  floats_tensor((1, 3, 3_2, 3_2)							,						rng=random.Random(A     )     )
      a	:     List[Any]							=  image.cpu().permute(0							,						2							,						3							,						1     )[0]
      a	:     Any							=  Image.fromarray(np.uinta(A     )     ).convert('RGB'     ).resize((3_2, 3_2)     )
      if str(A     ).startswith('mps'     ):
        a	:     int							=  torch.manual_seed(A     )
      else:
        a	:     Optional[Any]							=  torch.Generator(device=A     ).manual_seed(A     )
      a	:     str							=  {
          'image': image,
          'generator': generator,
          'num_inference_steps': 2,
          'guidance_scale': 6.0,
          'output_type': 'numpy',
      }
      return inputs





    def 							lowerCamelCase__ (				self      :   str     ):


      '''simple docstring'''




      a	:     str							=  'cpu'  # ensure determinism for the device-dependent torch.Generator
      a	:     List[Any]							=  self.get_dummy_components()
      a	:     Tuple							=  StableDiffusionImageVariationPipeline(**A     )
      a	:     List[str]							=  sd_pipe.to(A     )
      sd_pipe.set_progress_bar_config(disable=A     )

      a	:     List[Any]							=  self.get_dummy_inputs(A     )
      a	:     Optional[Any]							=  sd_pipe(**A     ).images
      a	:     Union[str, Any]							=  image[0, -3:, -3:, -1]

      assert image.shape == (1, 6_4, 6_4, 3)
      a	:     Optional[Any]							=  np.array([0.52_39, 0.57_23, 0.47_96, 0.50_49, 0.55_50, 0.46_85, 0.53_29, 0.48_91, 0.49_21]     )

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





    def 							lowerCamelCase__ (				self      :   List[Any]     ):


      '''simple docstring'''




      a	:     Optional[Any]							=  'cpu'  # ensure determinism for the device-dependent torch.Generator
      a	:     Union[str, Any]							=  self.get_dummy_components()
      a	:     List[str]							=  StableDiffusionImageVariationPipeline(**A     )
      a	:     Dict							=  sd_pipe.to(A     )
      sd_pipe.set_progress_bar_config(disable=A     )

      a	:     Tuple							=  self.get_dummy_inputs(A     )
      a	:     Dict							=  2 * [inputs['image']]
      a	:     Dict							=  sd_pipe(**A     )

      a	:     int							=  output.images

      a	:     Union[str, Any]							=  image[-1, -3:, -3:, -1]

      assert image.shape == (2, 6_4, 6_4, 3)
      a	:     Tuple							=  np.array([0.68_92, 0.56_37, 0.58_36, 0.57_71, 0.62_54, 0.64_09, 0.55_80, 0.55_69, 0.52_89]     )

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





    def 							lowerCamelCase__ (				self      :   Union[str, Any]     ):


      '''simple docstring'''




      super().test_inference_batch_single_identical(expected_max_diff=3E-3     )






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





    def 							lowerCamelCase__ (				self      :   Tuple     ):


      '''simple docstring'''




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





    def 							lowerCamelCase__ (				self      :   int							,						A      :   Optional[Any]							,						A      :   List[str]="cpu"							,						A      :   Union[str, Any]=torch.floataa							,						A      :   Tuple=0     ):


      '''simple docstring'''




      a	:     Tuple							=  torch.Generator(device=A     ).manual_seed(A     )
      a	:     Tuple							=  load_image(
          'https://huggingface.co/datasets/diffusers/test-arrays/resolve/main'
          '/stable_diffusion_imgvar/input_image_vermeer.png'     )
      a	:     Tuple							=  np.random.RandomState(A     ).standard_normal((1, 4, 6_4, 6_4)     )
      a	:     List[str]							=  torch.from_numpy(A     ).to(device=A							,						dtype=A     )
      a	:     Tuple							=  {
          'image': init_image,
          'latents': latents,
          'generator': generator,
          'num_inference_steps': 3,
          'guidance_scale': 7.5,
          'output_type': 'numpy',
      }
      return inputs





    def 							lowerCamelCase__ (				self      :   Any     ):


      '''simple docstring'''




      a	:     Any							=  StableDiffusionImageVariationPipeline.from_pretrained(
          'lambdalabs/sd-image-variations-diffusers'							,						safety_checker=A     )
      a	:     str							=  sd_pipe.to(A     )
      sd_pipe.set_progress_bar_config(disable=A     )

      a	:     Optional[int]							=  self.get_inputs(A     )
      a	:     int							=  sd_pipe(**A     ).images
      a	:     Dict							=  image[0, -3:, -3:, -1].flatten()

      assert image.shape == (1, 5_1_2, 5_1_2, 3)
      a	:     int							=  np.array([0.8_44_91, 0.9_07_89, 0.7_57_08, 0.7_87_34, 0.8_34_85, 0.7_00_99, 0.6_69_38, 0.6_87_27, 0.6_13_79]     )
      assert np.abs(image_slice - expected_slice     ).max() < 6E-3





    def 							lowerCamelCase__ (				self      :   str     ):


      '''simple docstring'''




      a	:     int							=  0

      def callback_fn(A      :   int							,						A      :   int							,						A      :   torch.FloatTensor     ) -> None:
        a	:     str							=  True
        nonlocal number_of_steps
        number_of_steps += 1
        if step == 1:
          a	:     Union[str, Any]							=  latents.detach().cpu().numpy()
          assert latents.shape == (1, 4, 6_4, 6_4)
          a	:     Union[str, Any]							=  latents[0, -3:, -3:, -1]
          a	:     List[Any]							=  np.array(
              [-0.16_21, 0.28_37, -0.79_79, -0.12_21, -1.30_57, 0.76_81, -2.11_91, 0.04_64, 1.63_09]     )

          assert np.abs(latents_slice.flatten() - expected_slice     ).max() < 5E-2
        elif step == 2:
          a	:     int							=  latents.detach().cpu().numpy()
          assert latents.shape == (1, 4, 6_4, 6_4)
          a	:     Optional[int]							=  latents[0, -3:, -3:, -1]
          a	:     List[Any]							=  np.array([0.62_99, 1.75_00, 1.19_92, -2.15_82, -1.89_94, 0.73_34, -0.70_90, 1.01_37, 1.52_73]     )

          assert np.abs(latents_slice.flatten() - expected_slice     ).max() < 5E-2

      a	:     Optional[int]							=  False

      a	:     Dict							=  StableDiffusionImageVariationPipeline.from_pretrained(
          'fusing/sd-image-variations-diffusers'							,						safety_checker=A							,						torch_dtype=torch.floataa							,						)
      pipe.to(A     )
      pipe.set_progress_bar_config(disable=A     )
      pipe.enable_attention_slicing()

      a	:     Optional[int]							=  self.get_inputs(A							,						dtype=torch.floataa     )
      pipe(**A							,						callback=A							,						callback_steps=1     )
      assert callback_fn.has_been_called
      assert number_of_steps == inputs["num_inference_steps"]





    def 							lowerCamelCase__ (				self      :   str     ):


      '''simple docstring'''




      torch.cuda.empty_cache()
      torch.cuda.reset_max_memory_allocated()
      torch.cuda.reset_peak_memory_stats()

      a	:     List[str]							=  'fusing/sd-image-variations-diffusers'
      a	:     List[Any]							=  StableDiffusionImageVariationPipeline.from_pretrained(
          A							,						safety_checker=A							,						torch_dtype=torch.floataa     )
      a	:     Union[str, Any]							=  pipe.to(A     )
      pipe.set_progress_bar_config(disable=A     )
      pipe.enable_attention_slicing(1     )
      pipe.enable_sequential_cpu_offload()

      a	:     Dict							=  self.get_inputs(A							,						dtype=torch.floataa     )
      a	:     str							=  pipe(**A     )

      a	:     int							=  torch.cuda.max_memory_allocated()
      # make sure that less than 2.6 GB is allocated
      assert mem_bytes < 2.6 * 1_0**9






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





    def 							lowerCamelCase__ (				self      :   List[Any]     ):


      '''simple docstring'''




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





    def 							lowerCamelCase__ (				self      :   Dict							,						A      :   str							,						A      :   List[Any]="cpu"							,						A      :   List[Any]=torch.floataa							,						A      :   Optional[int]=0     ):


      '''simple docstring'''




      a	:     List[Any]							=  torch.Generator(device=A     ).manual_seed(A     )
      a	:     List[Any]							=  load_image(
          'https://huggingface.co/datasets/diffusers/test-arrays/resolve/main'
          '/stable_diffusion_imgvar/input_image_vermeer.png'     )
      a	:     Dict							=  np.random.RandomState(A     ).standard_normal((1, 4, 6_4, 6_4)     )
      a	:     Optional[Any]							=  torch.from_numpy(A     ).to(device=A							,						dtype=A     )
      a	:     Optional[int]							=  {
          'image': init_image,
          'latents': latents,
          'generator': generator,
          'num_inference_steps': 5_0,
          'guidance_scale': 7.5,
          'output_type': 'numpy',
      }
      return inputs





    def 							lowerCamelCase__ (				self      :   List[Any]     ):


      '''simple docstring'''




      a	:     Any							=  StableDiffusionImageVariationPipeline.from_pretrained('fusing/sd-image-variations-diffusers'     )
      sd_pipe.to(A     )
      sd_pipe.set_progress_bar_config(disable=A     )

      a	:     str							=  self.get_inputs(A     )
      a	:     int							=  sd_pipe(**A     ).images[0]

      a	:     Optional[int]							=  load_numpy(
          'https://huggingface.co/datasets/diffusers/test-arrays/resolve/main'
          '/stable_diffusion_imgvar/lambdalabs_variations_pndm.npy'     )
      a	:     Dict							=  np.abs(expected_image - image     ).max()
      assert max_diff < 1E-3





    def 							lowerCamelCase__ (				self      :   str     ):


      '''simple docstring'''




      a	:     str							=  StableDiffusionImageVariationPipeline.from_pretrained('fusing/sd-image-variations-diffusers'     )
      a	:     Any							=  DPMSolverMultistepScheduler.from_config(sd_pipe.scheduler.config     )
      sd_pipe.to(A     )
      sd_pipe.set_progress_bar_config(disable=A     )

      a	:     Optional[Any]							=  self.get_inputs(A     )
      a	:     Any							=  2_5
      a	:     List[str]							=  sd_pipe(**A     ).images[0]

      a	:     Optional[Any]							=  load_numpy(
          'https://huggingface.co/datasets/diffusers/test-arrays/resolve/main'
          '/stable_diffusion_imgvar/lambdalabs_variations_dpm_multi.npy'     )
      a	:     str							=  np.abs(expected_image - image     ).max()
      assert max_diff < 1E-3












"""simple docstring"""


# 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


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




def      snake_case    (A_ :str ):







  '''simple docstring'''

  a	:     int							=  args.config_file
  if config_file is None and Path(A_ ).exists():
    a	:     List[Any]							=  default_config_file
  elif not Path(A_ ).exists():
    raise ValueError(f'''The passed config file located at {config_file} doesn\'t exist.''' )
  a	:     Dict							=  load_config_from_file(A_ )

  if config_file.endswith('.json' ):
    config.to_json_file(A_ )
  else:
    config.to_yaml_file(A_ )
  return config_file




def      snake_case    (A_ :List[str]      ,     A_ :Union[str, Any] ):







  '''simple docstring'''

  a	:     int							=  parser.add_parser('update'      ,     parents=A_      ,     help=A_      ,     formatter_class=A_ )
  parser.add_argument(
      '--config_file'      ,     default=A_      ,     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=A_ )
  return parser






def      snake_case    (A_ :List[str] ):







  '''simple docstring'''

  a	:     List[str]							=  update_config(A_ )
  print(f'''Sucessfully updated the configuration file at {config_file}.''' )




