

import copy
import logging
import random

import joblib
import numpy as np
import torch
import torch.nn as nn
from torch.utils.data import DataLoader
from tqdm import tqdm

from transformers import AdamW, GPTaLMHeadModel, get_linear_schedule_with_warmup


_SCREAMING_SNAKE_CASE          = logging.getLogger(__name__)





def    _snake_case		(snake_case_				:    Union[str, Any]     )       ->		str:
  random.seed(snake_case_     )
  np.random.seed(snake_case_     )
  torch.manual_seed(snake_case_     )
  torch.cuda.manual_seed_all(snake_case_     )





def    _snake_case		(snake_case_				:    List[Any]						,					snake_case_				:    str						,					snake_case_				:    int     )       ->		Any:
  model.eval()
  __UpperCAmelCase  :     int		=    next(model.parameters()     ).device
  __UpperCAmelCase  :     Optional[Any]		=    1
  __UpperCAmelCase  :     Union[str, Any]		=    torch.zeros((eval_batch_size, context_len)						,					dtype=torch.long						,					device=snake_case_     )
  __UpperCAmelCase  :     List[Any]		=    DataLoader(snake_case_						,					shuffle=snake_case_						,					batch_size=snake_case_     )
  __UpperCAmelCase  :     List[Any]		=    torch.zeros(1						,					device=snake_case_     )
  __UpperCAmelCase  :     Optional[int]		=    0
  for batch in eval_dataloader:
    batch.to(snake_case_     )
    # pad
    context.zero_()
    for i in range(snake_case_     ):
      __UpperCAmelCase  :     Union[str, Any]		=    batch[i]
    __UpperCAmelCase  :     Optional[int]		=    model(snake_case_						,					labels=snake_case_     )
    eval_loss += outputs[0].sum().item()
    nb_eval_examples += batch.size(0     )
  __UpperCAmelCase  :     Optional[int]		=    eval_loss / nb_eval_examples
  __UpperCAmelCase  :     List[str]		=    torch.exp(snake_case_     )
  model.train()
  return perplexity





def    _snake_case		(snake_case_				:    Optional[int]="gpt2"     )       ->		Optional[Any]:
  __UpperCAmelCase  :     Dict		=    GPTaLMHeadModel.from_pretrained(snake_case_						,					output_hidden_states=snake_case_     )
  torch.save(model.state_dict()						,					model_name + 'local.pt'     )
  return model





def    _snake_case		(snake_case_				:    Any						,					snake_case_				:    Optional[int]						,					snake_case_				:    List[str]     )       ->		Dict:
  __UpperCAmelCase  :     List[str]		=    copy.deepcopy(snake_case_     )
  model.to(snake_case_     )

  __UpperCAmelCase  :     Optional[int]		=    ['bias', 'LayerNorm.weight']
  __UpperCAmelCase  :     int		=    [
      {
          'params': [p for n, p in model.named_parameters() if not any(nd in n for nd in no_decay     )],
          'weight_decay': 0.0,
      },
      {'params': [p for n, p in model.named_parameters() if any(nd in n for nd in no_decay     )], 'weight_decay': 0.0},
  ]
  __UpperCAmelCase  :     List[Any]		=    AdamW(snake_case_						,					lr=5e-5						,					eps=1e-8     )
  __UpperCAmelCase  :     int		=    get_linear_schedule_with_warmup(snake_case_						,					0						,					snake_case_     )
  torch.cuda.empty_cache()
  return model, lm_optimizer, lm_scheduler





def    _snake_case		(snake_case_				:    List[str]						,					snake_case_				:    List[str]						,					snake_case_				:    Union[str, Any]						,					snake_case_				:    Tuple     )       ->		Tuple:
  __UpperCAmelCase  :     Optional[Any]		=    np.array(snake_case_     ).mean()
  __UpperCAmelCase  :     str		=    np.array(snake_case_     ).std()
  __UpperCAmelCase  :     List[str]		=    (real_perps - avg) / std
  __UpperCAmelCase  :     List[str]		=    list(zip(snake_case_						,					snake_case_						,					snake_case_     )     )
  joblib.dump(snake_case_						,					snake_case_     )





def    _snake_case		(snake_case_				:    str						,					snake_case_				:    Optional[Any]						,					snake_case_				:    Optional[int]						,					snake_case_				:    List[str]						,					snake_case_				:    Any						,					snake_case_				:    Dict						,					snake_case_				:    int						,					snake_case_				:    Optional[int]="dev.jbl"						,					snake_case_				:    Tuple=recopy_gpta						,					)       ->		Dict:
  __UpperCAmelCase  :     Dict		=    []
  __UpperCAmelCase  :     int		=    []
  __UpperCAmelCase  :     Optional[Any]		=    []

  # Initialize the transformer model
  __UpperCAmelCase  :     Tuple		=    copy.deepcopy(snake_case_     )
  orig_model.to(device='cpu'     )
  torch.cuda.empty_cache()

  # Compute perplexity of initial transformer model for comparison
  model.train()
  __UpperCAmelCase, __UpperCAmelCase, __UpperCAmelCase  :     str		=    recopy_model(snake_case_						,					snake_case_						,					snake_case_     )

  for step in tqdm(range(snake_case_     )     ):
    __UpperCAmelCase  :     Optional[int]		=    torch.zeros((1, context_len)						,					dtype=torch.long						,					device=snake_case_     )
    __UpperCAmelCase  :     List[Any]		=    random.choice(snake_case_     )
    __UpperCAmelCase  :     Optional[Any]		=    random.randint(0						,					len(story[0]     ) - context_len - 1     )
    __UpperCAmelCase  :     Union[str, Any]		=    story[0][start : start + context_len]
    lm_optimizer.zero_grad()
    __UpperCAmelCase  :     List[Any]		=    model(snake_case_						,					labels=snake_case_     )
    __UpperCAmelCase  :     Optional[int]		=    outputs[0]
    __UpperCAmelCase  :     Union[str, Any]		=    compute_perplexity(snake_case_						,					snake_case_						,					snake_case_     )
    model.train()
    lm_loss.backward()
    # Do LM backprop
    torch.nn.utils.clip_grad_norm_(model.parameters()						,					3.0     )
    lm_optimizer.step()
    lm_scheduler.step()  # Update learning rate schedule

    # Compute perplexity after back-propagating on the selected context
    __UpperCAmelCase  :     List[Any]		=    compute_perplexity(snake_case_						,					snake_case_						,					snake_case_     )

    # Periodically save the stored (X, IG(X)) pairs
    if step % 1000 == 0 and step > 1:
      intermittent_save(snake_case_						,					snake_case_						,					snake_case_						,					snake_case_     )

    # Reset the pretrained model to the original pretrained GPT-2 weights after each iteration
    __UpperCAmelCase, __UpperCAmelCase, __UpperCAmelCase  :     int		=    recopy_model(snake_case_						,					snake_case_						,					snake_case_     )

    past_perps.append(past_perp.item()     )
    real_perps.append(orig_perp - real_perp.item()     )
    contexts.append(np.array(context.cpu()     )     )

  intermittent_save(snake_case_						,					snake_case_						,					snake_case_						,					snake_case_     )





def    _snake_case		(snake_case_				:    str						,					snake_case_				:    int="data/tokenized_stories_train_wikitext103.jbl"						,					snake_case_				:    Tuple=100						,					snake_case_				:    str=1026						,					snake_case_				:    Union[str, Any]=True     )       ->		List[str]:
  __UpperCAmelCase  :     Union[str, Any]		=    joblib.load(snake_case_     )
  print('data loaded'     )
  __UpperCAmelCase  :     Any		=    []
  if trim:
    for i, example in enumerate(snake_case_     ):
      if len(example[0]     ) > min_len:
        __UpperCAmelCase  :     Optional[int]		=    random.randint(0						,					len(example[0]     ) - context_len - 1     )
        objective_set.append(example[0, start : start + context_len]     )
      if len(snake_case_     ) >= number:
        break
    __UpperCAmelCase  :     Union[str, Any]		=    []
    for j in range(i + 1						,					len(snake_case_     )     ):
      if len(data[j][0]     ) > min_len:
        train_data.append(data[j]     )
  else:
    __UpperCAmelCase  :     Any		=    data[0:number]
    __UpperCAmelCase  :     int		=    data[number:]

  joblib.dump(snake_case_						,					'objective_set.jbl'     )
  print('objective set saved'     )
  return train_data, objective_set





def    _snake_case		(snake_case_				:    Optional[Any]						,					snake_case_				:    Tuple						,					snake_case_				:    Optional[Any]						,					snake_case_				:    str						,					snake_case_				:    Dict=50						,					snake_case_				:    Any="secondary_learner.pt"     )       ->		Tuple:
  __UpperCAmelCase  :     int		=    torch.device('cuda:0' if torch.cuda.is_available() else 'cpu'     )
  # We will use the first 512 pairs from our dataset as a test set for
  # our secondary learner and the rest to train
  __UpperCAmelCase  :     int		=    train_dataset[:512]
  __UpperCAmelCase  :     Tuple		=    train_dataset[512:]
  __UpperCAmelCase  :     Any		=    DataLoader(snake_case_						,					shuffle=snake_case_						,					batch_size=snake_case_     )
  __UpperCAmelCase  :     str		=    DataLoader(snake_case_						,					shuffle=snake_case_						,					batch_size=snake_case_     )

  # secondary learner model set up
  __UpperCAmelCase  :     Any		=    nn.MSELoss()
  __UpperCAmelCase  :     Any		=    nn.MSELoss(reduction='sum'     )
  secondary_learner.to(snake_case_     )
  __UpperCAmelCase  :     Union[str, Any]		=    torch.optim.Adam(secondary_learner.parameters()						,					lr=0.00001     )
  secondary_learner.train()

  # TODO in original code this is written as number of actual batches seen
  # not number of items seen but other places it is number of items instead.
  # improve consistency! changed this to epochs for clarity
  __UpperCAmelCase  :     List[Any]		=    float('inf'     )
  # Iterate through batches until we've used max_steps batches
  for epoch in range(int(snake_case_     )     ):
    __UpperCAmelCase  :     int		=    0.0
    secondary_learner.train()
    for step, batch in enumerate(snake_case_     ):
      __UpperCAmelCase  :     Any		=    batch[0].to(snake_case_     )
      __UpperCAmelCase  :     Union[str, Any]		=    batch[1].to(snake_case_     )
      __UpperCAmelCase  :     List[Any]		=    secondary_learner(snake_case_     )
      q_optimizer.zero_grad()
      __UpperCAmelCase  :     Union[str, Any]		=    loss(snake_case_						,					real_q.float()     )
      q_loss.backward()
      q_optimizer.step()
      tr_q_loss += q_loss.item()

      # model trains fairly quickly so we won't wait for a full epoch
      # eval is triggered at eval_freq and end of epochs
      if (step % eval_freq == 0 and step > 0) or ((step + 1) == len(snake_case_     )):
        __UpperCAmelCase  :     Optional[int]		=    tr_q_loss / (step + 1)

        secondary_learner.eval()
        __UpperCAmelCase  :     List[str]		=    0.0
        __UpperCAmelCase  :     Tuple		=    0.0
        __UpperCAmelCase  :     int		=    []
        __UpperCAmelCase  :     Dict		=    []
        # Compute performance of the secondary learner after this batch
        for stepa, batcha in enumerate(snake_case_     ):
          __UpperCAmelCase  :     Optional[Any]		=    batcha[0].to(snake_case_     )
          __UpperCAmelCase  :     Optional[Any]		=    batcha[1].to(snake_case_     )
          __UpperCAmelCase  :     List[str]		=    secondary_learner(snake_case_     )
          q_lossa += test_loss(snake_case_						,					snake_case_     ).item()
          sum_qa += torch.sum(snake_case_     ).item()
          for ei, i in enumerate(predicted_qa.cpu().detach().numpy()     ):
            predicted.append(i.item()     )
          for ei, i in enumerate(real_qa.cpu().detach().numpy()     ):
            actual.append(i.item()     )

        q_lossa /= len(snake_case_     )
        print(
            'Epoch: '						,					snake_case_						,					'step: '						,					snake_case_						,					'Avg. q:'						,					sum_qa / len(snake_case_     )						,					'Train Loss: '						,					snake_case_						,					'Test Loss: '						,					snake_case_						,					)
        if q_lossa < best_test_loss:
          joblib.dump((predicted, actual)						,					'pred_vs_actual.jbl'     )
          torch.save(secondary_learner.state_dict()						,					snake_case_     )
          __UpperCAmelCase  :     int		=    q_lossa

      secondary_learner.train()
  return secondary_learner







class     lowerCamelCase_   ( nn.Module):




 def __init__( self     ,     _UpperCAmelCase       ):
   super(_UpperCAmelCase     ,     self       ).__init__()
   __UpperCAmelCase  :     List[str]		=    model.transformer.wte
   __UpperCAmelCase  :     Optional[int]		=    copy.deepcopy(model.transformer.wte.weight       )
   __UpperCAmelCase  :     str		=    nn.Convad(self.embeddings.weight.size(1       )     ,     256     ,     3     ,     padding=1       )
   __UpperCAmelCase  :     Dict		=    nn.Sequential(nn.Linear(256     ,     32       )     ,     nn.Dropout(p=0.1       )     ,     nn.Linear(32     ,     32       )     ,     nn.Linear(32     ,     1       )       )




 def 						a					( self     ,     _UpperCAmelCase       ):
   __UpperCAmelCase  :     Union[str, Any]		=    torch.max(self.conv(self.embeddings(_UpperCAmelCase       ).squeeze(1       ).transpose(1     ,     2       )       )     ,     2       )[0]
   __UpperCAmelCase  :     Dict		=    self.fc(_UpperCAmelCase       )
   return qs.squeeze(1       )




 @classmethod
 def 						a					( cls     ,     _UpperCAmelCase     ,     _UpperCAmelCase       ):
   __UpperCAmelCase  :     List[Any]		=    cls(_UpperCAmelCase       )  # this calls __init__
   __UpperCAmelCase  :     int		=    torch.load(_UpperCAmelCase       )
   secondary_learner.load_state_dict(_UpperCAmelCase       )
   __UpperCAmelCase  :     List[Any]		=    model.transformer.wte
   __UpperCAmelCase  :     List[Any]		=    copy.deepcopy(model.transformer.wte.weight       )
   return secondary_learner








import inspect
from typing import List, Optional, Union

import torch
from torch import nn
from torch.nn import functional as F
from torchvision import transforms
from transformers import CLIPImageProcessor, CLIPModel, CLIPTextModel, CLIPTokenizer

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


class     lowerCamelCase_   ( nn.Module):




 def __init__( self     ,     _UpperCAmelCase     ,     _UpperCAmelCase=1.0       ):
   super().__init__()

   __UpperCAmelCase  :     List[str]		=    cut_size
   __UpperCAmelCase  :     Dict		=    cut_power




 def 						a					( self     ,     _UpperCAmelCase     ,     _UpperCAmelCase       ):
   __UpperCAmelCase, __UpperCAmelCase  :     List[Any]		=    pixel_values.shape[2:4]
   __UpperCAmelCase  :     int		=    min(_UpperCAmelCase     ,     _UpperCAmelCase       )
   __UpperCAmelCase  :     int		=    min(_UpperCAmelCase     ,     _UpperCAmelCase     ,     self.cut_size       )
   __UpperCAmelCase  :     List[Any]		=    []
   for _ in range(_UpperCAmelCase       ):
     __UpperCAmelCase  :     str		=    int(torch.rand([]       ) ** self.cut_power * (max_size - min_size) + min_size       )
     __UpperCAmelCase  :     Optional[Any]		=    torch.randint(0     ,     sideX - size + 1     ,     ()       )
     __UpperCAmelCase  :     Any		=    torch.randint(0     ,     sideY - size + 1     ,     ()       )
     __UpperCAmelCase  :     int		=    pixel_values[:, :, offsety : offsety + size, offsetx : offsetx + size]
     cutouts.append(F.adaptive_avg_poolad(_UpperCAmelCase     ,     self.cut_size       )       )
   return torch.cat(_UpperCAmelCase       )





def    _snake_case		(snake_case_				:    Optional[int]						,					snake_case_				:    Optional[int]     )       ->		Dict:
  __UpperCAmelCase  :     Any		=    F.normalize(snake_case_						,					dim=-1     )
  __UpperCAmelCase  :     Any		=    F.normalize(snake_case_						,					dim=-1     )
  return (x - y).norm(dim=-1     ).div(2     ).arcsin().pow(2     ).mul(2     )





