
"""simple docstring"""





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_	(__a				,__a				,unittest.TestCase      ):




  '''simple docstring'''
  SCREAMING_SNAKE_CASE__			:						Tuple       =		StableDiffusionSAGPipeline
  SCREAMING_SNAKE_CASE__			:						List[str]       =		TEXT_TO_IMAGE_PARAMS
  SCREAMING_SNAKE_CASE__			:						Optional[int]       =		TEXT_TO_IMAGE_BATCH_PARAMS
  SCREAMING_SNAKE_CASE__			:						Union[str, Any]       =		TEXT_TO_IMAGE_IMAGE_PARAMS
  SCREAMING_SNAKE_CASE__			:						Any       =		TEXT_TO_IMAGE_IMAGE_PARAMS
  SCREAMING_SNAKE_CASE__			:						Any       =		False







  def        UpperCamelCase__					(		self  ):

     """simple docstring"""
     torch.manual_seed(0  )
     UpperCAmelCase_					:					str			     =    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_					:					int			     =    DDIMScheduler(
         beta_start=0.0_00_85				,   beta_end=0.0_12				,   beta_schedule="scaled_linear"				,   clip_sample=_UpperCamelCase				,   set_alpha_to_one=_UpperCamelCase				,   )
     torch.manual_seed(0  )
     UpperCAmelCase_					:					int			     =    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_					:					Any			     =    CLIPTextConfig(
         bos_token_id=0				,   eos_token_id=2				,   hidden_size=32				,   intermediate_size=37				,   layer_norm_eps=1E-0_5				,   num_attention_heads=4				,   num_hidden_layers=5				,   pad_token_id=1				,   vocab_size=1000				,   )
     UpperCAmelCase_					:					Union[str, Any]			     =    CLIPTextModel(_UpperCamelCase  )
     UpperCAmelCase_					:					Any			     =    CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip"  )

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







  def        UpperCamelCase__					(		self				,   lowercase_				,   lowercase_=0  ):

     """simple docstring"""
     if str(_UpperCamelCase  ).startswith("mps"  ):
        UpperCAmelCase_					:					Any			     =    torch.manual_seed(_UpperCamelCase  )
     else:
        UpperCAmelCase_					:					int			     =    torch.Generator(device=_UpperCamelCase  ).manual_seed(_UpperCamelCase  )
     UpperCAmelCase_					:					Union[str, Any]			     =    {
         "prompt": ".",
         "generator": generator,
         "num_inference_steps": 2,
         "guidance_scale": 1.0,
         "sag_scale": 1.0,
         "output_type": "numpy",
     }
     return inputs







  def        UpperCamelCase__					(		self  ):

     """simple docstring"""
     super().test_inference_batch_single_identical(expected_max_diff=3E-3  )




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




  '''simple docstring'''







  def        UpperCamelCase__					(		self  ):

     """simple docstring"""
     # clean up the VRAM after each test
     super().tearDown()
     gc.collect()
     torch.cuda.empty_cache()







  def        UpperCamelCase__					(		self  ):

     """simple docstring"""
     UpperCAmelCase_					:					Dict			     =    StableDiffusionSAGPipeline.from_pretrained("CompVis/stable-diffusion-v1-4"  )
     UpperCAmelCase_					:					List[Any]			     =    sag_pipe.to(_UpperCamelCase  )
     sag_pipe.set_progress_bar_config(disable=_UpperCamelCase  )

     UpperCAmelCase_					:					Tuple			     =    "."
     UpperCAmelCase_					:					Union[str, Any]			     =    torch.manual_seed(0  )
     UpperCAmelCase_					:					Any			     =    sag_pipe(
         [prompt]				,   generator=_UpperCamelCase				,   guidance_scale=7.5				,   sag_scale=1.0				,   num_inference_steps=20				,   output_type="np"  )

     UpperCAmelCase_					:					Union[str, Any]			     =    output.images

     UpperCAmelCase_					:					Optional[Any]			     =    image[0, -3:, -3:, -1]

     assert image.shape == (1, 512, 512, 3)
     UpperCAmelCase_					:					str			     =    np.array([0.15_68, 0.17_38, 0.16_95, 0.16_93, 0.15_07, 0.17_05, 0.15_47, 0.17_51, 0.19_49]  )

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







  def        UpperCamelCase__					(		self  ):

     """simple docstring"""
     UpperCAmelCase_					:					List[Any]			     =    StableDiffusionSAGPipeline.from_pretrained("stabilityai/stable-diffusion-2-1-base"  )
     UpperCAmelCase_					:					Any			     =    sag_pipe.to(_UpperCamelCase  )
     sag_pipe.set_progress_bar_config(disable=_UpperCamelCase  )

     UpperCAmelCase_					:					Optional[Any]			     =    "."
     UpperCAmelCase_					:					Tuple			     =    torch.manual_seed(0  )
     UpperCAmelCase_					:					Dict			     =    sag_pipe(
         [prompt]				,   generator=_UpperCamelCase				,   guidance_scale=7.5				,   sag_scale=1.0				,   num_inference_steps=20				,   output_type="np"  )

     UpperCAmelCase_					:					Any			     =    output.images

     UpperCAmelCase_					:					int			     =    image[0, -3:, -3:, -1]

     assert image.shape == (1, 512, 512, 3)
     UpperCAmelCase_					:					List[Any]			     =    np.array([0.34_59, 0.28_76, 0.25_37, 0.30_02, 0.26_71, 0.21_60, 0.30_26, 0.22_62, 0.23_71]  )

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







  def        UpperCamelCase__					(		self  ):

     """simple docstring"""
     UpperCAmelCase_					:					Optional[int]			     =    StableDiffusionSAGPipeline.from_pretrained("stabilityai/stable-diffusion-2-1-base"  )
     UpperCAmelCase_					:					List[Any]			     =    sag_pipe.to(_UpperCamelCase  )
     sag_pipe.set_progress_bar_config(disable=_UpperCamelCase  )

     UpperCAmelCase_					:					Dict			     =    "."
     UpperCAmelCase_					:					Any			     =    torch.manual_seed(0  )
     UpperCAmelCase_					:					Optional[int]			     =    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_					:					Optional[Any]			     =    output.images

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




from typing import TYPE_CHECKING

from ...utils import (
    OptionalDependencyNotAvailable,
    _LazyModule,
    is_flax_available,
    is_sentencepiece_available,
    is_tf_available,
    is_tokenizers_available,
    is_torch_available,
)


_A   =     {"configuration_xglm": ["XGLM_PRETRAINED_CONFIG_ARCHIVE_MAP", "XGLMConfig"]}

try:
			if not is_sentencepiece_available():
						raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
			pass
else:
			_A   =     ["XGLMTokenizer"]

try:
			if not is_tokenizers_available():
						raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
			pass
else:
			_A   =     ["XGLMTokenizerFast"]

try:
			if not is_torch_available():
						raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
			pass
else:
			_A   =     [
			    "XGLM_PRETRAINED_MODEL_ARCHIVE_LIST",
			    "XGLMForCausalLM",
			    "XGLMModel",
			    "XGLMPreTrainedModel",
			]


try:
			if not is_flax_available():
						raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
			pass
else:
			_A   =     [
			    "FlaxXGLMForCausalLM",
			    "FlaxXGLMModel",
			    "FlaxXGLMPreTrainedModel",
			]


try:
			if not is_tf_available():
						raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
			pass
else:
			_A   =     [
			    "TF_XGLM_PRETRAINED_MODEL_ARCHIVE_LIST",
			    "TFXGLMForCausalLM",
			    "TFXGLMModel",
			    "TFXGLMPreTrainedModel",
			]


if TYPE_CHECKING:
			from .configuration_xglm import XGLM_PRETRAINED_CONFIG_ARCHIVE_MAP, XGLMConfig

			try:
						if not is_sentencepiece_available():
									raise OptionalDependencyNotAvailable()
			except OptionalDependencyNotAvailable:
						pass
			else:
						from .tokenization_xglm import XGLMTokenizer

			try:
						if not is_tokenizers_available():
									raise OptionalDependencyNotAvailable()
			except OptionalDependencyNotAvailable:
						pass
			else:
						from .tokenization_xglm_fast import XGLMTokenizerFast

			try:
						if not is_torch_available():
									raise OptionalDependencyNotAvailable()
			except OptionalDependencyNotAvailable:
						pass
			else:
						from .modeling_xglm import XGLM_PRETRAINED_MODEL_ARCHIVE_LIST, XGLMForCausalLM, XGLMModel, XGLMPreTrainedModel

			try:
						if not is_flax_available():
									raise OptionalDependencyNotAvailable()
			except OptionalDependencyNotAvailable:
						pass
			else:
						from .modeling_flax_xglm import FlaxXGLMForCausalLM, FlaxXGLMModel, FlaxXGLMPreTrainedModel

			try:
						if not is_tf_available():
									raise OptionalDependencyNotAvailable()
			except OptionalDependencyNotAvailable:
						pass
			else:
						from .modeling_tf_xglm import (
						    TF_XGLM_PRETRAINED_MODEL_ARCHIVE_LIST,
						    TFXGLMForCausalLM,
						    TFXGLMModel,
						    TFXGLMPreTrainedModel,
						)


else:
			import sys

			_A   =     _LazyModule(__name__, globals()["__file__"], _import_structure)



