index.html

<!DOCTYPE html>
<html lang="en">
<head>
<meta charset="UTF-8">
<meta name="viewport" content="width=device-width, initial-scale=1.0">
<title>H-Space similarity explorer</title>
<link rel="shortcut icon" type="image/x-icon" href="favicon.ico?">
<style>
* {
  margin: 0;
  padding: 0;
  font-family: Arial, sans-serif;
}
#intro {
  margin: 1em auto 2em auto;
  width: 100%;
  max-width: 800px;
}
#intro p {
  margin: 1em;
  line-height: 1.3em;
}
.select-container {
  margin: .5em;
  display: flex;
  flex-wrap: wrap;
}
.select-group {
  margin: .5em;
}
.select-group span {
  margin-right: .3em;
}
#canvas-container {
  margin: 0 auto;
  margin-bottom: 2em;
  width: 100%;
  max-width: 800px;
}
#canvas-container h2 {
  text-align: center;
  margin: 2em 0 0 0;
}
#canvas-container h2 select {
  margin: 1em;
}
#canvas-container>p {
  text-align: center;
  color: grey;
}
#canvas-container div {
  width: 100%;
}
#canvas-container div div {
  display: inline-block;
  max-width: 44.9%;
  margin: 0 2.5%;
}
#canvas-container p {
  text-align: center;
  margin-bottom: .2em;
}
#canvas-container p input {
  margin: 0 .5em 0 .5em;
  height: .8em;
  max-width: 80%;
}
canvas {
  border: 1px solid #000;
  max-width: 100%;
  height: auto;
  display: block;
  touch-action: none;
}
footer {
  margin: 4em 1em 2em 1em;
}
</style>
</head>
<body>
  <div id="intro">
    <h1 style="text-align: center;">H-Space similarity explorer</h1>
    <p>
      This is a simple tool to explore the similarity between the representations of different concepts at different spatial positions in the h-space of diffusion models.
      The h-space is the output of the midblock of the diffusion unet, however, you can also explore the output from other positions in the unet.
      If you want to try this for your own prompts, check out the source code <a href="https://github.com/JonasLoos/h-space-similarity-explorer">here</a>.
    </p>
    <div class="select-container">
      <div class="select-group">
        <span>Model:</span>
        <select id="model-to-use"></select>
      </div>
      <div class="select-group">
        <span>Position:</span>
        <select id="position-to-use"></select>
      </div>
      <div class="select-group">
        <span>Similarity measure:</span>
        <select id="similarity-measure">
          <option value="cosine">Cosine</option>
          <option value="cosine_centered">Cosine (centered)</option>
          <!-- <option value="soft_cosine">Soft Cosine</option> -->
          <option value="manhattan">Manhattan (L1)</option>
          <option value="euclidean">Euclidean (L2)</option>
          <option value="chebyshev">Chebyshev (L∞)</option>
        </select>
      </div>
    </div>    
    <p id="model-description">
    </p>
    <p>
      All displayed similarities are relative to the currently hovered, or last clicked, tile.
    </p>
  </div>
  <div id="canvas-container"></div>
  <footer>
    <p style="text-align: center;">Made by <a href="https://github.com/JonasLoos">Jonas Loos</a> in 2024. The source code can be found on <a href="https://github.com/JonasLoos/h-space-similarity-explorer">Github</a>.</p>
  <script>

// initialize variables
let last_tile = null;  // remember the last tile to avoid redundant updates
let last_clicked = { concept_id: 0, col: 9, row: 6 };  // remember the last clicked tile
const available_models = [];
let current_model = null;
let available_positions = {};
let current_position = 'mid_block';  // assuming mid_block is always available
const repr_cache = {};
const available_concepts = [];
const concepts = [];
const initial_concepts = [
  { index: 0, initial_step: 4, },
  { index: 1, initial_step: 4, },
  { index: 1, initial_step: 2, },
];
let render_counter = 0;


// setup webworker for similarity computations
const worker = new Worker('worker/webworker.js', { type: 'module' });
worker.jobs = {};
worker.job_counter = 0;
worker.last_calc_promise = new Promise((resolve, reject) => {resolve()});
worker.onmessage = (e) => {
  const { id, data } = e.data;
  const job = worker.jobs[id];
  if (job) {
    job.resolve(data);
    delete worker.jobs[id];
  } else {
    console.warn('Received message for unknown job:', e.data);
  }
  if (data.status === 'error') {
    console.error('Error in worker:', data.msg);
  }
};
worker.onerror = (e) => {
  console.error('Error in worker:', e);
};

