204 lines
8.1 KiB
Python
204 lines
8.1 KiB
Python
import base64
|
||
import logging
|
||
import re
|
||
from pathlib import Path
|
||
from typing import Optional
|
||
|
||
import httpx
|
||
from config.settings import settings
|
||
from fastapi import HTTPException
|
||
|
||
from services.project_service import project_service
|
||
|
||
logger = logging.getLogger(__name__)
|
||
|
||
|
||
class ImageGenerator:
|
||
def __init__(self):
|
||
self.api_key = settings.OPENROUTER_API_KEY
|
||
self.base_url = getattr(
|
||
settings, "OPENROUTER_BASE_URL", "https://openrouter.ai/api/v1"
|
||
)
|
||
self.default_model = "qwen/qwen-image-2.1-turbo"
|
||
|
||
def _get_images_dir(self, project_id: str, user_id: str) -> Path:
|
||
"""Получить или создать директорию для изображений проекта"""
|
||
project_dir = project_service.get_project_dir(project_id, user_id)
|
||
images_dir = project_dir / "images"
|
||
images_dir.mkdir(exist_ok=True)
|
||
return images_dir
|
||
|
||
async def generate_image(
|
||
self,
|
||
prompt: str,
|
||
project_id: str,
|
||
user_id: str,
|
||
document_id: str,
|
||
document_type: str, # 'book' или 'chapter'
|
||
model: Optional[str] = None,
|
||
size: str = "1024x1024",
|
||
) -> dict:
|
||
"""
|
||
Генерирует изображение по промпту и сохраняет его в папку проекта.
|
||
Имя файла определяется по конвенции:
|
||
- document_type='book' -> cover.png (обложка книги)
|
||
- document_type='chapter' -> {document_id}.png (иллюстрация главы)
|
||
Возвращает информацию о сохранённом файле.
|
||
"""
|
||
if not self.api_key or self.api_key == "your_api_key_here":
|
||
raise ValueError("OpenRouter API key not configured")
|
||
|
||
model = model or self.default_model
|
||
|
||
# Генерация изображений использует OpenAI Images API endpoint,
|
||
# а не Chat Completions. См. официальную документацию модели.
|
||
payload = {
|
||
"model": model,
|
||
"prompt": prompt,
|
||
"n": 1,
|
||
"response_format": "b64_json",
|
||
}
|
||
|
||
# Мапим size в aspect_ratio для моделей, которые его поддерживают
|
||
aspect_ratio_map = {
|
||
"1024x1024": "1:1",
|
||
"1024x1792": "9:16",
|
||
"1792x1024": "16:9",
|
||
"512x512": "1:1",
|
||
}
|
||
if size in aspect_ratio_map:
|
||
payload["size"] = size
|
||
|
||
headers = {
|
||
"Authorization": f"Bearer {self.api_key}",
|
||
"Content-Type": "application/json",
|
||
"HTTP-Referer": "http://localhost:3000",
|
||
"X-Title": "Writer IDE",
|
||
}
|
||
|
||
try:
|
||
async with httpx.AsyncClient(timeout=120.0) as client:
|
||
response = await client.post(
|
||
f"{self.base_url}/images/generations",
|
||
json=payload,
|
||
headers=headers,
|
||
)
|
||
response.raise_for_status()
|
||
data = response.json()
|
||
|
||
# Проверяем наличие ошибки в ответе (некоторые провайдеры
|
||
# возвращают 200 с полем error вместо HTTP error status)
|
||
if "error" in data:
|
||
error_info = data["error"]
|
||
error_msg = (
|
||
error_info.get("message", str(error_info))
|
||
if isinstance(error_info, dict)
|
||
else str(error_info)
|
||
)
|
||
raise ValueError(f"API returned error: {error_msg}")
|
||
|
||
# OpenAI Images API формат: {"data": [{"b64_json": "..."}]}
|
||
images_data = data.get("data", [])
|
||
if not images_data:
|
||
raise ValueError(
|
||
f"No image data in response. Keys: {list(data.keys())}"
|
||
)
|
||
|
||
img_data = images_data[0]
|
||
image_url = None
|
||
image_base64 = None
|
||
|
||
if isinstance(img_data, dict):
|
||
image_base64 = img_data.get("b64_json")
|
||
image_url = img_data.get("url")
|
||
elif isinstance(img_data, str):
|
||
# Некоторые провайдеры могут вернуть URL напрямую в data
|
||
image_url = img_data if img_data.startswith("http") else None
|
||
|
||
# Фоллбэк: пробуем извлечь из text content (если провайдер
|
||
# использует chat/completions формат поверх images/generations)
|
||
if not image_url and not image_base64:
|
||
content = (
|
||
data.get("choices", [{}])[0].get("message", {}).get("content", "")
|
||
)
|
||
if content:
|
||
base64_match = re.search(
|
||
r"data:image/[^;]+;base64,([A-Za-z0-9+/=]+)", content
|
||
)
|
||
if base64_match:
|
||
image_base64 = base64_match.group(1)
|
||
|
||
if not image_url and not image_base64:
|
||
logger.warning(
|
||
"Could not extract image. Response keys: %s",
|
||
list(data.keys()),
|
||
)
|
||
raise ValueError(
|
||
f"Could not extract image from response. Available keys: {list(data.keys())}"
|
||
)
|
||
|
||
# Сохраняем изображение
|
||
images_dir = self._get_images_dir(project_id, user_id)
|
||
|
||
# Определяем имя файла по конвенции
|
||
if document_type == "book":
|
||
filename = "cover.png" # Обложка книги
|
||
else: # chapter
|
||
filename = f"{document_id}.png" # Изображение главы
|
||
|
||
filepath = images_dir / filename
|
||
|
||
if image_url:
|
||
async with httpx.AsyncClient(timeout=60.0) as client:
|
||
img_response = await client.get(image_url)
|
||
img_response.raise_for_status()
|
||
image_bytes = img_response.content
|
||
filepath.write_bytes(image_bytes)
|
||
data_url = f"data:image/png;base64,{base64.b64encode(image_bytes).decode('ascii')}"
|
||
elif image_base64:
|
||
image_bytes = base64.b64decode(image_base64)
|
||
filepath.write_bytes(image_bytes)
|
||
# Определяем MIME тип по первым байтам изображения
|
||
mime = "image/png" # default
|
||
if image_bytes[:2] == b"\xff\xd8":
|
||
mime = "image/jpeg"
|
||
elif image_bytes[:8] == b"\x89PNG\r\n\x1a\n":
|
||
mime = "image/png"
|
||
elif image_bytes[:6] in (b"GIF87a", b"GIF89a"):
|
||
mime = "image/gif"
|
||
elif image_bytes[:4] == b"RIFF" and image_bytes[8:12] == b"WEBP":
|
||
mime = "image/webp"
|
||
data_url = f"data:{mime};base64,{image_base64}"
|
||
else:
|
||
# Должно быть обработано выше, но на всякий случай
|
||
raise ValueError("No image data available to encode")
|
||
|
||
return {
|
||
"image_id": document_id,
|
||
"filename": filename,
|
||
"path": f"/api/images/{project_id}/{filename}",
|
||
"prompt": prompt,
|
||
"model": model,
|
||
"data_url": data_url,
|
||
"document_id": document_id,
|
||
"document_type": document_type,
|
||
}
|
||
|
||
except httpx.HTTPStatusError as e:
|
||
raise HTTPException(
|
||
status_code=502,
|
||
detail=f"Image generation API error: {e.response.status_code} - {e.response.text[:200]}",
|
||
)
|
||
except httpx.TimeoutException:
|
||
raise HTTPException(status_code=504, detail="Image generation timed out")
|
||
except HTTPException:
|
||
raise
|
||
except Exception as e:
|
||
raise HTTPException(
|
||
status_code=500, detail=f"Image generation failed: {str(e)}"
|
||
)
|
||
|
||
|
||
# Singleton
|
||
image_generator = ImageGenerator()
|