hao-ai-lab-FastVideo/fastvideo/entrypoints/openai/image_api.py

# Adapted from SGLang
# (https://github.com/sgl-project/sglang/blob/main/python/sglang/multimodal_gen/runtime/entrypoints/openai/image_api.py)

import asyncio
import base64
import os
import time

import aiofiles

from fastapi import (APIRouter, File, Form, HTTPException, Path, Query, UploadFile)
from fastapi.responses import FileResponse

from fastvideo.entrypoints.openai.state import (
    get_generator,
    get_output_dir,
)
from fastvideo.entrypoints.openai.protocol import (
    ImageGenerationsRequest,
    ImageResponse,
    ImageResponseData,
    generate_request_id,
)
from fastvideo.entrypoints.openai.stores import IMAGE_STORE
from fastvideo.entrypoints.openai.utils import (
    choose_image_ext,
    merge_image_input_list,
    parse_size,
    save_image_to_path,
)
from fastvideo.logger import init_logger

logger = init_logger(__name__)
router = APIRouter(prefix="/v1/images", tags=["images"])


def _build_generation_kwargs(
    request_id: str,
    prompt: str,
    n: int = 1,
    size: str | None = None,
    output_format: str | None = None,
    background: str | None = None,
    image_path: list[str] | None = None,
    seed: int | None = None,
    num_inference_steps: int | None = None,
    guidance_scale: float | None = None,
    true_cfg_scale: float | None = None,
    negative_prompt: str | None = None,
    enable_teacache: bool | None = None,
) -> dict:
    """Convert API request params to VideoGenerator.generate_video kwargs"""
    kwargs: dict = {"prompt": prompt}

    if size:
        w, h = parse_size(size)
        if w is not None and h is not None:
            kwargs["width"] = w
            kwargs["height"] = h

    ext = choose_image_ext(output_format, background)
    output_dir = os.path.join(get_output_dir(), "images")
    os.makedirs(output_dir, exist_ok=True)
    kwargs["output_path"] = os.path.join(output_dir, f"{request_id}.{ext}")

    # Image generation
    kwargs["num_frames"] = 1
    kwargs["save_video"] = True
    kwargs["num_videos_per_prompt"] = max(1, min(n, 10))

    if seed is not None:
        kwargs["seed"] = seed
    if num_inference_steps is not None:
        kwargs["num_inference_steps"] = num_inference_steps
    if guidance_scale is not None:
        kwargs["guidance_scale"] = guidance_scale
    if true_cfg_scale is not None:
        kwargs["true_cfg_scale"] = true_cfg_scale
    if negative_prompt is not None:
        kwargs["negative_prompt"] = negative_prompt
    if enable_teacache:
        kwargs["enable_teacache"] = True
    if image_path:
        kwargs["image_path"] = image_path[0] if len(image_path) == 1 else image_path

    return kwargs


@router.post("", response_model=ImageResponse)
async def generations(request: ImageGenerationsRequest):
    request_id = generate_request_id()
    generator = get_generator()
    loop = asyncio.get_running_loop()

    gen_kwargs = _build_generation_kwargs(
        request_id=request_id,
        prompt=request.prompt,
        n=request.n or 1,
        size=request.size,
        output_format=request.output_format,
        background=request.background,
        seed=request.seed,
        num_inference_steps=request.num_inference_steps,
        guidance_scale=request.guidance_scale,
        true_cfg_scale=request.true_cfg_scale,
        negative_prompt=request.negative_prompt,
        enable_teacache=request.enable_teacache,
    )

    start = time.perf_counter()
    try:
        await loop.run_in_executor(None, lambda: generator.generate_video(**gen_kwargs))
    except Exception as e:
        logger.error("Image generation failed: %s", e)
        raise HTTPException(status_code=500, detail=str(e)) from None
    elapsed = time.perf_counter() - start

    save_file_path = gen_kwargs["output_path"]

    resp_format = (request.response_format or "b64_json").lower()
    if resp_format == "b64_json":
        if not os.path.exists(save_file_path):
            raise HTTPException(status_code=500, detail="Image was not saved to disk")
        async with aiofiles.open(save_file_path, "rb") as f:
            b64_data = base64.b64encode(await f.read()).decode("utf-8")
        data = [ImageResponseData(b64_json=b64_data, revised_prompt=request.prompt)]
    elif resp_format == "url":
        data = [
            ImageResponseData(
                url=f"/v1/images/{request_id}/content",
                revised_prompt=request.prompt,
                file_path=os.path.abspath(save_file_path),
            )
        ]
    else:
        raise HTTPException(status_code=400, detail=f"response_format={resp_format} is not supported")

    await IMAGE_STORE.upsert(
        request_id,
        {
            "id": request_id,
            "created_at": int(time.time()),
            "file_path": save_file_path,
        },
    )

    return ImageResponse(
        id=request_id,
        data=data,
        inference_time_s=elapsed,
    )


