fix: Repair PyTorch Compatibility issue with torch.compile in flex_attention

[neverbiasu]

No description provided.

[spookandrey1]

Unfortunately, after replacing the qwen2_navit.py file, the error did not go away.

[ComfyUI-Manager] All startup tasks have been completed.
got prompt
Error loading BAGEL model: name 'Qwen2Config' is not defined
!!! Exception during processing !!! name 'Qwen2Config' is not defined
Traceback (most recent call last):
File "G:\ComfyUI_art\ComfyUI\execution.py", line 349, 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 "G:\ComfyUI_art\ComfyUI\execution.py", line 224, 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 "G:\ComfyUI_art\ComfyUI\execution.py", line 196, in _map_node_over_list
process_inputs(input_dict, i)
File "G:\ComfyUI_art\ComfyUI\execution.py", line 185, in process_inputs
results.append(getattr(obj, func)(**inputs))
^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "G:\ComfyUI_art\ComfyUI\custom_nodes\ComfyUI-BAGEL\node.py", line 342, in load_model
raise e
File "G:\ComfyUI_art\ComfyUI\custom_nodes\ComfyUI-BAGEL\node.py", line 259, in load_model
llm_config = Qwen2Config.from_json_file(
^^^^^^^^^^^
NameError: name 'Qwen2Config' is not defined

[neverbiasu]

@spookandrey1 actually I haven't got enough info to help u, the traceback contained little, while other reason we all discussed.
Just adding a try-catch, u could try it to locate the problem.

[HDANILO]

Hi, I'm moving the discussion to this PR to avoid the noise the issue can get, perhaps here we'll be able to communicate more efficiently, here's my latest response:

I suggest some problems come from flash_attention, but @spookandrey1 ur issue isn't, it probably came from the torch.compile mentioned by @HDANILO (thanks ur reproduction & debug). I create a PR #16 for it, plz copy the code for trying and feedback in that PR, thanks.

Thanks for the quick response! I checked out the branch q2cfg and tried running, and I still got the same issue, yesterday I left my PC up and running getting flash-attn again(it can take even couple hours to compile the thing), but turns out the problem persisted even with that. q2cfg still didn't help.

I still get

Error importing BAGEL modules: cannot import name 'check_supported_striding' from 'torch._inductor.kernel.mm_common' (/home/dragonrollstudio/miniconda3/lib/python3.12/site-packages/torch/_inductor/kernel/mm_common.py)
Please ensure BAGEL model files are properly installed and accessible in the Python path.

when running.

The one thing that seems different is that before, ComfyUI would point that the custom_node failed to load, but now it seems to be pointing to be properly loading:

   0.0 seconds: /home/dragonrollstudio/docker_opt/comfyui/custom_nodes/ComfyUI-BAGEL

This is the git log from my repo to show I'm in the right head:

commit b38df73f04f5f01d835a4c503701e185922a5e5e (HEAD -> q2cfg, origin/q2cfg)
Author: Fayche <1751162157@qq.com>
Date:   Thu May 29 00:24:13 2025 +0800

    fix: Refactor flex_attention compilation and improve error handling in BAGEL model

commit e6cf60f2e17205e02f3c3ff754e3387ee3176b82
Author: Fayche <1751162157@qq.com>
Date:   Wed May 28 15:42:22 2025 +0800

    fix: Repair PyTorch Compatibility issue with torch.compile in flex_attention

commit 70a0dddad5c0d416faffd72549909b693baa0113 (master)
Author: Fayche <1751162157@qq.com>
Date:   Mon May 26 15:24:27 2025 +0800

    fix: Update model download paths

commit 21a845cc5759511fedb0c85a5f5fb0aec18b6002
Merge: 95c984f e0d4460
Author: Faych Chen <90372299+neverbiasu@users.noreply.github.com>
Date:   Mon May 26 13:38:07 2025 +0800

    Merge pull request #1 from neverbiasu/registry

    feat: Add publish workflow and project metadata for ComfyUI-BAGEL

[HDANILO]

Easiest workaround the issue I found was to completely comment out the line 40 from qwen2_navit.py

I'm trying to get the model to run and I think it'll work because importing the flex_attention is no problem, the problem is torch.compiling it

I'm facing extra problems getting the model to run, I'll document that in a separated issue, but I think they might be easy to fix

[HDANILO]

And its working!

[HDANILO]

These are the diffs from my local to branch q2cfg

diff --git a/modeling/bagel/qwen2_navit.py b/modeling/bagel/qwen2_navit.py
index 4d873f6..b1b89f5 100644
--- a/modeling/bagel/qwen2_navit.py
+++ b/modeling/bagel/qwen2_navit.py
@@ -37,7 +37,7 @@ from modeling.qwen2.configuration_qwen2 import Qwen2Config as _Qwen2Config
 torch._dynamo.config.cache_size_limit = 512
 torch._dynamo.config.accumulated_cache_size_limit = 4096
 # flex_attention = torch.compile(flex_attention) # , dynamic=True, mode='max-autotune'
-flex_attention = torch.compile(flex_attention)
+# flex_attention = torch.compile(flex_attention)


 class Qwen2Config(_Qwen2Config):
diff --git a/node.py b/node.py
index 81d2ec7..3adec0d 100644
--- a/node.py
+++ b/node.py
@@ -38,6 +38,7 @@ except ImportError as e:
     print(
         "Please ensure BAGEL model files are properly installed and accessible in the Python path."
     )
+    raise e

 # Register the BAGEL model folder
 models_dir = os.path.join(os.getcwd(), "models")
@@ -259,7 +260,9 @@ class BagelModelLoader:

             # Load configuration files
             try:
-                llm_config = Qwen2Config.from_json_file("path/to/llm_config.json")
+                llm_config = Qwen2Config.from_json_file(
+                    os.path.join(local_model_dir, "llm_config.json")
+                )
             except Exception as e:
                 print(f"Error loading Qwen2Config: {e}")
                 raise

enumerating:

1. disable torch compilation of the flex_attention, perhaps this should be a checkbox of some sorts as it doesnt seem stable
2. raise the exception on the node import to clarify better why is it failing to load
3. changed llm_config import location from a non-sense folder "path/to/llm_config.json" to the actual folder where the model belongs: os.path.join(local_model_dir, "llm_config.json")

[neverbiasu]

@HDANILO good to hear that u solve it!

[neverbiasu]

@HDANILO actually I mismatch u two's problems together, the first commit in this PR is to solve urs while it cannot help @spookandrey1. Thus I notice the issue #7 is not same problems.