From dc756ce22a96554a7b764a70e6858f3436ef30d2 Mon Sep 17 00:00:00 2001 From: scottshireman <129428073+scottshireman@users.noreply.github.com> Date: Tue, 3 Dec 2024 08:47:28 -0500 Subject: [PATCH] Update sample_generator.py Changed {sample_index} to {sample_index:03} in file names and logs for sample generator to allow for proper sorting in wandb and other places when there are more than 10 samples as currently wandb sorts as follows: 0, 1, 10, 11, 2, etc. With this change it would be 000, 001, 002, ... , 010, 011, etc. --- utils/sample_generator.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/utils/sample_generator.py b/utils/sample_generator.py index ad40123..92058ca 100644 --- a/utils/sample_generator.py +++ b/utils/sample_generator.py @@ -272,15 +272,15 @@ class SampleGenerator: prompt = prompts[prompt_idx] clean_prompt = clean_filename(prompt) - result.save(f"{self.log_folder}/samples/gs{global_step:05}-{sample_index}-{extra_info}{clean_prompt[:100]}.jpg", format="JPEG", quality=95, optimize=True, progressive=False) - with open(f"{self.log_folder}/samples/gs{global_step:05}-{sample_index}-{extra_info}{clean_prompt[:100]}.txt", "w", encoding='utf-8') as f: + result.save(f"{self.log_folder}/samples/gs{global_step:05}-{sample_index:03}-{extra_info}{clean_prompt[:100]}.jpg", format="JPEG", quality=95, optimize=True, progressive=False) + with open(f"{self.log_folder}/samples/gs{global_step:05}-{sample_index:03}-{extra_info}{clean_prompt[:100]}.txt", "w", encoding='utf-8') as f: f.write(str(batch[prompt_idx])) tfimage = transforms.ToTensor()(result) if batch[prompt_idx].wants_random_caption: - self.log_writer.add_image(tag=f"sample_{sample_index}{extra_info}", img_tensor=tfimage, global_step=global_step) + self.log_writer.add_image(tag=f"sample_{sample_index:03}{extra_info}", img_tensor=tfimage, global_step=global_step) else: - self.log_writer.add_image(tag=f"sample_{sample_index}_{extra_info}{clean_prompt[:100]}", img_tensor=tfimage, global_step=global_step) + self.log_writer.add_image(tag=f"sample_{sample_index:03}_{extra_info}{clean_prompt[:100]}", img_tensor=tfimage, global_step=global_step) sample_index += 1 del result