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	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	2254	function loadNovelSchedulers() {
2240	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(/\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	4465	return secret_state[SECRET_KEYS.XAI];
4401	4466	case sources.google:
4402	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,aimlapizai">
 -                <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 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, 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);