



from ...configuration_utils import PretrainedConfig
from ...utils import logging


a_						:     str	     =		logging.get_logger(__name__)

a_						:     int	     =		{
    'caidas/swin2sr-classicalsr-x2-64': (
        'https://huggingface.co/caidas/swin2sr-classicalsr-x2-64/resolve/main/config.json'
    ),
}





class    lowerCamelCase__  (						UpperCAmelCase_):
     """simple docstring"""



     _A       =			'swin2sr'

     _A       =			{
         'hidden_size': 'embed_dim',
         'num_attention_heads': 'num_heads',
         'num_hidden_layers': 'num_layers',
     }







     def __init__(self				,			__a=64				,			__a=1				,			__a=3				,			__a=1_80				,			__a=[6, 6, 6, 6, 6, 6]				,			__a=[6, 6, 6, 6, 6, 6]				,			__a=8				,			__a=2.0				,			__a=True				,			__a=0.0				,			__a=0.0				,			__a=0.1				,			__a="gelu"				,			__a=False				,			__a=0.02				,			__a=1E-5				,			__a=2				,			__a=1.0				,			__a="1conv"				,			__a="pixelshuffle"				,			**__a				,			):






           '''simple docstring'''



           super().__init__(**__a    )

           lowerCamelCase								=		image_size
           lowerCamelCase								=		patch_size
           lowerCamelCase								=		num_channels
           lowerCamelCase								=		embed_dim
           lowerCamelCase								=		depths
           lowerCamelCase								=		len(__a    )
           lowerCamelCase								=		num_heads
           lowerCamelCase								=		window_size
           lowerCamelCase								=		mlp_ratio
           lowerCamelCase								=		qkv_bias
           lowerCamelCase								=		hidden_dropout_prob
           lowerCamelCase								=		attention_probs_dropout_prob
           lowerCamelCase								=		drop_path_rate
           lowerCamelCase								=		hidden_act
           lowerCamelCase								=		use_absolute_embeddings
           lowerCamelCase								=		layer_norm_eps
           lowerCamelCase								=		initializer_range
           lowerCamelCase								=		upscale
           lowerCamelCase								=		img_range
           lowerCamelCase								=		resi_connection
           lowerCamelCase								=		upsampler




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

from ...utils import logging


a_						:     Tuple	     =		logging.get_logger(__name__)





class    lowerCamelCase__  (						UpperCAmelCase_):
     """simple docstring"""



     _A       =			CLIPConfig

     _A       =			['CLIPEncoderLayer']







     def __init__(self				,			__a    ):






           '''simple docstring'''



           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    ):






           '''simple docstring'''



           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