Amendment to the stable diffusion notebook to do titling and generate /n/ images

diff.py

from PIL import Image
from PIL import ImageFont
from PIL import ImageDraw 
import math
#Self contain this for easy single clicking
from PIL import Image
def image_grid(imgs, rows, cols):
    #assert len(imgs) == rows*cols
    w, h = imgs[0].size
    grid = Image.new('RGB', size=(cols*w, rows*h))
    grid_w, grid_h = grid.size  
    for i, img in enumerate(imgs):
        grid.paste(img, box=(i%cols*w, i//cols*h))
    return grid

# grab the font
!test -f "/content/SansSerif.ttf" && echo "Skipping download" || wget -nc "https://github.com/PrincetonUniversity/COS333_Comet/blob/master/android/app/src/main/assets/fonts/Microsoft%20Sans%20Serif.ttf?raw=true" 
!test -f "/content/Microsoft Sans Serif.ttf?raw=true" && mv "/content/Microsoft Sans Serif.ttf?raw=true" "/content/SansSerif.ttf"
font = ImageFont.truetype("/content/SansSerif.ttf", 120)

# add a couple of options  
# can be any whole number above 0
manual_seed = 1024 #@param {type: 'integer'} 
# make it divisible by your column count and your row count
image_count = 9  #@param {type: 'integer'} 
# how may we add each time we do a new image
step_increase = 12 #@param {type: 'integer'} 
# jump over the first few steps to avoid NSFW flag
NSFW_skip_start_step = 10 #@param {type: 'integer'} 
height = 512 #@param {type: 'integer'} 
width = 1088 #@param {type: 'integer'} 
guidance_scale =8 #@param {type: 'integer'} 

# we are ready to prompt.
images=[]
prompt = "simple test aurora dramatic mountaintop, distant glowing figures" #@param {type: 'string'} 

# do the run
current_steps = NSFW_skip_start_step
for i in range(image_count):
  with autocast("cuda"):
        #eugh hack time  - skip the known bad step iterators.
    if(current_steps == 9):
      current_steps = 10
    if(current_steps >26 and current_steps<30):
      current_steps = 30
    if(current_steps >100 and current_steps<125):
      current_steps = 125
    if(current_steps >35 and current_steps<40):
      current_steps = 40
    if(current_steps >300 and current_steps<350):
      current_steps = 350  


    title = str(current_steps)
    generator = torch.Generator("cuda").manual_seed(manual_seed)
    try:
      img = pipe(prompt, generator=generator, height=height, width=width, 
                 num_inference_steps=current_steps, guidance_scale=guidance_scale )["sample"][0]
    finally:
      #document every fail
      current_steps
    draw = ImageDraw.Draw(img)
    draw.text((width/2, 50), title, (255,255,255),font=font)
    #img.save('/content/MyDrive/uniquename.jpg') TODO
    images.append(img)
    current_steps += step_increase 


#single, double, triple row et cetera - if indivisible will pad with black blank
rows = 3 #@param {type: 'integer'} 
cols = math.ceil((len(images)/rows))
grid = image_grid(images, rows=rows, cols=cols)

#now show it
grid

added automatic retrieval of font

updated font grab to do tests to save bandwidth

text will now just be placed 50% from the left automatically on any size image
draw.text((width/2, 50), title, (255,255,255),font=font) <--see to adjust this implementation

added a basic GUI, attempts to skip over "bad" step values silently.