




"""simple docstring"""

# This model implementation is heavily inspired by https://github.com/haofanwang/ControlNet-for-Diffusers/

import gc
import random
import tempfile
import unittest

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

from diffusers import (
    AutoencoderKL,
    ControlNetModel,
    DDIMScheduler,
    StableDiffusionControlNetImgaImgPipeline,
    UNetaDConditionModel,
)
from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_controlnet import MultiControlNetModel
from diffusers.utils import floats_tensor, load_image, load_numpy, randn_tensor, slow, torch_device
from diffusers.utils.import_utils import is_xformers_available
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu

from ..pipeline_params import (
    IMAGE_TO_IMAGE_IMAGE_PARAMS,
    TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS,
    TEXT_GUIDED_IMAGE_VARIATION_PARAMS,
)
from ..test_pipelines_common import (
    PipelineKarrasSchedulerTesterMixin,
    PipelineLatentTesterMixin,
    PipelineTesterMixin,
)


enable_full_determinism()





class     lowercase_	(					__lowerCAmelCase      ,     __lowerCAmelCase      ,     __lowerCAmelCase      ,     unittest.TestCase							):



 '''simple docstring'''
 UpperCAmelCase   :					Optional[int]				  =     StableDiffusionControlNetImgaImgPipeline
 UpperCAmelCase   :					Union[str, Any]				  =     TEXT_GUIDED_IMAGE_VARIATION_PARAMS - {'''height''', '''width'''}
 UpperCAmelCase   :					int				  =     TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS
 UpperCAmelCase   :					Optional[int]				  =     IMAGE_TO_IMAGE_IMAGE_PARAMS.union({'''control_image'''}							)
 UpperCAmelCase   :					int				  =     IMAGE_TO_IMAGE_IMAGE_PARAMS

 def       lowerCAmelCase_							(			self	:	str		):
       torch.manual_seed(0		)
       _A        =       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   ,							)
       torch.manual_seed(0		)
       _A        =       ControlNetModel(
           block_out_channels=(32, 64)   ,							layers_per_block=2   ,							in_channels=4   ,							down_block_types=('DownBlock2D', 'CrossAttnDownBlock2D')   ,							cross_attention_dim=32   ,							conditioning_embedding_out_channels=(16, 32)   ,							)
       torch.manual_seed(0		)
       _A        =       DDIMScheduler(
           beta_start=0.0_0085   ,							beta_end=0.012   ,							beta_schedule='scaled_linear'   ,							clip_sample=_UpperCAmelCase   ,							set_alpha_to_one=_UpperCAmelCase   ,							)
       torch.manual_seed(0		)
       _A        =       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		)
       _A        =       CLIPTextConfig(
           bos_token_id=0   ,							eos_token_id=2   ,							hidden_size=32   ,							intermediate_size=37   ,							layer_norm_eps=1E-0_5   ,							num_attention_heads=4   ,							num_hidden_layers=5   ,							pad_token_id=1   ,							vocab_size=1_000   ,							)
       _A        =       CLIPTextModel(_UpperCAmelCase		)
       _A        =       CLIPTokenizer.from_pretrained('hf-internal-testing/tiny-random-clip'		)

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

 def       lowerCAmelCase_							(			self	:	Optional[int]   ,							_UpperCAmelCase	:	List[Any]   ,							_UpperCAmelCase	:	Optional[int]=0		):
       if str(_UpperCAmelCase		).startswith('mps'		):
             _A        =       torch.manual_seed(_UpperCAmelCase		)
       else:
             _A        =       torch.Generator(device=_UpperCAmelCase		).manual_seed(_UpperCAmelCase		)

       _A        =       2
       _A        =       randn_tensor(
           (1, 3, 32 * controlnet_embedder_scale_factor, 32 * controlnet_embedder_scale_factor)   ,							generator=_UpperCAmelCase   ,							device=torch.device(_UpperCAmelCase		)   ,							)
       _A        =       floats_tensor(control_image.shape   ,							rng=random.Random(_UpperCAmelCase		)		).to(_UpperCAmelCase		)
       _A        =       image.cpu().permute(0   ,							2   ,							3   ,							1		)[0]
       _A        =       Image.fromarray(np.uinta(_UpperCAmelCase		)		).convert('RGB'		).resize((64, 64)		)
       _A        =       {
           'prompt': 'A painting of a squirrel eating a burger',
           'generator': generator,
           'num_inference_steps': 2,
           'guidance_scale': 6.0,
           'output_type': 'numpy',
           'image': image,
           'control_image': control_image,
       }

       return inputs

 def       lowerCAmelCase_							(			self	:	int		):
       return self._test_attention_slicing_forward_pass(expected_max_diff=2E-3		)

 @unittest.skipIf(
     torch_device != 'cuda' or not is_xformers_available()   ,							reason='XFormers attention is only available with CUDA and `xformers` installed'   ,							)
 def       lowerCAmelCase_							(			self	:	Optional[int]		):
       self._test_xformers_attention_forwardGenerator_pass(expected_max_diff=2E-3		)



 def       lowerCAmelCase_							(			self	:	Optional[Any]		):
       self._test_inference_batch_single_identical(expected_max_diff=2E-3		)





