feature complete; looks like ready for merge

This commit is contained in:
Lincoln Stein 2022-08-25 17:26:48 -04:00
parent 26dc05e0e0
commit b3e3b0e861
5 changed files with 61 additions and 42 deletions

View File

@ -2,6 +2,7 @@ Feature requests:
1. "gobig" mode - split image into strips, scale up, add detail using
img2img and reassemble with feathering. Issue #66.
See https://github.com/jquesnelle/txt2imghd
2. Port basujindal low VRAM optimizations. Issue #62

View File

@ -2,6 +2,7 @@
import os
import atexit
import re
from math import sqrt,floor,ceil
from PIL import Image,PngImagePlugin
# -------------------image generation utils-----
@ -24,7 +25,7 @@ class PngWriter:
print(e)
self.files_written.append([self.filepath,seed])
def unique_filename(self,seed,previouspath):
def unique_filename(self,seed,previouspath=None):
revision = 1
if previouspath is None:
@ -62,6 +63,22 @@ class PngWriter:
info.add_text("Dream",prompt)
image.save(path,"PNG",pnginfo=info)
def make_grid(self,image_list,rows=None,cols=None):
image_cnt = len(image_list)
if None in (rows,cols):
rows = floor(sqrt(image_cnt)) # try to make it square
cols = ceil(image_cnt/rows)
width = image_list[0].width
height = image_list[0].height
grid_img = Image.new('RGB',(width*cols,height*rows))
for r in range(0,rows):
for c in range (0,cols):
i = r*rows + c
grid_img.paste(image_list[i],(c*width,r*height))
return grid_img
class PromptFormatter():
def __init__(self,t2i,opt):
self.t2i = t2i
@ -80,8 +97,6 @@ class PromptFormatter():
switches.append(f'-H{opt.height or t2i.height}')
switches.append(f'-C{opt.cfg_scale or t2i.cfg_scale}')
switches.append(f'-m{t2i.sampler_name}')
if opt.variants:
switches.append(f'-v{opt.variants}')
if opt.init_img:
switches.append(f'-I{opt.init_img}')
if opt.strength and opt.init_img is not None:

View File

@ -20,7 +20,6 @@ from torch import autocast
from contextlib import contextmanager, nullcontext
import transformers
import time
import math
import re
import traceback

View File

@ -153,54 +153,60 @@ def main_loop(t2i,outdir,parser,log,infile):
continue
normalized_prompt = PromptFormatter(t2i,opt).normalize_prompt()
variants = None
individual_images = not opt.grid
try:
file_writer = PngWriter(outdir,normalized_prompt,opt.batch_size)
callback = file_writer.write_image
callback = file_writer.write_image if individual_images else None
t2i.prompt2image(image_callback=callback,
**vars(opt))
results = file_writer.files_written
image_list = t2i.prompt2image(image_callback=callback,**vars(opt))
results = file_writer.files_written if individual_images else image_list
if None not in (opt.variants,opt.init_img):
variants = generate_variants(t2i,outdir,opt,results)
if opt.grid and len(results) > 0:
grid_img = file_writer.make_grid([r[0] for r in results])
filename = file_writer.unique_filename(results[0][1])
seeds = [a[1] for a in results]
results = [[filename,seeds]]
metadata_prompt = f'{normalized_prompt} -S{results[0][1]}'
file_writer.save_image_and_prompt_to_png(grid_img,metadata_prompt,filename)
except AssertionError as e:
print(e)
continue
except OSError as e:
print(e)
continue
print("Outputs:")
write_log_message(t2i,normalized_prompt,results,log)
if variants is not None:
print('Variants:')
for vr in variants:
write_log_message(t2i,vr[0],vr[1],log)
print("goodbye!")
def generate_variants(t2i,outdir,opt,previous_gens):
variants = []
print(f"Generating {opt.variants} variant(s)...")
newopt = copy.deepcopy(opt)
newopt.iterations = 1
newopt.variants = None
for r in previous_gens:
newopt.init_img = r[0]
prompt = PromptFormatter(t2i,newopt).normalize_prompt()
print(f"] generating variant for {newopt.init_img}")
for j in range(0,opt.variants):
try:
file_writer = PngWriter(outdir,prompt,newopt.batch_size)
callback = file_writer.write_image
t2i.prompt2image(image_callback=callback,**vars(newopt))
results = file_writer.files_written
variants.append([prompt,results])
except AssertionError as e:
print(e)
continue
print(f'{opt.variants} variants generated')
return variants
# variant generation is going to be superseded by a generalized
# "prompt-morph" functionality
# def generate_variants(t2i,outdir,opt,previous_gens):
# variants = []
# print(f"Generating {opt.variants} variant(s)...")
# newopt = copy.deepcopy(opt)
# newopt.iterations = 1
# newopt.variants = None
# for r in previous_gens:
# newopt.init_img = r[0]
# prompt = PromptFormatter(t2i,newopt).normalize_prompt()
# print(f"] generating variant for {newopt.init_img}")
# for j in range(0,opt.variants):
# try:
# file_writer = PngWriter(outdir,prompt,newopt.batch_size)
# callback = file_writer.write_image
# t2i.prompt2image(image_callback=callback,**vars(newopt))
# results = file_writer.files_written
# variants.append([prompt,results])
# except AssertionError as e:
# print(e)
# continue
# print(f'{opt.variants} variants generated')
# return variants
def write_log_message(t2i,prompt,results,logfile):
@ -209,9 +215,6 @@ def write_log_message(t2i,prompt,results,logfile):
img_num = 1
seenit = {}
seeds = [a[1] for a in results]
seeds = f"(seeds for individual images: {seeds})"
for r in results:
seed = r[1]
log_message = (f'{r[0]}: {prompt} -S{seed}')
@ -275,7 +278,8 @@ def create_cmd_parser():
parser.add_argument('-i','--individual',action='store_true',help="generate individual files (default)")
parser.add_argument('-I','--init_img',type=str,help="path to input image for img2img mode (supersedes width and height)")
parser.add_argument('-f','--strength',default=0.75,type=float,help="strength for noising/unnoising. 0.0 preserves image exactly, 1.0 replaces it completely")
parser.add_argument('-v','--variants',type=int,help="in img2img mode, the first generated image will get passed back to img2img to generate the requested number of variants")
# variants is going to be superseded by a generalized "prompt-morph" function
# parser.add_argument('-v','--variants',type=int,help="in img2img mode, the first generated image will get passed back to img2img to generate the requested number of variants")
parser.add_argument('-x','--skip_normalize',action='store_true',help="skip subprompt weight normalization")
return parser

@ -1 +1 @@
Subproject commit db5799068749bf3a6d5845120ed32df16b7d883b
Subproject commit ef1bf07627c9a10ba9137e68a0206b844544a7d9