diff --git a/stable-diffusion.py b/stable-diffusion.py new file mode 100644 index 0000000..a8d02eb --- /dev/null +++ b/stable-diffusion.py @@ -0,0 +1,66 @@ +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 = { + "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 + ) + )