




import importlib.util
import os
import platform
from argparse import ArgumentParser

import huggingface_hub

from .. import __version__ as version
from ..utils import (
    is_accelerate_available,
    is_flax_available,
    is_safetensors_available,
    is_tf_available,
    is_torch_available,
)
from . import BaseTransformersCLICommand







def SCREAMING_SNAKE_CASE       ( lowercase_		)     -> Union[str, Any]:



       """simple docstring"""

       return EnvironmentCommand()







def SCREAMING_SNAKE_CASE       ( lowercase_		)     -> Dict:



       """simple docstring"""

       return EnvironmentCommand(args.accelerate_config_file		)


class        UpperCamelCase_					(				UpperCAmelCase__  ):

    '''simple docstring'''






    @staticmethod
    def  SCREAMING_SNAKE_CASE							(    UpperCAmelCase__   :			ArgumentParser)		->Union[str, Any]:





           '''simple docstring'''

           A__   				=					parser.add_parser('''env''')
           download_parser.set_defaults(func=UpperCAmelCase__)
           download_parser.add_argument(
               '''--accelerate-config_file'''      ,				default=UpperCAmelCase__      ,				help='''The accelerate config file to use for the default values in the launching script.'''      ,				)
           download_parser.set_defaults(func=UpperCAmelCase__)


    def __init__(    self   :			Optional[Any]      ,				UpperCAmelCase__   :			Optional[Any]      ,				*UpperCAmelCase__   :			Dict)		->None:





           '''simple docstring'''

           A__   				=					accelerate_config_file


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





           '''simple docstring'''

           A__   				=					'''not installed'''
           if is_safetensors_available():
                  import safetensors

                  A__   				=					safetensors.__version__
           elif importlib.util.find_spec('''safetensors''') is not None:
                  import safetensors

                  A__   				=					f"""{safetensors.__version__} but is ignored because of PyTorch version too old."""

           A__   				=					'''not installed'''
           A__   				=					A__   				=					'''not found'''
           if is_accelerate_available():
                  import accelerate
                  from accelerate.commands.config import default_config_file, load_config_from_file

                  A__   				=					accelerate.__version__
                  # Get the default from the config file.
                  if self._accelerate_config_file is not None or os.path.isfile(UpperCAmelCase__):
                         A__   				=					load_config_from_file(self._accelerate_config_file).to_dict()

                  A__   				=					(
                      '''\n'''.join([f"""\t- {prop}: {val}""" for prop, val in accelerate_config.items()])
                      if isinstance(UpperCAmelCase__      ,				UpperCAmelCase__)
                      else f"""\t{accelerate_config}"""
                  )

           A__   				=					'''not installed'''
           A__   				=					'''NA'''
           if is_torch_available():
                  import torch

                  A__   				=					torch.__version__
                  A__   				=					torch.cuda.is_available()

           A__   				=					'''not installed'''
           A__   				=					'''NA'''
           if is_tf_available():
                  import tensorflow as tf

                  A__   				=					tf.__version__
                  try:
                         # deprecated in v2.1
                         A__   				=					tf.test.is_gpu_available()
                  except AttributeError:
                         # returns list of devices, convert to bool
                         A__   				=					bool(tf.config.list_physical_devices('''GPU'''))

           A__   				=					'''not installed'''
           A__   				=					'''not installed'''
           A__   				=					'''not installed'''
           A__   				=					'''NA'''
           if is_flax_available():
                  import flax
                  import jax
                  import jaxlib

                  A__   				=					flax.__version__
                  A__   				=					jax.__version__
                  A__   				=					jaxlib.__version__
                  A__   				=					jax.lib.xla_bridge.get_backend().platform

           A__   				=					{
               '''`transformers` version''': version,
               '''Platform''': platform.platform(),
               '''Python version''': platform.python_version(),
               '''Huggingface_hub version''': huggingface_hub.__version__,
               '''Safetensors version''': f"""{safetensors_version}""",
               '''Accelerate version''': f"""{accelerate_version}""",
               '''Accelerate config''': f"""{accelerate_config_str}""",
               '''PyTorch version (GPU?)''': f"""{pt_version} ({pt_cuda_available})""",
               '''Tensorflow version (GPU?)''': f"""{tf_version} ({tf_cuda_available})""",
               '''Flax version (CPU?/GPU?/TPU?)''': f"""{flax_version} ({jax_backend})""",
               '''Jax version''': f"""{jax_version}""",
               '''JaxLib version''': f"""{jaxlib_version}""",
               '''Using GPU in script?''': '''<fill in>''',
               '''Using distributed or parallel set-up in script?''': '''<fill in>''',
           }

           print('''\nCopy-and-paste the text below in your GitHub issue and FILL OUT the two last points.\n''')
           print(self.format_dict(UpperCAmelCase__))

           return info





    @staticmethod
    def  SCREAMING_SNAKE_CASE							(    UpperCAmelCase__   :			str)		->Optional[int]:





           '''simple docstring'''

           return "\n".join([f"""- {prop}: {val}""" for prop, val in d.items()]) + "\n"








