code: def main():
      args = parse_args()
  
      logging.basicConfig(
          format="%(asctime)s - %(levelname)s - %(name)s - %(message)s",
          datefmt="%m/%d/%Y %H:%M:%S",
          level=logging.INFO,
      )
      # Setup logging, we only want one process per machine to log things on the screen.
      logger.setLevel(
          logging.INFO if jax.process_index() == 0 else logging.ERROR
      )
      if jax.process_index() == 0:
          transformers.utils.logging.set_verbosity_info()
      else:
          transformers.utils.logging.set_verbosity_error()
  
      if args.seed is not None:
          set_seed(args.seed)
  
      rng = jax.random.PRNGKey(args.seed)
  
      if args.with_prior_preservation:
          class_images_dir = Path(args.class_data_dir)
          if not class_images_dir.exists():
              class_images_dir.mkdir(parents=True)
          cur_class_images = len(list(class_images_dir.iterdir()))
  
          if cur_class_images < args.num_class_images:
              (
                  pipeline,
                  params,
              ) = FlaxStableDiffusionPipeline.from_pretrained(
                  args.pretrained_model_name_or_path,
                  safety_checker=None,
                  revision=args.revision,
              )
              pipeline.set_progress_bar_config(disable=True)
  
              num_new_images = args.num_class_images - cur_class_images
              logger.info(
                  f"Number of class images to sample: {num_new_images}."
              )
  
              sample_dataset = PromptDataset(
                  args.class_prompt, num_new_images
              )
              total_sample_batch_size = (
                  args.sample_batch_size * jax.local_device_count()
              )
              sample_dataloader = torch.utils.data.DataLoader(
                  sample_dataset, batch_size=total_sample_batch_size
              )
  
              for example in tqdm(
                  sample_dataloader,
                  desc="Generating class images",
                  disable=not jax.process_index() == 0,
              ):
                  prompt_ids = pipeline.prepare_inputs(
                      example["prompt"]
                  )
                  prompt_ids = shard(prompt_ids)
                  p_params = jax_utils.replicate(params)
                  rng = jax.random.split(rng)[0]
                  sample_rng = jax.random.split(rng, jax.device_count())
                  images = pipeline(
                      prompt_ids, p_params, sample_rng, jit=True
                  ).images
                  images = images.reshape(
                      (images.shape[0] * images.shape[1],)
                      + images.shape[-3:]
                  )
                  images = pipeline.numpy_to_pil(np.array(images))
  
                  for i, image in enumerate(images):
                      hash_image = hashlib.sha1(
                          image.tobytes()
                      ).hexdigest()
                      image_filename = (
                          class_images_dir
                          / f"{example['index'][i] + cur_class_images}-{hash_image}.jpg"
                      )
                      image.save(image_filename)
  
              del pipeline
  
      # Handle the repository creation
      if jax.process_index() == 0:
          if args.push_to_hub:
              if args.hub_model_id is None:
                  repo_name = get_full_repo_name(
                      Path(args.output_dir).name, token=args.hub_token
                  )
              else:
                  repo_name = args.hub_model_id
              create_repo(
                  repo_name, exist_ok=True, token=args.hub_token
              )
              repo = Repository(
                  args.output_dir,
                  clone_from=repo_name,
                  token=args.hub_token,
              )
  
              with open(
                  os.path.join(args.output_dir, ".gitignore"), "w+"
              ) as gitignore:
                  if "step_*" not in gitignore:
                      gitignore.write("step_*\n")
                  if "epoch_*" not in gitignore:
                      gitignore.write("epoch_*\n")
          elif args.output_dir is not None:
              os.makedirs(args.output_dir, exist_ok=True)
  
      # Load the tokenizer and add the placeholder token as a additional special token
      if args.tokenizer_name:
          tokenizer = CLIPTokenizer.from_pretrained(args.tokenizer_name)
      elif args.pretrained_model_name_or_path:
          tokenizer = CLIPTokenizer.from_pretrained(
              args.pretrained_model_name_or_path,
              subfolder="tokenizer",
              revision=args.revision,
          )
      else:
          raise NotImplementedError("No tokenizer specified!")
  
      train_dataset = DreamBoothDataset(
          instance_data_root=args.instance_data_dir,
          instance_prompt=args.instance_prompt,
          class_data_root=args.class_data_dir
          if args.with_prior_preservation
          else None,
          class_prompt=args.class_prompt,
          class_num=args.num_class_images,
          tokenizer=tokenizer,
          size=args.resolution,
          center_crop=args.center_crop,
      )
  
      def collate_fn(examples):
          input_ids = [
              example["instance_prompt_ids"] for example in examples
          ]
          pixel_values = [
              example["instance_images"] for example in examples
          ]
  
          # Concat class and instance examples for prior preservation.
          # We do this to avoid doing two forward passes.
          if args.with_prior_preservation:
              input_ids += [
                  example["class_prompt_ids"] for example in examples
              ]
              pixel_values += [
                  example["class_images"] for example in examples
              ]
  
          pixel_values = torch.stack(pixel_values)
          pixel_values = pixel_values.to(
              memory_format=torch.contiguous_format
          ).float()
  
          input_ids = tokenizer.pad(
              {"input_ids": input_ids},
              padding="max_length",
              max_length=tokenizer.model_max_length,
              return_tensors="pt",
          ).input_ids
  
          batch = {
              "input_ids": input_ids,
              "pixel_values": pixel_values,
          }
          batch = {k: v.numpy() for k, v in batch.items()}
          return batch
  
      total_train_batch_size = (
          args.train_batch_size * jax.local_device_count()
      )
      if len(train_dataset) < total_train_batch_size:
          raise ValueError(
              f"Training batch size is {total_train_batch_size}, but your dataset only contains"
              f" {len(train_dataset)} images. Please, use a larger dataset or reduce the effective batch size. Note that"
              f" there are {jax.local_device_count()} parallel devices, so your batch size can't be smaller than that."
          )
  
      train_dataloader = torch.utils.data.DataLoader(
          train_dataset,
          batch_size=total_train_batch_size,
          shuffle=True,
          collate_fn=collate_fn,
          drop_last=True,
      )
  
      weight_dtype = jnp.float32
      if args.mixed_precision == "fp16":
          weight_dtype = jnp.float16
      elif args.mixed_precision == "bf16":
          weight_dtype = jnp.bfloat16
  
      if args.pretrained_vae_name_or_path:
          # TODO(patil-suraj): Upload flax weights for the VAE
          vae_arg, vae_kwargs = (
              args.pretrained_vae_name_or_path,
              {"from_pt": True},
          )
      else:
          vae_arg, vae_kwargs = (
              args.pretrained_model_name_or_path,
              {"subfolder": "vae", "revision": args.revision},
          )
  
      # Load models and create wrapper for stable diffusion
      text_encoder = FlaxCLIPTextModel.from_pretrained(
          args.pretrained_model_name_or_path,
          subfolder="text_encoder",
          dtype=weight_dtype,
          revision=args.revision,
      )
      vae, vae_params = FlaxAutoencoderKL.from_pretrained(
          vae_arg,
          dtype=weight_dtype,
          **vae_kwargs,
      )
      unet, unet_params = FlaxUNet2DConditionModel.from_pretrained(
          args.pretrained_model_name_or_path,
          subfolder="unet",
          dtype=weight_dtype,
          revision=args.revision,
      )
  
      # Optimization
      if args.scale_lr:
          args.learning_rate = (
              args.learning_rate * total_train_batch_size
          )
  
      constant_scheduler = optax.constant_schedule(args.learning_rate)
  
      adamw = optax.adamw(
          learning_rate=constant_scheduler,
          b1=args.adam_beta1,
          b2=args.adam_beta2,
          eps=args.adam_epsilon,
          weight_decay=args.adam_weight_decay,
      )
  
      optimizer = optax.chain(
          optax.clip_by_global_norm(args.max_grad_norm),
          adamw,
      )
  
      unet_state = train_state.TrainState.create(
          apply_fn=unet.__call__, params=unet_params, tx=optimizer
      )
      text_encoder_state = train_state.TrainState.create(
          apply_fn=text_encoder.__call__,
          params=text_encoder.params,
          tx=optimizer,
      )
  
      noise_scheduler = FlaxDDPMScheduler(
          beta_start=0.00085,
          beta_end=0.012,
          beta_schedule="scaled_linear",
          num_train_timesteps=1000,
      )
      noise_scheduler_state = noise_scheduler.create_state()
  
      # Initialize our training
      train_rngs = jax.random.split(rng, jax.local_device_count())
  
      def train_step(
          unet_state, text_encoder_state, vae_params, batch, train_rng
      ):
          dropout_rng, sample_rng, new_train_rng = jax.random.split(
              train_rng, 3
          )
  
          if args.train_text_encoder:
              params = {
                  "text_encoder": text_encoder_state.params,
                  "unet": unet_state.params,
              }
          else:
              params = {"unet": unet_state.params}
  
          def compute_loss(params):
              # Convert images to latent space
              vae_outputs = vae.apply(
                  {"params": vae_params},
                  batch["pixel_values"],
                  deterministic=True,
                  method=vae.encode,
              )
              latents = vae_outputs.latent_dist.sample(sample_rng)
              # (NHWC) -> (NCHW)
              latents = jnp.transpose(latents, (0, 3, 1, 2))
              latents = latents * vae.config.scaling_factor
  
              # Sample noise that we'll add to the latents
              noise_rng, timestep_rng = jax.random.split(sample_rng)
              noise = jax.random.normal(noise_rng, latents.shape)
              # Sample a random timestep for each image
              bsz = latents.shape[0]
              timesteps = jax.random.randint(
                  timestep_rng,
                  (bsz,),
                  0,
                  noise_scheduler.config.num_train_timesteps,
              )
  
              # Add noise to the latents according to the noise magnitude at each timestep
              # (this is the forward diffusion process)
              noisy_latents = noise_scheduler.add_noise(
                  noise_scheduler_state, latents, noise, timesteps
              )
  
              # Get the text embedding for conditioning
              if args.train_text_encoder:
                  encoder_hidden_states = text_encoder_state.apply_fn(
                      batch["input_ids"],
                      params=params["text_encoder"],
                      dropout_rng=dropout_rng,
                      train=True,
                  )[0]
              else:
                  encoder_hidden_states = text_encoder(
                      batch["input_ids"],
                      params=text_encoder_state.params,
                      train=False,
                  )[0]
  
              # Predict the noise residual
              model_pred = unet.apply(
                  {"params": params["unet"]},
                  noisy_latents,
                  timesteps,
                  encoder_hidden_states,
                  train=True,
              ).sample
  
              # Get the target for loss depending on the prediction type
              if noise_scheduler.config.prediction_type == "epsilon":
                  target = noise
              elif (
                  noise_scheduler.config.prediction_type
                  == "v_prediction"
              ):
                  target = noise_scheduler.get_velocity(
                      noise_scheduler_state, latents, noise, timesteps
                  )
              else:
                  raise ValueError(
                      f"Unknown prediction type {noise_scheduler.config.prediction_type}"
                  )
  
              if args.with_prior_preservation:
                  # Chunk the noise and noise_pred into two parts and compute the loss on each part separately.
                  model_pred, model_pred_prior = jnp.split(
                      model_pred, 2, axis=0
                  )
                  target, target_prior = jnp.split(target, 2, axis=0)
  
                  # Compute instance loss
                  loss = (target - model_pred) ** 2
                  loss = loss.mean()
  
                  # Compute prior loss
                  prior_loss = (target_prior - model_pred_prior) ** 2
                  prior_loss = prior_loss.mean()
  
                  # Add the prior loss to the instance loss.
                  loss = loss + args.prior_loss_weight * prior_loss
              else:
                  loss = (target - model_pred) ** 2
                  loss = loss.mean()
  
              return loss
  
          grad_fn = jax.value_and_grad(compute_loss)
          loss, grad = grad_fn(params)
          grad = jax.lax.pmean(grad, "batch")
  
          new_unet_state = unet_state.apply_gradients(
              grads=grad["unet"]
          )
          if args.train_text_encoder:
              new_text_encoder_state = (
                  text_encoder_state.apply_gradients(
                      grads=grad["text_encoder"]
                  )
              )
          else:
              new_text_encoder_state = text_encoder_state
  
          metrics = {"loss": loss}
          metrics = jax.lax.pmean(metrics, axis_name="batch")
  
          return (
              new_unet_state,
              new_text_encoder_state,
              metrics,
              new_train_rng,
          )
  
      # Create parallel version of the train step
      p_train_step = jax.pmap(
          train_step, "batch", donate_argnums=(0, 1)
      )
  
      # Replicate the train state on each device
      unet_state = jax_utils.replicate(unet_state)
      text_encoder_state = jax_utils.replicate(text_encoder_state)
      vae_params = jax_utils.replicate(vae_params)
  
      # Train!
      num_update_steps_per_epoch = math.ceil(len(train_dataloader))
  
      # Scheduler and math around the number of training steps.
      if args.max_train_steps is None:
          args.max_train_steps = (
              args.num_train_epochs * num_update_steps_per_epoch
          )
  
      args.num_train_epochs = math.ceil(
          args.max_train_steps / num_update_steps_per_epoch
      )
  
      logger.info("***** Running training *****")
      logger.info(f"  Num examples = {len(train_dataset)}")
      logger.info(f"  Num Epochs = {args.num_train_epochs}")
      logger.info(
          f"  Instantaneous batch size per device = {args.train_batch_size}"
      )
      logger.info(
          f"  Total train batch size (w. parallel & distributed) = {total_train_batch_size}"
      )
      logger.info(
          f"  Total optimization steps = {args.max_train_steps}"
      )
  
      def checkpoint(step=None):
          # Create the pipeline using the trained modules and save it.
          scheduler, _ = FlaxPNDMScheduler.from_pretrained(
              "CompVis/stable-diffusion-v1-4", subfolder="scheduler"
          )
          safety_checker = (
              FlaxStableDiffusionSafetyChecker.from_pretrained(
                  "CompVis/stable-diffusion-safety-checker",
                  from_pt=True,
              )
          )
          pipeline = FlaxStableDiffusionPipeline(
              text_encoder=text_encoder,
              vae=vae,
              unet=unet,
              tokenizer=tokenizer,
              scheduler=scheduler,
              safety_checker=safety_checker,
              feature_extractor=CLIPImageProcessor.from_pretrained(
                  "openai/clip-vit-base-patch32"
              ),
          )
  
          outdir = (
              os.path.join(args.output_dir, str(step))
              if step
              else args.output_dir
          )
          pipeline.save_pretrained(
              outdir,
              params={
                  "text_encoder": get_params_to_save(
                      text_encoder_state.params
                  ),
                  "vae": get_params_to_save(vae_params),
                  "unet": get_params_to_save(unet_state.params),
                  "safety_checker": safety_checker.params,
              },
          )
  
          if args.push_to_hub:
              message = (
                  f"checkpoint-{step}"
                  if step is not None
                  else "End of training"
              )
              repo.push_to_hub(
                  commit_message=message,
                  blocking=False,
                  auto_lfs_prune=True,
              )
  
      global_step = 0
  
      epochs = tqdm(
          range(args.num_train_epochs), desc="Epoch ... ", position=0
      )
      for epoch in epochs:
          # ======================== Training ================================
  
          train_metrics = []
  
          steps_per_epoch = len(train_dataset) // total_train_batch_size
          train_step_progress_bar = tqdm(
              total=steps_per_epoch,
              desc="Training...",
              position=1,
              leave=False,
          )
          # train
          for batch in train_dataloader:
              batch = shard(batch)
              (
                  unet_state,
                  text_encoder_state,
                  train_metric,
                  train_rngs,
              ) = p_train_step(
                  unet_state,
                  text_encoder_state,
                  vae_params,
                  batch,
                  train_rngs,
              )
              train_metrics.append(train_metric)
  
              train_step_progress_bar.update(jax.local_device_count())
  
              global_step += 1
              if (
                  jax.process_index() == 0
                  and args.save_steps
                  and global_step % args.save_steps == 0
              ):
                  checkpoint(global_step)
              if global_step >= args.max_train_steps:
                  break
  
          train_metric = jax_utils.unreplicate(train_metric)
  
          train_step_progress_bar.close()
          epochs.write(
              f"Epoch... ({epoch + 1}/{args.num_train_epochs} | Loss: {train_metric['loss']})"
          )
  
      if jax.process_index() == 0:
          checkpoint()
  
imports: import argparse
  import hashlib
  import logging
  import math
  import os
  from pathlib import Path
  from typing import Optional
  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 HfFolder, Repository, create_repo, whoami
  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 4 for calling a global function with example usage: main() and returns: batch, new_unet_state,  new_text_encoder_state,  metrics,  new_train_rng, loss