class_docstr: A dataset to prepare the instance and class images with the prompts for fine-tuning the model.
  It pre-processes the images and the tokenizes prompts.
code: def __init__(
      self,
      instance_data_root,
      instance_prompt,
      tokenizer,
      class_data_root=None,
      class_prompt=None,
      class_num=None,
      size=512,
      center_crop=False,
  ):
      self.size = size
      self.center_crop = center_crop
      self.tokenizer = tokenizer
  
      self.instance_data_root = Path(instance_data_root)
      if not self.instance_data_root.exists():
          raise ValueError("Instance images root doesn't exists.")
  
      self.instance_images_path = list(
          Path(instance_data_root).iterdir()
      )
      self.num_instance_images = len(self.instance_images_path)
      self.instance_prompt = instance_prompt
      self._length = self.num_instance_images
  
      if class_data_root is not None:
          self.class_data_root = Path(class_data_root)
          self.class_data_root.mkdir(parents=True, exist_ok=True)
          self.class_images_path = list(self.class_data_root.iterdir())
          if class_num is not None:
              self.num_class_images = min(
                  len(self.class_images_path), class_num
              )
          else:
              self.num_class_images = len(self.class_images_path)
          self._length = max(
              self.num_class_images, self.num_instance_images
          )
          self.class_prompt = class_prompt
      else:
          self.class_data_root = None
  
      self.image_transforms = transforms.Compose(
          [
              transforms.Resize(
                  size,
                  interpolation=transforms.InterpolationMode.BILINEAR,
              ),
              transforms.CenterCrop(size)
              if center_crop
              else transforms.RandomCrop(size),
              transforms.ToTensor(),
              transforms.Normalize([0.5], [0.5]),
          ]
      )
  
imports: import argparse
  import logging
  import math
  import os
  from pathlib import Path
  import jax
  import jax.numpy
  import numpy
  import optax
  import torch
  import torch.utils.checkpoint
  import transformers
  from flax import jax_utils
  from flax.training import train_state
  from flax.training.common_utils import shard
  from huggingface_hub import create_repo, upload_folder
  from huggingface_hub.utils import insecure_hashlib
  from jax.experimental.compilation_cache import compilation_cache
  from PIL import Image
  from torch.utils.data import Dataset
  from torchvision import transforms
  from tqdm.auto import tqdm
  from transformers import CLIPImageProcessor, CLIPTokenizer, FlaxCLIPTextModel, set_seed
  from diffusers import FlaxAutoencoderKL, FlaxDDPMScheduler, FlaxPNDMScheduler, FlaxStableDiffusionPipeline, FlaxUNet2DConditionModel
  from diffusers.pipelines.stable_diffusion import FlaxStableDiffusionSafetyChecker
  from diffusers.utils import check_min_version
label_desc: Use image node_id 1 to create a new DreamBoothDataset object from inherited base classes: Dataset with example: obj = DreamBoothDataset(instance_data_root, instance_prompt, tokenizer, class_data_root, class_prompt, class_num, size, center_crop)