Adding new options to quality gate
This commit is contained in:
AIrjen
+101
-9
@@ -7,13 +7,17 @@ import sys, os
|
||||
from PIL import Image, PngImagePlugin
|
||||
from model_lists import *
|
||||
import time
|
||||
import random
|
||||
|
||||
def call_txt2img(passingprompt,ratio,upscale,debugmode,filename="",model = "currently selected model",samplingsteps = "40",cfg= "7",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",negativeprompt="",qualityhiresfix = False, qualitymode = "highest", qualitykeep="keep used"):
|
||||
|
||||
|
||||
def call_txt2img(passingprompt,ratio,upscale,debugmode,filename="",model = "currently selected model",samplingsteps = "40",cfg= "7",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",negativeprompt=""):
|
||||
|
||||
#set the prompt!
|
||||
prompt = passingprompt
|
||||
checkprompt = passingprompt.lower()
|
||||
|
||||
|
||||
#set the URL for the API
|
||||
url = apiurl
|
||||
|
||||
@@ -39,7 +43,7 @@ def call_txt2img(passingprompt,ratio,upscale,debugmode,filename="",model = "curr
|
||||
height = "512"
|
||||
#upscaler
|
||||
enable_hr = upscale
|
||||
if(debugmode==1):
|
||||
if(debugmode==1 or qualityhiresfix == True):
|
||||
enable_hr="False"
|
||||
|
||||
#defaults
|
||||
@@ -102,6 +106,15 @@ def call_txt2img(passingprompt,ratio,upscale,debugmode,filename="",model = "curr
|
||||
scoredeclist = []
|
||||
imagelist = []
|
||||
pnginfolist = []
|
||||
seedlist = []
|
||||
usedseed = -1
|
||||
imagethatiskept = ""
|
||||
|
||||
# flow things
|
||||
continuewithnextpart = True
|
||||
|
||||
# starting seed of -1
|
||||
seed = -1
|
||||
|
||||
|
||||
#call TXT2IMG
|
||||
@@ -117,7 +130,9 @@ def call_txt2img(passingprompt,ratio,upscale,debugmode,filename="",model = "curr
|
||||
"denoising_strength": denoising_strength,
|
||||
"hr_scale": hr_scale,
|
||||
"hr_upscaler": hr_upscaler,
|
||||
"hr_second_pass_steps": hr_second_pass_steps
|
||||
"hr_second_pass_steps": hr_second_pass_steps,
|
||||
"seed": seed,
|
||||
"hr_prompt": prompt
|
||||
}
|
||||
|
||||
if(model != "currently selected model"):
|
||||
@@ -125,6 +140,7 @@ def call_txt2img(passingprompt,ratio,upscale,debugmode,filename="",model = "curr
|
||||
|
||||
if(negativeprompt != ""):
|
||||
payload.update({"negative_prompt": negativeprompt})
|
||||
payload.update({"hr_negative_prompt": negativeprompt})
|
||||
|
||||
while Runs < MaxRuns:
|
||||
|
||||
@@ -134,6 +150,9 @@ def call_txt2img(passingprompt,ratio,upscale,debugmode,filename="",model = "curr
|
||||
outputTXT2IMGFull = '{}{}{}'.format(outputTXT2IMGfolder,filenamefull,outputTXT2IMGpng)
|
||||
|
||||
r = []
|
||||
# randomize the seed ( A number between 0 and 4,294,967,295 )
|
||||
seed = random.randrange(1, 4294967295)
|
||||
payload["seed"] = seed
|
||||
|
||||
# If we don't get an image back, we want to retry a few times. Max 3 times
|
||||
for i in range(4):
|
||||
@@ -181,6 +200,7 @@ def call_txt2img(passingprompt,ratio,upscale,debugmode,filename="",model = "curr
|
||||
score = image_scorer.get_score(image)
|
||||
scoredeclist.append(score)
|
||||
score = round(score,1)
|
||||
seedlist.append(seed)
|
||||
|
||||
scorelist.append(score)
|
||||
imagelist.append(outputTXT2IMGFull)
|
||||
@@ -193,6 +213,8 @@ def call_txt2img(passingprompt,ratio,upscale,debugmode,filename="",model = "curr
|
||||
else:
|
||||
runstodo = MaxRuns - Runs - 1
|
||||
print("Not a good result. Retrying for another " + str(runstodo) + " times or until the image is good enough.")
|
||||
# randomize the seed for the next run
|
||||
|
||||
|
||||
except ImportError:
|
||||
foundgood = True # just continue :)
|
||||
@@ -211,22 +233,92 @@ def call_txt2img(passingprompt,ratio,upscale,debugmode,filename="",model = "curr
|
||||
if(len(imagelist) > 0):
|
||||
|
||||
if(foundgood == True):
|
||||
print("Removing any other images generated this run (if any).")
|
||||
if(qualitykeep == "keep used"):
|
||||
print("Removing any other images generated this run (if any).")
|
||||
else:
|
||||
print("Stopped trying, keeping the best image we had so far.")
|
||||
if(qualitymode == "highest"):
|
||||
print("")
|
||||
print("Stopped trying, keeping the best image we had so far.")
|
||||
print("")
|
||||
else:
|
||||
print("")
|
||||
print("Eh, its all pretty bad. Not going forward with any image.")
|
||||
print("")
|
||||
|
||||
|
||||
# Get the index of the first occurrence of the maximum value in the list
|
||||
#if(qualitymode == "highest" or (qualitymode != "highest" and foundgood == True)):
|
||||
indexofimagetokeep = scoredeclist.index(max(scoredeclist))
|
||||
outputTXT2IMGFull = imagelist[indexofimagetokeep] #store the image to keep in here, so we can pass it along
|
||||
pnginfo = pnginfolist[indexofimagetokeep]
|
||||
imagelist.pop(indexofimagetokeep)
|
||||
usedseed = seedlist[indexofimagetokeep]
|
||||
imagethatiskept = imagelist[indexofimagetokeep]
|
||||
imagelist.pop(indexofimagetokeep)
|
||||
|
||||
|
||||
#remove all other images
|
||||
for imagelocation in imagelist:
|
||||
os.remove(imagelocation)
|
||||
if(qualitykeep == "keep used"):
|
||||
for imagelocation in imagelist:
|
||||
os.remove(imagelocation)
|
||||
|
||||
if(foundgood == False and qualitymode != "highest"):
|
||||
continuewithnextpart = False
|
||||
if(imagethatiskept != "" and qualitykeep == "keep used"):
|
||||
os.remove(imagethatiskept)
|
||||
|
||||
|
||||
# We have done everything, but if we want to run Hires fix from the quality gate, we are going to have to do it again. But this time a little easier.
|
||||
# We do have the check wether we want to run hiresfix first
|
||||
if(qualityhiresfix == True and upscale == False and continuewithnextpart == True):
|
||||
print("Quality Gate hires fix was enabled, but no hires fix settings were given.")
|
||||
if(qualityhiresfix == True and upscale == True and continuewithnextpart == True):
|
||||
print("Going to run the chosen image with hiresfix")
|
||||
|
||||
payload["seed"] = usedseed
|
||||
payload["enable_hr"] = "True"
|
||||
|
||||
# make the filename unique for hiresfix
|
||||
addrun = "_hiresfix"
|
||||
filenamefull = filename + addrun
|
||||
outputTXT2IMGFull = '{}{}{}'.format(outputTXT2IMGfolder,filenamefull,outputTXT2IMGpng)
|
||||
|
||||
# If we don't get an image back, we want to retry a few times. Max 3 times
|
||||
for i in range(4):
|
||||
response = requests.post(url=f'{url}/sdapi/v1/txt2img', json=payload)
|
||||
|
||||
r = response.json()
|
||||
if('images' in r):
|
||||
break # this means if we have the images object, then we "break" out of the for loop.
|
||||
else:
|
||||
if(i == 3):
|
||||
print("If this keeps happening: Is WebUI started with --api enabled?")
|
||||
print("")
|
||||
raise ValueError("API has not been responding after several retries. Stopped processing.")
|
||||
print("")
|
||||
print("We haven't received an image from the API. Maybe something went wrong. Will retry after waiting a bit.")
|
||||
|
||||
|
||||
time.sleep(10 * (i+1) ) # incremental waiting time
|
||||
|
||||
|
||||
|
||||
|
||||
for i in r['images']:
|
||||
image = Image.open(io.BytesIO(base64.b64decode(i.split(",",1)[0])))
|
||||
|
||||
png_payload = {
|
||||
"image": "data:image/png;base64," + i
|
||||
}
|
||||
response2 = requests.post(url=f'{url}/sdapi/v1/png-info', json=png_payload)
|
||||
|
||||
pnginfo = PngImagePlugin.PngInfo()
|
||||
pnginfo.add_text("parameters", response2.json().get("info"))
|
||||
image.save(outputTXT2IMGFull, pnginfo=pnginfo)
|
||||
|
||||
|
||||
|
||||
with open(outputTXT2IMGtxtFull,'w',encoding="utf8") as txt:
|
||||
json_object = json.dumps(payload, indent = 4)
|
||||
txt.write(json_object)
|
||||
|
||||
return [outputTXT2IMGFull,pnginfo]
|
||||
return [outputTXT2IMGFull,pnginfo,continuewithnextpart]
|
||||
@@ -13,7 +13,7 @@ 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, gender="all", chosensubjectsubtypeobject="all", chosensubjectsubtypehumanoid="all", chosensubjectsubtypeconcept="all", increasestability = False):
|
||||
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, gender="all", chosensubjectsubtypeobject="all", chosensubjectsubtypehumanoid="all", chosensubjectsubtypeconcept="all", increasestability = False, qualityhiresfix = False, qualitymode = "highest", qualitykeep="keep used"):
|
||||
loops = int(amount) # amount of images to generate
|
||||
steps = 0
|
||||
upscalefilelist=[]
|
||||
@@ -187,9 +187,10 @@ def generateimages(amount = 1, size = "all",model = "currently selected model",s
|
||||
|
||||
|
||||
|
||||
txt2img = call_txt2img(randomprompt, size ,hiresfix, 0, filenamecomplete,model ,samplingsteps,cfg, hiressteps, denoisestrength,samplingmethod, upscaler,hiresscale,apiurl,qualitygate,quality,runs,negativeprompt)
|
||||
txt2img = call_txt2img(randomprompt, size ,hiresfix, 0, filenamecomplete,model ,samplingsteps,cfg, hiressteps, denoisestrength,samplingmethod, upscaler,hiresscale,apiurl,qualitygate,quality,runs,negativeprompt, qualityhiresfix, qualitymode, qualitykeep)
|
||||
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
|
||||
continuewithnextpart = txt2img[2]
|
||||
|
||||
image = txt2img[0]
|
||||
else:
|
||||
@@ -206,7 +207,7 @@ def generateimages(amount = 1, size = "all",model = "currently selected model",s
|
||||
# upscale via img2img
|
||||
|
||||
img2imgloops = int(img2imgbatch)
|
||||
if(img2imgactivate == False): # If we dont want to run, turn it off
|
||||
if(img2imgactivate == False or continuewithnextpart == False): # If we dont want to run, turn it off
|
||||
img2imgloops = 0
|
||||
else:
|
||||
#Check if there is any random value we have to choose or not
|
||||
@@ -270,7 +271,7 @@ def generateimages(amount = 1, size = "all",model = "currently selected model",s
|
||||
|
||||
# upscale via extras upscaler next
|
||||
|
||||
if(enableextraupscale==True):
|
||||
if(enableextraupscale==True and continuewithnextpart == True):
|
||||
if(extrasupscaler1=="all"):
|
||||
extrasupscaler1 = random.choice(img2imgupscalerlist)
|
||||
print ("Going to upscale with upscaler 1 " + extrasupscaler1)
|
||||
@@ -281,7 +282,9 @@ def generateimages(amount = 1, size = "all",model = "currently selected model",s
|
||||
|
||||
image = call_extras(image, originalimage, originalpnginfo, apiurl, filenamecomplete,extrasupscaler1,extrasupscaler2 ,extrasupscaler2visiblity,extrasupscaler2gfpgan,extrasupscaler2codeformer,extrasupscaler2codeformerweight,extrasresize)
|
||||
|
||||
steps += 1
|
||||
if(continuewithnextpart == True):
|
||||
# only count images we actually fully processed
|
||||
steps += 1
|
||||
|
||||
|
||||
print("")
|
||||
|
||||
@@ -29,6 +29,9 @@ ANDtogglemode = ["none", "automatic", "prefix AND prompt + suffix", "prefix + pr
|
||||
seperatorlist = ["comma", "AND", "BREAK"]
|
||||
genders = ["all", "male", "female"]
|
||||
|
||||
qualitymodelist = ["highest", "gated"]
|
||||
qualitykeeplist = ["keep used","keep all"]
|
||||
|
||||
#for autorun and upscale
|
||||
sizelist = ["all", "portrait", "wide", "square", "ultrawide"]
|
||||
|
||||
@@ -672,7 +675,9 @@ class Script(scripts.Script):
|
||||
<font size="2">
|
||||
Uses aesthetic image scorer extension to check the quality of the image.
|
||||
|
||||
Once turned on, it will retry for n amount of times to get an image with the quality score. If not, it will take the best image so far and continue.
|
||||
Once turned on, it will retry for n amount of times to get an image with the quality score. If not, it will take the best image so far and continue or set it to gated to only take matching or higher quality images.
|
||||
|
||||
You can move Hiresfix to be applied afterwards. You can opt to keep all generated images.
|
||||
|
||||
Idea and inspiration by xKean.
|
||||
</font>
|
||||
@@ -682,6 +687,12 @@ class Script(scripts.Script):
|
||||
qualitygate = gr.Checkbox(label="Quality Gate", value=False)
|
||||
quality = gr.Slider(1, 10, value = "7.2", step=0.1, label="Quality", visible = False)
|
||||
runs = gr.Slider(1, 50, value = "5", step=1, label="Amount of tries", visible = False)
|
||||
with gr.Row():
|
||||
qualityhiresfix = gr.Checkbox(label="Move Hires fix afterwards", value=False, visible = False)
|
||||
qualitymode = gr.Dropdown(
|
||||
qualitymodelist, label= "Mode of operation", value="highest", visible = False)
|
||||
qualitykeep = gr.Dropdown(
|
||||
qualitykeeplist, label= "Images", value="keep used", visible = False)
|
||||
with gr.Row():
|
||||
gr.Markdown(
|
||||
"""
|
||||
@@ -774,7 +785,7 @@ class Script(scripts.Script):
|
||||
prompt4toworkflow.click(prompttoworkflowprompt, inputs=prompt4, outputs=workprompt)
|
||||
prompt5toworkflow.click(prompttoworkflowprompt, inputs=prompt5, outputs=workprompt)
|
||||
|
||||
startmain.click(generateimages, inputs=[amountofimages,size,model,samplingsteps,cfg,hiresfix,hiressteps,denoisestrength,samplingmethod, upscaler,hiresscale, apiurl, qualitygate, quality, runs,insanitylevel,subject, artist, imagetype, silentmode, workprompt, antistring, prefixprompt, suffixprompt,negativeprompt,promptcompounderlevel, seperator, img2imgbatch, img2imgsamplingsteps, img2imgcfg, img2imgsamplingmethod, img2imgupscaler, img2imgmodel,img2imgactivate, img2imgscale, img2imgpadding,img2imgdenoisestrength,ultimatesdupscale,usdutilewidth, usdutileheight, usdumaskblur, usduredraw, usduSeamsfix, usdusdenoise, usduswidth, usduspadding, usdusmaskblur, controlnetenabled, controlnetmodel,img2imgdenoisestrengthmod,enableextraupscale,controlnetblockymode,extrasupscaler1,extrasupscaler2,extrasupscaler2visiblity,extrasupscaler2gfpgan,extrasupscaler2codeformer,extrasupscaler2codeformerweight,extrasresize,onlyupscale,givensubject,smartsubject,giventypeofimage,imagemodechance, chosengender, chosensubjectsubtypeobject, chosensubjectsubtypehumanoid, chosensubjectsubtypeconcept, increasestability])
|
||||
startmain.click(generateimages, inputs=[amountofimages,size,model,samplingsteps,cfg,hiresfix,hiressteps,denoisestrength,samplingmethod, upscaler,hiresscale, apiurl, qualitygate, quality, runs,insanitylevel,subject, artist, imagetype, silentmode, workprompt, antistring, prefixprompt, suffixprompt,negativeprompt,promptcompounderlevel, seperator, img2imgbatch, img2imgsamplingsteps, img2imgcfg, img2imgsamplingmethod, img2imgupscaler, img2imgmodel,img2imgactivate, img2imgscale, img2imgpadding,img2imgdenoisestrength,ultimatesdupscale,usdutilewidth, usdutileheight, usdumaskblur, usduredraw, usduSeamsfix, usdusdenoise, usduswidth, usduspadding, usdusmaskblur, controlnetenabled, controlnetmodel,img2imgdenoisestrengthmod,enableextraupscale,controlnetblockymode,extrasupscaler1,extrasupscaler2,extrasupscaler2visiblity,extrasupscaler2gfpgan,extrasupscaler2codeformer,extrasupscaler2codeformerweight,extrasresize,onlyupscale,givensubject,smartsubject,giventypeofimage,imagemodechance, chosengender, chosensubjectsubtypeobject, chosensubjectsubtypehumanoid, chosensubjectsubtypeconcept, increasestability, qualityhiresfix, qualitymode, qualitykeep])
|
||||
|
||||
automatedoutputsfolderbutton.click(openfolder)
|
||||
|
||||
@@ -848,14 +859,18 @@ class Script(scripts.Script):
|
||||
|
||||
qualitygate: gr.update(visible=onlyupscale),
|
||||
quality: gr.update(visible=onlyupscale),
|
||||
runs: gr.update(visible=onlyupscale)
|
||||
runs: gr.update(visible=onlyupscale),
|
||||
qualityhiresfix: gr.update(visible=onlyupscale),
|
||||
qualitymode: gr.update(visible=onlyupscale),
|
||||
qualitykeep: gr.update(visible=onlyupscale)
|
||||
|
||||
|
||||
}
|
||||
|
||||
onlyupscale.change(
|
||||
onlyupscalevalues,
|
||||
[onlyupscale],
|
||||
[amountofimages,size,samplingsteps,cfg,hiresfix,hiressteps,hiresscale,denoisestrength,upscaler,model,samplingmethod,upscaler,qualitygate,quality,runs]
|
||||
[amountofimages,size,samplingsteps,cfg,hiresfix,hiressteps,hiresscale,denoisestrength,upscaler,model,samplingmethod,upscaler,qualitygate,quality,runs,qualityhiresfix,qualitymode,qualitykeep]
|
||||
)
|
||||
|
||||
|
||||
@@ -878,13 +893,16 @@ class Script(scripts.Script):
|
||||
def qgatechangevalues(qualitygate):
|
||||
return {
|
||||
quality: gr.update(visible=qualitygate),
|
||||
runs: gr.update(visible=qualitygate)
|
||||
runs: gr.update(visible=qualitygate),
|
||||
qualityhiresfix: gr.update(visible=qualitygate),
|
||||
qualitymode: gr.update(visible=qualitygate),
|
||||
qualitykeep: gr.update(visible=qualitygate)
|
||||
}
|
||||
|
||||
qualitygate.change(
|
||||
qgatechangevalues,
|
||||
[qualitygate],
|
||||
[quality,runs]
|
||||
[quality,runs,qualityhiresfix,qualitymode,qualitykeep]
|
||||
)
|
||||
|
||||
# Turn things off and on for USDU
|
||||
|
||||
Reference in New Issue
Block a user