

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"""




  UpperCAmelCase__      : Optional[Any]			       =StableDiffusionSAGPipeline
  UpperCAmelCase__      : Any			       =TEXT_TO_IMAGE_PARAMS
  UpperCAmelCase__      : Dict			       =TEXT_TO_IMAGE_BATCH_PARAMS
  UpperCAmelCase__      : List[str]			       =TEXT_TO_IMAGE_IMAGE_PARAMS
  UpperCAmelCase__      : int			       =TEXT_TO_IMAGE_IMAGE_PARAMS
  UpperCAmelCase__      : Optional[Any]			       =False


  def  _lowercase      (			self    : Dict		)					->str:
        """simple docstring"""







        torch.manual_seed(0		)
        SCREAMING_SNAKE_CASE :      str       =   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""")				,				cross_attention_dim=3_2				,				)
        SCREAMING_SNAKE_CASE :      Optional[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		)
        SCREAMING_SNAKE_CASE :      List[str]       =   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				,				)
        torch.manual_seed(0		)
        SCREAMING_SNAKE_CASE :      Tuple       =   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				,				)
        SCREAMING_SNAKE_CASE :      Optional[Any]       =   CLIPTextModel(UpperCAmelCase__		)
        SCREAMING_SNAKE_CASE :      List[Any]       =   CLIPTokenizer.from_pretrained("""hf-internal-testing/tiny-random-clip"""		)

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


  def  _lowercase      (			self    : int				,				UpperCAmelCase__    : Any				,				UpperCAmelCase__    : int=0		)					->Union[str, Any]:
        """simple docstring"""







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




  def  _lowercase      (			self    : Any		)					->str:
        """simple docstring"""







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







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







  """simple docstring"""




  def  _lowercase      (			self    : int		)					->List[str]:
        """simple docstring"""







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


  def  _lowercase      (			self    : int		)					->Any:
        """simple docstring"""







        SCREAMING_SNAKE_CASE :      Optional[int]       =   StableDiffusionSAGPipeline.from_pretrained("""CompVis/stable-diffusion-v1-4"""		)
        SCREAMING_SNAKE_CASE :      Union[str, Any]       =   sag_pipe.to(UpperCAmelCase__		)
        sag_pipe.set_progress_bar_config(disable=UpperCAmelCase__		)

        SCREAMING_SNAKE_CASE :      List[str]       =   """."""
        SCREAMING_SNAKE_CASE :      Optional[Any]       =   torch.manual_seed(0		)
        SCREAMING_SNAKE_CASE :      Optional[int]       =   sag_pipe(
            [prompt]				,				generator=UpperCAmelCase__				,				guidance_scale=7.5				,				sag_scale=1.0				,				num_inference_steps=2_0				,				output_type="""np"""		)

        SCREAMING_SNAKE_CASE :      Tuple       =   output.images

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

        assert image.shape == (1, 5_1_2, 5_1_2, 3)
        SCREAMING_SNAKE_CASE :      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  _lowercase      (			self    : List[str]		)					->int:
        """simple docstring"""







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

        SCREAMING_SNAKE_CASE :      Optional[int]       =   """."""
        SCREAMING_SNAKE_CASE :      int       =   torch.manual_seed(0		)
        SCREAMING_SNAKE_CASE :      Any       =   sag_pipe(
            [prompt]				,				generator=UpperCAmelCase__				,				guidance_scale=7.5				,				sag_scale=1.0				,				num_inference_steps=2_0				,				output_type="""np"""		)

        SCREAMING_SNAKE_CASE :      int       =   output.images

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

        assert image.shape == (1, 5_1_2, 5_1_2, 3)
        SCREAMING_SNAKE_CASE :      Optional[int]       =   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  _lowercase      (			self    : Optional[int]		)					->Dict:
        """simple docstring"""







        SCREAMING_SNAKE_CASE :      Dict       =   StableDiffusionSAGPipeline.from_pretrained("""stabilityai/stable-diffusion-2-1-base"""		)
        SCREAMING_SNAKE_CASE :      Optional[Any]       =   sag_pipe.to(UpperCAmelCase__		)
        sag_pipe.set_progress_bar_config(disable=UpperCAmelCase__		)

        SCREAMING_SNAKE_CASE :      List[str]       =   """."""
        SCREAMING_SNAKE_CASE :      Any       =   torch.manual_seed(0		)
        SCREAMING_SNAKE_CASE :      Tuple       =   sag_pipe(
            [prompt]				,				width=7_6_8				,				height=5_1_2				,				generator=UpperCAmelCase__				,				guidance_scale=7.5				,				sag_scale=1.0				,				num_inference_steps=2_0				,				output_type="""np"""				,				)

        SCREAMING_SNAKE_CASE :      Tuple       =   output.images

        assert image.shape == (1, 5_1_2, 7_6_8, 3)







