improve runai-streamer integration

Signed-off-by: Vladimir Mandic <mandic00@live.com>

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
 

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
 
 

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: