

import numpy as np
import torch
import torch.nn as nn
from transformers import CLIPConfig, CLIPVisionModelWithProjection, PreTrainedModel

from ...utils import logging


lowerCamelCase      				=     logging.get_logger(__name__)





class     snake_case_			(				_a							):


							"""simple docstring"""





							__UpperCAmelCase       =CLIPConfig

							__UpperCAmelCase       =['CLIPEncoderLayer']




							def __init__( self					,						_A							):
												super().__init__(_A							)

												__lowerCAmelCase					=							CLIPVisionModelWithProjection(config.vision_config							)

												__lowerCAmelCase					=							nn.Linear(config.vision_config.projection_dim					,						1							)
												__lowerCAmelCase					=							nn.Linear(config.vision_config.projection_dim					,						1							)




							@torch.no_grad()
							def        A__  ( self					,						_A					,						_A					,						_A=0.5					,						_A=0.5							):
												__lowerCAmelCase					=							self.vision_model(_A							)[0]

												__lowerCAmelCase					=							self.p_head(_A							)
												__lowerCAmelCase					=							nsfw_detected.flatten()
												__lowerCAmelCase					=							nsfw_detected > p_threshold
												__lowerCAmelCase					=							nsfw_detected.tolist()

												if any(_A							):
																	logger.warning(
																	    'Potential NSFW content was detected in one or more images. A black image will be returned instead.'
																	    ' Try again with a different prompt and/or seed.'							)

												for idx, nsfw_detected_ in enumerate(_A							):
																	if nsfw_detected_:
																						__lowerCAmelCase					=							np.zeros(images[idx].shape							)

												__lowerCAmelCase					=							self.w_head(_A							)
												__lowerCAmelCase					=							watermark_detected.flatten()
												__lowerCAmelCase					=							watermark_detected > w_threshold
												__lowerCAmelCase					=							watermark_detected.tolist()

												if any(_A							):
																	logger.warning(
																	    'Potential watermarked content was detected in one or more images. A black image will be returned instead.'
																	    ' Try again with a different prompt and/or seed.'							)

												for idx, watermark_detected_ in enumerate(_A							):
																	if watermark_detected_:
																						__lowerCAmelCase					=							np.zeros(images[idx].shape							)

												return images, nsfw_detected, watermark_detected









from dataclasses import dataclass
from typing import Dict, Optional, Tuple, Union

import torch
import torch.nn as nn

from ..configuration_utils import ConfigMixin, register_to_config
from ..utils import BaseOutput, apply_forward_hook
from .attention_processor import AttentionProcessor, AttnProcessor
from .modeling_utils import ModelMixin
from .vae import Decoder, DecoderOutput, DiagonalGaussianDistribution, Encoder





@dataclass
class     snake_case_			(				_a							):


							"""simple docstring"""





							__UpperCAmelCase       =42





