





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







					'''simple docstring'''

					def        SCREAMING_SNAKE_CASE							(						self    :  Optional[int]		):






						'''simple docstring'''







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

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







					'''simple docstring'''

					def        SCREAMING_SNAKE_CASE							(						self    :  Dict		):






						'''simple docstring'''







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

						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(__lowercase		)

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

						UpperCAmelCase_								=     pipeline(__lowercase ,					__lowercase ,					__lowercase ,					__lowercase ,					jit=__lowercase		).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.1514745		) < 1e-3
							assert np.abs(np.abs(__lowercase ,					dtype=np.floataa		).sum() - 49947.875		) < 5e-1

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





					def        SCREAMING_SNAKE_CASE							(						self    :  str		):






						'''simple docstring'''







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

						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(__lowercase		)

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

						UpperCAmelCase_								=     pipeline(__lowercase ,					__lowercase ,					__lowercase ,					__lowercase ,					jit=__lowercase		).images

						assert images.shape == (num_samples, 1, 5_12, 5_12, 3)
						if jax.device_count() == 8:
							assert np.abs((np.abs(images[0, 0, :2, :2, -2:] ,					dtype=np.floataa		).sum() - 0.05652401)		) < 1e-3
							assert np.abs((np.abs(__lowercase ,					dtype=np.floataa		).sum() - 2383808.2)		) < 5e-1





					def        SCREAMING_SNAKE_CASE							(						self    :  Optional[int]		):






						'''simple docstring'''







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

						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(__lowercase		)

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

						UpperCAmelCase_								=     pipeline(__lowercase ,					__lowercase ,					__lowercase ,					__lowercase ,					jit=__lowercase		).images

						assert images.shape == (num_samples, 1, 5_12, 5_12, 3)
						if jax.device_count() == 8:
							assert np.abs((np.abs(images[0, 0, :2, :2, -2:] ,					dtype=np.floataa		).sum() - 0.04003906)		) < 1e-3
							assert np.abs((np.abs(__lowercase ,					dtype=np.floataa		).sum() - 2373516.75)		) < 5e-1





					def        SCREAMING_SNAKE_CASE							(						self    :  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(__lowercase		)

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

						UpperCAmelCase_								=     pipeline(__lowercase ,					__lowercase ,					__lowercase ,					__lowercase ,					jit=__lowercase		).images

						assert images.shape == (num_samples, 1, 5_12, 5_12, 3)
						if jax.device_count() == 8:
							assert np.abs((np.abs(images[0, 0, :2, :2, -2:] ,					dtype=np.floataa		).sum() - 0.04003906)		) < 1e-3
							assert np.abs((np.abs(__lowercase ,					dtype=np.floataa		).sum() - 2373516.75)		) < 5e-1





					def        SCREAMING_SNAKE_CASE							(						self    :  Optional[Any]		):






						'''simple docstring'''







						UpperCAmelCase_								=     FlaxDDIMScheduler(
						    beta_start=0.00085 ,					beta_end=0.012 ,					beta_schedule="""scaled_linear""" ,					set_alpha_to_one=__lowercase ,					steps_offset=1 ,					)

						UpperCAmelCase_						,	UpperCAmelCase_								=     FlaxStableDiffusionPipeline.from_pretrained(
						    """CompVis/stable-diffusion-v1-4""" ,					revision="""bf16""" ,					dtype=jnp.bfloataa ,					scheduler=__lowercase ,					safety_checker=__lowercase ,					)
						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(__lowercase		)

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

						UpperCAmelCase_								=     pipeline(__lowercase ,					__lowercase ,					__lowercase ,					__lowercase ,					jit=__lowercase		).images

						assert images.shape == (num_samples, 1, 5_12, 5_12, 3)
						if jax.device_count() == 8:
							assert np.abs((np.abs(images[0, 0, :2, :2, -2:] ,					dtype=np.floataa		).sum() - 0.045043945)		) < 1e-3
							assert np.abs((np.abs(__lowercase ,					dtype=np.floataa		).sum() - 2347693.5)		) < 5e-1





					def        SCREAMING_SNAKE_CASE							(						self    :  Any		):






						'''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		) ,					__lowercase		)

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

						UpperCAmelCase_								=     replicate(__lowercase		)
						UpperCAmelCase_								=     pipeline.prepare_inputs(__lowercase		)
						UpperCAmelCase_								=     shard(__lowercase		)
						UpperCAmelCase_								=     pipeline(__lowercase ,					__lowercase ,					__lowercase ,					jit=__lowercase		).images
						assert images.shape == (num_samples, 1, 5_12, 5_12, 3)
						UpperCAmelCase_								=     images[2, 0, 2_56, 10:17, 1]

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

						UpperCAmelCase_								=     replicate(__lowercase		)
						UpperCAmelCase_								=     pipeline.prepare_inputs(__lowercase		)
						UpperCAmelCase_								=     shard(__lowercase		)
						UpperCAmelCase_								=     pipeline(__lowercase ,					__lowercase ,					__lowercase ,					jit=__lowercase		).images
						assert images_eff.shape == (num_samples, 1, 5_12, 5_12, 3)
						UpperCAmelCase_								=     images[2, 0, 2_56, 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









from argparse import ArgumentParser, Namespace

from ..utils import logging
from . import BaseTransformersCLICommand




def    A_(			A   ):
	return ConvertCommand(
	    args.model_type				,      args.tf_checkpoint				,      args.pytorch_dump_output				,      args.config				,      args.finetuning_task_name   )


UpperCamelCase__   : Any					  = """
transformers can only be used from the commandline to convert TensorFlow models in PyTorch, In that case, it requires
TensorFlow to be installed. Please see https://www.tensorflow.org/install/ for installation instructions.
"""







class   _UpperCamelCase  (       A_  ):







					'''simple docstring'''

					@staticmethod
					def        SCREAMING_SNAKE_CASE							(						__lowercase    :  ArgumentParser		):






						'''simple docstring'''







						UpperCAmelCase_								=     parser.add_parser(
						    """convert""" ,					help="""CLI tool to run convert model from original author checkpoints to Transformers PyTorch checkpoints.""" ,					)
						train_parser.add_argument("""--model_type""" ,					type=__lowercase ,					required=__lowercase ,					help="""Model's type."""		)
						train_parser.add_argument(
						    """--tf_checkpoint""" ,					type=__lowercase ,					required=__lowercase ,					help="""TensorFlow checkpoint path or folder."""		)
						train_parser.add_argument(
						    """--pytorch_dump_output""" ,					type=__lowercase ,					required=__lowercase ,					help="""Path to the PyTorch saved model output."""		)
						train_parser.add_argument("""--config""" ,					type=__lowercase ,					default="""""" ,					help="""Configuration file path or folder."""		)
						train_parser.add_argument(
						    """--finetuning_task_name""" ,					type=__lowercase ,					default=__lowercase ,					help="""Optional fine-tuning task name if the TF model was a finetuned model.""" ,					)
						train_parser.set_defaults(func=__lowercase		)





					def __init__(						self    :  Optional[int] ,					__lowercase    :  str ,					__lowercase    :  str ,					__lowercase    :  str ,					__lowercase    :  str ,					__lowercase    :  str ,					*__lowercase    :  Tuple ,					):






						'''simple docstring'''







						UpperCAmelCase_								=     logging.get_logger("""transformers-cli/converting"""		)

						self._logger.info(F"""Loading model {model_type}"""		)
						UpperCAmelCase_								=     model_type
						UpperCAmelCase_								=     tf_checkpoint
						UpperCAmelCase_								=     pytorch_dump_output
						UpperCAmelCase_								=     config
						UpperCAmelCase_								=     finetuning_task_name





					def        SCREAMING_SNAKE_CASE							(						self    :  List[Any]		):






						'''simple docstring'''







						if self._model_type == "albert":
							try:
								from ..models.albert.convert_albert_original_tf_checkpoint_to_pytorch import (
								    convert_tf_checkpoint_to_pytorch,
								)
							except ImportError:
								raise ImportError(__lowercase		)

							convert_tf_checkpoint_to_pytorch(self._tf_checkpoint ,					self._config ,					self._pytorch_dump_output		)
						elif self._model_type == "bert":
							try:
								from ..models.bert.convert_bert_original_tf_checkpoint_to_pytorch import (
								    convert_tf_checkpoint_to_pytorch,
								)
							except ImportError:
								raise ImportError(__lowercase		)

							convert_tf_checkpoint_to_pytorch(self._tf_checkpoint ,					self._config ,					self._pytorch_dump_output		)
						elif self._model_type == "funnel":
							try:
								from ..models.funnel.convert_funnel_original_tf_checkpoint_to_pytorch import (
								    convert_tf_checkpoint_to_pytorch,
								)
							except ImportError:
								raise ImportError(__lowercase		)

							convert_tf_checkpoint_to_pytorch(self._tf_checkpoint ,					self._config ,					self._pytorch_dump_output		)
						elif self._model_type == "t5":
							try:
								from ..models.ta.convert_ta_original_tf_checkpoint_to_pytorch import convert_tf_checkpoint_to_pytorch
							except ImportError:
								raise ImportError(__lowercase		)

							convert_tf_checkpoint_to_pytorch(self._tf_checkpoint ,					self._config ,					self._pytorch_dump_output		)
						elif self._model_type == "gpt":
							from ..models.openai.convert_openai_original_tf_checkpoint_to_pytorch import (
							    convert_openai_checkpoint_to_pytorch,
							)

							convert_openai_checkpoint_to_pytorch(self._tf_checkpoint ,					self._config ,					self._pytorch_dump_output		)
						elif self._model_type == "transfo_xl":
							try:
								from ..models.transfo_xl.convert_transfo_xl_original_tf_checkpoint_to_pytorch import (
								    convert_transfo_xl_checkpoint_to_pytorch,
								)
							except ImportError:
								raise ImportError(__lowercase		)

							if "ckpt" in self._tf_checkpoint.lower():
								UpperCAmelCase_								=     self._tf_checkpoint
								UpperCAmelCase_								=     """"""
							else:
								UpperCAmelCase_								=     self._tf_checkpoint
								UpperCAmelCase_								=     """"""
							convert_transfo_xl_checkpoint_to_pytorch(
							    __lowercase ,					self._config ,					self._pytorch_dump_output ,					__lowercase		)
						elif self._model_type == "gpt2":
							try:
								from ..models.gpta.convert_gpta_original_tf_checkpoint_to_pytorch import (
								    convert_gpta_checkpoint_to_pytorch,
								)
							except ImportError:
								raise ImportError(__lowercase		)

							convert_gpta_checkpoint_to_pytorch(self._tf_checkpoint ,					self._config ,					self._pytorch_dump_output		)
						elif self._model_type == "xlnet":
							try:
								from ..models.xlnet.convert_xlnet_original_tf_checkpoint_to_pytorch import (
								    convert_xlnet_checkpoint_to_pytorch,
								)
							except ImportError:
								raise ImportError(__lowercase		)

							convert_xlnet_checkpoint_to_pytorch(
							    self._tf_checkpoint ,					self._config ,					self._pytorch_dump_output ,					self._finetuning_task_name		)
						elif self._model_type == "xlm":
							from ..models.xlm.convert_xlm_original_pytorch_checkpoint_to_pytorch import (
							    convert_xlm_checkpoint_to_pytorch,
							)

							convert_xlm_checkpoint_to_pytorch(self._tf_checkpoint ,					self._pytorch_dump_output		)
						elif self._model_type == "lxmert":
							from ..models.lxmert.convert_lxmert_original_tf_checkpoint_to_pytorch import (
							    convert_lxmert_checkpoint_to_pytorch,
							)

							convert_lxmert_checkpoint_to_pytorch(self._tf_checkpoint ,					self._pytorch_dump_output		)
						elif self._model_type == "rembert":
							from ..models.rembert.convert_rembert_tf_checkpoint_to_pytorch import (
							    convert_rembert_tf_checkpoint_to_pytorch,
							)

							convert_rembert_tf_checkpoint_to_pytorch(self._tf_checkpoint ,					self._config ,					self._pytorch_dump_output		)
						else:
							raise ValueError(
							    """--model_type should be selected in the list [bert, gpt, gpt2, t5, transfo_xl, xlnet, xlm, lxmert]"""		)


