


'''simple docstring'''







import inspect
from typing import Optional, Union

import numpy as np
import PIL
import torch
from torch.nn import functional as F
from torchvision import transforms
from transformers import CLIPFeatureExtractor, CLIPModel, CLIPTextModel, CLIPTokenizer

from diffusers import (
    AutoencoderKL,
    DDIMScheduler,
    DiffusionPipeline,
    DPMSolverMultistepScheduler,
    LMSDiscreteScheduler,
    PNDMScheduler,
    UNetaDConditionModel,
)
from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion import StableDiffusionPipelineOutput
from diffusers.utils import (
    PIL_INTERPOLATION,
    randn_tensor,
)






def 				lowerCAmelCase__	(    lowerCamelCase					:      Optional[int]  ,lowerCamelCase					:      int  ,lowerCamelCase					:      List[Any]       ):
       if isinstance(lowerCamelCase  ,torch.Tensor       ):
              return image
       elif isinstance(lowerCamelCase  ,PIL.Image.Image       ):
              _A					:     Optional[int]     = [image]

       if isinstance(image[0]  ,PIL.Image.Image       ):
              _A					:     int     = [np.array(i.resize((w, h)  ,resample=PIL_INTERPOLATION['lanczos']       )       )[None, :] for i in image]
              _A					:     Optional[int]     = np.concatenate(lowerCamelCase  ,axis=0       )
              _A					:     Optional[int]     = np.array(lowerCamelCase       ).astype(np.floataa       ) / 255.0
              _A					:     List[Any]     = image.transpose(0  ,3  ,1  ,2       )
              _A					:     Union[str, Any]     = 2.0 * image - 1.0
              _A					:     Any     = torch.from_numpy(lowerCamelCase       )
       elif isinstance(image[0]  ,torch.Tensor       ):
              _A					:     List[Any]     = torch.cat(lowerCamelCase  ,dim=0       )
       return image






def 				lowerCAmelCase__	(    lowerCamelCase					:      List[str]  ,lowerCamelCase					:      Tuple  ,lowerCamelCase					:      List[str]  ,lowerCamelCase					:      List[str]=0.9_995       ):
       if not isinstance(lowerCamelCase  ,np.ndarray       ):
              _A					:     Optional[Any]     = True
              _A					:     List[str]     = va.device
              _A					:     List[Any]     = va.cpu().numpy()
              _A					:     List[Any]     = va.cpu().numpy()

       _A					:     int     = np.sum(va * va / (np.linalg.norm(lowerCamelCase       ) * np.linalg.norm(lowerCamelCase       ))       )
       if np.abs(lowerCamelCase       ) > DOT_THRESHOLD:
              _A					:     List[str]     = (1 - t) * va + t * va
       else:
              _A					:     Optional[Any]     = np.arccos(lowerCamelCase       )
              _A					:     Union[str, Any]     = np.sin(lowerCamelCase       )
              _A					:     Optional[int]     = theta_a * t
              _A					:     Optional[Any]     = np.sin(lowerCamelCase       )
              _A					:     int     = np.sin(theta_a - theta_t       ) / sin_theta_a
              _A					:     Optional[int]     = sin_theta_t / sin_theta_a
              _A					:     int     = sa * va + sa * va

       if inputs_are_torch:
              _A					:     List[str]     = torch.from_numpy(lowerCamelCase       ).to(lowerCamelCase       )

       return va






def 				lowerCAmelCase__	(    lowerCamelCase					:      int  ,lowerCamelCase					:      Union[str, Any]       ):
       _A					:     Union[str, Any]     = F.normalize(lowerCamelCase  ,dim=-1       )
       _A					:     Tuple     = F.normalize(lowerCamelCase  ,dim=-1       )
       return (x - y).norm(dim=-1       ).div(2       ).arcsin().pow(2       ).mul(2       )






def 				lowerCAmelCase__	(    lowerCamelCase					:      int  ,lowerCamelCase					:      Any       ):
       for param in model.parameters():
              _A					:     Dict     = value


class 			__lowerCamelCase			(	a_	):
    """simple docstring"""


    def __init__(  self      :							int				,							SCREAMING_SNAKE_CASE      :							AutoencoderKL				,							SCREAMING_SNAKE_CASE      :							CLIPTextModel				,							SCREAMING_SNAKE_CASE      :							CLIPModel				,							SCREAMING_SNAKE_CASE      :							CLIPTokenizer				,							SCREAMING_SNAKE_CASE      :							UNetaDConditionModel				,							SCREAMING_SNAKE_CASE      :							Union[PNDMScheduler, LMSDiscreteScheduler, DDIMScheduler, DPMSolverMultistepScheduler]				,							SCREAMING_SNAKE_CASE      :							CLIPFeatureExtractor				,							SCREAMING_SNAKE_CASE      :							Union[str, Any]=None				,							SCREAMING_SNAKE_CASE      :							str=None				,							SCREAMING_SNAKE_CASE      :							Optional[int]=None				,							):
           super().__init__()
           self.register_modules(
               vae=SCREAMING_SNAKE_CASE				,							text_encoder=SCREAMING_SNAKE_CASE				,							clip_model=SCREAMING_SNAKE_CASE				,							tokenizer=SCREAMING_SNAKE_CASE				,							unet=SCREAMING_SNAKE_CASE				,							scheduler=SCREAMING_SNAKE_CASE				,							feature_extractor=SCREAMING_SNAKE_CASE				,							coca_model=SCREAMING_SNAKE_CASE				,							coca_tokenizer=SCREAMING_SNAKE_CASE				,							coca_transform=SCREAMING_SNAKE_CASE				,							)
           _A					:     List[str]     = (
               feature_extractor.size
               if isinstance(feature_extractor.size				,							SCREAMING_SNAKE_CASE)
               else feature_extractor.size['shortest_edge']
           )
           _A					:     Union[str, Any]     = transforms.Normalize(mean=feature_extractor.image_mean				,							std=feature_extractor.image_std)
           set_requires_grad(self.text_encoder				,							SCREAMING_SNAKE_CASE)
           set_requires_grad(self.clip_model				,							SCREAMING_SNAKE_CASE)


    def        A   (  self      :							List[str]				,							SCREAMING_SNAKE_CASE      :							Optional[Union[str, int]] = "auto"):
           if slice_size == "auto":
                  # half the attention head size is usually a good trade-off between
                  # speed and memory
                  _A					:     Dict     = self.unet.config.attention_head_dim // 2
           self.unet.set_attention_slice(SCREAMING_SNAKE_CASE)


    def        A   (  self      :							Any):
           self.enable_attention_slicing(SCREAMING_SNAKE_CASE)


    def        A   (  self      :							Union[str, Any]):
           set_requires_grad(self.vae				,							SCREAMING_SNAKE_CASE)


    def        A   (  self      :							str):
           set_requires_grad(self.vae				,							SCREAMING_SNAKE_CASE)


    def        A   (  self      :							List[Any]):
           set_requires_grad(self.unet				,							SCREAMING_SNAKE_CASE)


    def        A   (  self      :							Any):
           set_requires_grad(self.unet				,							SCREAMING_SNAKE_CASE)


    def        A   (  self      :							Any				,							SCREAMING_SNAKE_CASE      :							Optional[int]				,							SCREAMING_SNAKE_CASE      :							Optional[int]				,							SCREAMING_SNAKE_CASE      :							Union[str, Any]):
           # get the original timestep using init_timestep
           _A					:     Optional[int]     = min(int(num_inference_steps * strength)				,							SCREAMING_SNAKE_CASE)

           _A					:     Any     = max(num_inference_steps - init_timestep				,							0)
           _A					:     str     = self.scheduler.timesteps[t_start:]

           return timesteps, num_inference_steps - t_start


    def        A   (  self      :							Union[str, Any]				,							SCREAMING_SNAKE_CASE      :							Any				,							SCREAMING_SNAKE_CASE      :							List[Any]				,							SCREAMING_SNAKE_CASE      :							Optional[int]				,							SCREAMING_SNAKE_CASE      :							Optional[Any]				,							SCREAMING_SNAKE_CASE      :							int				,							SCREAMING_SNAKE_CASE      :							Tuple=None):
           if not isinstance(SCREAMING_SNAKE_CASE				,							torch.Tensor):
                  raise ValueError(F'`image` has to be of type `torch.Tensor` but is {type(SCREAMING_SNAKE_CASE)}')

           _A					:     Any     = image.to(device=SCREAMING_SNAKE_CASE				,							dtype=SCREAMING_SNAKE_CASE)

           if isinstance(SCREAMING_SNAKE_CASE				,							SCREAMING_SNAKE_CASE):
                  _A					:     Any     = [
                      self.vae.encode(image[i : i + 1]).latent_dist.sample(generator[i]) for i in range(SCREAMING_SNAKE_CASE)
                  ]
                  _A					:     Any     = torch.cat(SCREAMING_SNAKE_CASE				,							dim=0)
           else:
                  _A					:     List[str]     = self.vae.encode(SCREAMING_SNAKE_CASE).latent_dist.sample(SCREAMING_SNAKE_CASE)

           # Hardcode 0.18215 because stable-diffusion-2-base has not self.vae.config.scaling_factor
           _A					:     Union[str, Any]     = 0.1_8215 * init_latents
           _A					:     List[Any]     = init_latents.repeat_interleave(SCREAMING_SNAKE_CASE				,							dim=0)

           _A					:     Optional[int]     = randn_tensor(init_latents.shape				,							generator=SCREAMING_SNAKE_CASE				,							device=SCREAMING_SNAKE_CASE				,							dtype=SCREAMING_SNAKE_CASE)

           # get latents
           _A					:     str     = self.scheduler.add_noise(SCREAMING_SNAKE_CASE				,							SCREAMING_SNAKE_CASE				,							SCREAMING_SNAKE_CASE)
           _A					:     List[Any]     = init_latents

           return latents


    def        A   (  self      :							Any				,							SCREAMING_SNAKE_CASE      :							Any):
           _A					:     Tuple     = self.coca_transform(SCREAMING_SNAKE_CASE).unsqueeze(0)
           with torch.no_grad(), torch.cuda.amp.autocast():
                  _A					:     str     = self.coca_model.generate(transformed_image.to(device=self.device				,							dtype=self.coca_model.dtype))
           _A					:     Union[str, Any]     = self.coca_tokenizer.decode(generated[0].cpu().numpy())
           return generated.split('<end_of_text>')[0].replace('<start_of_text>'				,							'').rstrip(' .,')


    def        A   (  self      :							Tuple				,							SCREAMING_SNAKE_CASE      :							List[Any]				,							SCREAMING_SNAKE_CASE      :							Dict):
           _A					:     Any     = self.feature_extractor.preprocess(SCREAMING_SNAKE_CASE)
           _A					:     str     = torch.from_numpy(clip_image_input['pixel_values'][0]).unsqueeze(0).to(self.device).half()
           _A					:     str     = self.clip_model.get_image_features(SCREAMING_SNAKE_CASE)
           _A					:     List[Any]     = image_embeddings_clip / image_embeddings_clip.norm(p=2				,							dim=-1				,							keepdim=SCREAMING_SNAKE_CASE)
           _A					:     Optional[Any]     = image_embeddings_clip.repeat_interleave(SCREAMING_SNAKE_CASE				,							dim=0)
           return image_embeddings_clip


    @torch.enable_grad()
    def        A   (  self      :							Dict				,							SCREAMING_SNAKE_CASE      :							Optional[Any]				,							SCREAMING_SNAKE_CASE      :							Tuple				,							SCREAMING_SNAKE_CASE      :							int				,							SCREAMING_SNAKE_CASE      :							List[str]				,							SCREAMING_SNAKE_CASE      :							Optional[Any]				,							SCREAMING_SNAKE_CASE      :							List[str]				,							SCREAMING_SNAKE_CASE      :							Any				,							):
           _A					:     str     = latents.detach().requires_grad_()

           _A					:     Optional[Any]     = self.scheduler.scale_model_input(SCREAMING_SNAKE_CASE				,							SCREAMING_SNAKE_CASE)

           # predict the noise residual
           _A					:     Union[str, Any]     = self.unet(SCREAMING_SNAKE_CASE				,							SCREAMING_SNAKE_CASE				,							encoder_hidden_states=SCREAMING_SNAKE_CASE).sample

           if isinstance(self.scheduler				,							(PNDMScheduler, DDIMScheduler, DPMSolverMultistepScheduler)):
                  _A					:     Optional[int]     = self.scheduler.alphas_cumprod[timestep]
                  _A					:     str     = 1 - alpha_prod_t
                  # compute predicted original sample from predicted noise also called
                  # "predicted x_0" of formula (12) from https://arxiv.org/pdf/2010.02502.pdf
                  _A					:     Any     = (latents - beta_prod_t ** 0.5 * noise_pred) / alpha_prod_t ** 0.5

                  _A					:     Dict     = torch.sqrt(SCREAMING_SNAKE_CASE)
                  _A					:     Tuple     = pred_original_sample * (fac) + latents * (1 - fac)
           elif isinstance(self.scheduler				,							SCREAMING_SNAKE_CASE):
                  _A					:     int     = self.scheduler.sigmas[index]
                  _A					:     Any     = latents - sigma * noise_pred
           else:
                  raise ValueError(F'scheduler type {type(self.scheduler)} not supported')

           # Hardcode 0.18215 because stable-diffusion-2-base has not self.vae.config.scaling_factor
           _A					:     Tuple     = 1 / 0.1_8215 * sample
           _A					:     Any     = self.vae.decode(SCREAMING_SNAKE_CASE).sample
           _A					:     Tuple     = (image / 2 + 0.5).clamp(0				,							1)

           _A					:     List[str]     = transforms.Resize(self.feature_extractor_size)(SCREAMING_SNAKE_CASE)
           _A					:     Optional[int]     = self.normalize(SCREAMING_SNAKE_CASE).to(latents.dtype)

           _A					:     List[str]     = self.clip_model.get_image_features(SCREAMING_SNAKE_CASE)
           _A					:     Any     = image_embeddings_clip / image_embeddings_clip.norm(p=2				,							dim=-1				,							keepdim=SCREAMING_SNAKE_CASE)

           _A					:     Dict     = spherical_dist_loss(SCREAMING_SNAKE_CASE				,							SCREAMING_SNAKE_CASE).mean() * clip_guidance_scale

           _A					:     Optional[Any]     = -torch.autograd.grad(SCREAMING_SNAKE_CASE				,							SCREAMING_SNAKE_CASE)[0]

           if isinstance(self.scheduler				,							SCREAMING_SNAKE_CASE):
                  _A					:     Dict     = latents.detach() + grads * (sigma**2)
                  _A					:     int     = noise_pred_original
           else:
                  _A					:     Any     = noise_pred_original - torch.sqrt(SCREAMING_SNAKE_CASE) * grads
           return noise_pred, latents


    @torch.no_grad()
    def __call__(  self      :							Dict				,							SCREAMING_SNAKE_CASE      :							Union[torch.FloatTensor, PIL.Image.Image]				,							SCREAMING_SNAKE_CASE      :							Union[torch.FloatTensor, PIL.Image.Image]				,							SCREAMING_SNAKE_CASE      :							Optional[str] = None				,							SCREAMING_SNAKE_CASE      :							Optional[str] = None				,							SCREAMING_SNAKE_CASE      :							Optional[int] = 512				,							SCREAMING_SNAKE_CASE      :							Optional[int] = 512				,							SCREAMING_SNAKE_CASE      :							float = 0.6				,							SCREAMING_SNAKE_CASE      :							Optional[int] = 50				,							SCREAMING_SNAKE_CASE      :							Optional[float] = 7.5				,							SCREAMING_SNAKE_CASE      :							Optional[int] = 1				,							SCREAMING_SNAKE_CASE      :							float = 0.0				,							SCREAMING_SNAKE_CASE      :							Optional[float] = 100				,							SCREAMING_SNAKE_CASE      :							Optional[torch.Generator] = None				,							SCREAMING_SNAKE_CASE      :							Optional[str] = "pil"				,							SCREAMING_SNAKE_CASE      :							bool = True				,							SCREAMING_SNAKE_CASE      :							float = 0.8				,							SCREAMING_SNAKE_CASE      :							float = 0.1				,							SCREAMING_SNAKE_CASE      :							float = 0.1				,							):
           if isinstance(SCREAMING_SNAKE_CASE				,							SCREAMING_SNAKE_CASE) and len(SCREAMING_SNAKE_CASE) != batch_size:
                  raise ValueError(F'You have passed {batch_size} batch_size, but only {len(SCREAMING_SNAKE_CASE)} generators.')

           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 isinstance(SCREAMING_SNAKE_CASE				,							torch.Generator) and batch_size > 1:
                  _A					:     int     = [generator] + [None] * (batch_size - 1)

           _A					:     str     = [
               ('model', self.coca_model is None),
               ('tokenizer', self.coca_tokenizer is None),
               ('transform', self.coca_transform is None),
           ]
           _A					:     str     = [x[0] for x in coca_is_none if x[1]]
           _A					:     List[str]     = ', '.join(SCREAMING_SNAKE_CASE)
           # generate prompts with coca model if prompt is None
           if content_prompt is None:
                  if len(SCREAMING_SNAKE_CASE):
                         raise ValueError(
                             F'Content prompt is None and CoCa [{coca_is_none_str}] is None.'
                             F'Set prompt or pass Coca [{coca_is_none_str}] to DiffusionPipeline.')
                  _A					:     Union[str, Any]     = self.get_image_description(SCREAMING_SNAKE_CASE)
           if style_prompt is None:
                  if len(SCREAMING_SNAKE_CASE):
                         raise ValueError(
                             F'Style prompt is None and CoCa [{coca_is_none_str}] is None.'
                             F' Set prompt or pass Coca [{coca_is_none_str}] to DiffusionPipeline.')
                  _A					:     Dict     = self.get_image_description(SCREAMING_SNAKE_CASE)

           # get prompt text embeddings for content and style
           _A					:     int     = self.tokenizer(
               SCREAMING_SNAKE_CASE				,							padding='max_length'				,							max_length=self.tokenizer.model_max_length				,							truncation=SCREAMING_SNAKE_CASE				,							return_tensors='pt'				,							)
           _A					:     Optional[Any]     = self.text_encoder(content_text_input.input_ids.to(self.device))[0]

           _A					:     int     = self.tokenizer(
               SCREAMING_SNAKE_CASE				,							padding='max_length'				,							max_length=self.tokenizer.model_max_length				,							truncation=SCREAMING_SNAKE_CASE				,							return_tensors='pt'				,							)
           _A					:     Optional[int]     = self.text_encoder(style_text_input.input_ids.to(self.device))[0]

           _A					:     int     = slerp(SCREAMING_SNAKE_CASE				,							SCREAMING_SNAKE_CASE				,							SCREAMING_SNAKE_CASE)

           # duplicate text embeddings for each generation per prompt
           _A					:     str     = text_embeddings.repeat_interleave(SCREAMING_SNAKE_CASE				,							dim=0)

           # set timesteps
           _A					:     Union[str, Any]     = 'offset' in set(inspect.signature(self.scheduler.set_timesteps).parameters.keys())
           _A					:     Union[str, Any]     = {}
           if accepts_offset:
                  _A					:     Optional[int]     = 1

           self.scheduler.set_timesteps(SCREAMING_SNAKE_CASE				,							**SCREAMING_SNAKE_CASE)
           # Some schedulers like PNDM have timesteps as arrays
           # It's more optimized to move all timesteps to correct device beforehand
           self.scheduler.timesteps.to(self.device)

           _A  ,			_A					:     Optional[Any]     = self.get_timesteps(SCREAMING_SNAKE_CASE				,							SCREAMING_SNAKE_CASE				,							self.device)
           _A					:     Dict     = timesteps[:1].repeat(SCREAMING_SNAKE_CASE)

           # Preprocess image
           _A					:     Optional[Any]     = preprocess(SCREAMING_SNAKE_CASE				,							SCREAMING_SNAKE_CASE				,							SCREAMING_SNAKE_CASE)
           _A					:     Optional[int]     = self.prepare_latents(
               SCREAMING_SNAKE_CASE				,							SCREAMING_SNAKE_CASE				,							SCREAMING_SNAKE_CASE				,							text_embeddings.dtype				,							self.device				,							SCREAMING_SNAKE_CASE)

           _A					:     Optional[Any]     = preprocess(SCREAMING_SNAKE_CASE				,							SCREAMING_SNAKE_CASE				,							SCREAMING_SNAKE_CASE)
           _A					:     Optional[int]     = self.prepare_latents(
               SCREAMING_SNAKE_CASE				,							SCREAMING_SNAKE_CASE				,							SCREAMING_SNAKE_CASE				,							text_embeddings.dtype				,							self.device				,							SCREAMING_SNAKE_CASE)

           _A					:     List[Any]     = slerp(SCREAMING_SNAKE_CASE				,							SCREAMING_SNAKE_CASE				,							SCREAMING_SNAKE_CASE)

           if clip_guidance_scale > 0:
                  _A					:     int     = self.get_clip_image_embeddings(SCREAMING_SNAKE_CASE				,							SCREAMING_SNAKE_CASE)
                  _A					:     str     = self.get_clip_image_embeddings(SCREAMING_SNAKE_CASE				,							SCREAMING_SNAKE_CASE)
                  _A					:     Union[str, Any]     = slerp(
                      SCREAMING_SNAKE_CASE				,							SCREAMING_SNAKE_CASE				,							SCREAMING_SNAKE_CASE)

           # 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					:     Any     = guidance_scale > 1.0
           # get unconditional embeddings for classifier free guidance
           if do_classifier_free_guidance:
                  _A					:     Tuple     = content_text_input.input_ids.shape[-1]
                  _A					:     List[Any]     = self.tokenizer(['']				,							padding='max_length'				,							max_length=SCREAMING_SNAKE_CASE				,							return_tensors='pt')
                  _A					:     str     = self.text_encoder(uncond_input.input_ids.to(self.device))[0]
                  # duplicate unconditional embeddings for each generation per prompt
                  _A					:     Optional[int]     = uncond_embeddings.repeat_interleave(SCREAMING_SNAKE_CASE				,							dim=0)

                  # 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					:     int     = 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					:     List[Any]     = (batch_size, self.unet.config.in_channels, height // 8, width // 8)
           _A					:     Any     = text_embeddings.dtype
           if latents is None:
                  if self.device.type == "mps":
                         # randn does not work reproducibly on mps
                         _A					:     Optional[Any]     = torch.randn(SCREAMING_SNAKE_CASE				,							generator=SCREAMING_SNAKE_CASE				,							device='cpu'				,							dtype=SCREAMING_SNAKE_CASE).to(
                             self.device)
                  else:
                         _A					:     Union[str, Any]     = torch.randn(SCREAMING_SNAKE_CASE				,							generator=SCREAMING_SNAKE_CASE				,							device=self.device				,							dtype=SCREAMING_SNAKE_CASE)
           else:
                  if latents.shape != latents_shape:
                         raise ValueError(F'Unexpected latents shape, got {latents.shape}, expected {latents_shape}')
                  _A					:     Any     = latents.to(self.device)

           # scale the initial noise by the standard deviation required by the scheduler
           _A					:     str     = 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					:     List[str]     = 'eta' in set(inspect.signature(self.scheduler.step).parameters.keys())
           _A					:     List[Any]     = {}
           if accepts_eta:
                  _A					:     Optional[int]     = eta

           # check if the scheduler accepts generator
           _A					:     Union[str, Any]     = 'generator' in set(inspect.signature(self.scheduler.step).parameters.keys())
           if accepts_generator:
                  _A					:     Any     = generator

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

                         # predict the noise residual
                         _A					:     Tuple     = self.unet(SCREAMING_SNAKE_CASE				,							SCREAMING_SNAKE_CASE				,							encoder_hidden_states=SCREAMING_SNAKE_CASE).sample

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

                         # perform clip guidance
                         if clip_guidance_scale > 0:
                                _A					:     str     = (
                                    text_embeddings.chunk(2)[1] if do_classifier_free_guidance else text_embeddings
                                )
                                _A  ,			_A					:     List[str]     = self.cond_fn(
                                    SCREAMING_SNAKE_CASE				,							SCREAMING_SNAKE_CASE				,							SCREAMING_SNAKE_CASE				,							SCREAMING_SNAKE_CASE				,							SCREAMING_SNAKE_CASE				,							SCREAMING_SNAKE_CASE				,							SCREAMING_SNAKE_CASE				,							)

                         # compute the previous noisy sample x_t -> x_t-1
                         _A					:     Dict     = self.scheduler.step(SCREAMING_SNAKE_CASE				,							SCREAMING_SNAKE_CASE				,							SCREAMING_SNAKE_CASE				,							**SCREAMING_SNAKE_CASE).prev_sample

        # Hardcode 0.18215 because stable-diffusion-2-base has not self.vae.config.scaling_factor
           _A					:     Optional[Any]     = 1 / 0.1_8215 * latents
           _A					:     Union[str, Any]     = self.vae.decode(SCREAMING_SNAKE_CASE).sample

           _A					:     Optional[Any]     = (image / 2 + 0.5).clamp(0				,							1)
           _A					:     int     = image.cpu().permute(0				,							2				,							3				,							1).numpy()

           if output_type == "pil":
                  _A					:     str     = self.numpy_to_pil(SCREAMING_SNAKE_CASE)

           if not return_dict:
                  return (image, None)

           return StableDiffusionPipelineOutput(images=SCREAMING_SNAKE_CASE				,							nsfw_content_detected=SCREAMING_SNAKE_CASE)










'''simple docstring'''







import gc
import importlib.metadata
import tempfile
import unittest

from packaging import version

from transformers import (
    AutoModel,
    AutoModelForCausalLM,
    AutoModelForSeqaSeqLM,
    AutoModelForSequenceClassification,
    AutoTokenizer,
    BitsAndBytesConfig,
    pipeline,
)
from transformers.testing_utils import (
    is_torch_available,
    require_accelerate,
    require_bitsandbytes,
    require_torch,
    require_torch_gpu,
    require_torch_multi_gpu,
    slow,
)






def 				lowerCAmelCase__	(    lowerCamelCase					:      Optional[int]       ):
       if model.config.model_type == "gpt2":
              return model.transformer.h[0].mlp.c_fc
       return model.transformer.h[0].mlp.dense_ah_to_h


if is_torch_available():
   import torch
   import torch.nn as nn


   class 			__lowerCamelCase			(	nn.Module	):
       """simple docstring"""


       def __init__(  self      :							int				,							SCREAMING_SNAKE_CASE      :							nn.Module				,							SCREAMING_SNAKE_CASE      :							int):
              super().__init__()
              _A					:     Tuple     = module
              _A					:     str     = nn.Sequential(
                  nn.Linear(module.in_features				,							SCREAMING_SNAKE_CASE				,							bias=SCREAMING_SNAKE_CASE)				,							nn.Linear(SCREAMING_SNAKE_CASE				,							module.out_features				,							bias=SCREAMING_SNAKE_CASE)				,							)
              _A					:     Any     = (2.0 / (5 * min(module.in_features				,							module.out_features))) ** 0.5
              nn.init.normal_(self.adapter[0].weight				,							std=SCREAMING_SNAKE_CASE)
              nn.init.zeros_(self.adapter[1].weight)
              self.adapter.to(module.weight.device)


       def        A   (  self      :							Any				,							SCREAMING_SNAKE_CASE      :							Any				,							*SCREAMING_SNAKE_CASE      :							List[str]				,							**SCREAMING_SNAKE_CASE      :							int):
              return self.module(SCREAMING_SNAKE_CASE				,							*SCREAMING_SNAKE_CASE				,							**SCREAMING_SNAKE_CASE) + self.adapter(SCREAMING_SNAKE_CASE)


@require_bitsandbytes
@require_accelerate
@require_torch
@require_torch_gpu
@slow
class 			__lowerCamelCase			(	unittest.TestCase	):
    """simple docstring"""
    # We keep the constants inside the init function and model loading inside setUp function

    # We need to test on relatively large models (aka >1b parameters otherwise the quantiztion may not work as expected)
    # Therefore here we use only bloom-1b3 to test our module
    a													=   "bigscience/bloom-1b7"

    # Constant values
    a													=   2.109_6595_5269_2574

    a													=   "Hello my name is"
    a													=   set()
    EXPECTED_OUTPUTS.add("Hello my name is John and I am a professional photographer. I"	)
    EXPECTED_OUTPUTS.add("Hello my name is John.\nI am a friend of your father.\n"	)
    EXPECTED_OUTPUTS.add("Hello my name is John Doe, I am a student at the University"	)
    a													=   10


    def        A   (  self      :							List[Any]):
           # Models and tokenizer
           _A					:     List[Any]     = AutoTokenizer.from_pretrained(self.model_name)




class 			__lowerCamelCase			(	a_	):
    """simple docstring"""


    def        A   (  self      :							List[str]):
           super().setUp()

           # Models and tokenizer
           _A					:     Tuple     = AutoModelForCausalLM.from_pretrained(
               self.model_name				,							torch_dtype=torch.floataa				,							device_map='auto')
           _A					:     str     = AutoModelForCausalLM.from_pretrained(self.model_name				,							load_in_abit=SCREAMING_SNAKE_CASE				,							device_map='auto')


    def        A   (  self      :							List[Any]):
           del self.model_fpaa
           del self.model_abit

           gc.collect()
           torch.cuda.empty_cache()


    def        A   (  self      :							str):
           _A					:     Any     = self.model_abit.config

           self.assertTrue(hasattr(SCREAMING_SNAKE_CASE				,							'quantization_config'))

           _A					:     str     = config.to_dict()
           _A					:     Dict     = config.to_diff_dict()

           _A					:     int     = config.to_json_string()


    def        A   (  self      :							List[str]):
           from bitsandbytes.nn import Paramsabit

           _A					:     List[str]     = self.model_fpaa.get_memory_footprint()
           _A					:     List[Any]     = self.model_abit.get_memory_footprint()

           self.assertAlmostEqual(mem_fpaa / mem_abit				,							self.EXPECTED_RELATIVE_DIFFERENCE)
           _A					:     str     = get_some_linear_layer(self.model_abit)
           self.assertTrue(linear.weight.__class__ == Paramsabit)


    def        A   (  self      :							List[Any]):
           from transformers import TaPreTrainedModel

           self.model_fpaa.get_memory_footprint()
           self.model_abit.get_memory_footprint()

           for name, module in self.model_abit.named_modules():
                  if isinstance(SCREAMING_SNAKE_CASE				,							torch.nn.Linear):
                         if name not in ["lm_head"] + TaPreTrainedModel._keep_in_fpaa_modules:
                                # 4-bit parameters are packed in uint8 variables
                                self.assertTrue(module.weight.dtype == torch.uinta)


    def        A   (  self      :							Any):
           _A					:     List[str]     = self.tokenizer(self.input_text				,							return_tensors='pt')
           _A					:     Any     = self.model_abit.generate(input_ids=encoded_input['input_ids'].to(0)				,							max_new_tokens=10)

           self.assertIn(self.tokenizer.decode(output_sequences[0]				,							skip_special_tokens=SCREAMING_SNAKE_CASE)				,							self.EXPECTED_OUTPUTS)


    def        A   (  self      :							List[str]):
           _A					:     List[Any]     = BitsAndBytesConfig()
           _A					:     Optional[int]     = True

           _A					:     Union[str, Any]     = AutoModelForCausalLM.from_pretrained(
               self.model_name				,							quantization_config=SCREAMING_SNAKE_CASE				,							device_map='auto')

           _A					:     int     = self.tokenizer(self.input_text				,							return_tensors='pt')
           _A					:     Any     = model_abit_from_config.generate(
               input_ids=encoded_input['input_ids'].to(0)				,							max_new_tokens=10)

           self.assertIn(self.tokenizer.decode(output_sequences[0]				,							skip_special_tokens=SCREAMING_SNAKE_CASE)				,							self.EXPECTED_OUTPUTS)


    def        A   (  self      :							Union[str, Any]):
           with self.assertRaises(SCREAMING_SNAKE_CASE), tempfile.TemporaryDirectory() as tmpdirname:
                  self.model_abit.save_pretrained(SCREAMING_SNAKE_CASE)


    def        A   (  self      :							int):
           _A					:     Any     = BitsAndBytesConfig()

           with self.assertRaises(SCREAMING_SNAKE_CASE):
                  _A					:     Optional[Any]     = AutoModelForCausalLM.from_pretrained(
                      self.model_name				,							quantization_config=SCREAMING_SNAKE_CASE				,							load_in_abit=SCREAMING_SNAKE_CASE				,							device_map='auto'				,							bnb_abit_quant_type='nf4'				,							)


    def        A   (  self      :							Optional[Any]):
           with self.assertRaises(SCREAMING_SNAKE_CASE):
                  # Tries with `str`
                  self.model_abit.to('cpu')

           with self.assertRaises(SCREAMING_SNAKE_CASE):
                  # Tries with a `dtype``
                  self.model_abit.to(torch.floataa)

           with self.assertRaises(SCREAMING_SNAKE_CASE):
                  # Tries with a `device`
                  self.model_abit.to(torch.device('cuda:0'))

           with self.assertRaises(SCREAMING_SNAKE_CASE):
                  # Tries with a `device`
                  self.model_abit.float()

           with self.assertRaises(SCREAMING_SNAKE_CASE):
                  # Tries with a `device`
                  self.model_abit.half()

           # Test if we did not break anything
           _A					:     str     = self.tokenizer(self.input_text				,							return_tensors='pt')

           _A					:     Optional[Any]     = self.model_fpaa.to(torch.floataa)
           _A					:     Optional[int]     = self.model_fpaa.generate(input_ids=encoded_input['input_ids'].to(0)				,							max_new_tokens=10)

           # Check this does not throw an error
           _A					:     Tuple     = self.model_fpaa.to('cpu')

           # Check this does not throw an error
           _A					:     Optional[int]     = self.model_fpaa.half()

           # Check this does not throw an error
           _A					:     int     = self.model_fpaa.float()


    def        A   (  self      :							Dict):
           _A					:     Any     = AutoModelForSeqaSeqLM.from_pretrained('t5-small'				,							load_in_abit=SCREAMING_SNAKE_CASE				,							device_map='auto')
           self.assertTrue(model.decoder.block[0].layer[2].DenseReluDense.wo.weight.dtype == torch.floataa)




@require_bitsandbytes
@require_accelerate
@require_torch
@require_torch_gpu
@slow
class 			__lowerCamelCase			(	unittest.TestCase	):
    """simple docstring"""


    @classmethod
    def        A   (  cls      :							Union[str, Any]):
           _A					:     Tuple     = 't5-small'
           _A					:     Union[str, Any]     = 'google/flan-t5-small'  # flan-t5 uses dense-act instead of dense-relu-dense
           _A					:     str     = AutoTokenizer.from_pretrained(cls.model_name)
           _A					:     Optional[int]     = 'Translate in German: Hello, my dog is cute'


    def        A   (  self      :							Any):
           gc.collect()
           torch.cuda.empty_cache()


    def        A   (  self      :							List[str]):
           from transformers import TaForConditionalGeneration

           _A					:     Union[str, Any]     = TaForConditionalGeneration._keep_in_fpaa_modules
           _A					:     List[str]     = None

           # test with `t5-small`
           _A					:     Optional[int]     = TaForConditionalGeneration.from_pretrained(self.model_name				,							load_in_abit=SCREAMING_SNAKE_CASE				,							device_map='auto')
           _A					:     int     = self.tokenizer(self.input_text				,							return_tensors='pt').to(0)
           _A					:     List[str]     = model.generate(**SCREAMING_SNAKE_CASE)

           # test with `flan-t5-small`
           _A					:     List[str]     = TaForConditionalGeneration.from_pretrained(
               self.dense_act_model_name				,							load_in_abit=SCREAMING_SNAKE_CASE				,							device_map='auto')
           _A					:     Dict     = self.tokenizer(self.input_text				,							return_tensors='pt').to(0)
           _A					:     Any     = model.generate(**SCREAMING_SNAKE_CASE)
           _A					:     List[str]     = modules


    def        A   (  self      :							Dict):
           import bitsandbytes as bnb

           from transformers import TaForConditionalGeneration

           # test with `t5-small`
           _A					:     str     = TaForConditionalGeneration.from_pretrained(self.model_name				,							load_in_abit=SCREAMING_SNAKE_CASE				,							device_map='auto')

           # there was a bug with decoders - this test checks that it is fixed
           self.assertTrue(isinstance(model.decoder.block[0].layer[0].SelfAttention.q				,							bnb.nn.Linearabit))

           _A					:     int     = self.tokenizer(self.input_text				,							return_tensors='pt').to(0)
           _A					:     Union[str, Any]     = model.generate(**SCREAMING_SNAKE_CASE)

           # test with `flan-t5-small`
           _A					:     Union[str, Any]     = TaForConditionalGeneration.from_pretrained(
               self.dense_act_model_name				,							load_in_abit=SCREAMING_SNAKE_CASE				,							device_map='auto')
           _A					:     str     = self.tokenizer(self.input_text				,							return_tensors='pt').to(0)
           _A					:     Optional[int]     = model.generate(**SCREAMING_SNAKE_CASE)




class 			__lowerCamelCase			(	a_	):
    """simple docstring"""


    def        A   (  self      :							Union[str, Any]):
           super().setUp()
           # model_name
           _A					:     Dict     = 'bigscience/bloom-560m'
           _A					:     List[str]     = 't5-small'

           # Different types of model

           _A					:     Optional[Any]     = AutoModel.from_pretrained(self.model_name				,							load_in_abit=SCREAMING_SNAKE_CASE				,							device_map='auto')
           # Sequence classification model
           _A					:     Tuple     = AutoModelForSequenceClassification.from_pretrained(
               self.model_name				,							load_in_abit=SCREAMING_SNAKE_CASE				,							device_map='auto')
           # CausalLM model
           _A					:     List[Any]     = AutoModelForCausalLM.from_pretrained(self.model_name				,							load_in_abit=SCREAMING_SNAKE_CASE				,							device_map='auto')
           # Seq2seq model
           _A					:     Dict     = AutoModelForSeqaSeqLM.from_pretrained(
               self.seq_to_seq_name				,							load_in_abit=SCREAMING_SNAKE_CASE				,							device_map='auto')


    def        A   (  self      :							str):
           del self.base_model
           del self.sequence_model
           del self.model_abit
           del self.seq_to_seq_model

           gc.collect()
           torch.cuda.empty_cache()


    def        A   (  self      :							Any):
           from bitsandbytes.nn import Paramsabit

           self.assertTrue(self.base_model.h[-1].mlp.dense_ah_to_h.weight.__class__ == Paramsabit)

           # Other heads should be nn.Parameter
           self.assertTrue(self.model_abit.lm_head.weight.__class__ == torch.nn.Parameter)
           self.assertTrue(self.sequence_model.score.weight.__class__ == torch.nn.Parameter)
           self.assertTrue(self.seq_to_seq_model.lm_head.weight.__class__ == torch.nn.Parameter)




class 			__lowerCamelCase			(	a_	):
    """simple docstring"""


    def        A   (  self      :							Optional[Any]):
           super().setUp()


    def        A   (  self      :							Any):
           del self.pipe

           gc.collect()
           torch.cuda.empty_cache()


    def        A   (  self      :							Optional[int]):
           _A					:     Dict     = pipeline(
               'text-generation'				,							model=self.model_name				,							model_kwargs={'device_map': 'auto', 'load_in_4bit': True, 'torch_dtype': torch.floataa}				,							max_new_tokens=self.MAX_NEW_TOKENS				,							)

           # Real second forward pass
           _A					:     Union[str, Any]     = self.pipe(self.input_text)
           self.assertIn(pipeline_output[0]['generated_text']				,							self.EXPECTED_OUTPUTS)




@require_torch_multi_gpu
class 			__lowerCamelCase			(	a_	):
    """simple docstring"""


    def        A   (  self      :							List[Any]):
           super().setUp()


    def        A   (  self      :							Any):
           _A					:     int     = AutoModelForCausalLM.from_pretrained(
               self.model_name				,							load_in_abit=SCREAMING_SNAKE_CASE				,							device_map='balanced')

           # Check correct device map
           self.assertEqual(set(model_parallel.hf_device_map.values())				,							{0, 1})

           # Check that inference pass works on the model
           _A					:     Tuple     = self.tokenizer(self.input_text				,							return_tensors='pt')

           # Second real batch
           _A					:     Union[str, Any]     = model_parallel.generate(input_ids=encoded_input['input_ids'].to(0)				,							max_new_tokens=10)
           self.assertIn(self.tokenizer.decode(output_parallel[0]				,							skip_special_tokens=SCREAMING_SNAKE_CASE)				,							self.EXPECTED_OUTPUTS)




class 			__lowerCamelCase			(	a_	):
    """simple docstring"""


    def        A   (  self      :							Optional[Any]):
           _A					:     Any     = 'facebook/opt-350m'
           super().setUp()


    def        A   (  self      :							List[Any]):
           if version.parse(importlib.metadata.version('bitsandbytes')) < version.parse('0.37.0'):
                  return

           # Step 1: freeze all parameters
           _A					:     Any     = AutoModelForCausalLM.from_pretrained(self.model_name				,							load_in_abit=SCREAMING_SNAKE_CASE)

           self.assertEqual(set(model.hf_device_map.values())				,							{torch.cuda.current_device()})

           for param in model.parameters():
                  _A					:     Optional[Any]     = False  # freeze the model - train adapters later
                  if param.ndim == 1:
                         # cast the small parameters (e.g. layernorm) to fp32 for stability
                         _A					:     Tuple     = param.data.to(torch.floataa)

        # Step 2: add adapters
           for _, module in model.named_modules():
                  if "OPTAttention" in repr(type(SCREAMING_SNAKE_CASE)):
                         _A					:     List[Any]     = LoRALayer(module.q_proj				,							rank=16)
                         _A					:     Union[str, Any]     = LoRALayer(module.k_proj				,							rank=16)
                         _A					:     str     = LoRALayer(module.v_proj				,							rank=16)

        # Step 3: dummy batch
           _A					:     Union[str, Any]     = self.tokenizer('Test batch '				,							return_tensors='pt').to(0)

           # Step 4: Check if the gradient is not None
           with torch.cuda.amp.autocast():
                  _A					:     Union[str, Any]     = model.forward(**SCREAMING_SNAKE_CASE)
                  out.logits.norm().backward()

           for module in model.modules():
                  if isinstance(SCREAMING_SNAKE_CASE				,							SCREAMING_SNAKE_CASE):
                         self.assertTrue(module.adapter[1].weight.grad is not None)
                         self.assertTrue(module.adapter[1].weight.grad.norm().item() > 0)
                  elif isinstance(SCREAMING_SNAKE_CASE				,							nn.Embedding):
                         self.assertTrue(module.weight.grad is None)




class 			__lowerCamelCase			(	a_	):
    """simple docstring"""
    a													=   "gpt2-xl"
    a													=   3.3191_8548_5415_2187






