Bump version to 1.0.2

Add image noise nodes with torchvision fallback

README.md

@@ -1,151 +1,119 @@
-# ComfyUI Flow Matching Inverter Nodes
+# ComfyUI Qwen Rectified Flow Inverter
 
-Custom nodes for ComfyUI that target rectified-flow models (for example, Qwen). The nodes cover three core tasks:
+Custom nodes for ComfyUI that target rectified-flow models (for example, Qwen). This pack focuses on inversion and diagnostics; the latent/image/conditioning noise nodes live in `Skoogeer-Noise` (see `../Skoogeer-Noise`).
 
-- integrate the model's velocity field to produce a partially inverted latent,
-- add noise using the same schedule as the sampler, and
-- perform small latent-domain cleanups or conditioning perturbations.
+## Installation
 
-Each node is available on its own, and the hybrid node chains them together for a single-drop workflow.
+1. Clone this repository into `ComfyUI/custom_nodes/ComfyUI-QwenRectifiedFlowInverter`.
+2. (Optional but recommended) Also install `Skoogeer-Noise` into `ComfyUI/custom_nodes/Skoogeer-Noise` for the forward-diffusion/noise/conditioning helpers.
+3. Restart ComfyUI.
 
-This node pack also includes a variety of latent and conditioning space noise nodes enabling you to do really
-interesting and insane things in the latent and conditioning space.
+## Data types
 
-## Conditioning primer
+### `LATENT`
 
-ComfyUI passes prompt information around as a Python list of `[embedding, metadata]` pairs. The `embedding` tensor stores the per-token hidden states produced by the text encoder (shape ≈ `batch × tokens × features`, or `tokens × features` when the batch dimension is implicit). The accompanying `metadata` dictionary carries auxiliary fields—most notably `pooled_output`, a single vector that summarises the whole prompt (CLS token, sentence average, etc.).
+ComfyUI latents are dictionaries containing a `"samples"` tensor.
 
-Whenever you tweak conditioning with the nodes below you are directly manipulating those tensors. Splitting, blurring, or scaling happens along the token axis, while pooled outputs are updated in lockstep so downstream samplers still receive a coherent summary.
+- Common SD-style latents: `samples` has shape `(B, C, H, W)`.
+- Video / flow-matching latents may be 5D: `(B, C, T, H, W)`.
 
-## Node overview
+### `CONDITIONING`
+
+ComfyUI represents prompt conditioning as a list of `[embedding, metadata]` entries. This pack reads the embedding tensor from `positive[0][0]` and uses it as the model context during inversion.
+
+## Nodes
 
 ### Qwen Rectified Flow Inverter
 
-Runs the rectified-flow ODE backwards to re-noise an encoded latent. Supports deterministic amplification, seeded stochastic perturbation, optional output normalisation, and shares the same interface as the sampler (`steps`, `conditioning`, etc.).
+Runs the rectified-flow integration to produce a partially inverted (re-noised) latent. Supports deterministic velocity amplification, seeded stochastic perturbation, and optional output normalization.
 
-### Latent Gaussian Blur
+- **Menu category:** `Qwen/Sampling`
+- **Returns:** `LATENT`
 
-Applies a Gaussian blur directly in latent space. `Spatial Only` blurs each channel independently; `Spatial and Channel` performs a joint 3D blur across channels for stronger smoothing.
+#### Inputs
 
-### Latent Frequency Split
+| Field | Type | Default | Notes |
+|------|------|---------|------|
+| `model` | `MODEL` | – | Diffusion model / UNet used to predict velocity. |
+| `latent_image` | `LATENT` | – | Source latent to invert. |
+| `positive` | `CONDITIONING` | – | Positive prompt conditioning used as context. |
+| `seed` | `INT` | `0` | Seed for the velocity perturbation noise. |
+| `steps` | `INT` | `20` | Integration budget; effective steps = `int(steps * inversion_strength)`. |
+| `inversion_strength` | `FLOAT` | `0.5` | Fraction of the schedule to traverse (0 = no-op). |
+| `velocity_amplification` | `FLOAT` | `0.0` | Multiplies predicted velocity by `(1 + velocity_amplification)` each step. |
+| `velocity_perturb_strength` | `FLOAT` | `0.0` | Adds seeded noise to velocity (scaled by velocity std). |
+| `internal_precision` | enum | `bfloat16` | Autocast dtype while running the UNet (`force_float32` is most stable). |
+| `normalize_output` | enum | `enable` | When enabled, normalizes each batch item to mean 0 / std 1 at the end. |
 
