






import gc
import math
import unittest

import torch

from diffusers import UNetaDModel
from diffusers.utils import floats_tensor, logging, slow, torch_all_close, torch_device
from diffusers.utils.testing_utils import enable_full_determinism

from .test_modeling_common import ModelTesterMixin, UNetTesterMixin


snake_case__   :						Dict													=   logging.get_logger(__name__)

enable_full_determinism()


class     _a							(	UpperCAmelCase__    ,							UpperCAmelCase__    ,							unittest.TestCase     ):






  """simple docstring"""





  A_						       = UNetaDModel
  A_						       = """sample"""




  @property
  def  _UpperCAmelCase  (  self			)      ->			Union[str, Any]:
     UpperCamelCase_						=  4
     UpperCamelCase_						=  3
     UpperCamelCase_						=  (32, 32)

     UpperCamelCase_						=  floats_tensor((batch_size, num_channels) + sizes			).to(_UpperCAmelCase			)
     UpperCamelCase_						=  torch.tensor([10]			).to(_UpperCAmelCase			)

     return {"sample": noise, "timestep": time_step}




  @property
  def  _UpperCAmelCase  (  self			)      ->			Union[str, Any]:
     return (3, 32, 32)




  @property
  def  _UpperCAmelCase  (  self			)      ->			Union[str, Any]:
     return (3, 32, 32)




  def  _UpperCAmelCase  (  self			)      ->			str:
     UpperCamelCase_						=  {
         'block_out_channels': (32, 64),
         'down_block_types': ('DownBlock2D', 'AttnDownBlock2D'),
         'up_block_types': ('AttnUpBlock2D', 'UpBlock2D'),
         'attention_head_dim': 3,
         'out_channels': 3,
         'in_channels': 3,
         'layers_per_block': 2,
         'sample_size': 32,
     }
     UpperCamelCase_						=  self.dummy_input
     return init_dict, inputs_dict




class     _a							(	UpperCAmelCase__    ,							UpperCAmelCase__    ,							unittest.TestCase     ):






  """simple docstring"""





  A_						       = UNetaDModel
  A_						       = """sample"""




  @property
  def  _UpperCAmelCase  (  self			)      ->			Union[str, Any]:
     UpperCamelCase_						=  4
     UpperCamelCase_						=  4
     UpperCamelCase_						=  (32, 32)

     UpperCamelCase_						=  floats_tensor((batch_size, num_channels) + sizes			).to(_UpperCAmelCase			)
     UpperCamelCase_						=  torch.tensor([10]			).to(_UpperCAmelCase			)

     return {"sample": noise, "timestep": time_step}




  @property
  def  _UpperCAmelCase  (  self			)      ->			Any:
     return (4, 32, 32)




  @property
  def  _UpperCAmelCase  (  self			)      ->			str:
     return (4, 32, 32)




  def  _UpperCAmelCase  (  self			)      ->			Any:
     UpperCamelCase_						=  {
         'sample_size': 32,
         'in_channels': 4,
         'out_channels': 4,
         'layers_per_block': 2,
         'block_out_channels': (32, 64),
         'attention_head_dim': 32,
         'down_block_types': ('DownBlock2D', 'DownBlock2D'),
         'up_block_types': ('UpBlock2D', 'UpBlock2D'),
     }
     UpperCamelCase_						=  self.dummy_input
     return init_dict, inputs_dict




  def  _UpperCAmelCase  (  self			)      ->			str:
     UpperCamelCase_		,   UpperCamelCase_						=  UNetaDModel.from_pretrained('fusing/unet-ldm-dummy-update'		,		output_loading_info=_UpperCAmelCase			)

     self.assertIsNotNone(_UpperCAmelCase			)
     self.assertEqual(len(loading_info['missing_keys']			)		,		0			)

     model.to(_UpperCAmelCase			)
     UpperCamelCase_						=  model(**self.dummy_input			).sample

     assert image is not None, "Make sure output is not None"




  @unittest.skipIf(torch_device != 'cuda'		,		'This test is supposed to run on GPU'			)
  def  _UpperCAmelCase  (  self			)      ->			Tuple:
     UpperCamelCase_		,   UpperCamelCase_						=  UNetaDModel.from_pretrained('fusing/unet-ldm-dummy-update'		,		output_loading_info=_UpperCAmelCase			)
     model.to(_UpperCAmelCase			)
     UpperCamelCase_						=  model(**self.dummy_input			).sample

     assert image is not None, "Make sure output is not None"




  @unittest.skipIf(torch_device != 'cuda'		,		'This test is supposed to run on GPU'			)
  def  _UpperCAmelCase  (  self			)      ->			Tuple:
     # by defautl model loading will use accelerate as `low_cpu_mem_usage=True`
     UpperCamelCase_		,   UpperCamelCase_						=  UNetaDModel.from_pretrained('fusing/unet-ldm-dummy-update'		,		output_loading_info=_UpperCAmelCase			)
     model_accelerate.to(_UpperCAmelCase			)
     model_accelerate.eval()

     UpperCamelCase_						=  torch.randn(
         1		,		model_accelerate.config.in_channels		,		model_accelerate.config.sample_size		,		model_accelerate.config.sample_size		,		generator=torch.manual_seed(0			)		,		)
     UpperCamelCase_						=  noise.to(_UpperCAmelCase			)
     UpperCamelCase_						=  torch.tensor([10] * noise.shape[0]			).to(_UpperCAmelCase			)

     UpperCamelCase_						=  model_accelerate(_UpperCAmelCase		,		_UpperCAmelCase			)['sample']

     # two models don't need to stay in the device at the same time
     del model_accelerate
     torch.cuda.empty_cache()
     gc.collect()

     UpperCamelCase_		,   UpperCamelCase_						=  UNetaDModel.from_pretrained(
         'fusing/unet-ldm-dummy-update'		,		output_loading_info=_UpperCAmelCase		,		low_cpu_mem_usage=_UpperCAmelCase			)
     model_normal_load.to(_UpperCAmelCase			)
     model_normal_load.eval()
     UpperCamelCase_						=  model_normal_load(_UpperCAmelCase		,		_UpperCAmelCase			)['sample']

     assert torch_all_close(_UpperCAmelCase		,		_UpperCAmelCase		,		rtol=1e-3			)




  def  _UpperCAmelCase  (  self			)      ->			List[str]:
     UpperCamelCase_						=  UNetaDModel.from_pretrained('fusing/unet-ldm-dummy-update'			)
     model.eval()
     model.to(_UpperCAmelCase			)

     UpperCamelCase_						=  torch.randn(
         1		,		model.config.in_channels		,		model.config.sample_size		,		model.config.sample_size		,		generator=torch.manual_seed(0			)		,		)
     UpperCamelCase_						=  noise.to(_UpperCAmelCase			)
     UpperCamelCase_						=  torch.tensor([10] * noise.shape[0]			).to(_UpperCAmelCase			)

     with torch.no_grad():
        UpperCamelCase_						=  model(_UpperCAmelCase		,		_UpperCAmelCase			).sample

     UpperCamelCase_						=  output[0, -1, -3:, -3:].flatten().cpu()
     # fmt: off
     UpperCamelCase_						=  torch.tensor([-1_3.3_2_5_8, -2_0.1_1_0_0, -1_5.9_8_7_3, -1_7.6_6_1_7, -2_3.0_5_9_6, -1_7.9_4_1_9, -1_3.3_6_7_5, -1_6.1_8_8_9, -1_2.3_8_0_0]			)
     # fmt: on

     self.assertTrue(torch_all_close(_UpperCAmelCase		,		_UpperCAmelCase		,		rtol=1e-3			)			)




