





# 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     __lowerCAmelCase	(   UpperCamelCase__						,    UpperCamelCase__						,    UpperCamelCase__						,    unittest.TestCase):
      _lowercase    : Dict            =      StableDiffusionControlNetImgaImgPipeline
      _lowercase    : Tuple            =      TEXT_GUIDED_IMAGE_VARIATION_PARAMS - {"""height""", """width"""}
      _lowercase    : Optional[Any]            =      TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS
      _lowercase    : str            =      IMAGE_TO_IMAGE_IMAGE_PARAMS.union({"""control_image"""})
      _lowercase    : Dict            =      IMAGE_TO_IMAGE_IMAGE_PARAMS







      def        _lowercase  (   self							)       ->						Dict:



        '''simple docstring'''




        torch.manual_seed(0							)
        a__     : Optional[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							)
        a__     : Optional[Any]  =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__     : List[str]  =DDIMScheduler(
            beta_start=0.0_00_85		,				beta_end=0.0_12		,				beta_schedule="scaled_linear"		,				clip_sample=lowerCAmelCase__		,				set_alpha_to_one=lowerCAmelCase__		,				)
        torch.manual_seed(0							)
        a__     : Optional[int]  =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__     : Optional[int]  =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__     : List[Any]  =CLIPTextModel(lowerCAmelCase__							)
        a__     : int  =CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip"							)

        a__     : Optional[Any]  ={
            "unet": unet,
            "controlnet": controlnet,
            "scheduler": scheduler,
            "vae": vae,
            "text_encoder": text_encoder,
            "tokenizer": tokenizer,
            "safety_checker": None,
            "feature_extractor": None,
        }
        return components







      def        _lowercase  (   self		,				lowerCAmelCase__		,				lowerCAmelCase__=0							)       ->						List[Any]:



        '''simple docstring'''




        if str(lowerCAmelCase__							).startswith("mps"							):
          a__     : Optional[Any]  =torch.manual_seed(lowerCAmelCase__							)
        else:
          a__     : Dict  =torch.Generator(device=lowerCAmelCase__							).manual_seed(lowerCAmelCase__							)

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



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



        '''simple docstring'''




        self._test_xformers_attention_forwardGenerator_pass(expected_max_diff=2E-3							)







      def        _lowercase  (   self							)       ->						Union[str, Any]:



        '''simple docstring'''




        self._test_inference_batch_single_identical(expected_max_diff=2E-3							)







