diff --git a/README.md b/README.md
index a6b1fca..dfd6827 100644
--- a/README.md
+++ b/README.md
@@ -8,7 +8,7 @@ For recent user-visible changes, please see the [ChangeLog](changelog.md).
* Better TAESD previews (see below).
* Visual previews for some audio models (currently only ACE-Steps).
-* Multi-frame video previews for most common video models (Wan 2.2, 2.1, Hunyuan, etc). See [the section on video encode/decode](#blehtaevideoencode-and-blehtaevideodecode).
+* Multi-frame video previews for most common video models (Wan 2.2, 2.1, Hunyuan, LTX 2.3, etc). See [the section on video encode/decode](#blehtaevideoencode-and-blehtaevideodecode).
* Allow setting seed, timestep range and step interval for HyperTile (look for the [`BlehHyperTile`](#blehhypertile) node).
* Allow applying Kohya Deep Shrink to multiple blocks, also allow gradually fading out the downscale factor (look for the [`BlehDeepShrink`](#blehdeepshrink) node).
* Allow discarding penultimate sigma (look for the `BlehDiscardPenultimateSigma` node). This can be useful if you find certain samplers are ruining your image by spewing a bunch of noise into it at the very end (usually only an issue with `dpm2 a` or SDE samplers).
@@ -22,17 +22,20 @@ For recent user-visible changes, please see the [ChangeLog](changelog.md).
## Configuration
-Copy either `blehconfig.yaml.example` or `blehconfig.json.example` to `blehconfig.yaml` or `blehconfig.json` respectively and edit the copy. When loading configuration, the YAML file will be prioritized if it exists and Python has YAML support.
+Copy either `blehconfig.yaml.example` to `blehconfig.yaml` and edit the copy. Bleh will also check for a `blehconfig.json` file when starting up if the YAML one doesn't exist. I recommend using YAML here.
Restart ComfyUI to apply any new changes.
### Better Previews
+There are links to the various TAE models used for high quality previewing near the bottom of this README.
+
* Supports setting max preview size (ComfyUI default is hardcoded to 512 max).
* Supports showing previews for more than the first latent in the batch.
* Supports throttling previews. Do you really need your expensive high quality preview to get updated 3 times a second?
+* Supports LTX (2.0 and 2.3). **Note**: You need to use the `BlehFixGuiderPreviewing` node for LTX. See the description of it below.
-The previewer can now show visual previews for ACE-Steps latents. If you want to disable that feature, you can add `aceaudio` to the
+The previewer can now show visual previews for ACE-Steps (1.0 and 1.5) latents. If you want to disable that feature, you can add `aceaudio` to the
`blacklist_formats` list. For example if you are using a YAML configuration file you could do: `blacklist_formats: ["aceaudio"]`
**General settings defaults:**
@@ -48,6 +51,7 @@ The previewer can now show visual previews for ACE-Steps latents. If you want to
|`throttle_secs`|`2`|Max frequency to decode the latents for previewing. `0.25` would be every quarter second, `2` would be once every two seconds|
|`maxed_batch_step_mode`|`false`|When `false`, you will see the first `max_batch` previews, when `true` you will see previews spread across the batch. Also applies to video frames.|
|`preview_device`|`null`|`null` (use the default device) or a string with a PyTorch device name like `"cpu"`, `"cuda:0"`, etc. Can be used to run TAESD previews on CPU or other available devices. Not recommended to change this unless you really need to, using the CPU device may prevent out of memory errors but will likely significantly slow down generation.|
+|`preview_dtype`|`bfloat16`|`null` (use the default device) or a string with a PyTorch dtype name like `float32`, `bfloat16`, `float16` (probably the only ones that will work).|
|`compile_previewer`|`false`|Controls whether the previewer gets compiled with `torch.compile`. May be a boolean or an object in which case the object will be used as argument to `torch.compile`. Note: May cause a delay/memory spike on the first preview.|
|`oom_fallback`|`latent2rgb`|May be set to `none` or `latent2rgb`. Controls what happens if trying to decode the preview runs out of memory.|
|`oom_retry`|`true`|If set to `false`, we will give up and use the `oom_fallback` behavior after hitting the first OOM. Otherwise, we'll attempt to decode with the normal previewer each time a preview is requested, even if that previously ran out of memory.|
@@ -71,11 +75,29 @@ More detailed explanation for skipping upscale layers: Latents (the thing you're
|`video_parallel`|`false`|Use parallel mode when decoding video latents. May actually use more memory than a full VAE decode.|
|`video_max_frames`|`-1`|Maximum frames to include in a preview. Frame limiting is treated like batch limiting. `-1` means unlimited.|
|`video_temporal_upscale_level`|`0`|Number of temporal upscale blocks to use, 0 will not do any temporal upscaling, 2 means full temporal upscaling.|
+|`animate_preview`|`none`|One of `none`, `video`, `batch` or `both`. Controls when animated previews are generated. When animating batch previews, batch items will be frames in the animation.|
These defaults are conservative. I would recommend setting `throttle_secs` to something relatively high (like 5-10) especially if you are generating batches at high resolution.
Slightly more detailed explanation for `maxed_batch_step_mode`: If max previews is set to `3` and the batch size is `15` you will see previews for indexes `0, 5, 10`. Or to put it a different way, it steps through the batch by `batch_size / max_previews` rounded up. This behavior may be useful for previewing generations with a high batch count like when using AnimateDiff.
+**Last preview endpoint settings defaults:**
+
+The last preview endpoint will serve up either a webpage that auto-updates the last preview `/bleh/last_preview.html` or
+a static image of the last preview `/bleh/last_preview`. It is disabled by default.
+
+|Key|Default|Description|
+|-|-|-|
+|`publish_last_preview`|`false`|Enables publishing last previews.|
+|`publish_last_preview_min_refresh`|`5`|When auto updating last previews, it will occur at most every this many seconds.|
+|`only_animate_last_preview`|`true`|When enabled, animated previews only appear in the last preview endpoint, not the sampler.|
+
+Last preview notes:
+
+* Leaving `only_animate_last_preview` on the default is reasonable. Unless you're using "Nodes 2.0", previews can't animate in the sampler so sending them just wastes resources.
+* If you turn on the last preview endpoint, anyone that can access your ComfyUI instance can see what you're generating. However external access to ComfyUI is disabled by default and if someone can access the instance they can do whatever they want, including seeing your previews and queue history so there isn't much of a risk to turning this on (to the best of my knowledge).
+
+***
**Note**: Other node packs that patch ComfyUI's previewer behavior may interfere with this feature. One I am aware of is [ComfyUI-VideoHelperSuite](https://github.com/Kosinkadink/ComfyUI-VideoHelperSuite) - if you have displaying animated previews turned on, it will overwrite Bleh's patched previewer. Or possibly, depending on the load order, Bleh will prevent it from working correctly.
@@ -279,6 +301,10 @@ just manually enter a list of sigmas. Note: Experimental, not well tested.
Ensures that Bleh's previewer is used. Generally not necessary unless some other custom node pack is overriding the default previewer. The node acts as a bridge for any input type.
+## BlehFixGuiderPreviewing
+
+Mostly only necessary for LTX previewing. You absolutely need to pass your guider through this for LTX (2.0, 2.3, 2.3 wide) previews to work. If you're generating videos with a FPS other than the video model's default (regardless of the video model type) then you can set `fps_override` to avoid your animated previews playing with the wrong speed. For LTX 2.3, you'll need to set `prefer_previewer` to `ltxav23` or `ltxav23wide` because there isn't a way for a previewer to detect whether the latent is in LTX 2.0 or LTX 2.3 format. I've been using the wide LTX 2.3 version (linked below) - it's better quality, but possibly somewhat higher resource usage.
+
### BlehTAEVideoEncode and BlehTAEVideoDecode
Fast video latent encoding/decoding with models from madebyollin (same person that made TAESD). Supports WAN 2.1, Hunyuan and Mochi. The node has a toggle for parallel mode which is faster but may use a lot of memory.
@@ -289,11 +315,18 @@ You will need to download the models and put them in `models/vae_approx`. Don't
* **WAN 2.1**: https://github.com/madebyollin/taehv/blob/main/taew2_1.pth
* **Hunyean**: https://github.com/madebyollin/taehv/blob/main/taehv.pth
* **Mochi**: https://github.com/madebyollin/taem1/blob/main/taem1.pth
+* **LTX 2.0**: https://github.com/madebyollin/taehv/blob/main/laeltx_2.pth
+* **LTX 2.3**: https://github.com/madebyollin/taehv/blob/main/laeltx2_3.pth
+* **LTX 2.3 wide**: https://github.com/madebyollin/taehv/blob/2026_03_11_taeltx23_wide/taeltx2_3_wide.pth
*Note*: If you run into issues it's probably a problem with my implementation and not the TAE video models or original inference code.
***
+## Blend Modes
+
+There is some incomplete documentation [here](blend_modes.md).
+
## Scaling Types
* bicubic: Generally the safe option.
diff --git a/__init__.py b/__init__.py
index d4f8a80..92a17d8 100644
--- a/__init__.py
+++ b/__init__.py
@@ -22,6 +22,12 @@ def blep_init():
sys.modules["_blepping_integrations"] = bi
nodes._blepping_integrations = bi # noqa: SLF001
samplers.add_sampler_presets()
+ if settings.SETTINGS.btp_publish_last_preview:
+ from .py.better_previews import last_preview # noqa: PLC0415
+
+ last_preview.init_routes(
+ min_refresh=settings.SETTINGS.btp_publish_last_preview_min_refresh,
+ )
blep_init()
diff --git a/blehconfig.example.json b/blehconfig.example.json
deleted file mode 100644
index cc6d3af..0000000
--- a/blehconfig.example.json
+++ /dev/null
@@ -1,20 +0,0 @@
-{
- "betterTaesdPreviews": {
- "enabled": true,
- "max_size": 768,
- "max_batch": 4,
- "max_batch_cols": 2,
- "throttle_secs": 1,
- "maxed_batch_step_mode": false,
- "preview_device": null,
- "skip_upscale_layers": 0,
- "compile_previewer": false,
- "oom_fallback": "latent2rgb",
- "oom_retry": true,
- "whitelist_formats": [],
- "blacklist_formats": [],
- "video_parallel": false,
- "video_max_frames": -1,
- "video_temporal_upscale_level": 0
- }
-}
diff --git a/blehconfig.example.yaml b/blehconfig.example.yaml
index 3594ad6..bae1e0a 100644
--- a/blehconfig.example.yaml
+++ b/blehconfig.example.yaml
@@ -1,5 +1,5 @@
# Copy this file to blehconfig.yaml
-betterTaesdPreviews:
+previews:
# If disabled, will use the old ComfyUI previewer.
enabled: true
@@ -36,7 +36,13 @@ betterTaesdPreviews:
# You may set it to a specific dtype: float32, float16, bfloat16
# Setting it to "keep" or null just leaves the dtype alone (which is probably float32).
# Setting it to "vae" will use whatever dtype ComfyUI is set to use for VAE.
- preview_dtype: null
+ preview_dtype: bfloat16
+
+ # Uses non-blocking transfers for previews when the device supports it.
+ # Not recommended as it is extremely likely to corrupt previews, especially if the previewer
+ # is relatively slow or the latent is large (video models, Chroma Radiance). However,
+ # it might decrease the performance impact of previewing.
+ preview_non_blocking: false
# Allows skipping upscale layers in the TAESD model, may increase performance when previewing large images or batches.
# May be set to -1 (conservative) or -2 (aggressive) to automatically calculate how many to skip. See README.md for details.
@@ -59,7 +65,7 @@ betterTaesdPreviews:
# When disabled, we use the fallback starting from the first OOM.
oom_retry: true
- # List of lowercase latent format names from https://github.com/comfyanonymous/ComfyUI/blob/master/comfy/latent_formats.py
+ # List of lowercase latent format names from https://github.com/Comfy-Org/ComfyUI/blob/master/comfy/latent_formats.py
# If the list is empty, this disables the whitelist. Otherwise, Bleh will
# only handle previewing for formats in the list.
whitelist_formats: []
@@ -78,3 +84,30 @@ betterTaesdPreviews:
# Number of temporal upscale blocks to run. 2 is the maximum and will fully decode
# the latent into image frames.
video_temporal_upscale_level: 0
+
+ # One of:
+ # none: Never use animated previews.
+ # video: Use animated previews for video generations.
+ # batch: Use animated previews for batch generations (batch items will be frames.)
+ # both: Use animated previews for video and batch generations.
+ animate_preview: none
+
+ # When enabled, this will add `/bleh/last_preview.html` and `/bleh/last_preview` paths
+ # under whatever URI you use to access ComfyUI. The HTML version will serve a webpage
+ # that shows the last preview and automatically updates. The non-HTML URI will serve
+ # an image (possibly animated) of the last preview.
+ # NOTE: When this is enabled, if someone can access your ComfyUI instance then they
+ # can see what you're generating, but they already could do that just less conveniently.
+ # Access from external hosts is not something ComfyUI enables by default.
+ publish_last_preview: false
+
+ # When publishing last previews, the HTML endpoint will only refresh the image at most
+ # once every this many seconds.
+ publish_last_preview_min_refresh: 5
+
+ # Only applies when you have a preview that would be animated and publish_last_preview
+ # is enabled. When this is turned on, the image sent to ComfyUI for the preview is one
+ # frame, while the image under the last preview may be animated. You probably want to
+ # leave this on the default because previews can only animated with "Nodes 2.0" which
+ # you probably don't want to use, so sending an animated preview to the sampler is pointless.
+ only_animate_last_preview: true
diff --git a/blend_modes.md b/blend_modes.md
new file mode 100644
index 0000000..ec87296
--- /dev/null
+++ b/blend_modes.md
@@ -0,0 +1,69 @@
+# Blend Modes
+
+There are so many blend modes now that it's probably pretty overwhelming. Also a lot of them are junk/don't work well but can't be removed without breaking existing workflows that might use them. So here is some incomplete, low-effort documentation!
+
+## Meta Modes
+
+Every blend function takes at least three parameters: `a`, `b` and `t` (the ratio to blend). `lerp(a, b, 0.25)` would mean `a * 0.75 + b * 0.25`.
+
+* `revWHATEVER` - Just flips the inputs, so if we were doing something like `lerp(a, b, 0.25)` (`a * 0.75 + b * 0.5`) the blend would be applied like `lerp(b, a, 0.25)`.
+* `normWHATEVER` - Tries to scale the input to -1...1 and then uses a simple LERP to find a range for the output. Pretty much garbage.
+
+Some modes also have suffixes. Not precisely a meta mode but this is probably the logical place to cover it.
+
+* `_d1`, `_d2`, etc - Indicates the mode will operate on that dimension. Dimension 1 in this case, which is typically channels in most latents.
+* `_copysign_a` - The blended result copies the sign from the `a` parameter.
+* `_avoidsign_a` - The blended result avoids the sign of the `a` parameter. In other words, negate A and then copy the sign from it to the result.
+* Digits like `_025` - Usually indicates a multiplier. `025` would stand for `0.25`.
+* `_base_a` - Mostly for CFG type blend modes. This means if `a` is cond and `b` is uncond, a ratio of 0 will give you cond. Normal CFG is `lerp(uncond, cond, ratio)` so if ratio is 0 you get pure uncond.
+
+## Custom Blend Parameters
+
+Nodes taking a blend mode use a selection but you can define modes as a string (whitespace is ignored so I recommend a multiline string widget) and force the connection to the mode parameter with the `BlehCast` node.
+
+Custom definitions use this syntax: `mode_name:param1=val1:param2=val2`. Integer and float values don't require any special handling. Boolean parameters use `true` and `false`. Nullable parameters can use `none`. An empty list is `()`. Lists are comma-separated and a singleton list is specified like `1,`. String literals (like mode names) should start with the caret, like `^whatever` otherwise they will be interpreted as a blend mode. Finally, some blend modes are wrappers for other blend modes. If it exists, a single trailing underscore will be stripped from the parameters when they are passed to the nested blend function. Not very convenient to use, but it allows specifying parameters where the names may clash. I.E. `some_mode:blend_mode=blah:blend_mode=other` would let `some_mode` use the `blend_mode` parameter and then pass the second to the `blah` blend mode handler.
+
+This is clunky/inconvenient and pretty limited but it does allow specifying custom parameters in most cases.
+
+All blend modes support some common parameters:
+
+* `rev` - boolean. Example `lerp:rev=true`. Flips the inputs.
+* `scale_multiplier` - float. Rescales the blend ratio. Example: `lerp:scale_multiplier=0.5`. `lerp(a, b, 1.0)` would result in `lerp(a, b, 0.5)`.
+* `invert_scale` - float. Adjusts the blend ratio by doing `scale_value - ratio`. Example: `lerp:invert_scale=1.0`. `lerp(a, b, 0.4)` would result in `lerp(a, b, 0.6)` (1 - 0.4 == 0.6).
+* `fork_rng` - boolean. Forks the random number generator state when calling the blend mode. Can be useful for probalistic blend modes like `problerp` which would perturb the RNG and affect stuff like noise used for ancestral sampling, changing your seed even if the ratio is tiny. Example: `problerp:fork_rng=true`
+
+Realistically, the unique parameters for custom blend modes will probably never get documented. Unfortunately, you will need to read the source in `latent_utils.py` to find out what the options are.
+
+## Generally Useful
+
+* `lerp` and friends. Linear interpolation, the most common blend mode. CFG is also just LERP.
+* `slerp` - Can sometimes be better than LERP for blending latents.
+* `inject` - Simple addition. `inject(a, b, 0.3)` is just `a + b * 0.3`. Useful for adding stuff like a CFG diff.
+*
+
+## Garbage/Redundant
+
+* `bislerp_wrong` - This is just LERP with the useless normalize.
+* `hslerp` (anything starting with HSLERP).
+* `colorize` - Literally just LERP.
+* `colordodge`, `difference`, `exclusion`, `glow`, `hardlight`, `linearlight`, `overlay`, `pinlight`, `reflect`, `screen`, `vividlight` - Photoshop filter type modes. They are designed for images and assume certain value ranges for the input so they are essentually useless for blending latents.
+* `linear_dodge` - Same as `inject`. This is just scaled addition.
+* `cosinesimilarity` - I actually like using this but it is roughly just a worse SLERP.
+
+## Experimental/Exotic Modes
+
+Note: Many of these are vibe-mathed, so the description explains what I was attempting to do and what I believe the mode does. I can't guarantee it is doing what it purports to because I don't always fully understand the math involved.
+
+* `ortho` - Orthogonal addition. You may want to specify the dimension parameters to control how the normalization happens. Example: `ortho:start_dim=2:end_dim=3` - For 4D latents, this would normalize over the height/width dimensions. The ortho blend function has many parameters, you will need to look at the source to see them.
+* `ortho_rescaled_lerpish` - Orthogonal blending means the parallel component gets thrown away, in other words you may end up adding less of something than you expected. This mode tries to adjust the result to target something like the result of a LERP.
+* `ortho_dyn_lerp` - Similar to the previous, except it calculates how much `b` got scaled down and LERPs to compensate.
+* `ortho_cfg` and `ortho_cfg_base_a` - Does ortho addition of the CFG diff.
+* `contrastive_ortho_cfg` (and `_base_a`) - Mostly useful for CFG. Let's say we're in base A mode and `a` is cond and `b` is uncond. The mode takes a `a_ortho_scale` (what is unique to cond) and `b_ortho_scale` (what is unique to uncond) parameters. `b_ortho` (AKA what is unique to uncond) gets _subtracted. So `contrastive_ortho_cfg:a_ortho_scale=0.0:b_ortho_scale=1.0` would mean only subtract what is unique to uncond but _don't_ scale up what is unique to cond. The reverse is also possible, only enhance what is unique to cond but don't subtract uncond's unique features. Compare this with normal CFG: `cond + (cond - uncond) * scale` or in other words `cond + cond * scale - uncond * scale`. If `cond` is 1 and uncond is 0 then the result would effectively be `cond + cond * scale` - we scale up cond since there isn't a value on the uncond side to cancel it out.
+* Modes starting with `slice` - Slices along a dimension. The ratio is the size of the dimension multiplied by the ratio. Example: `slice_d1` - slices dimension 1 (second dimension in zero-based dimension indexing). `slice_d1(a, b, 0.5)` would use the first 50% of channels from `a` and the remaining ones from `b`. Or maybe it's the other way around, I forget! These modes also have a `_flip` variant which would make it so the `a` result comes first or vice versa. The blend function supports various parameters from smoothing the result and controling the offset so if you wanted something like just the middle 25% of a dimension that is possible with custom parameters.
+* Modes starting with `wavelet`. `wavelet_b_hi_100_lo_0` means 100% of the high frequency components of `b` and none of the low frequency. `wavelet_b_hi_0_lo_100` is the reverse. Can be interesting as a CFG function. The blend function supports many parameters such as setting the wavelet type and ratios. This requires wavelet support from the `pytorch_wavelets` package which is unfortunately broken with recent Python versions and hasn't been updated in years. You can use my pull with a fix (or the repo it links to): https://github.com/fbcotter/pytorch_wavelets/pull/66
+* `pct_limited_025` - Limits the result to a maximum change (relative to `a` by default). This is a wrapper for other blend modes which defaults to LERP. Example: `pct_limited_025:diff_limit=0.25:blend_mode=inject` Let's say we do `blend(0.5, 100.0, 1.0)` Without limiting this is `0.5 + 100.0 * 1.0` which results in something like a 2,000% change. If we limit to adjusting `a` by 25% at most we get a limit of `0.125` (`0.5 * 0.25`) so the output is `0.5 + 0.125 == 0.625`. This blend function has various other parameters, for example to allow soft clamping, prevent sign flipping (which not limiting the actual value), calculating percentage change over a dimension rather than elementwise, etc.
+* `distro_aligned` - By default matches the distribution of `b` to `a`. `distro_aligned_result` aligns the output from the blend to the distribution of `a`. The blend function supports many options controling which gets aligned to what else, how it occurs, etc. You will need to look at the source.
+* `gaussian_aligned` and `gaussian_aligned_result`. The latter only aligns the result to the Gaussian distribution, the former aligns both `a` and `b` before blending and then also aligns the result. SDXL latents (and probably most latent formats?) should be in the Gaussian distribution (zero mean, std 1) so using this as a CFG function actually works pretty well.
+* `rms_interpolation_lerpsign` - By default, this is roughly `lerp(a**2, b**2, ratio)**0.5` and then copies the sign from `lerp(a, b, ratio)`. There are many options to control what blend mode is used, the power, whether the input gets converted to absolute values and how the sign is restored (since `whatever**2` will always be positive). `magnitude_interpolation_lerpsign` is just a preset for this with power 1 and using absolute inputs.
+* `moment_aligned` - Similar to `distro_aligned`, maybe just a worse version of it. Attempts to align the std and mean (of the input) to a reference. It does not align the result, though you could probably do that through nesting.
+* `pythagorean_lerp` - Like LERP but attempts to preserve the variance of the inputs. Let's say ratio is 0.4, so the multiplier for `a` would be `0.6` and `b` would be `0.4`. We'd calculate a variance division `(0.6**2 + 0.4 ** 2)**0.5` (roughly `0.7211`) and scale the weights (`0.6 / 0.7211 == 0.832`, `0.4 / 0.7211 == 0.554`). If the ratio was `0.1` you'd get something like `a * 0.993 + b * 0.11`.
diff --git a/changelog.md b/changelog.md
index d8202f7..4edbbad 100644
--- a/changelog.md
+++ b/changelog.md
@@ -2,6 +2,20 @@
Note, only relatively significant changes to user-visible functionality will be included here. Most recent changes at the top.
+## 2026TBD
+
+* Support LTX 2.0, LTX 2.3 and LTX 2.3 wide TAE models.
+* Added a `BlehFixGuiderPreviewing` node. Using this is necessary for LTX previews, it can also be used to set the LTX 2.3 wide previewer mode or override the default FPS when generating video previews.
+* ACE-Steps 1.5 latent visualization support in previews.
+* Default previewer datatype is now bfloat16. For the old behavior, you can set this to `null` in the config if you want.
+* Added `/bleh/last_preview.html` endpoint (defaults to disabled) which can serve the last preview in a separate page. Useful for video previews that don't normally animate in a ComfyUI workflow.
+* Many more blend modes.
+* Added a `BlehBlendConditioning` node. This allows blending multiple conditionings on both sides, unlike the built-in ComfyUI conditioning blend node.
+* `BlehBlockCFG` is better now? Maybe. I forget what I changed, but there are changes to it! (Note, still only works for diffusion models like SD 1.5 and SDXL.)
+* Added `BlehProcessLatentIn` and `BlehProcessLatentOut` nodes which can scale/unscale latents to what the model expects during sampling.
+* `previews` is now used for preview-related settings in the config file. `betterTaesdPreviews` will still work in existing config files.
+* Removed the example JSON configuration file. I don't recommend using JSON but it will still work if you really want to.
+
## 20250829
* Added support for Wan 2.2 video previews.
diff --git a/py/better_previews/base.py b/py/better_previews/base.py
index 04b3897..320d774 100644
--- a/py/better_previews/base.py
+++ b/py/better_previews/base.py
@@ -4,44 +4,109 @@
from comfy import latent_formats
+from .tae_vid import TAEVid, TAEVidBase, TAEVidLTX2, TAEVidLTX23Wide
+
if TYPE_CHECKING:
from pathlib import Path
class VideoModelInfo(NamedTuple):
+ name: str
latent_format: latent_formats.LatentFormat
- fps: int = 24
+ fps: int | float = 24
temporal_compression: int = 8
+ temporal_layers: int = 0
patch_size: int = 1
+ nested_tensor_index: int = 0
tae_model: str | Path | None = None
+ tae_class: TAEVidBase | None = TAEVid
VIDEO_FORMATS = {
- "mochi": VideoModelInfo(
- latent_formats.Mochi,
- temporal_compression=6,
- tae_model="taem1.pth",
- ),
- "hunyuanvideo": VideoModelInfo(
- latent_formats.HunyuanVideo,
- temporal_compression=4,
- tae_model="taehv.pth",
- ),
- "cosmos1cv8x8x8": VideoModelInfo(latent_formats.Cosmos1CV8x8x8),
- "wan21": VideoModelInfo(
- latent_formats.Wan21,
- fps=16,
- temporal_compression=4,
- tae_model="taew2_1.pth",
- ),
- "wan22": VideoModelInfo(
- latent_formats.Wan22,
- fps=24,
- temporal_compression=4,
- patch_size=2,
- tae_model="taew2_2.pth",
- ),
+ vmi.name: vmi
+ for vmi in (
+ VideoModelInfo(
+ "mochi",
+ latent_formats.Mochi,
+ temporal_compression=6,
+ tae_model="taem1.pth",
+ ),
+ VideoModelInfo(
+ "hunyuanvideo",
+ latent_formats.HunyuanVideo,
+ temporal_compression=4,
+ tae_model="taehv.pth",
+ ),
+ VideoModelInfo(
+ "hunyuanvideo15",
+ latent_formats.HunyuanVideo15,
+ temporal_compression=4,
+ patch_size=2,
+ tae_model="taehv1_5.pth",
+ ),
+ VideoModelInfo(
+ "cosmos1cv8x8x8",
+ latent_formats.Cosmos1CV8x8x8,
+ ),
+ VideoModelInfo(
+ "wan21",
+ latent_formats.Wan21,
+ fps=16,
+ temporal_compression=4,
+ temporal_layers=2,
+ tae_model="taew2_1.pth",
+ ),
+ VideoModelInfo(
+ "wan22",
+ latent_formats.Wan22,
+ fps=24,
+ temporal_compression=4,
+ temporal_layers=2,
+ patch_size=2,
+ tae_model="taew2_2.pth",
+ ),
+ VideoModelInfo(
+ "ltxv",
+ latent_formats.LTXV,
+ fps=24,
+ patch_size=4,
+ temporal_layers=3,
+ tae_model="taeltx_2.pth",
+ tae_class=TAEVidLTX2,
+ ),
+ VideoModelInfo(
+ "ltxav",
+ latent_formats.LTXV,
+ fps=24,
+ patch_size=4,
+ temporal_layers=3,
+ tae_model="taeltx_2.pth",
+ tae_class=TAEVidLTX2,
+ ),
+ VideoModelInfo(
+ "ltxav23",
+ latent_formats.LTXV,
+ fps=24,
+ patch_size=4,
+ temporal_layers=3,
+ tae_model="taeltx2_3.pth",
+ tae_class=TAEVidLTX2,
+ ),
+ VideoModelInfo(
+ "ltxav23wide",
+ latent_formats.LTXV,
+ fps=24,
+ patch_size=4,
+ temporal_layers=3,
+ tae_model="taeltx2_3_wide.pth",
+ tae_class=TAEVidLTX23Wide,
+ ),
+ )
}
+AMBIGUOUS_VIDEO_FORMATS = {
+ "ltxv": frozenset(("ltxav", "ltxav23", "ltxav23wide")),
+ "ltxav": frozenset(("ltxv", "ltxav23", "ltxav23wide")),
+}
-__all__ = ("VIDEO_FORMATS", "VideoModelInfo")
+__all__ = ("AMBIGUOUS_VIDEO_FORMATS", "VIDEO_FORMATS", "VideoModelInfo")
diff --git a/py/better_previews/last_preview.py b/py/better_previews/last_preview.py
new file mode 100644
index 0000000..e9fadfe
--- /dev/null
+++ b/py/better_previews/last_preview.py
@@ -0,0 +1,125 @@
+from time import time
+from typing import Any
+
+from aiohttp import web
+from server import PromptServer
+
+dum_page = """
+
+
+ bleh preview
+
+
+
+
+ ![]()
+
+
+
+
+"""
+
+
+class LastPreview:
+ image: bytes | None
+ stamp: float | None
+ content_type: str | None
+ duration: int
+
+ def __init__(self, *, min_refresh: int = 5):
+ self.min_refresh = min_refresh
+ self.image = None
+ self.stamp = None
+ self.content_type = None
+ self.duration = 10
+
+ def update(
+ self,
+ *,
+ image_bytes: bytes,
+ content_type: str,
+ stamp: float | None = None,
+ duration: int | None = None,
+ ):
+ self.image = image_bytes
+ self.stamp = time() if stamp is None else stamp
+ self.content_type = content_type
+ self.duration = int(
+ max(self.min_refresh, duration if duration is not None else 0)
+ )
+
+ async def __call__(self, request: web.Request):
+ if request.path.endswith(".html"):
+ return web.Response(body=dum_page, content_type="text/html")
+ if self.image is None or self.content_type is None:
+ raise web.HTTPNotFound(reason="OHNO")
+ return web.Response(
+ body=self.image,
+ content_type=self.content_type,
+ headers={"x-bleh-animation-duration": str(int(self.duration * 1000))},
+ )
+
+
+LAST_PREVIEW = None
+
+
+def init_routes(**kwargs: Any):
+ global LAST_PREVIEW # noqa: PLW0603
+ if LAST_PREVIEW is not None:
+ return
+ LAST_PREVIEW = LastPreview(**kwargs)
+ PromptServer.instance.routes.get("/bleh/last_preview")(LAST_PREVIEW)
+ PromptServer.instance.routes.get("/bleh/last_preview.html")(LAST_PREVIEW)
diff --git a/py/better_previews/previewer.py b/py/better_previews/previewer.py
index 45c2369..ab8dd24 100644
--- a/py/better_previews/previewer.py
+++ b/py/better_previews/previewer.py
@@ -1,12 +1,15 @@
from __future__ import annotations
import math
+from io import BytesIO
from time import time
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Any
+import comfy.utils as comfy_utils
import folder_paths
import latent_preview
import torch
+from comfy import latent_formats
from comfy.cli_args import LatentPreviewMethod
from comfy.cli_args import args as comfy_args
from comfy.model_management import device_supports_non_blocking, vae_dtype
@@ -15,18 +18,28 @@
from tqdm import tqdm
from ..settings import SETTINGS # noqa: TID252
-from .base import VIDEO_FORMATS, VideoModelInfo
+from . import last_preview
+from .base import AMBIGUOUS_VIDEO_FORMATS, VIDEO_FORMATS, VideoModelInfo
from .tae_vid import TAEVid
if TYPE_CHECKING:
+ from collections.abc import Callable
+
import numpy as np
from comfy import latent_formats
+
+class BlehPreviewerState:
+ last_latent_shapes: tuple | None = None
+ fps_override: float | None = None
+ prefer_previewer: str | None = None
+
+
+PREVIEWER_STATE = BlehPreviewerState()
+
_ORIG_PREVIEWER = latent_preview.TAESDPreviewerImpl
_ORIG_GET_PREVIEWER = latent_preview.get_previewer
-LAST_LATENT_FORMAT = None
-
# Referenced from https://github.com/learnables/learn2learn/blob/752200384c3ca8caeb8487b5dd1afd6568e8ec01/learn2learn/utils/__init__.py#L51
def clone_module(module, *, memo: dict | None = None) -> torch.nn.Module:
@@ -82,22 +95,57 @@ def normalize_to_scale(latent, target_min, target_max, *, dim=(-3, -2, -1)):
class ImageWrapper:
- def __init__(self, frames: tuple, frame_duration: int):
- self._frames = frames
+ def __init__(self, frames: tuple | Image, frame_duration: int = 250):
+ self._frames = (frames,) if not isinstance(frames, (tuple, list)) else frames
self._frame_duration = frame_duration
- def save(self, fp, format: str | None, **kwargs: dict): # noqa: A002
- if len(self._frames) == 1:
- return self._frames[0].save(fp, format, **kwargs)
- kwargs |= {
- "loop": 0,
- "save_all": True,
- "append_images": self._frames[1:],
- "duration": self._frame_duration,
- }
- return self._frames[0].save(fp, "webp", **kwargs)
-
- def resize(self, *args: list, **kwargs: dict) -> ImageWrapper:
+ def _save_image(
+ self,
+ frames: tuple[Image, ...],
+ *,
+ format: str | None, # noqa: A002
+ **kwargs: Any,
+ ) -> BytesIO:
+ buf = BytesIO()
+ extra_kwargs = (
+ {}
+ if len(frames) < 2
+ else {
+ "loop": 0,
+ "save_all": True,
+ "append_images": frames[1:],
+ "duration": self._frame_duration,
+ }
+ )
+ frames[0].save(buf, format, **extra_kwargs, **kwargs)
+ return buf
+
+ def save(self, fp, format: str | None, **kwargs: Any): # noqa: A002
+ frames = self._frames
+ publishing = last_preview.LAST_PREVIEW is not None
+ animated = len(frames) > 1
+ split_preview = (
+ animated and publishing and SETTINGS.btp_only_animate_last_preview
+ )
+ result_format = "webp" if animated else (format or "png")
+ result = self._save_image(frames, format=result_format, **kwargs).getvalue()
+ _preview_format, preview_result = (
+ (result_format, result)
+ if not split_preview
+ else (
+ format,
+ self._save_image(frames[:1], format=format, **kwargs).getvalue(),
+ )
+ )
+ if publishing:
+ last_preview.LAST_PREVIEW.update(
+ image_bytes=result,
+ content_type=f"image/{result_format}",
+ duration=2 + int(len(self._frames) / max(1, self._frame_duration)),
+ )
+ fp.write(preview_result)
+
+ def resize(self, *args: Any, **kwargs: Any) -> ImageWrapper:
return ImageWrapper(
tuple(frame.resize(*args, **kwargs) for frame in self._frames),
frame_duration=self._frame_duration,
@@ -123,6 +171,7 @@ def __init__(
self.device = device
raw_factors = latent_format.latent_rgb_factors
raw_bias = latent_format.latent_rgb_factors_bias
+ self.reshape_fun = getattr(latent_format, "latent_rgb_factors_reshape", None)
factors = torch.tensor(raw_factors, device=device, dtype=dtype).transpose(0, 1)
bias = (
torch.tensor(raw_bias, device=device, dtype=dtype)
@@ -144,6 +193,8 @@ def __init__(
@torch.no_grad()
def forward(self, x: torch.Tensor) -> torch.Tensor:
+ if self.reshape_fun is not None:
+ x = self.reshape_fun(x)
x = self.lin(x.movedim(1, -1)).movedim(-1, 1)
x = self.upsample(x).movedim(1, -1)
return x.add_(1.0).clamp_(0.0, 2.0).mul_(127.5).round_()
@@ -156,18 +207,26 @@ def __init__(
*,
dtype: torch.dtype,
device: torch.device,
+ height_factor: int = 4,
+ width_factor: int = 1,
normalize_dims: tuple = (-1,),
):
super().__init__()
self.dtype = dtype
self.device = device
self.normalize_dims = normalize_dims
+ self.height_factor = height_factor
+ self.width_factor = width_factor
@torch.no_grad()
def forward(self, x: torch.Tensor) -> torch.Tensor:
batch, temporal = x.shape[0], x.shape[-1]
x = normalize_to_scale(x, 0.0, 1.0, dim=self.normalize_dims) * 255.0
x = x.reshape(batch, -1, temporal)
+ if self.height_factor > 1:
+ x = x.repeat_interleave(dim=1, repeats=self.height_factor)
+ if self.width_factor > 1:
+ x = x.repeat_interleave(dim=1, repeats=self.width_factor)
return x[..., None].expand(*x.shape, 3)
@@ -179,7 +238,10 @@ def __init__(
latent_format: latent_formats.LatentFormat,
vid_info: VideoModelInfo | None = None,
):
- self.latent_format = latent_format
+ self.orig_latent_format = latent_format
+ self.latent_format = (
+ latent_format if vid_info is None else vid_info.latent_format
+ )
self.latent_format_name = (
"unknown"
if latent_format is None
@@ -406,7 +468,8 @@ def prepare_previewer(
return x0.to(
device=pdevice,
dtype=pdtype,
- non_blocking=device_supports_non_blocking(x0.device),
+ non_blocking=SETTINGS.btp_preview_non_blocking
+ and device_supports_non_blocking(x0.device),
)
def _decode_latent_taevid(self, x0: torch.Tensor) -> tuple[torch.Tensor, int, int]:
@@ -473,12 +536,25 @@ def calc_cols_rows(
rows = math.ceil(batch_size / cols)
return cols, rows
- @classmethod
- def decoded_to_animation(cls, samples: np.ndarray) -> ImageWrapper:
+ def decoded_to_animation(
+ self,
+ samples: np.ndarray,
+ video_frames: int,
+ ) -> ImageWrapper:
batch = samples.shape[0]
+ fps_override = PREVIEWER_STATE.fps_override
+ if self.vid_info is None or not video_frames:
+ frame_duration = 250 if not fps_override else 1000 / fps_override
+ else:
+ time_factor = self.vid_info.temporal_compression / max(
+ 1,
+ self.previewer_model.t_upscale,
+ )
+ ms_frame = 1000.0 / (fps_override or self.vid_info.fps)
+ frame_duration = ms_frame * time_factor
return ImageWrapper(
tuple(Image.fromarray(samples[idx]) for idx in range(batch)),
- frame_duration=250,
+ frame_duration=max(1, int(frame_duration)),
)
def decoded_to_image(
@@ -487,22 +563,23 @@ def decoded_to_image(
cols: int,
rows: int,
*,
- is_video=False,
+ video_frames: int = 0,
) -> Image | ImageWrapper:
batch, (height, width) = samples.shape[0], samples.shape[-3:-1]
samples = samples.to(
device="cpu",
dtype=torch.uint8,
- non_blocking=device_supports_non_blocking(samples.device),
+ non_blocking=SETTINGS.btp_preview_non_blocking
+ and device_supports_non_blocking(samples.device),
).numpy()
if batch == 1:
- self.cached = Image.fromarray(samples[0])
+ self.cached = ImageWrapper((Image.fromarray(samples[0]),))
return self.cached
if SETTINGS.btp_animate_preview == "both" or (
- is_video,
+ video_frames != 0,
SETTINGS.btp_animate_preview,
) in {(True, "video"), (False, "batch")}:
- return self.decoded_to_animation(samples)
+ return self.decoded_to_animation(samples, video_frames=video_frames)
cols, rows = self.calc_cols_rows(batch, width, height)
img_size = (width * cols, height * rows)
if self.cached is not None and self.cached.size == img_size:
@@ -514,10 +591,14 @@ def decoded_to_image(
Image.fromarray(samples[idx]),
box=((idx % cols) * width, ((idx // cols) % rows) * height),
)
- return result
+ return ImageWrapper((result,))
@torch.no_grad()
- def init_fallback_previewer(self, device: torch.device, dtype: torch.dtype) -> bool:
+ def init_fallback_previewer(
+ self,
+ device: torch.device | str,
+ dtype: torch.dtype,
+ ) -> bool:
if self.latent_format is None:
return False
if (
@@ -526,12 +607,14 @@ def init_fallback_previewer(self, device: torch.device, dtype: torch.dtype) -> b
and self.fallback_previewer_model.device == device
):
return True
- if self.latent_format_name == "aceaudio":
+ if self.latent_format_name in {"aceaudio", "aceaudio15"}:
self.fallback_previewer_model = ACEStepsPreviewerModel(
device=device,
dtype=dtype,
)
return True
+ if self.latent_format.latent_rgb_factors is None:
+ return False
self.fallback_previewer_model = FallbackPreviewerModel(
self.latent_format,
device=device,
@@ -563,14 +646,59 @@ def fallback_previewer(self, x0: torch.Tensor, *, quiet=False) -> Image:
except torch.OutOfMemoryError:
return self.blank
+ def ensure_x0_shape(self, x0: torch.Tensor) -> tuple[torch.Tensor, bool]: # noqa: PLR0911
+ expected_channels = self.latent_format.latent_channels
+ expected_ndim = 2 + self.latent_format.latent_dimensions
+ if x0.shape[0] == 0:
+ return x0, False
+ if self.latent_format_name == "aceaudio15" and x0.ndim == expected_ndim + 1:
+ expected_ndim += 1
+ if (
+ x0.ndim > 1
+ and x0.ndim == expected_ndim
+ and x0.shape[1] == expected_channels
+ ):
+ return x0, True
+ last_shapes = PREVIEWER_STATE.last_latent_shapes
+ if not last_shapes or not hasattr(comfy_utils, "unpack_latents"):
+ return x0, False
+ last_numel = sum(math.prod(tshape) for tshape in last_shapes)
+ if last_numel != x0.numel():
+ return x0, False
+ nest_idx = self.vid_info.nested_tensor_index if self.vid_info else 0
+ target_shape = None if len(last_shapes) <= nest_idx else last_shapes[nest_idx]
+ if (
+ # Have to have a nest shape
+ target_shape is None
+ # with at least a channel dimension,
+ or len(target_shape) < 2
+ # with the expected number of dims,
+ or len(target_shape) != expected_ndim
+ # And the correct number of channels.
+ or target_shape[1] != expected_channels
+ ):
+ return x0, False
+ unpacked_latents = comfy_utils.unpack_latents(x0, last_shapes)
+ target_latent = (
+ None if len(unpacked_latents) <= nest_idx else unpacked_latents[nest_idx]
+ )
+ if target_latent is None or target_latent.shape != target_shape:
+ return x0, False
+ return target_latent.reshape(*target_shape), True
+
def decode_latent_to_preview(self, x0: torch.Tensor) -> Image:
if self.check_use_cached():
return self.cached
- if x0.shape[0] == 0:
- return self.blank # Shouldn't actually be possible.
+ x0, can_preview = self.ensure_x0_shape(x0)
+ if not can_preview:
+ return self.blank
if (self.oom_count and not self.oom_retry) or self.previewer_model is None:
return self.fallback_previewer(x0, quiet=True)
is_video = x0.ndim == 5
+ if is_video:
+ # Who would be crazy enough to generate video batches?
+ # We'll just use the last item for now.
+ x0 = x0[-1:, ...]
used_fallback = False
start_time = time()
try:
@@ -579,7 +707,10 @@ def decode_latent_to_preview(self, x0: torch.Tensor) -> Image:
if is_video
else self._decode_latent_taesd(x0)
)
- result = self.decoded_to_image(*dargs, is_video=is_video)
+ result = self.decoded_to_image(
+ *dargs,
+ video_frames=x0.shape[2] if is_video else 0,
+ )
except torch.OutOfMemoryError:
used_fallback = True
result = self.fallback_previewer(x0)
@@ -593,56 +724,70 @@ def decode_latent_to_preview(self, x0: torch.Tensor) -> Image:
def bleh_get_previewer(
device,
latent_format: latent_formats.LatentFormat,
- *args: list,
- **kwargs: dict,
+ *args: Any,
+ **kwargs: Any,
) -> object | None:
def orig_get_previewer():
return _ORIG_GET_PREVIEWER(device, latent_format, *args, **kwargs)
preview_method = comfy_args.preview_method
- if preview_method == LatentPreviewMethod.NoPreviews:
+ if preview_method not in {
+ LatentPreviewMethod.TAESD,
+ LatentPreviewMethod.Auto,
+ LatentPreviewMethod.Latent2RGB,
+ }:
return orig_get_previewer()
format_name = latent_format.__class__.__name__.lower()
+ if PREVIEWER_STATE.prefer_previewer in AMBIGUOUS_VIDEO_FORMATS.get(
+ format_name,
+ frozenset(),
+ ):
+ format_name = PREVIEWER_STATE.prefer_previewer
if (
not SETTINGS.btp_enabled
or format_name in SETTINGS.btp_blacklist
or (SETTINGS.btp_whitelist and format_name not in SETTINGS.btp_whitelist)
):
return orig_get_previewer()
+ if format_name in {"aceaudio", "aceaudio15"}:
+ return BetterPreviewer(latent_format=latent_format)
+ vid_info = VIDEO_FORMATS.get(format_name)
+ eff_latent_format = (
+ vid_info.latent_format if vid_info is not None else latent_format
+ )
tae_model = None
if preview_method in {LatentPreviewMethod.TAESD, LatentPreviewMethod.Auto}:
- vid_info = VIDEO_FORMATS.get(format_name)
- if vid_info is not None and vid_info.tae_model is not None:
+ if (
+ vid_info is not None
+ and vid_info.tae_model is not None
+ and vid_info.tae_class is not None
+ ):
tae_model_path = folder_paths.get_full_path(
"vae_approx",
vid_info.tae_model,
)
- tupscale_limit = SETTINGS.btp_video_temporal_upscale_level
- decoder_time_upscale = tuple(
- i < tupscale_limit for i in range(TAEVid.temporal_upscale_blocks)
- )
tae_model = (
- TAEVid(
+ vid_info.tae_class(
checkpoint_path=tae_model_path,
vmi=vid_info,
device=torch.device("cpu"),
- decoder_time_upscale=decoder_time_upscale,
+ decoder_time_upscale_level=SETTINGS.btp_video_temporal_upscale_level,
)
if tae_model_path is not None
else None
)
- if tae_model is None and latent_format.taesd_decoder_name is not None:
+ elif vid_info is None and eff_latent_format.taesd_decoder_name is not None:
taesd_path = folder_paths.get_full_path(
"vae_approx",
- f"{latent_format.taesd_decoder_name}.pth",
+ f"{eff_latent_format.taesd_decoder_name}.pth",
)
tae_model = (
TAESD(
None,
taesd_path,
- latent_channels=latent_format.latent_channels,
+ latent_channels=eff_latent_format.latent_channels,
)
if taesd_path is not None
else None
@@ -650,10 +795,11 @@ def orig_get_previewer():
if tae_model is not None:
return BetterPreviewer(
taesd=tae_model,
- latent_format=latent_format,
+ latent_format=eff_latent_format,
vid_info=vid_info,
)
- if format_name == "aceaudio" or latent_format.latent_rgb_factors is not None:
+ # Using Latent2RGB either via setting or because no preview model.
+ if eff_latent_format.latent_rgb_factors is not None:
return BetterPreviewer(latent_format=latent_format)
return orig_get_previewer()
diff --git a/py/better_previews/tae_vid.py b/py/better_previews/tae_vid.py
index 325f2f0..f1affce 100644
--- a/py/better_previews/tae_vid.py
+++ b/py/better_previews/tae_vid.py
@@ -4,7 +4,7 @@
from __future__ import annotations
-from typing import TYPE_CHECKING, NamedTuple
+from typing import TYPE_CHECKING, Any, NamedTuple
import torch
from torch import nn
@@ -24,33 +24,75 @@ class TWorkItem(NamedTuple):
block_index: int
-def conv(n_in: int, n_out: int, **kwargs: dict) -> nn.Conv2d:
- return nn.Conv2d(n_in, n_out, 3, padding=1, **kwargs)
+def conv(
+ n_in: int,
+ n_out: int,
+ *,
+ kernel_size: int = 3,
+ stride: int = 1,
+ padding: int = 1,
+ **kwargs: Any,
+) -> nn.Conv2d:
+ return nn.Conv2d(
+ n_in,
+ n_out,
+ kernel_size=kernel_size,
+ stride=stride,
+ padding=padding,
+ **kwargs,
+ )
class Clamp(nn.Module):
@classmethod
def forward(cls, x: torch.Tensor) -> torch.Tensor:
- return torch.tanh(x / 3) * 3
+ return (x / 3.0).tanh_().mul_(3.0)
class MemBlock(nn.Module):
- def __init__(self, n_in, n_out):
+ def __init__(
+ self,
+ n_in: int,
+ n_out: int,
+ *,
+ wide: bool = False,
+ ):
super().__init__()
- self.conv = nn.Sequential(
- conv(n_in * 2, n_out),
- nn.ReLU(inplace=True),
- conv(n_out, n_out),
- nn.ReLU(inplace=True),
- conv(n_out, n_out),
- )
+ groups = max(1, n_out // 64) if wide else 1
+ if wide:
+ if n_out % groups != 0:
+ errstr = f"Bad n_out {n_out} parameter for wide MemBlock, must be divisible by 64"
+ raise ValueError(errstr)
+ self.conv = nn.Sequential(
+ conv(n_in * 2, n_out, kernel_size=1, padding=0),
+ nn.ReLU(inplace=True),
+ conv(n_out, n_out, groups=groups),
+ nn.ReLU(inplace=True),
+ conv(n_in, n_out, kernel_size=1, padding=0),
+ nn.ReLU(inplace=True),
+ conv(n_out, n_out, groups=groups),
+ )
+ else:
+ self.conv = nn.Sequential(
+ conv(n_in * 2, n_out),
+ nn.ReLU(inplace=True),
+ conv(n_out, n_out),
+ nn.ReLU(inplace=True),
+ conv(n_out, n_out),
+ )
self.skip = (
nn.Conv2d(n_in, n_out, 1, bias=False) if n_in != n_out else nn.Identity()
)
self.act = nn.ReLU(inplace=True)
def forward(self, x: torch.Tensor, past: torch.Tensor) -> torch.Tensor:
- return self.act(self.conv(torch.cat((x, past), 1)) + self.skip(x))
+ result = self.conv(torch.cat((x, past), 1))
+ result += self.skip(x)
+ return self.act(result)
+
+
+def make_memblocks(n: int, *, count: int = 3, **kwargs: Any) -> tuple[MemBlock, ...]:
+ return tuple(MemBlock(n, n, **kwargs) for _ in range(count))
class TPool(nn.Module):
@@ -183,8 +225,8 @@ def apply(self, x: torch.Tensor, *, show_progress=False) -> torch.Tensor:
return torch.stack(out, 1)
-class TAEVid(nn.Module):
- temporal_upscale_blocks = 2
+class TAEVidBase(nn.Module):
+ temporal_upscale_blocks = 3
spatial_upscale_blocks = 3
_nf = (256, 128, 64, 64)
@@ -195,67 +237,106 @@ def __init__(
vmi: VideoModelInfo,
image_channels: int = 3,
device="cpu",
- decoder_time_upscale=(True, True),
- decoder_space_upscale=(True, True, True),
+ encoder_time_downscale_level: int = 3,
+ decoder_time_upscale_level: int = 3,
+ decoder_space_upscale_level: int = 3,
):
- n_f = self._nf
super().__init__()
self.vmi = vmi
- self.latent_channels = vmi.latent_format.latent_channels
self.image_channels = image_channels
+ self.latent_channels = vmi.latent_format.latent_channels
self.patch_size = vmi.patch_size
- self.encoder = nn.Sequential(
- conv(image_channels * self.patch_size**2, 64),
- nn.ReLU(inplace=True),
- TPool(64, 2),
- conv(64, 64, stride=2, bias=False),
- MemBlock(64, 64),
- MemBlock(64, 64),
- MemBlock(64, 64),
- TPool(64, 2),
- conv(64, 64, stride=2, bias=False),
- MemBlock(64, 64),
- MemBlock(64, 64),
- MemBlock(64, 64),
- TPool(64, 1),
- conv(64, 64, stride=2, bias=False),
- MemBlock(64, 64),
- MemBlock(64, 64),
- MemBlock(64, 64),
- conv(64, vmi.latent_format.latent_channels),
+ encoder_time_downscale = self._get_encoder_flags(
+ time_level=encoder_time_downscale_level,
+ )
+ decoder_time_upscale, decoder_space_upscale = self._get_decoder_flags(
+ time_level=decoder_time_upscale_level,
+ space_level=decoder_space_upscale_level,
)
- self.frames_to_trim = 2 ** sum(decoder_time_upscale) - 1
- self.decoder = nn.Sequential(
+ encoder_strides = tuple(1 + int(flag) for flag in encoder_time_downscale)
+ decoder_strides = tuple(1 + int(flag) for flag in decoder_time_upscale)
+ decoder_scale_factors = tuple(1 + int(flag) for flag in decoder_space_upscale)
+ self.encoder = self._build_encoder(strides=encoder_strides)
+ self.decoder = self._build_decoder(
+ strides=decoder_strides,
+ scale_factors=decoder_scale_factors,
+ )
+ self.t_upscale = 2 ** sum(decoder_time_upscale)
+ self.t_downscale = 2 ** sum(encoder_time_downscale)
+ self.frames_to_trim = self.t_upscale - 1
+ if checkpoint_path is None:
+ return
+ sd = torch.load(checkpoint_path, map_location=device, weights_only=True)
+ self.load_state_dict(self.patch_tgrow_layers(sd))
+
+ def _get_decoder_flags(
+ self,
+ *,
+ time_level: int = 3,
+ space_level: int = 3,
+ ) -> tuple[tuple[bool, ...], tuple[bool, ...]]:
+ decoder_time_upscale = tuple(i < time_level for i in range(3))
+ decoder_space_upscale = tuple(i < space_level for i in range(3))
+ return decoder_time_upscale, decoder_space_upscale
+
+ def _get_encoder_flags(
+ self,
+ *,
+ time_level: int = 3,
+ ) -> tuple[bool, ...]:
+ return tuple(i < time_level for i in range(3))
+
+ def _build_decoder(
+ self,
+ *,
+ strides: tuple[int, ...],
+ scale_factors: tuple[int, ...],
+ memblock_kwargs: dict[str, Any] | None = None,
+ ) -> nn.Module:
+ if memblock_kwargs is None:
+ memblock_kwargs = {}
+ n_f = self._nf
+ return nn.Sequential(
Clamp(),
- conv(vmi.latent_format.latent_channels, n_f[0]),
+ conv(self.latent_channels, n_f[0]),
nn.ReLU(inplace=True),
- MemBlock(n_f[0], n_f[0]),
- MemBlock(n_f[0], n_f[0]),
- MemBlock(n_f[0], n_f[0]),
- nn.Upsample(scale_factor=2 if decoder_space_upscale[0] else 1),
- TGrow(n_f[0], 1),
+ *make_memblocks(n_f[0], **memblock_kwargs),
+ nn.Upsample(scale_factor=scale_factors[0]),
+ TGrow(n_f[0], strides[0]),
conv(n_f[0], n_f[1], bias=False),
- MemBlock(n_f[1], n_f[1]),
- MemBlock(n_f[1], n_f[1]),
- MemBlock(n_f[1], n_f[1]),
- nn.Upsample(scale_factor=2 if decoder_space_upscale[1] else 1),
- TGrow(n_f[1], 2 if decoder_time_upscale[0] else 1),
+ *make_memblocks(n_f[1], **memblock_kwargs),
+ nn.Upsample(scale_factor=scale_factors[1]),
+ TGrow(n_f[1], strides[1]),
conv(n_f[1], n_f[2], bias=False),
- MemBlock(n_f[2], n_f[2]),
- MemBlock(n_f[2], n_f[2]),
- MemBlock(n_f[2], n_f[2]),
- nn.Upsample(scale_factor=2 if decoder_space_upscale[2] else 1),
- TGrow(n_f[2], 2 if decoder_time_upscale[1] else 1),
+ *make_memblocks(n_f[2], **memblock_kwargs),
+ nn.Upsample(scale_factor=scale_factors[2]),
+ TGrow(n_f[2], strides[2]),
conv(n_f[2], n_f[3], bias=False),
nn.ReLU(inplace=True),
- conv(n_f[3], image_channels * self.patch_size**2),
+ conv(n_f[3], self.image_channels * self.patch_size**2),
)
- if checkpoint_path is None:
- return
- self.load_state_dict(
- self.patch_tgrow_layers(
- torch.load(checkpoint_path, map_location=device, weights_only=True),
- ),
+
+ def _build_encoder(
+ self,
+ *,
+ strides: tuple[int, ...],
+ memblock_kwargs: dict[str, Any] | None = None,
+ ) -> nn.Module:
+ if memblock_kwargs is None:
+ memblock_kwargs = {}
+ return nn.Sequential(
+ conv(self.image_channels * self.patch_size**2, 64),
+ nn.ReLU(inplace=True),
+ TPool(64, strides[0]),
+ conv(64, 64, stride=2, bias=False),
+ *make_memblocks(64, **memblock_kwargs),
+ TPool(64, strides[1]),
+ conv(64, 64, stride=2, bias=False),
+ *make_memblocks(64, **memblock_kwargs),
+ TPool(64, strides[2]),
+ conv(64, 64, stride=2, bias=False),
+ *make_memblocks(64, **memblock_kwargs),
+ conv(64, self.latent_channels),
)
def patch_tgrow_layers(self, sd: dict) -> dict:
@@ -304,8 +385,15 @@ def apply(
show_progress=False,
) -> torch.Tensor:
model = self.decoder if decode else self.encoder
- if not decode and self.vmi.patch_size > 1:
- x = F.pixel_unshuffle(x, self.patch_size)
+ if not decode:
+ if self.vmi.patch_size > 1:
+ x = F.pixel_unshuffle(x, self.patch_size)
+ # Pad handling copied from https://github.com/madebyollin
+ if x.shape[1] % self.t_downscale != 0:
+ # pad at end to multiple of self.t_downscale
+ n_pad = self.t_downscale - x.shape[1] % self.t_downscale
+ padding = x[:, -1:].repeat_interleave(n_pad, dim=1)
+ x = torch.cat((x, padding), 1)
if parallel:
result = self.apply_parallel(x, model, show_progress=show_progress)
else:
@@ -324,3 +412,61 @@ def encode(self, *args: list, **kwargs: dict) -> torch.Tensor:
def forward(self, x: torch.Tensor) -> torch.Tensor:
return self.c(x)
+
+
+class TAEVid(TAEVidBase):
+ def _get_decoder_flags(
+ self,
+ *,
+ time_level: int = 3,
+ space_level: int = 3,
+ ) -> tuple[tuple[bool, ...], tuple[bool, ...]]:
+ tu, su = super()._get_decoder_flags(
+ time_level=time_level,
+ space_level=space_level,
+ )
+ return (False, *tu[:2]), su
+
+ def _get_encoder_flags(
+ self,
+ *,
+ time_level: int = 3,
+ ) -> tuple[bool, ...]:
+ return (*super()._get_encoder_flags(time_level=time_level)[:2], False)
+
+
+class TAEVidLTX2(TAEVidBase):
+ def _get_decoder_flags(
+ self,
+ *,
+ time_level: int = 3,
+ space_level: int = 3,
+ ) -> tuple[tuple[bool, ...], tuple[bool, ...]]:
+ _tu, su = super()._get_decoder_flags(
+ time_level=time_level,
+ space_level=space_level,
+ )
+ return (True, True, True), su
+
+ def _get_encoder_flags(
+ self,
+ *,
+ time_level: int = 3, # noqa: ARG002
+ ) -> tuple[bool, ...]:
+ return (True, True, True)
+
+
+class TAEVidLTX23Wide(TAEVidLTX2):
+ def __init__(self, *args: Any, **kwargs: Any):
+ self._nf = (1024, 512, 256, 64)
+ super().__init__(*args, **kwargs)
+
+ def _build_decoder(
+ self,
+ *args: Any,
+ memblock_kwargs: dict[str, Any] | None = None,
+ **kwargs: Any,
+ ) -> nn.Module:
+ memblock_kwargs = {} if memblock_kwargs is None else memblock_kwargs.copy()
+ memblock_kwargs["wide"] = True
+ return super()._build_decoder(*args, memblock_kwargs=memblock_kwargs, **kwargs)
diff --git a/py/latent_utils.py b/py/latent_utils.py
index a75364d..73a8a35 100644
--- a/py/latent_utils.py
+++ b/py/latent_utils.py
@@ -5,18 +5,36 @@
import math
import os
from functools import partial
-from typing import ClassVar
+from typing import TYPE_CHECKING, Any, ClassVar, TypeVar
import kornia.filters as kf
import numpy as np
import torch
import torch.nn.functional as nnf
from torch import FloatTensor, LongTensor, fft
+from tqdm import tqdm
+
+from . import wavelet_functions as wavef
+
+if TYPE_CHECKING:
+ from collections.abc import Callable, Sequence
OVERRIDE_NO_SCALE = "COMFYUI_BLEH_OVERRIDE_NO_SCALE" in os.environ
USE_ORIG_NORMALIZE = "COMFYUI_BLEH_ORIG_NORMALIZE" in os.environ
+def pass_kwargs(*args: Any, **kwargs: Any) -> dict[str, Any]:
+ if args:
+ if not all(isinstance(a, dict) for a in args):
+ raise ValueError("Can only pass a single dict positionally")
+ a0 = args[0]
+ for a in args[1:]:
+ a0.update(a)
+ a0.update(kwargs)
+ kwargs = a0
+ return {k.removesuffix("_"): v for k, v in kwargs.items()}
+
+
def normalize_orig(latent, target_min=None, target_max=None, **_unused_kwargs: dict):
min_val = latent.min()
max_val = latent.max()
@@ -71,6 +89,160 @@ def normalize_to_scale(
)
+def soft_clamp(
+ t: torch.Tensor,
+ min_val: torch.Tensor | float = 0.0,
+ max_val: torch.Tensor | float = 1.0,
+ *,
+ # We define stiffness as a multiplier (beta) for the softplus function.
+ # Higher stiffness = sharper transition.
+ stiffness: float = 1.0,
+ safe: bool = True,
+) -> torch.Tensor:
+ if stiffness < 1e-04:
+ return t.clamp(min_val, max_val)
+ if not isinstance(min_val, torch.Tensor):
+ min_val = t.new_tensor(min_val)
+ if not isinstance(max_val, torch.Tensor):
+ max_val = t.new_tensor(max_val)
+
+ # Calculate how much we are exceeding the Max
+ # softplus(beta * x) / beta
+ upper_overshoot = nnf.softplus((t - max_val).mul_(stiffness)).div_(-stiffness)
+
+ # Calculate how much we are falling short of the Min
+ lower_undershoot = nnf.softplus((min_val - t).mul_(stiffness)).div_(stiffness)
+
+ # Apply corrections:
+ # Original - (Amount over max) + (Amount under min)
+ t = upper_overshoot.add_(t).add_(lower_undershoot)
+ return t.clamp_(min_val, max_val) if safe else t
+
+
+def force_gaussian_distribution(
+ t: torch.Tensor,
+ *,
+ start_dim: int = 1,
+ end_dim: int = -1,
+ # Invert the argsorts, option for crazy people. Not recommended.
+ invert1: bool = False,
+ invert2: bool = False,
+ eps: float = 1e-08,
+) -> torch.Tensor:
+ if start_dim < 0:
+ start_dim = t.ndim + start_dim
+ orig_shape = t.shape
+ t_flat = t.flatten(start_dim=start_dim, end_dim=end_dim).movedim(start_dim, -1)
+
+ # Get the rank of each element (0 to N-1)
+ # Double argsort safely returns the rank of the original elements
+ ranks = (
+ t_flat.argsort(dim=-1, descending=invert1)
+ .argsort(dim=-1, descending=invert2)
+ .to(t)
+ )
+
+ # Map ranks to a uniform distribution (0.0 to 1.0 exclusive)
+ # then to a Gaussian curve.
+ factor = max(eps, t_flat.shape[-1] / 2)
+ gaussian = ranks.div_(factor).add_(0.5 / factor - 1).erfinv_().mul_(2**0.5)
+
+ return gaussian.movedim(-1, start_dim).reshape(orig_shape)
+
+
+# Forces source to the distribution of reference.
+def match_distribution(
+ source: torch.Tensor,
+ *,
+ reference: torch.Tensor,
+ start_dim: int = 1,
+ end_dim: int = -1,
+ # Invert the sorts, option for crazy people. Not recommended.
+ invert1: bool = False,
+ invert2: bool = False,
+ invert3: bool = False,
+) -> torch.Tensor:
+ if source is reference:
+ return source.clone()
+ if start_dim < 0:
+ start_dim = source.ndim + start_dim
+ orig_shape = source.shape
+ s_flat = source.flatten(
+ start_dim=start_dim,
+ end_dim=end_dim,
+ ).movedim(start_dim, -1)
+ r_flat = reference.flatten(
+ start_dim=start_dim,
+ end_dim=end_dim,
+ ).movedim(start_dim, -1)
+
+ r_sorted = r_flat.sort(dim=-1, descending=invert1).values
+ s_ranks = s_flat.argsort(
+ dim=-1,
+ descending=invert2,
+ ).argsort(dim=-1, descending=invert3)
+
+ # 4. Give the source elements the values from the reference.
+ return (
+ r_sorted.gather(dim=-1, index=s_ranks)
+ .movedim(-1, start_dim)
+ .reshape(orig_shape)
+ )
+
+
+# Scales the source tensor to match the median and variance of the reference.
+def robust_scale_match(
+ source: torch.Tensor,
+ *,
+ reference: torch.Tensor | None = None,
+ # Default MAD if the reference is not passed. Targets the Gaussian distribution.
+ mad: float = 0.6745,
+ start_dim: int = 1,
+ end_dim: int = -1,
+ eps: float = 1e-8,
+) -> torch.Tensor:
+ if start_dim < 0:
+ start_dim = source.ndim + start_dim
+ orig_shape = source.shape
+ source = source.flatten(start_dim=start_dim, end_dim=end_dim).movedim(start_dim, -1)
+ # Find the median and spread (MAD) of the source
+ src_sub_median = source - source.median(dim=-1, keepdim=True).values
+ s_mad = (
+ src_sub_median.abs()
+ .median(
+ dim=-1,
+ keepdim=True,
+ )
+ .values.clamp_min_(eps)
+ )
+
+ # If no reference, target a Standard Gaussian scale.
+ # (A standard Gaussian has a median of 0 and a MAD of ~0.6745)
+ if reference is None:
+ mad = min(-eps, mad) if mad < 0 else max(eps, mad)
+ return (
+ src_sub_median.mul_(s_mad.reciprocal_().mul_(mad))
+ .movedim(-1, start_dim)
+ .reshape(orig_shape)
+ )
+ reference = reference.flatten(
+ start_dim=start_dim,
+ end_dim=end_dim,
+ ).movedim(start_dim, -1)
+
+ # Find the reference median and spread
+ r_median = reference.median(dim=-1, keepdim=True).values
+ r_mad = (reference - r_median).abs_().median(dim=-1, keepdim=True).values
+
+ # Stretch the source to match the reference
+ return (
+ src_sub_median.mul_(r_mad.div_(s_mad))
+ .add_(r_median)
+ .movedim(-1, start_dim)
+ .reshape(orig_shape)
+ )
+
+
def hslerp(a, b, t):
if a.shape != b.shape:
raise ValueError("Input tensors a and b must have the same shape.")
@@ -88,7 +260,9 @@ def hslerp(a, b, t):
interpolation_tensor[0, 0, 0, 0] = 1.0 if t < 0.5 else -1.0
result = (1 - t) * a + t * b
- result += (torch.norm(b - a, dim=1, keepdim=True) / 6) * interpolation_tensor
+ result += (
+ torch.linalg.vector_norm(b - a, dim=1, keepdim=True) / 6
+ ) * interpolation_tensor
return result
@@ -105,7 +279,7 @@ def hslerp_alt(a, b, t):
)
interp[0, 0] = 1.0
result = (1 - t) * a + t * b
- norm = (torch.norm(b - a, dim=1, keepdim=True) / 6) * interp
+ norm = (torch.linalg.vector_norm(b - a, dim=1, keepdim=True) / 6) * interp
norm[t.broadcast_to(norm.shape) < 0.5] *= -1
return result.add_(norm)
@@ -121,7 +295,7 @@ def hslerp_alt2(a, b, t, *, sign_order=(1.0, -1.0), sign_threshold=0.5):
((1 - t) * a)
.add_(t * b)
.add_(
- torch.norm(b - a, dim=1, keepdim=True).div_(6)
+ torch.linalg.vector_norm(b - a, dim=1, keepdim=True).div_(6)
* torch.where(t_expanded.abs() < sign_threshold, *sign_order),
)
)
@@ -132,8 +306,8 @@ def slerp_orig(b1, b2, r):
c = b1.shape[-1]
# norms
- b1_norms = torch.norm(b1, dim=-1, keepdim=True)
- b2_norms = torch.norm(b2, dim=-1, keepdim=True)
+ b1_norms = torch.linalg.vector_norm(b1, dim=-1, keepdim=True)
+ b2_norms = torch.linalg.vector_norm(b2, dim=-1, keepdim=True)
# normalize
b1_normalized = b1 / b1_norms
@@ -163,7 +337,7 @@ def slerp_orig(b1, b2, r):
# From https://gist.github.com/Birch-san/230ac46f99ec411ed5907b0a3d728efa
-def altslerp( # noqa: PLR0914
+def altslerp(
v0: FloatTensor,
v1: FloatTensor,
t: float | FloatTensor,
@@ -228,6 +402,7 @@ def stochasistic_blend(
fuzz=0.1,
clamp_t: bool | tuple = True,
blend=torch.lerp,
+ **kwargs: Any,
):
if not isinstance(t, torch.Tensor):
t = torch.tensor((t,), dtype=a.dtype, device=a.device)
@@ -248,7 +423,7 @@ def stochasistic_blend(
elif clamp_t:
tmin, tmax = t_orig.aminmax()
tadj = tadj.clamp_(min(0, tmin), max(1.0, tmax))
- return blend(a, b, tadj)
+ return blend(a, b, tadj, **pass_kwargs(kwargs))
def gaussian_smoothing(
@@ -283,9 +458,9 @@ def gaussian_smoothing(
elif ndim != 4:
raise ValueError("Can't handle tensor shape")
if len(kernel_size) == 1:
- kernel_size = kernel_size * 2 # noqa: PLR6104
+ kernel_size = kernel_size * 2
if len(sigma) == 1:
- sigma = sigma * 2 # noqa: PLR6104
+ sigma = sigma * 2
result = kf.gaussian_blur2d(t, kernel_size, sigma)
if ndim == 5:
return result.reshape(*ts)
@@ -307,7 +482,7 @@ def __call__(
*,
cpu=False,
collapse_dims=(),
- **kwargs: dict,
+ **kwargs: Any,
):
t_device = torch.device("cpu") if cpu else a.device
if not isinstance(t, torch.Tensor):
@@ -316,7 +491,7 @@ def __call__(
t = t.detach().clone().to(t_device)
tmin, tmax = t.aminmax()
tmin, tmax = min(tmin, 0.0), max(tmax, 1.0)
- t = t - tmin # noqa: PLR6104
+ t = t - tmin
tdiv = tmax - tmin
if tdiv != 0:
t /= tdiv
@@ -330,7 +505,7 @@ def __call__(
else:
prob_shape = a.shape
t = torch.bernoulli(t.clamp_(0, 1).broadcast_to(prob_shape)).to(a)
- return self.output(a, b, t, **kwargs)
+ return self.output(a, b, t, **pass_kwargs(kwargs))
class ProbBlendSmoothed(ProbBlend):
@@ -363,6 +538,7 @@ def gradient_blend_(
dim=-1,
scaling_constant=0.9,
blend_function=torch.lerp,
+ **kwargs: Any,
) -> torch.Tensor:
dim = max(0, min(a.ndim - 1, a.ndim + dim if dim < 0 else dim))
if not isinstance(t, torch.Tensor):
@@ -379,7 +555,7 @@ def gradient_blend_(
if scaling_constant != 1:
ratios *= scaling_constant
ratios = ratios.view(tuple(1 if i != dim else -1 for i in range(a.ndim)))
- return blend_function(a, b, ratios)
+ return blend_function(a, b, ratios, **pass_kwargs(kwargs))
def gradient_blend(
@@ -390,6 +566,7 @@ def gradient_blend(
flatten_start_dim=1,
scaling_constant=0.9,
blend_function=torch.lerp,
+ **kwargs: Any,
) -> torch.Tensor:
shape = a.shape
# print("\nBLEND:", t)
@@ -407,7 +584,7 @@ def gradient_blend(
)
if scaling_constant != 1:
ratios *= scaling_constant
- result = blend_function(a, b, ratios)
+ result = blend_function(a, b, ratios, **pass_kwargs(kwargs))
if result.shape != shape:
return result.reshape(*shape).contiguous()
return result
@@ -451,7 +628,7 @@ def slice_blend(
return result.reshape(orig_shape)
-def slice_blend_smooth( # noqa: PLR0914
+def slice_blend_smooth(
a: torch.Tensor,
b: torch.Tensor,
t: float | torch.Tensor,
@@ -465,6 +642,7 @@ def slice_blend_smooth( # noqa: PLR0914
b_blend_max: float = 1.0,
invert: bool = False, # Doesn't work propertly at the moment.
blend_function=torch.lerp,
+ **kwargs: Any,
) -> torch.Tensor:
if isinstance(t, torch.Tensor):
t = t.mean().clamp(0, 1)
@@ -525,7 +703,7 @@ def slice_blend_smooth( # noqa: PLR0914
blend_mask = blend_mask.view(
tuple(dim_els if d == dim else 1 for d in range(a.ndim)),
)
- return blend_function(a, b, blend_mask).reshape(orig_shape)
+ return blend_function(a, b, blend_mask, **pass_kwargs(kwargs)).reshape(orig_shape)
def lop_lerp(
@@ -541,7 +719,7 @@ def lop_lerp(
return (a_ratio - t.clamp(max=a_ratio)).mul(a).add_(b * (t * b_ratio))
-# # Thanks, ChatGPT though you did get the ratio reversed.
+# Thanks, ChatGPT though you did get the ratio reversed.
def cosine_similarity_blend_chatgpt_orig(
b: torch.Tensor,
a: torch.Tensor,
@@ -657,8 +835,8 @@ def cosine_similarity_blend_deepseek( # noqa: PLR0914
ratio = a.new_tensor(ratio)
# Compute magnitudes of a and b along the specified dimension
- mag_a = torch.norm(a, p=2, dim=dim, keepdim=True).add_(eps)
- mag_b = torch.norm(b, p=2, dim=dim, keepdim=True).add_(eps)
+ mag_a = torch.linalg.vector_norm(a, p=2, dim=dim, keepdim=True).add_(eps)
+ mag_b = torch.linalg.vector_norm(b, p=2, dim=dim, keepdim=True).add_(eps)
# Avoid division by zero during normalization
a_norm = a / mag_a
@@ -717,9 +895,9 @@ def cosine_similarity_blend(
a: torch.Tensor,
b: torch.Tensor,
t: float | torch.Tensor,
- *args: list,
+ *args: Any,
backend=DEFAULT_COSINE_SIMILARITY_BLEND_BACKEND,
- **kwargs: dict,
+ **kwargs: Any,
) -> torch.Tensor:
fun = COSINE_SIMILARITY_BLEND_BACKENDS.get(backend)
if fun is None:
@@ -798,6 +976,1717 @@ def blend_blend(
)
+def ortho_blend(
+ a: torch.Tensor,
+ b: torch.Tensor,
+ t: torch.Tensor | float,
+ *,
+ blend_mode: str | Callable | None = None,
+ proj_scale: float = -1.0,
+ ortho_scale: float = 1.0,
+ start_dim: int = 1,
+ end_dim: int = -1,
+ rescale_limit: float = 0.0,
+ # a, b, blend or None
+ rescale_result_mode: str | None = None,
+ # When rescale_target mode is blend, will use blend_mode if None.
+ rescale_result_blend_mode: str | Callable | None = None,
+ # LERP if None.
+ dyn_result_blend_mode: str | Callable | None = None,
+ dyn_ortho_mode: bool = False,
+ dyn_min_scale: float = 0.0,
+ dyn_max_scale: float = 1.0,
+ # Can only be used when the flattened tensor has 4 dimensions left.
+ smooth_factor_kernel_size: int | tuple[int, ...] = 0,
+ ortho_verbose: bool = False,
+ eps: float = 1e-06,
+) -> torch.Tensor:
+ orig_shape = a.shape
+ if not isinstance(smooth_factor_kernel_size, tuple):
+ smooth_factor_kernel_size = (smooth_factor_kernel_size | 1,)
+ else:
+ smooth_factor_kernel_size = tuple(sz | 1 for sz in smooth_factor_kernel_size)
+ ndim = a.ndim
+ if start_dim < 0:
+ start_dim = max(0, min(ndim + start_dim, ndim - 1))
+ if end_dim < 0:
+ end_dim = max(0, min(ndim + end_dim, ndim - 1))
+ if start_dim > end_dim:
+ start_dim, end_dim = end_dim, start_dim
+ if not isinstance(t, torch.Tensor):
+ t = a.new_tensor(t)
+ sync_t = False
+ else:
+ t = t.broadcast_to(a.shape)
+ sync_t = True
+ if sync_t:
+ t = t.flatten(start_dim=start_dim, end_dim=end_dim)
+ a = a.flatten(start_dim=start_dim, end_dim=end_dim)
+ b = b.flatten(start_dim=start_dim, end_dim=end_dim)
+ if end_dim != ndim - 1:
+ a = a.movedim(start_dim, -1)
+ b = b.movedim(start_dim, -1)
+ if sync_t:
+ t = t.movedim(start_dim, -1)
+ if start_dim == 0:
+ a = a.unsqueeze(0)
+ b = b.unsqueeze(0)
+ if sync_t:
+ t = t.unsqueeze(0)
+ b_normed = b.norm(dim=-1, keepdim=True) if rescale_limit else None
+ dot_ba = (b * a).sum(dim=-1, keepdim=True)
+ dot_aa = (a**2).sum(dim=-1, keepdim=True)
+ proj = (dot_ba / (dot_aa + eps)) * a
+ proj *= proj_scale
+ b_ortho = proj.add_(b if ortho_scale == 1.0 else b * ortho_scale)
+ if b_normed is not None:
+ rescale_limit = abs(rescale_limit)
+ if rescale_limit == 1:
+ rescale_limit += eps
+ b_ortho_normed = b_ortho.norm(dim=-1, keepdim=True)
+ b_ortho_normed += eps
+ b_normed /= b_ortho_normed
+ b_normed = b_normed.clamp_(-rescale_limit, rescale_limit)
+ b_ortho *= b_normed
+ if blend_mode is None:
+
+ def blend_function(a, b, t):
+ return (b * t).add_(a)
+ else:
+ blend_function = (
+ BLENDING_MODES[blend_mode] if isinstance(blend_mode, str) else blend_mode
+ )
+ ortho_result = blend_function(a, b_ortho, t)
+ rr_blend_function = None
+ if rescale_result_mode == "a":
+ rescale_result_target = a
+ elif rescale_result_mode == "b":
+ rescale_result_target = b
+ elif rescale_result_mode == "blend":
+ rr_blend_function = (
+ blend_function
+ if rescale_result_blend_mode is None
+ else (
+ BLENDING_MODES[rescale_result_blend_mode]
+ if isinstance(rescale_result_blend_mode, str)
+ else rescale_result_blend_mode
+ )
+ )
+ rescale_result_target = rr_blend_function(a, b, t)
+ else:
+ rescale_result_target = None
+ if rr_blend_function is None:
+ rr_blend_function = blend_function
+ if rescale_result_target is not None:
+ result_norm = ortho_result.norm(dim=-1, keepdim=True).add_(eps)
+ target_norm = rescale_result_target.norm(dim=-1, keepdim=True)
+ target_norm /= result_norm
+ ortho_result *= target_norm
+ if b_normed is not None and dyn_ortho_mode:
+ vanilla_result = (
+ rr_blend_function(a, b, t)
+ if rescale_result_mode != "blend"
+ else rescale_result_target
+ )
+ dyn_blend_function = (
+ torch.lerp
+ if dyn_result_blend_mode is None
+ else (
+ BLENDING_MODES[dyn_result_blend_mode]
+ if isinstance(dyn_result_blend_mode, str)
+ else dyn_result_blend_mode
+ )
+ )
+ ortho_factor = (
+ (1.0 - ((b_normed - 1.0) / (rescale_limit - 1.0)).clamp_(0.0, 1.0))
+ .add_(dyn_min_scale)
+ .mul_(dyn_max_scale - dyn_min_scale)
+ )
+ if smooth_factor_kernel_size not in {0, 1}:
+ if ortho_factor.ndim < 3:
+ raise ValueError(
+ f"Can't use smooth_factor_kernel_size when ortho_factor has less than 3 dimensions. It has shape: {ortho_factor.shape}",
+ )
+ pad_sizes = tuple((sz - 1) // 2 for sz in smooth_factor_kernel_size)
+ target_dim = -3 if ortho_factor.ndim > 3 else -2
+ pool_fun = (
+ torch.nn.functional.avg_pool2d
+ if target_dim == -3
+ else torch.nn.functional.avg_pool1d
+ )
+ ortho_factor = (
+ pool_fun(
+ ortho_factor.movedim(-1, target_dim),
+ kernel_size=smooth_factor_kernel_size,
+ stride=1,
+ padding=pad_sizes,
+ )
+ .movedim(target_dim, -1)
+ .clamp_(dyn_min_scale, dyn_max_scale)
+ )
+ ortho_result = dyn_blend_function(vanilla_result, ortho_result, ortho_factor)
+ if ortho_verbose:
+ tqdm.write(
+ f"ORTHO BLEND: b_norm min/max={b_normed.aminmax()}, avg: {ortho_factor.mean().item():.5f}, min: {ortho_factor.min().item():.5f}, max: {ortho_factor.max().item():.5f}",
+ )
+ if end_dim != ndim - 1:
+ ortho_result = ortho_result.movedim(-1, start_dim)
+ return ortho_result.reshape(orig_shape)
+
+
+def symmetric_ortho_blend(
+ a: torch.Tensor,
+ b: torch.Tensor,
+ t: torch.Tensor,
+ *,
+ symmetric_strength: float = 1.0,
+ symmetric_deduce_mode: bool = False,
+ **kwargs: Any,
+) -> torch.Tensor:
+ blended = ortho_blend(a, b, t, **kwargs)
+ if symmetric_strength == 0.0:
+ return blended
+ b_ortho = blended.sub_(a)
+ if symmetric_deduce_mode:
+ b_proj = b - b_ortho
+ # Projection would theoretically be the same for both, in the simple case at least?
+ # Actually, probably not. Oh well, this is here as an option now.
+ a_ortho = a - b_proj
+ else:
+ a_ortho = ortho_blend(b, a, a.new_tensor(1.0), **kwargs) - b
+ a_proj = a - a_ortho
+ return a_proj.mul_(1.0 - symmetric_strength).add_(a_ortho).add_(b_ortho)
+
+
+def contrastive_ortho_cfg_base_a(
+ a: torch.Tensor,
+ b: torch.Tensor,
+ t: float | torch.Tensor = 1.0,
+ *,
+ a_blend: float = 1.0,
+ b_blend: float = 1.0,
+ b_blend_sub_a_scale: float = 1.0,
+ a_blend_sub_b_scale: float = 1.0,
+ a_ortho_scale: torch.Tensor | float = 1.0,
+ b_ortho_scale: torch.Tensor | float = 1.0,
+ a_blend_kwargs: dict | None = None,
+ b_blend_kwargs: dict | None = None,
+ # If a is cond and b is cond - uncond (CFG diff), you could
+ # set this to -1.0 to get b=uncond
+ b_from_a_scale: float = 0.0,
+ # One of: add, lerp, add_t, lerp_t
+ # When using _t modes, swaps t and b_from_a_scale.
+ b_from_a_mode: str = "add",
+ output_base_mode: str = "a",
+ # a, b, None or a tensor reference
+ final_rescale_target: str | torch.Tensor | None = None,
+ final_rescale_energy: float = 1.0,
+ # Should be left at 0.
+ final_rescale_ortho_blend: float = 0.0,
+ final_rescale_strength: float = 1.0,
+ final_rescale_kwargs: dict | None = None,
+ diff_only: bool = False,
+ **kwargs: Any,
+) -> torch.Tensor:
+ kwargs = pass_kwargs(kwargs)
+ t_is_tensor = isinstance(t, torch.Tensor)
+ if not t_is_tensor:
+ if t == 0.0:
+ return a
+ else:
+ t = t.broadcast_to(a.shape)
+
+ if b_from_a_scale != 0:
+ if b_from_a_mode.endswith("_t"):
+ b_from_a_mode = b_from_a_mode[:-2]
+ t, b_from_a_scale = b_from_a_scale, t
+ t_is_tensor = isinstance(t, torch.Tensor)
+ if isinstance(b_from_a_scale, torch.Tensor):
+ b_from_a_scale = b_from_a_scale.broadcast_to(a.shape)
+ if b_from_a_mode == "add":
+ b = (b * b_from_a_scale).add_(a)
+ elif b_from_a_mode == "lerp":
+ b = a.lerp(b, b_from_a_scale)
+ else:
+ raise ValueError("Bad b_from_a_mode")
+
+ # Extract features unique to b (b orthogonal to a)
+ b_blend_kwargs = kwargs if b_blend_kwargs is None else kwargs | b_blend_kwargs
+ b_ortho = ortho_blend(a, b, a.new_tensor(b_blend), **b_blend_kwargs).sub_(
+ a if b_blend_sub_a_scale == 1.0 else a * b_blend_sub_a_scale,
+ )
+ if isinstance(b_ortho_scale, torch.Tensor) or b_ortho_scale != 1.0:
+ b_ortho *= (
+ b_ortho_scale.broadcast_to(b_ortho)
+ if isinstance(b_ortho_scale, torch.Tensor)
+ else b_ortho_scale
+ )
+
+ # Extract features unique to a (a orthogonal to b)
+ a_blend_kwargs = kwargs if a_blend_kwargs is None else kwargs | a_blend_kwargs
+ a_ortho = ortho_blend(b, a, a.new_tensor(a_blend), **a_blend_kwargs).sub_(
+ b if a_blend_sub_b_scale == 1.0 else b * a_blend_sub_b_scale,
+ )
+ if isinstance(a_ortho_scale, torch.Tensor) or a_ortho_scale != 1.0:
+ a_ortho *= (
+ a_ortho_scale.broadcast_to(a_ortho)
+ if isinstance(a_ortho_scale, torch.Tensor)
+ else a_ortho_scale
+ )
+
+ # Create the contrastive guidance vector
+ # Push towards the a-unique features, pull away from the b-unique features
+ guidance = a_ortho.sub_(b_ortho)
+
+ if t_is_tensor or t != 1.0:
+ guidance *= t
+ output_base = a if output_base_mode == "a" else b
+ if (
+ not isinstance(final_rescale_target, torch.Tensor)
+ and final_rescale_target not in {"a", "b", "mid"}
+ ) or final_rescale_strength == 0.0:
+ return guidance if diff_only else guidance.add_(output_base)
+ result = guidance.add_(output_base)
+ final_rescale_kwargs = (
+ kwargs | final_rescale_kwargs
+ if final_rescale_kwargs is not None
+ else kwargs.copy()
+ )
+ final_rescale_kwargs["rescale_result_mode"] = "b"
+ if "rescale_limit" not in final_rescale_kwargs:
+ final_rescale_kwargs["rescale_limit"] = 2.0
+ if isinstance(final_rescale_target, torch.Tensor):
+ target_b = final_rescale_target.broadcast_to(result.shape)
+ elif final_rescale_target == "b":
+ target_b = b
+ elif final_rescale_target == "mid":
+ target_b = a.lerp(b, 0.5)
+ else:
+ target_b = a
+ if final_rescale_energy != 1.0:
+ target_b = target_b * final_rescale_energy
+ final_result = ortho_blend(
+ result,
+ target_b,
+ result.new_tensor(final_rescale_ortho_blend),
+ **final_rescale_kwargs,
+ )
+ if final_rescale_strength != 1.0:
+ final_result = (final_result - result).mul_(final_rescale_strength).add_(result)
+ return final_result - output_base if diff_only else final_result
+
+
+class WaveletBlend:
+ wavelet: wavef.Wavelet | None = None
+ use_float64: bool = False
+
+ def __init__(
+ self,
+ *,
+ device: str | torch.device | None = None,
+ use_float64: bool = False,
+ **kwargs: Any,
+ ):
+ self.device = device
+ self.wavelet_kwargs = kwargs
+ self.use_float64 = use_float64
+
+ def get_wavelet(self, *, device: str | torch.device | None = None) -> wavef.Wavelet:
+ if self.wavelet is None:
+ self.wavelet = wavef.Wavelet(
+ device=device if device is not None else self.device,
+ **self.wavelet_kwargs,
+ ).to(dtype=torch.float64 if self.use_float64 else torch.float32)
+ self.device = device
+ return self.wavelet
+ if device is not None and self.wavelet.device != device:
+ self.wavelet = self.wavelet.to(device=device)
+ self.device = device
+ return self.wavelet
+
+ @staticmethod
+ def maybe_offset(
+ yl: torch.Tensor,
+ yh: Sequence[torch.Tensor],
+ offset_yl: float | torch.Tensor | None,
+ offset_yh: float | Sequence[float | Sequence[float]] | None,
+ *,
+ in_place: bool = False,
+ ) -> tuple[torch.Tensor, tuple[torch.Tensor, ...]]:
+ if offset_yl in {None, 1.0} and offset_yh in {None, 1.0}:
+ return (yl, tuple(yh))
+ return wavef.wavelet_scaling(
+ yl,
+ yh,
+ yl_scale=offset_yl if offset_yl is not None else 1.0,
+ yh_scales=offset_yh,
+ in_place=in_place,
+ )
+
+ def wavelet_blend(
+ self,
+ a: torch.Tensor,
+ b: torch.Tensor,
+ t: float | torch.Tensor,
+ *,
+ blend_mode_yl: str | Callable = torch.lerp,
+ blend_mode_yh: str | Callable | None = None,
+ a_offset_yl: float | torch.Tensor | None = None,
+ a_offset_yh: float | Sequence[float | Sequence[float]] | None = None,
+ b_offset_yl: float | torch.Tensor | None = None,
+ b_offset_yh: float | Sequence[float | Sequence[float]] | None = None,
+ out_offset_yl: float | torch.Tensor | None = None,
+ out_offset_yh: float | Sequence[float | Sequence[float]] | None = None,
+ blend_yl_offset: float = 1.0,
+ blend_yh_offset: float | torch.Tensor = 1.0,
+ two_step_inverse: bool = False,
+ in_place_offset: bool = True,
+ ) -> torch.Tensor:
+ if isinstance(blend_mode_yl, str):
+ blend_mode_yl = BLENDING_MODES[blend_mode_yl]
+ if blend_mode_yh is None:
+ blend_mode_yh = blend_mode_yl
+ elif isinstance(blend_mode_yh, str):
+ blend_mode_yh = BLENDING_MODES[blend_mode_yh]
+ wavelet = self.get_wavelet(device=a.device)
+ dtype = a.dtype
+ if a.ndim != b.ndim:
+ raise ValueError(
+ f"Tensor a ndim ({a.ndim}) must match tensor b ndim ({b.ndim})"
+ )
+ orig_shape = a.shape
+ # FIXME: This reshaping logic is almost certainly not reliable.
+ if a.ndim > 4:
+ a = a.reshape(a.shape[0], -1, *a.shape[-2:])
+ if b.ndim > 4:
+ b = a.reshape(b.shape[0], -1, *b.shape[-2:])
+ a = a.to(dtype=torch.float64 if self.use_float64 else torch.float32)
+ b = b.to(a)
+ t = a.new_tensor(t) if not isinstance(t, torch.Tensor) else t.to(a)
+ if t.ndim > 4:
+ t = a.reshape(t.shape[0], -1, *t.shape[-2:])
+ aw_l, aw_h = self.maybe_offset(
+ *wavelet.forward(a),
+ a_offset_yl,
+ a_offset_yh,
+ in_place=in_place_offset,
+ )
+ bw_l, bw_h = self.maybe_offset(
+ *wavelet.forward(b),
+ b_offset_yl,
+ b_offset_yh,
+ in_place=in_place_offset,
+ )
+ blend_yl_offset = t if blend_yl_offset == 1 else t * blend_yl_offset
+ blend_yh_offset = t if blend_yh_offset == 1 else t * blend_yh_offset
+ outw_l, outw_h = self.maybe_offset(
+ *wavef.wavelet_blend(
+ (aw_l, aw_h),
+ (bw_l, bw_h),
+ yl_factor=blend_yl_offset,
+ yh_factor=blend_yh_offset,
+ blend_function=blend_mode_yl,
+ yh_blend_function=blend_mode_yh,
+ ),
+ offset_yl=out_offset_yl,
+ offset_yh=out_offset_yh,
+ in_place=in_place_offset,
+ )
+ result = wavelet.inverse(outw_l, outw_h, two_step_inverse=two_step_inverse)
+ result = result[tuple(slice(None, dsize) for dsize in a.shape)]
+ return result.to(dtype=dtype).reshape(orig_shape)
+
+
+WAVELET_BLEND_CACHE: dict[frozenset[tuple[str, Any]], WaveletBlend] = {}
+
+
+def wavelet_blend(
+ a: torch.Tensor,
+ b: torch.Tensor,
+ t: float | torch.Tensor,
+ *,
+ blend_mode_yl: str | Callable = torch.lerp,
+ blend_mode_yh: str | Callable | None = None,
+ **kwargs: Any,
+) -> torch.Tensor:
+ if isinstance(blend_mode_yl, str):
+ blend_mode_yl = BLENDING_MODES[blend_mode_yl]
+ if blend_mode_yh is None:
+ blend_mode_yh = blend_mode_yl
+ _ = kwargs.pop("device", None)
+ wavelet_kwargs = {
+ k: kwargs.pop(k)
+ for k in (
+ "wave",
+ "level",
+ "mode",
+ "use_1d_dwt",
+ "use_dtcwt",
+ "biort",
+ "qshift",
+ "inv_wave",
+ "inv_mode",
+ "inv_biort",
+ "inv_qshift",
+ "two_step_inverse",
+ "use_float64",
+ )
+ if k in kwargs
+ }
+ cache_key = frozenset(
+ (
+ wavelet_kwargs
+ | {"blend_mode_yl": blend_mode_yl, "blend_mode_yh": blend_mode_yh}
+ ).items(),
+ )
+ print(f"\nWAVELET BLEND: cache key: {cache_key}")
+ wb = WAVELET_BLEND_CACHE.get(cache_key)
+ if wb is None:
+ wb = WaveletBlend(device=a.device, **wavelet_kwargs)
+ WAVELET_BLEND_CACHE[cache_key] = wb
+ return wb.wavelet_blend(
+ a,
+ b,
+ t,
+ blend_mode_yl=blend_mode_yl,
+ blend_mode_yh=blend_mode_yh,
+ **kwargs,
+ )
+
+
+class TieredBlendWrapper:
+ def __init__(
+ self,
+ blend_function: Callable,
+ *,
+ tiers: int = 16,
+ start_dim: int = 1,
+ end_dim: int = -1,
+ descending: bool = True,
+ abs_mode: bool = False,
+ # a, b, add, sub, lerp (50% LERP)
+ sort_target: str = "a",
+ pad_value: float = 0.0,
+ ):
+ """Wraps any blending function to operate on 'Probability Tiers'.
+
+ :param blend_function: Callable with signature (a, b, blend_ratio)
+ :param tiers: Number of fake 'channels' or tiers to divide the data into.
+ :param start_dim: The first dimension to flatten.
+ :param end_dim: The last dimension to flatten.
+ :param descending: Sort highest-to-lowest (True) or lowest-to-highest (False).
+ :param pad_value: Value to pad with if the flattened size isn't divisible by tiers.
+ (For logits, -math.inf might be better, but 0.0 is safe for latents).
+ """
+ self.blend_function = blend_function
+ self.tiers = tiers
+ self.start_dim = start_dim
+ self.end_dim = end_dim
+ self.descending = descending
+ self.abs_mode = abs_mode
+ self.sort_target = sort_target
+ self.pad_value = pad_value
+
+ def get_target(self, a_flat: torch.Tensor, b_flat: torch.Tensor) -> torch.Tensor:
+ starget = self.sort_target
+ if starget == "a":
+ return a_flat
+ if starget == "b":
+ return b_flat
+ if starget == "add":
+ return a_flat + b_flat
+ if starget == "sub":
+ return a_flat - b_flat
+ if starget == "lerp":
+ return a_flat.lerp(b_flat, 0.5)
+ raise ValueError("Invalid sort target")
+
+ def __call__(
+ self,
+ a: torch.Tensor,
+ b: torch.Tensor | float,
+ t: torch.Tensor | float,
+ **kwargs: Any,
+ ) -> torch.Tensor:
+ if self.tiers < 1:
+ return self.blend_function(a, b, t, **kwargs)
+
+ orig_shape = a.shape
+
+ start_dim, end_dim = (
+ d if d >= 0 else a.ndim + d for d in (self.start_dim, self.end_dim)
+ )
+ if any(d < 0 or d > a.ndim for d in (start_dim, end_dim)):
+ raise ValueError("Dimension out of range")
+
+ b = (
+ b.broadcast_to(orig_shape)
+ if isinstance(b, torch.Tensor)
+ else torch.full_like(a, fill_value=b)
+ )
+
+ a_flat = a.flatten(start_dim=start_dim, end_dim=end_dim)
+ b_flat = b.flatten(start_dim=start_dim, end_dim=end_dim)
+
+ t_is_tensor = isinstance(t, torch.Tensor) and t.numel() > 1
+ if t_is_tensor:
+ t_flat = t.broadcast_to(orig_shape).flatten(
+ start_dim=start_dim,
+ end_dim=end_dim,
+ )
+ else:
+ t_flat = t
+
+ a_flat = a_flat.transpose(start_dim, -1)
+ b_flat = b_flat.transpose(start_dim, -1)
+ if t_is_tensor:
+ t_flat = t_flat.transpose(start_dim, -1)
+
+ length = a_flat.shape[-1]
+
+ target = self.get_target(a_flat, b_flat)
+ if self.abs_mode:
+ target = target.abs()
+
+ target_vals, indices = torch.sort(target, dim=-1, descending=self.descending)
+ del target
+ if not self.abs_mode and self.sort_target == "a":
+ a_vals = target_vals
+ else:
+ del target_vals
+ a_vals = torch.gather(a_flat, dim=-1, index=indices)
+ b_vals = torch.gather(b_flat, dim=-1, index=indices)
+ if t_is_tensor:
+ t_vals = torch.gather(t_flat, dim=-1, index=indices)
+
+ # PAD (If length is not divisible by tiers)
+ pad_len = (self.tiers - (length % self.tiers)) % self.tiers
+ if pad_len > 0:
+ a_vals = nnf.pad(a_vals, (0, pad_len), value=self.pad_value)
+ b_vals = nnf.pad(b_vals, (0, pad_len), value=self.pad_value)
+ if t_is_tensor:
+ t_vals = nnf.pad(t_vals, (0, pad_len), value=self.pad_value)
+
+ # Reshape into tiers (e.g., [..., length] -> [..., tiers, features])
+ new_shape = (*a_vals.shape[:-1], self.tiers, -1)
+ a_tiered = a_vals.reshape(new_shape)
+ b_tiered = b_vals.reshape(new_shape)
+ effective_t = t_vals.reshape(new_shape) if t_is_tensor else t
+
+ blended_tiered = self.blend_function(a_tiered, b_tiered, effective_t, **kwargs)
+
+ blended_flat = blended_tiered.reshape(a_vals.shape)
+
+ if pad_len > 0:
+ blended_flat = blended_flat[..., :-pad_len]
+
+ # Scatter back to original element positions
+ result_flat = torch.empty_like(blended_flat)
+ result_flat.scatter_(dim=-1, index=indices, src=blended_flat)
+
+ return result_flat.transpose(start_dim, -1).reshape(orig_shape)
+
+
+def tiered_blend(
+ a: torch.Tensor,
+ b: torch.Tensor,
+ t: torch.Tensor | float,
+ *,
+ tiers_blend_mode: str | Callable = "lerp",
+ tiers_blend_kwargs: dict | None = None,
+ **kwargs: Any,
+) -> torch.Tensor:
+ if isinstance(tiers_blend_mode, str):
+ tiers_blend_mode = BLENDING_MODES[tiers_blend_mode]
+ tw_kwargs_map = {
+ "tiers": "tiers",
+ "tiers_start_dim": "start_dim",
+ "tiers_end_dim": "end_dim",
+ "tiers_descending": "descending",
+ "tiers_abs_mode": "abs_mode",
+ "tiers_sort_target": "sort_target",
+ "tiers_pad_value": "pad_value",
+ }
+ tw_kwargs = {tk: kwargs.pop(k) for k, tk in tw_kwargs_map.items() if k in kwargs}
+ wrapped_blend_function = TieredBlendWrapper(tiers_blend_mode, **tw_kwargs)
+ if tiers_blend_kwargs is not None:
+ kwargs = kwargs | tiers_blend_kwargs
+ return wrapped_blend_function(a, b, t, **pass_kwargs(kwargs))
+
+
+# Shortest path circular interpolation (with the default params)
+def sp_circular_interpolation(
+ a: torch.Tensor,
+ b: torch.Tensor,
+ t: torch.Tensor | float,
+ *,
+ period: float | None = 2.0 * torch.pi,
+ period_scale: float = 1.0,
+ # Optional: lock the final output inside the bounds
+ wrap_result: bool = False,
+ # Only needed if wrap_result is True
+ lower_bound: float | None = None,
+ lower_bound_scale: float = 1.0,
+ start_dim: int = 1,
+ end_dim: int | None = None,
+ elementwise: bool = False,
+ minimize_range: bool = False,
+ diff_preserve_sign: bool = False,
+ result_preserve_sign: bool = False,
+ eps: float = 1e-08,
+) -> torch.Tensor:
+ if period is None:
+ if start_dim < 1:
+ start_dim = a.ndim + start_dim
+ end_dim = (
+ a.ndim
+ if end_dim is None
+ else (a.ndim + end_dim if end_dim < 0 else end_dim)
+ ) + 1
+ dims = tuple(range(start_dim, end_dim))
+ period = (torch.minimum if minimize_range else torch.maximum)(
+ a.abs()
+ if elementwise
+ else a.abs().amax(
+ dim=dims,
+ keepdim=True,
+ ),
+ b.abs()
+ if elementwise
+ else b.abs().amax(
+ dim=dims,
+ keepdim=True,
+ ),
+ )
+ period = period.mul_(2.0).clamp_min_(eps)
+ else:
+ period = max(eps, abs(period))
+ if period_scale != 1.0:
+ period = period * period_scale
+ if wrap_result and lower_bound is None:
+ lower_bound = period * (-0.5 * lower_bound_scale)
+ diff_orig = diff = b - a
+ half_period = period * 0.5
+
+ # Wrap the difference to the shortest path around the "circle"
+ diff = diff + half_period
+ diff %= period
+ diff -= half_period
+ if diff_preserve_sign:
+ diff = diff.copysign_(diff_orig)
+
+ diff *= t.broadcast_to(a.shape) if isinstance(t, torch.Tensor) else t
+ result = diff.add_(a)
+
+ if wrap_result:
+ result_orig = result
+ result = result - lower_bound
+ result %= period
+ result += lower_bound
+ if result_preserve_sign:
+ result = result.copysign_(result_orig)
+ return result
+
+
+# Computes the matrix logarithm using complex eigendecomposition.
+# Safe for batched Rotation matrices (Vh) and Covariance matrices.
+# Matrice must be diagonalizable.
+def matrix_log(
+ m: torch.Tensor,
+ *,
+ eps: float = 1e-06,
+ ieps: complex = 1e-08j,
+ dtype: torch.dtype | None = torch.complex128,
+ keep_dtype: bool = True,
+) -> torch.Tensor:
+ if m.ndim not in {2, 3} or m.shape[-2] != m.shape[-1]:
+ raise ValueError("matrix_log only supports diagonalizable square matrices")
+ # 1. Add a tiny diagonal epsilon to prevent singular matrix errors / log(0)
+ eye = torch.eye(
+ m.shape[-1],
+ device=m.device,
+ dtype=m.dtype if dtype is None else dtype,
+ ).mul_(eps)
+ m_safe = m.to(dtype=eye.dtype) + eye
+
+ # 2. Eigendecomposition
+ # L = Eigenvalues, V = Eigenvectors
+ el, ev = torch.linalg.eig(m_safe)
+
+ # 3. Take the natural logarithm of the complex eigenvalues
+ # With a tiny complex epsilon to prevent log(0+0j) NaNs
+ log_l = el.add_(ieps).log_()
+
+ # 4. Reconstruct the matrix: V @ diag(log_L) @ V^-1
+ v_inv = torch.linalg.solve(
+ ev,
+ torch.eye(ev.shape[-1], device=ev.device, dtype=ev.dtype),
+ )
+ log_m = (ev * log_l.unsqueeze(-2)) @ v_inv
+
+ # 5. The result should be purely real (the imaginary parts cancel out to ~0)
+ return log_m.real if keep_dtype else log_m.real.to(dtype=m.dtype)
+
+
+# Suitable for blending coordinate spaces like the Vh component of SVD, covariance, etc.
+# Something like 0.5 is similar to LERP.
+# Values above 1 should be similar to CFG, but for for coordinate spaces.
+def geodesic_square_matrix(
+ a: torch.Tensor,
+ b: torch.Tensor,
+ t: torch.Tensor | float,
+ *,
+ use_pinv: bool = True,
+ # Allow non square matrices (but only operate on the square subset)
+ lax: bool = True,
+) -> torch.Tensor:
+ if a.shape != b.shape:
+ errstr = f"Input shape mismatch, A {a.shape} != B {b.shape}"
+ raise ValueError(errstr)
+ if a.ndim == 2:
+ x, y = a.shape
+ elif a.ndim == 3:
+ x, y = a.shape[1:]
+ else:
+ x = y = None
+ if not (x and y) or (not lax and x != y):
+ raise ValueError(
+ "geodesic_square_matrix only supports square matrices (batch dimension optional)",
+ )
+ a_orig = a
+ if x != y:
+ minsz = min(x, y)
+ a = a[..., :minsz, :minsz]
+ b = b[..., :minsz, :minsz]
+ inv_op = torch.linalg.pinv if use_pinv else torch.linalg.inv
+ mlg = matrix_log(b @ inv_op(a))
+ velocity = torch.linalg.matrix_exp(
+ mlg * (t.to(dtype=mlg.dtype) if isinstance(t, torch.Tensor) else t),
+ )
+ result = (velocity @ a.to(velocity.dtype)).to(dtype=a.dtype)
+ if x == y:
+ return result
+ a_orig = a_orig.clone()
+ a_orig[..., :minsz, :minsz] = result
+ return a_orig
+
+
+def fft_blend(
+ a: torch.Tensor,
+ b: torch.Tensor,
+ t: torch.Tensor | float,
+ *,
+ fft_dims: Sequence[int] = (-2, -1),
+ avg_t_frequency_dims: bool = True,
+ a_phase_offset: float = 0.0,
+ a_magnitude_multiplier: float = 1.0,
+ a_magnitude_power: float = 0.0,
+ b_phase_offset: float = 0.0,
+ b_magnitude_multiplier: float = 1.0,
+ b_magnitude_power: float = 0.0,
+ blended_phase_offset: float = 0.0,
+ blended_magnitude_multiplier: float = 1.0,
+ blended_magnitude_power: float = 0.0,
+ phase_blend_mode: str | Callable = sp_circular_interpolation,
+ phase_blend_kwargs: dict | None = None,
+ phase_blend_multiplier: float = 1.0,
+ phase_blend_offset: float = 0.0,
+ magnitude_blend_mode: str | Callable = torch.lerp,
+ magnitude_blend_kwargs: dict | None = None,
+ magnitude_blend_multiplier: float = 1.0,
+ magnitude_blend_offset: float = 0.0,
+ magnitude_eps: float = 1e-08,
+ **kwargs: Any,
+) -> torch.Tensor:
+ fft_dims = tuple(fft_dims)
+ if a.ndim < 3:
+ raise ValueError("fft blend can only handle tensors with 3+ dimensions.")
+ kwargs = pass_kwargs(kwargs)
+ p_blend = (
+ BLENDING_MODES[phase_blend_mode]
+ if isinstance(phase_blend_mode, str)
+ else phase_blend_mode
+ )
+ m_blend = (
+ BLENDING_MODES[magnitude_blend_mode]
+ if isinstance(magnitude_blend_mode, str)
+ else magnitude_blend_mode
+ )
+ a_f = torch.fft.rfftn(a, dim=fft_dims)
+ a_phase = torch.atan2(a_f.imag, a_f.real)
+ a_mag = a_f.abs()
+ if a_phase_offset != 0.0:
+ a_phase += a_phase_offset
+ if a_magnitude_multiplier != 1.0:
+ a_mag *= a_magnitude_multiplier
+ if a_magnitude_power != 0.0:
+ a_mag = a_mag.add_(magnitude_eps).pow_(a_magnitude_power)
+ b_f = torch.fft.rfftn(b, dim=fft_dims)
+ b_phase = torch.atan2(b_f.imag, b_f.real)
+ b_mag = b_f.abs()
+ if b_phase_offset != 0.0:
+ b_phase += b_phase_offset
+ if b_magnitude_multiplier != 1.0:
+ b_mag *= b_magnitude_multiplier
+ if b_magnitude_power != 0.0:
+ b_mag = b_mag.add_(magnitude_eps).pow_(b_magnitude_power)
+ if not isinstance(t, torch.Tensor):
+ t = a.new_tensor(t)
+ elif t.ndim > 1:
+ # FIXME: This probably doesn't work.
+ t = t.broadcast_to(a_f.shape)
+ if avg_t_frequency_dims:
+ t = t.mean(dim=fft_dims, keepdim=True)
+ magnitude_blend_kwargs = kwargs | (
+ magnitude_blend_kwargs if magnitude_blend_kwargs is not None else {}
+ )
+ phase_blend_kwargs = kwargs | (
+ phase_blend_kwargs if phase_blend_kwargs is not None else {}
+ )
+ t_mag = t if magnitude_blend_multiplier == 1.0 else t * magnitude_blend_multiplier
+ if magnitude_blend_offset != 0.0:
+ t_mag = magnitude_blend_offset + t_mag
+ t_phase = t if phase_blend_multiplier == 1.0 else t * phase_blend_multiplier
+ if phase_blend_offset != 0.0:
+ t_phase = phase_blend_offset + t_phase
+ blended_mag = m_blend(a_mag, b_mag, t_mag, **magnitude_blend_kwargs).abs()
+ blended_phase = p_blend(a_phase, b_phase, t_phase, **phase_blend_kwargs)
+ if blended_phase_offset != 0.0:
+ blended_phase += blended_phase_offset
+ if blended_magnitude_multiplier != 1.0:
+ blended_mag *= blended_magnitude_multiplier
+ if blended_magnitude_power != 0.0:
+ blended_mag = blended_mag.add_(magnitude_eps).pow_(blended_magnitude_power)
+ return torch.fft.irfftn(
+ torch.polar(blended_mag, blended_phase),
+ s=tuple(a.shape[d] for d in fft_dims),
+ dim=fft_dims,
+ )
+
+
+class DecompBlend:
+ @staticmethod
+ def decomp(
+ t: torch.Tensor,
+ *,
+ mode: str = "svd",
+ **kwargs: Any,
+ ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+ tqdm.write(f"DECOMP: shape1={t.shape}")
+ if t.ndim > 3:
+ raise ValueError(
+ "Can only handle batched or unbatched matrices (2 or 3 dimensions)",
+ )
+ tqdm.write(f"DECOMP: shape={t.shape}")
+ if mode == "svd":
+ return torch.linalg.svd(t, full_matrices=False)
+ if mode == "svd_lowrank":
+ q = kwargs.pop("q", None)
+ u, s, v = torch.svd_lowrank(
+ t,
+ q=q if q is not None else t.shape[-1],
+ **kwargs,
+ )
+ return u, s, v.mT
+ if mode == "qr":
+ u, r_mat = torch.linalg.qr(t)
+ s = r_mat.diagonal(dim1=-2, dim2=-1)
+ vh = (1.0 / s).masked_fill_(s == 0.0, 1.0 / 1e-08).unsqueeze(-1) * r_mat
+ return u, s, vh
+ raise ValueError("Bad decomp mode")
+
+ @staticmethod
+ def align(
+ left: tuple[torch.Tensor, torch.Tensor, torch.Tensor],
+ right: tuple[torch.Tensor, torch.Tensor, torch.Tensor],
+ *,
+ align_mode: str = "joint",
+ invert: bool = False,
+ eps: float = 1e-06,
+ ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+ align_mode = align_mode.strip().lower()
+ if align_mode not in {"u", "v", "vh", "joint"}:
+ raise ValueError("Align mode must be one of u, v[h], or joint")
+ align_mode = align_mode[0]
+
+ ul, sl, vl = left
+ ur, sr, vr = right
+
+ def get_sim(t_right: torch.Tensor, t_left: torch.Tensor) -> torch.Tensor:
+ norm_right, norm_left = (
+ torch.linalg.vector_norm(t, dim=-2, keepdim=True).clamp_min_(eps)
+ for t in (t_right, t_left)
+ )
+ return (t_right / norm_right).mT @ (t_left / norm_left)
+
+ sim_u = get_sim(ur, ul) if align_mode in "uj" else None
+ sim_vh = get_sim(vr.mT, vl.mT) if align_mode in "vj" else None
+ sim = (
+ sim_u * sim_vh
+ if align_mode == "j"
+ else (sim_u if sim_u is not None else sim_vh)
+ )
+
+ match_idx = (sim.abs().argmin if invert else sim.abs().argmax)(dim=-1)
+ signs = (
+ (sim if align_mode != "j" else sim_u)
+ .gather(
+ dim=-1,
+ index=match_idx.unsqueeze(-1),
+ )
+ .sign_()
+ )
+ return (
+ ul.gather(
+ dim=-1,
+ index=match_idx.unsqueeze(-2).expand(*ul.shape[:-1], sr.shape[-1]),
+ )
+ * signs.mT,
+ sl.gather(dim=-1, index=match_idx),
+ vl.gather(
+ dim=-2,
+ index=match_idx.unsqueeze(-1).expand(
+ *vl.shape[:-2],
+ sr.shape[-1],
+ vl.shape[-1],
+ ),
+ )
+ * signs,
+ )
+
+ @staticmethod
+ def get_size_with_offset(
+ *,
+ rank: int,
+ size: float,
+ offset: float = -1,
+ ) -> tuple[int, int]:
+ size = int(size) if abs(size) >= 1.0 else math.ceil(size * rank)
+ if size < 0:
+ size = rank + size
+ elif size == 0:
+ size = rank
+ offset = int(offset) if abs(offset) >= 1.0 else math.ceil(offset * rank)
+ if offset < 0:
+ offset = rank + offset
+ size = max(0, min(rank, size))
+ offset = min(rank - size, max(0, offset))
+ return (size, offset)
+
+ @classmethod
+ def rank_slice_blend(
+ cls,
+ a: torch.Tensor,
+ b: torch.Tensor,
+ # Negative values count from the end. Values > -1.0, < 1.0 are interpreted
+ # as percentage of ranks. Values outside of that range are truncated and treated
+ # as absolute rank indexes. Offsets are specified the same way.
+ t: float | torch.Tensor,
+ *,
+ feature_dim: int = 1,
+ # You will always get the exact size you specify and the offset will be adjusted
+ # if there isn't enough space.
+ rank_offset: float = -1.0,
+ decomp_mode: str = "svd",
+ align_mode: str = "joint",
+ align_invert: bool = False,
+ blend_components: str = "usv",
+ n_iter: int = 6,
+ q: int | None = None,
+ ) -> torch.Tensor:
+ if a.ndim < 2 or b.ndim < 2:
+ raise ValueError("Can only only handle 2+ dimensional tensors")
+ t = t.mean().detach().cpu().item() if isinstance(t, torch.Tensor) else float(t)
+
+ tqdm.write(f"ORIG SHAPE: {a.shape}")
+ a = a.movedim(feature_dim, -1)
+ flat_start_dim = 1 if a.ndim > 2 else 0
+ adj_shape = a.shape
+ a = a.flatten(start_dim=flat_start_dim, end_dim=-2)
+ b = b.movedim(feature_dim, -1).flatten(start_dim=flat_start_dim, end_dim=-2)
+
+ dl = cls.decomp(a, mode=decomp_mode, niter=n_iter, q=q)
+ dr = cls.decomp(b, mode=decomp_mode, niter=n_iter, q=q)
+
+ size, offset = cls.get_size_with_offset(
+ rank=dl[1].shape[-1],
+ size=t,
+ offset=rank_offset,
+ )
+ rs = slice(offset, offset + size)
+
+ if (align_mode := align_mode.strip().lower()) in {"joint", "u", "v", "vh"}:
+ dl = cls.align(
+ dl,
+ dr,
+ align_mode=align_mode,
+ invert=align_invert,
+ )
+ ((ua, sa, vha), (ub, sb, vhb)) = dl, dr
+ blend_components = blend_components.strip().lower()
+ if "u" in blend_components:
+ ua[..., rs] = ub[..., rs]
+ if "s" in blend_components:
+ sa[..., rs] = sb[..., rs]
+ if "v" in blend_components:
+ vha[..., rs, :] = vhb[..., rs, :]
+ result = ua @ sa.diag_embed() @ vha
+ return result.reshape(adj_shape).movedim(-1, feature_dim).contiguous()
+
+ @staticmethod
+ def normalizing_in(
+ t: torch.Tensor,
+ *,
+ centering_strength: float,
+ centering_restore_strength: float,
+ variance_normalizing: bool,
+ aug_scale: float,
+ dim: int | Sequence[int],
+ orig_features: int,
+ in_place: bool = True,
+ trim_features: bool = True,
+ eps: float = 1e-08,
+ ) -> tuple[torch.Tensor, torch.Tensor | None, torch.Tensor | None]:
+ if not in_place:
+ t = t.clone()
+ if not isinstance(dim, int):
+ dim = tuple(dim)
+ mean = (
+ t.mean(dim=dim, keepdim=True)
+ if centering_strength != 0.0 or centering_restore_strength != 0.0
+ else None
+ )
+ if centering_strength != 0 and mean is not None:
+ t -= mean * centering_strength if centering_strength != 1 else mean
+ if mean is not None and trim_features:
+ mean = mean[..., :orig_features]
+ if variance_normalizing:
+ std = t.std(dim=dim, keepdim=True).clamp_min_(eps)
+ t /= std
+ if trim_features:
+ std = std[..., :orig_features]
+ else:
+ std = None
+ if aug_scale != 0 and t.shape[-1] > orig_features:
+ t[..., orig_features:] *= aug_scale
+ return t, mean, std
+
+ @staticmethod
+ def normalizing_out(
+ t: torch.Tensor,
+ *,
+ mean: torch.Tensor | None = None,
+ std: torch.Tensor | None = None,
+ centering_restore_strength: float,
+ in_place: bool = True,
+ ) -> torch.Tensor:
+ if mean is None and std is None:
+ return t
+ if not in_place:
+ t = t.clone()
+ t_slices = tuple(slice(None, sz) for sz in t.shape)
+ mean, std = (None if temp is None else temp[t_slices] for temp in (mean, std))
+ if std is not None:
+ t *= std
+ if centering_restore_strength != 0 and mean is not None:
+ t += (
+ mean.mul_(centering_restore_strength)
+ if centering_restore_strength != 1
+ else mean
+ )
+ return t
+
+ @classmethod
+ def rank_blend(
+ cls,
+ a: torch.Tensor,
+ b: torch.Tensor,
+ t: float | torch.Tensor,
+ *,
+ # Controls whether changes get applied to a or b.
+ base_a: bool = True,
+ blend_mode: Callable | str = torch.lerp,
+ # If None, uses addition (inject) in diff mode or LERP in blend mode.
+ result_blend_mode: Callable | str | None = None,
+ # One of diff, slice, blend
+ blend_strategy: str = "diff",
+ # Negative values count from the end. Values > -1.0, < 1.0 are interpreted
+ # as percentage of ranks. Values outside of that range are truncated and treated
+ # as absolute rank indexes. Offsets are specified the same way.
+ # Since the ranks parameter is a size, 0 means all ranks, -2 to means total_ranks - 2, etc.
+ ranks: float = 0.5,
+ # You will always get the exact size you specify and the offset will be adjusted
+ # if there isn't enough space.
+ rank_offset: float = -1.0,
+ rank_start_scale: float = 1.0,
+ rank_end_scale: float = 1.0,
+ rank_ramp_power: float = 0.0,
+ use_log_rank_scales: bool = False,
+ # rmula, rmulb, amulb, rroll, rmulroll, rmulrollc
+ feature_augmentations: Sequence[str] = ("rmula", "rmulb", "amulb"),
+ feature_augmentation_scale: float = 0.0,
+ feature_dim: int = 1,
+ # Flattening occurs after the feature dim is moved to the end.
+ flatten_start_dim: int = 1,
+ flatten_end_dim: int = -2,
+ decomp_mode: str = "svd",
+ centering_strength: float = 0.0,
+ centering_restore_strength: float = 0.0,
+ result_scale: float = 1.0,
+ variance_normalizing: bool = False,
+ # Alignment only applies to the slice blend strategy.
+ align_mode: str = "joint",
+ align_invert: bool = False,
+ align_base: bool = True,
+ decomp_kwargs: dict[str, Any] | None = None,
+ **kwargs: Any,
+ ) -> torch.Tensor:
+ if blend_strategy not in {"diff", "blend", "slice"}:
+ raise ValueError("Bad blend_strategy")
+ if a.ndim < 2 or b.ndim < 2:
+ raise ValueError("Can only only handle 2+ dimensional tensors")
+ if isinstance(blend_mode, str):
+ blend_mode = BLENDING_MODES[blend_mode]
+ if isinstance(result_blend_mode, str):
+ result_blend_mode = BLENDING_MODES[result_blend_mode]
+ elif result_blend_mode is None and blend_strategy == "blend":
+ result_blend_mode = torch.lerp
+
+ kwargs = pass_kwargs(kwargs)
+ feature_augmentations = tuple(feature_augmentations)
+
+ use_aug = feature_augmentations and feature_augmentation_scale != 0
+ base = a if base_a else b
+ blend_result: torch.Tensor = blend_mode(a, b, t, **kwargs)
+ if blend_strategy == "diff":
+ blend_result -= base
+ blend_result = blend_result.movedim(feature_dim, -1)
+ adj_shape = blend_result.shape
+ orig_features = adj_shape[-1]
+ blend_result = blend_result.flatten(
+ start_dim=flatten_start_dim,
+ end_dim=flatten_end_dim,
+ )
+ if use_aug:
+ aug_list = [blend_result]
+ flat_a = a.movedim(feature_dim, -1).flatten(
+ start_dim=flatten_start_dim,
+ end_dim=flatten_end_dim,
+ )
+ flat_b = b.movedim(feature_dim, -1).flatten(
+ start_dim=flatten_start_dim,
+ end_dim=flatten_end_dim,
+ )
+ aug = None
+ for augtype in feature_augmentations:
+ if augtype == "rmula":
+ aug = blend_result * flat_a
+ elif augtype == "rmulb":
+ aug = blend_result * flat_b
+ elif augtype == "amulb":
+ aug = flat_a * flat_b
+ elif augtype == "rroll":
+ aug = blend_result.roll(shifts=1, dims=-2)
+ elif augtype == "rmulroll":
+ aug = blend_result.roll(shifts=1, dims=-2).mul_(blend_result)
+ elif augtype == "rmulrollc":
+ aug = blend_result.roll(shifts=1, dims=-1).mul_(blend_result)
+ else:
+ errstr = f"Unknown augmentation type: {augtype}"
+ raise ValueError(errstr)
+ aug_list.append(aug)
+ if len(aug_list) > 1:
+ blend_result = torch.cat(aug_list, dim=-1)
+ del aug, aug_list, flat_a, flat_b
+ norm_in = partial(
+ cls.normalizing_in,
+ centering_strength=centering_strength,
+ centering_restore_strength=centering_restore_strength,
+ variance_normalizing=variance_normalizing,
+ aug_scale=feature_augmentation_scale,
+ dim=flatten_start_dim,
+ orig_features=orig_features,
+ )
+ if blend_strategy == "slice":
+ flat_base = blend_result.clone()
+ flat_base[..., :orig_features] = base.movedim(feature_dim, -1).flatten(
+ start_dim=flatten_start_dim,
+ end_dim=flatten_end_dim,
+ )
+ flat_base, base_mean, base_std = norm_in(flat_base)
+ else:
+ flat_base = base_mean = base_std = None
+ blend_result, mean, std = norm_in(blend_result)
+ dr = cls.decomp(blend_result, mode=decomp_mode, **(decomp_kwargs or {}))
+ if blend_strategy == "slice" and flat_base is not None:
+ db = cls.decomp(flat_base, mode=decomp_mode, **(decomp_kwargs or {}))
+ if (align_mode := align_mode.strip().lower()) in {"joint", "u", "v", "vh"}:
+ temp = cls.align(
+ db if align_base else dr,
+ dr if align_base else db,
+ align_mode=align_mode,
+ invert=align_invert,
+ )
+ db, dr = (temp, dr) if align_base else (db, temp)
+ del temp
+
+ u, s, vh = dr
+ size, offset = cls.get_size_with_offset(
+ rank=s.shape[-1],
+ size=ranks,
+ offset=rank_offset,
+ )
+ if rank_start_scale != rank_end_scale:
+ if rank_offset < 0:
+ rank_start_scale, rank_end_scale = rank_end_scale, rank_start_scale
+ rank_scales = torch.linspace(
+ rank_start_scale,
+ rank_end_scale,
+ steps=size,
+ dtype=a.dtype,
+ device=a.device,
+ ).unsqueeze(0)
+ if rank_ramp_power != 0:
+ rank_scales = torch.where(
+ rank_scales == 0,
+ 0,
+ rank_scales.abs().pow_(rank_ramp_power).copysign_(rank_scales),
+ )
+ else:
+ rank_scales = rank_start_scale
+ rs = slice(offset, offset + size)
+ tqdm.write(
+ f"RANK BLEND: slice={rs}, scales={rank_scales}, shape={blend_result.shape}, adj={adj_shape}",
+ )
+ u, s, vh = u[..., rs], s[..., rs], vh[..., rs, :]
+ if not isinstance(rank_scales, float) or rank_scales != 1.0:
+ if use_log_rank_scales:
+ s = s.abs().log1p_().copysign_(s)
+ s *= rank_scales
+ if use_log_rank_scales:
+ s = s.abs().expm1_().copysign_(s)
+ result = (u @ s.diag_embed() @ vh)[..., :orig_features]
+ if result_scale != 1.0 and result_blend_mode is None:
+ if std is None:
+ result *= result_scale
+ else:
+ std *= result_scale
+ result = cls.normalizing_out(
+ result,
+ mean=mean,
+ std=std,
+ centering_restore_strength=centering_restore_strength,
+ )
+ if blend_strategy == "slice" and db is not None:
+ bu, bs, bvh = db
+ size_before = rs.start
+ size_after = bs.shape[-1] - rs.stop
+ base_result = (
+ (
+ bu[..., :size_before]
+ @ bs[..., :size_before].diag_embed()
+ @ bvh[..., :size_before, :]
+ )
+ if size_before > 0
+ else None
+ )
+ if size_after > 0:
+ temp = (
+ bu[..., rs.stop :]
+ @ bs[..., rs.stop :].diag_embed()
+ @ bvh[..., rs.stop :, :]
+ )
+ base_result = (
+ base_result.add_(temp) if base_result is not None else temp
+ )
+ if base_result is not None:
+ base_result = cls.normalizing_out(
+ base_result[..., :orig_features],
+ mean=base_mean,
+ std=base_std,
+ centering_restore_strength=centering_restore_strength,
+ )
+ result += base_result
+
+ result = result.reshape(adj_shape).movedim(-1, feature_dim).contiguous()
+ if blend_strategy == "slice":
+ return result
+ if result_blend_mode is not None:
+ return result_blend_mode(base, result, result_scale, **kwargs)
+ if blend_strategy == "diff":
+ return result.add_(base)
+ raise RuntimeError("Unhandled blend_strategy")
+
+
+def chain_blend(
+ a: torch.Tensor,
+ b: torch.Tensor,
+ *args: Any,
+ chain_iterations: int = 1,
+ chain_blend_mode: str | Callable = torch.lerp,
+ **kwargs: Any,
+) -> torch.Tensor:
+ if chain_iterations < 1:
+ return a.clone()
+ fun = (
+ BLENDING_MODES[chain_blend_mode]
+ if isinstance(chain_blend_mode, str)
+ else chain_blend_mode
+ )
+ kwargs = pass_kwargs(kwargs)
+ for _ in range(chain_iterations):
+ b = fun(a, b, *args, **kwargs)
+ return b
+
+
+def pct_limit_blend(
+ a: torch.Tensor,
+ b: torch.Tensor,
+ *args: Any,
+ base_a: bool = True,
+ diff_limit: float = 0.25,
+ eps: float = 1e-07,
+ blend_mode: str | Callable = torch.lerp,
+ dim: int | Sequence[int] | None = None,
+ # Only applies in elementwise mode (dim=None)
+ prevent_sign_flip: bool = False,
+ # Negative values disable soft clamp for that specific constraint.
+ # Higher stiffness -> approach the limit more closely before values warp.
+ pct_clamp_stiffness: float = 10.0,
+ sign_clamp_stiffness: float = 10.0,
+ **kwargs: Any,
+) -> torch.Tensor:
+ if a.ndim < 2:
+ raise ValueError("Blend function requires 2+ dimensions")
+ if diff_limit < 0:
+ raise ValueError("diff_limit must be positive")
+
+ blend_function = (
+ BLENDING_MODES[blend_mode] if isinstance(blend_mode, str) else blend_mode
+ )
+
+ br = blend_function(a, b, *args, **pass_kwargs(kwargs))
+ base = a if base_a else b
+
+ if diff_limit == 0:
+ return base.clone()
+
+ if dim is not None:
+ diff = br.sub_(base)
+ diff_norm = torch.linalg.vector_norm(diff, dim=dim, keepdim=True)
+ max_norm = (
+ torch.linalg.vector_norm(base, dim=dim, keepdim=True)
+ .mul_(diff_limit)
+ .clamp_min_(eps)
+ )
+
+ if pct_clamp_stiffness >= 0:
+ # Soft clamp the magnitude of the difference
+ soft_diff_norm = soft_clamp(
+ diff_norm,
+ min_val=0.0,
+ max_val=max_norm,
+ stiffness=pct_clamp_stiffness,
+ )
+ scale = soft_diff_norm.div_(diff_norm.clamp_min_(eps))
+ else:
+ # Hard clamp the scale
+ scale = max_norm.div_(diff_norm.clamp_min_(eps)).clamp_max_(1.0)
+
+ return diff.mul_(scale).add_(base)
+
+ # Elementwise handling.
+ if prevent_sign_flip:
+ # Create bounds using infinity so we ONLY restrict the zero-crossing
+ mask = base >= 0
+ sign_lower = torch.where(mask, 0.0, -torch.inf)
+ sign_upper = torch.where(mask, torch.inf, 0.0)
+
+ if sign_clamp_stiffness >= 0:
+ br = soft_clamp(br, sign_lower, sign_upper, stiffness=sign_clamp_stiffness)
+ else:
+ br = br.clamp_(min=sign_lower, max=sign_upper)
+
+ max_diff = base.abs().mul_(diff_limit).clamp_min_(eps)
+ lower_bound = base - max_diff
+ upper_bound = base + max_diff
+
+ if pct_clamp_stiffness < 0:
+ return br.clamp_(min=lower_bound, max=upper_bound)
+ return soft_clamp(br, lower_bound, upper_bound, stiffness=pct_clamp_stiffness)
+
+
+def moment_aligned_blend(
+ a: torch.Tensor,
+ b: torch.Tensor,
+ *args: Any,
+ blend_mode: str | Callable = torch.lerp,
+ base_a: bool = True,
+ mean_scale: float = 1.0,
+ std_scale: float = 1.0,
+ dim: int | Sequence[int] | None = None,
+ eps: float = 1e-07,
+ **kwargs: Any,
+) -> torch.Tensor:
+ blend_function = (
+ BLENDING_MODES[blend_mode] if isinstance(blend_mode, str) else blend_mode
+ )
+ kwargs = pass_kwargs(kwargs)
+
+ if mean_scale == 0.0 and std_scale == 0.0:
+ return blend_function(a, b, *args, **kwargs)
+
+ if dim is None:
+ dim = tuple(range(1, a.ndim))
+
+ target, base = (b, a) if base_a else (a, b)
+ mean_target = target.mean(dim=dim, keepdim=True)
+
+ # Centering here is always necessary. We will add the mean back if mean_scale is 0.
+ aligned = target - mean_target
+ if std_scale != 0.0:
+ std_target = target.std(dim=dim, keepdim=True).clamp_min_(eps)
+ std_base = base.std(dim=dim, keepdim=True).clamp_min_(eps)
+ std_goal = std_base if std_scale == 1 else std_target.lerp(std_base, std_scale)
+ aligned *= std_goal.div_(std_target)
+ del std_target, std_base, std_goal
+ if mean_scale != 0.0:
+ mean_base = base.mean(dim=dim, keepdim=True)
+ mean_goal = (
+ mean_base if mean_scale == 1 else mean_target.lerp_(mean_base, mean_scale)
+ )
+ del mean_base
+ else:
+ mean_goal = mean_target
+ aligned += mean_goal
+ del mean_goal, mean_target, base, target
+
+ return blend_function(
+ a if base_a else aligned,
+ aligned if base_a else b,
+ *args,
+ **kwargs,
+ )
+
+
+def distro_aligned_blend(
+ a: torch.Tensor,
+ b: torch.Tensor,
+ *args: Any,
+ blend_mode: str | Callable = torch.lerp,
+ reference_blend_mode: str | Callable = torch.lerp,
+ # When using a reference and unset, will use the original ratio.
+ reference_blend: float | torch.Tensor | None = None,
+ reference: torch.Tensor | None = None,
+ # One of da, db, dr(eference), sa, sb, sr, dg(aussian), sg
+ # 's' vs 'd' determines whether distro or robust scale matching is used.
+ align_a: str | None = None,
+ align_b: str | None = None,
+ align_result: str | None = None,
+ start_dim: int = 1,
+ end_dim: int = -1,
+ scale_match_mad: float = 0.6745,
+ **kwargs: Any,
+) -> torch.Tensor:
+ if len(args) == 0:
+ raise ValueError("Missing ratio positional parameter")
+ kwargs = pass_kwargs(kwargs)
+ blend_function = (
+ BLENDING_MODES[blend_mode] if isinstance(blend_mode, str) else blend_mode
+ )
+ if align_a is None and align_b is None and align_result is None:
+ return blend_function(a, b, *args, **kwargs)
+ need_ref = reference is None and any(
+ val in {"dr", "sr"} for val in (align_a, align_b, align_result)
+ )
+ if need_ref:
+ ref_blend_function = (
+ BLENDING_MODES[reference_blend_mode]
+ if isinstance(reference_blend_mode, str)
+ else reference_blend_mode
+ )
+ ref_args = (
+ reference_blend if reference_blend is not None else args[0],
+ *args[1:],
+ )
+ reference = ref_blend_function(a, b, *ref_args, **kwargs)
+
+ align_targets = {
+ "da": a,
+ "db": b,
+ "dr": reference,
+ "sa": a,
+ "sb": b,
+ "sr": reference,
+ }
+ if any(
+ val not in {None, "dg", "sg"} and align_targets.get(val) is None
+ for val in (align_a, align_b, align_result)
+ ):
+ raise ValueError("Invalid align target")
+
+ def do_align(t: torch.Tensor, amode: str | None) -> torch.Tensor:
+ if amode is None:
+ return t
+ if amode == "dg":
+ return force_gaussian_distribution(t, start_dim=start_dim, end_dim=end_dim)
+ if amode == "sg":
+ return robust_scale_match(
+ t,
+ start_dim=start_dim,
+ end_dim=end_dim,
+ mad=scale_match_mad,
+ )
+ if amode.startswith("d"):
+ return match_distribution(
+ t,
+ reference=align_targets[amode],
+ start_dim=start_dim,
+ end_dim=end_dim,
+ )
+ return robust_scale_match(
+ t,
+ reference=align_targets[amode],
+ start_dim=start_dim,
+ end_dim=end_dim,
+ )
+
+ a, b = (do_align(item, amode) for item, amode in ((a, align_a), (b, align_b)))
+
+ return do_align(blend_function(a, b, *args, **kwargs), align_result)
+
+
+# Standard LERP, but the weights are forced to preserve a variance of 1.
+def pythagorean_lerp(
+ a: torch.Tensor,
+ b: torch.Tensor,
+ t: float | torch.Tensor,
+ *,
+ eps: float = 1e-08,
+) -> torch.Tensor:
+ w_a = 1.0 - t
+ w_b = t
+
+ # Calculate how much the variance would shrink.
+ if isinstance(w_a, torch.Tensor):
+ variance_shrink = (w_a**2).add_(w_b**2).sqrt_().clamp_min_(eps)
+ else:
+ variance_shrink = max(eps, (w_a**2 + w_b**2) ** 0.5)
+
+ # Then scale the weights to compensate.
+ return a.mul(w_a / variance_shrink).add_(b * (w_b / variance_shrink))
+
+
+def rms_interpolation(
+ a: torch.Tensor,
+ b: torch.Tensor,
+ t: float | torch.Tensor,
+ *,
+ # To only use magnitude, enable this and set power to 1.
+ abs_inputs: bool = False,
+ # One of:
+ # blend, reference (reference must be supplied), a, b, leave,
+ sign_mode: str = "blend",
+ reference: torch.Tensor | None = None,
+ blend_mode: str | Callable = torch.lerp,
+ sign_blend_mode: str | Callable | None = None,
+ power: float | torch.Tensor = 2.0,
+ inv_power: float | torch.Tensor | None = None,
+ **kwargs: Any,
+) -> torch.Tensor:
+ kwargs = pass_kwargs(kwargs)
+ blend_function = (
+ BLENDING_MODES[blend_mode] if isinstance(blend_mode, str) else blend_mode
+ )
+ a_orig, b_orig = a, b
+ if abs_inputs:
+ a, b = a.abs(), b.abs()
+ if power != 1.0:
+ if inv_power is None:
+ inv_power = 1 / power
+ a, b = a**power, b**power
+ blend_result = blend_function(a, b, t, **kwargs)
+ if power != 1.0:
+ if sign_mode == "leave":
+ return blend_result.abs().pow_(inv_power).copysign_(blend_result)
+ blend_result = blend_result.abs_().pow_(inv_power)
+ elif sign_mode == "leave":
+ return blend_result
+ if sign_mode == "blend":
+ sign_blend_function = (
+ (
+ BLENDING_MODES[sign_blend_mode]
+ if isinstance(sign_blend_mode, str)
+ else sign_blend_mode
+ )
+ if sign_blend_mode is not None
+ else blend_function
+ )
+ reference = sign_blend_function(a_orig, b_orig, t, **kwargs)
+ return blend_result.copysign_(reference)
+ if sign_mode == "reference":
+ if reference is None:
+ raise ValueError("sign mode reference requires a reference to be supplied")
+ return blend_result.copysign_(reference.to(blend_result))
+ if sign_mode == "a":
+ return blend_result.copysign_(a_orig)
+ if sign_mode == "b":
+ return blend_result.copysign_(b_orig)
+ errstr = f"Unhandled sign mode: {sign_mode}"
+ raise ValueError(errstr)
+
+
+def orbit_blend(
+ a: torch.Tensor,
+ b: torch.Tensor,
+ t: float | torch.Tensor,
+ *,
+ # One of leave, clamp, bsuba, rsuba, rsubb
+ excess_mode: str = "leave",
+ invert_excess: bool = False,
+) -> torch.Tensor:
+ t_orig = t
+ if excess_mode == "clamp":
+ t = t.clamp(-1, 1) if isinstance(t, torch.Tensor) else max(-1.0, min(1.0, t))
+ angle = t * (math.pi / 2.0)
+ if isinstance(angle, torch.Tensor):
+ orbit = a.mul(angle.cos()).add_(b.mul(angle.sin()))
+ else:
+ orbit = (a * math.cos(angle)).add_(b * math.sin(angle))
+
+ if excess_mode not in {"leave", "clamp"}:
+ have_excess = (
+ torch.any(t_orig.abs() > 1).detach().cpu().item()
+ if isinstance(t_orig, torch.Tensor)
+ else abs(t_orig) > 1
+ )
+ else:
+ have_excess = False
+
+ if not have_excess:
+ return orbit
+
+ if excess_mode == "bsuba":
+ tangent_vector = b - a
+ elif excess_mode == "rsuba":
+ tangent_vector = orbit - a
+ elif excess_mode == "rsubb":
+ tangent_vector = orbit - b
+ else:
+ errstr = f"Unknown excess mode: {excess_mode}"
+ raise ValueError(errstr)
+ if invert_excess:
+ tangent_vector = tangent_vector.neg_()
+
+ # 2. How far out of bounds are we? (0 if t <= 1)
+ if isinstance(t_orig, torch.Tensor):
+ excess = (t_orig.abs() - 1.0).clamp_min_(0.0).copysign_(t_orig)
+ else:
+ excess = math.copysign(max(0.0, abs(t_orig) - 1.0), t_orig)
+
+ return orbit.add_(tangent_vector.mul_(excess))
+
+
class BlendMode:
__slots__ = (
"allow_scale",
@@ -805,24 +2694,36 @@ class BlendMode:
"f_kwargs",
"f_raw",
"force_rescale",
+ "fork_rng",
+ "invert_scale",
"norm",
"norm_dims",
"rescale_dims",
+ "rescale_max",
+ "rescale_min",
"rev",
+ "scale_multiplier",
+ "visible",
)
class _Empty:
pass
- def __init__( # noqa: PLR0917
+ def __init__(
self,
f,
norm=None,
- norm_dims=(-3, -2, -1),
- rev=False,
- allow_scale=True,
- rescale_dims=(-3, -2, -1),
- force_rescale=False,
+ norm_dims: tuple = (-3, -2, -1),
+ rev: bool = False,
+ allow_scale: bool = True,
+ rescale_dims: tuple = (-3, -2, -1),
+ rescale_min: float = 0.0,
+ rescale_max: float = 1.0,
+ force_rescale: bool = False,
+ fork_rng: bool = False,
+ invert_scale: float | None = None,
+ scale_multiplier: float = 1.0,
+ visible: bool = True,
**kwargs: dict,
):
self.f_raw = f
@@ -837,37 +2738,35 @@ def __init__( # noqa: PLR0917
self.rev = rev
self.allow_scale = allow_scale
self.rescale_dims = rescale_dims
+ self.rescale_min = rescale_min
+ self.rescale_max = rescale_max
self.force_rescale = force_rescale
+ self.fork_rng = fork_rng
+ self.invert_scale = invert_scale
+ self.scale_multiplier = scale_multiplier
+ self.visible = visible
- def edited(
- self,
- *,
- f=_Empty,
- norm=_Empty,
- norm_dims=_Empty,
- rev=_Empty,
- allow_scale=_Empty,
- rescale_dims=_Empty,
- force_rescale=_Empty,
- preserve_kwargs=True,
- **kwargs: dict,
- ) -> object:
+ def edited(self, *, f=_Empty, preserve_kwargs=True, **kwargs: dict) -> BlendMode:
empty = self._Empty
kwargs = (self.f_kwargs | kwargs) if preserve_kwargs else kwargs
- return self.__class__(
- f if f is not empty else self.f_raw,
- norm=norm if norm is not empty else self.norm,
- norm_dims=norm_dims if norm_dims is not empty else self.norm_dims,
- rev=rev if rev is not empty else self.rev,
- allow_scale=allow_scale if allow_scale is not empty else self.allow_scale,
- rescale_dims=rescale_dims
- if rescale_dims is not empty
- else self.rescale_dims,
- force_rescale=force_rescale
- if force_rescale is not empty
- else self.force_rescale,
- **kwargs,
- )
+ kwargs |= {
+ k: v if (v := kwargs.get(k, empty)) is not empty else getattr(self, k)
+ for k in (
+ "norm",
+ "norm_dims",
+ "rev",
+ "allow_scale",
+ "rescale_dims",
+ "rescale_min",
+ "rescale_max",
+ "force_rescale",
+ "fork_rng",
+ "invert_scale",
+ "scale_multiplier",
+ "visible",
+ )
+ }
+ return self.__class__(f if f is not empty else self.f_raw, **kwargs)
def rescale(self, t, *, rescale_dims=_Empty):
if t.ndim > 2:
@@ -879,26 +2778,82 @@ def rescale(self, t, *, rescale_dims=_Empty):
rescale_dims = -1
tmin = torch.amin(t, keepdim=True, dim=rescale_dims)
tmax = torch.amax(t, keepdim=True, dim=rescale_dims)
- return (t - tmin).div_(tmax - tmin).clamp_(0, 1), tmin, tmax
+ return (
+ (t - tmin).div_(tmax - tmin).clamp_(self.rescale_min, self.rescale_max),
+ tmin,
+ tmax,
+ )
- def __call__(self, a, b, t, *, norm_dims=_Empty) -> torch.Tensor:
+ def _blend_internal(
+ self,
+ a: torch.Tensor,
+ b: torch.Tensor,
+ t: torch.Tensor,
+ *,
+ norm_dims=_Empty,
+ rescale_min_blend: float = 0.5,
+ rescale_max_blend: float = 0.5,
+ rescale_min_blend_function: Callable = torch.lerp,
+ rescale_max_blend_function: Callable = torch.lerp,
+ **kwargs: dict,
+ ) -> torch.Tensor:
+ if self.force_rescale:
+ a, amin, amax = self.rescale(a)
+ b, bmin, bmax = self.rescale(b)
+ with torch.random.fork_rng(devices=(a.device, b.device), enabled=self.fork_rng):
+ result = self.__call__internal(a, b, t, norm_dims=norm_dims, **kwargs)
if not self.force_rescale:
- return self.__call__internal(a, b, t, norm_dims=norm_dims)
- a, amin, amax = self.rescale(a)
- b, bmin, bmax = self.rescale(b)
- result = self.__call__internal(a, b, t, norm_dims=norm_dims)
- del a, b
- rmin, rmax = torch.lerp(amin, bmin, 0.5), torch.lerp(amax, bmax, 0.5)
- del amin, amax, bmin, bmax
+ return result
+ rmin = rescale_min_blend_function(amin, bmin, rescale_min_blend)
+ rmax = rescale_max_blend_function(amax, bmax, rescale_max_blend)
return result.mul_(rmax.sub_(rmin)).add_(rmin)
- def __call__internal(self, a, b, t, *, norm_dims=_Empty) -> torch.Tensor:
- if not isinstance(t, torch.Tensor) and isinstance(a, torch.Tensor):
- t = a.new_full((1,), t)
+ _AT = TypeVar("_AT", torch.Tensor, float)
+
+ def __call__(
+ self,
+ a: _AT,
+ b: torch.Tensor | float,
+ t: torch.Tensor | float,
+ *,
+ norm_dims=_Empty,
+ **kwargs: Any,
+ ) -> _AT:
+ float_a = not isinstance(a, torch.Tensor)
+ if float_a:
+ a = torch.tensor(a, dtype=torch.float64, device="cpu")
+ if not isinstance(b, torch.Tensor):
+ b = a.new_tensor(b)
+ if b.ndim > 1:
+ b = b.broadcast_to(a.shape)
+ if not isinstance(t, torch.Tensor):
+ t = a.new_tensor(t)
+ if t.ndim > 1:
+ t = t.broadcast_to(a.shape)
+ if float_a and (b.numel() != 1 or t.numel() != 1):
+ raise ValueError(
+ "When passing the 'a' parameter as a float, 'b' and 't' must either be float or 1-element tensors.",
+ )
+ result = self._blend_internal(a, b, t, norm_dims=norm_dims, **kwargs)
+ return result.mean().detach().cpu().item() if float_a else result
+
+ def __call__internal(
+ self,
+ a: torch.Tensor,
+ b: torch.Tensor,
+ t: torch.Tensor | float,
+ *,
+ norm_dims=_Empty,
+ **kwargs: dict,
+ ) -> torch.Tensor:
if self.rev:
a, b = b, a
+ if self.invert_scale is not None:
+ t = self.invert_scale - t
+ if self.scale_multiplier != 1.0:
+ t = t * self.scale_multiplier
if self.norm is None:
- return self.f(a, b, t)
+ return self.f(a, b, t, **kwargs)
return self.norm(
self.f(a, b, t),
reference_latent=torch.lerp(a, b, t),
@@ -907,11 +2862,36 @@ def __call__internal(self, a, b, t, *, norm_dims=_Empty) -> torch.Tensor:
class BlendingModes:
+ BLEH = True
+
def __init__(self, builtins=None):
self.builtins = {} if builtins is None else builtins
self.cache = {}
- def get(self, k: str, default=None):
+ def get_dict_key(self, k: dict):
+ ds = frozenset(k.items())
+ cached = self.cache.get(ds)
+ if cached is not None:
+ return cached
+ name = k.get("name")
+ if name is None:
+ raise ValueError(
+ "When passing a blend mode key as dict, a string 'name' key must exist."
+ )
+ name = name.strip()
+ base_bm = self.builtins.get(name)
+ if base_bm is None:
+ errstr = f"Unknown mode {name} for extended blend specification"
+ raise ValueError(errstr)
+ bm_kwargs = k.copy()
+ del bm_kwargs["name"]
+ bm = base_bm.edited(**bm_kwargs)
+ self.cache[k] = bm
+ return bm
+
+ def get(self, k: str | dict, default=None):
+ if isinstance(k, dict):
+ return self.get_dict_key(k)
result = self.builtins.get(k)
if result is not None:
return result
@@ -995,23 +2975,23 @@ def try_extended(self, k: str, default=None) -> object:
return bm
def items(self):
- return self.builtins.items()
+ return ((k, v) for k, v in self.builtins.items() if v.visible)
def values(self):
- return self.builtins.values()
+ return (v for v in self.builtins.values() if v.visible)
def __contains__(self, k: str) -> bool:
return self.get(k) is not None
def __iter__(self):
- return self.builtins.__iter__()
+ return (k for k, _v in self.items())
keys = __iter__
- def __setitem__(self, k: str, v) -> str:
+ def __setitem__(self, k: str, v) -> None:
self.builtins[k] = v if isinstance(v, BlendMode) else BlendMode(v)
- def __getitem__(self, k: str):
+ def __getitem__(self, k: str) -> BlendMode:
result = self.get(k)
if result is None:
raise KeyError(k)
@@ -1054,22 +3034,20 @@ def copy(self):
"a_only": BlendMode(lambda a, _b, t: a * t, allow_scale=False),
"b_only": BlendMode(lambda _a, b, t: b * t, allow_scale=False),
# Interpolates between tensors a and b using normalized linear interpolation.
- "bislerp": BlendMode(
- lambda a, b, t: ((1 - t) * a).add_(t * b),
- normalize,
- ),
- # "nbislerp": BlendMode(lambda a, b, t: (1 - t) * a + t * b, normalize),
+ # This definitely isn't biSLERP.
+ "bislerp_wrong": BlendMode(torch.lerp, normalize),
+ # "^"bislerp": BlendMode(lambda a, b, t: (1 - t) * a + t * b, normalize),
"slerp": BlendMode(altslerp),
# Transfer the color from `b` to `a` by t` factor
- "colorize": BlendMode(lambda a, b, t: (b - a).mul_(t).add_(a)),
+ "colorize": BlendMode(torch.lerp),
# Interpolates between tensors a and b using cosine interpolation.
"cosinterp": BlendMode(
- lambda a, b, t: (
- (a + b).sub_((a - b).mul_(torch.cos(t * torch.tensor(math.pi))))
- ).div_(2),
+ lambda a, b, t: ((a + b).sub_((a - b).mul_((t * torch.pi).cos()))).div_(2),
),
# Interpolates between tensors a and b using cubic interpolation.
- "cuberp": BlendMode(lambda a, b, t: (b - a).mul_(3 * t**2 - 2 * t**3).add_(a)),
+ "cuberp": BlendMode(
+ lambda a, b, t: (b - a).mul_((3 * t**2).sub_(2 * t**3)).add_(a),
+ ),
# Interpolates between tensors a and b using normalized linear interpolation,
# with a twist when t is greater than or equal to 0.5.
"hslerp": BlendMode(hslerp),
@@ -1118,30 +3096,36 @@ def copy(self):
"inject_copysign_b": BlendMode(lambda a, b, t: (b * t).add_(a).copysign_(b)),
"inject_avoidsign_a": BlendMode(lambda a, b, t: (b * t).add_(a).copysign_(a.neg())),
"inject_avoidsign_b": BlendMode(lambda a, b, t: (b * t).add_(a).copysign_(b.neg())),
+ "cfg": BlendMode(torch.lerp),
+ "cfg_base_a": BlendMode(lambda a, b, t: (a - b).mul_(t).add_(a)),
# Interpolates between tensors a and b using linear interpolation.
- # "lerp": BlendMode(lambda a, b, t: ((1.0 - t) * a).add_(t * b)),
"lerp": BlendMode(torch.lerp),
- "lerp050x": BlendMode(lambda a, b, t: (((1 - t) * a).add_(t * b)).mul_(0.5)),
- "lerp075x": BlendMode(lambda a, b, t: (((1 - t) * a).add_(t * b)).mul_(0.75)),
- "lerp110x": BlendMode(lambda a, b, t: (((1 - t) * a).add_(t * b)).mul_(1.1)),
- "lerp125x": BlendMode(lambda a, b, t: (((1 - t) * a).add_(t * b)).mul_(1.25)),
- "lerp150x": BlendMode(lambda a, b, t: (((1 - t) * a).add_(t * b)).mul_(1.5)),
+ "lerp050x": BlendMode(lambda a, b, t: a.lerp(b, t).mul_(0.5)),
+ "lerp075x": BlendMode(lambda a, b, t: a.lerp(b, t).mul_(0.75)),
+ "lerp110x": BlendMode(lambda a, b, t: a.lerp(b, t).mul_(1.1)),
+ "lerp125x": BlendMode(lambda a, b, t: a.lerp(b, t).mul_(1.25)),
+ "lerp150x": BlendMode(lambda a, b, t: a.lerp(b, t).mul_(1.5)),
"lerp_copysign_a": BlendMode(
- lambda a, b, t: ((1.0 - t) * a).add_(t * b).copysign_(a),
+ lambda a, b, t: a.lerp(b, t).copysign_(a),
),
"lerp_copysign_b": BlendMode(
- lambda a, b, t: ((1.0 - t) * a).add_(t * b).copysign_(b),
+ lambda a, b, t: a.lerp(b, t).copysign_(b),
),
"lerp_avoidsign_a": BlendMode(
- lambda a, b, t: ((1.0 - t) * a).add_(t * b).copysign_(a.neg()),
+ lambda a, b, t: a.lerp(b, t).copysign_(a.neg()),
),
"lerp_avoidsign_b": BlendMode(
- lambda a, b, t: ((1.0 - t) * a).add_(t * b).copysign_(b.neg()),
+ lambda a, b, t: a.lerp(b, t).copysign_(b.neg()),
+ ),
+ "weighted_average": BlendMode(
+ lambda a, b, t: (b * t).add_(a) / (1.0 + abs(t)),
),
# Simulates a brightening effect by adding tensor b to tensor a, scaled by t.
"lineardodge": BlendMode(lambda a, b, t: (b * t).add_(a)),
"copysign": BlendMode(lambda a, b, _t: torch.copysign(a, b)),
- "probcopysign": BlendMode(lambda a, b, t: torch.copysign(a, prob_blend(a, b, t))),
+ "probcopysign": BlendMode(
+ lambda a, b, t: prob_blend(a, b, t).copysign_(a),
+ ),
"slice_flat_d1": BlendMode(slice_blend, dim=1, flatten=True),
"slice_flat_d2": BlendMode(slice_blend, dim=2, flatten=True),
"slice_d1": BlendMode(slice_blend, dim=1, flatten=False),
@@ -1235,10 +3219,12 @@ def copy(self):
),
"hardlight": BlendMode(
lambda a, b, t: (
- 2 * a * b * (a < 0.5).float()
- + (1 - 2 * (1 - a) * (1 - b)) * (a >= 0.5).float()
- )
- * t,
+ (
+ 2 * a * b * (a < 0.5).float()
+ + (1 - 2 * (1 - a) * (1 - b)) * (a >= 0.5).float()
+ )
+ * t
+ ),
allow_scale=False,
force_rescale=True,
),
@@ -1251,10 +3237,16 @@ def copy(self):
normalize,
allow_scale=False,
),
+ "multiply_by_b": BlendMode(
+ lambda a, b, _t: a * b,
+ allow_scale=False,
+ ),
"overlay": BlendMode(
- lambda a, b, t: (2 * a * b + a**2 - 2 * a * b * a) * t
- if torch.all(b < 0.5)
- else (1 - 2 * (1 - a) * (1 - b)) * t,
+ lambda a, b, t: (
+ (2 * a * b + a**2 - 2 * a * b * a) * t
+ if torch.all(b < 0.5)
+ else (1 - 2 * (1 - a) * (1 - b)) * t
+ ),
allow_scale=False,
force_rescale=True,
),
@@ -1296,6 +3288,109 @@ def copy(self):
allow_scale=False,
force_rescale=True,
),
+ "wavelet_b_hi_100_lo_0": BlendMode(
+ f=wavelet_blend,
+ blend_yl_offset=0.0,
+ blend_yh_offset=1.0,
+ wave="db4",
+ level=8,
+ ),
+ "wavelet_b_hi_0_lo_100": BlendMode(
+ f=wavelet_blend,
+ blend_yl_offset=1.0,
+ blend_yh_offset=0.0,
+ wave="db4",
+ level=8,
+ ),
+ "ortho": BlendMode(ortho_blend),
+ "ortho_rescaled": BlendMode(ortho_blend, rescale_limit=2.0),
+ "ortho_rescaled_lerpish": BlendMode(
+ ortho_blend,
+ rescale_limit=2.0,
+ rescale_result_blend_mode="lerp",
+ rescale_result_mode="blend",
+ ),
+ "ortho_lerp": BlendMode(ortho_blend, blend_mode="lerp"),
+ "ortho_dyn_lerp": BlendMode(
+ ortho_blend,
+ blend_mode="lerp",
+ rescale_result_mode="blend",
+ rescale_limit=4.0,
+ dyn_ortho_mode=True,
+ ),
+ "ortho_dyn_lerp_inverted": BlendMode(
+ ortho_blend,
+ blend_mode="lerp",
+ rescale_result_mode="blend",
+ rescale_limit=2.0,
+ dyn_ortho_mode=True,
+ rev=True,
+ invert_scale=1.0,
+ ),
+ "ortho_lerp_rescaled": BlendMode(
+ ortho_blend,
+ blend_mode="lerp",
+ rescale_result_mode="blend",
+ rescale_limit=2.0,
+ ),
+ "ortho_cfg": BlendMode(
+ lambda a, b, t, **kwargs: ortho_blend(b, a - b, t, **kwargs),
+ ),
+ "ortho_cfg_base_a": BlendMode(
+ lambda a, b, t, **kwargs: ortho_blend(a, a - b, t, **kwargs),
+ ),
+ "symmetric_ortho": BlendMode(symmetric_ortho_blend),
+ "symmetric_ortho_rescaled": BlendMode(symmetric_ortho_blend, rescale_limit=2.0),
+ "contrastive_ortho_cfg": BlendMode(
+ lambda a, b, t, **kwargs: contrastive_ortho_cfg_base_a(b, a, t, **kwargs),
+ ),
+ "contrastive_ortho_cfg_base_a": BlendMode(contrastive_ortho_cfg_base_a),
+ # These next two probably aren't actually useful.
+ "symmetric_ortho_cfg": BlendMode(
+ lambda a, b, t, **kwargs: symmetric_ortho_blend(b, a - b, t, **kwargs),
+ visible=False,
+ ),
+ "symmetric_ortho_cfg_base_a": BlendMode(
+ lambda a, b, t, **kwargs: symmetric_ortho_blend(a, a - b, t, **kwargs),
+ visible=False,
+ ),
+ "tiered_blend": BlendMode(
+ tiered_blend,
+ visible=False,
+ ),
+ "sp_circular_interpolation": BlendMode(sp_circular_interpolation),
+ "geodesic_square_matrix": BlendMode(geodesic_square_matrix),
+ "fft_blend": BlendMode(fft_blend),
+ "fft_phase_blend": BlendMode(partial(fft_blend, magnitude_blend_multiplier=0.0)),
+ "fft_magnitude_blend": BlendMode(partial(fft_blend, phase_blend_multiplier=0.0)),
+ "decomp_rank_blend": BlendMode(DecompBlend.rank_slice_blend),
+ "decomp_diff": BlendMode(DecompBlend.rank_blend),
+ "chain": BlendMode(chain_blend, visible=False),
+ "pct_limited_025": BlendMode(partial(pct_limit_blend, diff_limit=0.25)),
+ "moment_aligned": BlendMode(moment_aligned_blend),
+ "distro_aligned": BlendMode(partial(distro_aligned_blend, align_b="da")),
+ "distro_aligned_result": BlendMode(
+ partial(distro_aligned_blend, align_result="da"),
+ ),
+ "gaussian_aligned_result": BlendMode(
+ partial(distro_aligned_blend, align_result="dg"),
+ ),
+ "gaussian_aligned": BlendMode(
+ partial(
+ distro_aligned_blend,
+ align_a="dg",
+ align_b="dg",
+ align_result="dg",
+ ),
+ ),
+ "magnitude_interpolation_lerpsign": BlendMode(
+ partial(rms_interpolation, sign_mode="blend", abs_inputs=True, power=1.0),
+ ),
+ "rms_interpolation_lerpsign": BlendMode(
+ partial(rms_interpolation, sign_mode="blend"),
+ ),
+ "pythagorean_lerp": BlendMode(pythagorean_lerp),
+ "orbit": BlendMode(orbit_blend),
}
BLENDING_MODES |= {
@@ -1319,10 +3414,10 @@ def copy(self):
"bislerp": slerp_orig,
"altbislerp": altslerp,
"revaltbislerp": lambda a, b, t: altslerp(b, a, t),
- "bibislerp": BLENDING_MODES["bislerp"].edited(norm_dims=0),
+ "bibislerp": BLENDING_MODES["bislerp_wrong"].edited(norm_dims=0),
"revhslerp": lambda a, b, t: hslerp_alt(b, a, t),
"revbislerp": lambda a, b, t: slerp_orig(b, a, t),
- "revbibislerp": BLENDING_MODES["revbislerp"].edited(norm_dims=0),
+ "revbibislerp": BLENDING_MODES["revbislerp_wrong"].edited(norm_dims=0),
}
FILTER_PRESETS = {
@@ -1562,7 +3657,7 @@ def scale_samples(
).tolist()
mode, mode_h = (
m if mode not in {"random", "randomaa"} else RAND_UPSCALE_METHODS[ridx]
- for ridx, m in zip(ridxs, (mode, mode_h))
+ for ridx, m in zip(ridxs, (mode, mode_h), strict=True)
)
mode_h = mode
if mode in {"bicubic", "nearest-exact", "bilinear", "area"}:
@@ -1606,12 +3701,14 @@ def scale_samples(
# Modified from ComfyUI
-def biderp(samples, width, height, mode="bislerp", mode_h=None): # noqa: PLR0914
+def biderp(samples, width, height, mode="bislerp", mode_h=None):
if mode_h is None:
mode_h = mode
derp_w = (BIDERP_MODES if ":" not in mode else BLENDING_MODES).get(mode, slerp_orig)
- derp_h = (BIDERP_MODES if ":" not in mode_h else BLENDING_MODES).get(mode_h, slerp_orig)
+ derp_h = (BIDERP_MODES if ":" not in mode_h else BLENDING_MODES).get(
+ mode_h, slerp_orig
+ )
def generate_bilinear_data(length_old, length_new, device):
coords_1 = torch.arange(length_old, dtype=torch.float32, device=device).reshape(
diff --git a/py/nodes/__init__.py b/py/nodes/__init__.py
index f8358ec..bf4d4a7 100644
--- a/py/nodes/__init__.py
+++ b/py/nodes/__init__.py
@@ -36,10 +36,15 @@
"BlehPlug": misc.BlehPlug,
"BlehRefinerAfter": refinerAfter.BlehRefinerAfter,
"BlehSageAttentionSampler": sageAttention.BlehSageAttentionSampler,
+ "BlehAdvancedAttentionSampler": sageAttention.BlehAdvancedAttentionSampler,
"BlehSetSamplerPreset": samplers.BlehSetSamplerPreset,
"BlehSetSigmas": misc.BlehSetSigmas,
"BlehTAEVideoDecode": taevid.TAEVideoDecode,
"BlehTAEVideoEncode": taevid.TAEVideoEncode,
+ "BlehModelProcessLatentIn": misc.BlehModelProcessLatentIn,
+ "BlehModelProcessLatentOut": misc.BlehModelProcessLatentOut,
+ "BlehFixGuiderPreviewing": misc.BlehFixGuiderPreviewing,
+ "BlehBlendConditioning": misc.BlehBlendConditioning,
}
NODE_DISPLAY_NAME_MAPPINGS = {
diff --git a/py/nodes/blockCFG.py b/py/nodes/blockCFG.py
index 164769e..5918e2a 100644
--- a/py/nodes/blockCFG.py
+++ b/py/nodes/blockCFG.py
@@ -1,4 +1,60 @@
-from functools import partial
+import math
+from dataclasses import dataclass
+from enum import Enum, auto
+from functools import partial, reduce
+
+import torch
+from tqdm import tqdm
+
+
+class BlockType(Enum):
+ INPUT = auto()
+ OUTPUT = auto()
+ MIDDLE = auto()
+ ATTN_Q = auto()
+ ATTN_K = auto()
+ ATTN_V = auto()
+ ATTN = auto()
+
+
+class BlendType(Enum):
+ DIFF = auto()
+ RESULT = auto()
+
+
+class CondType(Enum):
+ COND = auto()
+ UNCOND = auto()
+ BOTH = auto()
+
+
+@dataclass
+class BlockCFGItem:
+ start_sigma: float
+ end_sigma: float
+ block_type: BlockType
+ target_type: BlendType
+ cond_type: CondType
+ block_num: int
+ scale: float
+ skip_mode: bool
+
+
+class BlockCFG:
+ start_sigma: float
+ end_sigma: float
+ block_types: frozenset[BlockType]
+ verbose: bool = True
+
+ def __init__(self, items: tuple[BlockCFGItem, ...]):
+ self.start_sigma, self.end_sigma = reduce(
+ lambda old, new: (min(old[0], new[0]), max(old[1], new[1])),
+ ((i.start_sigma, i.end_sigma) for i in items),
+ (math.inf, math.inf * -1),
+ )
+ self.block_types = frozenset(i.block_type for i in items)
+
+ # def check_applies(self,
class BlockCFGBleh:
@@ -123,6 +179,7 @@ def patch(
reverse = apply_to != "cond"
def check_applies(block_list, transformer_options):
+ tqdm.write(f"* BLOCKCFG: tf={transformer_options}")
cond_or_uncond = transformer_options["cond_or_uncond"]
if (
not (0 in cond_or_uncond and 1 in cond_or_uncond)
@@ -140,6 +197,13 @@ def check_applies(block_list, transformer_options):
return -1 in block_list
return block_def in {-1, transformer_options.get("transformer_index")}
+ def apply_cfg_fun_(tensor: torch.Tensor, primary_offset: int) -> torch.Tensor:
+ full_batch = tensor.shape[0]
+ if full_batch % 2:
+ raise RuntimeError("Batch size must be multiple of 2")
+ batch = full_batch // 2
+ diff = tensor[:batch, ...] - tensor[batch:, ...]
+
def apply_cfg_fun(tensor, primary_offset):
secondary_offset = 0 if primary_offset == 1 else 1
if reverse:
@@ -156,12 +220,15 @@ def apply_cfg_fun(tensor, primary_offset):
).mul_(scale)
return result
+ mid_patch = None
+
def non_output_block_patch(h, transformer_options, *, block_list):
+ nonlocal mid_patch
+ # print("\nSET????", mid_patch)
+ if mid_patch is not None:
+ mid_patch._bleh_set_topts(transformer_options)
cond_or_uncond = transformer_options["cond_or_uncond"]
- if not check_applies(
- block_list,
- transformer_options,
- ):
+ if not block_list or not check_applies(block_list, transformer_options):
return h
return apply_cfg_fun(h, cond_or_uncond.index(0))
@@ -180,7 +247,55 @@ def output_block_patch(h, hsp, transformer_options, *, block_list):
)
m = model.clone()
- if input_blocks:
+
+ if middle_blocks:
+ # print("******** MIDDLE")
+ try:
+ mb = model.get_model_object("diffusion_model.middle_block.0")
+ except AttributeError:
+ mb = None
+ orig_forward = getattr(mb, "forward", None)
+ if mb is None or orig_forward is None:
+ raise ValueError("Could not get middle block or forward")
+
+ class MBForward:
+ def __init__(self, orig_forward):
+ real_orig_forward = orig_forward
+ while temp := getattr(
+ real_orig_forward, "_bleh_orig_forward", None
+ ):
+ real_orig_forward = temp
+ orig_forward = real_orig_forward
+ self._bleh_orig_forward = orig_forward
+ self._bleh_topts = None
+
+ def _bleh_set_topts(self, transformer_options: dict) -> None:
+ # if self._bleh_topts:
+ # return
+ cond_or_uncond = transformer_options["cond_or_uncond"]
+ self._bleh_topts = {
+ "cond_or_uncond": cond_or_uncond.clone()
+ if isinstance(cond_or_uncond, torch.Tensor)
+ else cond_or_uncond,
+ "sigmas": transformer_options["sigmas"].clone(),
+ "block": ("middle", 0),
+ }
+
+ def __call__(self, *args: list, **kwargs: dict) -> torch.Tensor:
+ result = self._bleh_orig_forward(*args, **kwargs)
+ try:
+ return non_output_block_patch(
+ result,
+ self._bleh_topts,
+ block_list=middle_blocks,
+ )
+ finally:
+ self._bleh_topts = None
+
+ mid_patch = MBForward(orig_forward)
+ m.add_object_patch("diffusion_model.middle_block.0.forward", mid_patch)
+
+ if input_blocks or middle_blocks:
(
m.set_model_input_block_patch_after_skip
if skip_mode
@@ -188,11 +303,7 @@ def output_block_patch(h, hsp, transformer_options, *, block_list):
)(
partial(non_output_block_patch, block_list=input_blocks),
)
- if middle_blocks:
- m.set_model_patch(
- partial(non_output_block_patch, block_list=middle_blocks),
- "middle_block_patch",
- )
+
if output_blocks:
m.set_model_output_block_patch(
partial(output_block_patch, block_list=output_blocks),
diff --git a/py/nodes/misc.py b/py/nodes/misc.py
index 61c55f4..d098f5d 100644
--- a/py/nodes/misc.py
+++ b/py/nodes/misc.py
@@ -2,17 +2,24 @@
from __future__ import annotations
import contextlib
+import copy
+import math
import operator
import random
from decimal import Decimal
from functools import partial
+from itertools import pairwise
+from typing import TYPE_CHECKING, Any
import torch
from comfy import model_management
from comfy.model_management import throw_exception_if_processing_interrupted
-from ..better_previews.previewer import ensure_previewer
-from ..latent_utils import normalize_to_scale
+from .. import latent_utils
+from ..better_previews.previewer import PREVIEWER_STATE, ensure_previewer
+
+if TYPE_CHECKING:
+ from collections.abc import Callable
class DiscardPenultimateSigma:
@@ -405,7 +412,7 @@ def go(
)[..., :4]
if values_mode == "clamp":
return (image.clamp(0.0, 1.0),)
- image = normalize_to_scale(
+ image = latent_utils.normalize_to_scale(
image,
0.0,
1.0,
@@ -482,3 +489,558 @@ def output_block_patch(h, hsp, _transformer_options):
m.set_model_output_block_patch(output_block_patch)
return (m,)
+
+
+class BlehModelProcessLatentIn:
+ DESCRIPTION = "Advanced node that can be used to scale a raw latent for model input. Generally only needed if you're doing something that bypasses the normal latent input mechanisms."
+ RETURN_TYPES = ("LATENT",)
+ FUNCTION = "go"
+ CATEGORY = "latent/advanced"
+
+ @classmethod
+ def INPUT_TYPES(cls) -> dict:
+ return {
+ "required": {
+ "model": ("MODEL",),
+ "latent": ("LATENT",),
+ },
+ }
+
+ @classmethod
+ def go(cls, *, model, latent: dict) -> tuple[dict]:
+ latent_format = model.model.latent_format
+ samples = (
+ latent["samples"].detach().to(device="cpu", dtype=torch.float32, copy=True)
+ )
+ return (latent | {"samples": latent_format.process_in(samples)},)
+
+
+class BlehModelProcessLatentOut:
+ DESCRIPTION = "Advanced node that can be used to scale a latent to the correct range for output. Generally only needed if you're doing something that bypasses the normal latent output mechanisms."
+ RETURN_TYPES = ("LATENT",)
+ FUNCTION = "go"
+ CATEGORY = "latent/advanced"
+
+ @classmethod
+ def INPUT_TYPES(cls) -> dict:
+ return {
+ "required": {
+ "model": ("MODEL",),
+ "latent": ("LATENT",),
+ },
+ }
+
+ @classmethod
+ def go(cls, *, model, latent: dict) -> tuple[dict]:
+ latent_format = model.model.latent_format
+ samples = (
+ latent["samples"].detach().to(device="cpu", dtype=torch.float32, copy=True)
+ )
+ return (latent | {"samples": latent_format.process_out(samples)},)
+
+
+class PreviewFixGuider:
+ def __init__(self, guider, **kwargs: Any):
+ self.__guider = guider
+ self.__state_overrides = {k: v for k, v in kwargs.items() if v is not None}
+
+ def __getattr__(self, k):
+ return getattr(self.__guider, k)
+
+ def sample(self, noise, latent_image, *args: Any, **kwargs: Any):
+ latent_shapes = (
+ (tuple(latent_image.shape),)
+ if not latent_image.is_nested
+ else tuple(tuple(t.shape) for t in latent_image.unbind())
+ )
+ PREVIEWER_STATE.last_latent_shapes = latent_shapes
+ soverrides = self.__state_overrides
+ saved_state = {}
+ if soverrides:
+ saved_state |= {k: getattr(PREVIEWER_STATE, k, None) for k in soverrides}
+ for k, v in soverrides.items():
+ setattr(PREVIEWER_STATE, k, v)
+ try:
+ return self.__guider.sample(noise, latent_image, *args, **kwargs)
+ finally:
+ PREVIEWER_STATE.last_latent_shapes = None
+ if saved_state:
+ for k, v in saved_state.items():
+ setattr(PREVIEWER_STATE, k, v)
+
+
+class BlehFixGuiderPreviewing:
+ DESCRIPTION = "Wraps a guider to give the Bleh previewing system a hint about the latent shapes. Only necessary for models like LTX-2 which use nested tensors."
+ FUNCTION = "go"
+ OUTPUT_NODE = False
+ CATEGORY = "hacks"
+
+ RETURN_TYPES = ("GUIDER",)
+
+ @classmethod
+ def INPUT_TYPES(cls):
+ return {
+ "required": {
+ "guider": ("GUIDER",),
+ "fps_override": (
+ "FLOAT",
+ {
+ "default": 0.0,
+ "min": 0.0,
+ "max": 9999.0,
+ "tooltip": "Can be used to override the FPS when previewing with video models. Disabled if set to 0.",
+ },
+ ),
+ },
+ "optional": {
+ "prefer_previewer": (
+ ("default", "ltxav", "ltxav23", "ltxav23wide"),
+ {
+ "default": "default",
+ "tooltip": "This is mostly only useful for LTX 2.3 since there isn't a way for the internal logic to know what latent format is being used. Set this to ltxav23 for LTX 2.3 (ltxav23wide for the wide previewer model), otherwise leave on the default. Note: This option sets global state.",
+ },
+ ),
+ },
+ }
+
+ @classmethod
+ def go(
+ cls,
+ *,
+ guider,
+ fps_override: float | None = None,
+ prefer_previewer: str | None = None,
+ ) -> tuple:
+ if prefer_previewer:
+ PREVIEWER_STATE.prefer_previewer = prefer_previewer
+ return (
+ PreviewFixGuider(
+ guider,
+ fps_override=fps_override if fps_override != 0.0 else None,
+ ),
+ )
+
+
+class ConditioningBlender:
+ """Base class for complex chronological conditioning blending."""
+
+ def __call__(self, *args: Any, **kwargs: Any):
+ return self.blend(*args, **kwargs)
+
+ @staticmethod
+ def get_bounds(meta_dict: dict) -> tuple[float, float]:
+ start = meta_dict.get("start_percent", 0.0)
+ end = meta_dict.get("end_percent", 1.0)
+ if start >= end:
+ end = math.nextafter(start, 2.0)
+ return start, end
+
+ def get_boundaries(self, cond1: list, cond2: list) -> list[float]:
+ boundaries = {0.0, 1.0}
+ for item in cond1 + cond2:
+ start, end = self.get_bounds(item[1])
+ boundaries.add(start)
+ boundaries.add(end)
+ return sorted(boundaries)
+
+ def get_active_items(self, cond_list: list, t_start: float, t_end: float) -> list:
+ active = []
+ for item in cond_list:
+ c_start, c_end = self.get_bounds(item[1])
+ if c_start <= t_start and c_end >= t_end:
+ active.append(item)
+ return active
+
+ def simple_blend_wrapper(
+ self,
+ blend_function: Callable,
+ cond1: list,
+ cond2: list,
+ strength: float,
+ *,
+ new_start_percent: float,
+ new_end_percent: float,
+ **kwargs: Any,
+ ) -> list:
+ blended_tensor = blend_function(cond1[0], cond2[0], strength, **kwargs)
+ new_meta = copy.deepcopy(cond1[1])
+ new_meta |= {
+ "start_percent": new_start_percent,
+ "end_percent": new_end_percent,
+ }
+ return [blended_tensor, new_meta]
+
+ def blend(
+ self,
+ cond1: list,
+ cond2: list,
+ blend_func: Callable,
+ *,
+ strength: float = 0.5,
+ blend_full_items: bool = False,
+ **kwargs: Any,
+ ) -> list:
+ boundaries = self.get_boundaries(cond1, cond2)
+ if not blend_full_items:
+ blend_func = partial(self.simple_blend_wrapper, blend_func)
+
+ result_conditioning = []
+
+ for t_start, t_end in pairwise(boundaries):
+ if t_start >= t_end:
+ continue
+
+ active1 = self.get_active_items(cond1, t_start, t_end)
+ active2 = self.get_active_items(cond2, t_start, t_end)
+ first_active_cond = active1 or active2
+
+ if not first_active_cond:
+ continue
+
+ adjusted_end = math.nextafter(t_end, -1.0) if t_end < 1.0 else t_end
+
+ if not (active1 and active2):
+ for c in first_active_cond:
+ new_meta = copy.deepcopy(c[1])
+ new_meta |= {"start_percent": t_start, "end_percent": adjusted_end}
+ result_conditioning.append([c[0].clone(), new_meta])
+ continue
+
+ for c1 in active1:
+ for c2 in active2:
+ blend_result = blend_func(
+ c1,
+ c2,
+ strength,
+ new_start_percent=t_start,
+ new_end_percent=adjusted_end,
+ **kwargs,
+ )
+ result_conditioning.append(blend_result)
+
+ return result_conditioning
+
+
+class BlehConditioningBlender(ConditioningBlender):
+ """Extended Blender capable of N-dimensional dynamic slicing, padding, and alignment."""
+
+ def __init__(
+ self,
+ *,
+ size_mismatch_strategy: str = "zero",
+ size_mismatch_alignment_mode: str = "left",
+ **kwargs: Any,
+ ):
+ super().__init__(**kwargs)
+ self.strategy = size_mismatch_strategy
+ self.alignment = size_mismatch_alignment_mode
+
+ def get_slices(
+ self,
+ shape: torch.Size | tuple[int, ...],
+ target_shape: torch.Size | tuple[int, ...],
+ ) -> tuple[tuple[slice, ...], tuple[slice, ...]]:
+ """Generates N-dimensional slice objects for both source and target."""
+ src_slices, tgt_slices = [], []
+ for sz, tgt_sz in zip(shape, target_shape, strict=True):
+ # Slice start, slice end, target start, target end
+ ss = se = ts = te = None
+ if sz < tgt_sz: # Padding source into target
+ if self.alignment == "left":
+ ts = 0
+ elif self.alignment == "right":
+ ts = tgt_sz - sz
+ else: # center
+ ts = (tgt_sz - sz) // 2
+ te = ts + sz
+ elif sz > tgt_sz: # Cropping source down to target
+ if self.alignment == "left":
+ ss = 0
+ elif self.alignment == "right":
+ ss = sz - tgt_sz
+ else: # center
+ ss = (sz - tgt_sz) // 2
+ se = ss + tgt_sz
+ src_slices.append(slice(ss, se))
+ tgt_slices.append(slice(ts, te))
+ return tuple(src_slices), tuple(tgt_slices)
+
+ def process_tensor(
+ self,
+ t: torch.Tensor,
+ other_t: torch.Tensor,
+ *,
+ orig_t1: torch.Tensor,
+ orig_t2: torch.Tensor,
+ target_shape: tuple[int, ...],
+ ) -> torch.Tensor:
+ if t.shape == target_shape:
+ return t
+
+ if self.strategy == "replicate":
+ repeats = []
+ for sz, tgt_sz in zip(t.shape, target_shape, strict=True):
+ if tgt_sz % sz != 0:
+ errstr = f"Cannot replicate conditioning tensor: target size {tgt_sz} is not evenly divisible by source size {sz}."
+ raise ValueError(errstr)
+ repeats.append(tgt_sz // sz)
+ return t.repeat(*repeats)
+
+ src_slices, tgt_slices = self.get_slices(t.shape, target_shape)
+
+ if self.strategy == "smaller":
+ return t[src_slices]
+
+ # --- Padding Strategies ---
+ if self.strategy == "mean_cond_1":
+ # Always uses orig_t1, regardless of which tensor is currently 't'
+ dims = list(range(1, orig_t1.ndim)) if orig_t1.ndim > 1 else [0]
+ mean_val = orig_t1.mean(dim=dims, keepdim=True)
+ out = mean_val.expand(target_shape).clone()
+ elif self.strategy == "mean_cond_2":
+ # Always uses orig_t2, regardless of which tensor is currently 't'
+ dims = list(range(1, orig_t2.ndim)) if orig_t2.ndim > 1 else [0]
+ mean_val = orig_t2.mean(dim=dims, keepdim=True)
+ out = mean_val.expand(target_shape).clone()
+ elif self.strategy in {"larger", "match_cond_1", "match_cond_2"}:
+ out = other_t.expand(target_shape).clone()
+ else: # "zero"
+ out = t.new_zeros(target_shape)
+
+ out[tgt_slices] = t[src_slices]
+ return out
+
+ def align_tensors(
+ self,
+ t1: torch.Tensor,
+ t2: torch.Tensor,
+ ) -> tuple[torch.Tensor, torch.Tensor]:
+ """Aligns two arbitrary N-dimensional tensors based on the instance's strategy and alignment modes."""
+ if t1.shape == t2.shape:
+ return t1, t2
+
+ if t1.ndim != t2.ndim:
+ errstr = f"Can't align tensors with different numbers of dimensions. t1 shape: {t1.shape} ({t1.ndim}), t2 shape: {t2.shape} ({t2.ndim})"
+ raise ValueError(errstr)
+
+ if self.strategy == "error":
+ errstr = (
+ f"Size mismatch: {t1.shape} vs {t2.shape} (Strategy set to 'error')"
+ )
+ raise ValueError(errstr)
+
+ if self.strategy == "match_cond_1":
+ target_shape = t1.shape
+ elif self.strategy == "match_cond_2":
+ target_shape = t2.shape
+ else:
+ size_op: Callable = min if self.strategy == "smaller" else max
+ target_shape = tuple(
+ size_op(s1, s2) for s1, s2 in zip(t1.shape, t2.shape, strict=True)
+ )
+
+ # Pass t1 and t2 as the strict originals to preserve the mean math!
+ out1 = self.process_tensor(
+ t=t1,
+ other_t=t2,
+ orig_t1=t1,
+ orig_t2=t2,
+ target_shape=target_shape,
+ )
+ out2 = self.process_tensor(
+ t=t2,
+ other_t=t1,
+ orig_t1=t1,
+ orig_t2=t2,
+ target_shape=target_shape,
+ )
+ return out1, out2
+
+
+class BlehBlendConditioning:
+ DESCRIPTION = "Blends conditioning_1 with conditioning_2. Unlike the ComfyUI builtin node, this can handle multiple conditioning items on both side and respects time ranges but not other potential ranges (I.E. regional conditioning). Note however that conditionings are generally encoded as sequences of tokens, so something like 'a cute dog' blended with 'sketch of a cat' is likely to be doing something like blending 'a' with 'sketch', 'cute' with 'of', etc."
+ FUNCTION = "go"
+ OUTPUT_NODE = False
+ CATEGORY = "conditioning"
+ RETURN_TYPES = ("CONDITIONING",)
+
+ @classmethod
+ def INPUT_TYPES(cls):
+ return {
+ "required": {
+ "conditioning_1": ("CONDITIONING",),
+ "conditioning_2": ("CONDITIONING",),
+ "blend_mode": (
+ tuple(latent_utils.BLENDING_MODES.keys()),
+ {"default": "lerp"},
+ ),
+ "strength": (
+ "FLOAT",
+ {
+ "min": -99999.0,
+ "max": 999999.0,
+ "default": 0.5,
+ },
+ ),
+ "blend_cond_tensor": (
+ "BOOLEAN",
+ {
+ "default": True,
+ "tooltip": "Controls whether the base conditioning tensor is blended. When disabled, this will use the conditioning tensor from conditioning_1.",
+ },
+ ),
+ "blend_tensors": (
+ "STRING",
+ {
+ "default": "pooled_output",
+ "tooltip": "Comma-separated list of other tensors to blend if they exist (for example pooled_output)",
+ },
+ ),
+ "metadata_base_mode": (
+ ("cond_1", "cond_2", "prefer_cond_1", "prefer_cond_2", "empty"),
+ {
+ "default": "prefer_cond_1",
+ "tooltip": "Controls what metadata ends up in the blended conditioning item. Items in the blend_tensors list will always be included. Possible values:\ncond_1 - Uses conditioning_1\ncond_2 - See above\nprefer_cond_1 - Uses the key from conditioning_1 if it exists, otherwise conditioning_2\nprefer_cond_2 - See above.\nempty - Starts with empty metadata.",
+ },
+ ),
+ "size_mismatch_strategy": (
+ (
+ "zero",
+ "mean_cond_1",
+ "mean_cond_2",
+ "match_cond_1",
+ "match_cond_2",
+ "replicate",
+ "larger",
+ "smaller",
+ "error",
+ ),
+ {
+ "default": "zero",
+ "tooltip": "Handles the case of size mismatches between items to be blended.\nzero - Uses the larger size, fills extra elements with zero (how ComfyUI usually handles it)\nmean_cond_1 - Same as above, uses the mean from conditioning_1.\nmean_cond_2 - See above.\nmatch_cond_1 - Matches the size to conditioning_1. Similar to choosing larger/smaller mode per-dimension (so alignment may apply).\nmatch_cond_2 - Same as above aside from using conditioning_2's sizes.\nreplicate - Replicates the smaller size to match the larger. It is an error if the sizes aren't evenly divisible. May be use for conditioning types like CLIP where mismatches will be increments of the CLIP max tokens size (77).\nlarger - Uses the values from the item with the larger size.\nsmaller - Prunes the tensor to the smaller size.\nerror - Size mismatches are an error.",
+ },
+ ),
+ "size_mismatch_alignment_mode": (
+ ("left", "right", "center"),
+ {
+ "default": "left",
+ "tooltip": "Only applies to size mismatch strategies that use the larger size and pad.\nleft - Aligns existing values to lower indexes, padding will apply after them. This is ComfyUI's normal behavior, the other options are probably quite weird.\nright - Padding will apply to lower indexes.\ncenter - Tries to center (left-biased when the size isn't divisible by 2) populated values in the space.",
+ },
+ ),
+ },
+ }
+
+ @classmethod
+ def go(
+ cls,
+ *,
+ conditioning_1: list,
+ conditioning_2: list,
+ blend_mode: str,
+ strength: float,
+ blend_cond_tensor: bool,
+ blend_tensors: str,
+ metadata_base_mode: str,
+ size_mismatch_strategy: str,
+ size_mismatch_alignment_mode: str,
+ ) -> tuple:
+ blend_tensor_list = (s.strip() for s in blend_tensors.split(","))
+ blend_tensor_set = {s for s in blend_tensor_list if s}
+ blend_function = latent_utils.BLENDING_MODES[blend_mode]
+ mdbm = metadata_base_mode
+
+ # Initialize our stateful blender
+ blender = BlehConditioningBlender(
+ size_mismatch_strategy=size_mismatch_strategy,
+ size_mismatch_alignment_mode=size_mismatch_alignment_mode,
+ )
+
+ def blend_wrapper(
+ cond1: list,
+ cond2: list,
+ strength: float,
+ *,
+ new_start_percent: float,
+ new_end_percent: float,
+ **kwargs: Any,
+ ) -> list:
+ c1t, c1d = cond1[:2]
+ c2t, c2d = cond2[:2]
+ cond1_extra = cond1[2:]
+
+ if blend_cond_tensor:
+ c2t = c2t.to(c1t)
+ c1t, c2t = blender.align_tensors(c1t, c2t)
+
+ def is_blendable(t: torch.Tensor):
+ return isinstance(t, torch.Tensor) and t.dtype.is_floating_point
+
+ c1_tensor_keys = {
+ k
+ for k, v in c1d.items()
+ if k in blend_tensor_set and isinstance(v, torch.Tensor)
+ }
+ c2_tensor_keys = {
+ k
+ for k, v in c2d.items()
+ if k in blend_tensor_set and isinstance(v, torch.Tensor)
+ }
+ blendable_tensor_keys = c1_tensor_keys.intersection(c2_tensor_keys)
+
+ if mdbm == "empty":
+ md = {}
+ elif mdbm == "cond_1":
+ md = c1d
+ elif mdbm == "cond_2":
+ md = c2d
+ elif mdbm == "prefer_cond_1":
+ md = c2d | c1d
+ elif mdbm == "prefer_cond_2":
+ md = c1d | c2d
+ else:
+ raise ValueError("Bad metadata_base_mode")
+
+ if not blend_cond_tensor:
+ ct = c1t.clone()
+ else:
+ ct = blend_function(c1t, c2t, strength, **kwargs)
+
+ if ct is c1t or ct is c2t:
+ ct = ct.clone()
+
+ bmd = {}
+ for bk in blendable_tensor_keys:
+ mt1, mt2 = c1d[bk], c2d[bk]
+ if mt1.ndim != mt2.ndim:
+ continue
+ mt2 = mt2.to(mt1)
+
+ # Apply the exact same logic to pooled_output / metadata tensors
+ mt1, mt2 = blender.align_tensors(mt1, mt2)
+
+ bmresult = blend_function(mt1, mt2, strength, **kwargs)
+ if bmresult is mt1 or bmresult is mt2:
+ bmresult = bmresult.clone()
+ bmd[bk] = bmresult
+
+ if bmd:
+ md = {k: v for k, v in md.items() if k not in bmd}
+
+ md = copy.deepcopy(md)
+ md |= {
+ "start_percent": new_start_percent,
+ "end_percent": new_end_percent,
+ **bmd,
+ }
+ if cond1_extra:
+ cond1_extra = copy.deepcopy(cond1_extra)
+ return [ct, md, *cond1_extra]
+
+ result = blender.blend(
+ conditioning_1,
+ conditioning_2,
+ blend_wrapper,
+ strength=strength,
+ blend_full_items=True,
+ )
+ return (result,)
diff --git a/py/nodes/sageAttention.py b/py/nodes/sageAttention.py
index af1aa46..3067b49 100644
--- a/py/nodes/sageAttention.py
+++ b/py/nodes/sageAttention.py
@@ -3,12 +3,18 @@
import contextlib
import importlib
-from typing import TYPE_CHECKING
+import math
+from enum import Enum, auto
+from functools import partial, update_wrapper
+from typing import TYPE_CHECKING, Any, NamedTuple
import comfy.ldm.modules.attention as comfyattn
import torch
import yaml
from comfy.samplers import KSAMPLER
+from tqdm import tqdm
+
+from ..latent_utils import BLENDING_MODES
try:
import sageattention
@@ -51,6 +57,22 @@
HAVE_ATTN_OVERRIDE = hasattr(comfyattn, "register_attention_function")
+# class AttnsConfig(NamedTuple):
+# name: str
+# version: str
+# supported_head_sizes: collections.abc.Collection
+
+# class AttentionRule(NamedTuple):
+
+
+# class AttentionRules(NamedTuple):
+# orig_attn: Callable
+# start_sigma: float = math.inf
+# end_sigma: float = 0.0
+# verbose: bool = False
+# rules: tuple[AttentionRule, ...] = ()
+
+
def attention_bleh( # noqa: PLR0914
q: torch.Tensor,
k: torch.Tensor,
@@ -182,10 +204,14 @@ def make_attn_wrapper(
raise ValueError(
"SpargeAttention is not available, make sure you have the spas_sage_attn Python package installed",
)
- if sageattn_function == "sparge":
+ if sageattn_function in {"sparge", "sparge2"}:
sageattn_function = spas_sage_attn.spas_sage2_attn_meansim_cuda
+ elif sageattn_function == "sparge2_topk":
+ sageattn_function = spas_sage_attn.spas_sage2_attn_meansim_topk_cuda
elif sageattn_function == "sparge1":
sageattn_function = spas_sage_attn.spas_sage_attn_meansim_cuda
+ elif sageattn_function == "sparge1_topk":
+ sageattn_function = spas_sage_attn.spas_sage_attn_meansim_topk_cuda
else:
sageattn_function = getattr(sageattention, sageattn_function)
@@ -331,6 +357,19 @@ def go(
return (model,)
+class TimeMode(Enum):
+ PERCENT = auto()
+ SIGMA = auto()
+
+
+class SageAttnOptions(NamedTuple):
+ sampler: object
+ attn_kwargs: dict[str, Any]
+ start_time: float = math.inf
+ end_time: float = 0.0
+ time_mode: TimeMode = TimeMode.PERCENT
+
+
class BlehModelWrapper:
def __init__(self, model: object, model_call: Callable):
self.__bleh_model = model
@@ -344,20 +383,25 @@ def __getattr__(self, k: str):
def sageattn_sampler(
+ config: SageAttnOptions,
model: object,
x: torch.Tensor,
sigmas: torch.Tensor,
- *,
- sageattn_sampler_options: tuple,
+ # *,
+ # sageattn_sampler_options: tuple,
**kwargs: dict,
) -> torch.Tensor:
- sampler, start_percent, end_percent, sageattn_kwargs = sageattn_sampler_options
- ms = model.inner_model.inner_model.model_sampling
- start_sigma, end_sigma = (
- round(ms.percent_to_sigma(start_percent), 4),
- round(ms.percent_to_sigma(end_percent), 4),
- )
- del ms
+ # sampler, start_percent, end_percent, sageattn_kwargs = sageattn_sampler_options
+ if config.time_mode == TimeMode.PERCENT:
+ ms = model.inner_model.inner_model.model_sampling
+ start_sigma, end_sigma = (
+ round(ms.percent_to_sigma(config.start_time), 4),
+ round(ms.percent_to_sigma(config.end_time), 4),
+ )
+ del ms
+ else:
+ start_sigma = config.start_time
+ end_sigma = config.end_time
def model_call(
model: object,
@@ -369,7 +413,7 @@ def model_call(
enabled = end_sigma <= sigma_float <= start_sigma
with sageattn_context(
enabled=enabled,
- **sageattn_kwargs,
+ **config.attn_kwargs,
) as attn_override:
if enabled and HAVE_ATTN_OVERRIDE:
model_options = kwargs.pop("model_options", {}).copy()
@@ -381,12 +425,12 @@ def model_call(
kwargs["model_options"] = model_options
return model(x, sigma, **kwargs)
- return sampler.sampler_function(
+ return config.sampler.sampler_function(
BlehModelWrapper(model, model_call),
x,
sigmas,
**kwargs,
- **sampler.extra_options,
+ **config.sampler.extra_options,
)
@@ -450,14 +494,383 @@ def go(
)
return (
KSAMPLER(
- sageattn_sampler,
- extra_options={
- "sageattn_sampler_options": (
- sampler,
- start_percent,
- end_percent,
- get_yaml_parameters(yaml_parameters),
+ update_wrapper(
+ partial(
+ sageattn_sampler,
+ SageAttnOptions(
+ start_time=start_percent,
+ end_time=end_percent,
+ time_mode=TimeMode.PERCENT,
+ sampler=sampler,
+ attn_kwargs=get_yaml_parameters(yaml_parameters),
+ ),
+ ),
+ sampler.sampler_function,
+ ),
+ ),
+ )
+
+
+class AdvancedAttnRule(NamedTuple):
+ attn_function: Callable | None
+ attn_kwargs: dict[str, Any]
+ blend_function: Callable | None = None
+ start_sigma: float = math.inf
+ end_sigma: float = 0.0
+ check_nan: bool = False
+ q_multiplier: float = 1.0
+ k_multiplier: float = 1.0
+ v_multiplier: float = 1.0
+ output_multiplier: float = 1.0
+ device: torch.device | str | None = None
+ dtype: torch.dtype | str | None = None
+ blend: float = 1.0
+ op_q: str | None = None
+ op_k: str | None = None
+ op_v: str | None = None
+ op_current_result_preblend: str | None = None
+ op_result_preblend: str | None = None
+ op_result_postblend: str | None = None
+ op_result_postblend_diff: str | None = None
+
+ @classmethod
+ def build(
+ cls,
+ *,
+ attn_function: str | Callable | None = None,
+ blend_mode: str | None = None,
+ device=None,
+ dtype=None,
+ **kwargs,
+ ) -> NamedTuple:
+ blend_function = BLENDING_MODES[blend_mode] if blend_mode is not None else None
+ my_kwargs = {k: kwargs.pop(k) for k in cls._fields if k in kwargs}
+ if attn_function == "default":
+ attn_function = None
+ elif isinstance(attn_function, str):
+ kwargs["sageattn_function"] = attn_function
+ attn_function = make_attn_wrapper(orig_attn=None, **kwargs)
+ if isinstance(dtype, str):
+ dtype = {
+ "float32": torch.float32,
+ "float16": torch.float16,
+ "bfloat16": torch.bfloat16,
+ "float64": torch.float64,
+ }.get(dtype)
+ return cls(
+ device=device,
+ dtype=dtype,
+ blend_function=blend_function,
+ attn_function=attn_function,
+ attn_kwargs=kwargs,
+ **my_kwargs,
+ )
+
+
+class AdvancedAttnConfig(NamedTuple):
+ sampler: object
+ verbose: bool = False
+ rules: tuple[AdvancedAttnRule, ...] = ()
+ start_time: float = math.inf
+ end_time: float = 0.0
+ call_indexes: frozenset[float] = frozenset()
+ time_mode: TimeMode = TimeMode.PERCENT
+ min_cond_batch: int = 0
+ batch_slice: tuple | str | None = None
+ max_idx: int = -1
+ delegate_override: bool = True
+ op_result: str | None = None
+ latent_ops: dict[str, Callable] = {}
+
+ @classmethod
+ def build(
+ cls,
+ *,
+ rules=(),
+ time_mode: str | TimeMode | None = None,
+ call_indexes=(),
+ **kwargs,
+ ) -> NamedTuple:
+ fs = frozenset(cls._fields)
+ rules = tuple(AdvancedAttnRule.build(**r) for r in rules)
+ if isinstance(time_mode, str):
+ time_mode = getattr(TimeMode, time_mode.strip().upper())
+ call_indexes = frozenset(
+ i
+ if math.isnan(i) or i == math.inf or not isinstance(i, float)
+ else int(i) + 0.5
+ for i in call_indexes
+ )
+ kwargs = {k: v for k, v in kwargs.items() if k in fs}
+ return cls(
+ rules=rules,
+ time_mode=time_mode,
+ call_indexes=call_indexes,
+ **kwargs,
+ )
+
+ def call_op(
+ self, op_key: str | None, t: torch.Tensor, *, sigma: float
+ ) -> torch.Tensor:
+ op = None if op_key is None else self.latent_ops.get(op_key)
+ if op is None:
+ return t
+ return (
+ op(t)
+ if not hasattr(op, "EXTENDED_LATENT_OPERATION")
+ else op(t, sigma=sigma)
+ )
+
+ def attn_wrapper(
+ self,
+ sigma: float,
+ currattncall: CurrAttnCall,
+ old_override: Callable | None,
+ comfy_orig_attn: Callable,
+ q: torch.Tensor,
+ k: torch.Tensor,
+ v: torch.Tensor,
+ *args,
+ **kwargs,
+ ) -> torch.Tensor:
+ fallback_attn = (
+ partial(old_override, comfy_orig_attn)
+ if old_override is not None and self.delegate_override
+ else comfy_orig_attn
+ )
+ result = None
+ rules = self.rules
+ call_idx = currattncall.idx
+ rev_idx = (
+ -abs(currattncall.max_idx - currattncall.idx)
+ if currattncall.max_idx >= 0
+ else math.nan
+ )
+ ci = self.call_indexes
+ fallthrough_match = math.inf in ci
+ exclude = (call_idx + 0.5) in ci or (
+ not math.isnan(rev_idx) and (rev_idx + 0.5) in ci
+ )
+ matched = not exclude and (fallthrough_match or call_idx in ci or rev_idx in ci)
+ if self.verbose:
+ tqdm.write(
+ f"[BLEH] AdvancedAttn wrapper({currattncall.idx:<3}): sigma={sigma:.4f}, max_idx={currattncall.max_idx:<3}, rev_idx={rev_idx:<3}, matched={matched}, exclude={exclude}, q.shape={q.shape}",
+ )
+ currattncall.idx += 1
+ rules = self.rules if matched else ()
+ for rule in rules:
+ if not rule.end_sigma <= sigma <= rule.start_sigma:
+ continue
+ currq, currk, currv = q, k, v
+ if rule.q_multiplier != 1:
+ currq = currq * rule.q_multiplier
+ if rule.k_multiplier != 1:
+ currk = currk * rule.k_multiplier
+ if rule.v_multiplier != 1:
+ currv = currv * rule.v_multiplier
+ if rule.dtype is not None or rule.device is not None:
+ currq = currq.to(device=rule.device, dtype=rule.dtype)
+ currk = currk.to(device=rule.device, dtype=rule.dtype)
+ currv = currv.to(device=rule.device, dtype=rule.dtype)
+ currq = self.call_op(rule.op_q, currq, sigma=sigma)
+ currk = self.call_op(rule.op_k, currk, sigma=sigma)
+ currv = self.call_op(rule.op_v, currv, sigma=sigma)
+ attn_function = (
+ partial(rule.attn_function, fallback_attn)
+ if rule.attn_function
+ else fallback_attn
+ )
+ curr_result = attn_function(currq, currk, currv, *args, **kwargs)
+ del currq, currk, currv
+ if rule.check_nan and curr_result.isnan().any():
+ del curr_result
+ continue
+ if rule.output_multiplier != 1:
+ curr_result *= rule.output_multiplier
+ curr_result = self.call_op(
+ rule.op_current_result_preblend,
+ curr_result,
+ sigma=sigma,
+ )
+ if curr_result.dtype != q.dtype or curr_result.device != q.device:
+ curr_result = curr_result.to(q)
+ if result is not None:
+ result = self.call_op(rule.op_result_preblend, result, sigma=sigma)
+ if result is None or rule.blend_function is None:
+ result = curr_result
+ del curr_result
+ continue
+ prev_result = result
+ result = rule.blend_function(result, curr_result, rule.blend)
+ if rule.op_result_postblend_diff is not None:
+ result = prev_result + self.call_op(
+ rule.op_result_postblend_diff,
+ result - prev_result,
+ sigma=sigma,
+ )
+ del curr_result, prev_result
+ result = self.call_op(rule.op_result_postblend, result, sigma=sigma)
+ if result is None:
+ result = fallback_attn(q, k, v, *args, **kwargs)
+ return self.call_op(self.op_result, result, sigma=sigma)
+
+
+class CurrAttnCall:
+ def __init__(self, idx: int = 0, max_idx: int = -1):
+ self.idx = idx
+ self.max_idx = max_idx
+
+
+def advancedattn_sampler(
+ config: AdvancedAttnConfig,
+ model: object,
+ x: torch.Tensor,
+ sigmas: torch.Tensor,
+ **kwargs: dict,
+) -> torch.Tensor:
+ if config.time_mode == TimeMode.PERCENT:
+ ms = model.inner_model.inner_model.model_sampling
+ start_sigma, end_sigma = (
+ round(ms.percent_to_sigma(config.start_time), 4),
+ round(ms.percent_to_sigma(config.end_time), 4),
+ )
+ del ms
+ else:
+ start_sigma = config.start_time
+ end_sigma = config.end_time
+
+ max_idx = -1
+
+ def model_call(
+ model: object,
+ x: torch.Tensor,
+ sigma: torch.Tensor,
+ **kwargs: dict[str, Any],
+ ) -> torch.Tensor:
+ nonlocal max_idx
+ sigma_float = float(sigma.max().detach().cpu())
+ enabled = end_sigma <= sigma_float <= start_sigma
+ if not enabled:
+ return model(x, sigma, **kwargs)
+ calltracker = CurrAttnCall(idx=0, max_idx=max_idx)
+ if config.verbose:
+ tqdm.write(f"[BLEH] AdvancedAttn: Config: {config}")
+ model_options = kwargs.pop("model_options", {}).copy()
+ transformer_options = model_options.pop("transformer_options", {}).copy()
+ old_override = transformer_options.pop("optimized_attention_override", None)
+ attn_override = partial(
+ config.attn_wrapper,
+ sigma_float,
+ calltracker,
+ old_override,
+ )
+ transformer_options["optimized_attention_override"] = attn_override
+ model_options["transformer_options"] = transformer_options
+ kwargs["model_options"] = model_options
+ result = model(x, sigma, **kwargs)
+ max_idx = max(max_idx, calltracker.idx)
+ return result
+
+ return config.sampler.sampler_function(
+ BlehModelWrapper(model, model_call),
+ x,
+ sigmas,
+ **kwargs,
+ **config.sampler.extra_options,
+ )
+
+
+class BlehAdvancedAttentionSampler:
+ DESCRIPTION = "TBD"
+ CATEGORY = "sampling/custom_sampling/samplers"
+ RETURN_TYPES = ("SAMPLER",)
+ FUNCTION = "go"
+
+ DEFAULT_YAML_PARAMS = """verbose: false
+start_time: 0.0
+end_time: 1.0
+# One of: percent, sigma
+time_mode: percent
+# .inf means match everything. Whole float values exclude an index. I.E 2.0
+# Call index as in the Nth call to attention this model evaluation.
+# Negative indexes count from the end but can only match after a pass through the model.
+call_indexes: [.inf]
+# Can be set to null (everything), cond, uncond or a list.
+batch_slice: null
+# Requires cond batch information to be passed and at least this many items.
+min_cond_batch: 0
+rules:
+ # Passed as sageattn_function unless set to default.
+ # Keys not in this list are passed through like with the SageAttention node:
+ # attn_function, blend_mode, blend
+ - attn_function: default
+ # You can set whatever other keys you want here.
+ - attn_function: sageattn
+ # Blends target the last attention result and are ignored
+ # if it's missing.
+ # The default blend means:
+ # sageattn + (defaultattn - sageattn) * 2
+ blend_mode: cfg
+ blend: 2.0
+"""
+
+ @classmethod
+ def INPUT_TYPES(cls) -> dict:
+ return {
+ "required": {
+ "sampler": ("SAMPLER",),
+ "yaml_parameters": (
+ "STRING",
+ {
+ "default": cls.DEFAULT_YAML_PARAMS,
+ "tooltip": "Allows specifying custom parameters via YAML. These are mostly passed directly to the SageAttention function with no error checking. Must be empty or a YAML object.",
+ "dynamicPrompts": False,
+ "multiline": True,
+ "defaultInput": True,
+ },
+ ),
+ },
+ "optional": {
+ "op_0": ("LATENT_OPERATION",),
+ "op_1": ("LATENT_OPERATION",),
+ "op_2": ("LATENT_OPERATION",),
+ "op_3": ("LATENT_OPERATION",),
+ "op_4": ("LATENT_OPERATION",),
+ "op_5": ("LATENT_OPERATION",),
+ "op_6": ("LATENT_OPERATION",),
+ "op_7": ("LATENT_OPERATION",),
+ "op_8": ("LATENT_OPERATION",),
+ "op_9": ("LATENT_OPERATION",),
+ },
+ }
+
+ @classmethod
+ def go(
+ cls,
+ sampler: object,
+ yaml_parameters: str,
+ **kwargs: dict,
+ ) -> tuple:
+ if sageattention is None:
+ raise RuntimeError(
+ "sageattention not installed to Python environment: SageAttention feature unavailable",
+ )
+ if not HAVE_ATTN_OVERRIDE:
+ raise RuntimeError(
+ "This node only supports recent ComfyUI versions that support attention overrides.",
+ )
+ params = get_yaml_parameters(yaml_parameters)
+ params["latent_ops"] = {
+ k: v for k, v in kwargs.items() if k.startswith("op_") and v is not None
+ }
+ return (
+ KSAMPLER(
+ update_wrapper(
+ partial(
+ advancedattn_sampler,
+ AdvancedAttnConfig.build(sampler=sampler, **params),
),
- },
+ sampler.sampler_function,
+ ),
),
)
diff --git a/py/nodes/taevid.py b/py/nodes/taevid.py
index 0892342..4d3d6b3 100644
--- a/py/nodes/taevid.py
+++ b/py/nodes/taevid.py
@@ -14,16 +14,31 @@ class TAEVideoNodeBase:
FUNCTION = "go"
CATEGORY = "latent"
+ _download_map = { # noqa: RUF012
+ "hunyuanvideo": ("taehv.pth", "taehv"),
+ "ltxv": ("taeltx_2.pth", "taehv"),
+ "ltxv23": ("taeltx2_3.pth", "taehv"),
+ "ltxv23wide": ("taeltx2_3_wide.pth", "taehv"),
+ "mochi": ("taem1.pth", "taem1"),
+ "wan21": ("taew2_1.pth", "taehv"),
+ "wan22": ("taew2_2.pth", "taehv"),
+ }
+
@classmethod
def INPUT_TYPES(cls) -> dict:
return {
"required": {
- "latent_type": (("wan21", "wan22", "hunyuanvideo", "mochi"),),
+ "latent_type": (
+ tuple(cls._download_map),
+ {
+ "tooltip": "Use ltxv for LTX-2 AV.",
+ },
+ ),
"parallel_mode": (
"BOOLEAN",
{
"default": False,
- "tooltip": "Parallel mode is faster but requires more memory.",
+ "tooltip": "Parallel mode may be faster but requires more memory.",
},
),
},
@@ -39,15 +54,15 @@ def get_taevid_model(
raise ValueError("Bad latent type")
tae_model_path = folder_paths.get_full_path("vae_approx", vmi.tae_model)
if tae_model_path is None:
- if latent_type == "wan21":
- model_src = "taew2_1.pth from https://github.com/madebyollin/taehv"
- elif latent_type == "wan22":
- model_src = "taew2_2.pth from https://github.com/madebyollin/taehv"
- elif latent_type == "hunyuanvideo":
- model_src = "taehv.pth from https://github.com/madebyollin/taehv"
+ dl_info = cls._download_map.get(latent_type)
+ if dl_info is None:
+ err_string = (
+ f"Unexpected latent type {latent_type}, no information available"
+ )
else:
- model_src = "taem1.pth from https://github.com/madebyollin/taem1"
- err_string = f"Missing TAE video model. Download {model_src} and place it in the models/vae_approx directory"
+ filename, reponame = dl_info
+ model_src = f"{filename} from https://github.com/madebyollin/{reponame}"
+ err_string = f"Missing TAE video model. Download {model_src} and place it in the models/vae_approx directory"
raise RuntimeError(err_string)
device = model_management.vae_device()
dtype = model_management.vae_dtype(device=device)
@@ -66,7 +81,7 @@ def go(cls, *, latent, latent_type: str, parallel_mode: bool) -> tuple:
class TAEVideoDecode(TAEVideoNodeBase):
RETURN_TYPES = ("IMAGE",)
CATEGORY = "latent"
- DESCRIPTION = "Fast decoding of Wan, Hunyuan and Mochi video latents with the video equivalent of TAESD."
+ DESCRIPTION = "Fast decoding of Wan, Hunyuan, Mochi and LTX video latents with the video equivalent of TAESD."
@classmethod
def INPUT_TYPES(cls) -> dict:
@@ -100,7 +115,7 @@ def go(cls, *, latent: dict, latent_type: str, parallel_mode: bool) -> tuple:
class TAEVideoEncode(TAEVideoNodeBase):
RETURN_TYPES = ("LATENT",)
CATEGORY = "latent"
- DESCRIPTION = "Fast encoding of Wan, Hunyuan and Mochi video latents with the video equivalent of TAESD."
+ DESCRIPTION = "Fast encoding of Wan, Hunyuan, Mochi and LTX video latents with the video equivalent of TAESD."
@classmethod
def INPUT_TYPES(cls) -> dict:
diff --git a/py/settings.py b/py/settings.py
index 327e93d..d8a0aab 100644
--- a/py/settings.py
+++ b/py/settings.py
@@ -1,4 +1,7 @@
+from __future__ import annotations
+
from pathlib import Path
+from typing import NamedTuple
class Settings:
@@ -7,6 +10,8 @@ def __init__(self):
def update(self, obj):
btp = obj.get("betterTaesdPreviews", None)
+ if btp is None:
+ btp = obj.get("previews", None)
self.btp_enabled = btp is not None and btp.get("enabled", True) is True
if not self.btp_enabled:
return
@@ -22,7 +27,8 @@ def update(self, obj):
self.btp_skip_upscale_layers = btp.get("skip_upscale_layers", 0)
self.btp_preview_device = btp.get("preview_device")
# default, keep, float32, float16, bfloat16
- self.btp_preview_dtype = btp.get("preview_dtype")
+ self.btp_preview_dtype = btp.get("preview_dtype", "bfloat16")
+ self.btp_preview_non_blocking = bool(btp.get("preview_non_blocking", False))
self.btp_maxed_batch_step_mode = btp.get("maxed_batch_step_mode", False)
self.btp_compile_previewer = btp.get("compile_previewer", False)
self.btp_oom_fallback = btp.get("oom_fallback", "latent2rgb")
@@ -37,6 +43,12 @@ def update(self, obj):
)
self.btp_animate_preview = btp.get("animate_preview", "none")
self.btp_verbose = btp.get("verbose", False)
+ self.btp_publish_last_preview = btp.get("publish_last_preview", False)
+ self.btp_publish_last_preview_min_refresh = max(
+ 1,
+ btp.get("publish_last_preview_min_refresh", 5),
+ )
+ self.btp_only_animate_last_preview = btp.get("only_animate_last_preview", True)
@staticmethod
def get_cfg_path(filename) -> Path:
diff --git a/py/wavelet_functions.py b/py/wavelet_functions.py
new file mode 100644
index 0000000..b993b41
--- /dev/null
+++ b/py/wavelet_functions.py
@@ -0,0 +1,244 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Callable
+
+import torch
+
+if TYPE_CHECKING:
+ from collections.abc import Sequence
+
+try:
+ import pytorch_wavelets as ptwav
+ import pywt
+
+ HAVE_WAVELETS = True
+except ImportError:
+ ptwav = None
+ pywt = None
+ HAVE_WAVELETS = False
+
+
+def fallback[V, D](val: V | D, default: D = None) -> V | D:
+ return val if val is not None else default
+
+
+class Wavelet:
+ DEFAULT_MODE = "symmetric"
+ DEFAULT_LEVEL = 3
+ DEFAULT_WAVE = "db4"
+ DEFAULT_USE_1D_DWT = False
+ DEFAULT_USE_DTCWT = False
+ DEFAULT_QSHIFT = "qshift_a"
+ DEFAULT_BIORT = "near_sym_a"
+
+ def __init__(
+ self,
+ *,
+ wave: str = DEFAULT_WAVE,
+ level: int = DEFAULT_LEVEL,
+ mode: str = DEFAULT_MODE,
+ use_1d_dwt: bool = DEFAULT_USE_1D_DWT,
+ use_dtcwt: bool = DEFAULT_USE_DTCWT,
+ biort: str = DEFAULT_BIORT,
+ qshift: str = DEFAULT_QSHIFT,
+ inv_wave: str | None = None,
+ inv_mode: str | None = None,
+ inv_biort: str | None = None,
+ inv_qshift=None,
+ device: str | torch.device | None = None,
+ ):
+ if not HAVE_WAVELETS:
+ raise RuntimeError(
+ "Wavelet use requires the pytorch_wavelets package to be installed in your Python environment",
+ )
+ inv_wave = fallback(inv_wave, wave)
+ inv_mode = fallback(inv_mode, mode)
+ inv_biort = fallback(inv_biort, biort)
+ inv_qshift = fallback(inv_qshift, qshift)
+ if use_dtcwt:
+ fwdfun, invfun = ptwav.DTCWTForward, ptwav.DTCWTInverse
+ elif use_1d_dwt:
+ fwdfun, invfun = ptwav.DWT1DForward, ptwav.DWT1DInverse
+ else:
+ fwdfun, invfun = ptwav.DWTForward, ptwav.DWTInverse
+ if use_dtcwt:
+ self._wavelet_forward = fwdfun(
+ J=level,
+ mode=mode,
+ biort=biort,
+ qshift=qshift,
+ )
+ self._wavelet_inverse = invfun(
+ mode=inv_mode,
+ biort=inv_biort,
+ qshift=inv_qshift,
+ )
+ else:
+ self._wavelet_forward = fwdfun(J=level, wave=wave, mode=mode)
+ self._wavelet_inverse = invfun(wave=inv_wave, mode=inv_mode)
+ self.device = device
+ if device is not None:
+ self._wavelet_forward = self._wavelet_forward.to(device=device)
+ self._wavelet_inverse = self._wavelet_inverse.to(device=device)
+
+ def forward(
+ self,
+ t: torch.Tensor,
+ *,
+ forward_function: Callable | None = None,
+ ) -> tuple[torch.Tensor, tuple[torch.Tensor, ...]]:
+ return fallback(forward_function, self._wavelet_forward)(t)
+
+ def inverse(
+ self,
+ yl: torch.Tensor,
+ yh: tuple[torch.Tensor, ...],
+ *,
+ inverse_function: Callable | None = None,
+ two_step_inverse: bool = False,
+ ) -> torch.Tensor:
+ inverse_function = fallback(inverse_function, self._wavelet_inverse)
+ if not two_step_inverse:
+ return inverse_function((yl, yh))
+ result = inverse_function((torch.zeros_like(yl), yh))
+ result += inverse_function(
+ (
+ yl,
+ tuple(torch.zeros_like(yh_band) for yh_band in yh),
+ )
+ )
+ return result
+
+ def to(self, *args: list, copy: bool = False, **kwargs: dict) -> Wavelet:
+ o = Wavelet.__new__(Wavelet) if copy else self
+ o._wavelet_forward = self._wavelet_forward.to(*args, **kwargs) # noqa: SLF001
+ o._wavelet_inverse = self._wavelet_inverse.to(*args, **kwargs) # noqa: SLF001
+ o.device = kwargs.get("device")
+ return o
+
+ @staticmethod
+ def wavelist() -> tuple:
+ return tuple(pywt.wavelist()) if HAVE_WAVELETS else ()
+
+ @staticmethod
+ def biortlist() -> tuple:
+ return (
+ ("near_sym_a", "near_sym_b", "antonini", "legall") if HAVE_WAVELETS else ()
+ )
+
+ @staticmethod
+ def qshiftlist() -> tuple:
+ return (
+ ("qshift_a", "qshift_b", "qshift_c", "qshift_d", "qshift_06")
+ if HAVE_WAVELETS
+ else ()
+ )
+
+ @staticmethod
+ def modelist() -> tuple:
+ return (
+ (
+ "symmetric",
+ "zero",
+ "reflect",
+ "replicate",
+ "periodization",
+ "periodic",
+ "constant",
+ )
+ if HAVE_WAVELETS
+ else ()
+ )
+
+
+def expand_yh_scales(
+ yh: Sequence,
+ *,
+ yh_scales: float | Sequence = 1.0,
+) -> float | tuple:
+ yhlen = len(yh)
+ yh_shape = yh[0].shape
+ # Doesn't make sense to target orientations for 1D DWD (3D here).
+ olen = yh_shape[2] if len(yh_shape) > 3 else 1
+ # print(f"\nSIZES: yhlen={yhlen}, olen={olen}, yh_shape={yh[0].shape}")
+ if isinstance(yh_scales, (float, int)):
+ return ((float(yh_scales),) * olen,) * yhlen
+ otemplate = (1.0,) * olen
+ yh_scales = tuple(
+ (float(band),) * olen
+ if isinstance(band, (float, int))
+ else (
+ (
+ *(float(i) for i in band[:olen]),
+ *otemplate[: olen - len(band[:olen])],
+ )
+ if isinstance(band, (tuple, list))
+ else band
+ )
+ for band in yh_scales
+ )
+ if "fill" in yh_scales:
+ fillidx = yh_scales.index("fill")
+ if "fill" in yh_scales[fillidx + 1 :]:
+ raise ValueError("Only one fill allowed.")
+ if fillidx == 0 or len(yh_scales) < 2:
+ raise ValueError(
+ "Invalid fill value, cannot be in the first position or the only item.",
+ )
+ yhslen = len(yh_scales)
+ if yhslen - 1 < yhlen:
+ # Need to pad.
+ fill = (yh_scales[fillidx - 1],) * (yhlen - (len(yh_scales) - 1))
+ yh_scales = (*yh_scales[:fillidx], *fill, *yh_scales[fillidx + 1 :])
+ else:
+ # Just remove the "fill".
+ yh_scales = (*yh_scales[:fillidx], *yh_scales[fillidx + 1 :])
+ return yh_scales[:yhlen]
+
+
+def wavelet_scaling(
+ yl: torch.Tensor,
+ yh: Sequence[torch.Tensor],
+ yl_scale: float | torch.Tensor,
+ yh_scales: float | Sequence[float | Sequence[float]] | None,
+ *,
+ in_place: bool = False,
+) -> tuple[torch.Tensor, tuple[torch.Tensor, ...]]:
+ if not in_place:
+ yl = yl.clone()
+ yh = tuple(yhband.clone() for yhband in yh)
+ if yl_scale != 1.0:
+ yl *= yl_scale
+ yh_scales = expand_yh_scales(
+ yh,
+ yh_scales=yh_scales if yh_scales is not None else 1.0,
+ )
+ for hscale, ht in zip(yh_scales, yh):
+ if isinstance(hscale, (int, float)):
+ ht *= hscale # noqa: PLW2901
+ continue
+ for lidx in range(min(ht.shape[2], len(hscale))):
+ ht[:, :, lidx] *= hscale[lidx]
+ return (yl, yh)
+
+
+def wavelet_blend(
+ a: tuple,
+ b: tuple,
+ *,
+ yl_factor: torch.Tensor | float,
+ blend_function: Callable,
+ yh_factor: torch.Tensor | float | None = None,
+ yh_blend_function: Callable | None = None,
+) -> tuple:
+ if not isinstance(yl_factor, torch.Tensor):
+ yl_factor = a[0].new_full((1,), yl_factor)
+ if yh_factor is None:
+ yh_factor = yl_factor
+ elif not isinstance(yh_factor, torch.Tensor):
+ yh_factor = a[0].new_full((1,), yh_factor)
+ yh_blend_function = fallback(yh_blend_function, blend_function)
+ return (
+ blend_function(a[0], b[0], yl_factor),
+ tuple(yh_blend_function(ta, tb, yh_factor) for ta, tb in zip(a[1], b[1])),
+ )
diff --git a/ruff.toml b/ruff.toml
index 3fafb71..c284f53 100644
--- a/ruff.toml
+++ b/ruff.toml
@@ -7,6 +7,7 @@ ignore = [
"ANN202",
"ANN204",
"ANN206",
+ "ANN401",
"C901",
"CPY001",
"D100",
@@ -16,8 +17,10 @@ ignore = [
"D104",
"D105",
"D107",
+ "D203",
"D211",
"D213",
+ "D401",
"E402",
"E501",
"EM101",