p5+Dino-Diffusion

// TLDR: 

// To RUN the program: Press the Play button at the top of the screen.

// Press RETURN to start or re-start the AI process.

// Press SPACE to clear the canvas and start over.

// Press a key (and then RETURN) to see a letter. 

// Draw on the canvas with the mouse to provide input to the AI.

//

// Dino Diffusion: Bare-bones Diffusion in p5.js

// This project is heavily based on "Dino Diffusion: Bare-bones Diffusion Models"

// by Ollin Boer Bohan: https://madebyoll.in/posts/dino_diffusion/

// Uses p5.js v.1.11.1 and ONNX Runtime Web v1.18.0: https://onnxruntime.ai/

// Requires https://cdn.jsdelivr.net/npm/onnxruntime-web/dist/ort.min.js

//

// Uses graphics drawn into a 64x64 input image (in generateInputImage())

// to guide a real-time, progressively refined stable diffusion process.

// The program loads an 8MB model trained on historic drawings of plants. 

​

const modelFileName = "network.onnx"; 

const nSteps = 20; // How many steps to take in the diffusion process

let inputGraphics; // This 64x64-px offscreen buffer contains the user's "guiding" image.

let outputGraphics; // This 512x512px offscreen buffer contains the AI diffusion output.

​

//----------------------------------------------------

function setup() {

  createCanvas(1024, 512);

  noSmooth();

​

  // The AI ingests the pixel data from inputGraphics, a 64x64 buffer.

  inputGraphics = createGraphics(64,64); 

  inputGraphics.pixelDensity(1);

​

  // The AI synthesizes an image in outputGraphics (512x512).

  outputGraphics = createGraphics(512,512);

  outputGraphics.pixelDensity(1);

  outputGraphics.background(255);

​

  initAI(); 

  generateInputImage("A"); 

}

​

//----------------------------------------------------

function generateInputImage(ch){

  // Replace this function with your own graphics, 

  // drawing them into the inputGraphics offscreen buffer. 

  // Just remember that inputGraphics is tiny (64x64)!

  // You're drawing into `inputGraphics`, an offscreen buffer,

  // so all drawing commands must start with `inputGraphics.`

  // Press RETURN to generate the result. 

  inputGraphics.background(255);

  inputGraphics.fill(0); 

  inputGraphics.noStroke(); 

  inputGraphics.textAlign(CENTER, CENTER);

  inputGraphics.textFont('Garamond', 52);

  inputGraphics.text(ch, 32,32); 

  /*

  // here's a simple alternate idea: a random line

  inputGraphics.background(255); 

  inputGraphics.stroke(0); 

  inputGraphics.strokeWeight(2); 

  inputGraphics.line (random(64),15, random(64),50); 

  */

}

​

//----------------------------------------------------

function draw() {

  background(200);

​

  // Display the inputGraphics on the main canvas, magnified.

  image(inputGraphics,0,0,512,512);

  // If the mouse is pressed, draw the user's scribble in inputGraphics

  if (mouseIsPressed) {

    inputGraphics.stroke(0); 

    let x0 = mouseX/ratio; // scale down the mouse coordinates

    let y0 = mouseY/ratio; // from a 512x512 image to 64x64.

    let x1 = pmouseX/ratio;

    let y1 = pmouseY/ratio;

    inputGraphics.line(x0,y0, x1,y1);

  } 

  if (mouseIsPressed || keyIsPressed){ // show a message

    text("Press RETURN to send new image to AI", 20,30);

  }

​

  // Draw the AI-synthesized on the right half of the canvas.

  displayOutputGraphics(512,0,512,512); 

}

​

//----------------------------------------------------

function keyPressed() {

  if (keyCode == RETURN) {

    // Transmit the inputGraphics buffer to the AI!

    sendImageData();

  } else {

    background(255);

    generateInputImage(key); 

  } 

}