113 lines
3.1 KiB
Python
113 lines
3.1 KiB
Python
import requests
|
|
import json
|
|
from PIL import Image
|
|
import io
|
|
import base64
|
|
from typing import Optional
|
|
|
|
|
|
def send_request(api_endpoint:str, auth:Optional[str], arguments:dict) -> list[Image.Image]:
|
|
session = requests.Session()
|
|
if auth:
|
|
session.auth = (auth.split(":")[0], auth.split(":")[1])
|
|
api_endpoint = api_endpoint.rstrip("/") + "/sdapi/v1/txt2img"
|
|
response = session.post(api_endpoint, json=arguments)
|
|
response.raise_for_status()
|
|
response_json = response.json()
|
|
if "images" in response_json.keys():
|
|
images = [Image.open(io.BytesIO(base64.b64decode(i))) for i in response_json["images"]]
|
|
elif "image" in response_json.keys():
|
|
images = [Image.open(io.BytesIO(base64.b64decode(response_json["image"])))]
|
|
else:
|
|
raise ValueError("No image data in response")
|
|
return images
|
|
|
|
def construct_args(
|
|
prompt:str,
|
|
seed:int=-1,
|
|
negative_prompt:Optional[str] = None,
|
|
steps:int = 28,
|
|
width:int = 1024,
|
|
height:int = 1024,
|
|
hr_scale:float = 1.5,
|
|
hr_upscale:str = "Latent",
|
|
enable_hr:bool = False,
|
|
cfg_scale:int = 7,
|
|
):
|
|
arguments = {
|
|
"prompt": prompt,
|
|
"seed": seed,
|
|
"steps": steps,
|
|
"width": width,
|
|
"height": height,
|
|
"hr_scale": hr_scale,
|
|
"hr_upscale": hr_upscale,
|
|
"enable_hr": enable_hr,
|
|
"cfg_scale": cfg_scale,
|
|
}
|
|
if negative_prompt:
|
|
arguments["negative_prompt"] = negative_prompt
|
|
else:
|
|
arguments["negative_prompt"] = ""
|
|
return arguments
|
|
|
|
def get_image_from_prompt(
|
|
prompt:str,
|
|
api_endpoint:str,
|
|
auth:Optional[str]=None,
|
|
seed:int=-1,
|
|
negative_prompt:Optional[str] = None,
|
|
steps:int = 28,
|
|
width:int = 1024,
|
|
height:int = 1024,
|
|
hr_scale:float = 1.5,
|
|
hr_upscale:str = "Latent",
|
|
enable_hr:bool = False,
|
|
cfg_scale:int = 7,
|
|
):
|
|
arguments = construct_args(
|
|
prompt=prompt,
|
|
seed=seed,
|
|
negative_prompt=negative_prompt,
|
|
steps=steps,
|
|
width=width,
|
|
height=height,
|
|
hr_scale=hr_scale,
|
|
hr_upscale=hr_upscale,
|
|
enable_hr=enable_hr,
|
|
cfg_scale=cfg_scale,
|
|
)
|
|
return send_request(api_endpoint, auth, arguments)
|
|
|
|
def get_image_from_prompt_fallback(
|
|
prompt:str,
|
|
api_endpoint:str,
|
|
auth:Optional[str]=None,
|
|
seed:int=-1,
|
|
negative_prompt:Optional[str] = None,
|
|
steps:int = 28,
|
|
width:int = 1024,
|
|
height:int = 1024,
|
|
hr_scale:float = 1.5,
|
|
hr_upscale:str = "Latent",
|
|
enable_hr:bool = False,
|
|
cfg_scale:int = 7,
|
|
):
|
|
arguments = construct_args(
|
|
prompt=prompt,
|
|
seed=seed,
|
|
negative_prompt=negative_prompt,
|
|
steps=steps,
|
|
width=width,
|
|
height=height,
|
|
hr_scale=hr_scale,
|
|
hr_upscale=hr_upscale,
|
|
enable_hr=enable_hr,
|
|
cfg_scale=cfg_scale,
|
|
)
|
|
try:
|
|
return send_request(api_endpoint, auth, arguments)
|
|
except Exception as e:
|
|
# create blank image
|
|
return [Image.new("RGB", (width, height), (255, 255, 255))]
|