Initial commit: Freepik REST API
FastAPI async wrapper for Freepik cloud AI API supporting image generation (Mystic, Flux Dev/Pro, SeedReam), video generation (Kling, MiniMax, Seedance), image editing (upscale, relight, style transfer, expand, inpaint), and utilities (background removal, classifier, audio isolation). Includes async task tracking with polling, Docker containerization, and Gitea CI/CD workflow. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
This commit is contained in:
0
app/__init__.py
Normal file
0
app/__init__.py
Normal file
25
app/config.py
Normal file
25
app/config.py
Normal file
@@ -0,0 +1,25 @@
|
||||
from pydantic_settings import BaseSettings
|
||||
|
||||
|
||||
class Settings(BaseSettings):
|
||||
model_config = {'env_prefix': 'FP_'}
|
||||
|
||||
# Freepik API
|
||||
freepik_api_key: str
|
||||
freepik_base_url: str = 'https://api.freepik.com'
|
||||
|
||||
# Paths
|
||||
output_dir: str = '/data/outputs'
|
||||
temp_dir: str = '/data/temp'
|
||||
|
||||
# Limits
|
||||
max_upload_size_mb: int = 50
|
||||
task_poll_interval_seconds: int = 5
|
||||
task_poll_timeout_seconds: int = 600
|
||||
auto_cleanup_hours: int = 24
|
||||
|
||||
# Webhook
|
||||
webhook_secret: str = ''
|
||||
|
||||
|
||||
settings = Settings()
|
||||
42
app/main.py
Normal file
42
app/main.py
Normal file
@@ -0,0 +1,42 @@
|
||||
import logging
|
||||
from contextlib import asynccontextmanager
|
||||
|
||||
from fastapi import FastAPI
|
||||
|
||||
from app.routers import image_editing, image_generation, system, tasks, utilities, video_generation
|
||||
from app.services import file_manager, freepik_client
|
||||
from app.services.webhook import router as webhook_router
|
||||
|
||||
logging.basicConfig(
|
||||
level=logging.INFO,
|
||||
format='%(asctime)s %(levelname)s %(name)s: %(message)s',
|
||||
)
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
@asynccontextmanager
|
||||
async def lifespan(app: FastAPI):
|
||||
logger.info('Starting Freepik API...')
|
||||
file_manager.ensure_directories()
|
||||
await freepik_client.create_client()
|
||||
logger.info('Freepik API ready')
|
||||
yield
|
||||
logger.info('Shutting down...')
|
||||
await freepik_client.close_client()
|
||||
|
||||
|
||||
app = FastAPI(
|
||||
title='Freepik API',
|
||||
version='1.0.0',
|
||||
description='REST API wrapping Freepik cloud AI services',
|
||||
lifespan=lifespan,
|
||||
)
|
||||
|
||||
app.include_router(image_generation.router)
|
||||
app.include_router(video_generation.router)
|
||||
app.include_router(image_editing.router)
|
||||
app.include_router(utilities.router)
|
||||
app.include_router(utilities.icon_router)
|
||||
app.include_router(tasks.router)
|
||||
app.include_router(system.router)
|
||||
app.include_router(webhook_router)
|
||||
0
app/routers/__init__.py
Normal file
0
app/routers/__init__.py
Normal file
120
app/routers/image_editing.py
Normal file
120
app/routers/image_editing.py
Normal file
@@ -0,0 +1,120 @@
|
||||
import asyncio
|
||||
import logging
|
||||
from datetime import datetime, timezone
|
||||
|
||||
from fastapi import APIRouter, HTTPException, Query
|
||||
|
||||
from app.schemas.common import TaskDetail, TaskResponse, TaskStatus
|
||||
from app.schemas.image_editing import (
|
||||
ExpandRequest,
|
||||
InpaintRequest,
|
||||
RelightRequest,
|
||||
StyleTransferRequest,
|
||||
UpscaleCreativeRequest,
|
||||
UpscalePrecisionRequest,
|
||||
)
|
||||
from app.services import freepik_client, task_tracker
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
router = APIRouter(prefix='/api/v1/edit', tags=['image-editing'])
|
||||
|
||||
|
||||
async def _submit_and_respond(result: dict, sync: bool, metadata: dict) -> TaskResponse | TaskDetail:
|
||||
data = result.get('data', result)
|
||||
freepik_task_id = str(data.get('task_id') or data.get('id', ''))
|
||||
if not freepik_task_id:
|
||||
raise HTTPException(status_code=502, detail='No task_id in Freepik response')
|
||||
|
||||
internal_id = task_tracker.submit(freepik_task_id, metadata)
|
||||
|
||||
if not sync:
|
||||
return TaskResponse(
|
||||
task_id=internal_id,
|
||||
status=TaskStatus.pending,
|
||||
created_at=datetime.now(timezone.utc),
|
||||
)
|
||||
|
||||
from app.config import settings
|
||||
elapsed = 0
|
||||
while elapsed < settings.task_poll_timeout_seconds:
|
||||
await asyncio.sleep(2)
|
||||
elapsed += 2
|
||||
task = task_tracker.get_task(internal_id)
|
||||
if not task:
|
||||
raise HTTPException(status_code=500, detail='Task disappeared')
|
||||
if task['status'] == TaskStatus.completed:
|
||||
return TaskDetail(
|
||||
task_id=internal_id,
|
||||
status=TaskStatus.completed,
|
||||
created_at=task['created_at'],
|
||||
updated_at=task['updated_at'],
|
||||
progress=1.0,
|
||||
result_url=f'/api/v1/tasks/{internal_id}/result',
|
||||
)
|
||||
if task['status'] == TaskStatus.failed:
|
||||
raise HTTPException(status_code=502, detail=task.get('error', 'Task failed'))
|
||||
|
||||
raise HTTPException(status_code=504, detail='Task did not complete in time')
|
||||
|
||||
|
||||
@router.post('/upscale/creative', response_model=TaskResponse)
|
||||
async def upscale_creative(request: UpscaleCreativeRequest, sync: bool = Query(False)):
|
||||
result = await freepik_client.upscale_creative(
|
||||
image=request.image,
|
||||
prompt=request.prompt,
|
||||
scale=request.scale,
|
||||
creativity=request.creativity,
|
||||
resemblance=request.resemblance,
|
||||
)
|
||||
return await _submit_and_respond(result, sync, {'operation': 'upscale-creative'})
|
||||
|
||||
|
||||
@router.post('/upscale/precision', response_model=TaskResponse)
|
||||
async def upscale_precision(request: UpscalePrecisionRequest, sync: bool = Query(False)):
|
||||
result = await freepik_client.upscale_precision(
|
||||
image=request.image,
|
||||
scale=request.scale,
|
||||
)
|
||||
return await _submit_and_respond(result, sync, {'operation': 'upscale-precision'})
|
||||
|
||||
|
||||
@router.post('/relight', response_model=TaskResponse)
|
||||
async def relight(request: RelightRequest, sync: bool = Query(False)):
|
||||
result = await freepik_client.relight_image(
|
||||
image=request.image,
|
||||
prompt=request.prompt,
|
||||
light_source=request.light_source,
|
||||
intensity=request.intensity,
|
||||
)
|
||||
return await _submit_and_respond(result, sync, {'operation': 'relight'})
|
||||
|
||||
|
||||
@router.post('/style-transfer', response_model=TaskResponse)
|
||||
async def style_transfer(request: StyleTransferRequest, sync: bool = Query(False)):
|
||||
result = await freepik_client.style_transfer(
|
||||
image=request.image,
|
||||
style_reference=request.style_reference,
|
||||
strength=request.strength,
|
||||
)
|
||||
return await _submit_and_respond(result, sync, {'operation': 'style-transfer'})
|
||||
|
||||
|
||||
@router.post('/expand', response_model=TaskResponse)
|
||||
async def expand(request: ExpandRequest, sync: bool = Query(False)):
|
||||
result = await freepik_client.expand_image(
|
||||
image=request.image,
|
||||
prompt=request.prompt,
|
||||
direction=request.direction,
|
||||
)
|
||||
return await _submit_and_respond(result, sync, {'operation': 'expand'})
|
||||
|
||||
|
||||
@router.post('/inpaint', response_model=TaskResponse)
|
||||
async def inpaint(request: InpaintRequest, sync: bool = Query(False)):
|
||||
result = await freepik_client.inpaint(
|
||||
image=request.image,
|
||||
mask=request.mask,
|
||||
prompt=request.prompt,
|
||||
)
|
||||
return await _submit_and_respond(result, sync, {'operation': 'inpaint'})
|
||||
110
app/routers/image_generation.py
Normal file
110
app/routers/image_generation.py
Normal file
@@ -0,0 +1,110 @@
|
||||
import asyncio
|
||||
import logging
|
||||
from datetime import datetime, timezone
|
||||
|
||||
from fastapi import APIRouter, HTTPException, Query
|
||||
|
||||
from app.schemas.common import TaskDetail, TaskResponse, TaskStatus
|
||||
from app.schemas.image_generation import (
|
||||
FluxDevRequest,
|
||||
FluxProRequest,
|
||||
MysticRequest,
|
||||
SeedreamRequest,
|
||||
)
|
||||
from app.services import freepik_client, task_tracker
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
router = APIRouter(prefix='/api/v1/generate/image', tags=['image-generation'])
|
||||
|
||||
|
||||
async def _submit_and_respond(
|
||||
result: dict,
|
||||
sync: bool,
|
||||
metadata: dict | None = None,
|
||||
) -> TaskResponse | TaskDetail:
|
||||
"""Extract task_id from Freepik response, track it, optionally wait."""
|
||||
data = result.get('data', result)
|
||||
freepik_task_id = str(data.get('task_id') or data.get('id', ''))
|
||||
if not freepik_task_id:
|
||||
raise HTTPException(status_code=502, detail='No task_id in Freepik response')
|
||||
|
||||
internal_id = task_tracker.submit(freepik_task_id, metadata)
|
||||
|
||||
if not sync:
|
||||
return TaskResponse(
|
||||
task_id=internal_id,
|
||||
status=TaskStatus.pending,
|
||||
created_at=datetime.now(timezone.utc),
|
||||
)
|
||||
|
||||
# Sync mode: wait for completion
|
||||
from app.config import settings
|
||||
elapsed = 0
|
||||
while elapsed < settings.task_poll_timeout_seconds:
|
||||
await asyncio.sleep(2)
|
||||
elapsed += 2
|
||||
task = task_tracker.get_task(internal_id)
|
||||
if not task:
|
||||
raise HTTPException(status_code=500, detail='Task disappeared')
|
||||
if task['status'] == TaskStatus.completed:
|
||||
return TaskDetail(
|
||||
task_id=internal_id,
|
||||
status=TaskStatus.completed,
|
||||
created_at=task['created_at'],
|
||||
updated_at=task['updated_at'],
|
||||
progress=1.0,
|
||||
result_url=f'/api/v1/tasks/{internal_id}/result',
|
||||
)
|
||||
if task['status'] == TaskStatus.failed:
|
||||
raise HTTPException(status_code=502, detail=task.get('error', 'Task failed'))
|
||||
|
||||
raise HTTPException(status_code=504, detail='Task did not complete in time')
|
||||
|
||||
|
||||
@router.post('/mystic', response_model=TaskResponse)
|
||||
async def generate_mystic(request: MysticRequest, sync: bool = Query(False)):
|
||||
result = await freepik_client.generate_mystic(
|
||||
prompt=request.prompt,
|
||||
negative_prompt=request.negative_prompt,
|
||||
resolution=request.resolution,
|
||||
styling=request.styling,
|
||||
seed=request.seed,
|
||||
num_images=request.num_images,
|
||||
)
|
||||
return await _submit_and_respond(result, sync, {'model': 'mystic'})
|
||||
|
||||
|
||||
@router.post('/flux-dev', response_model=TaskResponse)
|
||||
async def generate_flux_dev(request: FluxDevRequest, sync: bool = Query(False)):
|
||||
result = await freepik_client.generate_flux_dev(
|
||||
prompt=request.prompt,
|
||||
image=request.image,
|
||||
guidance_scale=request.guidance_scale,
|
||||
num_images=request.num_images,
|
||||
seed=request.seed,
|
||||
)
|
||||
return await _submit_and_respond(result, sync, {'model': 'flux-dev'})
|
||||
|
||||
|
||||
@router.post('/flux-pro', response_model=TaskResponse)
|
||||
async def generate_flux_pro(request: FluxProRequest, sync: bool = Query(False)):
|
||||
result = await freepik_client.generate_flux_pro(
|
||||
prompt=request.prompt,
|
||||
image=request.image,
|
||||
guidance_scale=request.guidance_scale,
|
||||
seed=request.seed,
|
||||
)
|
||||
return await _submit_and_respond(result, sync, {'model': 'flux-pro'})
|
||||
|
||||
|
||||
@router.post('/seedream', response_model=TaskResponse)
|
||||
async def generate_seedream(request: SeedreamRequest, sync: bool = Query(False)):
|
||||
result = await freepik_client.generate_seedream(
|
||||
prompt=request.prompt,
|
||||
image=request.image,
|
||||
aspect_ratio=request.aspect_ratio,
|
||||
num_images=request.num_images,
|
||||
seed=request.seed,
|
||||
)
|
||||
return await _submit_and_respond(result, sync, {'model': 'seedream'})
|
||||
31
app/routers/system.py
Normal file
31
app/routers/system.py
Normal file
@@ -0,0 +1,31 @@
|
||||
import logging
|
||||
import os
|
||||
|
||||
import psutil
|
||||
from fastapi import APIRouter
|
||||
|
||||
from app.schemas.system import HealthResponse, SystemInfoResponse
|
||||
from app.services import freepik_client, task_tracker
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
router = APIRouter(prefix='/api/v1', tags=['system'])
|
||||
|
||||
|
||||
@router.get('/health', response_model=HealthResponse)
|
||||
async def health_check():
|
||||
return HealthResponse()
|
||||
|
||||
|
||||
@router.get('/system', response_model=SystemInfoResponse)
|
||||
async def system_info():
|
||||
api_key_valid = await freepik_client.check_api_key()
|
||||
mem = psutil.virtual_memory()
|
||||
|
||||
return SystemInfoResponse(
|
||||
api_key_valid=api_key_valid,
|
||||
active_tasks=task_tracker.active_count(),
|
||||
cpu_count=os.cpu_count() or 1,
|
||||
memory_total=mem.total,
|
||||
memory_available=mem.available,
|
||||
)
|
||||
75
app/routers/tasks.py
Normal file
75
app/routers/tasks.py
Normal file
@@ -0,0 +1,75 @@
|
||||
import logging
|
||||
from typing import Optional
|
||||
|
||||
from fastapi import APIRouter, HTTPException, Query
|
||||
from fastapi.responses import FileResponse
|
||||
|
||||
from app.schemas.common import TaskDetail, TaskListResponse, TaskStatus
|
||||
from app.services import task_tracker
|
||||
from app.services.file_manager import get_result_path
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
router = APIRouter(prefix='/api/v1/tasks', tags=['tasks'])
|
||||
|
||||
|
||||
@router.get('', response_model=TaskListResponse)
|
||||
async def list_tasks(
|
||||
status: Optional[TaskStatus] = Query(None),
|
||||
limit: int = Query(20, ge=1, le=100),
|
||||
offset: int = Query(0, ge=0),
|
||||
):
|
||||
tasks, total = task_tracker.list_tasks(status=status, limit=limit, offset=offset)
|
||||
return TaskListResponse(
|
||||
tasks=[_to_detail(t) for t in tasks],
|
||||
total=total,
|
||||
)
|
||||
|
||||
|
||||
@router.get('/{task_id}', response_model=TaskDetail)
|
||||
async def get_task(task_id: str):
|
||||
task = task_tracker.get_task(task_id)
|
||||
if not task:
|
||||
raise HTTPException(status_code=404, detail='Task not found')
|
||||
return _to_detail(task)
|
||||
|
||||
|
||||
@router.get('/{task_id}/result')
|
||||
async def get_task_result(task_id: str):
|
||||
task = task_tracker.get_task(task_id)
|
||||
if not task:
|
||||
raise HTTPException(status_code=404, detail='Task not found')
|
||||
if task['status'] != TaskStatus.completed:
|
||||
raise HTTPException(status_code=409, detail=f'Task status is {task["status"].value}')
|
||||
|
||||
path = task.get('local_path') or get_result_path(task_id)
|
||||
if not path:
|
||||
# Fall back to redirect if we have a remote URL
|
||||
if task.get('result_url'):
|
||||
from fastapi.responses import RedirectResponse
|
||||
return RedirectResponse(url=task['result_url'])
|
||||
raise HTTPException(status_code=404, detail='Result file not found')
|
||||
|
||||
return FileResponse(path)
|
||||
|
||||
|
||||
@router.delete('/{task_id}')
|
||||
async def delete_task(task_id: str):
|
||||
if not task_tracker.delete_task(task_id):
|
||||
raise HTTPException(status_code=404, detail='Task not found')
|
||||
return {'status': 'deleted', 'task_id': task_id}
|
||||
|
||||
|
||||
def _to_detail(task: dict) -> TaskDetail:
|
||||
result_url = None
|
||||
if task['status'] == TaskStatus.completed:
|
||||
result_url = f'/api/v1/tasks/{task["task_id"]}/result'
|
||||
return TaskDetail(
|
||||
task_id=task['task_id'],
|
||||
status=task['status'],
|
||||
created_at=task['created_at'],
|
||||
updated_at=task['updated_at'],
|
||||
progress=task.get('progress'),
|
||||
result_url=result_url,
|
||||
error=task.get('error'),
|
||||
)
|
||||
89
app/routers/utilities.py
Normal file
89
app/routers/utilities.py
Normal file
@@ -0,0 +1,89 @@
|
||||
import base64
|
||||
import logging
|
||||
from datetime import datetime, timezone
|
||||
|
||||
from fastapi import APIRouter, HTTPException
|
||||
|
||||
from app.schemas.common import TaskResponse, TaskStatus
|
||||
from app.schemas.utilities import ClassificationResponse, IconRequest
|
||||
from app.services import freepik_client, task_tracker
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
router = APIRouter(prefix='/api/v1/util', tags=['utilities'])
|
||||
|
||||
|
||||
@router.post('/remove-background')
|
||||
async def remove_background(request: dict):
|
||||
"""Remove background from an image. Returns processed image as base64."""
|
||||
image = request.get('image')
|
||||
if not image:
|
||||
raise HTTPException(status_code=400, detail='image field is required')
|
||||
|
||||
result_bytes = await freepik_client.remove_background(image)
|
||||
return {
|
||||
'image': base64.b64encode(result_bytes).decode(),
|
||||
'content_type': 'image/png',
|
||||
}
|
||||
|
||||
|
||||
@router.post('/classify', response_model=ClassificationResponse)
|
||||
async def classify_image(request: dict):
|
||||
"""Classify whether an image is AI-generated."""
|
||||
image = request.get('image')
|
||||
if not image:
|
||||
raise HTTPException(status_code=400, detail='image field is required')
|
||||
|
||||
result = await freepik_client.classify_image(image)
|
||||
data = result.get('data', result)
|
||||
return ClassificationResponse(
|
||||
is_ai_generated=data.get('is_ai_generated', False),
|
||||
ai_probability=data.get('ai_probability', 0.0),
|
||||
human_probability=data.get('human_probability', 0.0),
|
||||
)
|
||||
|
||||
|
||||
@router.post('/audio-isolate', response_model=TaskResponse)
|
||||
async def audio_isolate(request: dict):
|
||||
"""Isolate audio tracks from an audio file."""
|
||||
audio = request.get('audio')
|
||||
if not audio:
|
||||
raise HTTPException(status_code=400, detail='audio field is required')
|
||||
|
||||
result = await freepik_client.isolate_audio(audio)
|
||||
data = result.get('data', result)
|
||||
freepik_task_id = str(data.get('task_id') or data.get('id', ''))
|
||||
if not freepik_task_id:
|
||||
raise HTTPException(status_code=502, detail='No task_id in Freepik response')
|
||||
|
||||
internal_id = task_tracker.submit(freepik_task_id, {'operation': 'audio-isolate'})
|
||||
return TaskResponse(
|
||||
task_id=internal_id,
|
||||
status=TaskStatus.pending,
|
||||
created_at=datetime.now(timezone.utc),
|
||||
)
|
||||
|
||||
|
||||
# Icon generation lives under /api/v1/generate/ but is simple enough to keep here
|
||||
icon_router = APIRouter(prefix='/api/v1/generate', tags=['utilities'])
|
||||
|
||||
|
||||
@icon_router.post('/icon', response_model=TaskResponse)
|
||||
async def generate_icon(request: IconRequest):
|
||||
result = await freepik_client.generate_icon(
|
||||
prompt=request.prompt,
|
||||
color=request.color,
|
||||
shape=request.shape,
|
||||
style=request.style,
|
||||
)
|
||||
data = result.get('data', result)
|
||||
freepik_task_id = str(data.get('task_id') or data.get('id', ''))
|
||||
if not freepik_task_id:
|
||||
raise HTTPException(status_code=502, detail='No task_id in Freepik response')
|
||||
|
||||
internal_id = task_tracker.submit(freepik_task_id, {'operation': 'icon'})
|
||||
return TaskResponse(
|
||||
task_id=internal_id,
|
||||
status=TaskStatus.pending,
|
||||
created_at=datetime.now(timezone.utc),
|
||||
)
|
||||
70
app/routers/video_generation.py
Normal file
70
app/routers/video_generation.py
Normal file
@@ -0,0 +1,70 @@
|
||||
import logging
|
||||
from datetime import datetime, timezone
|
||||
|
||||
from fastapi import APIRouter, HTTPException
|
||||
|
||||
from app.schemas.common import TaskResponse, TaskStatus
|
||||
from app.schemas.video_generation import KlingRequest, MinimaxRequest, SeedanceRequest
|
||||
from app.services import freepik_client, task_tracker
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
router = APIRouter(prefix='/api/v1/generate/video', tags=['video-generation'])
|
||||
|
||||
|
||||
def _extract_task_id(result: dict) -> str:
|
||||
data = result.get('data', result)
|
||||
freepik_task_id = str(data.get('task_id') or data.get('id', ''))
|
||||
if not freepik_task_id:
|
||||
raise HTTPException(status_code=502, detail='No task_id in Freepik response')
|
||||
return freepik_task_id
|
||||
|
||||
|
||||
@router.post('/kling', response_model=TaskResponse)
|
||||
async def generate_video_kling(request: KlingRequest):
|
||||
result = await freepik_client.generate_video_kling(
|
||||
image=request.image,
|
||||
prompt=request.prompt,
|
||||
duration=request.duration,
|
||||
aspect_ratio=request.aspect_ratio,
|
||||
)
|
||||
freepik_id = _extract_task_id(result)
|
||||
internal_id = task_tracker.submit(freepik_id, {'model': 'kling'})
|
||||
return TaskResponse(
|
||||
task_id=internal_id,
|
||||
status=TaskStatus.pending,
|
||||
created_at=datetime.now(timezone.utc),
|
||||
)
|
||||
|
||||
|
||||
@router.post('/minimax', response_model=TaskResponse)
|
||||
async def generate_video_minimax(request: MinimaxRequest):
|
||||
result = await freepik_client.generate_video_minimax(
|
||||
prompt=request.prompt,
|
||||
first_frame_image=request.first_frame_image,
|
||||
subject_reference=request.subject_reference,
|
||||
)
|
||||
freepik_id = _extract_task_id(result)
|
||||
internal_id = task_tracker.submit(freepik_id, {'model': 'minimax'})
|
||||
return TaskResponse(
|
||||
task_id=internal_id,
|
||||
status=TaskStatus.pending,
|
||||
created_at=datetime.now(timezone.utc),
|
||||
)
|
||||
|
||||
|
||||
@router.post('/seedance', response_model=TaskResponse)
|
||||
async def generate_video_seedance(request: SeedanceRequest):
|
||||
result = await freepik_client.generate_video_seedance(
|
||||
prompt=request.prompt,
|
||||
image=request.image,
|
||||
duration=request.duration,
|
||||
resolution=request.resolution,
|
||||
)
|
||||
freepik_id = _extract_task_id(result)
|
||||
internal_id = task_tracker.submit(freepik_id, {'model': 'seedance'})
|
||||
return TaskResponse(
|
||||
task_id=internal_id,
|
||||
status=TaskStatus.pending,
|
||||
created_at=datetime.now(timezone.utc),
|
||||
)
|
||||
0
app/schemas/__init__.py
Normal file
0
app/schemas/__init__.py
Normal file
39
app/schemas/common.py
Normal file
39
app/schemas/common.py
Normal file
@@ -0,0 +1,39 @@
|
||||
from datetime import datetime
|
||||
from enum import Enum
|
||||
from typing import Optional
|
||||
|
||||
from pydantic import BaseModel
|
||||
|
||||
|
||||
class TaskStatus(str, Enum):
|
||||
pending = 'pending'
|
||||
processing = 'processing'
|
||||
completed = 'completed'
|
||||
failed = 'failed'
|
||||
|
||||
|
||||
class TaskResponse(BaseModel):
|
||||
task_id: str
|
||||
status: TaskStatus
|
||||
created_at: datetime
|
||||
|
||||
|
||||
class TaskDetail(BaseModel):
|
||||
task_id: str
|
||||
status: TaskStatus
|
||||
created_at: datetime
|
||||
updated_at: datetime
|
||||
progress: Optional[float] = None
|
||||
result_url: Optional[str] = None
|
||||
error: Optional[str] = None
|
||||
|
||||
|
||||
class TaskListResponse(BaseModel):
|
||||
tasks: list[TaskDetail]
|
||||
total: int
|
||||
|
||||
|
||||
class ErrorResponse(BaseModel):
|
||||
error: str
|
||||
detail: Optional[str] = None
|
||||
status_code: int
|
||||
41
app/schemas/image_editing.py
Normal file
41
app/schemas/image_editing.py
Normal file
@@ -0,0 +1,41 @@
|
||||
from typing import Optional
|
||||
|
||||
from pydantic import BaseModel, Field
|
||||
|
||||
|
||||
class UpscaleCreativeRequest(BaseModel):
|
||||
image: str = Field(..., description='Base64-encoded image')
|
||||
prompt: Optional[str] = None
|
||||
scale: Optional[int] = Field(None, ge=2, le=4)
|
||||
creativity: Optional[float] = Field(None, ge=0.0, le=1.0)
|
||||
resemblance: Optional[float] = Field(None, ge=0.0, le=1.0)
|
||||
|
||||
|
||||
class UpscalePrecisionRequest(BaseModel):
|
||||
image: str = Field(..., description='Base64-encoded image')
|
||||
scale: Optional[int] = Field(None, ge=2, le=4)
|
||||
|
||||
|
||||
class RelightRequest(BaseModel):
|
||||
image: str = Field(..., description='Base64-encoded image')
|
||||
prompt: Optional[str] = None
|
||||
light_source: Optional[str] = None
|
||||
intensity: Optional[float] = Field(None, ge=0.0, le=1.0)
|
||||
|
||||
|
||||
class StyleTransferRequest(BaseModel):
|
||||
image: str = Field(..., description='Base64-encoded image')
|
||||
style_reference: str = Field(..., description='Base64-encoded style reference image')
|
||||
strength: Optional[float] = Field(None, ge=0.0, le=1.0)
|
||||
|
||||
|
||||
class ExpandRequest(BaseModel):
|
||||
image: str = Field(..., description='Base64-encoded image')
|
||||
prompt: Optional[str] = None
|
||||
direction: Optional[str] = Field(None, description='Expansion direction')
|
||||
|
||||
|
||||
class InpaintRequest(BaseModel):
|
||||
image: str = Field(..., description='Base64-encoded image')
|
||||
mask: str = Field(..., description='Base64-encoded mask image')
|
||||
prompt: str = Field(..., min_length=1)
|
||||
35
app/schemas/image_generation.py
Normal file
35
app/schemas/image_generation.py
Normal file
@@ -0,0 +1,35 @@
|
||||
from typing import Optional
|
||||
|
||||
from pydantic import BaseModel, Field
|
||||
|
||||
|
||||
class MysticRequest(BaseModel):
|
||||
prompt: str = Field(..., min_length=1, max_length=4000)
|
||||
negative_prompt: Optional[str] = None
|
||||
resolution: Optional[str] = None
|
||||
styling: Optional[dict] = None
|
||||
seed: Optional[int] = None
|
||||
num_images: Optional[int] = Field(None, ge=1, le=4)
|
||||
|
||||
|
||||
class FluxDevRequest(BaseModel):
|
||||
prompt: str = Field(..., min_length=1, max_length=4000)
|
||||
image: Optional[str] = Field(None, description='Base64-encoded image for img2img')
|
||||
guidance_scale: Optional[float] = Field(None, ge=1.0, le=20.0)
|
||||
num_images: Optional[int] = Field(None, ge=1, le=4)
|
||||
seed: Optional[int] = None
|
||||
|
||||
|
||||
class FluxProRequest(BaseModel):
|
||||
prompt: str = Field(..., min_length=1, max_length=4000)
|
||||
image: Optional[str] = Field(None, description='Base64-encoded image for img2img')
|
||||
guidance_scale: Optional[float] = Field(None, ge=1.0, le=20.0)
|
||||
seed: Optional[int] = None
|
||||
|
||||
|
||||
class SeedreamRequest(BaseModel):
|
||||
prompt: str = Field(..., min_length=1, max_length=4000)
|
||||
image: Optional[str] = Field(None, description='Base64-encoded image for img2img')
|
||||
aspect_ratio: Optional[str] = None
|
||||
num_images: Optional[int] = Field(None, ge=1, le=4)
|
||||
seed: Optional[int] = None
|
||||
13
app/schemas/system.py
Normal file
13
app/schemas/system.py
Normal file
@@ -0,0 +1,13 @@
|
||||
from pydantic import BaseModel
|
||||
|
||||
|
||||
class HealthResponse(BaseModel):
|
||||
status: str = 'ok'
|
||||
|
||||
|
||||
class SystemInfoResponse(BaseModel):
|
||||
api_key_valid: bool
|
||||
active_tasks: int
|
||||
cpu_count: int
|
||||
memory_total: int
|
||||
memory_available: int
|
||||
16
app/schemas/utilities.py
Normal file
16
app/schemas/utilities.py
Normal file
@@ -0,0 +1,16 @@
|
||||
from typing import Optional
|
||||
|
||||
from pydantic import BaseModel, Field
|
||||
|
||||
|
||||
class ClassificationResponse(BaseModel):
|
||||
is_ai_generated: bool
|
||||
ai_probability: float
|
||||
human_probability: float
|
||||
|
||||
|
||||
class IconRequest(BaseModel):
|
||||
prompt: str = Field(..., min_length=1, max_length=4000)
|
||||
color: Optional[str] = None
|
||||
shape: Optional[str] = None
|
||||
style: Optional[str] = None
|
||||
23
app/schemas/video_generation.py
Normal file
23
app/schemas/video_generation.py
Normal file
@@ -0,0 +1,23 @@
|
||||
from typing import Optional
|
||||
|
||||
from pydantic import BaseModel, Field
|
||||
|
||||
|
||||
class KlingRequest(BaseModel):
|
||||
image: str = Field(..., description='Base64-encoded image')
|
||||
prompt: Optional[str] = None
|
||||
duration: Optional[str] = Field(None, description='5 or 10 seconds')
|
||||
aspect_ratio: Optional[str] = None
|
||||
|
||||
|
||||
class MinimaxRequest(BaseModel):
|
||||
prompt: str = Field(..., min_length=1, max_length=4000)
|
||||
first_frame_image: Optional[str] = Field(None, description='Base64-encoded image')
|
||||
subject_reference: Optional[str] = Field(None, description='Base64-encoded reference image')
|
||||
|
||||
|
||||
class SeedanceRequest(BaseModel):
|
||||
prompt: str = Field(..., min_length=1, max_length=4000)
|
||||
image: Optional[str] = Field(None, description='Base64-encoded image')
|
||||
duration: Optional[str] = None
|
||||
resolution: Optional[str] = None
|
||||
0
app/services/__init__.py
Normal file
0
app/services/__init__.py
Normal file
67
app/services/file_manager.py
Normal file
67
app/services/file_manager.py
Normal file
@@ -0,0 +1,67 @@
|
||||
import logging
|
||||
import os
|
||||
import time
|
||||
from pathlib import Path
|
||||
|
||||
import httpx
|
||||
|
||||
from app.config import settings
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def ensure_directories():
|
||||
os.makedirs(settings.output_dir, exist_ok=True)
|
||||
os.makedirs(settings.temp_dir, exist_ok=True)
|
||||
|
||||
|
||||
async def download_result(task_id: str, url: str) -> str:
|
||||
"""Download a result file from a URL and store it locally."""
|
||||
task_dir = os.path.join(settings.output_dir, task_id)
|
||||
os.makedirs(task_dir, exist_ok=True)
|
||||
|
||||
# Derive filename from URL
|
||||
filename = url.rsplit('/', 1)[-1].split('?')[0] or 'result'
|
||||
if '.' not in filename:
|
||||
filename += '.png'
|
||||
|
||||
output_path = os.path.join(task_dir, filename)
|
||||
|
||||
async with httpx.AsyncClient(timeout=120.0) as client:
|
||||
response = await client.get(url)
|
||||
response.raise_for_status()
|
||||
with open(output_path, 'wb') as f:
|
||||
f.write(response.content)
|
||||
|
||||
logger.info(f'Downloaded result for {task_id}: {output_path}')
|
||||
return output_path
|
||||
|
||||
|
||||
def get_result_path(task_id: str) -> str | None:
|
||||
"""Get the path to a task's result file."""
|
||||
task_dir = os.path.join(settings.output_dir, task_id)
|
||||
if not os.path.isdir(task_dir):
|
||||
return None
|
||||
files = os.listdir(task_dir)
|
||||
if not files:
|
||||
return None
|
||||
return os.path.join(task_dir, files[0])
|
||||
|
||||
|
||||
def cleanup_old_outputs():
|
||||
"""Remove output directories older than auto_cleanup_hours."""
|
||||
cutoff = time.time() - (settings.auto_cleanup_hours * 3600)
|
||||
output_dir = Path(settings.output_dir)
|
||||
if not output_dir.exists():
|
||||
return
|
||||
|
||||
removed = 0
|
||||
for task_dir in output_dir.iterdir():
|
||||
if task_dir.is_dir() and task_dir.stat().st_mtime < cutoff:
|
||||
for f in task_dir.iterdir():
|
||||
f.unlink()
|
||||
task_dir.rmdir()
|
||||
removed += 1
|
||||
|
||||
if removed:
|
||||
logger.info(f'Cleaned up {removed} old output directories')
|
||||
377
app/services/freepik_client.py
Normal file
377
app/services/freepik_client.py
Normal file
@@ -0,0 +1,377 @@
|
||||
import asyncio
|
||||
import logging
|
||||
from typing import Any, Optional
|
||||
|
||||
import httpx
|
||||
|
||||
from app.config import settings
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
_client: Optional[httpx.AsyncClient] = None
|
||||
|
||||
|
||||
def get_client() -> httpx.AsyncClient:
|
||||
if _client is None:
|
||||
raise RuntimeError('Freepik client not initialized')
|
||||
return _client
|
||||
|
||||
|
||||
async def create_client() -> httpx.AsyncClient:
|
||||
global _client
|
||||
_client = httpx.AsyncClient(
|
||||
base_url=settings.freepik_base_url,
|
||||
headers={
|
||||
'x-freepik-api-key': settings.freepik_api_key,
|
||||
'Accept': 'application/json',
|
||||
},
|
||||
timeout=httpx.Timeout(60.0, connect=10.0),
|
||||
)
|
||||
return _client
|
||||
|
||||
|
||||
async def close_client():
|
||||
global _client
|
||||
if _client:
|
||||
await _client.aclose()
|
||||
_client = None
|
||||
|
||||
|
||||
async def _request(
|
||||
method: str,
|
||||
path: str,
|
||||
*,
|
||||
json: Optional[dict] = None,
|
||||
data: Optional[dict] = None,
|
||||
files: Optional[dict] = None,
|
||||
params: Optional[dict] = None,
|
||||
max_retries: int = 3,
|
||||
) -> dict[str, Any]:
|
||||
"""Make an authenticated request to Freepik API with retry on 429."""
|
||||
client = get_client()
|
||||
for attempt in range(max_retries):
|
||||
kwargs: dict[str, Any] = {}
|
||||
if json is not None:
|
||||
kwargs['json'] = json
|
||||
if data is not None:
|
||||
kwargs['data'] = data
|
||||
if files is not None:
|
||||
kwargs['files'] = files
|
||||
if params is not None:
|
||||
kwargs['params'] = params
|
||||
|
||||
response = await client.request(method, path, **kwargs)
|
||||
|
||||
if response.status_code == 429:
|
||||
retry_after = int(response.headers.get('Retry-After', 2 ** attempt))
|
||||
logger.warning(f'Rate limited, retrying in {retry_after}s (attempt {attempt + 1})')
|
||||
await asyncio.sleep(retry_after)
|
||||
continue
|
||||
|
||||
response.raise_for_status()
|
||||
if response.headers.get('content-type', '').startswith('application/json'):
|
||||
return response.json()
|
||||
return {'raw': response.content}
|
||||
|
||||
raise httpx.HTTPStatusError(
|
||||
'Rate limit exceeded after retries',
|
||||
request=response.request,
|
||||
response=response,
|
||||
)
|
||||
|
||||
|
||||
async def _request_raw(
|
||||
method: str,
|
||||
path: str,
|
||||
*,
|
||||
json: Optional[dict] = None,
|
||||
data: Optional[dict] = None,
|
||||
files: Optional[dict] = None,
|
||||
max_retries: int = 3,
|
||||
) -> bytes:
|
||||
"""Make a request and return raw bytes (for binary responses)."""
|
||||
client = get_client()
|
||||
for attempt in range(max_retries):
|
||||
kwargs: dict[str, Any] = {}
|
||||
if json is not None:
|
||||
kwargs['json'] = json
|
||||
if data is not None:
|
||||
kwargs['data'] = data
|
||||
if files is not None:
|
||||
kwargs['files'] = files
|
||||
|
||||
response = await client.request(method, path, **kwargs)
|
||||
|
||||
if response.status_code == 429:
|
||||
retry_after = int(response.headers.get('Retry-After', 2 ** attempt))
|
||||
logger.warning(f'Rate limited, retrying in {retry_after}s (attempt {attempt + 1})')
|
||||
await asyncio.sleep(retry_after)
|
||||
continue
|
||||
|
||||
response.raise_for_status()
|
||||
return response.content
|
||||
|
||||
raise httpx.HTTPStatusError(
|
||||
'Rate limit exceeded after retries',
|
||||
request=response.request,
|
||||
response=response,
|
||||
)
|
||||
|
||||
|
||||
def _strip_none(d: dict) -> dict:
|
||||
"""Remove None values from a dict for clean API payloads."""
|
||||
return {k: v for k, v in d.items() if v is not None}
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Image generation
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
async def generate_mystic(
|
||||
prompt: str,
|
||||
negative_prompt: Optional[str] = None,
|
||||
resolution: Optional[str] = None,
|
||||
styling: Optional[dict] = None,
|
||||
seed: Optional[int] = None,
|
||||
num_images: Optional[int] = None,
|
||||
) -> dict:
|
||||
payload = _strip_none({
|
||||
'prompt': prompt,
|
||||
'negative_prompt': negative_prompt,
|
||||
'resolution': resolution,
|
||||
'styling': styling,
|
||||
'seed': seed,
|
||||
'num_images': num_images,
|
||||
})
|
||||
return await _request('POST', '/v1/ai/mystic', json=payload)
|
||||
|
||||
|
||||
async def generate_flux_dev(
|
||||
prompt: str,
|
||||
image: Optional[str] = None,
|
||||
guidance_scale: Optional[float] = None,
|
||||
num_images: Optional[int] = None,
|
||||
seed: Optional[int] = None,
|
||||
) -> dict:
|
||||
payload = _strip_none({
|
||||
'prompt': prompt,
|
||||
'image': image,
|
||||
'guidance_scale': guidance_scale,
|
||||
'num_images': num_images,
|
||||
'seed': seed,
|
||||
})
|
||||
return await _request('POST', '/v1/ai/text-to-image/flux-dev', json=payload)
|
||||
|
||||
|
||||
async def generate_flux_pro(
|
||||
prompt: str,
|
||||
image: Optional[str] = None,
|
||||
guidance_scale: Optional[float] = None,
|
||||
seed: Optional[int] = None,
|
||||
) -> dict:
|
||||
payload = _strip_none({
|
||||
'prompt': prompt,
|
||||
'image': image,
|
||||
'guidance_scale': guidance_scale,
|
||||
'seed': seed,
|
||||
})
|
||||
return await _request('POST', '/v1/ai/text-to-image/flux-pro-1.1', json=payload)
|
||||
|
||||
|
||||
async def generate_seedream(
|
||||
prompt: str,
|
||||
image: Optional[str] = None,
|
||||
aspect_ratio: Optional[str] = None,
|
||||
num_images: Optional[int] = None,
|
||||
seed: Optional[int] = None,
|
||||
) -> dict:
|
||||
payload = _strip_none({
|
||||
'prompt': prompt,
|
||||
'image': image,
|
||||
'aspect_ratio': aspect_ratio,
|
||||
'num_images': num_images,
|
||||
'seed': seed,
|
||||
})
|
||||
return await _request('POST', '/v1/ai/text-to-image/seedream', json=payload)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Video generation
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
async def generate_video_kling(
|
||||
image: str,
|
||||
prompt: Optional[str] = None,
|
||||
duration: Optional[str] = None,
|
||||
aspect_ratio: Optional[str] = None,
|
||||
) -> dict:
|
||||
payload = _strip_none({
|
||||
'image': image,
|
||||
'prompt': prompt,
|
||||
'duration': duration,
|
||||
'aspect_ratio': aspect_ratio,
|
||||
})
|
||||
return await _request('POST', '/v1/ai/image-to-video/kling', json=payload)
|
||||
|
||||
|
||||
async def generate_video_minimax(
|
||||
prompt: str,
|
||||
first_frame_image: Optional[str] = None,
|
||||
subject_reference: Optional[str] = None,
|
||||
) -> dict:
|
||||
payload = _strip_none({
|
||||
'prompt': prompt,
|
||||
'first_frame_image': first_frame_image,
|
||||
'subject_reference': subject_reference,
|
||||
})
|
||||
return await _request('POST', '/v1/ai/image-to-video/minimax', json=payload)
|
||||
|
||||
|
||||
async def generate_video_seedance(
|
||||
prompt: str,
|
||||
image: Optional[str] = None,
|
||||
duration: Optional[str] = None,
|
||||
resolution: Optional[str] = None,
|
||||
) -> dict:
|
||||
payload = _strip_none({
|
||||
'prompt': prompt,
|
||||
'image': image,
|
||||
'duration': duration,
|
||||
'resolution': resolution,
|
||||
})
|
||||
return await _request('POST', '/v1/ai/image-to-video/seedance', json=payload)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Image editing
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
async def upscale_creative(
|
||||
image: str,
|
||||
prompt: Optional[str] = None,
|
||||
scale: Optional[int] = None,
|
||||
creativity: Optional[float] = None,
|
||||
resemblance: Optional[float] = None,
|
||||
) -> dict:
|
||||
payload = _strip_none({
|
||||
'image': image,
|
||||
'prompt': prompt,
|
||||
'scale': scale,
|
||||
'creativity': creativity,
|
||||
'resemblance': resemblance,
|
||||
})
|
||||
return await _request('POST', '/v1/ai/upscale/creative', json=payload)
|
||||
|
||||
|
||||
async def upscale_precision(
|
||||
image: str,
|
||||
scale: Optional[int] = None,
|
||||
) -> dict:
|
||||
payload = _strip_none({
|
||||
'image': image,
|
||||
'scale': scale,
|
||||
})
|
||||
return await _request('POST', '/v1/ai/upscale/precision', json=payload)
|
||||
|
||||
|
||||
async def relight_image(
|
||||
image: str,
|
||||
prompt: Optional[str] = None,
|
||||
light_source: Optional[str] = None,
|
||||
intensity: Optional[float] = None,
|
||||
) -> dict:
|
||||
payload = _strip_none({
|
||||
'image': image,
|
||||
'prompt': prompt,
|
||||
'light_source': light_source,
|
||||
'intensity': intensity,
|
||||
})
|
||||
return await _request('POST', '/v1/ai/relight', json=payload)
|
||||
|
||||
|
||||
async def style_transfer(
|
||||
image: str,
|
||||
style_reference: str,
|
||||
strength: Optional[float] = None,
|
||||
) -> dict:
|
||||
payload = _strip_none({
|
||||
'image': image,
|
||||
'style_reference': style_reference,
|
||||
'strength': strength,
|
||||
})
|
||||
return await _request('POST', '/v1/ai/style-transfer', json=payload)
|
||||
|
||||
|
||||
async def expand_image(
|
||||
image: str,
|
||||
prompt: Optional[str] = None,
|
||||
direction: Optional[str] = None,
|
||||
) -> dict:
|
||||
payload = _strip_none({
|
||||
'image': image,
|
||||
'prompt': prompt,
|
||||
'direction': direction,
|
||||
})
|
||||
return await _request('POST', '/v1/ai/expand', json=payload)
|
||||
|
||||
|
||||
async def inpaint(
|
||||
image: str,
|
||||
mask: str,
|
||||
prompt: str,
|
||||
) -> dict:
|
||||
payload = {
|
||||
'image': image,
|
||||
'mask': mask,
|
||||
'prompt': prompt,
|
||||
}
|
||||
return await _request('POST', '/v1/ai/inpaint', json=payload)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Utilities
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
async def remove_background(image: str) -> bytes:
|
||||
return await _request_raw('POST', '/v1/ai/remove-background', json={'image': image})
|
||||
|
||||
|
||||
async def classify_image(image: str) -> dict:
|
||||
return await _request('POST', '/v1/ai/classifier', json={'image': image})
|
||||
|
||||
|
||||
async def isolate_audio(audio: str) -> dict:
|
||||
return await _request('POST', '/v1/ai/audio-isolate', json={'audio': audio})
|
||||
|
||||
|
||||
async def generate_icon(
|
||||
prompt: str,
|
||||
color: Optional[str] = None,
|
||||
shape: Optional[str] = None,
|
||||
style: Optional[str] = None,
|
||||
) -> dict:
|
||||
payload = _strip_none({
|
||||
'prompt': prompt,
|
||||
'color': color,
|
||||
'shape': shape,
|
||||
'style': style,
|
||||
})
|
||||
return await _request('POST', '/v1/ai/icon', json=payload)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Task management
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
async def get_task_status(task_id: str) -> dict:
|
||||
return await _request('GET', f'/v1/ai/tasks/{task_id}')
|
||||
|
||||
|
||||
async def check_api_key() -> bool:
|
||||
"""Verify the API key is valid by making a lightweight request."""
|
||||
try:
|
||||
client = get_client()
|
||||
response = await client.get('/v1/ai/tasks', params={'limit': 1})
|
||||
return response.status_code != 401
|
||||
except Exception:
|
||||
return False
|
||||
179
app/services/task_tracker.py
Normal file
179
app/services/task_tracker.py
Normal file
@@ -0,0 +1,179 @@
|
||||
import asyncio
|
||||
import logging
|
||||
import uuid
|
||||
from datetime import datetime, timezone
|
||||
from typing import Optional
|
||||
|
||||
from app.config import settings
|
||||
from app.schemas.common import TaskStatus
|
||||
from app.services import freepik_client
|
||||
from app.services.file_manager import download_result
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
_tasks: dict[str, dict] = {}
|
||||
_poll_tasks: dict[str, asyncio.Task] = {}
|
||||
|
||||
|
||||
def submit(freepik_task_id: str, metadata: Optional[dict] = None) -> str:
|
||||
"""Register a Freepik task and start background polling."""
|
||||
internal_id = str(uuid.uuid4())
|
||||
now = datetime.now(timezone.utc)
|
||||
_tasks[internal_id] = {
|
||||
'task_id': internal_id,
|
||||
'freepik_task_id': freepik_task_id,
|
||||
'status': TaskStatus.pending,
|
||||
'created_at': now,
|
||||
'updated_at': now,
|
||||
'progress': None,
|
||||
'result_url': None,
|
||||
'local_path': None,
|
||||
'error': None,
|
||||
'metadata': metadata or {},
|
||||
}
|
||||
_poll_tasks[internal_id] = asyncio.create_task(_poll_loop(internal_id))
|
||||
return internal_id
|
||||
|
||||
|
||||
def get_task(task_id: str) -> Optional[dict]:
|
||||
return _tasks.get(task_id)
|
||||
|
||||
|
||||
def list_tasks(
|
||||
status: Optional[TaskStatus] = None,
|
||||
limit: int = 20,
|
||||
offset: int = 0,
|
||||
) -> tuple[list[dict], int]:
|
||||
tasks = list(_tasks.values())
|
||||
if status:
|
||||
tasks = [t for t in tasks if t['status'] == status]
|
||||
tasks.sort(key=lambda t: t['created_at'], reverse=True)
|
||||
total = len(tasks)
|
||||
return tasks[offset:offset + limit], total
|
||||
|
||||
|
||||
def delete_task(task_id: str) -> bool:
|
||||
if task_id not in _tasks:
|
||||
return False
|
||||
# Cancel polling if active
|
||||
poll = _poll_tasks.pop(task_id, None)
|
||||
if poll and not poll.done():
|
||||
poll.cancel()
|
||||
_tasks.pop(task_id, None)
|
||||
return True
|
||||
|
||||
|
||||
def active_count() -> int:
|
||||
return sum(
|
||||
1 for t in _tasks.values()
|
||||
if t['status'] in (TaskStatus.pending, TaskStatus.processing)
|
||||
)
|
||||
|
||||
|
||||
async def _poll_loop(internal_id: str):
|
||||
"""Poll Freepik API until the task completes or times out."""
|
||||
task = _tasks.get(internal_id)
|
||||
if not task:
|
||||
return
|
||||
|
||||
freepik_id = task['freepik_task_id']
|
||||
elapsed = 0
|
||||
|
||||
try:
|
||||
while elapsed < settings.task_poll_timeout_seconds:
|
||||
await asyncio.sleep(settings.task_poll_interval_seconds)
|
||||
elapsed += settings.task_poll_interval_seconds
|
||||
|
||||
try:
|
||||
result = await freepik_client.get_task_status(freepik_id)
|
||||
except Exception as exc:
|
||||
logger.warning(f'Poll error for {internal_id}: {exc}')
|
||||
continue
|
||||
|
||||
data = result.get('data', result)
|
||||
fp_status = data.get('status', '')
|
||||
|
||||
task['updated_at'] = datetime.now(timezone.utc)
|
||||
|
||||
if fp_status in ('IN_PROGRESS', 'PROCESSING', 'processing'):
|
||||
task['status'] = TaskStatus.processing
|
||||
task['progress'] = data.get('progress')
|
||||
continue
|
||||
|
||||
if fp_status in ('COMPLETED', 'completed', 'done'):
|
||||
task['status'] = TaskStatus.completed
|
||||
task['progress'] = 1.0
|
||||
# Extract result URL from various response shapes
|
||||
result_url = (
|
||||
data.get('result_url')
|
||||
or data.get('output', {}).get('url')
|
||||
or _extract_first_url(data)
|
||||
)
|
||||
task['result_url'] = result_url
|
||||
if result_url:
|
||||
try:
|
||||
task['local_path'] = await download_result(
|
||||
internal_id, result_url
|
||||
)
|
||||
except Exception as exc:
|
||||
logger.error(f'Download failed for {internal_id}: {exc}')
|
||||
logger.info(f'Task {internal_id} completed')
|
||||
return
|
||||
|
||||
if fp_status in ('FAILED', 'failed', 'error'):
|
||||
task['status'] = TaskStatus.failed
|
||||
task['error'] = data.get('error', data.get('message', 'Unknown error'))
|
||||
logger.warning(f'Task {internal_id} failed: {task["error"]}')
|
||||
return
|
||||
|
||||
# Timeout
|
||||
task['status'] = TaskStatus.failed
|
||||
task['error'] = f'Polling timed out after {settings.task_poll_timeout_seconds}s'
|
||||
logger.warning(f'Task {internal_id} timed out')
|
||||
|
||||
except asyncio.CancelledError:
|
||||
logger.info(f'Polling cancelled for {internal_id}')
|
||||
except Exception as exc:
|
||||
task['status'] = TaskStatus.failed
|
||||
task['error'] = str(exc)
|
||||
logger.error(f'Unexpected error polling {internal_id}: {exc}')
|
||||
finally:
|
||||
_poll_tasks.pop(internal_id, None)
|
||||
|
||||
|
||||
def _extract_first_url(data: dict) -> Optional[str]:
|
||||
"""Try to extract the first URL from common Freepik response shapes."""
|
||||
# Some endpoints return {"data": {"images": [{"url": "..."}]}}
|
||||
for key in ('images', 'videos', 'results', 'outputs'):
|
||||
items = data.get(key, [])
|
||||
if isinstance(items, list) and items:
|
||||
first = items[0]
|
||||
if isinstance(first, dict) and 'url' in first:
|
||||
return first['url']
|
||||
if isinstance(first, str) and first.startswith('http'):
|
||||
return first
|
||||
# Direct URL field
|
||||
if 'url' in data and isinstance(data['url'], str):
|
||||
return data['url']
|
||||
return None
|
||||
|
||||
|
||||
def handle_webhook_completion(freepik_task_id: str, result_data: dict):
|
||||
"""Called when a webhook notification arrives for a completed task."""
|
||||
for task in _tasks.values():
|
||||
if task['freepik_task_id'] == freepik_task_id:
|
||||
task['status'] = TaskStatus.completed
|
||||
task['progress'] = 1.0
|
||||
task['updated_at'] = datetime.now(timezone.utc)
|
||||
result_url = (
|
||||
result_data.get('result_url')
|
||||
or result_data.get('output', {}).get('url')
|
||||
or _extract_first_url(result_data)
|
||||
)
|
||||
task['result_url'] = result_url
|
||||
# Cancel polling since webhook already notified us
|
||||
poll = _poll_tasks.pop(task['task_id'], None)
|
||||
if poll and not poll.done():
|
||||
poll.cancel()
|
||||
logger.info(f'Task {task["task_id"]} completed via webhook')
|
||||
break
|
||||
39
app/services/webhook.py
Normal file
39
app/services/webhook.py
Normal file
@@ -0,0 +1,39 @@
|
||||
import hashlib
|
||||
import hmac
|
||||
import logging
|
||||
|
||||
from fastapi import APIRouter, Header, HTTPException, Request
|
||||
|
||||
from app.config import settings
|
||||
from app.services.task_tracker import handle_webhook_completion
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
router = APIRouter(prefix='/api/v1', tags=['webhooks'])
|
||||
|
||||
|
||||
@router.post('/webhooks/task-complete')
|
||||
async def webhook_task_complete(
|
||||
request: Request,
|
||||
x_webhook_signature: str | None = Header(None),
|
||||
):
|
||||
body = await request.body()
|
||||
|
||||
if settings.webhook_secret:
|
||||
if not x_webhook_signature:
|
||||
raise HTTPException(status_code=401, detail='Missing webhook signature')
|
||||
expected = hmac.new(
|
||||
settings.webhook_secret.encode(),
|
||||
body,
|
||||
hashlib.sha256,
|
||||
).hexdigest()
|
||||
if not hmac.compare_digest(expected, x_webhook_signature):
|
||||
raise HTTPException(status_code=401, detail='Invalid webhook signature')
|
||||
|
||||
data = await request.json()
|
||||
task_id = data.get('task_id') or data.get('id')
|
||||
if not task_id:
|
||||
raise HTTPException(status_code=400, detail='Missing task_id in webhook payload')
|
||||
|
||||
handle_webhook_completion(str(task_id), data)
|
||||
return {'status': 'ok'}
|
||||
Reference in New Issue
Block a user