def    _snake_case		(snake_case_				:    Optional[int]						,					snake_case_				:    Optional[Any]     )       ->		Tuple:
  for param in model.parameters():
    __UpperCAmelCase  :     int		=    value







class     lowerCamelCase_   ( _lowerCAmelCase):




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

   __UpperCAmelCase  :     int		=    transforms.Normalize(mean=feature_extractor.image_mean     ,     std=feature_extractor.image_std       )
   __UpperCAmelCase  :     Dict		=    (
       feature_extractor.size
       if isinstance(feature_extractor.size     ,     _UpperCAmelCase       )
       else feature_extractor.size['shortest_edge']
   )
   __UpperCAmelCase  :     Optional[Any]		=    MakeCutouts(self.cut_out_size       )

   set_requires_grad(self.text_encoder     ,     _UpperCAmelCase       )
   set_requires_grad(self.clip_model     ,     _UpperCAmelCase       )




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




 def 						a					( self       ):
   self.enable_attention_slicing(_UpperCAmelCase       )




 def 						a					( self       ):
   set_requires_grad(self.vae     ,     _UpperCAmelCase       )




 def 						a					( self       ):
   set_requires_grad(self.vae     ,     _UpperCAmelCase       )




 def 						a					( self       ):
   set_requires_grad(self.unet     ,     _UpperCAmelCase       )




 def 						a					( self       ):
   set_requires_grad(self.unet     ,     _UpperCAmelCase       )




 @torch.enable_grad()
 def 						a					( self     ,     _UpperCAmelCase     ,     _UpperCAmelCase     ,     _UpperCAmelCase     ,     _UpperCAmelCase     ,     _UpperCAmelCase     ,     _UpperCAmelCase     ,     _UpperCAmelCase     ,     _UpperCAmelCase     ,     _UpperCAmelCase=True     ,     ):
   __UpperCAmelCase  :     List[str]		=    latents.detach().requires_grad_()

   __UpperCAmelCase  :     Any		=    self.scheduler.scale_model_input(_UpperCAmelCase     ,     _UpperCAmelCase       )

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

   if isinstance(self.scheduler     ,     (PNDMScheduler, DDIMScheduler, DPMSolverMultistepScheduler)       ):
     __UpperCAmelCase  :     Any		=    self.scheduler.alphas_cumprod[timestep]
     __UpperCAmelCase  :     int		=    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
     __UpperCAmelCase  :     Optional[int]		=    (latents - beta_prod_t ** 0.5 * noise_pred) / alpha_prod_t ** 0.5

     __UpperCAmelCase  :     Optional[int]		=    torch.sqrt(_UpperCAmelCase       )
     __UpperCAmelCase  :     Optional[int]		=    pred_original_sample * (fac) + latents * (1 - fac)
   elif isinstance(self.scheduler     ,     _UpperCAmelCase       ):
     __UpperCAmelCase  :     List[str]		=    self.scheduler.sigmas[index]
     __UpperCAmelCase  :     int		=    latents - sigma * noise_pred
   else:
     raise ValueError(F"scheduler type {type(self.scheduler       )} not supported"       )

   __UpperCAmelCase  :     int		=    1 / self.vae.config.scaling_factor * sample
   __UpperCAmelCase  :     Union[str, Any]		=    self.vae.decode(_UpperCAmelCase       ).sample
   __UpperCAmelCase  :     Any		=    (image / 2 + 0.5).clamp(0     ,     1       )

   if use_cutouts:
     __UpperCAmelCase  :     str		=    self.make_cutouts(_UpperCAmelCase     ,     _UpperCAmelCase       )
   else:
     __UpperCAmelCase  :     List[str]		=    transforms.Resize(self.cut_out_size       )(_UpperCAmelCase       )
   __UpperCAmelCase  :     Any		=    self.normalize(_UpperCAmelCase       ).to(latents.dtype       )

   __UpperCAmelCase  :     Union[str, Any]		=    self.clip_model.get_image_features(_UpperCAmelCase       )
   __UpperCAmelCase  :     str		=    image_embeddings_clip / image_embeddings_clip.norm(p=2     ,     dim=-1     ,     keepdim=_UpperCAmelCase       )

   if use_cutouts:
     __UpperCAmelCase  :     Optional[Any]		=    spherical_dist_loss(_UpperCAmelCase     ,     _UpperCAmelCase       )
     __UpperCAmelCase  :     Dict		=    dists.view([num_cutouts, sample.shape[0], -1]       )
     __UpperCAmelCase  :     Any		=    dists.sum(2       ).mean(0       ).sum() * clip_guidance_scale
   else:
     __UpperCAmelCase  :     Union[str, Any]		=    spherical_dist_loss(_UpperCAmelCase     ,     _UpperCAmelCase       ).mean() * clip_guidance_scale

   __UpperCAmelCase  :     Dict		=    -torch.autograd.grad(_UpperCAmelCase     ,     _UpperCAmelCase       )[0]

   if isinstance(self.scheduler     ,     _UpperCAmelCase       ):
     __UpperCAmelCase  :     Tuple		=    latents.detach() + grads * (sigma**2)
     __UpperCAmelCase  :     int		=    noise_pred_original
   else:
     __UpperCAmelCase  :     Tuple		=    noise_pred_original - torch.sqrt(_UpperCAmelCase       ) * grads
   return noise_pred, latents




 @torch.no_grad()
 def __call__( self     ,     _UpperCAmelCase     ,     _UpperCAmelCase = 512     ,     _UpperCAmelCase = 512     ,     _UpperCAmelCase = 50     ,     _UpperCAmelCase = 7.5     ,     _UpperCAmelCase = 1     ,     _UpperCAmelCase = 0.0     ,     _UpperCAmelCase = 100     ,     _UpperCAmelCase = None     ,     _UpperCAmelCase = 4     ,     _UpperCAmelCase = True     ,     _UpperCAmelCase = None     ,     _UpperCAmelCase = None     ,     _UpperCAmelCase = "pil"     ,     _UpperCAmelCase = True     ,     ):
   if isinstance(_UpperCAmelCase     ,     _UpperCAmelCase       ):
     __UpperCAmelCase  :     str		=    1
   elif isinstance(_UpperCAmelCase     ,     _UpperCAmelCase       ):
     __UpperCAmelCase  :     List[str]		=    len(_UpperCAmelCase       )
   else:
     raise ValueError(F"`prompt` has to be of type `str` or `list` but is {type(_UpperCAmelCase       )}"       )

   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}."       )

   # get prompt text embeddings
   __UpperCAmelCase  :     Any		=    self.tokenizer(
       _UpperCAmelCase     ,     padding='max_length'     ,     max_length=self.tokenizer.model_max_length     ,     truncation=_UpperCAmelCase     ,     return_tensors='pt'     ,     )
   __UpperCAmelCase  :     str		=    self.text_encoder(text_input.input_ids.to(self.device       )       )[0]
   # duplicate text embeddings for each generation per prompt
   __UpperCAmelCase  :     Optional[int]		=    text_embeddings.repeat_interleave(_UpperCAmelCase     ,     dim=0       )

   if clip_guidance_scale > 0:
     if clip_prompt is not None:
       __UpperCAmelCase  :     Dict		=    self.tokenizer(
           _UpperCAmelCase     ,     padding='max_length'     ,     max_length=self.tokenizer.model_max_length     ,     truncation=_UpperCAmelCase     ,     return_tensors='pt'     ,     ).input_ids.to(self.device       )
     else:
       __UpperCAmelCase  :     Optional[Any]		=    text_input.input_ids.to(self.device       )
     __UpperCAmelCase  :     Optional[Any]		=    self.clip_model.get_text_features(_UpperCAmelCase       )
     __UpperCAmelCase  :     Any		=    text_embeddings_clip / text_embeddings_clip.norm(p=2     ,     dim=-1     ,     keepdim=_UpperCAmelCase       )
     # duplicate text embeddings clip for each generation per prompt
     __UpperCAmelCase  :     int		=    text_embeddings_clip.repeat_interleave(_UpperCAmelCase     ,     dim=0       )

   # 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.
   __UpperCAmelCase  :     List[str]		=    guidance_scale > 1.0
   # get unconditional embeddings for classifier free guidance
   if do_classifier_free_guidance:
     __UpperCAmelCase  :     Union[str, Any]		=    text_input.input_ids.shape[-1]
     __UpperCAmelCase  :     Tuple		=    self.tokenizer(['']     ,     padding='max_length'     ,     max_length=_UpperCAmelCase     ,     return_tensors='pt'       )
     __UpperCAmelCase  :     Dict		=    self.text_encoder(uncond_input.input_ids.to(self.device       )       )[0]
     # duplicate unconditional embeddings for each generation per prompt
     __UpperCAmelCase  :     List[str]		=    uncond_embeddings.repeat_interleave(_UpperCAmelCase     ,     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
     __UpperCAmelCase  :     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`.
   __UpperCAmelCase  :     List[Any]		=    (batch_size * num_images_per_prompt, self.unet.config.in_channels, height // 8, width // 8)
   __UpperCAmelCase  :     Dict		=    text_embeddings.dtype
   if latents is None:
     if self.device.type == "mps":
       # randn does not work reproducibly on mps
       __UpperCAmelCase  :     Tuple		=    torch.randn(_UpperCAmelCase     ,     generator=_UpperCAmelCase     ,     device='cpu'     ,     dtype=_UpperCAmelCase       ).to(
           self.device       )
     else:
       __UpperCAmelCase  :     Tuple		=    torch.randn(_UpperCAmelCase     ,     generator=_UpperCAmelCase     ,     device=self.device     ,     dtype=_UpperCAmelCase       )
   else:
     if latents.shape != latents_shape:
       raise ValueError(F"Unexpected latents shape, got {latents.shape}, expected {latents_shape}"       )
     __UpperCAmelCase  :     str		=    latents.to(self.device       )

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

   self.scheduler.set_timesteps(_UpperCAmelCase     ,     **_UpperCAmelCase       )

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

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

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

   for i, t in enumerate(self.progress_bar(_UpperCAmelCase       )       ):
     # expand the latents if we are doing classifier free guidance
     __UpperCAmelCase  :     Optional[int]		=    torch.cat([latents] * 2       ) if do_classifier_free_guidance else latents
     __UpperCAmelCase  :     Union[str, Any]		=    self.scheduler.scale_model_input(_UpperCAmelCase     ,     _UpperCAmelCase       )

     # predict the noise residual
     __UpperCAmelCase  :     Optional[int]		=    self.unet(_UpperCAmelCase     ,     _UpperCAmelCase     ,     encoder_hidden_states=_UpperCAmelCase       ).sample

     # perform classifier free guidance
     if do_classifier_free_guidance:
       __UpperCAmelCase, __UpperCAmelCase  :     Tuple		=    noise_pred.chunk(2       )
       __UpperCAmelCase  :     int		=    noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)

     # perform clip guidance
     if clip_guidance_scale > 0:
       __UpperCAmelCase  :     int		=    (
           text_embeddings.chunk(2       )[1] if do_classifier_free_guidance else text_embeddings
       )
       __UpperCAmelCase, __UpperCAmelCase  :     Optional[int]		=    self.cond_fn(
           _UpperCAmelCase     ,     _UpperCAmelCase     ,     _UpperCAmelCase     ,     _UpperCAmelCase     ,     _UpperCAmelCase     ,     _UpperCAmelCase     ,     _UpperCAmelCase     ,     _UpperCAmelCase     ,     _UpperCAmelCase     ,     )

     # compute the previous noisy sample x_t -> x_t-1
     __UpperCAmelCase  :     Optional[Any]		=    self.scheduler.step(_UpperCAmelCase     ,     _UpperCAmelCase     ,     _UpperCAmelCase     ,     **_UpperCAmelCase       ).prev_sample

   # scale and decode the image latents with vae
   __UpperCAmelCase  :     Optional[Any]		=    1 / self.vae.config.scaling_factor * latents
   __UpperCAmelCase  :     List[str]		=    self.vae.decode(_UpperCAmelCase       ).sample

   __UpperCAmelCase  :     str		=    (image / 2 + 0.5).clamp(0     ,     1       )
   __UpperCAmelCase  :     str		=    image.cpu().permute(0     ,     2     ,     3     ,     1       ).numpy()

   if output_type == "pil":
     __UpperCAmelCase  :     Tuple		=    self.numpy_to_pil(_UpperCAmelCase       )

   if not return_dict:
     return (image, None)

   return StableDiffusionPipelineOutput(images=_UpperCAmelCase     ,     nsfw_content_detected=_UpperCAmelCase       )





