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`.
code: def __init__(
      self,
      vae: FlaxAutoencoderKL,
      text_encoder: FlaxCLIPTextModel,
      tokenizer: CLIPTokenizer,
      unet: FlaxUNet2DConditionModel,
      scheduler: Union[
          FlaxDDIMScheduler,
          FlaxPNDMScheduler,
          FlaxLMSDiscreteScheduler,
          FlaxDPMSolverMultistepScheduler,
      ],
      safety_checker: FlaxStableDiffusionSafetyChecker,
      feature_extractor: CLIPImageProcessor,
      dtype: jnp.dtype = jnp.float32,
  ):
      super().__init__()
      self.dtype = dtype
  
      if safety_checker is None:
          logger.warning(
              f"You have disabled the safety checker for {self.__class__} by passing `safety_checker=None`. Ensure"
              " that you abide to the conditions of the Stable Diffusion license and do not expose unfiltered"
              " results in services or applications open to the public. Both the diffusers team and Hugging Face"
              " strongly recommend to keep the safety filter enabled in all public facing circumstances, disabling"
              " it only for use-cases that involve analyzing network behavior or auditing its results. For more"
              " information, please have a look at https://github.com/huggingface/diffusers/pull/254 ."
          )
  
      is_unet_version_less_0_9_0 = hasattr(
          unet.config, "_diffusers_version"
      ) and version.parse(
          version.parse(unet.config._diffusers_version).base_version
      ) < version.parse(
          "0.9.0.dev0"
      )
      is_unet_sample_size_less_64 = (
          hasattr(unet.config, "sample_size")
          and unet.config.sample_size < 64
      )
      if is_unet_version_less_0_9_0 and is_unet_sample_size_less_64:
          deprecation_message = (
              "The configuration file of the unet has set the default `sample_size` to smaller than"
              " 64 which seems highly unlikely .If you're checkpoint is a fine-tuned version of any of the"
              " following: \n- CompVis/stable-diffusion-v1-4 \n- CompVis/stable-diffusion-v1-3 \n-"
              " CompVis/stable-diffusion-v1-2 \n- CompVis/stable-diffusion-v1-1 \n- runwayml/stable-diffusion-v1-5"
              " \n- runwayml/stable-diffusion-inpainting \n you should change 'sample_size' to 64 in the"
              " configuration file. Please make sure to update the config accordingly as leaving `sample_size=32`"
              " in the config might lead to incorrect results in future versions. If you have downloaded this"
              " checkpoint from the Hugging Face Hub, it would be very nice if you could open a Pull request for"
              " the `unet/config.json` file"
          )
          deprecate(
              "sample_size<64",
              "1.0.0",
              deprecation_message,
              standard_warn=False,
          )
          new_config = dict(unet.config)
          new_config["sample_size"] = 64
          unet._internal_dict = FrozenDict(new_config)
  
      self.register_modules(
          vae=vae,
          text_encoder=text_encoder,
          tokenizer=tokenizer,
          unet=unet,
          scheduler=scheduler,
          safety_checker=safety_checker,
          feature_extractor=feature_extractor,
      )
      self.vae_scale_factor = 2 ** (
          len(self.vae.config.block_out_channels) - 1
      )
  
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 1 to create a new FlaxStableDiffusionPipeline object from inherited base classes: FlaxDiffusionPipeline with example: obj = FlaxStableDiffusionPipeline(vae, text_encoder, tokenizer, unet, scheduler, safety_checker, feature_extractor, dtype)