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):
30 prompt_abbrev = prompt[:200]
32 base = "{}_{}_{}.png".format(seed, prompt_abbrev, counter)
34 base = "{}_{}.png".format(seed, prompt_abbrev)
35 base = normalize_filename(base)
36 filename = "{}{}/{}".format(DOCUMENT_ROOT, IMAGES_PATH, base)
37 if not os.path.exists(filename):
41 img = Image.open(io.BytesIO(artifact.binary))
43 return {"code": seed, "prompt": prompt,
44 "filename": filename.removeprefix(DOCUMENT_ROOT)}
47 os.environ['STABILITY_HOST'] = 'grpc.stability.ai:443'
49 if not 'STABILITY_KEY' in os.environ:
50 print("Error: STABILITY_KEY must be set with a valid key")
53 parser = argparse.ArgumentParser(
54 usage="%(prog)s [--seed=SEED] Prompt string here...",
55 description="Generate text-from-image using Stable Diffusion (via dreamstudio)"
57 parser.add_argument("--seed", nargs='?', type=int, required=False, default=None,
58 help="seed value to use (defaults to random)")
59 parser.add_argument("prompt", nargs='+',
60 help="text-to-image prompt")
62 args = parser.parse_args()
64 stability_api = client.StabilityInference(
65 key=os.environ['STABILITY_KEY'],
67 engine="stable-diffusion-768-v2-1"
70 prompt_string = " ".join(args.prompt)
73 answers = stability_api.generate(width=768, height=768,
74 seed=args.seed, prompt=prompt_string)
76 answers = stability_api.generate(width=768, height=768,
81 for artifact in resp.artifacts:
82 if artifact.finish_reason == generation.FILTER:
83 warnings.warn("Request tripped over API safety filters")
84 elif artifact.type == generation.ARTIFACT_IMAGE:
85 results.append(save_artifact(artifact, prompt_string))
86 print(json.dumps(results))
88 if __name__ == "__main__":