Fast latent extraction with Nemo Curator

[sippycoder]

Throughput on H100:
256 resolution: 69M per day
512 resolution: 18M per day

data:
  data_path: /mnt/pollux/nemo/sample
  output_path: /mnt/pollux/nemo/sample_latents_512/
  id_col: key
  batch_size: 256
  num_threads_per_worker: 16
  image_size: 512
  image_latent_column: image_latent_512

For the above configuration the latents are stored in cfg.data.output_path in parquet format.

ToDo:

• Wandb tracking
• Fix memory leak

[HaozheLiu-ST]

Can we extract the latent code of the 256/512 resolutions simultaneously? Since the model on 512 resolution will be trained with the pair data (caption,latentcode_256, latentcode_512). 256 is for the planning model, and 512 is for the generation part.