

import gc
import unittest

import numpy as np
import torch
from transformers import CLIPTextConfig, CLIPTextModel, CLIPTokenizer

from diffusers import (
    AutoencoderKL,
    DDIMParallelScheduler,
    DDPMParallelScheduler,
    StableDiffusionParadigmsPipeline,
    UNetaDConditionModel,
)
from diffusers.utils import slow, torch_device
from diffusers.utils.testing_utils import (
    enable_full_determinism,
    require_torch_gpu,
)

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()






class     __magic_name__     (							_snake_case						,    _snake_case						,    unittest.TestCase     ):
 UpperCAmelCase				     =    StableDiffusionParadigmsPipeline
 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		:			Tuple     )     ->   List[str]:
   torch.manual_seed(0     )
   UpperCAmelCase          =	UNetaDConditionModel(
       block_out_channels=(3_2, 6_4)      ,    layers_per_block=2      ,    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      ,    attention_head_dim=(2, 4)      ,    use_linear_projection=lowerCAmelCase__      ,    )
   UpperCAmelCase          =	DDIMParallelScheduler(
       beta_start=0.00_085      ,    beta_end=0.012      ,    beta_schedule="scaled_linear"      ,    clip_sample=lowerCAmelCase__      ,    set_alpha_to_one=lowerCAmelCase__      ,    )
   torch.manual_seed(0     )
   UpperCAmelCase          =	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      ,    sample_size=1_2_8      ,    )
   torch.manual_seed(0     )
   UpperCAmelCase          =	CLIPTextConfig(
       bos_token_id=0      ,    eos_token_id=2      ,    hidden_size=3_2      ,    intermediate_size=3_7      ,    layer_norm_eps=1e-0_5      ,    num_attention_heads=4      ,    num_hidden_layers=5      ,    pad_token_id=1      ,    vocab_size=1_0_0_0      ,    hidden_act="gelu"      ,    projection_dim=5_1_2      ,    )
   UpperCAmelCase          =	CLIPTextModel(lowerCAmelCase__     )
   UpperCAmelCase          =	CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip"     )

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


 def _UpperCamelCase						(    self		:			Optional[int]      ,    lowerCAmelCase__		:			Tuple      ,    lowerCAmelCase__		:			Dict=0     )     ->   Optional[int]:
   if str(lowerCAmelCase__     ).startswith("mps"     ):
     UpperCAmelCase          =	torch.manual_seed(lowerCAmelCase__     )
   else:
     UpperCAmelCase          =	torch.Generator(device=lowerCAmelCase__     ).manual_seed(lowerCAmelCase__     )
   UpperCAmelCase          =	{
       "prompt": "a photograph of an astronaut riding a horse",
       "generator": generator,
       "num_inference_steps": 1_0,
       "guidance_scale": 6.0,
       "output_type": "numpy",
       "parallel": 3,
       "debug": True,
   }
   return inputs


 def _UpperCamelCase						(    self		:			int     )     ->   int:
   UpperCAmelCase          =	"cpu"  # ensure determinism for the device-dependent torch.Generator
   UpperCAmelCase          =	self.get_dummy_components()
   UpperCAmelCase          =	StableDiffusionParadigmsPipeline(**lowerCAmelCase__     )
   UpperCAmelCase          =	sd_pipe.to(lowerCAmelCase__     )
   sd_pipe.set_progress_bar_config(disable=lowerCAmelCase__     )

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

   UpperCAmelCase          =	np.array([0.4_773, 0.5_417, 0.4_723, 0.4_925, 0.5_631, 0.4_752, 0.5_240, 0.4_935, 0.5_023]     )

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


 def _UpperCamelCase						(    self		:			Any     )     ->   Any:
   UpperCAmelCase          =	"cpu"  # ensure determinism for the device-dependent torch.Generator
   UpperCAmelCase          =	self.get_dummy_components()
   torch.manual_seed(0     )
   UpperCAmelCase          =	DDPMParallelScheduler()
   torch.manual_seed(0     )
   UpperCAmelCase          =	StableDiffusionParadigmsPipeline(**lowerCAmelCase__     )
   UpperCAmelCase          =	sd_pipe.to(lowerCAmelCase__     )
   sd_pipe.set_progress_bar_config(disable=lowerCAmelCase__     )

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

   UpperCAmelCase          =	np.array([0.3_573, 0.4_420, 0.4_960, 0.4_799, 0.3_796, 0.3_879, 0.4_819, 0.4_365, 0.4_468]     )

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


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


 def _UpperCamelCase						(    self		:			Any     )     ->   Tuple:
   super().test_inference_batch_single_identical(batch_size=2      ,    expected_max_diff=3e-3     )


 def _UpperCamelCase						(    self		:			str     )     ->   Any:
   UpperCAmelCase          =	"cpu"  # ensure determinism for the device-dependent torch.Generator
   UpperCAmelCase          =	self.get_dummy_components()
   UpperCAmelCase          =	StableDiffusionParadigmsPipeline(**lowerCAmelCase__     )
   UpperCAmelCase          =	sd_pipe.to(lowerCAmelCase__     )
   sd_pipe.set_progress_bar_config(disable=lowerCAmelCase__     )

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

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

   UpperCAmelCase          =	np.array([0.4_771, 0.5_420, 0.4_683, 0.4_918, 0.5_636, 0.4_725, 0.5_230, 0.4_923, 0.5_015]     )

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






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


 def _UpperCamelCase						(    self		:			Optional[Any]     )     ->   Dict:
   super().tearDown()
   gc.collect()
   torch.cuda.empty_cache()


 def _UpperCamelCase						(    self		:			Any      ,    lowerCAmelCase__		:			Union[str, Any]=0     )     ->   Any:
   UpperCAmelCase          =	torch.Generator(device=lowerCAmelCase__     ).manual_seed(lowerCAmelCase__     )
   UpperCAmelCase          =	{
       "prompt": "a photograph of an astronaut riding a horse",
       "generator": generator,
       "num_inference_steps": 1_0,
       "guidance_scale": 7.5,
       "output_type": "numpy",
       "parallel": 3,
       "debug": True,
   }
   return inputs


 def _UpperCamelCase						(    self		:			Any     )     ->   Optional[Any]:
   UpperCAmelCase          =	"stabilityai/stable-diffusion-2-base"
   UpperCAmelCase          =	DDIMParallelScheduler.from_pretrained(lowerCAmelCase__      ,    subfolder="scheduler"     )
   UpperCAmelCase          =	StableDiffusionParadigmsPipeline.from_pretrained(lowerCAmelCase__      ,    scheduler=lowerCAmelCase__      ,    safety_checker=lowerCAmelCase__     )
   pipe.to(lowerCAmelCase__     )
   pipe.set_progress_bar_config(disable=lowerCAmelCase__     )
   pipe.enable_attention_slicing()

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

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

   UpperCAmelCase          =	np.array([0.9_622, 0.9_602, 0.9_748, 0.9_591, 0.9_630, 0.9_691, 0.9_661, 0.9_631, 0.9_741]     )

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






