|
|
import base64
|
|
|
import datetime
|
|
|
import json
|
|
|
import os
|
|
|
|
|
|
import requests
|
|
|
|
|
|
# 文档
|
|
|
# https://github.com/AUTOMATIC1111/stable-diffusion-webui/wiki/API
|
|
|
|
|
|
# 参考
|
|
|
# https://www.bilibili.com/read/cv28519209/
|
|
|
|
|
|
# 主机地址
|
|
|
HOST = "http://10.10.21.21:7860"
|
|
|
|
|
|
|
|
|
def submit_post(url: str, data: dict):
|
|
|
return requests.post(url, data=json.dumps(data))
|
|
|
|
|
|
|
|
|
def save_encoded_image(b64_image: str, output_path: str):
|
|
|
# 判断当前目录下是否存在 output 文件夹，如果不存在则创建
|
|
|
if not os.path.exists("output"):
|
|
|
os.mkdir("output")
|
|
|
timestamp = datetime.datetime.now().strftime("%Y%m%d%H%M%S")
|
|
|
output_path = f"{output_path}_{timestamp}" + ".png"
|
|
|
# 将文件放入当前目录下的 output 文件夹中
|
|
|
output_path = f"output/{output_path}"
|
|
|
with open(output_path, "wb") as f:
|
|
|
f.write(base64.b64decode(b64_image))
|
|
|
|
|
|
|
|
|
def save_json_file(data: dict, output_path: str):
|
|
|
# 忽略 data 中的 images 字段
|
|
|
data.pop('images')
|
|
|
# 将 data 中的 info 字段转为 json 字符串，info 当前数据需要转义
|
|
|
data['info'] = json.loads(data['info'])
|
|
|
|
|
|
# 输出 data.info.infotexts
|
|
|
info_texts = data['info']['infotexts']
|
|
|
for info_text in info_texts:
|
|
|
print(info_text)
|
|
|
|
|
|
# 判断当前目录下是否存在 output 文件夹，如果不存在则创建
|
|
|
if not os.path.exists("output"):
|
|
|
os.mkdir("output")
|
|
|
timestamp = datetime.datetime.now().strftime("%Y%m%d%H%M%S")
|
|
|
output_path = f"{output_path}_{timestamp}" + ".json"
|
|
|
# 将文件放入当前目录下的 output 文件夹中
|
|
|
output_path = f"output/{output_path}"
|
|
|
with open(output_path, "w") as f:
|
|
|
json.dump(data, f, indent=4, ensure_ascii=False)
|
|
|
|
|
|
|
|
|
def getQueue():
|
|
|
# 探测队列中任务数量
|
|
|
url = HOST + "/queue/status"
|
|
|
|
|
|
# 发送GET请求
|
|
|
response = requests.get(url)
|
|
|
|
|
|
# 检查请求是否成功
|
|
|
if response.status_code == 200:
|
|
|
# 获取响应内容
|
|
|
data = response.json() # 假设返回的数据是JSON格式
|
|
|
print(data)
|
|
|
else:
|
|
|
print('请求失败，状态码：', response.status_code)
|
|
|
|
|
|
|
|
|
# 文生图
|
|
|
def txt2img():
|
|
|
txt2img_url = HOST + "/sdapi/v1/txt2img" # 服务器地址
|
|
|
prompt = input("请输入提示词：")
|
|
|
negative_prompt = input("请输入反面提示词：")
|
|
|
data = {'prompt': prompt, 'negative_prompt': negative_prompt, "cfg_scale": 7, "height": 1024, "width": 1024,
|
|
|
"sampler_name": "DPM++ 2M Karras", "samples": 1, "steps": 30}
|
|
|
|
|
|
# 将 data.prompt 中的文本，删除文件名非法字符，已下划线分隔，作为文件名
|
|
|
output_path = data['prompt'].replace(" ", "_").replace("/", "_").replace("\\", "_").replace(":", "_").replace("\"",
|
|
|
"_").replace(
|
|
|
"<", "_").replace(">", "_").replace("|", "_")
|
|
|
response = submit_post(txt2img_url, data)
|
|
|
save_encoded_image(response.json()['images'][0], output_path)
|
|
|
save_json_file(response.json(), output_path)
|
|
|
|
|
|
|
|
|
if __name__ == '__main__':
|
|
|
# 查询运行队列情况
|
|
|
getQueue()
|
|
|
# 文生图
|
|
|
txt2img()
|
