Refactor some codes.

This commit is contained in:
hodanov 2023-06-05 10:16:09 +09:00
parent 41d94ee060
commit 053b142a87

View File

@ -162,18 +162,13 @@ def entrypoint(
inputs["max_embeddings_multiples"] = util.count_token(p=prompt, n=n_prompt) inputs["max_embeddings_multiples"] = util.count_token(p=prompt, n=n_prompt)
directory = util.make_directory() directory = util.make_directory()
util.save_prompts(inputs)
sd = StableDiffusion() sd = StableDiffusion()
for i in range(samples): for i in range(samples):
start_time = time.time() start_time = time.time()
images = sd.run_inference.call(inputs) images = sd.run_inference.call(inputs)
for j, image_bytes in enumerate(images): util.save_images(directory, images, i)
formatted_time = time.strftime("%Y%m%d%H%M%S", time.localtime(time.time()))
output_path = directory / f"{formatted_time}_{i}_{j}.png"
print(f"Saving it to {output_path}")
with open(output_path, "wb") as file:
file.write(image_bytes)
total_time = time.time() - start_time total_time = time.time() - start_time
print(f"Sample {i} took {total_time:.3f}s ({(total_time)/len(images):.3f}s / image).") print(f"Sample {i} took {total_time:.3f}s ({(total_time)/len(images):.3f}s / image).")
util.save_prompts(inputs)