diff --git a/installer.py b/installer.py index 480c54e83..5957570e6 100644 --- a/installer.py +++ b/installer.py @@ -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' diff --git a/modules/sd_hijack_safetensors.py b/modules/sd_hijack_safetensors.py index 20e85f40d..a9ca5ca78 100644 --- a/modules/sd_hijack_safetensors.py +++ b/modules/sd_hijack_safetensors.py @@ -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 diff --git a/modules/sd_models.py b/modules/sd_models.py index a8d501105..66048d58e 100644 --- a/modules/sd_models.py +++ b/modules/sd_models.py @@ -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: diff --git a/modules/video_models/video_load.py b/modules/video_models/video_load.py index 4ad4a3925..375861840 100644 --- a/modules/video_models/video_load.py +++ b/modules/video_models/video_load.py @@ -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'] diff --git a/pipelines/generic.py b/pipelines/generic.py index be59bd75e..911cdc3c6 100644 --- a/pipelines/generic.py +++ b/pipelines/generic.py @@ -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: