Minor bug fixes: mask list

This commit is contained in:
Sameer Kankute 2025-12-19 22:52:23 +05:30
parent 5feb6d5388
commit 248eb4ea49
8 changed files with 324 additions and 60 deletions

View file

@ -16,7 +16,7 @@ LiteLLM provides image editing functionality that maps to OpenAI's `/images/edit
| Supported operations | Create image edits | Single and multiple images supported |
| Supported LiteLLM SDK Versions | 1.63.8+ | Gemini support requires 1.79.3+ |
| Supported LiteLLM Proxy Versions | 1.71.1+ | Gemini support requires 1.79.3+ |
| Supported LLM providers | **OpenAI**, **Gemini (Google AI Studio)**, **Vertex AI** | Gemini supports the new `gemini-2.5-flash-image` family. Vertex AI supports both Gemini and Imagen models. |
| Supported LLM providers | **OpenAI**, **Gemini (Google AI Studio)**, **Vertex AI**, **Stability AI**, **AWS Bedrock (Stability)** | Gemini supports the new `gemini-2.5-flash-image` family. Vertex AI supports both Gemini and Imagen models. Stability AI and Bedrock Stability support various image editing operations. |
#### ⚡See all supported models and providers at [models.litellm.ai](https://models.litellm.ai/)

View file

