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__     (				SCREAMING_SNAKE_CASE_     ,SCREAMING_SNAKE_CASE_     ,unittest.TestCase			):


		'''simple docstring'''






		lowerCAmelCase_     =   StableDiffusionPanoramaPipeline
		lowerCAmelCase_     =   TEXT_TO_IMAGE_PARAMS
		lowerCAmelCase_     =   TEXT_TO_IMAGE_BATCH_PARAMS
		lowerCAmelCase_     =   TEXT_TO_IMAGE_IMAGE_PARAMS
		lowerCAmelCase_     =   TEXT_TO_IMAGE_IMAGE_PARAMS
		def   _snake_case       (       self     :     Any		)      ->	Optional[Any]:
				torch.manual_seed(0		)
				_lowerCamelCase				   =  UNetaDConditionModel(
				    block_out_channels=(3_2, 6_4)     ,  layers_per_block=1     ,  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     ,  )
				_lowerCamelCase				   =  DDIMScheduler()
				torch.manual_seed(0		)
				_lowerCamelCase				   =  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		)
				_lowerCamelCase				   =  CLIPTextConfig(
				    bos_token_id=0     ,  eos_token_id=2     ,  hidden_size=3_2     ,  intermediate_size=3_7     ,  layer_norm_eps=1e-05     ,  num_attention_heads=4     ,  num_hidden_layers=5     ,  pad_token_id=1     ,  vocab_size=1_0_0_0     ,  )
				_lowerCamelCase				   =  CLIPTextModel(snake_case__		)
				_lowerCamelCase				   =  CLIPTokenizer.from_pretrained('hf-internal-testing/tiny-random-clip'		)

				_lowerCamelCase				   =  {
				    'unet': unet,
				    'scheduler': scheduler,
				    'vae': vae,
				    'text_encoder': text_encoder,
				    'tokenizer': tokenizer,
				    'safety_checker': None,
				    'feature_extractor': None,
				}
				return components
		def   _snake_case       (       self     :     Tuple     ,  snake_case__     :     Union[str, Any]     ,  snake_case__     :     Dict=0		)      ->	str:
				_lowerCamelCase				   =  torch.manual_seed(snake_case__		)
				_lowerCamelCase				   =  {
				    '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     :     Union[str, Any]		)      ->	Optional[Any]:
				_lowerCamelCase				   =  'cpu'  # ensure determinism for the device-dependent torch.Generator
				_lowerCamelCase				   =  self.get_dummy_components()
				_lowerCamelCase				   =  StableDiffusionPanoramaPipeline(**snake_case__		)
				_lowerCamelCase				   =  sd_pipe.to(snake_case__		)
				sd_pipe.set_progress_bar_config(disable=snake_case__		)

				_lowerCamelCase				   =  self.get_dummy_inputs(snake_case__		)
				_lowerCamelCase				   =  sd_pipe(**snake_case__		).images
				_lowerCamelCase				   =  image[0, -3:, -3:, -1]
				assert image.shape == (1, 6_4, 6_4, 3)

				_lowerCamelCase				   =  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     :     Tuple		)      ->	Union[str, Any]:
				super().test_inference_batch_consistent(batch_sizes=[1, 2]		)
		def   _snake_case       (       self     :     str		)      ->	int:
				super().test_inference_batch_single_identical(batch_size=2     ,  expected_max_diff=3.25e-3		)
		def   _snake_case       (       self     :     Optional[int]		)      ->	Any:
				_lowerCamelCase				   =  'cpu'  # ensure determinism for the device-dependent torch.Generator
				_lowerCamelCase				   =  self.get_dummy_components()
				_lowerCamelCase				   =  StableDiffusionPanoramaPipeline(**snake_case__		)
				_lowerCamelCase				   =  sd_pipe.to(snake_case__		)
				sd_pipe.set_progress_bar_config(disable=snake_case__		)

				_lowerCamelCase				   =  self.get_dummy_inputs(snake_case__		)
				_lowerCamelCase				   =  'french fries'
				_lowerCamelCase				   =  sd_pipe(**snake_case__     ,  negative_prompt=snake_case__		)
				_lowerCamelCase				   =  output.images
				_lowerCamelCase				   =  image[0, -3:, -3:, -1]

				assert image.shape == (1, 6_4, 6_4, 3)

				_lowerCamelCase				   =  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     :     Union[str, Any]		)      ->	int:
				_lowerCamelCase				   =  'cpu'  # ensure determinism for the device-dependent torch.Generator
				_lowerCamelCase				   =  self.get_dummy_components()
				_lowerCamelCase				   =  StableDiffusionPanoramaPipeline(**snake_case__		)
				_lowerCamelCase				   =  sd_pipe.to(snake_case__		)
				sd_pipe.set_progress_bar_config(disable=snake_case__		)

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

				assert image.shape == (1, 6_4, 6_4, 3)

				_lowerCamelCase				   =  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     :     List[Any]		)      ->	Optional[Any]:
				_lowerCamelCase				   =  'cpu'  # ensure determinism for the device-dependent torch.Generator
				_lowerCamelCase				   =  self.get_dummy_components()
				_lowerCamelCase				   =  EulerAncestralDiscreteScheduler(
				    beta_start=0.0_0085     ,  beta_end=0.012     ,  beta_schedule='scaled_linear'		)
				_lowerCamelCase				   =  StableDiffusionPanoramaPipeline(**snake_case__		)
				_lowerCamelCase				   =  sd_pipe.to(snake_case__		)
				sd_pipe.set_progress_bar_config(disable=snake_case__		)

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

				assert image.shape == (1, 6_4, 6_4, 3)

				_lowerCamelCase				   =  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]		)      ->	str:
				_lowerCamelCase				   =  'cpu'  # ensure determinism for the device-dependent torch.Generator
				_lowerCamelCase				   =  self.get_dummy_components()
				_lowerCamelCase				   =  PNDMScheduler(
				    beta_start=0.0_0085     ,  beta_end=0.012     ,  beta_schedule='scaled_linear'     ,  skip_prk_steps=snake_case__		)
				_lowerCamelCase				   =  StableDiffusionPanoramaPipeline(**snake_case__		)
				_lowerCamelCase				   =  sd_pipe.to(snake_case__		)
				sd_pipe.set_progress_bar_config(disable=snake_case__		)

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

				assert image.shape == (1, 6_4, 6_4, 3)

				_lowerCamelCase				   =  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			):


		'''simple docstring'''






		def   _snake_case       (       self     :     List[str]		)      ->	Optional[int]:
				super().tearDown()
				gc.collect()
				torch.cuda.empty_cache()
		def   _snake_case       (       self     :     Tuple     ,  snake_case__     :     Optional[Any]=0		)      ->	Tuple:
				_lowerCamelCase				   =  torch.manual_seed(snake_case__		)
				_lowerCamelCase				   =  {
				    '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     :     Optional[int]		)      ->	List[str]:
				_lowerCamelCase				   =  'stabilityai/stable-diffusion-2-base'
				_lowerCamelCase				   =  DDIMScheduler.from_pretrained(snake_case__     ,  subfolder='scheduler'		)
				_lowerCamelCase				   =  StableDiffusionPanoramaPipeline.from_pretrained(snake_case__     ,  scheduler=snake_case__     ,  safety_checker=snake_case__		)
				pipe.to(snake_case__		)
				pipe.set_progress_bar_config(disable=snake_case__		)
				pipe.enable_attention_slicing()

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

				assert image.shape == (1, 5_1_2, 2_0_4_8, 3)

				_lowerCamelCase				   =  np.array(
				    [
				        0.3696_8392,
				        0.2702_5372,
				        0.3244_6766,
				        0.2837_9387,
				        0.3636_3274,
				        0.3073_3347,
				        0.2710_0027,
				        0.2705_4125,
				        0.2553_6096,
				    ]		)

				assert np.abs(expected_slice - image_slice		).max() < 1e-2
		def   _snake_case       (       self     :     List[str]		)      ->	Tuple:
				_lowerCamelCase				   =  StableDiffusionPanoramaPipeline.from_pretrained(
				    'stabilityai/stable-diffusion-2-base'     ,  safety_checker=snake_case__		)
				_lowerCamelCase				   =  LMSDiscreteScheduler.from_config(pipe.scheduler.config		)
				pipe.to(snake_case__		)
				pipe.set_progress_bar_config(disable=snake_case__		)
				pipe.enable_attention_slicing()

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

				assert image.shape == (1, 5_1_2, 2_0_4_8, 3)

				_lowerCamelCase				   =  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     :     Optional[int]		)      ->	str:
				_lowerCamelCase				   =  0

				def callback_fn(snake_case__     :     int     ,  snake_case__     :     int     ,  snake_case__     :     torch.FloatTensor		) -> None:
						_lowerCamelCase				   =  True
						nonlocal number_of_steps
						number_of_steps += 1
						if step == 1:
								_lowerCamelCase				   =  latents.detach().cpu().numpy()
								assert latents.shape == (1, 4, 6_4, 2_5_6)
								_lowerCamelCase				   =  latents[0, -3:, -3:, -1]

								_lowerCamelCase				   =  np.array(
								    [
								        0.1868_1869,
								        0.3390_7816,
								        0.536_1276,
								        0.1443_2865,
								        -0.0285_6611,
								        -0.7394_1123,
								        0.2339_7987,
								        0.4732_2682,
								        -0.3782_3164,
								    ]		)
								assert np.abs(latents_slice.flatten() - expected_slice		).max() < 5e-2
						elif step == 2:
								_lowerCamelCase				   =  latents.detach().cpu().numpy()
								assert latents.shape == (1, 4, 6_4, 2_5_6)
								_lowerCamelCase				   =  latents[0, -3:, -3:, -1]

								_lowerCamelCase				   =  np.array(
								    [
								        0.1853_9645,
								        0.3398_7248,
								        0.537_8559,
								        0.1443_7142,
								        -0.0245_5261,
								        -0.733_8317,
								        0.2399_0755,
								        0.4735_6272,
								        -0.378_6505,
								    ]		)

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

				_lowerCamelCase				   =  False

				_lowerCamelCase				   =  'stabilityai/stable-diffusion-2-base'
				_lowerCamelCase				   =  DDIMScheduler.from_pretrained(snake_case__     ,  subfolder='scheduler'		)
				_lowerCamelCase				   =  StableDiffusionPanoramaPipeline.from_pretrained(snake_case__     ,  scheduler=snake_case__     ,  safety_checker=snake_case__		)
				_lowerCamelCase				   =  pipe.to(snake_case__		)
				pipe.set_progress_bar_config(disable=snake_case__		)
				pipe.enable_attention_slicing()

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



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

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

				_lowerCamelCase				   =  self.get_inputs()
				_lowerCamelCase				   =  pipe(**snake_case__		)

				_lowerCamelCase				   =  torch.cuda.max_memory_allocated()
				# make sure that less than 5.2 GB is allocated
				assert mem_bytes < 5.5 * 1_0**9
