Error while processing rearrange-reduction pattern "B (T S) N C -> (B S) N T C". Input tensor shape: torch.Size([1, 34320, 40, 128]). Additional info: {'T': 21, 'S': 1634, 'N': 40, 'C': 128}. Shape mismatch, 34320 != 34314

### Custom Node Testing

- [x] I have tried disabling custom nodes and the issue persists (see [how to disable custom nodes](https://docs.comfy.org/troubleshooting/custom-node-issues#step-1%3A-test-with-all-custom-nodes-disabled) if you need help)

### Expected Behavior

# ComfyUI Error Report
## Error Details
- **Node ID:** 149
- **Node Type:** KSampler
- **Exception Type:** einops.EinopsError
- **Exception Message:** Error while processing rearrange-reduction pattern "B (T S) N C -> (B S) N T C".
Input tensor shape: torch.Size([1, 34320, 40, 128]). Additional info: {'T': 21, 'S': 1634, 'N': 40, 'C': 128}.
Shape mismatch, 34320 != 34314

## Stack Trace
```
File "D:\ComfyUI\execution.py", line 361, in execute
output_data, output_ui, has_subgraph = get_output_data(obj, input_data_all, execution_block_cb=execution_block_cb, pre_execute_cb=pre_execute_cb)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^

File "D:\ComfyUI\execution.py", line 236, in get_output_data
return_values = _map_node_over_list(obj, input_data_all, obj.FUNCTION, allow_interrupt=True, execution_block_cb=execution_block_cb, pre_execute_cb=pre_execute_cb)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^

File "D:\ComfyUI\execution.py", line 208, in _map_node_over_list
process_inputs(input_dict, i)

File "D:\ComfyUI\execution.py", line 197, in process_inputs
results.append(getattr(obj, func)(**inputs))
^^^^^^^^^^^^^^^^^^^^^^^^^^^^

File "D:\ComfyUI\nodes.py", line 1516, in sample
return common_ksampler(model, seed, steps, cfg, sampler_name, scheduler, positive, negative, latent_image, denoise=denoise)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^

File "D:\ComfyUI\nodes.py", line 1483, in common_ksampler
samples = comfy.sample.sample(model, noise, steps, cfg, sampler_name, scheduler, positive, negative, latent_image,
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^

File "D:\ComfyUI\comfy\sample.py", line 45, in sample
samples = sampler.sample(noise, positive, negative, cfg=cfg, latent_image=latent_image, start_step=start_step, last_step=last_step, force_full_denoise=force_full_denoise, denoise_mask=noise_mask, sigmas=sigmas, callback=callback, disable_pbar=disable_pbar, seed=seed)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^

File "D:\ComfyUI\comfy\samplers.py", line 1143, in sample
return sample(self.model, noise, positive, negative, cfg, self.device, sampler, sigmas, self.model_options, latent_image=latent_image, denoise_mask=denoise_mask, callback=callback, disable_pbar=disable_pbar, seed=seed)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^

File "D:\ComfyUI\comfy\samplers.py", line 1033, in sample
return cfg_guider.sample(noise, latent_image, sampler, sigmas, denoise_mask, callback, disable_pbar, seed)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^

File "D:\ComfyUI\comfy\samplers.py", line 1018, in sample
output = executor.execute(noise, latent_image, sampler, sigmas, denoise_mask, callback, disable_pbar, seed)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^

File "D:\ComfyUI\comfy\patcher_extension.py", line 111, in execute
return self.original(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^

File "D:\ComfyUI\comfy\samplers.py", line 986, in outer_sample
output = self.inner_sample(noise, latent_image, device, sampler, sigmas, denoise_mask, callback, disable_pbar, seed)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^

File "D:\ComfyUI\comfy\samplers.py", line 969, in inner_sample
samples = executor.execute(self, sigmas, extra_args, callback, noise, latent_image, denoise_mask, disable_pbar)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^

File "D:\ComfyUI\comfy\patcher_extension.py", line 111, in execute
return self.original(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^

File "D:\ComfyUI\comfy\samplers.py", line 748, in sample
samples = self.sampler_function(model_k, noise, sigmas, extra_args=extra_args, callback=k_callback, disable=disable_pbar, **self.extra_options)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^

File "D:\ComfyUI\comfy\extra_samplers\uni_pc.py", line 868, in sample_unipc
x = uni_pc.sample(noise, timesteps=timesteps, skip_type="time_uniform", method="multistep", order=order, lower_order_final=True, callback=callback, disable_pbar=disable)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^

File "D:\ComfyUI\comfy\extra_samplers\uni_pc.py", line 715, in sample
model_prev_list = [self.model_fn(x, vec_t)]
^^^^^^^^^^^^^^^^^^^^^^^

File "D:\ComfyUI\comfy\extra_samplers\uni_pc.py", line 410, in model_fn
return self.data_prediction_fn(x, t)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^

File "D:\ComfyUI\comfy\extra_samplers\uni_pc.py", line 394, in data_prediction_fn
noise = self.noise_prediction_fn(x, t)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^

File "D:\ComfyUI\comfy\extra_samplers\uni_pc.py", line 388, in noise_prediction_fn
return self.model(x, t)
^^^^^^^^^^^^^^^^

File "D:\ComfyUI\comfy\extra_samplers\uni_pc.py", line 329, in model_fn
return noise_pred_fn(x, t_continuous)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^

File "D:\ComfyUI\comfy\extra_samplers\uni_pc.py", line 297, in noise_pred_fn
output = model(x, t_input, **model_kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^

File "D:\ComfyUI\comfy\extra_samplers\uni_pc.py", line 859, in
lambda input, sigma, **kwargs: predict_eps_sigma(model, input, sigma, **kwargs),
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^

File "D:\ComfyUI\comfy\extra_samplers\uni_pc.py", line 843, in predict_eps_sigma
return (input - model(input, sigma_in, **kwargs)) / sigma
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^

File "D:\ComfyUI\comfy\samplers.py", line 400, in __call__
out = self.inner_model(x, sigma, model_options=model_options, seed=seed)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^

File "D:\ComfyUI\comfy\samplers.py", line 949, in __call__
return self.predict_noise(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^

File "D:\ComfyUI\comfy\samplers.py", line 952, in predict_noise
return sampling_function(self.inner_model, x, timestep, self.conds.get("negative", None), self.conds.get("positive", None), self.cfg, model_options=model_options, seed=seed)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^

File "D:\ComfyUI\comfy\samplers.py", line 380, in sampling_function
out = calc_cond_batch(model, conds, x, timestep, model_options)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^

File "D:\ComfyUI\comfy\samplers.py", line 206, in calc_cond_batch
return executor.execute(model, conds, x_in, timestep, model_options)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^

File "D:\ComfyUI\comfy\patcher_extension.py", line 111, in execute
return self.original(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^

File "D:\ComfyUI\comfy\samplers.py", line 325, in _calc_cond_batch
output = model.apply_model(input_x, timestep_, **c).chunk(batch_chunks)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^

File "D:\ComfyUI\comfy\model_base.py", line 152, in apply_model
return comfy.patcher_extension.WrapperExecutor.new_class_executor(
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^

File "D:\ComfyUI\comfy\patcher_extension.py", line 111, in execute
return self.original(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^

File "D:\ComfyUI\comfy\model_base.py", line 190, in _apply_model
model_output = self.diffusion_model(xc, t, context=context, control=control, transformer_options=transformer_options, **extra_conds).float()
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^

File "D:\ComfyUI\venv\Lib\site-packages\torch\nn\modules\module.py", line 1751, in _wrapped_call_impl
return self._call_impl(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^

File "D:\ComfyUI\venv\Lib\site-packages\torch\nn\modules\module.py", line 1762, in _call_impl
return forward_call(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^

File "D:\ComfyUI\comfy\ldm\wan\model.py", line 563, in forward
return self.forward_orig(x, timestep, context, clip_fea=clip_fea, freqs=freqs, transformer_options=transformer_options, **kwargs)[:, :, :t, :h, :w]
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^

File "D:\ComfyUI\comfy\ldm\wan\model.py", line 533, in forward_orig
x = block(x, e=e0, freqs=freqs, context=context, context_img_len=context_img_len)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^

File "D:\ComfyUI\venv\Lib\site-packages\torch\nn\modules\module.py", line 1751, in _wrapped_call_impl
return self._call_impl(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^

File "D:\ComfyUI\venv\Lib\site-packages\torch\nn\modules\module.py", line 1762, in _call_impl
return forward_call(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^

File "D:\ComfyUI\comfy\ldm\wan\model.py", line 209, in forward
y = self.self_attn(
^^^^^^^^^^^^^^^

File "D:\ComfyUI\venv\Lib\site-packages\torch\nn\modules\module.py", line 1751, in _wrapped_call_impl
return self._call_impl(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^

File "D:\ComfyUI\venv\Lib\site-packages\torch\nn\modules\module.py", line 1762, in _call_impl
return forward_call(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^

File "D:\ComfyUI\custom_nodes\ComfyUI-KJNodes\nodes\model_optimization_nodes.py", line 1410, in wrapped_attention
return modified_wan_self_attention_forward(self_module, *args, **kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^

File "D:\ComfyUI\custom_nodes\ComfyUI-KJNodes\nodes\model_optimization_nodes.py", line 1341, in modified_wan_self_attention_forward
feta_scores = get_feta_scores(q, k, self.num_frames, self.enhance_weight)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^

File "D:\ComfyUI\custom_nodes\ComfyUI-KJNodes\nodes\model_optimization_nodes.py", line 1364, in get_feta_scores
query_image = rearrange(
^^^^^^^^^^

File "D:\ComfyUI\venv\Lib\site-packages\einops\einops.py", line 600, in rearrange
return reduce(tensor, pattern, reduction="rearrange", **axes_lengths)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^

File "D:\ComfyUI\venv\Lib\site-packages\einops\einops.py", line 542, in reduce
raise EinopsError(message + "\n {}".format(e))

```
## System Information
- **ComfyUI Version:** 0.3.43
- **Arguments:** D:\ComfyUI\main.py --auto-launch --preview-method auto --use-sage-attention --disable-cuda-malloc --fast
- **OS:** nt
- **Python Version:** 3.12.10 (tags/v3.12.10:0cc8128, Apr 8 2025, 12:21:36) [MSC v.1943 64 bit (AMD64)]
- **Embedded Python:** false
- **PyTorch Version:** 2.7.0+cu128
## Devices

- **Name:** cuda:0 NVIDIA GeForce RTX 4070 Ti : cudaMallocAsync
- **Type:** cuda
- **VRAM Total:** 12878086144
- **VRAM Free:** 8001112312
- **Torch VRAM Total:** 3623878656
- **Torch VRAM Free:** 124209400

### Actual Behavior

...

### Steps to Reproduce

### Debug Logs

```powershell
...
```

### Other

_No response_