commit 7aa734a482743c198d97b932b0c5918b14ad28de Author: Jaret Burkett Date: Fri Mar 14 12:14:21 2025 -0600 Initial commit diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..68bc17f --- /dev/null +++ b/.gitignore @@ -0,0 +1,160 @@ +# Byte-compiled / optimized / DLL files +__pycache__/ +*.py[cod] +*$py.class + +# C extensions +*.so + +# Distribution / packaging +.Python +build/ +develop-eggs/ +dist/ +downloads/ +eggs/ +.eggs/ +lib/ +lib64/ +parts/ +sdist/ +var/ +wheels/ +share/python-wheels/ +*.egg-info/ +.installed.cfg +*.egg +MANIFEST + +# PyInstaller +# Usually these files are written by a python script from a template +# before PyInstaller builds the exe, so as to inject date/other infos into it. +*.manifest +*.spec + +# Installer logs +pip-log.txt +pip-delete-this-directory.txt + +# Unit test / coverage reports +htmlcov/ +.tox/ +.nox/ +.coverage +.coverage.* +.cache +nosetests.xml +coverage.xml +*.cover +*.py,cover +.hypothesis/ +.pytest_cache/ +cover/ + +# Translations +*.mo +*.pot + +# Django stuff: +*.log +local_settings.py +db.sqlite3 +db.sqlite3-journal + +# Flask stuff: +instance/ +.webassets-cache + +# Scrapy stuff: +.scrapy + +# Sphinx documentation +docs/_build/ + +# PyBuilder +.pybuilder/ +target/ + +# Jupyter Notebook +.ipynb_checkpoints + +# IPython +profile_default/ +ipython_config.py + +# pyenv +# For a library or package, you might want to ignore these files since the code is +# intended to run in multiple environments; otherwise, check them in: +# .python-version + +# pipenv +# According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control. +# However, in case of collaboration, if having platform-specific dependencies or dependencies +# having no cross-platform support, pipenv may install dependencies that don't work, or not +# install all needed dependencies. +#Pipfile.lock + +# poetry +# Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control. +# This is especially recommended for binary packages to ensure reproducibility, and is more +# commonly ignored for libraries. +# https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control +#poetry.lock + +# pdm +# Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control. +#pdm.lock +# pdm stores project-wide configurations in .pdm.toml, but it is recommended to not include it +# in version control. +# https://pdm.fming.dev/#use-with-ide +.pdm.toml + +# PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm +__pypackages__/ + +# Celery stuff +celerybeat-schedule +celerybeat.pid + +# SageMath parsed files +*.sage.py + +# Environments +.env +.venv +env/ +venv/ +ENV/ +env.bak/ +venv.bak/ + +# Spyder project settings +.spyderproject +.spyproject + +# Rope project settings +.ropeproject + +# mkdocs documentation +/site + +# mypy +.mypy_cache/ +.dmypy.json +dmypy.json + +# Pyre type checker +.pyre/ + +# pytype static type analyzer +.pytype/ + +# Cython debug symbols +cython_debug/ + +# PyCharm +# JetBrains specific template is maintained in a separate JetBrains.gitignore that can +# be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore +# and can be added to the global gitignore or merged into this file. For a more nuclear +# option (not recommended) you can uncomment the following to ignore the entire idea folder. +#.idea/ diff --git a/LICENSE b/LICENSE new file mode 100644 index 0000000..305cdac --- /dev/null +++ b/LICENSE @@ -0,0 +1,21 @@ +MIT License + +Copyright (c) 2023 Ostris, LLC + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +SOFTWARE. diff --git a/README.md b/README.md new file mode 100644 index 0000000..8863591 --- /dev/null +++ b/README.md @@ -0,0 +1,13 @@ +# Flex.1 tools + +Some tools to help with [Flex.1-alpha](https://huggingface.co/ostris/Flex.1-alpha) inference on Comfy UI. + +## Installation + +Clone this repo into your `custom_nodes` directory. + +## Nodes + +- **Flex Guidance**: Allows you to set the guidance for the Flex.1 guidance embedder, or bypass it completly to use true CFG. +- **Flex LoRA Loader**: Loads LoRAs and automatically prunes them to Flex.1 layers. It will not be perfect as Flex is heavily diverged from Flux dev and is not a direct ancenstor of it, but it should be good enough for most purposes. +- **Flex LoRA Loader (Model Only)**: Same as Flex LoRA Loader, but only loads the model and not the text encoder. Most Flux LoRAs do not train the text encoder. \ No newline at end of file diff --git a/__init__.py b/__init__.py new file mode 100644 index 0000000..47ebf54 --- /dev/null +++ b/__init__.py @@ -0,0 +1,163 @@ +import folder_paths +import node_helpers +import comfy.sd +import comfy.utils + + +class FlexGuidance: + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "conditioning": ("CONDITIONING", ), + "guidance": ("FLOAT", {"default": 3.5, "min": 0.0, "max": 100.0, "step": 0.1}), + "bypass_guidance_embedder": (["yes", "no"], {"default": "no"}), + } + } + + RETURN_TYPES = ("CONDITIONING",) + FUNCTION = "do_it" + + CATEGORY = "advanced/conditioning/flux" + + def do_it(self, conditioning, guidance, bypass_guidance_embedder): + bypass_guidance_embedder = bypass_guidance_embedder == "yes" + guidance_value = guidance + if bypass_guidance_embedder: + guidance_value = None + cond = node_helpers.conditioning_set_values( + conditioning, {"guidance": guidance_value} + ) + return (cond, ) + + +class FlexLoraLoader: + def __init__(self): + self.loaded_lora = None + + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "model": ("MODEL", {"tooltip": "The diffusion model the LoRA will be applied to."}), + "clip": ("CLIP", {"tooltip": "The CLIP model the LoRA will be applied to."}), + "lora_name": (folder_paths.get_filename_list("loras"), {"tooltip": "The name of the LoRA."}), + "strength_model": ("FLOAT", {"default": 1.0, "min": -100.0, "max": 100.0, "step": 0.01, "tooltip": "How strongly to modify the diffusion model. This value can be negative."}), + "strength_clip": ("FLOAT", {"default": 1.0, "min": -100.0, "max": 100.0, "step": 0.01, "tooltip": "How strongly to modify the CLIP model. This value can be negative."}), + } + } + + RETURN_TYPES = ("MODEL", "CLIP") + OUTPUT_TOOLTIPS = ("The modified diffusion model.", + "The modified CLIP model.") + FUNCTION = "load_lora" + + CATEGORY = "loaders" + DESCRIPTION = "Loads Loras and automatically converts Flux loras to Flex loras." + + def load_lora(self, model, clip, lora_name, strength_model, strength_clip): + if strength_model == 0 and strength_clip == 0: + return (model, clip) + + lora_path = folder_paths.get_full_path_or_raise("loras", lora_name) + lora = None + if self.loaded_lora is not None: + if self.loaded_lora[0] == lora_path: + lora = self.loaded_lora[1] + else: + self.loaded_lora = None + + if lora is None: + lora = comfy.utils.load_torch_file(lora_path, safe_load=True) + # convert it to Flex LoRA + # the pruning squashed double idx 5-15 into idx 4 + # making idx 16, 17, 18 become 5, 6, 7 + # we will drop double blocks with idx 5-15 + # and move idx 16, 17, 18 to 5, 6, 7 + # it is best to drop idx 4 as well since it is so divergent due to pruning + + # loras have different naming patterns, the ones I know about are below + block_test_targets = [ + "double_blocks.{idx}.", + "transformer.transformer_blocks.{idx}.", + "lora_unet_double_blocks_{idx}_", + "lycoris_unet_double_blocks_{idx}_", + "lycoris_transformer_blocks_{idx}_", + "lora_transformer_blocks_{idx}_", + ] + + # we trained the guidance embedder from scratch, the weights will not match at all + # loras will destroy it, so we will ignore it + ignore_if_contains = [ + "guidance_in", + "guidance_embedder" + ] + + # check if any of the keys start with the block_test_targets with idx 8-18, + # if they do, then this it is a Flux lora + + is_flux_lora = False + for idx in range(8, 19): + for target in block_test_targets: + if any(k.startswith(target.format(idx=idx)) for k in lora.keys()): + is_flux_lora = True + break + if is_flux_lora: + break + + if is_flux_lora: + flex_lora = {} + drop_idxs = list(range(4, 16)) + move_idxs = {16: 5, 17: 6, 18: 7} + for k, v in lora.items(): + if any(k.startswith(target.format(idx=idx)) for target in block_test_targets for idx in drop_idxs): + # drop it + continue + if any(target in k for target in ignore_if_contains): + continue + for old_idx, new_idx in move_idxs.items(): + replaced = False + for target in block_test_targets: + formatted_target = target.format(idx=old_idx) + if k.startswith(formatted_target): + k = k.replace(formatted_target, + target.format(idx=new_idx)) + replaced = True + break + if replaced: + break + flex_lora[k] = v + lora = flex_lora + + self.loaded_lora = (lora_path, lora) + + model_lora, clip_lora = comfy.sd.load_lora_for_models( + model, clip, lora, strength_model, strength_clip) + return (model_lora, clip_lora) + + +class FlexLoraLoaderModelOnly(FlexLoraLoader): + @classmethod + def INPUT_TYPES(s): + return {"required": {"model": ("MODEL",), + "lora_name": (folder_paths.get_filename_list("loras"), ), + "strength_model": ("FLOAT", {"default": 1.0, "min": -100.0, "max": 100.0, "step": 0.01}), + }} + RETURN_TYPES = ("MODEL",) + FUNCTION = "load_lora_model_only" + + def load_lora_model_only(self, model, lora_name, strength_model): + return (self.load_lora(model, None, lora_name, strength_model, 0)[0],) + + +NODE_CLASS_MAPPINGS = { + "FlexGuidance": FlexGuidance, + "FlexLoraLoader": FlexLoraLoader, + "FlexLoraLoaderModelOnly": FlexLoraLoaderModelOnly, +} + +NODE_DISPLAY_NAME_MAPPINGS = { + "FlexGuidance": "Flex Guidance", + "FlexLoraLoader": "Flex LoRA Loader", + "FlexLoraLoaderModelOnly": "Flex LoRA Loader (Model Only)", +}