This commit is contained in:
admin committed
commit
46af290db3
100 files changed
+16663
No files matched your search
@@ -0,0 +1,203 @@
|
||||
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()
|
||||
Reference in new issue
Block a user