from typing import TYPE_CHECKING, Any, ClassVar, Literal from invokeai.app.invocations.baseinvocation import BaseInvocation, BaseInvocationOutput, invocation from invokeai.app.invocations.fields import ( FieldDescriptions, ImageField, InputField, MetadataField, WithBoard, WithMetadata, ) from invokeai.app.invocations.model import ModelIdentifierField from invokeai.app.invocations.primitives import ImageCollectionOutput from invokeai.app.services.external_generation.external_generation_common import ( ExternalGenerationRequest, ExternalGenerationResult, ExternalReferenceImage, ) from invokeai.app.services.shared.invocation_context import InvocationContext from invokeai.backend.model_manager.configs.external_api import ExternalApiModelConfig, ExternalGenerationMode from invokeai.backend.model_manager.taxonomy import BaseModelType, ModelFormat, ModelType if TYPE_CHECKING: from invokeai.app.services.invocation_services import InvocationServices class BaseExternalImageGenerationInvocation(BaseInvocation, WithMetadata, WithBoard): """Generate images using an external provider.""" provider_id: ClassVar[str | None] = None model: ModelIdentifierField = InputField( description=FieldDescriptions.main_model, ui_model_base=[BaseModelType.External], ui_model_type=[ModelType.ExternalImageGenerator], ui_model_format=[ModelFormat.ExternalApi], ) mode: ExternalGenerationMode = InputField( default="txt2img", description="Generation mode. Not all modes are supported by every model; unsupported modes raise at runtime.", ) prompt: str = InputField(description="Prompt") seed: int | None = InputField(default=None, description=FieldDescriptions.seed) num_images: int = InputField(default=1, gt=0, description="Number of images to generate") width: int = InputField(default=1024, gt=0, description=FieldDescriptions.width) height: int = InputField(default=1024, gt=0, description=FieldDescriptions.height) image_size: str | None = InputField(default=None, description="Image size preset (e.g. 1K, 2K, 4K)") init_image: ImageField | None = InputField(default=None, description="Init image for img2img/inpaint") mask_image: ImageField | None = InputField(default=None, description="Mask image for inpaint") reference_images: list[ImageField] = InputField(default=[], description="Reference images") def _build_provider_options(self) -> dict[str, Any] | None: """Override in provider-specific subclasses to pass extra options.""" return None def invoke(self, context: InvocationContext) -> ImageCollectionOutput: model_config = context.models.get_config(self.model) if not isinstance(model_config, ExternalApiModelConfig): raise ValueError("Selected model is not an external API model") if self.provider_id is not None and model_config.provider_id != self.provider_id: raise ValueError( f"Selected model provider '{model_config.provider_id}' does not match node provider '{self.provider_id}'" ) init_image = None if self.init_image is not None: init_image = context.images.get_pil(self.init_image.image_name, mode="RGB") mask_image = None if self.mask_image is not None: mask_image = context.images.get_pil(self.mask_image.image_name, mode="L") reference_images: list[ExternalReferenceImage] = [] for image_field in self.reference_images: reference_image = context.images.get_pil(image_field.image_name, mode="RGB") reference_images.append(ExternalReferenceImage(image=reference_image)) request = ExternalGenerationRequest( model=model_config, mode=self.mode, prompt=self.prompt, seed=self.seed, num_images=self.num_images, width=self.width, height=self.height, image_size=self.image_size, init_image=init_image, mask_image=mask_image, reference_images=reference_images, metadata=self._build_request_metadata(), provider_options=self._build_provider_options(), ) result = context._services.external_generation.generate(request) outputs: list[ImageField] = [] for generated in result.images: metadata = self._build_output_metadata(model_config, result, generated.seed) image_dto = context.images.save(image=generated.image, metadata=metadata) outputs.append(ImageField(image_name=image_dto.image_name)) return ImageCollectionOutput(collection=outputs) def invoke_internal(self, context: InvocationContext, services: "InvocationServices") -> BaseInvocationOutput: """Override default cache behavior so cache hits produce new gallery entries. The standard invocation cache returns the cached output (with stale image_name references) without re-running invoke(), which means no new images are saved to the gallery on repeat invokes. For external API nodes — where the API call is the expensive part — we want cache hits to skip the API call but still produce fresh gallery entries by copying the cached images. """ if services.configuration.node_cache_size == 0 or not self.use_cache: return super().invoke_internal(context, services) key = services.invocation_cache.create_key(self) cached_value = services.invocation_cache.get(key) if cached_value is None: services.logger.debug(f'Invocation cache miss for type "{self.get_type()}": {self.id}') output = self.invoke(context) services.invocation_cache.save(key, output) return output services.logger.debug(f'Invocation cache hit for type "{self.get_type()}": {self.id}, duplicating images') if not isinstance(cached_value, ImageCollectionOutput): return cached_value outputs: list[ImageField] = [] for image_field in cached_value.collection: cached_image = context.images.get_pil(image_field.image_name, mode="RGB") image_dto = context.images.save(image=cached_image) outputs.append(ImageField(image_name=image_dto.image_name)) return ImageCollectionOutput(collection=outputs) def _build_request_metadata(self) -> dict[str, Any] | None: if self.metadata is None: return None return self.metadata.root def _build_output_metadata( self, model_config: ExternalApiModelConfig, result: ExternalGenerationResult, image_seed: int | None, ) -> MetadataField | None: metadata: dict[str, Any] = {} if self.metadata is not None: metadata.update(self.metadata.root) metadata.update( { "external_provider": model_config.provider_id, "external_model_id": model_config.provider_model_id, } ) if self.image_size is not None: metadata["image_size"] = self.image_size provider_request_id = getattr(result, "provider_request_id", None) if provider_request_id: metadata["external_request_id"] = provider_request_id provider_metadata = getattr(result, "provider_metadata", None) if provider_metadata: metadata["external_provider_metadata"] = provider_metadata if image_seed is not None: metadata["external_seed"] = image_seed metadata.update(self._build_output_provider_metadata()) if not metadata: return None return MetadataField(root=metadata) def _build_output_provider_metadata(self) -> dict[str, Any]: """Override in provider-specific subclasses to add recall-relevant fields to the image metadata.""" return {} @invocation( "openai_image_generation", title="OpenAI Image Generation", tags=["external", "generation", "openai"], category="image", version="1.0.0", ) class OpenAIImageGenerationInvocation(BaseExternalImageGenerationInvocation): """Generate images using an OpenAI-hosted external model.""" provider_id = "openai" model: ModelIdentifierField = InputField( description=FieldDescriptions.main_model, ui_model_base=[BaseModelType.External], ui_model_type=[ModelType.ExternalImageGenerator], ui_model_format=[ModelFormat.ExternalApi], ui_model_provider_id=["openai"], ) # OpenAI's API has no img2img/inpaint distinction — the edits endpoint is used # automatically when reference images are provided. Hide mode and init_image # (init_image is functionally identical to a reference image), and hide # mask_image since no OpenAI model supports inpainting. mode: ExternalGenerationMode = InputField(default="txt2img", description="Generation mode.", ui_hidden=True) init_image: ImageField | None = InputField( default=None, description="Init image (use reference_images instead)", ui_hidden=True ) mask_image: ImageField | None = InputField(default=None, description="Mask image for inpaint", ui_hidden=True) quality: Literal["auto", "high", "medium", "low"] = InputField(default="auto", description="Output image quality") background: Literal["auto", "transparent", "opaque"] = InputField( default="auto", description="Background transparency handling" ) input_fidelity: Literal["low", "high"] | None = InputField( default=None, description="Fidelity to source images (edits only)" ) def _build_provider_options(self) -> dict[str, Any]: options: dict[str, Any] = { "quality": self.quality, "background": self.background, } if self.input_fidelity is not None: options["input_fidelity"] = self.input_fidelity return options def _build_output_provider_metadata(self) -> dict[str, Any]: metadata: dict[str, Any] = { "openai_quality": self.quality, "openai_background": self.background, } if self.input_fidelity is not None: metadata["openai_input_fidelity"] = self.input_fidelity return metadata @invocation( "gemini_image_generation", title="Gemini Image Generation", tags=["external", "generation", "gemini"], category="image", version="1.0.0", ) class GeminiImageGenerationInvocation(BaseExternalImageGenerationInvocation): """Generate images using a Gemini-hosted external model.""" provider_id = "gemini" model: ModelIdentifierField = InputField( description=FieldDescriptions.main_model, ui_model_base=[BaseModelType.External], ui_model_type=[ModelType.ExternalImageGenerator], ui_model_format=[ModelFormat.ExternalApi], ui_model_provider_id=["gemini"], ) # Gemini only supports txt2img — hide mode/init_image/mask_image fields # that are inherited from the base class but not usable with any Gemini model. mode: ExternalGenerationMode = InputField(default="txt2img", description="Generation mode.", ui_hidden=True) init_image: ImageField | None = InputField( default=None, description="Init image for img2img/inpaint", ui_hidden=True ) mask_image: ImageField | None = InputField(default=None, description="Mask image for inpaint", ui_hidden=True) temperature: float | None = InputField(default=None, ge=0.0, le=2.0, description="Sampling temperature") thinking_level: Literal["minimal", "high"] | None = InputField( default=None, description="Thinking level for image generation" ) def _build_provider_options(self) -> dict[str, Any] | None: options: dict[str, Any] = {} if self.temperature is not None: options["temperature"] = self.temperature if self.thinking_level is not None: options["thinking_level"] = self.thinking_level return options or None def _build_output_provider_metadata(self) -> dict[str, Any]: metadata: dict[str, Any] = {} if self.temperature is not None: metadata["gemini_temperature"] = self.temperature if self.thinking_level is not None: metadata["gemini_thinking_level"] = self.thinking_level return metadata @invocation( "seedream_image_generation", title="Seedream Image Generation", tags=["external", "generation", "seedream"], category="image", version="1.1.0", ) class SeedreamImageGenerationInvocation(BaseExternalImageGenerationInvocation): """Generate images using a BytePlus Seedream model.""" provider_id = "seedream" model: ModelIdentifierField = InputField( description=FieldDescriptions.main_model, ui_model_base=[BaseModelType.External], ui_model_type=[ModelType.ExternalImageGenerator], ui_model_format=[ModelFormat.ExternalApi], ui_model_provider_id=["seedream"], ) # Seedream's API has only one endpoint and no inpaint support — mode is implicit # from inputs (img2img happens automatically when init_image or reference_images # are provided). Hide mode and mask_image since they have no effect. mode: ExternalGenerationMode = InputField(default="txt2img", description="Generation mode.", ui_hidden=True) mask_image: ImageField | None = InputField(default=None, description="Mask image for inpaint", ui_hidden=True) watermark: bool = InputField(default=False, description="Add watermark to generated images") optimize_prompt: bool = InputField(default=False, description="Let the model optimize the prompt before generation") def _build_provider_options(self) -> dict[str, Any]: return { "watermark": self.watermark, "optimize_prompt": self.optimize_prompt, } def _build_output_provider_metadata(self) -> dict[str, Any]: return { "seedream_watermark": self.watermark, "seedream_optimize_prompt": self.optimize_prompt, } @invocation( "alibabacloud_image_generation", title="Alibaba Cloud DashScope Image Generation", tags=["external", "generation", "alibabacloud", "dashscope"], category="image", version="1.0.0", ) class AlibabaCloudImageGenerationInvocation(BaseExternalImageGenerationInvocation): """Generate images using an Alibaba Cloud DashScope external model.""" provider_id = "alibabacloud" model: ModelIdentifierField = InputField( description=FieldDescriptions.main_model, ui_model_base=[BaseModelType.External], ui_model_type=[ModelType.ExternalImageGenerator], ui_model_format=[ModelFormat.ExternalApi], ui_model_provider_id=["alibabacloud"], )