






'''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


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







class 	UpperCAmelCase	(		__a):


 '''simple docstring'''




 def __init__(   self							,  lowerCAmelCase_							,  lowerCAmelCase_							,  lowerCAmelCase_							,  lowerCAmelCase_							,  lowerCAmelCase_							,  lowerCAmelCase_							,  lowerCAmelCase_							,  ) ->      List[Any]:



      """simple docstring"""







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




 def       lowercase_  (   self							,  lowerCAmelCase_ = "auto") ->      Dict:



      """simple docstring"""







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




 def       lowercase_  (   self) ->      Tuple:



      """simple docstring"""







      self.enable_attention_slicing(lowerCAmelCase_)




 @torch.no_grad()
 def __call__(   self							,  lowerCAmelCase_							,  lowerCAmelCase_ = 5_1_2							,  lowerCAmelCase_ = 5_1_2							,  lowerCAmelCase_ = 5_0							,  lowerCAmelCase_ = 7.5							,  lowerCAmelCase_ = None							,  lowerCAmelCase_ = 1							,  lowerCAmelCase_ = 0.0							,  lowerCAmelCase_ = None							,  lowerCAmelCase_ = None							,  lowerCAmelCase_ = "pil"							,  lowerCAmelCase_ = True							,  lowerCAmelCase_ = None							,  lowerCAmelCase_ = 1							,  lowerCAmelCase_ = None							,  **lowerCAmelCase_							,  ) ->      List[str]:



      """simple docstring"""







      if isinstance(lowerCAmelCase_							,  lowerCAmelCase_):
           a_       		=1
      elif isinstance(lowerCAmelCase_							,  lowerCAmelCase_):
           a_       		=len(lowerCAmelCase_)
      else:
           raise ValueError(f"""`prompt` has to be of type `str` or `list` but is {type(lowerCAmelCase_)}""")

      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(lowerCAmelCase_							,  lowerCAmelCase_) or callback_steps <= 0)
      ):
           raise ValueError(
               f"""`callback_steps` has to be a positive integer but is {callback_steps} of type"""
               f""" {type(lowerCAmelCase_)}.""")

      # get prompt text embeddings
      a_       		=self.tokenizer(
          lowerCAmelCase_							,  padding="max_length"							,  max_length=self.tokenizer.model_max_length							,  return_tensors="pt"							,  )
      a_       		=text_inputs.input_ids

      if text_input_ids.shape[-1] > self.tokenizer.model_max_length:
           a_       		=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_       		=text_input_ids[:, : self.tokenizer.model_max_length]

      if text_embeddings is None:
           a_       		=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_       		=text_embeddings.shape
      a_       		=text_embeddings.repeat(1							,  lowerCAmelCase_							,  1)
      a_       		=text_embeddings.view(bs_embed * num_images_per_prompt							,  lowerCAmelCase_							,  -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_       		=guidance_scale > 1.0
      # get unconditional embeddings for classifier free guidance
      if do_classifier_free_guidance:
           a_       		=4_2
           if negative_prompt is None:
                a_       		=[""]
           elif type(lowerCAmelCase_) is not type(lowerCAmelCase_):
                raise TypeError(
                    f"""`negative_prompt` should be the same type to `prompt`, but got {type(lowerCAmelCase_)} !="""
                    f""" {type(lowerCAmelCase_)}.""")
           elif isinstance(lowerCAmelCase_							,  lowerCAmelCase_):
                a_       		=[negative_prompt]
           elif batch_size != len(lowerCAmelCase_):
                raise ValueError(
                    f"""`negative_prompt`: {negative_prompt} has batch size {len(lowerCAmelCase_)}, but `prompt`:"""
                    f""" {prompt} has batch size {batch_size}. Please make sure that passed `negative_prompt` matches"""
                    " the batch size of `prompt`.")
           else:
                a_       		=negative_prompt

           a_       		=text_input_ids.shape[-1]
           a_       		=self.tokenizer(
               lowerCAmelCase_							,  padding="max_length"							,  max_length=lowerCAmelCase_							,  truncation=lowerCAmelCase_							,  return_tensors="pt"							,  )
           a_       		=self.text_encoder(uncond_input.input_ids.to(self.device))[0]

           # duplicate unconditional embeddings for each generation per prompt, using mps friendly method
           a_       		=uncond_embeddings.shape[1]
           a_       		=uncond_embeddings.repeat(lowerCAmelCase_							,  lowerCAmelCase_							,  1)
           a_       		=uncond_embeddings.view(batch_size * num_images_per_prompt							,  lowerCAmelCase_							,  -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_       		=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_       		=(batch_size * num_images_per_prompt, self.unet.config.in_channels, height // 8, width // 8)
      a_       		=(batch_size * num_images_per_prompt, self.unet.config.in_channels, 6_4, 6_4)
      a_       		=text_embeddings.dtype
      if latents is None:
           if self.device.type == "mps":
                # randn does not exist on mps
                a_       		=torch.randn(
                    lowerCAmelCase_							,  generator=lowerCAmelCase_							,  device="cpu"							,  dtype=lowerCAmelCase_).to(self.device)
                a_       		=torch.randn(lowerCAmelCase_							,  generator=lowerCAmelCase_							,  device="cpu"							,  dtype=lowerCAmelCase_).to(
                    self.device)
           else:
                a_       		=torch.randn(
                    lowerCAmelCase_							,  generator=lowerCAmelCase_							,  device=self.device							,  dtype=lowerCAmelCase_)
                a_       		=torch.randn(lowerCAmelCase_							,  generator=lowerCAmelCase_							,  device=self.device							,  dtype=lowerCAmelCase_)
      else:
           if latents_reference.shape != latents_shape:
                raise ValueError(f"""Unexpected latents shape, got {latents.shape}, expected {latents_shape}""")
           a_       		=latents_reference.to(self.device)
           a_       		=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_       		=(latents_shape[3] - latents_shape_reference[3]) // 2
      a_       		=(latents_shape[2] - latents_shape_reference[2]) // 2
      a_       		=latents_shape_reference[3] if dx >= 0 else latents_shape_reference[3] + 2 * dx
      a_       		=latents_shape_reference[2] if dy >= 0 else latents_shape_reference[2] + 2 * dy
      a_       		=0 if dx < 0 else dx
      a_       		=0 if dy < 0 else dy
      a_       		=max(-dx							,  0)
      a_       		=max(-dy							,  0)
      # import pdb
      # pdb.set_trace()
      a_       		=latents_reference[:, :, dy : dy + h, dx : dx + w]

      # set timesteps
      self.scheduler.set_timesteps(lowerCAmelCase_)

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

      # scale the initial noise by the standard deviation required by the scheduler
      a_       		=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_       		="eta" in set(inspect.signature(self.scheduler.step).parameters.keys())
      a_       		={}
      if accepts_eta:
           a_       		=eta

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

           # predict the noise residual
           a_       		=self.unet(lowerCAmelCase_							,  lowerCAmelCase_							,  encoder_hidden_states=lowerCAmelCase_).sample

           # perform guidance
           if do_classifier_free_guidance:
                a_		,						a_       		=noise_pred.chunk(2)
                a_       		=noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)

           # compute the previous noisy sample x_t -> x_t-1
           a_       		=self.scheduler.step(lowerCAmelCase_							,  lowerCAmelCase_							,  lowerCAmelCase_							,  **lowerCAmelCase_).prev_sample

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

      a_       		=1 / 0.1_8_2_1_5 * latents
      a_       		=self.vae.decode(lowerCAmelCase_).sample

      a_       		=(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_       		=image.cpu().permute(0							,  2							,  3							,  1).float().numpy()

      if self.safety_checker is not None:
           a_       		=self.feature_extractor(self.numpy_to_pil(lowerCAmelCase_)							,  return_tensors="pt").to(
               self.device)
           a_		,						a_       		=self.safety_checker(
               images=lowerCAmelCase_							,  clip_input=safety_checker_input.pixel_values.to(text_embeddings.dtype))
      else:
           a_       		=None

      if output_type == "pil":
           a_       		=self.numpy_to_pil(lowerCAmelCase_)

      if not return_dict:
           return (image, has_nsfw_concept)

      return StableDiffusionPipelineOutput(images=lowerCAmelCase_							,  nsfw_content_detected=lowerCAmelCase_)







'''simple docstring'''


import torch

from diffusers import StableDiffusionPipeline


lowercase							       =			'''path-to-your-trained-model'''
lowercase							       =			StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.floataa).to('''cuda''')

lowercase							       =			'''A photo of sks dog in a bucket'''
lowercase							       =			pipe(prompt, num_inference_steps=50, guidance_scale=7.5).images[0]

image.save('''dog-bucket.png''')