// send a task to the webworker and return a promise
function callWasm(task, data) {
  const id = worker.job_counter++;
  const job = {id, task};
  const promise = new Promise((resolve, reject) => {
    job.resolve = resolve;
    job.reject = reject;
  });
  if (task === 'calc_similarities') {
    // wait for the previous calculation task to finish
    const render_counter_backup = render_counter;
    worker.last_calc_promise = worker.last_calc_promise.then(() => {
      if (render_counter_backup !== render_counter) {
        return 'outdated';
      }
      worker.jobs[id] = job;
      worker.postMessage({ id, task, data });
      return promise;
    });
    return worker.last_calc_promise;
  } else if (task === 'fetch_repr') {
    worker.jobs[id] = job;
    worker.postMessage({ id, task, data });
    return promise;
  } else {
    console.error('Unknown task:', task);
  }
}


// create an html element with attributes and append it to a parent
function createElem(tag, attrs, parent) {
  const elem = document.createElement(tag);
  for (const [key, value] of Object.entries(attrs)) {
    elem[key] = value;
  }
  parent.appendChild(elem);
  return elem;
}


// setup a concept with its canvas elements and event listeners
function setupConcept({ index, initial_step }, i) {
  const { name, prompt } = available_concepts[index];
  const concept = { id: i, name: name, step: initial_step, repr_loading_started: null };
  // create html elements
  const canvas_container = document.getElementById('canvas-container');
  const title_elem = createElem('h2', {}, canvas_container)
  const title_text = createElem('span', { textContent: name }, title_elem);
  const title_select = createElem('select', { id: `concept-${i}-select` }, title_elem);
  available_concepts.forEach((concept, i) => {
    const option = createElem('option', { value: concept.name, textContent: concept.name }, title_select);
    if (concept.name === name) option.selected = true;
  });
  const prompt_p = createElem('p', { textContent: `Prompt: ${prompt}` }, canvas_container);
  const canvas_row = createElem('div', {}, canvas_container);
  const image_canvas_div = createElem('div', {}, canvas_row);
  const canvas_size = image_canvas_div.clientWidth * 2;
  const title_p = createElem('p', {}, image_canvas_div);
  createElem('span', { textContent: `Image at step:` }, title_p);
  const slider = createElem('input', { type: 'range', min: '1', max: current_model.steps, value: concept.step }, title_p);
  const slider_value = createElem('span', { textContent: concept.step }, title_p);
  concept.image_canvas = createElem('canvas', { width: canvas_size, height: canvas_size }, image_canvas_div);
  const tile_canvas_div = createElem('div', {}, canvas_row);
  const text_p = createElem('p', {}, tile_canvas_div);
  createElem('span', { textContent: `Similarities, avg: ` }, text_p);
  concept.text = createElem('span', { textContent: `?` }, text_p);
  concept.tile_canvas = createElem('canvas', { width: canvas_size, height: canvas_size }, tile_canvas_div);

  // get convas contexts
  concept.image_ctx = concept.image_canvas.getContext('2d');
  concept.tile_ctx = concept.tile_canvas.getContext('2d');

  concept.getUrl = () => `${window.location.pathname.split('/').slice(0, -1).join('/')}/representations/${current_model.short}/${concept.name}/repr-${current_position}-${concept.step-1}.bin`;

  // helper function to load the representation
  const getRepr = () => {
    concept.repr_loading_started = Date.now();
    const { steps, n, m } = current_model.getShapes();
    callWasm('fetch_repr', { url: concept.getUrl(), n, m })
      .then(() => {
        console.log(`Fetched repr for ${concept.name} from ${concept.getUrl()}`);
        updateCanvasesWithLastClicked();
      })
      .catch(error => {
        console.error(`Error while fetching repr for ${concept.name}:`, error);
        self.repr_loading_started = null;
      });
  };

  // update
  concept.update = () => {
    // setup prompt related elements
    const prompt = available_concepts.find(x => x.name === concept.name).prompt;
    title_text.textContent = concept.name;
    prompt_p.textContent = prompt;

    // update slider
    concept.step = Math.max(Math.round(concept.step / slider.max * current_model.steps), 1)  // update step to match new model
    slider.max = current_model.steps;
    slider.value = concept.step;
    slider_value.textContent = concept.step;

    // update image und representation
    concept.img.src = '';
    concept.img.src = `representations/${current_model.short}/${concept.name}/${concept.step}.jpg`;
    updateCanvasesWithLastClicked();
    getRepr();
  }

  // load image
  concept.img = new Image();
  concept.img.src = `representations/${current_model.short}/${concept.name}/${concept.step}.jpg`;
  concept.img.onload = function() {
    concept.image_ctx.globalCompositeOperation = 'destination-over';
    concept.image_ctx.drawImage(concept.img, 0, 0, concept.image_ctx.canvas.width, concept.image_ctx.canvas.height);
    concept.image_ctx.globalCompositeOperation = 'source-over';
    updateCanvasesWithLastClicked();
  }

  // load representation
  getRepr();

  // setup event listeners  
  [concept.image_canvas, concept.tile_canvas].forEach(canvas => {
    const moveHandler = function(event) {
      const { col, row } = getMousePos(canvas, event);
      updateCanvases(concept, col, row);
    };
    const leaveHandler = function(event) {
      updateCanvases(concepts[last_clicked.concept_id], last_clicked.col, last_clicked.row);
    };
    const clickHandler = function(event) {
      const { col, row } = getMousePos(canvas, event);
      last_clicked.concept_id = concept.id;
      last_clicked.col = col;
      last_clicked.row = row;
      last_tile = null;  // force update
      updateCanvases(concept, col, row);

      // log clicked tile
      if (!concept.repr) return;
      const { m, n } = current_model.getShapes();
      const block_repr = [];
      const offset = (concept.step-1)*m*n*n + row*n*m + col*m;
      for (let i = 0; i < m; i++) {
        block_repr.push(concept.repr[offset + i]);
      }
      console.log(`Clicked on ${concept.name} at (${col},${row})`, block_repr);
    };
    canvas.addEventListener('click', clickHandler);
    canvas.addEventListener('mousemove', moveHandler);
    canvas.addEventListener('mouseleave', leaveHandler);
    canvas.addEventListener('touchstart', clickHandler);
    canvas.addEventListener('touchmove', clickHandler);
  });
  
  // slider event listener
  slider.addEventListener('input', event => {
    concept.step = event.target.value;
    concept.update();
  });

  // concept name select event listener
  title_select.addEventListener('change', event => {
    const value = event.target.value;
    concept.name = value;
    concept.update();
  });

  // concept is ready
  return concept;
}


