10 from stability_sdk import client
11 import stability_sdk.interfaces.gooseai.generation.generation_pb2 as generation
13 DOCUMENT_ROOT = "/srv/cworth.org/zombocom"
14 IMAGES_PATH = "/images"
16 def normalize_filename(fn):
20 if str.isalpha(c) or str.isdigit(c) or (c in valid):
26 def save_artifact(artifact, prompt):
31 base = "{}_{}_{}.png".format(seed, prompt, counter)
33 base = "{}_{}.png".format(seed, prompt)
34 base = normalize_filename(base)
35 filename = "{}{}/{}".format(DOCUMENT_ROOT, IMAGES_PATH, base)
36 if not os.path.exists(filename):
40 img = Image.open(io.BytesIO(artifact.binary))
42 return {"seed": seed, "prompt": prompt,
43 "filename": filename.removeprefix(DOCUMENT_ROOT)}
46 os.environ['STABILITY_HOST'] = 'grpc.stability.ai:443'
48 if not 'STABILITY_KEY' in os.environ:
49 print("Error: STABILITY_KEY must be set with a valid key")
52 parser = argparse.ArgumentParser(
53 usage="%(prog)s [--seed=SEED] Prompt string here...",
54 description="Generate text-from-image using Stable Diffusion (via dreamstudio)"
56 parser.add_argument("--seed", nargs='?', type=int, required=False, default=None,
57 help="seed value to use (defaults to random)")
58 parser.add_argument("prompt", nargs='+',
59 help="text-to-image prompt")
61 args = parser.parse_args()
63 stability_api = client.StabilityInference(
64 key=os.environ['STABILITY_KEY'],
66 engine="stable-diffusion-512-v2-0"
69 prompt_string = " ".join(args.prompt)
72 answers = stability_api.generate(seed=args.seed, prompt=prompt_string)
74 answers = stability_api.generate(prompt=prompt_string)
78 for artifact in resp.artifacts:
79 if artifact.finish_reason == generation.FILTER:
80 warnings.warn("Request tripped over API safety filters")
81 elif artifact.type == generation.ARTIFACT_IMAGE:
82 results.append(save_artifact(artifact, prompt_string))
83 print(json.dumps(results))
85 if __name__ == "__main__":