Browse Source
Move generateWithReplicate() into providers/replicate.ts. It imports schema utilities from schemaUtils.ts. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>handoff-20260429-1057
shrimbly
5 months ago
2 changed files with 281 additions and 268 deletions
@ -0,0 +1,280 @@ |
|||
/** |
|||
* Replicate Provider for Generate API Route |
|||
* |
|||
* Handles image/video generation using Replicate's prediction API. |
|||
*/ |
|||
|
|||
import { GenerationInput, GenerationOutput } from "@/lib/providers/types"; |
|||
import { |
|||
getParameterTypesFromSchema, |
|||
coerceParameterTypes, |
|||
getInputMappingFromSchema, |
|||
} from "../schemaUtils"; |
|||
|
|||
/** |
|||
* Generate image using Replicate API |
|||
*/ |
|||
export async function generateWithReplicate( |
|||
requestId: string, |
|||
apiKey: string, |
|||
input: GenerationInput |
|||
): Promise<GenerationOutput> { |
|||
console.log(`[API:${requestId}] Replicate generation - Model: ${input.model.id}, Images: ${input.images?.length || 0}, Prompt: ${input.prompt.length} chars`); |
|||
|
|||
const REPLICATE_API_BASE = "https://api.replicate.com/v1"; |
|||
|
|||
// Get the latest version of the model
|
|||
const modelId = input.model.id; |
|||
const [owner, name] = modelId.split("/"); |
|||
|
|||
// First, get the model to find the latest version
|
|||
const modelResponse = await fetch( |
|||
`${REPLICATE_API_BASE}/models/${owner}/${name}`, |
|||
{ |
|||
headers: { |
|||
Authorization: `Bearer ${apiKey}`, |
|||
}, |
|||
} |
|||
); |
|||
|
|||
if (!modelResponse.ok) { |
|||
return { |
|||
success: false, |
|||
error: `Failed to get model info: ${modelResponse.status}`, |
|||
}; |
|||
} |
|||
|
|||
const modelData = await modelResponse.json(); |
|||
const version = modelData.latest_version?.id; |
|||
|
|||
if (!version) { |
|||
return { |
|||
success: false, |
|||
error: "Model has no available version", |
|||
}; |
|||
} |
|||
|
|||
const hasDynamicInputs = input.dynamicInputs && Object.keys(input.dynamicInputs).length > 0; |
|||
console.log(`[API:${requestId}] Model version: ${version}, Dynamic inputs: ${hasDynamicInputs ? Object.keys(input.dynamicInputs!).join(", ") : "none"}`); |
|||
|
|||
// Get schema for type coercion and input mapping
|
|||
const schema = modelData.latest_version?.openapi_schema as Record<string, unknown> | undefined; |
|||
const parameterTypes = getParameterTypesFromSchema(schema); |
|||
|
|||
// Build input for the prediction, coercing parameter types from schema
|
|||
const predictionInput: Record<string, unknown> = { |
|||
...coerceParameterTypes(input.parameters, parameterTypes), |
|||
}; |
|||
|
|||
// Add dynamic inputs if provided (these come from schema-mapped connections)
|
|||
if (hasDynamicInputs) { |
|||
const { schemaArrayParams } = getInputMappingFromSchema(schema); |
|||
|
|||
// Apply array wrapping based on schema type
|
|||
for (const [key, value] of Object.entries(input.dynamicInputs!)) { |
|||
if (value !== null && value !== undefined && value !== '') { |
|||
if (schemaArrayParams.has(key) && !Array.isArray(value)) { |
|||
predictionInput[key] = [value]; // Wrap in array
|
|||
} else { |
|||
predictionInput[key] = value; |
|||
} |
|||
} |
|||
} |
|||
} else { |
|||
// Fallback: use schema to map generic input names to model-specific parameter names
|
|||
const { paramMap, arrayParams } = getInputMappingFromSchema(schema); |
|||
|
|||
// Map prompt input
|
|||
if (input.prompt) { |
|||
const promptParam = paramMap.prompt || "prompt"; |
|||
predictionInput[promptParam] = input.prompt; |
|||
} |
|||
|
|||
// Map image input - use array or string format based on schema
|
|||
if (input.images && input.images.length > 0) { |
|||
const imageParam = paramMap.image || "image"; |
|||
if (arrayParams.has("image")) { |
|||
predictionInput[imageParam] = input.images; |
|||
} else { |
|||
predictionInput[imageParam] = input.images[0]; |
|||
} |
|||
} |
|||
|
|||
// Map any parameters that might need renaming (use coerced values)
|
|||
const coercedParams = coerceParameterTypes(input.parameters, parameterTypes); |
|||
for (const [key, value] of Object.entries(coercedParams)) { |
|||
const mappedKey = paramMap[key] || key; |
|||
predictionInput[mappedKey] = value; |
|||
} |
|||
} |
|||
|
|||
// Create a prediction
|
|||
const createResponse = await fetch(`${REPLICATE_API_BASE}/predictions`, { |
|||
method: "POST", |
|||
headers: { |
|||
Authorization: `Bearer ${apiKey}`, |
|||
"Content-Type": "application/json", |
|||
}, |
|||
body: JSON.stringify({ |
|||
version, |
|||
input: predictionInput, |
|||
}), |
|||
}); |
|||
|
|||
if (!createResponse.ok) { |
|||
const errorText = await createResponse.text(); |
|||
let errorDetail = errorText; |
|||
try { |
|||
const errorJson = JSON.parse(errorText); |
|||
errorDetail = errorJson.detail || errorJson.message || errorJson.error || errorText; |
|||
} catch { |
|||
// Keep original text if not JSON
|
|||
} |
|||
|
|||
// Handle rate limits
|
|||
if (createResponse.status === 429) { |
|||
return { |
|||
success: false, |
|||
error: `${input.model.name}: Rate limit exceeded. Try again in a moment.`, |
|||
}; |
|||
} |
|||
|
|||
return { |
|||
success: false, |
|||
error: `${input.model.name}: ${errorDetail}`, |
|||
}; |
|||
} |
|||
|
|||
const prediction = await createResponse.json(); |
|||
console.log(`[API:${requestId}] Prediction created: ${prediction.id}`); |
|||
|
|||
// Poll for completion
|
|||
const maxWaitTime = 5 * 60 * 1000; // 5 minutes
|
|||
const pollInterval = 1000; // 1 second
|
|||
const startTime = Date.now(); |
|||
|
|||
let currentPrediction = prediction; |
|||
let lastStatus = ""; |
|||
|
|||
while ( |
|||
currentPrediction.status !== "succeeded" && |
|||
currentPrediction.status !== "failed" && |
|||
currentPrediction.status !== "canceled" |
|||
) { |
|||
if (Date.now() - startTime > maxWaitTime) { |
|||
return { |
|||
success: false, |
|||
error: `${input.model.name}: Generation timed out after 5 minutes. Video models may take longer - try again.`, |
|||
}; |
|||
} |
|||
|
|||
await new Promise((resolve) => setTimeout(resolve, pollInterval)); |
|||
|
|||
const pollResponse = await fetch( |
|||
`${REPLICATE_API_BASE}/predictions/${currentPrediction.id}`, |
|||
{ |
|||
headers: { |
|||
Authorization: `Bearer ${apiKey}`, |
|||
}, |
|||
} |
|||
); |
|||
|
|||
if (!pollResponse.ok) { |
|||
return { |
|||
success: false, |
|||
error: `Failed to poll prediction: ${pollResponse.status}`, |
|||
}; |
|||
} |
|||
|
|||
currentPrediction = await pollResponse.json(); |
|||
if (currentPrediction.status !== lastStatus) { |
|||
console.log(`[API:${requestId}] Prediction status: ${currentPrediction.status}`); |
|||
lastStatus = currentPrediction.status; |
|||
} |
|||
} |
|||
|
|||
if (currentPrediction.status === "failed") { |
|||
const failureReason = currentPrediction.error || "Prediction failed"; |
|||
return { |
|||
success: false, |
|||
error: `${input.model.name}: ${failureReason}`, |
|||
}; |
|||
} |
|||
|
|||
if (currentPrediction.status === "canceled") { |
|||
return { |
|||
success: false, |
|||
error: "Prediction was canceled", |
|||
}; |
|||
} |
|||
|
|||
// Extract output
|
|||
const output = currentPrediction.output; |
|||
if (!output) { |
|||
return { |
|||
success: false, |
|||
error: "No output from prediction", |
|||
}; |
|||
} |
|||
|
|||
// Output can be a single URL string or an array of URLs
|
|||
const outputUrls: string[] = Array.isArray(output) ? output : [output]; |
|||
|
|||
if (outputUrls.length === 0) { |
|||
return { |
|||
success: false, |
|||
error: "No output from prediction", |
|||
}; |
|||
} |
|||
|
|||
// Fetch the first output and convert to base64
|
|||
const mediaUrl = outputUrls[0]; |
|||
console.log(`[API:${requestId}] Fetching output from: ${mediaUrl.substring(0, 80)}...`); |
|||
const mediaResponse = await fetch(mediaUrl); |
|||
|
|||
if (!mediaResponse.ok) { |
|||
return { |
|||
success: false, |
|||
error: `Failed to fetch output: ${mediaResponse.status}`, |
|||
}; |
|||
} |
|||
|
|||
// Determine MIME type from response
|
|||
const contentType = mediaResponse.headers.get("content-type") || "image/png"; |
|||
const isVideo = contentType.startsWith("video/"); |
|||
|
|||
const mediaArrayBuffer = await mediaResponse.arrayBuffer(); |
|||
const mediaSizeBytes = mediaArrayBuffer.byteLength; |
|||
const mediaSizeMB = mediaSizeBytes / (1024 * 1024); |
|||
|
|||
console.log(`[API:${requestId}] Output: ${contentType}, ${mediaSizeMB.toFixed(2)}MB`); |
|||
|
|||
// For very large videos (>20MB), return URL directly instead of base64
|
|||
if (isVideo && mediaSizeMB > 20) { |
|||
console.log(`[API:${requestId}] SUCCESS - Returning URL for large video`); |
|||
return { |
|||
success: true, |
|||
outputs: [ |
|||
{ |
|||
type: "video", |
|||
data: mediaUrl, // Return URL directly for very large videos
|
|||
url: mediaUrl, |
|||
}, |
|||
], |
|||
}; |
|||
} |
|||
|
|||
const mediaBase64 = Buffer.from(mediaArrayBuffer).toString("base64"); |
|||
console.log(`[API:${requestId}] SUCCESS - Returning ${isVideo ? "video" : "image"}`); |
|||
|
|||
return { |
|||
success: true, |
|||
outputs: [ |
|||
{ |
|||
type: isVideo ? "video" : "image", |
|||
data: `data:${contentType};base64,${mediaBase64}`, |
|||
url: mediaUrl, |
|||
}, |
|||
], |
|||
}; |
|||
} |
|||
Loading…
Reference in new issue