Z.AI: Add image generation and web search (#4895) * Z.AI: Add image generation * Z.AI: Add web search endpoint * Add image URL validation * Combine validation conditions

bb3ea08b5a3872d9cb4fcf5dab9fef4b80fcd8be

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

Signed

Parents: 000fbccda950466ceb55b4c2ba1cc96f57c8f682

public/scripts/extensions/stable-diffusion/index.js+70 -1 Viewed

     falai: 'falai',
     xai: 'xai',
     google: 'google',
+    zai: 'zai',
 };
 const initiators = {

         sources.xai,
         sources.google,
         sources.chutes,
+        sources.zai,
     ];
     if (cloudSources.includes(extension_settings.sd.source)) {

         case sources.bfl:
             samplers = ['N/A'];
             break;
+        case sources.falai:
+            samplers = ['N/A'];
+            break;
         case sources.xai:
             samplers = ['N/A'];
             break;
         case sources.google:
             samplers = ['N/A'];
             break;
+        case sources.zai:
+            samplers = ['N/A'];
+            break;
+    }
     for (const sampler of samplers) {

         case sources.google:
             models = await loadGoogleModels();
             break;
+        case sources.zai:
+            models = await loadZaiModels();
+            break;
+    }
     if (extension_settings.sd.source === sources.electronhub) {

@@ -2236,6 +2247,10 @@ async function loadGoogleModels() {
2236	].map(name => ({ value: name, text: name }));	2247	].map(name => ({ value: name, text: name }));
2237	}	2248	}
2238		2249
		2250	async function loadZaiModels() {
		2251	return ['cogview-4-250304'].map(name => ({ value: name, text: name }));
		2252	}
		2253
2239	function loadNovelSchedulers() {	2254	function loadNovelSchedulers() {
2240	return ['karras', 'native', 'exponential', 'polyexponential'];	2255	return ['karras', 'native', 'exponential', 'polyexponential'];
2241	}	2256	}

         case sources.google:
             schedulers = ['N/A'];
             break;