from ...configuration_utils import PretrainedConfig
from ...utils import logging


UpperCAmelCase__ :     Dict				    =    logging.get_logger(__name__)

UpperCAmelCase__ :     int				    =    {
    """abeja/gpt-neox-japanese-2.7b""": """https://huggingface.co/abeja/gpt-neox-japanese-2.7b/resolve/main/config.json""",
}







class        a__       (   UpperCAmelCase       ):







  """simple docstring"""




  UpperCAmelCase__      : Any			       ="""gpt_neox_japanese"""


  def __init__(			self    : Any				,				UpperCAmelCase__    : Any=3_2_0_0_0				,				UpperCAmelCase__    : Dict=2_5_6_0				,				UpperCAmelCase__    : List[str]=3_2				,				UpperCAmelCase__    : Optional[int]=3_2				,				UpperCAmelCase__    : List[str]=4				,				UpperCAmelCase__    : List[str]="gelu"				,				UpperCAmelCase__    : Optional[int]=1.00				,				UpperCAmelCase__    : List[Any]=1_0_0_0_0				,				UpperCAmelCase__    : Tuple=2_0_4_8				,				UpperCAmelCase__    : Optional[int]=0.02				,				UpperCAmelCase__    : Dict=1e-5				,				UpperCAmelCase__    : Dict=True				,				UpperCAmelCase__    : Tuple=3_1_9_9_6				,				UpperCAmelCase__    : Tuple=3_1_9_9_9				,				UpperCAmelCase__    : Dict=0.1				,				UpperCAmelCase__    : Union[str, Any]=0.0				,				**UpperCAmelCase__    : Optional[Any]				,				)					->Optional[Any]:
        """simple docstring"""







        super().__init__(bos_token_id=UpperCAmelCase__				,				eos_token_id=UpperCAmelCase__				,				**UpperCAmelCase__		)
        SCREAMING_SNAKE_CASE :      List[str]       =   vocab_size
        SCREAMING_SNAKE_CASE :      List[str]       =   max_position_embeddings
        SCREAMING_SNAKE_CASE :      Tuple       =   hidden_size
        SCREAMING_SNAKE_CASE :      Union[str, Any]       =   num_hidden_layers
        SCREAMING_SNAKE_CASE :      Optional[int]       =   num_attention_heads
        SCREAMING_SNAKE_CASE :      Any       =   intermediate_multiple_size
        SCREAMING_SNAKE_CASE :      Tuple       =   hidden_act
        SCREAMING_SNAKE_CASE :      Any       =   rotary_pct
        SCREAMING_SNAKE_CASE :      Tuple       =   rotary_emb_base
        SCREAMING_SNAKE_CASE :      Union[str, Any]       =   initializer_range
        SCREAMING_SNAKE_CASE :      Optional[Any]       =   layer_norm_eps
        SCREAMING_SNAKE_CASE :      str       =   use_cache
        SCREAMING_SNAKE_CASE :      Any       =   attention_dropout
        SCREAMING_SNAKE_CASE :      str       =   hidden_dropout




