Public/automatic
1
0
Fork 0
mirror of https://github.com/vladmandic/automatic synced 2026-09-20 01:31:13 +02:00
Code Issues Packages Projects Releases Wiki Activity

Merge pull request #4436 from CalamitousFelicitousness/runai-update

Browse Source 
Update runai-model-streamer logging integration
This commit is contained in:
Vladimir Mandic
2025-12-02 03:59:38 -05:00
committed by GitHub
parent 1673380b94 55c089ae48
commit 3b4f909862
5 changed files with 55 additions and 44 deletions
Download Patch File Download Diff File Expand all files Collapse all files
installer.py
+1
View File
@@ -1392,6 +1392,7 @@ def set_environment():
os.environ.setdefault('UVICORN_TIMEOUT_KEEP_ALIVE', '60')
os.environ.setdefault('RUNAI_STREAMER_CHUNK_BYTESIZE', '2097152')
os.environ.setdefault('RUNAI_STREAMER_MEMORY_LIMIT', '-1')
os.environ.setdefault('RUNAI_STREAMER_LOG_LEVEL', 'DEBUG' if os.environ.get('SD_LOAD_DEBUG') else 'WARNING')
allocator = f'garbage_collection_threshold:{opts.get("torch_gc_threshold", 80)/100:0.2f},max_split_size_mb:512'
if opts.get("torch_malloc", "native") == 'cudaMallocAsync':
allocator += ',backend:cudaMallocAsync'
modules/sd_hijack_safetensors.py
+18 -29
View File
@@ -1,6 +1,3 @@
import io
import os
import contextlib
import safetensors.torch
import transformers
from installer import install, log
@@ -15,21 +12,17 @@ def hijacked_load_file(checkpoint_file, device="cpu"):
if not checkpoint_file.endswith('.safetensors'):
return orig_load_file(checkpoint_file, device=device)
install('runai_model_streamer')
log.debug(f'Loader: method=runai type=file chunk={os.environ["RUNAI_STREAMER_CHUNK_BYTESIZE"]} limit={os.environ["RUNAI_STREAMER_MEMORY_LIMIT"]} device={device}')
install('runai_model_streamer>=0.15.1')
state_dict = {}
stdout = io.StringIO()
from runai_model_streamer import SafetensorsStreamer
with contextlib.redirect_stdout(stdout):
try:
with SafetensorsStreamer() as streamer:
streamer.stream_file(checkpoint_file)
for key, tensor in streamer.get_tensors():
state_dict[key] = tensor.to(device)
except Exception as e:
log.error(f'Loader: {e}')
log.error(stdout.getvalue())
errors.display(e, 'runai')
try:
with SafetensorsStreamer() as streamer:
streamer.stream_file(checkpoint_file)
for key, tensor in streamer.get_tensors():
state_dict[key] = tensor.to(device)
except Exception as e:
log.error(f'Loader: {e}')
errors.display(e, 'runai')
return state_dict
@@ -37,21 +30,17 @@ def hijacked_load_state_dict(checkpoint_file, is_quantized: bool = False, map_lo
if not checkpoint_file.endswith(".safetensors"):
return orig_load_state_dict(checkpoint_file=checkpoint_file, is_quantized=is_quantized, map_location=map_location, weights_only=weights_only)
install('runai_model_streamer')
log.trace(f'Loader: method=runai type=dict chunk={os.environ["RUNAI_STREAMER_CHUNK_BYTESIZE"]} limit={os.environ["RUNAI_STREAMER_MEMORY_LIMIT"]} device={map_location} quantized={is_quantized}')
install('runai_model_streamer>=0.15.1')
state_dict = {}
stdout = io.StringIO()
from runai_model_streamer import SafetensorsStreamer
with contextlib.redirect_stdout(stdout):
try:
with SafetensorsStreamer() as streamer:
streamer.stream_file(checkpoint_file)
for key, tensor in streamer.get_tensors():
state_dict[key] = tensor.to(map_location) if map_location != "meta" else tensor
except Exception as e:
log.error(f'Loader: {e}')
log.error(stdout.getvalue())
errors.display(e, 'runai')
try:
with SafetensorsStreamer() as streamer:
streamer.stream_file(checkpoint_file)
for key, tensor in streamer.get_tensors():
state_dict[key] = tensor.to(map_location) if map_location != "meta" else tensor
except Exception as e:
log.error(f'Loader: {e}')
errors.display(e, 'runai')
return state_dict
modules/sd_models.py
+3 -2
View File
@@ -65,6 +65,8 @@ def set_huggingface_options():
else:
sd_hijack_accelerate.restore_accelerate()
if (shared.opts.runai_streamer_diffusers or shared.opts.runai_streamer_transformers) and (sys.platform == 'linux'):
import os
log.debug(f'Loader: runai enabled chunk={os.environ["RUNAI_STREAMER_CHUNK_BYTESIZE"]} limit={os.environ["RUNAI_STREAMER_MEMORY_LIMIT"]}')
sd_hijack_safetensors.hijack_safetensors(shared.opts.runai_streamer_diffusers, shared.opts.runai_streamer_transformers)
else:
sd_hijack_safetensors.restore_safetensors()
@@ -629,8 +631,7 @@ def load_sdnq_model(checkpoint_info, pipeline, diffusers_load_config, op):
if shared.opts.runai_streamer_diffusers and (sys.platform == 'linux'):
load_method = 'streamer'
from installer import install
install('runai_model_streamer')
shared.log.trace(f'Loader: method={load_method} chunk={os.environ["RUNAI_STREAMER_CHUNK_BYTESIZE"]} limit={os.environ["RUNAI_STREAMER_MEMORY_LIMIT"]}')
install('runai_model_streamer>=0.15.1')
elif shared.opts.sd_parallel_load:
load_method = 'threaded'
else:
modules/video_models/video_load.py
+13 -3
View File
@@ -1,4 +1,5 @@
import os
import sys
import copy
import time
import transformers # pylint: disable=unused-import
@@ -7,6 +8,15 @@ from modules import shared, errors, sd_models, sd_checkpoint, model_quant, devic
from modules.video_models import models_def, video_utils, video_overrides, video_cache
def _loader(component):
"""Return loader type for log messages."""
if sys.platform != 'linux':
return 'default'
if component == 'diffusers':
return 'runai' if shared.opts.runai_streamer_diffusers else 'default'
return 'runai' if shared.opts.runai_streamer_transformers else 'default'
loaded_model = None
@@ -60,7 +70,7 @@ def load_model(selected: models_def.Model):
selected.te_folder = 'text_encoder'
selected.te_revision = None
shared.log.debug(f'Video load: module=te repo="{selected.te or selected.repo}" folder="{selected.te_folder}" cls={selected.te_cls.__name__} quant={model_quant.get_quant_type(quant_args)}')
shared.log.debug(f'Video load: module=te repo="{selected.te or selected.repo}" folder="{selected.te_folder}" cls={selected.te_cls.__name__} quant={model_quant.get_quant_type(quant_args)} loader={_loader("transformers")}')
kwargs["text_encoder"] = selected.te_cls.from_pretrained(
pretrained_model_name_or_path=selected.te or selected.repo,
subfolder=selected.te_folder,
@@ -80,7 +90,7 @@ def load_model(selected: models_def.Model):
if dit_folder is not None and dit_folder not in kwargs:
# get a new quant arg on every loop to prevent the quant config classes getting entangled
load_args, quant_args = model_quant.get_dit_args({}, module='Model', device_map=True)
shared.log.debug(f'Video load: module=transformer repo="{selected.dit or selected.repo}" module="{dit_folder}" folder="{dit_folder}" cls={selected.dit_cls.__name__} quant={model_quant.get_quant_type(quant_args)}')
shared.log.debug(f'Video load: module=transformer repo="{selected.dit or selected.repo}" module="{dit_folder}" folder="{dit_folder}" cls={selected.dit_cls.__name__} quant={model_quant.get_quant_type(quant_args)} loader={_loader("diffusers")}')
kwargs[dit_folder] = selected.dit_cls.from_pretrained(
pretrained_model_name_or_path=selected.dit or selected.repo,
subfolder=dit_folder,
@@ -91,7 +101,7 @@ def load_model(selected: models_def.Model):
**offline_args,
)
else:
shared.log.debug(f'Video load: module=transformer repo="{selected.dit or selected.repo}" module="{dit_folder}" folder="{dit_folder}" cls={selected.dit_cls.__name__} skip')
shared.log.debug(f'Video load: module=transformer repo="{selected.dit or selected.repo}" module="{dit_folder}" folder="{dit_folder}" cls={selected.dit_cls.__name__} loader={_loader("diffusers")} skip')
if selected.dit_folder is None:
selected.dit_folder = ['transformer']
pipelines/generic.py
+20 -10
View File
@@ -1,4 +1,5 @@
import os
import sys
import json
import diffusers
import transformers
@@ -8,6 +9,15 @@ from modules import shared, devices, errors, sd_models, model_quant
debug = os.environ.get('SD_LOAD_DEBUG', None) is not None
def _loader(component):
"""Return loader type for log messages."""
if sys.platform != 'linux':
return 'default'
if component == 'diffusers':
return 'runai' if shared.opts.runai_streamer_diffusers else 'default'
return 'runai' if shared.opts.runai_streamer_transformers else 'default'
def load_transformer(repo_id, cls_name, load_config=None, subfolder="transformer", allow_quant=True, variant=None, dtype=None, modules_to_not_convert=None, modules_dtype_dict=None):
transformer = None
if load_config is None:
@@ -31,7 +41,7 @@ def load_transformer(repo_id, cls_name, load_config=None, subfolder="transformer
local_file = sd_unet.unet_dict[shared.opts.sd_unet]
if local_file is not None and local_file.lower().endswith('.gguf'):
shared.log.debug(f'Load model: transformer="{local_file}" cls={cls_name.__name__} quant="{quant_type}" args={load_args}')
shared.log.debug(f'Load model: transformer="{local_file}" cls={cls_name.__name__} quant="{quant_type}" loader={_loader("diffusers")} args={load_args}')
from modules import ggml
ggml.install_gguf()
loader = cls_name.from_single_file if hasattr(cls_name, 'from_single_file') else cls_name.from_pretrained
@@ -43,7 +53,7 @@ def load_transformer(repo_id, cls_name, load_config=None, subfolder="transformer
)
transformer = model_quant.do_post_load_quant(transformer, allow=quant_type is not None)
elif local_file is not None and local_file.lower().endswith('.safetensors'):
shared.log.debug(f'Load model: transformer="{local_file}" cls={cls_name.__name__} quant="{quant_type}" args={load_args}')
shared.log.debug(f'Load model: transformer="{local_file}" cls={cls_name.__name__} quant="{quant_type}" loader={_loader("diffusers")} args={load_args}')
if dtype is not None:
load_args['torch_dtype'] = dtype
loader = cls_name.from_single_file if hasattr(cls_name, 'from_single_file') else cls_name.from_pretrained
@@ -54,7 +64,7 @@ def load_transformer(repo_id, cls_name, load_config=None, subfolder="transformer
**quant_args,
)
else:
shared.log.debug(f'Load model: transformer="{repo_id}" cls={cls_name.__name__} subfolder={subfolder} quant="{quant_type}" args={load_args}')
shared.log.debug(f'Load model: transformer="{repo_id}" cls={cls_name.__name__} subfolder={subfolder} quant="{quant_type}" loader={_loader("diffusers")} args={load_args}')
if 'sdnq-' in repo_id.lower():
quant_args = {}
if dtype is not None:
@@ -115,7 +125,7 @@ def load_text_encoder(repo_id, cls_name, load_config=None, subfolder="text_encod
# load from local file gguf
if local_file is not None and local_file.lower().endswith('.gguf'):
shared.log.debug(f'Load model: text_encoder="{local_file}" cls={cls_name.__name__} quant="{quant_type}"')
shared.log.debug(f'Load model: text_encoder="{local_file}" cls={cls_name.__name__} quant="{quant_type}" loader={_loader("transformers")}')
"""
from modules import ggml
ggml.install_gguf()
@@ -132,7 +142,7 @@ def load_text_encoder(repo_id, cls_name, load_config=None, subfolder="text_encod
# load from local file safetensors
elif local_file is not None and local_file.lower().endswith('.safetensors'):
shared.log.debug(f'Load model: text_encoder="{local_file}" cls={cls_name.__name__} quant="{quant_type}"')
shared.log.debug(f'Load model: text_encoder="{local_file}" cls={cls_name.__name__} quant="{quant_type}" loader={_loader("transformers")}')
from modules import model_te
text_encoder = model_te.load_t5(local_file)
text_encoder = model_quant.do_post_load_quant(text_encoder, allow=quant_type is not None)
@@ -143,7 +153,7 @@ def load_text_encoder(repo_id, cls_name, load_config=None, subfolder="text_encod
import nunchaku
repo_id = 'nunchaku-tech/nunchaku-t5/awq-int4-flux.1-t5xxl.safetensors'
cls_name = nunchaku.NunchakuT5EncoderModel
shared.log.debug(f'Load model: text_encoder="{repo_id}" cls={cls_name.__name__} quant="SVDQuant"')
shared.log.debug(f'Load model: text_encoder="{repo_id}" cls={cls_name.__name__} quant="SVDQuant" loader={_loader("transformers")}')
text_encoder = nunchaku.NunchakuT5EncoderModel.from_pretrained(
repo_id,
torch_dtype=dtype,
@@ -157,7 +167,7 @@ def load_text_encoder(repo_id, cls_name, load_config=None, subfolder="text_encod
repo_id = 'Disty0/t5-xxl'
with open(os.path.join('configs', 'flux', 'text_encoder_2', 'config.json'), encoding='utf8') as f:
load_args['config'] = transformers.T5Config(**json.load(f))
shared.log.debug(f'Load model: text_encoder="{repo_id}" cls={cls_name.__name__} quant="{quant_type}" shared={shared.opts.te_shared_t5}')
shared.log.debug(f'Load model: text_encoder="{repo_id}" cls={cls_name.__name__} quant="{quant_type}" loader={_loader("transformers")} shared={shared.opts.te_shared_t5}')
text_encoder = cls_name.from_pretrained(
repo_id,
cache_dir=shared.opts.hfcache_dir,
@@ -170,7 +180,7 @@ def load_text_encoder(repo_id, cls_name, load_config=None, subfolder="text_encod
else:
repo_id = 'Wan-AI/Wan2.1-T2V-1.3B-Diffusers'
subfolder = 'text_encoder'
shared.log.debug(f'Load model: text_encoder="{repo_id}" cls={cls_name.__name__} quant="{quant_type}" shared={shared.opts.te_shared_t5}')
shared.log.debug(f'Load model: text_encoder="{repo_id}" cls={cls_name.__name__} quant="{quant_type}" loader={_loader("transformers")} shared={shared.opts.te_shared_t5}')
text_encoder = cls_name.from_pretrained(
repo_id,
cache_dir=shared.opts.hfcache_dir,
@@ -181,7 +191,7 @@ def load_text_encoder(repo_id, cls_name, load_config=None, subfolder="text_encod
elif cls_name == transformers.Qwen2_5_VLForConditionalGeneration and allow_shared and shared.opts.te_shared_t5:
repo_id = 'hunyuanvideo-community/HunyuanImage-2.1-Diffusers'
subfolder = 'text_encoder'
shared.log.debug(f'Load model: text_encoder="{repo_id}" cls={cls_name.__name__} quant="{quant_type}" shared={shared.opts.te_shared_t5}')
shared.log.debug(f'Load model: text_encoder="{repo_id}" cls={cls_name.__name__} quant="{quant_type}" loader={_loader("transformers")} shared={shared.opts.te_shared_t5}')
text_encoder = cls_name.from_pretrained(
repo_id,
cache_dir=shared.opts.hfcache_dir,
@@ -192,7 +202,7 @@ def load_text_encoder(repo_id, cls_name, load_config=None, subfolder="text_encod
# load from repo
if text_encoder is None:
shared.log.debug(f'Load model: text_encoder="{repo_id}" cls={cls_name.__name__} quant="{quant_type}" shared={shared.opts.te_shared_t5}')
shared.log.debug(f'Load model: text_encoder="{repo_id}" cls={cls_name.__name__} quant="{quant_type}" loader={_loader("transformers")} shared={shared.opts.te_shared_t5}')
if subfolder is not None:
load_args['subfolder'] = subfolder
if variant is not None: