freepik-api/app/routers/image_editing.py

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
from app.services.freepik_client import TaskResult, _extract_task_id

logger = logging.getLogger(__name__)

router = APIRouter(prefix='/api/v1/edit', tags=['image-editing'])


async def _submit_and_respond(result: TaskResult, sync: bool, metadata: dict) -> TaskResponse | TaskDetail:
    freepik_task_id = _extract_task_id(result.data)
    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, result.status_path, 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_factor=request.scale_factor,
        creativity=request.creativity,
        resemblance=request.resemblance,
        optimized_for=request.optimized_for,
    )
    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_factor=request.scale_factor,
    )
    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,
        transfer_light_from_reference_image=request.transfer_light_from_reference_image,
        light_transfer_strength=request.light_transfer_strength,
    )
    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,
        reference_image=request.reference_image,
        prompt=request.prompt,
        style_strength=request.style_strength,
        structure_strength=request.structure_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,
    )
    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'})