import sys
from collections.abc import Mapping

import numpy as np
import pyarrow as pa

from .. import config
from ..utils.py_utils import map_nested
from .formatting import TensorFormatter






class     __magic_name__     (							TensorFormatter[Mapping, np.ndarray, Mapping]     ):


 def __init__(    self		:			Optional[int]      ,    lowerCAmelCase__		:			str=None      ,    **lowerCAmelCase__		:			int     )     ->   Optional[Any]:
   super().__init__(features=lowerCAmelCase__     )
   UpperCAmelCase          =	np_array_kwargs


 def _UpperCamelCase						(    self		:			Any      ,    lowerCAmelCase__		:			Dict     )     ->   Union[str, Any]:
   if isinstance(lowerCAmelCase__      ,    lowerCAmelCase__     ):
     if column and all(
         isinstance(lowerCAmelCase__      ,    np.ndarray     ) and x.shape == column[0].shape and x.dtype == column[0].dtype for x in column     ):
       return np.stack(lowerCAmelCase__     )
     else:
       # don't use np.array(column, dtype=object)
       # since it fails in certain cases
       # see https://stackoverflow.com/q/51005699
       UpperCAmelCase          =	np.empty(len(lowerCAmelCase__     )      ,    dtype=lowerCAmelCase__     )
       UpperCAmelCase          =	column
       return out
   return column


 def _UpperCamelCase						(    self		:			Optional[Any]      ,    lowerCAmelCase__		:			Union[str, Any]     )     ->   List[Any]:
   if isinstance(lowerCAmelCase__      ,    (str, bytes, type(lowerCAmelCase__     ))     ):
     return value
   elif isinstance(lowerCAmelCase__      ,    (np.character, np.ndarray)     ) and np.issubdtype(value.dtype      ,    np.character     ):
     return value
   elif isinstance(lowerCAmelCase__      ,    np.number     ):
     return value

   UpperCAmelCase          =	{}

   if isinstance(lowerCAmelCase__      ,    np.ndarray     ) and np.issubdtype(value.dtype      ,    np.integer     ):
     UpperCAmelCase          =	{"dtype": np.intaa}
   elif isinstance(lowerCAmelCase__      ,    np.ndarray     ) and np.issubdtype(value.dtype      ,    np.floating     ):
     UpperCAmelCase          =	{"dtype": np.floataa}
   elif config.PIL_AVAILABLE and "PIL" in sys.modules:
     import PIL.Image

     if isinstance(lowerCAmelCase__      ,    PIL.Image.Image     ):
       return np.asarray(lowerCAmelCase__      ,    **self.np_array_kwargs     )

   return np.asarray(lowerCAmelCase__      ,    **{**default_dtype, **self.np_array_kwargs}     )


 def _UpperCamelCase						(    self		:			List[Any]      ,    lowerCAmelCase__		:			str     )     ->   Optional[int]:
   # support for torch, tf, jax etc.
   if config.TORCH_AVAILABLE and "torch" in sys.modules:
     import torch

     if isinstance(lowerCAmelCase__      ,    torch.Tensor     ):
       return self._tensorize(data_struct.detach().cpu().numpy()[()]     )
   if hasattr(lowerCAmelCase__      ,    "__array__"     ) and not isinstance(lowerCAmelCase__      ,    (np.ndarray, np.character, np.number)     ):
     UpperCAmelCase          =	data_struct.__array__()
   # support for nested types like struct of list of struct
   if isinstance(lowerCAmelCase__      ,    np.ndarray     ):
     if data_struct.dtype == object:
       return self._consolidate([self.recursive_tensorize(lowerCAmelCase__     ) for substruct in data_struct]     )
   if isinstance(lowerCAmelCase__      ,    (list, tuple)     ):
     return self._consolidate([self.recursive_tensorize(lowerCAmelCase__     ) for substruct in data_struct]     )
   return self._tensorize(lowerCAmelCase__     )


 def _UpperCamelCase						(    self		:			Union[str, Any]      ,    lowerCAmelCase__		:			dict     )     ->   Union[str, Any]:
   return map_nested(self._recursive_tensorize      ,    lowerCAmelCase__      ,    map_list=lowerCAmelCase__     )


 def _UpperCamelCase						(    self		:			int      ,    lowerCAmelCase__		:			pa.Table     )     ->   Mapping:
   UpperCAmelCase          =	self.numpy_arrow_extractor().extract_row(lowerCAmelCase__     )
   UpperCAmelCase          =	self.python_features_decoder.decode_row(lowerCAmelCase__     )
   return self.recursive_tensorize(lowerCAmelCase__     )


 def _UpperCamelCase						(    self		:			List[str]      ,    lowerCAmelCase__		:			pa.Table     )     ->   np.ndarray:
   UpperCAmelCase          =	self.numpy_arrow_extractor().extract_column(lowerCAmelCase__     )
   UpperCAmelCase          =	self.python_features_decoder.decode_column(lowerCAmelCase__      ,    pa_table.column_names[0]     )
   UpperCAmelCase          =	self.recursive_tensorize(lowerCAmelCase__     )
   UpperCAmelCase          =	self._consolidate(lowerCAmelCase__     )
   return column


 def _UpperCamelCase						(    self		:			List[Any]      ,    lowerCAmelCase__		:			pa.Table     )     ->   Mapping:
   UpperCAmelCase          =	self.numpy_arrow_extractor().extract_batch(lowerCAmelCase__     )
   UpperCAmelCase          =	self.python_features_decoder.decode_batch(lowerCAmelCase__     )
   UpperCAmelCase          =	self.recursive_tensorize(lowerCAmelCase__     )
   for column_name in batch:
     UpperCAmelCase          =	self._consolidate(batch[column_name]     )
   return batch



