Z.AI: Add video generation models

ad860f4447c77d8893a19481c789a2f7129bfaa3

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

Parents: 3047045d3722bb7f6e59d7f98f1e6f3c1b33c59b

public/scripts/extensions/stable-diffusion/index.js+103 -49 Viewed

@@ -2297,7 +2297,7 @@ async function loadGoogleModels() {
2297	2297	}
2298	2298
2299	2299	async function loadZaiModels() {
2300	2300	return ['cogview-4-250304', 'cogvideox-3', 'viduq1-text'].map(name => ({ value: name, text: name }));
2301	2301	}
2302	2302
2303	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'|'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, sizes = []) {
 -    const response = await fetch('/api/sd/electronhub/sizes', {
 +    const sizesData = [];
 -        method: 'POST',
++
 -        headers: getRequestHeaders(),
 +    if (Array.isArray(sizes) && sizes.length > 0) {
 -        body: JSON.stringify({
 +        sizesData.push(...sizes);
 -            model: extension_settings.sd.model,
 +    } else if (extension_settings.sd.source === sources.electronhub) {
 -        }),
 +        const response = await fetch('/api/sd/electronhub/sizes', {
 -    });
 +            method: 'POST',
 -    if (!response.ok) {
 +            headers: getRequestHeaders(),
 -        const text = await response.text();
 +            body: JSON.stringify({
 -        throw new Error(text);
 +                model: extension_settings.sd.model,
 +            }),
 +        });
 +        if (!response.ok) {
 +            const text = await response.text();
 +            throw new Error(text);
 +        }
 +        const result = await response.json();
 +        sizesData.push(...result.sizes);
 +    } else {
 +        return null;
 +    }
++
 +    const targetWidth = Number(width);
 +    const targetHeight = Number(height);
++
 +    if (isNaN(targetWidth) || isNaN(targetHeight)) {
 +        return null;
+    }
 -    const result = await response.json();
++
     const sizesDatatargetAspect = result.sizestargetWidth / 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)) {
             return closest;
+        }
         const sizeAreaaspectDiff = Math.abs((sizeWidth */ sizeHeight) - targetAspect) / targetAspect;
         const targetArearesolutionDiff = targetWidthMath.abs(sizeWidth * targetHeightsizeHeight - targetResolution) / targetResolution;
         const diff = Math.abs(sizeAreaaspectDiff -+ targetArea)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) {
 -    // Round width and height to nearest multiple of 16, and clamp to 512-2048 range
 +    if (/(cogvideox|vidu)/.test(extension_settings.sd.model)) {
 -    let width = clamp(Math.round(extension_settings.sd.width / 16) * 16, 512, 2048);
 +        const videoParams = {};
 -    let height = clamp(Math.round(extension_settings.sd.height / 16) * 16, 512, 2048);
 +        if (/cogvideox/.test(extension_settings.sd.model)) {
+-
 +            const cogVideoSizes = ['1280x720', '720x1280', '1024x1024', '1080x1920', '2048x1080', '3840x2160'];
 -    // Make sure the pixel count does not exceed 2^21px
 +            videoParams.quality = extension_settings.sd.openai_quality === 'hd' ? 'quality' : 'speed';
 -    while ((width * height) > Math.pow(2, 21)) {
 +            videoParams.size = await getClosestSize(extension_settings.sd.width, extension_settings.sd.height, cogVideoSizes);
 -        if (width >= height) {
 +        }
 -            width -= 16;
 +        if (/vidu/.test(extension_settings.sd.model)) {
 -        } else {
 +            videoParams.aspect_ratio = getClosestAspectRatio(extension_settings.sd.width, extension_settings.sd.height, 'zai');
 -            height -= 16;
+        }
 -    }
             const resultvideoResult = await fetch('/api/sd/zai/generate-video', {
                     method: 'POST',
                     headers: getRequestHeaders(),
                     signal: signal,
                     body: JSON.stringify({
                             prompt: prompt,
                             model: extension_settings.sd.model,
 -            quality: extension_settings.sd.openai_quality,
 +                ...videoParams,
 -            size: `${width}x${height}`,
 +            }),
         }),;
 -    });
             if (resultvideoResult.ok) {
                     const data = await resultvideoResult.json();
                     return { format: data.format, data: data.imagevideo };
+            }
             const text = await resultvideoResult.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);
++
 +        // Make sure the pixel count does not exceed 2^21px
 +        while ((width * height) > Math.pow(2, 21)) {
 +            if (width >= height) {
 +                width -= 16;
 +            } else {
 +                height -= 16;
 +            }
 +        }
++
 +        const result = await fetch('/api/sd/zai/generate', {
 +            method: 'POST',
 +            headers: getRequestHeaders(),
 +            signal: signal,
 +            body: JSON.stringify({
 +                prompt: prompt,
 +                model: extension_settings.sd.model,
 +                quality: extension_settings.sd.openai_quality,
 +                size: `${width}x${height}`,
 +            }),
 +        });
++
 +        if (result.ok) {
 +            const data = await result.json();
 +            return { format: data.format, data: data.image };
 +        }
++
 +        const text = await result.text();
 +        throw new Error(text);
 +    }
+}
 /**

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,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);