





'''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 			_lowercase						(       __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     UpperCamelCase			( self    ) ->	List[str]:
      torch.manual_seed(0    )
      snake_case				=			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    ,   )
      snake_case				=			DDIMScheduler()
      torch.manual_seed(0    )
      snake_case				=			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    )
      snake_case				=			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    ,   )
      snake_case				=			CLIPTextModel(A__    )
      snake_case				=			CLIPTokenizer.from_pretrained('''hf-internal-testing/tiny-random-clip'''    )

      snake_case				=			{
          '''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    ) ->	Optional[Any]:
      snake_case				=			torch.manual_seed(A__    )
      snake_case				=			{
          '''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    ) ->	List[str]:
      snake_case				=			'''cpu'''  # ensure determinism for the device-dependent torch.Generator
      snake_case				=			self.get_dummy_components()
      snake_case				=			StableDiffusionPanoramaPipeline(**A__    )
      snake_case				=			sd_pipe.to(A__    )
      sd_pipe.set_progress_bar_config(disable=A__    )

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

      snake_case				=			np.array([0.6_1_8_6, 0.5_3_7_4, 0.4_9_1_5, 0.4_1_3_5, 0.4_1_1_4, 0.4_5_6_3, 0.5_1_2_8, 0.4_9_7_7, 0.4_7_5_7]    )

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







     def     UpperCamelCase			( self    ) ->	List[Any]:
      super().test_inference_batch_consistent(batch_sizes=[1, 2]    )







     def     UpperCamelCase			( self    ) ->	Optional[int]:
      super().test_inference_batch_single_identical(batch_size=2    ,   expected_max_diff=3.25e-3    )







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

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

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

      snake_case				=			np.array([0.6_1_8_7, 0.5_3_7_5, 0.4_9_1_5, 0.4_1_3_6, 0.4_1_1_4, 0.4_5_6_3, 0.5_1_2_8, 0.4_9_7_6, 0.4_7_5_7]    )

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







     def     UpperCamelCase			( self    ) ->	str:
      snake_case				=			'''cpu'''  # ensure determinism for the device-dependent torch.Generator
      snake_case				=			self.get_dummy_components()
      snake_case				=			StableDiffusionPanoramaPipeline(**A__    )
      snake_case				=			sd_pipe.to(A__    )
      sd_pipe.set_progress_bar_config(disable=A__    )

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

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

      snake_case				=			np.array([0.6_1_8_7, 0.5_3_7_5, 0.4_9_1_5, 0.4_1_3_6, 0.4_1_1_4, 0.4_5_6_3, 0.5_1_2_8, 0.4_9_7_6, 0.4_7_5_7]    )

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







     def     UpperCamelCase			( self    ) ->	Union[str, Any]:
      snake_case				=			'''cpu'''  # ensure determinism for the device-dependent torch.Generator
      snake_case				=			self.get_dummy_components()
      snake_case				=			EulerAncestralDiscreteScheduler(
          beta_start=0.0_0_0_8_5    ,   beta_end=0.0_1_2    ,   beta_schedule='''scaled_linear'''    )
      snake_case				=			StableDiffusionPanoramaPipeline(**A__    )
      snake_case				=			sd_pipe.to(A__    )
      sd_pipe.set_progress_bar_config(disable=A__    )

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

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

      snake_case				=			np.array([0.4_0_2_4, 0.6_5_1_0, 0.4_9_0_1, 0.5_3_7_8, 0.5_8_1_3, 0.5_6_2_2, 0.4_7_9_5, 0.4_4_6_7, 0.4_9_5_2]    )

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







     def     UpperCamelCase			( self    ) ->	int:
      snake_case				=			'''cpu'''  # ensure determinism for the device-dependent torch.Generator
      snake_case				=			self.get_dummy_components()
      snake_case				=			PNDMScheduler(
          beta_start=0.0_0_0_8_5    ,   beta_end=0.0_1_2    ,   beta_schedule='''scaled_linear'''    ,   skip_prk_steps=A__    )
      snake_case				=			StableDiffusionPanoramaPipeline(**A__    )
      snake_case				=			sd_pipe.to(A__    )
      sd_pipe.set_progress_bar_config(disable=A__    )

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

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

      snake_case				=			np.array([0.6_3_9_1, 0.6_2_9_1, 0.4_8_6_1, 0.5_1_3_4, 0.5_5_5_2, 0.4_5_7_8, 0.5_0_3_2, 0.5_0_2_3, 0.4_5_3_9]    )

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




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







     def     UpperCamelCase			( self    ) ->	Dict:
      super().tearDown()
      gc.collect()
      torch.cuda.empty_cache()







     def     UpperCamelCase			( self    ,   A__=0    ) ->	str:
      snake_case				=			torch.manual_seed(A__    )
      snake_case				=			{
          '''prompt''': '''a photo of the dolomites''',
          '''generator''': generator,
          '''num_inference_steps''': 3,
          '''guidance_scale''': 7.5,
          '''output_type''': '''numpy''',
      }
      return inputs







     def     UpperCamelCase			( self    ) ->	Any:
      snake_case				=			'''stabilityai/stable-diffusion-2-base'''
      snake_case				=			DDIMScheduler.from_pretrained(A__    ,   subfolder='''scheduler'''    )
      snake_case				=			StableDiffusionPanoramaPipeline.from_pretrained(A__    ,   scheduler=A__    ,   safety_checker=A__    )
      pipe.to(A__    )
      pipe.set_progress_bar_config(disable=A__    )
      pipe.enable_attention_slicing()

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

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

      snake_case				=			np.array(
          [
              0.3_6_9_6_8_3_9_2,
              0.2_7_0_2_5_3_7_2,
              0.3_2_4_4_6_7_6_6,
              0.2_8_3_7_9_3_8_7,
              0.3_6_3_6_3_2_7_4,
              0.3_0_7_3_3_3_4_7,
              0.2_7_1_0_0_0_2_7,
              0.2_7_0_5_4_1_2_5,
              0.2_5_5_3_6_0_9_6,
          ]    )

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







     def     UpperCamelCase			( self    ) ->	Tuple:
      snake_case				=			StableDiffusionPanoramaPipeline.from_pretrained(
          '''stabilityai/stable-diffusion-2-base'''    ,   safety_checker=A__    )
      snake_case				=			LMSDiscreteScheduler.from_config(pipe.scheduler.config    )
      pipe.to(A__    )
      pipe.set_progress_bar_config(disable=A__    )
      pipe.enable_attention_slicing()

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

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

      snake_case				=			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    ) ->	Union[str, Any]:
      snake_case				=			0

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

        snake_case				=			np.array(
            [
                0.1_8_6_8_1_8_6_9,
                0.3_3_9_0_7_8_1_6,
                0.5_3_6_1_2_7_6,
                0.1_4_4_3_2_8_6_5,
                -0.0_2_8_5_6_6_1_1,
                -0.7_3_9_4_1_1_2_3,
                0.2_3_3_9_7_9_8_7,
                0.4_7_3_2_2_6_8_2,
                -0.3_7_8_2_3_1_6_4,
            ]    )
        assert np.abs(latents_slice.flatten() - expected_slice    ).max() < 5e-2
       elif step == 2:
        snake_case				=			latents.detach().cpu().numpy()
        assert latents.shape == (1, 4, 64, 2_56)
        snake_case				=			latents[0, -3:, -3:, -1]

        snake_case				=			np.array(
            [
                0.1_8_5_3_9_6_4_5,
                0.3_3_9_8_7_2_4_8,
                0.5_3_7_8_5_5_9,
                0.1_4_4_3_7_1_4_2,
                -0.0_2_4_5_5_2_6_1,
                -0.7_3_3_8_3_1_7,
                0.2_3_9_9_0_7_5_5,
                0.4_7_3_5_6_2_7_2,
                -0.3_7_8_6_5_0_5,
            ]    )

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

      snake_case				=			False

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

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







     def     UpperCamelCase			( self    ) ->	Union[str, Any]:
      torch.cuda.empty_cache()
      torch.cuda.reset_max_memory_allocated()
      torch.cuda.reset_peak_memory_stats()

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

      snake_case				=			self.get_inputs()
      snake_case				=			pipe(**A__    )

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








'''simple docstring'''







import inspect
import re
from hashlib import shaaaa
from typing import Dict, List

from .arrow import arrow
from .audiofolder import audiofolder
from .csv import csv
from .imagefolder import imagefolder
from .json import json
from .pandas import pandas
from .parquet import parquet
from .sql import sql  # noqa F401
from .text import text

def    __UpperCamelCase	(	a    : List[str]     ) ->str:
 snake_case				=			[]
 for line in lines:
  snake_case				=			re.sub(R'''#.*'''					,   ''''''					,   a     )  # remove comments
  if line:
   filtered_lines.append(a     )
 snake_case				=			'''\n'''.join(a     )

 # Make a hash from all this code
 snake_case				=			full_str.encode('''utf-8'''     )
 return shaaaa(a     ).hexdigest()


# get importable module names and hash for caching
_lowercase       						=			{
    'csv': (csv.__name__, _hash_python_lines(inspect.getsource(csv).splitlines())),
    'json': (json.__name__, _hash_python_lines(inspect.getsource(json).splitlines())),
    'pandas': (pandas.__name__, _hash_python_lines(inspect.getsource(pandas).splitlines())),
    'parquet': (parquet.__name__, _hash_python_lines(inspect.getsource(parquet).splitlines())),
    'arrow': (arrow.__name__, _hash_python_lines(inspect.getsource(arrow).splitlines())),
    'text': (text.__name__, _hash_python_lines(inspect.getsource(text).splitlines())),
    'imagefolder': (imagefolder.__name__, _hash_python_lines(inspect.getsource(imagefolder).splitlines())),
    'audiofolder': (audiofolder.__name__, _hash_python_lines(inspect.getsource(audiofolder).splitlines())),
}

# Used to infer the module to use based on the data files extensions
_lowercase       						=			{
    '.csv': ('csv', {}),
    '.tsv': ('csv', {'sep': '\t'}),
    '.json': ('json', {}),
    '.jsonl': ('json', {}),
    '.parquet': ('parquet', {}),
    '.arrow': ('arrow', {}),
    '.txt': ('text', {}),
}
_EXTENSION_TO_MODULE.update({ext: ('imagefolder', {}) for ext in imagefolder.ImageFolder.EXTENSIONS})
_EXTENSION_TO_MODULE.update({ext.upper(): ('imagefolder', {}) for ext in imagefolder.ImageFolder.EXTENSIONS})
_EXTENSION_TO_MODULE.update({ext: ('audiofolder', {}) for ext in audiofolder.AudioFolder.EXTENSIONS})
_EXTENSION_TO_MODULE.update({ext.upper(): ('audiofolder', {}) for ext in audiofolder.AudioFolder.EXTENSIONS})
_lowercase       						=			{'imagefolder', 'audiofolder'}

# Used to filter data files based on extensions given a module name
_lowercase       						=			{}
for _ext, (_module, _) in _EXTENSION_TO_MODULE.items():
   _MODULE_TO_EXTENSIONS.setdefault(_module, []).append(_ext)

_MODULE_TO_EXTENSIONS["imagefolder"].append('.zip')






_MODULE_TO_EXTENSIONS["audiofolder"].append('.zip')

