chore: import upstream snapshot with attribution

invokeai/app/services/external_generation/external_generation_default.py

@@ -0,0 +1,369 @@
+from __future__ import annotations
+
+import dataclasses
+import time
+from logging import Logger
+from typing import TYPE_CHECKING
+
+from PIL import Image
+from PIL.Image import Image as PILImageType
+
+from invokeai.app.services.external_generation.errors import (
+    ExternalProviderCapabilityError,
+    ExternalProviderNotConfiguredError,
+    ExternalProviderNotFoundError,
+    ExternalProviderRateLimitError,
+)
+from invokeai.app.services.external_generation.external_generation_base import (
+    ExternalGenerationServiceBase,
+    ExternalProvider,
+)
+from invokeai.app.services.external_generation.external_generation_common import (
+    ExternalGeneratedImage,
+    ExternalGenerationRequest,
+    ExternalGenerationResult,
+    ExternalProviderStatus,
+)
+from invokeai.backend.model_manager.configs.external_api import ExternalApiModelConfig, ExternalImageSize
+from invokeai.backend.model_manager.starter_models import STARTER_MODELS
+
+if TYPE_CHECKING:
+    from invokeai.app.services.model_records import ModelRecordServiceBase
+
+
+class ExternalGenerationService(ExternalGenerationServiceBase):
+    def __init__(
+        self,
+        providers: dict[str, ExternalProvider],
+        logger: Logger,
+        record_store: ModelRecordServiceBase | None = None,
+    ) -> None:
+        self._providers = providers
+        self._logger = logger
+        self._record_store = record_store
+
+    def generate(self, request: ExternalGenerationRequest) -> ExternalGenerationResult:
+        provider = self._providers.get(request.model.provider_id)
+        if provider is None:
+            raise ExternalProviderNotFoundError(f"No external provider registered for '{request.model.provider_id}'")
+
+        if not provider.is_configured():
+            raise ExternalProviderNotConfiguredError(f"Provider '{request.model.provider_id}' is missing credentials")
+
+        request = self._refresh_model_capabilities(request)
+        resize_to_original_inpaint_size = _get_resize_target_for_inpaint(request)
+        request = self._bucket_request(request)
+        request = self._drop_unsupported_capabilities(request)
+
+        self._validate_request(request)
+        result = self._generate_with_retry(provider, request)
+
+        if resize_to_original_inpaint_size is None:
+            return result
+
+        width, height = resize_to_original_inpaint_size
+        return _resize_result_images(result, width, height)
+
+    _MAX_RETRIES = 3
+    _DEFAULT_RETRY_DELAY = 10.0
+    _MAX_RETRY_DELAY = 60.0
+
+    def _generate_with_retry(
+        self, provider: ExternalProvider, request: ExternalGenerationRequest
+    ) -> ExternalGenerationResult:
+        for attempt in range(self._MAX_RETRIES):
+            try:
+                return provider.generate(request)
+            except ExternalProviderRateLimitError as exc:
+                if attempt == self._MAX_RETRIES - 1:
+                    raise
+                delay = min(exc.retry_after or self._DEFAULT_RETRY_DELAY, self._MAX_RETRY_DELAY)
+                self._logger.warning(
+                    "Rate limited by %s (attempt %d/%d), retrying in %.0fs",
+                    request.model.provider_id,
+                    attempt + 1,
+                    self._MAX_RETRIES,
+                    delay,
+                )
+                time.sleep(delay)
+        raise ExternalProviderRateLimitError("Rate limit exceeded after all retries")
+
+    def get_provider_statuses(self) -> dict[str, ExternalProviderStatus]:
+        return {provider_id: provider.get_status() for provider_id, provider in self._providers.items()}
+
+    def _validate_request(self, request: ExternalGenerationRequest) -> None:
+        capabilities = request.model.capabilities
+
+        self._logger.debug(
+            "Validating external request provider=%s model=%s mode=%s supported=%s",
+            request.model.provider_id,
+            request.model.provider_model_id,
+            request.mode,
+            capabilities.modes,
+        )
+
+        if request.mode not in capabilities.modes:
+            raise ExternalProviderCapabilityError(f"Mode '{request.mode}' is not supported by {request.model.name}")
+
+        if request.reference_images and not capabilities.supports_reference_images:
+            raise ExternalProviderCapabilityError(f"Reference images are not supported by {request.model.name}")
+
+        if capabilities.max_reference_images is not None:
+            if len(request.reference_images) > capabilities.max_reference_images:
+                raise ExternalProviderCapabilityError(
+                    f"{request.model.name} supports at most {capabilities.max_reference_images} reference images"
+                )
+
+        if capabilities.max_images_per_request is not None and request.num_images > capabilities.max_images_per_request:
+            raise ExternalProviderCapabilityError(
+                f"{request.model.name} supports at most {capabilities.max_images_per_request} images per request"
+            )
+
+        if capabilities.max_image_size is not None:
+            if request.width > capabilities.max_image_size.width or request.height > capabilities.max_image_size.height:
+                raise ExternalProviderCapabilityError(
+                    f"{request.model.name} supports a maximum size of {capabilities.max_image_size.width}x{capabilities.max_image_size.height}"
+                )
+
+        if capabilities.allowed_aspect_ratios:
+            aspect_ratio = _format_aspect_ratio(request.width, request.height)
+            if aspect_ratio not in capabilities.allowed_aspect_ratios:
+                size_ratio = None
+                if capabilities.aspect_ratio_sizes:
+                    size_ratio = _ratio_for_size(request.width, request.height, capabilities.aspect_ratio_sizes)
+                if size_ratio is None or size_ratio not in capabilities.allowed_aspect_ratios:
+                    ratio_label = size_ratio or aspect_ratio
+                    raise ExternalProviderCapabilityError(
+                        f"{request.model.name} does not support aspect ratio {ratio_label}"
+                    )
+
+        required_modes = capabilities.input_image_required_for or ["img2img", "inpaint"]
+        if request.mode in required_modes and request.init_image is None:
+            raise ExternalProviderCapabilityError(
+                f"Mode '{request.mode}' requires an init image for {request.model.name}"
+            )
+
+        if request.mode == "inpaint" and request.mask_image is None:
+            raise ExternalProviderCapabilityError(
+                f"Mode '{request.mode}' requires a mask image for {request.model.name}"
+            )
+
+    def _drop_unsupported_capabilities(self, request: ExternalGenerationRequest) -> ExternalGenerationRequest:
+        """Silently drop request fields the selected model does not support so workflow-editor runs don't fail
+        when users wire them in regardless."""
+        capabilities = request.model.capabilities
+        updates: dict[str, object] = {}
+
+        if request.seed is not None and not capabilities.supports_seed:
+            self._logger.debug(
+                "Dropping seed for %s: model does not support seed control",
+                request.model.name,
+            )
+            updates["seed"] = None
+
+        if updates:
+            return dataclasses.replace(request, **updates)
+        return request
+
+    def _refresh_model_capabilities(self, request: ExternalGenerationRequest) -> ExternalGenerationRequest:
+        if self._record_store is None:
+            return request
+
+        try:
+            record = self._record_store.get_model(request.model.key)
+        except Exception:
+            record = None
+
+        if not isinstance(record, ExternalApiModelConfig):
+            return request
+
+        if record.key != request.model.key:
+            return request
+
+        if record.provider_id != request.model.provider_id:
+            return request
+
+        if record.provider_model_id != request.model.provider_model_id:
+            return request
+
+        record = _apply_starter_overrides(record)
+
+        if record == request.model:
+            return request
+
+        return ExternalGenerationRequest(
+            model=record,
+            mode=request.mode,
+            prompt=request.prompt,
+            seed=request.seed,
+            num_images=request.num_images,
+            width=request.width,
+            height=request.height,
+            image_size=request.image_size,
+            init_image=request.init_image,
+            mask_image=request.mask_image,
+            reference_images=request.reference_images,
+            metadata=request.metadata,
+            provider_options=request.provider_options,
+        )
+
+    def _bucket_request(self, request: ExternalGenerationRequest) -> ExternalGenerationRequest:
+        capabilities = request.model.capabilities
+        if not capabilities.allowed_aspect_ratios:
+            return request
+
+        aspect_ratio = _format_aspect_ratio(request.width, request.height)
+        size = None
+        if capabilities.aspect_ratio_sizes:
+            size = capabilities.aspect_ratio_sizes.get(aspect_ratio)
+
+        if size is not None:
+            if request.width == size.width and request.height == size.height:
+                return request
+            return self._bucket_to_size(request, size.width, size.height, aspect_ratio)
+
+        if aspect_ratio in capabilities.allowed_aspect_ratios:
+            return request
+
+        if not capabilities.aspect_ratio_sizes:
+            return request
+
+        closest = _select_closest_ratio(
+            request.width,
+            request.height,
+            capabilities.allowed_aspect_ratios,
+        )
+        if closest is None:
+            return request
+
+        size = capabilities.aspect_ratio_sizes.get(closest)
+        if size is None:
+            return request
+
+        return self._bucket_to_size(request, size.width, size.height, closest)
+
+    def _bucket_to_size(
+        self,
+        request: ExternalGenerationRequest,
+        width: int,
+        height: int,
+        ratio: str,
+    ) -> ExternalGenerationRequest:
+        self._logger.info(
+            "Bucketing external request provider=%s model=%s %sx%s -> %sx%s (ratio %s)",
+            request.model.provider_id,
+            request.model.provider_model_id,
+            request.width,
+            request.height,
+            width,
+            height,
+            ratio,
+        )
+
+        return ExternalGenerationRequest(
+            model=request.model,
+            mode=request.mode,
+            prompt=request.prompt,
+            seed=request.seed,
+            num_images=request.num_images,
+            width=width,
+            height=height,
+            image_size=request.image_size,
+            init_image=_resize_image(request.init_image, width, height, "RGB"),
+            mask_image=_resize_image(request.mask_image, width, height, "L"),
+            reference_images=request.reference_images,
+            metadata=request.metadata,
+            provider_options=request.provider_options,
+        )
+
+
+def _format_aspect_ratio(width: int, height: int) -> str:
+    divisor = _gcd(width, height)
+    return f"{width // divisor}:{height // divisor}"
+
+
+def _select_closest_ratio(width: int, height: int, ratios: list[str]) -> str | None:
+    ratio = width / height
+    parsed: list[tuple[str, float]] = []
+    for value in ratios:
+        parsed_ratio = _parse_ratio(value)
+        if parsed_ratio is not None:
+            parsed.append((value, parsed_ratio))
+    if not parsed:
+        return None
+    return min(parsed, key=lambda item: abs(item[1] - ratio))[0]
+
+
+def _ratio_for_size(width: int, height: int, sizes: dict[str, ExternalImageSize]) -> str | None:
+    for ratio, size in sizes.items():
+        if size.width == width and size.height == height:
+            return ratio
+    return None
+
+
+def _parse_ratio(value: str) -> float | None:
+    if ":" not in value:
+        return None
+    left, right = value.split(":", 1)
+    try:
+        numerator = float(left)
+        denominator = float(right)
+    except ValueError:
+        return None
+    if denominator == 0:
+        return None
+    return numerator / denominator
+
+
+def _gcd(a: int, b: int) -> int:
+    while b:
+        a, b = b, a % b
+    return a
+
+
+def _resize_image(image: PILImageType | None, width: int, height: int, mode: str) -> PILImageType | None:
+    if image is None:
+        return None
+    if image.width == width and image.height == height:
+        return image
+    return image.convert(mode).resize((width, height), Image.Resampling.LANCZOS)
+
+
+def _get_resize_target_for_inpaint(request: ExternalGenerationRequest) -> tuple[int, int] | None:
+    if request.mode != "inpaint" or request.init_image is None:
+        return None
+    return request.init_image.width, request.init_image.height
+
+
+def _resize_result_images(result: ExternalGenerationResult, width: int, height: int) -> ExternalGenerationResult:
+    resized_images = [
+        ExternalGeneratedImage(
+            image=generated.image
+            if generated.image.width == width and generated.image.height == height
+            else generated.image.resize((width, height), Image.Resampling.LANCZOS),
+            seed=generated.seed,
+        )
+        for generated in result.images
+    ]
+    return ExternalGenerationResult(
+        images=resized_images,
+        seed_used=result.seed_used,
+        provider_request_id=result.provider_request_id,
+        provider_metadata=result.provider_metadata,
+        content_filters=result.content_filters,
+    )
+
+
+def _apply_starter_overrides(model: ExternalApiModelConfig) -> ExternalApiModelConfig:
+    source = model.source or f"external://{model.provider_id}/{model.provider_model_id}"
+    starter_match = next((starter for starter in STARTER_MODELS if starter.source == source), None)
+    if starter_match is None:
+        return model
+    updates: dict[str, object] = {}
+    if starter_match.capabilities is not None:
+        updates["capabilities"] = starter_match.capabilities
+    if starter_match.default_settings is not None:
+        updates["default_settings"] = starter_match.default_settings
+    if not updates:
+        return model
+    return model.model_copy(update=updates)