


import gc
import random
import tempfile
import unittest

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

from diffusers import AutoencoderKL, DDIMScheduler, LMSDiscreteScheduler, PNDMScheduler, UNetaDConditionModel
from diffusers.pipelines.stable_diffusion_safe import StableDiffusionPipelineSafe as StableDiffusionPipeline
from diffusers.utils import floats_tensor, nightly, torch_device
from diffusers.utils.testing_utils import require_torch_gpu



class 	__lowerCamelCase			(unittest.TestCase						):







       def 		snake_case_	(			self: str					):






            '''simple docstring'''




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







       @property
       def 		snake_case_	(			self: Optional[int]					):






            '''simple docstring'''




            __UpperCamelCase =			1
            __UpperCamelCase =			3
            __UpperCamelCase =			(32, 32)

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







       @property
       def 		snake_case_	(			self: str					):






            '''simple docstring'''




            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,)
            return model







       @property
       def 		snake_case_	(			self: int					):






            '''simple docstring'''




            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,)
            return model







       @property
       def 		snake_case_	(			self: Any					):






            '''simple docstring'''




            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,)
            return CLIPTextModel(A_					)







       @property
       def 		snake_case_	(			self: Dict					):






            '''simple docstring'''




            def extract(*A_: Any,**A_: Dict					):



                 class 	__lowerCamelCase			:







                        def __init__(			self: Optional[int]					):






                             '''simple docstring'''




                             __UpperCamelCase =			torch.ones([0]					)







                        def 		snake_case_	(			self: Optional[Any],A_: int					):






                             '''simple docstring'''




                             self.pixel_values.to(A_					)
                             return self



                 return Out()

            return extract







       def 		snake_case_	(			self: int					):






            '''simple docstring'''




            __UpperCamelCase =			'cpu'  # ensure determinism for the device-dependent torch.Generator
            __UpperCamelCase =			self.dummy_cond_unet
            __UpperCamelCase =			DDIMScheduler(
                beta_start=0.0_0_0_8_5,beta_end=0.0_1_2,beta_schedule='scaled_linear',clip_sample=A_,set_alpha_to_one=A_,)

            __UpperCamelCase =			self.dummy_vae
            __UpperCamelCase =			self.dummy_text_encoder
            __UpperCamelCase =			CLIPTokenizer.from_pretrained('hf-internal-testing/tiny-random-clip'					)

            # make sure here that pndm scheduler skips prk
            __UpperCamelCase =			StableDiffusionPipeline(
                unet=A_,scheduler=A_,vae=A_,text_encoder=A_,tokenizer=A_,safety_checker=A_,feature_extractor=self.dummy_extractor,)
            __UpperCamelCase =			sd_pipe.to(A_					)
            sd_pipe.set_progress_bar_config(disable=A_					)

            __UpperCamelCase =			'A painting of a squirrel eating a burger'

            __UpperCamelCase =			torch.Generator(device=A_					).manual_seed(0					)
            __UpperCamelCase =			sd_pipe([prompt],generator=A_,guidance_scale=6.0,num_inference_steps=2,output_type='np'					)
            __UpperCamelCase =			output.images

            __UpperCamelCase =			torch.Generator(device=A_					).manual_seed(0					)
            __UpperCamelCase =			sd_pipe(
                [prompt],generator=A_,guidance_scale=6.0,num_inference_steps=2,output_type='np',return_dict=A_,)[0]

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

            assert image.shape == (1, 64, 64, 3)
            __UpperCamelCase =			np.array([0.5_7_5_6, 0.6_1_1_8, 0.5_0_0_5, 0.5_0_4_1, 0.5_4_7_1, 0.4_7_2_6, 0.4_9_7_6, 0.4_8_6_5, 0.4_8_6_4]					)

            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 		snake_case_	(			self: Optional[Any]					):






            '''simple docstring'''




            __UpperCamelCase =			'cpu'  # ensure determinism for the device-dependent torch.Generator
            __UpperCamelCase =			self.dummy_cond_unet
            __UpperCamelCase =			PNDMScheduler(skip_prk_steps=A_					)
            __UpperCamelCase =			self.dummy_vae
            __UpperCamelCase =			self.dummy_text_encoder
            __UpperCamelCase =			CLIPTokenizer.from_pretrained('hf-internal-testing/tiny-random-clip'					)

            # make sure here that pndm scheduler skips prk
            __UpperCamelCase =			StableDiffusionPipeline(
                unet=A_,scheduler=A_,vae=A_,text_encoder=A_,tokenizer=A_,safety_checker=A_,feature_extractor=self.dummy_extractor,)
            __UpperCamelCase =			sd_pipe.to(A_					)
            sd_pipe.set_progress_bar_config(disable=A_					)

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

            __UpperCamelCase =			output.images

            __UpperCamelCase =			torch.Generator(device=A_					).manual_seed(0					)
            __UpperCamelCase =			sd_pipe(
                [prompt],generator=A_,guidance_scale=6.0,num_inference_steps=2,output_type='np',return_dict=A_,)[0]

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

            assert image.shape == (1, 64, 64, 3)
            __UpperCamelCase =			np.array([0.5_1_2_5, 0.5_7_1_6, 0.4_8_2_8, 0.5_0_6_0, 0.5_6_5_0, 0.4_7_6_8, 0.5_1_8_5, 0.4_8_9_5, 0.4_9_9_3]					)

            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 		snake_case_	(			self: Tuple					):






            '''simple docstring'''




            __UpperCamelCase =			StableDiffusionPipeline.from_pretrained(
                'hf-internal-testing/tiny-stable-diffusion-lms-pipe',safety_checker=A_					)
            assert isinstance(A_,A_					)
            assert isinstance(pipe.scheduler,A_					)
            assert pipe.safety_checker is None

            __UpperCamelCase =			pipe('example prompt',num_inference_steps=2					).images[0]
            assert image is not None

            # check that there's no error when saving a pipeline with one of the models being None
            with tempfile.TemporaryDirectory() as tmpdirname:
                 pipe.save_pretrained(A_					)
                 __UpperCamelCase =			StableDiffusionPipeline.from_pretrained(A_					)

            # sanity check that the pipeline still works
            assert pipe.safety_checker is None
            __UpperCamelCase =			pipe('example prompt',num_inference_steps=2					).images[0]
            assert image is not None







       @unittest.skipIf(torch_device != 'cuda','This test requires a GPU'					)
       def 		snake_case_	(			self: Tuple					):






            '''simple docstring'''




            __UpperCamelCase =			self.dummy_cond_unet
            __UpperCamelCase =			PNDMScheduler(skip_prk_steps=A_					)
            __UpperCamelCase =			self.dummy_vae
            __UpperCamelCase =			self.dummy_text_encoder
            __UpperCamelCase =			CLIPTokenizer.from_pretrained('hf-internal-testing/tiny-random-clip'					)

            # put models in fp16
            __UpperCamelCase =			unet.half()
            __UpperCamelCase =			vae.half()
            __UpperCamelCase =			bert.half()

            # make sure here that pndm scheduler skips prk
            __UpperCamelCase =			StableDiffusionPipeline(
                unet=A_,scheduler=A_,vae=A_,text_encoder=A_,tokenizer=A_,safety_checker=A_,feature_extractor=self.dummy_extractor,)
            __UpperCamelCase =			sd_pipe.to(A_					)
            sd_pipe.set_progress_bar_config(disable=A_					)

            __UpperCamelCase =			'A painting of a squirrel eating a burger'
            __UpperCamelCase =			sd_pipe([prompt],num_inference_steps=2,output_type='np'					).images

            assert image.shape == (1, 64, 64, 3)



