from app import BZHStableSignatureDemo
import sys
from pathlib import Path

def generate_all(captions, outdir='out'):
    backend = BZHStableSignatureDemo()
    decoders = list(backend.decoders.keys())
    outdir = Path(outdir)
    outdir.mkdir(parents=True, exist_ok=True)

    for decoder in decoders:
        # TODO: optimize by patching only once..
        # TEMP
        if decoder != 'medium':
            continue
        outdirdec = Path(outdir) / Path(decoder)
        outdirdec.mkdir(parents=True, exist_ok = True)
        for i, prompt in enumerate(captions):
            img = backend.generate(decoder, 0, prompt)
            img.save(outdirdec / Path('%04d.png' % i), format='PNG')

if __name__ == '__main__':
    infile = sys.argv[1]
    outdir = sys.argv[2]
    with open(infile) as f:
        generate_all([line.strip() for line in f.readlines()], outdir=outdir)