-Separates a latent into low-pass structure and high-frequency detail bands using a Gaussian crossover so you can process each side differently before recombining.
+#### Notes
 
-### Add Latent Noise (Seeded)
+- If you see `NaN detected at step ...`, increase precision (`internal_precision = force_float32`) and/or reduce `velocity_amplification` / `velocity_perturb_strength`.
+- `normalize_output = enable` is strongly recommended at higher `inversion_strength`.
 
-Adds seeded Gaussian noise scaled by the input latent's standard deviation. Useful whenever a latent needs a controlled bump in noise without touching the schedule.
-
-### Latent Perlin Fractal Noise
-
-Generates smooth octave-based Perlin noise and adds it to the latent. Tune the base frequency, octaves, persistence, and lacunarity to decide whether you want broad undulations or fine stippling, and choose between shared or per-channel fields for subtle colour-channel offsets.
-
-### Latent Simplex Noise
-
-Produces layered simplex noise—Perlin's successor—for isotropic, artifact-free texture. Frequency, octaves, persistence, and lacunarity work just like the Perlin node, while `temporal_mode` lets you keep the same pattern across video frames (`locked`) or reseed each frame (`animated`).
-
-### Latent Worley Noise
-
-Synthesises cellular Worley (Voronoi) patterns by scattering seeded feature points. Adjust the point count, distance metric, and jitter to morph between cracked clay, reptilian scales, or bubbly foams. Combine with animated mode to introduce evolving cellular motion in video latents.
-
-### Latent Reaction-Diffusion
-
-Runs a Gray–Scott reaction-diffusion simulation in latent space to generate self-organising Turing patterns. Feed, kill, and diffusion controls expose the classic Gray–Scott parameter space; `temporal_mode` toggles between a single locked pattern and re-simulating per frame.
-
-### Latent Fractal Brownian Motion
-
-Wraps any of the base noise primitives (simplex, Perlin, or Worley) in an fBm stack to build rich multi-scale detail. Combine frequency, octaves, persistence, and lacunarity to decide how aggressively the layers accumulate, and pick Worley for cellular fBm, Perlin for smooth ridges, or simplex for organic grain.
-
-Under the hood the node keeps resampling the chosen base noise at higher frequency while shrinking the amplitude (classic fBm). Each octave is normalized before it is added, so persistence becomes a clean gain control instead of inheriting whatever variance the base noise happens to produce. Worley uses the `feature_points` count scaled by frequency to keep the cell density consistent as lacunarity climbs. Once all octaves are summed, the composite field is normalized again, multiplied by the latent's standard deviation, and finally scaled by `strength` before being added back. `channel_mode` decides whether every channel shares the same Brownian field or gets an independent reseeded stack, and `temporal_mode` controls whether video batches reuse the field (`locked`) or generate fresh motion on every frame (`animated`).
-
-### Latent Swirl Noise
-
-Warps the latent around one or more seeded vortices using `grid_sample`. Choose between a global grid (shared across channels) or per-channel grids, set the fraction of channels to affect, and control vortex count, rotation, falloff radius, centre jitter, and blend to inject painterly whirlpools or gentle twisting motion where you want it.
-
-### Conditioning (Add Noise)
-
-Adds seeded Gaussian noise to the CLIP conditioning embeddings (and pooled output when present). Great for introducing gentle prompt variation without rewriting text.
-
-Sensible values of `strength` range from 0.0 to 1.5. Above this range, the Qwen model's output tends to get silly, with Chinese characters overlaid on the image.
-
-### Conditioning (Gaussian Blur)
-
-Smooths the token embeddings with a Gaussian kernel along the prompt sequence, softening sharp emphasis changes while preserving the overall prompt content.
-
-### Conditioning (Frequency Split)
-
-Generates low- and high-frequency prompt embeddings so you can tame the overall narrative while separately shaping high-energy emphasis tokens.
-
-Sensible values of `sigma` range from 0.1 to 1.3. The higher you raise `sigma`, the more of the conditioning information that will be shoved into the `low_pass` conditioning output. Values well above 1.0 are really equivalent to knocking out the high pass signal entirely. But strange and interesting results can be obtained by using a very high sigma in combination with a very high value of `gain` for the `high_pass` conditioning in the `Conditioning (Frequency Merge)` node.
-
-### Conditioning (Frequency Merge)
-
-Recombines low/high conditioning bands with adjustable gains so you can dial detail back in after sculpting each band independently. Boosting the low band reinforces the prompt's broad narrative and stability, while boosting the high band pushes sharp emphasis changes and punctuation to the forefront.
-
-Feed the output of the `Conditioning (Frequency Split)` node into the input of this node and then play around with the gain knobs for interesting effects. Or peel off the high or low-pass conditioning and then add noise to one or the other, or combine the high pass conditionings of two different `CLIP Text Encode` outputs... the combinations are endless.
-
-### Conditioning (Scale)
-
-Multiplies the conditioning embeddings and pooled outputs by a user-defined factor so you can mute (0.0), keep (1.0), or amplify (>1.0) prompt influence without editing text.
-
-### Forward Diffusion (Add Scheduled Noise)
-
-Uses the model's noise schedule (via `KSampler`) to add the amount of noise that corresponds to a given progress value. This is the "anchor" latent that samplers expect when starting from a given denoise level.
+---
 
 ### Latent Hybrid Inverter (Qwen)
 