@nightly
@require_torch_gpu
class 	__lowerCamelCase			(unittest.TestCase						):







       def 		snake_case_	(			self: List[Any]					):






            '''simple docstring'''




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







       def 		snake_case_	(			self: int					):






            '''simple docstring'''




            __UpperCamelCase =			StableDiffusionPipeline.from_pretrained('runwayml/stable-diffusion-v1-5',safety_checker=A_					)
            __UpperCamelCase =			LMSDiscreteScheduler.from_config(sd_pipe.scheduler.config					)
            __UpperCamelCase =			sd_pipe.to(A_					)
            sd_pipe.set_progress_bar_config(disable=A_					)

            __UpperCamelCase =			(
                'portrait of girl with smokey eyes makeup in abandoned hotel, grange clothes, redshift, wide high angle'
                ' coloured polaroid photograph with flash, kodak film, hyper real, stunning moody cinematography, with'
                ' anamorphic lenses, by maripol, fallen angels by wong kar - wai, style of suspiria and neon demon and'
                ' children from bahnhof zoo, detailed '
            )
            __UpperCamelCase =			40_0366_0346
            __UpperCamelCase =			7

            # without safety guidance (sld_guidance_scale = 0)
            __UpperCamelCase =			torch.manual_seed(A_					)
            __UpperCamelCase =			sd_pipe(
                [prompt],generator=A_,guidance_scale=A_,num_inference_steps=50,output_type='np',width=512,height=512,sld_guidance_scale=0,)

            __UpperCamelCase =			output.images
            __UpperCamelCase =			image[0, -3:, -3:, -1]
            __UpperCamelCase =			[0.2_2_7_8, 0.2_2_3_1, 0.2_2_4_9, 0.2_3_3_3, 0.2_3_0_3, 0.1_8_8_5, 0.2_2_7_3, 0.2_1_4_4, 0.2_1_7_6]

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

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

            # without safety guidance (strong configuration)
            __UpperCamelCase =			torch.manual_seed(A_					)
            __UpperCamelCase =			sd_pipe(
                [prompt],generator=A_,guidance_scale=A_,num_inference_steps=50,output_type='np',width=512,height=512,sld_guidance_scale=2000,sld_warmup_steps=7,sld_threshold=0.0_2_5,sld_momentum_scale=0.5,sld_mom_beta=0.7,)

            __UpperCamelCase =			output.images
            __UpperCamelCase =			image[0, -3:, -3:, -1]
            __UpperCamelCase =			[0.2_3_8_3, 0.2_2_7_6, 0.2_3_6, 0.2_1_9_2, 0.2_1_8_6, 0.2_0_5_3, 0.1_9_7_1, 0.1_9_0_1, 0.1_7_1_9]

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

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







       def 		snake_case_	(			self: List[Any]					):






            '''simple docstring'''




            __UpperCamelCase =			StableDiffusionPipeline.from_pretrained('runwayml/stable-diffusion-v1-5',safety_checker=A_					)
            __UpperCamelCase =			LMSDiscreteScheduler.from_config(sd_pipe.scheduler.config					)
            __UpperCamelCase =			sd_pipe.to(A_					)
            sd_pipe.set_progress_bar_config(disable=A_					)

            __UpperCamelCase =			'padme amidala taking a bath artwork, safe for work, no nudity'
            __UpperCamelCase =			27_3497_1755
            __UpperCamelCase =			7

            __UpperCamelCase =			torch.manual_seed(A_					)
            __UpperCamelCase =			sd_pipe(
                [prompt],generator=A_,guidance_scale=A_,num_inference_steps=50,output_type='np',width=512,height=512,sld_guidance_scale=0,)

            __UpperCamelCase =			output.images
            __UpperCamelCase =			image[0, -3:, -3:, -1]
            __UpperCamelCase =			[0.3_5_0_2, 0.3_6_2_2, 0.3_3_9_6, 0.3_6_4_2, 0.3_4_7_8, 0.3_3_1_8, 0.3_5, 0.3_3_4_8, 0.3_2_9_7]

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

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

            __UpperCamelCase =			torch.manual_seed(A_					)
            __UpperCamelCase =			sd_pipe(
                [prompt],generator=A_,guidance_scale=A_,num_inference_steps=50,output_type='np',width=512,height=512,sld_guidance_scale=2000,sld_warmup_steps=7,sld_threshold=0.0_2_5,sld_momentum_scale=0.5,sld_mom_beta=0.7,)

            __UpperCamelCase =			output.images
            __UpperCamelCase =			image[0, -3:, -3:, -1]
            __UpperCamelCase =			[0.5_5_3_1, 0.5_2_0_6, 0.4_8_9_5, 0.5_1_5_6, 0.5_1_8_2, 0.4_7_5_1, 0.4_8_0_2, 0.4_8_0_3, 0.4_4_4_3]

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

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







       def 		snake_case_	(			self: List[str]					):






            '''simple docstring'''




            __UpperCamelCase =			StableDiffusionPipeline.from_pretrained('runwayml/stable-diffusion-v1-5'					)
            __UpperCamelCase =			sd_pipe.to(A_					)
            sd_pipe.set_progress_bar_config(disable=A_					)

            __UpperCamelCase =			(
                'the four horsewomen of the apocalypse, painting by tom of finland, gaston bussiere, craig mullins, j. c.'
                ' leyendecker'
            )
            __UpperCamelCase =			10_4435_5234
            __UpperCamelCase =			12

            __UpperCamelCase =			torch.manual_seed(A_					)
            __UpperCamelCase =			sd_pipe(
                [prompt],generator=A_,guidance_scale=A_,num_inference_steps=50,output_type='np',width=512,height=512,sld_guidance_scale=0,)

            __UpperCamelCase =			output.images
            __UpperCamelCase =			image[0, -3:, -3:, -1]
            __UpperCamelCase =			np.array([0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]					)

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

            assert np.abs(image_slice.flatten() - expected_slice					).max() < 1E-7

            __UpperCamelCase =			torch.manual_seed(A_					)
            __UpperCamelCase =			sd_pipe(
                [prompt],generator=A_,guidance_scale=A_,num_inference_steps=50,output_type='np',width=512,height=512,sld_guidance_scale=2000,sld_warmup_steps=7,sld_threshold=0.0_2_5,sld_momentum_scale=0.5,sld_mom_beta=0.7,)

            __UpperCamelCase =			output.images
            __UpperCamelCase =			image[0, -3:, -3:, -1]
            __UpperCamelCase =			np.array([0.5_8_1_8, 0.6_2_8_5, 0.6_8_3_5, 0.6_0_1_9, 0.6_2_5, 0.6_7_5_4, 0.6_0_9_6, 0.6_3_3_4, 0.6_5_6_1]					)
            assert image.shape == (1, 512, 512, 3)

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






