mirror of
https://github.com/easydiffusion/easydiffusion.git
synced 2025-04-04 14:01:15 +02:00
872c09c220
The old code had a 1.6% chance of name collisions (AAb/345 and AAB3/45 would generate the same filename) on Linux and 39% on Windows (base64 is case sensitive, Windows isn't). This code uses base36 (0-9, A-Z) to avoid case issues. To avoid collisions on fast computers or multi GPU computers, the time resolution of the timestamp is changed from seconds to 0.1ms, and the image number isn't added to the timestamp but appended as an extra character. Due to the limitation of the timestamps to 7 characters, the timecode will repeat every 90 days. This shouldn't be an issue since most sessions will not last this long.
132 lines
4.6 KiB
Python
132 lines
4.6 KiB
Python
import os
|
|
import time
|
|
import re
|
|
|
|
from easydiffusion.types import TaskData, GenerateImageRequest
|
|
|
|
from sdkit.utils import save_images, save_dicts
|
|
from numpy import base_repr
|
|
|
|
filename_regex = re.compile("[^a-zA-Z0-9._-]")
|
|
|
|
# keep in sync with `ui/media/js/dnd.js`
|
|
TASK_TEXT_MAPPING = {
|
|
"prompt": "Prompt",
|
|
"width": "Width",
|
|
"height": "Height",
|
|
"seed": "Seed",
|
|
"num_inference_steps": "Steps",
|
|
"guidance_scale": "Guidance Scale",
|
|
"prompt_strength": "Prompt Strength",
|
|
"use_face_correction": "Use Face Correction",
|
|
"use_upscale": "Use Upscaling",
|
|
"upscale_amount": "Upscale By",
|
|
"sampler_name": "Sampler",
|
|
"negative_prompt": "Negative Prompt",
|
|
"use_stable_diffusion_model": "Stable Diffusion model",
|
|
"use_vae_model": "VAE model",
|
|
"use_hypernetwork_model": "Hypernetwork model",
|
|
"hypernetwork_strength": "Hypernetwork Strength",
|
|
}
|
|
|
|
|
|
def save_images_to_disk(images: list, filtered_images: list, req: GenerateImageRequest, task_data: TaskData):
|
|
now = time.time()
|
|
save_dir_path = os.path.join(task_data.save_to_disk_path, filename_regex.sub("_", task_data.session_id))
|
|
metadata_entries = get_metadata_entries_for_request(req, task_data)
|
|
make_filename = make_filename_callback(req, now=now)
|
|
|
|
if task_data.show_only_filtered_image or filtered_images is images:
|
|
save_images(
|
|
filtered_images,
|
|
save_dir_path,
|
|
file_name=make_filename,
|
|
output_format=task_data.output_format,
|
|
output_quality=task_data.output_quality,
|
|
)
|
|
if task_data.metadata_output_format.lower() in ["json", "txt", "embed"]:
|
|
save_dicts(
|
|
metadata_entries,
|
|
save_dir_path,
|
|
file_name=make_filename,
|
|
output_format=task_data.metadata_output_format,
|
|
file_format=task_data.output_format,
|
|
)
|
|
else:
|
|
make_filter_filename = make_filename_callback(req, now=now, suffix="filtered")
|
|
|
|
save_images(
|
|
images,
|
|
save_dir_path,
|
|
file_name=make_filename,
|
|
output_format=task_data.output_format,
|
|
output_quality=task_data.output_quality,
|
|
)
|
|
save_images(
|
|
filtered_images,
|
|
save_dir_path,
|
|
file_name=make_filter_filename,
|
|
output_format=task_data.output_format,
|
|
output_quality=task_data.output_quality,
|
|
)
|
|
if task_data.metadata_output_format.lower() in ["json", "txt", "embed"]:
|
|
save_dicts(
|
|
metadata_entries,
|
|
save_dir_path,
|
|
file_name=make_filter_filename,
|
|
output_format=task_data.metadata_output_format,
|
|
file_format=task_data.output_format,
|
|
)
|
|
|
|
|
|
def get_metadata_entries_for_request(req: GenerateImageRequest, task_data: TaskData):
|
|
metadata = get_printable_request(req)
|
|
metadata.update(
|
|
{
|
|
"use_stable_diffusion_model": task_data.use_stable_diffusion_model,
|
|
"use_vae_model": task_data.use_vae_model,
|
|
"use_hypernetwork_model": task_data.use_hypernetwork_model,
|
|
"use_face_correction": task_data.use_face_correction,
|
|
"use_upscale": task_data.use_upscale,
|
|
}
|
|
)
|
|
if metadata["use_upscale"] is not None:
|
|
metadata["upscale_amount"] = task_data.upscale_amount
|
|
if task_data.use_hypernetwork_model is None:
|
|
del metadata["hypernetwork_strength"]
|
|
|
|
# if text, format it in the text format expected by the UI
|
|
is_txt_format = task_data.metadata_output_format.lower() == "txt"
|
|
if is_txt_format:
|
|
metadata = {TASK_TEXT_MAPPING[key]: val for key, val in metadata.items() if key in TASK_TEXT_MAPPING}
|
|
|
|
entries = [metadata.copy() for _ in range(req.num_outputs)]
|
|
for i, entry in enumerate(entries):
|
|
entry["Seed" if is_txt_format else "seed"] = req.seed + i
|
|
|
|
return entries
|
|
|
|
|
|
def get_printable_request(req: GenerateImageRequest):
|
|
metadata = req.dict()
|
|
del metadata["init_image"]
|
|
del metadata["init_image_mask"]
|
|
if req.init_image is None:
|
|
del metadata["prompt_strength"]
|
|
return metadata
|
|
|
|
|
|
def make_filename_callback(req: GenerateImageRequest, suffix=None, now=None):
|
|
if now is None:
|
|
now = time.time()
|
|
|
|
def make_filename(i):
|
|
img_id = base_repr(int(now * 10000), 36)[-7:] + base_repr(int(i),36) # Base 36 conversion, 0-9, A-Z
|
|
|
|
prompt_flattened = filename_regex.sub("_", req.prompt)[:50]
|
|
name = f"{prompt_flattened}_{img_id}"
|
|
name = name if suffix is None else f"{name}_{suffix}"
|
|
return name
|
|
|
|
return make_filename
|