mirror of
https://github.com/oobabooga/text-generation-webui.git
synced 2025-12-06 07:12:10 +01:00
Make --image-model load the model on startup
This commit is contained in:
parent
977d08de14
commit
e328748bc8
17
server.py
17
server.py
|
|
@ -5,6 +5,7 @@ from pathlib import Path
|
|||
|
||||
from modules import shared
|
||||
from modules.block_requests import OpenMonkeyPatch, RequestBlocker
|
||||
from modules.image_models import load_image_model
|
||||
from modules.logging_colors import logger
|
||||
from modules.prompts import load_prompt
|
||||
|
||||
|
|
@ -320,6 +321,22 @@ if __name__ == "__main__":
|
|||
if shared.args.lora:
|
||||
add_lora_to_model(shared.args.lora)
|
||||
|
||||
# Load image model if specified via CLI
|
||||
if shared.args.image_model:
|
||||
logger.info(f"Loading image model: {shared.args.image_model}")
|
||||
result = load_image_model(
|
||||
shared.args.image_model,
|
||||
dtype=shared.settings.get('image_dtype', 'bfloat16'),
|
||||
attn_backend=shared.settings.get('image_attn_backend', 'sdpa'),
|
||||
cpu_offload=shared.settings.get('image_cpu_offload', False),
|
||||
compile_model=shared.settings.get('image_compile', False),
|
||||
quant_method=shared.settings.get('image_quant', 'none')
|
||||
)
|
||||
if result is not None:
|
||||
shared.image_model_name = shared.args.image_model
|
||||
else:
|
||||
logger.error(f"Failed to load image model: {shared.args.image_model}")
|
||||
|
||||
shared.generation_lock = Lock()
|
||||
|
||||
if shared.args.idle_timeout > 0:
|
||||
|
|
|
|||
Loading…
Reference in a new issue