OneButtonPrompt/call_extras.py

80 lines
2.9 KiB
Python

import os
import requests
import io
import base64
import uuid
from PIL import Image, PngImagePlugin
def call_extras(imagelocation,originalimage, originalpnginfo ="", apiurl="http://127.0.0.1:7860",filename="",extrasupscaler1 = "all",extrasupscaler2 ="all",extrasupscaler2visiblity="0.5",extrasupscaler2gfpgan="0",extrasupscaler2codeformer="0.15",extrasupscaler2codeformerweight="0.1",extrasresize="2"):
#rest of prompt things
upscaling_resize = extrasresize
upscaler_1 = extrasupscaler1
upscaler_2 = extrasupscaler2
with open(imagelocation, "rb") as image_file:
encoded_string = base64.b64encode(image_file.read())
encodedstring2 = encoded_string.decode('utf-8')
#params to stay the same
url = apiurl
script_dir = os.path.dirname(os.path.abspath(__file__)) # Script directory
outputextrasfolder = os.path.join(script_dir, "./automated_outputs/extras/" )
if(filename==""):
filename = str(uuid.uuid4())
outputextrasilename = filename
outputextraspng = '.png'
outputextrasFull = '{}{}{}'.format(outputextrasfolder,outputextrasilename,outputextraspng)
payload = {
"upscaling_resize": float(upscaling_resize),
"upscaler_1": upscaler_1,
"image": encodedstring2,
"resize_mode": 0,
"show_extras_results": "false",
"gfpgan_visibility": extrasupscaler2gfpgan ,
"codeformer_visibility": extrasupscaler2visiblity,
"codeformer_weight": extrasupscaler2codeformerweight,
"upscaling_crop": "false",
"upscaler_2": upscaler_2,
"extras_upscaler_2_visibility": extrasupscaler2visiblity,
"upscale_first": "true",
"rb_enabled": "false", # the remove backgrounds plugin is automatically turned on, need to turn it off
"models": "None" # the remove backgrounds plugin is automatically turned on, need to turn it off
}
response = requests.post(url=f'{url}/sdapi/v1/extra-single-image', json=payload)
image = Image.open(io.BytesIO(base64.b64decode(response.json().get("image"))))
# when using just upscale, we somehow can't get the png info. Unless we do IMG2IMG or TXT2IMG first, then it is added.
# minor issue, so this solves it for now
#if(originalpnginfo==""):
# png_payload = {
# "image": "data:image/png;base64," + image[0]
# }
#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(outputextrasFull, pnginfo=originalpnginfo)
return outputextrasFull