From 6d4d19948b28dcc7fdd5b814d82f6c92bad851dc Mon Sep 17 00:00:00 2001
From: shrimbly <falloon.willie@gmail.com>
Date: Wed, 18 Feb 2026 15:24:00 +1300
Subject: [PATCH] feat: add audio model detection to Replicate & WaveSpeed
 providers

Audio models (music, TTS, speech) were misclassified as text-to-image
because capability inference only checked for 3D and video keywords.
This caused audio models to not appear in GenerateAudio's model browser
and to create nanoBanana nodes instead of generateAudio nodes.

- Add audio keyword detection to inferReplicateCapabilities()
- Add audio keyword detection to inferWaveSpeedCapabilities()
- Add audio output handling to replicate.ts provider
- Add audio output handling to wavespeed.ts provider

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
---
 src/app/api/generate/providers/replicate.ts | 18 +++++++++++
 src/app/api/generate/providers/wavespeed.ts | 22 +++++++++++--
 src/app/api/models/route.ts                 | 35 +++++++++++++++++++++
 3 files changed, 73 insertions(+), 2 deletions(-)

diff --git a/src/app/api/generate/providers/replicate.ts b/src/app/api/generate/providers/replicate.ts
index 482d4d78..63c75edc 100644
--- a/src/app/api/generate/providers/replicate.ts
+++ b/src/app/api/generate/providers/replicate.ts
@@ -280,6 +280,8 @@ export async function generateWithReplicate(
   // Determine MIME type from response
   const contentType = mediaResponse.headers.get("content-type") || "image/png";
   const isVideo = contentType.startsWith("video/");
+  const isAudio = contentType.startsWith("audio/") ||
+    input.model.capabilities.some(c => c.includes("audio"));
 
   const mediaArrayBuffer = await mediaResponse.arrayBuffer();
   const mediaSizeBytes = mediaArrayBuffer.byteLength;
@@ -303,6 +305,22 @@ export async function generateWithReplicate(
   }
 
   const mediaBase64 = Buffer.from(mediaArrayBuffer).toString("base64");
+
+  if (isAudio) {
+    const audioContentType = contentType.startsWith("audio/") ? contentType : "audio/mpeg";
+    console.log(`[API:${requestId}] SUCCESS - Returning audio`);
+    return {
+      success: true,
+      outputs: [
+        {
+          type: "audio",
+          data: `data:${audioContentType};base64,${mediaBase64}`,
+          url: mediaUrl,
+        },
+      ],
+    };
+  }
+
   console.log(`[API:${requestId}] SUCCESS - Returning ${isVideo ? "video" : "image"}`);
 
   return {
diff --git a/src/app/api/generate/providers/wavespeed.ts b/src/app/api/generate/providers/wavespeed.ts
index 82febf22..1a98daa8 100644
--- a/src/app/api/generate/providers/wavespeed.ts
+++ b/src/app/api/generate/providers/wavespeed.ts
@@ -373,10 +373,11 @@ export async function generateWithWaveSpeed(
   const outputSizeMB = outputArrayBuffer.byteLength / (1024 * 1024);
 
   const rawContentType = outputResponse.headers.get("content-type");
+  const isAudioModel = input.model.capabilities.some(c => c.includes("audio"));
   const contentType =
-    (rawContentType && (rawContentType.startsWith("video/") || rawContentType.startsWith("image/")))
+    (rawContentType && (rawContentType.startsWith("video/") || rawContentType.startsWith("image/") || rawContentType.startsWith("audio/")))
       ? rawContentType
-      : (isVideoModel ? "video/mp4" : "image/png");
+      : (isVideoModel ? "video/mp4" : isAudioModel ? "audio/mpeg" : "image/png");
 
   console.log(`[API:${requestId}] Output: ${contentType}, ${outputSizeMB.toFixed(2)}MB`);
 
@@ -396,6 +397,23 @@ export async function generateWithWaveSpeed(
   }
 
   const outputBase64 = Buffer.from(outputArrayBuffer).toString("base64");
+  const isAudio = contentType.startsWith("audio/") || isAudioModel;
+
+  if (isAudio) {
+    const audioContentType = contentType.startsWith("audio/") ? contentType : "audio/mpeg";
+    console.log(`[API:${requestId}] SUCCESS - Returning audio`);
+    return {
+      success: true,
+      outputs: [
+        {
+          type: "audio",
+          data: `data:${audioContentType};base64,${outputBase64}`,
+          url: outputUrl,
+        },
+      ],
+    };
+  }
+
   console.log(`[API:${requestId}] SUCCESS - Returning ${isVideoModel ? "video" : "image"}`);
 
   return {
diff --git a/src/app/api/models/route.ts b/src/app/api/models/route.ts
index 35952835..9563d188 100644
--- a/src/app/api/models/route.ts
+++ b/src/app/api/models/route.ts
@@ -491,6 +491,23 @@ function inferReplicateCapabilities(model: ReplicateModel): ModelCapability[] {
     return capabilities;
   }
 
+  // Check for audio-related keywords
+  const isAudioModel =
+    searchText.includes("music") ||
+    searchText.includes("audio") ||
+    searchText.includes("tts") ||
+    searchText.includes("text-to-speech") ||
+    searchText.includes("speech") ||
+    searchText.includes("sound effect") ||
+    searchText.includes("voice") ||
+    searchText.includes("bark") ||
+    searchText.includes("xtts");
+
+  if (isAudioModel) {
+    capabilities.push("text-to-audio");
+    return capabilities;
+  }
+
   // Check for video-related keywords
   const isVideoModel =
     searchText.includes("video") ||
@@ -651,6 +668,24 @@ function inferWaveSpeedCapabilities(model: WaveSpeedModel): ModelCapability[] {
     return capabilities;
   }
 
+  // Check for audio-related keywords
+  const isAudioModel =
+    searchText.includes("music") ||
+    searchText.includes("audio") ||
+    searchText.includes("tts") ||
+    searchText.includes("text-to-speech") ||
+    searchText.includes("speech") ||
+    searchText.includes("sound effect") ||
+    searchText.includes("voice") ||
+    category.includes("audio") ||
+    category.includes("music") ||
+    category.includes("speech");
+
+  if (isAudioModel) {
+    capabilities.push("text-to-audio");
+    return capabilities;
+  }
+
   // Check for video-related keywords
   const isVideoModel =
     searchText.includes("video") ||