"""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 	UpperCAmelCase_      (       _a					,				_a					,				unittest.TestCase):
      lowerCamelCase__						:       int					    =    StableDiffusionPanoramaPipeline
      lowerCamelCase__						:       Any					    =    TEXT_TO_IMAGE_PARAMS
      lowerCamelCase__						:       str					    =    TEXT_TO_IMAGE_BATCH_PARAMS
      lowerCamelCase__						:       Union[str, Any]					    =    TEXT_TO_IMAGE_IMAGE_PARAMS
      lowerCamelCase__						:       Tuple					    =    TEXT_TO_IMAGE_IMAGE_PARAMS


      def 						_UpperCAmelCase			(  self				)			->   Dict:
             torch.manual_seed(0				)
             lowercase__     :						Tuple         =      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						,     )
             lowercase__     :						int         =      DDIMScheduler()
             torch.manual_seed(0				)
             lowercase__     :						Dict         =      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				)
             lowercase__     :						Union[str, Any]         =      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						,     )
             lowercase__     :						str         =      CLIPTextModel(lowercase__				)
             lowercase__     :						Optional[Any]         =      CLIPTokenizer.from_pretrained('hf-internal-testing/tiny-random-clip'				)

             lowercase__     :						List[str]         =      {
                 'unet': unet,
                 'scheduler': scheduler,
                 'vae': vae,
                 'text_encoder': text_encoder,
                 'tokenizer': tokenizer,
                 'safety_checker': None,
                 'feature_extractor': None,
             }
             return components


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

             lowercase__     :						Optional[Any]         =      self.get_dummy_inputs(lowercase__				)
             lowercase__     :						Optional[Any]         =      sd_pipe(**lowercase__				).images
             lowercase__     :						Optional[int]         =      image[0, -3:, -3:, -1]
             assert image.shape == (1, 6_4, 6_4, 3)

             lowercase__     :						str         =      np.array([0.6_186, 0.5_374, 0.4_915, 0.4_135, 0.4_114, 0.4_563, 0.5_128, 0.4_977, 0.4_757]				)

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


      def 						_UpperCAmelCase			(  self				)			->   Dict:
             super().test_inference_batch_consistent(batch_sizes=[1, 2]				)


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


      def 						_UpperCAmelCase			(  self				)			->   Optional[Any]:
             lowercase__     :						Union[str, Any]         =      'cpu'  # ensure determinism for the device-dependent torch.Generator
             lowercase__     :						List[Any]         =      self.get_dummy_components()
             lowercase__     :						Any         =      StableDiffusionPanoramaPipeline(**lowercase__				)
             lowercase__     :						Tuple         =      sd_pipe.to(lowercase__				)
             sd_pipe.set_progress_bar_config(disable=lowercase__				)

             lowercase__     :						str         =      self.get_dummy_inputs(lowercase__				)
             lowercase__     :						Optional[Any]         =      'french fries'
             lowercase__     :						Dict         =      sd_pipe(**lowercase__						,     negative_prompt=lowercase__				)
             lowercase__     :						List[str]         =      output.images
             lowercase__     :						Any         =      image[0, -3:, -3:, -1]

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

             lowercase__     :						Optional[Any]         =      np.array([0.6_187, 0.5_375, 0.4_915, 0.4_136, 0.4_114, 0.4_563, 0.5_128, 0.4_976, 0.4_757]				)

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


      def 						_UpperCAmelCase			(  self				)			->   Optional[Any]:
             lowercase__     :						str         =      'cpu'  # ensure determinism for the device-dependent torch.Generator
             lowercase__     :						List[Any]         =      self.get_dummy_components()
             lowercase__     :						Optional[Any]         =      StableDiffusionPanoramaPipeline(**lowercase__				)
             lowercase__     :						Dict         =      sd_pipe.to(lowercase__				)
             sd_pipe.set_progress_bar_config(disable=lowercase__				)

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

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

             lowercase__     :						Any         =      np.array([0.6_187, 0.5_375, 0.4_915, 0.4_136, 0.4_114, 0.4_563, 0.5_128, 0.4_976, 0.4_757]				)

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


      def 						_UpperCAmelCase			(  self				)			->   Optional[int]:
             lowercase__     :						Union[str, Any]         =      'cpu'  # ensure determinism for the device-dependent torch.Generator
             lowercase__     :						str         =      self.get_dummy_components()
             lowercase__     :						int         =      EulerAncestralDiscreteScheduler(
                 beta_start=0.00_085						,     beta_end=0.012						,     beta_schedule='scaled_linear'				)
             lowercase__     :						List[str]         =      StableDiffusionPanoramaPipeline(**lowercase__				)
             lowercase__     :						Any         =      sd_pipe.to(lowercase__				)
             sd_pipe.set_progress_bar_config(disable=lowercase__				)

             lowercase__     :						Any         =      self.get_dummy_inputs(lowercase__				)
             lowercase__     :						int         =      sd_pipe(**lowercase__				).images
             lowercase__     :						List[str]         =      image[0, -3:, -3:, -1]

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

             lowercase__     :						Dict         =      np.array([0.4_024, 0.6_510, 0.4_901, 0.5_378, 0.5_813, 0.5_622, 0.4_795, 0.4_467, 0.4_952]				)

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


      def 						_UpperCAmelCase			(  self				)			->   Optional[int]:
             lowercase__     :						Optional[Any]         =      'cpu'  # ensure determinism for the device-dependent torch.Generator
             lowercase__     :						str         =      self.get_dummy_components()
             lowercase__     :						Optional[Any]         =      PNDMScheduler(
                 beta_start=0.00_085						,     beta_end=0.012						,     beta_schedule='scaled_linear'						,     skip_prk_steps=lowercase__				)
             lowercase__     :						Union[str, Any]         =      StableDiffusionPanoramaPipeline(**lowercase__				)
             lowercase__     :						Tuple         =      sd_pipe.to(lowercase__				)
             sd_pipe.set_progress_bar_config(disable=lowercase__				)

             lowercase__     :						List[str]         =      self.get_dummy_inputs(lowercase__				)
             lowercase__     :						Tuple         =      sd_pipe(**lowercase__				).images
             lowercase__     :						Optional[int]         =      image[0, -3:, -3:, -1]

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

             lowercase__     :						Optional[Any]         =      np.array([0.6_391, 0.6_291, 0.4_861, 0.5_134, 0.5_552, 0.4_578, 0.5_032, 0.5_023, 0.4_539]				)

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





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


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


      def 						_UpperCAmelCase			(  self						,     a=0				)			->   Union[str, Any]:
             lowercase__     :						Union[str, Any]         =      torch.manual_seed(lowercase__				)
             lowercase__     :						Tuple         =      {
                 'prompt': 'a photo of the dolomites',
                 'generator': generator,
                 'num_inference_steps': 3,
                 'guidance_scale': 7.5,
                 'output_type': 'numpy',
             }
             return inputs


      def 						_UpperCAmelCase			(  self				)			->   Optional[Any]:
             lowercase__     :						Tuple         =      'stabilityai/stable-diffusion-2-base'
             lowercase__     :						Tuple         =      DDIMScheduler.from_pretrained(lowercase__						,     subfolder='scheduler'				)
             lowercase__     :						int         =      StableDiffusionPanoramaPipeline.from_pretrained(lowercase__						,     scheduler=lowercase__						,     safety_checker=lowercase__				)
             pipe.to(lowercase__				)
             pipe.set_progress_bar_config(disable=lowercase__				)
             pipe.enable_attention_slicing()

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

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

             lowercase__     :						Dict         =      np.array(
                 [
                     0.36_968_392,
                     0.27_025_372,
                     0.32_446_766,
                     0.28_379_387,
                     0.36_363_274,
                     0.30_733_347,
                     0.27_100_027,
                     0.27_054_125,
                     0.25_536_096,
                 ]				)

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


      def 						_UpperCAmelCase			(  self				)			->   List[Any]:
             lowercase__     :						int         =      StableDiffusionPanoramaPipeline.from_pretrained(
                 'stabilityai/stable-diffusion-2-base'						,     safety_checker=lowercase__				)
             lowercase__     :						str         =      LMSDiscreteScheduler.from_config(pipe.scheduler.config				)
             pipe.to(lowercase__				)
             pipe.set_progress_bar_config(disable=lowercase__				)
             pipe.enable_attention_slicing()

             lowercase__     :						Optional[int]         =      self.get_inputs()
             lowercase__     :						Dict         =      pipe(**lowercase__				).images
             lowercase__     :						Union[str, Any]         =      image[0, -3:, -3:, -1].flatten()

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

             lowercase__     :						Any         =      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 						_UpperCAmelCase			(  self				)			->   int:
             lowercase__     :						Optional[Any]         =      0

             def callback_fn(a						,     a						,     a				) -> None:
                    lowercase__     :						str         =      True
                    nonlocal number_of_steps
                    number_of_steps += 1
                    if step == 1:
                           lowercase__     :						Union[str, Any]         =      latents.detach().cpu().numpy()
                           assert latents.shape == (1, 4, 6_4, 2_5_6)
                           lowercase__     :						str         =      latents[0, -3:, -3:, -1]

                           lowercase__     :						Any         =      np.array(
                               [
                                   0.18_681_869,
                                   0.33_907_816,
                                   0.5_361_276,
                                   0.14_432_865,
                                   -0.02_856_611,
                                   -0.73_941_123,
                                   0.23_397_987,
                                   0.47_322_682,
                                   -0.37_823_164,
                               ]				)
                           assert np.abs(latents_slice.flatten() - expected_slice				).max() < 5e-2
                    elif step == 2:
                           lowercase__     :						Union[str, Any]         =      latents.detach().cpu().numpy()
                           assert latents.shape == (1, 4, 6_4, 2_5_6)
                           lowercase__     :						List[Any]         =      latents[0, -3:, -3:, -1]

                           lowercase__     :						Optional[int]         =      np.array(
                               [
                                   0.18_539_645,
                                   0.33_987_248,
                                   0.5_378_559,
                                   0.14_437_142,
                                   -0.02_455_261,
                                   -0.7_338_317,
                                   0.23_990_755,
                                   0.47_356_272,
                                   -0.3_786_505,
                               ]				)

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

             lowercase__     :						Dict         =      False

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

             lowercase__     :						str         =      self.get_inputs()
             pipe(**lowercase__						,     callback=lowercase__						,     callback_steps=1				)
             assert callback_fn.has_been_called
             assert number_of_steps == 3


      def 						_UpperCAmelCase			(  self				)			->   Tuple:
             torch.cuda.empty_cache()
             torch.cuda.reset_max_memory_allocated()
             torch.cuda.reset_peak_memory_stats()

             lowercase__     :						Tuple         =      'stabilityai/stable-diffusion-2-base'
             lowercase__     :						List[str]         =      DDIMScheduler.from_pretrained(lowercase__						,     subfolder='scheduler'				)
             lowercase__     :						Dict         =      StableDiffusionPanoramaPipeline.from_pretrained(lowercase__						,     scheduler=lowercase__						,     safety_checker=lowercase__				)
             lowercase__     :						Optional[int]         =      pipe.to(lowercase__				)
             pipe.set_progress_bar_config(disable=lowercase__				)
             pipe.enable_attention_slicing(1				)
             pipe.enable_sequential_cpu_offload()

             lowercase__     :						Dict         =      self.get_inputs()
             lowercase__     :						Dict         =      pipe(**lowercase__				)

             lowercase__     :						Any         =      torch.cuda.max_memory_allocated()
             # make sure that less than 5.2 GB is allocated
             assert mem_bytes < 5.5 * 1_0**9




import os
import re
import shutil
import sys
import tempfile
import unittest

import black


_lowerCamelCase      :					Union[str, Any]					    =      os.path.abspath(os.path.dirname(os.path.dirname(os.path.dirname(__file__))))
sys.path.append(os.path.join(git_repo_path, '''utils'''))

import check_copies  # noqa: E402


# This is the reference code that will be used in the tests.
# If DDPMSchedulerOutput is changed in scheduling_ddpm.py, this code needs to be manually updated.
_lowerCamelCase      :					str					    =      '''    \"""
    Output class for the scheduler\'s step function output.

    Args:
        prev_sample (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)` for images):
            Computed sample (x_{t-1}) of previous timestep. `prev_sample` should be used as next model input in the
            denoising loop.
        pred_original_sample (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)` for images):
            The predicted denoised sample (x_{0}) based on the model output from the current timestep.
            `pred_original_sample` can be used to preview progress or for guidance.
    \"""

    prev_sample: torch.FloatTensor
    pred_original_sample: Optional[torch.FloatTensor] = None
'''





class        lowerCAmelCase__  (					unittest.TestCase			):



       '''simple docstring'''







       def     __UpperCamelCase     (							self     ):







             '''simple docstring'''



             __A    =tempfile.mkdtemp()
             os.makedirs(os.path.join(self.diffusers_dir ,   '''schedulers/'''     )     )
             __A    =self.diffusers_dir
             shutil.copy(
                 os.path.join(lowercase__ ,   '''src/diffusers/schedulers/scheduling_ddpm.py'''     ) ,   os.path.join(self.diffusers_dir ,   '''schedulers/scheduling_ddpm.py'''     ) ,   )

       def     __UpperCamelCase     (							self     ):







             '''simple docstring'''



             __A    ='''src/diffusers'''
             shutil.rmtree(self.diffusers_dir     )

       def     __UpperCamelCase     (							self ,   lowercase__ ,   lowercase__ ,   lowercase__ ,   lowercase__=None     ):







             '''simple docstring'''



             __A    =comment + f'''\nclass {class_name}(nn.Module):\n''' + class_code
             if overwrite_result is not None:
                   __A    =comment + f'''\nclass {class_name}(nn.Module):\n''' + overwrite_result
             __A    =black.Mode(target_versions={black.TargetVersion.PYaa} ,   line_length=1_1_9     )
             __A    =black.format_str(lowercase__ ,   mode=lowercase__     )
             __A    =os.path.join(self.diffusers_dir ,   '''new_code.py'''     )
             with open(lowercase__ ,   '''w''' ,   newline='''\n'''     ) as f:
                   f.write(lowercase__     )
             if overwrite_result is None:
                   self.assertTrue(len(check_copies.is_copy_consistent(lowercase__     )     ) == 0     )
             else:
                   check_copies.is_copy_consistent(f.name ,   overwrite=lowercase__     )
                   with open(lowercase__ ,   '''r'''     ) as f:
                         self.assertTrue(f.read() ,   lowercase__     )

       def     __UpperCamelCase     (							self     ):







             '''simple docstring'''



             __A    =check_copies.find_code_in_diffusers('''schedulers.scheduling_ddpm.DDPMSchedulerOutput'''     )
             self.assertEqual(lowercase__ ,   lowercase__     )


       def     __UpperCamelCase     (							self     ):







             '''simple docstring'''



             self.check_copy_consistency(
                 '''# Copied from diffusers.schedulers.scheduling_ddpm.DDPMSchedulerOutput''' ,   '''DDPMSchedulerOutput''' ,   REFERENCE_CODE + '''\n''' ,   )

             # With no empty line at the end
             self.check_copy_consistency(
                 '''# Copied from diffusers.schedulers.scheduling_ddpm.DDPMSchedulerOutput''' ,   '''DDPMSchedulerOutput''' ,   lowercase__ ,   )

             # Copy consistency with rename
             self.check_copy_consistency(
                 '''# Copied from diffusers.schedulers.scheduling_ddpm.DDPMSchedulerOutput with DDPM->Test''' ,   '''TestSchedulerOutput''' ,   re.sub('''DDPM''' ,   '''Test''' ,   lowercase__     ) ,   )

             # Copy consistency with a really long name
             __A    ='''TestClassWithAReallyLongNameBecauseSomePeopleLikeThatForSomeReason'''
             self.check_copy_consistency(
                 f'''# Copied from diffusers.schedulers.scheduling_ddpm.DDPMSchedulerOutput with DDPM->{long_class_name}''' ,   f'''{long_class_name}SchedulerOutput''' ,   re.sub('''Bert''' ,   lowercase__ ,   lowercase__     ) ,   )

             # Copy consistency with overwrite
             self.check_copy_consistency(
                 '''# Copied from diffusers.schedulers.scheduling_ddpm.DDPMSchedulerOutput with DDPM->Test''' ,   '''TestSchedulerOutput''' ,   lowercase__ ,   overwrite_result=re.sub('''DDPM''' ,   '''Test''' ,   lowercase__     ) ,   )