+        case sources.zai:
+            schedulers = ['N/A'];
+            break;
+    }
     for (const scheduler of schedulers) {

         case sources.google:
             vaes = ['N/A'];
             break;
+        case sources.zai:
+            vaes = ['N/A'];
+            break;
+    }
     for (const vae of vaes) {

     str = str.normalize('NFD');
     // Strip out non-alphanumeric characters barring model syntax exceptions
-    str = str.replace(/[^a-zA-Z0-9.,:_(){}<>[\]\-'|#]+/g, ' ');
+    str = str.replace(/[^a-zA-Z0-9.,:_(){}<>[\]/\-'|#]+/g, ' ');
     str = str.replace(/\s+/g, ' '); // Collapse multiple whitespaces into one
     str = str.trim();

             case sources.google:
                 result = await generateGoogleImage(prefixedPrompt, negativePrompt, signal);
                 break;
+            case sources.zai:
+                result = await generateZaiImage(prefixedPrompt, signal);
+                break;
+        }
         if (!result.data) {

+    }
+}
+/**
+ * Generates an image using the Z.AI API.
+ * @param {string} prompt The main instruction used to guide the image generation.
+ * @param {AbortSignal} signal An AbortSignal object that can be used to cancel the request.
+ * @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
+    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);
+}
 async function onComfyOpenWorkflowEditorClick() {
     let workflow = await (await fetch('/api/sd/comfy/workflow', {
         method: 'POST',

@@ -4400,6 +4465,10 @@ function isValidState() {
4400	return secret_state[SECRET_KEYS.XAI];	4465	return secret_state[SECRET_KEYS.XAI];
4401	case sources.google:	4466	case sources.google:
4402	return secret_state[SECRET_KEYS.MAKERSUITE] \|\| secret_state[SECRET_KEYS.VERTEXAI] \|\| secret_state[SECRET_KEYS.VERTEXAI_SERVICE_ACCOUNT];	4467	return secret_state[SECRET_KEYS.MAKERSUITE] \|\| secret_state[SECRET_KEYS.VERTEXAI] \|\| secret_state[SECRET_KEYS.VERTEXAI_SERVICE_ACCOUNT];
		4468	case sources.zai:
		4469	return secret_state[SECRET_KEYS.ZAI];
		4470	default:
		4471	return false;
4403	}	4472	}
4404	}	4473	}
4405		4474

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

                 <option value="horde">Stable Horde</option>
                 <option value="togetherai">TogetherAI</option>
                 <option value="xai">xAI (Grok)</option>
+                <option value="zai">Z.AI (CogView)</option>
             </select>
             <div data-sd-source="auto">
                 <label for="sd_auto_url">SD Web UI URL</label>

                     </div>
                 </div>
             </div>
-            <div data-sd-source="openai,aimlapi">
+            <div data-sd-source="zai">
-                <div data-sd-model="dall-e-3" class="flex-container">
+                <b>Will use Common API. Coding API is not supported!</b>
-                    <div class="flex1">
+            </div>
+            <div data-sd-source="openai,aimlapi,zai">
+                <div class="flex-container">
+                    <div data-sd-model="dall-e-3" class="flex1">
                         <label for="sd_openai_style" data-i18n="Image Style">Image Style</label>
                         <select id="sd_openai_style">
                             <option value="vivid">Vivid</option>
                             <option value="natural">Natural</option>
                         </select>
                     </div>
-                    <div class="flex1">
+                    <div data-sd-model="dall-e-3,cogview-4-250304" 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/openai.js+1 -0 Viewed

         if (request.body.api === 'zai') {
             key = readSecret(request.user.directories, SECRET_KEYS.ZAI);
+            bodyParams.max_tokens = 4096; // default is 1024
+        }
         const noKeyTypes = ['custom', 'ooba', 'koboldcpp', 'vllm', 'llamacpp', 'pollinations'];

src/endpoints/search.js+46 -0 Viewed

+    }
 });
+router.post('/zai', async (request, response) => {
+    try {
+        const key = readSecret(request.user.directories, SECRET_KEYS.ZAI);
+        if (!key) {
+            console.error('No Z.AI key found');
+            return response.sendStatus(400);
+        }
+        const { query } = request.body;
+        if (!query) {
+            console.error('No query provided for /zai');
+            return response.sendStatus(400);
+        }
+        console.debug('Z.AI web search query', query);
+        const result = await fetch('https://api.z.ai/api/paas/v4/web_search', {
+            method: 'POST',
+            headers: {
+                'Content-Type': 'application/json',
+                'Authorization': `Bearer ${key}`,
+            },
+            body: JSON.stringify({
+                // TODO: There's only one engine option for now
+                search_engine: 'search-prime',
+                search_query: query,
+            }),
+        });
+        if (!result.ok) {
+            const text = await result.text();
+            console.error('Z.AI request failed', result.statusText, text);
+            return response.status(500).send(text);
+        }
+        const data = await result.json();
+        console.debug('Z.AI web search response', data);
+        return response.json(data);
+    } catch (error) {
+        console.error(error);
+        return response.sendStatus(500);
+    }
+});
 router.post('/visit', async (request, response) => {
     try {
         const url = request.body.url;

src/endpoints/stable-diffusion.js+62 -1 Viewed

 import urlJoin from 'url-join';
 import _ from 'lodash';
-import { delay, getBasicAuthHeader, tryParse } from '../util.js';
+import { delay, getBasicAuthHeader, isValidUrl, tryParse } from '../util.js';
 import { readSecret, SECRET_KEYS } from './secrets.js';
 import { AIMLAPI_HEADERS } from '../constants.js';

+    }
 });
+const zai = express.Router();
+zai.post('/generate', async (request, response) => {
+    try {
+        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 image request:', request.body);
+        const generateResponse = await fetch('https://api.z.ai/api/paas/v4/images/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,
+            }),
+        });
+        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 image response:', data);
+        const url = data?.data?.[0]?.url;
+        if (!url || !isValidUrl(url) || !new URL(url).hostname.endsWith('.z.ai')) {
+            console.warn('Z.AI returned an invalid image URL.');
+            return response.sendStatus(500);
+        }
+        const imageResponse = await fetch(url);
+        if (!imageResponse.ok) {
+            console.warn('Z.AI image fetch returned an error.');
+            return response.sendStatus(500);
+        }
+        const buffer = await imageResponse.arrayBuffer();
+        const image = Buffer.from(buffer).toString('base64');
+        const format = path.extname(url).substring(1).toLowerCase() || 'png';
+        return response.send({ image, format });
+    } catch (error) {
+        console.error(error);
+        return response.sendStatus(500);
+    }
+});
 router.use('/comfy', comfy);
 router.use('/together', together);
 router.use('/drawthings', drawthings);

 router.use('/falai', falai);
 router.use('/xai', xai);
 router.use('/aimlapi', aimlapi);
+router.use('/zai', zai);