File size: 1,673 Bytes
923696f |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 |
import io
import os.path
import uuid
from base64 import b64decode
import openai
import requests
from PIL import Image
from autogpt.config import Config
cfg = Config()
working_directory = "auto_gpt_workspace"
def generate_image(prompt):
filename = str(uuid.uuid4()) + ".jpg"
# DALL-E
if cfg.image_provider == "dalle":
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)
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
elif cfg.image_provider == "sd":
API_URL = (
"https://api-inference.huggingface.co/models/CompVis/stable-diffusion-v1-4"
)
if cfg.huggingface_api_token is None:
raise ValueError(
"You need to set your Hugging Face API token in the config file."
)
headers = {"Authorization": "Bearer " + cfg.huggingface_api_token}
response = requests.post(
API_URL,
headers=headers,
json={
"inputs": prompt,
},
)
image = Image.open(io.BytesIO(response.content))
print("Image Generated for prompt:" + prompt)
image.save(os.path.join(working_directory, filename))
return "Saved to disk:" + filename
else:
return "No Image Provider Set"
|