import base64
import json
import os
from pathlib import Path
from urllib.request import Request, urlopen

payload = {
    "model": "grok-imagine-image",
    "prompt": "A minimalist perfume campaign, morning light, soft mint backdrop, no text",
    "n": 1,
    "response_format": "b64_json",
}
request = Request(
    "https://api.modelkiwi.com/v1/images/generations",
    data=json.dumps(payload).encode("utf-8"),
    headers={
        "Authorization": "Bearer " + os.environ["IMAGE_GEN_API_KEY"],
        "Content-Type": "application/json",
    },
    method="POST",
)
with urlopen(request, timeout=300) as response:
    result = json.load(response)
image = base64.b64decode(result["data"][0]["b64_json"], validate=True)
if image.startswith(b"\x89PNG\r\n\x1a\n"):
    extension = "png"
elif image.startswith(b"\xff\xd8\xff"):
    extension = "jpg"
elif image[:4] == b"RIFF" and image[8:12] == b"WEBP":
    extension = "webp"
else:
    raise ValueError("Unexpected image format; inspect the response")
output = Path("generated-image." + extension)
output.write_bytes(image)
print(output.resolve())