|
import json |
|
from os import path, mkdir |
|
from typing import Any |
|
from uuid import uuid4 |
|
from backend.models.lcmdiffusion_setting import LCMDiffusionSetting |
|
from utils import get_image_file_extension |
|
|
|
|
|
def get_exclude_keys(): |
|
exclude_keys = { |
|
"init_image": True, |
|
"generated_images": True, |
|
"lora": { |
|
"models_dir": True, |
|
"path": True, |
|
}, |
|
"dirs": True, |
|
"controlnet": { |
|
"adapter_path": True, |
|
}, |
|
} |
|
return exclude_keys |
|
|
|
|
|
class ImageSaver: |
|
@staticmethod |
|
def save_images( |
|
output_path: str, |
|
images: Any, |
|
folder_name: str = "", |
|
format: str = "PNG", |
|
jpeg_quality: int = 90, |
|
lcm_diffusion_setting: LCMDiffusionSetting = None, |
|
) -> list[str]: |
|
gen_id = uuid4() |
|
image_ids = [] |
|
|
|
if images: |
|
image_seeds = [] |
|
|
|
for index, image in enumerate(images): |
|
|
|
image_seed = image.info.get('image_seed') |
|
if image_seed is not None: |
|
image_seeds.append(image_seed) |
|
|
|
if not path.exists(output_path): |
|
mkdir(output_path) |
|
|
|
if folder_name: |
|
out_path = path.join( |
|
output_path, |
|
folder_name, |
|
) |
|
else: |
|
out_path = output_path |
|
|
|
if not path.exists(out_path): |
|
mkdir(out_path) |
|
image_extension = get_image_file_extension(format) |
|
image_file_name = f"{gen_id}-{index+1}{image_extension}" |
|
image_ids.append(image_file_name) |
|
image.save(path.join(out_path, image_file_name), quality = jpeg_quality) |
|
if lcm_diffusion_setting: |
|
data = lcm_diffusion_setting.model_dump(exclude=get_exclude_keys()) |
|
if image_seeds: |
|
data['image_seeds'] = image_seeds |
|
with open(path.join(out_path, f"{gen_id}.json"), "w") as json_file: |
|
json.dump( |
|
data, |
|
json_file, |
|
indent=4, |
|
) |
|
return image_ids |
|
|
|
|
