




'''simple docstring'''


import contextlib
import importlib
import io
import unittest

import transformers

# Try to import everything from transformers to ensure every object can be loaded.
from transformers import *  # noqa F406
from transformers.testing_utils import DUMMY_UNKNOWN_IDENTIFIER, require_flax, require_tf, require_torch
from transformers.utils import ContextManagers, find_labels, is_flax_available, is_tf_available, is_torch_available


if is_torch_available():
				from transformers import BertForPreTraining, BertForQuestionAnswering, BertForSequenceClassification

if is_tf_available():
				from transformers import TFBertForPreTraining, TFBertForQuestionAnswering, TFBertForSequenceClassification

if is_flax_available():
				from transformers import FlaxBertForPreTraining, FlaxBertForQuestionAnswering, FlaxBertForSequenceClassification


_snake_case : List[str]  				=						DUMMY_UNKNOWN_IDENTIFIER
# An actual model hosted on huggingface.co

_snake_case : Union[str, Any]  				=						'main'
# Default branch name
_snake_case : Dict  				=						'f2c752cfc5c0ab6f4bdec59acea69eefbee381c2'
# One particular commit (not the top of `main`)
_snake_case : Dict  				=						'aaaaaaa'
# This commit does not exist, so we should 404.

_snake_case : str  				=						'd9e9f15bc825e4b2c9249e9578f884bbcb5e3684'
# Sha-1 of config.json on the top of `main`, for checking purposes
_snake_case : int  				=						'4b243c475af8d0a7754e87d7d096c92e5199ec2fe168a2ee7998e3b8e9bcb1d3'
@contextlib.contextmanager
def 						snake_case_       ():





							'''simple docstring'''


							print('''Welcome!''' )
							yield
							print('''Bye!''' )
@contextlib.contextmanager
def 						snake_case_       ():





							'''simple docstring'''


							print('''Bonjour!''' )
							yield
							print('''Au revoir!''' )

class      A       (						unittest.TestCase       ):







							def   __lowerCAmelCase     (    self			:    Dict   )     ->				Union[str, Any]:


														"""simple docstring"""


														assert transformers.__spec__ is not None
														assert importlib.util.find_spec('''transformers'''   ) is not None



class      A       (						unittest.TestCase       ):







							@unittest.mock.patch('''sys.stdout'''				,     new_callable=io.StringIO   )
							def   __lowerCAmelCase     (    self			:    str				,     lowerCAmelCase_			:    str   )     ->				int:


														"""simple docstring"""


														with ContextManagers([]   ):
																					print('''Transformers are awesome!'''   )
														# The print statement adds a new line at the end of the output
														self.assertEqual(mock_stdout.getvalue()				,     '''Transformers are awesome!\n'''   )







							@unittest.mock.patch('''sys.stdout'''				,     new_callable=io.StringIO   )
							def   __lowerCAmelCase     (    self			:    Any				,     lowerCAmelCase_			:    Optional[int]   )     ->				Any:


														"""simple docstring"""


														with ContextManagers([context_en()]   ):
																					print('''Transformers are awesome!'''   )
														# The output should be wrapped with an English welcome and goodbye
														self.assertEqual(mock_stdout.getvalue()				,     '''Welcome!\nTransformers are awesome!\nBye!\n'''   )







							@unittest.mock.patch('''sys.stdout'''				,     new_callable=io.StringIO   )
							def   __lowerCAmelCase     (    self			:    Optional[Any]				,     lowerCAmelCase_			:    List[Any]   )     ->				str:


														"""simple docstring"""


														with ContextManagers([context_fr(), context_en()]   ):
																					print('''Transformers are awesome!'''   )
														# The output should be wrapped with an English and French welcome and goodbye
														self.assertEqual(mock_stdout.getvalue()				,     '''Bonjour!\nWelcome!\nTransformers are awesome!\nBye!\nAu revoir!\n'''   )







							@require_torch
							def   __lowerCAmelCase     (    self			:    int   )     ->				Optional[int]:


														"""simple docstring"""


														self.assertEqual(find_labels(lowerCAmelCase_   )				,     ['''labels''']   )
														self.assertEqual(find_labels(lowerCAmelCase_   )				,     ['''labels''', '''next_sentence_label''']   )
														self.assertEqual(find_labels(lowerCAmelCase_   )				,     ['''start_positions''', '''end_positions''']   )

														class      A       (						_a       ):
																					pass



														self.assertEqual(find_labels(lowerCAmelCase_   )				,     ['''labels''']   )







							@require_tf
							def   __lowerCAmelCase     (    self			:    Tuple   )     ->				Optional[Any]:


														"""simple docstring"""


														self.assertEqual(find_labels(lowerCAmelCase_   )				,     ['''labels''']   )
														self.assertEqual(find_labels(lowerCAmelCase_   )				,     ['''labels''', '''next_sentence_label''']   )
														self.assertEqual(find_labels(lowerCAmelCase_   )				,     ['''start_positions''', '''end_positions''']   )

														class      A       (						_a       ):
																					pass



														self.assertEqual(find_labels(lowerCAmelCase_   )				,     ['''labels''']   )







							@require_flax
							def   __lowerCAmelCase     (    self			:    int   )     ->				Dict:


														"""simple docstring"""


														self.assertEqual(find_labels(lowerCAmelCase_   )				,     []   )
														self.assertEqual(find_labels(lowerCAmelCase_   )				,     []   )
														self.assertEqual(find_labels(lowerCAmelCase_   )				,     []   )

														class      A       (						_a       ):
																					pass



														self.assertEqual(find_labels(lowerCAmelCase_   )				,     []   )




'''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