-Convenience node that calls the inverter and forward diffusion nodes internally, then blends their outputs using spherical interpolation. It keeps separate seeds for the creative and anchor paths and exposes the blend weight so you can bias towards either latent.
+Runs the inverter and an internal forward-diffusion anchor pass, then blends both latents via spherical interpolation (SLERP).
 
-## Creative play
+- **Menu category:** `Qwen/Sampling`
+- **Returns:** `LATENT`
 
-Image exploration rarely means chasing a single "correct" answer. These nodes are designed to be mixed and layered so you can steer a latent in different, often surprising, directions:
+#### Inputs
 
-- Lead with `Latent Hybrid Inverter` to recover structure, then nudge that structure by adding `Conditioning (Add Noise)` at a subtle strength (≈0.05–0.15). The latent stays coherent while the prompt wanders just enough to spark new ideas.
-- Follow the noise with `Conditioning (Gaussian Blur)` to soften abrupt emphasis changes or to merge multi-prompt blends into a single vibe. Blurring after noise often produces dreamlike, painterly shifts instead of chaotic drift.
-- Use `Conditioning (Scale)` as a volume knob while iterating: dial the factor down to 0.3–0.5 when you want the image to respond mostly to the recovered latent, then crank it above 1.5 when the textual guidance should take the lead.
-- Pair `Latent Gaussian Blur` with `Conditioning (Gaussian Blur)` for holistic smoothing—latent blur calms texture while conditioning blur calms prompt pacing. Reintroduce energy by sprinkling in `Add Latent Noise` or ramping the conditioning scale afterwards.
-- Sculpt structured detail by layering `Latent Perlin Fractal Noise` (shared mode for broad waves, per-channel for colour shifts) before blending back with `Latent Swirl Noise`—increase `vortices` for multi-spiral motion, switch to per-channel grids, or drop `channel_fraction` to swirl only a handful of feature maps.
-- Split a latent with `Latent Frequency Split`, sharpen or noise the high band while blurring the low band, then recombine via `Latent Mixer`/`Add` nodes to weave sharp detail onto soft composition.
-- Carve prompts with `Conditioning (Frequency Split)` and return them through `Conditioning (Frequency Merge)`: keep low-band text steady for the story while remapping the high band through noise, blur, or scale to emphasize only certain fragments.
-- Swap seeds between the inverter, forward diffusion, and noise nodes as you iterate. Small tweaks here can shift the interplay between latent detail and prompt guidance, rewarding playful experimentation.
-- Chase naturalistic textures by stacking `Latent Simplex Noise` or the fBm node before sampling; switch `temporal_mode` to `animated` when generating video to keep organic flicker under control. For harder cellular structures, blend in `Latent Worley Noise`, and when you need life-like stripes or spots, seed `Latent Reaction-Diffusion` with slightly different feed/kill rates to explore the Gray–Scott pattern atlas.
+| Field | Type | Default | Notes |
+|------|------|---------|------|
+| `model` | `MODEL` | – | Model used for both passes. |
+| `latent_image` | `LATENT` | – | Source latent. |
+| `positive` | `CONDITIONING` | – | Context used by the inverter pass. |
+| `steps` | `INT` | `20` | Schedule/integration steps. |
+| `strength` | `FLOAT` | `0.5` | Target noise level (also used as the forward-diffusion strength). |
+| `blend_factor` | `FLOAT` | `0.5` | `0` = 100% forward diffusion (stable), `1` = 100% inverter (creative). |
+| `inverter_seed` | `INT` | `0` | Seed for the inverter’s velocity perturbation. |
+| `forward_diffusion_seed` | `INT` | `1` | Seed for the forward-diffusion anchor noise. |
+| `velocity_amplification` | `FLOAT` | `1.0` | Passed through to the inverter node. |
+| `velocity_perturb_strength` | `FLOAT` | `0.2` | Passed through to the inverter node. |
+| `internal_precision` | enum | `bfloat16` | Passed through to the inverter node. |
+| `normalize_output` | enum | `enable` | Passed through to the inverter node. |
 
