Skip to content

Instantly share code, notes, and snippets.

@janzheng
Forked from mattdesl/cli.js
Created September 13, 2022 23:01
Show Gist options
  • Save janzheng/e5724eed4dc1c606ddadf38a7633be98 to your computer and use it in GitHub Desktop.
Save janzheng/e5724eed4dc1c606ddadf38a7633be98 to your computer and use it in GitHub Desktop.
colour palette from text prompt using Stable Diffusion https://twitter.com/mattdesl/status/1569457645182152705
/**
* General-purpose NodeJS CLI/API wrapping the Stable-Diffusion python scripts.
*
* Note that this uses an older fork of stable-diffusion
* with the 'txt2img.py' script, and that script was modified to
* support the --outfile command.
*/
var { spawn, exec } = require("child_process");
var path = require("path");
const argv = require("minimist")(process.argv.slice(2), {
alias: {
seed: "S",
width: "W",
height: "H",
},
});
const cwd = process.cwd();
// You might need to edit this
const STABLE_DIFFUSION_REPO_PATH = path.resolve(cwd, "../stable-diffusion");
function resolvePath(p) {
if (path.isAbsolute(p)) return p;
else return path.resolve(cwd, p);
}
module.exports = generate;
function generate(opts = {}) {
return new Promise((resolve, reject) => {
console.time("exec");
const outfile = opts.outfile ? opts.outfile : Date.now() + "_x.png";
const outdir = opts.outdir
? opts.outdir
: path.join(process.cwd(), "output");
const prompt = opts.prompt;
console.log("Prompt:", JSON.stringify(prompt));
const width = opts.width || 512;
const height = opts.height || 512;
const seed = opts.seed;
const steps = opts.steps || 50;
const inputImage = opts.image ? resolvePath(opts.image) : "";
const strength = opts.strength != null ? opts.strength : 0.75;
const scale = opts.scale != null ? opts.scale : 7;
// const ddim_eta = opts.ddim_eta != null ? opts.ddim_eta : "0.0";
const txt2imgArgs = [
"scripts/txt2img.py",
"--prompt",
JSON.stringify(prompt),
// "--plms",
"--outdir",
outdir,
"--outfile",
outfile,
"--ckpt",
"sd-v1-4.ckpt",
"--W",
width,
"--H",
height,
// `--ddim_eta=${ddim_eta}`,
"--skip_grid",
`--scale=${scale}`,
`--n_samples=1`,
"--n_rows=1",
"--n_iter=1",
"--fixed_code",
seed != null ? `--seed=${seed}` : null,
`--ddim_steps=${steps}`,
].filter((d) => d != null);
const img2imgArgs = [
"scripts/img2img.py",
"--prompt",
JSON.stringify(prompt),
"--outdir",
outdir,
"--outfile",
outfile,
"--ckpt",
"sd-v1-4.ckpt",
"--skip_grid",
`--n_samples=1`,
"--n_rows=1",
"--n_iter=1",
"--fixed_code",
seed != null ? `--seed=${seed}` : null,
`--ddim_steps=${steps}`,
`--scale=${scale}`,
`--strength=${strength}`,
`--init-img=${JSON.stringify(inputImage)}`,
].filter((d) => d != null);
const cliArgs = inputImage ? img2imgArgs : txt2imgArgs;
const environmentName = "ldm";
const command = `conda run -n ${environmentName} --no-capture-output python ${cliArgs.join(
" "
)}`;
const cp = exec(
command,
{
stdio: "inherit",
cwd: STABLE_DIFFUSION_REPO_PATH,
},
(err, stdout, stderr) => {
console.timeEnd("exec");
if (err) reject(err);
else resolve();
}
);
cp.stdout.pipe(process.stdout);
cp.stderr.pipe(process.stderr);
});
}
if (!module.parent) {
(async () => {
const opts = {
...argv,
prompt: argv._.join(" "),
};
const n_samples = argv.n || 1;
for (let i = 0; i < n_samples; i++) {
await generate(opts);
}
})();
}
/**
* Node CLI tool that generates a palette from text prompt, using Stable Diffusion.
*/
const path = require("path");
const fs = require("fs/promises");
const generate = require("./cli");
const { promisify } = require("util");
const getPixels = promisify(require("get-pixels"));
const { createCanvas } = require("canvas");
const Color = require("canvas-sketch-util/color");
const { quantize } = require("gifenc");
const argv = require("minimist")(process.argv.slice(2), {
alias: {
seed: "S",
width: "W",
height: "H",
},
});
(async () => {
const outdir = path.resolve(process.cwd(), "output");
const outfile = `${Date.now()}.png`;
const basename = path.basename(outfile, path.extname(outfile));
const prompt = [
...argv._,
", photograph, hd hq 8k 4k, canon, leica, 35mm, octane render",
].join(" ");
await generate({
...argv,
prompt,
outdir,
outfile,
});
const pixels = await getPixels(path.resolve(outdir, outfile));
const n = 5;
const palette = quantize(pixels.data, n);
const canvas = createCanvas(512, 128);
const tileWidth = canvas.width / n;
const context = canvas.getContext("2d");
for (let i = 0; i < palette.length; i++) {
const color = palette[i];
const hex = Color.parse(color).hex;
context.fillStyle = hex;
context.fillRect(i * tileWidth, 0, tileWidth, canvas.height);
}
const buf = canvas.toBuffer();
await fs.writeFile(path.join(outdir, `${basename}_palette.png`), buf);
})();
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment