class_docstr: Flax-based pipeline for text-to-image generation using Stable Diffusion with ControlNet Guidance.
  
  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.
      controlnet ([`FlaxControlNetModel`]:
          Provides additional conditioning to the `unet` during the denoising process.
      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-to-image","generation","using","Stable","Diffusion","with","ControlNet","Guidance",".","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",".","controlnet","(","[","`","FlaxControlNetModel","`","]",":","Provides","additional","conditioning","to","the","`","unet","`","during","the","denoising","process",".","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_ids (`jnp.ndarray`):
          The prompt or prompts to guide the image generation.
      image (`jnp.ndarray`):
          Array representing the ControlNet input condition to provide guidance to the `unet` for generation.
      params (`Dict` or `FrozenDict`):
          Dictionary containing the model parameters/weights.
      prng_seed (`jax.Array`):
          Array containing random number generator key.
      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`.
      controlnet_conditioning_scale (`float` or `jnp.ndarray`, *optional*, defaults to 1.0):
          The outputs of the ControlNet are multiplied by `controlnet_conditioning_scale` before they are added
          to the residual in the original `unet`.
      return_dict (`bool`, *optional*, defaults to `True`):
          Whether or not to return a [`~pipelines.stable_diffusion.FlaxStableDiffusionPipelineOutput`] instead of
          a plain tuple.
      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>
  
  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_ids","(","`","jnp.ndarray","`",")",":","The","prompt","or","prompts","to","guide","the","image","generation",".","image","(","`","jnp.ndarray","`",")",":","Array","representing","the","ControlNet","input","condition","to","provide","guidance","to","the","`","unet","`","for","generation",".","params","(","`","Dict","`","or","`","FrozenDict","`",")",":","Dictionary","containing","the","model","parameters\/weights",".","prng_seed","(","`","jax.Array","`",")",":","Array","containing","random","number","generator","key",".","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","`",".","controlnet_conditioning_scale","(","`","float","`","or","`","jnp.ndarray","`",",","*","optional","*",",","defaults","to","1.0",")",":","The","outputs","of","the","ControlNet","are","multiplied","by","`","controlnet_conditioning_scale","`","before","they","are","added","to","the","residual","in","the","original","`","unet","`",".","return_dict","(","`","bool","`",",","*","optional","*",",","defaults","to","`","True","`",")",":","Whether","or","not","to","return","a","[","`","~pipelines.stable_diffusion.FlaxStableDiffusionPipelineOutput","`","]","instead","of","a","plain","tuple",".","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",">","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.ndarray,
      image: jnp.ndarray,
      params: Union[Dict, FrozenDict],
      prng_seed: jax.Array,
      num_inference_steps: int = 50,
      guidance_scale: Union[float, jnp.ndarray] = 7.5,
      latents: jnp.ndarray = None,
      neg_prompt_ids: jnp.ndarray = None,
      controlnet_conditioning_scale: Union[float, jnp.ndarray] = 1.0,
      return_dict: bool = True,
      jit: bool = False,
  ):
      r"""
      The call function to the pipeline for generation.
  
      Args:
          prompt_ids (`jnp.ndarray`):
              The prompt or prompts to guide the image generation.
          image (`jnp.ndarray`):
              Array representing the ControlNet input condition to provide guidance to the `unet` for generation.
          params (`Dict` or `FrozenDict`):
              Dictionary containing the model parameters/weights.
          prng_seed (`jax.Array`):
              Array containing random number generator key.
          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`.
          controlnet_conditioning_scale (`float` or `jnp.ndarray`, *optional*, defaults to 1.0):
              The outputs of the ControlNet are multiplied by `controlnet_conditioning_scale` before they are added
              to the residual in the original `unet`.
          return_dict (`bool`, *optional*, defaults to `True`):
              Whether or not to return a [`~pipelines.stable_diffusion.FlaxStableDiffusionPipelineOutput`] instead of
              a plain tuple.
          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>
  
      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.
      """
  
      height, width = image.shape[-2:]
  
      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 isinstance(controlnet_conditioning_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.
          controlnet_conditioning_scale = jnp.array(
              [controlnet_conditioning_scale] * prompt_ids.shape[0]
          )
          if len(prompt_ids.shape) > 2:
              # Assume sharded
              controlnet_conditioning_scale = (
                  controlnet_conditioning_scale[:, None]
              )
  
      if jit:
          images = _p_generate(
              self,
              prompt_ids,
              image,
              params,
              prng_seed,
              num_inference_steps,
              guidance_scale,
              latents,
              neg_prompt_ids,
              controlnet_conditioning_scale,
          )
      else:
          images = self._generate(
              prompt_ids,
              image,
              params,
              prng_seed,
              num_inference_steps,
              guidance_scale,
              latents,
              neg_prompt_ids,
              controlnet_conditioning_scale,
          )
  
      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.array(images)
  
          # block images
          if any(has_nsfw_concept):
              for i, is_nsfw in enumerate(has_nsfw_concept):
                  if is_nsfw:
                      images[i] = 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 PIL import Image
  from transformers import CLIPFeatureExtractor, CLIPTokenizer, FlaxCLIPTextModel
  from ...models import FlaxAutoencoderKL, FlaxControlNetModel, FlaxUNet2DConditionModel
  from ...schedulers import FlaxDDIMScheduler, FlaxDPMSolverMultistepScheduler, FlaxLMSDiscreteScheduler, FlaxPNDMScheduler
  from ...utils import PIL_INTERPOLATION, logging, replace_example_docstring
  from ..pipeline_flax_utils import FlaxDiffusionPipeline
  from ..stable_diffusion import FlaxStableDiffusionPipelineOutput
  from ..stable_diffusion.safety_checker_flax import FlaxStableDiffusionSafetyChecker
label_desc: Use image node_id 7 for calling the FlaxStableDiffusionControlNetPipeline obj's underlying member method code with example usage: obj.__call__(prompt_ids, image, params, prng_seed, num_inference_steps, guidance_scale, latents, neg_prompt_ids, controlnet_conditioning_scale, return_dict, jit) and returns: FlaxStableDiffusionPipelineOutput, images,  has_nsfw_concept