invoke-ai--invokeai/invokeai/app/services/external_generation/providers/seedream.py

from __future__ import annotations

import requests

from invokeai.app.services.external_generation.errors import (
    ExternalProviderCapabilityError,
    ExternalProviderRateLimitError,
    ExternalProviderRequestError,
)
from invokeai.app.services.external_generation.external_generation_base import ExternalProvider
from invokeai.app.services.external_generation.external_generation_common import (
    ExternalGeneratedImage,
    ExternalGenerationRequest,
    ExternalGenerationResult,
)
from invokeai.app.services.external_generation.image_utils import decode_image_base64, encode_image_base64

_SEEDREAM_BATCH_PREFIXES = (
    "seedream-5",
    "seedream-4.5",
    "seedream-4.0",
    "seedream-4-5",
    "seedream-4-0",
    "seedream-5-0",
)

# Seedream batch endpoint accepts up to 15 total images counting both inputs (reference + init)
# and outputs combined. Hitting this only after the API call wastes a request and produces a
# confusing 400, so we enforce it locally for batch-capable models.
_SEEDREAM_BATCH_MAX_TOTAL_IMAGES = 15


class SeedreamProvider(ExternalProvider):
    provider_id = "seedream"

    def is_configured(self) -> bool:
        return bool(self._app_config.external_seedream_api_key)

    def generate(self, request: ExternalGenerationRequest) -> ExternalGenerationResult:
        api_key = self._app_config.external_seedream_api_key
        if not api_key:
            raise ExternalProviderRequestError("Seedream API key is not configured")

        base_url = (self._app_config.external_seedream_base_url or "https://ark.ap-southeast.bytepluses.com").rstrip(
            "/"
        )
        endpoint = f"{base_url}/api/v3/images/generations"
        headers = {"Authorization": f"Bearer {api_key}"}

        model_id = request.model.provider_model_id
        is_batch_model = any(model_id.startswith(prefix) for prefix in _SEEDREAM_BATCH_PREFIXES)

        if is_batch_model:
            input_image_count = len(request.reference_images) + (1 if request.init_image is not None else 0)
            total_images = input_image_count + request.num_images
            if total_images > _SEEDREAM_BATCH_MAX_TOTAL_IMAGES:
                raise ExternalProviderCapabilityError(
                    f"{request.model.name} supports at most {_SEEDREAM_BATCH_MAX_TOTAL_IMAGES} images total "
                    f"(reference + init + output), got {total_images}"
                )

        opts = request.provider_options or {}

        payload: dict[str, object] = {
            "model": model_id,
            "prompt": request.prompt,
            "size": f"{request.width}x{request.height}",
            "response_format": "b64_json",
            "watermark": opts.get("watermark", False),
        }

        if opts.get("optimize_prompt"):
            payload["optimize_prompt_options"] = {"optimize_prompt": True}

        # Seed and guidance_scale are only supported on 3.0 models
        if not is_batch_model and request.seed is not None and request.seed >= 0:
            payload["seed"] = request.seed
        if not is_batch_model and opts.get("guidance_scale") is not None:
            payload["guidance_scale"] = opts["guidance_scale"]

        # Batch generation for 4.x/5.x models
        if is_batch_model:
            if request.num_images > 1:
                payload["sequential_image_generation"] = "auto"
                payload["sequential_image_generation_options"] = {"max_images": request.num_images}
            else:
                payload["sequential_image_generation"] = "disabled"

        # Image input: init_image for img2img, reference images for 4.x
        images_b64: list[str] = []
        if request.init_image is not None:
            images_b64.append(f"data:image/png;base64,{encode_image_base64(request.init_image)}")
        for reference in request.reference_images:
            images_b64.append(f"data:image/png;base64,{encode_image_base64(reference.image)}")

        if images_b64:
            payload["image"] = images_b64 if len(images_b64) > 1 else images_b64[0]

        response = requests.post(endpoint, headers=headers, json=payload, timeout=120)

        if not response.ok:
            if response.status_code == 429:
                retry_after = _parse_retry_after(response.headers.get("retry-after"))
                raise ExternalProviderRateLimitError(
                    f"Seedream rate limit exceeded. {f'Retry after {retry_after:.0f}s.' if retry_after else 'Please try again later.'}",
                    retry_after=retry_after,
                )
            raise ExternalProviderRequestError(
                f"Seedream request failed with status {response.status_code}: {response.text}"
            )

        body = response.json()
        if not isinstance(body, dict):
            raise ExternalProviderRequestError("Seedream response payload was not a JSON object")

        generated_images: list[ExternalGeneratedImage] = []
        item_errors: list[dict[str, object]] = []
        data_items = body.get("data")
        if not isinstance(data_items, list):
            raise ExternalProviderRequestError("Seedream response payload missing image data")

        for item in data_items:
            if not isinstance(item, dict):
                continue
            # Items may be error objects for failed images in batch — collect rather than discard
            # so partial-failure causes (e.g., content filter) are visible to the caller.
            if "error" in item:
                error_payload = item["error"]
                item_errors.append(
                    error_payload if isinstance(error_payload, dict) else {"message": str(error_payload)}
                )
                continue
            encoded = item.get("b64_json")
            if not encoded:
                continue
            image = decode_image_base64(encoded)
            generated_images.append(ExternalGeneratedImage(image=image, seed=request.seed))

        if not generated_images:
            if item_errors:
                first = item_errors[0]
                message = first.get("message") if isinstance(first, dict) else None
                raise ExternalProviderRequestError(
                    f"Seedream returned no images. Provider reported: {message or item_errors}"
                )
            raise ExternalProviderRequestError("Seedream response contained no images")

        provider_metadata: dict[str, object] = {"model": model_id}
        if item_errors:
            provider_metadata["partial_failures"] = item_errors
            self._logger.warning(
                "Seedream returned %d image(s) with %d partial failure(s): %s",
                len(generated_images),
                len(item_errors),
                item_errors,
            )

        return ExternalGenerationResult(
            images=generated_images,
            seed_used=request.seed,
            provider_metadata=provider_metadata,
        )


def _parse_retry_after(value: str | None) -> float | None:
    if not value:
        return None
    try:
        return float(value)
    except ValueError:
        return None