class     lowercase_	(					__lowerCAmelCase      ,     __lowerCAmelCase      ,     unittest.TestCase							):



 '''simple docstring'''
 UpperCAmelCase   :					Dict				  =     StableDiffusionControlNetImgaImgPipeline
 UpperCAmelCase   :					Dict				  =     TEXT_GUIDED_IMAGE_VARIATION_PARAMS - {'''height''', '''width'''}
 UpperCAmelCase   :					List[str]				  =     TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS
 UpperCAmelCase   :					str				  =     frozenset([]							)  # TO_DO: add image_params once refactored VaeImageProcessor.preprocess

 def       lowerCAmelCase_							(			self	:	List[str]		):
       torch.manual_seed(0		)
       _A        =       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   ,							)
       torch.manual_seed(0		)

       def init_weights(_UpperCAmelCase	:	Union[str, Any]		):
             if isinstance(_UpperCAmelCase   ,							torch.nn.Convad		):
                   torch.nn.init.normal(m.weight		)
                   m.bias.data.fill_(1.0		)

       _A        =       ControlNetModel(
           block_out_channels=(32, 64)   ,							layers_per_block=2   ,							in_channels=4   ,							down_block_types=('DownBlock2D', 'CrossAttnDownBlock2D')   ,							cross_attention_dim=32   ,							conditioning_embedding_out_channels=(16, 32)   ,							)
       controlneta.controlnet_down_blocks.apply(_UpperCAmelCase		)

       torch.manual_seed(0		)
       _A        =       ControlNetModel(
           block_out_channels=(32, 64)   ,							layers_per_block=2   ,							in_channels=4   ,							down_block_types=('DownBlock2D', 'CrossAttnDownBlock2D')   ,							cross_attention_dim=32   ,							conditioning_embedding_out_channels=(16, 32)   ,							)
       controlneta.controlnet_down_blocks.apply(_UpperCAmelCase		)

       torch.manual_seed(0		)
       _A        =       DDIMScheduler(
           beta_start=0.0_0085   ,							beta_end=0.012   ,							beta_schedule='scaled_linear'   ,							clip_sample=_UpperCAmelCase   ,							set_alpha_to_one=_UpperCAmelCase   ,							)
       torch.manual_seed(0		)
       _A        =       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		)
       _A        =       CLIPTextConfig(
           bos_token_id=0   ,							eos_token_id=2   ,							hidden_size=32   ,							intermediate_size=37   ,							layer_norm_eps=1E-0_5   ,							num_attention_heads=4   ,							num_hidden_layers=5   ,							pad_token_id=1   ,							vocab_size=1_000   ,							)
       _A        =       CLIPTextModel(_UpperCAmelCase		)
       _A        =       CLIPTokenizer.from_pretrained('hf-internal-testing/tiny-random-clip'		)

       _A        =       MultiControlNetModel([controlneta, controlneta]		)

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

 def       lowerCAmelCase_							(			self	:	Dict   ,							_UpperCAmelCase	:	str   ,							_UpperCAmelCase	:	int=0		):
       if str(_UpperCAmelCase		).startswith('mps'		):
             _A        =       torch.manual_seed(_UpperCAmelCase		)
       else:
             _A        =       torch.Generator(device=_UpperCAmelCase		).manual_seed(_UpperCAmelCase		)

       _A        =       2

       _A        =       [
           randn_tensor(
               (1, 3, 32 * controlnet_embedder_scale_factor, 32 * controlnet_embedder_scale_factor)   ,							generator=_UpperCAmelCase   ,							device=torch.device(_UpperCAmelCase		)   ,							),
           randn_tensor(
               (1, 3, 32 * controlnet_embedder_scale_factor, 32 * controlnet_embedder_scale_factor)   ,							generator=_UpperCAmelCase   ,							device=torch.device(_UpperCAmelCase		)   ,							),
       ]

       _A        =       floats_tensor(control_image[0].shape   ,							rng=random.Random(_UpperCAmelCase		)		).to(_UpperCAmelCase		)
       _A        =       image.cpu().permute(0   ,							2   ,							3   ,							1		)[0]
       _A        =       Image.fromarray(np.uinta(_UpperCAmelCase		)		).convert('RGB'		).resize((64, 64)		)
       _A        =       {
           'prompt': 'A painting of a squirrel eating a burger',
           'generator': generator,
           'num_inference_steps': 2,
           'guidance_scale': 6.0,
           'output_type': 'numpy',
           'image': image,
           'control_image': control_image,
       }

       return inputs

 def       lowerCAmelCase_							(			self	:	Any		):
       _A        =       self.get_dummy_components()
       _A        =       self.pipeline_class(**_UpperCAmelCase		)
       pipe.to(_UpperCAmelCase		)

       _A        =       10.0
       _A        =       4

       _A        =       self.get_dummy_inputs(_UpperCAmelCase		)
       _A        =       steps
       _A        =       scale
       _A        =       pipe(**_UpperCAmelCase		)[0]

       _A        =       self.get_dummy_inputs(_UpperCAmelCase		)
       _A        =       steps
       _A        =       scale
       _A        =       pipe(**_UpperCAmelCase   ,							control_guidance_start=0.1   ,							control_guidance_end=0.2		)[0]

       _A        =       self.get_dummy_inputs(_UpperCAmelCase		)
       _A        =       steps
       _A        =       scale
       _A        =       pipe(**_UpperCAmelCase   ,							control_guidance_start=[0.1, 0.3]   ,							control_guidance_end=[0.2, 0.7]		)[0]

       _A        =       self.get_dummy_inputs(_UpperCAmelCase		)
       _A        =       steps
       _A        =       scale
       _A        =       pipe(**_UpperCAmelCase   ,							control_guidance_start=0.4   ,							control_guidance_end=[0.5, 0.8]		)[0]

       # make sure that all outputs are different
       assert np.sum(np.abs(output_a - output_a		)		) > 1E-3
       assert np.sum(np.abs(output_a - output_a		)		) > 1E-3
       assert np.sum(np.abs(output_a - output_a		)		) > 1E-3

 def       lowerCAmelCase_							(			self	:	List[str]		):
       return self._test_attention_slicing_forward_pass(expected_max_diff=2E-3		)

 @unittest.skipIf(
     torch_device != 'cuda' or not is_xformers_available()   ,							reason='XFormers attention is only available with CUDA and `xformers` installed'   ,							)
 def       lowerCAmelCase_							(			self	:	List[Any]		):
       self._test_xformers_attention_forwardGenerator_pass(expected_max_diff=2E-3		)

 def       lowerCAmelCase_							(			self	:	Union[str, Any]		):
       self._test_inference_batch_single_identical(expected_max_diff=2E-3		)



 def       lowerCAmelCase_							(			self	:	str		):
       _A        =       self.get_dummy_components()
       _A        =       self.pipeline_class(**_UpperCAmelCase		)
       pipe.to(_UpperCAmelCase		)
       pipe.set_progress_bar_config(disable=_UpperCAmelCase		)
       with tempfile.TemporaryDirectory() as tmpdir:
             try:
                   # save_pretrained is not implemented for Multi-ControlNet
                   pipe.save_pretrained(_UpperCAmelCase		)
             except NotImplementedError:
                   pass





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



 '''simple docstring'''

 def       lowerCAmelCase_							(			self	:	Tuple		):
       super().tearDown()
       gc.collect()
       torch.cuda.empty_cache()



 def       lowerCAmelCase_							(			self	:	List[str]		):
       _A        =       ControlNetModel.from_pretrained('lllyasviel/sd-controlnet-canny'		)

       _A        =       StableDiffusionControlNetImgaImgPipeline.from_pretrained(
           'runwayml/stable-diffusion-v1-5'   ,							safety_checker=_UpperCAmelCase   ,							controlnet=_UpperCAmelCase		)
       pipe.enable_model_cpu_offload()
       pipe.set_progress_bar_config(disable=_UpperCAmelCase		)

       _A        =       torch.Generator(device='cpu'		).manual_seed(0		)
       _A        =       'evil space-punk bird'
       _A        =       load_image(
           'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/sd_controlnet/bird_canny.png'		).resize((512, 512)		)
       _A        =       load_image(
           'https://huggingface.co/lllyasviel/sd-controlnet-canny/resolve/main/images/bird.png'		).resize((512, 512)		)

       _A        =       pipe(
           _UpperCAmelCase   ,							_UpperCAmelCase   ,							control_image=_UpperCAmelCase   ,							generator=_UpperCAmelCase   ,							output_type='np'   ,							num_inference_steps=50   ,							strength=0.6   ,							)

       _A        =       output.images[0]

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

       _A        =       load_numpy(
           'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/sd_controlnet/img2img.npy'		)

       assert np.abs(expected_image - image		).max() < 9E-2







