






import gc
import unittest

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

from diffusers import (
    AutoencoderKL,
    DDIMScheduler,
    StableDiffusionSAGPipeline,
    UNetaDConditionModel,
)
from diffusers.utils import slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu

from ..pipeline_params import TEXT_TO_IMAGE_BATCH_PARAMS, TEXT_TO_IMAGE_IMAGE_PARAMS, TEXT_TO_IMAGE_PARAMS
from ..test_pipelines_common import PipelineLatentTesterMixin, PipelineTesterMixin


enable_full_determinism()


class 						_a       (  UpperCAmelCase__						,       UpperCAmelCase__						,       unittest.TestCase    ):






  """simple docstring"""





  A_						=   StableDiffusionSAGPipeline
  A_						=   TEXT_TO_IMAGE_PARAMS
  A_						=   TEXT_TO_IMAGE_BATCH_PARAMS
  A_						=   TEXT_TO_IMAGE_IMAGE_PARAMS
  A_						=   TEXT_TO_IMAGE_IMAGE_PARAMS
  A_						=   False




  def    _UpperCAmelCase     (     self )			-> int:
     torch.manual_seed(0 )
     UpperCamelCase_									=     UNetaDConditionModel(
         block_out_channels=(32, 64)		,		layers_per_block=2		,		sample_size=32		,		in_channels=4		,		out_channels=4		,		down_block_types=('DownBlock2D', 'CrossAttnDownBlock2D')		,		up_block_types=('CrossAttnUpBlock2D', 'UpBlock2D')		,		cross_attention_dim=32		,		)
     UpperCamelCase_									=     DDIMScheduler(
         beta_start=0.0_0_0_8_5		,		beta_end=0.0_1_2		,		beta_schedule='scaled_linear'		,		clip_sample=_UpperCAmelCase		,		set_alpha_to_one=_UpperCAmelCase		,		)
     torch.manual_seed(0 )
     UpperCamelCase_									=     AutoencoderKL(
         block_out_channels=[32, 64]		,		in_channels=3		,		out_channels=3		,		down_block_types=['DownEncoderBlock2D', 'DownEncoderBlock2D']		,		up_block_types=['UpDecoderBlock2D', 'UpDecoderBlock2D']		,		latent_channels=4		,		)
     torch.manual_seed(0 )
     UpperCamelCase_									=     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=1000		,		)
     UpperCamelCase_									=     CLIPTextModel(_UpperCAmelCase )
     UpperCamelCase_									=     CLIPTokenizer.from_pretrained('hf-internal-testing/tiny-random-clip' )

     UpperCamelCase_									=     {
         'unet': unet,
         'scheduler': scheduler,
         'vae': vae,
         'text_encoder': text_encoder,
         'tokenizer': tokenizer,
         'safety_checker': None,
         'feature_extractor': None,
     }
     return components




  def    _UpperCAmelCase     (     self		,		_UpperCAmelCase		,		_UpperCAmelCase=0 )			-> Any:
     if str(_UpperCAmelCase ).startswith('mps' ):
        UpperCamelCase_									=     torch.manual_seed(_UpperCAmelCase )
     else:
        UpperCamelCase_									=     torch.Generator(device=_UpperCAmelCase ).manual_seed(_UpperCAmelCase )
     UpperCamelCase_									=     {
         'prompt': '.',
         'generator': generator,
         'num_inference_steps': 2,
         'guidance_scale': 1.0,
         'sag_scale': 1.0,
         'output_type': 'numpy',
     }
     return inputs




  def    _UpperCAmelCase     (     self )			-> Dict:
     super().test_inference_batch_single_identical(expected_max_diff=3e-3 )




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






  """simple docstring"""





  def    _UpperCAmelCase     (     self )			-> int:
     # clean up the VRAM after each test
     super().tearDown()
     gc.collect()
     torch.cuda.empty_cache()




  def    _UpperCAmelCase     (     self )			-> Optional[Any]:
     UpperCamelCase_									=     StableDiffusionSAGPipeline.from_pretrained('CompVis/stable-diffusion-v1-4' )
     UpperCamelCase_									=     sag_pipe.to(_UpperCAmelCase )
     sag_pipe.set_progress_bar_config(disable=_UpperCAmelCase )

     UpperCamelCase_									=     '.'
     UpperCamelCase_									=     torch.manual_seed(0 )
     UpperCamelCase_									=     sag_pipe(
         [prompt]		,		generator=_UpperCAmelCase		,		guidance_scale=7.5		,		sag_scale=1.0		,		num_inference_steps=20		,		output_type='np' )

     UpperCamelCase_									=     output.images

     UpperCamelCase_									=     image[0, -3:, -3:, -1]

     assert image.shape == (1, 512, 512, 3)
     UpperCamelCase_									=     np.array([0.1_5_6_8, 0.1_7_3_8, 0.1_6_9_5, 0.1_6_9_3, 0.1_5_0_7, 0.1_7_0_5, 0.1_5_4_7, 0.1_7_5_1, 0.1_9_4_9] )

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




  def    _UpperCAmelCase     (     self )			-> List[str]:
     UpperCamelCase_									=     StableDiffusionSAGPipeline.from_pretrained('stabilityai/stable-diffusion-2-1-base' )
     UpperCamelCase_									=     sag_pipe.to(_UpperCAmelCase )
     sag_pipe.set_progress_bar_config(disable=_UpperCAmelCase )

     UpperCamelCase_									=     '.'
     UpperCamelCase_									=     torch.manual_seed(0 )
     UpperCamelCase_									=     sag_pipe(
         [prompt]		,		generator=_UpperCAmelCase		,		guidance_scale=7.5		,		sag_scale=1.0		,		num_inference_steps=20		,		output_type='np' )

     UpperCamelCase_									=     output.images

     UpperCamelCase_									=     image[0, -3:, -3:, -1]

     assert image.shape == (1, 512, 512, 3)
     UpperCamelCase_									=     np.array([0.3_4_5_9, 0.2_8_7_6, 0.2_5_3_7, 0.3_0_0_2, 0.2_6_7_1, 0.2_1_6_0, 0.3_0_2_6, 0.2_2_6_2, 0.2_3_7_1] )

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




  def    _UpperCAmelCase     (     self )			-> Tuple:
     UpperCamelCase_									=     StableDiffusionSAGPipeline.from_pretrained('stabilityai/stable-diffusion-2-1-base' )
     UpperCamelCase_									=     sag_pipe.to(_UpperCAmelCase )
     sag_pipe.set_progress_bar_config(disable=_UpperCAmelCase )

     UpperCamelCase_									=     '.'
     UpperCamelCase_									=     torch.manual_seed(0 )
     UpperCamelCase_									=     sag_pipe(
         [prompt]		,		width=768		,		height=512		,		generator=_UpperCAmelCase		,		guidance_scale=7.5		,		sag_scale=1.0		,		num_inference_steps=20		,		output_type='np'		,		)

     UpperCamelCase_									=     output.images

     assert image.shape == (1, 512, 768, 3)













