


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 						_UpperCAmelCase      (     lowercase		,		lowercase		,		unittest.TestCase							):
		lowerCamelCase_	:			Dict       			=  StableDiffusionPanoramaPipeline
		lowerCamelCase_	:			Union[str, Any]       			=  TEXT_TO_IMAGE_PARAMS
		lowerCamelCase_	:			List[Any]       			=  TEXT_TO_IMAGE_BATCH_PARAMS
		lowerCamelCase_	:			List[str]       			=  TEXT_TO_IMAGE_IMAGE_PARAMS
		lowerCamelCase_	:			int       			=  TEXT_TO_IMAGE_IMAGE_PARAMS






		def _snake_case		(						self				:     str):
				torch.manual_seed(0)
				SCREAMING_SNAKE_CASE_				:Union[str, 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   ,  )
				SCREAMING_SNAKE_CASE_				:Union[str, Any]				     =						DDIMScheduler()
				torch.manual_seed(0)
				SCREAMING_SNAKE_CASE_				:Optional[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)
				SCREAMING_SNAKE_CASE_				: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=10_00   ,  )
				SCREAMING_SNAKE_CASE_				:str				     =						CLIPTextModel(UpperCAmelCase)
				SCREAMING_SNAKE_CASE_				:Union[str, Any]				     =						CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip")

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






		def _snake_case		(						self				:     str   ,  UpperCAmelCase				:     List[str]   ,  UpperCAmelCase				:     Dict=0):
				SCREAMING_SNAKE_CASE_				:List[str]				     =						torch.manual_seed(UpperCAmelCase)
				SCREAMING_SNAKE_CASE_				:Dict				     =						{
				    "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):
				SCREAMING_SNAKE_CASE_				:str				     =						"cpu"  # ensure determinism for the device-dependent torch.Generator
				SCREAMING_SNAKE_CASE_				:int				     =						self.get_dummy_components()
				SCREAMING_SNAKE_CASE_				:int				     =						StableDiffusionPanoramaPipeline(**UpperCAmelCase)
				SCREAMING_SNAKE_CASE_				:List[Any]				     =						sd_pipe.to(UpperCAmelCase)
				sd_pipe.set_progress_bar_config(disable=UpperCAmelCase)

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

				SCREAMING_SNAKE_CASE_				:Any				     =						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				:     Dict):
				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.25E-3)






		def _snake_case		(						self				:     Any):
				SCREAMING_SNAKE_CASE_				:List[str]				     =						"cpu"  # ensure determinism for the device-dependent torch.Generator
				SCREAMING_SNAKE_CASE_				:Union[str, Any]				     =						self.get_dummy_components()
				SCREAMING_SNAKE_CASE_				:List[Any]				     =						StableDiffusionPanoramaPipeline(**UpperCAmelCase)
				SCREAMING_SNAKE_CASE_				:Union[str, Any]				     =						sd_pipe.to(UpperCAmelCase)
				sd_pipe.set_progress_bar_config(disable=UpperCAmelCase)

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

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

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

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

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

				SCREAMING_SNAKE_CASE_				:List[str]				     =						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[Any]):
				SCREAMING_SNAKE_CASE_				:Union[str, Any]				     =						"cpu"  # ensure determinism for the device-dependent torch.Generator
				SCREAMING_SNAKE_CASE_				:Optional[int]				     =						self.get_dummy_components()
				SCREAMING_SNAKE_CASE_				:Union[str, Any]				     =						EulerAncestralDiscreteScheduler(
				    beta_start=0.00085   ,  beta_end=0.012   ,  beta_schedule="scaled_linear")
				SCREAMING_SNAKE_CASE_				:List[str]				     =						StableDiffusionPanoramaPipeline(**UpperCAmelCase)
				SCREAMING_SNAKE_CASE_				:List[str]				     =						sd_pipe.to(UpperCAmelCase)
				sd_pipe.set_progress_bar_config(disable=UpperCAmelCase)

				SCREAMING_SNAKE_CASE_				:str				     =						self.get_dummy_inputs(UpperCAmelCase)
				SCREAMING_SNAKE_CASE_				:Tuple				     =						sd_pipe(**UpperCAmelCase).images
				SCREAMING_SNAKE_CASE_				:Any				     =						image[0, -3:, -3:, -1]

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

				SCREAMING_SNAKE_CASE_				:Optional[int]				     =						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[int]):
				SCREAMING_SNAKE_CASE_				:Union[str, Any]				     =						"cpu"  # ensure determinism for the device-dependent torch.Generator
				SCREAMING_SNAKE_CASE_				:str				     =						self.get_dummy_components()
				SCREAMING_SNAKE_CASE_				:Optional[Any]				     =						PNDMScheduler(
				    beta_start=0.00085   ,  beta_end=0.012   ,  beta_schedule="scaled_linear"   ,  skip_prk_steps=UpperCAmelCase)
				SCREAMING_SNAKE_CASE_				:Optional[int]				     =						StableDiffusionPanoramaPipeline(**UpperCAmelCase)
				SCREAMING_SNAKE_CASE_				:str				     =						sd_pipe.to(UpperCAmelCase)
				sd_pipe.set_progress_bar_config(disable=UpperCAmelCase)

				SCREAMING_SNAKE_CASE_				:int				     =						self.get_dummy_inputs(UpperCAmelCase)
				SCREAMING_SNAKE_CASE_				:Tuple				     =						sd_pipe(**UpperCAmelCase).images
				SCREAMING_SNAKE_CASE_				:Optional[Any]				     =						image[0, -3:, -3:, -1]

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

				SCREAMING_SNAKE_CASE_				:Any				     =						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 						_UpperCAmelCase      (     unittest.TestCase							):






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






		def _snake_case		(						self				:     int   ,  UpperCAmelCase				:     Union[str, Any]=0):
				SCREAMING_SNAKE_CASE_				:Any				     =						torch.manual_seed(UpperCAmelCase)
				SCREAMING_SNAKE_CASE_				:Dict				     =						{
				    "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				:     str):
				SCREAMING_SNAKE_CASE_				:Optional[int]				     =						"stabilityai/stable-diffusion-2-base"
				SCREAMING_SNAKE_CASE_				:str				     =						DDIMScheduler.from_pretrained(UpperCAmelCase   ,  subfolder="scheduler")
				SCREAMING_SNAKE_CASE_				:List[str]				     =						StableDiffusionPanoramaPipeline.from_pretrained(UpperCAmelCase   ,  scheduler=UpperCAmelCase   ,  safety_checker=UpperCAmelCase)
				pipe.to(UpperCAmelCase)
				pipe.set_progress_bar_config(disable=UpperCAmelCase)
				pipe.enable_attention_slicing()

				SCREAMING_SNAKE_CASE_				:int				     =						self.get_inputs()
				SCREAMING_SNAKE_CASE_				:Dict				     =						pipe(**UpperCAmelCase).images
				SCREAMING_SNAKE_CASE_				:str				     =						image[0, -3:, -3:, -1].flatten()

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

				SCREAMING_SNAKE_CASE_				:Optional[Any]				     =						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]):
				SCREAMING_SNAKE_CASE_				:List[Any]				     =						StableDiffusionPanoramaPipeline.from_pretrained(
				    "stabilityai/stable-diffusion-2-base"   ,  safety_checker=UpperCAmelCase)
				SCREAMING_SNAKE_CASE_				:Tuple				     =						LMSDiscreteScheduler.from_config(pipe.scheduler.config)
				pipe.to(UpperCAmelCase)
				pipe.set_progress_bar_config(disable=UpperCAmelCase)
				pipe.enable_attention_slicing()

				SCREAMING_SNAKE_CASE_				:Dict				     =						self.get_inputs()
				SCREAMING_SNAKE_CASE_				:List[str]				     =						pipe(**UpperCAmelCase).images
				SCREAMING_SNAKE_CASE_				:Dict				     =						image[0, -3:, -3:, -1].flatten()

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

				SCREAMING_SNAKE_CASE_				:Tuple				     =						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[Any]):
				SCREAMING_SNAKE_CASE_				:List[Any]				     =						0

				def callback_fn(UpperCAmelCase				:     int   ,  UpperCAmelCase				:     int   ,  UpperCAmelCase				:     torch.FloatTensor) -> None:
						SCREAMING_SNAKE_CASE_				:Any				     =						True
						nonlocal number_of_steps
						number_of_steps += 1
						if step == 1:
								SCREAMING_SNAKE_CASE_				:str				     =						latents.detach().cpu().numpy()
								assert latents.shape == (1, 4, 64, 2_56)
								SCREAMING_SNAKE_CASE_				:Union[str, Any]				     =						latents[0, -3:, -3:, -1]

								SCREAMING_SNAKE_CASE_				:Tuple				     =						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:
								SCREAMING_SNAKE_CASE_				:int				     =						latents.detach().cpu().numpy()
								assert latents.shape == (1, 4, 64, 2_56)
								SCREAMING_SNAKE_CASE_				:Tuple				     =						latents[0, -3:, -3:, -1]

								SCREAMING_SNAKE_CASE_				:str				     =						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

				SCREAMING_SNAKE_CASE_				:List[Any]				     =						False

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

				SCREAMING_SNAKE_CASE_				:int				     =						self.get_inputs()
				pipe(**UpperCAmelCase   ,  callback=UpperCAmelCase   ,  callback_steps=1)
				assert callback_fn.has_been_called
				assert number_of_steps == 3







		def _snake_case		(						self				:     str):
				torch.cuda.empty_cache()
				torch.cuda.reset_max_memory_allocated()
				torch.cuda.reset_peak_memory_stats()

				SCREAMING_SNAKE_CASE_				:Union[str, Any]				     =						"stabilityai/stable-diffusion-2-base"
				SCREAMING_SNAKE_CASE_				:str				     =						DDIMScheduler.from_pretrained(UpperCAmelCase   ,  subfolder="scheduler")
				SCREAMING_SNAKE_CASE_				:Any				     =						StableDiffusionPanoramaPipeline.from_pretrained(UpperCAmelCase   ,  scheduler=UpperCAmelCase   ,  safety_checker=UpperCAmelCase)
				SCREAMING_SNAKE_CASE_				:Union[str, Any]				     =						pipe.to(UpperCAmelCase)
				pipe.set_progress_bar_config(disable=UpperCAmelCase)
				pipe.enable_attention_slicing(1)
				pipe.enable_sequential_cpu_offload()

				SCREAMING_SNAKE_CASE_				:int				     =						self.get_inputs()
				SCREAMING_SNAKE_CASE_				:Optional[int]				     =						pipe(**UpperCAmelCase)

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










