





"""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  _UpperCAmelCase  (   lowerCamelCase_			,   lowerCamelCase_			,   lowerCamelCase_			,   unittest.TestCase		):



      '''simple docstring'''
      lowercase_      : List[Any]	    =   StableDiffusionControlNetImgaImgPipeline
      lowercase_      : Optional[int]	    =   TEXT_GUIDED_IMAGE_VARIATION_PARAMS - {"""height""", """width"""}
      lowercase_      : Optional[int]	    =   TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS
      lowercase_      : Dict	    =   IMAGE_TO_IMAGE_IMAGE_PARAMS.union({"""control_image"""}		)
      lowercase_      : Any	    =   IMAGE_TO_IMAGE_IMAGE_PARAMS

      def 					lowerCamelCase_			(     self  ):

        """simple docstring"""

        torch.manual_seed(0  )
        A_  :       Tuple										=  UNetaDConditionModel(
            block_out_channels=(3_2, 6_4)					,      layers_per_block=2					,      sample_size=3_2					,      in_channels=4					,      out_channels=4					,      down_block_types=('DownBlock2D', 'CrossAttnDownBlock2D')					,      up_block_types=('CrossAttnUpBlock2D', 'UpBlock2D')					,      cross_attention_dim=3_2					,      )
        torch.manual_seed(0  )
        A_  :       Dict										=  ControlNetModel(
            block_out_channels=(3_2, 6_4)					,      layers_per_block=2					,      in_channels=4					,      down_block_types=('DownBlock2D', 'CrossAttnDownBlock2D')					,      cross_attention_dim=3_2					,      conditioning_embedding_out_channels=(1_6, 3_2)					,      )
        torch.manual_seed(0  )
        A_  :       Dict										=  DDIMScheduler(
            beta_start=0.0_00_85					,      beta_end=0.0_12					,      beta_schedule='scaled_linear'					,      clip_sample=_lowercase					,      set_alpha_to_one=_lowercase					,      )
        torch.manual_seed(0  )
        A_  :       List[str]										=  AutoencoderKL(
            block_out_channels=[3_2, 6_4]					,      in_channels=3					,      out_channels=3					,      down_block_types=['DownEncoderBlock2D', 'DownEncoderBlock2D']					,      up_block_types=['UpDecoderBlock2D', 'UpDecoderBlock2D']					,      latent_channels=4					,      )
        torch.manual_seed(0  )
        A_  :       List[str]										=  CLIPTextConfig(
            bos_token_id=0					,      eos_token_id=2					,      hidden_size=3_2					,      intermediate_size=3_7					,      layer_norm_eps=1E-05					,      num_attention_heads=4					,      num_hidden_layers=5					,      pad_token_id=1					,      vocab_size=1_0_0_0					,      )
        A_  :       Tuple										=  CLIPTextModel(_lowercase  )
        A_  :       Optional[int]										=  CLIPTokenizer.from_pretrained('hf-internal-testing/tiny-random-clip'  )

        A_  :       List[str]										=  {
            '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					,      snake_case_					,      snake_case_=0  ):

        """simple docstring"""

        if str(_lowercase  ).startswith('mps'  ):
          A_  :       Optional[Any]										=  torch.manual_seed(_lowercase  )
        else:
          A_  :       Union[str, Any]										=  torch.Generator(device=_lowercase  ).manual_seed(_lowercase  )

        A_  :       Optional[Any]										=  2
        A_  :       Tuple										=  randn_tensor(
            (1, 3, 3_2 * controlnet_embedder_scale_factor, 3_2 * controlnet_embedder_scale_factor)					,      generator=_lowercase					,      device=torch.device(_lowercase  )					,      )
        A_  :       Union[str, Any]										=  floats_tensor(control_image.shape					,      rng=random.Random(_lowercase  )  ).to(_lowercase  )
        A_  :       Dict										=  image.cpu().permute(0					,      2					,      3					,      1  )[0]
        A_  :       List[Any]										=  Image.fromarray(np.uinta(_lowercase  )  ).convert('RGB'  ).resize((6_4, 6_4)  )
        A_  :       Any										=  {
            '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  ):

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

        """simple docstring"""

        self._test_xformers_attention_forwardGenerator_pass(expected_max_diff=2E-3  )

      def 					lowerCamelCase_			(     self  ):

        """simple docstring"""

        self._test_inference_batch_single_identical(expected_max_diff=2E-3  )


class  _UpperCAmelCase  (   lowerCamelCase_			,   lowerCamelCase_			,   unittest.TestCase		):



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

      def 					lowerCamelCase_			(     self  ):

        """simple docstring"""

        torch.manual_seed(0  )
        A_  :       Union[str, Any]										=  UNetaDConditionModel(
            block_out_channels=(3_2, 6_4)					,      layers_per_block=2					,      sample_size=3_2					,      in_channels=4					,      out_channels=4					,      down_block_types=('DownBlock2D', 'CrossAttnDownBlock2D')					,      up_block_types=('CrossAttnUpBlock2D', 'UpBlock2D')					,      cross_attention_dim=3_2					,      )
        torch.manual_seed(0  )

        def init_weights(snake_case_  ):
          if isinstance(_lowercase					,      torch.nn.Convad  ):
            torch.nn.init.normal(m.weight  )
            m.bias.data.fill_(1.0  )

        A_  :       Dict										=  ControlNetModel(
            block_out_channels=(3_2, 6_4)					,      layers_per_block=2					,      in_channels=4					,      down_block_types=('DownBlock2D', 'CrossAttnDownBlock2D')					,      cross_attention_dim=3_2					,      conditioning_embedding_out_channels=(1_6, 3_2)					,      )
        controlneta.controlnet_down_blocks.apply(_lowercase  )

        torch.manual_seed(0  )
        A_  :       str										=  ControlNetModel(
            block_out_channels=(3_2, 6_4)					,      layers_per_block=2					,      in_channels=4					,      down_block_types=('DownBlock2D', 'CrossAttnDownBlock2D')					,      cross_attention_dim=3_2					,      conditioning_embedding_out_channels=(1_6, 3_2)					,      )
        controlneta.controlnet_down_blocks.apply(_lowercase  )

        torch.manual_seed(0  )
        A_  :       Optional[Any]										=  DDIMScheduler(
            beta_start=0.0_00_85					,      beta_end=0.0_12					,      beta_schedule='scaled_linear'					,      clip_sample=_lowercase					,      set_alpha_to_one=_lowercase					,      )
        torch.manual_seed(0  )
        A_  :       Dict										=  AutoencoderKL(
            block_out_channels=[3_2, 6_4]					,      in_channels=3					,      out_channels=3					,      down_block_types=['DownEncoderBlock2D', 'DownEncoderBlock2D']					,      up_block_types=['UpDecoderBlock2D', 'UpDecoderBlock2D']					,      latent_channels=4					,      )
        torch.manual_seed(0  )
        A_  :       List[Any]										=  CLIPTextConfig(
            bos_token_id=0					,      eos_token_id=2					,      hidden_size=3_2					,      intermediate_size=3_7					,      layer_norm_eps=1E-05					,      num_attention_heads=4					,      num_hidden_layers=5					,      pad_token_id=1					,      vocab_size=1_0_0_0					,      )
        A_  :       Optional[Any]										=  CLIPTextModel(_lowercase  )
        A_  :       List[str]										=  CLIPTokenizer.from_pretrained('hf-internal-testing/tiny-random-clip'  )

        A_  :       List[str]										=  MultiControlNetModel([controlneta, controlneta]  )

        A_  :       List[str]										=  {
            '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					,      snake_case_					,      snake_case_=0  ):

        """simple docstring"""

        if str(_lowercase  ).startswith('mps'  ):
          A_  :       Optional[int]										=  torch.manual_seed(_lowercase  )
        else:
          A_  :       Optional[int]										=  torch.Generator(device=_lowercase  ).manual_seed(_lowercase  )

        A_  :       int										=  2

        A_  :       List[str]										=  [
            randn_tensor(
                (1, 3, 3_2 * controlnet_embedder_scale_factor, 3_2 * controlnet_embedder_scale_factor)					,      generator=_lowercase					,      device=torch.device(_lowercase  )					,      ),
            randn_tensor(
                (1, 3, 3_2 * controlnet_embedder_scale_factor, 3_2 * controlnet_embedder_scale_factor)					,      generator=_lowercase					,      device=torch.device(_lowercase  )					,      ),
        ]

        A_  :       Any										=  floats_tensor(control_image[0].shape					,      rng=random.Random(_lowercase  )  ).to(_lowercase  )
        A_  :       str										=  image.cpu().permute(0					,      2					,      3					,      1  )[0]
        A_  :       Optional[int]										=  Image.fromarray(np.uinta(_lowercase  )  ).convert('RGB'  ).resize((6_4, 6_4)  )
        A_  :       int										=  {
            '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  ):

        """simple docstring"""

        A_  :       str										=  self.get_dummy_components()
        A_  :       Dict										=  self.pipeline_class(**_lowercase  )
        pipe.to(_lowercase  )

        A_  :       int										=  10.0
        A_  :       List[Any]										=  4

        A_  :       List[Any]										=  self.get_dummy_inputs(_lowercase  )
        A_  :       Any										=  steps
        A_  :       Union[str, Any]										=  scale
        A_  :       Union[str, Any]										=  pipe(**_lowercase  )[0]

        A_  :       Dict										=  self.get_dummy_inputs(_lowercase  )
        A_  :       Optional[Any]										=  steps
        A_  :       Any										=  scale
        A_  :       List[str]										=  pipe(**_lowercase					,      control_guidance_start=0.1					,      control_guidance_end=0.2  )[0]

        A_  :       Optional[int]										=  self.get_dummy_inputs(_lowercase  )
        A_  :       int										=  steps
        A_  :       List[str]										=  scale
        A_  :       Dict										=  pipe(**_lowercase					,      control_guidance_start=[0.1, 0.3]					,      control_guidance_end=[0.2, 0.7]  )[0]

        A_  :       int										=  self.get_dummy_inputs(_lowercase  )
        A_  :       Optional[int]										=  steps
        A_  :       str										=  scale
        A_  :       Union[str, Any]										=  pipe(**_lowercase					,      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  ):

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

        """simple docstring"""

        self._test_xformers_attention_forwardGenerator_pass(expected_max_diff=2E-3  )

      def 					lowerCamelCase_			(     self  ):

        """simple docstring"""

        self._test_inference_batch_single_identical(expected_max_diff=2E-3  )

      def 					lowerCamelCase_			(     self  ):

        """simple docstring"""

        A_  :       Tuple										=  self.get_dummy_components()
        A_  :       Union[str, Any]										=  self.pipeline_class(**_lowercase  )
        pipe.to(_lowercase  )
        pipe.set_progress_bar_config(disable=_lowercase  )
        with tempfile.TemporaryDirectory() as tmpdir:
          try:
            # save_pretrained is not implemented for Multi-ControlNet
            pipe.save_pretrained(_lowercase  )
          except NotImplementedError:
            pass


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



      '''simple docstring'''

      def 					lowerCamelCase_			(     self  ):

        """simple docstring"""

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

      def 					lowerCamelCase_			(     self  ):

        """simple docstring"""

        A_  :       Optional[int]										=  ControlNetModel.from_pretrained('lllyasviel/sd-controlnet-canny'  )

        A_  :       List[Any]										=  StableDiffusionControlNetImgaImgPipeline.from_pretrained(
            'runwayml/stable-diffusion-v1-5'					,      safety_checker=_lowercase					,      controlnet=_lowercase  )
        pipe.enable_model_cpu_offload()
        pipe.set_progress_bar_config(disable=_lowercase  )

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

        A_  :       List[str]										=  pipe(
            _lowercase					,      _lowercase					,      control_image=_lowercase					,      generator=_lowercase					,      output_type='np'					,      num_inference_steps=5_0					,      strength=0.6					,      )

        A_  :       List[Any]										=  output.images[0]

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

        A_  :       Union[str, Any]										=  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
from datetime import datetime

import requests
from bsa import BeautifulSoup

if __name__ == "__main__":
   __lowerCamelCase			: str							   =		input('''Enter image url: ''').strip()
   print(F"""Downloading image from {url} ...""")
   __lowerCamelCase			: Any							   =		BeautifulSoup(requests.get(url).content, '''html.parser''')
   # The image URL is in the content field of the first meta tag with property og:image
   __lowerCamelCase			: List[Any]							   =		soup.find('''meta''', {'''property''': '''og:image'''})['''content''']
   __lowerCamelCase			: Tuple							   =		requests.get(image_url).content
   __lowerCamelCase			: Union[str, Any]							   =		F"""{datetime.now():%Y-%m-%d_%H:%M:%S}.jpg"""
   with open(file_name, '''wb''') as fp:
      fp.write(image_data)
   print(F"""Done. Image saved to disk as {file_name}.""")