import inspect
import os
import unittest
from pathlib import Path

import torch

import accelerate
from accelerate.test_utils import execute_subprocess_async
from accelerate.test_utils.testing import run_command


class 						_a       (  unittest.TestCase    ):






  """simple docstring"""





  A_						=   inspect.getfile(accelerate.test_utils    )
  A_						=   os.path.sep.join(mod_file.split(os.path.sep    )[:-1] + ["""scripts""", """test_cli.py"""]    )

  A_						=   ["""accelerate""", """launch"""]
  A_						=   Path.home() / """.cache/huggingface/accelerate"""
  A_						=   """default_config.yaml"""
  A_						=   config_folder / config_file
  A_						=   config_folder / """_default_config.yaml"""

  A_						=   Path("""tests/test_configs"""    )




  @classmethod
  def    _UpperCAmelCase     (     cls )			-> Union[str, Any]:
     if cls.config_path.is_file():
        cls.config_path.rename(cls.changed_path )




  @classmethod
  def    _UpperCAmelCase     (     cls )			-> List[str]:
     if cls.changed_path.is_file():
        cls.changed_path.rename(cls.config_path )




  def    _UpperCAmelCase     (     self )			-> Optional[Any]:
     UpperCamelCase_									=     self.base_cmd
     if torch.cuda.is_available() and (torch.cuda.device_count() > 1):
        cmd += ["--multi_gpu"]
     execute_subprocess_async(cmd + [self.test_file_path]		,		env=os.environ.copy() )




  def    _UpperCAmelCase     (     self )			-> Optional[int]:
     for config in sorted(self.test_config_path.glob('**/*.yaml' ) ):
        with self.subTest(config_file=_UpperCAmelCase ):
           execute_subprocess_async(
               self.base_cmd + ['--config_file', str(_UpperCAmelCase ), self.test_file_path]		,		env=os.environ.copy() )




  def    _UpperCAmelCase     (     self )			-> Tuple:
     execute_subprocess_async(['accelerate', 'test']		,		env=os.environ.copy() )




