add grok to cloud models

Signed-off-by: Vladimir Mandic <mandic00@live.com>

CHANGELOG.md

@@ -7,9 +7,10 @@
 All-about-optimizations: 
 - improved LoRA performance and quality, especially with quantized models
 - newly structured attention mechanisms
-- modular pipelines with new guidance
+- modular pipelines with new guidance methods
 - support for different caching stacks
 - compute updates across the board
+- enhanced cloud model support
 
 ### Details for 2026-09-07
 
@@ -51,7 +52,7 @@ All-about-optimizations:
   - implement progress and preview
   - intercept and profiling hooks
   - on-demand convert standard model on-demand 
-- **Google**  
+- **Cloud**
   - updated support for google models in text, image and video workflows  
     *note*: requires google api key  
   - [Google Veo](https://ai.google.dev/gemini-api/docs/veo) in *preview*, *fast* and *lite* variants  
@@ -62,6 +63,10 @@ All-about-optimizations:
     workflows: *caption*  
   - [Google Gemini](https://ai.google.dev/gemini-api/docs/models/gemini-3.8-flash) in *flash* and *pro* variants  
     workflows: *caption, prompt-enhance*  
+  - added support for xai grok models  
+  - [X.AI Grok](https://x.ai/grok) in *3*, *3 fast*, *3 mini* and *3 mini fast* variants  
+    workflows: *caption, prompt-enhance*  
+    *note*: requires grok api key  
 - **Compute**
   - cuda: update `torch==2.14.0` with `cuda==13.2`
   - openvino: update `openvino==2026.3.1` with `torch==2.13.0`
@@ -75,7 +80,9 @@ All-about-optimizations:
   - update `numpy` and `scipy` frozen requirements as required by new compute drivers  
     *note*: this may break compatibility with some legacy packages, so report any finidings  
 - **Other**
-  - Video Preview: TAESD support for **MiniMax**
+  - video preview: TAESD support for **MiniMax**
+  - support `xai grok` for prompt enhance workflows  
+    *note*: requires grok api key  
   - remove `/redocs` as `/docs` are primary api docs  
   - rebuild docs site index  
 - **Wiki/Docs**:
@@ -99,6 +106,7 @@ All-about-optimizations:
   - vae: fetch scale factor from the model
   - todo: remove dead code, thanks @Anai-Guo
   - offline: honor offline mode for more models, thanks @ryanmeador
+  - prompt enhance: cloud models use correct system prompt  
 
 ## Update for 2026-08-26
 

modules/caption/gemini.py

@@ -3,6 +3,7 @@ import os
 from modules import shared
 from modules.logger import log
 
+
 debug_enabled = os.environ.get('SD_CAPTION_DEBUG', None) is not None
 debug_log = log.trace if debug_enabled else lambda *args, **kwargs: None
 
@@ -63,11 +64,13 @@ class GoogleGeminiPipeline():
             config['temperature'] = kwargs['temperature']
         if 'max_output_tokens' in kwargs:
             config['max_output_tokens'] = kwargs['max_output_tokens']
-        debug_log(f'Gemini config: {config}')
+        debug_log(f'LLM config: {config}')
+        debug_log(f'LLM instructions: "{instructions}"')
+        debug_log(f'LLM image: {image}')
         question = question.replace('<', '').replace('>', '').replace('_', ' ')
         if prefill:
             question += prefill
-        debug_log(f'Gemini question: "{question}"')
+        debug_log(f'LLM question: "{question}"')
 
         if image:
             data = io.BytesIO()
@@ -80,9 +83,9 @@ class GoogleGeminiPipeline():
         answer = ''
         try:
             response = self.client.models.generate_content(
-                model=model,
-                contents=contents,
-                config=config,
+                model = model,
+                contents = contents,
+                config = config,
             )
             debug_log(f'Gemini response: {response}')
             answer = response.text
@@ -94,8 +97,8 @@ class GoogleGeminiPipeline():
 
 ai = None
 
-def predict(question, image, vqa_model, system_prompt, model_name, prefill, thinking, gen_kwargs):
+def predict(question, image, model_name, system_prompt, prefill, thinking, gen_kwargs):
     global ai # pylint: disable=global-statement
     if ai is None:
         ai = GoogleGeminiPipeline(model_name)
-    return ai(question, image, vqa_model, system_prompt, prefill, thinking, gen_kwargs)
+    return ai(question, image, model_name, system_prompt, prefill, thinking, gen_kwargs)

modules/caption/grok.py

@@ -0,0 +1,99 @@
+import io
+import os
+import base64
+from modules import shared
+from modules.logger import log
+
+
+debug_enabled = os.environ.get('SD_CAPTION_DEBUG', None) is not None
+debug_log = log.trace if debug_enabled else lambda *args, **kwargs: None
+
+
+class XAIGrokPipeline():
+    def __init__(self, model_name: str):
+        self.url = 'https://api.x.ai/v1'
+        self.model = model_name.split(' (')[0].replace('xai/', '')
+        from installer import install
+        install('openai')
+        from openai import OpenAI # pylint: disable=no-name-in-module
+        args = self.get_args()
+        if not args:
+            return
+        self.client = OpenAI(**args)
+        log.debug(f'Load model: type=XAIGrok model="{self.model}"')
+
+    def get_args(self):
+        from modules.shared import opts
+        # Use UI settings only - env vars are intentionally ignored
+        api_key = opts.xai_api_key
+        has_api_key = api_key and len(api_key) > 0
+        if not has_api_key: # Gemini Developer API: api_key only
+            log.error(f'Cloud: model="{self.model}" API key not provided')
+            return None
+        args = {
+            'api_key': api_key,
+            'base_url': self.url,
+        }
+        # Debug logging
+        args_log = args.copy()
+        if args_log.get('api_key'):
+            args_log['api_key'] = '...' + args_log['api_key'][-4:]
+        log.debug(f'Cloud: model="{self.model}" args={args_log}')
+        return args
+
+    def __call__(self, question, image, model, instructions, prefill, thinking, kwargs):
+        question = question.replace('<', '').replace('>', '').replace('_', ' ')
+        if prefill:
+            question += prefill
+        debug_log(f'LLM instructions: "{instructions}"')
+        debug_log(f'LLM question: "{question}"')
+        debug_log(f'LLM image: {image}')
+        answer = ''
+        temperature = kwargs.get('temperature', 0.0)
+        try:
+            if image is not None:
+                image_data = image.convert('RGB')
+                image_bytes = io.BytesIO()
+                image_data.save(image_bytes, format='JPEG')
+                image_bytes.seek(0)
+                content = [
+                    {
+                        "type": "image_url",
+                        "image_url": {
+                            "url": f"data:image/jpeg;base64,{base64.b64encode(image_bytes.getvalue()).decode('utf-8')}",
+                            "detail": "high",
+                        },
+                    },
+                    {
+                        "type": "text",
+                        "text": question,
+                    },
+                ]
+            else:
+                content = question
+            response = self.client.chat.completions.create(
+                model = self.model,
+                messages = [
+                    {"role": "system", "content": instructions or shared.opts.caption_vlm_system},
+                    {"role": "user", "content": content},
+                ],
+                stream = False,
+                temperature = temperature,
+                reasoning_effort = "high" if thinking else "low"
+            )
+            text = (response.choices[0].message.content or "").strip()
+            debug_log(f'Grok response: {response}')
+            answer = text
+        except Exception as e:
+            log.error(f'Grok: {e}')
+            answer = f'Error: {e}'
+        return answer
+
+
+ai = None
+
+def predict(question, image, model_name, system_prompt, prefill, thinking, gen_kwargs):
+    global ai # pylint: disable=global-statement
+    if ai is None:
+        ai = XAIGrokPipeline(model_name)
+    return ai(question, image, model_name, system_prompt, prefill, thinking, gen_kwargs)

modules/caption/models_def.py

@@ -81,6 +81,10 @@ vlm_models = {
     f"Google Gemini 3.5 Flash Lite {ui_symbols.cloud}": "gemini-3.5-flash-lite",
     f"Google Gemini 3.1 Flash Lite {ui_symbols.cloud}": "gemini-3.1-flash-lite",
     f"Google Gemini 3.1 Pro {ui_symbols.cloud}": "gemini-3.1-pro-preview",
+    f"X.AI Grok 3 {ui_symbols.cloud}": "grok-3-latest",
+    f"X.AI Grok 3 Fast {ui_symbols.cloud}": "grok-3-fast-latest",
+    f"X.AI Grok 3 Mini {ui_symbols.cloud}": "grok-3-mini-latest",
+    f"X.AI Grok 3 Mini Fast {ui_symbols.cloud}": "grok-3-mini-fast-latest",
 }
 
 # Default model
@@ -224,5 +228,10 @@ Summary:
 
 def get_vlm_repo(display_name: str) -> str:
     """Look up repo ID from display name, stripping any trailing symbols."""
+    from modules.logger import log
     name = display_name.strip()
-    return vlm_models.get(name, name)
+    model = vlm_models.get(name, None)
+    if model is None:
+        log.warning(f"Model '{name}' not found")
+        return name
+    return model

modules/caption/vqa.py

@@ -1597,7 +1597,12 @@ class VQA:
                 handler = 'gemini'
                 gen_kwargs = get_kwargs(self.model)
                 from modules.caption import gemini
-                answer = gemini.predict(question, image, vqa_model, system_prompt, model_name, prefill, thinking_mode, gen_kwargs)
+                answer = gemini.predict(question, image, vqa_model, system_prompt, prefill, thinking_mode, gen_kwargs)
+            elif 'grok' in vqa_model.lower():
+                handler = 'grok'
+                gen_kwargs = get_kwargs(self.model)
+                from modules.caption import grok
+                answer = grok.predict(question, image, vqa_model, system_prompt, prefill, thinking_mode, gen_kwargs)
             else:
                 answer = 'unknown model'
         except Exception as e:

modules/ui_definitions.py

@@ -113,6 +113,8 @@ def create_settings(cmd_opts):
         "google_api_key": OptionInfo("", "Google cloud API key", gr.Textbox, secret=True, env_var='GOOGLE_API_KEY'),
         "google_project_id": OptionInfo("", "Google Cloud project ID", gr.Textbox, secret=True, env_var='GOOGLE_PROJECT_ID'),
         "google_location_id": OptionInfo("", "Google Cloud location ID", gr.Textbox),
+        "model_xai_sep": OptionInfo("<h2>XAI Grok</h2>", "", gr.HTML),
+        "xai_api_key": OptionInfo("", "XAI Grok API key", gr.Textbox, secret=True, env_var='XAI_API_KEY'),
         "model_krea2_sep": OptionInfo("<h2>Krea 2</h2>", "", gr.HTML),
         "model_krea2_dense": OptionInfo(False, "Use dense masking"),
         "model_sd3_sep": OptionInfo("<h2>Stable Diffusion 3.x</h2>", "", gr.HTML),

scripts/prompt_enhance/options.py

@@ -43,6 +43,10 @@ class Options:
         'google/gemini-3.5-flash-lite',
         'google/gemini-3.1-flash-lite',
         'google/gemini-3.1-pro-preview',
+        'xai/grok-3-latest',
+        'xai/grok-3-fast-latest',
+        'xai/grok-3-mini-latest',
+        'xai/grok-3-mini-fast-latest',
     ]
     models = {
         # Gemma
@@ -90,6 +94,11 @@ class Options:
         'google/gemini-3.5-flash-lite': {},
         'google/gemini-3.1-flash-lite': {},
         'google/gemini-3.1-pro-preview': {},
+        # Grok
+        'xai/grok-3-latest': {},
+        'xai/grok-3-fast-latest': {},
+        'xai/grok-3-mini-latest': {},
+        'xai/grok-3-mini-fast-latest': {},
         # SmolLM
         'HuggingFaceTB/SmolLM2-135M-Instruct': {},
         'HuggingFaceTB/SmolLM2-360M-Instruct': {},

scripts/prompt_enhance/template.py

@@ -9,13 +9,24 @@ debug_enabled = os.environ.get('SD_LLM_DEBUG', None) is not None
 debug_log = log.trace if debug_enabled else lambda *args, **kwargs: None
 
 
-def get_text_template(system, prompt, options, nsfw, has_system, has_prompt, has_processor, is_video, _image) -> list[dict]:
-    if not has_system:
-        system = options.t2v_prompt if is_video else options.t2i_prompt
-        system += options.nsfw_ok if nsfw else options.nsfw_no
-        system += options.details_prompt
-        system += options.details_format
-        debug_log(f'Prompt enhance: system="{system}"')
+def get_system_prompt(system: str | None, options: Options, nsfw: bool, has_prompt: bool, is_video: bool, is_image: bool) -> str:
+    if system is not None and len(system) > 4:
+        return system
+    if is_video:
+        system = options.t2v_prompt if has_prompt else options.t2v_noprompt
+    elif is_image:
+        system = options.i2i_prompt if has_prompt else options.i2i_noprompt
+    else:
+        system = options.t2i_prompt if has_prompt else options.t2i_noprompt
+    system += options.nsfw_ok if nsfw else options.nsfw_no
+    system += options.details_prompt
+    system += options.details_format
+    debug_log(f'Prompt enhance: system="{system}"')
+    return system
+
+
+def get_text_template(system, prompt, options, nsfw, has_prompt, has_processor, is_video, _image) -> list[dict]:
+    system = get_system_prompt(system, options, nsfw, has_prompt, is_video, is_image=False)
     if not has_prompt:
         prompt = 'be creative!'
     if not has_processor:
@@ -35,16 +46,8 @@ def get_text_template(system, prompt, options, nsfw, has_system, has_prompt, has
     return chat_template
 
 
-def get_image_template(system, prompt, options, nsfw, has_system, has_prompt, _has_processor, is_video, image) -> list[dict]:
-    if not has_system:
-        if is_video:
-            system = options.i2v_prompt if has_prompt else options.i2v_noprompt
-        else:
-            system = options.i2i_prompt if has_prompt else options.i2i_noprompt
-        system += options.nsfw_ok if nsfw else options.nsfw_no
-        system += options.details_prompt
-        system += options.details_format
-        debug_log(f'Prompt enhance: system="{system}"')
+def get_image_template(system, prompt, options, nsfw, has_prompt, _has_processor, is_video, image) -> list[dict]:
+    system = get_system_prompt(system, options, nsfw, has_prompt, is_video, is_image=True)
     if has_prompt:
         chat_template = [
             { "role": "system", "content": [
@@ -78,12 +81,11 @@ def set_template(
         module: str | None = None,
 ) -> list[dict] | str:
     chat_template = []
-    has_system = system is not None and len(system) > 4
     has_prompt = prompt is not None and len(prompt) > 4
     has_image = image is not None and isinstance(image, Image.Image)
     is_video = module == 'video'
 
-    debug_log(f'Prompt enhance template: module={module} system={has_system} prompt={has_prompt} image={has_image} video={is_video} model="{model}" nsfw={nsfw} processor={has_processor}')
+    debug_log(f'Prompt enhance template: module={module} prompt={has_prompt} image={has_image} video={is_video} model="{model}" nsfw={nsfw} processor={has_processor}')
 
     if has_image:
         if is_cloud_model(model):
@@ -93,8 +95,8 @@ def set_template(
             return prompt if prompt is not None else '' # Return original text part if image cannot be processed
 
     if has_image:
-        chat_template = get_image_template(system, prompt, options, nsfw, has_system, has_prompt, has_processor, is_video, image)
+        chat_template = get_image_template(system, prompt, options, nsfw, has_prompt, has_processor, is_video, image)
     else:
-        chat_template = get_text_template(system, prompt, options, nsfw, has_system, has_prompt, has_processor, is_video, image)
+        chat_template = get_text_template(system, prompt, options, nsfw, has_prompt, has_processor, is_video, image)
 
     return chat_template

scripts/prompt_enhance_ext.py

@@ -14,7 +14,7 @@ from modules.caption.logits import LogitsParser
 from modules.caption import helpers
 from scripts.prompt_enhance.options import Options
 from scripts.prompt_enhance.helpers import is_cloud_model, is_vision_model, is_thinking_model, get_model_repo_from_display
-from scripts.prompt_enhance.template import set_template
+from scripts.prompt_enhance.template import set_template, get_system_prompt
 
 
 debug_enabled = os.environ.get('SD_LLM_DEBUG', None) is not None
@@ -389,6 +389,45 @@ class PromptEnhanceScript(scripts_manager.Script):
                     current_image = current_image.convert('RGB')
                     debug_log('Prompt enhance: Converted image to RGB mode')
 
+        # Prepare prefill (VQA approach: string concatenation, not assistant message)
+        prefill_text = (prefill or '').strip()
+
+        t0 = time.time()
+        self.busy = True
+
+        if is_cloud_model(model):
+            has_prompt = prompt_text is not None and len(prompt_text) > 4
+            system = get_system_prompt(system, self.options, nsfw, has_prompt=has_prompt, is_video=self.parent=='video', is_image=current_image is not None)
+            if 'gemini' in model:
+                from modules.caption import gemini
+                kwargs = {
+                    'temperature': temperature,
+                    'min_output_tokens': min_tokens,
+                    'max_output_tokens': max_tokens,
+                }
+                model_name = model.replace('google/', '')
+                response = gemini.predict(prompt_text, current_image, model_name, system, prefill_text, thinking, kwargs)
+                t1 = time.time()
+                log.info(f'Prompt enhance: model="{model}" nsfw={nsfw} time={t1-t0:.2f} prefill="{prefill_text[:20] if prefill_text else None}" response={len(response)}')
+                debug_log(f'Prompt enhance: response="{response}"')
+                self.busy = False
+                return response
+            elif 'grok' in model:
+                from modules.caption import grok
+                kwargs = {
+                    'temperature': temperature,
+                }
+                model_name = model.replace('xai/', '')
+                response = grok.predict(prompt_text, current_image, model_name, system, prefill_text, thinking, kwargs)
+                t1 = time.time()
+                log.info(f'Prompt enhance: model="{model}" nsfw={nsfw} time={t1-t0:.2f} prefill="{prefill_text[:20] if prefill_text else None}" response={len(response)}')
+                debug_log(f'Prompt enhance: response="{response}"')
+                self.busy = False
+                return response
+            else:
+                self.busy = False
+                return 'Model not recognized'
+
         chat_template = set_template(
             system=system,
             prompt=prompt_text,
@@ -400,35 +439,11 @@ class PromptEnhanceScript(scripts_manager.Script):
             module=self.parent,
         )
 
-        # Prepare prefill (VQA approach: string concatenation, not assistant message)
-        prefill_text = (prefill or '').strip()
         use_prefill = len(prefill_text) > 0
         is_thinking = is_thinking_model(model)
-
         debug_log(f'Prompt enhance: system="{system}"')
         debug_log(f'Prompt enhance: prompt="{prompt_text}"')
         debug_log(f'Prompt template: roles={[msg["role"] for msg in chat_template]} thinking={is_thinking}:{thinking} prefill={use_prefill}')
-        t0 = time.time()
-        self.busy = True
-
-        if is_cloud_model(model):
-            if 'gemini' in model:
-                from modules.caption import gemini
-                kwargs = {
-                    'temperature': temperature,
-                    'min_output_tokens': min_tokens,
-                    'max_output_tokens': max_tokens,
-                }
-                model_name = model.replace('google/', '')
-                response = gemini.predict(prompt_text, current_image, model_name, system, model, prefill_text, thinking, kwargs)
-                t1 = time.time()
-                log.info(f'Prompt enhance: model="{model}" nsfw={nsfw} time={t1-t0:.2f} prefill="{prefill_text[:20] if prefill_text else None}" response={len(response)}')
-                debug_log(f'Prompt enhance: response="{response}"')
-                self.busy = False
-                return response
-
-            else:
-                return 'Model not recognized'
 
         try:
             # Qwen3.5 uses native enable_thinking parameter in the chat template