// Function to draw the grid
function drawGrid(ctx) {
  const { n } = current_model.getShapes();
  const tile_size = ctx.canvas.width / n;
  ctx.strokeStyle = 'black';
  ctx.lineWidth = 2;
  ctx.beginPath();

  // Draw the vertical lines
  for (let i = 1; i < n; i++) {
    ctx.moveTo(i * tile_size, 0);
    ctx.lineTo(i * tile_size, ctx.canvas.height);
  }

  // Draw the horizontal lines
  for (let i = 1; i < n; i++) {
    ctx.moveTo(0, i * tile_size);
    ctx.lineTo(ctx.canvas.width, i * tile_size);
  }

  ctx.stroke();
}


function drawSimilarities(concept, similarities) {
  const { n, m } = current_model.getShapes();
  const tile_ctx = concept.tile_ctx;
  const tile_size = tile_ctx.canvas.width / n;
  const { width, height } = tile_ctx.canvas;
  tile_ctx.clearRect(0, 0, width, height);
  for (let i = 0; i < n; i++) {
    for (let j = 0; j < n; j++) {
      const similarity = similarities[i + j * n];
      // draw orange for positive similarity and blue for negative similarity
      tile_ctx.fillStyle = similarity > 0 ? `rgba(255, 165, 0, ${similarity})` : `rgba(0, 165, 255, ${-similarity})`;
      tile_ctx.fillRect(i * tile_size, j * tile_size, tile_size, tile_size);
    }
  }
  // calculate and update average similarity
  const averageSimilarity = similarities.reduce((a, b) => a + b, 0) / similarities.length;
  const stdDev = Math.sqrt(similarities.reduce((a, b) => a + (b - averageSimilarity) ** 2, 0) / similarities.length);
  concept.text.textContent = `${averageSimilarity.toFixed(3)}±${stdDev.toFixed(3)}`;
}

function drawError(concept, msg) {
  const tile_ctx = concept.tile_ctx;
  const { width, height } = tile_ctx.canvas;
  tile_ctx.clearRect(0, 0, width, height);
  tile_ctx.fillStyle = 'black';
  tile_ctx.font = '30px Arial';
  tile_ctx.textAlign = 'center';
  tile_ctx.textBaseline = 'middle';
  const loading_dots_count = Math.floor((Date.now() % 1000) / 250);
  const loading_dots = '.'.repeat(loading_dots_count) + ' '.repeat(3 - loading_dots_count);
  msg = msg.replace('...', loading_dots);
  const lines = msg.split('\n');
  for (let i = 0; i < lines.length; i++) {
    tile_ctx.fillText(lines[i], width / 2, height / 2 - (lines.length-1) * 20 + i * 40);
  }
  concept.text.textContent = `?`;
}


