#!/usr/bin/env python3 # Copyright (c) 2022 Lincoln D. Stein (https://github.com/lstein) import argparse import shlex import os import sys import copy from ldm.dream_util import Completer,PngWriter,PromptFormatter debugging = False def main(): ''' Initialize command-line parsers and the diffusion model ''' arg_parser = create_argv_parser() opt = arg_parser.parse_args() if opt.laion400m: # defaults suitable to the older latent diffusion weights width = 256 height = 256 config = "configs/latent-diffusion/txt2img-1p4B-eval.yaml" weights = "models/ldm/text2img-large/model.ckpt" else: # some defaults suitable for stable diffusion weights width = 512 height = 512 config = "configs/stable-diffusion/v1-inference.yaml" weights = "models/ldm/stable-diffusion-v1/model.ckpt" print("* Initializing, be patient...\n") sys.path.append('.') from pytorch_lightning import logging from ldm.simplet2i import T2I # these two lines prevent a horrible warning message from appearing # when the frozen CLIP tokenizer is imported import transformers transformers.logging.set_verbosity_error() # creating a simple text2image object with a handful of # defaults passed on the command line. # additional parameters will be added (or overriden) during # the user input loop t2i = T2I(width=width, height=height, sampler_name=opt.sampler_name, weights=weights, full_precision=opt.full_precision, config=config, latent_diffusion_weights=opt.laion400m, # this is solely for recreating the prompt embedding_path=opt.embedding_path, device=opt.device ) # make sure the output directory exists if not os.path.exists(opt.outdir): os.makedirs(opt.outdir) # gets rid of annoying messages about random seed logging.getLogger("pytorch_lightning").setLevel(logging.ERROR) infile = None try: if opt.infile is not None: infile = open(opt.infile,'r') except FileNotFoundError as e: print(e) exit(-1) # preload the model t2i.load_model() # load GFPGAN if requested if opt.use_gfpgan: print("\n* --gfpgan was specified, loading gfpgan...") try: model_path = os.path.join(opt.gfpgan_dir, opt.gfpgan_model_path) if not os.path.isfile(model_path): raise Exception("GFPGAN model not found at path "+model_path) sys.path.append(os.path.abspath(opt.gfpgan_dir)) from gfpgan import GFPGANer bg_upsampler = None if opt.gfpgan_bg_upsampler is not None: bg_upsampler = load_gfpgan_bg_upsampler(opt.gfpgan_bg_upsampler, opt.gfpgan_bg_tile) t2i.gfpgan = GFPGANer(model_path=model_path, upscale=opt.gfpgan_upscale, arch='clean', channel_multiplier=2, bg_upsampler=bg_upsampler) except Exception: import traceback print("Error loading GFPGAN:", file=sys.stderr) print(traceback.format_exc(), file=sys.stderr) print("\n* Initialization done! Awaiting your command (-h for help, 'q' to quit, 'cd' to change output dir, 'pwd' to print output dir)...") log_path = os.path.join(opt.outdir,'dream_log.txt') with open(log_path,'a') as log: cmd_parser = create_cmd_parser() main_loop(t2i,opt.outdir,cmd_parser,log,infile) log.close() if infile: infile.close() def main_loop(t2i,outdir,parser,log,infile): ''' prompt/read/execute loop ''' done = False while not done: try: command = infile.readline() if infile else input("dream> ") except EOFError: done = True break if infile and len(command)==0: done = True break if command.startswith(('#','//')): continue # before splitting, escape single quotes so as not to mess # up the parser command = command.replace("'","\\'") try: elements = shlex.split(command) except ValueError as e: print(str(e)) continue if len(elements)==0: continue if elements[0]=='q': done = True break if elements[0]=='cd' and len(elements)>1: if os.path.exists(elements[1]): print(f"setting image output directory to {elements[1]}") outdir=elements[1] else: print(f"directory {elements[1]} does not exist") continue if elements[0]=='pwd': print(f"current output directory is {outdir}") continue if elements[0].startswith('!dream'): # in case a stored prompt still contains the !dream command elements.pop(0) # rearrange the arguments to mimic how it works in the Dream bot. switches = [''] switches_started = False for el in elements: if el[0]=='-' and not switches_started: switches_started = True if switches_started: switches.append(el) else: switches[0] += el switches[0] += ' ' switches[0] = switches[0][:len(switches[0])-1] try: opt = parser.parse_args(switches) except SystemExit: parser.print_help() continue if len(opt.prompt)==0: print("Try again with a prompt!") continue normalized_prompt = PromptFormatter(t2i,opt).normalize_prompt() individual_images = not opt.grid try: file_writer = PngWriter(outdir,normalized_prompt,opt.batch_size) callback = file_writer.write_image if individual_images else None image_list = t2i.prompt2image(image_callback=callback,**vars(opt)) results = file_writer.files_written if individual_images else image_list 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) print("goodbye!") def load_gfpgan_bg_upsampler(bg_upsampler, bg_tile=400): import torch if bg_upsampler == 'realesrgan': if not torch.cuda.is_available(): # CPU import warnings warnings.warn('The unoptimized RealESRGAN is slow on CPU. We do not use it. ' 'If you really want to use it, please modify the corresponding codes.') bg_upsampler = None else: from basicsr.archs.rrdbnet_arch import RRDBNet from realesrgan import RealESRGANer model = RRDBNet(num_in_ch=3, num_out_ch=3, num_feat=64, num_block=23, num_grow_ch=32, scale=2) bg_upsampler = RealESRGANer( scale=2, model_path='https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.1/RealESRGAN_x2plus.pth', model=model, tile=bg_tile, tile_pad=10, pre_pad=0, half=True) # need to set False in CPU mode else: bg_upsampler = None return bg_upsampler # 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): ''' logs the name of the output image, its prompt and seed to the terminal, log file, and a Dream text chunk in the PNG metadata''' last_seed = None img_num = 1 seenit = {} for r in results: seed = r[1] log_message = (f'{r[0]}: {prompt} -S{seed}') print(log_message) logfile.write(log_message+"\n") logfile.flush() def create_argv_parser(): parser = argparse.ArgumentParser(description="Parse script's command line args") parser.add_argument("--laion400m", "--latent_diffusion", "-l", dest='laion400m', action='store_true', help="fallback to the latent diffusion (laion400m) weights and config") parser.add_argument("--from_file", dest='infile', type=str, help="if specified, load prompts from this file") parser.add_argument('-n','--iterations', type=int, default=1, help="number of images to generate") parser.add_argument('-F','--full_precision', dest='full_precision', action='store_true', help="use slower full precision math for calculations") parser.add_argument('--sampler','-m', dest="sampler_name", choices=['ddim', 'k_dpm_2_a', 'k_dpm_2', 'k_euler_a', 'k_euler', 'k_heun', 'k_lms', 'plms'], default='k_lms', help="which sampler to use (k_lms) - can only be set on command line") parser.add_argument('--outdir', '-o', type=str, default="outputs/img-samples", help="directory in which to place generated images and a log of prompts and seeds") parser.add_argument('--embedding_path', type=str, help="Path to a pre-trained embedding manager checkpoint - can only be set on command line") parser.add_argument('--device', '-d', type=str, default="cuda", help="device to run stable diffusion on. defaults to cuda `torch.cuda.current_device()` if avalible") # GFPGAN related args parser.add_argument('--gfpgan', dest='use_gfpgan', action='store_true', help="load gfpgan for use in the dreambot. Note: Enabling GFPGAN will require more GPU memory") parser.add_argument("--gfpgan_upscale", type=int, default=2, help="The final upsampling scale of the image. Default: 2. Only used if --gfpgan is specified") parser.add_argument("--gfpgan_bg_upsampler", type=str, default='realesrgan', help="Background upsampler. Default: None. Options: realesrgan, none. Only used if --gfpgan is specified") parser.add_argument("--gfpgan_bg_tile", type=int, default=400, help="Tile size for background sampler, 0 for no tile during testing. Default: 400. Only used if --gfpgan is specified") parser.add_argument("--gfpgan_model_path", type=str, default='experiments/pretrained_models/GFPGANv1.3.pth', help="indicates the path to the GFPGAN model, relative to --gfpgan_dir. Only used if --gfpgan is specified") parser.add_argument("--gfpgan_dir", type=str, default='../gfpgan', help="indicates the directory containing the GFPGAN code. Only used if --gfpgan is specified") return parser def create_cmd_parser(): parser = argparse.ArgumentParser(description='Example: dream> a fantastic alien landscape -W1024 -H960 -s100 -n12') parser.add_argument('prompt') parser.add_argument('-s','--steps',type=int,help="number of steps") parser.add_argument('-S','--seed',type=int,help="image seed") parser.add_argument('-n','--iterations',type=int,default=1,help="number of samplings to perform (slower, but will provide seeds for individual images)") parser.add_argument('-b','--batch_size',type=int,default=1,help="number of images to produce per sampling (will not provide seeds for individual images!)") parser.add_argument('-W','--width',type=int,help="image width, multiple of 64") parser.add_argument('-H','--height',type=int,help="image height, multiple of 64") parser.add_argument('-C','--cfg_scale',default=7.5,type=float,help="prompt configuration scale") parser.add_argument('-g','--grid',action='store_true',help="generate a grid") 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('-G','--gfpgan_strength', default=0.5, type=float, help="The strength at which to apply the GFPGAN model to the result, in order to improve faces.") # 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 if __name__ == "__main__": main()