diff --git a/example_workflows/nunchaku-flux.1-dev.jpg b/example_workflows/nunchaku-flux.1-dev.jpg new file mode 100644 index 00000000..0e8dfaad Binary files /dev/null and b/example_workflows/nunchaku-flux.1-dev.jpg differ diff --git a/example_workflows/nunchaku-flux.1-dev.json b/example_workflows/nunchaku-flux.1-dev.json new file mode 100644 index 00000000..5e72d2f2 --- /dev/null +++ b/example_workflows/nunchaku-flux.1-dev.json @@ -0,0 +1 @@ +{"id":"f253212e-0ec7-40c5-9671-bafc52d66023","revision":0,"last_node_id":56,"last_link_id":146,"nodes":[{"id":26,"type":"FluxGuidance","pos":[533.9339599609375,118.7322998046875],"size":[317.4000244140625,58],"flags":{},"order":16,"mode":0,"inputs":[{"localized_name":"conditioning","name":"conditioning","type":"CONDITIONING","link":41},{"localized_name":"guidance","name":"guidance","type":"FLOAT","widget":{"name":"guidance"},"link":null}],"outputs":[{"localized_name":"CONDITIONING","name":"CONDITIONING","type":"CONDITIONING","slot_index":0,"links":[42]}],"properties":{"cnr_id":"comfy-core","ver":"0.3.24","Node name for S&R":"FluxGuidance"},"widgets_values":[3.5],"color":"#233","bgcolor":"#355"},{"id":16,"type":"KSamplerSelect","pos":[813.7789306640625,847.5765991210938],"size":[315,58],"flags":{},"order":0,"mode":0,"inputs":[{"localized_name":"sampler_name","name":"sampler_name","type":"COMBO","widget":{"name":"sampler_name"},"link":null}],"outputs":[{"localized_name":"SAMPLER","name":"SAMPLER","type":"SAMPLER","links":[19]}],"properties":{"cnr_id":"comfy-core","ver":"0.3.24","Node name for S&R":"KSamplerSelect"},"widgets_values":["euler"]},{"id":22,"type":"BasicGuider","pos":[691.55615234375,303.7747497558594],"size":[222.3482666015625,46],"flags":{},"order":17,"mode":0,"inputs":[{"localized_name":"model","name":"model","type":"MODEL","link":54},{"localized_name":"conditioning","name":"conditioning","type":"CONDITIONING","link":42}],"outputs":[{"localized_name":"GUIDER","name":"GUIDER","type":"GUIDER","slot_index":0,"links":[30]}],"properties":{"cnr_id":"comfy-core","ver":"0.3.24","Node name for S&R":"BasicGuider"},"widgets_values":[]},{"id":27,"type":"EmptySD3LatentImage","pos":[825.0562133789062,526.280029296875],"size":[315,126],"flags":{},"order":9,"mode":0,"inputs":[{"localized_name":"width","name":"width","type":"INT","widget":{"name":"width"},"link":112},{"localized_name":"height","name":"height","type":"INT","widget":{"name":"height"},"link":113},{"localized_name":"batch_size","name":"batch_size","type":"INT","widget":{"name":"batch_size"},"link":null}],"outputs":[{"localized_name":"LATENT","name":"LATENT","type":"LATENT","slot_index":0,"links":[116]}],"properties":{"cnr_id":"comfy-core","ver":"0.3.24","Node name for S&R":"EmptySD3LatentImage"},"widgets_values":[832,1216,1]},{"id":13,"type":"SamplerCustomAdvanced","pos":[867.9028930664062,126.43718719482422],"size":[272.3617858886719,124.53733825683594],"flags":{},"order":18,"mode":0,"inputs":[{"localized_name":"noise","name":"noise","type":"NOISE","link":37},{"localized_name":"guider","name":"guider","type":"GUIDER","link":30},{"localized_name":"sampler","name":"sampler","type":"SAMPLER","link":19},{"localized_name":"sigmas","name":"sigmas","type":"SIGMAS","link":20},{"localized_name":"latent_image","name":"latent_image","type":"LATENT","link":116}],"outputs":[{"localized_name":"output","name":"output","type":"LATENT","slot_index":0,"links":[24]},{"localized_name":"denoised_output","name":"denoised_output","type":"LATENT","links":null}],"properties":{"cnr_id":"comfy-core","ver":"0.3.24","Node name for S&R":"SamplerCustomAdvanced"},"widgets_values":[]},{"id":30,"type":"ModelSamplingFlux","pos":[879.5872802734375,1107.096923828125],"size":[210,130],"flags":{},"order":11,"mode":0,"inputs":[{"localized_name":"model","name":"model","type":"MODEL","link":143},{"localized_name":"max_shift","name":"max_shift","type":"FLOAT","widget":{"name":"max_shift"},"link":null},{"localized_name":"base_shift","name":"base_shift","type":"FLOAT","widget":{"name":"base_shift"},"link":null},{"localized_name":"width","name":"width","type":"INT","widget":{"name":"width"},"link":115},{"localized_name":"height","name":"height","type":"INT","widget":{"name":"height"},"link":114}],"outputs":[{"localized_name":"MODEL","name":"MODEL","type":"MODEL","slot_index":0,"links":[54,55]}],"properties":{"cnr_id":"comfy-core","ver":"0.3.24","Node name for S&R":"ModelSamplingFlux"},"widgets_values":[1.15,0.5,832,1216]},{"id":17,"type":"BasicScheduler","pos":[818.5535888671875,948.8093872070312],"size":[315,106],"flags":{},"order":13,"mode":0,"inputs":[{"localized_name":"model","name":"model","type":"MODEL","link":55},{"localized_name":"scheduler","name":"scheduler","type":"COMBO","widget":{"name":"scheduler"},"link":null},{"localized_name":"steps","name":"steps","type":"INT","widget":{"name":"steps"},"link":null},{"localized_name":"denoise","name":"denoise","type":"FLOAT","widget":{"name":"denoise"},"link":null}],"outputs":[{"localized_name":"SIGMAS","name":"SIGMAS","type":"SIGMAS","links":[20]}],"properties":{"cnr_id":"comfy-core","ver":"0.3.24","Node name for S&R":"BasicScheduler"},"widgets_values":["simple",8,1]},{"id":10,"type":"VAELoader","pos":[421.0907897949219,767.467041015625],"size":[311.81634521484375,60.429901123046875],"flags":{},"order":1,"mode":0,"inputs":[{"localized_name":"vae_name","name":"vae_name","type":"COMBO","widget":{"name":"vae_name"},"link":null}],"outputs":[{"localized_name":"VAE","name":"VAE","type":"VAE","slot_index":0,"links":[12]}],"properties":{"cnr_id":"comfy-core","ver":"0.3.24","Node name for S&R":"VAELoader"},"widgets_values":["flux1\\ae.safetensors"]},{"id":8,"type":"VAEDecode","pos":[938.151611328125,299.3258056640625],"size":[210,46],"flags":{},"order":19,"mode":0,"inputs":[{"localized_name":"samples","name":"samples","type":"LATENT","link":24},{"localized_name":"vae","name":"vae","type":"VAE","link":12}],"outputs":[{"localized_name":"IMAGE","name":"IMAGE","type":"IMAGE","slot_index":0,"links":[135]}],"properties":{"cnr_id":"comfy-core","ver":"0.3.24","Node name for S&R":"VAEDecode"},"widgets_values":[]},{"id":44,"type":"NunchakuTextEncoderLoader","pos":[-291.055908203125,303.4891662597656],"size":[352.79998779296875,178],"flags":{},"order":2,"mode":0,"inputs":[{"localized_name":"model_type","name":"model_type","type":"COMBO","widget":{"name":"model_type"},"link":null},{"localized_name":"text_encoder1","name":"text_encoder1","type":"COMBO","widget":{"name":"text_encoder1"},"link":null},{"localized_name":"text_encoder2","name":"text_encoder2","type":"COMBO","widget":{"name":"text_encoder2"},"link":null},{"localized_name":"t5_min_length","name":"t5_min_length","type":"INT","widget":{"name":"t5_min_length"},"link":null},{"localized_name":"use_4bit_t5","name":"use_4bit_t5","type":"COMBO","widget":{"name":"use_4bit_t5"},"link":null},{"localized_name":"int4_model","name":"int4_model","type":"COMBO","widget":{"name":"int4_model"},"link":null}],"outputs":[{"localized_name":"CLIP","name":"CLIP","type":"CLIP","links":[127]}],"properties":{"Node name for S&R":"NunchakuTextEncoderLoader"},"widgets_values":["flux","t5xxl_fp8_e4m3fn.safetensors","clip_l.safetensors",512,"disable","none"]},{"id":34,"type":"PrimitiveNode","pos":[696.5196533203125,397.441650390625],"size":[210,82],"flags":{},"order":3,"mode":0,"inputs":[],"outputs":[{"name":"INT","type":"INT","widget":{"name":"width"},"slot_index":0,"links":[112,115]}],"title":"width","properties":{"Run widget replace on values":false},"widgets_values":[832,"fixed"],"color":"#323","bgcolor":"#535"},{"id":35,"type":"PrimitiveNode","pos":[937.396484375,392.3027038574219],"size":[210,86.4900131225586],"flags":{},"order":4,"mode":0,"inputs":[],"outputs":[{"name":"INT","type":"INT","widget":{"name":"height"},"slot_index":0,"links":[113,114]}],"title":"height","properties":{"Run widget replace on values":false},"widgets_values":[1216,"fixed"],"color":"#323","bgcolor":"#535"},{"id":25,"type":"RandomNoise","pos":[819.1885986328125,709.9674072265625],"size":[315,82],"flags":{},"order":5,"mode":0,"inputs":[{"localized_name":"noise_seed","name":"noise_seed","type":"INT","widget":{"name":"noise_seed"},"link":null}],"outputs":[{"localized_name":"NOISE","name":"NOISE","type":"NOISE","links":[37]}],"properties":{"cnr_id":"comfy-core","ver":"0.3.24","Node name for S&R":"RandomNoise"},"widgets_values":[924280468,"fixed"],"color":"#2a363b","bgcolor":"#3f5159"},{"id":49,"type":"PreviewImage","pos":[1170.8551025390625,128.45352172851562],"size":[897.65087890625,1284.666748046875],"flags":{},"order":20,"mode":0,"inputs":[{"localized_name":"images","name":"images","type":"IMAGE","link":135}],"outputs":[],"properties":{"cnr_id":"comfy-core","ver":"0.3.43","Node name for S&R":"PreviewImage"},"widgets_values":[]},{"id":45,"type":"NunchakuFluxDiTLoader","pos":[-548.04833984375,990.260498046875],"size":[315,202],"flags":{},"order":6,"mode":0,"inputs":[{"localized_name":"model_path","name":"model_path","type":"COMBO","widget":{"name":"model_path"},"link":null},{"localized_name":"cache_threshold","name":"cache_threshold","type":"FLOAT","widget":{"name":"cache_threshold"},"link":null},{"localized_name":"attention","name":"attention","type":"COMBO","widget":{"name":"attention"},"link":null},{"localized_name":"cpu_offload","name":"cpu_offload","type":"COMBO","widget":{"name":"cpu_offload"},"link":null},{"localized_name":"device_id","name":"device_id","type":"INT","widget":{"name":"device_id"},"link":null},{"localized_name":"data_type","name":"data_type","type":"COMBO","widget":{"name":"data_type"},"link":null},{"localized_name":"i2f_mode","name":"i2f_mode","shape":7,"type":"COMBO","widget":{"name":"i2f_mode"},"link":null}],"outputs":[{"localized_name":"MODEL","name":"MODEL","type":"MODEL","links":[133]}],"properties":{"Node name for S&R":"NunchakuFluxDiTLoader"},"widgets_values":["svdq-int4-flux.1-dev",0,"nunchaku-fp16","auto",0,"bfloat16","enabled"]},{"id":48,"type":"Lora Loader (LoraManager)","pos":[-52.50986099243164,909.7772827148438],"size":[403.350341796875,481.60552978515625],"flags":{},"order":10,"mode":0,"inputs":[{"localized_name":"model","name":"model","type":"MODEL","link":133},{"localized_name":"text","name":"text","type":"STRING","widget":{"name":"text"},"link":null},{"name":"clip","shape":7,"type":"CLIP","link":null},{"name":"lora_stack","shape":7,"type":"LORA_STACK","link":144}],"outputs":[{"localized_name":"MODEL","name":"MODEL","type":"MODEL","links":[143]},{"localized_name":"CLIP","name":"CLIP","type":"CLIP","links":null},{"localized_name":"trigger_words","name":"trigger_words","type":"STRING","links":[136]},{"localized_name":"loaded_loras","name":"loaded_loras","type":"STRING","links":null}],"properties":{"cnr_id":"comfyui-lora-manager","ver":"4043846767d4339bcf6bbb2eb3407c7fac94615d","Node name for S&R":"Lora Loader (LoraManager)"},"widgets_values":[" ",[{"name":"Mezzotint_Artstyle_for_Flux_-_by_Ethanar","strength":0.2,"active":true,"clipStrength":0.2,"expanded":false},{"name":"RetroAnimeFluxV1","strength":0.8,"active":true,"clipStrength":0.8,"expanded":false},{"name":"flux.1_lora_flyway_ink-dynamic","strength":0.8,"active":true,"clipStrength":0.8,"expanded":false},{"name":"ck-painterly-fantasy-000017","strength":0.2,"active":true,"clipStrength":0.2,"expanded":false},{"name":"FluxMythG0thicL1nes","strength":0.4,"active":true,"clipStrength":0.4,"expanded":false},{"name":"1","strength":0.75,"active":true,"clipStrength":0.75,"expanded":false}]],"color":"#323","bgcolor":"#535"},{"id":50,"type":"TriggerWord Toggle (LoraManager)","pos":[390.2729797363281,926.0449829101562],"size":[318.5542907714844,260],"flags":{},"order":12,"mode":0,"inputs":[{"localized_name":"group_mode","name":"group_mode","type":"BOOLEAN","widget":{"name":"group_mode"},"link":null},{"localized_name":"default_active","name":"default_active","type":"BOOLEAN","widget":{"name":"default_active"},"link":null},{"name":"trigger_words","shape":7,"type":"string","link":136}],"outputs":[{"localized_name":"filtered_trigger_words","name":"filtered_trigger_words","type":"STRING","links":[137]}],"properties":{"cnr_id":"comfyui-lora-manager","ver":"4043846767d4339bcf6bbb2eb3407c7fac94615d","Node name for S&R":"TriggerWord Toggle (LoraManager)"},"widgets_values":[true,true,[{"text":"G0thicL1nes","active":true},{"text":"amanoer","active":true}],"G0thicL1nes,, amanoer"],"color":"#323","bgcolor":"#535"},{"id":55,"type":"Lora Stacker (LoraManager)","pos":[-573.553955078125,586.1253662109375],"size":[414.74151611328125,332],"flags":{},"order":7,"mode":0,"inputs":[{"localized_name":"text","name":"text","type":"STRING","widget":{"name":"text"},"link":null},{"name":"lora_stack","shape":7,"type":"LORA_STACK","link":null}],"outputs":[{"localized_name":"LORA_STACK","name":"LORA_STACK","type":"LORA_STACK","links":[144]},{"localized_name":"trigger_words","name":"trigger_words","type":"STRING","links":null},{"localized_name":"active_loras","name":"active_loras","type":"STRING","links":null}],"properties":{"cnr_id":"comfyui-lora-manager","ver":"4043846767d4339bcf6bbb2eb3407c7fac94615d","Node name for S&R":"Lora Stacker (LoraManager)"},"widgets_values":["",[{"name":"FLUX.1-Turbo-Alpha","strength":1,"active":true,"clipStrength":1,"expanded":false}]],"color":"#233","bgcolor":"#355"},{"id":56,"type":"StringConstantMultiline","pos":[-60.967350006103516,584.9976196289062],"size":[400,200],"flags":{},"order":8,"mode":0,"inputs":[{"localized_name":"string","name":"string","type":"STRING","widget":{"name":"string"},"link":null},{"localized_name":"strip_newlines","name":"strip_newlines","type":"BOOLEAN","widget":{"name":"strip_newlines"},"link":null}],"outputs":[{"localized_name":"STRING","name":"STRING","type":"STRING","links":[146]}],"properties":{"cnr_id":"comfyui-kjnodes","ver":"1.1.2","Node name for S&R":"StringConstantMultiline"},"widgets_values":["a futuristic, cyberpunk-styled female warrior. She stands in the center of the image, exuding a commanding presence. Her skin is pale and her long, flowing white hair contrasts starkly with the dark, metallic armor that covers her body. The armor is intricately detailed with sharp, angular lines and glowing blue accents, particularly around her right arm, which emits a vibrant, ethereal blue light.\nShe wields a large, double-headed scythe that is almost as tall as she is. The scythe has a metallic sheen, with one blade being a sharp, white, almost crystalline material, and the other blade emitting a similar blue glow. Her expression is intense and focused, with piercing blue eyes that match the glow from her arm.\nBehind her, a massive, serpentine dragon with a similarly pale, almost white, scales is coiled protectively. The dragon's eyes are glowing blue, matching her arm and adding to the overall mystical and powerful aura of the scene. The background is a dark, industrial setting with blurred, metallic structures and neon lights, enhancing the cyberpunk feel. The artwork is highly detailed and rendered in a realistic yet fantastical style, typical of high-quality digital art.",false],"color":"#232","bgcolor":"#353"},{"id":51,"type":"JoinStrings","pos":[416.567138671875,598.2608032226562],"size":[270,78],"flags":{},"order":14,"mode":0,"inputs":[{"localized_name":"string1","name":"string1","type":"STRING","link":137},{"localized_name":"string2","name":"string2","type":"STRING","link":146},{"localized_name":"delimiter","name":"delimiter","type":"STRING","widget":{"name":"delimiter"},"link":null}],"outputs":[{"localized_name":"STRING","name":"STRING","type":"STRING","links":[139]}],"properties":{"cnr_id":"comfyui-kjnodes","ver":"1.1.2","Node name for S&R":"JoinStrings"},"widgets_values":[", "],"color":"#232","bgcolor":"#353"},{"id":6,"type":"CLIPTextEncode","pos":[240.22817993164062,363.7594909667969],"size":[263.2095642089844,88],"flags":{"collapsed":false},"order":15,"mode":0,"inputs":[{"localized_name":"clip","name":"clip","type":"CLIP","link":127},{"localized_name":"text","name":"text","type":"STRING","widget":{"name":"text"},"link":139}],"outputs":[{"localized_name":"CONDITIONING","name":"CONDITIONING","type":"CONDITIONING","slot_index":0,"links":[41]}],"title":"CLIP Text Encode (Positive Prompt)","properties":{"cnr_id":"comfy-core","ver":"0.3.24","Node name for S&R":"CLIPTextEncode"},"widgets_values":["a dynamic and dramatic digital artwork featuring a stylized anthropomorphic white tiger with striking yellow eyes. The tiger is depicted in a powerful stance, wielding a katana with one hand raised above its head. Its fur is detailed with black stripes, and its mane flows wildly, blending with the stormy background. The scene is set amidst swirling dark clouds and flashes of lightning, enhancing the sense of movement and energy. The composition is vertical, with the tiger positioned centrally, creating a sense of depth and intensity. The color palette is dominated by shades of blue, gray, and white, with bright highlights from the lightning. The overall style is reminiscent of fantasy or manga art, with a focus on dynamic action and dramatic lighting."],"color":"#232","bgcolor":"#353"}],"links":[[12,10,0,8,1,"VAE"],[19,16,0,13,2,"SAMPLER"],[20,17,0,13,3,"SIGMAS"],[24,13,0,8,0,"LATENT"],[30,22,0,13,1,"GUIDER"],[37,25,0,13,0,"NOISE"],[41,6,0,26,0,"CONDITIONING"],[42,26,0,22,1,"CONDITIONING"],[54,30,0,22,0,"MODEL"],[55,30,0,17,0,"MODEL"],[112,34,0,27,0,"INT"],[113,35,0,27,1,"INT"],[114,35,0,30,4,"INT"],[115,34,0,30,3,"INT"],[116,27,0,13,4,"LATENT"],[127,44,0,6,0,"CLIP"],[133,45,0,48,0,"MODEL"],[135,8,0,49,0,"IMAGE"],[136,48,2,50,2,"string"],[137,50,0,51,0,"STRING"],[139,51,0,6,1,"STRING"],[143,48,0,30,0,"MODEL"],[144,55,0,48,3,"LORA_STACK"],[146,56,0,51,1,"STRING"]],"groups":[],"config":{},"extra":{"ds":{"scale":0.9849732675807675,"offset":[486.518281834556,-142.01476750486412]},"groupNodes":{"EmptyLatentImage":{"nodes":[{"type":"PrimitiveNode","pos":[432,480],"size":{"0":210,"1":82},"flags":{},"order":6,"mode":0,"outputs":[{"name":"INT","type":"INT","links":[],"widget":{"name":"height"},"slot_index":0}],"title":"height","properties":{"Run widget replace on values":false},"color":"#323","bgcolor":"#535","index":0},{"type":"PrimitiveNode","pos":[672,480],"size":{"0":210,"1":82},"flags":{},"order":7,"mode":0,"outputs":[{"name":"INT","type":"INT","links":[],"slot_index":0,"widget":{"name":"width"}}],"title":"width","properties":{"Run widget replace on values":false},"color":"#323","bgcolor":"#535","index":1},{"type":"EmptySD3LatentImage","pos":[480,624],"size":{"0":315,"1":106},"flags":{},"order":10,"mode":0,"inputs":[{"name":"width","type":"INT","link":null,"widget":{"name":"width"}},{"name":"height","type":"INT","link":null,"widget":{"name":"height"}}],"outputs":[{"name":"LATENT","type":"LATENT","links":[],"shape":3,"slot_index":0}],"properties":{"Node name for S&R":"EmptySD3LatentImage"},"widgets_values":[1024,1024,1],"index":2}],"links":[[1,0,2,0,34,"INT"],[0,0,2,1,35,"INT"]],"external":[[0,0,"INT"],[1,0,"INT"],[2,0,"LATENT"]],"config":{"0":{"output":{"0":{"name":"height"}},"input":{"value":{"visible":true}}},"1":{"output":{"0":{"name":"width"}},"input":{"value":{"visible":true}}},"2":{"input":{"width":{"visible":false},"height":{"visible":false}}}}}},"node_versions":{"comfy-core":"0.3.24"}},"version":0.4} \ No newline at end of file diff --git a/py/nodes/lora_loader.py b/py/nodes/lora_loader.py index 58a9aa89..8d45e19c 100644 --- a/py/nodes/lora_loader.py +++ b/py/nodes/lora_loader.py @@ -2,14 +2,14 @@ import logging from nodes import LoraLoader from comfy.comfy_types import IO # type: ignore import asyncio -from .utils import FlexibleOptionalInputType, any_type, get_lora_info, extract_lora_name, get_loras_list +from .utils import FlexibleOptionalInputType, any_type, get_lora_info, extract_lora_name, get_loras_list, nunchaku_load_lora logger = logging.getLogger(__name__) class LoraManagerLoader: NAME = "Lora Loader (LoraManager)" CATEGORY = "Lora Manager/loaders" - + @classmethod def INPUT_TYPES(cls): return { @@ -37,19 +37,39 @@ class LoraManagerLoader: clip = kwargs.get('clip', None) lora_stack = kwargs.get('lora_stack', None) + + # Check if model is a Nunchaku Flux model - simplified approach + is_nunchaku_model = False + + try: + model_wrapper = model.model.diffusion_model + # Check if model is a Nunchaku Flux model using only class name + if model_wrapper.__class__.__name__ == "ComfyFluxWrapper": + is_nunchaku_model = True + logger.info("Detected Nunchaku Flux model") + except (AttributeError, TypeError): + # Not a model with the expected structure + pass + # First process lora_stack if available if lora_stack: for lora_path, model_strength, clip_strength in lora_stack: - # Apply the LoRA using the provided path and strengths - model, clip = LoraLoader().load_lora(model, clip, lora_path, model_strength, clip_strength) + # Apply the LoRA using the appropriate loader + if is_nunchaku_model: + # Use our custom function for Flux models + model = nunchaku_load_lora(model, lora_path, model_strength) + # clip remains unchanged for Nunchaku models + else: + # Use default loader for standard models + model, clip = LoraLoader().load_lora(model, clip, lora_path, model_strength, clip_strength) # Extract lora name for trigger words lookup lora_name = extract_lora_name(lora_path) _, trigger_words = asyncio.run(get_lora_info(lora_name)) all_trigger_words.extend(trigger_words) - # Add clip strength to output if different from model strength - if abs(model_strength - clip_strength) > 0.001: + # Add clip strength to output if different from model strength (except for Nunchaku models) + if not is_nunchaku_model and abs(model_strength - clip_strength) > 0.001: loaded_loras.append(f"{lora_name}: {model_strength},{clip_strength}") else: loaded_loras.append(f"{lora_name}: {model_strength}") @@ -68,11 +88,17 @@ class LoraManagerLoader: # Get lora path and trigger words lora_path, trigger_words = asyncio.run(get_lora_info(lora_name)) - # Apply the LoRA using the resolved path with separate strengths - model, clip = LoraLoader().load_lora(model, clip, lora_path, model_strength, clip_strength) + # Apply the LoRA using the appropriate loader + if is_nunchaku_model: + # For Nunchaku models, use our custom function + model = nunchaku_load_lora(model, lora_path, model_strength) + # clip remains unchanged + else: + # Use default loader for standard models + model, clip = LoraLoader().load_lora(model, clip, lora_path, model_strength, clip_strength) - # Include clip strength in output if different from model strength - if abs(model_strength - clip_strength) > 0.001: + # Include clip strength in output if different from model strength and not a Nunchaku model + if not is_nunchaku_model and abs(model_strength - clip_strength) > 0.001: loaded_loras.append(f"{lora_name}: {model_strength},{clip_strength}") else: loaded_loras.append(f"{lora_name}: {model_strength}") diff --git a/py/nodes/utils.py b/py/nodes/utils.py index 1feb1a77..33a3c972 100644 --- a/py/nodes/utils.py +++ b/py/nodes/utils.py @@ -35,7 +35,12 @@ any_type = AnyType("*") # Common methods extracted from lora_loader.py and lora_stacker.py import os import logging -import asyncio +import copy +import folder_paths +import torch +import safetensors.torch +from diffusers.utils.state_dict_utils import convert_unet_state_dict_to_peft +from diffusers.loaders import FluxLoraLoaderMixin from ..services.lora_scanner import LoraScanner from ..config import config @@ -81,4 +86,64 @@ def get_loras_list(kwargs): # Unexpected format else: logger.warning(f"Unexpected loras format: {type(loras_data)}") - return [] \ No newline at end of file + return [] + +def load_state_dict_in_safetensors(path, device="cpu", filter_prefix=""): + """Simplified version of load_state_dict_in_safetensors that just loads from a local path""" + state_dict = {} + with safetensors.torch.safe_open(path, framework="pt", device=device) as f: + for k in f.keys(): + if filter_prefix and not k.startswith(filter_prefix): + continue + state_dict[k.removeprefix(filter_prefix)] = f.get_tensor(k) + return state_dict + +def to_diffusers(input_lora): + """Simplified version of to_diffusers for Flux LoRA conversion""" + if isinstance(input_lora, str): + tensors = load_state_dict_in_safetensors(input_lora, device="cpu") + else: + tensors = {k: v for k, v in input_lora.items()} + + # Convert FP8 tensors to BF16 + for k, v in tensors.items(): + if v.dtype not in [torch.float64, torch.float32, torch.bfloat16, torch.float16]: + tensors[k] = v.to(torch.bfloat16) + + new_tensors = FluxLoraLoaderMixin.lora_state_dict(tensors) + new_tensors = convert_unet_state_dict_to_peft(new_tensors) + + return new_tensors + +def nunchaku_load_lora(model, lora_name, lora_strength): + """Load a Flux LoRA for Nunchaku model""" + model_wrapper = model.model.diffusion_model + transformer = model_wrapper.model + + # Save the transformer temporarily + model_wrapper.model = None + ret_model = copy.deepcopy(model) # copy everything except the model + ret_model_wrapper = ret_model.model.diffusion_model + + # Restore the model and set it for the copy + model_wrapper.model = transformer + ret_model_wrapper.model = transformer + + # Get full path to the LoRA file + lora_path = folder_paths.get_full_path("loras", lora_name) + ret_model_wrapper.loras.append((lora_path, lora_strength)) + + # Convert the LoRA to diffusers format + sd = to_diffusers(lora_path) + + # Handle embedding adjustment if needed + if "transformer.x_embedder.lora_A.weight" in sd: + new_in_channels = sd["transformer.x_embedder.lora_A.weight"].shape[1] + assert new_in_channels % 4 == 0 + new_in_channels = new_in_channels // 4 + + old_in_channels = ret_model.model.model_config.unet_config["in_channels"] + if old_in_channels < new_in_channels: + ret_model.model.model_config.unet_config["in_channels"] = new_in_channels + + return ret_model \ No newline at end of file diff --git a/pyproject.toml b/pyproject.toml index a586d61c..746ec8b6 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -7,6 +7,7 @@ dependencies = [ "aiohttp", "jinja2", "safetensors", + "diffusers", "watchdog", "beautifulsoup4", "piexif",