// Update the canvas highlightings based on the mouse position
function updateCanvases(base_concept, col, row) {
  const { n } = current_model.getShapes();
  // Check if the mouse is inside the canvas
  if (col < 0 || col >= n || row < 0 || row >= n) return;

  // Check if the tile has changed
  const curr_tile = `${base_concept.id}-${col}-${row}`;
  if (curr_tile === last_tile) return;
  last_tile = curr_tile;
  render_counter++;

  // Update the canvases
  const func = document.getElementById('similarity-measure').value;
  const concepts_sorted = concepts.slice().sort((a, b) => b == base_concept);  // sort concepts so that the base concept is drawn first
  for (const concept of concepts_sorted) {
    const tile_size = concept.image_ctx.canvas.width / n;
    // update image canvas
    const img_ctx = concept.image_ctx;
    img_ctx.clearRect(0, 0, img_ctx.canvas.width, img_ctx.canvas.height);
    try {
      img_ctx.drawImage(concept.img, 0, 0, img_ctx.canvas.width, img_ctx.canvas.height);      
    } catch (error) {
      console.warn('Error while drawing image: ', error.toString());
    }
    drawGrid(img_ctx);
    // highlight the currently selected tile
    if (concept === base_concept) {
      img_ctx.strokeStyle = 'rgba(255, 165, 0, 1)';
      img_ctx.lineWidth = 2;
      img_ctx.strokeRect(col * tile_size, row * tile_size, tile_size, tile_size);
      img_ctx.fillStyle = 'rgba(255, 165, 0, 0.7)';
      img_ctx.fillRect(col * tile_size, row * tile_size, tile_size, tile_size);
    }
    // highlight the last clicked tile (might be confusing)
    // if (last_clicked.concept_id === concept.id && !(last_clicked.col === col && last_clicked.row === row && concept === base_concept)) {
    //   img_ctx.fillStyle = 'rgba(255, 165, 165, 0.7)';
    //   img_ctx.fillRect(last_clicked.col * tile_size, last_clicked.row * tile_size, tile_size, tile_size);
    // }
    // update tile canvas
    callWasm('calc_similarities', { func, repr1_str: base_concept.getUrl(), repr2_str: concept.getUrl(), row, col })
      .then((similarities) => {
        if (similarities === 'loading') {
          drawError(concept, 'Loading...');
          updateCanvasesSoon(base_concept, col, row);  // schedule update to check if representations are loaded and animate loading text
          return;
        } else if (similarities === 'outdated') {
          return;
        } else {
          drawSimilarities(concept, similarities);
          // highlight the most similar tile
          if (concept !== base_concept) {
            const max_sim = Math.max(...similarities);
            const max_idx = similarities.indexOf(max_sim);
            const max_col = max_idx % n;
            const max_row = Math.floor(max_idx / n);
            img_ctx.fillStyle = 'rgba(165, 255, 0, 0.7)';
            img_ctx.fillRect(max_col * tile_size, max_row * tile_size, tile_size, tile_size);
          }
        }
      })
      .catch((error) => {
        // console.error('Error while calculating similarities:', error);
        drawError(concept, 'Error calculating similarities...\nMaybe try to change settings.');
        updateCanvasesSoon(base_concept, col, row);  // schedule update to check if representations are loaded and animate loading text
      });
  }
}


// force update of all canvases with the last clicked tile as the base
function updateCanvasesWithLastClicked() {
  last_tile = null;  // force update
  updateCanvases(concepts[last_clicked.concept_id], last_clicked.col, last_clicked.row);
}


// delayed canvas update
function updateCanvasesSoon(concept, col, row) {
  // force update
  last_tile = null;
  // update now if the last update was more than 100ms ago, to ensure updates even when this function is called often
  if (self.timer && self.last_updated && self.last_updated + 200 < Date.now()) {
    self.last_updated = Date.now();
    updateCanvases(concept, col, row);
  }
  // scehdule update in 200ms
  clearTimeout(self.timer);
  self.timer = setTimeout(() => {
    self.last_updated = Date.now();
    updateCanvases(concept, col, row);
  }, 200);
}



function updatePositionSelector() {
  if (current_model.representations[current_position] === undefined || !current_model.representations[current_position].available)
    current_position = 'mid_block';  // fallback to mid_block, if current position is not available
  document.getElementById('position-to-use').innerHTML = '';  // clear old options
  Object.keys(current_model.representations).forEach(position => {
    // setup available positions (where the representations are extracted from)
    createElem('option', {
      value: position,
      textContent: position == 'mid_block' ? 'mid_block (h-space)' : position,
      disabled: !current_model.representations[position].available,
      title: !current_model.representations[position].available ? 'Not available due to large representation size' : undefined,
      selected: position == current_position
    }, document.getElementById('position-to-use'));
  });
}


