
import gc
import unittest

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

from diffusers import (
    AutoencoderKL,
    DDIMScheduler,
    EulerAncestralDiscreteScheduler,
    LMSDiscreteScheduler,
    PNDMScheduler,
    StableDiffusionPanoramaPipeline,
    UNetaDConditionModel,
)
from diffusers.utils import slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu, skip_mps

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()





@skip_mps
class  _snake_case   (	a_							,		a_							,		unittest.TestCase							):
			UpperCamelCase__        =       StableDiffusionPanoramaPipeline
			UpperCamelCase__        =       TEXT_TO_IMAGE_PARAMS
			UpperCamelCase__        =       TEXT_TO_IMAGE_BATCH_PARAMS
			UpperCamelCase__        =       TEXT_TO_IMAGE_IMAGE_PARAMS
			UpperCamelCase__        =       TEXT_TO_IMAGE_IMAGE_PARAMS







			def      SCREAMING_SNAKE_CASE       (							self						):
				torch.manual_seed(0						)
				__magic_name__						:						List[Any]        =	UNetaDConditionModel(
				    block_out_channels=(32, 64)		,  layers_per_block=1		,  sample_size=32		,  in_channels=4		,  out_channels=4		,  down_block_types=("DownBlock2D", "CrossAttnDownBlock2D")		,  up_block_types=("CrossAttnUpBlock2D", "UpBlock2D")		,  cross_attention_dim=32		,  )
				__magic_name__						:						Dict        =	DDIMScheduler()
				torch.manual_seed(0						)
				__magic_name__						:						List[Any]        =	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		,  )
				torch.manual_seed(0						)
				__magic_name__						:						Optional[Any]        =	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=1_000		,  )
				__magic_name__						:						Optional[Any]        =	CLIPTextModel(__SCREAMING_SNAKE_CASE						)
				__magic_name__						:						List[Any]        =	CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip"						)

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







			def      SCREAMING_SNAKE_CASE       (							self		,  _a		,  _a=0						):
				__magic_name__						:						List[str]        =	torch.manual_seed(__SCREAMING_SNAKE_CASE						)
				__magic_name__						:						str        =	{
				    "prompt": "a photo of the dolomites",
				    "generator": generator,
				    # Setting height and width to None to prevent OOMs on CPU.
				    "height": None,
				    "width": None,
				    "num_inference_steps": 1,
				    "guidance_scale": 6.0,
				    "output_type": "numpy",
				}
				return inputs







			def      SCREAMING_SNAKE_CASE       (							self						):
				__magic_name__						:						int        =	"cpu"  # ensure determinism for the device-dependent torch.Generator
				__magic_name__						:						Tuple        =	self.get_dummy_components()
				__magic_name__						:						Any        =	StableDiffusionPanoramaPipeline(**__SCREAMING_SNAKE_CASE						)
				__magic_name__						:						Tuple        =	sd_pipe.to(__SCREAMING_SNAKE_CASE						)
				sd_pipe.set_progress_bar_config(disable=__SCREAMING_SNAKE_CASE						)

				__magic_name__						:						Tuple        =	self.get_dummy_inputs(__SCREAMING_SNAKE_CASE						)
				__magic_name__						:						Union[str, Any]        =	sd_pipe(**__SCREAMING_SNAKE_CASE						).images
				__magic_name__						:						Union[str, Any]        =	image[0, -3:, -3:, -1]
				assert image.shape == (1, 64, 64, 3)

				__magic_name__						:						Dict        =	np.array([0.61_86, 0.53_74, 0.49_15, 0.41_35, 0.41_14, 0.45_63, 0.51_28, 0.49_77, 0.47_57]						)

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







			def      SCREAMING_SNAKE_CASE       (							self						):
				super().test_inference_batch_consistent(batch_sizes=[1, 2]						)







			def      SCREAMING_SNAKE_CASE       (							self						):
				super().test_inference_batch_single_identical(batch_size=2		,  expected_max_diff=3.25e-3						)







			def      SCREAMING_SNAKE_CASE       (							self						):
				__magic_name__						:						Dict        =	"cpu"  # ensure determinism for the device-dependent torch.Generator
				__magic_name__						:						Tuple        =	self.get_dummy_components()
				__magic_name__						:						Tuple        =	StableDiffusionPanoramaPipeline(**__SCREAMING_SNAKE_CASE						)
				__magic_name__						:						Union[str, Any]        =	sd_pipe.to(__SCREAMING_SNAKE_CASE						)
				sd_pipe.set_progress_bar_config(disable=__SCREAMING_SNAKE_CASE						)

				__magic_name__						:						Union[str, Any]        =	self.get_dummy_inputs(__SCREAMING_SNAKE_CASE						)
				__magic_name__						:						Dict        =	"french fries"
				__magic_name__						:						Optional[Any]        =	sd_pipe(**__SCREAMING_SNAKE_CASE		,  negative_prompt=__SCREAMING_SNAKE_CASE						)
				__magic_name__						:						Tuple        =	output.images
				__magic_name__						:						int        =	image[0, -3:, -3:, -1]

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

				__magic_name__						:						Optional[Any]        =	np.array([0.61_87, 0.53_75, 0.49_15, 0.41_36, 0.41_14, 0.45_63, 0.51_28, 0.49_76, 0.47_57]						)

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







			def      SCREAMING_SNAKE_CASE       (							self						):
				__magic_name__						:						Optional[int]        =	"cpu"  # ensure determinism for the device-dependent torch.Generator
				__magic_name__						:						Any        =	self.get_dummy_components()
				__magic_name__						:						List[str]        =	StableDiffusionPanoramaPipeline(**__SCREAMING_SNAKE_CASE						)
				__magic_name__						:						List[str]        =	sd_pipe.to(__SCREAMING_SNAKE_CASE						)
				sd_pipe.set_progress_bar_config(disable=__SCREAMING_SNAKE_CASE						)

				__magic_name__						:						Any        =	self.get_dummy_inputs(__SCREAMING_SNAKE_CASE						)
				__magic_name__						:						Optional[int]        =	sd_pipe(**__SCREAMING_SNAKE_CASE		,  view_batch_size=2						)
				__magic_name__						:						Union[str, Any]        =	output.images
				__magic_name__						:						List[str]        =	image[0, -3:, -3:, -1]

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

				__magic_name__						:						str        =	np.array([0.61_87, 0.53_75, 0.49_15, 0.41_36, 0.41_14, 0.45_63, 0.51_28, 0.49_76, 0.47_57]						)

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







			def      SCREAMING_SNAKE_CASE       (							self						):
				__magic_name__						:						List[str]        =	"cpu"  # ensure determinism for the device-dependent torch.Generator
				__magic_name__						:						Union[str, Any]        =	self.get_dummy_components()
				__magic_name__						:						Dict        =	EulerAncestralDiscreteScheduler(
				    beta_start=0.0_00_85		,  beta_end=0.0_12		,  beta_schedule="scaled_linear"						)
				__magic_name__						:						Optional[Any]        =	StableDiffusionPanoramaPipeline(**__SCREAMING_SNAKE_CASE						)
				__magic_name__						:						Optional[int]        =	sd_pipe.to(__SCREAMING_SNAKE_CASE						)
				sd_pipe.set_progress_bar_config(disable=__SCREAMING_SNAKE_CASE						)

				__magic_name__						:						Tuple        =	self.get_dummy_inputs(__SCREAMING_SNAKE_CASE						)
				__magic_name__						:						Union[str, Any]        =	sd_pipe(**__SCREAMING_SNAKE_CASE						).images
				__magic_name__						:						List[Any]        =	image[0, -3:, -3:, -1]

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

				__magic_name__						:						Dict        =	np.array([0.40_24, 0.65_10, 0.49_01, 0.53_78, 0.58_13, 0.56_22, 0.47_95, 0.44_67, 0.49_52]						)

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







			def      SCREAMING_SNAKE_CASE       (							self						):
				__magic_name__						:						Optional[Any]        =	"cpu"  # ensure determinism for the device-dependent torch.Generator
				__magic_name__						:						str        =	self.get_dummy_components()
				__magic_name__						:						List[str]        =	PNDMScheduler(
				    beta_start=0.0_00_85		,  beta_end=0.0_12		,  beta_schedule="scaled_linear"		,  skip_prk_steps=__SCREAMING_SNAKE_CASE						)
				__magic_name__						:						Optional[int]        =	StableDiffusionPanoramaPipeline(**__SCREAMING_SNAKE_CASE						)
				__magic_name__						:						Optional[Any]        =	sd_pipe.to(__SCREAMING_SNAKE_CASE						)
				sd_pipe.set_progress_bar_config(disable=__SCREAMING_SNAKE_CASE						)

				__magic_name__						:						Tuple        =	self.get_dummy_inputs(__SCREAMING_SNAKE_CASE						)
				__magic_name__						:						Union[str, Any]        =	sd_pipe(**__SCREAMING_SNAKE_CASE						).images
				__magic_name__						:						List[Any]        =	image[0, -3:, -3:, -1]

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

				__magic_name__						:						Tuple        =	np.array([0.63_91, 0.62_91, 0.48_61, 0.51_34, 0.55_52, 0.45_78, 0.50_32, 0.50_23, 0.45_39]						)

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







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







			def      SCREAMING_SNAKE_CASE       (							self						):
				super().tearDown()
				gc.collect()
				torch.cuda.empty_cache()







			def      SCREAMING_SNAKE_CASE       (							self		,  _a=0						):
				__magic_name__						:						List[str]        =	torch.manual_seed(__SCREAMING_SNAKE_CASE						)
				__magic_name__						:						Tuple        =	{
				    "prompt": "a photo of the dolomites",
				    "generator": generator,
				    "num_inference_steps": 3,
				    "guidance_scale": 7.5,
				    "output_type": "numpy",
				}
				return inputs







			def      SCREAMING_SNAKE_CASE       (							self						):
				__magic_name__						:						Tuple        =	"stabilityai/stable-diffusion-2-base"
				__magic_name__						:						Dict        =	DDIMScheduler.from_pretrained(__SCREAMING_SNAKE_CASE		,  subfolder="scheduler"						)
				__magic_name__						:						Tuple        =	StableDiffusionPanoramaPipeline.from_pretrained(__SCREAMING_SNAKE_CASE		,  scheduler=__SCREAMING_SNAKE_CASE		,  safety_checker=__SCREAMING_SNAKE_CASE						)
				pipe.to(__SCREAMING_SNAKE_CASE						)
				pipe.set_progress_bar_config(disable=__SCREAMING_SNAKE_CASE						)
				pipe.enable_attention_slicing()

				__magic_name__						:						str        =	self.get_inputs()
				__magic_name__						:						List[Any]        =	pipe(**__SCREAMING_SNAKE_CASE						).images
				__magic_name__						:						Union[str, Any]        =	image[0, -3:, -3:, -1].flatten()

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

				__magic_name__						:						Tuple        =	np.array(
				    [
				        0.36_96_83_92,
				        0.27_02_53_72,
				        0.32_44_67_66,
				        0.28_37_93_87,
				        0.36_36_32_74,
				        0.30_73_33_47,
				        0.27_10_00_27,
				        0.27_05_41_25,
				        0.25_53_60_96,
				    ]						)

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







			def      SCREAMING_SNAKE_CASE       (							self						):
				__magic_name__						:						Optional[Any]        =	StableDiffusionPanoramaPipeline.from_pretrained(
				    "stabilityai/stable-diffusion-2-base"		,  safety_checker=__SCREAMING_SNAKE_CASE						)
				__magic_name__						:						Any        =	LMSDiscreteScheduler.from_config(pipe.scheduler.config						)
				pipe.to(__SCREAMING_SNAKE_CASE						)
				pipe.set_progress_bar_config(disable=__SCREAMING_SNAKE_CASE						)
				pipe.enable_attention_slicing()

				__magic_name__						:						Dict        =	self.get_inputs()
				__magic_name__						:						Dict        =	pipe(**__SCREAMING_SNAKE_CASE						).images
				__magic_name__						:						Optional[int]        =	image[0, -3:, -3:, -1].flatten()

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

				__magic_name__						:						Dict        =	np.array(
				    [
				        [
				            0.0,
				            0.0,
				            0.0,
				            0.0,
				            0.0,
				            0.0,
				            0.0,
				            0.0,
				            0.0,
				        ]
				    ]						)

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







			def      SCREAMING_SNAKE_CASE       (							self						):
				__magic_name__						:						Any        =	0

				def callback_fn(_a		,  _a		,  _a						) -> None:
					__magic_name__						:						Dict        =	True
					nonlocal number_of_steps
					number_of_steps += 1
					if step == 1:
						__magic_name__						:						Optional[Any]        =	latents.detach().cpu().numpy()
						assert latents.shape == (1, 4, 64, 256)
						__magic_name__						:						Union[str, Any]        =	latents[0, -3:, -3:, -1]

						__magic_name__						:						Optional[Any]        =	np.array(
						    [
						        0.18_68_18_69,
						        0.33_90_78_16,
						        0.5_36_12_76,
						        0.14_43_28_65,
						        -0.02_85_66_11,
						        -0.73_94_11_23,
						        0.23_39_79_87,
						        0.47_32_26_82,
						        -0.37_82_31_64,
						    ]						)
						assert np.abs(latents_slice.flatten() - expected_slice						).max() < 5e-2
					elif step == 2:
						__magic_name__						:						Union[str, Any]        =	latents.detach().cpu().numpy()
						assert latents.shape == (1, 4, 64, 256)
						__magic_name__						:						Any        =	latents[0, -3:, -3:, -1]

						__magic_name__						:						Optional[int]        =	np.array(
						    [
						        0.18_53_96_45,
						        0.33_98_72_48,
						        0.5_37_85_59,
						        0.14_43_71_42,
						        -0.02_45_52_61,
						        -0.7_33_83_17,
						        0.23_99_07_55,
						        0.47_35_62_72,
						        -0.3_78_65_05,
						    ]						)

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

				__magic_name__						:						Dict        =	False

				__magic_name__						:						str        =	"stabilityai/stable-diffusion-2-base"
				__magic_name__						:						Dict        =	DDIMScheduler.from_pretrained(__SCREAMING_SNAKE_CASE		,  subfolder="scheduler"						)
				__magic_name__						:						Any        =	StableDiffusionPanoramaPipeline.from_pretrained(__SCREAMING_SNAKE_CASE		,  scheduler=__SCREAMING_SNAKE_CASE		,  safety_checker=__SCREAMING_SNAKE_CASE						)
				__magic_name__						:						List[Any]        =	pipe.to(__SCREAMING_SNAKE_CASE						)
				pipe.set_progress_bar_config(disable=__SCREAMING_SNAKE_CASE						)
				pipe.enable_attention_slicing()

				__magic_name__						:						List[str]        =	self.get_inputs()
				pipe(**__SCREAMING_SNAKE_CASE		,  callback=__SCREAMING_SNAKE_CASE		,  callback_steps=1						)
				assert callback_fn.has_been_called
				assert number_of_steps == 3







			def      SCREAMING_SNAKE_CASE       (							self						):
				torch.cuda.empty_cache()
				torch.cuda.reset_max_memory_allocated()
				torch.cuda.reset_peak_memory_stats()

				__magic_name__						:						List[Any]        =	"stabilityai/stable-diffusion-2-base"
				__magic_name__						:						Dict        =	DDIMScheduler.from_pretrained(__SCREAMING_SNAKE_CASE		,  subfolder="scheduler"						)
				__magic_name__						:						Dict        =	StableDiffusionPanoramaPipeline.from_pretrained(__SCREAMING_SNAKE_CASE		,  scheduler=__SCREAMING_SNAKE_CASE		,  safety_checker=__SCREAMING_SNAKE_CASE						)
				__magic_name__						:						Any        =	pipe.to(__SCREAMING_SNAKE_CASE						)
				pipe.set_progress_bar_config(disable=__SCREAMING_SNAKE_CASE						)
				pipe.enable_attention_slicing(1						)
				pipe.enable_sequential_cpu_offload()

				__magic_name__						:						Union[str, Any]        =	self.get_inputs()
				__magic_name__						:						int        =	pipe(**__SCREAMING_SNAKE_CASE						)

				__magic_name__						:						int        =	torch.cuda.max_memory_allocated()
				# make sure that less than 5.2 GB is allocated
				assert mem_bytes < 5.5 * 10**9

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

from argparse import ArgumentParser

from accelerate.commands.config import get_config_parser
from accelerate.commands.env import env_command_parser
from accelerate.commands.launch import launch_command_parser
from accelerate.commands.test import test_command_parser
from accelerate.commands.tpu import tpu_command_parser



def 	__A		(      ):



 """simple docstring"""


 __a							=	ArgumentParser("Accelerate CLI tool"				,					usage="accelerate <command> [<args>]"				,					allow_abbrev=_A       )
 __a							=	parser.add_subparsers(help="accelerate command helpers"       )

 # Register commands
 get_config_parser(subparsers=_A       )
 env_command_parser(subparsers=_A       )
 launch_command_parser(subparsers=_A       )
 tpu_command_parser(subparsers=_A       )
 test_command_parser(subparsers=_A       )

 # Let's go
 __a							=	parser.parse_args()

 if not hasattr(_A				,					"func"       ):
  parser.print_help()
  exit(1       )

 # Run
 args.func(_A       )


if __name__ == "__main__":
  main()


