diff --git a/plugins/ai-gateway/index.ts b/plugins/ai-gateway/index.ts
index d74c55c1..3d32d6fb 100644
--- a/plugins/ai-gateway/index.ts
+++ b/plugins/ai-gateway/index.ts
@@ -119,14 +119,28 @@ const aiGatewayPlugin: IntegrationPlugin = {
           key: "imageModel",
           label: "Model",
           type: "select",
-          defaultValue: "google/imagen-4.0-generate",
+          defaultValue: "google/imagen-4.0-generate-001",
           options: [
             {
-              value: "google/imagen-4.0-generate",
-              label: "Imagen 4.0 (Google)",
+              value: "google/imagen-4.0-generate-001",
+              label: "Imagen 4",
+            },
+            {
+              value: "google/imagen-4.0-fast-generate-001",
+              label: "Imagen 4 Fast",
+            },
+            {
+              value: "google/imagen-4.0-ultra-generate-001",
+              label: "Imagen 4 Ultra",
+            },
+            {
+              value: "bfl/flux-kontext-pro",
+              label: "FLUX.1 Kontext Pro",
+            },
+            {
+              value: "bfl/flux-kontext-max",
+              label: "FLUX.1 Kontext Max",
             },
-            { value: "openai/dall-e-3", label: "DALL-E 3 (OpenAI)" },
-            { value: "openai/dall-e-2", label: "DALL-E 2 (OpenAI)" },
           ],
         },
         {
diff --git a/plugins/ai-gateway/steps/generate-image.ts b/plugins/ai-gateway/steps/generate-image.ts
index 0217dd8b..4d1d83fd 100644
--- a/plugins/ai-gateway/steps/generate-image.ts
+++ b/plugins/ai-gateway/steps/generate-image.ts
@@ -40,7 +40,7 @@ async function stepHandler(
     };
   }
 
-  const modelId = input.imageModel || "google/imagen-4.0-generate";
+  const modelId = input.imageModel || "google/imagen-4.0-generate-001";
   const promptText = input.imagePrompt || "";
 
   if (!promptText || promptText.trim() === "") {