


'''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      snake_case    (    unittest.TestCase			):





  """simple docstring"""




  def   snake_case		(    self     ):







   """simple docstring"""






   with tempfile.TemporaryDirectory() as tmpdirname:
    # pipeline has Flax weights
    lowerCamelCase_      =  FlaxDiffusionPipeline.from_pretrained(
        "hf-internal-testing/tiny-stable-diffusion-pipe"			,  safety_checker=_UpperCAmelCase			,  cache_dir=_UpperCAmelCase     )

    lowerCamelCase_      =  [t[-1] for t in os.walk(os.path.join(_UpperCAmelCase			,  os.listdir(_UpperCAmelCase     )[0]			,  "snapshots"     )     )]
    lowerCamelCase_      =  [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      snake_case    (    unittest.TestCase			):





  """simple docstring"""




  def   snake_case		(    self     ):







   """simple docstring"""
   lowerCamelCase_      =  FlaxStableDiffusionPipeline.from_pretrained(
       "hf-internal-testing/tiny-stable-diffusion-pipe"			,  safety_checker=_UpperCAmelCase     )

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

   lowerCamelCase_      =  jax.random.PRNGKey(0     )
   lowerCamelCase_      =  4

   lowerCamelCase_      =  jax.device_count()
   lowerCamelCase_      =  num_samples * [prompt]
   lowerCamelCase_      =  pipeline.prepare_inputs(_UpperCAmelCase     )

   # shard inputs and rng
   lowerCamelCase_      =  replicate(_UpperCAmelCase     )
   lowerCamelCase_      =  jax.random.split(_UpperCAmelCase			,  _UpperCAmelCase     )
   lowerCamelCase_      =  shard(_UpperCAmelCase     )

   lowerCamelCase_      =  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.1_514_745     ) < 1e-3
    assert np.abs(np.abs(_UpperCAmelCase			,  dtype=np.floataa     ).sum() - 49_947.875     ) < 5e-1

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






  def   snake_case		(    self     ):







   """simple docstring"""
   lowerCamelCase_      =  FlaxStableDiffusionPipeline.from_pretrained(
       "CompVis/stable-diffusion-v1-4"			,  revision="flax"			,  safety_checker=_UpperCAmelCase     )

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

   lowerCamelCase_      =  jax.random.PRNGKey(0     )
   lowerCamelCase_      =  50

   lowerCamelCase_      =  jax.device_count()
   lowerCamelCase_      =  num_samples * [prompt]
   lowerCamelCase_      =  pipeline.prepare_inputs(_UpperCAmelCase     )

   # shard inputs and rng
   lowerCamelCase_      =  replicate(_UpperCAmelCase     )
   lowerCamelCase_      =  jax.random.split(_UpperCAmelCase			,  _UpperCAmelCase     )
   lowerCamelCase_      =  shard(_UpperCAmelCase     )

   lowerCamelCase_      =  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.05_652_401)     ) < 1e-3
    assert np.abs((np.abs(_UpperCAmelCase			,  dtype=np.floataa     ).sum() - 2_383_808.2)     ) < 5e-1






  def   snake_case		(    self     ):







   """simple docstring"""
   lowerCamelCase_      =  FlaxStableDiffusionPipeline.from_pretrained(
       "CompVis/stable-diffusion-v1-4"			,  revision="bf16"			,  dtype=jnp.bfloataa			,  safety_checker=_UpperCAmelCase     )

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

   lowerCamelCase_      =  jax.random.PRNGKey(0     )
   lowerCamelCase_      =  50

   lowerCamelCase_      =  jax.device_count()
   lowerCamelCase_      =  num_samples * [prompt]
   lowerCamelCase_      =  pipeline.prepare_inputs(_UpperCAmelCase     )

   # shard inputs and rng
   lowerCamelCase_      =  replicate(_UpperCAmelCase     )
   lowerCamelCase_      =  jax.random.split(_UpperCAmelCase			,  _UpperCAmelCase     )
   lowerCamelCase_      =  shard(_UpperCAmelCase     )

   lowerCamelCase_      =  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.04_003_906)     ) < 1e-3
    assert np.abs((np.abs(_UpperCAmelCase			,  dtype=np.floataa     ).sum() - 2_373_516.75)     ) < 5e-1






  def   snake_case		(    self     ):







   """simple docstring"""
   lowerCamelCase_      =  FlaxStableDiffusionPipeline.from_pretrained(
       "CompVis/stable-diffusion-v1-4"			,  revision="bf16"			,  dtype=jnp.bfloataa     )

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

   lowerCamelCase_      =  jax.random.PRNGKey(0     )
   lowerCamelCase_      =  50

   lowerCamelCase_      =  jax.device_count()
   lowerCamelCase_      =  num_samples * [prompt]
   lowerCamelCase_      =  pipeline.prepare_inputs(_UpperCAmelCase     )

   # shard inputs and rng
   lowerCamelCase_      =  replicate(_UpperCAmelCase     )
   lowerCamelCase_      =  jax.random.split(_UpperCAmelCase			,  _UpperCAmelCase     )
   lowerCamelCase_      =  shard(_UpperCAmelCase     )

   lowerCamelCase_      =  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.04_003_906)     ) < 1e-3
    assert np.abs((np.abs(_UpperCAmelCase			,  dtype=np.floataa     ).sum() - 2_373_516.75)     ) < 5e-1






  def   snake_case		(    self     ):







   """simple docstring"""
   lowerCamelCase_      =  FlaxDDIMScheduler(
       beta_start=0.00_085			,  beta_end=0.012			,  beta_schedule="scaled_linear"			,  set_alpha_to_one=_UpperCAmelCase			,  steps_offset=1			,  )

   lowerCamelCase_      =  FlaxStableDiffusionPipeline.from_pretrained(
       "CompVis/stable-diffusion-v1-4"			,  revision="bf16"			,  dtype=jnp.bfloataa			,  scheduler=_UpperCAmelCase			,  safety_checker=_UpperCAmelCase			,  )
   lowerCamelCase_      =  scheduler.create_state()

   lowerCamelCase_      =  scheduler_state

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

   lowerCamelCase_      =  jax.random.PRNGKey(0     )
   lowerCamelCase_      =  50

   lowerCamelCase_      =  jax.device_count()
   lowerCamelCase_      =  num_samples * [prompt]
   lowerCamelCase_      =  pipeline.prepare_inputs(_UpperCAmelCase     )

   # shard inputs and rng
   lowerCamelCase_      =  replicate(_UpperCAmelCase     )
   lowerCamelCase_      =  jax.random.split(_UpperCAmelCase			,  _UpperCAmelCase     )
   lowerCamelCase_      =  shard(_UpperCAmelCase     )

   lowerCamelCase_      =  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.045_043_945)     ) < 1e-3
    assert np.abs((np.abs(_UpperCAmelCase			,  dtype=np.floataa     ).sum() - 2_347_693.5)     ) < 5e-1







  def   snake_case		(    self     ):







   """simple docstring"""
   lowerCamelCase_      =  (
       'A cinematic film still of Morgan Freeman starring as Jimi Hendrix, portrait, 40mm lens, shallow depth of'
       ' field, close up, split lighting, cinematic'
   )

   lowerCamelCase_      =  jax.device_count()
   lowerCamelCase_      =  num_samples * [prompt]
   lowerCamelCase_      =  jax.random.split(jax.random.PRNGKey(0     )			,  _UpperCAmelCase     )

   lowerCamelCase_      =  FlaxStableDiffusionPipeline.from_pretrained(
       "CompVis/stable-diffusion-v1-4"			,  revision="bf16"			,  dtype=jnp.bfloataa			,  safety_checker=_UpperCAmelCase			,  )

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

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

   lowerCamelCase_      =  replicate(_UpperCAmelCase     )
   lowerCamelCase_      =  pipeline.prepare_inputs(_UpperCAmelCase     )
   lowerCamelCase_      =  shard(_UpperCAmelCase     )
   lowerCamelCase_      =  pipeline(_UpperCAmelCase			,  _UpperCAmelCase			,  _UpperCAmelCase			,  jit=_UpperCAmelCase     ).images
   assert images_eff.shape == (num_samples, 1, 512, 512, 3)
   lowerCamelCase_      =  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










'''simple docstring'''




import shutil
import tempfile
import unittest

import numpy as np
import pytest

from transformers.testing_utils import require_vision
from transformers.utils import is_vision_available


if is_vision_available():
  from PIL import Image

  from transformers import AutoProcessor, BertTokenizer, BlipImageProcessor, BlipProcessor, PreTrainedTokenizerFast


@require_vision
class  SCREAMING_SNAKE_CASE    (unittest.TestCase							):


     def 							SCREAMING_SNAKE_CASE  (						self):

          '''simple docstring'''






          __A						: Dict      =       tempfile.mkdtemp()

          __A						: Optional[int]      =       BlipImageProcessor()
          __A						: List[str]      =       BertTokenizer.from_pretrained('hf-internal-testing/tiny-random-BertModel')

          __A						: Any      =       BlipProcessor(_UpperCAmelCase						,	_UpperCAmelCase)

          processor.save_pretrained(self.tmpdirname)


     def 							SCREAMING_SNAKE_CASE  (						self						,	**_UpperCAmelCase):

          '''simple docstring'''






          return AutoProcessor.from_pretrained(self.tmpdirname						,	**_UpperCAmelCase).tokenizer


     def 							SCREAMING_SNAKE_CASE  (						self						,	**_UpperCAmelCase):

          '''simple docstring'''






          return AutoProcessor.from_pretrained(self.tmpdirname						,	**_UpperCAmelCase).image_processor


     def 							SCREAMING_SNAKE_CASE  (						self):

          '''simple docstring'''






          shutil.rmtree(self.tmpdirname)


     def 							SCREAMING_SNAKE_CASE  (						self):

          '''simple docstring'''






          __A						: int      =       [np.random.randint(255						,	size=(3, 30, 400)						,	dtype=np.uinta)]

          __A						: Union[str, Any]      =       [Image.fromarray(np.moveaxis(_UpperCAmelCase						,	0						,	-1)) for x in image_inputs]

          return image_inputs


     def 							SCREAMING_SNAKE_CASE  (						self):

          '''simple docstring'''






          __A						: int      =       BlipProcessor(tokenizer=self.get_tokenizer()						,	image_processor=self.get_image_processor())
          processor.save_pretrained(self.tmpdirname)

          __A						: int      =       self.get_tokenizer(bos_token='(BOS)'						,	eos_token='(EOS)')
          __A						: List[Any]      =       self.get_image_processor(do_normalize=_UpperCAmelCase						,	padding_value=1.0)

          __A						: Tuple      =       BlipProcessor.from_pretrained(
              self.tmpdirname						,	bos_token='(BOS)'						,	eos_token='(EOS)'						,	do_normalize=_UpperCAmelCase						,	padding_value=1.0)

          self.assertEqual(processor.tokenizer.get_vocab()						,	tokenizer_add_kwargs.get_vocab())
          self.assertIsInstance(processor.tokenizer						,	_UpperCAmelCase)

          self.assertEqual(processor.image_processor.to_json_string()						,	image_processor_add_kwargs.to_json_string())
          self.assertIsInstance(processor.image_processor						,	_UpperCAmelCase)


     def 							SCREAMING_SNAKE_CASE  (						self):

          '''simple docstring'''






          __A						: List[str]      =       self.get_image_processor()
          __A						: str      =       self.get_tokenizer()

          __A						: int      =       BlipProcessor(tokenizer=_UpperCAmelCase						,	image_processor=_UpperCAmelCase)

          __A						: List[str]      =       self.prepare_image_inputs()

          __A						: str      =       image_processor(_UpperCAmelCase						,	return_tensors='np')
          __A						: str      =       processor(images=_UpperCAmelCase						,	return_tensors='np')

          for key in input_feat_extract.keys():
               self.assertAlmostEqual(input_feat_extract[key].sum()						,	input_processor[key].sum()						,	delta=1e-2)


     def 							SCREAMING_SNAKE_CASE  (						self):

          '''simple docstring'''






          __A						: Any      =       self.get_image_processor()
          __A						: Tuple      =       self.get_tokenizer()

          __A						: Union[str, Any]      =       BlipProcessor(tokenizer=_UpperCAmelCase						,	image_processor=_UpperCAmelCase)

          __A						: str      =       'lower newer'

          __A						: Dict      =       processor(text=_UpperCAmelCase)

          __A						: Tuple      =       tokenizer(_UpperCAmelCase						,	return_token_type_ids=_UpperCAmelCase)

          for key in encoded_tok.keys():
               self.assertListEqual(encoded_tok[key]						,	encoded_processor[key])


     def 							SCREAMING_SNAKE_CASE  (						self):

          '''simple docstring'''






          __A						: List[Any]      =       self.get_image_processor()
          __A						: Optional[Any]      =       self.get_tokenizer()

          __A						: Optional[Any]      =       BlipProcessor(tokenizer=_UpperCAmelCase						,	image_processor=_UpperCAmelCase)

          __A						: Dict      =       'lower newer'
          __A						: int      =       self.prepare_image_inputs()

          __A						: List[Any]      =       processor(text=_UpperCAmelCase						,	images=_UpperCAmelCase)

          self.assertListEqual(list(inputs.keys())						,	['pixel_values', 'input_ids', 'attention_mask'])

          # test if it raises when no input is passed
          with pytest.raises(_UpperCAmelCase):
               processor()


     def 							SCREAMING_SNAKE_CASE  (						self):

          '''simple docstring'''






          __A						: Union[str, Any]      =       self.get_image_processor()
          __A						: Optional[Any]      =       self.get_tokenizer()

          __A						: Union[str, Any]      =       BlipProcessor(tokenizer=_UpperCAmelCase						,	image_processor=_UpperCAmelCase)

          __A						: Tuple      =       [[1, 4, 5, 8, 1, 0, 8], [3, 4, 3, 1, 1, 8, 9]]

          __A						: List[str]      =       processor.batch_decode(_UpperCAmelCase)
          __A						: List[str]      =       tokenizer.batch_decode(_UpperCAmelCase)

          self.assertListEqual(_UpperCAmelCase						,	_UpperCAmelCase)


     def 							SCREAMING_SNAKE_CASE  (						self):

          '''simple docstring'''






          __A						: Any      =       self.get_image_processor()
          __A						: Union[str, Any]      =       self.get_tokenizer()

          __A						: Optional[Any]      =       BlipProcessor(tokenizer=_UpperCAmelCase						,	image_processor=_UpperCAmelCase)

          __A						: str      =       'lower newer'
          __A						: str      =       self.prepare_image_inputs()

          __A						: Optional[int]      =       processor(text=_UpperCAmelCase						,	images=_UpperCAmelCase)

          # For now the processor supports only ['pixel_values', 'input_ids', 'attention_mask']
          self.assertListEqual(list(inputs.keys())						,	['pixel_values', 'input_ids', 'attention_mask'])