Skip to content

Commit 9645e69

Browse files
committed
feat: Add ModelsLab image generation provider
- Add ModelsLab as new image generation provider (alongside OpenAI) - Support Flux, Juggernaut XL, RealVisXL, DreamShaper XL, SDXL models - Server-side route handles async polling for image generation - Restructure image_generator block with provider selector - Add conditional model/options based on provider selection
1 parent 115f04e commit 9645e69

File tree

6 files changed

+450
-14
lines changed

6 files changed

+450
-14
lines changed
Lines changed: 197 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,197 @@
1+
import { createLogger } from '@sim/logger'
2+
import { type NextRequest, NextResponse } from 'next/server'
3+
import { checkInternalAuth } from '@/lib/auth/hybrid'
4+
import { getInternalApiBaseUrl } from '@/lib/core/utils/urls'
5+
import type { ImageGenerationRequestBody } from '@/tools/image/types'
6+
7+
const logger = createLogger('ImageGenerateAPI')
8+
9+
export const dynamic = 'force-dynamic'
10+
export const maxDuration = 300 // 5 minutes for image generation with polling
11+
12+
export async function POST(request: NextRequest) {
13+
const requestId = crypto.randomUUID()
14+
logger.info(`[${requestId}] Image generation request started`)
15+
16+
try {
17+
const authResult = await checkInternalAuth(request, { requireWorkflowId: false })
18+
if (!authResult.success) {
19+
return NextResponse.json({ error: 'Unauthorized' }, { status: 401 })
20+
}
21+
22+
const body: ImageGenerationRequestBody = await request.json()
23+
const { provider, apiKey, model, prompt, width, height, negativePrompt } = body
24+
25+
if (!provider || !apiKey || !prompt) {
26+
return NextResponse.json(
27+
{ error: 'Missing required fields: provider, apiKey, and prompt' },
28+
{ status: 400 }
29+
)
30+
}
31+
32+
if (provider !== 'modelslab') {
33+
return NextResponse.json(
34+
{ error: `Unsupported provider: ${provider}. Currently supports: modelslab` },
35+
{ status: 400 }
36+
)
37+
}
38+
39+
if (prompt.length < 3 || prompt.length > 2000) {
40+
return NextResponse.json(
41+
{ error: 'Prompt must be between 3 and 2000 characters' },
42+
{ status: 400 }
43+
)
44+
}
45+
46+
const resolvedWidth = width && width > 0 ? width : 1024
47+
const resolvedHeight = height && height > 0 ? height : 1024
48+
49+
logger.info(`[${requestId}] Generating image with ModelsLab, model: ${model || 'flux'}`)
50+
51+
const result = await generateWithModelsLab(
52+
apiKey,
53+
prompt,
54+
model || 'flux',
55+
resolvedWidth,
56+
resolvedHeight,
57+
negativePrompt,
58+
requestId
59+
)
60+
61+
// Fetch the image and convert to base64 via existing image proxy
62+
let imageFile: string | undefined
63+
if (result.imageUrl) {
64+
try {
65+
const baseUrl = getInternalApiBaseUrl()
66+
const proxyUrl = new URL('/api/tools/image', baseUrl)
67+
proxyUrl.searchParams.append('url', result.imageUrl)
68+
69+
const { generateInternalToken } = await import('@/lib/auth/internal')
70+
const token = await generateInternalToken()
71+
72+
const imageResponse = await fetch(proxyUrl.toString(), {
73+
headers: {
74+
Accept: 'image/*, */*',
75+
Authorization: `Bearer ${token}`,
76+
},
77+
cache: 'no-store',
78+
})
79+
80+
if (imageResponse.ok) {
81+
const arrayBuffer = await imageResponse.arrayBuffer()
82+
if (arrayBuffer.byteLength > 0) {
83+
imageFile = Buffer.from(arrayBuffer).toString('base64')
84+
}
85+
}
86+
} catch (error) {
87+
logger.warn(`[${requestId}] Failed to fetch image for base64 conversion:`, error)
88+
// Non-fatal: still return the URL
89+
}
90+
}
91+
92+
logger.info(`[${requestId}] Image generation complete`)
93+
94+
return NextResponse.json({
95+
imageUrl: result.imageUrl,
96+
imageFile,
97+
model: model || 'flux',
98+
provider: 'modelslab',
99+
})
100+
} catch (error) {
101+
const errorMessage = error instanceof Error ? error.message : String(error)
102+
logger.error(`[${requestId}] Image generation error:`, { error: errorMessage })
103+
return NextResponse.json({ error: errorMessage }, { status: 500 })
104+
}
105+
}
106+
107+
async function generateWithModelsLab(
108+
apiKey: string,
109+
prompt: string,
110+
model: string,
111+
width: number,
112+
height: number,
113+
negativePrompt: string | undefined,
114+
requestId: string
115+
): Promise<{ imageUrl: string }> {
116+
logger.info(`[${requestId}] Calling ModelsLab text2img, model: ${model}`)
117+
118+
const requestBody: Record<string, unknown> = {
119+
key: apiKey,
120+
model_id: model,
121+
prompt,
122+
width,
123+
height,
124+
samples: 1,
125+
safety_checker: false,
126+
enhance_prompt: false,
127+
}
128+
129+
if (negativePrompt) {
130+
requestBody.negative_prompt = negativePrompt
131+
}
132+
133+
const createResponse = await fetch('https://modelslab.com/api/v6/images/text2img', {
134+
method: 'POST',
135+
headers: { 'Content-Type': 'application/json' },
136+
body: JSON.stringify(requestBody),
137+
})
138+
139+
if (!createResponse.ok) {
140+
const errText = await createResponse.text()
141+
throw new Error(`ModelsLab API error: ${createResponse.status} - ${errText}`)
142+
}
143+
144+
const createData = await createResponse.json()
145+
logger.info(`[${requestId}] ModelsLab response status: ${createData.status}`)
146+
147+
// Immediate success
148+
if (createData.status === 'success' && createData.output?.length > 0) {
149+
return { imageUrl: createData.output[0] }
150+
}
151+
152+
// Async processing — poll fetch endpoint
153+
if (createData.status === 'processing' && createData.id) {
154+
const jobId = String(createData.id)
155+
const maxAttempts = 40 // 40 × 5s = 200s max
156+
let attempts = 0
157+
158+
while (attempts < maxAttempts) {
159+
await new Promise((resolve) => setTimeout(resolve, 5000))
160+
161+
const fetchResponse = await fetch(
162+
`https://modelslab.com/api/v6/images/fetch/${jobId}`,
163+
{
164+
method: 'POST',
165+
headers: { 'Content-Type': 'application/json' },
166+
body: JSON.stringify({ key: apiKey }),
167+
}
168+
)
169+
170+
if (!fetchResponse.ok) {
171+
throw new Error(`ModelsLab fetch error: ${fetchResponse.status}`)
172+
}
173+
174+
const fetchData = await fetchResponse.json()
175+
logger.info(`[${requestId}] Poll ${attempts + 1}: status=${fetchData.status}`)
176+
177+
if (fetchData.status === 'success' && fetchData.output?.length > 0) {
178+
return { imageUrl: fetchData.output[0] }
179+
}
180+
181+
if (fetchData.status === 'error' || fetchData.status === 'failed') {
182+
throw new Error(`ModelsLab image generation failed: ${fetchData.message || 'Unknown error'}`)
183+
}
184+
185+
attempts++
186+
}
187+
188+
throw new Error('ModelsLab image generation timed out after 200 seconds')
189+
}
190+
191+
// Error response
192+
if (createData.status === 'error' || createData.error) {
193+
throw new Error(`ModelsLab API error: ${createData.message || createData.error || 'Unknown error'}`)
194+
}
195+
196+
throw new Error(`ModelsLab unexpected response: ${JSON.stringify(createData)}`)
197+
}

0 commit comments

Comments
 (0)