Output Grid Code Block for Stable Diffusion Colab

sb-output-grid.py

#@title 🌌 Run to start dreaming.{ vertical-output: true, display-mode: "form" }

import IPython
import base64
from io import BytesIO

all_images = []

# Clear sample output
!rm /content/stable-diffusion/outputs/txt2img-samples/samples/*
!rm /content/stable-diffusion/outputs/txt2img-samples/grid*

# Python really needs objects like JS
class Option:
  def __init__(self, seed):
    self.value = seed
  def __str__(self):
    return self.value

# Clear code output
clear_output()

### HTML ELEMENTS
# Create an HTML grid to put images in
def createGrid():
  display('create grid..')
  display(IPython.display.Javascript("""
    var seeds = []

    const outputFooter = document.querySelector("#output-footer")
    const grid = document.createElement('div')
    grid.classList.add('stream')
    grid.id = 'image-grid'
    grid.style = `
      display: grid;
      grid-template-columns: repeat(6, 1fr);
      grid-gap: 10px;
      padding: 3em 2em 5em 2em;
    `
    outputFooter.appendChild(grid)

    const showSeeds = document.createElement('div')
    showSeeds.style = `
      border: 2px solid grey;
      background: #333;
      margin: 0 0 2em 0;
      padding: 2em 1em;
      cursor: pointer;
      width: 100%;
    `
    showSeeds.onmouseover = showSeeds.onfocus = () => {{
      showSeeds.style.background = '#111'
    }}
    showSeeds.onmouseout = showSeeds.onblur = () => {{
      showSeeds.style.background = '#333'
    }}
    showSeeds.innerText = 'Click to copy seeds: ' + seeds.join(' ')
    showSeeds.onclick = () => {{
      navigator.clipboard.writeText(seeds.join(' ')).then(() => {{
        console.log('seeds copied to clipboard!')
      }}, () => {{
        console.log('failed to copy to clipboard')
      }})
    }}
    outputFooter.appendChild(showSeeds)

    // Listen for the "updateSeeds" event
    window.addEventListener('updateSeeds', function (e) {{
      seeds = e.detail
      showSeeds.innerText = 'Click to copy seeds: ' + seeds.join(' ')
    }})
  """))
# Add image to grid
def addToGrid(img_base64, seedStr):
  display(IPython.display.Javascript("""
    const grid = document.querySelector('#image-grid')
    const container = document.createElement('div')
    container.style = `
      display: flex;
      flex-direction: column;
    `
    const img = document.createElement('img')
    img.src = `{img_base64}`
    img.style = `
      width: 100%;
      height: 100%;
      object-fit: cover;
    `
    const button = document.createElement('button')
    let active = false
    button.style = `
      transition: all 80ms ease-out;
      border: 4px solid grey;
      border-radius: 4px;
      filter: drop-shadow(0px 0px 0px black);
      padding: 0;
      cursor: pointer;
    `
    button.onmouseover = button.onfocus = () => {{
      button.style.transform = active ? 'scale(1)' : 'scale(1.2)';
      button.style.filter = 'drop-shadow(0px 2px 12px black)';
      button.style.zIndex = '10';
    }}
    button.onmouseout = button.onblur = () => {{
      button.style.transform = active ? 'scale(0.8)' : 'scale(1)';
      button.style.filter = 'drop-shadow(0px 0px 0px black)';
      button.style.zIndex = 'unset';
    }}
    button.onclick = () => {{
      active = !active

      if (active) {{
        seeds.push('{seed}')
        button.style.opacity = '0.5'
        button.style.border = '4px solid #4747f5';
        button.style.transform = 'scale(0.8)';
      }} else {{
        seeds.splice(seeds.indexOf('{seed}'), 1)
        button.style.border = '4px solid grey';
        button.style.opacity = '1'
        button.style.transform = 'scale(1)';
      }}

      // Create new custom event to share the "seeds" array object
      var event = new CustomEvent('updateSeeds', {{detail: seeds}})
      window.dispatchEvent(event)
    }}
    button.appendChild(img)
    container.appendChild(button)
    grid.appendChild(container)
  """.format(img_base64=img_base64, seed=seedStr)))

# Create image grid
createGrid()

### Prepare config
# Get iterable list from widget seedsopt string
widgetDict = get_widget_extractor(widget_opt)
# Remove hanging empty value if it ends with a comma
seeds = [s for s in widgetDict['seeds'].value.strip().replace('  ', ' ').split(' ') if s]
iterCount = widgetDict['n_iter'].value
widgetDict['n_iter'] = Option(1)
display(IPython.display.Javascript('if (typeof seeds === "undefined") { var seeds = [] } else { seeds = [] };'))

# Define run
def doRun(seedStr):
  # Run inference
  run(widgetDict)
  # Update image grid
  buffered = BytesIO()
  all_images.pop().save(buffered, format="JPEG")
  img_str = base64.b64encode(buffered.getvalue())
  img_base64 = (bytes("data:image/jpeg;base64,", encoding='utf-8') + img_str).decode("utf-8") 
  addToGrid(img_base64, seedStr)

# Runs
for baseSeed in seeds:
  for iter in range(iterCount):
    # Seed
    seed = int(baseSeed) + iter
    seedStr = str(seed)
    # Add seed to dict manually
    widgetDict['seed'] = Option(seed)
    print('Seed:', seedStr)
    doRun(seedStr)

# Get Seeds button at bottom
"""display(IPython.display.Javascript('''
  console.log('create get seeds button')
  const button = document.createElement('button')
  button.innerText = 'Get seeds'
  button.onclick = () => {{
    console.log('seeds:', seeds)
    const streams = document.querySelectorAll(".stream")
    streams[streams.length - 1].appendChild(document.createTextNode(`Seeds: ${seeds.join(' ')}`))
  }}
  const streams = document.querySelectorAll(".stream")
  streams[streams.length - 1].appendChild(button)
'''))

for img in all_images:
  print('show image in grid..')
  buffered = BytesIO()
  img.save(buffered, format="JPEG")
  img_str = base64.b64encode(buffered.getvalue())
  img_base64 = (bytes("data:image/jpeg;base64,", encoding='utf-8') + img_str).decode("utf-8") 
  addToGrid(img_base64)
"""

print('Batch complete!')

# Clear VRAM
torch.cuda.empty_cache()








#@title Main one-time Setup Code (replaces txt2img.py from SD repo) { display-mode: "form" }

# Special installs
!pip install diffusers

# Slightly modified version of: https://github.com/CompVis/stable-diffusion/blob/main/scripts/txt2img.py
import argparse, os, sys, glob    
import torch    
import numpy as np
import datetime

from omegaconf import OmegaConf    
from PIL import Image
from PIL.PngImagePlugin import PngInfo    
#from tqdm.auto import tqdm, trange  # NOTE: updated for notebook
from tqdm import tqdm, trange  # NOTE: updated for notebook
from itertools import islice    
from einops import rearrange    
from torchvision.utils import make_grid    
import time    
import rich
from pytorch_lightning import seed_everything    
from torch import autocast    
from contextlib import contextmanager, nullcontext    
    
from ldm.util import instantiate_from_config    
from ldm.models.diffusion.ddim import DDIMSampler    
from ldm.models.diffusion.plms import PLMSSampler
from scripts.txt2img import chunk, load_model_from_config

from IPython.display import clear_output

# Code to turn kwargs into Jupyter widgets
import ipywidgets as widgets
from collections import OrderedDict


def load_model(opt):
    """Seperates the loading of the model from the inference"""
    
    if opt.laion400m:
        print("Falling back to LAION 400M model...")
        opt.config = "configs/latent-diffusion/txt2img-1p4B-eval.yaml"
        opt.ckpt = "models/ldm/text2img-large/model.ckpt"
        opt.outdir = "outputs/txt2img-samples-laion400m"

    config = OmegaConf.load(f"{opt.config}")
    model = load_model_from_config(config, f"{opt.ckpt}")

    if torch.cuda.is_available():
        device = torch.device("cuda")
    else:
        print("Warning - running in CPU mode!")
        device = torch.device("cpu")
    model = model.to(device)

    return model

all_images = []

def run_inference(opt, model):
    """Seperates the loading of the model from the inference
    
    Additionally, slightly modified to display generated images inline
    """
    seed_everything(opt.seed)

    if opt.plms:
        sampler = PLMSSampler(model)
    else:
        sampler = DDIMSampler(model)

    os.makedirs(opt.outdir, exist_ok=True)
    outpath = opt.outdir

    batch_size = opt.n_samples
    n_rows = opt.n_rows if opt.n_rows > 0 else batch_size

    if not opt.from_file:
        prompt = opt.prompt
        assert prompt is not None
        data = [batch_size * [prompt]]
    else:
        #print(f"reading prompts from {opt.from_file}")
        with open(opt.from_file, "r") as f:
            data = f.read().splitlines()
            data = list(chunk(data, batch_size))

    # * Variables for saved image filenames
    # date + time
    datetimeStr = datetime.datetime.now().isoformat()
    # Filename-safe prompt string
    slugPrompt = "".join(c if c.isalnum() else "_" for c in opt.prompt)

    sample_path = os.path.join(outpath, "samples")
    os.makedirs(sample_path, exist_ok=True)
    base_count = len(os.listdir(sample_path))
    grid_count = len(os.listdir(outpath)) - 1

    start_code = None
    if opt.fixed_code:
        start_code = torch.randn([opt.n_samples, opt.C, opt.H // opt.f, opt.W // opt.f], device=device)

    precision_scope = autocast if opt.precision=="autocast" else nullcontext
    with torch.no_grad():
        with precision_scope("cuda"):
            with model.ema_scope():
                tic = time.time()
                all_samples = list()
                for n in range(opt.n_iter): # trange(opt.n_iter, desc="Sampling"):
                    for prompts in data: # tqdm(data, desc="data"):
                        uc = None
                        if opt.scale != 1.0:
                            uc = model.get_learned_conditioning(batch_size * [""])
                        if isinstance(prompts, tuple):
                            prompts = list(prompts)
                        c = model.get_learned_conditioning(prompts)
                        shape = [opt.C, opt.H // opt.f, opt.W // opt.f]
                        samples_ddim, _ = sampler.sample(S=opt.ddim_steps,
                                                         conditioning=c,
                                                         batch_size=opt.n_samples,
                                                         shape=shape,
                                                         verbose=False,
                                                         unconditional_guidance_scale=opt.scale,
                                                         unconditional_conditioning=uc,
                                                         eta=opt.ddim_eta,
                                                         x_T=start_code)

                        x_samples_ddim = model.decode_first_stage(samples_ddim)
                        x_samples_ddim = torch.clamp((x_samples_ddim + 1.0) / 2.0, min=0.0, max=1.0)

                        if not opt.skip_save:
                            for x_sample in x_samples_ddim:
                                x_sample = 255. * rearrange(x_sample.cpu().numpy(), 'c h w -> h w c')
                                imgPath = f'{slugPrompt[:150]}_{opt.seed}_{datetimeStr}_{base_count:05}.png'
                                #Image.fromarray(x_sample.astype(np.uint8)).save(
                                #    os.path.join(sample_path, imgPath))
                                
                                img = Image.fromarray(x_sample.astype(np.uint8)) # Image.open(imgPath)
                                all_images.append(img)
                                metadata = PngInfo()
                                metadata.add_text("artist", 'Chris Hayes')
                                metadata.add_text("copyright", 'Public Domain')
                                metadata.add_text("software", "Stable Diffusion 1.4")
                                metadata.add_text("title", opt.prompt)
                                config = f"prompt: {opt.prompt}, seed: {seed}, steps: {opt.ddim_steps}, CGS: {opt.scale}"
                                metadata.add_text("config", config)
                                img.save(os.path.join(sample_path, imgPath), pnginfo=metadata)

                                base_count += 1

                        if not opt.skip_grid:
                            all_samples.append(x_samples_ddim)

                if not opt.skip_grid:
                    # additionally, save as grid
                    grid = torch.stack(all_samples, 0)
                    grid = rearrange(grid, 'n b c h w -> (n b) c h w')
                    grid = make_grid(grid, nrow=n_rows)

                    # to image
                    grid = 255. * rearrange(grid, 'c h w -> h w c').cpu().numpy()
                    #Image.fromarray(grid.astype(np.uint8)).save(os.path.join(outpath, f'grid-{slugPrompt[:150]}_{opt.seed}_{datetimeStr}_{grid_count:04}.png'))
                    grid_count += 1
                    
                    # display
                    #if opt.display_inline:
                        #clear_output()
                        #display(Image.fromarray(grid.astype(np.uint8)))

                toc = time.time()

    #print(f"Your samples have been saved to: \n{outpath} \n"
    #      f" \nEnjoy.")


def run(opt):
    """If the model parameters changed, reload the model, otherwise, just do inference"""

    #print(f"Creating image ({opt.H},{opt.W}) from prompt:\n\"{opt.prompt}\"\n")
    
    # FIXME global hack
    global last_config
    global last_ckpt
    global model

    if (opt.config != last_config) or (opt.ckpt != last_ckpt):
        model = load_model(opt)
        # FIXME global hack
        last_config = opt.config
        last_ckpt   = opt.ckpt

    run_inference(opt, model)


# FIXME global hack
last_config = ""
last_ckpt   = ""


####### Widget GUI code #######

def get_widget_extractor(widget_dict):
    # allows accessing after setting, this is to reduce the diff against the argparse code
    class WidgetDict(OrderedDict):
        def __getattr__(self,val):
            return self[val].value
    return WidgetDict(widget_dict)

# Allows long widget descriptions
style  = {'description_width': 'initial'}
# Force widget width to max
layout = widgets.Layout(width='100%')

# args from argparse converted to widgets:
# https://github.com/CompVis/stable-diffusion/blob/main/scripts/txt2img.py#L48-L177
widget_opt = OrderedDict()

widget_opt['outdir'] = widgets.Text(
    layout=layout, style=style,
    description='dir to write results to',
    value="outputs/txt2img-samples",
    disabled=False
)
widget_opt['skip_grid'] = widgets.Checkbox(
    layout=layout, style=style,
    value=False,
    description='do not save a grid, only individual samples. Helpful when evaluating lots of samples',
    indent=False,
    disabled=False
)
widget_opt['skip_save'] = widgets.Checkbox(
    layout=layout, style=style,
    value=False,
    description='do not save individual samples. For speed measurements.',
    indent=False,
    disabled=False
)
widget_opt['plms'] = widgets.Checkbox(
    layout=layout, style=style,
    value=False,
    description='use plms sampling (not checked = ddim)',
    indent=False,
    disabled=False
)
widget_opt['laion400m'] = widgets.Checkbox(
    layout=layout, style=style,
    value=False,
    description='uses the LAION400M model',
    indent=False,
    disabled=False
)
widget_opt['fixed_code'] = widgets.Checkbox(
    layout=layout, style=style,
    value=False,
    description='if enabled, uses the same starting code across samples',
    indent=False,
    disabled=False
)
widget_opt['ddim_eta'] = widgets.FloatText(
    layout=layout, style=style,
    description='ddim eta (eta=0.0 corresponds to deterministic sampling',
    value=0.0,
    disabled=False
)
widget_opt['C'] = widgets.IntText(
    layout=layout, style=style,
    description='latent channels',
    value=4,
    disabled=False
)
widget_opt['f'] = widgets.IntText(
    layout=layout, style=style,
    description='downsampling factor',
    value=8,
    disabled=False
)
widget_opt['n_rows'] = widgets.IntText(
    layout=layout, style=style,
    description='rows in the grid (default: n_samples)',
    value=0,
    disabled=False
)
widget_opt['scale'] = widgets.FloatText(
    layout=layout, style=style,
    description='unconditional guidance scale: eps = eps(x, empty) + scale * (eps(x, cond) - eps(x, empty))',
    value=7.5,
    disabled=False
)
widget_opt['from_file'] = widgets.Text(
    layout=layout, style=style,
    description='if specified, load prompts from this file',
    value=None,
    disabled=False
)
widget_opt['config'] = widgets.Text(
    layout=layout, style=style,
    description='path to config which constructs model',
    value="configs/stable-diffusion/v1-inference.yaml",
    disabled=False
)
widget_opt['ckpt'] = widgets.Text(
    layout=layout, style=style,
    description='path to checkpoint of model',
    value="models/ldm/stable-diffusion-v1/model.ckpt",
    disabled=False
)
widget_opt['precision'] = widgets.Combobox(
    layout=layout, style=style,
    description='evaluate at this precision',
    value="autocast",
    options=["full", "autocast"],
    disabled=False
)
# Extra option for the notebook
widget_opt['display_inline'] = widgets.Checkbox(
    layout=layout, style=style,
    value=True,
    description='display output images inline (in addition to saving them)',
    indent=False,
    disabled=False
)

# Common
widget_opt['n_iter'] = widgets.IntText(
    layout=layout, style=style,
    description='sample this often',
    value=1,
    disabled=False
)
widget_opt['n_samples'] = widgets.IntText(
    layout=layout, style=style,
    description='how many samples to produce for each given prompt. A.k.a. batch size',
    value=1,
    disabled=False
)
widget_opt['H'] = widgets.IntText(
    layout=layout, style=style,
    description='image height, in pixel space',
    value=512,
    disabled=False
)
widget_opt['W'] = widgets.IntText(
    layout=layout, style=style,
    description='image width, in pixel space',
    value=512,
    disabled=False
)
widget_opt['prompt'] = widgets.Text(
    layout=layout, style=style,
    description='the prompt to render',
    #value="a painting of a virus monster playing guitar",  # script default
    value="a photograph of an astronaut riding a horse",  # README default
    disabled=False
)
widget_opt['ddim_steps'] = widgets.IntText(
    layout=layout, style=style,
    description='number of ddim sampling steps',
    value=50,
    disabled=False
)
widget_opt['seeds'] = widgets.Text(
    layout=layout, style=style,
    description='multiple seeds for batch runs (separate with a space)',
    value='42',
    disabled=False
)

# Button that runs the 
# Alternatively, you can just run the following in a new cell:
# run(get_widget_extractor(widget_opt))
run_button = widgets.Button(
    description='CLICK TO DREAM',
    disabled=False,
    button_style='', # 'success', 'info', 'warning', 'danger' or ''
    tooltip='Click to run (settings will update automatically)',
    icon='check'
)

run_button_out = widgets.Output()
# this doesn't get used
def on_run_button_click(b):
    with run_button_out:
        widgetDict = get_widget_extractor(widget_opt)
        for seed in widgetDict['seeds'].split(','):
            #clear_output()
            widgetDict['seed'] = seed
            run(widgetDict)

run_button.on_click(on_run_button_click)

# Package into box and render
#primary_options = ['prompt', 'outdir']  # options to put up top
#secondary_options = [k for k in widget_opt.keys() if k not in primary_options]  # rest, ordered by insertion

load_options = ['config', 'ckpt']
inference_options = [k for k in widget_opt.keys() if k not in load_options]  # rest, ordered by insertion
assert all([k in inference_options + load_options for k in widget_opt.keys()])  # make sure we didn't miss any options

# Package into box for rendering
gui = widgets.VBox(
    [widget_opt[k] for k in inference_options] + [widget_opt[k] for k in load_options]  # + [run_button, run_button_out]
)