gork-imagine Closes #5216

3f8b9998ca5e2bb5081b0f94943e7c2de9caf68c

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

public/scripts/extensions/stable-diffusion/index.js+32 -2 Viewed

 async function loadXAIModels() {
     return [
+        { value: 'grok-imagine-image', text: 'grok-imagine-image' },
+        { value: 'grok-imagine-image-pro', text: 'grok-imagine-image-pro' },
         { value: 'grok-2-image-1212', text: 'grok-2-image-1212' },
     ];
+}

  * 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
+ * @param {'google'|'stability'|'zai'|'xai'} source Source of the request, used to determine aspect ratio
  * @returns {string} Closest aspect ratio as a string
  */
 function getClosestAspectRatio(width, height, source) {

                     '16:9': 16 / 9,
                     '9:16': 9 / 16,
                 };
+            case 'xai':
+                return {
+                    '1:1': 1,
+                    '3:4': 3 / 4,
+                    '4:3': 4 / 3,
+                    '9:16': 9 / 16,
+                    '16:9': 16 / 9,
+                    '2:3': 2 / 3,
+                    '3:2': 3 / 2,
+                    '9:19.5': 9 / 19.5,
+                    '19.5:9': 19.5 / 9,
+                    '9:20': 9 / 20,
+                    '20:9': 20 / 9,
+                    '1:2': 1 / 2,
+                    '2:1': 2 / 1,
+                };
             default:
                 console.warn(`Unknown source "${source}" for aspect ratio calculation.`);
                 return null;

  * @returns {Promise<{format: string, data: string}>} A promise that resolves when the image generation and processing are complete.
  */
 async function generateXAIImage(prompt, _negativePrompt, signal) {
+    let aspectRatio;
+    let resolution;
+    if (/grok-imagine/.test(extension_settings.sd.model)) {
+        const resolutionThreshold = 1296 * 864;
+        const use2kResolution = (extension_settings.sd.width * extension_settings.sd.height) > resolutionThreshold;
+        aspectRatio = getClosestAspectRatio(extension_settings.sd.width, extension_settings.sd.height, 'xai');
+        resolution = use2kResolution ? '2k' : '1k';
+    }
     const result = await fetch('/api/sd/xai/generate', {
         method: 'POST',
         headers: getRequestHeaders(),

         body: JSON.stringify({
             prompt: prompt,
             model: extension_settings.sd.model,
+            aspect_ratio: aspectRatio,
+            resolution: resolution,
         }),
     });
     if (result.ok) {
         const data = await result.json();
-        return { format: 'jpg', data: data.image };
+        return { format: data.format, data: data.image };
     } else {
         const text = await result.text();
         throw new Error(text);

src/endpoints/stable-diffusion.js+11 -3 Viewed

         const requestBody = {
             prompt: request.body.prompt,
             model: request.body.model,
+            aspect_ratio: request.body.aspect_ratio,
+            resolution: request.body.resolution,
             response_format: 'b64_json',
         };

         /** @type {any} */
         const data = await result.json();
-        const image = data?.data?.[0]?.b64_json;
+        // Can either be a base64 buffer (always JPEG) or a data URL (with MIME type)
-        if (!image) {
+        const encodedImage = String(data?.data?.[0]?.b64_json || '');
+        if (!encodedImage) {
             console.warn('xAI returned invalid data.');
             return response.sendStatus(500);
+        }
-        return response.send({ image });
+        const dataUrlMatch = encodedImage.match(/^data:(.+);base64,(.+)$/);
+        const mimeType = dataUrlMatch?.[1] || 'image/jpeg';
+        const format = mime.extension(mimeType) || 'jpg';
+        const image = dataUrlMatch?.[2] || encodedImage;
+        return response.send({ image, format });
     } catch (error) {
         console.error('Error communicating with xAI', error);
         return response.sendStatus(500);