light_finder.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Finding Light\n",
    "\n",
    "In this notebook we show how to find the stylespace channel for light.\n",
    "\n",
    "First, we create a stylegan for bedrooms and show a few generated images."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "from rewrite_utils import zdataset, show, labwidget, paintwidget, renormalize, nethook, imgviz, pbar\n",
    "from torchvision.transforms import Compose, Resize, CenterCrop, ToTensor, Normalize\n",
    "from PIL import Image\n",
    "from rewrite_utils.stylegan2 import load_seq_stylegan\n",
    "#import clip\n",
    "\n",
    "\n",
    "gan_model = load_seq_stylegan('bedroom', mconv='seq', truncation=0.90)\n",
    "nethook.set_requires_grad(False, gan_model)\n",
    "zds = zdataset.z_dataset_for_model(gan_model, size=1000)\n",
    "\n",
    "image_numbers = range(100)#  [46] [77] [59] [95]\n",
    "def pilim(idata):\n",
    "    return renormalize.as_image(idata)\n",
    "show([\n",
    "    [pilim(gan_model(zds[i][0][None].cuda())[0]), f'Image {i}']\n",
    "    for i in image_numbers])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Notice that image number 41 has really nice light patterns cast from the lamps.\n",
    "\n",
    "Let's find out how the model represents those light patterns.\n",
    "\n",
    "First step is, we will manually make a mask over the light patterns we're interested in:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "image_number = 41\n",
    "\n",
    "PW = paintwidget.PaintWidget(image=renormalize.as_url(\n",
    "    gan_model(zds[image_number][0][None].cuda())[0]))\n",
    "def do_reset():\n",
    "    PW.mask = ''\n",
    "RB = labwidget.Button('reset').on('click', do_reset)\n",
    "show([PW, RB])\n",
    "\n",
    "# I'll pre-draw a mask I drew for the experiment.\n",
    "PW.mask = 'data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAQAAAAEACAYAAABccqhmAAAMfElEQVR4Xu3dv48cZxkH8HcWI3FpsJsonosUR7RInEsQ2EtPgUUBZST+AKCjAlNRQnokiwIJKjv/wVxAtLZEQ0cs5dZWhGRHQo4U2zdozmd0Nvdjf8yz+/jez3WRd5593s/z3jc7e7M7TfFDgEC1Ak21K7dwAgSKALAJCFQsIAAqHr6lExAA9gCBigUEQMXDt3QCAsAeIFCxgACoePiWTkAA2AMEKhYQABUP39IJCAB7gEDFAgKg4uFbOgEBYA8QqFhAAFQ8fEsnIADsAQIVCwiAiodv6QQEgD1AoGIBAVDx8C2dgACwBwhULCAAKh6+pRMQAPYAgYoFBEDFw7d0AgLAHiBQsYAAqHj4lk5AANgDBCoWEAAVD9/SCQgAe4BAxQICoOLhWzoBAWAPEKhYQABUPHxLJyAA7AECFQsIgIqHb+kEBIA9QKBiAQFQ8fAtnYAAsAcIVCwgACoevqUTEAD2AIGKBQRAxcO3dAICwB4gULGAAKh4+JZOQADYAwQqFhAAFQ/f0gkIAHuAQMUCAqDi4Vs6AQFgDxCoWEAAVDx8SycgAOwBAhULCICKh2/pBASAPUCgYgEBUPHwLZ2AALAHCFQsIAAqHr6lExAA9gCBigUEQMXDt3QCAsAeIFCxgACoePiWTkAA2AMEKhYQABUP39IJCIAF9sCjR9OLly51jxc4xEMJpBYQAKeMZzab/rzvy7dKKTtNU3YOH/pJKeWTpildKeWPly93w3/7IfBGCgiAY8b22WfTnadPy60jv/THDrfvy+OmKTfbtvvwjZy+pqsXEACvbYG9venNpim/XnBndFtb5cZYpwdONRbU9/ClBQTAEbrhJX8p5XdLanZt231/mWOdaiyj5pgxBATAoeKDB9Mr+/vlbtOUiyvA/qJtu9/Pe7xTjXmlPC5KQAAcys5m0+FNveurQA/vCUwm5eo8bwxmONVYZa2OPR8CAqCUMvyf+NmzcneMkfZ9+c32dnfztFqbOtUYY31qnC8BAVBKefBg+kHfl1sjjXa3bbvpSbU2caox0rqUOYcCAqCUMptNh/P2n40x3+E0YHu7u3RSrXWfaoyxJjXOr4AAeBEAK5//H90iTVPeP+59gHWfapzfbWtlYwkIgIAAaNvuWNd1nmqMtUHUOd8CAqCUsuQ78iftjPtt21057h/Xeapxvret1Y0lIABevAL4YSnl9kioH7VtN9T7v591nWqMtA5lKhAQAC/+CnCl78u/Rpr3iRcDjR0AJ51qjLQOZSoQEACHQx7p5fn9ra2yc9JnAtZ1qlHBvrXEkQQEwCHk8AGcJ09K1zQHH/9d6ufChXL17be7eycdvK5TjaWad1CVAgLgyNhX+TPdPFcArutUo8qdbNFLCQiA19gOQ+BOKeW9OUU/L+XgOwHm+hDQOk415uzbwwgUAXDMJhhOB774ogzX8591deBu05QP5vnwz8unWcephn1NYF4BAXCK1OGrgeEvBMNXgr28vn/4CrB7TVPuXb7cDVcQLvwTfaqxcEMOqFZAAGxo9NGnGhtalqd9wwQEwAYHFnmqscFleeo3SEAAJBhW1KlGgqVpIbmAAEg+IO0RiBQQAJG6ahNILiAAkg9IewQiBQRApK7aBJILCIDkA9IegUgBARCpu0LtR4++896lS3+/v0IJhxI4U0AAnEm0ngfs7V2/XUrzzVL67aZptg6f9Vnfl/9MJuVu35dftW33t/V041lqERAAG570w4fXfvz8eXPryC/9CR31/fCtRW27+6MNt+zpz5GAANjgMD/9dNpNJovdjWj42vG33vpyx+nBBgd3jp5aAGxomMNL/qZpjv3uwLNaOuveA2cd798JvBQQAKfshajbdM9m0++W0n9cSrO0f9/3d7a3d2/YygRWEVh6A67ypFmPXddtumez6eNSytdXcxjeE2iueWNwNcXajxYAhzcHffq03GqasnPahhheejfNwbf/fLjsxhne9Nvfn/x52eOPHre/X3bffffk+xCO8RxqnG+B6gNgyW/q7ba2yo2Tvv33tC0zm03/UEr56Ujb6vO27S6OVEuZCgWqDoBN3KZ7b+/63aZpTn2lMf8+7Pu23Z3M/3iPJPCqQLUBsKnbdI9z/v/KEL/nfQC/1ssKVBsAY9ylZ3hPYDIpVxf5UtCxA8DdgZbd+o4bBKoMgFW+lPP1bTPP/QCOHrPMxT+nbNXnbdtdsJUJLCtQZQBs8jbde3vXfts0k18uO7BXj+sftu3u5XFqqVKjQJUBMNLNOQ72y6JX5b24CKj8dYzN5mKgMRTrrlFrAAzf5399rNE3TXl/kfcBRvpLwPOtrS+/4TMBY02xzjoCYIS5L/pG3PBZ/ydPvvrPpmm+tuzTTyb7P3nnnY//suzxjiNQ7ZuAS178c9KOud+23ZVFt9MqVwS6AnBRbY8/SaDWVwDDp/Buj7QtPmrbbqlP9R2GwJ9KKV+Zp5f+4CsBig8BzYPlMXMJVBkAmW7TfXg6cGeOqwOHuxD/wEU/c+1rD5pToMoAGGxG+kvA/a2tsrPMZwJen8+LbwYqO30/+fZk8r8PJf277/t/NE2zO+/tx+ecu4cROBCoNgDcpttvAIGKA2AY/ipXBC56BaDNRiCjQLWvAF4Ow226M25LPa1LoPoAGKDdpntd283zZBMQAEcm4jbd2banfqIFBEC0sPoEEgsIgMTD0RqBaAEBEC2sPoHEAgIg8XC0RiBaQABEC6tPILGAAEg8HK0RiBYQANHC6hNILCAAEg9HawSiBQRAtLD6BBILCIDEw9EagWgBARAtrD6BxAICIPFwtEYgWkAARAurTyCxgABIPBytEYgWEADRwuoTSCwgABIPR2sEogUEQLSw+gQSCwiAxMPRGoFoAQEQLaw+gcQCAiDxcLRGIFpAAEQLq08gsYAASDwcrRGIFhAA0cLqE0gsIAASD0drBKIFBEC0sPoEEgsIgMTD0RqBaAEBEC2sPoHEAgIg8XC0RiBaQABEC6tPILGAAEg8HK0RiBYQANHC6hNILCAAEg9HawSiBQRAtLD6BBILCIDEw9EagWgBARAtrD6BxAICIPFwtEYgWkAARAurTyCxgABIPBytEYgWEADRwuoTSCwgABIPR2sEogUEQLSw+gQSCwiAxMPRGoFoAQEQLaw+gcQCAiDxcLRGIFpAAEQLq08gsYAASDwcrRGIFhAA0cLqE0gsIAASD0drBKIFBEC0sPoEEgsIgMTD0RqBaAEBEC2sPoHEAgIg8XC0RiBaQABEC6tPILGAAEg8HK0RiBYQANHC6hNILCAAEg9HawSiBQRAtLD6BBILCIDEw9EagWgBARAtrD6BxAICIPFwtEYgWkAARAurTyCxgABIPBytEYgWEADRwuoTSCwgABIPR2sEogUEQLSw+gQSCwiAxMPRGoFoAQEQLaw+gcQCAiDxcLRGIFpAAEQLq08gsYAASDwcrRGIFhAA0cLqE0gsIAASD0drBKIFBEC0sPoEEgsIgMTD0RqBaAEBEC2sPoHEAgIg8XC0RiBaQABEC6tPILGAAEg8HK0RiBYQANHC6hNILCAAEg9HawSiBQRAtLD6BBILCIDEw9EagWgBARAtrD6BxAICIPFwtEYgWkAARAurTyCxgABIPBytEYgWEADRwuoTSCwgABIPR2sEogUEQLSw+gQSCwiAxMPRGoFoAQEQLaw+gcQCAiDxcLRGIFpAAEQLq08gsYAASDwcrRGIFhAA0cLqE0gsIAASD0drBKIFBEC0sPoEEgsIgMTD0RqBaAEBEC2sPoHEAgIg8XC0RiBaQABEC6tPILGAAEg8HK0RiBYQANHC6hNILCAAEg9HawSiBQRAtLD6BBILCIDEw9EagWgBARAtrD6BxAICIPFwtEYgWkAARAurTyCxgABIPBytEYgWEADRwuoTSCwgABIPR2sEogUEQLSw+gQSCwiAxMPRGoFoAQEQLaw+gcQCAiDxcLRGIFpAAEQLq08gsYAASDwcrRGIFhAA0cLqE0gsIAASD0drBKIFBEC0sPoEEgsIgMTD0RqBaAEBEC2sPoHEAgIg8XC0RiBaQABEC6tPILGAAEg8HK0RiBYQANHC6hNILCAAEg9HawSiBQRAtLD6BBILCIDEw9EagWgBARAtrD6BxAICIPFwtEYgWkAARAurTyCxgABIPBytEYgWEADRwuoTSCwgABIPR2sEogUEQLSw+gQSCwiAxMPRGoFoAQEQLaw+gcQCAiDxcLRGIFpAAEQLq08gsYAASDwcrRGIFhAA0cLqE0gs8F99KF8fqMSP9QAAAABJRU5ErkJggg=='"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "PW.image = PW.mask"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#PW.image"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we can hunt through all the layers and identify S channels that, when zeroed, alter the masked pixels the most.  We will go through many layers and units."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "device = 'cuda'\n",
    "mask = renormalize.from_url(PW.mask, target='pt')[0].to(device)\n",
    "z = zds[image_number][0][None].cuda()\n",
    "test_result = []\n",
    "for layername in pbar([f'layer{l}' for l in range(3, 14)]): # or 3, 14 for all the layers\n",
    "    with nethook.Trace(gan_model, f'{layername}.sconv.mconv.modulation') as slayer:\n",
    "        baseline = gan_model(z)[0]\n",
    "    for u in range(slayer.output.style.shape[1]):\n",
    "        def edit_style(output):\n",
    "            output.style[:, u] = 0\n",
    "            return output\n",
    "        with nethook.Trace(gan_model, f'{layername}.sconv.mconv.modulation', edit_output=edit_style):\n",
    "            changed = gan_model(z)[0]\n",
    "            delta = ((baseline - changed).abs() * mask).sum()\n",
    "            test_result.append((delta.item(), (layername, u)))\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Sort through the tested units and visualize the best 12."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "top = list(reversed(sorted(test_result)))\n",
    "out = []\n",
    "for rank, (score, (layername, u)) in enumerate(top[:20]):\n",
    "    with nethook.Trace(gan_model, f'{layername}.sconv.mconv.modulation', edit_output=edit_style):\n",
    "        changed = gan_model(zds[image_number][0][None].cuda())[0]\n",
    "        delta = (baseline - changed)\n",
    "        out.append([pilim(delta).resize((180, 180)), f'#{rank+1}: {layername} neuron {u}'])\n",
    "show(out)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    " * layer8 265 looks like it switches lights cleanly, both lampshade and reflectiion\n",
    " * layer11 398 controls reflected light only, without lampshade\n",
    " * layer11 112 controls yellowness of reflected light\n",
    " * layer9 252 specifically controls umbra shadows from light\n",
    " \n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's make a simple widget to control one of these."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from importlib import reload\n",
    "\n",
    "reload(labwidget)\n",
    "rw = labwidget.Range()\n",
    "im = labwidget.Image()\n",
    "cp = labwidget.ColorPicker('#ffffff', desc='color: ')\n",
    "num = labwidget.Textbox(desc='imgnum: ')\n",
    "lay = labwidget.Textbox(desc='layername: ')\n",
    "uni = labwidget.Textbox(desc='unitnum: ')\n",
    "\n",
    "num.value = image_number\n",
    "lay.value = 'layer8'\n",
    "uni.value = 265\n",
    "\n",
    "\n",
    "def get_lit_scene(imgnum, amount, layername, unitnum):\n",
    "    def change_light(output):\n",
    "        output.style[:, int(unitnum)] = 10 * amount\n",
    "        return output\n",
    "    with nethook.Trace(gan_model, f'{layername}.sconv.mconv.modulation', edit_output=change_light):\n",
    "        return gan_model(zds[imgnum][0][None].cuda())[0]\n",
    "\n",
    "baseline = get_lit_scene(int(num.value), -1.0, lay.value, uni.value)\n",
    "im.render(pilim(baseline))\n",
    "\n",
    "def readcolor():\n",
    "    try:\n",
    "        floatcolor = [float(int(cp.value[i:i+2], 16))/255.0  for i in [1,3,5]]\n",
    "        color = torch.tensor(floatcolor).float()\n",
    "        if len(color) == 3:\n",
    "            return color\n",
    "    except:\n",
    "        pass\n",
    "    return torch.tensor([1.0, 1.0, 1.0]).float()\n",
    "\n",
    "def newimage():\n",
    "    frac = ((float(rw.value) * 2 - 100) / 100.0)\n",
    "    lit = get_lit_scene(int(num.value), frac, lay.value, uni.value)\n",
    "    baseline = get_lit_scene(int(num.value), -1.0, lay.value, uni.value)\n",
    "    light = lit - baseline\n",
    "    color = readcolor().cuda()[:,None,None]\n",
    "    colored = baseline + (light * color)\n",
    "    im.render(pilim(colored))\n",
    "\n",
    "show([[num, '', lay, '', uni, '', rw, im]])\n",
    "\n",
    "num.on('value', newimage)\n",
    "cp.on('value', newimage)\n",
    "rw.on('value', newimage)\n",
    "None"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}