


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





      def        _lowercase				(		self						:					List[Any]       )    ->   Dict:






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





      def        _lowercase				(		self						:					List[str]       )    ->   Union[str, Any]:






         """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, 6_4, 6_4, 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(__lowerCamelCase					,       dtype=np.floataa       ).sum() - 4_9_9_4_7.8_7_5       ) < 5e-1

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





      def        _lowercase				(		self						:					Tuple       )    ->   Dict:






         """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							=     5_0

         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, 5_1_2, 5_1_2, 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(__lowerCamelCase					,       dtype=np.floataa       ).sum() - 2_3_8_3_8_0_8.2)       ) < 5e-1





      def        _lowercase				(		self						:					int       )    ->   List[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							=     5_0

         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, 5_1_2, 5_1_2, 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(__lowerCamelCase					,       dtype=np.floataa       ).sum() - 2_3_7_3_5_1_6.7_5)       ) < 5e-1





      def        _lowercase				(		self						:					Optional[Any]       )    ->   List[str]:






         """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							=     5_0

         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, 5_1_2, 5_1_2, 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(__lowerCamelCase					,       dtype=np.floataa       ).sum() - 2_3_7_3_5_1_6.7_5)       ) < 5e-1





      def        _lowercase				(		self						:					int       )    ->   List[str]:






         """simple docstring"""







         UpperCAmelCase							=     FlaxDDIMScheduler(
             beta_start=0.00_085					,       beta_end=0.012					,       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							=     5_0

         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, 5_1_2, 5_1_2, 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(__lowerCamelCase					,       dtype=np.floataa       ).sum() - 2_3_4_7_6_9_3.5)       ) < 5e-1





      def        _lowercase				(		self						:					Dict       )    ->   List[str]:






         """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, 5_1_2, 5_1_2, 3)
         UpperCAmelCase							=     images[2, 0, 2_5_6, 1_0:1_7, 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, 5_1_2, 5_1_2, 3)
         UpperCAmelCase							=     images[2, 0, 2_5_6, 1_0:1_7, 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







import math

import qiskit
def   _UpperCamelCase				(						lowerCAmelCase_ = 1  ,						lowerCAmelCase_ = 1  ,						lowerCAmelCase_ = 1   )						->qiskit.result.counts.Counts:
   if (
       isinstance(lowerCAmelCase_  ,						lowerCAmelCase_   )
       or isinstance(lowerCAmelCase_  ,						lowerCAmelCase_   )
       or isinstance(lowerCAmelCase_  ,						lowerCAmelCase_   )
   ):
      raise TypeError("""inputs must be integers."""   )

   if (input_a < 0) or (input_a < 0) or (carry_in < 0):
      raise ValueError("""inputs must be positive."""   )

   if (
       (math.floor(lowerCAmelCase_   ) != input_a)
       or (math.floor(lowerCAmelCase_   ) != input_a)
       or (math.floor(lowerCAmelCase_   ) != carry_in)
   ):
      raise ValueError("""inputs must be exact integers."""   )

   if (input_a > 2) or (input_a > 2) or (carry_in > 2):
      raise ValueError("""inputs must be less or equal to 2."""   )

   # build registers
   UpperCAmelCase							=     qiskit.QuantumRegister(4  ,						"""qr"""   )
   UpperCAmelCase							=     qiskit.ClassicalRegister(2  ,						"""cr"""   )
   # list the entries
   UpperCAmelCase							=     [input_a, input_a, carry_in]

   UpperCAmelCase							=     qiskit.QuantumCircuit(lowerCAmelCase_  ,						lowerCAmelCase_   )

   for i in range(0  ,						3   ):
      if entry[i] == 2:
         quantum_circuit.h(lowerCAmelCase_   )  # for hadamard entries
      elif entry[i] == 1:
         quantum_circuit.x(lowerCAmelCase_   )  # for 1 entries
      elif entry[i] == 0:
         quantum_circuit.i(lowerCAmelCase_   )  # for 0 entries

    # build the circuit
   quantum_circuit.ccx(0  ,						1  ,						3   )  # ccx = toffoli gate
   quantum_circuit.cx(0  ,						1   )
   quantum_circuit.ccx(1  ,						2  ,						3   )
   quantum_circuit.cx(1  ,						2   )
   quantum_circuit.cx(0  ,						1   )

   quantum_circuit.measure([2, 3]  ,						lowerCAmelCase_   )  # measure the last two qbits

   UpperCAmelCase							=     qiskit.Aer.get_backend("""aer_simulator"""   )
   UpperCAmelCase							=     qiskit.execute(lowerCAmelCase_  ,						lowerCAmelCase_  ,						shots=1_0_0_0   )

   return job.result().get_counts(lowerCAmelCase_   )


if __name__ == "__main__":
 print(F"""Total sum count for state is: {quantum_full_adder(1, 1, 1)}""")



