
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__			    =			StableDiffusionSAGPipeline
  SCREAMING_SNAKE_CASE__			    =			TEXT_TO_IMAGE_PARAMS
  SCREAMING_SNAKE_CASE__			    =			TEXT_TO_IMAGE_BATCH_PARAMS
  SCREAMING_SNAKE_CASE__			    =			TEXT_TO_IMAGE_IMAGE_PARAMS
  SCREAMING_SNAKE_CASE__			    =			TEXT_TO_IMAGE_IMAGE_PARAMS
  SCREAMING_SNAKE_CASE__			    =			False







  def        UpperCAmelCase_					(self						):
       torch.manual_seed(0						)
       UpperCamelCase__				=				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__				=				DDIMScheduler(
           beta_start=0.0_0085       ,			beta_end=0.012       ,			beta_schedule="""scaled_linear"""       ,			clip_sample=_UpperCAmelCase       ,			set_alpha_to_one=_UpperCAmelCase       ,			)
       torch.manual_seed(0						)
       UpperCamelCase__				=				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__				=				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=10_00       ,			)
       UpperCamelCase__				=				CLIPTextModel(_UpperCAmelCase						)
       UpperCamelCase__				=				CLIPTokenizer.from_pretrained("""hf-internal-testing/tiny-random-clip"""						)

       UpperCamelCase__				=				{
           'unet': unet,
           'scheduler': scheduler,
           'vae': vae,
           'text_encoder': text_encoder,
           'tokenizer': tokenizer,
           'safety_checker': None,
           'feature_extractor': None,
       }
       return components







  def        UpperCAmelCase_					(self       ,			SCREAMING_SNAKE_CASE_       ,			SCREAMING_SNAKE_CASE_=0						):






       if str(_UpperCAmelCase						).startswith("""mps"""						):
            UpperCamelCase__				=				torch.manual_seed(_UpperCAmelCase						)
       else:
            UpperCamelCase__				=				torch.Generator(device=_UpperCAmelCase						).manual_seed(_UpperCAmelCase						)
       UpperCamelCase__				=				{
           'prompt': '.',
           'generator': generator,
           'num_inference_steps': 2,
           'guidance_scale': 1.0,
           'sag_scale': 1.0,
           'output_type': 'numpy',
       }
       return inputs







  def        UpperCAmelCase_					(self						):
       super().test_inference_batch_single_identical(expected_max_diff=3E-3						)






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

  """simple docstring"""







  def        UpperCAmelCase_					(self						):
       super().tearDown()
       gc.collect()
       torch.cuda.empty_cache()







  def        UpperCAmelCase_					(self						):
       UpperCamelCase__				=				StableDiffusionSAGPipeline.from_pretrained("""CompVis/stable-diffusion-v1-4"""						)
       UpperCamelCase__				=				sag_pipe.to(_UpperCAmelCase						)
       sag_pipe.set_progress_bar_config(disable=_UpperCAmelCase						)

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

       UpperCamelCase__				=				output.images

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

       assert image.shape == (1, 5_12, 5_12, 3)
       UpperCamelCase__				=				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        UpperCAmelCase_					(self						):
       UpperCamelCase__				=				StableDiffusionSAGPipeline.from_pretrained("""stabilityai/stable-diffusion-2-1-base"""						)
       UpperCamelCase__				=				sag_pipe.to(_UpperCAmelCase						)
       sag_pipe.set_progress_bar_config(disable=_UpperCAmelCase						)

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

       UpperCamelCase__				=				output.images

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

       assert image.shape == (1, 5_12, 5_12, 3)
       UpperCamelCase__				=				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        UpperCAmelCase_					(self						):
       UpperCamelCase__				=				StableDiffusionSAGPipeline.from_pretrained("""stabilityai/stable-diffusion-2-1-base"""						)
       UpperCamelCase__				=				sag_pipe.to(_UpperCAmelCase						)
       sag_pipe.set_progress_bar_config(disable=_UpperCAmelCase						)

       UpperCamelCase__				=				'.'
       UpperCamelCase__				=				torch.manual_seed(0						)
       UpperCamelCase__				=				sag_pipe(
           [prompt]       ,			width=7_68       ,			height=5_12       ,			generator=_UpperCAmelCase       ,			guidance_scale=7.5       ,			sag_scale=1.0       ,			num_inference_steps=20       ,			output_type="""np"""       ,			)

       UpperCamelCase__				=				output.images

       assert image.shape == (1, 5_12, 7_68, 3)












'''simple docstring'''




import os
import tempfile
import unittest

import numpy as np

from diffusers.utils import is_flax_available
from diffusers.utils.testing_utils import require_flax, slow


if is_flax_available():
  import jax
  import jax.numpy as jnp
  from flax.jax_utils import replicate
  from flax.training.common_utils import shard

  from diffusers import FlaxDDIMScheduler, FlaxDiffusionPipeline, FlaxStableDiffusionPipeline


@require_flax
class  SCREAMING_SNAKE_CASE    (unittest.TestCase							):


     def 							SCREAMING_SNAKE_CASE  (						self):

          '''simple docstring'''






          with tempfile.TemporaryDirectory() as tmpdirname:
               # pipeline has Flax weights
               __A						: Union[str, Any]      =       FlaxDiffusionPipeline.from_pretrained(
                   'hf-internal-testing/tiny-stable-diffusion-pipe'						,	safety_checker=_UpperCAmelCase						,	cache_dir=_UpperCAmelCase)

               __A						: Optional[Any]      =       [t[-1] for t in os.walk(os.path.join(_UpperCAmelCase						,	os.listdir(_UpperCAmelCase)[0]						,	'snapshots'))]
               __A						: int      =       [item for sublist in all_root_files for item in sublist]

               # None of the downloaded files should be a PyTorch file even if we have some here:
               # https://huggingface.co/hf-internal-testing/tiny-stable-diffusion-pipe/blob/main/unet/diffusion_pytorch_model.bin
               assert not any(f.endswith('.bin') for f in files)




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


     def 							SCREAMING_SNAKE_CASE  (						self):

          '''simple docstring'''






          __A ,__A						: Optional[int]      =       FlaxStableDiffusionPipeline.from_pretrained(
              'hf-internal-testing/tiny-stable-diffusion-pipe'						,	safety_checker=_UpperCAmelCase)

          __A						: Dict      =       (
              'A cinematic film still of Morgan Freeman starring as Jimi Hendrix, portrait, 40mm lens, shallow depth of'
              ' field, close up, split lighting, cinematic'
          )

          __A						: Optional[Any]      =       jax.random.PRNGKey(0)
          __A						: int      =       4

          __A						: Tuple      =       jax.device_count()
          __A						: Union[str, Any]      =       num_samples * [prompt]
          __A						: Tuple      =       pipeline.prepare_inputs(_UpperCAmelCase)

          # shard inputs and rng
          __A						: str      =       replicate(_UpperCAmelCase)
          __A						: Tuple      =       jax.random.split(_UpperCAmelCase						,	_UpperCAmelCase)
          __A						: Union[str, Any]      =       shard(_UpperCAmelCase)

          __A						: Union[str, Any]      =       pipeline(_UpperCAmelCase						,	_UpperCAmelCase						,	_UpperCAmelCase						,	_UpperCAmelCase						,	jit=_UpperCAmelCase).images

          assert images.shape == (num_samples, 1, 64, 64, 3)
          if jax.device_count() == 8:
               assert np.abs(np.abs(images[0, 0, :2, :2, -2:]						,	dtype=np.floataa).sum() - 4.1514745) < 1e-3
               assert np.abs(np.abs(_UpperCAmelCase						,	dtype=np.floataa).sum() - 49947.875) < 5e-1

          __A						: List[str]      =       pipeline.numpy_to_pil(np.asarray(images.reshape((num_samples,) + images.shape[-3:])))
          assert len(_UpperCAmelCase) == num_samples


     def 							SCREAMING_SNAKE_CASE  (						self):

          '''simple docstring'''






          __A ,__A						: str      =       FlaxStableDiffusionPipeline.from_pretrained(
              'CompVis/stable-diffusion-v1-4'						,	revision='flax'						,	safety_checker=_UpperCAmelCase)

          __A						: List[Any]      =       (
              'A cinematic film still of Morgan Freeman starring as Jimi Hendrix, portrait, 40mm lens, shallow depth of'
              ' field, close up, split lighting, cinematic'
          )

          __A						: Tuple      =       jax.random.PRNGKey(0)
          __A						: Any      =       50

          __A						: str      =       jax.device_count()
          __A						: Union[str, Any]      =       num_samples * [prompt]
          __A						: List[str]      =       pipeline.prepare_inputs(_UpperCAmelCase)

          # shard inputs and rng
          __A						: Dict      =       replicate(_UpperCAmelCase)
          __A						: Optional[Any]      =       jax.random.split(_UpperCAmelCase						,	_UpperCAmelCase)
          __A						: int      =       shard(_UpperCAmelCase)

          __A						: Tuple      =       pipeline(_UpperCAmelCase						,	_UpperCAmelCase						,	_UpperCAmelCase						,	_UpperCAmelCase						,	jit=_UpperCAmelCase).images

          assert images.shape == (num_samples, 1, 512, 512, 3)
          if jax.device_count() == 8:
               assert np.abs((np.abs(images[0, 0, :2, :2, -2:]						,	dtype=np.floataa).sum() - 0.05652401)) < 1e-3
               assert np.abs((np.abs(_UpperCAmelCase						,	dtype=np.floataa).sum() - 2383808.2)) < 5e-1


     def 							SCREAMING_SNAKE_CASE  (						self):

          '''simple docstring'''






          __A ,__A						: List[str]      =       FlaxStableDiffusionPipeline.from_pretrained(
              'CompVis/stable-diffusion-v1-4'						,	revision='bf16'						,	dtype=jnp.bfloataa						,	safety_checker=_UpperCAmelCase)

          __A						: List[Any]      =       (
              'A cinematic film still of Morgan Freeman starring as Jimi Hendrix, portrait, 40mm lens, shallow depth of'
              ' field, close up, split lighting, cinematic'
          )

          __A						: str      =       jax.random.PRNGKey(0)
          __A						: Any      =       50

          __A						: Optional[int]      =       jax.device_count()
          __A						: int      =       num_samples * [prompt]
          __A						: Optional[int]      =       pipeline.prepare_inputs(_UpperCAmelCase)

          # shard inputs and rng
          __A						: Optional[int]      =       replicate(_UpperCAmelCase)
          __A						: List[str]      =       jax.random.split(_UpperCAmelCase						,	_UpperCAmelCase)
          __A						: Dict      =       shard(_UpperCAmelCase)

          __A						: str      =       pipeline(_UpperCAmelCase						,	_UpperCAmelCase						,	_UpperCAmelCase						,	_UpperCAmelCase						,	jit=_UpperCAmelCase).images

          assert images.shape == (num_samples, 1, 512, 512, 3)
          if jax.device_count() == 8:
               assert np.abs((np.abs(images[0, 0, :2, :2, -2:]						,	dtype=np.floataa).sum() - 0.04003906)) < 1e-3
               assert np.abs((np.abs(_UpperCAmelCase						,	dtype=np.floataa).sum() - 2373516.75)) < 5e-1


     def 							SCREAMING_SNAKE_CASE  (						self):

          '''simple docstring'''






          __A ,__A						: Optional[Any]      =       FlaxStableDiffusionPipeline.from_pretrained(
              'CompVis/stable-diffusion-v1-4'						,	revision='bf16'						,	dtype=jnp.bfloataa)

          __A						: Union[str, Any]      =       (
              'A cinematic film still of Morgan Freeman starring as Jimi Hendrix, portrait, 40mm lens, shallow depth of'
              ' field, close up, split lighting, cinematic'
          )

          __A						: Any      =       jax.random.PRNGKey(0)
          __A						: List[str]      =       50

          __A						: Optional[int]      =       jax.device_count()
          __A						: List[Any]      =       num_samples * [prompt]
          __A						: List[Any]      =       pipeline.prepare_inputs(_UpperCAmelCase)

          # shard inputs and rng
          __A						: Union[str, Any]      =       replicate(_UpperCAmelCase)
          __A						: Optional[Any]      =       jax.random.split(_UpperCAmelCase						,	_UpperCAmelCase)
          __A						: List[str]      =       shard(_UpperCAmelCase)

          __A						: int      =       pipeline(_UpperCAmelCase						,	_UpperCAmelCase						,	_UpperCAmelCase						,	_UpperCAmelCase						,	jit=_UpperCAmelCase).images

          assert images.shape == (num_samples, 1, 512, 512, 3)
          if jax.device_count() == 8:
               assert np.abs((np.abs(images[0, 0, :2, :2, -2:]						,	dtype=np.floataa).sum() - 0.04003906)) < 1e-3
               assert np.abs((np.abs(_UpperCAmelCase						,	dtype=np.floataa).sum() - 2373516.75)) < 5e-1


     def 							SCREAMING_SNAKE_CASE  (						self):

          '''simple docstring'''






          __A						: List[Any]      =       FlaxDDIMScheduler(
              beta_start=0.00085						,	beta_end=0.012						,	beta_schedule='scaled_linear'						,	set_alpha_to_one=_UpperCAmelCase						,	steps_offset=1						,	)

          __A ,__A						: Any      =       FlaxStableDiffusionPipeline.from_pretrained(
              'CompVis/stable-diffusion-v1-4'						,	revision='bf16'						,	dtype=jnp.bfloataa						,	scheduler=_UpperCAmelCase						,	safety_checker=_UpperCAmelCase						,	)
          __A						: Optional[Any]      =       scheduler.create_state()

          __A						: Any      =       scheduler_state

          __A						: List[str]      =       (
              'A cinematic film still of Morgan Freeman starring as Jimi Hendrix, portrait, 40mm lens, shallow depth of'
              ' field, close up, split lighting, cinematic'
          )

          __A						: Union[str, Any]      =       jax.random.PRNGKey(0)
          __A						: Optional[int]      =       50

          __A						: Optional[Any]      =       jax.device_count()
          __A						: Any      =       num_samples * [prompt]
          __A						: Optional[Any]      =       pipeline.prepare_inputs(_UpperCAmelCase)

          # shard inputs and rng
          __A						: int      =       replicate(_UpperCAmelCase)
          __A						: Any      =       jax.random.split(_UpperCAmelCase						,	_UpperCAmelCase)
          __A						: Tuple      =       shard(_UpperCAmelCase)

          __A						: Union[str, Any]      =       pipeline(_UpperCAmelCase						,	_UpperCAmelCase						,	_UpperCAmelCase						,	_UpperCAmelCase						,	jit=_UpperCAmelCase).images

          assert images.shape == (num_samples, 1, 512, 512, 3)
          if jax.device_count() == 8:
               assert np.abs((np.abs(images[0, 0, :2, :2, -2:]						,	dtype=np.floataa).sum() - 0.045043945)) < 1e-3
               assert np.abs((np.abs(_UpperCAmelCase						,	dtype=np.floataa).sum() - 2347693.5)) < 5e-1


     def 							SCREAMING_SNAKE_CASE  (						self):

          '''simple docstring'''






          __A						: Dict      =       (
              'A cinematic film still of Morgan Freeman starring as Jimi Hendrix, portrait, 40mm lens, shallow depth of'
              ' field, close up, split lighting, cinematic'
          )

          __A						: int      =       jax.device_count()
          __A						: List[Any]      =       num_samples * [prompt]
          __A						: List[Any]      =       jax.random.split(jax.random.PRNGKey(0)						,	_UpperCAmelCase)

          __A ,__A						: Tuple      =       FlaxStableDiffusionPipeline.from_pretrained(
              'CompVis/stable-diffusion-v1-4'						,	revision='bf16'						,	dtype=jnp.bfloataa						,	safety_checker=_UpperCAmelCase						,	)

          __A						: str      =       replicate(_UpperCAmelCase)
          __A						: str      =       pipeline.prepare_inputs(_UpperCAmelCase)
          __A						: str      =       shard(_UpperCAmelCase)
          __A						: int      =       pipeline(_UpperCAmelCase						,	_UpperCAmelCase						,	_UpperCAmelCase						,	jit=_UpperCAmelCase).images
          assert images.shape == (num_samples, 1, 512, 512, 3)
          __A						: Any      =       images[2, 0, 256, 10:17, 1]

          # With memory efficient attention
          __A ,__A						: str      =       FlaxStableDiffusionPipeline.from_pretrained(
              'CompVis/stable-diffusion-v1-4'						,	revision='bf16'						,	dtype=jnp.bfloataa						,	safety_checker=_UpperCAmelCase						,	use_memory_efficient_attention=_UpperCAmelCase						,	)

          __A						: Any      =       replicate(_UpperCAmelCase)
          __A						: List[Any]      =       pipeline.prepare_inputs(_UpperCAmelCase)
          __A						: Optional[Any]      =       shard(_UpperCAmelCase)
          __A						: List[Any]      =       pipeline(_UpperCAmelCase						,	_UpperCAmelCase						,	_UpperCAmelCase						,	jit=_UpperCAmelCase).images
          assert images_eff.shape == (num_samples, 1, 512, 512, 3)
          __A						: List[Any]      =       images[2, 0, 256, 10:17, 1]

          # I checked the results visually and they are very similar. However, I saw that the max diff is `1` and the `sum`
          # over the 8 images is exactly `256`, which is very suspicious. Testing a random slice for now.
          assert abs(slice_eff - slice).max() < 1e-2