1+ import asyncio
12from enum import Enum
23from typing import Literal
34
1920 SchemaField ,
2021)
2122from backend .integrations .providers import ProviderName
22- from backend .util .file import MediaFileType
23+ from backend .util .file import MediaFileType , store_media_file
2324
2425
2526class 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+
3045class 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