mirror of
https://github.com/vladmandic/automatic
synced 2026-09-19 01:04:32 +02:00
add grok to cloud models
Signed-off-by: Vladimir Mandic <mandic00@live.com>
This commit is contained in:
@@ -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': {},
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user