Z.AI: Add video generation models

ad860f4447c77d8893a19481c789a2f7129bfaa3

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

Parents: 3047045d3722bb7f6e59d7f98f1e6f3c1b33c59b

public/scripts/extensions/stable-diffusion/index.js+64 -10 Viewed

@@ -2297,7 +2297,7 @@ async function loadGoogleModels() {
2297	}	2297	}
2298		2298
2299	async function loadZaiModels() {	2299	async function loadZaiModels() {
2300	return ['cogview-4-250304'].map(name => ({ value: name, text: name }));	2300	return ['cogview-4-250304', 'cogvideox-3', 'viduq1-text'].map(name => ({ value: name, text: name }));
2301	}	2301	}
2302		2302
2303	async function loadOpenRouterModels() {	2303	async function loadOpenRouterModels() {

  * Gets an aspect ratio for Stability that is the closest to the given width and height.
  * @param {number} width Target width
  * @param {number} height Target height
- * @param {'google'|'stability'} source Source of the request, used to determine aspect ratio
+ * @param {'google'|'stability'|'zai'} source Source of the request, used to determine aspect ratio
  * @returns {string} Closest aspect ratio as a string
  */
 function getClosestAspectRatio(width, height, source) {

                     '4:3': 4 / 3,
                     '3:4': 3 / 4,
                 };
+            case 'zai':
+                return {
+                    '1:1': 1,
+                    '16:9': 16 / 9,
+                    '9:16': 9 / 16,
+                };
             default:
                 console.warn(`Unknown source "${source}" for aspect ratio calculation.`);
                 return null;

  * Get closest size for Electron Hub
  * @param {number} width - The width of the image
  * @param {number} height - The height of the image
+ * @param {string[]} sizes - Available sizes
  * @returns {Promise<string>} - The closest size
  */
-async function getClosestSize(width, height) {
+async function getClosestSize(width, height, sizes = []) {
+    const sizesData = [];
+    if (Array.isArray(sizes) && sizes.length > 0) {
+        sizesData.push(...sizes);
+    } else if (extension_settings.sd.source === sources.electronhub) {
         const response = await fetch('/api/sd/electronhub/sizes', {
             method: 'POST',
             headers: getRequestHeaders(),

             throw new Error(text);
+        }
         const result = await response.json();
-    const sizesData = result.sizes;
+        sizesData.push(...result.sizes);
+    } else {
+        return null;
+    }
+    const targetWidth = Number(width);
+    const targetHeight = Number(height);
+    if (isNaN(targetWidth) || isNaN(targetHeight)) {
+        return null;
+    }
+    const targetAspect = targetWidth / targetHeight;
+    const targetResolution = targetWidth * targetHeight;
     const closestSize = sizesData.reduce((closest, size) => {
         if (!size || typeof size !== 'string') {

         const sizeWidth = Number(sizeParts[0]);
         const sizeHeight = Number(sizeParts[1]);
-        const targetWidth = Number(width);
-        const targetHeight = Number(height);
-        if (isNaN(sizeWidth) || isNaN(sizeHeight) || isNaN(targetWidth) || isNaN(targetHeight)) {
+        if (isNaN(sizeWidth) || isNaN(sizeHeight)) {
             return closest;
+        }
-        const sizeArea = sizeWidth * sizeHeight;
+        const aspectDiff = Math.abs((sizeWidth / sizeHeight) - targetAspect) / targetAspect;
-        const targetArea = targetWidth * targetHeight;
+        const resolutionDiff = Math.abs(sizeWidth * sizeHeight - targetResolution) / targetResolution;
-        const diff = Math.abs(sizeArea - targetArea);
+        const diff = aspectDiff + resolutionDiff;
         return diff < closest.diff ? { size, diff } : closest;
     }, { size: null, diff: Infinity });

  * @returns {Promise<{format: string, data: string}>} A promise that resolves when the image generation and processing are complete.
  */
 async function generateZaiImage(prompt, signal) {
+    if (/(cogvideox|vidu)/.test(extension_settings.sd.model)) {
+        const videoParams = {};
+        if (/cogvideox/.test(extension_settings.sd.model)) {
+            const cogVideoSizes = ['1280x720', '720x1280', '1024x1024', '1080x1920', '2048x1080', '3840x2160'];
+            videoParams.quality = extension_settings.sd.openai_quality === 'hd' ? 'quality' : 'speed';
+            videoParams.size = await getClosestSize(extension_settings.sd.width, extension_settings.sd.height, cogVideoSizes);
+        }
+        if (/vidu/.test(extension_settings.sd.model)) {
+            videoParams.aspect_ratio = getClosestAspectRatio(extension_settings.sd.width, extension_settings.sd.height, 'zai');
+        }
+        const videoResult = await fetch('/api/sd/zai/generate-video', {
+            method: 'POST',
+            headers: getRequestHeaders(),
+            signal: signal,
+            body: JSON.stringify({
+                prompt: prompt,
+                model: extension_settings.sd.model,
+                ...videoParams,
+            }),
+        });
+        if (videoResult.ok) {
+            const data = await videoResult.json();
+            return { format: data.format, data: data.video };
+        }
+        const text = await videoResult.text();
+        throw new Error(text);
+    } else {
         // Round width and height to nearest multiple of 16, and clamp to 512-2048 range
         let width = clamp(Math.round(extension_settings.sd.width / 16) * 16, 512, 2048);
         let height = clamp(Math.round(extension_settings.sd.height / 16) * 16, 512, 2048);

         const text = await result.text();
         throw new Error(text);
+    }
+}
 /**
  * Generates an image using the OpenRouter API.

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

                             <option value="high" data-i18n="High">High</option>
                         </select>
                     </div>
-                    <div data-sd-model="dall-e-3,cogview-4" class="flex1">
+                    <div data-sd-model="dall-e-3,cogview-4,cogvideox" class="flex1">
                         <label for="sd_openai_quality" data-i18n="Image Quality">Image Quality</label>
                         <select id="sd_openai_quality">
                             <option value="standard" data-i18n="Standard">Standard</option>

src/endpoints/stable-diffusion.js+101 -0 Viewed

+    }
 });
+zai.post('/generate-video', async (request, response) => {
+    try {
+        const controller = new AbortController();
+        request.socket.removeAllListeners('close');
+        request.socket.on('close', function () {
+            controller.abort();
+        });
+        const key = readSecret(request.user.directories, SECRET_KEYS.ZAI);
+        if (!key) {
+            console.warn('Z.AI key not found.');
+            return response.sendStatus(400);
+        }
+        console.debug('Z.AI video request:', request.body);
+        const generateResponse = await fetch('https://api.z.ai/api/paas/v4/videos/generations', {
+            method: 'POST',
+            headers: {
+                'Content-Type': 'application/json',
+                'Authorization': `Bearer ${key}`,
+            },
+            body: JSON.stringify({
+                prompt: request.body.prompt,
+                model: request.body.model,
+                quality: request.body.quality,
+                size: request.body.size,
+                aspect_ratio: request.body.aspect_ratio,
+            }),
+            signal: controller.signal,
+        });
+        if (!generateResponse.ok) {
+            const text = await generateResponse.text();
+            console.warn('Z.AI returned an error.', text);
+            return response.sendStatus(500);
+        }
+        /** @type {any} */
+        const data = await generateResponse.json();
+        console.debug('Z.AI video response:', data);
+        // Poll for video generation completion
+        for (let attempt = 0; attempt < 30; attempt++) {
+            if (controller.signal.aborted) {
+                console.info('Z.AI video generation aborted by client');
+                return response.status(500).send('Video generation aborted by client');
+            }
+            await delay(5000 + attempt * 1000);
+            console.debug(`Polling Z.AI video job ${data.id}, attempt ${attempt + 1}`);
+            const pollResponse = await fetch(`https://api.z.ai/api/paas/v4/async-result/${data.id}`, {
+                method: 'GET',
+                headers: {
+                    'Authorization': `Bearer ${key}`,
+                },
+            });
+            if (!pollResponse.ok) {
+                const text = await pollResponse.text();
+                console.warn('Z.AI video job polling failed', pollResponse.statusText, text);
+                return response.status(500).send(text);
+            }
+            /** @type {any} */
+            const pollResult = await pollResponse.json();
+            console.debug(`Z.AI video job status: ${pollResult.task_status}`);
+            if (pollResult.task_status === 'FAIL') {
+                console.warn('Z.AI video generation failed', pollResult);
+                return response.status(500).send('Video generation failed');
+            }
+            if (pollResult.task_status === 'SUCCESS') {
+                console.debug('Z.AI video generation succeeded', pollResult);
+                const url = pollResult?.video_result?.[0]?.url;
+                if (!url || !isValidUrl(url)) {
+                    console.warn('Z.AI returned an invalid video URL.');
+                    return response.sendStatus(500);
+                }
+                const contentResponse = await fetch(url);
+                if (!contentResponse.ok) {
+                    const text = await contentResponse.text();
+                    console.warn('Z.AI video content fetch failed', contentResponse.statusText, text);
+                    return response.status(500).send(text);
+                }
+                const contentBuffer = await contentResponse.arrayBuffer();
+                return response.send({ format: 'mp4', video: Buffer.from(contentBuffer).toString('base64') });
+            }
+        }
+    } catch (error) {
+        console.error(error);
+        return response.sendStatus(500);
+    }
+});
 router.use('/comfy', comfy);
 router.use('/comfyrunpod', comfyRunPod);
 router.use('/together', together);