feat(ImageGen): add DALL-E support

scripts/image_gen.py

@@ -4,6 +4,8 @@ import os.path
 from PIL import Image
 from config import Config
 import uuid
+import openai
+from base64 import b64decode
 
 cfg = Config()
 
@@ -13,16 +15,36 @@ API_URL = "https://api-inference.huggingface.co/models/CompVis/stable-diffusion-
 headers = {"Authorization": "Bearer " + cfg.huggingface_api_token}
 
 def generate_image(prompt):
+
+    filename = str(uuid.uuid4()) + ".jpg"
+
+    # DALL-E
+    openai.api_key = cfg.openai_api_key
+
+    response = openai.Image.create(
+        prompt=prompt,
+        n=1,
+        size="256x256",
+        response_format="b64_json",
+    )
+
+    print("Image Generated for prompt:" + prompt)
+    print(response["data"][0]["b64_json"][:50])
+
+    image_data = b64decode(response["data"][0]["b64_json"])
+    with open(working_directory + "/" + filename, mode="wb") as png:
+        png.write(image_data)
+
+    return "Saved to disk:" + filename
+
+    # STABLE DIFFUSION
     response = requests.post(API_URL, headers=headers, json={
         "inputs": prompt,
     })
     image = Image.open(io.BytesIO(response.content))
     print("Image Generated for prompt:" + prompt)
 
-    filename = str(uuid.uuid4()) + ".jpg"
-
     image.save(os.path.join(working_directory, filename))
-
     print("Saved to disk:" + filename)
 
     return str("Image " + filename + " saved to disk for prompt: " + prompt)