Google: Add Veo models (#4742) * Google: Add Veo models * Combine retry into a singe loop * Improve error logging and reporting * Add client abort handling for video generation endpoint

ad72fa806eb0e82d61feec890332f7a0218cd6b3

Cohee <18619528+Cohee1207@users.noreply.github.com>

Signed

Parents: 5a2b325ff7b4fb23a548912ef50e13b03f608569

public/scripts/extensions/stable-diffusion/index.js+48 -0 Viewed

     // Google settings
     google_api: 'makersuite',
     google_enhance: true,
+    google_duration: 6,
 };
 const writePromptFieldsDebounced = debounce(writePromptFields, debounce_timeout.relaxed);

     $('#sd_bfl_upsampling').prop('checked', extension_settings.sd.bfl_upsampling);
     $('#sd_google_api').val(extension_settings.sd.google_api);
     $('#sd_google_enhance').prop('checked', extension_settings.sd.google_enhance);
+    $('#sd_google_duration').val(extension_settings.sd.google_duration);
     for (const style of extension_settings.sd.styles) {
         const option = document.createElement('option');

 async function loadGoogleModels() {
     return [
+        'imagen-4.0-generate-001',
+        'imagen-4.0-ultra-generate-001',
+        'imagen-4.0-fast-generate-001',
         'imagen-4.0-generate-preview-06-06',
         'imagen-4.0-fast-generate-preview-06-06',
         'imagen-4.0-ultra-generate-preview-06-06',

         'imagegeneration@006',
         'imagegeneration@005',
         'imagegeneration@002',
+        'veo-3.1-generate-preview',
+        'veo-3.1-fast-generate-preview',
+        'veo-3.0-generate-001',
+        'veo-3.0-fast-generate-001',
+        'veo-2.0-generate-001',
+        'veo-2.0-generate-exp',
+        'veo-2.0-generate-preview',
     ].map(name => ({ value: name, text: name }));
+}

  * @returns {Promise<{format: string, data: string}>} A promise that resolves when the image generation and processing are complete.
  */
 async function generateGoogleImage(prompt, negativePrompt, signal) {
+    const isVeo = /^veo-/.test(extension_settings.sd.model);
+    if (isVeo) {
+        const aspectRatio = extension_settings.sd.width / extension_settings.sd.height;
+        const maxPromptLength = 3000; // 1024 tokens approx.
+        const videoResult = await fetch('/api/google/generate-video', {
+            method: 'POST',
+            headers: getRequestHeaders(),
+            signal: signal,
+            body: JSON.stringify({
+                prompt: prompt.slice(0, maxPromptLength),
+                aspect_ratio: aspectRatio >= 1 ? '16:9' : '9:16',
+                seconds: extension_settings.sd.google_duration,
+                negative_prompt: negativePrompt,
+                model: extension_settings.sd.model,
+                api: extension_settings.sd.google_api || 'makersuite',
+                seed: extension_settings.sd.seed >= 0 ? extension_settings.sd.seed : undefined,
+                vertexai_auth_mode: oai_settings.vertexai_auth_mode,
+                vertexai_region: oai_settings.vertexai_region,
+                vertexai_express_project_id: oai_settings.vertexai_express_project_id,
+            }),
+        });
+        if (!videoResult.ok) {
+            const text = await videoResult.text();
+            throw new Error(text);
+        }
+        const data = await videoResult.json();
+        return { format: 'mp4', data: data.video };
+    }
     const result = await fetch('/api/google/generate-image', {
         method: 'POST',
         headers: getRequestHeaders(),

         extension_settings.sd.google_enhance = $(this).prop('checked');
         saveSettingsDebounced();
     });
+    $('#sd_google_duration').on('input', function () {
+        extension_settings.sd.google_duration = Number($(this).val());
+        saveSettingsDebounced();
+    });
     $('#sd_electronhub_quality').on('change', function () {
         extension_settings.sd.electronhub_quality = String($(this).val());
         saveSettingsDebounced();

public/scripts/extensions/stable-diffusion/settings.html+9 -1 Viewed

                         </select>
                     </div>
                 </div>
-                <div class="flex-container">
+                <div class="flex-container alignItemsCenter">
                     <label class="flex1 checkbox_label" for="sd_google_enhance" data-i18n="[title]Enables prompt enhancing (passes prompts through an LLM to add detail)." title="Enables prompt enhancing (passes prompts through an LLM to add detail).">
                         <input id="sd_google_enhance" type="checkbox" />
                         <span data-i18n="Enhance">
                             Enhance
                         </span>
                     </label>
+                    <div class="flex1">
+                        <label for="sd_google_duration" data-i18n="Duration (Veo)">Duration (Veo)</label>
+                        <select id="sd_google_duration">
+                            <option value="4">Short (4 seconds)</option>
+                            <option value="6">Medium (6 seconds)</option>
+                            <option value="8">Long (8 seconds)</option>
+                        </select>
+                    </div>
                 </div>
             </div>

src/endpoints/google.js+169 -15 Viewed

 import express from 'express';
 import { speak, languages } from 'google-translate-api-x';
 import crypto from 'node:crypto';
+import util from 'node:util';
+import urlJoin from 'url-join';
+import lodash from 'lodash';
 import { readSecret, SECRET_KEYS } from './secrets.js';
 import { GEMINI_SAFETY } from '../constants.js';
-import { getConfigValue, trimTrailingSlash } from '../util.js';
+import { delay, getConfigValue, trimTrailingSlash } from '../util.js';
 const API_MAKERSUITE = 'https://generativelanguage.googleapis.com';
 const API_VERTEX_AI = 'https://us-central1-aiplatform.googleapis.com';

  * @param {express.Request} request Express request object
  * @param {string} model Model name to use
  * @param {string} endpoint API endpoint (default: 'generateContent')
- * @returns {Promise<{url: string, headers: object, apiName: string}>} URL, headers, and API name
+ * @returns {Promise<{url: string, headers: object, apiName: string, baseUrl: string}>} URL, headers, and API name
  */
 export async function getGoogleApiConfig(request, model, endpoint = 'generateContent') {
     const useVertexAi = request.body.api === 'vertexai';

     const apiName = useVertexAi ? 'Google Vertex AI' : 'Google AI Studio';
     let url;
+    let baseUrl;
     let headers = {
         'Content-Type': 'application/json',
     };

             // Express mode: use API key parameter
             const keyParam = authHeader.replace('Bearer ', '');
             const projectId = request.body.vertexai_express_project_id;
-            const baseUrl = region === 'global'
+            baseUrl = region === 'global'
-                ? 'https://aiplatform.googleapis.com'
+                ? 'https://aiplatform.googleapis.com/v1'
-                : `https://${region}-aiplatform.googleapis.com`;
+                : `https://${region}-aiplatform.googleapis.com/v1`;
             url = projectId
-                ? `https://aiplatform.googleapis.com/v1/projects/${projectId}/locations/${region}/publishers/google/models/${model}:${endpoint}?key=${keyParam}`
+                ? `${baseUrl}/projects/${projectId}/locations/${region}/publishers/google/models/${model}:${endpoint}`
-                : `${baseUrl}/v1/publishers/google/models/${model}:${endpoint}?key=${keyParam}`;
+                : `${baseUrl}/publishers/google/models/${model}:${endpoint}`;
+            headers['x-goog-api-key'] = keyParam;
         } else if (authType === 'full') {
             // Full mode: use project-specific URL with Authorization header
             // Get project ID from Service Account JSON

                 throw new Error('Failed to extract project ID from Service Account JSON.');
+            }
             // Handle global region differently - no region prefix in hostname
-            url = region === 'global'
+            baseUrl = region === 'global'
-                ? `https://aiplatform.googleapis.com/v1/projects/${projectId}/locations/${region}/publishers/google/models/${model}:${endpoint}`
+                ? 'https://aiplatform.googleapis.com/v1'
-                : `https://${region}-aiplatform.googleapis.com/v1/projects/${projectId}/locations/${region}/publishers/google/models/${model}:${endpoint}`;
+                : `https://${region}-aiplatform.googleapis.com/v1`;
+            url = `${baseUrl}/projects/${projectId}/locations/${region}/publishers/google/models/${model}:${endpoint}`;
             headers['Authorization'] = authHeader;
         } else {
             // Proxy mode: use Authorization header
             const apiUrl = trimTrailingSlash(request.body.reverse_proxy || API_VERTEX_AI);
-            url = `${apiUrl}/v1/publishers/google/models/${model}:${endpoint}`;
+            baseUrl = `${apiUrl}/v1`;
+            url = `${baseUrl}/publishers/google/models/${model}:${endpoint}`;
             headers['Authorization'] = authHeader;
+        }
     } else {

         const apiKey = request.body.reverse_proxy ? request.body.proxy_password : readSecret(request.user.directories, SECRET_KEYS.MAKERSUITE);
         const apiUrl = trimTrailingSlash(request.body.reverse_proxy || API_MAKERSUITE);
         const apiVersion = getConfigValue('gemini.apiVersion', 'v1beta');
-        url = `${apiUrl}/${apiVersion}/models/${model}:${endpoint}?key=${apiKey}`;
+        baseUrl = `${apiUrl}/${apiVersion}`;
+        url = `${baseUrl}/models/${model}:${endpoint}`;
+        headers['x-goog-api-key'] = apiKey;
+    }
-    return { url, headers, apiName };
+    return { url, headers, apiName, baseUrl };
+}
 export const router = express.Router();

         if (!result.ok) {
             const errorText = await result.text();
             console.warn(`${apiName} image generation error: ${result.status} ${result.statusText}`, errorText);
-            return response.sendStatus(500);
+            return response.status(500).send('Image generation request failed');
+        }
         /** @type {any} */

         if (!imagePart) {
             console.warn(`${apiName} image generation error: No image data found in response`);
-            return response.sendStatus(500);
+            return response.status(500).send('No image data found in response');
+        }
         return response.send({ image: imagePart });

         return response.end();
+    }
 });
