





"""simple docstring"""





import unittest

import torch

from diffusers import VQModel
from diffusers.utils import floats_tensor, torch_device
from diffusers.utils.testing_utils import enable_full_determinism

from .test_modeling_common import ModelTesterMixin, UNetTesterMixin


enable_full_determinism()





class        _lowerCamelCase      (       _lowercase  ,			_lowercase  ,			unittest.TestCase ):
						UpperCAmelCase_ =			VQModel
						UpperCAmelCase_ =			"sample"


						@property
						def     snake_case_   (self   ,							__a=(32, 32)		) ->   Union[str, Any]:
								UpperCamelCase   =	4
								UpperCamelCase   =	3

								UpperCamelCase   =	floats_tensor((batch_size, num_channels) + sizes		).to(__a		)

								return {"sample": image}


						@property
						def     snake_case_   (self		) ->   List[str]:
								return (3, 32, 32)


						@property
						def     snake_case_   (self		) ->   Any:
								return (3, 32, 32)


						def     snake_case_   (self		) ->   int:
								UpperCamelCase   =	{
								    "block_out_channels": [32, 64],
								    "in_channels": 3,
								    "out_channels": 3,
								    "down_block_types": ["DownEncoderBlock2D", "DownEncoderBlock2D"],
								    "up_block_types": ["UpDecoderBlock2D", "UpDecoderBlock2D"],
								    "latent_channels": 3,
								}
								UpperCamelCase   =	self.dummy_input
								return init_dict, inputs_dict


						def     snake_case_   (self		) ->   Dict:
								pass


						def     snake_case_   (self		) ->   Optional[Any]:
								pass


						def     snake_case_   (self		) ->   List[Any]:
								UpperCamelCase   ,				UpperCamelCase   =	VQModel.from_pretrained("fusing/vqgan-dummy"   ,							output_loading_info=__a		)
								self.assertIsNotNone(__a		)
								self.assertEqual(len(loading_info["missing_keys"]		)   ,							0		)

								model.to(__a		)
								UpperCamelCase   =	model(**self.dummy_input		)

								assert image is not None, "Make sure output is not None"




						def     snake_case_   (self		) ->   Optional[Any]:
								UpperCamelCase   =	VQModel.from_pretrained("fusing/vqgan-dummy"		)
								model.to(__a		).eval()

								torch.manual_seed(0		)
								if torch.cuda.is_available():
										torch.cuda.manual_seed_all(0		)

								UpperCamelCase   =	torch.randn(1   ,							model.config.in_channels   ,							model.config.sample_size   ,							model.config.sample_size		)
								UpperCamelCase   =	image.to(__a		)
								with torch.no_grad():
										UpperCamelCase   =	model(__a		).sample

								UpperCamelCase   =	output[0, -1, -3:, -3:].flatten().cpu()
								# fmt: off
								UpperCamelCase   =	torch.tensor([-0.0153, -0.4044, -0.1880, -0.5161, -0.2418, -0.4072, -0.1612, -0.0633, -0.0143]		)
								# fmt: on
								self.assertTrue(torch.allclose(__a   ,							__a   ,							atol=1e-3		)		)