from __future__ import annotations

from collections import deque
from collections.abc import Iterator
from dataclasses import dataclass





@dataclass
class     lowerCAmelCase__     :


		'''simple docstring'''






		lowerCAmelCase_     =   42
		lowerCAmelCase_     =   42





class     lowerCAmelCase__     :


		'''simple docstring'''






		def __init__(       self     :     int     ,  snake_case__     :     int		)      ->	Any:
				_lowerCamelCase				   =  [[] for _ in range(snake_case__		)]
				_lowerCamelCase				   =  size
		def __getitem__(       self     :     Dict     ,  snake_case__     :     int		)      ->	Iterator[Edge]:
				return iter(self._graph[vertex]		)
		@property
		def   _snake_case       (       self     :     str		)      ->	int:
				return self._size
		def   _snake_case       (       self     :     List[Any]     ,  snake_case__     :     int     ,  snake_case__     :     int     ,  snake_case__     :     int		)      ->	Tuple:
				if weight not in (0, 1):
						raise ValueError('Edge weight must be either 0 or 1.'		)

				if to_vertex < 0 or to_vertex >= self.size:
						raise ValueError('Vertex indexes must be in [0; size).'		)

				self._graph[from_vertex].append(Edge(snake_case__     ,  snake_case__		)		)



		def   _snake_case       (       self     :     Dict     ,  snake_case__     :     int     ,  snake_case__     :     int		)      ->	int | None:
				_lowerCamelCase				   =  deque([start_vertex]		)
				_lowerCamelCase				   =  [None] * self.size
				_lowerCamelCase				   =  0

				while queue:
						_lowerCamelCase				   =  queue.popleft()
						_lowerCamelCase				   =  distances[current_vertex]
						if current_distance is None:
								continue

						for edge in self[current_vertex]:
								_lowerCamelCase				   =  current_distance + edge.weight
								_lowerCamelCase				   =  distances[edge.destination_vertex]
								if (
								    isinstance(snake_case__     ,  snake_case__		)
								    and new_distance >= dest_vertex_distance
								):
										continue
								_lowerCamelCase				   =  new_distance
								if edge.weight == 0:
										queue.appendleft(edge.destination_vertex		)
								else:
										queue.append(edge.destination_vertex		)

				if distances[finish_vertex] is None:
						raise ValueError('No path from start_vertex to finish_vertex.'		)

				return distances[finish_vertex]


if __name__ == "__main__":
				import doctest

				doctest.testmod()