1. Fix Normalization Bug:
- The previous implementation normalized each token immediately after scaling, effectively cancelling the emphasis weight.
- Updated logic to calculate global mean before scaling and restore it after scaling all tokens, correctly implementing A1111's 'EmphasisOriginal' behavior.
2. Improve Model Compatibility:
- Modified token handling to preserve all keys in the `tokens` dictionary (e.g., `t5xxl` for Flux/SD3) instead of reconstructing a partial dictionary. This prevents silent failures for non-SDXL/SD1.5 models.
- Added a safeguard in `_apply_direct_scaling_sdxl` to check if the embedding dimension matches the expected SDXL size (2048). If not (e.g., SD3), it falls back to generic scaling and logs a warning instead of producing incorrect output.
3. Note on hooks.py:
- The `hooks.py` file remains unused as per current implementation status (placeholder for future sampler integration).
The previous implementation of normalization in `_apply_direct_scaling` applied normalization per-token immediately after scaling. This effectively cancelled out the applied weight, rendering the emphasis feature non-functional when normalization was enabled.
This change fixes the logic to:
1. Calculate the global mean of the conditioning tensor before scaling.
2. Apply weights to all tokens.
3. Normalize the entire tensor to restore the original global mean.
This correctly implements A1111-style 'EmphasisOriginal' behavior where relative weights are preserved while maintaining global magnitude.
google-labs-jules[bot]