+router.post('/generate-video', async (request, response) => {
+    try {
+        const controller = new AbortController();
+        request.socket.removeAllListeners('close');
+        request.socket.on('close', function () {
+            controller.abort();
+        });
+        const model = request.body.model || 'veo-3.1-generate-preview';
+        const { url, headers, apiName, baseUrl } = await getGoogleApiConfig(request, model, 'predictLongRunning');
+        const useVertexAi = request.body.api === 'vertexai';
+        const isVeo3 = /veo-3/.test(model);
+        const lowerBound = isVeo3 ? 4 : 5;
+        const upperBound = isVeo3 ? 8 : 8;
+        const requestBody = {
+            instances: [{
+                prompt: String(request.body.prompt || ''),
+            }],
+            parameters: {
+                negativePrompt: String(request.body.negative_prompt || ''),
+                durationSeconds: lodash.clamp(Number(request.body.seconds || 6), lowerBound, upperBound),
+                aspectRatio: String(request.body.aspect_ratio || '16:9'),
+                personGeneration: 'allow_all',
+                seed: isVeo3 ? Number(request.body.seed ?? Math.floor(Math.random() * 1000000)) : undefined,
+            },
+        };
+        console.debug(`${apiName} video generation request:`, model, requestBody);
+        const videoJobResponse = await fetch(url, {
+            method: 'POST',
+            headers: headers,
+            body: JSON.stringify(requestBody),
+        });
+        if (!videoJobResponse.ok) {
+            const errorText = await videoJobResponse.text();
+            console.warn(`${apiName} video generation error: ${videoJobResponse.status} ${videoJobResponse.statusText}`, errorText);
+            return response.status(500).send('Video generation request failed');
+        }
+        /** @type {any} */
+        const videoJobData = await videoJobResponse.json();
+        const videoJobName = videoJobData?.name;
+        if (!videoJobName) {
+            console.warn(`${apiName} video generation error: No job name found in response`);
+            return response.status(500).send('No video job name found in response');
+        }
+        console.debug(`${apiName} video job name:`, videoJobName);
+        for (let attempt = 0; attempt < 30; attempt++) {
+            if (controller.signal.aborted) {
+                console.info(`${apiName} video generation aborted by client`);
+                return response.status(500).send('Video generation aborted by client');
+            }
+            await delay(5000 + attempt * 1000);
+            if (useVertexAi) {
+                const { url: pollUrl, headers: pollHeaders } = await getGoogleApiConfig(request, model, 'fetchPredictOperation');
+                const pollResponse = await fetch(pollUrl, {
+                    method: 'POST',
+                    headers: pollHeaders,
+                    body: JSON.stringify({ operationName: videoJobName }),
+                });
+                if (!pollResponse.ok) {
+                    const errorText = await pollResponse.text();
+                    console.warn(`${apiName} video job status error: ${pollResponse.status} ${pollResponse.statusText}`, errorText);
+                    return response.status(500).send('Video job status request failed');
+                }
+                /** @type {any} */
+                const pollData = await pollResponse.json();
+                const jobDone = pollData?.done;
+                console.debug(`${apiName} video job status attempt ${attempt + 1}: ${jobDone ? 'done' : 'running'}`);
+                if (jobDone) {
+                    const videoData = pollData?.response?.videos?.[0]?.bytesBase64Encoded;
+                    if (!videoData) {
+                        const pollDataLog = util.inspect(pollData, { depth: 5, colors: true, maxStringLength: 500 });
+                        console.warn(`${apiName} video generation error: No video data found in response`, pollDataLog);
+                        return response.status(500).send('No video data found in response');
+                    }
+                    return response.send({ video: videoData });
+                }
+            } else {
+                const pollUrl = urlJoin(baseUrl, videoJobName);
+                const pollResponse = await fetch(pollUrl, {
+                    method: 'GET',
+                    headers: headers,
+                });
+                if (!pollResponse.ok) {
+                    const errorText = await pollResponse.text();
+                    console.warn(`${apiName} video job status error: ${pollResponse.status} ${pollResponse.statusText}`, errorText);
+                    return response.status(500).send('Video job status request failed');
+                }
+                /** @type {any} */
+                const pollData = await pollResponse.json();
+                const jobDone = pollData?.done;
+                console.debug(`${apiName} video job status attempt ${attempt + 1}: ${jobDone ? 'done' : 'running'}`);
+                if (jobDone) {
+                    const videoUri = pollData?.response?.generateVideoResponse?.generatedSamples?.[0]?.video?.uri;
+                    console.debug(`${apiName} video URI:`, videoUri);
+                    if (!videoUri) {
+                        const pollDataLog = util.inspect(pollData, { depth: 5, colors: true, maxStringLength: 500 });
+                        console.warn(`${apiName} video generation error: No video URI found in response`, pollDataLog);
+                        return response.status(500).send('No video URI found in response');
+                    }
+                    const videoResponse = await fetch(videoUri, {
+                        method: 'GET',
+                        headers: headers,
+                    });
+                    if (!videoResponse.ok) {
+                        console.warn(`${apiName} video fetch error: ${videoResponse.status} ${videoResponse.statusText}`);
+                        return response.status(500).send('Video fetch request failed');
+                    }
+                    const videoData = await videoResponse.arrayBuffer();
+                    const videoBase64 = Buffer.from(videoData).toString('base64');
+                    return response.send({ video: videoBase64 });
+                }
+            }
+        }
+        console.warn(`${apiName} video generation error: Job timed out after multiple attempts`);
+        return response.status(500).send('Video generation timed out');
+    } catch (error) {
+        console.error('Google Video generation failed:', error);
+        return response.sendStatus(500);
+    }
+});