function updateModelDescription() {
  const { steps, m, n} = current_model.getShapes();
  document.getElementById('model-description').innerHTML = `The examples below are the output for <a href="https://huggingface.co/${current_model.name}">${current_model.name}</a> at position ${current_position} in the diffusion unet, which has a dimension of (${steps},${m},${n},${n}), i.e. ${steps} steps, ${m} channels, and ${n}x${n} spatial resolution.`;
};


// get mouse/touch position on canvas
function getMousePos(canvas, event) {
  const { n } = current_model.getShapes();
  if (event.touches) event = event.touches[0];  // handle touch events for mobile
  const rect = canvas.getBoundingClientRect();
  const x = event.clientX - rect.left;
  const y = event.clientY - rect.top;
  const col = Math.floor(2 * x / canvas.width * n);
  const row = Math.floor(2 * y / canvas.width * n);
  return { col, row };
}


function init() {
  // load available prompts/concepts
  const prompts_promise = fetch('representations/prompts.json')
    .then(response => response.json())
    .then(data => {
      Object.entries(data).forEach(([name, prompt], i) => {
        available_concepts.push({ name, prompt});
      });
    })
    .catch(error => console.error('Error while fetching prompts info:', error));
  
  // load available models
  const models_promise = fetch('representations/models.json')
    .then(response => response.json())
    .then(data => {
      available_models.push(...data);
      available_models.forEach(model => {
        createElem('option', { value: model.short, textContent: model.short }, document.getElementById('model-to-use'));
        model.getShapes = () => ({ steps: model.steps, m: model.representations[current_position].channels, n: model.representations[current_position].spatial })
      });
      current_model = available_models[0];
      updatePositionSelector();
      updateModelDescription();
    })
    .catch(error => console.error('Error while fetching models info:', error));
  
  // initialize concepts
  Promise.all([prompts_promise, models_promise])
    .then(() => {
      // check if there are any concepts
      if (available_concepts.length === 0) {
        createElem('p', { textContent: 'Error loading Data' }, document.getElementById('canvas-container'));
        return;
      }
      // setup concepts
      initial_concepts.forEach((index, i) => {
        concepts.push(setupConcept(index, i));
      });
    })
    .then(() => {
      // update canvases on load with a delay
      setTimeout(updateCanvasesWithLastClicked, 100);
    });
  
  // setup model change event listener
  document.getElementById('model-to-use').addEventListener('change', event => {
    const prev_hspace_spatial = current_model.getShapes().n;

    // update current model, position selector, and description
    current_model = available_models.find(x => x.short === event.target.value);
    updatePositionSelector();
    updateModelDescription();

    // update last clicked if spatial resolution changed
    const { n } = current_model.getShapes();
    if (n !== prev_hspace_spatial) {
      last_clicked.col = Math.floor(last_clicked.col * n / prev_hspace_spatial);
      last_clicked.row = Math.floor(last_clicked.row * n / prev_hspace_spatial);
    }

    // update concepts and canvases
    concepts.forEach(concept => concept.update());
    updateCanvasesWithLastClicked();
  });

  // setup position change event listener
  document.getElementById('position-to-use').addEventListener('change', event => {
    const prev_hspace_spatial = current_model.getShapes().n;

    // update current position and description
    current_position = event.target.value;
    updateModelDescription();

    // update last clicked if spatial resolution changed
    const { n } = current_model.getShapes();
    if (n !== prev_hspace_spatial) {
      last_clicked.col = Math.floor(last_clicked.col * n / prev_hspace_spatial);
      last_clicked.row = Math.floor(last_clicked.row * n / prev_hspace_spatial);
    }

    // update concepts and canvases
    concepts.forEach(concept => concept.update());
    updateCanvasesWithLastClicked();
  });

  // setup similarity measure change event listener
  document.getElementById('similarity-measure').addEventListener('change', updateCanvasesWithLastClicked);

  // reload on resize to update canvas size
  window.addEventListener('resize', () => {
      concepts.forEach(concept => {
        const canvas_size = concept.image_ctx.canvas.parentElement.clientWidth * 2;
        concept.image_canvas.width = canvas_size;
        concept.image_canvas.height = canvas_size;
        concept.tile_canvas.width = canvas_size;
        concept.tile_canvas.height = canvas_size;
      });
      updateCanvasesWithLastClicked();
  });
}


// start the app
init();

</script>

</body>
</html>