add grok to cloud models

Signed-off-by: Vladimir Mandic <mandic00@live.com>
This commit is contained in:
Vladimir Mandic
2026-09-07 09:20:58 +02:00
parent 0a94b2d8ae
commit fe03d7b38e
9 changed files with 210 additions and 58 deletions
+9
View File
@@ -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': {},
+23 -21
View File
@@ -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
+40 -25
View File
@@ -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