class_docstr: LIT model wrapper for the Dalle-Mini Text-to-Image model.
  
  The model wrapper consists of a few already known works connected in an
  interesting way to generate images from the text:
  
  * VQGAN: The generative image model.
  * BART: A sequence-to-sequence autoencoder used to reconstruct input text.
  * CLIP: A scoring model that assesses the alingment of an image and a prompt.
  
  The basic flow within this model wrapper's predict() function is:
  
  1.  The BART encoder is fed the prompt text.
  2.  The BART decoder is sampled multiple times to generate candidates.
  3.  Each candidate is passed to VQGAN, which generates images.
  4.  CLIP then scores each generated image against the prompt.
class_docstr_tok: ["LIT","model","wrapper","for","the","Dalle-Mini","Text-to-Image","model",".","The","model","wrapper","consists","of","a","few","already","known","works","connected","in","an","interesting","way","to","generate","images","from","the","text",":","*","VQGAN",":","The","generative","image","model",".","*","BART",":","A","sequence-to-sequence","autoencoder","used","to","reconstruct","input","text",".","*","CLIP",":","A","scoring","model","that","assesses","the","alingment","of","an","image","and","a","prompt",".","The","basic","flow","within","this","model","wrapper","'s","predict","(",")","function","is",":","1",".","The","BART","encoder","is","fed","the","prompt","text",".","2",".","The","BART","decoder","is","sampled","multiple","times","to","generate","candidates",".","3",".","Each","candidate","is","passed","to","VQGAN",",","which","generates","images",".","4",".","CLIP","then","scores","each","generated","image","against","the","prompt","."]
code: def predict(
      self, inputs: Iterable[_JsonDict], **unused_kw
  ) -> Iterable[_JsonDict]:
      # Model prediction based on code pipeline in doc
      # https://github.com/borisdayma/dalle-mini
      #
      # Three models are required for prediction and it's in the following stages:
      #
      # 1.  First the prompt(text from which image is generated) is fed into
      #     BART encoder model (dalle_bart_model) then the BART decoder is sampled
      #     multiple times to generate candidates.
      #
      # 2.  Each candidate is passed to VQGAN which generates images. The
      #     generated images are stored in two arrays. One images[], which
      #     contains images in ImageBytes format second pil_images[] which
      #     contains the same images in PIL format which is used to generate CLIP
      #     score using the CLIP model.
      #
      # 3.  CLIP is another model by openai, which takes in an image and a prompt
      #     and tells how well they match by generating a score. At the last, the
      #     final output is formated such that it includes the CLIP score along
      #     with the generated images.
  
      # BART model inference function
      @functools.partial(jax.pmap, axis_name="batch")
      def p_generate(
          tokenized_prompt,
          key,
          params,
          top_k: Optional[int] = None,
          top_p: Optional[float] = None,
          temperature: Optional[float] = None,
          condition_scale: float = 10.0,
      ):
          return self.dalle_bart_model.generate(
              **tokenized_prompt,
              prng_key=key,
              params=params,
              top_k=top_k,
              top_p=top_p,
              temperature=temperature,
              condition_scale=condition_scale,
          )
  
      # VQGAN image decoder function
      @functools.partial(jax.pmap, axis_name="batch")
      def p_decode(indices, params):
          return self.vqgan.decode_code(indices, params=params)
  
      # CLIP <image, prompt> scoring function
      @functools.partial(jax.pmap, axis_name="batch")
      def p_clip(inputs, params):
          logits = self.clip(params=params, **inputs).logits_per_image
          return logits
  
      # create a random key
      seed = random.randint(0, 2**32 - 1)
      key = jax.random.PRNGKey(seed)
      prompts = [ex["prompt"] for ex in inputs]
      tokenized_prompts = self.processor(prompts)
      tokenized_prompt = _flax_replicate(tokenized_prompts)
  
      # images has all generated ImageByste
      # pil_images has images in pil format for clip score
      images = []
      pil_images = []
      for _ in _trange(
          max(self.n_predictions // jax.device_count(), 1)
      ):
          # Get a new key; passed to the model on each device to generate unique
          # inference.
          key, subkey = jax.random.split(key)
          # generate images
          encoded_images = p_generate(
              tokenized_prompt,
              _flax_shard_prng_key(subkey),
              self.params,
          )
  
          # remove BOS
          encoded_images = encoded_images.sequences[..., 1:]
          # decode images
          decoded_images = p_decode(encoded_images, self.vqgan_params)
          decoded_images = decoded_images.clip(0.0, 1.0).reshape(
              (-1, 256, 256, 3)
          )
          for decoded_img in decoded_images:
              img = Image.fromarray(
                  np.asarray(decoded_img * 255, dtype=np.uint8)
              )
              # need pil format images too to generate CLIP score
              pil_images.append(img)
              # convert to ImageBytes
              image_str = image_utils.convert_pil_to_image_str(img)
              images.append(image_str)
  
      # get CLIP scores
      clip_inputs = self.clip_processor(
          text=prompts * jax.device_count(),
          images=pil_images,
          return_tensors="np",
          padding="max_length",
          max_length=77,
          truncation=True,
      ).data
  
      # array containing clip score for all images
      p = len(prompts)
      logits = p_clip(_flax_shard(clip_inputs), self.clip_params)
      logits = np.asarray(
          [logits[:, i::p, i] for i in range(p)]
      ).squeeze()
  
      # Organize VQGAN images and CLIP scores per prompt.
      #
      # Until now, images contains data structured as:
      # [
      #     prompt1_prediction_1,
      #     prompt2_prediction_1,
      #     prompt1_prediction_2,
      #     prompt2_prediction_2,
      # ]
      #
      # Below, the structure changes to...
      # [
      #     {'image': [prompt1_prediction_1, prompt1_prediction_2]},
      #     {'image': [prompt2_prediction_1, prompt2_prediction_2]},
      # ]
      #
      # ...and also adds clip score to check image accuracy
      final_images = []
      clip_score = []
      images_per_prompt = []
  
      # Add images as per prompt [prompt1_prediction_1,prompt1_prediction_2]
      # and also add clip_score
      for i in range(p):
          # logits shape is different if n_predictions is only 1, vs more than one
          clip_prompts = (
              logits if logits.ndim == 0 else logits[i].argsort()[::-1]
          )
          if logits.ndim == 0 or len(clip_prompts) == 1:
              my_loop = [0] * self.n_predictions
          else:
              my_loop = clip_prompts
          for idx in range(len(my_loop)):
              images_per_prompt.append(images[idx * p + i])
              # Append CLIP score depending on size hence the condition
              if logits.ndim == 0:
                  clip_score.append((str(logits), None))
              elif len(np.shape(logits)) == 1:
                  clip_score.append((str(logits[idx]), None))
              else:
                  clip_score.append((str(logits[i][idx]), None))
          # Append to final list[JsonDict]
          final_images.append(
              {
                  "image": images_per_prompt,
                  "clip_score": clip_score,
              }
          )
          # Reset images & clip score for new prompt
          images_per_prompt = []
          clip_score = []
  
      return final_images
  
imports: from collections.abc import Iterable
  import functools
  import random
  from typing import Optional
  import dalle_mini
  import flax
  from flax.training import common_utils
  import jax
  from jax import numpy
  from lit_nlp.api import model
  from lit_nlp.api import types
  from lit_nlp.lib import image_utils
  import numpy
  from PIL import Image
  import tqdm.notebook
  import transformers
  from vqgan_jax import modeling_flax_vqgan
label_desc: Use image node_id 3 for calling the DalleModel obj's underlying member method code with example usage: obj.predict(inputs) and returns: final_images, self, self, logits