diff --git a/CHANGELOG.md b/CHANGELOG.md
index 2a9b72955..73e4a5b2a 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -30,6 +30,7 @@
   - **NNCF** update to 2.16.0  
     major refactoring of NNCF quantization code  
     new quant types: `INT8_SYM` (new default), `INT4` and `INT4_SYM`  
+    pre-load quantization support  
   - **HiDream-I1** optimized offloading and prompt-encode caching  
     it now works in 12GB VRAM / 26GB RAM!  
   - **CogView3** and **CogView4** model loader optimizations  
diff --git a/modules/processing_args.py b/modules/processing_args.py
index b65047b57..a7bb4831e 100644
--- a/modules/processing_args.py
+++ b/modules/processing_args.py
@@ -299,7 +299,9 @@ def set_pipeline_args(p, model, prompts:list, negative_prompts:list, prompts_2:t
     elif 'callback_on_step_end' in possible:
         args['callback_on_step_end'] = diffusers_callback
         if 'callback_on_step_end_tensor_inputs' in possible:
-            if 'prompt_embeds' in possible and 'negative_prompt_embeds' in possible and hasattr(model, '_callback_tensor_inputs'):
+            if 'HiDreamImage' in model.__class__.__name__: # uses prompt_embeds_t5 and prompt_embeds_llama3 instead
+                args['callback_on_step_end_tensor_inputs'] = model._callback_tensor_inputs # pylint: disable=protected-access
+            elif 'prompt_embeds' in possible and 'negative_prompt_embeds' in possible and hasattr(model, '_callback_tensor_inputs'):
                 args['callback_on_step_end_tensor_inputs'] = model._callback_tensor_inputs # pylint: disable=protected-access
             else:
                 args['callback_on_step_end_tensor_inputs'] = ['latents']