


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 _lowerCamelCase(		unittest.TestCase     ):
       def     UpperCamelCase  (      self) ->     int:




           """simple docstring"""

           super().tearDown()
           gc.collect()
       def     UpperCamelCase  (      self) ->     str:




           """simple docstring"""

           _lowercase :							Tuple   	=	FlaxControlNetModel.from_pretrained(
               'lllyasviel/sd-controlnet-canny',   from_pt=lowerCamelCase,   dtype=jnp.bfloataa)
           _lowercase :							Any   	=	FlaxStableDiffusionControlNetPipeline.from_pretrained(
               'runwayml/stable-diffusion-v1-5',   controlnet=lowerCamelCase,   from_pt=lowerCamelCase,   dtype=jnp.bfloataa)
           _lowercase :							List[str]   	=	controlnet_params

           _lowercase :							Union[str, Any]   	=	"bird"
           _lowercase :							Optional[int]   	=	jax.device_count()
           _lowercase :							Tuple   	=	pipe.prepare_text_inputs([prompts] * num_samples)

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

           _lowercase :							List[str]   	=	jax.random.PRNGKey(0)
           _lowercase :							str   	=	jax.random.split(lowerCamelCase,   jax.device_count())

           _lowercase :							int   	=	replicate(lowerCamelCase)
           _lowercase :							Any   	=	shard(lowerCamelCase)
           _lowercase :							Any   	=	shard(lowerCamelCase)

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

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

           _lowercase :							Union[str, Any]   	=	jnp.asarray(jax.device_get(image_slice.flatten()))
           _lowercase :							List[Any]   	=	jnp.array(
               [0.1_6_7_9_6_9, 0.1_1_6_6_9_9, 0.0_8_1_5_4_3, 0.1_5_4_2_9_7, 0.1_3_2_8_1_2, 0.1_0_8_8_8_7, 0.1_6_9_9_2_2, 0.1_6_9_9_2_2, 0.2_0_5_0_7_8])
           print(F'''output_slice: {output_slice}''')
           assert jnp.abs(output_slice - expected_slice).max() < 1E-2




       def     UpperCamelCase  (      self) ->     Optional[int]:




           """simple docstring"""

           _lowercase :							List[str]   	=	FlaxControlNetModel.from_pretrained(
               'lllyasviel/sd-controlnet-openpose',   from_pt=lowerCamelCase,   dtype=jnp.bfloataa)
           _lowercase :							Optional[Any]   	=	FlaxStableDiffusionControlNetPipeline.from_pretrained(
               'runwayml/stable-diffusion-v1-5',   controlnet=lowerCamelCase,   from_pt=lowerCamelCase,   dtype=jnp.bfloataa)
           _lowercase :							str   	=	controlnet_params

           _lowercase :							int   	=	"Chef in the kitchen"
           _lowercase :							List[Any]   	=	jax.device_count()
           _lowercase :							Dict   	=	pipe.prepare_text_inputs([prompts] * num_samples)

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

           _lowercase :							List[str]   	=	jax.random.PRNGKey(0)
           _lowercase :							Any   	=	jax.random.split(lowerCamelCase,   jax.device_count())

           _lowercase :							Dict   	=	replicate(lowerCamelCase)
           _lowercase :							Tuple   	=	shard(lowerCamelCase)
           _lowercase :							Optional[int]   	=	shard(lowerCamelCase)

           _lowercase :							Optional[Any]   	=	pipe(
               prompt_ids=lowerCamelCase,   image=lowerCamelCase,   params=lowerCamelCase,   prng_seed=lowerCamelCase,   num_inference_steps=50,   jit=lowerCamelCase,   ).images
           assert images.shape == (jax.device_count(), 1, 7_68, 5_12, 3)

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

           _lowercase :							Tuple   	=	jnp.asarray(jax.device_get(image_slice.flatten()))
           _lowercase :							List[str]   	=	jnp.array(
               [[0.2_7_1_4_8_4, 0.2_6_1_7_1_9, 0.2_7_5_3_9_1, 0.2_7_7_3_4_4, 0.2_7_9_2_9_7, 0.2_9_1_0_1_6, 0.2_9_4_9_2_2, 0.3_0_2_7_3_4, 0.3_0_2_7_3_4]])
           print(F'''output_slice: {output_slice}''')
           assert jnp.abs(output_slice - expected_slice).max() < 1E-2






import subprocess
import sys

from transformers import BertConfig, BertModel, BertTokenizer, pipeline
from transformers.testing_utils import TestCasePlus, require_torch






class _lowerCamelCase(		_a     ):
       @require_torch
       def     UpperCamelCase  (      self) ->     int:




           """simple docstring"""

           _lowercase :							Optional[Any]   	=	'\nfrom transformers import BertConfig, BertModel, BertTokenizer, pipeline\n        '

           _lowercase :							Union[str, Any]   	=	'\nmname = "hf-internal-testing/tiny-random-bert"\nBertConfig.from_pretrained(mname)\nBertModel.from_pretrained(mname)\nBertTokenizer.from_pretrained(mname)\npipe = pipeline(task="fill-mask", model=mname)\nprint("success")\n        '

           _lowercase :							List[str]   	=	'\nimport socket\ndef offline_socket(*args, **kwargs): raise RuntimeError("Offline mode is enabled, we shouldn\'t access internet")\nsocket.socket = offline_socket\n        '

           # Force fetching the files so that we can use the cache
           _lowercase :							Union[str, Any]   	=	'hf-internal-testing/tiny-random-bert'
           BertConfig.from_pretrained(lowerCamelCase)
           BertModel.from_pretrained(lowerCamelCase)
           BertTokenizer.from_pretrained(lowerCamelCase)
           pipeline(task='fill-mask',   model=lowerCamelCase)

           # baseline - just load from_pretrained with normal network
           _lowercase :							Any   	=	[sys.executable, '-c', '\n'.join([load, run, mock])]

           # should succeed
           _lowercase :							List[str]   	=	self.get_env()
           # should succeed as TRANSFORMERS_OFFLINE=1 tells it to use local files
           _lowercase :							Any   	=	'1'
           _lowercase :							List[Any]   	=	subprocess.run(lowerCamelCase,   env=lowerCamelCase,   check=lowerCamelCase,   capture_output=lowerCamelCase)
           self.assertEqual(result.returncode,   0,   result.stderr)
           self.assertIn('success',   result.stdout.decode())
       @require_torch
       def     UpperCamelCase  (      self) ->     List[Any]:




           """simple docstring"""

           _lowercase :							int   	=	'\nfrom transformers import BertConfig, BertModel, BertTokenizer, pipeline\n        '

           _lowercase :							Union[str, Any]   	=	'\nmname = "hf-internal-testing/tiny-random-bert"\nBertConfig.from_pretrained(mname)\nBertModel.from_pretrained(mname)\nBertTokenizer.from_pretrained(mname)\npipe = pipeline(task="fill-mask", model=mname)\nprint("success")\n        '

           _lowercase :							List[Any]   	=	'\nimport socket\ndef offline_socket(*args, **kwargs): raise socket.error("Faking flaky internet")\nsocket.socket = offline_socket\n        '

           # Force fetching the files so that we can use the cache
           _lowercase :							Any   	=	'hf-internal-testing/tiny-random-bert'
           BertConfig.from_pretrained(lowerCamelCase)
           BertModel.from_pretrained(lowerCamelCase)
           BertTokenizer.from_pretrained(lowerCamelCase)
           pipeline(task='fill-mask',   model=lowerCamelCase)

           # baseline - just load from_pretrained with normal network
           _lowercase :							List[str]   	=	[sys.executable, '-c', '\n'.join([load, run, mock])]

           # should succeed
           _lowercase :							Any   	=	self.get_env()
           _lowercase :							str   	=	subprocess.run(lowerCamelCase,   env=lowerCamelCase,   check=lowerCamelCase,   capture_output=lowerCamelCase)
           self.assertEqual(result.returncode,   0,   result.stderr)
           self.assertIn('success',   result.stdout.decode())
       @require_torch
       def     UpperCamelCase  (      self) ->     Union[str, Any]:




           """simple docstring"""

           _lowercase :							Tuple   	=	'\nfrom transformers import BertConfig, BertModel, BertTokenizer\n        '

           _lowercase :							int   	=	'\nmname = "hf-internal-testing/tiny-random-bert-sharded"\nBertConfig.from_pretrained(mname)\nBertModel.from_pretrained(mname)\nprint("success")\n        '

           _lowercase :							List[Any]   	=	'\nimport socket\ndef offline_socket(*args, **kwargs): raise ValueError("Offline mode is enabled")\nsocket.socket = offline_socket\n        '

           # baseline - just load from_pretrained with normal network
           _lowercase :							int   	=	[sys.executable, '-c', '\n'.join([load, run])]

           # should succeed
           _lowercase :							List[str]   	=	self.get_env()
           _lowercase :							Dict   	=	subprocess.run(lowerCamelCase,   env=lowerCamelCase,   check=lowerCamelCase,   capture_output=lowerCamelCase)
           self.assertEqual(result.returncode,   0,   result.stderr)
           self.assertIn('success',   result.stdout.decode())

           # next emulate no network
           _lowercase :							Union[str, Any]   	=	[sys.executable, '-c', '\n'.join([load, mock, run])]

           # Doesn't fail anymore since the model is in the cache due to other tests, so commenting this.
           # env["TRANSFORMERS_OFFLINE"] = "0"
           # result = subprocess.run(cmd, env=env, check=False, capture_output=True)
           # self.assertEqual(result.returncode, 1, result.stderr)

           # should succeed as TRANSFORMERS_OFFLINE=1 tells it to use local files
           _lowercase :							Optional[Any]   	=	'1'
           _lowercase :							Optional[Any]   	=	subprocess.run(lowerCamelCase,   env=lowerCamelCase,   check=lowerCamelCase,   capture_output=lowerCamelCase)
           self.assertEqual(result.returncode,   0,   result.stderr)
           self.assertIn('success',   result.stdout.decode())
       @require_torch
       def     UpperCamelCase  (      self) ->     Optional[Any]:




           """simple docstring"""

           _lowercase :							Optional[Any]   	=	'\nfrom transformers import pipeline\n        '
           _lowercase :							Dict   	=	'\nmname = "hf-internal-testing/tiny-random-bert"\npipe = pipeline(model=mname)\n        '

           _lowercase :							Optional[Any]   	=	'\nimport socket\ndef offline_socket(*args, **kwargs): raise socket.error("Offline mode is enabled")\nsocket.socket = offline_socket\n        '
           _lowercase :							Tuple   	=	self.get_env()
           _lowercase :							Tuple   	=	'1'
           _lowercase :							Union[str, Any]   	=	[sys.executable, '-c', '\n'.join([load, mock, run])]
           _lowercase :							Tuple   	=	subprocess.run(lowerCamelCase,   env=lowerCamelCase,   check=lowerCamelCase,   capture_output=lowerCamelCase)
           self.assertEqual(result.returncode,   1,   result.stderr)
           self.assertIn(
               'You cannot infer task automatically within `pipeline` when using offline mode',   result.stderr.decode().replace('\n',   ''),   )




       @require_torch
       def     UpperCamelCase  (      self) ->     Tuple:




           """simple docstring"""

           _lowercase :							Dict   	=	'\nfrom transformers import AutoModel\n        '
           _lowercase :							int   	=	'\nmname = "hf-internal-testing/test_dynamic_model"\nAutoModel.from_pretrained(mname, trust_remote_code=True)\nprint("success")\n        '

           # baseline - just load from_pretrained with normal network
           _lowercase :							Optional[int]   	=	[sys.executable, '-c', '\n'.join([load, run])]

           # should succeed
           _lowercase :							int   	=	self.get_env()
           _lowercase :							List[str]   	=	subprocess.run(lowerCamelCase,   env=lowerCamelCase,   check=lowerCamelCase,   capture_output=lowerCamelCase)
           self.assertEqual(result.returncode,   0,   result.stderr)
           self.assertIn('success',   result.stdout.decode())

           # should succeed as TRANSFORMERS_OFFLINE=1 tells it to use local files
           _lowercase :							Tuple   	=	'1'
           _lowercase :							Dict   	=	subprocess.run(lowerCamelCase,   env=lowerCamelCase,   check=lowerCamelCase,   capture_output=lowerCamelCase)
           self.assertEqual(result.returncode,   0,   result.stderr)
           self.assertIn('success',   result.stdout.decode())