#
# This a `torch.distributed` diagnostics script that checks that all GPUs in the cluster (one or
# many nodes) can talk to each other via nccl and allocate gpu memory.
#
# To run first adjust the number of processes and nodes:
#
# python -m torch.distributed.run --nproc_per_node 2 --nnodes 1 torch-distributed-gpu-test.py
#
# You may need to add --master_addr $MASTER_ADDR --master_port $MASTER_PORT if using a custom addr:port
#
# You can also use the rdzv API: --rdzv_endpoint $MASTER_ADDR:$MASTER_PORT --rdzv_backend c10d
#
# use torch.distributed.launch instead of torch.distributed.run for torch < 1.9
#
# If you get a hanging in `barrier` calls you have some network issues, you may try to debug this with:
#
# NCCL_DEBUG=INFO python -m torch.distributed.run --nproc_per_node 2 --nnodes 1 torch-distributed-gpu-test.py
#
# which should tell you what's going on behind the scenes.
#
#
# This script can be run via `srun` in the SLURM environment as well. Here is a SLURM script that
# runs on 2 nodes of 4 gpus per node:
#
# #SBATCH --job-name=test-nodes        # name
# #SBATCH --nodes=2                    # nodes
# #SBATCH --ntasks-per-node=1          # crucial - only 1 task per dist per node!
# #SBATCH --cpus-per-task=10           # number of cores per tasks
# #SBATCH --gres=gpu:4                 # number of gpus
# #SBATCH --time 0:05:00               # maximum execution time (HH:MM:SS)
# #SBATCH --output=%x-%j.out           # output file name
#
# GPUS_PER_NODE=4
# MASTER_ADDR=$(scontrol show hostnames $SLURM_JOB_NODELIST | head -n 1)
# MASTER_PORT=6000
#
# srun --jobid $SLURM_JOBID bash -c 'python -m torch.distributed.run \
# --nproc_per_node $GPUS_PER_NODE --nnodes $SLURM_NNODES --node_rank $SLURM_PROCID \
# --master_addr $MASTER_ADDR --master_port $MASTER_PORT \
# torch-distributed-gpu-test.py'
#

