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'|'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.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);