Switch to new 2.1 model of Stable Diffusion and also 768x768 pixels

[zombocom-ai] / generate-image.py

diff --git a/generate-image.py b/generate-image.py
index 3a2d3f3f65593fdd03a4b308af3ebd412295abf0..408fe21bebfe8086accbf4f373176970fa6dcaaa 100755 (executable)
--- a/generate-image.py
+++ b/generate-image.py
@@ -63,15 +63,17 @@ def main() -> None:
     stability_api = client.StabilityInference(
         key=os.environ['STABILITY_KEY'],
         verbose=True,
-        engine="stable-diffusion-512-v2-0"
+        engine="stable-diffusion-768-v2-1"
     )
 
     prompt_string = " ".join(args.prompt)
 
     if args.seed:
-        answers = stability_api.generate(seed=args.seed, prompt=prompt_string)
+        answers = stability_api.generate(width=768, height=768,
+                                         seed=args.seed, prompt=prompt_string)
     else:
-        answers = stability_api.generate(prompt=prompt_string)
+        answers = stability_api.generate(width=768, height=768,
+                                         prompt=prompt_string)
 
     results = []
     for resp in answers: