





'''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  SCREAMING_SNAKE_CASE    (a__					,		a__					,		unittest.TestCase							):
     lowerCAmelCase					=       StableDiffusionSAGPipeline
     lowerCAmelCase					=       TEXT_TO_IMAGE_PARAMS
     lowerCAmelCase					=       TEXT_TO_IMAGE_BATCH_PARAMS
     lowerCAmelCase					=       TEXT_TO_IMAGE_IMAGE_PARAMS
     lowerCAmelCase					=       TEXT_TO_IMAGE_IMAGE_PARAMS
     lowerCAmelCase					=       False


     def 							SCREAMING_SNAKE_CASE  (						self):

          '''simple docstring'''






          torch.manual_seed(0)
          __A						: Optional[int]      =       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						,	)
          __A						: Any      =       DDIMScheduler(
              beta_start=0.00085						,	beta_end=0.012						,	beta_schedule='scaled_linear'						,	clip_sample=_UpperCAmelCase						,	set_alpha_to_one=_UpperCAmelCase						,	)
          torch.manual_seed(0)
          __A						: Optional[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)
          __A						: 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						,	)
          __A						: Optional[Any]      =       CLIPTextModel(_UpperCAmelCase)
          __A						: Optional[Any]      =       CLIPTokenizer.from_pretrained('hf-internal-testing/tiny-random-clip')

          __A						: Optional[Any]      =       {
              'unet': unet,
              'scheduler': scheduler,
              'vae': vae,
              'text_encoder': text_encoder,
              'tokenizer': tokenizer,
              'safety_checker': None,
              'feature_extractor': None,
          }
          return components


     def 							SCREAMING_SNAKE_CASE  (						self						,	_UpperCAmelCase						,	_UpperCAmelCase=0):

          '''simple docstring'''






          if str(_UpperCAmelCase).startswith('mps'):
               __A						: Optional[Any]      =       torch.manual_seed(_UpperCAmelCase)
          else:
               __A						: Optional[int]      =       torch.Generator(device=_UpperCAmelCase).manual_seed(_UpperCAmelCase)
          __A						: Any      =       {
              'prompt': '.',
              'generator': generator,
              'num_inference_steps': 2,
              'guidance_scale': 1.0,
              'sag_scale': 1.0,
              'output_type': 'numpy',
          }
          return inputs


     def 							SCREAMING_SNAKE_CASE  (						self):

          '''simple docstring'''






          super().test_inference_batch_single_identical(expected_max_diff=3e-3)




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


     def 							SCREAMING_SNAKE_CASE  (						self):

          '''simple docstring'''






          super().tearDown()
          gc.collect()
          torch.cuda.empty_cache()


     def 							SCREAMING_SNAKE_CASE  (						self):

          '''simple docstring'''






          __A						: str      =       StableDiffusionSAGPipeline.from_pretrained('CompVis/stable-diffusion-v1-4')
          __A						: Optional[Any]      =       sag_pipe.to(_UpperCAmelCase)
          sag_pipe.set_progress_bar_config(disable=_UpperCAmelCase)

          __A						: Optional[Any]      =       '.'
          __A						: Tuple      =       torch.manual_seed(0)
          __A						: Optional[int]      =       sag_pipe(
              [prompt]						,	generator=_UpperCAmelCase						,	guidance_scale=7.5						,	sag_scale=1.0						,	num_inference_steps=20						,	output_type='np')

          __A						: Dict      =       output.images

          __A						: Dict      =       image[0, -3:, -3:, -1]

          assert image.shape == (1, 512, 512, 3)
          __A						: Optional[Any]      =       np.array([0.1568, 0.1738, 0.1695, 0.1693, 0.1507, 0.1705, 0.1547, 0.1751, 0.1949])

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


     def 							SCREAMING_SNAKE_CASE  (						self):

          '''simple docstring'''






          __A						: Optional[int]      =       StableDiffusionSAGPipeline.from_pretrained('stabilityai/stable-diffusion-2-1-base')
          __A						: str      =       sag_pipe.to(_UpperCAmelCase)
          sag_pipe.set_progress_bar_config(disable=_UpperCAmelCase)

          __A						: int      =       '.'
          __A						: Tuple      =       torch.manual_seed(0)
          __A						: List[str]      =       sag_pipe(
              [prompt]						,	generator=_UpperCAmelCase						,	guidance_scale=7.5						,	sag_scale=1.0						,	num_inference_steps=20						,	output_type='np')

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

          __A						: str      =       image[0, -3:, -3:, -1]

          assert image.shape == (1, 512, 512, 3)
          __A						: List[str]      =       np.array([0.3459, 0.2876, 0.2537, 0.3002, 0.2671, 0.2160, 0.3026, 0.2262, 0.2371])

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


     def 							SCREAMING_SNAKE_CASE  (						self):

          '''simple docstring'''






          __A						: Union[str, Any]      =       StableDiffusionSAGPipeline.from_pretrained('stabilityai/stable-diffusion-2-1-base')
          __A						: List[Any]      =       sag_pipe.to(_UpperCAmelCase)
          sag_pipe.set_progress_bar_config(disable=_UpperCAmelCase)

          __A						: List[str]      =       '.'
          __A						: List[Any]      =       torch.manual_seed(0)
          __A						: Tuple      =       sag_pipe(
              [prompt]						,	width=768						,	height=512						,	generator=_UpperCAmelCase						,	guidance_scale=7.5						,	sag_scale=1.0						,	num_inference_steps=20						,	output_type='np'						,	)

          __A						: List[str]      =       output.images

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






'''simple docstring'''




from typing import TYPE_CHECKING

from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_torch_available


lowercase__    :      List[str]						=    {
    '''configuration_git''': ['''GIT_PRETRAINED_CONFIG_ARCHIVE_MAP''', '''GitConfig''', '''GitVisionConfig'''],
    '''processing_git''': ['''GitProcessor'''],
}

try:
  if not is_torch_available():
    raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
  pass
else:
  lowercase__    :      Any						=    [
      '''GIT_PRETRAINED_MODEL_ARCHIVE_LIST''',
      '''GitForCausalLM''',
      '''GitModel''',
      '''GitPreTrainedModel''',
      '''GitVisionModel''',
  ]

if TYPE_CHECKING:
  from .configuration_git import GIT_PRETRAINED_CONFIG_ARCHIVE_MAP, GitConfig, GitVisionConfig
  from .processing_git import GitProcessor

  try:
    if not is_torch_available():
      raise OptionalDependencyNotAvailable()
  except OptionalDependencyNotAvailable:
    pass
  else:
    from .modeling_git import (
        GIT_PRETRAINED_MODEL_ARCHIVE_LIST,
        GitForCausalLM,
        GitModel,
        GitPreTrainedModel,
        GitVisionModel,
    )

else:
  import sys

  lowercase__    :      List[str]						=    _LazyModule(__name__, globals()['''__file__'''], _import_structure, module_spec=__spec__)