import gc
import random
import unittest

import numpy as np
import torch
from PIL import Image
from transformers import CLIPTextConfig, CLIPTextModel, CLIPTokenizer

from diffusers import AutoencoderKL, DDIMScheduler, DDPMScheduler, StableDiffusionUpscalePipeline, UNetaDConditionModel
from diffusers.utils import floats_tensor, load_image, load_numpy, slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu


enable_full_determinism()


class        UpperCamelCase_					(				unittest.TestCase  ):

    '''simple docstring'''






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





           '''simple docstring'''

           super().tearDown()
           gc.collect()
           torch.cuda.empty_cache()


    @property
    def  SCREAMING_SNAKE_CASE							(    self   :			Union[str, Any])		->Dict:





           '''simple docstring'''

           A__   				=					1
           A__   				=					3
           A__   				=					(32, 32)

           A__   				=					floats_tensor((batch_size, num_channels) + sizes      ,				rng=random.Random(0)).to(UpperCAmelCase__)
           return image


    @property
    def  SCREAMING_SNAKE_CASE							(    self   :			Any)		->str:





           '''simple docstring'''

           torch.manual_seed(0)
           A__   				=					UNetaDConditionModel(
               block_out_channels=(32, 32, 64)      ,				layers_per_block=2      ,				sample_size=32      ,				in_channels=7      ,				out_channels=4      ,				down_block_types=('''DownBlock2D''', '''CrossAttnDownBlock2D''', '''CrossAttnDownBlock2D''')      ,				up_block_types=('''CrossAttnUpBlock2D''', '''CrossAttnUpBlock2D''', '''UpBlock2D''')      ,				cross_attention_dim=32      ,				attention_head_dim=8      ,				use_linear_projection=UpperCAmelCase__      ,				only_cross_attention=(True, True, False)      ,				num_class_embeds=100      ,				)
           return model


    @property
    def  SCREAMING_SNAKE_CASE							(    self   :			str)		->Dict:





           '''simple docstring'''

           torch.manual_seed(0)
           A__   				=					AutoencoderKL(
               block_out_channels=[32, 32, 64]      ,				in_channels=3      ,				out_channels=3      ,				down_block_types=['''DownEncoderBlock2D''', '''DownEncoderBlock2D''', '''DownEncoderBlock2D''']      ,				up_block_types=['''UpDecoderBlock2D''', '''UpDecoderBlock2D''', '''UpDecoderBlock2D''']      ,				latent_channels=4      ,				)
           return model


    @property
    def  SCREAMING_SNAKE_CASE							(    self   :			List[str])		->Dict:





           '''simple docstring'''

           torch.manual_seed(0)
           A__   				=					CLIPTextConfig(
               bos_token_id=0      ,				eos_token_id=2      ,				hidden_size=32      ,				intermediate_size=37      ,				layer_norm_eps=1e-05      ,				num_attention_heads=4      ,				num_hidden_layers=5      ,				pad_token_id=1      ,				vocab_size=1_000      ,				hidden_act='''gelu'''      ,				projection_dim=512      ,				)
           return CLIPTextModel(UpperCAmelCase__)


    def  SCREAMING_SNAKE_CASE							(    self   :			Union[str, Any])		->Tuple:





           '''simple docstring'''

           A__   				=					'''cpu'''  # ensure determinism for the device-dependent torch.Generator
           A__   				=					self.dummy_cond_unet_upscale
           A__   				=					DDPMScheduler()
           A__   				=					DDIMScheduler(prediction_type='''v_prediction''')
           A__   				=					self.dummy_vae
           A__   				=					self.dummy_text_encoder
           A__   				=					CLIPTokenizer.from_pretrained('''hf-internal-testing/tiny-random-clip''')

           A__   				=					self.dummy_image.cpu().permute(0      ,				2      ,				3      ,				1)[0]
           A__   				=					Image.fromarray(np.uinta(UpperCAmelCase__)).convert('''RGB''').resize((64, 64))

           # make sure here that pndm scheduler skips prk
           A__   				=					StableDiffusionUpscalePipeline(
               unet=UpperCAmelCase__      ,				low_res_scheduler=UpperCAmelCase__      ,				scheduler=UpperCAmelCase__      ,				vae=UpperCAmelCase__      ,				text_encoder=UpperCAmelCase__      ,				tokenizer=UpperCAmelCase__      ,				max_noise_level=350      ,				)
           A__   				=					sd_pipe.to(UpperCAmelCase__)
           sd_pipe.set_progress_bar_config(disable=UpperCAmelCase__)

           A__   				=					'''A painting of a squirrel eating a burger'''
           A__   				=					torch.Generator(device=UpperCAmelCase__).manual_seed(0)
           A__   				=					sd_pipe(
               [prompt]      ,				image=UpperCAmelCase__      ,				generator=UpperCAmelCase__      ,				guidance_scale=6.0      ,				noise_level=20      ,				num_inference_steps=2      ,				output_type='''np'''      ,				)

           A__   				=					output.images

           A__   				=					torch.Generator(device=UpperCAmelCase__).manual_seed(0)
           A__   				=					sd_pipe(
               [prompt]      ,				image=UpperCAmelCase__      ,				generator=UpperCAmelCase__      ,				guidance_scale=6.0      ,				noise_level=20      ,				num_inference_steps=2      ,				output_type='''np'''      ,				return_dict=UpperCAmelCase__      ,				)[0]

           A__   				=					image[0, -3:, -3:, -1]
           A__   				=					image_from_tuple[0, -3:, -3:, -1]

           A__   				=					low_res_image.size[0] * 4
           assert image.shape == (1, expected_height_width, expected_height_width, 3)
           A__   				=					np.array([0.3113, 0.3910, 0.4272, 0.4859, 0.5061, 0.4652, 0.5362, 0.5715, 0.5661])

           assert np.abs(image_slice.flatten() - expected_slice).max() < 1e-2
           assert np.abs(image_from_tuple_slice.flatten() - expected_slice).max() < 1e-2


    def  SCREAMING_SNAKE_CASE							(    self   :			Union[str, Any])		->int:





           '''simple docstring'''

           A__   				=					'''cpu'''  # ensure determinism for the device-dependent torch.Generator
           A__   				=					self.dummy_cond_unet_upscale
           A__   				=					DDPMScheduler()
           A__   				=					DDIMScheduler(prediction_type='''v_prediction''')
           A__   				=					self.dummy_vae
           A__   				=					self.dummy_text_encoder
           A__   				=					CLIPTokenizer.from_pretrained('''hf-internal-testing/tiny-random-clip''')

           A__   				=					self.dummy_image.cpu().permute(0      ,				2      ,				3      ,				1)[0]
           A__   				=					Image.fromarray(np.uinta(UpperCAmelCase__)).convert('''RGB''').resize((64, 64))

           # make sure here that pndm scheduler skips prk
           A__   				=					StableDiffusionUpscalePipeline(
               unet=UpperCAmelCase__      ,				low_res_scheduler=UpperCAmelCase__      ,				scheduler=UpperCAmelCase__      ,				vae=UpperCAmelCase__      ,				text_encoder=UpperCAmelCase__      ,				tokenizer=UpperCAmelCase__      ,				max_noise_level=350      ,				)
           A__   				=					sd_pipe.to(UpperCAmelCase__)
           sd_pipe.set_progress_bar_config(disable=UpperCAmelCase__)

           A__   				=					'''A painting of a squirrel eating a burger'''
           A__   				=					sd_pipe(
               2 * [prompt]      ,				image=2 * [low_res_image]      ,				guidance_scale=6.0      ,				noise_level=20      ,				num_inference_steps=2      ,				output_type='''np'''      ,				)
           A__   				=					output.images
           assert image.shape[0] == 2

           A__   				=					torch.Generator(device=UpperCAmelCase__).manual_seed(0)
           A__   				=					sd_pipe(
               [prompt]      ,				image=UpperCAmelCase__      ,				generator=UpperCAmelCase__      ,				num_images_per_prompt=2      ,				guidance_scale=6.0      ,				noise_level=20      ,				num_inference_steps=2      ,				output_type='''np'''      ,				)
           A__   				=					output.images
           assert image.shape[0] == 2





    @unittest.skipIf(torch_device != '''cuda'''      ,				'''This test requires a GPU''')
    def  SCREAMING_SNAKE_CASE							(    self   :			Any)		->str:





           '''simple docstring'''

           A__   				=					self.dummy_cond_unet_upscale
           A__   				=					DDPMScheduler()
           A__   				=					DDIMScheduler(prediction_type='''v_prediction''')
           A__   				=					self.dummy_vae
           A__   				=					self.dummy_text_encoder
           A__   				=					CLIPTokenizer.from_pretrained('''hf-internal-testing/tiny-random-clip''')

           A__   				=					self.dummy_image.cpu().permute(0      ,				2      ,				3      ,				1)[0]
           A__   				=					Image.fromarray(np.uinta(UpperCAmelCase__)).convert('''RGB''').resize((64, 64))

           # put models in fp16, except vae as it overflows in fp16
           A__   				=					unet.half()
           A__   				=					text_encoder.half()

           # make sure here that pndm scheduler skips prk
           A__   				=					StableDiffusionUpscalePipeline(
               unet=UpperCAmelCase__      ,				low_res_scheduler=UpperCAmelCase__      ,				scheduler=UpperCAmelCase__      ,				vae=UpperCAmelCase__      ,				text_encoder=UpperCAmelCase__      ,				tokenizer=UpperCAmelCase__      ,				max_noise_level=350      ,				)
           A__   				=					sd_pipe.to(UpperCAmelCase__)
           sd_pipe.set_progress_bar_config(disable=UpperCAmelCase__)

           A__   				=					'''A painting of a squirrel eating a burger'''
           A__   				=					torch.manual_seed(0)
           A__   				=					sd_pipe(
               [prompt]      ,				image=UpperCAmelCase__      ,				generator=UpperCAmelCase__      ,				num_inference_steps=2      ,				output_type='''np'''      ,				).images

           A__   				=					low_res_image.size[0] * 4
           assert image.shape == (1, expected_height_width, expected_height_width, 3)


