OneButtonPrompt/call_img2img.py

206 lines
7.9 KiB
Python

import os
import requests
import io
import base64
import uuid
from PIL import Image, PngImagePlugin
from modules import shared
from model_lists import *
def call_img2img(imagelocation,originalimage, originalpnginfo ="", apiurl="http://127.0.0.1:7860",filename="", prompt = "", negativeprompt = "", img2imgsamplingsteps = "20", img2imgcfg = "7", img2imgsamplingmethod = "DPM++ SDE Karras", img2imgupscaler = "R-ESRGAN 4x+", img2imgmodel = "currently selected model", denoising_strength = "0.3", scale = "2", padding = "64",upscalescript="SD upscale",usdutilewidth = "512", usdutileheight = "0", usdumaskblur = "8", usduredraw ="Linear", usduSeamsfix = "None", usdusdenoise = "0.35", usduswidth = "64", usduspadding ="32", usdusmaskblur = "8",controlnetenabled=False, controlnetmodel="",controlnetblockymode=False):
negativepromptfound = 0
#params to stay the same
url = apiurl
script_dir = os.path.dirname(os.path.abspath(__file__)) # Script directory
outputimg2imgfolder = os.path.join(script_dir, "./automated_outputs/img2img/" )
if(filename==""):
filename = str(uuid.uuid4())
outputimg2imgpng = '.png'
outputimg2imgFull = '{}{}{}'.format(outputimg2imgfolder,filename,outputimg2imgpng)
encodedstringlist = []
# need to convert the values to the correct index number for Ultimate SD Upscaler
redrawmodelist =["Linear","Chess","None"]
seamsfixmodelist = ["None","Band pass","Half tile offset pass","Half tile offset pass + intersections"]
usduredrawint = redrawmodelist.index(usduredraw)
seamsfixmodeint = seamsfixmodelist.index(usduSeamsfix)
#rest of prompt things
sampler_index = img2imgsamplingmethod
steps = img2imgsamplingsteps
cfg_scale = img2imgcfg
with open(imagelocation, "rb") as image_file:
encoded_string = base64.b64encode(image_file.read())
encodedstringlist.append(encoded_string.decode('utf-8'))
# If we don't have a prompt, get it from the original image file
# This is used when only_upscale is activated
if(prompt==""):
with open(originalimage, "rb") as originalimage_file:
originalencoded_string = base64.b64encode(originalimage_file.read())
encodedstring2 = originalencoded_string.decode('utf-8')
# get prompt from picture
png_payload = {
"image": encodedstring2
}
response3 = requests.post(url=f'{url}/sdapi/v1/png-info', json=png_payload)
pnginfo = str(response3.json().get("info"))
prompt = pnginfo[:pnginfo.rfind("Steps")]
if(prompt.rfind("Negative prompt") != -1):
prompt = prompt[:prompt.rfind("Negative prompt")]
negativepromptfound = 1
if(negativepromptfound == 1):
negativeprompt = pnginfo[:pnginfo.rfind("Steps")]
negativeprompt = negativeprompt.replace(prompt,"")
# set the automatic upscale
checkprompt = prompt.lower()
if(img2imgupscaler != "automatic"):
upscaler = img2imgupscaler
else:
upscalerlist = get_upscalers_for_img2img()
# on automatic, make some choices about what upscaler to use
# photos, prefer 4x ultrasharp
# anime, cartoon or drawing, go for R-ESRGAN 4x+ Anime6B
# else, R-ESRGAN 4x+"
if("hoto" in checkprompt and "4x-UltraSharp" in upscalerlist):
upscaler = "4x-UltraSharp"
elif("anime" in checkprompt or "cartoon" in checkprompt or "draw" in checkprompt or "vector" in checkprompt or "cel shad" in checkprompt or "visual novel" in checkprompt):
upscaler = "R-ESRGAN 4x+ Anime6B"
else:
upscaler = "R-ESRGAN 4x+"
if(upscaler== "4x-UltraSharp"):
denoising_strength = "0.35"
if(upscaler== "R-ESRGAN 4x+ Anime6B+"):
denoising_strength = "0.6" # 0.6 is fine for the anime upscaler
if(upscaler== "R-ESRGAN 4x+"):
denoising_strength = "0.5" # default 0.6 is a lot and changes a lot of details
#wierd blocky mode comes up when the treshold is set way too high and the denoising strenght is strong
if(controlnetblockymode==True):
treshold = int(padding)
if(float(denoising_strength) < 0.65):
denoising_strength = "0.65"
else:
treshold = 1
payload = {
"resize_mode": 0,
"denoising_strength": denoising_strength,
"sampler_index": sampler_index,
"batch_size": "1",
"n_iter": "1",
"prompt": prompt,
"negative_prompt": negativeprompt,
"steps": steps,
"cfg_scale": cfg_scale,
#"width": width,
#"height": height,
"include_init_images": "true",
"init_images": encodedstringlist,
}
if(img2imgmodel != "currently selected model"):
payload.update({"sd_model": img2imgmodel})
#
# https://github.com/Mikubill/sd-webui-controlnet/wiki/API
#
if(controlnetenabled==True and controlnetmodel!=""):
payload.update({"alwayson_scripts": {
"controlnet": {
"args": [
{
"module": "tile_resample",
"model": controlnetmodel, # control_v11f1e_sd15_tile [a371b31b]
#"input_image": encodedstringlist,
"control_mode": 2, #"ControlNet is more important" : the controlnet model has more impact than the prompt
#"resize_mode": 0
"threshold_a": treshold
}
]
}
}
})
if(upscalescript=="SD upscale"):
payload.update({"script_name": upscalescript})
payload.update({"script_args": ["",int(padding),upscaler,float(scale)]})
if(upscalescript=="Ultimate SD upscale"):
upscaler_index = [x.name.lower() for x in shared.sd_upscalers].index(upscaler.lower())
payload.update({"script_name": upscalescript})
payload.update({"script_args": ["",int(usdutilewidth),int(usdutileheight),int(usdumaskblur),int(padding), int(usduswidth), usdusdenoise,int(usduspadding),
upscaler_index,True,usduredrawint,False,int(usdusmaskblur),
seamsfixmodeint,2,"","",float(scale)]})
# Ultimate SD Upscale params:
#_, tile_width, tile_height, mask_blur, padding, seams_fix_width, seams_fix_denoise, seams_fix_padding,
# upscaler_index, save_upscaled_image, redraw_mode, save_seams_fix_image, seams_fix_mask_blur,
# seams_fix_type, target_size_type, custom_width, custom_height, custom_scale):
# target_size_type = 2
# custom_scale = 2
response = requests.post(url=f'{url}/sdapi/v1/img2img', json=payload)
r = response.json()
for i in r['images']:
image = Image.open(io.BytesIO(base64.b64decode(i.split(",",1)[0])))
if(originalpnginfo==""):
png_payload = {
"image": "data:image/png;base64," + i
}
#print("and here!")
#print(png_payload)
response2 = requests.post(url=f'{url}/sdapi/v1/png-info', json=png_payload)
#print("here!")
#print(response2)
pnginfo = PngImagePlugin.PngInfo()
pnginfo.add_text("parameters", response2.json().get("info"))
originalpnginfo = pnginfo
image.save(outputimg2imgFull, pnginfo=originalpnginfo)
return [outputimg2imgFull,originalpnginfo]