mirror of
https://github.com/MODSetter/SurfSense.git
synced 2026-04-25 16:56:22 +02:00
710 lines
24 KiB
Python
710 lines
24 KiB
Python
"""
|
|
Image Generation routes:
|
|
- CRUD for ImageGenerationConfig (user-created image model configs)
|
|
- Global image gen configs endpoint (from YAML)
|
|
- Image generation execution (calls litellm.aimage_generation())
|
|
- CRUD for ImageGeneration records (results)
|
|
- Image serving endpoint (serves b64_json images from DB, protected by signed tokens)
|
|
"""
|
|
|
|
import base64
|
|
import logging
|
|
|
|
from fastapi import APIRouter, Depends, HTTPException, Query
|
|
from fastapi.responses import Response
|
|
from litellm import aimage_generation
|
|
from sqlalchemy import select
|
|
from sqlalchemy.exc import SQLAlchemyError
|
|
from sqlalchemy.ext.asyncio import AsyncSession
|
|
|
|
from app.config import config
|
|
from app.db import (
|
|
ImageGeneration,
|
|
ImageGenerationConfig,
|
|
Permission,
|
|
SearchSpace,
|
|
SearchSpaceMembership,
|
|
User,
|
|
get_async_session,
|
|
)
|
|
from app.schemas import (
|
|
GlobalImageGenConfigRead,
|
|
ImageGenerationConfigCreate,
|
|
ImageGenerationConfigRead,
|
|
ImageGenerationConfigUpdate,
|
|
ImageGenerationCreate,
|
|
ImageGenerationListRead,
|
|
ImageGenerationRead,
|
|
)
|
|
from app.services.image_gen_router_service import (
|
|
IMAGE_GEN_AUTO_MODE_ID,
|
|
ImageGenRouterService,
|
|
is_image_gen_auto_mode,
|
|
)
|
|
from app.users import current_active_user
|
|
from app.utils.rbac import check_permission
|
|
from app.utils.signed_image_urls import verify_image_token
|
|
|
|
router = APIRouter()
|
|
logger = logging.getLogger(__name__)
|
|
|
|
# Provider mapping for building litellm model strings.
|
|
# Only includes providers that support image generation.
|
|
# See: https://docs.litellm.ai/docs/image_generation#supported-providers
|
|
_PROVIDER_MAP = {
|
|
"OPENAI": "openai",
|
|
"AZURE_OPENAI": "azure",
|
|
"GOOGLE": "gemini", # Google AI Studio
|
|
"VERTEX_AI": "vertex_ai",
|
|
"BEDROCK": "bedrock", # AWS Bedrock
|
|
"RECRAFT": "recraft",
|
|
"OPENROUTER": "openrouter",
|
|
"XINFERENCE": "xinference",
|
|
"NSCALE": "nscale",
|
|
}
|
|
|
|
|
|
def _get_global_image_gen_config(config_id: int) -> dict | None:
|
|
"""Get a global image generation configuration by ID (negative IDs)."""
|
|
if config_id == IMAGE_GEN_AUTO_MODE_ID:
|
|
return {
|
|
"id": IMAGE_GEN_AUTO_MODE_ID,
|
|
"name": "Auto (Fastest)",
|
|
"provider": "AUTO",
|
|
"model_name": "auto",
|
|
"is_auto_mode": True,
|
|
}
|
|
if config_id > 0:
|
|
return None
|
|
for cfg in config.GLOBAL_IMAGE_GEN_CONFIGS:
|
|
if cfg.get("id") == config_id:
|
|
return cfg
|
|
return None
|
|
|
|
|
|
def _build_model_string(
|
|
provider: str, model_name: str, custom_provider: str | None
|
|
) -> str:
|
|
"""Build a litellm model string from provider + model_name."""
|
|
if custom_provider:
|
|
return f"{custom_provider}/{model_name}"
|
|
prefix = _PROVIDER_MAP.get(provider.upper(), provider.lower())
|
|
return f"{prefix}/{model_name}"
|
|
|
|
|
|
async def _execute_image_generation(
|
|
session: AsyncSession,
|
|
image_gen: ImageGeneration,
|
|
search_space: SearchSpace,
|
|
) -> None:
|
|
"""
|
|
Call litellm.aimage_generation() with the appropriate config.
|
|
|
|
Resolution order:
|
|
1. Explicit image_generation_config_id on the request
|
|
2. Search space's image_generation_config_id preference
|
|
3. Falls back to Auto mode if available
|
|
"""
|
|
config_id = image_gen.image_generation_config_id
|
|
if config_id is None:
|
|
config_id = search_space.image_generation_config_id or IMAGE_GEN_AUTO_MODE_ID
|
|
image_gen.image_generation_config_id = config_id
|
|
|
|
# Build kwargs
|
|
gen_kwargs = {}
|
|
if image_gen.n is not None:
|
|
gen_kwargs["n"] = image_gen.n
|
|
if image_gen.quality is not None:
|
|
gen_kwargs["quality"] = image_gen.quality
|
|
if image_gen.size is not None:
|
|
gen_kwargs["size"] = image_gen.size
|
|
if image_gen.style is not None:
|
|
gen_kwargs["style"] = image_gen.style
|
|
if image_gen.response_format is not None:
|
|
gen_kwargs["response_format"] = image_gen.response_format
|
|
|
|
if is_image_gen_auto_mode(config_id):
|
|
if not ImageGenRouterService.is_initialized():
|
|
raise ValueError(
|
|
"Auto mode requested but Image Generation Router not initialized. "
|
|
"Ensure global_llm_config.yaml has global_image_generation_configs."
|
|
)
|
|
response = await ImageGenRouterService.aimage_generation(
|
|
prompt=image_gen.prompt, model="auto", **gen_kwargs
|
|
)
|
|
elif config_id < 0:
|
|
# Global config from YAML
|
|
cfg = _get_global_image_gen_config(config_id)
|
|
if not cfg:
|
|
raise ValueError(f"Global image generation config {config_id} not found")
|
|
|
|
model_string = _build_model_string(
|
|
cfg.get("provider", ""), cfg["model_name"], cfg.get("custom_provider")
|
|
)
|
|
gen_kwargs["api_key"] = cfg.get("api_key")
|
|
if cfg.get("api_base"):
|
|
gen_kwargs["api_base"] = cfg["api_base"]
|
|
if cfg.get("api_version"):
|
|
gen_kwargs["api_version"] = cfg["api_version"]
|
|
if cfg.get("litellm_params"):
|
|
gen_kwargs.update(cfg["litellm_params"])
|
|
|
|
# User model override
|
|
if image_gen.model:
|
|
model_string = image_gen.model
|
|
|
|
response = await aimage_generation(
|
|
prompt=image_gen.prompt, model=model_string, **gen_kwargs
|
|
)
|
|
else:
|
|
# Positive ID = DB ImageGenerationConfig
|
|
result = await session.execute(
|
|
select(ImageGenerationConfig).filter(ImageGenerationConfig.id == config_id)
|
|
)
|
|
db_cfg = result.scalars().first()
|
|
if not db_cfg:
|
|
raise ValueError(f"Image generation config {config_id} not found")
|
|
|
|
model_string = _build_model_string(
|
|
db_cfg.provider.value, db_cfg.model_name, db_cfg.custom_provider
|
|
)
|
|
gen_kwargs["api_key"] = db_cfg.api_key
|
|
if db_cfg.api_base:
|
|
gen_kwargs["api_base"] = db_cfg.api_base
|
|
if db_cfg.api_version:
|
|
gen_kwargs["api_version"] = db_cfg.api_version
|
|
if db_cfg.litellm_params:
|
|
gen_kwargs.update(db_cfg.litellm_params)
|
|
|
|
# User model override
|
|
if image_gen.model:
|
|
model_string = image_gen.model
|
|
|
|
response = await aimage_generation(
|
|
prompt=image_gen.prompt, model=model_string, **gen_kwargs
|
|
)
|
|
|
|
# Store response
|
|
image_gen.response_data = (
|
|
response.model_dump() if hasattr(response, "model_dump") else dict(response)
|
|
)
|
|
if not image_gen.model and hasattr(response, "_hidden_params"):
|
|
hidden = response._hidden_params
|
|
if isinstance(hidden, dict) and hidden.get("model"):
|
|
image_gen.model = hidden["model"]
|
|
|
|
|
|
# =============================================================================
|
|
# Global Image Generation Configs (from YAML)
|
|
# =============================================================================
|
|
|
|
|
|
@router.get(
|
|
"/global-image-generation-configs",
|
|
response_model=list[GlobalImageGenConfigRead],
|
|
)
|
|
async def get_global_image_gen_configs(
|
|
user: User = Depends(current_active_user),
|
|
):
|
|
"""Get all global image generation configs. API keys are hidden."""
|
|
try:
|
|
global_configs = config.GLOBAL_IMAGE_GEN_CONFIGS
|
|
safe_configs = []
|
|
|
|
if global_configs and len(global_configs) > 0:
|
|
safe_configs.append(
|
|
{
|
|
"id": 0,
|
|
"name": "Auto (Fastest)",
|
|
"description": "Automatically routes across available image generation providers.",
|
|
"provider": "AUTO",
|
|
"custom_provider": None,
|
|
"model_name": "auto",
|
|
"api_base": None,
|
|
"api_version": None,
|
|
"litellm_params": {},
|
|
"is_global": True,
|
|
"is_auto_mode": True,
|
|
}
|
|
)
|
|
|
|
for cfg in global_configs:
|
|
safe_configs.append(
|
|
{
|
|
"id": cfg.get("id"),
|
|
"name": cfg.get("name"),
|
|
"description": cfg.get("description"),
|
|
"provider": cfg.get("provider"),
|
|
"custom_provider": cfg.get("custom_provider"),
|
|
"model_name": cfg.get("model_name"),
|
|
"api_base": cfg.get("api_base") or None,
|
|
"api_version": cfg.get("api_version") or None,
|
|
"litellm_params": cfg.get("litellm_params", {}),
|
|
"is_global": True,
|
|
}
|
|
)
|
|
|
|
return safe_configs
|
|
except Exception as e:
|
|
logger.exception("Failed to fetch global image generation configs")
|
|
raise HTTPException(
|
|
status_code=500, detail=f"Failed to fetch configs: {e!s}"
|
|
) from e
|
|
|
|
|
|
# =============================================================================
|
|
# ImageGenerationConfig CRUD
|
|
# =============================================================================
|
|
|
|
|
|
@router.post("/image-generation-configs", response_model=ImageGenerationConfigRead)
|
|
async def create_image_gen_config(
|
|
config_data: ImageGenerationConfigCreate,
|
|
session: AsyncSession = Depends(get_async_session),
|
|
user: User = Depends(current_active_user),
|
|
):
|
|
"""Create a new image generation config for a search space."""
|
|
try:
|
|
await check_permission(
|
|
session,
|
|
user,
|
|
config_data.search_space_id,
|
|
Permission.IMAGE_GENERATIONS_CREATE.value,
|
|
"You don't have permission to create image generation configs in this search space",
|
|
)
|
|
|
|
db_config = ImageGenerationConfig(**config_data.model_dump(), user_id=user.id)
|
|
session.add(db_config)
|
|
await session.commit()
|
|
await session.refresh(db_config)
|
|
return db_config
|
|
|
|
except HTTPException:
|
|
raise
|
|
except Exception as e:
|
|
await session.rollback()
|
|
logger.exception("Failed to create ImageGenerationConfig")
|
|
raise HTTPException(
|
|
status_code=500, detail=f"Failed to create config: {e!s}"
|
|
) from e
|
|
|
|
|
|
@router.get("/image-generation-configs", response_model=list[ImageGenerationConfigRead])
|
|
async def list_image_gen_configs(
|
|
search_space_id: int,
|
|
skip: int = 0,
|
|
limit: int = 100,
|
|
session: AsyncSession = Depends(get_async_session),
|
|
user: User = Depends(current_active_user),
|
|
):
|
|
"""List image generation configs for a search space."""
|
|
try:
|
|
await check_permission(
|
|
session,
|
|
user,
|
|
search_space_id,
|
|
Permission.IMAGE_GENERATIONS_READ.value,
|
|
"You don't have permission to view image generation configs in this search space",
|
|
)
|
|
|
|
result = await session.execute(
|
|
select(ImageGenerationConfig)
|
|
.filter(ImageGenerationConfig.search_space_id == search_space_id)
|
|
.order_by(ImageGenerationConfig.created_at.desc())
|
|
.offset(skip)
|
|
.limit(limit)
|
|
)
|
|
return result.scalars().all()
|
|
|
|
except HTTPException:
|
|
raise
|
|
except Exception as e:
|
|
logger.exception("Failed to list ImageGenerationConfigs")
|
|
raise HTTPException(
|
|
status_code=500, detail=f"Failed to fetch configs: {e!s}"
|
|
) from e
|
|
|
|
|
|
@router.get(
|
|
"/image-generation-configs/{config_id}", response_model=ImageGenerationConfigRead
|
|
)
|
|
async def get_image_gen_config(
|
|
config_id: int,
|
|
session: AsyncSession = Depends(get_async_session),
|
|
user: User = Depends(current_active_user),
|
|
):
|
|
"""Get a specific image generation config by ID."""
|
|
try:
|
|
result = await session.execute(
|
|
select(ImageGenerationConfig).filter(ImageGenerationConfig.id == config_id)
|
|
)
|
|
db_config = result.scalars().first()
|
|
if not db_config:
|
|
raise HTTPException(status_code=404, detail="Config not found")
|
|
|
|
await check_permission(
|
|
session,
|
|
user,
|
|
db_config.search_space_id,
|
|
Permission.IMAGE_GENERATIONS_READ.value,
|
|
"You don't have permission to view image generation configs in this search space",
|
|
)
|
|
return db_config
|
|
|
|
except HTTPException:
|
|
raise
|
|
except Exception as e:
|
|
logger.exception("Failed to get ImageGenerationConfig")
|
|
raise HTTPException(
|
|
status_code=500, detail=f"Failed to fetch config: {e!s}"
|
|
) from e
|
|
|
|
|
|
@router.put(
|
|
"/image-generation-configs/{config_id}", response_model=ImageGenerationConfigRead
|
|
)
|
|
async def update_image_gen_config(
|
|
config_id: int,
|
|
update_data: ImageGenerationConfigUpdate,
|
|
session: AsyncSession = Depends(get_async_session),
|
|
user: User = Depends(current_active_user),
|
|
):
|
|
"""Update an existing image generation config."""
|
|
try:
|
|
result = await session.execute(
|
|
select(ImageGenerationConfig).filter(ImageGenerationConfig.id == config_id)
|
|
)
|
|
db_config = result.scalars().first()
|
|
if not db_config:
|
|
raise HTTPException(status_code=404, detail="Config not found")
|
|
|
|
await check_permission(
|
|
session,
|
|
user,
|
|
db_config.search_space_id,
|
|
Permission.IMAGE_GENERATIONS_CREATE.value,
|
|
"You don't have permission to update image generation configs in this search space",
|
|
)
|
|
|
|
for key, value in update_data.model_dump(exclude_unset=True).items():
|
|
setattr(db_config, key, value)
|
|
|
|
await session.commit()
|
|
await session.refresh(db_config)
|
|
return db_config
|
|
|
|
except HTTPException:
|
|
raise
|
|
except Exception as e:
|
|
await session.rollback()
|
|
logger.exception("Failed to update ImageGenerationConfig")
|
|
raise HTTPException(
|
|
status_code=500, detail=f"Failed to update config: {e!s}"
|
|
) from e
|
|
|
|
|
|
@router.delete("/image-generation-configs/{config_id}", response_model=dict)
|
|
async def delete_image_gen_config(
|
|
config_id: int,
|
|
session: AsyncSession = Depends(get_async_session),
|
|
user: User = Depends(current_active_user),
|
|
):
|
|
"""Delete an image generation config."""
|
|
try:
|
|
result = await session.execute(
|
|
select(ImageGenerationConfig).filter(ImageGenerationConfig.id == config_id)
|
|
)
|
|
db_config = result.scalars().first()
|
|
if not db_config:
|
|
raise HTTPException(status_code=404, detail="Config not found")
|
|
|
|
await check_permission(
|
|
session,
|
|
user,
|
|
db_config.search_space_id,
|
|
Permission.IMAGE_GENERATIONS_DELETE.value,
|
|
"You don't have permission to delete image generation configs in this search space",
|
|
)
|
|
|
|
await session.delete(db_config)
|
|
await session.commit()
|
|
return {
|
|
"message": "Image generation config deleted successfully",
|
|
"id": config_id,
|
|
}
|
|
|
|
except HTTPException:
|
|
raise
|
|
except Exception as e:
|
|
await session.rollback()
|
|
logger.exception("Failed to delete ImageGenerationConfig")
|
|
raise HTTPException(
|
|
status_code=500, detail=f"Failed to delete config: {e!s}"
|
|
) from e
|
|
|
|
|
|
# =============================================================================
|
|
# Image Generation Execution + Results CRUD
|
|
# =============================================================================
|
|
|
|
|
|
@router.post("/image-generations", response_model=ImageGenerationRead)
|
|
async def create_image_generation(
|
|
data: ImageGenerationCreate,
|
|
session: AsyncSession = Depends(get_async_session),
|
|
user: User = Depends(current_active_user),
|
|
):
|
|
"""Create and execute an image generation request."""
|
|
try:
|
|
await check_permission(
|
|
session,
|
|
user,
|
|
data.search_space_id,
|
|
Permission.IMAGE_GENERATIONS_CREATE.value,
|
|
"You don't have permission to create image generations in this search space",
|
|
)
|
|
|
|
result = await session.execute(
|
|
select(SearchSpace).filter(SearchSpace.id == data.search_space_id)
|
|
)
|
|
search_space = result.scalars().first()
|
|
if not search_space:
|
|
raise HTTPException(status_code=404, detail="Search space not found")
|
|
|
|
db_image_gen = ImageGeneration(
|
|
prompt=data.prompt,
|
|
model=data.model,
|
|
n=data.n,
|
|
quality=data.quality,
|
|
size=data.size,
|
|
style=data.style,
|
|
response_format=data.response_format,
|
|
image_generation_config_id=data.image_generation_config_id,
|
|
search_space_id=data.search_space_id,
|
|
created_by_id=user.id,
|
|
)
|
|
session.add(db_image_gen)
|
|
await session.flush()
|
|
|
|
try:
|
|
await _execute_image_generation(session, db_image_gen, search_space)
|
|
except Exception as e:
|
|
logger.exception("Image generation call failed")
|
|
db_image_gen.error_message = str(e)
|
|
|
|
await session.commit()
|
|
await session.refresh(db_image_gen)
|
|
return db_image_gen
|
|
|
|
except HTTPException:
|
|
raise
|
|
except SQLAlchemyError:
|
|
await session.rollback()
|
|
raise HTTPException(
|
|
status_code=500, detail="Database error during image generation"
|
|
) from None
|
|
except Exception as e:
|
|
await session.rollback()
|
|
logger.exception("Failed to create image generation")
|
|
raise HTTPException(
|
|
status_code=500, detail=f"Image generation failed: {e!s}"
|
|
) from e
|
|
|
|
|
|
@router.get("/image-generations", response_model=list[ImageGenerationListRead])
|
|
async def list_image_generations(
|
|
search_space_id: int | None = None,
|
|
skip: int = 0,
|
|
limit: int = 50,
|
|
session: AsyncSession = Depends(get_async_session),
|
|
user: User = Depends(current_active_user),
|
|
):
|
|
"""List image generations."""
|
|
if skip < 0 or limit < 1:
|
|
raise HTTPException(status_code=400, detail="Invalid pagination parameters")
|
|
if limit > 100:
|
|
limit = 100
|
|
|
|
try:
|
|
if search_space_id is not None:
|
|
await check_permission(
|
|
session,
|
|
user,
|
|
search_space_id,
|
|
Permission.IMAGE_GENERATIONS_READ.value,
|
|
"You don't have permission to read image generations in this search space",
|
|
)
|
|
result = await session.execute(
|
|
select(ImageGeneration)
|
|
.filter(ImageGeneration.search_space_id == search_space_id)
|
|
.order_by(ImageGeneration.created_at.desc())
|
|
.offset(skip)
|
|
.limit(limit)
|
|
)
|
|
else:
|
|
result = await session.execute(
|
|
select(ImageGeneration)
|
|
.join(SearchSpace)
|
|
.join(SearchSpaceMembership)
|
|
.filter(SearchSpaceMembership.user_id == user.id)
|
|
.order_by(ImageGeneration.created_at.desc())
|
|
.offset(skip)
|
|
.limit(limit)
|
|
)
|
|
|
|
return [
|
|
ImageGenerationListRead.from_orm_with_count(img)
|
|
for img in result.scalars().all()
|
|
]
|
|
|
|
except HTTPException:
|
|
raise
|
|
except SQLAlchemyError:
|
|
raise HTTPException(
|
|
status_code=500, detail="Database error fetching image generations"
|
|
) from None
|
|
|
|
|
|
@router.get("/image-generations/{image_gen_id}", response_model=ImageGenerationRead)
|
|
async def get_image_generation(
|
|
image_gen_id: int,
|
|
session: AsyncSession = Depends(get_async_session),
|
|
user: User = Depends(current_active_user),
|
|
):
|
|
"""Get a specific image generation by ID."""
|
|
try:
|
|
result = await session.execute(
|
|
select(ImageGeneration).filter(ImageGeneration.id == image_gen_id)
|
|
)
|
|
image_gen = result.scalars().first()
|
|
if not image_gen:
|
|
raise HTTPException(status_code=404, detail="Image generation not found")
|
|
|
|
await check_permission(
|
|
session,
|
|
user,
|
|
image_gen.search_space_id,
|
|
Permission.IMAGE_GENERATIONS_READ.value,
|
|
"You don't have permission to read image generations in this search space",
|
|
)
|
|
return image_gen
|
|
|
|
except HTTPException:
|
|
raise
|
|
except SQLAlchemyError:
|
|
raise HTTPException(
|
|
status_code=500, detail="Database error fetching image generation"
|
|
) from None
|
|
|
|
|
|
@router.delete("/image-generations/{image_gen_id}", response_model=dict)
|
|
async def delete_image_generation(
|
|
image_gen_id: int,
|
|
session: AsyncSession = Depends(get_async_session),
|
|
user: User = Depends(current_active_user),
|
|
):
|
|
"""Delete an image generation record."""
|
|
try:
|
|
result = await session.execute(
|
|
select(ImageGeneration).filter(ImageGeneration.id == image_gen_id)
|
|
)
|
|
db_image_gen = result.scalars().first()
|
|
if not db_image_gen:
|
|
raise HTTPException(status_code=404, detail="Image generation not found")
|
|
|
|
await check_permission(
|
|
session,
|
|
user,
|
|
db_image_gen.search_space_id,
|
|
Permission.IMAGE_GENERATIONS_DELETE.value,
|
|
"You don't have permission to delete image generations in this search space",
|
|
)
|
|
|
|
await session.delete(db_image_gen)
|
|
await session.commit()
|
|
return {"message": "Image generation deleted successfully"}
|
|
|
|
except HTTPException:
|
|
raise
|
|
except SQLAlchemyError:
|
|
await session.rollback()
|
|
raise HTTPException(
|
|
status_code=500, detail="Database error deleting image generation"
|
|
) from None
|
|
|
|
|
|
# =============================================================================
|
|
# Image Serving (serves generated images from DB, protected by signed tokens)
|
|
# =============================================================================
|
|
|
|
|
|
@router.get("/image-generations/{image_gen_id}/image")
|
|
async def serve_generated_image(
|
|
image_gen_id: int,
|
|
token: str = Query(..., description="Signed access token"),
|
|
index: int = 0,
|
|
session: AsyncSession = Depends(get_async_session),
|
|
):
|
|
"""
|
|
Serve a generated image by ID, protected by a signed token.
|
|
|
|
The token is generated when the image URL is created by the generate_image
|
|
tool and encodes the image_gen_id, search_space_id, and an expiry timestamp.
|
|
This ensures only users with access to the search space can view images,
|
|
without requiring auth headers (which <img> tags cannot pass).
|
|
|
|
Args:
|
|
image_gen_id: The image generation record ID
|
|
token: HMAC-signed access token (included as query parameter)
|
|
index: Which image to serve if multiple were generated (default: 0)
|
|
"""
|
|
try:
|
|
result = await session.execute(
|
|
select(ImageGeneration).filter(ImageGeneration.id == image_gen_id)
|
|
)
|
|
image_gen = result.scalars().first()
|
|
if not image_gen:
|
|
raise HTTPException(status_code=404, detail="Image generation not found")
|
|
|
|
# Verify the access token against the one stored on the record
|
|
if not verify_image_token(image_gen.access_token, token):
|
|
raise HTTPException(status_code=403, detail="Invalid image access token")
|
|
|
|
if not image_gen.response_data:
|
|
raise HTTPException(status_code=404, detail="No image data available")
|
|
|
|
images = image_gen.response_data.get("data", [])
|
|
if not images or index >= len(images):
|
|
raise HTTPException(
|
|
status_code=404, detail="Image not found at the specified index"
|
|
)
|
|
|
|
image_entry = images[index]
|
|
|
|
# If there's a URL, redirect to it
|
|
if image_entry.get("url"):
|
|
from fastapi.responses import RedirectResponse
|
|
|
|
return RedirectResponse(url=image_entry["url"])
|
|
|
|
# If there's b64_json data, decode and serve it
|
|
if image_entry.get("b64_json"):
|
|
image_bytes = base64.b64decode(image_entry["b64_json"])
|
|
return Response(
|
|
content=image_bytes,
|
|
media_type="image/png",
|
|
headers={
|
|
"Cache-Control": "public, max-age=86400",
|
|
"Content-Disposition": f'inline; filename="generated-{image_gen_id}-{index}.png"',
|
|
},
|
|
)
|
|
|
|
raise HTTPException(status_code=404, detail="No displayable image data")
|
|
|
|
except HTTPException:
|
|
raise
|
|
except Exception as e:
|
|
logger.exception("Failed to serve generated image")
|
|
raise HTTPException(
|
|
status_code=500, detail=f"Failed to serve image: {e!s}"
|
|
) from e
|