-Treat each slider as a brushstroke: push a setting until it breaks, back off to the sweet spot, and capture the happy accidents along the way.
+---
+
+### Memory Diagnostics (Pass-Through)
+
+Passes any input through unchanged while printing basic GPU memory statistics to the console.
+
+- **Menu category:** `Qwen/Diagnostics`
+- **Returns:** same as input
+
+#### Inputs
+
+| Field | Type | Default | Notes |
+|------|------|---------|------|
+| `data` | any | – | Input payload (any ComfyUI type). |
+| `label` | `STRING` | `""` | Optional label included in the printed line. |
+| `synchronize` | enum | `enable` | When enabled, synchronizes GPU kernels before sampling memory. |
 
 ## Example workflows
 
-### Manual composition
+### Manual composition (recommended with Skoogeer-Noise)
 
 1. Encode an image to a latent.
-2. Feed the latent into `Qwen Rectified Flow Inverter` (pick strength/amplification/perturbation).
-3. Feed the clean latent into `Forward Diffusion` with the same strength but a different seed.
-4. Optionally clean either latent with `Latent Gaussian Blur` or `Add Latent Noise`.
-5. Blend the creative and anchor latents (e.g. `Latent Slerp`) and sample with `denoise` equal to the chosen strength.
+2. Run `Qwen Rectified Flow Inverter`.
+3. In `Skoogeer-Noise`, run `Forward Diffusion (Add Scheduled Noise)` with the same `steps` and `noise_strength` but a different seed (this gives an anchor latent samplers expect).
+4. Blend the two latents (e.g. `Latent Slerp`) and sample with `denoise` matching the strength you used.
 
 ### Hybrid node quickstart
 
 1. Encode an image to a latent.
-2. Run the latent through `Latent Hybrid Inverter` with your preferred strength and blend factor.
-3. Sample the result with `denoise` matching the passed strength.
-
-### Combination example
-
-Check `Examples/Combination Example.png` for a full workflow that combines the inverter, simplex and Worley noise, reaction-diffusion, and fBm nodes. The diagram illustrates how to stage noise injection around the hybrid inverter and how the new `temporal_mode`/`channel_mode` options keep video latents coherent while still layering rich structure.
-
-## Installation
-
-1. Clone or download this repository into `ComfyUI/custom_nodes/`.
-2. Restart ComfyUI to load the nodes (look under `Qwen/Sampling` and `Latent`).
+2. Run `Latent Hybrid Inverter (Qwen)`.
+3. Sample the output latent with `denoise = strength`.
 
 ## Running the tests
 
-The repository ships with a helper script that prepares a virtual environment and
-executes the pytest suite. From the project root run:
+Run:
 
 ```bash
 ./run_tests.sh
 ```
 
-The script creates (or reuses) a `.venv` folder, installs the dependencies
-from `requirements.txt`, and then launches `pytest`. Pass additional arguments
-to forward them to pytest—for example `./run_tests.sh -k noise` runs only the
-tests whose names include `noise`.
+The script creates (or reuses) a `.venv` folder, installs the dependencies from `requirements.txt`, and launches `pytest`.

pyproject.toml

@@ -1,7 +1,7 @@
 [project]
 name = "flowmatching-inverter"
