text: # @title Define imports and utility functions.
  import collections
  import functools
  import time
  from pathlib import Path
  
  import gin
  import jax
  import jax.numpy as jnp
  import numpy as np
  import PIL.Image
  import skvideo.io
  import tensorflow as tf
  from absl import app, flags, logging
  from flax import jax_utils
  from flax import linen as nn
  from flax import optim
  from flax.training import checkpoints
  from jax import random
  
  from conerf import (
      configs,
      datasets,
      evaluation,
      gpath,
      image_utils,
      model_utils,
      models,
      schedules,
      utils,
  )
  from conerf import visualization as viz
  
  flags.DEFINE_enum(
      "mode",
      None,
      ["jax_cpu", "jax_gpu", "jax_tpu"],
      "Distributed strategy approach.",
  )
  
  flags.DEFINE_string("base_folder", None, "where to store ckpts and logs")
  flags.mark_flag_as_required("base_folder")
  flags.DEFINE_multi_string("gin_bindings", None, "Gin parameter bindings.")
  flags.DEFINE_multi_string("gin_configs", (), "Gin config files.")
  flags.DEFINE_string("betas", None, "Indices from the dataset to be used")
  flags.DEFINE_bool(
      "use_lr", False, "Whether use linear regression to predict attributes"
  )
  FLAGS = flags.FLAGS
  
  
  def main(argv):
      jax.config.parse_flags_with_absl()
      tf.config.experimental.set_visible_devices([], "GPU")
      del argv
      logging.info("*** Starting generating video")
      # Assume G3 path for config files when running locally.
      gin_configs = FLAGS.gin_configs
  
      logging.info("*** Loading Gin configs from: %s", str(gin_configs))
      gin.parse_config_files_and_bindings(
          config_files=gin_configs,
          bindings=FLAGS.gin_bindings,
          skip_unknown=True,
      )
      exp_config = configs.ExperimentConfig()
      train_config = configs.TrainConfig()
      eval_config = configs.EvalConfig()
      dummy_model = models.NerfModel({}, 0, 0, 0)
      # Get directory information.
      exp_dir = gpath.GPath(FLAGS.base_folder)
      if exp_config.subname:
          exp_dir = exp_dir / exp_config.subname
      checkpoint_dir = exp_dir / "checkpoints"
      # Log and create directories if this is the main process.
      if jax.process_index() == 0:
          logging.info("exp_dir = %s", exp_dir)
          if not exp_dir.exists():
              exp_dir.mkdir(parents=True, exist_ok=True)
  
          logging.info("checkpoint_dir = %s", checkpoint_dir)
          if not checkpoint_dir.exists():
              checkpoint_dir.mkdir(parents=True, exist_ok=True)
  
      logging.info(
          "Starting process %d. There are %d processes.",
          jax.process_index(),
          jax.process_count(),
      )
      logging.info(
          "Found %d accelerator devices: %s.",
          jax.local_device_count(),
          str(jax.local_devices()),
      )
      logging.info(
          "Found %d total devices: %s.", jax.device_count(), str(jax.devices())
      )
  
      datasource = exp_config.datasource_cls(
          image_scale=exp_config.image_scale,
          random_seed=exp_config.random_seed,
          # Enable metadata based on model needs.
          use_warp_id=dummy_model.use_warp,
          use_appearance_id=(
              dummy_model.nerf_embed_key == "appearance"
              or dummy_model.hyper_embed_key == "appearance"
          ),
          use_camera_id=dummy_model.nerf_embed_key == "camera",
          use_time=dummy_model.warp_embed_key == "time",
      )
  
      rng = random.PRNGKey(exp_config.random_seed)
      np.random.seed(exp_config.random_seed + jax.process_index())
      devices_to_use = jax.devices()
  
      learning_rate_sched = schedules.from_config(train_config.lr_schedule)
      nerf_alpha_sched = schedules.from_config(train_config.nerf_alpha_schedule)
      warp_alpha_sched = schedules.from_config(train_config.warp_alpha_schedule)
      hyper_alpha_sched = schedules.from_config(
          train_config.hyper_alpha_schedule
      )
      hyper_sheet_alpha_sched = schedules.from_config(
          train_config.hyper_sheet_alpha_schedule
      )
  
      rng, key = random.split(rng)
      params = {}
      model, params["model"] = models.construct_nerf(
          key,
          batch_size=train_config.batch_size,
          embeddings_dict=datasource.embeddings_dict,
          near=datasource.near,
          far=datasource.far,
          num_attributes=datasource.num_attributes,
      )
  
      optimizer_def = optim.Adam(learning_rate_sched(0))
      optimizer = optimizer_def.create(params)
  
      state = model_utils.TrainState(
          optimizer=optimizer,
          nerf_alpha=nerf_alpha_sched(0),
          warp_alpha=warp_alpha_sched(0),
          hyper_alpha=hyper_alpha_sched(0),
          hyper_sheet_alpha=hyper_sheet_alpha_sched(0),
      )
  
      optimizer_def = optim.Adam(0.0)
      if train_config.use_weight_norm:
          optimizer_def = optim.WeightNorm(optimizer_def)
      optimizer = optimizer_def.create(params)
      init_state = model_utils.TrainState(optimizer=optimizer)
      del params
  
      # @title Define pmapped render function.
  
      devices = jax.devices()
  
      def _model_fn(key_0, key_1, params, rays_dict, extra_params):
          out = model.apply(
              {"params": params},
              rays_dict,
              extra_params=extra_params,
              rngs={"coarse": key_0, "fine": key_1},
              mutable=False,
              metadata_encoded=FLAGS.use_lr,
          )
          return jax.lax.all_gather(out, axis_name="batch")
  
      pmodel_fn = jax.pmap(
          # Note rng_keys are useless in eval mode since there's no randomness.
          _model_fn,
          in_axes=(0, 0, 0, 0, 0),  # Only distribute the data input.
          devices=devices_to_use,
          axis_name="batch",
      )
  
      render_fn = functools.partial(
          evaluation.render_image,
          model_fn=pmodel_fn,
          device_count=len(devices),
          chunk=eval_config.chunk,
      )
  
      camera_path = Path("camera-paths") / "orbit-mild"
  
      camera_dir = Path(datasource.data_dir, camera_path)
      logging.info(f"Loading cameras from {camera_dir}")
      test_camera_paths = datasource.glob_cameras(camera_dir)
      test_cameras = utils.parallel_map(
          datasource.load_camera, test_camera_paths, show_pbar=True
      )
  
      rng = rng + jax.process_index()  # Make random seed separate across hosts.
      _ = random.split(rng, len(devices))
      warp_lr_params = None
      hyper_lr_params = None
      if (exp_dir / "warp_lr.npy").exists():
          warp_lr_params = jnp.load(exp_dir / "warp_lr.npy")
      if (exp_dir / "hyper_lr.npy").exists():
          hyper_lr_params = jnp.load(exp_dir / "hyper_lr.npy")
  
      last_step = 0
      while True:
          if not checkpoint_dir.exists():
              logging.info("No checkpoints yet.")
              time.sleep(10)
              continue
  
          state = checkpoints.restore_checkpoint(checkpoint_dir, init_state)
          state = jax_utils.replicate(state, devices=devices_to_use)
          step = int(state.optimizer.state.step[0])
          if step <= last_step:
              logging.info("No new checkpoints (%d <= %d).", step, last_step)
              time.sleep(10)
              continue
  
          results = []
  
          renders_dir = exp_dir / "eval_renders"
  
          if FLAGS.betas is not None:
              betas = map(int, FLAGS.betas.split(","))
          else:
              betas = [0]
          for beta in betas:
              if model.use_attribute_conditioning or FLAGS.use_lr:
                  num_attributes = datasource.num_attributes
                  attribute_values = [-1.0, 0.0, 1.0]
                  attributes_meshgrid = jnp.stack(
                      jnp.meshgrid(*([attribute_values] * num_attributes)),
                      axis=-1,
                  ).reshape((-1, num_attributes))
                  logging.info(
                      f"{len(attributes_meshgrid)} sequences to generate"
                  )
  
                  for attribute_set in attributes_meshgrid:
                      frames = collections.defaultdict(list)
                      for i in range(len(test_cameras)):
                          suffix = (
                              f"_beta_{beta}" if FLAGS.betas is not None else ""
                          )
                          out_dir = renders_dir / (
                              "_".join([f"{attr:.2f}" for attr in attribute_set])
                              + f"_camera_{i}"
                              + suffix
                          )
                          out_dir.mkdir(parents=True, exist_ok=True)
  
                          logging.info(
                              f"Rendering frame {i+1}/{len(test_cameras)}"
                          )
                          camera = test_cameras[i]
                          batch = datasets.camera_to_rays(camera)
                          batch["metadata"] = {
                              "appearance": jnp.zeros_like(
                                  batch["origins"][..., 0, jnp.newaxis],
                                  jnp.uint32,
                              )
                              + beta,
                              "warp": jnp.zeros_like(
                                  batch["origins"][..., 0, jnp.newaxis],
                                  jnp.uint32,
                              )
                              + beta,
                              "hyper_embed": (
                                  attribute_set[jnp.newaxis, jnp.newaxis]
                                  .repeat(
                                      repeats=batch["origins"].shape[0], axis=0
                                  )
                                  .repeat(
                                      repeats=batch["origins"].shape[1], axis=1
                                  )
                              ),
                          }
  
                          if FLAGS.use_lr:
                              encodings = {}
                              if warp_lr_params is not None:
                                  encodings["encoded_warp"] = (
                                      batch["metadata"]["hyper_embed"]
                                      @ warp_lr_params
                                  )
                              if hyper_lr_params is not None:
                                  encodings["encoded_hyper"] = (
                                      batch["metadata"]["hyper_embed"]
                                      @ hyper_lr_params
                                  )
                              elif model.hyper_use_warp_embed:
                                  encodings["encoded_hyper"] = (
                                      batch["metadata"]["hyper_embed"]
                                      @ warp_lr_params
                                  )
                              batch["metadata"] = encodings
  
                          render = render_fn(state, batch, rng=rng)
                          rgb = np.array(render["rgb"])
                          depth_med = np.array(render["med_depth"])
                          results.append((rgb, depth_med))
                          depth_viz = viz.colorize(
                              depth_med.squeeze(),
                              cmin=datasource.near,
                              cmax=datasource.far,
                              invert=True,
                          )
  
                          frames["rgb"].append(image_utils.image_to_uint8(rgb))
                          frames["depth"].append(
                              image_utils.image_to_uint8(depth_viz)
                          )
                          if "attribute_rgb" in render:
                              attribute = np.array(
                                  nn.sigmoid(render["attribute_rgb"])
                              )
                              mask = np.concatenate(
                                  np.split(
                                      attribute,
                                      indices_or_sections=num_attributes,
                                      axis=-1,
                                  ),
                                  axis=1,
                              ).repeat(repeats=3, axis=-1)
                              overlay = np.clip(
                                  (
                                      attribute.max(axis=-1, keepdims=True) * 0.5
                                      + rgb * 0.5
                                  ),
                                  a_min=0.0,
                                  a_max=1.0,
                              )
                              frames["mask"].append(
                                  image_utils.image_to_uint8(mask)
                              )
                              frames["overlay"].append(
                                  image_utils.image_to_uint8(overlay)
                              )
  
                          for key, vals in frames.items():
                              PIL.Image.fromarray(vals[-1]).save(
                                  (
                                      out_dir / "{}_{:04d}.png".format(key, i)
                                  ).as_posix()
                              )
                      fps = "30"
                      for key, vals in frames.items():
                          skvideo.io.vwrite(
                              (out_dir / f"{key}.mp4").as_posix(),
                              vals,
                              inputdict={"-r": fps},
                              outputdict={"-r": fps},
                          )
              else:
                  frames = collections.defaultdict(list)
                  suffix = f"_beta_{beta}" if FLAGS.betas is not None else ""
                  for i in range(len(test_cameras)):
                      out_dir = renders_dir / (f"camera_{i}" + suffix)
                      out_dir.mkdir(parents=True, exist_ok=True)
  
                      logging.info(f"Rendering frame {i+1}/{len(test_cameras)}")
                      camera = test_cameras[i]
                      batch = datasets.camera_to_rays(camera)
                      batch["metadata"] = {
                          "appearance": jnp.zeros_like(
                              batch["origins"][..., 0, jnp.newaxis], jnp.uint32
                          )
                          + 78,
                          "warp": jnp.zeros_like(
                              batch["origins"][..., 0, jnp.newaxis], jnp.uint32
                          )
                          + 78,
                      }
  
                      render = render_fn(state, batch, rng=rng)
                      rgb = np.array(render["rgb"])
                      depth_med = np.array(render["med_depth"])
                      results.append((rgb, depth_med))
                      depth_viz = viz.colorize(
                          depth_med.squeeze(),
                          cmin=datasource.near,
                          cmax=datasource.far,
                          invert=True,
                      )
  
                      frames["rgb"].append(image_utils.image_to_uint8(rgb))
                      frames["depth"].append(
                          image_utils.image_to_uint8(depth_viz)
                      )
  
                      for key, vals in frames.items():
                          PIL.Image.fromarray(vals[-1]).save(
                              (
                                  out_dir / "{}_{:04d}.png".format(key, i)
                              ).as_posix()
                          )
  
                  fps = "30"
                  for key, vals in frames.items():
                      skvideo.io.vwrite(
                          (out_dir / f"{key}.mp4").as_posix(),
                          vals,
                          inputdict={"-r": fps},
                          outputdict={"-r": fps},
                      )
          if step >= train_config.max_steps:
              break
          last_step = step
  
  
  if __name__ == "__main__":
      app.run(main)
  