stable-diffusion-storage-extension/scripts/storage.py

import base64, uuid
from io import BytesIO
import os,json
import re
import modules.scripts as scripts
import gradio as gr

from ots_agent import create_table, insert_row
from oss_util import upload_to_oss, upload_bytes_to_oss

pk, user_id, instance_id = "pk", "user_id", "instance_id"

create_table()

class Scripts(scripts.Script):
    def title(self):
        return "Aliyun Storage"

    def show(self, is_img2img):
        return scripts.AlwaysVisible

    def ui(self, is_img2img):
        checkbox_save_to_db = gr.inputs.Checkbox(label="Save to DB", default=True)
        return [checkbox_save_to_db, ]

    def postprocess(self, p, processed,checkbox_save_to_db):
        # Extract image information
        regex = r"Steps:.*$"
        info = re.findall(regex, processed.info, re.M)[0]
        input_dict = dict(item.split(": ") for item in str(info).split(", "))
        steps = int(input_dict["Steps"])
        seed = int(input_dict["Seed"])
        sampler = input_dict["Sampler"]
        cfg_scale = float(input_dict["CFG scale"])
        size = tuple(map(int, input_dict["Size"].split("x")))
        model_hash = input_dict["Model hash"]
        model = input_dict["Model"]
            
        task_id = str(uuid.uuid4())
        image_list = []
        
        for i in range(len(processed.images)):
            seed = processed.seed
            prompt = processed.prompt
            neg_prompt = processed.negative_prompt


            image = processed.images[i] # PIL.Image
            buffer = BytesIO()
            image.save(buffer, "png")
            image_bytes = buffer.getvalue()
            
            oss_file_name = f"sd_images/{user_id}/{instance_id}/{task_id}/{i}.png"
            
            upload_bytes_to_oss(image_bytes, oss_file_name)
            image_list.append(oss_file_name)

        parameters = {
                "prompt": prompt, 
                "negative_prompt": neg_prompt, 
                "steps": int(steps), 
                "seed": int(seed), 
                "sampler": sampler,
                "cfg_scale": float(cfg_scale), 
                "size": size, 
                "model_hash": model_hash, 
                "model": model
        }
        
        print("insert result of task_id: ", task_id, pk, user_id, instance_id, task_id, prompt, ','.join(image_list) , parameters)
        insert_row(pk, user_id, instance_id, task_id, prompt, json.dumps(image_list, ensure_ascii=False) , request_params=json.dumps(parameters, ensure_ascii=False))
        return True