From 9057beecd72ae5df2673e98ef5b9b0f236ae5319 Mon Sep 17 00:00:00 2001 From: soulteary Date: Sat, 10 Dec 2022 05:54:53 +0800 Subject: [PATCH] fix: #52 https://github.com/apple/ml-stable-diffusion/issues/52 --- python_coreml_stable_diffusion/pipeline.py | 26 +++++++++++++++------- 1 file changed, 18 insertions(+), 8 deletions(-) diff --git a/python_coreml_stable_diffusion/pipeline.py b/python_coreml_stable_diffusion/pipeline.py index 79751ca..ed759a0 100644 --- a/python_coreml_stable_diffusion/pipeline.py +++ b/python_coreml_stable_diffusion/pipeline.py @@ -437,21 +437,30 @@ def get_coreml_pipe(pytorch_pipe, return coreml_pipe -def get_image_path(args, **override_kwargs): +def get_image_path(args): """ mkdir output folder and encode metadata in the filename """ - out_folder = os.path.join(args.o, "_".join(args.prompt.replace("/", "_").rsplit(" "))) + current = time.strftime('%Y-%m-%d %H-%M-%S',time.localtime(time.time())) + out_folder = os.path.join(args.o, current) os.makedirs(out_folder, exist_ok=True) + return os.path.join(out_folder, current + ".png") - out_fname = f"randomSeed_{override_kwargs.get('seed', None) or args.seed}" - out_fname += f"_computeUnit_{override_kwargs.get('compute_unit', None) or args.compute_unit}" - out_fname += f"_modelVersion_{override_kwargs.get('model_version', None) or args.model_version.replace('/', '_')}" + +def save_image_args(args, image_fname, **override_kwargs): + """ save encode metadata in text file which filename same to image file + """ + content = "_".join(args.prompt.replace("/", "_").rsplit(" ")) + "\n" + content += f"randomSeed_{override_kwargs.get('seed', None) or args.seed}" + "\n" + content += f"_computeUnit_{override_kwargs.get('compute_unit', None) or args.compute_unit}" + "\n" + content += f"_modelVersion_{override_kwargs.get('model_version', None) or args.model_version.replace('/', '_')}" + "\n" if args.scheduler is not None: - out_fname += f"_customScheduler_{override_kwargs.get('scheduler', None) or args.scheduler}" - out_fname += f"_numInferenceSteps{override_kwargs.get('num_inference_steps', None) or args.num_inference_steps}" + content += f"_customScheduler_{override_kwargs.get('scheduler', None) or args.scheduler}" + "\n" + content += f"_numInferenceSteps{override_kwargs.get('num_inference_steps', None) or args.num_inference_steps}" + "\n" - return os.path.join(out_folder, out_fname + ".png") + with open(image_fname.replace(".png", ".txt"), 'w') as f: + f.write(content) + return def main(args): @@ -483,6 +492,7 @@ def main(args): ) out_path = get_image_path(args) + save_image_args(args, out_path) logger.info(f"Saving generated image to {out_path}") image["images"][0].save(out_path)