"""simple docstring"""

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



 '''simple docstring'''
 UpperCAmelCase   :					str				  =     inspect.getfile(accelerate.test_utils							)
 UpperCAmelCase   :					Optional[Any]				  =     os.path.sep.join(mod_file.split(os.path.sep							)[:-1] + ['''scripts''', '''test_cli.py''']							)

 UpperCAmelCase   :					List[Any]				  =     ['''accelerate''', '''launch''']
 UpperCAmelCase   :					Dict				  =     Path.home() / '''.cache/huggingface/accelerate'''
 UpperCAmelCase   :					Union[str, Any]				  =     '''default_config.yaml'''
 UpperCAmelCase   :					Union[str, Any]				  =     config_folder / config_file
 UpperCAmelCase   :					Union[str, Any]				  =     config_folder / '''_default_config.yaml'''

 UpperCAmelCase   :					List[Any]				  =     Path('''tests/test_configs'''							)

 @classmethod
 def       lowerCAmelCase_							(			cls	:	List[Any]		):
       if cls.config_path.is_file():
             cls.config_path.rename(cls.changed_path		)

 @classmethod
 def       lowerCAmelCase_							(			cls	:	Tuple		):
       if cls.changed_path.is_file():
             cls.changed_path.rename(cls.config_path		)

 def       lowerCAmelCase_							(			self	:	List[Any]		):
       _A        =       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       lowerCAmelCase_							(			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       lowerCAmelCase_							(			self	:	Any		):
       execute_subprocess_async(['accelerate', 'test']   ,							env=os.environ.copy()		)





class     lowercase_	(					unittest.TestCase							):



 '''simple docstring'''
 UpperCAmelCase   :					Dict				  =     '''test-tpu'''
 UpperCAmelCase   :					Optional[int]				  =     '''us-central1-a'''
 UpperCAmelCase   :					List[str]				  =     '''ls'''
 UpperCAmelCase   :					str				  =     ['''accelerate''', '''tpu-config''']
 UpperCAmelCase   :					Optional[Any]				  =     '''cd /usr/share'''
 UpperCAmelCase   :					Optional[Any]				  =     '''tests/test_samples/test_command_file.sh'''
 UpperCAmelCase   :					str				  =     '''Running gcloud compute tpus tpu-vm ssh'''

 def       lowerCAmelCase_							(			self	:	Any		):
       _A        =       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       lowerCAmelCase_							(			self	:	Dict		):
       _A        =       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       lowerCAmelCase_							(			self	:	Optional[int]		):
       _A        =       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       lowerCAmelCase_							(			self	:	str		):
       _A        =       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       lowerCAmelCase_							(			self	:	List[str]		):
       _A        =       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       lowerCAmelCase_							(			self	:	str		):
       _A        =       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       lowerCAmelCase_							(			self	:	List[Any]		):
       _A        =       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       lowerCAmelCase_							(			self	:	int		):
       _A        =       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       lowerCAmelCase_							(			self	:	Optional[int]		):
       _A        =       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   ,							)

