




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     UpperCamelCase_   (			UpperCAmelCase__		, UpperCAmelCase__		, unittest.TestCase					):

    '''simple docstring'''






    UpperCAmelCase__							  =				StableDiffusionSAGPipeline
    UpperCAmelCase__							  =				TEXT_TO_IMAGE_PARAMS
    UpperCAmelCase__							  =				TEXT_TO_IMAGE_BATCH_PARAMS
    UpperCAmelCase__							  =				TEXT_TO_IMAGE_IMAGE_PARAMS
    UpperCAmelCase__							  =				TEXT_TO_IMAGE_IMAGE_PARAMS
    UpperCAmelCase__							  =				False


    def 	SCREAMING_SNAKE_CASE							(      self				:     Dict)						->str:





        '''simple docstring'''

        torch.manual_seed(0)
        A__							=   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__							=   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__							=   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__							=   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=1_000 ,			)
        A__							=   CLIPTextModel(UpperCAmelCase__)
        A__							=   CLIPTokenizer.from_pretrained('''hf-internal-testing/tiny-random-clip''')

        A__							=   {
            '''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				:     Any ,			UpperCAmelCase__				:     List[Any] ,			UpperCAmelCase__				:     List[Any]=0)						->Optional[int]:





        '''simple docstring'''

        if str(UpperCAmelCase__).startswith('''mps'''):
            A__							=   torch.manual_seed(UpperCAmelCase__)
        else:
            A__							=   torch.Generator(device=UpperCAmelCase__).manual_seed(UpperCAmelCase__)
        A__							=   {
            '''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				:     Tuple)						->str:





        '''simple docstring'''

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


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

    '''simple docstring'''






    def 	SCREAMING_SNAKE_CASE							(      self				:     Optional[Any])						->Optional[Any]:





        '''simple docstring'''

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


    def 	SCREAMING_SNAKE_CASE							(      self				:     int)						->Optional[int]:





        '''simple docstring'''

        A__							=   StableDiffusionSAGPipeline.from_pretrained('''CompVis/stable-diffusion-v1-4''')
        A__							=   sag_pipe.to(UpperCAmelCase__)
        sag_pipe.set_progress_bar_config(disable=UpperCAmelCase__)

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

        A__							=   output.images

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

        assert image.shape == (1, 512, 512, 3)
        A__							=   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				:     List[Any])						->int:





        '''simple docstring'''

        A__							=   StableDiffusionSAGPipeline.from_pretrained('''stabilityai/stable-diffusion-2-1-base''')
        A__							=   sag_pipe.to(UpperCAmelCase__)
        sag_pipe.set_progress_bar_config(disable=UpperCAmelCase__)

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

        A__							=   output.images

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

        assert image.shape == (1, 512, 512, 3)
        A__							=   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				:     Dict)						->Dict:





        '''simple docstring'''

        A__							=   StableDiffusionSAGPipeline.from_pretrained('''stabilityai/stable-diffusion-2-1-base''')
        A__							=   sag_pipe.to(UpperCAmelCase__)
        sag_pipe.set_progress_bar_config(disable=UpperCAmelCase__)

        A__							=   '''.'''
        A__							=   torch.manual_seed(0)
        A__							=   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__							=   output.images

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








from typing import TYPE_CHECKING

# rely on isort to merge the imports
from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_torch_available


_lowerCamelCase     :		str								=       {"""configuration_focalnet""": ["""FOCALNET_PRETRAINED_CONFIG_ARCHIVE_MAP""", """FocalNetConfig"""]}


try:
  if not is_torch_available():
    raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
  pass
else:
  _lowerCamelCase     :		Dict								=       [
      """FOCALNET_PRETRAINED_MODEL_ARCHIVE_LIST""",
      """FocalNetForImageClassification""",
      """FocalNetForMaskedImageModeling""",
      """FocalNetBackbone""",
      """FocalNetModel""",
      """FocalNetPreTrainedModel""",
  ]

if TYPE_CHECKING:
  from .configuration_focalnet import FOCALNET_PRETRAINED_CONFIG_ARCHIVE_MAP, FocalNetConfig

  try:
    if not is_torch_available():
      raise OptionalDependencyNotAvailable()
  except OptionalDependencyNotAvailable:
    pass
  else:
    from .modeling_focalnet import (
        FOCALNET_PRETRAINED_MODEL_ARCHIVE_LIST,
        FocalNetBackbone,
        FocalNetForImageClassification,
        FocalNetForMaskedImageModeling,
        FocalNetModel,
        FocalNetPreTrainedModel,
    )

else:
  import sys

  _lowerCamelCase     :		Tuple								=       _LazyModule(__name__, globals()["""__file__"""], _import_structure, module_spec=__spec__)