import fcntl
import os
import socket

import torch
import torch.distributed as dist






def 		_A		(		*_lowercase )					-> Tuple:


     """simple docstring"""


     with open(_lowercase					,  'r' ) as fh:
          fcntl.flock(_lowercase					,  fcntl.LOCK_EX )
          try:
               print(*_lowercase )
          finally:
               fcntl.flock(_lowercase					,  fcntl.LOCK_UN )


__snake_case   							=						int(os.environ['''LOCAL_RANK'''])
torch.cuda.set_device(local_rank)
__snake_case   							=						torch.device('''cuda''', local_rank)
__snake_case   							=						socket.gethostname()

__snake_case   							=						f"""[{hostname}-{local_rank}]"""

try:
    # test distributed
    dist.init_process_group('''nccl''')
    dist.all_reduce(torch.ones(1).to(device), op=dist.ReduceOp.SUM)
    dist.barrier()

    # test cuda is available and can allocate memory
    torch.cuda.is_available()
    torch.ones(1).cuda(local_rank)

    # global rank
    __snake_case   							=						dist.get_rank()
    __snake_case   							=						dist.get_world_size()

    printflock(f"""{gpu} is OK (global rank: {rank}/{world_size})""")

    dist.barrier()
    if rank == 0:
        printflock(f"""pt={torch.__version__}, cuda={torch.version.cuda}, nccl={torch.cuda.nccl.version()}""")

except Exception:
    printflock(f"""{gpu} is broken""")
    raise


