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	return secret_state[SECRET_KEYS.BFL];	3835	return secret_state[SECRET_KEYS.BFL];
3783	case sources.falai:	3836	case sources.falai:
3784	return secret_state[SECRET_KEYS.FALAI];	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);