(Deprecated) AWS: Stable Diffusion 1.0: Generate Image

const MODEL = "stability.stable-diffusion-xl-v1";
const SERVICE = "bedrock";

function main() {
    ////// == 工程コンフィグ・ワークフローデータの参照 / Config & Data Retrieving ==
    const awsKey = configs.getObject("conf_AccessKey").getToken();
    const awsSecret = configs.getObject("conf_SecretKey").getToken();
    const region = retrieveRegion();
    const payload = buildPayload();
    const fileName = configs.get("conf_FileName");
    if (fileName === "") {
        throw new Error("File name is blank.");
    }

    ////// == 演算 / Calculating ==
    const { generatedFile, seed } = generateImage(awsKey, awsSecret, region, payload, fileName);

    ////// == ワークフローデータへの代入 / Data Updating ==
    saveFileData('conf_File', generatedFile);
    saveData('conf_Seed', seed);
}

/**
 * config からリージョンコードを読み出す
 * リージョンコードの形式として不正な場合はエラー
 * @return {String}
 */
const retrieveRegion = () => {
    const region = configs.get("conf_Region");
    // 今後リージョンが増えることも考えて、中央の最大文字数には余裕をみている
    const reg = new RegExp("^[a-z]{2}-[a-z]{4,16}-[1-9]$");
    if (!reg.test(region)) {
        throw new Error("Region Code is invalid.");
    }
    return region;
}

/**
 * config から API リクエストのボティにセットするパラメータを読み出し、
 * ペイロードを構築する
 * @return {Object}
 */
const buildPayload = () => {
    // cfg_scale / steps は stability.ai のデフォルト値で行く
    const payload = {
        text_prompts: retrieveTextPrompts(),
        seed: retrieveSeed(),
        style_preset: retrieveStyle()
    };

    const initImage = retrieveInitImage();
    const size = retrieveImageSize();
    if (initImage !== null) {
        Object.assign(payload, {
            init_image: initImage
        });
    } else if (size !== null) { // init_image がある場合は size は無視される
        Object.assign(payload, {
            width: size.width,
            height: size.height
        });
    }
    return payload;
}

/**
 * config からテキストプロンプトとネガティブプロンプトを読み出す
 * @return {Array}
 */
const retrieveTextPrompts = () => {
    const message = configs.get('conf_Message');
    if (message === '') {
        throw new Error('Prompt is blank.');
    }
    const negative = configs.get('conf_Negative');

    const prompts = [{
        text: message,
        weight: 1.0
    }];
    if (negative !== '') {
        prompts.push({
            text: negative,
            weight: -0.5
        });
    }
    return prompts;
};

/**
 * config から画像サイズを読み出す
 * @return {Object} {width, height}
 */
const retrieveImageSize = () => {
    const size = configs.get("conf_Size");
    if (size === "") {
        return null;
    }
    const width = Number.parseInt(size.split("x")[0]);
    const height = Number.parseInt(size.split("x")[1]);
    return {
        width,
        height
    };
};

/**
 * config からベースとする画像を読み出し、base64 エンコードして返す
 * @return {String} base64 エンコードされた画像
 */
const retrieveInitImage = () => {
    const imageDef = configs.getObject("conf_InitImage");
    if (imageDef === null) {
        return null;
    }
    const images = engine.findData(imageDef);
    if (images === null) {
        return null;
    }
    if (images.size() !== 1) {
        throw new Error("More than one base image files attached.");
    }
    const image = images.get(0);
    if (!image.getContentType().startsWith('image/')) {
        throw new Error("The base image file is not an image.");
    }
    // 幅、高さがそれぞれ 64 の倍数である必要があるが、事前に確認するのは難しいので、API に任せる
    return base64.encodeToString(fileRepository.readFile(image));
};

/**
 * config から画像スタイルを読み出す
 * @return {String}
 */
const retrieveStyle = () => {
    const style = configs.get("conf_Style");
    if (style === "") {
        return undefined;
    }
    return style;
};

const MAX_SEED = 4294967295;

/**
 * config からシードを読み出す
 * @return {Number}
 */
const retrieveSeed = () => {
    const seedDef = configs.getObject("conf_Seed");
    if (seedDef === null) {
        return undefined;
    }
    const seed = engine.findData(seedDef);
    if (seed === null) {
        return undefined;
    }
    const reg = new RegExp('^(0|[1-9][0-9]{0,9})$');
    const errorMsg = `Seed must be an integer from 0 to ${MAX_SEED}.`;
    if (!reg.test(seed)) {
        throw new Error(errorMsg);
    }
    const seedNum = Number.parseInt(seed);
    if (seedNum > MAX_SEED) {
        throw new Error(errorMsg);
    }
    return seedNum;
};

/**
 * 画像を生成する API リクエストを送信
 * @param awsKey
 * @param awsSecret
 * @param region
 * @param payload
 * @param fileName
 * @return {Object} {generatedFile: NewQfile, seed: String}
 */
const generateImage = (awsKey, awsSecret, region, payload, fileName) => {
    const URL = `https://bedrock-runtime.${region}.amazonaws.com/model/${MODEL}/invoke`;
    const response = httpClient.begin()
        .awsSignV4(awsKey, awsSecret, region, SERVICE)
        .body(JSON.stringify(payload), "application/json")
        .post(URL);

    const status = response.getStatusCode();
    const respTxt = response.getResponseAsString();
    if (status !== 200) {
        engine.log(respTxt);
        throw new Error(`Failed to invoke model. status: ${status}`);
    }

    const {artifacts} = JSON.parse(respTxt);
    engine.log(JSON.stringify(Object.assign({}, artifacts[0], { // base64 以外の値をログ出力
        base64: undefined
    })));
    const image = base64.decodeFromStringToByteArray(artifacts[0].base64);
    const generatedFile = new com.questetra.bpms.core.event.scripttask.NewQfile(
        fileName,
        "image/png",
        image
    );
    const seed = String(artifacts[0].seed)
    return { generatedFile, seed };
};

/**
 * データ項目への保存
 * @param configName
 * @param newFile
 */
const saveFileData = (configName, newFile) => {
    const def = configs.getObject(configName);
    if (def === null) {
        return;
    }
    let files = engine.findData(def);
    if (files === null) {
        files = new java.util.ArrayList();
    }
    files.add(newFile);
    engine.setData(def, files);
};

/**
 * データ項目への保存
 * @param configName
 * @param data
 */
const saveData = (configName, data) => {
    const def = configs.getObject(configName);
    if (def === null) {
        return;
    }
    engine.setData(def, data);
};