class     _a							(	UpperCAmelCase__    ,							UpperCAmelCase__    ,							unittest.TestCase     ):






  """simple docstring"""





  A_						       = UNetaDModel
  A_						       = """sample"""




  @property
  def  _UpperCAmelCase  (  self		,		_UpperCAmelCase=(32, 32)			)      ->			str:
     UpperCamelCase_						=  4
     UpperCamelCase_						=  3

     UpperCamelCase_						=  floats_tensor((batch_size, num_channels) + sizes			).to(_UpperCAmelCase			)
     UpperCamelCase_						=  torch.tensor(batch_size * [10]			).to(dtype=torch.intaa		,		device=_UpperCAmelCase			)

     return {"sample": noise, "timestep": time_step}




  @property
  def  _UpperCAmelCase  (  self			)      ->			Dict:
     return (3, 32, 32)




  @property
  def  _UpperCAmelCase  (  self			)      ->			Optional[int]:
     return (3, 32, 32)




  def  _UpperCAmelCase  (  self			)      ->			int:
     UpperCamelCase_						=  {
         'block_out_channels': [32, 64, 64, 64],
         'in_channels': 3,
         'layers_per_block': 1,
         'out_channels': 3,
         'time_embedding_type': 'fourier',
         'norm_eps': 1e-6,
         'mid_block_scale_factor': math.sqrt(2.0			),
         'norm_num_groups': None,
         'down_block_types': [
             'SkipDownBlock2D',
             'AttnSkipDownBlock2D',
             'SkipDownBlock2D',
             'SkipDownBlock2D',
         ],
         'up_block_types': [
             'SkipUpBlock2D',
             'SkipUpBlock2D',
             'AttnSkipUpBlock2D',
             'SkipUpBlock2D',
         ],
     }
     UpperCamelCase_						=  self.dummy_input
     return init_dict, inputs_dict




  @slow
  def  _UpperCAmelCase  (  self			)      ->			int:
     UpperCamelCase_		,   UpperCamelCase_						=  UNetaDModel.from_pretrained('google/ncsnpp-celebahq-256'		,		output_loading_info=_UpperCAmelCase			)
     self.assertIsNotNone(_UpperCAmelCase			)
     self.assertEqual(len(loading_info['missing_keys']			)		,		0			)

     model.to(_UpperCAmelCase			)
     UpperCamelCase_						=  self.dummy_input
     UpperCamelCase_						=  floats_tensor((4, 3) + (256, 256)			).to(_UpperCAmelCase			)
     UpperCamelCase_						=  noise
     UpperCamelCase_						=  model(**_UpperCAmelCase			)

     assert image is not None, "Make sure output is not None"




  @slow
  def  _UpperCAmelCase  (  self			)      ->			Optional[int]:
     UpperCamelCase_						=  UNetaDModel.from_pretrained('google/ncsnpp-celebahq-256'			)
     model.to(_UpperCAmelCase			)

     UpperCamelCase_						=  4
     UpperCamelCase_						=  3
     UpperCamelCase_						=  (256, 256)

     UpperCamelCase_						=  torch.ones((batch_size, num_channels) + sizes			).to(_UpperCAmelCase			)
     UpperCamelCase_						=  torch.tensor(batch_size * [1e-4]			).to(_UpperCAmelCase			)

     with torch.no_grad():
        UpperCamelCase_						=  model(_UpperCAmelCase		,		_UpperCAmelCase			).sample

     UpperCamelCase_						=  output[0, -3:, -3:, -1].flatten().cpu()
     # fmt: off
     UpperCamelCase_						=  torch.tensor([-4_8_4_2.8_6_9_1, -6_4_9_9.6_6_3_1, -3_8_0_0.1_9_5_3, -7_9_7_8.2_6_8_6, -1_0_9_8_0.7_1_2_9, -2_0_0_2_8.8_5_3_5, 8_1_4_8.2_8_2_2, 2_3_4_2.2_9_0_5, 5_6_7.7_6_0_8]			)
     # fmt: on

     self.assertTrue(torch_all_close(_UpperCAmelCase		,		_UpperCAmelCase		,		rtol=1e-2			)			)




  def  _UpperCAmelCase  (  self			)      ->			Optional[Any]:
     UpperCamelCase_						=  UNetaDModel.from_pretrained('fusing/ncsnpp-ffhq-ve-dummy-update'			)
     model.to(_UpperCAmelCase			)

     UpperCamelCase_						=  4
     UpperCamelCase_						=  3
     UpperCamelCase_						=  (32, 32)

     UpperCamelCase_						=  torch.ones((batch_size, num_channels) + sizes			).to(_UpperCAmelCase			)
     UpperCamelCase_						=  torch.tensor(batch_size * [1e-4]			).to(_UpperCAmelCase			)

     with torch.no_grad():
        UpperCamelCase_						=  model(_UpperCAmelCase		,		_UpperCAmelCase			).sample

     UpperCamelCase_						=  output[0, -3:, -3:, -1].flatten().cpu()
     # fmt: off
     UpperCamelCase_						=  torch.tensor([-0.0_3_2_5, -0.0_9_0_0, -0.0_8_6_9, -0.0_3_3_2, -0.0_7_2_5, -0.0_2_7_0, -0.0_1_0_1, 0.0_2_2_7, 0.0_2_5_6]			)
     # fmt: on

     self.assertTrue(torch_all_close(_UpperCAmelCase		,		_UpperCAmelCase		,		rtol=1e-2			)			)




  def  _UpperCAmelCase  (  self			)      ->			Dict:
     # not required for this model
     pass













