From fb07233d32b98e5a56bb326c5239ff3149753e44 Mon Sep 17 00:00:00 2001 From: AIrjen Date: Fri, 21 Jul 2023 12:04:30 +0200 Subject: [PATCH] Adding new options to quality gate --- call_txt2img.py | 110 ++++++++++++++++++++++++++++++++++--- main.py | 13 +++-- scripts/onebuttonprompt.py | 30 ++++++++-- 3 files changed, 133 insertions(+), 20 deletions(-) diff --git a/call_txt2img.py b/call_txt2img.py index 0dbb765..12cb196 100644 --- a/call_txt2img.py +++ b/call_txt2img.py @@ -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] \ No newline at end of file + return [outputTXT2IMGFull,pnginfo,continuewithnextpart] \ No newline at end of file diff --git a/main.py b/main.py index aae87c4..1b9aaca 100644 --- a/main.py +++ b/main.py @@ -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("") diff --git a/scripts/onebuttonprompt.py b/scripts/onebuttonprompt.py index 53e669d..ab57f16 100644 --- a/scripts/onebuttonprompt.py +++ b/scripts/onebuttonprompt.py @@ -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): 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. @@ -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