
'''simple docstring'''



import random
import unittest

import numpy as np
import torch
from transformers import CLIPTextConfig, CLIPTextModel, CLIPTokenizer

from diffusers import (
    AutoencoderKL,
    DDIMScheduler,
    UNetaDConditionModel,
    VideoToVideoSDPipeline,
)
from diffusers.utils import floats_tensor, is_xformers_available, skip_mps
from diffusers.utils.testing_utils import enable_full_determinism, slow, torch_device

from ..pipeline_params import (
    TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS,
    TEXT_GUIDED_IMAGE_VARIATION_PARAMS,
)
from ..test_pipelines_common import PipelineTesterMixin


enable_full_determinism()

@skip_mps
class  _a					(snake_case_	,       unittest.TestCase):





       """simple docstring"""




       SCREAMING_SNAKE_CASE   =	VideoToVideoSDPipeline
       SCREAMING_SNAKE_CASE   =	TEXT_GUIDED_IMAGE_VARIATION_PARAMS.union({'video'}) - {'image', 'width', 'height'}
       SCREAMING_SNAKE_CASE   =	TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS.union({'video'}) - {'image'}
       SCREAMING_SNAKE_CASE   =	PipelineTesterMixin.required_optional_params - {'latents'}
       SCREAMING_SNAKE_CASE   =	False

       # No `output_type`.
       SCREAMING_SNAKE_CASE   =	frozenset(
           [
               'num_inference_steps',
               'generator',
               'latents',
               'return_dict',
               'callback',
               'callback_steps',
           ])






       def 							UpperCamelCase					(				self   )     ->      Any:
          torch.manual_seed(0   )
          _SCREAMING_SNAKE_CASE  			=					UNetaDConditionModel(
              block_out_channels=(32, 64, 64, 64)							,				layers_per_block=2							,				sample_size=32							,				in_channels=4							,				out_channels=4							,				down_block_types=("""CrossAttnDownBlock3D""", """CrossAttnDownBlock3D""", """CrossAttnDownBlock3D""", """DownBlock3D""")							,				up_block_types=("""UpBlock3D""", """CrossAttnUpBlock3D""", """CrossAttnUpBlock3D""", """CrossAttnUpBlock3D""")							,				cross_attention_dim=32							,				attention_head_dim=4							,				)
          _SCREAMING_SNAKE_CASE  			=					DDIMScheduler(
              beta_start=0.0_0085							,				beta_end=0.012							,				beta_schedule="""scaled_linear"""							,				clip_sample=lowerCAmelCase__							,				set_alpha_to_one=lowerCAmelCase__							,				)
          torch.manual_seed(0   )
          _SCREAMING_SNAKE_CASE  			=					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							,				sample_size=1_28							,				)
          torch.manual_seed(0   )
          _SCREAMING_SNAKE_CASE  			=					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=10_00							,				hidden_act="""gelu"""							,				projection_dim=5_12							,				)
          _SCREAMING_SNAKE_CASE  			=					CLIPTextModel(lowerCAmelCase__   )
          _SCREAMING_SNAKE_CASE  			=					CLIPTokenizer.from_pretrained("""hf-internal-testing/tiny-random-clip"""   )

          _SCREAMING_SNAKE_CASE  			=					{
              '''unet''': unet,
              '''scheduler''': scheduler,
              '''vae''': vae,
              '''text_encoder''': text_encoder,
              '''tokenizer''': tokenizer,
          }
          return components






       def 							UpperCamelCase					(				self							,				A__							,				A__=0   )     ->      str:
          _SCREAMING_SNAKE_CASE  			=					floats_tensor((1, 3, 3, 32, 32)							,				rng=random.Random(lowerCAmelCase__   )   ).to(lowerCAmelCase__   )

          if str(lowerCAmelCase__   ).startswith("""mps"""   ):
             _SCREAMING_SNAKE_CASE  			=					torch.manual_seed(lowerCAmelCase__   )
          else:
             _SCREAMING_SNAKE_CASE  			=					torch.Generator(device=lowerCAmelCase__   ).manual_seed(lowerCAmelCase__   )
          _SCREAMING_SNAKE_CASE  			=					{
              '''prompt''': '''A painting of a squirrel eating a burger''',
              '''video''': video,
              '''generator''': generator,
              '''num_inference_steps''': 2,
              '''guidance_scale''': 6.0,
              '''output_type''': '''pt''',
          }
          return inputs






       def 							UpperCamelCase					(				self   )     ->      Dict:
          _SCREAMING_SNAKE_CASE  			=					'''cpu'''  # ensure determinism for the device-dependent torch.Generator
          _SCREAMING_SNAKE_CASE  			=					self.get_dummy_components()
          _SCREAMING_SNAKE_CASE  			=					VideoToVideoSDPipeline(**lowerCAmelCase__   )
          _SCREAMING_SNAKE_CASE  			=					sd_pipe.to(lowerCAmelCase__   )
          sd_pipe.set_progress_bar_config(disable=lowerCAmelCase__   )

          _SCREAMING_SNAKE_CASE  			=					self.get_dummy_inputs(lowerCAmelCase__   )
          _SCREAMING_SNAKE_CASE  			=					'''np'''
          _SCREAMING_SNAKE_CASE  			=					sd_pipe(**lowerCAmelCase__   ).frames
          _SCREAMING_SNAKE_CASE  			=					frames[0][-3:, -3:, -1]

          assert frames[0].shape == (32, 32, 3)
          _SCREAMING_SNAKE_CASE  			=					np.array([1_06, 1_17, 1_13, 1_74, 1_37, 1_12, 1_48, 1_51, 1_31]   )

          assert np.abs(image_slice.flatten() - expected_slice   ).max() < 1E-2






       @unittest.skipIf(
           torch_device != """cuda""" or not is_xformers_available()							,				reason="""XFormers attention is only available with CUDA and `xformers` installed"""							,				)
       def 							UpperCamelCase					(				self   )     ->      Union[str, Any]:
          self._test_xformers_attention_forwardGenerator_pass(test_mean_pixel_difference=lowerCAmelCase__							,				expected_max_diff=5E-3   )






       @unittest.skip(reason="""Batching needs to be properly figured out first for this pipeline."""   )
       def 							UpperCamelCase					(				self   )     ->      List[str]:
          pass






       @unittest.skip(reason="""Batching needs to be properly figured out first for this pipeline."""   )
       def 							UpperCamelCase					(				self   )     ->      Optional[Any]:
          pass






       @unittest.skip(reason="""`num_images_per_prompt` argument is not supported for this pipeline."""   )
       def 							UpperCamelCase					(				self   )     ->      List[Any]:
          pass






       def 							UpperCamelCase					(				self   )     ->      List[Any]:
          return super().test_progress_bar()






@slow
@skip_mps
class  _a					(unittest.TestCase):





       """simple docstring"""




       def 							UpperCamelCase					(				self   )     ->      Tuple:
          _SCREAMING_SNAKE_CASE  			=					VideoToVideoSDPipeline.from_pretrained("""cerspense/zeroscope_v2_XL"""							,				torch_dtype=torch.floataa   )
          pipe.enable_model_cpu_offload()

          # 10 frames
          _SCREAMING_SNAKE_CASE  			=					torch.Generator(device="""cpu"""   ).manual_seed(0   )
          _SCREAMING_SNAKE_CASE  			=					torch.randn((1, 10, 3, 10_24, 5_76)							,				generator=lowerCAmelCase__   )
          _SCREAMING_SNAKE_CASE  			=					video.to("""cuda"""   )

          _SCREAMING_SNAKE_CASE  			=					'''Spiderman is surfing'''

          _SCREAMING_SNAKE_CASE  			=					pipe(lowerCAmelCase__							,				video=lowerCAmelCase__							,				generator=lowerCAmelCase__							,				num_inference_steps=3							,				output_type="""pt"""   ).frames

          _SCREAMING_SNAKE_CASE  			=					np.array([-1.045_8984, -1.127_9297, -0.966_3086, -0.9150_3906, -0.7509_7656]   )
          assert np.abs(video_frames.cpu().numpy()[0, 0, 0, 0, -5:] - expected_array   ).sum() < 1E-2











'''simple docstring'''



import torch
import torch.nn as nn
from transformers import CLIPConfig, CLIPVisionModel, PreTrainedModel

from ...utils import logging


lowerCAmelCase_  :							int             =    logging.get_logger(__name__)
def        UpperCAmelCase (    A	:					List[Any]       ,		A	:					Optional[int]					):
							SCREAMING_SNAKE_CASE       :   Any         =      nn.functional.normalize(A					)
							SCREAMING_SNAKE_CASE       :   List[str]         =      nn.functional.normalize(A					)
							return torch.mm(A       ,		normalized_text_embeds.t()					)