class     snake_case_			(				_a			, _a							):


							"""simple docstring"""





							__UpperCAmelCase       =True




							@register_to_config
							def __init__( self					,						_A = 3					,						_A = 3					,						_A = ("DownEncoderBlock2D",)					,						_A = ("UpDecoderBlock2D",)					,						_A = (6_4,)					,						_A = 1					,						_A = "silu"					,						_A = 4					,						_A = 3_2					,						_A = 3_2					,						_A = 0.1_8215					,						):
												super().__init__()

												# pass init params to Encoder
												__lowerCAmelCase					=							Encoder(
												    in_channels=_A					,						out_channels=_A					,						down_block_types=_A					,						block_out_channels=_A					,						layers_per_block=_A					,						act_fn=_A					,						norm_num_groups=_A					,						double_z=_A					,						)

												# pass init params to Decoder
												__lowerCAmelCase					=							Decoder(
												    in_channels=_A					,						out_channels=_A					,						up_block_types=_A					,						block_out_channels=_A					,						layers_per_block=_A					,						norm_num_groups=_A					,						act_fn=_A					,						)

												__lowerCAmelCase					=							nn.Convad(2 * latent_channels					,						2 * latent_channels					,						1							)
												__lowerCAmelCase					=							nn.Convad(_A					,						_A					,						1							)

												__lowerCAmelCase					=							False
												__lowerCAmelCase					=							False

												# only relevant if vae tiling is enabled
												__lowerCAmelCase					=							self.config.sample_size
												__lowerCAmelCase					=							(
												    self.config.sample_size[0]
												    if isinstance(self.config.sample_size					,						(list, tuple)							)
												    else self.config.sample_size
												)
												__lowerCAmelCase					=							int(sample_size / (2 ** (len(self.config.block_out_channels							) - 1))							)
												__lowerCAmelCase					=							0.25




							def        A__  ( self					,						_A					,						_A=False							):
												if isinstance(_A					,						(Encoder, Decoder)							):
																	__lowerCAmelCase					=							value




							def        A__  ( self					,						_A = True							):
												__lowerCAmelCase					=							use_tiling




							def        A__  ( self							):
												self.enable_tiling(_A							)




							def        A__  ( self							):
												__lowerCAmelCase					=							True




							def        A__  ( self							):
												__lowerCAmelCase					=							False




							@property
							# Copied from diffusers.models.unet_2d_condition.UNet2DConditionModel.attn_processors
							def        A__  ( self							):
												__lowerCAmelCase					=							{}

												def fn_recursive_add_processors(_A					,						_A					,						_A							):
																	if hasattr(_A					,						'set_processor'							):
																						__lowerCAmelCase					=							module.processor

																	for sub_name, child in module.named_children():
																						fn_recursive_add_processors(F"""{name}.{sub_name}"""					,						_A					,						_A							)

																	return processors

												for name, module in self.named_children():
																	fn_recursive_add_processors(_A					,						_A					,						_A							)

												return processors




							def        A__  ( self					,						_A							):
												__lowerCAmelCase					=							len(self.attn_processors.keys()							)

												if isinstance(_A					,						_A							) and len(_A							) != count:
																	raise ValueError(
																	    F"""A dict of processors was passed, but the number of processors {len(_A							)} does not match the"""
																	    F""" number of attention layers: {count}. Please make sure to pass {count} processor classes."""							)

												def fn_recursive_attn_processor(_A					,						_A					,						_A							):
																	if hasattr(_A					,						'set_processor'							):
																						if not isinstance(_A					,						_A							):
																											module.set_processor(_A							)
																						else:
																											module.set_processor(processor.pop(F"""{name}.processor"""							)							)

																	for sub_name, child in module.named_children():
																						fn_recursive_attn_processor(F"""{name}.{sub_name}"""					,						_A					,						_A							)

												for name, module in self.named_children():
																	fn_recursive_attn_processor(_A					,						_A					,						_A							)




							def        A__  ( self							):
												self.set_attn_processor(AttnProcessor()							)




							@apply_forward_hook
							def        A__  ( self					,						_A					,						_A = True							):
												if self.use_tiling and (x.shape[-1] > self.tile_sample_min_size or x.shape[-2] > self.tile_sample_min_size):
																	return self.tiled_encode(_A					,						return_dict=_A							)

												if self.use_slicing and x.shape[0] > 1:
																	__lowerCAmelCase					=							[self.encoder(_A							) for x_slice in x.split(1							)]
																	__lowerCAmelCase					=							torch.cat(_A							)
												else:
																	__lowerCAmelCase					=							self.encoder(_A							)

												__lowerCAmelCase					=							self.quant_conv(_A							)
												__lowerCAmelCase					=							DiagonalGaussianDistribution(_A							)

												if not return_dict:
																	return (posterior,)

												return AutoencoderKLOutput(latent_dist=_A							)




							def        A__  ( self					,						_A					,						_A = True							):
												if self.use_tiling and (z.shape[-1] > self.tile_latent_min_size or z.shape[-2] > self.tile_latent_min_size):
																	return self.tiled_decode(_A					,						return_dict=_A							)

												__lowerCAmelCase					=							self.post_quant_conv(_A							)
												__lowerCAmelCase					=							self.decoder(_A							)

												if not return_dict:
																	return (dec,)

												return DecoderOutput(sample=_A							)




							@apply_forward_hook
							def        A__  ( self					,						_A					,						_A = True							):
												if self.use_slicing and z.shape[0] > 1:
																	__lowerCAmelCase					=							[self._decode(_A							).sample for z_slice in z.split(1							)]
																	__lowerCAmelCase					=							torch.cat(_A							)
												else:
																	__lowerCAmelCase					=							self._decode(_A							).sample

												if not return_dict:
																	return (decoded,)

												return DecoderOutput(sample=_A							)




							def        A__  ( self					,						_A					,						_A					,						_A							):
												__lowerCAmelCase					=							min(a.shape[2]					,						b.shape[2]					,						_A							)
												for y in range(_A							):
																	__lowerCAmelCase					=							a[:, :, -blend_extent + y, :] * (1 - y / blend_extent) + b[:, :, y, :] * (y / blend_extent)
												return b




							def        A__  ( self					,						_A					,						_A					,						_A							):
												__lowerCAmelCase					=							min(a.shape[3]					,						b.shape[3]					,						_A							)
												for x in range(_A							):
																	__lowerCAmelCase					=							a[:, :, :, -blend_extent + x] * (1 - x / blend_extent) + b[:, :, :, x] * (x / blend_extent)
												return b




							def        A__  ( self					,						_A					,						_A = True							):
												__lowerCAmelCase					=							int(self.tile_sample_min_size * (1 - self.tile_overlap_factor)							)
												__lowerCAmelCase					=							int(self.tile_latent_min_size * self.tile_overlap_factor							)
												__lowerCAmelCase					=							self.tile_latent_min_size - blend_extent

												# Split the image into 512x512 tiles and encode them separately.
												__lowerCAmelCase					=							[]
												for i in range(0					,						x.shape[2]					,						_A							):
																	__lowerCAmelCase					=							[]
																	for j in range(0					,						x.shape[3]					,						_A							):
																						__lowerCAmelCase					=							x[:, :, i : i + self.tile_sample_min_size, j : j + self.tile_sample_min_size]
																						__lowerCAmelCase					=							self.encoder(_A							)
																						__lowerCAmelCase					=							self.quant_conv(_A							)
																						row.append(_A							)
																	rows.append(_A							)
												__lowerCAmelCase					=							[]
												for i, row in enumerate(_A							):
																	__lowerCAmelCase					=							[]
																	for j, tile in enumerate(_A							):
																						# blend the above tile and the left tile
																						# to the current tile and add the current tile to the result row
																						if i > 0:
																											__lowerCAmelCase					=							self.blend_v(rows[i - 1][j]					,						_A					,						_A							)
																						if j > 0:
																											__lowerCAmelCase					=							self.blend_h(row[j - 1]					,						_A					,						_A							)
																						result_row.append(tile[:, :, :row_limit, :row_limit]							)
																	result_rows.append(torch.cat(_A					,						dim=3							)							)

												__lowerCAmelCase					=							torch.cat(_A					,						dim=2							)
												__lowerCAmelCase					=							DiagonalGaussianDistribution(_A							)

												if not return_dict:
																	return (posterior,)

												return AutoencoderKLOutput(latent_dist=_A							)




							def        A__  ( self					,						_A					,						_A = True							):
												__lowerCAmelCase					=							int(self.tile_latent_min_size * (1 - self.tile_overlap_factor)							)
												__lowerCAmelCase					=							int(self.tile_sample_min_size * self.tile_overlap_factor							)
												__lowerCAmelCase					=							self.tile_sample_min_size - blend_extent

												# Split z into overlapping 64x64 tiles and decode them separately.
												# The tiles have an overlap to avoid seams between tiles.
												__lowerCAmelCase					=							[]
												for i in range(0					,						z.shape[2]					,						_A							):
																	__lowerCAmelCase					=							[]
																	for j in range(0					,						z.shape[3]					,						_A							):
																						__lowerCAmelCase					=							z[:, :, i : i + self.tile_latent_min_size, j : j + self.tile_latent_min_size]
																						__lowerCAmelCase					=							self.post_quant_conv(_A							)
																						__lowerCAmelCase					=							self.decoder(_A							)
																						row.append(_A							)
																	rows.append(_A							)
												__lowerCAmelCase					=							[]
												for i, row in enumerate(_A							):
																	__lowerCAmelCase					=							[]
																	for j, tile in enumerate(_A							):
																						# blend the above tile and the left tile
																						# to the current tile and add the current tile to the result row
																						if i > 0:
																											__lowerCAmelCase					=							self.blend_v(rows[i - 1][j]					,						_A					,						_A							)
																						if j > 0:
																											__lowerCAmelCase					=							self.blend_h(row[j - 1]					,						_A					,						_A							)
																						result_row.append(tile[:, :, :row_limit, :row_limit]							)
																	result_rows.append(torch.cat(_A					,						dim=3							)							)

												__lowerCAmelCase					=							torch.cat(_A					,						dim=2							)
												if not return_dict:
																	return (dec,)

												return DecoderOutput(sample=_A							)




							def        A__  ( self					,						_A					,						_A = False					,						_A = True					,						_A = None					,						):
												__lowerCAmelCase					=							sample
												__lowerCAmelCase					=							self.encode(_A							).latent_dist
												if sample_posterior:
																	__lowerCAmelCase					=							posterior.sample(generator=_A							)
												else:
																	__lowerCAmelCase					=							posterior.mode()
												__lowerCAmelCase					=							self.decode(_A							).sample

												if not return_dict:
																	return (dec,)

												return DecoderOutput(sample=_A							)






