79 lines
2.1 KiB
Python
79 lines
2.1 KiB
Python
import re
|
|
from os import getenv
|
|
from pathlib import Path
|
|
|
|
import requests
|
|
from openai import OpenAI
|
|
|
|
API_KEY = getenv("RECRAFT_API_TOKEN")
|
|
client = OpenAI(base_url="https://external.api.recraft.ai/v1", api_key=API_KEY)
|
|
|
|
styles = {"nature-simplism": "ee2bb119-a790-409f-888e-b4b7b73f0aa0"}
|
|
|
|
|
|
def generate_image_v3(prompt: str, file_path: Path, style: str):
|
|
response = client.images.generate(
|
|
prompt=prompt,
|
|
size="1024x1024",
|
|
model="recraftv3",
|
|
extra_body={
|
|
"style_id": styles[style],
|
|
},
|
|
)
|
|
image_url = response.data[0].url
|
|
if not image_url:
|
|
raise ValueError("No image URL returned by the API.")
|
|
|
|
img_response = requests.get(image_url)
|
|
if img_response.status_code != 200:
|
|
raise IOError(f"Failed to download image from {image_url}")
|
|
|
|
with open(file_path, "wb") as f:
|
|
f.write(img_response.content)
|
|
|
|
|
|
def generate_image_v2(prompt: str, file_path: Path):
|
|
response = client.images.generate(
|
|
prompt=prompt,
|
|
size="1024x1024",
|
|
model="recraftv2",
|
|
style="icon",
|
|
)
|
|
image_url = response.data[0].url
|
|
if not image_url:
|
|
raise ValueError("No image URL returned by the API.")
|
|
|
|
img_response = requests.get(image_url)
|
|
if img_response.status_code != 200:
|
|
raise IOError(f"Failed to download image from {image_url}")
|
|
|
|
with open(file_path, "wb") as f:
|
|
f.write(img_response.content)
|
|
|
|
|
|
def sanitize_filename(name: str) -> str:
|
|
# Define illegal characters (union of Windows + Unix)
|
|
illegal_chars = r'[\\\/:*?"<>|]'
|
|
|
|
# Replace spaces and illegal characters with "_"
|
|
name = re.sub(r"\s", "_", name)
|
|
name = re.sub(illegal_chars, "_", name)
|
|
|
|
# Windows also doesn't allow names ending in . or space
|
|
name = name.rstrip(" .")
|
|
|
|
# Windows reserved device names
|
|
reserved_names = {
|
|
"CON",
|
|
"PRN",
|
|
"AUX",
|
|
"NUL",
|
|
*(f"COM{i}" for i in range(1, 10)),
|
|
*(f"LPT{i}" for i in range(1, 10)),
|
|
}
|
|
name_upper = name.upper().split(".")[0]
|
|
if name_upper in reserved_names:
|
|
name = f"_{name}"
|
|
|
|
return name
|