src/prompt-modules/list.js

/*
    Copyright 2022 juenbug12851

    Licensed under the Apache License, Version 2.0 (the "License");
    you may not use this file except in compliance with the License.
    You may obtain a copy of the License at

        http://www.apache.org/licenses/LICENSE-2.0

    Unless required by applicable law or agreed to in writing, software
    distributed under the License is distributed on an "AS IS" BASIS,
    WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    See the License for the specific language governing permissions and
    limitations under the License.
*/

/**
 * @file
 * @brief Pipeline stage: replace {name} with a random list line plus emphasis / editing / alternating randomization. Notes: notes/reference/prompt-dsl.md.
 */

import fs from "node:fs";
import _ from "lodash";

// Bring in helper functions
import randomEmphasis from "../helpers/randomEmphasis.js";
import randomEditing from "../helpers/randomEditing.js";
import randomAlternating from "../helpers/randomAlternating.js";
import listFiles from "../helpers/listFiles.js";

// List of all prompt funcs
let promptFuncsSd = [randomEmphasis, randomEditing, randomAlternating];

let promptFuncsNai = [randomEmphasis, randomAlternating];

let promptFuncsMdj = [randomEmphasis, randomAlternating];

// List to give every prompt func a turn
let promptFuncsTmp = [];

/**
 * Refill the per-pull pool of randomizers so each gets a turn before any repeats.
 * @param {Function[]} list The mode-appropriate randomizer set to clone into the pool.
 * @returns {void}
 */
function reloadPromptFunc(list) {
  promptFuncsTmp = _.clone(list);
}

/**
 * Pull one random entry from list `name`, optionally applying a single randomizer
 * (emphasis / editing / alternating, chosen without replacement) and resolving any
 * nested `{list}` tokens the randomizer leaves behind.
 * @param {string} name The list name to pull from.
 * @param {object} settings The merged generation settings.
 * @param {boolean} [emphasis=true] Whether this keyword is eligible for randomization.
 * @returns {string} The (possibly randomized) keyword.
 */
// Pulls a random line from a list file
function sampleFile(name, settings, emphasis) {
  // If emphasis is not set, default to true, otherwise, convert to boolean
  emphasis = emphasis === undefined ? true : emphasis == true;

  if (!emphasis || _.random(0.0, 1.0, true) > settings.emphasisChance)
    return listFiles.pull(settings, name);

  // Set correct prompt func for target AI Generator
  let targList = promptFuncsSd;

  if (settings.mode == "NovelAI") targList = promptFuncsNai;
  else if (settings.mode == "Midjourney") targList = promptFuncsMdj;

  // Start list over if depleted
  if (promptFuncsTmp.length == 0) reloadPromptFunc(targList);

  // Shuffle funcs
  promptFuncsTmp = _.shuffle(promptFuncsTmp);

  // Put back in braces
  name = `{${name}}`;

  // Process and save
  name = promptFuncsTmp[0](settings, name).keyword;

  // Remove used entry from tmp list
  promptFuncsTmp.splice(0, 1);

  // Expand
  name = name.replaceAll(/\{(.*?)\}/gm, function (match, p1) {
    return listFiles.pull(settings, p1);
  });

  // Convert to NovelAI if it's enabled
  if (settings.mode == "NovelAI") {
    name = name.replaceAll("(", "{");
    name = name.replaceAll(")", "}");
  }

  return name;
}

/**
 * List pipeline stage: replace every `{name}` token with a random line from that
 * list, applying emphasis/editing/alternating to non-artist keywords.
 * @param {string} prompt The prompt after the dynamic-prompt and salt stages.
 * @param {object} settings The merged generation settings.
 * @param {object} [imageSettings] Image settings (unused; stage-signature parity).
 * @param {object} [upscaleSettings] Upscale settings (unused).
 * @returns {string} The prompt with all list tokens resolved.
 */
export default function (prompt, settings, imageSettings, upscaleSettings) {
  // Process prompt, 2nd pass, expand list keywords into random items from list
  // also include random prompt if requested
  prompt = prompt.replaceAll(/\{(.*?)\}/gm, function (match, p1) {
    // If from the artist file, then pcik a random artist but do not emphasize
    // them
    if (p1 == settings.artistFilename || p1.includes("artist"))
      return sampleFile(p1, settings, false);
    // Otherwise, pull from the file and follow normal emphasis settings
    else return sampleFile(p1, settings, settings.keywordEmphasis);
  });

  // Return prompt
  return prompt;
}