oobabooga/text-generation-webui
1
0
Fork 0
mirror of https://github.com/oobabooga/text-generation-webui.git synced 2025-12-06 07:12:10 +01:00
Code Issues Actions Packages Projects Releases Wiki Activity

Make --image-model load the model on startup

Browse source
This commit is contained in:
oobabooga 2025-12-02 06:15:21 -08:00
parent 977d08de14
commit e328748bc8
1 changed files with 17 additions and 0 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

17
server.py
View file

@ -5,6 +5,7 @@ from pathlib import Path
from modules import shared from modules import shared
from modules.block_requests import OpenMonkeyPatch, RequestBlocker from modules.block_requests import OpenMonkeyPatch, RequestBlocker
from modules.image_models import load_image_model
from modules.logging_colors import logger from modules.logging_colors import logger
from modules.prompts import load_prompt from modules.prompts import load_prompt
@ -320,6 +321,22 @@ if __name__ == "__main__":
if shared.args.lora: if shared.args.lora:
add_lora_to_model(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() shared.generation_lock = Lock()
if shared.args.idle_timeout > 0: if shared.args.idle_timeout > 0: