Spaces:
Runtime error
Runtime error
#!/bin/env python | |
import io | |
import json | |
import base64 | |
import argparse | |
import requests | |
from PIL import Image | |
options = { | |
"negative_prompt": "", | |
"steps": 20, | |
"batch_size": 1, | |
"n_iter": 1, | |
"seed": -1, | |
"sampler_name": "UniPC", | |
"cfg_scale": 6, | |
"width": 512, | |
"height": 512, | |
"save_images": False, | |
"send_images": True, | |
} | |
styles = [] | |
def pil_to_b64(img: Image, size: int, quality: int): | |
img = img.convert('RGB') | |
img = img.resize((size, size)) | |
buffer = io.BytesIO() | |
img.save(buffer, format="JPEG", quality=quality) | |
b64encoded = base64.b64encode(buffer.getvalue()).decode("utf-8") | |
return f'data:image/jpeg;base64,{b64encoded}' | |
def post(endpoint: str, dct: dict = None): | |
req = requests.post(endpoint, json = dct, timeout=300, verify=False) | |
if req.status_code != 200: | |
return { 'error': req.status_code, 'reason': req.reason, 'url': req.url } | |
else: | |
return req.json() | |
if __name__ == '__main__': | |
parser = argparse.ArgumentParser(description = 'gen-styles.py') | |
parser.add_argument('--input', type=str, required=True, help="input text file with one line per prompt") | |
parser.add_argument('--output', type=str, required=True, help="output json file") | |
parser.add_argument('--nopreviews', default=False, action='store_true', help = 'generate previews') | |
parser.add_argument('--prompt', type=str, required=False, default='girl walking in a city', help="applied prompt when generating previews") | |
parser.add_argument('--size', type=int, default=128, help="image size for previews") | |
parser.add_argument('--quality', type=int, default=35, help="image quality for previews") | |
parser.add_argument('--url', type=str, required=False, default='http://127.0.0.1:7860', help="sd.next server url") | |
args = parser.parse_args() | |
with open(args.input, encoding='utf-8') as f: | |
lines = f.readlines() | |
for line in lines: | |
line = line.strip().replace('\n', '') | |
if len(line) == 0: | |
continue | |
print(f'processing: {line}') | |
if not args.nopreviews: | |
options['prompt'] = f'{line} {args.prompt}' | |
data = post(f'{args.url}/sdapi/v1/txt2img', options) | |
if 'error' in data: | |
print(f'error: {data}') | |
continue | |
b64str = data['images'][0].split(',',1)[0] | |
image = Image.open(io.BytesIO(base64.b64decode(b64str))) | |
else: | |
image = None | |
styles.append({ | |
'name': line, | |
'prompt': line + ' {prompt}', | |
'negative': '', | |
'extra': '', | |
'preview': pil_to_b64(image, args.size, args.quality) if image is not None else '', | |
}) | |
with open(args.output, 'w', encoding='utf-8') as outfile: | |
json.dump(styles, outfile, indent=2) | |