@ -8,7 +8,7 @@ https://stability.ai/
| Description | Stability AI creates open AI models for image, video, audio, and 3D generation. Known for Stable Diffusion. |
| Provider Route on LiteLLM | `stability/` |
| Link to Provider Doc | [Stability AI API ↗](https://platform.stability.ai/docs/api-reference) |
| Supported Operations | [`/images/generations`](#image-generation) |
| Supported Operations | [`/images/generations`](#image-generation), [`/images/edits`](#image-editing) |
LiteLLM supports Stability AI Image Generation calls via the Stability AI REST API (not via Bedrock).
@ -169,13 +169,285 @@ Stability AI returns images in base64 format. The response is OpenAI-compatible:
}
```
## Comparing with Bedrock
## Image Editing
Stability AI supports various image editing operations including inpainting, upscaling, outpainting, background removal, and more.
### Usage - LiteLLM Python SDK
#### Inpainting (Edit with Mask)
```python showLineNumbers
from litellm import image_edit
import os
os.environ['STABILITY_API_KEY'] = "your-api-key"
# Inpainting - edit specific areas using a mask
response = image_edit(
model="stability/stable-image-inpaint-v1:0",
image=open("original_image.png", "rb"),
mask=open("mask_image.png", "rb"),
prompt="Add a beautiful sunset in the masked area",
size="1024x1024",
)
print(response)
```
#### Image Upscaling
```python showLineNumbers
from litellm import image_edit
import os
os.environ['STABILITY_API_KEY'] = "your-api-key"
# Conservative upscaling - preserves details
response = image_edit(
model="stability/stable-conservative-upscale-v1:0",
image=open("low_res_image.png", "rb"),
prompt="Upscale this image while preserving details",
)
# Creative upscaling - adds creative details
response = image_edit(
model="stability/stable-creative-upscale-v1:0",
image=open("low_res_image.png", "rb"),
prompt="Upscale and enhance with creative details",
creativity=0.3, # 0-0.35, higher = more creative
)
# Fast upscaling - quick upscaling
response = image_edit(
model="stability/stable-fast-upscale-v1:0",
image=open("low_res_image.png", "rb"),
prompt="Quickly upscale this image",
)
print(response)
```
#### Image Outpainting
```python showLineNumbers
from litellm import image_edit
import os
os.environ['STABILITY_API_KEY'] = "your-api-key"
# Extend image beyond its borders
response = image_edit(
model="stability/stable-outpaint-v1:0",
image=open("original_image.png", "rb"),
prompt="Extend this landscape with mountains",
left=100, # Pixels to extend on the left
right=100, # Pixels to extend on the right
up=50, # Pixels to extend on top
down=50, # Pixels to extend on bottom
)
print(response)
```
#### Background Removal
```python showLineNumbers
from litellm import image_edit
import os
os.environ['STABILITY_API_KEY'] = "your-api-key"
# Remove background from image
response = image_edit(
model="stability/stable-image-remove-background-v1:0",
image=open("portrait.png", "rb"),
prompt="Remove the background",
)
print(response)
```
#### Search and Replace
```python showLineNumbers
from litellm import image_edit
import os
os.environ['STABILITY_API_KEY'] = "your-api-key"
# Search and replace objects in image
response = image_edit(
model="stability/stable-image-search-replace-v1:0",
image=open("scene.png", "rb"),
prompt="A red sports car",
search_prompt="blue sedan", # What to replace
)
# Search and recolor
response = image_edit(
model="stability/stable-image-search-recolor-v1:0",
image=open("scene.png", "rb"),
prompt="Make it golden yellow",
select_prompt="the car", # What to recolor
)
print(response)
```
#### Image Control (Sketch/Structure)
```python showLineNumbers
from litellm import image_edit
import os
os.environ['STABILITY_API_KEY'] = "your-api-key"
# Control with sketch
response = image_edit(
model="stability/stable-image-control-sketch-v1:0",
image=open("sketch.png", "rb"),
prompt="Turn this sketch into a realistic photo",
control_strength=0.7, # 0-1, higher = more control
)
# Control with structure
response = image_edit(
model="stability/stable-image-control-structure-v1:0",
image=open("structure_reference.png", "rb"),
prompt="Generate image following this structure",
control_strength=0.7,
)
print(response)
```
#### Erase Objects
```python showLineNumbers
from litellm import image_edit
import os
os.environ['STABILITY_API_KEY'] = "your-api-key"
# Erase objects from image
response = image_edit(
model="stability/stable-image-erase-object-v1:0",
image=open("scene.png", "rb"),
mask=open("object_mask.png", "rb"), # Mask the object to erase
prompt="Remove the object",
)
print(response)
```
### Supported Image Edit Models
| Model Name | Function Call | Description |
|------------|---------------|-------------|
| stable-image-inpaint-v1:0 | `image_edit(model="stability/stable-image-inpaint-v1:0", ...)` | Inpainting with mask |
| stable-conservative-upscale-v1:0 | `image_edit(model="stability/stable-conservative-upscale-v1:0", ...)` | Conservative upscaling |
| stable-creative-upscale-v1:0 | `image_edit(model="stability/stable-creative-upscale-v1:0", ...)` | Creative upscaling |
| stable-fast-upscale-v1:0 | `image_edit(model="stability/stable-fast-upscale-v1:0", ...)` | Fast upscaling |
| stable-outpaint-v1:0 | `image_edit(model="stability/stable-outpaint-v1:0", ...)` | Extend image borders |
| stable-image-remove-background-v1:0 | `image_edit(model="stability/stable-image-remove-background-v1:0", ...)` | Remove background |
| stable-image-search-replace-v1:0 | `image_edit(model="stability/stable-image-search-replace-v1:0", ...)` | Search and replace objects |
| stable-image-search-recolor-v1:0 | `image_edit(model="stability/stable-image-search-recolor-v1:0", ...)` | Search and recolor |
| stable-image-control-sketch-v1:0 | `image_edit(model="stability/stable-image-control-sketch-v1:0", ...)` | Control with sketch |
| stable-image-control-structure-v1:0 | `image_edit(model="stability/stable-image-control-structure-v1:0", ...)` | Control with structure |
| stable-image-erase-object-v1:0 | `image_edit(model="stability/stable-image-erase-object-v1:0", ...)` | Erase objects |
| stable-image-style-guide-v1:0 | `image_edit(model="stability/stable-image-style-guide-v1:0", ...)` | Apply style guide |
| stable-style-transfer-v1:0 | `image_edit(model="stability/stable-style-transfer-v1:0", ...)` | Transfer style |
### Usage - LiteLLM Proxy Server
#### 1. Setup config.yaml
```yaml showLineNumbers
model_list:
- model_name: stability-inpaint
litellm_params:
model: stability/stable-image-inpaint-v1:0
api_key: os.environ/STABILITY_API_KEY
model_info:
mode: image_edit
- model_name: stability-upscale
litellm_params:
model: stability/stable-conservative-upscale-v1:0
api_key: os.environ/STABILITY_API_KEY
model_info:
mode: image_edit
general_settings:
master_key: sk-1234
```
#### 2. Start the proxy
```bash showLineNumbers
litellm --config config.yaml
# RUNNING on http://0.0.0.0:4000
```
#### 3. Test it
```bash showLineNumbers
curl -X POST "http://0.0.0.0:4000/v1/images/edits" \
-H "Authorization: Bearer sk-1234" \
-F "model=stability-inpaint" \
-F "image=@original_image.png" \
-F "mask=@mask_image.png" \
-F "prompt=Add a beautiful garden in the masked area"
```
## AWS Bedrock (Stability)
LiteLLM also supports Stability AI models via AWS Bedrock. This is useful if you're already using AWS infrastructure.
### Usage - Bedrock Stability
```python showLineNumbers
from litellm import image_edit
import os
# Set AWS credentials
os.environ["AWS_ACCESS_KEY_ID"] = "your-access-key"
os.environ["AWS_SECRET_ACCESS_KEY"] = "your-secret-key"
os.environ["AWS_REGION_NAME"] = "us-east-1"
# Bedrock Stability inpainting
response = image_edit(
model="bedrock/us.stability.stable-image-inpaint-v1:0",
image=open("original_image.png", "rb"),
mask=open("mask_image.png", "rb"),
prompt="Add flowers in the masked area",
size="1024x1024",
)
print(response)
```
### Supported Bedrock Stability Models
All Stability AI image edit models are available via Bedrock with the `bedrock/` prefix:
| Direct API Model | Bedrock Model | Description |
|------------------|---------------|-------------|
| stability/stable-image-inpaint-v1:0 | bedrock/us.stability.stable-image-inpaint-v1:0 | Inpainting |
| stability/stable-conservative-upscale-v1:0 | bedrock/stability.stable-conservative-upscale-v1:0 | Conservative upscaling |
| stability/stable-creative-upscale-v1:0 | bedrock/stability.stable-creative-upscale-v1:0 | Creative upscaling |
| stability/stable-fast-upscale-v1:0 | bedrock/stability.stable-fast-upscale-v1:0 | Fast upscaling |
| stability/stable-outpaint-v1:0 | bedrock/stability.stable-outpaint-v1:0 | Outpainting |
| stability/stable-image-remove-background-v1:0 | bedrock/stability.stable-image-remove-background-v1:0 | Remove background |
| stability/stable-image-search-replace-v1:0 | bedrock/stability.stable-image-search-replace-v1:0 | Search and replace |
| stability/stable-image-search-recolor-v1:0 | bedrock/stability.stable-image-search-recolor-v1:0 | Search and recolor |
| stability/stable-image-control-sketch-v1:0 | bedrock/stability.stable-image-control-sketch-v1:0 | Control with sketch |
| stability/stable-image-control-structure-v1:0 | bedrock/stability.stable-image-control-structure-v1:0 | Control with structure |
| stability/stable-image-erase-object-v1:0 | bedrock/stability.stable-image-erase-object-v1:0 | Erase objects |
**Note:** Bedrock model IDs may use `us.stability.*` or `stability.*` prefix depending on the region and model.
## Comparing Routes
LiteLLM supports Stability AI models via two routes:
| Route | Provider | Use Case |
|-------|----------|----------|
| `stability/` | Stability AI Direct API | Direct access, all latest models |
| `bedrock/stability.*` | AWS Bedrock | AWS integration, enterprise features |
| Route | Provider | Use Case | Image Generation | Image Editing |
|-------|----------|----------|------------------|---------------|
| `stability/` | Stability AI Direct API | Direct access, all latest models | ✅ | ✅ |
| `bedrock/stability.*` | AWS Bedrock | AWS integration, enterprise features | ✅ | ✅ |
Use `stability/` for direct API access. Use `bedrock/stability.*` if you're already using AWS Bedrock.

View file

@ -892,6 +892,7 @@ BEDROCK_INVOKE_PROVIDERS_LITERAL = Literal[
"qwen2",
"twelvelabs",
"openai",
"stability",
]
BEDROCK_EMBEDDING_PROVIDERS_LITERAL = Literal[

View file

@ -839,6 +839,7 @@ def image_edit(
if custom_llm_provider == "bedrock":
if model is None:
raise Exception("Model needs to be set for bedrock")
image_edit_request_params.update(non_default_params)
return bedrock_image_edit.image_edit( # type: ignore
model=model,
image=images,

View file

@ -365,6 +365,10 @@ class BaseAWSLLM:
model_id = BaseAWSLLM._get_model_id_from_model_with_spec(
model_id, spec="qwen3"
)
elif provider == "stability" and "stability/" in model_id:
model_id = BaseAWSLLM._get_model_id_from_model_with_spec(
model_id, spec="stability"
)
return model_id
@staticmethod

View file

@ -300,8 +300,8 @@ class BedrockImageEdit(BaseAWSLLM):
config_class = self.get_config_class(model=model)
config_instance = config_class()
config_instance.transform_image_edit_response(
model_response=model_response,
model_response = config_instance.transform_image_edit_response(
model=model,
raw_response=response,
logging_obj=logging_obj,
)

View file

@ -153,43 +153,6 @@ class BedrockStabilityImageEditConfig(BaseImageEditConfig):
return mapped_params
def _get_model_operation(self, model: str) -> str:
"""
Get the operation type for a given model.
"""
model_lower = model.lower()
if "upscale" in model_lower:
if "fast" in model_lower:
return "fast"
elif "conservative" in model_lower:
return "conservative"
elif "creative" in model_lower:
return "creative"
elif "outpaint" in model_lower:
return "outpaint"
elif "inpaint" in model_lower:
return "inpaint"
elif "erase" in model_lower:
return "erase"
elif "remove-background" in model_lower:
return "remove-background"
elif "search-recolor" in model_lower:
return "search-and-recolor"
elif "search-replace" in model_lower:
return "search-and-replace"
elif "control-sketch" in model_lower:
return "sketch"
elif "control-structure" in model_lower:
return "structure"
elif "style-guide" in model_lower:
return "style"
elif "style-transfer" in model_lower:
return "style-transfer"
# Default to inpaint
return "inpaint"
def transform_image_edit_request(
self,
model: str,
@ -234,17 +197,29 @@ class BedrockStabilityImageEditConfig(BaseImageEditConfig):
if key.startswith("_") or value is None:
continue
# File-like optional param
if key == "mask":
if hasattr(value, 'read'):
mask_bytes = value.read()
# File-like optional params (mask, init_image, style_image, etc.)
if key in ["mask", "init_image", "style_image"]:
# Handle case where value might be in a list
file_value = value
if isinstance(value, list) and len(value) > 0:
file_value = value[0]
if hasattr(file_value, 'read'):
file_bytes = file_value.read()
elif isinstance(file_value, bytes):
file_bytes = file_value
elif isinstance(file_value, str):
# Already a base64 string
data[key] = file_value
continue
else:
mask_bytes = value
if isinstance(mask_bytes, bytes):
mask_b64 = base64.b64encode(mask_bytes).decode('utf-8')
file_bytes = file_value
if isinstance(file_bytes, bytes):
file_b64 = base64.b64encode(file_bytes).decode('utf-8')
else:
mask_b64 = mask_bytes
data["mask"] = mask_b64
file_b64 = file_bytes
data[key] = file_b64
continue
# Supported text fields
@ -270,8 +245,6 @@ class BedrockStabilityImageEditConfig(BaseImageEditConfig):
"composition_fidelity",
"style_strength",
"change_strength",
"init_image",
"style_image",
]:
data[key] = value # type: ignore
@ -293,6 +266,8 @@ class BedrockStabilityImageEditConfig(BaseImageEditConfig):
"""
try:
response_data = raw_response.json()
with open("response_data.json", "w") as f:
json.dump(response_data, f)
except Exception as e:
raise self.get_error_class(
error_message=f"Error parsing Bedrock Stability response: {e}",

View file

@ -201,7 +201,20 @@ class StabilityImageEditConfig(BaseImageEditConfig):
# File-like optional param
if key == "mask":
files["mask"] = value # type: ignore
# Handle case where mask might be in a list
mask_value = value
if isinstance(value, list) and len(value) > 0:
mask_value = value[0]
files["mask"] = mask_value # type: ignore
continue
# File-like optional params (init_image, style_image, etc.)
if key in ["init_image", "style_image"]:
# Handle case where value might be in a list
file_value = value
if isinstance(value, list) and len(value) > 0:
file_value = value[0]
files[key] = file_value # type: ignore
continue
# Supported text fields
@ -221,8 +234,6 @@ class StabilityImageEditConfig(BaseImageEditConfig):
"grow_mask",
"select_prompt",
"control_strength",
"init_image",
"style_image",
"composition_fidelity",
"change_strength"
]: