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  A_			(							a_      ,							a_      ,							unittest.TestCase			):
  _SCREAMING_SNAKE_CASE     					=      StableDiffusionPanoramaPipeline
  _SCREAMING_SNAKE_CASE     					=      TEXT_TO_IMAGE_PARAMS
  _SCREAMING_SNAKE_CASE     					=      TEXT_TO_IMAGE_BATCH_PARAMS
  _SCREAMING_SNAKE_CASE     					=      TEXT_TO_IMAGE_IMAGE_PARAMS
  _SCREAMING_SNAKE_CASE     					=      TEXT_TO_IMAGE_IMAGE_PARAMS

  def 		_UpperCAmelCase (		self			:   List[str]							):
   torch.manual_seed(0							)
   __a							=	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  ,				)
   __a							=	DDIMScheduler()
   torch.manual_seed(0							)
   __a							=	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							)
   __a							=	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  ,				)
   __a							=	CLIPTextModel(__SCREAMING_SNAKE_CASE							)
   __a							=	CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip"							)

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

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

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

   __a							=	np.array([0.61_86, 0.53_74, 0.49_15, 0.41_35, 0.41_14, 0.45_63, 0.51_28, 0.49_77, 0.47_57]							)

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

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

  def 		_UpperCAmelCase (		self			:   Tuple							):
   super().test_inference_batch_single_identical(batch_size=2  ,				expected_max_diff=3.25E-3							)

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

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

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

   __a							=	np.array([0.61_87, 0.53_75, 0.49_15, 0.41_36, 0.41_14, 0.45_63, 0.51_28, 0.49_76, 0.47_57]							)

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

  def 		_UpperCAmelCase (		self			:   Optional[int]							):
   __a							=	"cpu"  # ensure determinism for the device-dependent torch.Generator
   __a							=	self.get_dummy_components()
   __a							=	StableDiffusionPanoramaPipeline(**__SCREAMING_SNAKE_CASE							)
   __a							=	sd_pipe.to(__SCREAMING_SNAKE_CASE							)
   sd_pipe.set_progress_bar_config(disable=__SCREAMING_SNAKE_CASE							)

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

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

   __a							=	np.array([0.61_87, 0.53_75, 0.49_15, 0.41_36, 0.41_14, 0.45_63, 0.51_28, 0.49_76, 0.47_57]							)

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

  def 		_UpperCAmelCase (		self			:   List[str]							):
   __a							=	"cpu"  # ensure determinism for the device-dependent torch.Generator
   __a							=	self.get_dummy_components()
   __a							=	EulerAncestralDiscreteScheduler(
       beta_start=0.0_00_85  ,				beta_end=0.0_12  ,				beta_schedule="scaled_linear"							)
   __a							=	StableDiffusionPanoramaPipeline(**__SCREAMING_SNAKE_CASE							)
   __a							=	sd_pipe.to(__SCREAMING_SNAKE_CASE							)
   sd_pipe.set_progress_bar_config(disable=__SCREAMING_SNAKE_CASE							)

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

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

   __a							=	np.array([0.40_24, 0.65_10, 0.49_01, 0.53_78, 0.58_13, 0.56_22, 0.47_95, 0.44_67, 0.49_52]							)

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

  def 		_UpperCAmelCase (		self			:   List[str]							):
   __a							=	"cpu"  # ensure determinism for the device-dependent torch.Generator
   __a							=	self.get_dummy_components()
   __a							=	PNDMScheduler(
       beta_start=0.0_00_85  ,				beta_end=0.0_12  ,				beta_schedule="scaled_linear"  ,				skip_prk_steps=__SCREAMING_SNAKE_CASE							)
   __a							=	StableDiffusionPanoramaPipeline(**__SCREAMING_SNAKE_CASE							)
   __a							=	sd_pipe.to(__SCREAMING_SNAKE_CASE							)
   sd_pipe.set_progress_bar_config(disable=__SCREAMING_SNAKE_CASE							)

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

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

   __a							=	np.array([0.63_91, 0.62_91, 0.48_61, 0.51_34, 0.55_52, 0.45_78, 0.50_32, 0.50_23, 0.45_39]							)

   assert np.abs(image_slice.flatten() - expected_slice							).max() < 1E-2
@slow
@require_torch_gpu
class  A_			(							unittest.TestCase			):

  def 		_UpperCAmelCase (		self			:   int							):
   super().tearDown()
   gc.collect()
   torch.cuda.empty_cache()

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

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

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

   __a							=	np.array(
       [
           0.36_96_83_92,
           0.27_02_53_72,
           0.32_44_67_66,
           0.28_37_93_87,
           0.36_36_32_74,
           0.30_73_33_47,
           0.27_10_00_27,
           0.27_05_41_25,
           0.25_53_60_96,
       ]							)

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

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

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

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

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

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

     __a							=	np.array(
         [
             0.18_68_18_69,
             0.33_90_78_16,
             0.5_36_12_76,
             0.14_43_28_65,
             -0.02_85_66_11,
             -0.73_94_11_23,
             0.23_39_79_87,
             0.47_32_26_82,
             -0.37_82_31_64,
         ]							)
     assert np.abs(latents_slice.flatten() - expected_slice							).max() < 5E-2
    elif step == 2:
     __a							=	latents.detach().cpu().numpy()
     assert latents.shape == (1, 4, 64, 2_56)
     __a							=	latents[0, -3:, -3:, -1]

     __a							=	np.array(
         [
             0.18_53_96_45,
             0.33_98_72_48,
             0.5_37_85_59,
             0.14_43_71_42,
             -0.02_45_52_61,
             -0.7_33_83_17,
             0.23_99_07_55,
             0.47_35_62_72,
             -0.3_78_65_05,
         ]							)

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

   __a							=	False

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

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

  def 		_UpperCAmelCase (		self			:   Dict							):
   torch.cuda.empty_cache()
   torch.cuda.reset_max_memory_allocated()
   torch.cuda.reset_peak_memory_stats()

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

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

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



import os
from collections import namedtuple

import pytest

from datasets import ClassLabel, Features, Sequence, Value
from datasets.commands.test import TestCommand
from datasets.info import DatasetInfo, DatasetInfosDict


SCREAMING_SNAKE_CASE			:						Union[str, Any]				  =		namedtuple(
    """_TestCommandArgs""",
    [
        """dataset""",
        """name""",
        """cache_dir""",
        """data_dir""",
        """all_configs""",
        """save_infos""",
        """ignore_verifications""",
        """force_redownload""",
        """clear_cache""",
    ],
    defaults=[None, None, None, False, False, False, False, False],
)



def 	__A		(      _A				,					_A       ):



 """simple docstring"""


 return (abs(source - target       ) / target) < 0.01






@pytest.mark.integration
def 	__A		(      _A       ):



 """simple docstring"""


 __a							=	_TestCommandArgs(dataset=_A				,					all_configs=_A				,					save_infos=_A       )
 __a							=	TestCommand(*_A       )
 test_command.run()
 __a							=	os.path.join(_A				,					"README.md"       )
 assert os.path.exists(_A       )
 __a							=	DatasetInfosDict.from_directory(_A       )
 __a							=	DatasetInfosDict(
     {
         "default": DatasetInfo(
             features=Features(
                 {
                     "tokens": Sequence(Value("string"       )       ),
                     "ner_tags": Sequence(
                         ClassLabel(names=["O", "B-PER", "I-PER", "B-ORG", "I-ORG", "B-LOC", "I-LOC"]       )       ),
                     "langs": Sequence(Value("string"       )       ),
                     "spans": Sequence(Value("string"       )       ),
                 }       )				,					splits=[
                 {
                     "name": "train",
                     "num_bytes": 235_1563,
                     "num_examples": 1_0000,
                 },
                 {
                     "name": "validation",
                     "num_bytes": 23_8418,
                     "num_examples": 1000,
                 },
             ]				,					download_size=394_0680				,					dataset_size=258_9981				,					)
     }       )
 assert dataset_infos.keys() == expected_dataset_infos.keys()
 for key in DatasetInfo._INCLUDED_INFO_IN_YAML:
  __a							,  __a							=	getattr(dataset_infos["default"]				,					_A       ), getattr(expected_dataset_infos["default"]				,					_A       )
  if key == "num_bytes":
   assert is_apercent_close(_A				,					_A       )
  elif key == "splits":
   assert list(_A       ) == list(_A       )
   for split in result:
    assert result[split].name == expected[split].name
    assert result[split].num_examples == expected[split].num_examples
    assert is_apercent_close(result[split].num_bytes				,					expected[split].num_bytes       )
  else:
   result == expected


