





"""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       a							(   a__   , a__   , a__   , unittest.TestCase  ):
 snake_case__             =       StableDiffusionControlNetImgaImgPipeline
 snake_case__             =       TEXT_GUIDED_IMAGE_VARIATION_PARAMS - {'''height''', '''width'''}
 snake_case__             =       TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS
 snake_case__             =       IMAGE_TO_IMAGE_IMAGE_PARAMS.union({'''control_image'''}  )
 snake_case__             =       IMAGE_TO_IMAGE_IMAGE_PARAMS



 def    UpperCamelCase__  (     self     ):







  """simple docstring"""






  torch.manual_seed(0     )
  lowerCAmelCase    				= 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     )
  lowerCAmelCase    				= 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     )
  lowerCAmelCase    				= DDIMScheduler(
      beta_start=0.00_085				,       beta_end=0.012				,       beta_schedule='scaled_linear'				,       clip_sample=_snake_case				,       set_alpha_to_one=_snake_case				,       )
  torch.manual_seed(0     )
  lowerCAmelCase    				= 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     )
  lowerCAmelCase    				= 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=10_00				,       )
  lowerCAmelCase    				= CLIPTextModel(_snake_case     )
  lowerCAmelCase    				= CLIPTokenizer.from_pretrained('hf-internal-testing/tiny-random-clip'     )

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



 def    UpperCamelCase__  (     self				,       _snake_case				,       _snake_case=0     ):







  """simple docstring"""






  if str(_snake_case     ).startswith('mps'     ):
   lowerCAmelCase    				= torch.manual_seed(_snake_case     )
  else:
   lowerCAmelCase    				= torch.Generator(device=_snake_case     ).manual_seed(_snake_case     )

  lowerCAmelCase    				= 2
  lowerCAmelCase    				= randn_tensor(
      (1, 3, 32 * controlnet_embedder_scale_factor, 32 * controlnet_embedder_scale_factor)				,       generator=_snake_case				,       device=torch.device(_snake_case     )				,       )
  lowerCAmelCase    				= floats_tensor(control_image.shape				,       rng=random.Random(_snake_case     )     ).to(_snake_case     )
  lowerCAmelCase    				= image.cpu().permute(0				,       2				,       3				,       1     )[0]
  lowerCAmelCase    				= Image.fromarray(np.uinta(_snake_case     )     ).convert('RGB'     ).resize((64, 64)     )
  lowerCAmelCase    				= {
      '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    UpperCamelCase__  (     self     ):







  """simple docstring"""






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







  """simple docstring"""






  self._test_xformers_attention_forwardGenerator_pass(expected_max_diff=2E-3     )



 def    UpperCamelCase__  (     self     ):







  """simple docstring"""






  self._test_inference_batch_single_identical(expected_max_diff=2E-3     )







class       a							(   a__   , a__   , unittest.TestCase  ):
 snake_case__             =       StableDiffusionControlNetImgaImgPipeline
 snake_case__             =       TEXT_GUIDED_IMAGE_VARIATION_PARAMS - {'''height''', '''width'''}
 snake_case__             =       TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS
 snake_case__             =       frozenset([]  )  # TO_DO: add image_params once refactored VaeImageProcessor.preprocess



 def    UpperCamelCase__  (     self     ):







  """simple docstring"""






  torch.manual_seed(0     )
  lowerCAmelCase    				= 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(_snake_case     ):
   if isinstance(_snake_case				,       torch.nn.Convad     ):
    torch.nn.init.normal(m.weight     )
    m.bias.data.fill_(1.0     )

  lowerCAmelCase    				= 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(_snake_case     )

  torch.manual_seed(0     )
  lowerCAmelCase    				= 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(_snake_case     )

  torch.manual_seed(0     )
  lowerCAmelCase    				= DDIMScheduler(
      beta_start=0.00_085				,       beta_end=0.012				,       beta_schedule='scaled_linear'				,       clip_sample=_snake_case				,       set_alpha_to_one=_snake_case				,       )
  torch.manual_seed(0     )
  lowerCAmelCase    				= 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     )
  lowerCAmelCase    				= 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=10_00				,       )
  lowerCAmelCase    				= CLIPTextModel(_snake_case     )
  lowerCAmelCase    				= CLIPTokenizer.from_pretrained('hf-internal-testing/tiny-random-clip'     )

  lowerCAmelCase    				= MultiControlNetModel([controlneta, controlneta]     )

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



 def    UpperCamelCase__  (     self				,       _snake_case				,       _snake_case=0     ):







  """simple docstring"""






  if str(_snake_case     ).startswith('mps'     ):
   lowerCAmelCase    				= torch.manual_seed(_snake_case     )
  else:
   lowerCAmelCase    				= torch.Generator(device=_snake_case     ).manual_seed(_snake_case     )

  lowerCAmelCase    				= 2

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

  lowerCAmelCase    				= floats_tensor(control_image[0].shape				,       rng=random.Random(_snake_case     )     ).to(_snake_case     )
  lowerCAmelCase    				= image.cpu().permute(0				,       2				,       3				,       1     )[0]
  lowerCAmelCase    				= Image.fromarray(np.uinta(_snake_case     )     ).convert('RGB'     ).resize((64, 64)     )
  lowerCAmelCase    				= {
      '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    UpperCamelCase__  (     self     ):







  """simple docstring"""






  lowerCAmelCase    				= self.get_dummy_components()
  lowerCAmelCase    				= self.pipeline_class(**_snake_case     )
  pipe.to(_snake_case     )

  lowerCAmelCase    				= 10.0
  lowerCAmelCase    				= 4

  lowerCAmelCase    				= self.get_dummy_inputs(_snake_case     )
  lowerCAmelCase    				= steps
  lowerCAmelCase    				= scale
  lowerCAmelCase    				= pipe(**_snake_case     )[0]

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

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

  lowerCAmelCase    				= self.get_dummy_inputs(_snake_case     )
  lowerCAmelCase    				= steps
  lowerCAmelCase    				= scale
  lowerCAmelCase    				= pipe(**_snake_case				,       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    UpperCamelCase__  (     self     ):







  """simple docstring"""






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







  """simple docstring"""






  self._test_xformers_attention_forwardGenerator_pass(expected_max_diff=2E-3     )



 def    UpperCamelCase__  (     self     ):







  """simple docstring"""






  self._test_inference_batch_single_identical(expected_max_diff=2E-3     )



 def    UpperCamelCase__  (     self     ):







  """simple docstring"""






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







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



 def    UpperCamelCase__  (     self     ):







  """simple docstring"""






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



 def    UpperCamelCase__  (     self     ):







  """simple docstring"""






  lowerCAmelCase    				= ControlNetModel.from_pretrained('lllyasviel/sd-controlnet-canny'     )

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

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

  lowerCAmelCase    				= pipe(
      _snake_case				,       _snake_case				,       control_image=_snake_case				,       generator=_snake_case				,       output_type='np'				,       num_inference_steps=50				,       strength=0.6				,       )

  lowerCAmelCase    				= output.images[0]

  assert image.shape == (5_12, 5_12, 3)

  lowerCAmelCase    				= 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 argparse
import math
import os

import torch
from neural_compressor.utils.pytorch import load
from PIL import Image
from transformers import CLIPTextModel, CLIPTokenizer

from diffusers import AutoencoderKL, StableDiffusionPipeline, UNetaDConditionModel

def 		_SCREAMING_SNAKE_CASE			():
 lowerCAmelCase    				= argparse.ArgumentParser()
 parser.add_argument(
     '-m'  ,     '--pretrained_model_name_or_path'  ,     type=_UpperCAmelCase  ,     default=_UpperCAmelCase  ,     required=_UpperCAmelCase  ,     help='Path to pretrained model or model identifier from huggingface.co/models.'  ,     )
 parser.add_argument(
     '-c'  ,     '--caption'  ,     type=_UpperCAmelCase  ,     default='robotic cat with wings'  ,     help='Text used to generate images.'  ,     )
 parser.add_argument(
     '-n'  ,     '--images_num'  ,     type=_UpperCAmelCase  ,     default=4  ,     help='How much images to generate.'  ,     )
 parser.add_argument(
     '-s'  ,     '--seed'  ,     type=_UpperCAmelCase  ,     default=42  ,     help='Seed for random process.'  ,     )
 parser.add_argument(
     '-ci'  ,     '--cuda_id'  ,     type=_UpperCAmelCase  ,     default=0  ,     help='cuda_id.'  ,     )
 lowerCAmelCase    				= parser.parse_args()
 return args

def 		_SCREAMING_SNAKE_CASE			(_UpperCAmelCase					:    int  ,     _UpperCAmelCase					:    List[Any]  ,     _UpperCAmelCase					:    Optional[int]			):
 if not len(_UpperCAmelCase			) == rows * cols:
  raise ValueError('The specified number of rows and columns are not correct.'			)

 lowerCAmelCase				,lowerCAmelCase    				= imgs[0].size
 lowerCAmelCase    				= Image.new('RGB'  ,     size=(cols * w, rows * h)			)
 lowerCAmelCase				,lowerCAmelCase    				= grid.size

 for i, img in enumerate(_UpperCAmelCase			):
  grid.paste(_UpperCAmelCase  ,     box=(i % cols * w, i // cols * h)			)
 return grid

def 		_SCREAMING_SNAKE_CASE			(_UpperCAmelCase					:    Any  ,     _UpperCAmelCase					:    Union[str, Any]="robotic cat with wings"  ,     _UpperCAmelCase					:    Optional[int]=7.5  ,     _UpperCAmelCase					:    Dict=50  ,     _UpperCAmelCase					:    Tuple=1  ,     _UpperCAmelCase					:    int=42  ,     ):
 lowerCAmelCase    				= torch.Generator(pipeline.device			).manual_seed(_UpperCAmelCase			)
 lowerCAmelCase    				= pipeline(
     _UpperCAmelCase  ,     guidance_scale=_UpperCAmelCase  ,     num_inference_steps=_UpperCAmelCase  ,     generator=_UpperCAmelCase  ,     num_images_per_prompt=_UpperCAmelCase  ,     ).images
 lowerCAmelCase    				= int(math.sqrt(_UpperCAmelCase			)			)
 lowerCAmelCase    				= image_grid(_UpperCAmelCase  ,     rows=_rows  ,     cols=num_images_per_prompt // _rows			)
 return grid, images


__UpperCamelCase     :						Optional[Any]						       =  parse_args()
# Load models and create wrapper for stable diffusion
__UpperCamelCase     :						List[Any]						       =  CLIPTokenizer.from_pretrained(args.pretrained_model_name_or_path, subfolder='''tokenizer''')
__UpperCamelCase     :						str						       =  CLIPTextModel.from_pretrained(args.pretrained_model_name_or_path, subfolder='''text_encoder''')
__UpperCamelCase     :						Optional[int]						       =  AutoencoderKL.from_pretrained(args.pretrained_model_name_or_path, subfolder='''vae''')
__UpperCamelCase     :						List[str]						       =  UNetaDConditionModel.from_pretrained(args.pretrained_model_name_or_path, subfolder='''unet''')

__UpperCamelCase     :						Tuple						       =  StableDiffusionPipeline.from_pretrained(
    args.pretrained_model_name_or_path, text_encoder=text_encoder, vae=vae, unet=unet, tokenizer=tokenizer
)
__UpperCamelCase     :						Union[str, Any]						       =  lambda images, clip_input: (images, False)
if os.path.exists(os.path.join(args.pretrained_model_name_or_path, '''best_model.pt''')):
       __UpperCamelCase     :						Dict						       =  load(args.pretrained_model_name_or_path, model=unet)
       unet.eval()
       setattr(pipeline, '''unet''', unet)
else:
       __UpperCamelCase     :						Dict						       =  unet.to(torch.device('''cuda''', args.cuda_id))
__UpperCamelCase     :						Optional[Any]						       =  pipeline.to(unet.device)
__UpperCamelCase       ,__UpperCamelCase     :						List[Any]						       =  generate_images(pipeline, prompt=args.caption, num_images_per_prompt=args.images_num, seed=args.seed)
grid.save(os.path.join(args.pretrained_model_name_or_path, '''{}.png'''.format('''_'''.join(args.caption.split()))))
__UpperCamelCase     :						int						       =  os.path.join(args.pretrained_model_name_or_path, '''_'''.join(args.caption.split()))
os.makedirs(dirname, exist_ok=True)
for idx, image in enumerate(images):
       image.save(os.path.join(dirname, '''{}.png'''.format(idx + 1)))





