OneButtonPrompt/main.py

226 lines
11 KiB
Python

import sys, os
import random
import uuid
import re
from datetime import datetime
sys.path.append(os.path.abspath(".."))
from call_txt2img import *
from call_img2img import *
from build_dynamic_prompt import *
from call_extras import *
from model_lists import *
def generateimages(amount = 1, size = "all",model = "currently selected model",samplingsteps = "40",cfg= "7",hiresfix = True,hiressteps ="0",denoisestrength="0.6",samplingmethod="DPM++ SDE Karras", upscaler="R-ESRGAN 4x+", hiresscale="2",apiurl="http://127.0.0.1:7860",qualitygate=False,quality="7.6",runs="5",insanitylevel="5",subject="all", artist="all", imagetype="all",silentmode=False, workprompt="", antistring="",prefixprompt="", suffixprompt="", negativeprompt="",promptcompounderlevel = "1", seperator="comma", img2imgbatch = "1", img2imgsamplingsteps = "20", img2imgcfg = "7", img2imgsamplingmethod = "DPM++ SDE Karras", img2imgupscaler = "R-ESRGAN 4x+", img2imgmodel = "currently selected model", img2imgactivate = False, img2imgscale = "2", img2imgpadding = "64",img2imgdenoisestrength="0.3",ultimatesdupscale=False,usdutilewidth = "512", usdutileheight = "0", usdumaskblur = "8", usduredraw ="Linear", usduSeamsfix = "None", usdusdenoise = "0.35", usduswidth = "64", usduspadding ="32", usdusmaskblur = "8",controlnetenabled=False, controlnetmodel="",img2imgdenoisestrengthmod="-0.05",enableextraupscale = False,controlnetblockymode = False,extrasupscaler1 = "all",extrasupscaler2 ="all",extrasupscaler2visiblity="0.5",extrasupscaler2gfpgan="0",extrasupscaler2codeformer="0.15",extrasupscaler2codeformerweight="0.1",extrasresize="2",onlyupscale="false",givensubject="",smartsubject=True,giventypeofimage="",imagemodechance=20):
loops = int(amount) # amount of images to generate
steps = 0
upscalefilelist=[]
originalimage = ""
originalpnginfo =""
randomprompt = ""
filename=""
originalsize=size
originalmodel = model
originalimg2imgmodel = img2imgmodel
insanitylevel = int(insanitylevel)
if(onlyupscale==True):
script_dir = os.path.dirname(os.path.abspath(__file__)) # Script directory
inputupscalemefolder = os.path.join(script_dir, "./automated_outputs/upscale_me/" )
for upscalefilename in os.listdir(inputupscalemefolder):
f = os.path.join(inputupscalemefolder, upscalefilename)
# checking if it is a file
if os.path.isfile(f):
if(f[-3:]!="txt"):
upscalefilelist.append(f)
loops = len(upscalefilelist)
if(loops==0):
print('No files to upscale found! Please place images in //upscale_me// folder')
else:
print("")
print("Found and upscaling files")
print("")
modellist=get_models()
samplerlist=get_samplers()
upscalerlist=get_upscalers()
img2imgupscalerlist=get_upscalers_for_img2img()
img2imgsamplerlist=get_samplers_for_img2img()
if(ultimatesdupscale==False):
upscalescript="SD upscale"
else:
upscalescript="Ultimate SD upscale"
while steps < loops:
# build prompt
if(silentmode==True and workprompt == ""):
print("Trying to use provided workflow prompt, but is empty. Generating a random prompt instead.")
if(onlyupscale==False): # only do txt2img when onlyupscale is False
if(silentmode==True and workprompt != ""):
randomprompt = workprompt
print("Using provided workflow prompt")
print(workprompt)
else:
randomprompt = build_dynamic_prompt(insanitylevel,subject,artist,imagetype, False,antistring,prefixprompt,suffixprompt,promptcompounderlevel, seperator,givensubject,smartsubject,giventypeofimage,imagemodechance)
# make the filename, from from a to the first comma
start_index = randomprompt.find("of a ") + len("of a ")
# find the index of the first comma after "of a" or end of the prompt
end_index = randomprompt.find(",", start_index)
if(end_index == -1):
end_index=len(randomprompt)
# extract the desired substring using slicing
filename = randomprompt[start_index:end_index]
# cleanup some unsafe things in the filename
filename = filename.replace("\"", "")
filename = filename.replace("[", "")
filename = filename.replace("|", "")
filename = filename.replace("]", "")
filename = filename.replace("<", "")
filename = filename.replace(">", "")
filename = filename.replace(":", "_")
filename = re.sub(r'[0-9]+', '', filename)
if(filename==""):
filename = str(uuid.uuid4())
# create a datetime object for the current date and time
now = datetime.now()
filenamecomplete = now.strftime("%Y%m%d%H%M%S") + "_" + filename.replace(" ", "_").strip()
# prompt + size
if(originalsize == "all"):
sizelist = ["portrait", "wide", "square"]
size = random.choice(sizelist)
#Check if there is any random value we have to choose or not
if(originalmodel=="all"):
model = random.choice(modellist)
#lets not do inpainting models
while "inpaint" in model:
model = random.choice(modellist)
print("Going to run with model " + model)
# set the model here
if(originalmodel!="currently selected model"):
option_payload = {
"sd_model_checkpoint": model
}
response = requests.post(url=f'{apiurl}/sdapi/v1/options', json=option_payload)
if(samplingmethod=="all"):
samplingmethod = random.choice(samplerlist)
print ("Going to run with sampling method " + samplingmethod)
if(upscaler=="all" and hiresfix == True):
upscaler = random.choice(upscalerlist)
print ("Going to run with upscaler " + upscaler)
# WebUI fix for PLMS and UniPC and hiresfix
if(samplingmethod in ['PLMS', 'UniPC']): # PLMS/UniPC do not support hirefix so we just silently switch to DDIM
samplingmethod = 'DDIM'
txt2img = call_txt2img(randomprompt, size ,hiresfix, 0, filenamecomplete,model ,samplingsteps,cfg, hiressteps, denoisestrength,samplingmethod, upscaler,hiresscale,apiurl,qualitygate,quality,runs,negativeprompt)
originalimage = txt2img[0] #Set this for later use
originalpnginfo = txt2img[1] #Sort of hacky way of bringing this forward. But if it works, it works
image = txt2img[0]
else:
if(filename==""):
filename = str(uuid.uuid4())
# create a datetime object for the current date and time
now = datetime.now()
filenamecomplete = now.strftime("%Y%m%d%H%M%S") + "_" + filename.replace(" ", "_").strip()
image = upscalefilelist[steps] # else we get the image from the upscale file list
originalimage = image # this is also the original image file
# upscale via img2img
img2imgloops = int(img2imgbatch)
if(img2imgactivate == False): # If we dont want to run, turn it off
img2imgloops = 0
img2imgsteps = 0
# start the batching!
while img2imgsteps < img2imgloops:
#Check if there is any random value we have to choose or not
if(originalimg2imgmodel=="all"):
img2imgmodel = random.choice(modellist)
#lets not do inpainting models
while "inpaint" in model:
img2imgmodel = random.choice(modellist)
print("Going to upscale with model " + img2imgmodel)
# set the model here
if(originalimg2imgmodel!="currently selected model"):
option_payload = {
"sd_model_checkpoint": img2imgmodel
}
response = requests.post(url=f'{apiurl}/sdapi/v1/options', json=option_payload)
if(img2imgsamplingmethod=="all"):
img2imgsamplingmethod = random.choice(img2imgsamplerlist)
print ("Going to upscale with sampling method " + img2imgsamplingmethod)
if(img2imgupscaler=="all"):
img2imgupscaler = random.choice(img2imgupscalerlist)
print ("Going to run with upscaler " + img2imgupscaler)
# WebUI fix for PLMS and UniPC and img2img
if(img2imgsamplingmethod in ['PLMS', 'UniPC']): # PLMS/UniPC do not support img2img so we just silently switch to DDIM
img2imgsamplingmethod = 'DDIM'
img2img = call_img2img(image, originalimage, originalpnginfo, apiurl, filenamecomplete, randomprompt,negativeprompt,img2imgsamplingsteps, img2imgcfg, img2imgsamplingmethod, img2imgupscaler, img2imgmodel, img2imgdenoisestrength, img2imgscale, img2imgpadding,upscalescript,usdutilewidth, usdutileheight, usdumaskblur, usduredraw, usduSeamsfix, usdusdenoise, usduswidth, usduspadding, usdusmaskblur,controlnetenabled, controlnetmodel,controlnetblockymode)
image = img2img[0]
if(originalpnginfo==""):
originalpnginfo = img2img[1]
img2imgdenoisestrength = str(float(img2imgdenoisestrength) + float(img2imgdenoisestrengthmod)) # lower or increase the denoise strength for each batch
img2imgpadding = int(int(img2imgpadding) * float(img2imgscale)) # also increase padding by scale
if(int(img2imgpadding)>256): # but not overdo it :D
img2imgpadding="256"
img2imgsteps += 1
# upscale via extras upscaler next
if(enableextraupscale==True):
if(extrasupscaler1=="all"):
extrasupscaler1 = random.choice(img2imgupscalerlist)
print ("Going to upscale with upscaler 1 " + extrasupscaler1)
if(extrasupscaler2=="all"):
extrasupscaler2 = random.choice(img2imgupscalerlist)
print ("Going to upscale with upscaler 2 " + extrasupscaler2)
image = call_extras(image, originalimage, originalpnginfo, apiurl, filenamecomplete,extrasupscaler1,extrasupscaler2 ,extrasupscaler2visiblity,extrasupscaler2gfpgan,extrasupscaler2codeformer,extrasupscaler2codeformerweight,extrasresize)
steps += 1
print("")
print("All done!")