class     __lowerCAmelCase	(   UpperCamelCase__						,    UpperCamelCase__						,    unittest.TestCase):
      _lowercase    : Any            =      StableDiffusionControlNetImgaImgPipeline
      _lowercase    : int            =      TEXT_GUIDED_IMAGE_VARIATION_PARAMS - {"""height""", """width"""}
      _lowercase    : Any            =      TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS
      _lowercase    : Optional[Any]            =      frozenset([])  # TO_DO: add image_params once refactored VaeImageProcessor.preprocess







      def        _lowercase  (   self							)       ->						Optional[int]:



        '''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(lowerCAmelCase__							):
          if isinstance(lowerCAmelCase__		,				torch.nn.Convad							):
            torch.nn.init.normal(m.weight							)
            m.bias.data.fill_(1.0							)

        a__     : Tuple  =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(lowerCAmelCase__							)

        torch.manual_seed(0							)
        a__     : List[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(lowerCAmelCase__							)

        torch.manual_seed(0							)
        a__     : List[Any]  =DDIMScheduler(
            beta_start=0.0_00_85		,				beta_end=0.0_12		,				beta_schedule="scaled_linear"		,				clip_sample=lowerCAmelCase__		,				set_alpha_to_one=lowerCAmelCase__		,				)
        torch.manual_seed(0							)
        a__     : Tuple  =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__     : 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__     : Union[str, Any]  =CLIPTextModel(lowerCAmelCase__							)
        a__     : int  =CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip"							)

        a__     : List[Any]  =MultiControlNetModel([controlneta, controlneta]							)

        a__     : List[Any]  ={
            "unet": unet,
            "controlnet": controlnet,
            "scheduler": scheduler,
            "vae": vae,
            "text_encoder": text_encoder,
            "tokenizer": tokenizer,
            "safety_checker": None,
            "feature_extractor": None,
        }
        return components







      def        _lowercase  (   self		,				lowerCAmelCase__		,				lowerCAmelCase__=0							)       ->						Optional[Any]:



        '''simple docstring'''




        if str(lowerCAmelCase__							).startswith("mps"							):
          a__     : Union[str, Any]  =torch.manual_seed(lowerCAmelCase__							)
        else:
          a__     : Optional[Any]  =torch.Generator(device=lowerCAmelCase__							).manual_seed(lowerCAmelCase__							)

        a__     : Optional[Any]  =2

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

        a__     : Optional[Any]  =floats_tensor(control_image[0].shape		,				rng=random.Random(lowerCAmelCase__							)							).to(lowerCAmelCase__							)
        a__     : Optional[int]  =image.cpu().permute(0		,				2		,				3		,				1							)[0]
        a__     : int  =Image.fromarray(np.uinta(lowerCAmelCase__							)							).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        _lowercase  (   self							)       ->						Union[str, Any]:



        '''simple docstring'''




        a__     : str  =self.get_dummy_components()
        a__     : List[Any]  =self.pipeline_class(**lowerCAmelCase__							)
        pipe.to(lowerCAmelCase__							)

        a__     : Union[str, Any]  =10.0
        a__     : Union[str, Any]  =4

        a__     : str  =self.get_dummy_inputs(lowerCAmelCase__							)
        a__     : List[str]  =steps
        a__     : Tuple  =scale
        a__     : Optional[int]  =pipe(**lowerCAmelCase__							)[0]

        a__     : Optional[int]  =self.get_dummy_inputs(lowerCAmelCase__							)
        a__     : str  =steps
        a__     : Optional[Any]  =scale
        a__     : Union[str, Any]  =pipe(**lowerCAmelCase__		,				control_guidance_start=0.1		,				control_guidance_end=0.2							)[0]

        a__     : List[Any]  =self.get_dummy_inputs(lowerCAmelCase__							)
        a__     : Optional[Any]  =steps
        a__     : int  =scale
        a__     : Tuple  =pipe(**lowerCAmelCase__		,				control_guidance_start=[0.1, 0.3]		,				control_guidance_end=[0.2, 0.7]							)[0]

        a__     : Dict  =self.get_dummy_inputs(lowerCAmelCase__							)
        a__     : Any  =steps
        a__     : List[Any]  =scale
        a__     : Union[str, Any]  =pipe(**lowerCAmelCase__		,				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        _lowercase  (   self							)       ->						Dict:



        '''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        _lowercase  (   self							)       ->						Union[str, Any]:



        '''simple docstring'''




        self._test_xformers_attention_forwardGenerator_pass(expected_max_diff=2E-3							)







      def        _lowercase  (   self							)       ->						Any:



        '''simple docstring'''




        self._test_inference_batch_single_identical(expected_max_diff=2E-3							)







      def        _lowercase  (   self							)       ->						Optional[int]:



        '''simple docstring'''




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







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







      def        _lowercase  (   self							)       ->						Union[str, Any]:



        '''simple docstring'''




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







      def        _lowercase  (   self							)       ->						int:



        '''simple docstring'''




        a__     : int  =ControlNetModel.from_pretrained("lllyasviel/sd-controlnet-canny"							)

        a__     : Any  =StableDiffusionControlNetImgaImgPipeline.from_pretrained(
            "runwayml/stable-diffusion-v1-5"		,				safety_checker=lowerCAmelCase__		,				controlnet=lowerCAmelCase__							)
        pipe.enable_model_cpu_offload()
        pipe.set_progress_bar_config(disable=lowerCAmelCase__							)

        a__     : Optional[int]  =torch.Generator(device="cpu"							).manual_seed(0							)
        a__     : int  ="evil space-punk bird"
        a__     : List[Any]  =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[str]  =load_image(
            "https://huggingface.co/lllyasviel/sd-controlnet-canny/resolve/main/images/bird.png"							).resize((5_1_2, 5_1_2)							)

        a__     : List[Any]  =pipe(
            lowerCAmelCase__		,				lowerCAmelCase__		,				control_image=lowerCAmelCase__		,				generator=lowerCAmelCase__		,				output_type="np"		,				num_inference_steps=5_0		,				strength=0.6		,				)

        a__     : Optional[int]  =output.images[0]

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

        a__     : List[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













import copy
from dataclasses import dataclass, field
from typing import ClassVar, Dict

from ..features import ClassLabel, Features, Image
from .base import TaskTemplate







@dataclass(frozen=UpperCamelCase__)
class     __lowerCAmelCase	(   UpperCamelCase__):
      _lowercase    : str            =      field(default="""image-classification"""						,    metadata={"""include_in_asdict_even_if_is_default""": True})
      _lowercase    : ClassVar[Features]            =      Features({"""image""": Image()})
      _lowercase    : ClassVar[Features]            =      Features({"""labels""": ClassLabel})
      _lowercase    : str            =      "image"
      _lowercase    : str            =      "labels"







      def        _lowercase  (   self		,				lowerCAmelCase__							)       ->						List[Any]:



        '''simple docstring'''




        if self.label_column not in features:
          raise ValueError(F'''Column {self.label_column} is not present in features.'''							)
        if not isinstance(features[self.label_column]		,				lowerCAmelCase__							):
          raise ValueError(F'''Column {self.label_column} is not a ClassLabel.'''							)
        a__     : int  =copy.deepcopy(self							)
        a__     : Dict  =self.label_schema.copy()
        a__     : int  =features[self.label_column]
        a__     : str  =label_schema
        return task_template







      @property
      def        _lowercase  (   self							)       ->						Dict[str, str]:



        '''simple docstring'''




        return {
            self.image_column: "image",
            self.label_column: "labels",
        }






