90 lines
3.3 KiB
Python
90 lines
3.3 KiB
Python
import base64
|
|
import struct
|
|
|
|
import requests
|
|
|
|
from config import PIAPI_API_KEY, PIAPI_IMAGE_MODEL, get_proxies
|
|
|
|
PIAPI_IMAGE_URL = "https://api.piapi.ai/v1/images/generations"
|
|
|
|
# Доступные размеры (gpt-image-1)
|
|
IMAGE_SIZES = {
|
|
"square": ("1024x1024", "Квадрат 1024x1024"),
|
|
"wide": ("1536x1024", "Горизонтальный 1536x1024"),
|
|
"tall": ("1024x1536", "Вертикальный 1024x1536"),
|
|
}
|
|
|
|
# Пресеты стилей - дописываются к промпту
|
|
STYLE_PRESETS = {
|
|
"none": ("Без стиля", ""),
|
|
"photo": ("Фотореализм", ", ultra realistic photo, professional photography, natural lighting"),
|
|
"illustration": ("Иллюстрация", ", digital illustration, vibrant colors, detailed"),
|
|
"minimal": ("Минимализм", ", minimalist style, clean background, simple composition"),
|
|
"3d": ("3D-рендер", ", 3D render, soft shadows, modern design"),
|
|
"flat": ("Плоская графика", ", flat vector illustration, bold shapes, social media style"),
|
|
}
|
|
|
|
|
|
def detect_ext(data: bytes) -> str:
|
|
"""Определяет расширение файла по magic-байтам."""
|
|
if data[:8] == b"\x89PNG\r\n\x1a\n":
|
|
return ".png"
|
|
if data[:2] == b"\xff\xd8":
|
|
return ".jpg"
|
|
if data[:4] == b"RIFF" and data[8:12] == b"WEBP":
|
|
return ".webp"
|
|
if data[:4] == b"GIF8":
|
|
return ".gif"
|
|
return ".png"
|
|
|
|
|
|
class ImageGenerator:
|
|
"""Генерация изображений через PiAPI (OpenAI-совместимый /images/generations)."""
|
|
|
|
def __init__(self, model: str | None = None):
|
|
self.model = model or PIAPI_IMAGE_MODEL
|
|
self.proxies = get_proxies()
|
|
|
|
def generate(self, prompt: str, size: str = "square") -> bytes:
|
|
"""Генерирует картинку и возвращает её байты. Поднимает исключение при ошибке."""
|
|
if size not in IMAGE_SIZES:
|
|
size = "square"
|
|
size_value = IMAGE_SIZES[size][0]
|
|
|
|
response = requests.post(
|
|
PIAPI_IMAGE_URL,
|
|
headers={
|
|
"Content-Type": "application/json",
|
|
"Authorization": f"Bearer {PIAPI_API_KEY}",
|
|
},
|
|
json={
|
|
"model": self.model,
|
|
"prompt": prompt,
|
|
"n": 1,
|
|
"size": size_value,
|
|
},
|
|
proxies=self.proxies,
|
|
timeout=240,
|
|
)
|
|
response.raise_for_status()
|
|
|
|
data = response.json().get("data") or []
|
|
if not data:
|
|
raise RuntimeError("Пустой ответ от API картинок")
|
|
|
|
item = data[0]
|
|
if item.get("b64_json"):
|
|
return base64.b64decode(item["b64_json"])
|
|
if item.get("url"):
|
|
img_response = requests.get(item["url"], proxies=self.proxies, timeout=120)
|
|
img_response.raise_for_status()
|
|
return img_response.content
|
|
|
|
raise RuntimeError("Неожиданный формат ответа от API картинок")
|
|
|
|
def build_prompt(self, prompt: str, style: str = "none") -> str:
|
|
"""Дополняет промпт пресетом стиля."""
|
|
if style in STYLE_PRESETS:
|
|
return prompt.strip() + STYLE_PRESETS[style][1]
|
|
return prompt.strip()
|