@slow
@require_torch_gpu
class        UpperCamelCase_					(				unittest.TestCase  ):

    '''simple docstring'''






    def  SCREAMING_SNAKE_CASE							(    self   :			Dict)		->str:





           '''simple docstring'''

           super().tearDown()
           gc.collect()
           torch.cuda.empty_cache()


    def  SCREAMING_SNAKE_CASE							(    self   :			str)		->Any:





           '''simple docstring'''

           A__   				=					load_image(
               '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'''
               '''/sd2-upscale/low_res_cat.png''')
           A__   				=					load_numpy(
               '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/sd2-upscale'''
               '''/upsampled_cat.npy''')

           A__   				=					'''stabilityai/stable-diffusion-x4-upscaler'''
           A__   				=					StableDiffusionUpscalePipeline.from_pretrained(UpperCAmelCase__)
           pipe.to(UpperCAmelCase__)
           pipe.set_progress_bar_config(disable=UpperCAmelCase__)
           pipe.enable_attention_slicing()

           A__   				=					'''a cat sitting on a park bench'''

           A__   				=					torch.manual_seed(0)
           A__   				=					pipe(
               prompt=UpperCAmelCase__      ,				image=UpperCAmelCase__      ,				generator=UpperCAmelCase__      ,				output_type='''np'''      ,				)
           A__   				=					output.images[0]

           assert image.shape == (512, 512, 3)
           assert np.abs(expected_image - image).max() < 1e-3


    def  SCREAMING_SNAKE_CASE							(    self   :			int)		->str:





           '''simple docstring'''

           A__   				=					load_image(
               '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'''
               '''/sd2-upscale/low_res_cat.png''')
           A__   				=					load_numpy(
               '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/sd2-upscale'''
               '''/upsampled_cat_fp16.npy''')

           A__   				=					'''stabilityai/stable-diffusion-x4-upscaler'''
           A__   				=					StableDiffusionUpscalePipeline.from_pretrained(
               UpperCAmelCase__      ,				torch_dtype=torch.floataa      ,				)
           pipe.to(UpperCAmelCase__)
           pipe.set_progress_bar_config(disable=UpperCAmelCase__)
           pipe.enable_attention_slicing()

           A__   				=					'''a cat sitting on a park bench'''

           A__   				=					torch.manual_seed(0)
           A__   				=					pipe(
               prompt=UpperCAmelCase__      ,				image=UpperCAmelCase__      ,				generator=UpperCAmelCase__      ,				output_type='''np'''      ,				)
           A__   				=					output.images[0]

           assert image.shape == (512, 512, 3)
           assert np.abs(expected_image - image).max() < 5e-1





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





           '''simple docstring'''

           torch.cuda.empty_cache()
           torch.cuda.reset_max_memory_allocated()
           torch.cuda.reset_peak_memory_stats()

           A__   				=					load_image(
               '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'''
               '''/sd2-upscale/low_res_cat.png''')

           A__   				=					'''stabilityai/stable-diffusion-x4-upscaler'''
           A__   				=					StableDiffusionUpscalePipeline.from_pretrained(
               UpperCAmelCase__      ,				torch_dtype=torch.floataa      ,				)
           pipe.to(UpperCAmelCase__)
           pipe.set_progress_bar_config(disable=UpperCAmelCase__)
           pipe.enable_attention_slicing(1)
           pipe.enable_sequential_cpu_offload()

           A__   				=					'''a cat sitting on a park bench'''

           A__   				=					torch.manual_seed(0)
           A__   				=					pipe(
               prompt=UpperCAmelCase__      ,				image=UpperCAmelCase__      ,				generator=UpperCAmelCase__      ,				num_inference_steps=5      ,				output_type='''np'''      ,				)

           A__   				=					torch.cuda.max_memory_allocated()
           # make sure that less than 2.9 GB is allocated
           assert mem_bytes < 2.9 * 10**9


