UniPC and `linalg_solve_ex` with AuraFlow and ComfyUI

exdysa: <p>Hopefully the right place to be here…<br>
I’ve been encountering an error while using Pytorch in conjunction with AuraFlow and UniPC which brought me to the support matrix. I noticed there that <code>linalg_solve_triangular</code> is marked as green and enabled, but I don’t see <code>linalg_solve_ex</code> listed there. The instructions<br>
suggest checking the matrix first before commenting on the GitHub issue.</p>
<p>My questions are:</p>
<ol>
<li>Are <code>solve_ex</code> and <code>solve_triangular</code> different operations, or are they both part of the same <code>linalg_solve</code> functionality?</li>
<li>Should I comment on the GitHub issue to request support for <code>linalg_solve_ex</code>, even though <code>linalg_solve_triangular</code> is already supported?</li>
</ol>
<p>Exact Error</p>
<pre><code class="lang-auto">got prompt
  8%|██████████▉                                                                                                                              | 2/25 [00:15&lt;02:52,  7.52s/it]
!!! Exception during processing !!! The operator 'aten::_linalg_solve_ex.result' is not currently implemented for the MPS device. If you want this op to be considered for addition please comment on https://github.com/pytorch/pytorch/issues/141287 and mention use-case, that resulted in missing op as well as commit hash 3dd667797a7f46f6c241e492ad1586058acc5174. As a temporary fix, you can set the environment variable `PYTORCH_ENABLE_MPS_FALLBACK=1` to use the CPU as a fallback for this op. WARNING: this will be slower than running natively on MPS.
Traceback (most recent call last):
  File "/Users/unauthorized/Documents/GitHub/src/comfyui/execution.py", line 324, 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 "/Users/unauthorized/Documents/GitHub/src/comfyui/execution.py", line 199, 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 "/Users/unauthorized/Documents/GitHub/src/comfyui/execution.py", line 170, in _map_node_over_list
    process_inputs(input_dict, i)
  File "/Users/unauthorized/Documents/GitHub/src/comfyui/execution.py", line 159, in process_inputs
    results.append(getattr(obj, func)(**inputs))
                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/unauthorized/Documents/GitHub/src/comfyui/nodes.py", line 1467, in sample
    return common_ksampler(model, seed, steps, cfg, sampler_name, scheduler, positive, negative, latent_image, denoise=denoise)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/unauthorized/Documents/GitHub/src/comfyui/nodes.py", line 1434, in common_ksampler
    samples = comfy.sample.sample(model, noise, steps, cfg, sampler_name, scheduler, positive, negative, latent_image,
              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/unauthorized/Documents/GitHub/src/comfyui/comfy/sample.py", line 43, 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 "/Users/unauthorized/Documents/GitHub/src/comfyui/comfy/samplers.py", line 1020, 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 "/Users/unauthorized/Documents/GitHub/src/comfyui/comfy/samplers.py", line 918, in sample
    return cfg_guider.sample(noise, latent_image, sampler, sigmas, denoise_mask, callback, disable_pbar, seed)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/unauthorized/Documents/GitHub/src/comfyui/comfy/samplers.py", line 904, in sample
    output = executor.execute(noise, latent_image, sampler, sigmas, denoise_mask, callback, disable_pbar, seed)
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/unauthorized/Documents/GitHub/src/comfyui/comfy/patcher_extension.py", line 110, in execute
    return self.original(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/unauthorized/Documents/GitHub/src/comfyui/comfy/samplers.py", line 873, in outer_sample
    output = self.inner_sample(noise, latent_image, device, sampler, sigmas, denoise_mask, callback, disable_pbar, seed)
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/unauthorized/Documents/GitHub/src/comfyui/comfy/samplers.py", line 857, in inner_sample
    samples = executor.execute(self, sigmas, extra_args, callback, noise, latent_image, denoise_mask, disable_pbar)
              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/unauthorized/Documents/GitHub/src/comfyui/comfy/patcher_extension.py", line 110, in execute
    return self.original(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/unauthorized/Documents/GitHub/src/comfyui/comfy/samplers.py", line 714, in sample
    samples = self.sampler_function(model_k, noise, sigmas, extra_args=extra_args, callback=k_callback, disable=disable_pbar, **self.extra_options)
              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/unauthorized/Documents/GitHub/src/comfyui/comfy/extra_samplers/uni_pc.py", line 870, 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 "/Users/unauthorized/Documents/GitHub/src/comfyui/comfy/extra_samplers/uni_pc.py", line 724, in sample
    x, model_x = self.multistep_uni_pc_update(x, model_prev_list, t_prev_list, vec_t, init_order, use_corrector=True)
                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/unauthorized/Documents/GitHub/src/comfyui/comfy/extra_samplers/uni_pc.py", line 472, in multistep_uni_pc_update
    return self.multistep_uni_pc_bh_update(x, model_prev_list, t_prev_list, t, order, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/unauthorized/Documents/GitHub/src/comfyui/comfy/extra_samplers/uni_pc.py", line 654, in multistep_uni_pc_bh_update
    rhos_c = torch.linalg.solve(R, b)
             ^^^^^^^^^^^^^^^^^^^^^^^^
</code></pre>
<p>M4 Max 16 Nov 2024 Sequoia 128gb 15.1.1<br>
Python 3.12.7 [Clang 18.1.8 ] on darwin<br>
Using Python 3.12.7 environment .virtualenvs/comfyui<br>
Package            Version</p>
<hr>
<p>aiohappyeyeballs   2.4.4<br>
aiohttp            3.11.10<br>
aiosignal          1.3.1<br>
attrs              24.2.0<br>
certifi            2024.8.30<br>
cffi               1.17.1<br>
charset-normalizer 3.4.0<br>
einops             0.8.0<br>
filelock           3.16.1<br>
frozenlist         1.5.0<br>
fsspec             2024.10.0<br>
huggingface-hub    0.26.5<br>
idna               3.10<br>
jinja2             3.1.4<br>
kornia             0.7.4<br>
kornia-rs          0.1.7<br>
markupsafe         2.1.5<br>
mpmath             1.3.0<br>
multidict          6.1.0<br>
networkx           3.4.2<br>
numpy              2.1.2<br>
packaging          24.2<br>
pillow             11.0.0<br>
propcache          0.2.1<br>
psutil             6.1.0<br>
pycparser          2.22<br>
pyyaml             6.0.2<br>
regex              2024.11.6<br>
requests           2.32.3<br>
safetensors        0.4.5<br>
scipy              1.14.1<br>
sentencepiece      0.2.0<br>
setuptools         70.2.0<br>
soundfile          0.12.1<br>
spandrel           0.4.0<br>
sympy              1.13.1<br>
tokenizers         0.21.0<br>
torch              2.6.0.dev20241206<br>
torchaudio         2.5.0.dev20241206<br>
torchsde           0.2.6<br>
torchvision        0.20.0.dev20241206<br>
tqdm               4.67.1<br>
trampoline         0.1.2<br>
transformers       4.47.0<br>
typing-extensions  4.12.2<br>
urllib3            2.2.3<br>
yarl               1.18.3</p>
<p>Any insights or suggestions would be greatly appreciated!</p>
<p>Thanks</p>