-description = "A set of nodes for reversing the sampling process in flow matching models such as Qwen, plus various latent and conditioning noise nodes"
-version = "1.0.0"
+description = "A set of nodes for reversing the sampling process in flow matching models such as Qwen"
+version = "1.0.2"
 license = { text = "MIT License" }
 classifiers = [
     "Operating System :: OS Independent",
@@ -20,4 +20,3 @@ DisplayName = "ComfyUI-FlowMatching-Inverter"
 Icon = "https://github.com/ttulttul/ComfyUI-FlowMatching-Inverter/raw/refs/heads/main/flowmatching-inverter-icon.png"
 includes = [] 
 # "requires-comfyui" = ">=1.0.0"  # ComfyUI version compatibility
-

requirements.txt

@@ -2,5 +2,4 @@
 # are provided by ComfyUI natively when the node is used within
 # ComfyUI's environment.
 pytest>=8.0.0,<9.0.0
-torch>=2.1.0,<2.3.0
-torchvision>=0.17.0,<0.18.0
+torch>=2.1.0

run_tests.sh

@@ -20,4 +20,4 @@ else
     "${PIP_BIN}" install -r "${ROOT_DIR}/requirements.txt"
 fi
 
-"${PYTHON_BIN}" -m pytest "$@"
+"${PYTHON_BIN}" -m pytest "$@"

setup_venv.sh

@@ -0,0 +1,41 @@
+#!/bin/bash
+
+# The name of the virtual environment directory
+VENV_DIR="comfy-env"
+
+# The Python version to use
+PYTHON_VERSION="3.13"
+
+# --- Script Start ---
+
+echo "Checking for Python ${PYTHON_VERSION}..."
+
+# Attempt to find the python3.13 executable
+PYTHON_EXEC=$(command -v "python${PYTHON_VERSION}")
+
+if [ -z "$PYTHON_EXEC" ]; then
+    echo "Error: python${PYTHON_VERSION} not found in your PATH." >&2
+    echo "Please install Python ${PYTHON_VERSION} or make sure it is in your system's PATH." >&2
+    exit 1
+fi
+
+echo "Found Python at: ${PYTHON_EXEC}"
+
+# Check if the virtual environment directory already exists
+if [ -d "$VENV_DIR" ]; then
+    echo "Virtual environment '${VENV_DIR}' already exists."
+else
+    echo "Creating virtual environment '${VENV_DIR}' with ${PYTHON_EXEC}..."
+    # Create the virtual environment using the specified Python version
+    "$PYTHON_EXEC" -m venv "$VENV_DIR"
+    if [ $? -ne 0 ]; then
+        echo "Error: Failed to create the virtual environment." >&2
+        exit 1
+    fi
+    echo "Virtual environment created successfully."
+fi
+
+# Activate the virtual environment
+echo "Activating virtual environment: ${VENV_DIR}"
+echo "To activate, run 'source "${VENV_DIR}/bin/activate"' in your terminal."
+echo "To deactivate, simply type 'deactivate' in your terminal."

tests/conftest.py

@@ -25,6 +25,19 @@ def _ensure_comfy_stubs():
     utils_module = types.ModuleType("comfy.utils")
     utils_module.ProgressBar = _DummyProgressBar
 
+    class _AnyType(str):
+        def __new__(cls, name="*"):
+            return str.__new__(cls, name)
+
+        def __ne__(self, other):
+            return False
+
+    def _any_type(name="*"):
+        return _AnyType(name)
+
+    utils_module.AnyType = _AnyType
+    utils_module.any_type = _any_type
+
     samplers_module = types.ModuleType("comfy.samplers")
 
     class _DummyKSampler:

tests/test_nodes.py

@@ -6,11 +6,6 @@ import torch
 
 MODULE_PATH = Path(__file__).resolve().parents[1] / "__init__.py"
 
-FLOW_MATCHING_CHANNELS = 16
-FLOW_MATCHING_TEMPORAL = 1
-FLOW_MATCHING_WIDTH = 8
-FLOW_MATCHING_HEIGHT = 8
-
 
 def load_module():
     spec = importlib.util.spec_from_file_location("flowmatching_inverter", MODULE_PATH)
@@ -19,105 +14,31 @@ def load_module():
     return module
 
 
-def make_flow_matching_latent(batch_size=1, width=FLOW_MATCHING_WIDTH, height=FLOW_MATCHING_HEIGHT):
-    samples = torch.randn(batch_size, FLOW_MATCHING_CHANNELS, FLOW_MATCHING_TEMPORAL, width, height)
-    return {"samples": samples}
-
-
-def assert_flow_matching_shape(tensor, batch_size=1, width=FLOW_MATCHING_WIDTH, height=FLOW_MATCHING_HEIGHT):
-    expected = (batch_size, FLOW_MATCHING_CHANNELS, FLOW_MATCHING_TEMPORAL, width, height)
-    assert tensor.shape == expected
-
-
-def test_latent_gaussian_blur_modifies_values():
+def test_node_mappings_only_expose_inverter_nodes():
     module = load_module()
-    node = module.LatentGaussianBlur()
-    latent = make_flow_matching_latent()
 
-    (result,) = node.blur_latent(latent, sigma=1.5, blur_mode="Spatial Only")
+    assert "QwenRectifiedFlowInverter" in module.NODE_CLASS_MAPPINGS
+    assert "LatentHybridInverter" in module.NODE_CLASS_MAPPINGS
+    assert "MemoryDiagnosticsPassThrough" in module.NODE_CLASS_MAPPINGS
 
-    original = latent["samples"]
-    blurred = result["samples"]
-
-    assert_flow_matching_shape(original)
-    assert_flow_matching_shape(blurred)
-    assert not torch.allclose(blurred, original)
+    # moved to Skoogeer-Noise
+    assert "LatentGaussianBlur" not in module.NODE_CLASS_MAPPINGS
+    assert "LatentAddNoise" not in module.NODE_CLASS_MAPPINGS
+    assert "ConditioningAddNoise" not in module.NODE_CLASS_MAPPINGS
 
 
-def test_latent_add_noise_reproducible_with_seed():
+def test_memory_diagnostics_pass_through_logs_and_returns_input(capsys):
     module = load_module()
-    node = module.LatentAddNoise()
-    latent = make_flow_matching_latent()
+    node = module.MemoryDiagnosticsPassThrough()
 
-    (first,) = node.add_noise(latent, seed=123, strength=0.8)
-    (second,) = node.add_noise(latent, seed=123, strength=0.8)
-    (third,) = node.add_noise(latent, seed=321, strength=0.8)
+    payload = {"samples": torch.zeros(1)}
 
-    assert_flow_matching_shape(first["samples"])
-    assert_flow_matching_shape(second["samples"])
-    assert_flow_matching_shape(third["samples"])
-    assert torch.allclose(first["samples"], second["samples"])
-    assert not torch.allclose(first["samples"], third["samples"])
+    (result,) = node.log_memory(payload, label="unit", synchronize="disable")
 
+    assert result is payload
 
-def test_perlin_noise_strength_affects_latent():
-    module = load_module()
-    node = module.LatentPerlinFractalNoise()
-    latent = make_flow_matching_latent()
+    output = capsys.readouterr().out
 
-    (result,) = node.add_perlin_noise(
-        latent,
-        seed=7,
-        frequency=1.5,
-        octaves=2,
-        persistence=0.5,
-        lacunarity=2.0,
-        strength=0.5,
-        channel_mode="shared",
-    )
-
-    assert_flow_matching_shape(result["samples"])
-    assert not torch.allclose(result["samples"], latent["samples"])
-
-
-def test_simplex_noise_zero_strength_is_noop():
-    module = load_module()
-    node = module.LatentSimplexNoise()
-    latent = make_flow_matching_latent()
-
-    (result,) = node.add_simplex_noise(
-        latent,
-        seed=42,
-        frequency=1.0,
-        octaves=2,
-        persistence=0.5,
-        lacunarity=2.0,
-        strength=0.0,
-        channel_mode="shared",
-        temporal_mode="locked",
-    )
-
-    assert_flow_matching_shape(result["samples"])
-    assert torch.allclose(result["samples"], latent["samples"])
-
-
-def test_conditioning_nodes_modify_embeddings_and_metadata():
-    module = load_module()
-    noise_node = module.ConditioningAddNoise()
-    blur_node = module.ConditioningGaussianBlur()
-
-    embedding = torch.randn(2, 4, 6)
-    pooled = torch.randn(4)
-    conditioning = [[embedding.clone(), {"pooled_output": pooled.clone()}]]
-
-    (noised,) = noise_node.add_noise(conditioning, seed=0, strength=0.5)
-    (blurred,) = blur_node.blur(noised, sigma=1.0)
-
-    noised_embedding = noised[0][0]
-    noised_pooled = noised[0][1]["pooled_output"]
-
-    assert not torch.allclose(noised_embedding, embedding)
-    assert not torch.allclose(noised_pooled, pooled)
-
-    blurred_embedding = blurred[0][0]
-    assert blurred_embedding.shape == embedding.shape
+    assert "[MemoryDiagnostics]" in output
+    assert "label=unit" in output
+    assert "device=" in output