"""simple docstring"""





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        _lowerCamelCase      (       _lowercase  ,			_lowercase  ,			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     snake_case_   (self		) ->   int:
								torch.manual_seed(0		)
								UpperCamelCase   =	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   ,							)
								UpperCamelCase   =	DDIMScheduler()
								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   ,							)
								torch.manual_seed(0		)
								UpperCamelCase   =	CLIPTextConfig(
								    bos_token_id=0   ,							eos_token_id=2   ,							hidden_size=32   ,							intermediate_size=37   ,							layer_norm_eps=1e-0_5   ,							num_attention_heads=4   ,							num_hidden_layers=5   ,							pad_token_id=1   ,							vocab_size=10_00   ,							)
								UpperCamelCase   =	CLIPTextModel(__a		)
								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     snake_case_   (self   ,							__a   ,							__a=0		) ->   List[str]:
								UpperCamelCase   =	torch.manual_seed(__a		)
								UpperCamelCase   =	{
								    "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     snake_case_   (self		) ->   int:
								UpperCamelCase   =	"cpu"  # ensure determinism for the device-dependent torch.Generator
								UpperCamelCase   =	self.get_dummy_components()
								UpperCamelCase   =	StableDiffusionPanoramaPipeline(**__a		)
								UpperCamelCase   =	sd_pipe.to(__a		)
								sd_pipe.set_progress_bar_config(disable=__a		)

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

								UpperCamelCase   =	np.array([0.6186, 0.5374, 0.4915, 0.4135, 0.4114, 0.4563, 0.5128, 0.4977, 0.4757]		)

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


						def     snake_case_   (self		) ->   Optional[int]:
								super().test_inference_batch_consistent(batch_sizes=[1, 2]		)


						def     snake_case_   (self		) ->   Tuple:
								super().test_inference_batch_single_identical(batch_size=2   ,							expected_max_diff=3.2_5e-3		)


						def     snake_case_   (self		) ->   Tuple:
								UpperCamelCase   =	"cpu"  # ensure determinism for the device-dependent torch.Generator
								UpperCamelCase   =	self.get_dummy_components()
								UpperCamelCase   =	StableDiffusionPanoramaPipeline(**__a		)
								UpperCamelCase   =	sd_pipe.to(__a		)
								sd_pipe.set_progress_bar_config(disable=__a		)

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

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

								UpperCamelCase   =	np.array([0.6187, 0.5375, 0.4915, 0.4136, 0.4114, 0.4563, 0.5128, 0.4976, 0.4757]		)

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


						def     snake_case_   (self		) ->   Tuple:
								UpperCamelCase   =	"cpu"  # ensure determinism for the device-dependent torch.Generator
								UpperCamelCase   =	self.get_dummy_components()
								UpperCamelCase   =	StableDiffusionPanoramaPipeline(**__a		)
								UpperCamelCase   =	sd_pipe.to(__a		)
								sd_pipe.set_progress_bar_config(disable=__a		)

								UpperCamelCase   =	self.get_dummy_inputs(__a		)
								UpperCamelCase   =	sd_pipe(**__a   ,							view_batch_size=2		)
								UpperCamelCase   =	output.images
								UpperCamelCase   =	image[0, -3:, -3:, -1]

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

								UpperCamelCase   =	np.array([0.6187, 0.5375, 0.4915, 0.4136, 0.4114, 0.4563, 0.5128, 0.4976, 0.4757]		)

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


						def     snake_case_   (self		) ->   Optional[int]:
								UpperCamelCase   =	"cpu"  # ensure determinism for the device-dependent torch.Generator
								UpperCamelCase   =	self.get_dummy_components()
								UpperCamelCase   =	EulerAncestralDiscreteScheduler(
								    beta_start=0.00085   ,							beta_end=0.012   ,							beta_schedule="scaled_linear"		)
								UpperCamelCase   =	StableDiffusionPanoramaPipeline(**__a		)
								UpperCamelCase   =	sd_pipe.to(__a		)
								sd_pipe.set_progress_bar_config(disable=__a		)

								UpperCamelCase   =	self.get_dummy_inputs(__a		)
								UpperCamelCase   =	sd_pipe(**__a		).images
								UpperCamelCase   =	image[0, -3:, -3:, -1]

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

								UpperCamelCase   =	np.array([0.4024, 0.6510, 0.4901, 0.5378, 0.5813, 0.5622, 0.4795, 0.4467, 0.4952]		)

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




						def     snake_case_   (self		) ->   Optional[Any]:
								UpperCamelCase   =	"cpu"  # ensure determinism for the device-dependent torch.Generator
								UpperCamelCase   =	self.get_dummy_components()
								UpperCamelCase   =	PNDMScheduler(
								    beta_start=0.00085   ,							beta_end=0.012   ,							beta_schedule="scaled_linear"   ,							skip_prk_steps=__a		)
								UpperCamelCase   =	StableDiffusionPanoramaPipeline(**__a		)
								UpperCamelCase   =	sd_pipe.to(__a		)
								sd_pipe.set_progress_bar_config(disable=__a		)

								UpperCamelCase   =	self.get_dummy_inputs(__a		)
								UpperCamelCase   =	sd_pipe(**__a		).images
								UpperCamelCase   =	image[0, -3:, -3:, -1]

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

								UpperCamelCase   =	np.array([0.6391, 0.6291, 0.4861, 0.5134, 0.5552, 0.4578, 0.5032, 0.5023, 0.4539]		)

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





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


						def     snake_case_   (self		) ->   List[str]:
								super().tearDown()
								gc.collect()
								torch.cuda.empty_cache()


						def     snake_case_   (self   ,							__a=0		) ->   List[str]:
								UpperCamelCase   =	torch.manual_seed(__a		)
								UpperCamelCase   =	{
								    "prompt": "a photo of the dolomites",
								    "generator": generator,
								    "num_inference_steps": 3,
								    "guidance_scale": 7.5,
								    "output_type": "numpy",
								}
								return inputs


						def     snake_case_   (self		) ->   Union[str, Any]:
								UpperCamelCase   =	"stabilityai/stable-diffusion-2-base"
								UpperCamelCase   =	DDIMScheduler.from_pretrained(__a   ,							subfolder="scheduler"		)
								UpperCamelCase   =	StableDiffusionPanoramaPipeline.from_pretrained(__a   ,							scheduler=__a   ,							safety_checker=__a		)
								pipe.to(__a		)
								pipe.set_progress_bar_config(disable=__a		)
								pipe.enable_attention_slicing()

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

								assert image.shape == (1, 5_12, 20_48, 3)

								UpperCamelCase   =	np.array(
								    [
								        0.36968392,
								        0.27025372,
								        0.32446766,
								        0.28379387,
								        0.36363274,
								        0.30733347,
								        0.27100027,
								        0.27054125,
								        0.25536096,
								    ]		)

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


						def     snake_case_   (self		) ->   Union[str, Any]:
								UpperCamelCase   =	StableDiffusionPanoramaPipeline.from_pretrained(
								    "stabilityai/stable-diffusion-2-base"   ,							safety_checker=__a		)
								UpperCamelCase   =	LMSDiscreteScheduler.from_config(pipe.scheduler.config		)
								pipe.to(__a		)
								pipe.set_progress_bar_config(disable=__a		)
								pipe.enable_attention_slicing()

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

								assert image.shape == (1, 5_12, 20_48, 3)

								UpperCamelCase   =	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     snake_case_   (self		) ->   Tuple:
								UpperCamelCase   =	0

								def callback_fn(__a   ,							__a   ,							__a		) -> None:
										UpperCamelCase   =	True
										nonlocal number_of_steps
										number_of_steps += 1
										if step == 1:
												UpperCamelCase   =	latents.detach().cpu().numpy()
												assert latents.shape == (1, 4, 64, 2_56)
												UpperCamelCase   =	latents[0, -3:, -3:, -1]

												UpperCamelCase   =	np.array(
												    [
												        0.18681869,
												        0.33907816,
												        0.5361276,
												        0.14432865,
												        -0.02856611,
												        -0.73941123,
												        0.23397987,
												        0.47322682,
												        -0.37823164,
												    ]		)
												assert np.abs(latents_slice.flatten() - expected_slice		).max() < 5e-2
										elif step == 2:
												UpperCamelCase   =	latents.detach().cpu().numpy()
												assert latents.shape == (1, 4, 64, 2_56)
												UpperCamelCase   =	latents[0, -3:, -3:, -1]

												UpperCamelCase   =	np.array(
												    [
												        0.18539645,
												        0.33987248,
												        0.5378559,
												        0.14437142,
												        -0.02455261,
												        -0.7338317,
												        0.23990755,
												        0.47356272,
												        -0.3786505,
												    ]		)

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

								UpperCamelCase   =	False

								UpperCamelCase   =	"stabilityai/stable-diffusion-2-base"
								UpperCamelCase   =	DDIMScheduler.from_pretrained(__a   ,							subfolder="scheduler"		)
								UpperCamelCase   =	StableDiffusionPanoramaPipeline.from_pretrained(__a   ,							scheduler=__a   ,							safety_checker=__a		)
								UpperCamelCase   =	pipe.to(__a		)
								pipe.set_progress_bar_config(disable=__a		)
								pipe.enable_attention_slicing()

								UpperCamelCase   =	self.get_inputs()
								pipe(**__a   ,							callback=__a   ,							callback_steps=1		)
								assert callback_fn.has_been_called
								assert number_of_steps == 3




						def     snake_case_   (self		) ->   List[Any]:
								torch.cuda.empty_cache()
								torch.cuda.reset_max_memory_allocated()
								torch.cuda.reset_peak_memory_stats()

								UpperCamelCase   =	"stabilityai/stable-diffusion-2-base"
								UpperCamelCase   =	DDIMScheduler.from_pretrained(__a   ,							subfolder="scheduler"		)
								UpperCamelCase   =	StableDiffusionPanoramaPipeline.from_pretrained(__a   ,							scheduler=__a   ,							safety_checker=__a		)
								UpperCamelCase   =	pipe.to(__a		)
								pipe.set_progress_bar_config(disable=__a		)
								pipe.enable_attention_slicing(1		)
								pipe.enable_sequential_cpu_offload()

								UpperCamelCase   =	self.get_inputs()
								UpperCamelCase   =	pipe(**__a		)

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






