






"""simple docstring"""



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__      (							snake_case__,       snake_case__,       snake_case__,       unittest.TestCase       ):
			_UpperCAmelCase		:Dict             =			StableDiffusionControlNetImgaImgPipeline
			_UpperCAmelCase		:Optional[int]             =			TEXT_GUIDED_IMAGE_VARIATION_PARAMS - {"height", "width"}
			_UpperCAmelCase		:Tuple             =			TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS
			_UpperCAmelCase		:Tuple             =			IMAGE_TO_IMAGE_IMAGE_PARAMS.union({"control_image"}       )
			_UpperCAmelCase		:Union[str, Any]             =			IMAGE_TO_IMAGE_IMAGE_PARAMS


			def      UpperCAmelCase__						(				self					:  Dict ):
										torch.manual_seed(0 )
										lowerCamelCase_	:							Tuple									=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_	:							int									=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_	:							List[str]									=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_	:							str									=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_	:							Dict									=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       ,	)
										lowerCamelCase_	:							List[str]									=CLIPTextModel(snake_case__ )
										lowerCamelCase_	:							List[Any]									=CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip" )

										lowerCamelCase_	:							Dict									={
										    "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					:  Tuple       ,	snake_case__					:  Union[str, Any]       ,	snake_case__					:  List[str]=0 ):
										if str(snake_case__ ).startswith("mps" ):
																	lowerCamelCase_	:							Any									=torch.manual_seed(snake_case__ )
										else:
																	lowerCamelCase_	:							List[str]									=torch.Generator(device=snake_case__ ).manual_seed(snake_case__ )

										lowerCamelCase_	:							Any									=2
										lowerCamelCase_	:							Optional[Any]									=randn_tensor(
										    (1, 3, 32 * controlnet_embedder_scale_factor, 32 * controlnet_embedder_scale_factor)       ,	generator=snake_case__       ,	device=torch.device(snake_case__ )       ,	)
										lowerCamelCase_	:							str									=floats_tensor(control_image.shape       ,	rng=random.Random(snake_case__ ) ).to(snake_case__ )
										lowerCamelCase_	:							List[str]									=image.cpu().permute(0       ,	2       ,	3       ,	1 )[0]
										lowerCamelCase_	:							Optional[int]									=Image.fromarray(np.uinta(snake_case__ ) ).convert("RGB" ).resize((64, 64) )
										lowerCamelCase_	:							Optional[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      UpperCAmelCase__						(				self					:  Dict ):
										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					:  Tuple ):
										self._test_xformers_attention_forwardGenerator_pass(expected_max_diff=2E-3 )


			def      UpperCAmelCase__						(				self					:  str ):
										self._test_inference_batch_single_identical(expected_max_diff=2E-3 )







class        lowercase__      (							snake_case__,       snake_case__,       unittest.TestCase       ):
			_UpperCAmelCase		:List[str]             =			StableDiffusionControlNetImgaImgPipeline
			_UpperCAmelCase		:str             =			TEXT_GUIDED_IMAGE_VARIATION_PARAMS - {"height", "width"}
			_UpperCAmelCase		:Dict             =			TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS
			_UpperCAmelCase		:Optional[int]             =			frozenset([]       )  # TO_DO: add image_params once refactored VaeImageProcessor.preprocess


			def      UpperCAmelCase__						(				self					:  str ):
										torch.manual_seed(0 )
										lowerCamelCase_	:							Optional[int]									=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__					:  Optional[int] ):
																	if isinstance(snake_case__       ,	torch.nn.Convad ):
																								torch.nn.init.normal(m.weight )
																								m.bias.data.fill_(1.0 )

										lowerCamelCase_	:							List[str]									=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_	:							Any									=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_	:							Optional[int]									=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_	:							Any									=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_	:							List[str]									=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       ,	)
										lowerCamelCase_	:							Tuple									=CLIPTextModel(snake_case__ )
										lowerCamelCase_	:							str									=CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip" )

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

										lowerCamelCase_	:							str									={
										    "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					:  int       ,	snake_case__					:  Dict       ,	snake_case__					:  Optional[int]=0 ):
										if str(snake_case__ ).startswith("mps" ):
																	lowerCamelCase_	:							Optional[int]									=torch.manual_seed(snake_case__ )
										else:
																	lowerCamelCase_	:							Optional[int]									=torch.Generator(device=snake_case__ ).manual_seed(snake_case__ )

										lowerCamelCase_	:							List[str]									=2

										lowerCamelCase_	:							Any									=[
										    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_	:							Any									=floats_tensor(control_image[0].shape       ,	rng=random.Random(snake_case__ ) ).to(snake_case__ )
										lowerCamelCase_	:							int									=image.cpu().permute(0       ,	2       ,	3       ,	1 )[0]
										lowerCamelCase_	:							Optional[Any]									=Image.fromarray(np.uinta(snake_case__ ) ).convert("RGB" ).resize((64, 64) )
										lowerCamelCase_	:							Dict									={
										    "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					:  Union[str, Any] ):
										lowerCamelCase_	:							str									=self.get_dummy_components()
										lowerCamelCase_	:							Any									=self.pipeline_class(**snake_case__ )
										pipe.to(snake_case__ )

										lowerCamelCase_	:							int									=10.0
										lowerCamelCase_	:							Optional[int]									=4

										lowerCamelCase_	:							Optional[int]									=self.get_dummy_inputs(snake_case__ )
										lowerCamelCase_	:							str									=steps
										lowerCamelCase_	:							List[Any]									=scale
										lowerCamelCase_	:							Any									=pipe(**snake_case__ )[0]

										lowerCamelCase_	:							Tuple									=self.get_dummy_inputs(snake_case__ )
										lowerCamelCase_	:							Union[str, Any]									=steps
										lowerCamelCase_	:							Union[str, Any]									=scale
										lowerCamelCase_	:							Optional[Any]									=pipe(**snake_case__       ,	control_guidance_start=0.1       ,	control_guidance_end=0.2 )[0]

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

										lowerCamelCase_	:							str									=self.get_dummy_inputs(snake_case__ )
										lowerCamelCase_	:							Optional[int]									=steps
										lowerCamelCase_	:							Union[str, Any]									=scale
										lowerCamelCase_	:							List[str]									=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					:  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      UpperCAmelCase__						(				self					:  Optional[Any] ):
										self._test_xformers_attention_forwardGenerator_pass(expected_max_diff=2E-3 )


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


			def      UpperCAmelCase__						(				self					:  Optional[int] ):
										lowerCamelCase_	:							List[str]									=self.get_dummy_components()
										lowerCamelCase_	:							str									=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        lowercase__      (							unittest.TestCase       ):


			def      UpperCAmelCase__						(				self					:  int ):
										super().tearDown()
										gc.collect()
										torch.cuda.empty_cache()


			def      UpperCAmelCase__						(				self					:  Union[str, Any] ):
										lowerCamelCase_	:							List[str]									=ControlNetModel.from_pretrained("lllyasviel/sd-controlnet-canny" )

										lowerCamelCase_	:							Union[str, Any]									=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_	:							Union[str, Any]									=torch.Generator(device="cpu" ).manual_seed(0 )
										lowerCamelCase_	:							Dict									="evil space-punk bird"
										lowerCamelCase_	:							Any									=load_image(
										    "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/sd_controlnet/bird_canny.png" ).resize((512, 512) )
										lowerCamelCase_	:							List[str]									=load_image(
										    "https://huggingface.co/lllyasviel/sd-controlnet-canny/resolve/main/images/bird.png" ).resize((512, 512) )

										lowerCamelCase_	:							Tuple									=pipe(
										    snake_case__       ,	snake_case__       ,	control_image=snake_case__       ,	generator=snake_case__       ,	output_type="np"       ,	num_inference_steps=50       ,	strength=0.6       ,	)

										lowerCamelCase_	:							str									=output.images[0]

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

										lowerCamelCase_	:							str									=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 itertools
from dataclasses import dataclass
from typing import List, Optional

import pyarrow as pa
import pyarrow.parquet as pq

import datasets
from datasets.table import table_cast


A__	:    List[str]			 = datasets.utils.logging.get_logger(__name__)







@dataclass
class        lowercase__      (							datasets.BuilderConfig       ):
			_UpperCAmelCase		:int             =			10000
			_UpperCAmelCase		:Optional[List[str]]             =			None
			_UpperCAmelCase		:Optional[datasets.Features]             =			None







class        lowercase__      (							datasets.ArrowBasedBuilder       ):
			_UpperCAmelCase		:Optional[int]             =			ParquetConfig


			def      UpperCAmelCase__						(				self					:  Optional[int] ):
										return datasets.DatasetInfo(features=self.config.features )


			def      UpperCAmelCase__						(				self					:  List[Any]       ,	snake_case__					:  List[Any] ):
										if not self.config.data_files:
																	raise ValueError(F"""At least one data file must be specified, but got data_files={self.config.data_files}""" )
										lowerCamelCase_	:							str									=dl_manager.download_and_extract(self.config.data_files )
										if isinstance(snake_case__       ,	(str, list, tuple) ):
																	lowerCamelCase_	:							Dict									=data_files
																	if isinstance(snake_case__       ,	snake_case__ ):
																								lowerCamelCase_	:							List[Any]									=[files]
																	# Use `dl_manager.iter_files` to skip hidden files in an extracted archive
																	lowerCamelCase_	:							Any									=[dl_manager.iter_files(snake_case__ ) for file in files]
																	return [datasets.SplitGenerator(name=datasets.Split.TRAIN       ,	gen_kwargs={"files": files} )]
										lowerCamelCase_	:							Optional[int]									=[]
										for split_name, files in data_files.items():
																	if isinstance(snake_case__       ,	snake_case__ ):
																								lowerCamelCase_	:							Optional[int]									=[files]
																	# Use `dl_manager.iter_files` to skip hidden files in an extracted archive
																	lowerCamelCase_	:							int									=[dl_manager.iter_files(snake_case__ ) for file in files]
																	# Infer features is they are stoed in the arrow schema
																	if self.info.features is None:
																								for file in itertools.chain.from_iterable(snake_case__ ):
																															with open(snake_case__       ,	"rb" ) as f:
																																						lowerCamelCase_	:							List[Any]									=datasets.Features.from_arrow_schema(pq.read_schema(snake_case__ ) )
																															break
																	splits.append(datasets.SplitGenerator(name=snake_case__       ,	gen_kwargs={"files": files} ) )
										return splits


			def      UpperCAmelCase__						(				self					:  int       ,	snake_case__					:  pa.Table ):
										if self.info.features is not None:
																	# more expensive cast to support nested features with keys in a different order
																	# allows str <-> int/float or str to Audio for example
																	lowerCamelCase_	:							List[str]									=table_cast(snake_case__       ,	self.info.features.arrow_schema )
										return pa_table


			def      UpperCAmelCase__						(				self					:  int       ,	snake_case__					:  Optional[Any] ):
										lowerCamelCase_	:							Tuple									=self.info.features.arrow_schema if self.info.features is not None else None
										if self.info.features is not None and self.config.columns is not None:
																	if sorted(field.name for field in schema ) != sorted(self.config.columns ):
																								raise ValueError(
																								    F"""Tried to load parquet data with columns '{self.config.columns}' with mismatching features '{self.info.features}'""" )
										for file_idx, file in enumerate(itertools.chain.from_iterable(snake_case__ ) ):
																	with open(snake_case__       ,	"rb" ) as f:
																								lowerCamelCase_	:							List[str]									=pq.ParquetFile(snake_case__ )
																								try:
																															for batch_idx, record_batch in enumerate(
																															    parquet_file.iter_batches(batch_size=self.config.batch_size       ,	columns=self.config.columns ) ):
																																						lowerCamelCase_	:							Union[str, Any]									=pa.Table.from_batches([record_batch] )
																																						# Uncomment for debugging (will print the Arrow table size and elements)
																																						# logger.warning(f"pa_table: {pa_table} num rows: {pa_table.num_rows}")
																																						# logger.warning('\n'.join(str(pa_table.slice(i, 1).to_pydict()) for i in range(pa_table.num_rows)))
																																						yield F"""{file_idx}_{batch_idx}""", self._cast_table(snake_case__ )
																								except ValueError as e:
																															logger.error(F"""Failed to read file '{file}' with error {type(snake_case__ )}: {e}""" )
																															raise



