





import inspect
import unittest

from transformers import ViTHybridConfig
from transformers.testing_utils import require_accelerate, require_torch, require_vision, slow, torch_device
from transformers.utils import cached_property, is_torch_available, is_vision_available

from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, _config_zero_init, floats_tensor, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
     import torch
     from torch import nn

     from transformers import ViTHybridForImageClassification, ViTHybridImageProcessor, ViTHybridModel
     from transformers.models.vit_hybrid.modeling_vit_hybrid import VIT_HYBRID_PRETRAINED_MODEL_ARCHIVE_LIST


if is_vision_available():
     from PIL import Image

class      A       :



      def __init__(	self:						str		,  _lowerCAmelCase:						Dict		,  _lowerCAmelCase:						List[str]=13		,  _lowerCAmelCase:						List[Any]=64		,  _lowerCAmelCase:						Dict=2		,  _lowerCAmelCase:						List[str]=3		,  _lowerCAmelCase:						List[Any]=True		,  _lowerCAmelCase:						List[Any]=True		,  _lowerCAmelCase:						str=32		,  _lowerCAmelCase:						Optional[Any]=5		,  _lowerCAmelCase:						str=4		,  _lowerCAmelCase:						Any=37		,  _lowerCAmelCase:						Optional[int]="gelu"		,  _lowerCAmelCase:						List[str]=0.1		,  _lowerCAmelCase:						int=0.1		,  _lowerCAmelCase:						List[Any]=10		,  _lowerCAmelCase:						Optional[int]=0.02		,  _lowerCAmelCase:						Optional[int]=[1, 16, 4, 4]		,  _lowerCAmelCase:						int=None		,  )						->    List[Any]:
         '''simple docstring'''

         UpperCAmelCase_											=parent
         UpperCAmelCase_											=batch_size
         UpperCAmelCase_											=image_size
         UpperCAmelCase_											=patch_size
         UpperCAmelCase_											=num_channels
         UpperCAmelCase_											=is_training
         UpperCAmelCase_											=use_labels
         UpperCAmelCase_											=hidden_size
         UpperCAmelCase_											=num_hidden_layers
         UpperCAmelCase_											=num_attention_heads
         UpperCAmelCase_											=intermediate_size
         UpperCAmelCase_											=hidden_act
         UpperCAmelCase_											=hidden_dropout_prob
         UpperCAmelCase_											=attention_probs_dropout_prob
         UpperCAmelCase_											=type_sequence_label_size
         UpperCAmelCase_											=initializer_range
         UpperCAmelCase_											=scope
         UpperCAmelCase_											=backbone_featmap_shape

         # in ViT hybrid, the seq length equals the number of patches + 1 (we add 1 for the [CLS] token)
         # the number of patches is based on the feature map of the backbone, which by default uses an output stride
         # of 32, which means that the feature map has a spatial resolution of 1/32 of the input image size
         UpperCAmelCase_											=(self.image_size // 32) ** 2
         UpperCAmelCase_											=num_patches + 1



      def 	lowerCAmelCase__    (	self:						Tuple       )						->    Tuple:
         '''simple docstring'''

         UpperCAmelCase_											=floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]       )

         UpperCAmelCase_											=None
         if self.use_labels:
            UpperCAmelCase_											=ids_tensor([self.batch_size]		,  self.type_sequence_label_size       )

         UpperCAmelCase_											=self.get_config()

         return config, pixel_values, labels



      def 	lowerCAmelCase__    (	self:						Any       )						->    List[Any]:
         '''simple docstring'''

         UpperCAmelCase_											={
             "global_padding": "same",
             "layer_type": "bottleneck",
             "depths": [3, 4, 9],
             "out_features": ["stage1", "stage2", "stage3"],
             "embedding_dynamic_padding": True,
             "hidden_sizes": [4, 8, 16, 32],
             "num_groups": 2,
         }

         return ViTHybridConfig(
             image_size=self.image_size		,  patch_size=self.patch_size		,  num_channels=self.num_channels		,  hidden_size=self.hidden_size		,  num_hidden_layers=self.num_hidden_layers		,  num_attention_heads=self.num_attention_heads		,  intermediate_size=self.intermediate_size		,  hidden_act=self.hidden_act		,  hidden_dropout_prob=self.hidden_dropout_prob		,  attention_probs_dropout_prob=self.attention_probs_dropout_prob		,  is_decoder=_lowerCAmelCase		,  initializer_range=self.initializer_range		,  backbone_featmap_shape=self.backbone_featmap_shape		,  backbone_config=_lowerCAmelCase		,  )



      def 	lowerCAmelCase__    (	self:						Optional[Any]		,  _lowerCAmelCase:						List[str]		,  _lowerCAmelCase:						str		,  _lowerCAmelCase:						Optional[Any]       )						->    List[Any]:
         '''simple docstring'''

         UpperCAmelCase_											=ViTHybridModel(config=_lowerCAmelCase       )
         model.to(_lowerCAmelCase       )
         model.eval()
         UpperCAmelCase_											=model(_lowerCAmelCase       )
         self.parent.assertEqual(result.last_hidden_state.shape		,  (self.batch_size, self.seq_length, self.hidden_size)       )



      def 	lowerCAmelCase__    (	self:						Dict		,  _lowerCAmelCase:						int		,  _lowerCAmelCase:						Union[str, Any]		,  _lowerCAmelCase:						List[str]       )						->    Dict:
         '''simple docstring'''

         UpperCAmelCase_											=self.type_sequence_label_size
         UpperCAmelCase_											=ViTHybridForImageClassification(_lowerCAmelCase       )
         model.to(_lowerCAmelCase       )
         model.eval()
         UpperCAmelCase_											=model(_lowerCAmelCase		,  labels=_lowerCAmelCase       )
         self.parent.assertEqual(result.logits.shape		,  (self.batch_size, self.type_sequence_label_size)       )






      def 	lowerCAmelCase__    (	self:						List[Any]       )						->    str:
         '''simple docstring'''

         UpperCAmelCase_											=self.prepare_config_and_inputs()
         UpperCAmelCase_						,       UpperCAmelCase_						,       UpperCAmelCase_											=config_and_inputs
         UpperCAmelCase_											={"pixel_values": pixel_values}
         return config, inputs_dict

@require_torch
class      A       (	__lowercase		,      __lowercase		,      unittest.TestCase						):
      _snake_case							     =(ViTHybridModel, ViTHybridForImageClassification) if is_torch_available() else ()
      _snake_case							     =(
          {'''feature-extraction''': ViTHybridModel, '''image-classification''': ViTHybridForImageClassification}
          if is_torch_available()
          else {}
      )
      _snake_case							     =False
      _snake_case							     =False
      _snake_case							     =False



      def 	lowerCAmelCase__    (	self:						Optional[Any]       )						->    Optional[Any]:
         '''simple docstring'''

         UpperCAmelCase_											=ViTHybridModelTester(self       )
         UpperCAmelCase_											=ConfigTester(self		,  config_class=_lowerCAmelCase		,  has_text_modality=_lowerCAmelCase		,  hidden_size=37       )



      def 	lowerCAmelCase__    (	self:						Tuple       )						->    Dict:
         '''simple docstring'''

         self.config_tester.run_common_tests()



      @unittest.skip(reason="ViT does not use inputs_embeds"       )
      def 	lowerCAmelCase__    (	self:						Optional[int]       )						->    str:
         '''simple docstring'''

         pass



      def 	lowerCAmelCase__    (	self:						List[str]       )						->    List[Any]:
         '''simple docstring'''

         UpperCAmelCase_						,       UpperCAmelCase_											=self.model_tester.prepare_config_and_inputs_for_common()

         for model_class in self.all_model_classes:
            UpperCAmelCase_											=model_class(_lowerCAmelCase       )
            self.assertIsInstance(model.get_input_embeddings()		,  (nn.Module)       )
            UpperCAmelCase_											=model.get_output_embeddings()
            self.assertTrue(x is None or isinstance(_lowerCAmelCase		,  nn.Linear       )       )



      def 	lowerCAmelCase__    (	self:						Optional[Any]       )						->    Any:
         '''simple docstring'''

         UpperCAmelCase_						,       UpperCAmelCase_											=self.model_tester.prepare_config_and_inputs_for_common()

         for model_class in self.all_model_classes:
            UpperCAmelCase_											=model_class(_lowerCAmelCase       )
            UpperCAmelCase_											=inspect.signature(model.forward       )
            # signature.parameters is an OrderedDict => so arg_names order is deterministic
            UpperCAmelCase_											=[*signature.parameters.keys()]

            UpperCAmelCase_											=["pixel_values"]
            self.assertListEqual(arg_names[:1]		,  _lowerCAmelCase       )



      def 	lowerCAmelCase__    (	self:						Optional[Any]       )						->    List[str]:
         '''simple docstring'''

         UpperCAmelCase_											=self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_model(*_lowerCAmelCase       )



      def 	lowerCAmelCase__    (	self:						Optional[int]       )						->    Optional[int]:
         '''simple docstring'''

         UpperCAmelCase_											=self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_for_image_classification(*_lowerCAmelCase       )



      def 	lowerCAmelCase__    (	self:						Optional[int]       )						->    Optional[int]:
         '''simple docstring'''

         UpperCAmelCase_						,       UpperCAmelCase_											=self.model_tester.prepare_config_and_inputs_for_common()

         UpperCAmelCase_											=_config_zero_init(_lowerCAmelCase       )
         for model_class in self.all_model_classes:
            UpperCAmelCase_											=model_class(config=_lowerCAmelCase       )
            # Skip the check for the backbone
            for name, module in model.named_modules():
               if module.__class__.__name__ == "ViTHybridPatchEmbeddings":
                  UpperCAmelCase_											=[F'{name}.{key}' for key in module.state_dict().keys()]
                  break

            for name, param in model.named_parameters():
               if param.requires_grad:
                  if name in backbone_params:
                     continue
                  self.assertIn(
                      ((param.data.mean() * 1e9).round() / 1e9).item()		,  [0.0, 1.0]		,  msg=F'Parameter {name} of model {model_class} seems not properly initialized'		,  )






      @slow
      def 	lowerCAmelCase__    (	self:						List[str]       )						->    List[str]:
         '''simple docstring'''

         for model_name in VIT_HYBRID_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
            UpperCAmelCase_											=ViTHybridModel.from_pretrained(_lowerCAmelCase       )
            self.assertIsNotNone(_lowerCAmelCase       )




def     a__ (					):






   '''simple docstring'''






   UpperCAmelCase_											=Image.open("./tests/fixtures/tests_samples/COCO/000000039769.png"  )
   return image





@require_torch
@require_vision
class      A       (	unittest.TestCase						):



      @cached_property
      def 	lowerCAmelCase__    (	self:						Tuple       )						->    Optional[Any]:
         '''simple docstring'''

         return (
             ViTHybridImageProcessor.from_pretrained(VIT_HYBRID_PRETRAINED_MODEL_ARCHIVE_LIST[0]       )
             if is_vision_available()
             else None
         )



      @slow
      def 	lowerCAmelCase__    (	self:						List[Any]       )						->    List[str]:
         '''simple docstring'''

         UpperCAmelCase_											=ViTHybridForImageClassification.from_pretrained(VIT_HYBRID_PRETRAINED_MODEL_ARCHIVE_LIST[0]       ).to(
             _lowerCAmelCase       )

         UpperCAmelCase_											=self.default_image_processor
         UpperCAmelCase_											=prepare_img()
         UpperCAmelCase_											=image_processor(images=_lowerCAmelCase		,  return_tensors="pt"       ).to(_lowerCAmelCase       )

         # forward pass
         with torch.no_grad():
            UpperCAmelCase_											=model(**_lowerCAmelCase       )

         # verify the logits
         UpperCAmelCase_											=torch.Size((1, 1000)       )
         self.assertEqual(outputs.logits.shape		,  _lowerCAmelCase       )

         UpperCAmelCase_											=torch.tensor([-1.90_90, -0.49_93, -0.23_89]       ).to(_lowerCAmelCase       )

         self.assertTrue(torch.allclose(outputs.logits[0, :3]		,  _lowerCAmelCase		,  atol=1e-4       )       )






      @slow
      @require_accelerate
      def 	lowerCAmelCase__    (	self:						Optional[Any]       )						->    Any:
         '''simple docstring'''

         UpperCAmelCase_											=ViTHybridImageProcessor.from_pretrained("google/vit-hybrid-base-bit-384"       )
         UpperCAmelCase_											=ViTHybridForImageClassification.from_pretrained("google/vit-hybrid-base-bit-384"		,  device_map="auto"       )

         UpperCAmelCase_											=prepare_img()

         UpperCAmelCase_											=image_processor(images=_lowerCAmelCase		,  return_tensors="pt"       )
         UpperCAmelCase_											=model(**_lowerCAmelCase       )
         UpperCAmelCase_											=outputs.logits
         # model predicts one of the 1000 ImageNet classes
         UpperCAmelCase_											=logits.argmax(-1       ).item()

         self.assertTrue(model.config.idalabel[predicted_class_idx]		,  "tabby, tabby cat"       )











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



      def 	lowerCAmelCase__    (	self:						str       )						->    Union[str, Any]:
         '''simple docstring'''

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

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



      def 	lowerCAmelCase__    (	self:						Union[str, Any]       )						->    Tuple:
         '''simple docstring'''

         UpperCAmelCase_						,       UpperCAmelCase_											=FlaxStableDiffusionPipeline.from_pretrained(
             "hf-internal-testing/tiny-stable-diffusion-pipe"		,  safety_checker=_lowerCAmelCase       )

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

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

         UpperCAmelCase_											=jax.device_count()
         UpperCAmelCase_											=num_samples * [prompt]
         UpperCAmelCase_											=pipeline.prepare_inputs(_lowerCAmelCase       )

         # shard inputs and rng
         UpperCAmelCase_											=replicate(_lowerCAmelCase       )
         UpperCAmelCase_											=jax.random.split(_lowerCAmelCase		,  _lowerCAmelCase       )
         UpperCAmelCase_											=shard(_lowerCAmelCase       )

         UpperCAmelCase_											=pipeline(_lowerCAmelCase		,  _lowerCAmelCase		,  _lowerCAmelCase		,  _lowerCAmelCase		,  jit=_lowerCAmelCase       ).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_51_47_45       ) < 1e-3
            assert np.abs(np.abs(_lowerCAmelCase		,  dtype=np.floataa       ).sum() - 4_99_47.8_75       ) < 5e-1

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



      def 	lowerCAmelCase__    (	self:						List[Any]       )						->    Optional[Any]:
         '''simple docstring'''

         UpperCAmelCase_						,       UpperCAmelCase_											=FlaxStableDiffusionPipeline.from_pretrained(
             "CompVis/stable-diffusion-v1-4"		,  revision="flax"		,  safety_checker=_lowerCAmelCase       )

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

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

         UpperCAmelCase_											=jax.device_count()
         UpperCAmelCase_											=num_samples * [prompt]
         UpperCAmelCase_											=pipeline.prepare_inputs(_lowerCAmelCase       )

         # shard inputs and rng
         UpperCAmelCase_											=replicate(_lowerCAmelCase       )
         UpperCAmelCase_											=jax.random.split(_lowerCAmelCase		,  _lowerCAmelCase       )
         UpperCAmelCase_											=shard(_lowerCAmelCase       )

         UpperCAmelCase_											=pipeline(_lowerCAmelCase		,  _lowerCAmelCase		,  _lowerCAmelCase		,  _lowerCAmelCase		,  jit=_lowerCAmelCase       ).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_65_24_01)       ) < 1e-3
            assert np.abs((np.abs(_lowerCAmelCase		,  dtype=np.floataa       ).sum() - 2_38_38_08.2)       ) < 5e-1



      def 	lowerCAmelCase__    (	self:						int       )						->    Union[str, Any]:
         '''simple docstring'''

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

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

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

         UpperCAmelCase_											=jax.device_count()
         UpperCAmelCase_											=num_samples * [prompt]
         UpperCAmelCase_											=pipeline.prepare_inputs(_lowerCAmelCase       )

         # shard inputs and rng
         UpperCAmelCase_											=replicate(_lowerCAmelCase       )
         UpperCAmelCase_											=jax.random.split(_lowerCAmelCase		,  _lowerCAmelCase       )
         UpperCAmelCase_											=shard(_lowerCAmelCase       )

         UpperCAmelCase_											=pipeline(_lowerCAmelCase		,  _lowerCAmelCase		,  _lowerCAmelCase		,  _lowerCAmelCase		,  jit=_lowerCAmelCase       ).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_00_39_06)       ) < 1e-3
            assert np.abs((np.abs(_lowerCAmelCase		,  dtype=np.floataa       ).sum() - 2_37_35_16.75)       ) < 5e-1



      def 	lowerCAmelCase__    (	self:						Optional[int]       )						->    Tuple:
         '''simple docstring'''

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

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

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

         UpperCAmelCase_											=jax.device_count()
         UpperCAmelCase_											=num_samples * [prompt]
         UpperCAmelCase_											=pipeline.prepare_inputs(_lowerCAmelCase       )

         # shard inputs and rng
         UpperCAmelCase_											=replicate(_lowerCAmelCase       )
         UpperCAmelCase_											=jax.random.split(_lowerCAmelCase		,  _lowerCAmelCase       )
         UpperCAmelCase_											=shard(_lowerCAmelCase       )

         UpperCAmelCase_											=pipeline(_lowerCAmelCase		,  _lowerCAmelCase		,  _lowerCAmelCase		,  _lowerCAmelCase		,  jit=_lowerCAmelCase       ).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_00_39_06)       ) < 1e-3
            assert np.abs((np.abs(_lowerCAmelCase		,  dtype=np.floataa       ).sum() - 2_37_35_16.75)       ) < 5e-1



      def 	lowerCAmelCase__    (	self:						Optional[Any]       )						->    str:
         '''simple docstring'''

         UpperCAmelCase_											=FlaxDDIMScheduler(
             beta_start=0.0_00_85		,  beta_end=0.0_12		,  beta_schedule="scaled_linear"		,  set_alpha_to_one=_lowerCAmelCase		,  steps_offset=1		,  )

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

         UpperCAmelCase_											=scheduler_state

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

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

         UpperCAmelCase_											=jax.device_count()
         UpperCAmelCase_											=num_samples * [prompt]
         UpperCAmelCase_											=pipeline.prepare_inputs(_lowerCAmelCase       )

         # shard inputs and rng
         UpperCAmelCase_											=replicate(_lowerCAmelCase       )
         UpperCAmelCase_											=jax.random.split(_lowerCAmelCase		,  _lowerCAmelCase       )
         UpperCAmelCase_											=shard(_lowerCAmelCase       )

         UpperCAmelCase_											=pipeline(_lowerCAmelCase		,  _lowerCAmelCase		,  _lowerCAmelCase		,  _lowerCAmelCase		,  jit=_lowerCAmelCase       ).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.0_45_04_39_45)       ) < 1e-3
            assert np.abs((np.abs(_lowerCAmelCase		,  dtype=np.floataa       ).sum() - 2_34_76_93.5)       ) < 5e-1






      def 	lowerCAmelCase__    (	self:						Optional[Any]       )						->    int:
         '''simple docstring'''

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

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

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

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

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

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