import unittest

from transformers import is_torch_available, is_vision_available
from transformers.testing_utils import require_torch, require_vision, slow, torch_device


if is_torch_available():
		import torch

		from transformers import AutoModelForImageClassification

if is_vision_available():
		from transformers import AutoImageProcessor



@require_torch
@require_vision
class 						_UpperCAmelCase      (     unittest.TestCase							):






		@slow
		def _snake_case		(						self				:     List[str]):
				SCREAMING_SNAKE_CASE_				:Tuple				     =						AutoImageProcessor.from_pretrained("microsoft/dit-base-finetuned-rvlcdip")
				SCREAMING_SNAKE_CASE_				:int				     =						AutoModelForImageClassification.from_pretrained("microsoft/dit-base-finetuned-rvlcdip")
				model.to(UpperCAmelCase)

				from datasets import load_dataset

				SCREAMING_SNAKE_CASE_				:Optional[Any]				     =						load_dataset("nielsr/rvlcdip-demo")

				SCREAMING_SNAKE_CASE_				:Dict				     =						dataset["train"][0]["image"].convert("RGB")

				SCREAMING_SNAKE_CASE_				:Optional[Any]				     =						image_processor(UpperCAmelCase   ,  return_tensors="pt").to(UpperCAmelCase)

				# forward pass
				with torch.no_grad():
						SCREAMING_SNAKE_CASE_				:Optional[int]				     =						model(**UpperCAmelCase)
						SCREAMING_SNAKE_CASE_				:List[str]				     =						outputs.logits

				SCREAMING_SNAKE_CASE_				:int				     =						torch.Size((1, 16))
				self.assertEqual(logits.shape   ,  UpperCAmelCase)

				SCREAMING_SNAKE_CASE_				:Optional[Any]				     =						torch.tensor(
				    [-0.4158, -0.4092, -0.4347]   ,  device=UpperCAmelCase   ,  dtype=torch.float   ,  )
				self.assertTrue(torch.allclose(logits[0, :3]   ,  UpperCAmelCase   ,  atol=1E-4))






