From d2caffa7b4f49cc85a5e0f423443a590c4f02c5f Mon Sep 17 00:00:00 2001 From: Vladimir Mandic Date: Sun, 2 Nov 2025 12:14:12 -0500 Subject: [PATCH] improve runai-streamer integration Signed-off-by: Vladimir Mandic --- CHANGELOG.md | 32 ++++++++++++++------------ modules/sd_hijack_safetensors.py | 39 ++++++++++++++++++++++---------- pipelines/generic.py | 12 ++++------ 3 files changed, 49 insertions(+), 34 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index a2520773a..29989d2c9 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -10,21 +10,25 @@ - **wildcards**: add inline processing using curly braces syntax - add setting to control `cudnn` enable/disable - change `vlm` beams to 1 by default for faster response + - update diffusers - **Fixes** - - fix: rocm possible endless loop during hip detection - - fix: rocm auto-disable miopen for gfx120x - - fix: better handling of detailer settings, thanks @awsr - - fix: cleanup `--optional` installer - - fix: guard against multi-controlnet in hires - - fix: update diffusers - - fix: inpaint handling - - fix: model type detection - - fix: version detection when cloned with `.git` suffix, thanks @awsr - - fix: init `sdnq` on video model load - - fix: add vae scale override for chrono - - fix: add tracing to model detection - - ui: fix full-screen image viewer buttons with non-standard ui theme - - ui: control tab show override section + - `chrono` transformers handling + - `chrono` extract last frame + - `chrono` add vae scale override, thanks @CalamitousFelicitousness + - `runai` improve streamer integration + - `transformers` dtype use new syntax + - `rocm` possible endless loop during hip detection + - `rocm` auto-disable miopen for gfx120x + - `detailer` better handling of settings, thanks @awsr + - `installer` cleanup `--optional` + - `hires` guard against multi-controlnet + - `inpaint` handling + - `version` detection when cloned with `.git` suffix, thanks @awsr + - `sdnq` init on video model load + - `model type` detection + - `model type` add tracing to model detection + - `ui` fix full-screen image viewer buttons with non-standard ui theme + - `ui` control tab show override section ## Update for 2025-10-31 diff --git a/modules/sd_hijack_safetensors.py b/modules/sd_hijack_safetensors.py index e8f775c5b..20e85f40d 100644 --- a/modules/sd_hijack_safetensors.py +++ b/modules/sd_hijack_safetensors.py @@ -1,7 +1,10 @@ +import io import os +import contextlib import safetensors.torch import transformers from installer import install, log +from modules import errors orig_load_file = safetensors.torch.load_file @@ -13,14 +16,20 @@ def hijacked_load_file(checkpoint_file, device="cpu"): return orig_load_file(checkpoint_file, device=device) install('runai_model_streamer') - log.trace(f'Loader: method=runai chunk={os.environ["RUNAI_STREAMER_CHUNK_BYTESIZE"]} limit={os.environ["RUNAI_STREAMER_MEMORY_LIMIT"]} device={device}') + log.debug(f'Loader: method=runai type=file chunk={os.environ["RUNAI_STREAMER_CHUNK_BYTESIZE"]} limit={os.environ["RUNAI_STREAMER_MEMORY_LIMIT"]} device={device}') state_dict = {} + stdout = io.StringIO() from runai_model_streamer import SafetensorsStreamer - with SafetensorsStreamer() as streamer: - streamer.stream_file(checkpoint_file) - for key, tensor in streamer.get_tensors(): - state_dict[key] = tensor.to(device) - + 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') return state_dict @@ -29,14 +38,20 @@ def hijacked_load_state_dict(checkpoint_file, is_quantized: bool = False, map_lo 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 chunk={os.environ["RUNAI_STREAMER_CHUNK_BYTESIZE"]} limit={os.environ["RUNAI_STREAMER_MEMORY_LIMIT"]} device={map_location} quantized={is_quantized}') + 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}') state_dict = {} + stdout = io.StringIO() from runai_model_streamer import SafetensorsStreamer - 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 - + 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') return state_dict diff --git a/pipelines/generic.py b/pipelines/generic.py index 40a951379..896d3c5c2 100644 --- a/pipelines/generic.py +++ b/pipelines/generic.py @@ -101,7 +101,9 @@ def load_text_encoder(repo_id, cls_name, load_config=None, subfolder="text_encod from modules import sdnq # pylint: disable=unused-import # register to diffusers and transformers load_args, quant_args = model_quant.get_dit_args(load_config, module='TE', device_map=True, allow_quant=allow_quant, modules_to_not_convert=modules_to_not_convert, modules_dtype_dict=modules_dtype_dict) quant_type = model_quant.get_quant_type(quant_args) + load_args.pop('torch_dtype', None) dtype = dtype or devices.dtype + load_args['dtype'] = dtype # load from local file if specified local_file = None @@ -128,12 +130,14 @@ def load_text_encoder(repo_id, cls_name, load_config=None, subfolder="text_encod """ text_encoder = model_te.load_t5(local_file) text_encoder = model_quant.do_post_load_quant(text_encoder, allow=quant_type is not None) + # 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}"') 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) + # use shared t5 if possible elif cls_name == transformers.T5EncoderModel and allow_shared and shared.opts.te_shared_t5: if model_quant.check_nunchaku('TE'): @@ -155,8 +159,6 @@ def load_text_encoder(repo_id, cls_name, load_config=None, subfolder="text_encod 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}') - if dtype is not None: - load_args['torch_dtype'] = dtype text_encoder = cls_name.from_pretrained( repo_id, cache_dir=shared.opts.hfcache_dir, @@ -170,8 +172,6 @@ def load_text_encoder(repo_id, cls_name, load_config=None, subfolder="text_encod 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}') - if dtype is not None: - load_args['torch_dtype'] = dtype text_encoder = cls_name.from_pretrained( repo_id, cache_dir=shared.opts.hfcache_dir, @@ -183,8 +183,6 @@ def load_text_encoder(repo_id, cls_name, load_config=None, subfolder="text_encod 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}') - if dtype is not None: - load_args['torch_dtype'] = dtype text_encoder = cls_name.from_pretrained( repo_id, cache_dir=shared.opts.hfcache_dir, @@ -196,8 +194,6 @@ 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}') - if dtype is not None: - load_args['torch_dtype'] = dtype if subfolder is not None: load_args['subfolder'] = subfolder if variant is not None: