class_docstr: Flax-based pipeline for text-to-image generation using Stable Diffusion.
  
  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`.
docstr: The call function to the pipeline for generation.
  
  Args:
      prompt (`str` or `List[str]`, *optional*):
          The prompt or prompts to guide image generation.
      height (`int`, *optional*, defaults to `self.unet.config.sample_size * self.vae_scale_factor`):
          The height in pixels of the generated image.
      width (`int`, *optional*, defaults to `self.unet.config.sample_size * self.vae_scale_factor`):
          The width in pixels of the generated image.
      num_inference_steps (`int`, *optional*, defaults to 50):
          The number of denoising steps. More denoising steps usually lead to a higher quality image at the
          expense of slower inference.
      guidance_scale (`float`, *optional*, defaults to 7.5):
          A higher guidance scale value encourages the model to generate images closely linked to the text
          `prompt` at the expense of lower image quality. Guidance scale is enabled when `guidance_scale > 1`.
      latents (`jnp.ndarray`, *optional*):
          Pre-generated noisy latents sampled from a Gaussian distribution, to be used as inputs for image
          generation. Can be used to tweak the same generation with different prompts. If not provided, a latents
          array is generated by sampling using the supplied random `generator`.
      jit (`bool`, defaults to `False`):
          Whether to run `pmap` versions of the generation and safety scoring functions.
  
              <Tip warning={true}>
  
              This argument exists because `__call__` is not yet end-to-end pmap-able. It will be removed in a
              future release.
  
              </Tip>
  
      return_dict (`bool`, *optional*, defaults to `True`):
          Whether or not to return a [`~pipelines.stable_diffusion.FlaxStableDiffusionPipelineOutput`] instead of
          a plain tuple.
  
  Examples:
  
  Returns:
      [`~pipelines.stable_diffusion.FlaxStableDiffusionPipelineOutput`] or `tuple`:
          If `return_dict` is `True`, [`~pipelines.stable_diffusion.FlaxStableDiffusionPipelineOutput`] is
          returned, otherwise a `tuple` is returned where the first element is a list with the generated images
          and the second element is a list of `bool`s indicating whether the corresponding generated image
          contains "not-safe-for-work" (nsfw) content.
docstr_tok: ["The","call","function","to","the","pipeline","for","generation",".","Args",":","prompt","(","`","str","`","or","`","List","[","str","]","`",",","*","optional","*",")",":","The","prompt","or","prompts","to","guide","image","generation",".","height","(","`","int","`",",","*","optional","*",",","defaults","to","`","self.unet.config.sample_size","*","self.vae_scale_factor","`",")",":","The","height","in","pixels","of","the","generated","image",".","width","(","`","int","`",",","*","optional","*",",","defaults","to","`","self.unet.config.sample_size","*","self.vae_scale_factor","`",")",":","The","width","in","pixels","of","the","generated","image",".","num_inference_steps","(","`","int","`",",","*","optional","*",",","defaults","to","50",")",":","The","number","of","denoising","steps",".","More","denoising","steps","usually","lead","to","a","higher","quality","image","at","the","expense","of","slower","inference",".","guidance_scale","(","`","float","`",",","*","optional","*",",","defaults","to","7.5",")",":","A","higher","guidance","scale","value","encourages","the","model","to","generate","images","closely","linked","to","the","text","`","prompt","`","at","the","expense","of","lower","image","quality",".","Guidance","scale","is","enabled","when","`","guidance_scale",">","1","`",".","latents","(","`","jnp.ndarray","`",",","*","optional","*",")",":","Pre-generated","noisy","latents","sampled","from","a","Gaussian","distribution",",","to","be","used","as","inputs","for","image","generation",".","Can","be","used","to","tweak","the","same","generation","with","different","prompts",".","If","not","provided",",","a","latents","array","is","generated","by","sampling","using","the","supplied","random","`","generator","`",".","jit","(","`","bool","`",",","defaults","to","`","False","`",")",":","Whether","to","run","`","pmap","`","versions","of","the","generation","and","safety","scoring","functions",".","<","Tip","warning=","{","true","}",">","This","argument","exists","because","`","__call__","`","is","not","yet","end-to-end","pmap-able",".","It","will","be","removed","in","a","future","release",".","<","\/Tip",">","return_dict","(","`","bool","`",",","*","optional","*",",","defaults","to","`","True","`",")",":","Whether","or","not","to","return","a","[","`","~pipelines.stable_diffusion.FlaxStableDiffusionPipelineOutput","`","]","instead","of","a","plain","tuple",".","Examples",":","Returns",":","[","`","~pipelines.stable_diffusion.FlaxStableDiffusionPipelineOutput","`","]","or","`","tuple","`",":","If","`","return_dict","`","is","`","True","`",",","[","`","~pipelines.stable_diffusion.FlaxStableDiffusionPipelineOutput","`","]","is","returned",",","otherwise","a","`","tuple","`","is","returned","where","the","first","element","is","a","list","with","the","generated","images","and","the","second","element","is","a","list","of","`","bool","`","s","indicating","whether","the","corresponding","generated","image","contains","``","not-safe-for-work","''","(","nsfw",")","content","."]
code: def __call__(
      self,
      prompt_ids: jnp.array,
      params: Union[Dict, FrozenDict],
      prng_seed: jax.Array,
      num_inference_steps: int = 50,
      height: Optional[int] = None,
      width: Optional[int] = None,
      guidance_scale: Union[float, jnp.ndarray] = 7.5,
      latents: jnp.ndarray = None,
      neg_prompt_ids: jnp.ndarray = None,
      return_dict: bool = True,
      jit: bool = False,
  ):
      r"""
      The call function to the pipeline for generation.
  
      Args:
          prompt (`str` or `List[str]`, *optional*):
              The prompt or prompts to guide image generation.
          height (`int`, *optional*, defaults to `self.unet.config.sample_size * self.vae_scale_factor`):
              The height in pixels of the generated image.
          width (`int`, *optional*, defaults to `self.unet.config.sample_size * self.vae_scale_factor`):
              The width in pixels of the generated image.
          num_inference_steps (`int`, *optional*, defaults to 50):
              The number of denoising steps. More denoising steps usually lead to a higher quality image at the
              expense of slower inference.
          guidance_scale (`float`, *optional*, defaults to 7.5):
              A higher guidance scale value encourages the model to generate images closely linked to the text
              `prompt` at the expense of lower image quality. Guidance scale is enabled when `guidance_scale > 1`.
          latents (`jnp.ndarray`, *optional*):
              Pre-generated noisy latents sampled from a Gaussian distribution, to be used as inputs for image
              generation. Can be used to tweak the same generation with different prompts. If not provided, a latents
              array is generated by sampling using the supplied random `generator`.
          jit (`bool`, defaults to `False`):
              Whether to run `pmap` versions of the generation and safety scoring functions.
  
                  <Tip warning={true}>
  
                  This argument exists because `__call__` is not yet end-to-end pmap-able. It will be removed in a
                  future release.
  
                  </Tip>
  
          return_dict (`bool`, *optional*, defaults to `True`):
              Whether or not to return a [`~pipelines.stable_diffusion.FlaxStableDiffusionPipelineOutput`] instead of
              a plain tuple.
  
      Examples:
  
      Returns:
          [`~pipelines.stable_diffusion.FlaxStableDiffusionPipelineOutput`] or `tuple`:
              If `return_dict` is `True`, [`~pipelines.stable_diffusion.FlaxStableDiffusionPipelineOutput`] is
              returned, otherwise a `tuple` is returned where the first element is a list with the generated images
              and the second element is a list of `bool`s indicating whether the corresponding generated image
              contains "not-safe-for-work" (nsfw) content.
      """
      # 0. Default height and width to unet
      height = (
          height or self.unet.config.sample_size * self.vae_scale_factor
      )
      width = (
          width or self.unet.config.sample_size * self.vae_scale_factor
      )
  
      if isinstance(guidance_scale, float):
          # Convert to a tensor so each device gets a copy. Follow the prompt_ids for
          # shape information, as they may be sharded (when `jit` is `True`), or not.
          guidance_scale = jnp.array(
              [guidance_scale] * prompt_ids.shape[0]
          )
          if len(prompt_ids.shape) > 2:
              # Assume sharded
              guidance_scale = guidance_scale[:, None]
  
      if jit:
          images = _p_generate(
              self,
              prompt_ids,
              params,
              prng_seed,
              num_inference_steps,
              height,
              width,
              guidance_scale,
              latents,
              neg_prompt_ids,
          )
      else:
          images = self._generate(
              prompt_ids,
              params,
              prng_seed,
              num_inference_steps,
              height,
              width,
              guidance_scale,
              latents,
              neg_prompt_ids,
          )
  
      if self.safety_checker is not None:
          safety_params = params["safety_checker"]
          images_uint8_casted = (images * 255).round().astype("uint8")
          num_devices, batch_size = images.shape[:2]
  
          images_uint8_casted = np.asarray(images_uint8_casted).reshape(
              num_devices * batch_size, height, width, 3
          )
          (
              images_uint8_casted,
              has_nsfw_concept,
          ) = self._run_safety_checker(
              images_uint8_casted, safety_params, jit
          )
          images = np.asarray(images).copy()
  
          # block images
          if any(has_nsfw_concept):
              for i, is_nsfw in enumerate(has_nsfw_concept):
                  if is_nsfw:
                      images[i, 0] = np.asarray(images_uint8_casted[i])
  
          images = images.reshape(
              num_devices, batch_size, height, width, 3
          )
      else:
          images = np.asarray(images)
          has_nsfw_concept = False
  
      if not return_dict:
          return (images, has_nsfw_concept)
  
      return FlaxStableDiffusionPipelineOutput(
          images=images, nsfw_content_detected=has_nsfw_concept
      )
  
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 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 6 for calling the FlaxStableDiffusionPipeline obj's underlying member method code with example usage: obj.__call__(prompt_ids, params, prng_seed, num_inference_steps, height, width, guidance_scale, latents, neg_prompt_ids, return_dict, jit) and returns: FlaxStableDiffusionPipelineOutput, images,  has_nsfw_concept