Skip to content

Commit e983d5c

Browse files
Torantulinomajdyzclaudentindle
authored
fix(backend): Implement passed uploaded media support for AI image customizer block (#11441)
- Added `store_media_file` utility to convert local file paths to Data URIs for image processing. - Updated `AIImageCustomizerBlock` to utilize processed images in model execution, improving compatibility with Replicate API. - Added optional Aspect ratio input to AIImageCustomizerBlock This change enhances the image handling capabilities of the AI image customizer, ensuring that images are properly formatted for external processing. <!-- Clearly explain the need for these changes: --> ### Changes 🏗️ <!-- Concisely describe all of the changes made in this pull request: --> ### Checklist 📋 #### For code changes: - [x] I have clearly listed my changes in the PR description - [x] I have made a test plan - [x] I have tested my changes according to the test plan: <!-- Put your test plan here: --> - [x] Created agent using AI Image Customizer block attached to agent file input - [x] Run agent, confirmed block is working - [x] Confirm block is still working in original direct file upload setup. ### Testing Results #### Before (dev cloud): <img width="836" height="592" alt="image" src="https://github.com/user-attachments/assets/88c75668-c5c9-44bb-bec5-6554088a0cb7" /> #### After (local): <img width="827" height="587" alt="image" src="https://github.com/user-attachments/assets/04fea431-70a5-4173-bc84-d354c03d7174" /> <!-- CURSOR_SUMMARY --> --- > [!NOTE] > Preprocesses input images to data URIs and adds an `aspect_ratio` option, wiring both through to Replicate in `AIImageCustomizerBlock`. > > - **Backend** > - **`backend/blocks/ai_image_customizer.py`**: > - Preprocesses input images via `store_media_file(..., return_content=True)` to Data URIs before invoking Replicate. > - Adds `AspectRatio` enum and `aspect_ratio` input; passed through `run_model` and included in Replicate input. > - Updates block test input accordingly. > > <sup>Written by [Cursor Bugbot](https://cursor.com/dashboard?tab=bugbot) for commit 4116cf8. This will update automatically on new commits. Configure [here](https://cursor.com/dashboard?tab=bugbot).</sup> <!-- /CURSOR_SUMMARY --> --------- Co-authored-by: Zamil Majdy <[email protected]> Co-authored-by: Claude <[email protected]> Co-authored-by: Nicholas Tindle <[email protected]>
1 parent bdb94a3 commit e983d5c

File tree

1 file changed

+39
-2
lines changed

1 file changed

+39
-2
lines changed

autogpt_platform/backend/backend/blocks/ai_image_customizer.py

Lines changed: 39 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,4 @@
1+
import asyncio
12
from enum import Enum
23
from typing import Literal
34

@@ -19,14 +20,28 @@
1920
SchemaField,
2021
)
2122
from backend.integrations.providers import ProviderName
22-
from backend.util.file import MediaFileType
23+
from backend.util.file import MediaFileType, store_media_file
2324

2425

2526
class GeminiImageModel(str, Enum):
2627
NANO_BANANA = "google/nano-banana"
2728
NANO_BANANA_PRO = "google/nano-banana-pro"
2829

2930

31+
class AspectRatio(str, Enum):
32+
MATCH_INPUT_IMAGE = "match_input_image"
33+
ASPECT_1_1 = "1:1"
34+
ASPECT_2_3 = "2:3"
35+
ASPECT_3_2 = "3:2"
36+
ASPECT_3_4 = "3:4"
37+
ASPECT_4_3 = "4:3"
38+
ASPECT_4_5 = "4:5"
39+
ASPECT_5_4 = "5:4"
40+
ASPECT_9_16 = "9:16"
41+
ASPECT_16_9 = "16:9"
42+
ASPECT_21_9 = "21:9"
43+
44+
3045
class OutputFormat(str, Enum):
3146
JPG = "jpg"
3247
PNG = "png"
@@ -69,6 +84,11 @@ class Input(BlockSchemaInput):
6984
default=[],
7085
title="Input Images",
7186
)
87+
aspect_ratio: AspectRatio = SchemaField(
88+
description="Aspect ratio of the generated image",
89+
default=AspectRatio.MATCH_INPUT_IMAGE,
90+
title="Aspect Ratio",
91+
)
7292
output_format: OutputFormat = SchemaField(
7393
description="Format of the output image",
7494
default=OutputFormat.PNG,
@@ -92,6 +112,7 @@ def __init__(self):
92112
"prompt": "Make the scene more vibrant and colorful",
93113
"model": GeminiImageModel.NANO_BANANA,
94114
"images": [],
115+
"aspect_ratio": AspectRatio.MATCH_INPUT_IMAGE,
95116
"output_format": OutputFormat.JPG,
96117
"credentials": TEST_CREDENTIALS_INPUT,
97118
},
@@ -116,11 +137,25 @@ async def run(
116137
**kwargs,
117138
) -> BlockOutput:
118139
try:
140+
# Convert local file paths to Data URIs (base64) so Replicate can access them
141+
processed_images = await asyncio.gather(
142+
*(
143+
store_media_file(
144+
graph_exec_id=graph_exec_id,
145+
file=img,
146+
user_id=user_id,
147+
return_content=True,
148+
)
149+
for img in input_data.images
150+
)
151+
)
152+
119153
result = await self.run_model(
120154
api_key=credentials.api_key,
121155
model_name=input_data.model.value,
122156
prompt=input_data.prompt,
123-
images=input_data.images,
157+
images=processed_images,
158+
aspect_ratio=input_data.aspect_ratio.value,
124159
output_format=input_data.output_format.value,
125160
)
126161
yield "image_url", result
@@ -133,12 +168,14 @@ async def run_model(
133168
model_name: str,
134169
prompt: str,
135170
images: list[MediaFileType],
171+
aspect_ratio: str,
136172
output_format: str,
137173
) -> MediaFileType:
138174
client = ReplicateClient(api_token=api_key.get_secret_value())
139175

140176
input_params: dict = {
141177
"prompt": prompt,
178+
"aspect_ratio": aspect_ratio,
142179
"output_format": output_format,
143180
}
144181

0 commit comments

Comments
 (0)