| 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 |
| |
|
|
