





import gc
import random
import tempfile
import unittest

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

from diffusers import AutoencoderKL, DDIMScheduler, LMSDiscreteScheduler, PNDMScheduler, UNetaDConditionModel
from diffusers.pipelines.stable_diffusion_safe import StableDiffusionPipelineSafe as StableDiffusionPipeline
from diffusers.utils import floats_tensor, nightly, torch_device
from diffusers.utils.testing_utils import require_torch_gpu







class   __lowerCAmelCase			(					unittest.TestCase  ):





   def 	_snake_case ( self      )				->  Dict:
      """simple docstring"""





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





   @property
   def 	_snake_case ( self      )				->  Dict:
      """simple docstring"""





      a__	:						Any      		=  1
      a__	:						Dict      		=  3
      a__	:						Union[str, Any]      		=  (32, 32)

      a__	:						Union[str, Any]      		=  floats_tensor((batch_size, num_channels) + sizes					,			rng=random.Random(0      )      ).to(__A      )
      return image





   @property
   def 	_snake_case ( self      )				->  str:
      """simple docstring"""





      torch.manual_seed(0      )
      a__	:						List[Any]      		=  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					,			)
      return model





   @property
   def 	_snake_case ( self      )				->  int:
      """simple docstring"""





      torch.manual_seed(0      )
      a__	:						Dict      		=  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					,			)
      return model





   @property
   def 	_snake_case ( self      )				->  str:
      """simple docstring"""





      torch.manual_seed(0      )
      a__	:						List[str]      		=  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					,			)
      return CLIPTextModel(__A      )





   @property
   def 	_snake_case ( self      )				->  str:
      """simple docstring"""





      def extract(*snake_case					,			**snake_case      ):







         class   __lowerCAmelCase			:





            def __init__( self      )				->  Union[str, Any]:
               """simple docstring"""





               a__	:						str      		=  torch.ones([0]      )





            def 	_snake_case ( self					,			snake_case      )				->  Tuple:
               """simple docstring"""





               self.pixel_values.to(__A      )
               return self







         return Out()

      return extract





   def 	_snake_case ( self      )				->  str:
      """simple docstring"""





      a__	:						Optional[Any]      		=  "cpu"  # ensure determinism for the device-dependent torch.Generator
      a__	:						Optional[Any]      		=  self.dummy_cond_unet
      a__	:						Optional[Any]      		=  DDIMScheduler(
          beta_start=0.00_085					,			beta_end=0.012					,			beta_schedule="scaled_linear"					,			clip_sample=__A					,			set_alpha_to_one=__A					,			)

      a__	:						Dict      		=  self.dummy_vae
      a__	:						Union[str, Any]      		=  self.dummy_text_encoder
      a__	:						Dict      		=  CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip"      )

      # make sure here that pndm scheduler skips prk
      a__	:						Dict      		=  StableDiffusionPipeline(
          unet=__A					,			scheduler=__A					,			vae=__A					,			text_encoder=__A					,			tokenizer=__A					,			safety_checker=__A					,			feature_extractor=self.dummy_extractor					,			)
      a__	:						Any      		=  sd_pipe.to(__A      )
      sd_pipe.set_progress_bar_config(disable=__A      )

      a__	:						str      		=  "A painting of a squirrel eating a burger"

      a__	:						Optional[int]      		=  torch.Generator(device=__A      ).manual_seed(0      )
      a__	:						Optional[Any]      		=  sd_pipe([prompt]					,			generator=__A					,			guidance_scale=6.0					,			num_inference_steps=2					,			output_type="np"      )
      a__	:						int      		=  output.images

      a__	:						Dict      		=  torch.Generator(device=__A      ).manual_seed(0      )
      a__	:						int      		=  sd_pipe(
          [prompt]					,			generator=__A					,			guidance_scale=6.0					,			num_inference_steps=2					,			output_type="np"					,			return_dict=__A					,			)[0]

      a__	:						Tuple      		=  image[0, -3:, -3:, -1]
      a__	:						Optional[Any]      		=  image_from_tuple[0, -3:, -3:, -1]

      assert image.shape == (1, 64, 64, 3)
      a__	:						int      		=  np.array([0.5_756, 0.6_118, 0.5_005, 0.5_041, 0.5_471, 0.4_726, 0.4_976, 0.4_865, 0.4_864]      )

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





   def 	_snake_case ( self      )				->  str:
      """simple docstring"""





      a__	:						List[Any]      		=  "cpu"  # ensure determinism for the device-dependent torch.Generator
      a__	:						Any      		=  self.dummy_cond_unet
      a__	:						List[Any]      		=  PNDMScheduler(skip_prk_steps=__A      )
      a__	:						List[Any]      		=  self.dummy_vae
      a__	:						Tuple      		=  self.dummy_text_encoder
      a__	:						Optional[Any]      		=  CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip"      )

      # make sure here that pndm scheduler skips prk
      a__	:						Dict      		=  StableDiffusionPipeline(
          unet=__A					,			scheduler=__A					,			vae=__A					,			text_encoder=__A					,			tokenizer=__A					,			safety_checker=__A					,			feature_extractor=self.dummy_extractor					,			)
      a__	:						List[str]      		=  sd_pipe.to(__A      )
      sd_pipe.set_progress_bar_config(disable=__A      )

      a__	:						List[str]      		=  "A painting of a squirrel eating a burger"
      a__	:						Union[str, Any]      		=  torch.Generator(device=__A      ).manual_seed(0      )
      a__	:						Optional[int]      		=  sd_pipe([prompt]					,			generator=__A					,			guidance_scale=6.0					,			num_inference_steps=2					,			output_type="np"      )

      a__	:						Tuple      		=  output.images

      a__	:						int      		=  torch.Generator(device=__A      ).manual_seed(0      )
      a__	:						str      		=  sd_pipe(
          [prompt]					,			generator=__A					,			guidance_scale=6.0					,			num_inference_steps=2					,			output_type="np"					,			return_dict=__A					,			)[0]

      a__	:						List[str]      		=  image[0, -3:, -3:, -1]
      a__	:						Optional[int]      		=  image_from_tuple[0, -3:, -3:, -1]

      assert image.shape == (1, 64, 64, 3)
      a__	:						Dict      		=  np.array([0.5_125, 0.5_716, 0.4_828, 0.5_060, 0.5_650, 0.4_768, 0.5_185, 0.4_895, 0.4_993]      )

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





   def 	_snake_case ( self      )				->  Optional[int]:
      """simple docstring"""





      a__	:						Dict      		=  StableDiffusionPipeline.from_pretrained(
          "hf-internal-testing/tiny-stable-diffusion-lms-pipe"					,			safety_checker=__A      )
      assert isinstance(__A					,			__A      )
      assert isinstance(pipe.scheduler					,			__A      )
      assert pipe.safety_checker is None

      a__	:						List[str]      		=  pipe("example prompt"					,			num_inference_steps=2      ).images[0]
      assert image is not None

      # check that there's no error when saving a pipeline with one of the models being None
      with tempfile.TemporaryDirectory() as tmpdirname:
         pipe.save_pretrained(__A      )
         a__	:						Optional[int]      		=  StableDiffusionPipeline.from_pretrained(__A      )

      # sanity check that the pipeline still works
      assert pipe.safety_checker is None
      a__	:						Dict      		=  pipe("example prompt"					,			num_inference_steps=2      ).images[0]
      assert image is not None





   @unittest.skipIf(torch_device != "cuda"					,			"This test requires a GPU"      )
   def 	_snake_case ( self      )				->  Optional[int]:
      """simple docstring"""





      a__	:						Any      		=  self.dummy_cond_unet
      a__	:						List[str]      		=  PNDMScheduler(skip_prk_steps=__A      )
      a__	:						List[Any]      		=  self.dummy_vae
      a__	:						Union[str, Any]      		=  self.dummy_text_encoder
      a__	:						str      		=  CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip"      )

      # put models in fp16
      a__	:						List[Any]      		=  unet.half()
      a__	:						Dict      		=  vae.half()
      a__	:						str      		=  bert.half()

      # make sure here that pndm scheduler skips prk
      a__	:						Dict      		=  StableDiffusionPipeline(
          unet=__A					,			scheduler=__A					,			vae=__A					,			text_encoder=__A					,			tokenizer=__A					,			safety_checker=__A					,			feature_extractor=self.dummy_extractor					,			)
      a__	:						Tuple      		=  sd_pipe.to(__A      )
      sd_pipe.set_progress_bar_config(disable=__A      )

      a__	:						List[str]      		=  "A painting of a squirrel eating a burger"
      a__	:						Optional[Any]      		=  sd_pipe([prompt]					,			num_inference_steps=2					,			output_type="np"      ).images

      assert image.shape == (1, 64, 64, 3)







@nightly
@require_torch_gpu
class   __lowerCAmelCase			(					unittest.TestCase  ):





   def 	_snake_case ( self      )				->  Optional[int]:
      """simple docstring"""





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





   def 	_snake_case ( self      )				->  Any:
      """simple docstring"""





      a__	:						Dict      		=  StableDiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5"					,			safety_checker=__A      )
      a__	:						Union[str, Any]      		=  LMSDiscreteScheduler.from_config(sd_pipe.scheduler.config      )
      a__	:						Dict      		=  sd_pipe.to(__A      )
      sd_pipe.set_progress_bar_config(disable=__A      )

      a__	:						Optional[int]      		=  (
          "portrait of girl with smokey eyes makeup in abandoned hotel, grange clothes, redshift, wide high angle"
          " coloured polaroid photograph with flash, kodak film, hyper real, stunning moody cinematography, with"
          " anamorphic lenses, by maripol, fallen angels by wong kar - wai, style of suspiria and neon demon and"
          " children from bahnhof zoo, detailed "
      )
      a__	:						Optional[int]      		=  4_003_660_346
      a__	:						Tuple      		=  7

      # without safety guidance (sld_guidance_scale = 0)
      a__	:						Optional[Any]      		=  torch.manual_seed(__A      )
      a__	:						Optional[int]      		=  sd_pipe(
          [prompt]					,			generator=__A					,			guidance_scale=__A					,			num_inference_steps=50					,			output_type="np"					,			width=512					,			height=512					,			sld_guidance_scale=0					,			)

      a__	:						Dict      		=  output.images
      a__	:						List[Any]      		=  image[0, -3:, -3:, -1]
      a__	:						Tuple      		=  [0.2_278, 0.2_231, 0.2_249, 0.2_333, 0.2_303, 0.1_885, 0.2_273, 0.2_144, 0.2_176]

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

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

      # without safety guidance (strong configuration)
      a__	:						Any      		=  torch.manual_seed(__A      )
      a__	:						Dict      		=  sd_pipe(
          [prompt]					,			generator=__A					,			guidance_scale=__A					,			num_inference_steps=50					,			output_type="np"					,			width=512					,			height=512					,			sld_guidance_scale=2_000					,			sld_warmup_steps=7					,			sld_threshold=0.025					,			sld_momentum_scale=0.5					,			sld_mom_beta=0.7					,			)

      a__	:						List[str]      		=  output.images
      a__	:						Any      		=  image[0, -3:, -3:, -1]
      a__	:						str      		=  [0.2_383, 0.2_276, 0.236, 0.2_192, 0.2_186, 0.2_053, 0.1_971, 0.1_901, 0.1_719]

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

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





   def 	_snake_case ( self      )				->  Union[str, Any]:
      """simple docstring"""





      a__	:						Any      		=  StableDiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5"					,			safety_checker=__A      )
      a__	:						Any      		=  LMSDiscreteScheduler.from_config(sd_pipe.scheduler.config      )
      a__	:						List[str]      		=  sd_pipe.to(__A      )
      sd_pipe.set_progress_bar_config(disable=__A      )

      a__	:						Optional[Any]      		=  "padme amidala taking a bath artwork, safe for work, no nudity"
      a__	:						Any      		=  2_734_971_755
      a__	:						List[Any]      		=  7

      a__	:						int      		=  torch.manual_seed(__A      )
      a__	:						Any      		=  sd_pipe(
          [prompt]					,			generator=__A					,			guidance_scale=__A					,			num_inference_steps=50					,			output_type="np"					,			width=512					,			height=512					,			sld_guidance_scale=0					,			)

      a__	:						List[Any]      		=  output.images
      a__	:						Optional[int]      		=  image[0, -3:, -3:, -1]
      a__	:						Optional[Any]      		=  [0.3_502, 0.3_622, 0.3_396, 0.3_642, 0.3_478, 0.3_318, 0.35, 0.3_348, 0.3_297]

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

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

      a__	:						Dict      		=  torch.manual_seed(__A      )
      a__	:						str      		=  sd_pipe(
          [prompt]					,			generator=__A					,			guidance_scale=__A					,			num_inference_steps=50					,			output_type="np"					,			width=512					,			height=512					,			sld_guidance_scale=2_000					,			sld_warmup_steps=7					,			sld_threshold=0.025					,			sld_momentum_scale=0.5					,			sld_mom_beta=0.7					,			)

      a__	:						Optional[int]      		=  output.images
      a__	:						Dict      		=  image[0, -3:, -3:, -1]
      a__	:						Any      		=  [0.5_531, 0.5_206, 0.4_895, 0.5_156, 0.5_182, 0.4_751, 0.4_802, 0.4_803, 0.4_443]

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

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





   def 	_snake_case ( self      )				->  Optional[Any]:
      """simple docstring"""





      a__	:						str      		=  StableDiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5"      )
      a__	:						Any      		=  sd_pipe.to(__A      )
      sd_pipe.set_progress_bar_config(disable=__A      )

      a__	:						Union[str, Any]      		=  (
          "the four horsewomen of the apocalypse, painting by tom of finland, gaston bussiere, craig mullins, j. c."
          " leyendecker"
      )
      a__	:						Any      		=  1_044_355_234
      a__	:						List[str]      		=  12

      a__	:						Dict      		=  torch.manual_seed(__A      )
      a__	:						List[str]      		=  sd_pipe(
          [prompt]					,			generator=__A					,			guidance_scale=__A					,			num_inference_steps=50					,			output_type="np"					,			width=512					,			height=512					,			sld_guidance_scale=0					,			)

      a__	:						List[Any]      		=  output.images
      a__	:						Any      		=  image[0, -3:, -3:, -1]
      a__	:						str      		=  np.array([0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]      )

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

      assert np.abs(image_slice.flatten() - expected_slice      ).max() < 1E-7

      a__	:						Tuple      		=  torch.manual_seed(__A      )
      a__	:						Any      		=  sd_pipe(
          [prompt]					,			generator=__A					,			guidance_scale=__A					,			num_inference_steps=50					,			output_type="np"					,			width=512					,			height=512					,			sld_guidance_scale=2_000					,			sld_warmup_steps=7					,			sld_threshold=0.025					,			sld_momentum_scale=0.5					,			sld_mom_beta=0.7					,			)

      a__	:						Optional[Any]      		=  output.images
      a__	:						Union[str, Any]      		=  image[0, -3:, -3:, -1]
      a__	:						Any      		=  np.array([0.5_818, 0.6_285, 0.6_835, 0.6_019, 0.625, 0.6_754, 0.6_096, 0.6_334, 0.6_561]      )
      assert image.shape == (1, 512, 512, 3)

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








# Copyright 2022 The HuggingFace Team and The OpenBMB Team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
from typing import TYPE_CHECKING

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


SCREAMING_SNAKE_CASE__       :			int			      =						{
    """configuration_cpmant""": ["""CPMANT_PRETRAINED_CONFIG_ARCHIVE_MAP""", """CpmAntConfig"""],
    """tokenization_cpmant""": ["""CpmAntTokenizer"""],
}

try:
  if not is_torch_available():
    raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
  pass
else:
  SCREAMING_SNAKE_CASE__       :			List[str]			      =						[
      """CPMANT_PRETRAINED_MODEL_ARCHIVE_LIST""",
      """CpmAntForCausalLM""",
      """CpmAntModel""",
      """CpmAntPreTrainedModel""",
  ]


if TYPE_CHECKING:
  from .configuration_cpmant import CPMANT_PRETRAINED_CONFIG_ARCHIVE_MAP, CpmAntConfig
  from .tokenization_cpmant import CpmAntTokenizer

  try:
    if not is_torch_available():
      raise OptionalDependencyNotAvailable()
  except OptionalDependencyNotAvailable:
    pass
  else:
    from .modeling_cpmant import (
        CPMANT_PRETRAINED_MODEL_ARCHIVE_LIST,
        CpmAntForCausalLM,
        CpmAntModel,
        CpmAntPreTrainedModel,
    )


else:
  import sys

  SCREAMING_SNAKE_CASE__       :			Union[str, Any]			      =						_LazyModule(__name__, globals()["""__file__"""], _import_structure, module_spec=__spec__)

