class_docstr: Flax-based pipeline for text-guided image inpainting using Stable Diffusion.
  
  <Tip warning={true}>
  
  🧪 This is an experimental feature!
  
  </Tip>
  
  This model inherits from [`FlaxDiffusionPipeline`]. Check the superclass documentation for the generic methods
  implemented for all pipelines (downloading, saving, running on a particular device, etc.).
  
  Args:
      vae ([`FlaxAutoencoderKL`]):
          Variational Auto-Encoder (VAE) model to encode and decode images to and from latent representations.
      text_encoder ([`~transformers.FlaxCLIPTextModel`]):
          Frozen text-encoder ([clip-vit-large-patch14](https://huggingface.co/openai/clip-vit-large-patch14)).
      tokenizer ([`~transformers.CLIPTokenizer`]):
          A `CLIPTokenizer` to tokenize text.
      unet ([`FlaxUNet2DConditionModel`]):
          A `FlaxUNet2DConditionModel` to denoise the encoded image latents.
      scheduler ([`SchedulerMixin`]):
          A scheduler to be used in combination with `unet` to denoise the encoded image latents. Can be one of
          [`FlaxDDIMScheduler`], [`FlaxLMSDiscreteScheduler`], [`FlaxPNDMScheduler`], or
          [`FlaxDPMSolverMultistepScheduler`].
      safety_checker ([`FlaxStableDiffusionSafetyChecker`]):
          Classification module that estimates whether generated images could be considered offensive or harmful.
          Please refer to the [model card](https://huggingface.co/runwayml/stable-diffusion-v1-5) for more details
          about a model's potential harms.
      feature_extractor ([`~transformers.CLIPImageProcessor`]):
          A `CLIPImageProcessor` to extract features from generated images; used as inputs to the `safety_checker`.
class_docstr_tok: ["Flax-based","pipeline","for","text-guided","image","inpainting","using","Stable","Diffusion",".","<","Tip","warning=","{","true","}",">","\ud83e\uddea","This","is","an","experimental","feature","!","<","\/Tip",">","This","model","inherits","from","[","`","FlaxDiffusionPipeline","`","]",".","Check","the","superclass","documentation","for","the","generic","methods","implemented","for","all","pipelines","(","downloading",",","saving",",","running","on","a","particular","device",",","etc.",")",".","Args",":","vae","(","[","`","FlaxAutoencoderKL","`","]",")",":","Variational","Auto-Encoder","(","VAE",")","model","to","encode","and","decode","images","to","and","from","latent","representations",".","text_encoder","(","[","`","~transformers.FlaxCLIPTextModel","`","]",")",":","Frozen","text-encoder","(","[","clip-vit-large-patch14","]","(","https",":","\/\/huggingface.co\/openai\/clip-vit-large-patch14",")",")",".","tokenizer","(","[","`","~transformers.CLIPTokenizer","`","]",")",":","A","`","CLIPTokenizer","`","to","tokenize","text",".","unet","(","[","`","FlaxUNet2DConditionModel","`","]",")",":","A","`","FlaxUNet2DConditionModel","`","to","denoise","the","encoded","image","latents",".","scheduler","(","[","`","SchedulerMixin","`","]",")",":","A","scheduler","to","be","used","in","combination","with","`","unet","`","to","denoise","the","encoded","image","latents",".","Can","be","one","of","[","`","FlaxDDIMScheduler","`","]",",","[","`","FlaxLMSDiscreteScheduler","`","]",",","[","`","FlaxPNDMScheduler","`","]",",","or","[","`","FlaxDPMSolverMultistepScheduler","`","]",".","safety_checker","(","[","`","FlaxStableDiffusionSafetyChecker","`","]",")",":","Classification","module","that","estimates","whether","generated","images","could","be","considered","offensive","or","harmful",".","Please","refer","to","the","[","model","card","]","(","https",":","\/\/huggingface.co\/runwayml\/stable-diffusion-v1-5",")","for","more","details","about","a","model","'s","potential","harms",".","feature_extractor","(","[","`","~transformers.CLIPImageProcessor","`","]",")",":","A","`","CLIPImageProcessor","`","to","extract","features","from","generated","images",";","used","as","inputs","to","the","`","safety_checker","`","."]
code: def _generate(
      self,
      prompt_ids: jnp.ndarray,
      mask: jnp.ndarray,
      masked_image: jnp.ndarray,
      params: Union[Dict, FrozenDict],
      prng_seed: jax.Array,
      num_inference_steps: int,
      height: int,
      width: int,
      guidance_scale: float,
      latents: Optional[jnp.ndarray] = None,
      neg_prompt_ids: Optional[jnp.ndarray] = None,
  ):
      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
      prompt_embeds = self.text_encoder(
          prompt_ids, params=params["text_encoder"]
      )[0]
  
      # TODO: currently it is assumed `do_classifier_free_guidance = guidance_scale > 1.0`
      # implement this conditional `do_classifier_free_guidance = guidance_scale > 1.0`
      batch_size = prompt_ids.shape[0]
  
      max_length = prompt_ids.shape[-1]
  
      if neg_prompt_ids is None:
          uncond_input = self.tokenizer(
              [""] * batch_size,
              padding="max_length",
              max_length=max_length,
              return_tensors="np",
          ).input_ids
      else:
          uncond_input = neg_prompt_ids
      negative_prompt_embeds = self.text_encoder(
          uncond_input, params=params["text_encoder"]
      )[0]
      context = jnp.concatenate([negative_prompt_embeds, prompt_embeds])
  
      latents_shape = (
          batch_size,
          self.vae.config.latent_channels,
          height // self.vae_scale_factor,
          width // self.vae_scale_factor,
      )
      if latents is None:
          latents = jax.random.normal(
              prng_seed, shape=latents_shape, dtype=self.dtype
          )
      else:
          if latents.shape != latents_shape:
              raise ValueError(
                  f"Unexpected latents shape, got {latents.shape}, expected {latents_shape}"
              )
  
      prng_seed, mask_prng_seed = jax.random.split(prng_seed)
  
      masked_image_latent_dist = self.vae.apply(
          {"params": params["vae"]},
          masked_image,
          method=self.vae.encode,
      ).latent_dist
      masked_image_latents = masked_image_latent_dist.sample(
          key=mask_prng_seed
      ).transpose((0, 3, 1, 2))
      masked_image_latents = (
          self.vae.config.scaling_factor * masked_image_latents
      )
      del mask_prng_seed
  
      mask = jax.image.resize(
          mask,
          (*mask.shape[:-2], *masked_image_latents.shape[-2:]),
          method="nearest",
      )
  
      # 8. Check that sizes of mask, masked image and latents match
      num_channels_latents = self.vae.config.latent_channels
      num_channels_mask = mask.shape[1]
      num_channels_masked_image = masked_image_latents.shape[1]
      if (
          num_channels_latents
          + num_channels_mask
          + num_channels_masked_image
          != self.unet.config.in_channels
      ):
          raise ValueError(
              f"Incorrect configuration settings! The config of `pipeline.unet`: {self.unet.config} expects"
              f" {self.unet.config.in_channels} but received `num_channels_latents`: {num_channels_latents} +"
              f" `num_channels_mask`: {num_channels_mask} + `num_channels_masked_image`: {num_channels_masked_image}"
              f" = {num_channels_latents+num_channels_masked_image+num_channels_mask}. Please verify the config of"
              " `pipeline.unet` or your `mask_image` or `image` input."
          )
  
      def loop_body(step, args):
          latents, mask, masked_image_latents, scheduler_state = args
          # 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
          latents_input = jnp.concatenate([latents] * 2)
          mask_input = jnp.concatenate([mask] * 2)
          masked_image_latents_input = jnp.concatenate(
              [masked_image_latents] * 2
          )
  
          t = jnp.array(scheduler_state.timesteps, dtype=jnp.int32)[
              step
          ]
          timestep = jnp.broadcast_to(t, latents_input.shape[0])
  
          latents_input = self.scheduler.scale_model_input(
              scheduler_state, latents_input, t
          )
          # concat latents, mask, masked_image_latents in the channel dimension
          latents_input = jnp.concatenate(
              [latents_input, mask_input, masked_image_latents_input],
              axis=1,
          )
  
          # predict the noise residual
          noise_pred = self.unet.apply(
              {"params": params["unet"]},
              jnp.array(latents_input),
              jnp.array(timestep, dtype=jnp.int32),
              encoder_hidden_states=context,
          ).sample
          # perform guidance
          noise_pred_uncond, noise_prediction_text = jnp.split(
              noise_pred, 2, axis=0
          )
          noise_pred = noise_pred_uncond + guidance_scale * (
              noise_prediction_text - noise_pred_uncond
          )
  
          # compute the previous noisy sample x_t -> x_t-1
          latents, scheduler_state = self.scheduler.step(
              scheduler_state, noise_pred, t, latents
          ).to_tuple()
          return latents, mask, masked_image_latents, scheduler_state
  
      scheduler_state = self.scheduler.set_timesteps(
          params["scheduler"],
          num_inference_steps=num_inference_steps,
          shape=latents.shape,
      )
  
      # scale the initial noise by the standard deviation required by the scheduler
      latents = latents * params["scheduler"].init_noise_sigma
  
      if DEBUG:
          # run with python for loop
          for i in range(num_inference_steps):
              (
                  latents,
                  mask,
                  masked_image_latents,
                  scheduler_state,
              ) = loop_body(
                  i,
                  (
                      latents,
                      mask,
                      masked_image_latents,
                      scheduler_state,
                  ),
              )
      else:
          latents, _, _, _ = jax.lax.fori_loop(
              0,
              num_inference_steps,
              loop_body,
              (latents, mask, masked_image_latents, scheduler_state),
          )
  
      # scale and decode the image latents with vae
      latents = 1 / self.vae.config.scaling_factor * latents
      image = self.vae.apply(
          {"params": params["vae"]}, latents, method=self.vae.decode
      ).sample
  
      image = (image / 2 + 0.5).clip(0, 1).transpose(0, 2, 3, 1)
      return image
  
imports: import warnings
  from functools import partial
  from typing import Dict, List, Optional, Union
  import jax
  import jax.numpy
  import numpy
  from flax.core.frozen_dict import FrozenDict
  from flax.jax_utils import unreplicate
  from flax.training.common_utils import shard
  from packaging import version
  from PIL import Image
  from transformers import CLIPImageProcessor, CLIPTokenizer, FlaxCLIPTextModel
  from ...models import FlaxAutoencoderKL, FlaxUNet2DConditionModel
  from ...schedulers import FlaxDDIMScheduler, FlaxDPMSolverMultistepScheduler, FlaxLMSDiscreteScheduler, FlaxPNDMScheduler
  from ...utils import PIL_INTERPOLATION, deprecate, logging, replace_example_docstring
  from ..pipeline_flax_utils import FlaxDiffusionPipeline
  from .pipeline_output import FlaxStableDiffusionPipelineOutput
  from .safety_checker_flax import FlaxStableDiffusionSafetyChecker
label_desc: Use image node_id 5 for calling the FlaxStableDiffusionInpaintPipeline obj's underlying member method code with example usage: obj._generate(prompt_ids, mask, masked_image, params, prng_seed, num_inference_steps, height, width, guidance_scale, latents, neg_prompt_ids) and returns: image, latents,  mask,  masked_image_latents,  scheduler_state