




import gc
import unittest

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

from diffusers import (
    AutoencoderKL,
    DiffusionPipeline,
    EulerDiscreteScheduler,
    StableDiffusionXLPipeline,
    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 							__lowercase   (			_snake_case	,		_snake_case	,		unittest.TestCase    ):
     _lowerCAmelCase		       =						StableDiffusionXLPipeline
     _lowerCAmelCase		       =						TEXT_TO_IMAGE_PARAMS
     _lowerCAmelCase		       =						TEXT_TO_IMAGE_BATCH_PARAMS
     _lowerCAmelCase		       =						TEXT_TO_IMAGE_IMAGE_PARAMS
     _lowerCAmelCase		       =						TEXT_TO_IMAGE_IMAGE_PARAMS



     def 				__magic_name__		(  self   :							Tuple       ):
       torch.manual_seed(0       )
       a_		    =  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''')      ,		attention_head_dim=(2, 4)      ,		use_linear_projection=lowerCAmelCase__      ,		addition_embed_type='''text_time'''      ,		addition_time_embed_dim=8      ,		transformer_layers_per_block=(1, 2)      ,		projection_class_embeddings_input_dim=8_0      ,		cross_attention_dim=6_4      ,		)
       a_		    =  EulerDiscreteScheduler(
           beta_start=0.0_0085      ,		beta_end=0.012      ,		steps_offset=1      ,		beta_schedule='''scaled_linear'''      ,		timestep_spacing='''leading'''      ,		)
       torch.manual_seed(0       )
       a_		    =  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       )
       a_		    =  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      ,		hidden_act='''gelu'''      ,		projection_dim=3_2      ,		)
       a_		    =  CLIPTextModel(lowerCAmelCase__       )
       a_		    =  CLIPTokenizer.from_pretrained('''hf-internal-testing/tiny-random-clip'''      ,		local_files_only=lowerCAmelCase__       )

       a_		    =  CLIPTextModelWithProjection(lowerCAmelCase__       )
       a_		    =  CLIPTokenizer.from_pretrained('''hf-internal-testing/tiny-random-clip'''      ,		local_files_only=lowerCAmelCase__       )

       a_		    =  {
           '''unet''': unet,
           '''scheduler''': scheduler,
           '''vae''': vae,
           '''text_encoder''': text_encoder,
           '''tokenizer''': tokenizer,
           '''text_encoder_2''': text_encoder_a,
           '''tokenizer_2''': tokenizer_a,
           # "safety_checker": None,
           # "feature_extractor": None,
       }
       return components



     def 				__magic_name__		(  self   :							Optional[Any]      ,		lowercase__   :							Any      ,		lowercase__   :							Optional[Any]=0       ):
       if str(lowerCAmelCase__       ).startswith('''mps'''       ):
         a_		    =  torch.manual_seed(lowerCAmelCase__       )
       else:
         a_		    =  torch.Generator(device=lowerCAmelCase__       ).manual_seed(lowerCAmelCase__       )
       a_		    =  {
           '''prompt''': '''A painting of a squirrel eating a burger''',
           '''generator''': generator,
           '''num_inference_steps''': 2,
           '''guidance_scale''': 5.0,
           '''output_type''': '''numpy''',
       }
       return inputs



     def 				__magic_name__		(  self   :							Any       ):
       a_		    =  '''cpu'''  # ensure determinism for the device-dependent torch.Generator
       a_		    =  self.get_dummy_components()
       a_		    =  StableDiffusionXLPipeline(**lowerCAmelCase__       )
       a_		    =  sd_pipe.to(lowerCAmelCase__       )
       sd_pipe.set_progress_bar_config(disable=lowerCAmelCase__       )

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

       assert image.shape == (1, 6_4, 6_4, 3)
       a_		    =  np.array([0.5873, 0.6128, 0.4797, 0.5122, 0.5674, 0.4639, 0.5227, 0.5149, 0.4747]       )

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



     def 				__magic_name__		(  self   :							Union[str, Any]       ):
       a_		    =  self.get_dummy_components()
       a_		    =  StableDiffusionXLPipeline(**lowerCAmelCase__       )
       a_		    =  sd_pipe.to(lowerCAmelCase__       )
       a_		    =  sd_pipe.to(lowerCAmelCase__       )
       sd_pipe.set_progress_bar_config(disable=lowerCAmelCase__       )

       # forward without prompt embeds
       a_		    =  self.get_dummy_inputs(lowerCAmelCase__       )
       a_		    =  3 * ['''this is a negative prompt''']
       a_		    =  negative_prompt
       a_		    =  3 * [inputs['''prompt''']]

       a_		    =  sd_pipe(**lowerCAmelCase__       )
       a_		    =  output.images[0, -3:, -3:, -1]

       # forward with prompt embeds
       a_		    =  self.get_dummy_inputs(lowerCAmelCase__       )
       a_		    =  3 * ['''this is a negative prompt''']
       a_		    =  3 * [inputs.pop('''prompt'''       )]

       (
           (
           a_						
       )    , (
           a_						
       )    , (
           a_						
       )    , (
           a_						
       )    , 						
       )		    =  sd_pipe.encode_prompt(lowerCAmelCase__      ,		negative_prompt=lowerCAmelCase__       )

       a_		    =  sd_pipe(
           **lowerCAmelCase__      ,		prompt_embeds=lowerCAmelCase__      ,		negative_prompt_embeds=lowerCAmelCase__      ,		pooled_prompt_embeds=lowerCAmelCase__      ,		negative_pooled_prompt_embeds=lowerCAmelCase__      ,		)
       a_		    =  output.images[0, -3:, -3:, -1]

       # make sure that it's equal
       assert np.abs(image_slice_a.flatten() - image_slice_a.flatten()       ).max() < 1e-4



     def 				__magic_name__		(  self   :							List[Any]       ):
       super().test_attention_slicing_forward_pass(expected_max_diff=3e-3       )



     def 				__magic_name__		(  self   :							Optional[Any]       ):
       super().test_inference_batch_single_identical(expected_max_diff=3e-3       )



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



     def 				__magic_name__		(  self   :							Optional[int]       ):
       super().tearDown()
       gc.collect()
       torch.cuda.empty_cache()



     def 				__magic_name__		(  self   :							List[Any]      ,		lowercase__   :							int      ,		lowercase__   :							List[Any]="cpu"      ,		lowercase__   :							Dict=torch.floataa      ,		lowercase__   :							Tuple=0       ):
       a_		    =  torch.Generator(device=lowerCAmelCase__       ).manual_seed(lowerCAmelCase__       )
       a_		    =  np.random.RandomState(lowerCAmelCase__       ).standard_normal((1, 4, 6_4, 6_4)       )
       a_		    =  torch.from_numpy(lowerCAmelCase__       ).to(device=lowerCAmelCase__      ,		dtype=lowerCAmelCase__       )
       a_		    =  {
           '''prompt''': '''a photograph of an astronaut riding a horse''',
           '''latents''': latents,
           '''generator''': generator,
           '''num_inference_steps''': 3,
           '''guidance_scale''': 7.5,
           '''output_type''': '''numpy''',
       }
       return inputs



     def 				__magic_name__		(  self   :							Optional[Any]       ):
       a_		    =  DiffusionPipeline.from_pretrained('''stabilityai/stable-diffusion-2-base'''       )
       pipe.to(lowerCAmelCase__       )
       pipe.set_progress_bar_config(disable=lowerCAmelCase__       )

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

       assert image.shape == (1, 5_1_2, 5_1_2, 3)
       a_		    =  np.array([0.4_9493, 0.4_7896, 0.4_0798, 0.5_4214, 0.5_3212, 0.4_8202, 0.4_7656, 0.4_6329, 0.4_8506]       )
       assert np.abs(image_slice - expected_slice       ).max() < 7e-3



import json
import os
from collections.abc import Mapping
from typing import Dict, List, Optional, Union

import numpy as np
import torch

from .imports import is_safetensors_available
def   _lowerCAmelCase(						__A ,							__A ,							__A ,							__A=None      ):
  UpperCAmelCase          =	None
  # Check the string instead of the dtype to be compatible with versions of PyTorch that don't have bfloat16.
  if str(weight.dtype      ) == "torch.bfloat16":
    # Need to reinterpret the underlined data as int16 since NumPy does not handle bfloat16s.
    UpperCAmelCase          =	weight.view(torch.intaa      )
    UpperCAmelCase          =	"bfloat16"
  UpperCAmelCase          =	weight.cpu().numpy()
  UpperCAmelCase          =	os.path.join(__A ,							F"{weight_name}.dat"      )
  if index is not None:
    if dtype is None:
      UpperCAmelCase          =	str(array.dtype      )
    UpperCAmelCase          =	{"dtype": dtype, "shape": list(array.shape      )}
  if array.ndim == 0:
    UpperCAmelCase          =	array[None]
  UpperCAmelCase          =	np.memmap(__A ,							dtype=array.dtype ,							mode="w+" ,							shape=array.shape      )
  UpperCAmelCase          =	array[:]
  file_array.flush()
  return index
def   _lowerCAmelCase(						__A ,							__A      ):
  UpperCAmelCase          =	tuple(weight_info["shape"]      )
  if shape == ():
    # NumPy memory-mapped arrays can't have 0 dims so it was saved as 1d tensor
    UpperCAmelCase          =	(1,)

  UpperCAmelCase          =	weight_info["dtype"]
  if dtype == "bfloat16":
    # NumPy does not support bfloat16 so this was saved as a int16
    UpperCAmelCase          =	"int16"

  UpperCAmelCase          =	np.memmap(__A ,							dtype=__A ,							shape=__A ,							mode="r"      )

  if len(weight_info["shape"]      ) == 0:
    UpperCAmelCase          =	weight[0]
  UpperCAmelCase          =	torch.tensor(__A      )
  if weight_info["dtype"] == "bfloat16":
    UpperCAmelCase          =	weight.view(torch.bfloataa      )

  return weight
def   _lowerCAmelCase(						__A ,							__A      ):
  if index is None or len(__A      ) == 0:
    # Nothing to save
    return

  UpperCAmelCase          =	os.path.join(__A ,							"index.json"      )
  if os.path.isfile(__A      ):
    with open(__A ,							"r" ,							encoding="utf-8"      ) as f:
      UpperCAmelCase          =	json.load(__A      )
  else:
    UpperCAmelCase          =	{}
  current_index.update(__A      )

  with open(__A ,							"w" ,							encoding="utf-8"      ) as f:
    json.dump(__A ,							__A ,							indent=2      )
def   _lowerCAmelCase(						__A ,							__A      ):
  os.makedirs(__A ,							exist_ok=__A      )
  UpperCAmelCase          =	{}
  for name, parameter in state_dict.items():
    UpperCAmelCase          =	offload_weight(__A ,							__A ,							__A ,							index=__A      )

  # Update index
  save_offload_index(__A ,							__A      )






class     __magic_name__     (							_snake_case     ):


 def __init__(    self		:			Tuple      ,    lowerCAmelCase__		:			Mapping      ,    lowerCAmelCase__		:			str     )     ->   Tuple:
   UpperCAmelCase          =	dataset
   UpperCAmelCase          =	prefix


 def __getitem__(    self		:			Dict      ,    lowerCAmelCase__		:			Any     )     ->   Optional[int]:
   return self.dataset[f"{self.prefix}{key}"]


 def __iter__(    self		:			Optional[int]     )     ->   Any:
   return iter([key for key in self.dataset if key.startswith(self.prefix     )]     )


 def __len__(    self		:			Union[str, Any]     )     ->   int:
   return len(self.dataset     )






class     __magic_name__     (							_snake_case     ):


 def __init__(    self		:			Optional[int]      ,    lowerCAmelCase__		:			Dict[str, torch.Tensor] = None      ,    lowerCAmelCase__		:			Optional[Union[str, os.PathLike]] = None      ,    lowerCAmelCase__		:			Mapping = None      ,    lowerCAmelCase__		:			Dict=None      ,    )     ->   Dict:
   if state_dict is None and save_folder is None:
     raise ValueError("Need either a `state_dict` or a `save_folder` containing offloaded weights."     )

   UpperCAmelCase          =	{} if state_dict is None else state_dict
   UpperCAmelCase          =	save_folder
   if index is None and save_folder is not None:
     with open(os.path.join(lowerCAmelCase__      ,    "index.json"     )     ) as f:
       UpperCAmelCase          =	json.load(lowerCAmelCase__     )
   UpperCAmelCase          =	{} if index is None else index
   UpperCAmelCase          =	list(self.state_dict.keys()     )
   self.all_keys.extend([key for key in self.index if key not in self.all_keys]     )
   UpperCAmelCase          =	device


 def __getitem__(    self		:			str      ,    lowerCAmelCase__		:			str     )     ->   Dict:
   # State dict gets priority
   if key in self.state_dict:
     return self.state_dict[key]
   UpperCAmelCase          =	self.index[key]
   if weight_info.get("safetensors_file"     ) is not None:
     if not is_safetensors_available():
       raise ImportError("These offloaded weights require the use of safetensors: `pip install safetensors`."     )

     from safetensors import safe_open

     UpperCAmelCase          =	"cpu" if self.device is None else self.device
     with safe_open(weight_info["safetensors_file"]      ,    framework="pt"      ,    device=lowerCAmelCase__     ) as f:
       UpperCAmelCase          =	f.get_tensor(weight_info.get("weight_name"      ,    lowerCAmelCase__     )     )

     if "dtype" in weight_info:
       return tensor.to(getattr(lowerCAmelCase__      ,    weight_info["dtype"]     )     )
     else:
       return tensor

   UpperCAmelCase          =	os.path.join(self.save_folder      ,    f"{key}.dat"     )
   return load_offloaded_weight(lowerCAmelCase__      ,    lowerCAmelCase__     )


 def __iter__(    self		:			str     )     ->   Dict:
   return iter(self.all_keys     )


 def __len__(    self		:			Optional[Any]     )     ->   List[Any]:
   return len(self.all_keys     )

def   _lowerCAmelCase(						__A ,							__A      ):
  UpperCAmelCase          =	{}
  for module_name in submodule_names:
    # We want to catch module_name parameter (module_name.xxx) or potentially module_name, but not any of the
    # submodules that could being like module_name (transformers.h.1 and transformers.h.10 for instance)
    result.update(
        {
            key: param
            for key, param in state_dict.items()
            if key == module_name or key.startswith(module_name + "."      )
        }      )
  return result



