Unverified Commit 4bf675f4 authored by Hanusz Leszek's avatar Hanusz Leszek Committed by GitHub
Browse files

Dreambooth class image generation: using unique names to avoid overwriting existing image (#847)

* Add an underscore to filename if it already exists

* Use sha1sum hash instead of adding underscores
parent 7674a36a
import argparse
import hashlib
import itertools
import math
import os
......@@ -362,7 +363,9 @@ def main():
images = pipeline(example["prompt"]).images
for i, image in enumerate(images):
image.save(class_images_dir / f"{example['index'][i] + cur_class_images}.jpg")
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
if torch.cuda.is_available():
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment