ComfyUI 0.21.1 + HiDream O1 Error

### Custom Node Testing

- [ ] 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

I can not run workflow from Kijai PR.

### Actual Behavior

.

### Steps to Reproduce

.

### Debug Logs

```powershell
------------------------------------------------------
!!! Exception during processing !!! Error while processing rearrange-reduction pattern "B C (H p1) (W p2) -> B (H W) (C p1 p2)".
Input tensor shape: torch.Size([1, 3, 2728, 1536]). Additional info: {'p1': 32, 'p2': 32}.
Shape mismatch, can't divide axis of length 2728 in chunks of 32
Traceback (most recent call last):
File "D:\_ComfyUI_windows_portable\ComfyUI\execution.py", line 535, in execute
output_data, output_ui, has_subgraph, has_pending_tasks = await get_output_data(prompt_id, unique_id, obj, input_data_all, execution_block_cb=execution_block_cb, pre_execute_cb=pre_execute_cb, v3_data=v3_data)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "D:\_ComfyUI_windows_portable\ComfyUI\execution.py", line 335, in get_output_data
return_values = await _async_map_node_over_list(prompt_id, unique_id, obj, input_data_all, obj.FUNCTION, allow_interrupt=True, execution_block_cb=execution_block_cb, pre_execute_cb=pre_execute_cb, v3_data=v3_data)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "D:\_ComfyUI_windows_portable\ComfyUI\execution.py", line 309, in _async_map_node_over_list
await process_inputs(input_dict, i)
File "D:\_ComfyUI_windows_portable\ComfyUI\execution.py", line 297, in process_inputs
result = f(**inputs)
File "D:\_ComfyUI_windows_portable\ComfyUI\comfy_api\internal\__init__.py", line 149, in wrapped_func
return method(locked_class, **inputs)
File "D:\_ComfyUI_windows_portable\ComfyUI\comfy_api\latest\_io.py", line 1833, in EXECUTE_NORMALIZED
to_return = cls.execute(*args, **kwargs)
File "D:\_ComfyUI_windows_portable\ComfyUI\comfy_extras\nodes_custom_sampler.py", line 769, in execute
samples = comfy.sample.sample_custom(model, noise, cfg, sampler, sigmas, positive, negative, latent_image, noise_mask=noise_mask, callback=callback, disable_pbar=disable_pbar, seed=noise_seed)
File "D:\_ComfyUI_windows_portable\ComfyUI\comfy\sample.py", line 71, in sample_custom
samples = comfy.samplers.sample(model, noise, positive, negative, cfg, model.load_device, sampler, sigmas, model_options=model.model_options, latent_image=latent_image, denoise_mask=noise_mask, callback=callback, disable_pbar=disable_pbar, seed=seed)
File "D:\_ComfyUI_windows_portable\ComfyUI\comfy\samplers.py", line 1070, in sample
return cfg_guider.sample(noise, latent_image, sampler, sigmas, denoise_mask, callback, disable_pbar, seed)
~~~~~~~~~~~~~~~~~^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "D:\_ComfyUI_windows_portable\ComfyUI\comfy\samplers.py", line 1052, in sample
output = executor.execute(noise, latent_image, sampler, sigmas, denoise_mask, callback, disable_pbar, seed, latent_shapes=latent_shapes)
File "D:\_ComfyUI_windows_portable\ComfyUI\comfy\patcher_extension.py", line 112, in execute
return self.original(*args, **kwargs)
~~~~~~~~~~~~~^^^^^^^^^^^^^^^^^
File "D:\_ComfyUI_windows_portable\ComfyUI\comfy\samplers.py", line 995, in outer_sample
output = self.inner_sample(noise, latent_image, device, sampler, sigmas, denoise_mask, callback, disable_pbar, seed, latent_shapes=latent_shapes)
File "D:\_ComfyUI_windows_portable\ComfyUI\comfy\samplers.py", line 981, in inner_sample
samples = executor.execute(self, sigmas, extra_args, callback, noise, latent_image, denoise_mask, disable_pbar)
File "D:\_ComfyUI_windows_portable\ComfyUI\comfy\patcher_extension.py", line 112, in execute
return self.original(*args, **kwargs)
~~~~~~~~~~~~~^^^^^^^^^^^^^^^^^
File "D:\_ComfyUI_windows_portable\ComfyUI\comfy\samplers.py", line 751, 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_windows_portable\python_embeded\Lib\site-packages\torch\utils\_contextlib.py", line 124, in decorate_context
return func(*args, **kwargs)
File "D:\_ComfyUI_windows_portable\ComfyUI\comfy\k_diffusion\sampling.py", line 971, in sample_dpmpp_2m_sde_gpu
return sample_dpmpp_2m_sde(model, x, sigmas, extra_args=extra_args, callback=callback, disable=disable, eta=eta, s_noise=s_noise, noise_sampler=noise_sampler, solver_type=solver_type)
File "D:\_ComfyUI_windows_portable\python_embeded\Lib\site-packages\torch\utils\_contextlib.py", line 124, in decorate_context
return func(*args, **kwargs)
File "D:\_ComfyUI_windows_portable\ComfyUI\comfy\k_diffusion\sampling.py", line 845, in sample_dpmpp_2m_sde
denoised = model(x, sigmas[i] * s_in, **extra_args)
File "D:\_ComfyUI_windows_portable\ComfyUI\comfy\samplers.py", line 400, in __call__
out = self.inner_model(x, sigma, model_options=model_options, seed=seed)
File "D:\_ComfyUI_windows_portable\ComfyUI\comfy\samplers.py", line 954, in __call__
return self.outer_predict_noise(*args, **kwargs)
~~~~~~~~~~~~~~~~~~~~~~~~^^^^^^^^^^^^^^^^^
File "D:\_ComfyUI_windows_portable\ComfyUI\comfy\samplers.py", line 961, in outer_predict_noise
).execute(x, timestep, model_options, seed)
~~~~~~~^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "D:\_ComfyUI_windows_portable\ComfyUI\comfy\patcher_extension.py", line 112, in execute
return self.original(*args, **kwargs)
~~~~~~~~~~~~~^^^^^^^^^^^^^^^^^
File "D:\_ComfyUI_windows_portable\ComfyUI\comfy\samplers.py", line 964, 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_windows_portable\ComfyUI\comfy\samplers.py", line 380, in sampling_function
out = calc_cond_batch(model, conds, x, timestep, model_options)
File "D:\_ComfyUI_windows_portable\ComfyUI\comfy\samplers.py", line 205, in calc_cond_batch
return _calc_cond_batch_outer(model, conds, x_in, timestep, model_options)
File "D:\_ComfyUI_windows_portable\ComfyUI\comfy\samplers.py", line 213, in _calc_cond_batch_outer
return executor.execute(model, conds, x_in, timestep, model_options)
~~~~~~~~~~~~~~~~^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "D:\_ComfyUI_windows_portable\ComfyUI\comfy\patcher_extension.py", line 112, in execute
return self.original(*args, **kwargs)
~~~~~~~~~~~~~^^^^^^^^^^^^^^^^^
File "D:\_ComfyUI_windows_portable\ComfyUI\comfy\samplers.py", line 325, in _calc_cond_batch
output = model.apply_model(input_x, timestep_, **c).chunk(batch_chunks)
~~~~~~~~~~~~~~~~~^^^^^^^^^^^^^^^^^^^^^^^^^
File "D:\_ComfyUI_windows_portable\ComfyUI\comfy\model_base.py", line 182, in apply_model
return comfy.patcher_extension.WrapperExecutor.new_class_executor(
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
...<2 lines>...
comfy.patcher_extension.get_all_wrappers(comfy.patcher_extension.WrappersMP.APPLY_MODEL, transformer_options)
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
).execute(x, t, c_concat, c_crossattn, control, transformer_options, **kwargs)
~~~~~~~~~^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "D:\_ComfyUI_windows_portable\ComfyUI\comfy\patcher_extension.py", line 112, in execute
return self.original(*args, **kwargs)
~~~~~~~~~~~~~^^^^^^^^^^^^^^^^^
File "D:\_ComfyUI_windows_portable\ComfyUI\comfy\model_base.py", line 226, in _apply_model
model_output = self.diffusion_model(xc, t, context=context, control=control, transformer_options=transformer_options, **extra_conds)
File "D:\_ComfyUI_windows_portable\python_embeded\Lib\site-packages\torch\nn\modules\module.py", line 1778, in _wrapped_call_impl
return self._call_impl(*args, **kwargs)
~~~~~~~~~~~~~~~^^^^^^^^^^^^^^^^^
File "D:\_ComfyUI_windows_portable\python_embeded\Lib\site-packages\torch\nn\modules\module.py", line 1789, in _call_impl
return forward_call(*args, **kwargs)
File "D:\_ComfyUI_windows_portable\ComfyUI\comfy\ldm\hidream_o1\model.py", line 137, in forward
return comfy.patcher_extension.WrapperExecutor.new_class_executor(
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
...<2 lines>...
comfy.patcher_extension.get_all_wrappers(comfy.patcher_extension.WrappersMP.DIFFUSION_MODEL, transformer_options)
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
).execute(x, timesteps, context, transformer_options, **kwargs)
~~~~~~~~~^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "D:\_ComfyUI_windows_portable\ComfyUI\comfy\patcher_extension.py", line 113, in execute
return self.wrappers[self.idx](self, *args, **kwargs)
~~~~~~~~~~~~~~~~~~~~~~~^^^^^^^^^^^^^^^^^^^^^^^
File "D:\_ComfyUI_windows_portable\ComfyUI\comfy_extras\nodes_hidream_o1.py", line 203, in smoothing_wrapper
pred = executor(*args, **kwargs)
File "D:\_ComfyUI_windows_portable\ComfyUI\comfy\patcher_extension.py", line 105, in __call__
return new_executor.execute(*args, **kwargs)
~~~~~~~~~~~~~~~~~~~~^^^^^^^^^^^^^^^^^
File "D:\_ComfyUI_windows_portable\ComfyUI\comfy\patcher_extension.py", line 112, in execute
return self.original(*args, **kwargs)
~~~~~~~~~~~~~^^^^^^^^^^^^^^^^^
File "D:\_ComfyUI_windows_portable\ComfyUI\comfy\ldm\hidream_o1\model.py", line 154, in _forward
z = einops.rearrange(
x, 'B C (H p1) (W p2) -> B (H W) (C p1 p2)',
p1=self.patch_size, p2=self.patch_size,
)
File "D:\_ComfyUI_windows_portable\python_embeded\Lib\site-packages\einops\einops.py", line 616, in rearrange
return reduce(tensor, pattern, reduction="rearrange", **axes_lengths)
File "D:\_ComfyUI_windows_portable\python_embeded\Lib\site-packages\einops\einops.py", line 550, in reduce
raise EinopsError(message + f"\n {e}") from None
einops.EinopsError: Error while processing rearrange-reduction pattern "B C (H p1) (W p2) -> B (H W) (C p1 p2)".
Input tensor shape: torch.Size([1, 3, 2728, 1536]). Additional info: {'p1': 32, 'p2': 32}.
Shape mismatch, can't divide axis of length 2728 in chunks of 32
```

### Other

_No response_