class 						_a       (  unittest.TestCase    ):






  """simple docstring"""





  A_						=   """test-tpu"""
  A_						=   """us-central1-a"""
  A_						=   """ls"""
  A_						=   ["""accelerate""", """tpu-config"""]
  A_						=   """cd /usr/share"""
  A_						=   """tests/test_samples/test_command_file.sh"""
  A_						=   """Running gcloud compute tpus tpu-vm ssh"""




  def    _UpperCAmelCase     (     self )			-> Any:
     UpperCamelCase_									=     run_command(
         self.cmd
         + ['--command', self.command, '--tpu_zone', self.tpu_zone, '--tpu_name', self.tpu_name, '--debug']		,		return_stdout=_UpperCAmelCase		,		)
     self.assertIn(
         f"""{self.gcloud} test-tpu --zone us-central1-a --command {self.base_output}; ls --worker all"""		,		_UpperCAmelCase		,		)




  def    _UpperCAmelCase     (     self )			-> Any:
     UpperCamelCase_									=     run_command(
         self.cmd
         + [
             '--config_file',
             'tests/test_configs/0_12_0.yaml',
             '--command',
             self.command,
             '--tpu_zone',
             self.tpu_zone,
             '--tpu_name',
             self.tpu_name,
             '--debug',
         ]		,		return_stdout=_UpperCAmelCase		,		)
     self.assertIn(
         f"""{self.gcloud} test-tpu --zone us-central1-a --command {self.base_output}; ls --worker all"""		,		_UpperCAmelCase		,		)




  def    _UpperCAmelCase     (     self )			-> Any:
     UpperCamelCase_									=     run_command(
         self.cmd + ['--config_file', 'tests/test_configs/latest.yaml', '--debug']		,		return_stdout=_UpperCAmelCase )
     self.assertIn(
         f"""{self.gcloud} test-tpu --zone us-central1-a --command {self.base_output}; echo \"hello world\"; echo \"this is a second command\" --worker all"""		,		_UpperCAmelCase		,		)




  def    _UpperCAmelCase     (     self )			-> Tuple:
     UpperCamelCase_									=     run_command(
         self.cmd + ['--config_file', 'tests/test_configs/latest.yaml', '--command', self.command, '--debug']		,		return_stdout=_UpperCAmelCase		,		)
     self.assertIn(
         f"""{self.gcloud} test-tpu --zone us-central1-a --command {self.base_output}; ls --worker all"""		,		_UpperCAmelCase		,		)




  def    _UpperCAmelCase     (     self )			-> Optional[Any]:
     UpperCamelCase_									=     run_command(
         self.cmd
         + [
             '--config_file',
             'tests/test_configs/latest.yaml',
             '--command',
             self.command,
             '--command',
             'echo "Hello World"',
             '--debug',
         ]		,		return_stdout=_UpperCAmelCase		,		)
     self.assertIn(
         f"""{self.gcloud} test-tpu --zone us-central1-a --command {self.base_output}; ls; echo \"Hello World\" --worker all"""		,		_UpperCAmelCase		,		)




  def    _UpperCAmelCase     (     self )			-> Dict:
     UpperCamelCase_									=     run_command(
         self.cmd
         + ['--config_file', 'tests/test_configs/latest.yaml', '--command_file', self.command_file, '--debug']		,		return_stdout=_UpperCAmelCase		,		)
     self.assertIn(
         f"""{self.gcloud} test-tpu --zone us-central1-a --command {self.base_output}; echo \"hello world\"; echo \"this is a second command\" --worker all"""		,		_UpperCAmelCase		,		)




  def    _UpperCAmelCase     (     self )			-> Dict:
     UpperCamelCase_									=     run_command(
         self.cmd
         + [
             '--config_file',
             'tests/test_configs/0_12_0.yaml',
             '--command_file',
             self.command_file,
             '--tpu_zone',
             self.tpu_zone,
             '--tpu_name',
             self.tpu_name,
             '--debug',
         ]		,		return_stdout=_UpperCAmelCase		,		)
     self.assertIn(
         f"""{self.gcloud} test-tpu --zone us-central1-a --command {self.base_output}; echo \"hello world\"; echo \"this is a second command\" --worker all"""		,		_UpperCAmelCase		,		)




  def    _UpperCAmelCase     (     self )			-> List[Any]:
     UpperCamelCase_									=     run_command(
         self.cmd + ['--config_file', 'tests/test_configs/latest.yaml', '--install_accelerate', '--debug']		,		return_stdout=_UpperCAmelCase		,		)
     self.assertIn(
         f"""{self.gcloud} test-tpu --zone us-central1-a --command {self.base_output}; pip install accelerate -U; echo \"hello world\"; echo \"this is a second command\" --worker all"""		,		_UpperCAmelCase		,		)




  def    _UpperCAmelCase     (     self )			-> Optional[Any]:
     UpperCamelCase_									=     run_command(
         self.cmd
         + [
             '--config_file',
             'tests/test_configs/latest.yaml',
             '--install_accelerate',
             '--accelerate_version',
             '12.0.0',
             '--debug',
         ]		,		return_stdout=_UpperCAmelCase		,		)
     self.assertIn(
         f"""{self.gcloud} test-tpu --zone us-central1-a --command {self.base_output}; pip install accelerate==12.0.0; echo \"hello world\"; echo \"this is a second command\" --worker all"""		,		_UpperCAmelCase		,		)