import gc
import unittest

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

from diffusers import (
    AutoencoderKL,
    DDIMScheduler,
    StableDiffusionSAGPipeline,
    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     _a							(	UpperCAmelCase__    ,							UpperCAmelCase__    ,							unittest.TestCase     ):






  """simple docstring"""





  A_						       = StableDiffusionSAGPipeline
  A_						       = TEXT_TO_IMAGE_PARAMS
  A_						       = TEXT_TO_IMAGE_BATCH_PARAMS
  A_						       = TEXT_TO_IMAGE_IMAGE_PARAMS
  A_						       = TEXT_TO_IMAGE_IMAGE_PARAMS
  A_						       = False




  def  _UpperCAmelCase  (  self			)      ->			Optional[Any]:
     torch.manual_seed(0			)
     UpperCamelCase_						=  UNetaDConditionModel(
         block_out_channels=(32, 64)		,		layers_per_block=2		,		sample_size=32		,		in_channels=4		,		out_channels=4		,		down_block_types=('DownBlock2D', 'CrossAttnDownBlock2D')		,		up_block_types=('CrossAttnUpBlock2D', 'UpBlock2D')		,		cross_attention_dim=32		,		)
     UpperCamelCase_						=  DDIMScheduler(
         beta_start=0.0_0_0_8_5		,		beta_end=0.0_1_2		,		beta_schedule='scaled_linear'		,		clip_sample=_UpperCAmelCase		,		set_alpha_to_one=_UpperCAmelCase		,		)
     torch.manual_seed(0			)
     UpperCamelCase_						=  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			)
     UpperCamelCase_						=  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=1000		,		)
     UpperCamelCase_						=  CLIPTextModel(_UpperCAmelCase			)
     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		,		_UpperCAmelCase		,		_UpperCAmelCase=0			)      ->			List[Any]:
     if str(_UpperCAmelCase			).startswith('mps'			):
        UpperCamelCase_						=  torch.manual_seed(_UpperCAmelCase			)
     else:
        UpperCamelCase_						=  torch.Generator(device=_UpperCAmelCase			).manual_seed(_UpperCAmelCase			)
     UpperCamelCase_						=  {
         'prompt': '.',
         'generator': generator,
         'num_inference_steps': 2,
         'guidance_scale': 1.0,
         'sag_scale': 1.0,
         'output_type': 'numpy',
     }
     return inputs




  def  _UpperCAmelCase  (  self			)      ->			Tuple:
     super().test_inference_batch_single_identical(expected_max_diff=3e-3			)




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






  """simple docstring"""





  def  _UpperCAmelCase  (  self			)      ->			Tuple:
     # clean up the VRAM after each test
     super().tearDown()
     gc.collect()
     torch.cuda.empty_cache()




  def  _UpperCAmelCase  (  self			)      ->			str:
     UpperCamelCase_						=  StableDiffusionSAGPipeline.from_pretrained('CompVis/stable-diffusion-v1-4'			)
     UpperCamelCase_						=  sag_pipe.to(_UpperCAmelCase			)
     sag_pipe.set_progress_bar_config(disable=_UpperCAmelCase			)

     UpperCamelCase_						=  '.'
     UpperCamelCase_						=  torch.manual_seed(0			)
     UpperCamelCase_						=  sag_pipe(
         [prompt]		,		generator=_UpperCAmelCase		,		guidance_scale=7.5		,		sag_scale=1.0		,		num_inference_steps=20		,		output_type='np'			)

     UpperCamelCase_						=  output.images

     UpperCamelCase_						=  image[0, -3:, -3:, -1]

     assert image.shape == (1, 512, 512, 3)
     UpperCamelCase_						=  np.array([0.1_5_6_8, 0.1_7_3_8, 0.1_6_9_5, 0.1_6_9_3, 0.1_5_0_7, 0.1_7_0_5, 0.1_5_4_7, 0.1_7_5_1, 0.1_9_4_9]			)

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




  def  _UpperCAmelCase  (  self			)      ->			List[Any]:
     UpperCamelCase_						=  StableDiffusionSAGPipeline.from_pretrained('stabilityai/stable-diffusion-2-1-base'			)
     UpperCamelCase_						=  sag_pipe.to(_UpperCAmelCase			)
     sag_pipe.set_progress_bar_config(disable=_UpperCAmelCase			)

     UpperCamelCase_						=  '.'
     UpperCamelCase_						=  torch.manual_seed(0			)
     UpperCamelCase_						=  sag_pipe(
         [prompt]		,		generator=_UpperCAmelCase		,		guidance_scale=7.5		,		sag_scale=1.0		,		num_inference_steps=20		,		output_type='np'			)

     UpperCamelCase_						=  output.images

     UpperCamelCase_						=  image[0, -3:, -3:, -1]

     assert image.shape == (1, 512, 512, 3)
     UpperCamelCase_						=  np.array([0.3_4_5_9, 0.2_8_7_6, 0.2_5_3_7, 0.3_0_0_2, 0.2_6_7_1, 0.2_1_6_0, 0.3_0_2_6, 0.2_2_6_2, 0.2_3_7_1]			)

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




  def  _UpperCAmelCase  (  self			)      ->			Dict:
     UpperCamelCase_						=  StableDiffusionSAGPipeline.from_pretrained('stabilityai/stable-diffusion-2-1-base'			)
     UpperCamelCase_						=  sag_pipe.to(_UpperCAmelCase			)
     sag_pipe.set_progress_bar_config(disable=_UpperCAmelCase			)

     UpperCamelCase_						=  '.'
     UpperCamelCase_						=  torch.manual_seed(0			)
     UpperCamelCase_						=  sag_pipe(
         [prompt]		,		width=768		,		height=512		,		generator=_UpperCAmelCase		,		guidance_scale=7.5		,		sag_scale=1.0		,		num_inference_steps=20		,		output_type='np'		,		)

     UpperCamelCase_						=  output.images

     assert image.shape == (1, 512, 768, 3)





