src/backend/image_saver.py

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