1
0
Fork 0
scripts/stable-diffusion.py

66 lines
2.2 KiB
Python

import asyncio
import argparse
import aiohttp
import base64
import os
async def api_request(url: str,
username: str, password: str, params: dict) -> dict:
headers = {
"Authorization": aiohttp.BasicAuth(username, password).encode(),
"Accept-Encoding": "identity",
}
async with aiohttp.ClientSession() as session:
async with session.post(f"{url}/sdapi/v1/txt2img",
headers=headers, json=params) as resp:
return await resp.json()
async def save_images(username: str, password: str,
prompt: str, negative_prompt: str,
destination: str, params_dict: dict[str, str]) -> None:
params = {
"enable_hr": True,
"sampler name": "DPM++ 2M Karras",
"width": 896,
"height": 896,
"steps": 30,
"cfg_scale": 12,
"n_iter": 3,
"batch_size": 1,
"denoising_strength": 0.6,
"prompt": prompt,
"negative_prompt": negative_prompt,
}
response = await api_request(username, password, params)
for i, image_data in enumerate(response["images"]):
image_bytes = base64.decodebytes(image_data.encode())
with open(os.path.join(destination, f"file{i + 1}.png"), "wb") as f:
f.write(image_bytes)
if __name__ == "__main__":
parser = argparse.ArgumentParser()
parser.add_argument(
"--username", default="root",
help="Username for basic authentication (default: root)")
parser.add_argument(
"--password", default="itsreallyme",
help="Password for basic authentication (default: itsreallyme)")
parser.add_argument("url", help="URL to use for API")
parser.add_argument("prompt", help="Prompt for image generation")
parser.add_argument("negative_prompt",
help="Negative prompt for image generation")
parser.add_argument("destination", help="Directory to save images to")
args = parser.parse_args()
asyncio.run(
save_images(
args.url,
args.username,
args.password,
args.prompt,
args.negative_prompt,
args.destination
)
)