scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-07 08:26:10 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 36

Add stability models for image edit

Browse source
This commit is contained in:
Sameer Kankute 2025-12-15 15:35:18 +05:30
parent e1d25670cd
commit b09ea08424
13 changed files with 533 additions and 16 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

6
litellm/__init__.py
View file

@ -1203,9 +1203,9 @@ from .llms.bedrock.chat.invoke_transformations.amazon_openai_transformation impo
AmazonBedrockOpenAIConfig,
)
from .llms.bedrock.image.amazon_stability1_transformation import AmazonStabilityConfig
from .llms.bedrock.image.amazon_stability3_transformation import AmazonStability3Config
from .llms.bedrock.image.amazon_nova_canvas_transformation import AmazonNovaCanvasConfig
from .llms.bedrock.image_generation.amazon_stability1_transformation import AmazonStabilityConfig
from .llms.bedrock.image_generation.amazon_stability3_transformation import AmazonStability3Config
from .llms.bedrock.image_generation.amazon_nova_canvas_transformation import AmazonNovaCanvasConfig
from .llms.bedrock.embed.amazon_titan_g1_transformation import AmazonTitanG1Config
from .llms.bedrock.embed.amazon_titan_multimodal_transformation import (
AmazonTitanMultimodalEmbeddingG1Config,

1
litellm/images/utils.py
View file

@ -82,7 +82,6 @@ class ImageEditRequestUtils:
filtered_params = {
k: v for k, v in params.items() if k in valid_keys and v is not None
}
return cast(ImageEditOptionalRequestParams, filtered_params)
@staticmethod

2
litellm/llms/bedrock/image/cost_calculator.py
View file

@ -1,6 +1,6 @@
from typing import Optional
from litellm.llms.bedrock.image.image_handler import BedrockImageGeneration
from litellm.llms.bedrock.image_generation.image_handler import BedrockImageGeneration
from litellm.types.utils import ImageResponse

6
litellm/llms/bedrock/image/image_handler.py
View file

@ -9,13 +9,13 @@ from pydantic import BaseModel
import litellm
from litellm._logging import verbose_logger
from litellm.litellm_core_utils.litellm_logging import Logging as LitellmLogging
from litellm.llms.bedrock.image.amazon_nova_canvas_transformation import (
from litellm.llms.bedrock.image_generation.amazon_nova_canvas_transformation import (
AmazonNovaCanvasConfig,
)
from litellm.llms.bedrock.image.amazon_stability3_transformation import (
from litellm.llms.bedrock.image_generation.amazon_stability3_transformation import (
AmazonStability3Config,
)
from litellm.llms.bedrock.image.amazon_titan_transformation import (
from litellm.llms.bedrock.image_generation.amazon_titan_transformation import (
AmazonTitanImageGenerationConfig,
)
from litellm.llms.custom_httpx.http_handler import (

2
litellm/llms/custom_httpx/llm_http_handler.py
View file

@ -3761,7 +3761,7 @@ class BaseLLMHTTPHandler:
input=prompt,
api_key="",
additional_args={
"complete_input_dict": data,
"complete_input_dict": files,
"api_base": api_base,
"headers": headers,
},

37
litellm/llms/stability/image_edit/__init__.py Normal file
View file

@ -0,0 +1,37 @@
"""
Stability AI Image Edit Module
Factory function for getting the appropriate config class.
"""
from litellm.llms.base_llm.image_edit.transformation import (
BaseImageEditConfig,
)
from .transformation import StabilityImageEditConfig
__all__ = [
"StabilityImageEditConfig",
"get_stability_image_edit_config",
]
def get_stability_image_edit_config(model: str) -> BaseImageEditConfig:
"""
Get the appropriate Stability AI config for the given model.
Currently all models use the same config class, but this factory
allows for model-specific configs in the future.
Args:
model: The model name (e.g., "stability/inpaint", "stability/outpaint")
Returns:
BaseImageEditConfig instance for Stability AI
"""
# For now, all models use the same config
# In the future, we could have model-specific configs:
# - StabilityInpaintConfig for Inpaint models
# - StabilityOutpaintConfig for Outpaint models
# - etc.
return StabilityImageEditConfig()

291
litellm/llms/stability/image_edit/transformations.py Normal file
View file

@ -0,0 +1,291 @@
"""
Stability AI Image Edit Config
Handles transformation between OpenAI-compatible format and Stability AI API format.
API Reference: https://platform.stability.ai/docs/api-reference
"""
from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple
import httpx
from httpx._types import RequestFiles
from litellm.llms.base_llm.image_edit.transformation import BaseImageEditConfig
from litellm.secret_managers.main import get_secret_str
from litellm.types.llms.openai import (
OpenAIImageEditOptionalParams,
)
from litellm.types.llms.stability import (
OPENAI_SIZE_TO_STABILITY_ASPECT_RATIO,
STABILITY_EDIT_ENDPOINTS,
StabilityImageEditRequest,
)
from litellm.types.utils import FileTypes, ImageObject, ImageResponse
from litellm.utils import get_model_info
if TYPE_CHECKING:
from litellm.litellm_core_utils.litellm_logging import Logging as _LiteLLMLoggingObj
LiteLLMLoggingObj = _LiteLLMLoggingObj
else:
LiteLLMLoggingObj = Any
class StabilityImageEditConfig(BaseImageEditConfig):
"""
Configuration for Stability AI image edit.
Supports:
- Stable Diffusion 3 (SD3, SD3.5) Image Edit
"""
DEFAULT_BASE_URL: str = "https://api.stability.ai"
def get_supported_openai_params(
self, model: str
) -> List[OpenAIImageEditOptionalParams]:
"""
Return list of OpenAI params supported by Stability AI.
https://platform.stability.ai/docs/api-reference
"""
return [
"n", # Number of images (Stability always returns 1, we can loop)
"size", # Maps to aspect_ratio
"response_format", # b64_json or url (Stability only returns b64)
"mask"
]
def map_openai_params(
self,
image_edit_optional_params: dict,
model: str,
drop_params: bool,
) -> dict:
"""
Map OpenAI parameters to Stability AI parameters.
OpenAI -> Stability mappings:
- size -> aspect_ratio
- n -> (handled separately, Stability returns 1 image per request)
"""
supported_params = self.get_supported_openai_params(model)
# Define mapping from OpenAI params to Stability params
param_mapping = {
"size": "aspect_ratio",
# "n" and "response_format" are handled separately
}
# Create a copy to not mutate original
mapped_params = image_edit_optional_params.copy()
for k, v in image_edit_optional_params.items():
if k in param_mapping:
# Map param if mapping exists and value is valid
if k == "size" and v in OPENAI_SIZE_TO_STABILITY_ASPECT_RATIO:
mapped_params[param_mapping[k]] = OPENAI_SIZE_TO_STABILITY_ASPECT_RATIO[v]
# Don't copy "size" itself to final dict
elif k == "n":
# Store for logic but do not add to outgoing params
mapped_params["_n"] = v
elif k == "response_format":
# Only b64 supported at Stability; store for postprocessing
mapped_params["_response_format"] = v
elif k not in supported_params:
if not drop_params:
raise ValueError(
f"Parameter {k} is not supported for model {model}. "
f"Supported parameters are {supported_params}. "
f"Set drop_params=True to drop unsupported parameters."
)
# Otherwise, param will simply be dropped
else:
# param is supported and not mapped, keep as-is
continue
# Remove OpenAI params that have been mapped unless they're in stability
for mapped in ["size", "n", "response_format"]:
if mapped in mapped_params:
del mapped_params[mapped]
return mapped_params
def _get_model_endpoint(self, model: str) -> str:
"""
Get the API endpoint for a given model.
"""
# Remove "stability/" prefix if present
model_name = model.lower()
if model_name.startswith("stability/"):
model_name = model_name[10:] # Remove "stability/" prefix
# Check if model is in our mapping
for key, endpoint in STABILITY_EDIT_ENDPOINTS.items():
if key in model_name:
return endpoint
# Default to SD3 endpoint
return "/v2beta/stable-image/edit/inpaint"
def get_complete_url(
self,
api_base: Optional[str],
model: str,
litellm_params: dict,
) -> str:
"""
Get the complete URL for the Stability AI API request.
"""
base_url: str = (
api_base
or get_secret_str("STABILITY_API_BASE")
or litellm_params.get("api_base", None)
or self.DEFAULT_BASE_URL
)
base_url = base_url.rstrip("/")
endpoint = self._get_model_endpoint(model)
return f"{base_url}{endpoint}"
def validate_environment(
self,
headers: dict,
model: str,
api_key: Optional[str] = None,
) -> dict:
"""
Validate environment and set up headers for Stability AI.
"""
final_api_key: Optional[str] = api_key or get_secret_str("STABILITY_API_KEY")
if not final_api_key:
raise ValueError(
"STABILITY_API_KEY is not set. "
"Please set it via environment variable or pass api_key parameter."
)
headers["Authorization"] = f"Bearer {final_api_key}"
headers["Accept"] = "application/json"
return headers
def transform_image_edit_request(
self,
model: str,
prompt: str,
image: FileTypes,
image_edit_optional_request_params: OpenAIImageEditOptionalParams,
litellm_params: dict,
headers: dict,
) -> Tuple[Dict, RequestFiles]:
"""
Transform OpenAI-style request to Stability AI request format.
Note: Stability AI uses multipart/form-data, but the HTTP handler
will handle the conversion from dict to form data.
"""
# Build Stability request
# Populate multipart form-data as separate text fields (data) and files.
# Stability expects prompt/output_format/etc. as normal form fields, not file parts.
data: StabilityImageEditRequest = {
"prompt": prompt,
"output_format": "png", # Default to PNG
}
files: Dict[str, Any] = {"image": image[0]}
# Add optional params (already mapped in map_openai_params)
for key, value in image_edit_optional_request_params.items():
# Skip internal params (prefixed with _)
if key.startswith("_") or value is None:
continue
# File-like optional param
if key == "mask":
files["mask"] = value # type: ignore
continue
# Supported text fields
if key in [
"negative_prompt",
"aspect_ratio",
"seed",
"output_format",
"model",
"mode",
"strength",
"style_preset",
]:
data[key] = value # type: ignore
return data, files
def transform_image_edit_response(
self,
model: str,
raw_response: httpx.Response,
logging_obj: LiteLLMLoggingObj,
api_key: Optional[str] = None,
json_mode: Optional[bool] = None,
) -> ImageResponse:
"""
Transform Stability AI response to OpenAI-compatible ImageResponse.
Stability returns: {"image": "base64...", "finish_reason": "SUCCESS", "seed": 123}
OpenAI expects: {"data": [{"b64_json": "base64..."}], "created": timestamp}
"""
try:
response_data = raw_response.json()
except Exception as e:
raise self.get_error_class(
error_message=f"Error parsing Stability AI response: {e}",
status_code=raw_response.status_code,
headers=raw_response.headers,
)
# Check for errors in response
if "errors" in response_data:
raise self.get_error_class(
error_message=f"Stability AI error: {response_data['errors']}",
status_code=raw_response.status_code,
headers=raw_response.headers,
)
# Check finish_reason
finish_reason = response_data.get("finish_reason", "")
if finish_reason == "CONTENT_FILTERED":
raise self.get_error_class(
error_message="Content was filtered by Stability AI safety systems",
status_code=400,
headers=raw_response.headers,
)
model_response = ImageResponse()
if not model_response.data:
model_response.data = []
# Extract image from response
image_b64 = response_data.get("image")
if image_b64:
model_response.data.append(
ImageObject(
b64_json=image_b64,
url=None,
revised_prompt=None,
)
)
if not hasattr(model_response, "_hidden_params"):
model_response._hidden_params = {}
if "additional_headers" not in model_response._hidden_params:
model_response._hidden_params["additional_headers"] = {}
# Override: fetch model-cost from model_cost map based on the provided model name
model_info = get_model_info(model, custom_llm_provider="stability")
cost_per_image = model_info.get("output_cost_per_image", 0)
model_response._hidden_params["additional_headers"]["llm_provider-x-litellm-response-cost"] = float(cost_per_image)
return model_response
def use_multipart_form_data(self) -> bool:
"""
Stability AI requires multipart/form-data for image generation.
"""
return True

2
litellm/main.py
View file

@ -165,7 +165,7 @@ from .llms.azure_ai.anthropic.handler import AzureAnthropicChatCompletion
from .llms.azure_ai.embed import AzureAIEmbedding
from .llms.bedrock.chat import BedrockConverseLLM, BedrockLLM
from .llms.bedrock.embed.embedding import BedrockEmbedding
from .llms.bedrock.image.image_handler import BedrockImageGeneration
from .llms.bedrock.image_generation.image_handler import BedrockImageGeneration
from .llms.bytez.chat.transformation import BytezChatConfig
from .llms.clarifai.chat.transformation import ClarifaiConfig
from .llms.codestral.completion.handler import CodestralTextCompletion

84
litellm/model_prices_and_context_window_backup.json
View file

@ -24483,6 +24483,90 @@
"output_cost_per_image": 0.08,
"supported_endpoints": ["/v1/images/generations"]
},
"stability/inpaint": {
"litellm_provider": "stability",
"mode": "image_edit",
"output_cost_per_image": 0.005,
"supported_endpoints": ["/v1/images/edits"]
},
"stability/outpaint": {
"litellm_provider": "stability",
"mode": "image_edit",
"output_cost_per_image": 0.004,
"supported_endpoints": ["/v1/images/edits"]
},
"stability/erase": {
"litellm_provider": "stability",
"mode": "image_edit",
"output_cost_per_image": 0.005,
"supported_endpoints": ["/v1/images/edits"]
},
"stability/search-and-replace": {
"litellm_provider": "stability",
"mode": "image_edit",
"output_cost_per_image": 0.005,
"supported_endpoints": ["/v1/images/edits"]
},
"stability/search-and-recolor": {
"litellm_provider": "stability",
"mode": "image_edit",
"output_cost_per_image": 0.005,
"supported_endpoints": ["/v1/images/edits"]
},
"stability/remove-background": {
"litellm_provider": "stability",
"mode": "image_edit",
"output_cost_per_image": 0.005,
"supported_endpoints": ["/v1/images/edits"]
},
"stability/replace-background-and-relight": {
"litellm_provider": "stability",
"mode": "image_edit",
"output_cost_per_image": 0.008,
"supported_endpoints": ["/v1/images/edits"]
},
"stability/sketch": {
"litellm_provider": "stability",
"mode": "image_edit",
"output_cost_per_image": 0.005,
"supported_endpoints": ["/v1/images/edits"]
},
"stability/structure": {
"litellm_provider": "stability",
"mode": "image_edit",
"output_cost_per_image": 0.005,
"supported_endpoints": ["/v1/images/edits"]
},
"stability/style": {
"litellm_provider": "stability",
"mode": "image_edit",
"output_cost_per_image": 0.005,
"supported_endpoints": ["/v1/images/edits"]
},
"stability/style-transfer": {
"litellm_provider": "stability",
"mode": "image_edit",
"output_cost_per_image": 0.008,
"supported_endpoints": ["/v1/images/edits"]
},
"stability/fast": {
"litellm_provider": "stability",
"mode": "image_edit",
"output_cost_per_image": 0.002,
"supported_endpoints": ["/v1/images/edits"]
},
"stability/conservative": {
"litellm_provider": "stability",
"mode": "image_edit",
"output_cost_per_image": 0.04,
"supported_endpoints": ["/v1/images/edits"]
},
"stability/creative": {
"litellm_provider": "stability",
"mode": "image_edit",
"output_cost_per_image": 0.06,
"supported_endpoints": ["/v1/images/edits"]
},
"stability/stable-image-core": {
"litellm_provider": "stability",
"mode": "image_generation",

13
litellm/types/llms/openai.py
View file

@ -1029,6 +1029,19 @@ OpenAIImageGenerationOptionalParams = Literal[
"user",
]
OpenAIImageEditOptionalParams = Literal[
"background",
"n",
"mask"
"output_compression",
"output_format",
"quality",
"partial_images",
"response_format",
"size",
"style",
"user",
]
class ComputerToolParam(TypedDict, total=False):
display_height: Required[float]

15
litellm/types/llms/stability.py
View file

@ -29,6 +29,13 @@ class StabilityImageGenerationRequest(TypedDict, total=False):
strength: Optional[float] # How much to transform the image (0-1)
style_preset: Optional[str] # Style preset name
class StabilityImageEditRequest(StabilityImageGenerationRequest):
"""
Request parameters for Stability AI image edit endpoint.
Endpoint: /v2beta/stable-image/edit/inpaint
"""
mask: Optional[str] # Base64-encoded mask (white = edit, black = keep)
class StabilityImageGenerationResponse(TypedDict, total=False):
"""
@ -197,16 +204,12 @@ STABILITY_EDIT_ENDPOINTS = {
"search-and-replace": "/v2beta/stable-image/edit/search-and-replace",
"search-and-recolor": "/v2beta/stable-image/edit/search-and-recolor",
"remove-background": "/v2beta/stable-image/edit/remove-background",
}
STABILITY_UPSCALE_ENDPOINTS = {
"replace-background-and-relight": "/v2beta/stable-image/edit/replace-background-and-relight",
"fast": "/v2beta/stable-image/upscale/fast",
"conservative": "/v2beta/stable-image/upscale/conservative",
"creative": "/v2beta/stable-image/upscale/creative",
}
STABILITY_CONTROL_ENDPOINTS = {
"sketch": "/v2beta/stable-image/control/sketch",
"structure": "/v2beta/stable-image/control/structure",
"style": "/v2beta/stable-image/control/style",
"style-transfer": "/v2beta/stable-image/control/style-transfer",
}

6
litellm/utils.py
View file

@ -7966,6 +7966,12 @@ class ProviderConfigManager:
)
return get_vertex_ai_image_edit_config(model)
elif LlmProviders.STABILITY == provider:
from litellm.llms.stability.image_edit import (
get_stability_image_edit_config,
)
return get_stability_image_edit_config(model)
return None
@staticmethod

84
model_prices_and_context_window.json
View file

@ -24483,6 +24483,90 @@
"output_cost_per_image": 0.08,
"supported_endpoints": ["/v1/images/generations"]
},
"stability/inpaint": {
"litellm_provider": "stability",
"mode": "image_edit",
"output_cost_per_image": 0.005,
"supported_endpoints": ["/v1/images/edits"]
},
"stability/outpaint": {
"litellm_provider": "stability",
"mode": "image_edit",
"output_cost_per_image": 0.004,
"supported_endpoints": ["/v1/images/edits"]
},
"stability/erase": {
"litellm_provider": "stability",
"mode": "image_edit",
"output_cost_per_image": 0.005,
"supported_endpoints": ["/v1/images/edits"]
},
"stability/search-and-replace": {
"litellm_provider": "stability",
"mode": "image_edit",
"output_cost_per_image": 0.005,
"supported_endpoints": ["/v1/images/edits"]
},
"stability/search-and-recolor": {
"litellm_provider": "stability",
"mode": "image_edit",
"output_cost_per_image": 0.005,
"supported_endpoints": ["/v1/images/edits"]
},
"stability/remove-background": {
"litellm_provider": "stability",
"mode": "image_edit",
"output_cost_per_image": 0.005,
"supported_endpoints": ["/v1/images/edits"]
},
"stability/replace-background-and-relight": {
"litellm_provider": "stability",
"mode": "image_edit",
"output_cost_per_image": 0.008,
"supported_endpoints": ["/v1/images/edits"]
},
"stability/sketch": {
"litellm_provider": "stability",
"mode": "image_edit",
"output_cost_per_image": 0.005,
"supported_endpoints": ["/v1/images/edits"]
},
"stability/structure": {
"litellm_provider": "stability",
"mode": "image_edit",
"output_cost_per_image": 0.005,
"supported_endpoints": ["/v1/images/edits"]
},
"stability/style": {
"litellm_provider": "stability",
"mode": "image_edit",
"output_cost_per_image": 0.005,
"supported_endpoints": ["/v1/images/edits"]
},
"stability/style-transfer": {
"litellm_provider": "stability",
"mode": "image_edit",
"output_cost_per_image": 0.008,
"supported_endpoints": ["/v1/images/edits"]
},
"stability/fast": {
"litellm_provider": "stability",
"mode": "image_edit",
"output_cost_per_image": 0.002,
"supported_endpoints": ["/v1/images/edits"]
},
"stability/conservative": {
"litellm_provider": "stability",
"mode": "image_edit",
"output_cost_per_image": 0.04,
"supported_endpoints": ["/v1/images/edits"]
},
"stability/creative": {
"litellm_provider": "stability",
"mode": "image_edit",
"output_cost_per_image": 0.06,
"supported_endpoints": ["/v1/images/edits"]
},
"stability/stable-image-core": {
"litellm_provider": "stability",
"mode": "image_generation",