Fixes Imagen sampling example and updates container

[terrykong]

1. Fix imagen sampling loop when prompt_ct is a multiple of batch_size // gen_per_prompt
2. Add comment explaining the invocation of sampling script of what exactly is expected due to implicit checkpoint dir requirements and quoting
3. add 2B base model generation gin configs
4. parametrize imagen sampling scripts
5. Updates the imagen image with the fix in (1); built as follows:

docker buildx build --push -t ghcr.io/nvidia/t5x:imagen-2023-10-02.v3 ./JAX-Toolbox -f - <<EOF
FROM ghcr.io/nvidia/t5x:imagen-2023-10-02
COPY rosetta/rosetta/projects/diffusion/common/set_gpu_xla_flags.sh /opt/rosetta/rosetta/projects/diffusion/common/set_gpu_xla_flags.sh
COPY rosetta/rosetta/projects/imagen/README.md /opt/rosetta/rosetta/projects/imagen/README.md
COPY rosetta/rosetta/projects/imagen/configs/imagen_1024_sample_2b.gin /opt/rosetta/rosetta/projects/imagen/configs/imagen_1024_sample_2b.gin
COPY rosetta/rosetta/projects/imagen/configs/imagen_256_sample_2b.gin /opt/rosetta/rosetta/projects/imagen/configs/imagen_256_sample_2b.gin
COPY rosetta/rosetta/projects/imagen/imagen_pipe.py /opt/rosetta/rosetta/projects/imagen/imagen_pipe.py
COPY rosetta/rosetta/projects/imagen/scripts/example_slurm_inf_train.sub /opt/rosetta/rosetta/projects/imagen/scripts/example_slurm_inf_train.sub
COPY rosetta/rosetta/projects/imagen/scripts/sample_imagen_1024.sh /opt/rosetta/rosetta/projects/imagen/scripts/sample_imagen_1024.sh
COPY rosetta/rosetta/projects/imagen/scripts/sample_imagen_256.sh /opt/rosetta/rosetta/projects/imagen/scripts/sample_imagen_256.sh
EOF