diff --git a/scripts/image_gen.py b/scripts/image_gen.py index bb3e7686..92cda290 100644 --- a/scripts/image_gen.py +++ b/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)