
from argparse import ArgumentParser

from .add_new_model import AddNewModelCommand
from .add_new_model_like import AddNewModelLikeCommand
from .convert import ConvertCommand
from .download import DownloadCommand
from .env import EnvironmentCommand
from .lfs import LfsCommands
from .pt_to_tf import PTtoTFCommand
from .run import RunCommand
from .serving import ServeCommand
from .user import UserCommands




def    __UpperCamelCase     (				):
 __a :  Any						=  ArgumentParser('''Transformers CLI tool'''  ,   usage='''transformers-cli <command> [<args>]'''  )
 __a :  int						=  parser.add_subparsers(help='''transformers-cli command helpers'''  )

 # Register commands
 ConvertCommand.register_subcommand(lowerCAmelCase__  )
 DownloadCommand.register_subcommand(lowerCAmelCase__  )
 EnvironmentCommand.register_subcommand(lowerCAmelCase__  )
 RunCommand.register_subcommand(lowerCAmelCase__  )
 ServeCommand.register_subcommand(lowerCAmelCase__  )
 UserCommands.register_subcommand(lowerCAmelCase__  )
 AddNewModelCommand.register_subcommand(lowerCAmelCase__  )
 AddNewModelLikeCommand.register_subcommand(lowerCAmelCase__  )
 LfsCommands.register_subcommand(lowerCAmelCase__  )
 PTtoTFCommand.register_subcommand(lowerCAmelCase__  )

 # Let's go
 __a :  List[str]						=  parser.parse_args()

 if not hasattr(lowerCAmelCase__  ,   '''func'''  ):
  parser.print_help()
  exit(1  )

 # Run
 __a :  Optional[int]						=  args.func(lowerCAmelCase__  )
 service.run()


if __name__ == "__main__":
      main()


import gc
import unittest

from diffusers import FlaxDPMSolverMultistepScheduler, FlaxStableDiffusionPipeline
from diffusers.utils import is_flax_available, 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 		UpperCamelCase__  (       unittest.TestCase    ):
   def       lowerCAmelCase			(self :					Optional[Any]   ):
    # clean up the VRAM after each test
    super().tearDown()
    gc.collect()
   def       lowerCAmelCase			(self :					int   ):
    __a				,       __a :  Optional[Any]						=  FlaxStableDiffusionPipeline.from_pretrained(
        '''stabilityai/stable-diffusion-2'''					,				revision='''bf16'''					,				dtype=jnp.bfloataa					,				)

    __a :  Any						=  '''A painting of a squirrel eating a burger'''
    __a :  Dict						=  jax.device_count()
    __a :  Optional[int]						=  num_samples * [prompt]
    __a :  Optional[Any]						=  sd_pipe.prepare_inputs(snake_case_   )

    __a :  Optional[Any]						=  replicate(snake_case_   )
    __a :  Optional[int]						=  shard(snake_case_   )

    __a :  int						=  jax.random.PRNGKey(0   )
    __a :  str						=  jax.random.split(snake_case_					,				jax.device_count()   )

    __a :  int						=  sd_pipe(snake_case_					,				snake_case_					,				snake_case_					,				num_inference_steps=2_5					,				jit=snake_case_   )[0]
    assert images.shape == (jax.device_count(), 1, 7_6_8, 7_6_8, 3)

    __a :  Optional[int]						=  images.reshape((images.shape[0] * images.shape[1],) + images.shape[-3:]   )
    __a :  List[Any]						=  images[0, 2_5_3:2_5_6, 2_5_3:2_5_6, -1]

    __a :  Dict						=  jnp.asarray(jax.device_get(image_slice.flatten()   )   )
    __a :  str						=  jnp.array([0.4238, 0.4414, 0.4395, 0.4453, 0.4629, 0.4590, 0.4531, 0.4_5508, 0.4512]   )
    print(f"output_slice: {output_slice}"   )
    assert jnp.abs(output_slice - expected_slice   ).max() < 1E-2







   def       lowerCAmelCase			(self :					Tuple   ):
    __a :  Optional[Any]						=  '''stabilityai/stable-diffusion-2'''
    __a				,       __a :  Dict						=  FlaxDPMSolverMultistepScheduler.from_pretrained(snake_case_					,				subfolder='''scheduler'''   )
    __a				,       __a :  List[Any]						=  FlaxStableDiffusionPipeline.from_pretrained(
        snake_case_					,				scheduler=snake_case_					,				revision='''bf16'''					,				dtype=jnp.bfloataa					,				)
    __a :  Union[str, Any]						=  scheduler_params

    __a :  List[Any]						=  '''A painting of a squirrel eating a burger'''
    __a :  Any						=  jax.device_count()
    __a :  Any						=  num_samples * [prompt]
    __a :  List[Any]						=  sd_pipe.prepare_inputs(snake_case_   )

    __a :  Tuple						=  replicate(snake_case_   )
    __a :  Dict						=  shard(snake_case_   )

    __a :  Dict						=  jax.random.PRNGKey(0   )
    __a :  Dict						=  jax.random.split(snake_case_					,				jax.device_count()   )

    __a :  str						=  sd_pipe(snake_case_					,				snake_case_					,				snake_case_					,				num_inference_steps=2_5					,				jit=snake_case_   )[0]
    assert images.shape == (jax.device_count(), 1, 7_6_8, 7_6_8, 3)

    __a :  List[Any]						=  images.reshape((images.shape[0] * images.shape[1],) + images.shape[-3:]   )
    __a :  Any						=  images[0, 2_5_3:2_5_6, 2_5_3:2_5_6, -1]

    __a :  List[Any]						=  jnp.asarray(jax.device_get(image_slice.flatten()   )   )
    __a :  List[Any]						=  jnp.array([0.4336, 0.4_2969, 0.4453, 0.4199, 0.4297, 0.4531, 0.4434, 0.4434, 0.4297]   )
    print(f"output_slice: {output_slice}"   )
    assert jnp.abs(output_slice - expected_slice   ).max() < 1E-2
