Add xAI for image generation extension

0c4c86ef0620dd72749dc7fd339c6f6b517b4626

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

Parents: c1544fb60cc547fddc31179478e5bec3daea29ce

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

     nanogpt: 'nanogpt',
     bfl: 'bfl',
     falai: 'falai',
 +    xai: 'xai',
 };
 const initiators = {

         sources.nanogpt,
         sources.bfl,
         sources.falai,
 +        sources.xai,
     ];
     if (cloudSources.includes(extension_settings.sd.source)) {

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

         case sources.falai:
             models = await loadFalaiModels();
             break;
 +        case sources.xai:
 +            models = await loadXAIModels();
 +            break;
+    }
     for (const model of models) {

     return [];
+}
 +async function loadXAIModels() {
 +    return [
 +        { value: 'grok-2-image-1212', text: 'grok-2-image-1212' },
 +    ];
 +}
++
 async function loadPollinationsModels() {
     const result = await fetch('/api/sd/pollinations/models', {
         method: 'POST',

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

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

             case sources.falai:
                 result = await generateFalaiImage(prefixedPrompt, negativePrompt, signal);
                 break;
 +            case sources.xai:
 +                result = await generateXAIImage(prefixedPrompt, negativePrompt, signal);
 +                break;
+        }
         if (!result.data) {

+}
 /**
 + * Generates an image using the xAI API.
 + * @param {string} prompt The main instruction used to guide the image generation.
 + * @param {string} _negativePrompt Negative prompt is not used in this API
 + * @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 generateXAIImage(prompt, _negativePrompt, signal) {
 +    const result = await fetch('/api/sd/xai/generate', {
 +        method: 'POST',
 +        headers: getRequestHeaders(),
 +        signal: signal,
 +        body: JSON.stringify({
 +            prompt: prompt,
 +            model: extension_settings.sd.model,
 +        }),
 +    });
++
 +    if (result.ok) {
 +        const data = await result.json();
 +        return { format: 'jpg', data: data.image };
 +    } else {
 +        const text = await result.text();
 +        throw new Error(text);
 +    }
 +}
++
 +/**
  * Generates an image using the FAL.AI API.
  * @param {string} prompt - The main instruction used to guide the image generation.
  * @param {string} negativePrompt - The negative prompt used to guide the image generation.

@@ -3782,6 +3835,8 @@ function isValidState() {
3782	3835	return secret_state[SECRET_KEYS.BFL];
3783	3836	case sources.falai:
3784	3837	return secret_state[SECRET_KEYS.FALAI];
	3838	+ case sources.xai:
	3839	+ return secret_state[SECRET_KEYS.XAI];
3785	3840	}
3786	3841	}
3787	3842

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

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

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

                         'Authorization': `Key ${key}`,
                     },
                 });
 +                /** @type {any} */
                 const resultData = await resultFetch.json();
                 if (resultData.detail !== null && resultData.detail !== undefined) {

+    }
 });
 +const xai = express.Router();
++
 +xai.post('/generate', async (request, response) => {
 +    try {
 +        const key = readSecret(request.user.directories, SECRET_KEYS.XAI);
++
 +        if (!key) {
 +            console.warn('xAI key not found.');
 +            return response.sendStatus(400);
 +        }
++
 +        const requestBody = {
 +            prompt: request.body.prompt,
 +            model: request.body.model,
 +            response_format: 'b64_json',
 +        };
++
 +        console.debug('xAI request:', requestBody);
++
 +        const result = await fetch('https://api.x.ai/v1/images/generations', {
 +            method: 'POST',
 +            body: JSON.stringify(requestBody),
 +            headers: {
 +                'Content-Type': 'application/json',
 +                'Authorization': `Bearer ${key}`,
 +            },
 +        });
++
 +        if (!result.ok) {
 +            const text = await result.text();
 +            console.warn('xAI returned an error.', text);
 +            return response.sendStatus(500);
 +        }
++
 +        /** @type {any} */
 +        const data = await result.json();
++
 +        const image = data?.data?.[0]?.b64_json;
 +        if (!image) {
 +            console.warn('xAI returned invalid data.');
 +            return response.sendStatus(500);
 +        }
++
 +        return response.send({ image });
 +    } catch (error) {
 +        console.error('Error communicating with xAI', error);
 +        return response.sendStatus(500);
 +    }
 +});
++
 router.use('/comfy', comfy);
 router.use('/together', together);
 router.use('/drawthings', drawthings);

 router.use('/nanogpt', nanogpt);
 router.use('/bfl', bfl);
 router.use('/falai', falai);
 +router.use('/xai', xai);