

'''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 				_snake_case					(   a_     , a_     , unittest.TestCase			):
  SCREAMING_SNAKE_CASE			:       Dict				=     StableDiffusionSAGPipeline
  SCREAMING_SNAKE_CASE			:       int				=     TEXT_TO_IMAGE_PARAMS
  SCREAMING_SNAKE_CASE			:       str				=     TEXT_TO_IMAGE_BATCH_PARAMS
  SCREAMING_SNAKE_CASE			:       Tuple				=     TEXT_TO_IMAGE_IMAGE_PARAMS
  SCREAMING_SNAKE_CASE			:       str				=     TEXT_TO_IMAGE_IMAGE_PARAMS
  SCREAMING_SNAKE_CASE			:       Any				=     False

  def   _SCREAMING_SNAKE_CASE				(				self	):





      '''simple docstring'''

      torch.manual_seed(0	)
      lowerCAmelCase 						=							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      , )
      lowerCAmelCase 						=							DDIMScheduler(
          beta_start=0.00_085      , beta_end=0.012      , beta_schedule='scaled_linear'      , clip_sample=_SCREAMING_SNAKE_CASE      , set_alpha_to_one=_SCREAMING_SNAKE_CASE      , )
      torch.manual_seed(0	)
      lowerCAmelCase 						=							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	)
      lowerCAmelCase 						=							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      , )
      lowerCAmelCase 						=							CLIPTextModel(_SCREAMING_SNAKE_CASE	)
      lowerCAmelCase 						=							CLIPTokenizer.from_pretrained('hf-internal-testing/tiny-random-clip'	)

      lowerCAmelCase 						=							{
          '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      , _SCREAMING_SNAKE_CASE      , _SCREAMING_SNAKE_CASE=0	):





      '''simple docstring'''

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

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

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

      lowerCAmelCase 						=							output.images

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

      assert image.shape == (1, 5_12, 5_12, 3)
      lowerCAmelCase 						=							np.array([0.1_568, 0.1_738, 0.1_695, 0.1_693, 0.1_507, 0.1_705, 0.1_547, 0.1_751, 0.1_949]	)

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

  def   _SCREAMING_SNAKE_CASE				(				self	):





      '''simple docstring'''

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

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

      lowerCAmelCase 						=							output.images

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

      assert image.shape == (1, 5_12, 5_12, 3)
      lowerCAmelCase 						=							np.array([0.3_459, 0.2_876, 0.2_537, 0.3_002, 0.2_671, 0.2_160, 0.3_026, 0.2_262, 0.2_371]	)

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


  def   _SCREAMING_SNAKE_CASE				(				self	):





      '''simple docstring'''

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

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

      lowerCAmelCase 						=							output.images

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








'''simple docstring'''



import gc
import unittest

from diffusers import FlaxControlNetModel, FlaxStableDiffusionControlNetPipeline
from diffusers.utils import is_flax_available, load_image, slow
from diffusers.utils.testing_utils import require_flax


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



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

  def   _SCREAMING_SNAKE_CASE				(				self	):





      '''simple docstring'''

      super().tearDown()
      gc.collect()

  def   _SCREAMING_SNAKE_CASE				(				self	):





      '''simple docstring'''

      lowerCAmelCase						,   lowerCAmelCase 						=							FlaxControlNetModel.from_pretrained(
          'lllyasviel/sd-controlnet-canny'      , from_pt=_SCREAMING_SNAKE_CASE      , dtype=jnp.bfloataa	)
      lowerCAmelCase						,   lowerCAmelCase 						=							FlaxStableDiffusionControlNetPipeline.from_pretrained(
          'runwayml/stable-diffusion-v1-5'      , controlnet=_SCREAMING_SNAKE_CASE      , from_pt=_SCREAMING_SNAKE_CASE      , dtype=jnp.bfloataa	)
      lowerCAmelCase 						=							controlnet_params

      lowerCAmelCase 						=							'bird'
      lowerCAmelCase 						=							jax.device_count()
      lowerCAmelCase 						=							pipe.prepare_text_inputs([prompts] * num_samples	)

      lowerCAmelCase 						=							load_image(
          'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/sd_controlnet/bird_canny.png'	)
      lowerCAmelCase 						=							pipe.prepare_image_inputs([canny_image] * num_samples	)

      lowerCAmelCase 						=							jax.random.PRNGKey(0	)
      lowerCAmelCase 						=							jax.random.split(_SCREAMING_SNAKE_CASE      , jax.device_count()	)

      lowerCAmelCase 						=							replicate(_SCREAMING_SNAKE_CASE	)
      lowerCAmelCase 						=							shard(_SCREAMING_SNAKE_CASE	)
      lowerCAmelCase 						=							shard(_SCREAMING_SNAKE_CASE	)

      lowerCAmelCase 						=							pipe(
          prompt_ids=_SCREAMING_SNAKE_CASE      , image=_SCREAMING_SNAKE_CASE      , params=_SCREAMING_SNAKE_CASE      , prng_seed=_SCREAMING_SNAKE_CASE      , num_inference_steps=50      , jit=_SCREAMING_SNAKE_CASE      , ).images
      assert images.shape == (jax.device_count(), 1, 7_68, 5_12, 3)

      lowerCAmelCase 						=							images.reshape((images.shape[0] * images.shape[1],) + images.shape[-3:]	)
      lowerCAmelCase 						=							images[0, 2_53:2_56, 2_53:2_56, -1]

      lowerCAmelCase 						=							jnp.asarray(jax.device_get(image_slice.flatten()	)	)
      lowerCAmelCase 						=							jnp.array(
          [0.167_969, 0.116_699, 0.081_543, 0.154_297, 0.132_812, 0.108_887, 0.169_922, 0.169_922, 0.205_078]	)
      print(F'output_slice: {output_slice}'	)
      assert jnp.abs(output_slice - expected_slice	).max() < 1e-2


  def   _SCREAMING_SNAKE_CASE				(				self	):





      '''simple docstring'''

      lowerCAmelCase						,   lowerCAmelCase 						=							FlaxControlNetModel.from_pretrained(
          'lllyasviel/sd-controlnet-openpose'      , from_pt=_SCREAMING_SNAKE_CASE      , dtype=jnp.bfloataa	)
      lowerCAmelCase						,   lowerCAmelCase 						=							FlaxStableDiffusionControlNetPipeline.from_pretrained(
          'runwayml/stable-diffusion-v1-5'      , controlnet=_SCREAMING_SNAKE_CASE      , from_pt=_SCREAMING_SNAKE_CASE      , dtype=jnp.bfloataa	)
      lowerCAmelCase 						=							controlnet_params

      lowerCAmelCase 						=							'Chef in the kitchen'
      lowerCAmelCase 						=							jax.device_count()
      lowerCAmelCase 						=							pipe.prepare_text_inputs([prompts] * num_samples	)

      lowerCAmelCase 						=							load_image(
          'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/sd_controlnet/pose.png'	)
      lowerCAmelCase 						=							pipe.prepare_image_inputs([pose_image] * num_samples	)

      lowerCAmelCase 						=							jax.random.PRNGKey(0	)
      lowerCAmelCase 						=							jax.random.split(_SCREAMING_SNAKE_CASE      , jax.device_count()	)

      lowerCAmelCase 						=							replicate(_SCREAMING_SNAKE_CASE	)
      lowerCAmelCase 						=							shard(_SCREAMING_SNAKE_CASE	)
      lowerCAmelCase 						=							shard(_SCREAMING_SNAKE_CASE	)

      lowerCAmelCase 						=							pipe(
          prompt_ids=_SCREAMING_SNAKE_CASE      , image=_SCREAMING_SNAKE_CASE      , params=_SCREAMING_SNAKE_CASE      , prng_seed=_SCREAMING_SNAKE_CASE      , num_inference_steps=50      , jit=_SCREAMING_SNAKE_CASE      , ).images
      assert images.shape == (jax.device_count(), 1, 7_68, 5_12, 3)

      lowerCAmelCase 						=							images.reshape((images.shape[0] * images.shape[1],) + images.shape[-3:]	)
      lowerCAmelCase 						=							images[0, 2_53:2_56, 2_53:2_56, -1]

      lowerCAmelCase 						=							jnp.asarray(jax.device_get(image_slice.flatten()	)	)
      lowerCAmelCase 						=							jnp.array(
          [[0.271_484, 0.261_719, 0.275_391, 0.277_344, 0.279_297, 0.291_016, 0.294_922, 0.302_734, 0.302_734]]	)
      print(F'output_slice: {output_slice}'	)
      assert jnp.abs(output_slice - expected_slice	).max() < 1e-2





