





import json
import os
import shutil
import tempfile
import unittest

import numpy as np
import pytest

from transformers import CLIPTokenizer, CLIPTokenizerFast
from transformers.models.clip.tokenization_clip import VOCAB_FILES_NAMES
from transformers.testing_utils import require_vision
from transformers.utils import IMAGE_PROCESSOR_NAME, is_vision_available


if is_vision_available():
 from PIL import Image

 from transformers import OwlViTImageProcessor, OwlViTProcessor


@require_vision
class     __magic_name__     (							unittest.TestCase							):





     """simple docstring"""



     def   SCREAMING_SNAKE_CASE						(							self				:int     ):





        '''simple docstring'''




        A_    : Union[str, Any]								=	tempfile.mkdtemp()

        # fmt: off
        A_    : int								=	["""""", """l""", """o""", """w""", """e""", """r""", """s""", """t""", """i""", """d""", """n""", """lo""", """l</w>""", """w</w>""", """r</w>""", """t</w>""", """low</w>""", """er</w>""", """lowest</w>""", """newer</w>""", """wider""", """<unk>""", """<|startoftext|>""", """<|endoftext|>"""]
        # fmt: on
        A_    : Dict								=	dict(zip(snake_case	,					range(len(snake_case     )     )     )     )
        A_    : Optional[int]								=	["""#version: 0.2""", """l o""", """lo w</w>""", """e r</w>""", """"""]
        A_    : List[str]								=	{"""unk_token""": """<unk>"""}

        A_    : Any								=	os.path.join(self.tmpdirname	,					VOCAB_FILES_NAMES["vocab_file"]     )
        A_    : Tuple								=	os.path.join(self.tmpdirname	,					VOCAB_FILES_NAMES["merges_file"]     )
        with open(self.vocab_file	,					"w"	,					encoding="utf-8"     ) as fp:
           fp.write(json.dumps(snake_case     ) + "\n"     )
        with open(self.merges_file	,					"w"	,					encoding="utf-8"     ) as fp:
           fp.write("\n".join(snake_case     )     )

        A_    : Optional[Any]								=	{
            """do_resize""": True,
            """size""": 20,
            """do_center_crop""": True,
            """crop_size""": 18,
            """do_normalize""": True,
            """image_mean""": [0.48145466, 0.4578275, 0.40821073],
            """image_std""": [0.26862954, 0.26130258, 0.27577711],
        }
        A_    : str								=	os.path.join(self.tmpdirname	,					snake_case     )
        with open(self.image_processor_file	,					"w"	,					encoding="utf-8"     ) as fp:
           json.dump(snake_case	,					snake_case     )



     def   SCREAMING_SNAKE_CASE						(							self				:Optional[Any]	,					**snake_case				:Optional[int]     ):





        '''simple docstring'''




        return CLIPTokenizer.from_pretrained(self.tmpdirname	,					pad_token="!"	,					**snake_case     )



     def   SCREAMING_SNAKE_CASE						(							self				:Optional[Any]	,					**snake_case				:Union[str, Any]     ):





        '''simple docstring'''




        return CLIPTokenizerFast.from_pretrained(self.tmpdirname	,					pad_token="!"	,					**snake_case     )



     def   SCREAMING_SNAKE_CASE						(							self				:Tuple	,					**snake_case				:List[str]     ):





        '''simple docstring'''




        return OwlViTImageProcessor.from_pretrained(self.tmpdirname	,					**snake_case     )



     def   SCREAMING_SNAKE_CASE						(							self				:Tuple     ):





        '''simple docstring'''




        shutil.rmtree(self.tmpdirname     )



     def   SCREAMING_SNAKE_CASE						(							self				:Optional[int]     ):





        '''simple docstring'''




        A_    : List[str]								=	[np.random.randint(255	,					size=(3, 30, 400)	,					dtype=np.uinta     )]

        A_    : int								=	[Image.fromarray(np.moveaxis(snake_case	,					0	,					-1     )     ) for x in image_inputs]

        return image_inputs



     def   SCREAMING_SNAKE_CASE						(							self				:List[str]     ):





        '''simple docstring'''




        A_    : Optional[int]								=	self.get_tokenizer()
        A_    : int								=	self.get_rust_tokenizer()
        A_    : Tuple								=	self.get_image_processor()

        A_    : Tuple								=	OwlViTProcessor(tokenizer=snake_case	,					image_processor=snake_case     )
        processor_slow.save_pretrained(self.tmpdirname     )
        A_    : Dict								=	OwlViTProcessor.from_pretrained(self.tmpdirname	,					use_fast=snake_case     )

        A_    : Dict								=	OwlViTProcessor(tokenizer=snake_case	,					image_processor=snake_case     )
        processor_fast.save_pretrained(self.tmpdirname     )
        A_    : Optional[int]								=	OwlViTProcessor.from_pretrained(self.tmpdirname     )

        self.assertEqual(processor_slow.tokenizer.get_vocab()	,					tokenizer_slow.get_vocab()     )
        self.assertEqual(processor_fast.tokenizer.get_vocab()	,					tokenizer_fast.get_vocab()     )
        self.assertEqual(tokenizer_slow.get_vocab()	,					tokenizer_fast.get_vocab()     )
        self.assertIsInstance(processor_slow.tokenizer	,					snake_case     )
        self.assertIsInstance(processor_fast.tokenizer	,					snake_case     )

        self.assertEqual(processor_slow.image_processor.to_json_string()	,					image_processor.to_json_string()     )
        self.assertEqual(processor_fast.image_processor.to_json_string()	,					image_processor.to_json_string()     )
        self.assertIsInstance(processor_slow.image_processor	,					snake_case     )
        self.assertIsInstance(processor_fast.image_processor	,					snake_case     )



     def   SCREAMING_SNAKE_CASE						(							self				:int     ):





        '''simple docstring'''




        A_    : Dict								=	OwlViTProcessor(tokenizer=self.get_tokenizer()	,					image_processor=self.get_image_processor()     )
        processor.save_pretrained(self.tmpdirname     )

        A_    : Dict								=	self.get_tokenizer(bos_token="(BOS)"	,					eos_token="(EOS)"     )
        A_    : Optional[Any]								=	self.get_image_processor(do_normalize=snake_case     )

        A_    : Tuple								=	OwlViTProcessor.from_pretrained(
            self.tmpdirname	,					bos_token="(BOS)"	,					eos_token="(EOS)"	,					do_normalize=snake_case     )

        self.assertEqual(processor.tokenizer.get_vocab()	,					tokenizer_add_kwargs.get_vocab()     )
        self.assertIsInstance(processor.tokenizer	,					snake_case     )

        self.assertEqual(processor.image_processor.to_json_string()	,					image_processor_add_kwargs.to_json_string()     )
        self.assertIsInstance(processor.image_processor	,					snake_case     )



     def   SCREAMING_SNAKE_CASE						(							self				:Any     ):





        '''simple docstring'''




        A_    : List[Any]								=	self.get_image_processor()
        A_    : int								=	self.get_tokenizer()

        A_    : Optional[int]								=	OwlViTProcessor(tokenizer=snake_case	,					image_processor=snake_case     )

        A_    : Dict								=	self.prepare_image_inputs()

        A_    : str								=	image_processor(snake_case	,					return_tensors="np"     )
        A_    : List[str]								=	processor(images=snake_case	,					return_tensors="np"     )

        for key in input_image_proc.keys():
           self.assertAlmostEqual(input_image_proc[key].sum()	,					input_processor[key].sum()	,					delta=1e-2     )



     def   SCREAMING_SNAKE_CASE						(							self				:List[Any]     ):





        '''simple docstring'''




        A_    : List[str]								=	self.get_image_processor()
        A_    : Optional[Any]								=	self.get_tokenizer()

        A_    : str								=	OwlViTProcessor(tokenizer=snake_case	,					image_processor=snake_case     )

        A_    : int								=	"""lower newer"""

        A_    : Optional[Any]								=	processor(text=snake_case	,					return_tensors="np"     )

        A_    : Dict								=	tokenizer(snake_case	,					return_tensors="np"     )

        for key in encoded_tok.keys():
           self.assertListEqual(encoded_tok[key][0].tolist()	,					encoded_processor[key][0].tolist()     )



     def   SCREAMING_SNAKE_CASE						(							self				:Optional[Any]     ):





        '''simple docstring'''




        A_    : str								=	self.get_image_processor()
        A_    : Optional[int]								=	self.get_tokenizer()

        A_    : List[Any]								=	OwlViTProcessor(tokenizer=snake_case	,					image_processor=snake_case     )

        A_    : Optional[Any]								=	"""lower newer"""
        A_    : Tuple								=	self.prepare_image_inputs()

        A_    : List[str]								=	processor(text=snake_case	,					images=snake_case     )

        self.assertListEqual(list(inputs.keys()     )	,					["input_ids", "attention_mask", "pixel_values"]     )

        # test if it raises when no input is passed
        with pytest.raises(snake_case     ):
           processor()



     def   SCREAMING_SNAKE_CASE						(							self				:str     ):





        '''simple docstring'''




        A_    : Dict								=	"""google/owlvit-base-patch32"""
        A_    : int								=	OwlViTProcessor.from_pretrained(snake_case     )

        A_    : Any								=	["""cat""", """nasa badge"""]
        A_    : List[Any]								=	processor(text=snake_case     )

        A_    : int								=	16
        self.assertListEqual(list(inputs.keys()     )	,					["input_ids", "attention_mask"]     )
        self.assertEqual(inputs["input_ids"].shape	,					(2, seq_length)     )

        # test if it raises when no input is passed
        with pytest.raises(snake_case     ):
           processor()



     def   SCREAMING_SNAKE_CASE						(							self				:Optional[int]     ):





        '''simple docstring'''




        A_    : Optional[int]								=	"""google/owlvit-base-patch32"""
        A_    : Optional[Any]								=	OwlViTProcessor.from_pretrained(snake_case     )

        A_    : Dict								=	[["""cat""", """nasa badge"""], ["""person"""]]
        A_    : Any								=	processor(text=snake_case     )

        A_    : Dict								=	16
        A_    : Union[str, Any]								=	len(snake_case     )
        A_    : Union[str, Any]								=	max([len(snake_case     ) for texts in input_texts]     )

        self.assertListEqual(list(inputs.keys()     )	,					["input_ids", "attention_mask"]     )
        self.assertEqual(inputs["input_ids"].shape	,					(batch_size * num_max_text_queries, seq_length)     )

        # test if it raises when no input is passed
        with pytest.raises(snake_case     ):
           processor()



     def   SCREAMING_SNAKE_CASE						(							self				:Optional[Any]     ):





        '''simple docstring'''




        A_    : Any								=	"""google/owlvit-base-patch32"""
        A_    : int								=	OwlViTProcessor.from_pretrained(snake_case     )

        A_    : List[str]								=	["""cat""", """nasa badge"""]
        A_    : Optional[Any]								=	processor(text=snake_case     )

        A_    : Any								=	16
        A_    : Tuple								=	inputs["""input_ids"""]
        A_    : Union[str, Any]								=	[
            [49_406, 2_368, 49_407, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
            [49_406, 6_841, 11_301, 49_407, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
        ]

        self.assertListEqual(list(inputs.keys()     )	,					["input_ids", "attention_mask"]     )
        self.assertEqual(inputs["input_ids"].shape	,					(2, seq_length)     )
        self.assertListEqual(list(input_ids[0]     )	,					predicted_ids[0]     )
        self.assertListEqual(list(input_ids[1]     )	,					predicted_ids[1]     )



     def   SCREAMING_SNAKE_CASE						(							self				:Tuple     ):





        '''simple docstring'''




        A_    : int								=	self.get_image_processor()
        A_    : str								=	self.get_tokenizer()

        A_    : Optional[Any]								=	OwlViTProcessor(tokenizer=snake_case	,					image_processor=snake_case     )

        A_    : Tuple								=	self.prepare_image_inputs()
        A_    : int								=	self.prepare_image_inputs()

        A_    : Any								=	processor(images=snake_case	,					query_images=snake_case     )

        self.assertListEqual(list(inputs.keys()     )	,					["query_pixel_values", "pixel_values"]     )

        # test if it raises when no input is passed
        with pytest.raises(snake_case     ):
           processor()





     def   SCREAMING_SNAKE_CASE						(							self				:Tuple     ):





        '''simple docstring'''




        A_    : Any								=	self.get_image_processor()
        A_    : List[str]								=	self.get_tokenizer()

        A_    : Any								=	OwlViTProcessor(tokenizer=snake_case	,					image_processor=snake_case     )

        A_    : Union[str, Any]								=	[[1, 4, 5, 8, 1, 0, 8], [3, 4, 3, 1, 1, 8, 9]]

        A_    : Dict								=	processor.batch_decode(snake_case     )
        A_    : List[str]								=	tokenizer.batch_decode(snake_case     )

        self.assertListEqual(snake_case	,					snake_case     )





'''simple docstring'''




import inspect
from typing import Callable, List, Optional, Union

import torch
from transformers import CLIPImageProcessor, CLIPTextModel, CLIPTokenizer

from diffusers import DiffusionPipeline
from diffusers.models import AutoencoderKL, UNetaDConditionModel
from diffusers.pipelines.stable_diffusion import StableDiffusionPipelineOutput
from diffusers.pipelines.stable_diffusion.safety_checker import StableDiffusionSafetyChecker
from diffusers.schedulers import DDIMScheduler, LMSDiscreteScheduler, PNDMScheduler
from diffusers.utils import logging


lowerCamelCase							:Any				    =					logging.get_logger(__name__)  # pylint: disable=invalid-name







class 	_lowerCAmelCase							(			__UpperCAmelCase ):


	def __init__(self       ,      lowercase       ,      lowercase       ,      lowercase       ,      lowercase       ,      lowercase       ,      lowercase       ,      lowercase       ,      ):
				super().__init__()
				self.register_modules(
				    vae=lowercase       ,      text_encoder=lowercase       ,      tokenizer=lowercase       ,      unet=lowercase       ,      scheduler=lowercase       ,      safety_checker=lowercase       ,      feature_extractor=lowercase       ,      )


	def     _a	(self       ,      lowercase = "auto"							):
				if slice_size == "auto":
							# half the attention head size is usually a good trade-off between
							# speed and memory
							A_	:						int	=     self.unet.config.attention_head_dim // 2
				self.unet.set_attention_slice(lowercase							)


	def     _a	(self							):
				self.enable_attention_slicing(lowercase							)


	@torch.no_grad()
	def __call__(self       ,      lowercase       ,      lowercase = 512       ,      lowercase = 512       ,      lowercase = 50       ,      lowercase = 7.5       ,      lowercase = None       ,      lowercase = 1       ,      lowercase = 0.0       ,      lowercase = None       ,      lowercase = None       ,      lowercase = "pil"       ,      lowercase = True       ,      lowercase = None       ,      lowercase = 1       ,      lowercase = None       ,      **lowercase       ,      ):

				if isinstance(lowercase       ,      lowercase							):
							A_	:						Union[str, Any]	=     1
				elif isinstance(lowercase       ,      lowercase							):
							A_	:						Any	=     len(lowercase							)
				else:
							raise ValueError(F'`prompt` has to be of type `str` or `list` but is {type(lowercase							)}'							)

				if height % 8 != 0 or width % 8 != 0:
							raise ValueError(F'`height` and `width` have to be divisible by 8 but are {height} and {width}.'							)

				if (callback_steps is None) or (
				    callback_steps is not None and (not isinstance(lowercase       ,      lowercase							) or callback_steps <= 0)
				):
							raise ValueError(
							    F'`callback_steps` has to be a positive integer but is {callback_steps} of type'
							    F' {type(lowercase							)}.'							)

				# get prompt text embeddings
				A_	:						Optional[Any]	=     self.tokenizer(
				    lowercase       ,      padding="""max_length"""       ,      max_length=self.tokenizer.model_max_length       ,      return_tensors="""pt"""       ,      )
				A_	:						Dict	=     text_inputs.input_ids

				if text_input_ids.shape[-1] > self.tokenizer.model_max_length:
							A_	:						Union[str, Any]	=     self.tokenizer.batch_decode(text_input_ids[:, self.tokenizer.model_max_length :]							)
							logger.warning(
							    """The following part of your input was truncated because CLIP can only handle sequences up to"""
							    F' {self.tokenizer.model_max_length} tokens: {removed_text}'							)
							A_	:						Dict	=     text_input_ids[:, : self.tokenizer.model_max_length]

				if text_embeddings is None:
							A_	:						Any	=     self.text_encoder(text_input_ids.to(self.device							)							)[0]

				# duplicate text embeddings for each generation per prompt, using mps friendly method
				A_,    A_,    A_	:						Tuple	=     text_embeddings.shape
				A_	:						Optional[Any]	=     text_embeddings.repeat(1       ,      lowercase       ,      1							)
				A_	:						Any	=     text_embeddings.view(bs_embed * num_images_per_prompt       ,      lowercase       ,      -1							)

				# here `guidance_scale` is defined analog to the guidance weight `w` of equation (2)
				# of the Imagen paper: https://arxiv.org/pdf/2205.11487.pdf . `guidance_scale = 1`
				# corresponds to doing no classifier free guidance.
				A_	:						List[str]	=     guidance_scale > 1.0
				# get unconditional embeddings for classifier free guidance
				if do_classifier_free_guidance:
							A_	:						List[str]
							if negative_prompt is None:
										A_	:						Optional[int]	=     [""""""]
							elif type(lowercase							) is not type(lowercase							):
										raise TypeError(
										    F'`negative_prompt` should be the same type to `prompt`, but got {type(lowercase							)} !='
										    F' {type(lowercase							)}.'							)
							elif isinstance(lowercase       ,      lowercase							):
										A_	:						Dict	=     [negative_prompt]
							elif batch_size != len(lowercase							):
										raise ValueError(
										    F'`negative_prompt`: {negative_prompt} has batch size {len(lowercase							)}, but `prompt`:'
										    F' {prompt} has batch size {batch_size}. Please make sure that passed `negative_prompt` matches'
										    """ the batch size of `prompt`."""							)
							else:
										A_	:						Dict	=     negative_prompt

							A_	:						int	=     text_input_ids.shape[-1]
							A_	:						List[Any]	=     self.tokenizer(
							    lowercase       ,      padding="""max_length"""       ,      max_length=lowercase       ,      truncation=lowercase       ,      return_tensors="""pt"""       ,      )
							A_	:						str	=     self.text_encoder(uncond_input.input_ids.to(self.device							)							)[0]

							# duplicate unconditional embeddings for each generation per prompt, using mps friendly method
							A_	:						Optional[Any]	=     uncond_embeddings.shape[1]
							A_	:						str	=     uncond_embeddings.repeat(lowercase       ,      lowercase       ,      1							)
							A_	:						List[str]	=     uncond_embeddings.view(batch_size * num_images_per_prompt       ,      lowercase       ,      -1							)

							# For classifier free guidance, we need to do two forward passes.
							# Here we concatenate the unconditional and text embeddings into a single batch
							# to avoid doing two forward passes
							A_	:						Dict	=     torch.cat([uncond_embeddings, text_embeddings]							)

				# get the initial random noise unless the user supplied it

				# Unlike in other pipelines, latents need to be generated in the target device
				# for 1-to-1 results reproducibility with the CompVis implementation.
				# However this currently doesn't work in `mps`.
				A_	:						int	=     (batch_size * num_images_per_prompt, self.unet.config.in_channels, height // 8, width // 8)
				A_	:						Dict	=     (batch_size * num_images_per_prompt, self.unet.config.in_channels, 64, 64)
				A_	:						Dict	=     text_embeddings.dtype
				if latents is None:
							if self.device.type == "mps":
										# randn does not exist on mps
										A_	:						Tuple	=     torch.randn(
										    lowercase       ,      generator=lowercase       ,      device="""cpu"""       ,      dtype=lowercase							).to(self.device							)
										A_	:						int	=     torch.randn(lowercase       ,      generator=lowercase       ,      device="""cpu"""       ,      dtype=lowercase							).to(
										    self.device							)
							else:
										A_	:						int	=     torch.randn(
										    lowercase       ,      generator=lowercase       ,      device=self.device       ,      dtype=lowercase							)
										A_	:						str	=     torch.randn(lowercase       ,      generator=lowercase       ,      device=self.device       ,      dtype=lowercase							)
				else:
							if latents_reference.shape != latents_shape:
										raise ValueError(F'Unexpected latents shape, got {latents.shape}, expected {latents_shape}'							)
							A_	:						str	=     latents_reference.to(self.device							)
							A_	:						Tuple	=     latents.to(self.device							)

				# This is the key part of the pipeline where we
				# try to ensure that the generated images w/ the same seed
				# but different sizes actually result in similar images
				A_	:						Optional[int]	=     (latents_shape[3] - latents_shape_reference[3]) // 2
				A_	:						Optional[int]	=     (latents_shape[2] - latents_shape_reference[2]) // 2
				A_	:						Optional[int]	=     latents_shape_reference[3] if dx >= 0 else latents_shape_reference[3] + 2 * dx
				A_	:						int	=     latents_shape_reference[2] if dy >= 0 else latents_shape_reference[2] + 2 * dy
				A_	:						Optional[Any]	=     0 if dx < 0 else dx
				A_	:						Optional[Any]	=     0 if dy < 0 else dy
				A_	:						Optional[int]	=     max(-dx       ,      0							)
				A_	:						List[str]	=     max(-dy       ,      0							)
				# import pdb
				# pdb.set_trace()
				A_	:						str	=     latents_reference[:, :, dy : dy + h, dx : dx + w]

				# set timesteps
				self.scheduler.set_timesteps(lowercase							)

				# Some schedulers like PNDM have timesteps as arrays
				# It's more optimized to move all timesteps to correct device beforehand
				A_	:						Any	=     self.scheduler.timesteps.to(self.device							)

				# scale the initial noise by the standard deviation required by the scheduler
				A_	:						Tuple	=     latents * self.scheduler.init_noise_sigma

				# prepare extra kwargs for the scheduler step, since not all schedulers have the same signature
				# eta (η) is only used with the DDIMScheduler, it will be ignored for other schedulers.
				# eta corresponds to η in DDIM paper: https://arxiv.org/abs/2010.02502
				# and should be between [0, 1]
				A_	:						Dict	=     """eta""" in set(inspect.signature(self.scheduler.step							).parameters.keys()							)
				A_	:						Any	=     {}
				if accepts_eta:
							A_	:						Optional[int]	=     eta

				for i, t in enumerate(self.progress_bar(lowercase							)							):
							# expand the latents if we are doing classifier free guidance
							A_	:						Tuple	=     torch.cat([latents] * 2							) if do_classifier_free_guidance else latents
							A_	:						Tuple	=     self.scheduler.scale_model_input(lowercase       ,      lowercase							)

							# predict the noise residual
							A_	:						List[str]	=     self.unet(lowercase       ,      lowercase       ,      encoder_hidden_states=lowercase							).sample

							# perform guidance
							if do_classifier_free_guidance:
										A_,    A_	:						str	=     noise_pred.chunk(2							)
										A_	:						Optional[Any]	=     noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)

							# compute the previous noisy sample x_t -> x_t-1
							A_	:						List[str]	=     self.scheduler.step(lowercase       ,      lowercase       ,      lowercase       ,      **lowercase							).prev_sample

							# call the callback, if provided
							if callback is not None and i % callback_steps == 0:
										callback(lowercase       ,      lowercase       ,      lowercase							)

				A_	:						List[str]	=     1 / 0.1_82_15 * latents
				A_	:						List[str]	=     self.vae.decode(lowercase							).sample

				A_	:						Optional[int]	=     (image / 2 + 0.5).clamp(0       ,      1							)

				# we always cast to float32 as this does not cause significant overhead and is compatible with bfloat16
				A_	:						Optional[int]	=     image.cpu().permute(0       ,      2       ,      3       ,      1							).float().numpy()

				if self.safety_checker is not None:
							A_	:						Union[str, Any]	=     self.feature_extractor(self.numpy_to_pil(lowercase							)       ,      return_tensors="""pt"""							).to(
							    self.device							)
							A_,    A_	:						Optional[int]	=     self.safety_checker(
							    images=lowercase       ,      clip_input=safety_checker_input.pixel_values.to(text_embeddings.dtype							)							)
				else:
							A_	:						Tuple	=     None

				if output_type == "pil":
							A_	:						Tuple	=     self.numpy_to_pil(lowercase							)

				if not return_dict:
							return (image, has_nsfw_concept)

				return StableDiffusionPipelineOutput(images=lowercase       ,      nsfw_content_detected=lowercase							)