@router.post("/edits", response_model=ImageResponse)
async def edits(
        image: list[UploadFile] | None = File(None),  # noqa: B008
        image_array: list[UploadFile] | None = File(  # noqa: B008
            None, alias="image[]"),
        url: list[str] | None = Form(None),  # noqa: B008
        url_array: list[str] | None = Form(None, alias="url[]"),  # noqa: B008
        prompt: str = Form(...),
        model: str | None = Form(None),
        n: int | None = Form(1),
        response_format: str | None = Form(None),
        size: str | None = Form(None),
        output_format: str | None = Form(None),
        background: str | None = Form("auto"),
        seed: int | None = Form(1024),
        negative_prompt: str | None = Form(None),
        guidance_scale: float | None = Form(None),
        true_cfg_scale: float | None = Form(None),
        num_inference_steps: int | None = Form(None),
        enable_teacache: bool | None = Form(False),
):
    request_id = generate_request_id()
    generator = get_generator()
    loop = asyncio.get_running_loop()

    images = image or image_array
    urls = url or url_array
    if (not images or len(images) == 0) and (not urls or len(urls) == 0):
        raise HTTPException(status_code=422, detail="Field 'image' or 'url' is required")

    # Save input images
    uploads_dir = os.path.join(get_output_dir(), "uploads")
    os.makedirs(uploads_dir, exist_ok=True)
    image_list = merge_image_input_list(images, urls)

    input_paths: list[str] = []
    try:
        for idx, img in enumerate(image_list):
            filename = getattr(img, "filename", f"image_{idx}")
            input_path = await save_image_to_path(img, os.path.join(uploads_dir, f"{request_id}_{idx}_{filename}"))
            input_paths.append(input_path)
    except Exception as e:
        raise HTTPException(status_code=400, detail=f"Failed to process image: {e}") from None

    gen_kwargs = _build_generation_kwargs(
        request_id=request_id,
        prompt=prompt,
        n=n or 1,
        size=size,
        output_format=output_format,
        background=background,
        image_path=input_paths,
        seed=seed,
        num_inference_steps=num_inference_steps,
        guidance_scale=guidance_scale,
        true_cfg_scale=true_cfg_scale,
        negative_prompt=negative_prompt,
        enable_teacache=enable_teacache,
    )

    start = time.perf_counter()
    try:
        await loop.run_in_executor(None, lambda: generator.generate_video(**gen_kwargs))
    except Exception as e:
        logger.error("Image edit failed: %s", e)
        raise HTTPException(status_code=500, detail=str(e)) from None
    elapsed = time.perf_counter() - start

    save_file_path = gen_kwargs["output_path"]

    resp_format = (response_format or "b64_json").lower()
    if resp_format == "b64_json":
        async with aiofiles.open(save_file_path, "rb") as f:
            b64_data = base64.b64encode(await f.read()).decode("utf-8")
        data = [
            ImageResponseData(
                b64_json=b64_data,
                revised_prompt=prompt,
                file_path=os.path.abspath(save_file_path),
            )
        ]
    else:
        data = [
            ImageResponseData(
                url=f"/v1/images/{request_id}/content",
                revised_prompt=prompt,
                file_path=os.path.abspath(save_file_path),
            )
        ]

    await IMAGE_STORE.upsert(
        request_id,
        {
            "id": request_id,
            "created_at": int(time.time()),
            "file_path": save_file_path,
        },
    )

    return ImageResponse(id=request_id, data=data, inference_time_s=elapsed)


@router.get("/{image_id}/content")
async def download_image_content(image_id: str = Path(...), variant: str | None = Query(None)):
    item = await IMAGE_STORE.get(image_id)
    if not item:
        raise HTTPException(status_code=404, detail="Image not found")

    file_path = item.get("file_path")
    if not file_path or not os.path.exists(file_path):
        raise HTTPException(status_code=404, detail="Image is still being generated")

    ext = os.path.splitext(file_path)[1].lower()
    media_type = {".png": "image/png", ".webp": "image/webp"}.get(ext, "image/jpeg")

    return FileResponse(path=file_path, media_type=media_type, filename=os.path.basename(file_path))