class_docstr: Base class for Flax-based pipelines.
  
  [`FlaxDiffusionPipeline`] stores all components (models, schedulers, and processors) for diffusion pipelines and
  provides methods for loading, downloading and saving models. It also includes methods to:
  
      - enable/disable the progress bar for the denoising iteration
  
  Class attributes:
  
      - **config_name** ([`str`]) -- The configuration filename that stores the class and module names of all the
        diffusion pipeline's components.
docstr: Save all saveable variables of the pipeline to a directory. A pipeline variable can be saved and loaded if its
  class implements both a save and loading method. The pipeline is easily reloaded using the
  [`~FlaxDiffusionPipeline.from_pretrained`] class method.
  
  Arguments:
      save_directory (`str` or `os.PathLike`):
          Directory to which to save. Will be created if it doesn't exist.
      push_to_hub (`bool`, *optional*, defaults to `False`):
          Whether or not to push your model to the Hugging Face model hub after saving it. You can specify the
          repository you want to push to with `repo_id` (will default to the name of `save_directory` in your
          namespace).
      kwargs (`Dict[str, Any]`, *optional*):
          Additional keyword arguments passed along to the [`~utils.PushToHubMixin.push_to_hub`] method.
docstr_tok: ["Save","all","saveable","variables","of","the","pipeline","to","a","directory",".","A","pipeline","variable","can","be","saved","and","loaded","if","its","class","implements","both","a","save","and","loading","method",".","The","pipeline","is","easily","reloaded","using","the","[","`","~FlaxDiffusionPipeline.from_pretrained","`","]","class","method",".","Arguments",":","save_directory","(","`","str","`","or","`","os.PathLike","`",")",":","Directory","to","which","to","save",".","Will","be","created","if","it","does","n't","exist",".","push_to_hub","(","`","bool","`",",","*","optional","*",",","defaults","to","`","False","`",")",":","Whether","or","not","to","push","your","model","to","the","Hugging","Face","model","hub","after","saving","it",".","You","can","specify","the","repository","you","want","to","push","to","with","`","repo_id","`","(","will","default","to","the","name","of","`","save_directory","`","in","your","namespace",")",".","kwargs","(","`","Dict","[","str",",","Any","]","`",",","*","optional","*",")",":","Additional","keyword","arguments","passed","along","to","the","[","`","~utils.PushToHubMixin.push_to_hub","`","]","method","."]
code: def save_pretrained(
      self,
      save_directory: Union[str, os.PathLike],
      params: Union[Dict, FrozenDict],
      push_to_hub: bool = False,
      **kwargs,
  ):
      # TODO: handle inference_state
      """
      Save all saveable variables of the pipeline to a directory. A pipeline variable can be saved and loaded if its
      class implements both a save and loading method. The pipeline is easily reloaded using the
      [`~FlaxDiffusionPipeline.from_pretrained`] class method.
  
      Arguments:
          save_directory (`str` or `os.PathLike`):
              Directory to which to save. Will be created if it doesn't exist.
          push_to_hub (`bool`, *optional*, defaults to `False`):
              Whether or not to push your model to the Hugging Face model hub after saving it. You can specify the
              repository you want to push to with `repo_id` (will default to the name of `save_directory` in your
              namespace).
          kwargs (`Dict[str, Any]`, *optional*):
              Additional keyword arguments passed along to the [`~utils.PushToHubMixin.push_to_hub`] method.
      """
      self.save_config(save_directory)
  
      model_index_dict = dict(self.config)
      model_index_dict.pop("_class_name")
      model_index_dict.pop("_diffusers_version")
      model_index_dict.pop("_module", None)
  
      if push_to_hub:
          commit_message = kwargs.pop("commit_message", None)
          private = kwargs.pop("private", False)
          create_pr = kwargs.pop("create_pr", False)
          token = kwargs.pop("token", None)
          repo_id = kwargs.pop(
              "repo_id", save_directory.split(os.path.sep)[-1]
          )
          repo_id = create_repo(
              repo_id, exist_ok=True, private=private, token=token
          ).repo_id
  
      for pipeline_component_name in model_index_dict.keys():
          sub_model = getattr(self, pipeline_component_name)
          if sub_model is None:
              # edge case for saving a pipeline with safety_checker=None
              continue
  
          model_cls = sub_model.__class__
  
          save_method_name = None
          # search for the model's base class in LOADABLE_CLASSES
          for library_name, library_classes in LOADABLE_CLASSES.items():
              library = importlib.import_module(library_name)
              for (
                  base_class,
                  save_load_methods,
              ) in library_classes.items():
                  class_candidate = getattr(library, base_class, None)
                  if class_candidate is not None and issubclass(
                      model_cls, class_candidate
                  ):
                      # if we found a suitable base class in LOADABLE_CLASSES then grab its save method
                      save_method_name = save_load_methods[0]
                      break
              if save_method_name is not None:
                  break
  
          save_method = getattr(sub_model, save_method_name)
          expects_params = "params" in set(
              inspect.signature(save_method).parameters.keys()
          )
  
          if expects_params:
              save_method(
                  os.path.join(save_directory, pipeline_component_name),
                  params=params[pipeline_component_name],
              )
          else:
              save_method(
                  os.path.join(save_directory, pipeline_component_name)
              )
  
          if push_to_hub:
              self._upload_folder(
                  save_directory,
                  repo_id,
                  token=token,
                  commit_message=commit_message,
                  create_pr=create_pr,
              )
  
imports: import importlib
  import inspect
  import os
  from typing import Any, Dict, List, Optional, Union
  import flax
  import numpy
  import PIL.Image
  from flax.core.frozen_dict import FrozenDict
  from huggingface_hub import create_repo, snapshot_download
  from PIL import Image
  from tqdm.auto import tqdm
  from ..configuration_utils import ConfigMixin
  from ..models.modeling_flax_utils import FLAX_WEIGHTS_NAME, FlaxModelMixin
  from ..schedulers.scheduling_utils_flax import SCHEDULER_CONFIG_NAME, FlaxSchedulerMixin
  from ..utils import CONFIG_NAME, DIFFUSERS_CACHE, BaseOutput, PushToHubMixin, http_user_agent, is_transformers_available, logging
label_desc: Use image node_id 2 for calling the FlaxDiffusionPipeline obj's underlying member method code with example usage: obj.save_pretrained(save_directory, params, push_to_hub) without return types