class 						lowerCamelCase_				(       snake_case_      ):
						_lowerCAmelCase  : Optional[Any]						  =					CLIPConfig

						_lowerCAmelCase  : List[str]						  =					['CLIPEncoderLayer']







						def __init__(	self :   int					,     lowerCAmelCase__ :   CLIPConfig					):

													"""simple docstring"""
													super().__init__(lowerCAmelCase__					)

													SCREAMING_SNAKE_CASE       :   Any         =      CLIPVisionModel(config.vision_config					)
													SCREAMING_SNAKE_CASE       :   List[str]         =      nn.Linear(config.vision_config.hidden_size					,     config.projection_dim					,     bias=lowerCAmelCase__					)

													SCREAMING_SNAKE_CASE       :   Any         =      nn.Parameter(torch.ones(17					,     config.projection_dim					)					,     requires_grad=lowerCAmelCase__					)
													SCREAMING_SNAKE_CASE       :   int         =      nn.Parameter(torch.ones(3					,     config.projection_dim					)					,     requires_grad=lowerCAmelCase__					)

													SCREAMING_SNAKE_CASE       :   Dict         =      nn.Parameter(torch.ones(17					)					,     requires_grad=lowerCAmelCase__					)
													SCREAMING_SNAKE_CASE       :   List[str]         =      nn.Parameter(torch.ones(3					)					,     requires_grad=lowerCAmelCase__					)







						@torch.no_grad()
						def        __lowercase     (	self :   List[str]					,     lowerCAmelCase__ :   List[str]					,     lowerCAmelCase__ :   Optional[Any]					):

													"""simple docstring"""
													SCREAMING_SNAKE_CASE       :   Any         =      self.vision_model(lowerCAmelCase__					)[1]  # pooled_output
													SCREAMING_SNAKE_CASE       :   Any         =      self.visual_projection(lowerCAmelCase__					)

													# we always cast to float32 as this does not cause significant overhead and is compatible with bfloat16
													SCREAMING_SNAKE_CASE       :   Union[str, Any]         =      cosine_distance(lowerCAmelCase__					,     self.special_care_embeds					).cpu().float().numpy()
													SCREAMING_SNAKE_CASE       :   List[str]         =      cosine_distance(lowerCAmelCase__					,     self.concept_embeds					).cpu().float().numpy()

													SCREAMING_SNAKE_CASE       :   Union[str, Any]         =      []
													SCREAMING_SNAKE_CASE       :   int         =      image_embeds.shape[0]
													for i in range(lowerCAmelCase__					):
																				SCREAMING_SNAKE_CASE       :   Any         =      {'''special_scores''': {}, '''special_care''': [], '''concept_scores''': {}, '''bad_concepts''': []}

																				# increase this value to create a stronger `nfsw` filter
																				# at the cost of increasing the possibility of filtering benign images
																				SCREAMING_SNAKE_CASE       :   int         =      0.0

																				for concept_idx in range(len(special_cos_dist[0]					)					):
																											SCREAMING_SNAKE_CASE       :   Optional[int]         =      special_cos_dist[i][concept_idx]
																											SCREAMING_SNAKE_CASE       :   int         =      self.special_care_embeds_weights[concept_idx].item()
																											SCREAMING_SNAKE_CASE       :   Union[str, Any]         =      round(concept_cos - concept_threshold + adjustment					,     3					)
																											if result_img["special_scores"][concept_idx] > 0:
																																		result_img["special_care"].append({concept_idx, result_img['''special_scores'''][concept_idx]}					)
																																		SCREAMING_SNAKE_CASE       :   Tuple         =      0.01

																				for concept_idx in range(len(cos_dist[0]					)					):
																											SCREAMING_SNAKE_CASE       :   Optional[Any]         =      cos_dist[i][concept_idx]
																											SCREAMING_SNAKE_CASE       :   Union[str, Any]         =      self.concept_embeds_weights[concept_idx].item()
																											SCREAMING_SNAKE_CASE       :   Dict         =      round(concept_cos - concept_threshold + adjustment					,     3					)
																											if result_img["concept_scores"][concept_idx] > 0:
																																		result_img["bad_concepts"].append(lowerCAmelCase__					)

																				result.append(lowerCAmelCase__					)

													SCREAMING_SNAKE_CASE       :   Dict         =      [len(res['''bad_concepts''']					) > 0 for res in result]

													return images, has_nsfw_concepts







						@torch.no_grad()
						def        __lowercase     (	self :   List[str]					,     lowerCAmelCase__ :   torch.FloatTensor					,     lowerCAmelCase__ :   torch.FloatTensor					):

													"""simple docstring"""
													SCREAMING_SNAKE_CASE       :   Union[str, Any]         =      self.vision_model(lowerCAmelCase__					)[1]  # pooled_output
													SCREAMING_SNAKE_CASE       :   List[Any]         =      self.visual_projection(lowerCAmelCase__					)

													SCREAMING_SNAKE_CASE       :   List[str]         =      cosine_distance(lowerCAmelCase__					,     self.special_care_embeds					)
													SCREAMING_SNAKE_CASE       :   List[str]         =      cosine_distance(lowerCAmelCase__					,     self.concept_embeds					)

													# increase this value to create a stronger `nsfw` filter
													# at the cost of increasing the possibility of filtering benign images
													SCREAMING_SNAKE_CASE       :   Optional[int]         =      0.0

													SCREAMING_SNAKE_CASE       :   Any         =      special_cos_dist - self.special_care_embeds_weights + adjustment
													# special_scores = special_scores.round(decimals=3)
													SCREAMING_SNAKE_CASE       :   List[Any]         =      torch.any(special_scores > 0					,     dim=1					)
													SCREAMING_SNAKE_CASE       :   Dict         =      special_care * 0.01
													SCREAMING_SNAKE_CASE       :   List[str]         =      special_adjustment.unsqueeze(1					).expand(-1					,     cos_dist.shape[1]					)

													SCREAMING_SNAKE_CASE       :   Optional[int]         =      (cos_dist - self.concept_embeds_weights) + special_adjustment
													# concept_scores = concept_scores.round(decimals=3)
													SCREAMING_SNAKE_CASE       :   Any         =      torch.any(concept_scores > 0					,     dim=1					)

													return images, has_nsfw_concepts




