From 08d6c98cfe6d0330f1f00f71ad8cad816edcc72e Mon Sep 17 00:00:00 2001 From: Maxed-Out-99 Date: Sat, 8 Aug 2026 05:53:35 -0700 Subject: [PATCH] Move Smart Model Loaders out to their own repo; bump to 3.0.0 Core ComfyUI's memory handling has improved enough that GGUF is no longer a default need, so the five Smart loader nodes (UNET + CLIP/Dual/Triple/Quad) move to the standalone Smart-Model-Loaders-MXD pack. 67 -> 62 nodes. Schema re-baselined intentionally. Side effect worth noting: those nodes carried the pack's only pip dependencies (gguf, sentencepiece, protobuf), so requirements.txt is deleted. ComfyUI-MaxedOut now installs with nothing beyond what ComfyUI itself ships. Version 2.0.0 -> 3.0.0. Major, because this release removes 15 nodes in total and any workflow using one will show it as missing. Co-Authored-By: Claude Opus 5 --- README.md | 9 +- __init__.py | 1 - pyproject.toml | 2 +- requirements.txt | 5 - scripts/node_schema_baseline.json | 263 +-------------- smart_loaders_mxd/LICENSE | 201 ------------ smart_loaders_mxd/UPSTREAM.md | 41 --- smart_loaders_mxd/__init__.py | 3 - smart_loaders_mxd/dequant.py | 303 ----------------- smart_loaders_mxd/loader.py | 508 ----------------------------- smart_loaders_mxd/nodes.py | 358 -------------------- smart_loaders_mxd/ops.py | 283 ---------------- smart_loaders_mxd/tools/convert.py | 369 --------------------- 13 files changed, 7 insertions(+), 2339 deletions(-) delete mode 100644 requirements.txt delete mode 100644 smart_loaders_mxd/LICENSE delete mode 100644 smart_loaders_mxd/UPSTREAM.md delete mode 100644 smart_loaders_mxd/__init__.py delete mode 100644 smart_loaders_mxd/dequant.py delete mode 100644 smart_loaders_mxd/loader.py delete mode 100644 smart_loaders_mxd/nodes.py delete mode 100644 smart_loaders_mxd/ops.py delete mode 100644 smart_loaders_mxd/tools/convert.py diff --git a/README.md b/README.md index b61b4bd..01c98c9 100644 --- a/README.md +++ b/README.md @@ -70,6 +70,7 @@ only what you want. Nothing below is required. | [Live-Preview-MXD](https://github.com/Maxed-Out-99/Live-Preview-MXD) | Watch video generations animate while they render, in a dockable panel. Also adds LTX 2.3 previews, which core can't do at all. | | [Prompt-Library-MXD](https://github.com/Maxed-Out-99/Prompt-Library-MXD) | Save named prompt snippets by category and pull them in by name, or let a wildcard pick one at random. | | [Spell-Check-MXD](https://github.com/Maxed-Out-99/Spell-Check-MXD) | The prompt spellchecker on its own, if you don't want the rest of this pack. Safe to install alongside — it won't double up. | +| [Smart-Model-Loaders-MXD](https://github.com/Maxed-Out-99/Smart-Model-Loaders-MXD) | Loader nodes that accept safetensors or GGUF in the same slot. Only needed if you actually use GGUF quants. | ## Free Workflows @@ -96,7 +97,7 @@ If you star this repo, definitely consider starring theirs too. ## License -MIT — see [LICENSE](LICENSE). Vendored dependencies keep their own licenses: -`smart_loaders_mxd/` (GGUF loading, from -[city96/ComfyUI-GGUF](https://github.com/city96/ComfyUI-GGUF)) and -`web/vendor/` each carry their license files alongside the code. +MIT — see [LICENSE](LICENSE). Vendored code under `web/vendor/` keeps its own +licenses, carried alongside it in that folder. + +This pack has no pip dependencies — everything it needs ships with ComfyUI. diff --git a/__init__.py b/__init__.py index 9ea68f3..af17295 100644 --- a/__init__.py +++ b/__init__.py @@ -22,7 +22,6 @@ NODE_DISPLAY_NAME_MAPPINGS = {} for _name in ( "nodes", "loraloader_mxd", - "smart_loaders_mxd", "system.model_paths", ): _mod = _safe_import(_name) diff --git a/pyproject.toml b/pyproject.toml index 617b396..68339f2 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,7 +1,7 @@ [project] name = "ComfyUI-MaxedOut" description = "Custom ComfyUI nodes used in Maxed Out workflows (SDXL, Flux, Wan 2.2, etc.)" -version = "2.0.0" +version = "3.0.0" license = {file = "LICENSE"} # classifiers = [ # # For OS-independent nodes (works on all operating systems) diff --git a/requirements.txt b/requirements.txt deleted file mode 100644 index 04cd9a9..0000000 --- a/requirements.txt +++ /dev/null @@ -1,5 +0,0 @@ -# Required for smart_loaders_mxd (GGUF model loading) -gguf>=0.13.0 -# Optional - GGUF tokenizer reconstruction -sentencepiece -protobuf diff --git a/scripts/node_schema_baseline.json b/scripts/node_schema_baseline.json index 864fc9b..df69546 100644 --- a/scripts/node_schema_baseline.json +++ b/scripts/node_schema_baseline.json @@ -1,5 +1,5 @@ { - "node_count": 67, + "node_count": 62, "nodes": { "BboxDetectorCombinedBatchMXD": { "category": "MXD/Detector", @@ -42,61 +42,6 @@ "MASK" ] }, - "CLIPLoaderUnified": { - "category": "advanced/loaders", - "class": "CLIPLoaderUnified", - "display_name": "Smart CLIP Loader MXD", - "function": "load_clip", - "input_types": { - "required": { - "clip_name1": { - "type": { - "combo": "nonempty", - "dynamic": true - } - }, - "type": { - "type": { - "combo": [ - "stable_diffusion", - "stable_cascade", - "sd3", - "stable_audio", - "mochi", - "ltxv", - "pixart", - "cosmos", - "lumina2", - "wan", - "hidream", - "chroma", - "ace", - "omnigen2", - "qwen_image", - "hunyuan_image", - "flux2", - "ovis", - "longcat_image", - "cogvideox", - "lens", - "pixeldit", - "ideogram4", - "boogu", - "krea2", - "joyimage", - "mage", - "minimax" - ] - } - } - } - }, - "output_node": false, - "return_names": null, - "return_types": [ - "CLIP" - ] - }, "CombineVideos_MXD": { "category": "MXD/video", "class": "CombineVideos_MXD", @@ -151,51 +96,6 @@ "IMAGE" ] }, - "DualCLIPLoaderUnified": { - "category": "advanced/loaders", - "class": "DualCLIPLoaderUnified", - "display_name": "Smart Dual CLIP Loader MXD", - "function": "load_clip", - "input_types": { - "required": { - "clip_name1": { - "type": { - "combo": "nonempty", - "dynamic": true - } - }, - "clip_name2": { - "type": { - "combo": "nonempty", - "dynamic": true - } - }, - "type": { - "type": { - "combo": [ - "sdxl", - "sd3", - "flux", - "hunyuan_video", - "hidream", - "hunyuan_image", - "hunyuan_video_15", - "kandinsky5", - "kandinsky5_image", - "ltxv", - "newbie", - "ace" - ] - } - } - } - }, - "output_node": false, - "return_names": null, - "return_types": [ - "CLIP" - ] - }, "Extract Workflow From Image MXD": { "category": "MXD/Image", "class": "ExtractWorkflowFromImageMXD", @@ -1763,79 +1663,6 @@ "CONDITIONING" ] }, - "QuadrupleCLIPLoaderUnified": { - "category": "advanced/loaders", - "class": "QuadrupleCLIPLoaderUnified", - "display_name": "Smart Quad CLIP Loader MXD", - "function": "load_clip", - "input_types": { - "required": { - "clip_name1": { - "type": { - "combo": "nonempty", - "dynamic": true - } - }, - "clip_name2": { - "type": { - "combo": "nonempty", - "dynamic": true - } - }, - "clip_name3": { - "type": { - "combo": "nonempty", - "dynamic": true - } - }, - "clip_name4": { - "type": { - "combo": "nonempty", - "dynamic": true - } - }, - "type": { - "type": { - "combo": [ - "stable_diffusion", - "stable_cascade", - "sd3", - "stable_audio", - "mochi", - "ltxv", - "pixart", - "cosmos", - "lumina2", - "wan", - "hidream", - "chroma", - "ace", - "omnigen2", - "qwen_image", - "hunyuan_image", - "flux2", - "ovis", - "longcat_image", - "cogvideox", - "lens", - "pixeldit", - "ideogram4", - "boogu", - "krea2", - "joyimage", - "mage", - "minimax" - ] - } - } - } - }, - "output_node": false, - "return_names": null, - "return_types": [ - "CLIP" - ] - }, "QwenImageEditSingleMXD": { "category": "MXD/conditioning", "class": "QwenImageEditSingleMXD", @@ -2472,94 +2299,6 @@ "IMAGE" ] }, - "TripleCLIPLoaderUnified": { - "category": "advanced/loaders", - "class": "TripleCLIPLoaderUnified", - "display_name": "Smart Triple CLIP Loader MXD", - "function": "load_clip", - "input_types": { - "required": { - "clip_name1": { - "type": { - "combo": "nonempty", - "dynamic": true - } - }, - "clip_name2": { - "type": { - "combo": "nonempty", - "dynamic": true - } - }, - "clip_name3": { - "type": { - "combo": "nonempty", - "dynamic": true - } - }, - "type": { - "type": { - "combo": [ - "stable_diffusion", - "stable_cascade", - "sd3", - "stable_audio", - "mochi", - "ltxv", - "pixart", - "cosmos", - "lumina2", - "wan", - "hidream", - "chroma", - "ace", - "omnigen2", - "qwen_image", - "hunyuan_image", - "flux2", - "ovis", - "longcat_image", - "cogvideox", - "lens", - "pixeldit", - "ideogram4", - "boogu", - "krea2", - "joyimage", - "mage", - "minimax" - ] - } - } - } - }, - "output_node": false, - "return_names": null, - "return_types": [ - "CLIP" - ] - }, - "UNETLoaderUnified": { - "category": "advanced/loaders", - "class": "UNETLoaderUnified", - "display_name": "Smart Model Loader MXD", - "function": "load_unet", - "input_types": { - "required": { - "unet_name": { - "type": { - "combo": "nonempty", - "dynamic": true - } - } - } - }, - "output_node": false, - "return_names": null, - "return_types": [ - "MODEL" - ] - }, "Video Comparer MXD": { "category": "MXD/video", "class": "MxdVideoComparer", diff --git a/smart_loaders_mxd/LICENSE b/smart_loaders_mxd/LICENSE deleted file mode 100644 index 261eeb9..0000000 --- a/smart_loaders_mxd/LICENSE +++ /dev/null @@ -1,201 +0,0 @@ - Apache License - Version 2.0, January 2004 - http://www.apache.org/licenses/ - - TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION - - 1. Definitions. - - "License" shall mean the terms and conditions for use, reproduction, - and distribution as defined by Sections 1 through 9 of this document. - - "Licensor" shall mean the copyright owner or entity authorized by - the copyright owner that is granting the License. - - "Legal Entity" shall mean the union of the acting entity and all - other entities that control, are controlled by, or are under common - control with that entity. For the purposes of this definition, - "control" means (i) the power, direct or indirect, to cause the - direction or management of such entity, whether by contract or - otherwise, or (ii) ownership of fifty percent (50%) or more of the - outstanding shares, or (iii) beneficial ownership of such entity. - - "You" (or "Your") shall mean an individual or Legal Entity - exercising permissions granted by this License. - - "Source" form shall mean the preferred form for making modifications, - including but not limited to software source code, documentation - source, and configuration files. - - "Object" form shall mean any form resulting from mechanical - transformation or translation of a Source form, including but - not limited to compiled object code, generated documentation, - and conversions to other media types. - - "Work" shall mean the work of authorship, whether in Source or - Object form, made available under the License, as indicated by a - copyright notice that is included in or attached to the work - (an example is provided in the Appendix below). - - "Derivative Works" shall mean any work, whether in Source or Object - form, that is based on (or derived from) the Work and for which the - editorial revisions, annotations, elaborations, or other modifications - represent, as a whole, an original work of authorship. For the purposes - of this License, Derivative Works shall not include works that remain - separable from, or merely link (or bind by name) to the interfaces of, - the Work and Derivative Works thereof. - - "Contribution" shall mean any work of authorship, including - the original version of the Work and any modifications or additions - to that Work or Derivative Works thereof, that is intentionally - submitted to Licensor for inclusion in the Work by the copyright owner - or by an individual or Legal Entity authorized to submit on behalf of - the copyright owner. For the purposes of this definition, "submitted" - means any form of electronic, verbal, or written communication sent - to the Licensor or its representatives, including but not limited to - communication on electronic mailing lists, source code control systems, - and issue tracking systems that are managed by, or on behalf of, the - Licensor for the purpose of discussing and improving the Work, but - excluding communication that is conspicuously marked or otherwise - designated in writing by the copyright owner as "Not a Contribution." - - "Contributor" shall mean Licensor and any individual or Legal Entity - on behalf of whom a Contribution has been received by Licensor and - subsequently incorporated within the Work. - - 2. Grant of Copyright License. Subject to the terms and conditions of - this License, each Contributor hereby grants to You a perpetual, - worldwide, non-exclusive, no-charge, royalty-free, irrevocable - copyright license to reproduce, prepare Derivative Works of, - publicly display, publicly perform, sublicense, and distribute the - Work and such Derivative Works in Source or Object form. - - 3. Grant of Patent License. Subject to the terms and conditions of - this License, each Contributor hereby grants to You a perpetual, - worldwide, non-exclusive, no-charge, royalty-free, irrevocable - (except as stated in this section) patent license to make, have made, - use, offer to sell, sell, import, and otherwise transfer the Work, - where such license applies only to those patent claims licensable - by such Contributor that are necessarily infringed by their - Contribution(s) alone or by combination of their Contribution(s) - with the Work to which such Contribution(s) was submitted. If You - institute patent litigation against any entity (including a - cross-claim or counterclaim in a lawsuit) alleging that the Work - or a Contribution incorporated within the Work constitutes direct - or contributory patent infringement, then any patent licenses - granted to You under this License for that Work shall terminate - as of the date such litigation is filed. - - 4. Redistribution. You may reproduce and distribute copies of the - Work or Derivative Works thereof in any medium, with or without - modifications, and in Source or Object form, provided that You - meet the following conditions: - - (a) You must give any other recipients of the Work or - Derivative Works a copy of this License; and - - (b) You must cause any modified files to carry prominent notices - stating that You changed the files; and - - (c) You must retain, in the Source form of any Derivative Works - that You distribute, all copyright, patent, trademark, and - attribution notices from the Source form of the Work, - excluding those notices that do not pertain to any part of - the Derivative Works; and - - (d) If the Work includes a "NOTICE" text file as part of its - distribution, then any Derivative Works that You distribute must - include a readable copy of the attribution notices contained - within such NOTICE file, excluding those notices that do not - pertain to any part of the Derivative Works, in at least one - of the following places: within a NOTICE text file distributed - as part of the Derivative Works; within the Source form or - documentation, if provided along with the Derivative Works; or, - within a display generated by the Derivative Works, if and - wherever such third-party notices normally appear. The contents - of the NOTICE file are for informational purposes only and - do not modify the License. You may add Your own attribution - notices within Derivative Works that You distribute, alongside - or as an addendum to the NOTICE text from the Work, provided - that such additional attribution notices cannot be construed - as modifying the License. - - You may add Your own copyright statement to Your modifications and - may provide additional or different license terms and conditions - for use, reproduction, or distribution of Your modifications, or - for any such Derivative Works as a whole, provided Your use, - reproduction, and distribution of the Work otherwise complies with - the conditions stated in this License. - - 5. Submission of Contributions. Unless You explicitly state otherwise, - any Contribution intentionally submitted for inclusion in the Work - by You to the Licensor shall be under the terms and conditions of - this License, without any additional terms or conditions. - Notwithstanding the above, nothing herein shall supersede or modify - the terms of any separate license agreement you may have executed - with Licensor regarding such Contributions. - - 6. Trademarks. This License does not grant permission to use the trade - names, trademarks, service marks, or product names of the Licensor, - except as required for reasonable and customary use in describing the - origin of the Work and reproducing the content of the NOTICE file. - - 7. Disclaimer of Warranty. Unless required by applicable law or - agreed to in writing, Licensor provides the Work (and each - Contributor provides its Contributions) on an "AS IS" BASIS, - WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or - implied, including, without limitation, any warranties or conditions - of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A - PARTICULAR PURPOSE. You are solely responsible for determining the - appropriateness of using or redistributing the Work and assume any - risks associated with Your exercise of permissions under this License. - - 8. Limitation of Liability. In no event and under no legal theory, - whether in tort (including negligence), contract, or otherwise, - unless required by applicable law (such as deliberate and grossly - negligent acts) or agreed to in writing, shall any Contributor be - liable to You for damages, including any direct, indirect, special, - incidental, or consequential damages of any character arising as a - result of this License or out of the use or inability to use the - Work (including but not limited to damages for loss of goodwill, - work stoppage, computer failure or malfunction, or any and all - other commercial damages or losses), even if such Contributor - has been advised of the possibility of such damages. - - 9. Accepting Warranty or Additional Liability. While redistributing - the Work or Derivative Works thereof, You may choose to offer, - and charge a fee for, acceptance of support, warranty, indemnity, - or other liability obligations and/or rights consistent with this - License. However, in accepting such obligations, You may act only - on Your own behalf and on Your sole responsibility, not on behalf - of any other Contributor, and only if You agree to indemnify, - defend, and hold each Contributor harmless for any liability - incurred by, or claims asserted against, such Contributor by reason - of your accepting any such warranty or additional liability. - - END OF TERMS AND CONDITIONS - - APPENDIX: How to apply the Apache License to your work. - - To apply the Apache License to your work, attach the following - boilerplate notice, with the fields enclosed by brackets "[]" - replaced with your own identifying information. (Don't include - the brackets!) The text should be enclosed in the appropriate - comment syntax for the file format. We also recommend that a - file or class name and description of purpose be included on the - same "printed page" as the copyright notice for easier - identification within third-party archives. - - Copyright [yyyy] [name of copyright owner] - - Licensed under the Apache License, Version 2.0 (the "License"); - you may not use this file except in compliance with the License. - You may obtain a copy of the License at - - http://www.apache.org/licenses/LICENSE-2.0 - - Unless required by applicable law or agreed to in writing, software - distributed under the License is distributed on an "AS IS" BASIS, - WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - See the License for the specific language governing permissions and - limitations under the License. diff --git a/smart_loaders_mxd/UPSTREAM.md b/smart_loaders_mxd/UPSTREAM.md deleted file mode 100644 index 48f534b..0000000 --- a/smart_loaders_mxd/UPSTREAM.md +++ /dev/null @@ -1,41 +0,0 @@ -# Upstream sync notes - -`dequant.py`, `loader.py`, `ops.py`, and `tools/convert.py` in this package are -vendored **verbatim** from [city96/ComfyUI-GGUF](https://github.com/city96/ComfyUI-GGUF) -(Apache-2.0 — see `LICENSE` in this folder). `nodes.py` is our own adaptation -(unified UNET/CLIP loaders that pick between safetensors and GGUF automatically) -and is not meant to track upstream 1:1. - -## Baseline - -Vendored from `Maxed-Out-99/ComfyUI-SmartModelLoaders-MXD` at commit -`7fd86b156afbe3bb453958900d5b6bd8c074ab8a` (2026-02-22, "Add Gemma3 GGUF support -and mmap/memory fixes"). That repo is a manual copy of city96/ComfyUI-GGUF -rather than a git fork, but as of the last check (2026-06-13) its -`dequant.py`/`loader.py`/`ops.py` were byte-identical to city96's `main` -branch — so this baseline should be treated as "current with upstream `main`" -as of that date. - -## How to sync when city96 ships an update - -1. Download the current `dequant.py`, `loader.py`, `ops.py`, and - `tools/convert.py` from - `https://github.com/city96/ComfyUI-GGUF/tree/main`. -2. Diff each one against the matching file in this folder (ignore the header - comment block added here). If the diff is clean (no local modifications - to reconcile — these 4 files are never hand-edited), just drop the new - versions in and re-add the vendoring header comment to the top of each. -3. Leave `nodes.py` alone — cross-check it against upstream's `nodes.py` only - if you want to pull in a new *loader node* (e.g. their `GGUFOps` - changes), and merge by hand since ours diverges intentionally (unified - loaders, MXD info-icon UI hooks). -4. Bump the baseline note above (commit/date) after syncing. - -## What was intentionally left out - -The standalone repo's `tools/` folder also has `fix_5d_tensors.py`, -`fix_lines_ending.py`, `lcpp.patch`, and `read_tensors.py` — one-off CLI/dev -scripts, not runtime dependencies. Only `convert.py` was vendored here, -because `loader.py`'s compatibility-mode fallback does -`from .tools.convert import detect_arch` at runtime for GGUF files missing -standard architecture metadata. diff --git a/smart_loaders_mxd/__init__.py b/smart_loaders_mxd/__init__.py deleted file mode 100644 index 39a8c6b..0000000 --- a/smart_loaders_mxd/__init__.py +++ /dev/null @@ -1,3 +0,0 @@ -from .nodes import NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS - -__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS"] diff --git a/smart_loaders_mxd/dequant.py b/smart_loaders_mxd/dequant.py deleted file mode 100644 index 0c32084..0000000 --- a/smart_loaders_mxd/dequant.py +++ /dev/null @@ -1,303 +0,0 @@ -# (c) City96 || Apache-2.0 (apache.org/licenses/LICENSE-2.0) -# Vendored verbatim from city96/ComfyUI-GGUF. See smart_loaders_mxd/UPSTREAM.md -# for the sync baseline and procedure. Do not hand-edit; port upstream diffs instead. -import gguf -import torch -from tqdm import tqdm - - -TORCH_COMPATIBLE_QTYPES = (None, gguf.GGMLQuantizationType.F32, gguf.GGMLQuantizationType.F16) - -def is_torch_compatible(tensor): - return tensor is None or getattr(tensor, "tensor_type", None) in TORCH_COMPATIBLE_QTYPES - -def is_quantized(tensor): - return not is_torch_compatible(tensor) - -def dequantize_tensor(tensor, dtype=None, dequant_dtype=None): - qtype = getattr(tensor, "tensor_type", None) - oshape = getattr(tensor, "tensor_shape", tensor.shape) - - if qtype in TORCH_COMPATIBLE_QTYPES: - return tensor.to(dtype) - elif qtype in dequantize_functions: - dequant_dtype = dtype if dequant_dtype == "target" else dequant_dtype - return dequantize(tensor.data, qtype, oshape, dtype=dequant_dtype).to(dtype) - else: - # this is incredibly slow - tqdm.write(f"Falling back to numpy dequant for qtype: {getattr(qtype, 'name', repr(qtype))}") - new = gguf.quants.dequantize(tensor.cpu().numpy(), qtype) - return torch.from_numpy(new).to(tensor.device, dtype=dtype) - -def dequantize(data, qtype, oshape, dtype=None): - """ - Dequantize tensor back to usable shape/dtype - """ - block_size, type_size = gguf.GGML_QUANT_SIZES[qtype] - dequantize_blocks = dequantize_functions[qtype] - - rows = data.reshape( - (-1, data.shape[-1]) - ).view(torch.uint8) - - n_blocks = rows.numel() // type_size - blocks = rows.reshape((n_blocks, type_size)) - blocks = dequantize_blocks(blocks, block_size, type_size, dtype) - return blocks.reshape(oshape) - -def to_uint32(x): - # no uint32 :( - x = x.view(torch.uint8).to(torch.int32) - return (x[:, 0] | x[:, 1] << 8 | x[:, 2] << 16 | x[:, 3] << 24).unsqueeze(1) - -def to_uint16(x): - x = x.view(torch.uint8).to(torch.int32) - return (x[:, 0] | x[:, 1] << 8).unsqueeze(1) - -def split_block_dims(blocks, *args): - n_max = blocks.shape[1] - dims = list(args) + [n_max - sum(args)] - return torch.split(blocks, dims, dim=1) - -# Full weights # -def dequantize_blocks_BF16(blocks, block_size, type_size, dtype=None): - return (blocks.view(torch.int16).to(torch.int32) << 16).view(torch.float32) - -# Legacy Quants # -def dequantize_blocks_Q8_0(blocks, block_size, type_size, dtype=None): - d, x = split_block_dims(blocks, 2) - d = d.view(torch.float16).to(dtype) - x = x.view(torch.int8) - return (d * x) - -def dequantize_blocks_Q5_1(blocks, block_size, type_size, dtype=None): - n_blocks = blocks.shape[0] - - d, m, qh, qs = split_block_dims(blocks, 2, 2, 4) - d = d.view(torch.float16).to(dtype) - m = m.view(torch.float16).to(dtype) - qh = to_uint32(qh) - - qh = qh.reshape((n_blocks, 1)) >> torch.arange(32, device=d.device, dtype=torch.int32).reshape(1, 32) - ql = qs.reshape((n_blocks, -1, 1, block_size // 2)) >> torch.tensor([0, 4], device=d.device, dtype=torch.uint8).reshape(1, 1, 2, 1) - qh = (qh & 1).to(torch.uint8) - ql = (ql & 0x0F).reshape((n_blocks, -1)) - - qs = (ql | (qh << 4)) - return (d * qs) + m - -def dequantize_blocks_Q5_0(blocks, block_size, type_size, dtype=None): - n_blocks = blocks.shape[0] - - d, qh, qs = split_block_dims(blocks, 2, 4) - d = d.view(torch.float16).to(dtype) - qh = to_uint32(qh) - - qh = qh.reshape(n_blocks, 1) >> torch.arange(32, device=d.device, dtype=torch.int32).reshape(1, 32) - ql = qs.reshape(n_blocks, -1, 1, block_size // 2) >> torch.tensor([0, 4], device=d.device, dtype=torch.uint8).reshape(1, 1, 2, 1) - - qh = (qh & 1).to(torch.uint8) - ql = (ql & 0x0F).reshape(n_blocks, -1) - - qs = (ql | (qh << 4)).to(torch.int8) - 16 - return (d * qs) - -def dequantize_blocks_Q4_1(blocks, block_size, type_size, dtype=None): - n_blocks = blocks.shape[0] - - d, m, qs = split_block_dims(blocks, 2, 2) - d = d.view(torch.float16).to(dtype) - m = m.view(torch.float16).to(dtype) - - qs = qs.reshape((n_blocks, -1, 1, block_size // 2)) >> torch.tensor([0, 4], device=d.device, dtype=torch.uint8).reshape(1, 1, 2, 1) - qs = (qs & 0x0F).reshape(n_blocks, -1) - - return (d * qs) + m - -def dequantize_blocks_Q4_0(blocks, block_size, type_size, dtype=None): - n_blocks = blocks.shape[0] - - d, qs = split_block_dims(blocks, 2) - d = d.view(torch.float16).to(dtype) - - qs = qs.reshape((n_blocks, -1, 1, block_size // 2)) >> torch.tensor([0, 4], device=d.device, dtype=torch.uint8).reshape((1, 1, 2, 1)) - qs = (qs & 0x0F).reshape((n_blocks, -1)).to(torch.int8) - 8 - return (d * qs) - -# K Quants # -QK_K = 256 -K_SCALE_SIZE = 12 - -def get_scale_min(scales): - n_blocks = scales.shape[0] - scales = scales.view(torch.uint8) - scales = scales.reshape((n_blocks, 3, 4)) - - d, m, m_d = torch.split(scales, scales.shape[-2] // 3, dim=-2) - - sc = torch.cat([d & 0x3F, (m_d & 0x0F) | ((d >> 2) & 0x30)], dim=-1) - min = torch.cat([m & 0x3F, (m_d >> 4) | ((m >> 2) & 0x30)], dim=-1) - - return (sc.reshape((n_blocks, 8)), min.reshape((n_blocks, 8))) - -def dequantize_blocks_Q6_K(blocks, block_size, type_size, dtype=None): - n_blocks = blocks.shape[0] - - ql, qh, scales, d, = split_block_dims(blocks, QK_K // 2, QK_K // 4, QK_K // 16) - - scales = scales.view(torch.int8).to(dtype) - d = d.view(torch.float16).to(dtype) - d = (d * scales).reshape((n_blocks, QK_K // 16, 1)) - - ql = ql.reshape((n_blocks, -1, 1, 64)) >> torch.tensor([0, 4], device=d.device, dtype=torch.uint8).reshape((1, 1, 2, 1)) - ql = (ql & 0x0F).reshape((n_blocks, -1, 32)) - qh = qh.reshape((n_blocks, -1, 1, 32)) >> torch.tensor([0, 2, 4, 6], device=d.device, dtype=torch.uint8).reshape((1, 1, 4, 1)) - qh = (qh & 0x03).reshape((n_blocks, -1, 32)) - q = (ql | (qh << 4)).to(torch.int8) - 32 - q = q.reshape((n_blocks, QK_K // 16, -1)) - - return (d * q).reshape((n_blocks, QK_K)) - -def dequantize_blocks_Q5_K(blocks, block_size, type_size, dtype=None): - n_blocks = blocks.shape[0] - - d, dmin, scales, qh, qs = split_block_dims(blocks, 2, 2, K_SCALE_SIZE, QK_K // 8) - - d = d.view(torch.float16).to(dtype) - dmin = dmin.view(torch.float16).to(dtype) - - sc, m = get_scale_min(scales) - - d = (d * sc).reshape((n_blocks, -1, 1)) - dm = (dmin * m).reshape((n_blocks, -1, 1)) - - ql = qs.reshape((n_blocks, -1, 1, 32)) >> torch.tensor([0, 4], device=d.device, dtype=torch.uint8).reshape((1, 1, 2, 1)) - qh = qh.reshape((n_blocks, -1, 1, 32)) >> torch.tensor([i for i in range(8)], device=d.device, dtype=torch.uint8).reshape((1, 1, 8, 1)) - ql = (ql & 0x0F).reshape((n_blocks, -1, 32)) - qh = (qh & 0x01).reshape((n_blocks, -1, 32)) - q = (ql | (qh << 4)) - - return (d * q - dm).reshape((n_blocks, QK_K)) - -def dequantize_blocks_Q4_K(blocks, block_size, type_size, dtype=None): - n_blocks = blocks.shape[0] - - d, dmin, scales, qs = split_block_dims(blocks, 2, 2, K_SCALE_SIZE) - d = d.view(torch.float16).to(dtype) - dmin = dmin.view(torch.float16).to(dtype) - - sc, m = get_scale_min(scales) - - d = (d * sc).reshape((n_blocks, -1, 1)) - dm = (dmin * m).reshape((n_blocks, -1, 1)) - - qs = qs.reshape((n_blocks, -1, 1, 32)) >> torch.tensor([0, 4], device=d.device, dtype=torch.uint8).reshape((1, 1, 2, 1)) - qs = (qs & 0x0F).reshape((n_blocks, -1, 32)) - - return (d * qs - dm).reshape((n_blocks, QK_K)) - -def dequantize_blocks_Q3_K(blocks, block_size, type_size, dtype=None): - n_blocks = blocks.shape[0] - - hmask, qs, scales, d = split_block_dims(blocks, QK_K // 8, QK_K // 4, 12) - d = d.view(torch.float16).to(dtype) - - lscales, hscales = scales[:, :8], scales[:, 8:] - lscales = lscales.reshape((n_blocks, 1, 8)) >> torch.tensor([0, 4], device=d.device, dtype=torch.uint8).reshape((1, 2, 1)) - lscales = lscales.reshape((n_blocks, 16)) - hscales = hscales.reshape((n_blocks, 1, 4)) >> torch.tensor([0, 2, 4, 6], device=d.device, dtype=torch.uint8).reshape((1, 4, 1)) - hscales = hscales.reshape((n_blocks, 16)) - scales = (lscales & 0x0F) | ((hscales & 0x03) << 4) - scales = (scales.to(torch.int8) - 32) - - dl = (d * scales).reshape((n_blocks, 16, 1)) - - ql = qs.reshape((n_blocks, -1, 1, 32)) >> torch.tensor([0, 2, 4, 6], device=d.device, dtype=torch.uint8).reshape((1, 1, 4, 1)) - qh = hmask.reshape(n_blocks, -1, 1, 32) >> torch.tensor([i for i in range(8)], device=d.device, dtype=torch.uint8).reshape((1, 1, 8, 1)) - ql = ql.reshape((n_blocks, 16, QK_K // 16)) & 3 - qh = (qh.reshape((n_blocks, 16, QK_K // 16)) & 1) ^ 1 - q = (ql.to(torch.int8) - (qh << 2).to(torch.int8)) - - return (dl * q).reshape((n_blocks, QK_K)) - -def dequantize_blocks_Q2_K(blocks, block_size, type_size, dtype=None): - n_blocks = blocks.shape[0] - - scales, qs, d, dmin = split_block_dims(blocks, QK_K // 16, QK_K // 4, 2) - d = d.view(torch.float16).to(dtype) - dmin = dmin.view(torch.float16).to(dtype) - - # (n_blocks, 16, 1) - dl = (d * (scales & 0xF)).reshape((n_blocks, QK_K // 16, 1)) - ml = (dmin * (scales >> 4)).reshape((n_blocks, QK_K // 16, 1)) - - shift = torch.tensor([0, 2, 4, 6], device=d.device, dtype=torch.uint8).reshape((1, 1, 4, 1)) - - qs = (qs.reshape((n_blocks, -1, 1, 32)) >> shift) & 3 - qs = qs.reshape((n_blocks, QK_K // 16, 16)) - qs = dl * qs - ml - - return qs.reshape((n_blocks, -1)) - -# IQ quants -KVALUES = torch.tensor([-127, -104, -83, -65, -49, -35, -22, -10, 1, 13, 25, 38, 53, 69, 89, 113], dtype=torch.int8) - -def dequantize_blocks_IQ4_NL(blocks, block_size, type_size, dtype=None): - n_blocks = blocks.shape[0] - - d, qs = split_block_dims(blocks, 2) - d = d.view(torch.float16).to(dtype) - - qs = qs.reshape((n_blocks, -1, 1, block_size//2)) >> torch.tensor([0, 4], device=d.device, dtype=torch.uint8).reshape((1, 1, 2, 1)) - qs = (qs & 0x0F).reshape((n_blocks, -1, 1)).to(torch.int64) - - kvalues = KVALUES.to(qs.device).expand(*qs.shape[:-1], 16) - qs = torch.gather(kvalues, dim=-1, index=qs).reshape((n_blocks, -1)) - del kvalues # should still be view, but just to be safe - - return (d * qs) - -def dequantize_blocks_IQ4_XS(blocks, block_size, type_size, dtype=None): - n_blocks = blocks.shape[0] - d, scales_h, scales_l, qs = split_block_dims(blocks, 2, 2, QK_K // 64) - d = d.view(torch.float16).to(dtype) - scales_h = to_uint16(scales_h) - - shift_a = torch.tensor([0, 4], device=d.device, dtype=torch.uint8).reshape((1, 1, 2)) - shift_b = torch.tensor([2 * i for i in range(QK_K // 32)], device=d.device, dtype=torch.uint8).reshape((1, -1, 1)) - - scales_l = scales_l.reshape((n_blocks, -1, 1)) >> shift_a.reshape((1, 1, 2)) - scales_h = scales_h.reshape((n_blocks, -1, 1)) >> shift_b.reshape((1, -1, 1)) - - scales_l = scales_l.reshape((n_blocks, -1)) & 0x0F - scales_h = scales_h.reshape((n_blocks, -1)).to(torch.uint8) & 0x03 - - scales = (scales_l | (scales_h << 4)).to(torch.int8) - 32 - dl = (d * scales.to(dtype)).reshape((n_blocks, -1, 1)) - - qs = qs.reshape((n_blocks, -1, 1, 16)) >> shift_a.reshape((1, 1, 2, 1)) - qs = qs.reshape((n_blocks, -1, 32, 1)) & 0x0F - - kvalues = KVALUES.to(qs.device).expand(*qs.shape[:-1], 16) - qs = torch.gather(kvalues, dim=-1, index=qs.to(torch.int64)).reshape((n_blocks, -1, 32)) - del kvalues # see IQ4_NL - del shift_a - del shift_b - - return (dl * qs).reshape((n_blocks, -1)) - -dequantize_functions = { - gguf.GGMLQuantizationType.BF16: dequantize_blocks_BF16, - gguf.GGMLQuantizationType.Q8_0: dequantize_blocks_Q8_0, - gguf.GGMLQuantizationType.Q5_1: dequantize_blocks_Q5_1, - gguf.GGMLQuantizationType.Q5_0: dequantize_blocks_Q5_0, - gguf.GGMLQuantizationType.Q4_1: dequantize_blocks_Q4_1, - gguf.GGMLQuantizationType.Q4_0: dequantize_blocks_Q4_0, - gguf.GGMLQuantizationType.Q6_K: dequantize_blocks_Q6_K, - gguf.GGMLQuantizationType.Q5_K: dequantize_blocks_Q5_K, - gguf.GGMLQuantizationType.Q4_K: dequantize_blocks_Q4_K, - gguf.GGMLQuantizationType.Q3_K: dequantize_blocks_Q3_K, - gguf.GGMLQuantizationType.Q2_K: dequantize_blocks_Q2_K, - gguf.GGMLQuantizationType.IQ4_NL: dequantize_blocks_IQ4_NL, - gguf.GGMLQuantizationType.IQ4_XS: dequantize_blocks_IQ4_XS, -} diff --git a/smart_loaders_mxd/loader.py b/smart_loaders_mxd/loader.py deleted file mode 100644 index 96ac363..0000000 --- a/smart_loaders_mxd/loader.py +++ /dev/null @@ -1,508 +0,0 @@ -# (c) City96 || Apache-2.0 (apache.org/licenses/LICENSE-2.0) -# Vendored verbatim from city96/ComfyUI-GGUF. See smart_loaders_mxd/UPSTREAM.md -# for the sync baseline and procedure. Do not hand-edit; port upstream diffs instead. -import warnings -import logging -import torch -import gguf -import re -import os - -from .ops import GGMLTensor -from .dequant import is_quantized, dequantize_tensor - -IMG_ARCH_LIST = {"flux", "sd1", "sdxl", "sd3", "aura", "hidream", "cosmos", "ltxv", "hyvid", "wan", "lumina2", "qwen_image"} -TXT_ARCH_LIST = {"t5", "t5encoder", "llama", "qwen2vl", "qwen3", "qwen3vl", "gemma3"} -VIS_TYPE_LIST = {"clip-vision", "mmproj"} - -def get_orig_shape(reader, tensor_name): - field_key = f"comfy.gguf.orig_shape.{tensor_name}" - field = reader.get_field(field_key) - if field is None: - return None - # Has original shape metadata, so we try to decode it. - if len(field.types) != 2 or field.types[0] != gguf.GGUFValueType.ARRAY or field.types[1] != gguf.GGUFValueType.INT32: - raise TypeError(f"Bad original shape metadata for {field_key}: Expected ARRAY of INT32, got {field.types}") - return torch.Size(tuple(int(field.parts[part_idx][0]) for part_idx in field.data)) - -def get_field(reader, field_name, field_type): - field = reader.get_field(field_name) - if field is None: - return None - elif field_type == str: - # extra check here as this is used for checking arch string - if len(field.types) != 1 or field.types[0] != gguf.GGUFValueType.STRING: - raise TypeError(f"Bad type for GGUF {field_name} key: expected string, got {field.types!r}") - return str(field.parts[field.data[-1]], encoding="utf-8") - elif field_type in [int, float, bool]: - return field_type(field.parts[field.data[-1]].item()) - else: - raise TypeError(f"Unknown field type {field_type}") - -def get_list_field(reader, field_name, field_type): - field = reader.get_field(field_name) - if field is None: - return None - elif field_type == str: - return tuple(str(field.parts[part_idx], encoding="utf-8") for part_idx in field.data) - elif field_type in [int, float, bool]: - return tuple(field_type(field.parts[part_idx][0]) for part_idx in field.data) - else: - raise TypeError(f"Unknown field type {field_type}") - -def get_gguf_metadata(reader): - """Extract all simple metadata fields like safetensors""" - metadata = {} - for field_name in reader.fields: - try: - field = reader.get_field(field_name) - if len(field.types) == 1: # Simple scalar fields only - if field.types[0] == gguf.GGUFValueType.STRING: - metadata[field_name] = str(field.parts[field.data[-1]], "utf-8") - elif field.types[0] == gguf.GGUFValueType.INT32: - metadata[field_name] = int(field.parts[field.data[-1]]) - elif field.types[0] == gguf.GGUFValueType.F32: - metadata[field_name] = float(field.parts[field.data[-1]]) - elif field.types[0] == gguf.GGUFValueType.BOOL: - metadata[field_name] = bool(field.parts[field.data[-1]]) - except: - continue - return metadata - -def gguf_sd_loader(path, handle_prefix="model.diffusion_model.", is_text_model=False): - """ - Read state dict as fake tensors - """ - reader = gguf.GGUFReader(path) - - # filter and strip prefix - has_prefix = False - if handle_prefix is not None: - prefix_len = len(handle_prefix) - tensor_names = set(tensor.name for tensor in reader.tensors) - has_prefix = any(s.startswith(handle_prefix) for s in tensor_names) - - tensors = [] - for tensor in reader.tensors: - sd_key = tensor_name = tensor.name - if has_prefix: - if not tensor_name.startswith(handle_prefix): - continue - sd_key = tensor_name[prefix_len:] - tensors.append((sd_key, tensor)) - - # detect and verify architecture - compat = None - arch_str = get_field(reader, "general.architecture", str) - type_str = get_field(reader, "general.type", str) - if arch_str in [None, "pig", "cow"]: - if is_text_model: - raise ValueError(f"This gguf file is incompatible with llama.cpp!\nConsider using safetensors or a compatible gguf file\n({path})") - compat = "sd.cpp" if arch_str is None else arch_str - # import here to avoid changes to convert.py breaking regular models - from .tools.convert import detect_arch - try: - arch_str = detect_arch(set(val[0] for val in tensors)).arch - except Exception as e: - raise ValueError(f"This model is not currently supported - ({e})") - elif arch_str not in TXT_ARCH_LIST and is_text_model: - if type_str not in VIS_TYPE_LIST: - raise ValueError(f"Unexpected text model architecture type in GGUF file: {arch_str!r}") - elif arch_str not in IMG_ARCH_LIST and not is_text_model: - raise ValueError(f"Unexpected architecture type in GGUF file: {arch_str!r}") - - if compat: - logging.warning(f"Warning: This gguf model file is loaded in compatibility mode '{compat}' [arch:{arch_str}]") - - # main loading loop - state_dict = {} - qtype_dict = {} - for sd_key, tensor in tensors: - tensor_name = tensor.name - # torch_tensor = torch.from_numpy(tensor.data) # mmap - - # NOTE: line above replaced with this block to avoid persistent numpy warning about mmap - with warnings.catch_warnings(): - warnings.filterwarnings("ignore", message="The given NumPy array is not writable") - torch_tensor = torch.from_numpy(tensor.data) # mmap - - shape = get_orig_shape(reader, tensor_name) - if shape is None: - shape = torch.Size(tuple(int(v) for v in reversed(tensor.shape))) - # Workaround for stable-diffusion.cpp SDXL detection. - if compat == "sd.cpp" and arch_str == "sdxl": - if any([tensor_name.endswith(x) for x in (".proj_in.weight", ".proj_out.weight")]): - while len(shape) > 2 and shape[-1] == 1: - shape = shape[:-1] - - # add to state dict - if tensor.tensor_type in {gguf.GGMLQuantizationType.F32, gguf.GGMLQuantizationType.F16}: - torch_tensor = torch_tensor.view(*shape) - state_dict[sd_key] = GGMLTensor(torch_tensor, tensor_type=tensor.tensor_type, tensor_shape=shape) - - # 1D tensors shouldn't be quantized, this is a fix for BF16 - if len(shape) <= 1 and tensor.tensor_type == gguf.GGMLQuantizationType.BF16: - state_dict[sd_key] = dequantize_tensor(state_dict[sd_key], dtype=torch.float32) - - # keep track of loaded tensor types - tensor_type_str = getattr(tensor.tensor_type, "name", repr(tensor.tensor_type)) - qtype_dict[tensor_type_str] = qtype_dict.get(tensor_type_str, 0) + 1 - - # print loaded tensor type counts - logging.info("gguf qtypes: " + ", ".join(f"{k} ({v})" for k, v in qtype_dict.items())) - - # mark largest tensor for vram estimation - qsd = {k:v for k,v in state_dict.items() if is_quantized(v)} - if len(qsd) > 0: - max_key = max(qsd.keys(), key=lambda k: qsd[k].numel()) - state_dict[max_key].is_largest_weight = True - - # extra info to return - extra = { - "arch_str": arch_str, - "metadata": get_gguf_metadata(reader) - } - return (state_dict, extra) - -# for remapping llama.cpp -> original key names -T5_SD_MAP = { - "enc.": "encoder.", - ".blk.": ".block.", - "token_embd": "shared", - "output_norm": "final_layer_norm", - "attn_q": "layer.0.SelfAttention.q", - "attn_k": "layer.0.SelfAttention.k", - "attn_v": "layer.0.SelfAttention.v", - "attn_o": "layer.0.SelfAttention.o", - "attn_norm": "layer.0.layer_norm", - "attn_rel_b": "layer.0.SelfAttention.relative_attention_bias", - "ffn_up": "layer.1.DenseReluDense.wi_1", - "ffn_down": "layer.1.DenseReluDense.wo", - "ffn_gate": "layer.1.DenseReluDense.wi_0", - "ffn_norm": "layer.1.layer_norm", -} - -LLAMA_SD_MAP = { - "blk.": "model.layers.", - "attn_norm": "input_layernorm", - "attn_q_norm.": "self_attn.q_norm.", - "attn_k_norm.": "self_attn.k_norm.", - "attn_v_norm.": "self_attn.v_norm.", - "attn_q": "self_attn.q_proj", - "attn_k": "self_attn.k_proj", - "attn_v": "self_attn.v_proj", - "attn_output": "self_attn.o_proj", - "ffn_up": "mlp.up_proj", - "ffn_down": "mlp.down_proj", - "ffn_gate": "mlp.gate_proj", - "ffn_norm": "post_attention_layernorm", - "token_embd": "model.embed_tokens", - "output_norm": "model.norm", - "output.weight": "lm_head.weight", -} - -GEMMA3_SD_MAP = LLAMA_SD_MAP.copy() -GEMMA3_SD_MAP.update({ - "ffn_norm": "pre_feedforward_layernorm", - "post_ffw_norm": "post_feedforward_layernorm", - "post_attention_norm": "post_attention_layernorm", -}) - -CLIP_VISION_SD_MAP = { - "mm.": "visual.merger.mlp.", - "v.post_ln.": "visual.merger.ln_q.", - "v.patch_embd": "visual.patch_embed.proj", - "v.blk.": "visual.blocks.", - "ffn_up": "mlp.up_proj", - "ffn_down": "mlp.down_proj", - "ffn_gate": "mlp.gate_proj", - "attn_out.": "attn.proj.", - "ln1.": "norm1.", - "ln2.": "norm2.", -} - -def sd_map_replace(raw_sd, key_map): - sd = {} - for k,v in raw_sd.items(): - for s,d in key_map.items(): - k = k.replace(s,d) - sd[k] = v - return sd - -def llama_permute(raw_sd, n_head, n_head_kv): - # Reverse version of LlamaModel.permute in llama.cpp convert script - sd = {} - permute = lambda x,h: x.reshape(h, x.shape[0] // h // 2, 2, *x.shape[1:]).swapaxes(1, 2).reshape(x.shape) - for k,v in raw_sd.items(): - if k.endswith(("q_proj.weight", "q_proj.bias")): - v.data = permute(v.data, n_head) - if k.endswith(("k_proj.weight", "k_proj.bias")): - v.data = permute(v.data, n_head_kv) - sd[k] = v - return sd - -def gemma3_norm_corrections(sd): - # Reverse change from Gemma3Model modify_tensors in llama.cpp convert script - norm_patterns = [ - "input_layernorm.weight", - "post_attention_layernorm.weight", - "pre_feedforward_layernorm.weight", - "post_feedforward_layernorm.weight", - "self_attn.q_norm.weight", - "self_attn.k_norm.weight", - "model.norm.weight" - ] - corrected = 0 - for key in list(sd.keys()): - if any(p in key for p in norm_patterns): - if is_quantized(sd[key]): - sd[key] = dequantize_tensor(sd[key], dtype=torch.float32) - 1.0 - else: - sd[key] = sd[key].float() - 1.0 - corrected += 1 - #logging.info(f"Gemma3: Applied -1 norm correction to {corrected} tensors") - return sd - -def strip_quant_suffix(name): - pattern = r"[-_]?(?:ud-)?i?q[0-9]_[a-z0-9_\-]{1,8}$" - match = re.search(pattern, name, re.IGNORECASE) - if match: - name = name[:match.start()] - return name - -def gguf_mmproj_loader(path): - # Reverse version of Qwen2VLVisionModel.modify_tensors - logging.info("Attenpting to find mmproj file for text encoder...") - - # get name to match w/o quant suffix - tenc_fname = os.path.basename(path) - tenc = os.path.splitext(tenc_fname)[0].lower() - tenc = strip_quant_suffix(tenc) - - # try and find matching mmproj - target = [] - root = os.path.dirname(path) - for fname in os.listdir(root): - name, ext = os.path.splitext(fname) - if ext.lower() != ".gguf": - continue - if "mmproj" not in name.lower(): - continue - if tenc in name.lower(): - target.append(fname) - - if len(target) == 0: - logging.error(f"Error: Can't find mmproj file for '{tenc_fname}' (matching:'{tenc}')! Qwen-Image-Edit will be broken!") - return {} - if len(target) > 1: - logging.error(f"Ambiguous mmproj for text encoder '{tenc_fname}', will use first match.") - - logging.info(f"Using mmproj '{target[0]}' for text encoder '{tenc_fname}'.") - target = os.path.join(root, target[0]) - vsd, _ = gguf_sd_loader(target, is_text_model=True) - - # concat 4D to 5D - if "v.patch_embd.weight.1" in vsd: - w1 = dequantize_tensor(vsd.pop("v.patch_embd.weight"), dtype=torch.float32) - w2 = dequantize_tensor(vsd.pop("v.patch_embd.weight.1"), dtype=torch.float32) - vsd["v.patch_embd.weight"] = torch.stack([w1, w2], dim=2) - - # run main replacement - vsd = sd_map_replace(vsd, CLIP_VISION_SD_MAP) - - # handle split Q/K/V - if "visual.blocks.0.attn_q.weight" in vsd: - attns = {} - # filter out attentions + group - for k,v in vsd.items(): - if any(x in k for x in ["attn_q", "attn_k", "attn_v"]): - k_attn, k_name = k.rsplit(".attn_", 1) - k_attn += ".attn.qkv." + k_name.split(".")[-1] - if k_attn not in attns: - attns[k_attn] = {} - attns[k_attn][k_name] = dequantize_tensor( - v, dtype=(torch.bfloat16 if is_quantized(v) else torch.float16) - ) - - # recombine - for k,v in attns.items(): - suffix = k.split(".")[-1] - vsd[k] = torch.cat([ - v[f"q.{suffix}"], - v[f"k.{suffix}"], - v[f"v.{suffix}"], - ], dim=0) - del attns - - return vsd - -def gguf_tokenizer_loader(path, temb_shape): - # convert gguf tokenizer to spiece - logging.info("Attempting to recreate sentencepiece tokenizer from GGUF file metadata...") - try: - from sentencepiece import sentencepiece_model_pb2 as model - except ImportError: - raise ImportError("Please make sure sentencepiece and protobuf are installed.\npip install sentencepiece protobuf") - spm = model.ModelProto() - - reader = gguf.GGUFReader(path) - - if get_field(reader, "tokenizer.ggml.model", str) == "t5": - if temb_shape == (256384, 4096): # probably UMT5 - spm.trainer_spec.model_type == 1 # Unigram (do we have a T5 w/ BPE?) - else: - raise NotImplementedError("Unknown model, can't set tokenizer!") - else: - raise NotImplementedError("Unknown model, can't set tokenizer!") - - spm.normalizer_spec.add_dummy_prefix = get_field(reader, "tokenizer.ggml.add_space_prefix", bool) - spm.normalizer_spec.remove_extra_whitespaces = get_field(reader, "tokenizer.ggml.remove_extra_whitespaces", bool) - - tokens = get_list_field(reader, "tokenizer.ggml.tokens", str) - scores = get_list_field(reader, "tokenizer.ggml.scores", float) - toktypes = get_list_field(reader, "tokenizer.ggml.token_type", int) - - for idx, (token, score, toktype) in enumerate(zip(tokens, scores, toktypes)): - # # These aren't present in the original? - # if toktype == 5 and idx >= temb_shape[0]%1000): - # continue - - piece = spm.SentencePiece() - piece.piece = token - piece.score = score - piece.type = toktype - spm.pieces.append(piece) - - # unsure if any of these are correct - spm.trainer_spec.byte_fallback = True - spm.trainer_spec.vocab_size = len(tokens) # split off unused? - spm.trainer_spec.max_sentence_length = 4096 - spm.trainer_spec.eos_id = get_field(reader, "tokenizer.ggml.eos_token_id", int) - spm.trainer_spec.pad_id = get_field(reader, "tokenizer.ggml.padding_token_id", int) - - logging.info(f"Created tokenizer with vocab size of {len(spm.pieces)}") - del reader - return torch.ByteTensor(list(spm.SerializeToString())) - -def gguf_tekken_tokenizer_loader(path, temb_shape): - # convert ggml (hf) tokenizer metadata to tekken/comfy data - logging.info("Attempting to recreate tekken tokenizer from GGUF file metadata...") - import json - import base64 - from transformers.convert_slow_tokenizer import bytes_to_unicode - - reader = gguf.GGUFReader(path) - - model_str = get_field(reader, "tokenizer.ggml.model", str) - if model_str == "gpt2": - if temb_shape == (131072, 5120): # probably Mistral - data = { - "config": {"num_vocab_tokens": 150000, "default_vocab_size": 131072}, - "vocab": [], - "special_tokens": [], - } - else: - raise NotImplementedError("Unknown model, can't set tokenizer!") - else: - raise NotImplementedError("Unknown model, can't set tokenizer!") - - tokens = get_list_field(reader, "tokenizer.ggml.tokens", str) - toktypes = get_list_field(reader, "tokenizer.ggml.token_type", int) - - decoder = {v: k for k, v in bytes_to_unicode().items()} - for idx, (token, toktype) in enumerate(zip(tokens, toktypes)): - if toktype == 3: - data["special_tokens"].append( - {'rank': idx, 'token_str': token, 'is_control': True} - ) - else: - tok = bytes([decoder[char] for char in token]) - data["vocab"].append({ - "rank": len(data["vocab"]), - "token_bytes": base64.b64encode(tok).decode("ascii"), - "token_str": tok.decode("utf-8", errors="replace") # ? - }) - - logging.info(f"Created tekken tokenizer with vocab size of {len(data['vocab'])} (+{len(data['special_tokens'])})") - del reader - return torch.ByteTensor(list(json.dumps(data).encode('utf-8'))) - -def gguf_gemma3_tokenizer_loader(path): - #TODO: merge into gguf_tokenizer_loader - logging.info("Attempting to recreate sentencepiece tokenizer from GGUF file metadata...") - try: - from sentencepiece import sentencepiece_model_pb2 as model - except ImportError: - raise ImportError("Please install sentencepiece and protobuf.\npip install sentencepiece protobuf") - spm = model.ModelProto() - reader = gguf.GGUFReader(path) - - spm.normalizer_spec.name = "identity" - spm.normalizer_spec.add_dummy_prefix = False - spm.trainer_spec.model_type = 2 - spm.trainer_spec.input_format = "tsv" - spm.trainer_spec.byte_fallback = True - spm.trainer_spec.max_sentence_length = 4192 - spm.trainer_spec.bos_piece = "" - - tokens = get_list_field(reader, "tokenizer.ggml.tokens", str) - scores = get_list_field(reader, "tokenizer.ggml.scores", float) - toktype = get_list_field(reader, "tokenizer.ggml.token_type", int) - - if not tokens or not scores or not toktype: - raise ValueError("Missing tokenizer metadata") - - for idx in range(len(tokens)): - piece = spm.SentencePiece() - piece.piece = tokens[idx] - if idx == 3: # UNK position - piece.type = 2 # UNK Token - piece.score = 0.0 # UNK Score - else: - piece.type = toktype[idx] - piece.score = scores[idx] - spm.pieces.append(piece) - - spm.trainer_spec.vocab_size = len(spm.pieces) - logging.info(f"Created tokenizer with vocab size of {len(spm.pieces)}") - - del reader - return torch.ByteTensor(list(spm.SerializeToString())) - -def gguf_clip_loader(path): - sd, extra = gguf_sd_loader(path, is_text_model=True) - arch = extra.get("arch_str", None) - if arch in {"t5", "t5encoder"}: - temb_key = "token_embd.weight" - if temb_key in sd and sd[temb_key].shape == (256384, 4096): - # non-standard Comfy-Org tokenizer - sd["spiece_model"] = gguf_tokenizer_loader(path, sd[temb_key].shape) - # TODO: dequantizing token embed here is janky but otherwise we OOM due to tensor being massive. - logging.warning(f"Dequantizing {temb_key} to prevent runtime OOM.") - sd[temb_key] = dequantize_tensor(sd[temb_key], dtype=torch.float16) - sd = sd_map_replace(sd, T5_SD_MAP) - elif arch in {"llama", "qwen2vl", "qwen3", "qwen3vl", "gemma3"}: - # TODO: pass model_options["vocab_size"] to loader somehow - temb_key = "token_embd.weight" - if temb_key in sd and sd[temb_key].shape[0] >= (64 * 1024): - if arch == "llama" and sd[temb_key].shape == (131072, 5120): - # non-standard Comfy-Org tokenizer - sd["tekken_model"] = gguf_tekken_tokenizer_loader(path, sd[temb_key].shape) - elif arch == "gemma3": - sd["spiece_model"] = gguf_gemma3_tokenizer_loader(path) - # See note above for T5. - logging.warning(f"Dequantizing {temb_key} to prevent runtime OOM.") - sd[temb_key] = dequantize_tensor(sd[temb_key], dtype=torch.float16) - if arch == "gemma3": - sd = sd_map_replace(sd, GEMMA3_SD_MAP) - sd = gemma3_norm_corrections(sd) - else: - sd = sd_map_replace(sd, LLAMA_SD_MAP) - if arch == "llama": - sd = llama_permute(sd, 32, 8) # L3 / Mistral - if arch == "qwen2vl": - vsd = gguf_mmproj_loader(path) - sd.update(vsd) - else: - pass - return sd diff --git a/smart_loaders_mxd/nodes.py b/smart_loaders_mxd/nodes.py deleted file mode 100644 index 999e72a..0000000 --- a/smart_loaders_mxd/nodes.py +++ /dev/null @@ -1,358 +0,0 @@ -# (c) City96 || Apache-2.0 (apache.org/licenses/LICENSE-2.0) -# Forked from https://github.com/city96/ComfyUI-GGUF -# Modified by Maxed-Out-99 -# -# Not vendored verbatim (see UPSTREAM.md) — this is our own unified -# loader layer on top of the vendored dequant/loader/ops engine files. -import logging -import collections -import inspect - -import comfy.sd -import comfy.lora -import comfy.float -import comfy.utils -import comfy.model_patcher -import comfy.model_management -import folder_paths -import nodes -import torch - -from .ops import GGMLOps, move_patch_to_device -from .loader import gguf_sd_loader, gguf_clip_loader -from .dequant import is_quantized, is_torch_compatible - -def update_folder_names_and_paths(key, targets=[]): - # check for existing key - base = folder_paths.folder_names_and_paths.get(key, ([], {})) - base = base[0] if isinstance(base[0], (list, set, tuple)) else [] - # find base key & add w/ fallback, sanity check + warning - target = next((x for x in targets if x in folder_paths.folder_names_and_paths), targets[0]) - orig, _ = folder_paths.folder_names_and_paths.get(target, ([], {})) - folder_paths.folder_names_and_paths[key] = (orig or base, {".gguf"}) - if base and base != orig: - logging.warning(f"Unknown file list already present on key {key}: {base}") - -# Add a custom keys for files ending in .gguf -update_folder_names_and_paths("unet_gguf", ["diffusion_models", "unet"]) -update_folder_names_and_paths("clip_gguf", ["text_encoders", "clip"]) - -class GGUFModelPatcher(comfy.model_patcher.ModelPatcher): - patch_on_device = False - - def patch_weight_to_device(self, key, device_to=None, inplace_update=False): - if key not in self.patches: - return - weight = comfy.utils.get_attr(self.model, key) - - patches = self.patches[key] - if is_quantized(weight): - out_weight = weight.to(device_to) - patches = move_patch_to_device(patches, self.load_device if self.patch_on_device else self.offload_device) - # TODO: do we ever have legitimate duplicate patches? (i.e. patch on top of patched weight) - out_weight.patches = [(patches, key)] - else: - inplace_update = self.weight_inplace_update or inplace_update - if key not in self.backup: - self.backup[key] = collections.namedtuple('Dimension', ['weight', 'inplace_update'])( - weight.to(device=self.offload_device, copy=inplace_update), inplace_update - ) - - if device_to is not None: - temp_weight = comfy.model_management.cast_to_device(weight, device_to, torch.float32, copy=True) - else: - temp_weight = weight.to(torch.float32, copy=True) - - out_weight = comfy.lora.calculate_weight(patches, temp_weight, key) - out_weight = comfy.float.stochastic_rounding(out_weight, weight.dtype) - - if inplace_update: - comfy.utils.copy_to_param(self.model, key, out_weight) - else: - comfy.utils.set_attr_param(self.model, key, out_weight) - - def unpatch_model(self, device_to=None, unpatch_weights=True): - if unpatch_weights: - for p in self.model.parameters(): - if is_torch_compatible(p): - continue - patches = getattr(p, "patches", []) - if len(patches) > 0: - p.patches = [] - # TODO: Find another way to not unload after patches - return super().unpatch_model(device_to=device_to, unpatch_weights=unpatch_weights) - - def pin_weight_to_device(self, key): - op_key = key.rsplit('.', 1)[0] - if not self.mmap_released and op_key in self.named_modules_to_munmap: - # TODO: possible to OOM, find better way to detach - self.named_modules_to_munmap[op_key].to(self.load_device).to(self.offload_device) - del self.named_modules_to_munmap[op_key] - super().pin_weight_to_device(key) - - mmap_released = False - named_modules_to_munmap = {} - - def load(self, *args, force_patch_weights=False, **kwargs): - if not self.mmap_released: - self.named_modules_to_munmap = dict(self.model.named_modules()) - - # always call `patch_weight_to_device` even for lowvram - super().load(*args, force_patch_weights=True, **kwargs) - - # make sure nothing stays linked to mmap after first load - if not self.mmap_released: - linked = [] - if kwargs.get("lowvram_model_memory", 0) > 0: - for n, m in self.named_modules_to_munmap.items(): - if hasattr(m, "weight"): - device = getattr(m.weight, "device", None) - if device == self.offload_device: - linked.append((n, m)) - continue - if hasattr(m, "bias"): - device = getattr(m.bias, "device", None) - if device == self.offload_device: - linked.append((n, m)) - continue - if linked and self.load_device != self.offload_device: - logging.info(f"Attempting to release mmap ({len(linked)})") - for n, m in linked: - # TODO: possible to OOM, find better way to detach - m.to(self.load_device).to(self.offload_device) - self.mmap_released = True - self.named_modules_to_munmap = {} - - def clone(self, *args, **kwargs): - src_cls = self.__class__ - self.__class__ = GGUFModelPatcher - n = super().clone(*args, **kwargs) - n.__class__ = GGUFModelPatcher - self.__class__ = src_cls - # GGUF specific clone values below - n.patch_on_device = getattr(self, "patch_on_device", False) - n.mmap_released = getattr(self, "mmap_released", False) - if src_cls != GGUFModelPatcher: - n.size = 0 # force recalc - return n - -class UNETLoaderUnified: - @classmethod - def INPUT_TYPES(s): - # Combine both lists: standard + gguf - model_list = folder_paths.get_filename_list("diffusion_models") + folder_paths.get_filename_list("unet_gguf") - return { - "required": { - "unet_name": (sorted(set(model_list)),), - } - } - - RETURN_TYPES = ("MODEL",) - FUNCTION = "load_unet" - CATEGORY = "advanced/loaders" - TITLE = "Smart Model Loader MXD" - - def load_unet(self, unet_name): - if unet_name.endswith(".gguf"): - return self.load_gguf_unet(unet_name) - else: - return self.load_standard_unet(unet_name) - - def load_standard_unet(self, unet_name): - unet_path = folder_paths.get_full_path_or_raise("diffusion_models", unet_name) - model = comfy.sd.load_diffusion_model(unet_path) - return (model,) - - def load_gguf_unet(self, unet_name): - unet_path = folder_paths.get_full_path("unet", unet_name) - ops = GGMLOps() - - # Load state dict from GGUF - loaded = gguf_sd_loader(unet_path) - if loaded is None: - raise RuntimeError(f"Failed to load GGUF model: {unet_path}") - - if isinstance(loaded, tuple) and len(loaded) == 2: - sd, extra = loaded - else: - sd, extra = loaded, {} - - kwargs = {} - valid_params = inspect.signature(comfy.sd.load_diffusion_model_state_dict).parameters - if "metadata" in valid_params: - kwargs["metadata"] = extra.get("metadata", {}) - - model = comfy.sd.load_diffusion_model_state_dict( - sd, model_options={"custom_operations": ops}, **kwargs - ) - if model is None: - raise RuntimeError(f"Could not detect GGUF model type for: {unet_path}") - model = GGUFModelPatcher.clone(model) - return (model,) - -########################################################################################################################################## - -class SmartCLIPLoaderBase: - def get_clip_paths(self, *clip_names): - return [folder_paths.get_full_path("clip", name) for name in clip_names] - - def load_clip_models(self, clip_paths, clip_type): - if all(path is not None and not path.endswith(".gguf") for path in clip_paths): - clip = comfy.sd.load_clip( - ckpt_paths=clip_paths, - embedding_directory=folder_paths.get_folder_paths("embeddings"), - clip_type=clip_type, - model_options={}, - ) - return (clip,) - - clip_data = [] - use_ggml_ops = False - - for path in clip_paths: - if path.endswith(".gguf"): - sd = gguf_clip_loader(path) - use_ggml_ops = True - else: - sd, metadata = comfy.utils.load_torch_file(path, safe_load=True, return_metadata=True) - sd, _metadata = comfy.utils.convert_old_quants(sd, model_prefix="", metadata=metadata) - clip_data.append(sd) - - model_options = { - "initial_device": comfy.model_management.text_encoder_offload_device(), - } - if use_ggml_ops: - model_options["custom_operations"] = GGMLOps - - clip = comfy.sd.load_text_encoder_state_dicts( - clip_type=clip_type, - state_dicts=clip_data, - model_options=model_options, - embedding_directory=folder_paths.get_folder_paths("embeddings"), - ) - - if use_ggml_ops and hasattr(clip, "patcher") and clip.patcher is not None: - clip.patcher = GGUFModelPatcher.clone(clip.patcher) - - return (clip,) - -########################################################################################################################################## - -class DualCLIPLoaderUnified(SmartCLIPLoaderBase): - @classmethod - def INPUT_TYPES(cls): - clip_files = folder_paths.get_filename_list("text_encoders") + folder_paths.get_filename_list("clip_gguf") - clip_options = (sorted(set(clip_files)),) - return { - "required": { - "clip_name1": clip_options, - "clip_name2": clip_options, - "type": nodes.DualCLIPLoader.INPUT_TYPES()["required"]["type"], - } - } - - RETURN_TYPES = ("CLIP",) - FUNCTION = "load_clip" - CATEGORY = "advanced/loaders" - TITLE = "Smart Dual CLIP Loader MXD" - - def load_clip(self, clip_name1, clip_name2, type): - clip_type = getattr(comfy.sd.CLIPType, type.upper(), comfy.sd.CLIPType.STABLE_DIFFUSION) - clip_paths = self.get_clip_paths(clip_name1, clip_name2) - return self.load_clip_models(clip_paths, clip_type) - -########################################################################################################################################## - -class TripleCLIPLoaderUnified(SmartCLIPLoaderBase): - @classmethod - def INPUT_TYPES(cls): - clip_files = folder_paths.get_filename_list("text_encoders") + folder_paths.get_filename_list("clip_gguf") - clip_options = (sorted(set(clip_files)),) - return { - "required": { - "clip_name1": clip_options, - "clip_name2": clip_options, - "clip_name3": clip_options, - "type": nodes.CLIPLoader.INPUT_TYPES()["required"]["type"], - } - } - - RETURN_TYPES = ("CLIP",) - FUNCTION = "load_clip" - CATEGORY = "advanced/loaders" - TITLE = "Smart Triple CLIP Loader MXD" - - def load_clip(self, clip_name1, clip_name2, clip_name3, type): - clip_type = getattr(comfy.sd.CLIPType, type.upper(), comfy.sd.CLIPType.STABLE_DIFFUSION) - clip_paths = self.get_clip_paths(clip_name1, clip_name2, clip_name3) - return self.load_clip_models(clip_paths, clip_type) - -########################################################################################################################################## - - -class QuadrupleCLIPLoaderUnified(SmartCLIPLoaderBase): - @classmethod - def INPUT_TYPES(cls): - clip_files = folder_paths.get_filename_list("text_encoders") + folder_paths.get_filename_list("clip_gguf") - clip_options = (sorted(set(clip_files)),) - return { - "required": { - "clip_name1": clip_options, - "clip_name2": clip_options, - "clip_name3": clip_options, - "clip_name4": clip_options, - "type": nodes.CLIPLoader.INPUT_TYPES()["required"]["type"], - } - } - - RETURN_TYPES = ("CLIP",) - FUNCTION = "load_clip" - CATEGORY = "advanced/loaders" - TITLE = "Smart Quad CLIP Loader MXD" - - def load_clip(self, clip_name1, clip_name2, clip_name3, clip_name4, type): - clip_type = getattr(comfy.sd.CLIPType, type.upper(), comfy.sd.CLIPType.STABLE_DIFFUSION) - clip_paths = self.get_clip_paths(clip_name1, clip_name2, clip_name3, clip_name4) - return self.load_clip_models(clip_paths, clip_type) - -########################################################################################################################################## - -class CLIPLoaderUnified(SmartCLIPLoaderBase): - @classmethod - def INPUT_TYPES(cls): - clip_files = folder_paths.get_filename_list("text_encoders") + folder_paths.get_filename_list("clip_gguf") - clip_options = (sorted(set(clip_files)),) - return { - "required": { - "clip_name1": clip_options, - "type": nodes.CLIPLoader.INPUT_TYPES()["required"]["type"], - } - } - - RETURN_TYPES = ("CLIP",) - FUNCTION = "load_clip" - CATEGORY = "advanced/loaders" - TITLE = "Smart CLIP Loader MXD" - - def load_clip(self, clip_name1, type): - clip_type = getattr(comfy.sd.CLIPType, type.upper(), comfy.sd.CLIPType.STABLE_DIFFUSION) - return self.load_clip_models(self.get_clip_paths(clip_name1), clip_type) - -########################################################################################################################################## - -NODE_CLASS_MAPPINGS = { - # Keep the original keys for backward compatibility with existing ComfyUI workflows. - "UNETLoaderUnified": UNETLoaderUnified, - "CLIPLoaderUnified": CLIPLoaderUnified, - "DualCLIPLoaderUnified": DualCLIPLoaderUnified, - "TripleCLIPLoaderUnified": TripleCLIPLoaderUnified, - "QuadrupleCLIPLoaderUnified": QuadrupleCLIPLoaderUnified, -} - -NODE_DISPLAY_NAME_MAPPINGS = { - "UNETLoaderUnified": "Smart Model Loader MXD", - "CLIPLoaderUnified": "Smart CLIP Loader MXD", - "DualCLIPLoaderUnified": "Smart Dual CLIP Loader MXD", - "TripleCLIPLoaderUnified": "Smart Triple CLIP Loader MXD", - "QuadrupleCLIPLoaderUnified": "Smart Quad CLIP Loader MXD", -} diff --git a/smart_loaders_mxd/ops.py b/smart_loaders_mxd/ops.py deleted file mode 100644 index bfafa13..0000000 --- a/smart_loaders_mxd/ops.py +++ /dev/null @@ -1,283 +0,0 @@ -# (c) City96 || Apache-2.0 (apache.org/licenses/LICENSE-2.0) -# Vendored verbatim from city96/ComfyUI-GGUF. See smart_loaders_mxd/UPSTREAM.md -# for the sync baseline and procedure. Do not hand-edit; port upstream diffs instead. -import gguf -import torch -import logging - -import comfy.ops -import comfy.lora -import comfy.model_management -from .dequant import dequantize_tensor, is_quantized - -def chained_hasattr(obj, chained_attr): - probe = obj - for attr in chained_attr.split('.'): - if hasattr(probe, attr): - probe = getattr(probe, attr) - else: - return False - return True - -# A bakcward and forward compatible way to get `torch.compiler.disable`. -def get_torch_compiler_disable_decorator(): - def dummy_decorator(*args, **kwargs): - def noop(x): - return x - return noop - - from packaging import version - - if not chained_hasattr(torch, "compiler.disable"): - logging.info("ComfyUI-GGUF: Torch too old for torch.compile - bypassing") - return dummy_decorator # torch too old - elif version.parse(torch.__version__) >= version.parse("2.8"): - logging.info("ComfyUI-GGUF: Allowing full torch compile") - return dummy_decorator # torch compile works - if chained_hasattr(torch, "_dynamo.config.nontraceable_tensor_subclasses"): - logging.info("ComfyUI-GGUF: Allowing full torch compile (nightly)") - return dummy_decorator # torch compile works, nightly before 2.8 release - else: - logging.info("ComfyUI-GGUF: Partial torch compile only, consider updating pytorch") - return torch.compiler.disable - -torch_compiler_disable = get_torch_compiler_disable_decorator() - -class GGMLTensor(torch.Tensor): - """ - Main tensor-like class for storing quantized weights - """ - def __init__(self, *args, tensor_type, tensor_shape, patches=[], **kwargs): - super().__init__() - self.tensor_type = tensor_type - self.tensor_shape = tensor_shape - self.patches = patches - - def __new__(cls, *args, tensor_type, tensor_shape, patches=[], **kwargs): - return super().__new__(cls, *args, **kwargs) - - def to(self, *args, **kwargs): - new = super().to(*args, **kwargs) - new.tensor_type = getattr(self, "tensor_type", None) - new.tensor_shape = getattr(self, "tensor_shape", new.data.shape) - new.patches = getattr(self, "patches", []).copy() - return new - - def clone(self, *args, **kwargs): - return self - - def detach(self, *args, **kwargs): - return self - - def copy_(self, *args, **kwargs): - # fixes .weight.copy_ in comfy/clip_model/CLIPTextModel - try: - return super().copy_(*args, **kwargs) - except Exception as e: - logging.warning(f"ignoring 'copy_' on tensor: {e}") - - def new_empty(self, size, *args, **kwargs): - # Intel Arc fix, ref#50 - new_tensor = super().new_empty(size, *args, **kwargs) - return GGMLTensor( - new_tensor, - tensor_type = getattr(self, "tensor_type", None), - tensor_shape = size, - patches = getattr(self, "patches", []).copy() - ) - - @property - def shape(self): - if not hasattr(self, "tensor_shape"): - self.tensor_shape = self.size() - return self.tensor_shape - -class GGMLLayer(torch.nn.Module): - """ - This (should) be responsible for de-quantizing on the fly - """ - comfy_cast_weights = True - dequant_dtype = None - patch_dtype = None - largest_layer = False - torch_compatible_tensor_types = {None, gguf.GGMLQuantizationType.F32, gguf.GGMLQuantizationType.F16} - - def is_ggml_quantized(self, *, weight=None, bias=None): - if weight is None: - weight = self.weight - if bias is None: - bias = self.bias - return is_quantized(weight) or is_quantized(bias) - - def _load_from_state_dict(self, state_dict, prefix, *args, **kwargs): - weight, bias = state_dict.get(f"{prefix}weight"), state_dict.get(f"{prefix}bias") - # NOTE: using modified load for linear due to not initializing on creation, see GGMLOps todo - if self.is_ggml_quantized(weight=weight, bias=bias) or isinstance(self, torch.nn.Linear): - return self.ggml_load_from_state_dict(state_dict, prefix, *args, **kwargs) - # Not strictly required, but fixes embedding shape mismatch. Threshold set in loader.py - if isinstance(self, torch.nn.Embedding) and self.weight.shape[0] >= (64 * 1024): - return self.ggml_load_from_state_dict(state_dict, prefix, *args, **kwargs) - return super()._load_from_state_dict(state_dict, prefix, *args, **kwargs) - - def ggml_load_from_state_dict(self, state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs): - prefix_len = len(prefix) - for k,v in state_dict.items(): - if k[prefix_len:] == "weight": - self.weight = torch.nn.Parameter(v, requires_grad=False) - elif k[prefix_len:] == "bias" and v is not None: - self.bias = torch.nn.Parameter(v, requires_grad=False) - else: - unexpected_keys.append(k) - - # For Linear layer with missing weight - if self.weight is None and isinstance(self, torch.nn.Linear): - v = torch.zeros(self.in_features, self.out_features) - self.weight = torch.nn.Parameter(v, requires_grad=False) - missing_keys.append(prefix+"weight") - - # for vram estimation (TODO: less fragile logic?) - if getattr(self.weight, "is_largest_weight", False): - self.largest_layer = True - - def _save_to_state_dict(self, *args, **kwargs): - if self.is_ggml_quantized(): - return self.ggml_save_to_state_dict(*args, **kwargs) - return super()._save_to_state_dict(*args, **kwargs) - - def ggml_save_to_state_dict(self, destination, prefix, keep_vars): - # This is a fake state dict for vram estimation - weight = torch.zeros_like(self.weight, device=torch.device("meta")) - destination[prefix + "weight"] = weight - if self.bias is not None: - bias = torch.zeros_like(self.bias, device=torch.device("meta")) - destination[prefix + "bias"] = bias - - # Take into account space required for dequantizing the largest tensor - if self.largest_layer: - shape = getattr(self.weight, "tensor_shape", self.weight.shape) - dtype = self.dequant_dtype if self.dequant_dtype and self.dequant_dtype != "target" else torch.float16 - temp = torch.empty(*shape, device=torch.device("meta"), dtype=dtype) - destination[prefix + "temp.weight"] = temp - - return - # This would return the dequantized state dict - destination[prefix + "weight"] = self.get_weight(self.weight) - if bias is not None: - destination[prefix + "bias"] = self.get_weight(self.bias) - - def get_weight(self, tensor, dtype): - if tensor is None: - return - - # consolidate and load patches to GPU in async - patch_list = [] - device = tensor.device - for patches, key in getattr(tensor, "patches", []): - patch_list += move_patch_to_device(patches, device) - - # dequantize tensor while patches load - weight = dequantize_tensor(tensor, dtype, self.dequant_dtype) - - # prevent propagating custom tensor class - if isinstance(weight, GGMLTensor): - weight = torch.Tensor(weight) - - # apply patches - if len(patch_list) > 0: - if self.patch_dtype is None: - weight = comfy.lora.calculate_weight(patch_list, weight, key) - else: - # for testing, may degrade image quality - patch_dtype = dtype if self.patch_dtype == "target" else self.patch_dtype - weight = comfy.lora.calculate_weight(patch_list, weight, key, patch_dtype) - return weight - - @torch_compiler_disable() - def cast_bias_weight(s, input=None, dtype=None, device=None, bias_dtype=None): - if input is not None: - if dtype is None: - dtype = getattr(input, "dtype", torch.float32) - if bias_dtype is None: - bias_dtype = dtype - if device is None: - device = input.device - - bias = None - non_blocking = comfy.model_management.device_supports_non_blocking(device) - if s.bias is not None: - bias = s.get_weight(s.bias.to(device), dtype) - bias = comfy.ops.cast_to(bias, bias_dtype, device, non_blocking=non_blocking, copy=False) - - weight = s.get_weight(s.weight.to(device), dtype) - weight = comfy.ops.cast_to(weight, dtype, device, non_blocking=non_blocking, copy=False) - return weight, bias - - def forward_comfy_cast_weights(self, input, *args, **kwargs): - if self.is_ggml_quantized(): - out = self.forward_ggml_cast_weights(input, *args, **kwargs) - else: - out = super().forward_comfy_cast_weights(input, *args, **kwargs) - - # non-ggml forward might still propagate custom tensor class - if isinstance(out, GGMLTensor): - out = torch.Tensor(out) - return out - - def forward_ggml_cast_weights(self, input): - raise NotImplementedError - -class GGMLOps(comfy.ops.manual_cast): - """ - Dequantize weights on the fly before doing the compute - """ - class Linear(GGMLLayer, comfy.ops.manual_cast.Linear): - def __init__(self, in_features, out_features, bias=True, device=None, dtype=None): - torch.nn.Module.__init__(self) - # TODO: better workaround for reserved memory spike on windows - # Issue is with `torch.empty` still reserving the full memory for the layer - # Windows doesn't over-commit memory so without this 24GB+ of pagefile is used - self.in_features = in_features - self.out_features = out_features - self.weight = None - self.bias = None - - def forward_ggml_cast_weights(self, input): - weight, bias = self.cast_bias_weight(input) - return torch.nn.functional.linear(input, weight, bias) - - class Conv2d(GGMLLayer, comfy.ops.manual_cast.Conv2d): - def forward_ggml_cast_weights(self, input): - weight, bias = self.cast_bias_weight(input) - return self._conv_forward(input, weight, bias) - - class Embedding(GGMLLayer, comfy.ops.manual_cast.Embedding): - def forward_ggml_cast_weights(self, input, out_dtype=None): - output_dtype = out_dtype - if self.weight.dtype == torch.float16 or self.weight.dtype == torch.bfloat16: - out_dtype = None - weight, _bias = self.cast_bias_weight(self, device=input.device, dtype=out_dtype) - return torch.nn.functional.embedding( - input, weight, self.padding_idx, self.max_norm, self.norm_type, self.scale_grad_by_freq, self.sparse - ).to(dtype=output_dtype) - - class LayerNorm(GGMLLayer, comfy.ops.manual_cast.LayerNorm): - def forward_ggml_cast_weights(self, input): - if self.weight is None: - return super().forward_comfy_cast_weights(input) - weight, bias = self.cast_bias_weight(input) - return torch.nn.functional.layer_norm(input, self.normalized_shape, weight, bias, self.eps) - - class GroupNorm(GGMLLayer, comfy.ops.manual_cast.GroupNorm): - def forward_ggml_cast_weights(self, input): - weight, bias = self.cast_bias_weight(input) - return torch.nn.functional.group_norm(input, self.num_groups, weight, bias, self.eps) - -def move_patch_to_device(item, device): - if isinstance(item, torch.Tensor): - return item.to(device, non_blocking=True) - elif isinstance(item, tuple): - return tuple(move_patch_to_device(x, device) for x in item) - elif isinstance(item, list): - return [move_patch_to_device(x, device) for x in item] - else: - return item diff --git a/smart_loaders_mxd/tools/convert.py b/smart_loaders_mxd/tools/convert.py deleted file mode 100644 index 69d636c..0000000 --- a/smart_loaders_mxd/tools/convert.py +++ /dev/null @@ -1,369 +0,0 @@ -# (c) City96 || Apache-2.0 (apache.org/licenses/LICENSE-2.0) -# Vendored verbatim from city96/ComfyUI-GGUF. See smart_loaders_mxd/UPSTREAM.md -# for the sync baseline and procedure. Do not hand-edit; port upstream diffs instead. -# Only detect_arch() is used at runtime (loader.py's compat-mode fallback); the -# rest of this module (CLI conversion script) is kept intact for parity with upstream. -import os -import gguf -import torch -import logging -import argparse -from tqdm import tqdm -from safetensors.torch import load_file, save_file - -QUANTIZATION_THRESHOLD = 1024 -REARRANGE_THRESHOLD = 512 -MAX_TENSOR_NAME_LENGTH = 127 -MAX_TENSOR_DIMS = 4 - -class ModelTemplate: - arch = "invalid" # string describing architecture - shape_fix = False # whether to reshape tensors - keys_detect = [] # list of lists to match in state dict - keys_banned = [] # list of keys that should mark model as invalid for conversion - keys_hiprec = [] # list of keys that need to be kept in fp32 for some reason - keys_ignore = [] # list of strings to ignore keys by when found - - def handle_nd_tensor(self, key, data): - raise NotImplementedError(f"Tensor detected that exceeds dims supported by C++ code! ({key} @ {data.shape})") - -class ModelFlux(ModelTemplate): - arch = "flux" - keys_detect = [ - ("transformer_blocks.0.attn.norm_added_k.weight",), - ("double_blocks.0.img_attn.proj.weight",), - ] - keys_banned = ["transformer_blocks.0.attn.norm_added_k.weight",] - -class ModelSD3(ModelTemplate): - arch = "sd3" - keys_detect = [ - ("transformer_blocks.0.attn.add_q_proj.weight",), - ("joint_blocks.0.x_block.attn.qkv.weight",), - ] - keys_banned = ["transformer_blocks.0.attn.add_q_proj.weight",] - -class ModelAura(ModelTemplate): - arch = "aura" - keys_detect = [ - ("double_layers.3.modX.1.weight",), - ("joint_transformer_blocks.3.ff_context.out_projection.weight",), - ] - keys_banned = ["joint_transformer_blocks.3.ff_context.out_projection.weight",] - -class ModelHiDream(ModelTemplate): - arch = "hidream" - keys_detect = [ - ( - "caption_projection.0.linear.weight", - "double_stream_blocks.0.block.ff_i.shared_experts.w3.weight" - ) - ] - keys_hiprec = [ - # nn.parameter, can't load from BF16 ver - ".ff_i.gate.weight", - "img_emb.emb_pos" - ] - -class CosmosPredict2(ModelTemplate): - arch = "cosmos" - keys_detect = [ - ( - "blocks.0.mlp.layer1.weight", - "blocks.0.adaln_modulation_cross_attn.1.weight", - ) - ] - keys_hiprec = ["pos_embedder"] - keys_ignore = ["_extra_state", "accum_"] - -class ModelHyVid(ModelTemplate): - arch = "hyvid" - keys_detect = [ - ( - "double_blocks.0.img_attn_proj.weight", - "txt_in.individual_token_refiner.blocks.1.self_attn_qkv.weight", - ) - ] - - def handle_nd_tensor(self, key, data): - # hacky but don't have any better ideas - path = f"./fix_5d_tensors_{self.arch}.safetensors" # TODO: somehow get a path here?? - if os.path.isfile(path): - raise RuntimeError(f"5D tensor fix file already exists! {path}") - fsd = {key: torch.from_numpy(data)} - tqdm.write(f"5D key found in state dict! Manual fix required! - {key} {data.shape}") - save_file(fsd, path) - -class ModelWan(ModelHyVid): - arch = "wan" - keys_detect = [ - ( - "blocks.0.self_attn.norm_q.weight", - "text_embedding.2.weight", - "head.modulation", - ) - ] - keys_hiprec = [ - ".modulation" # nn.parameter, can't load from BF16 ver - ] - -class ModelLTXV(ModelTemplate): - arch = "ltxv" - keys_detect = [ - ( - "adaln_single.emb.timestep_embedder.linear_2.weight", - "transformer_blocks.27.scale_shift_table", - "caption_projection.linear_2.weight", - ) - ] - keys_hiprec = [ - "scale_shift_table" # nn.parameter, can't load from BF16 base quant - ] - -class ModelSDXL(ModelTemplate): - arch = "sdxl" - shape_fix = True - keys_detect = [ - ("down_blocks.0.downsamplers.0.conv.weight", "add_embedding.linear_1.weight",), - ( - "input_blocks.3.0.op.weight", "input_blocks.6.0.op.weight", - "output_blocks.2.2.conv.weight", "output_blocks.5.2.conv.weight", - ), # Non-diffusers - ("label_emb.0.0.weight",), - ] - -class ModelSD1(ModelTemplate): - arch = "sd1" - shape_fix = True - keys_detect = [ - ("down_blocks.0.downsamplers.0.conv.weight",), - ( - "input_blocks.3.0.op.weight", "input_blocks.6.0.op.weight", "input_blocks.9.0.op.weight", - "output_blocks.2.1.conv.weight", "output_blocks.5.2.conv.weight", "output_blocks.8.2.conv.weight" - ), # Non-diffusers - ] - -class ModelLumina2(ModelTemplate): - arch = "lumina2" - keys_detect = [ - ("cap_embedder.1.weight", "context_refiner.0.attention.qkv.weight") - ] - -arch_list = [ModelFlux, ModelSD3, ModelAura, ModelHiDream, CosmosPredict2, - ModelLTXV, ModelHyVid, ModelWan, ModelSDXL, ModelSD1, ModelLumina2] - -def is_model_arch(model, state_dict): - # check if model is correct - matched = False - invalid = False - for match_list in model.keys_detect: - if all(key in state_dict for key in match_list): - matched = True - invalid = any(key in state_dict for key in model.keys_banned) - break - assert not invalid, "Model architecture not allowed for conversion! (i.e. reference VS diffusers format)" - return matched - -def detect_arch(state_dict): - model_arch = None - for arch in arch_list: - if is_model_arch(arch, state_dict): - model_arch = arch() - break - assert model_arch is not None, "Unknown model architecture!" - return model_arch - -def parse_args(): - parser = argparse.ArgumentParser(description="Generate F16 GGUF files from single UNET") - parser.add_argument("--src", required=True, help="Source model ckpt file.") - parser.add_argument("--dst", help="Output unet gguf file.") - args = parser.parse_args() - - if not os.path.isfile(args.src): - parser.error("No input provided!") - - return args - -def strip_prefix(state_dict): - # prefix for mixed state dict - prefix = None - for pfx in ["model.diffusion_model.", "model."]: - if any([x.startswith(pfx) for x in state_dict.keys()]): - prefix = pfx - break - - # prefix for uniform state dict - if prefix is None: - for pfx in ["net."]: - if all([x.startswith(pfx) for x in state_dict.keys()]): - prefix = pfx - break - - # strip prefix if found - if prefix is not None: - logging.info(f"State dict prefix found: '{prefix}'") - sd = {} - for k, v in state_dict.items(): - if prefix not in k: - continue - k = k.replace(prefix, "") - sd[k] = v - else: - logging.debug("State dict has no prefix") - sd = state_dict - - return sd - -def load_state_dict(path): - if any(path.endswith(x) for x in [".ckpt", ".pt", ".bin", ".pth"]): - state_dict = torch.load(path, map_location="cpu", weights_only=True) - for subkey in ["model", "module"]: - if subkey in state_dict: - state_dict = state_dict[subkey] - break - if len(state_dict) < 20: - raise RuntimeError(f"pt subkey load failed: {state_dict.keys()}") - else: - state_dict = load_file(path) - - return strip_prefix(state_dict) - -def handle_tensors(writer, state_dict, model_arch): - name_lengths = tuple(sorted( - ((key, len(key)) for key in state_dict.keys()), - key=lambda item: item[1], - reverse=True, - )) - if not name_lengths: - return - max_name_len = name_lengths[0][1] - if max_name_len > MAX_TENSOR_NAME_LENGTH: - bad_list = ", ".join(f"{key!r} ({namelen})" for key, namelen in name_lengths if namelen > MAX_TENSOR_NAME_LENGTH) - raise ValueError(f"Can only handle tensor names up to {MAX_TENSOR_NAME_LENGTH} characters. Tensors exceeding the limit: {bad_list}") - for key, data in tqdm(state_dict.items()): - old_dtype = data.dtype - - if any(x in key for x in model_arch.keys_ignore): - tqdm.write(f"Filtering ignored key: '{key}'") - continue - - if data.dtype == torch.bfloat16: - data = data.to(torch.float32).numpy() - # this is so we don't break torch 2.0.X - elif data.dtype in [getattr(torch, "float8_e4m3fn", "_invalid"), getattr(torch, "float8_e5m2", "_invalid")]: - data = data.to(torch.float16).numpy() - else: - data = data.numpy() - - n_dims = len(data.shape) - data_shape = data.shape - if old_dtype == torch.bfloat16: - data_qtype = gguf.GGMLQuantizationType.BF16 - # elif old_dtype == torch.float32: - # data_qtype = gguf.GGMLQuantizationType.F32 - else: - data_qtype = gguf.GGMLQuantizationType.F16 - - # The max no. of dimensions that can be handled by the quantization code is 4 - if len(data.shape) > MAX_TENSOR_DIMS: - model_arch.handle_nd_tensor(key, data) - continue # needs to be added back later - - # get number of parameters (AKA elements) in this tensor - n_params = 1 - for dim_size in data_shape: - n_params *= dim_size - - if old_dtype in (torch.float32, torch.bfloat16): - if n_dims == 1: - # one-dimensional tensors should be kept in F32 - # also speeds up inference due to not dequantizing - data_qtype = gguf.GGMLQuantizationType.F32 - - elif n_params <= QUANTIZATION_THRESHOLD: - # very small tensors - data_qtype = gguf.GGMLQuantizationType.F32 - - elif any(x in key for x in model_arch.keys_hiprec): - # tensors that require max precision - data_qtype = gguf.GGMLQuantizationType.F32 - - if (model_arch.shape_fix # NEVER reshape for models such as flux - and n_dims > 1 # Skip one-dimensional tensors - and n_params >= REARRANGE_THRESHOLD # Only rearrange tensors meeting the size requirement - and (n_params / 256).is_integer() # Rearranging only makes sense if total elements is divisible by 256 - and not (data.shape[-1] / 256).is_integer() # Only need to rearrange if the last dimension is not divisible by 256 - ): - orig_shape = data.shape - data = data.reshape(n_params // 256, 256) - writer.add_array(f"comfy.gguf.orig_shape.{key}", tuple(int(dim) for dim in orig_shape)) - - try: - data = gguf.quants.quantize(data, data_qtype) - except (AttributeError, gguf.QuantError) as e: - tqdm.write(f"falling back to F16: {e}") - data_qtype = gguf.GGMLQuantizationType.F16 - data = gguf.quants.quantize(data, data_qtype) - - new_name = key # do we need to rename? - - shape_str = f"{{{', '.join(str(n) for n in reversed(data.shape))}}}" - tqdm.write(f"{f'%-{max_name_len + 4}s' % f'{new_name}'} {old_dtype} --> {data_qtype.name}, shape = {shape_str}") - - writer.add_tensor(new_name, data, raw_dtype=data_qtype) - -def convert_file(path, dst_path=None, interact=True, overwrite=False): - # load & run model detection logic - state_dict = load_state_dict(path) - model_arch = detect_arch(state_dict) - logging.info(f"* Architecture detected from input: {model_arch.arch}") - - # detect & set dtype for output file - dtypes = [x.dtype for x in state_dict.values()] - dtypes = {x:dtypes.count(x) for x in set(dtypes)} - main_dtype = max(dtypes, key=dtypes.get) - - if main_dtype == torch.bfloat16: - ftype_name = "BF16" - ftype_gguf = gguf.LlamaFileType.MOSTLY_BF16 - # elif main_dtype == torch.float32: - # ftype_name = "F32" - # ftype_gguf = None - else: - ftype_name = "F16" - ftype_gguf = gguf.LlamaFileType.MOSTLY_F16 - - if dst_path is None: - dst_path = f"{os.path.splitext(path)[0]}-{ftype_name}.gguf" - elif "{ftype}" in dst_path: # lcpp logic - dst_path = dst_path.replace("{ftype}", ftype_name) - - if os.path.isfile(dst_path) and not overwrite: - if interact: - input("Output exists enter to continue or ctrl+c to abort!") - else: - raise OSError("Output exists and overwriting is disabled!") - - # handle actual file - writer = gguf.GGUFWriter(path=None, arch=model_arch.arch) - writer.add_quantization_version(gguf.GGML_QUANT_VERSION) - if ftype_gguf is not None: - writer.add_file_type(ftype_gguf) - - handle_tensors(writer, state_dict, model_arch) - writer.write_header_to_file(path=dst_path) - writer.write_kv_data_to_file() - writer.write_tensors_to_file(progress=True) - writer.close() - - fix = f"./fix_5d_tensors_{model_arch.arch}.safetensors" - if os.path.isfile(fix): - logging.warning(f"\n### Warning! Fix file found at '{fix}'") - logging.warning(" you most likely need to run 'fix_5d_tensors.py' after quantization.") - - return dst_path, model_arch - -if __name__ == "__main__": - args = parse_args() - convert_file(args.src, args.dst) -