Remove unused temp files
This commit is contained in:
ussoewwin
@@ -1,277 +0,0 @@
|
||||
"""
|
||||
This module provides the :class:`NunchakuFluxLoraLoader` node
|
||||
for applying LoRA weights to Nunchaku FLUX models within ComfyUI.
|
||||
"""
|
||||
|
||||
import logging
|
||||
import os
|
||||
|
||||
from nunchaku.lora.flux import to_diffusers
|
||||
|
||||
from ...wrappers.flux import ComfyFluxWrapper, copy_with_ctx
|
||||
from ..utils import get_filename_list, get_full_path_or_raise
|
||||
|
||||
# Get log level from environment variable (default to INFO)
|
||||
log_level = os.getenv("LOG_LEVEL", "INFO").upper()
|
||||
|
||||
# Configure logging
|
||||
logging.basicConfig(level=getattr(logging, log_level, logging.INFO), format="%(asctime)s - %(levelname)s - %(message)s")
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class NunchakuFluxLoraLoader:
|
||||
"""
|
||||
Node for loading and applying a LoRA to a Nunchaku FLUX model.
|
||||
|
||||
Attributes
|
||||
----------
|
||||
RETURN_TYPES : tuple
|
||||
The return type of the node ("MODEL",).
|
||||
OUTPUT_TOOLTIPS : tuple
|
||||
Tooltip for the output.
|
||||
FUNCTION : str
|
||||
The function to call ("load_lora").
|
||||
TITLE : str
|
||||
Node title.
|
||||
CATEGORY : str
|
||||
Node category.
|
||||
DESCRIPTION : str
|
||||
Node description.
|
||||
"""
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
"""
|
||||
Defines the input types and tooltips for the node.
|
||||
|
||||
Returns
|
||||
-------
|
||||
dict
|
||||
A dictionary specifying the required inputs and their descriptions for the node interface.
|
||||
"""
|
||||
return {
|
||||
"required": {
|
||||
"model": (
|
||||
"MODEL",
|
||||
{
|
||||
"tooltip": "The diffusion model the LoRA will be applied to. "
|
||||
"Make sure the model is loaded by `Nunchaku FLUX DiT Loader`."
|
||||
},
|
||||
),
|
||||
"lora_name": (
|
||||
get_filename_list("loras"),
|
||||
{"tooltip": "The file name of the LoRA."},
|
||||
),
|
||||
"lora_strength": (
|
||||
"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.",
|
||||
},
|
||||
),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("MODEL",)
|
||||
OUTPUT_TOOLTIPS = ("The modified diffusion model.",)
|
||||
FUNCTION = "load_lora"
|
||||
TITLE = "Nunchaku FLUX LoRA Loader"
|
||||
|
||||
CATEGORY = "Nunchaku"
|
||||
DESCRIPTION = (
|
||||
"LoRAs are used to modify the diffusion model, "
|
||||
"altering the way in which latents are denoised such as applying styles. "
|
||||
"You can link multiple LoRA nodes."
|
||||
)
|
||||
|
||||
def load_lora(self, model, lora_name: str, lora_strength: float):
|
||||
"""
|
||||
Apply a LoRA to a Nunchaku FLUX diffusion model.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
model : object
|
||||
The diffusion model to modify.
|
||||
lora_name : str
|
||||
The name of the LoRA to apply.
|
||||
lora_strength : float
|
||||
The strength with which to apply the LoRA.
|
||||
|
||||
Returns
|
||||
-------
|
||||
tuple
|
||||
A tuple containing the modified diffusion model.
|
||||
"""
|
||||
if abs(lora_strength) < 1e-5:
|
||||
return (model,) # If the strength is too small, return the original model
|
||||
|
||||
model_wrapper = model.model.diffusion_model
|
||||
assert isinstance(model_wrapper, ComfyFluxWrapper)
|
||||
|
||||
lora_path = get_full_path_or_raise("loras", lora_name)
|
||||
|
||||
ret_model_wrapper, ret_model = copy_with_ctx(model_wrapper)
|
||||
|
||||
ret_model_wrapper.loras = [*model_wrapper.loras, (lora_path, lora_strength)]
|
||||
sd = to_diffusers(lora_path)
|
||||
|
||||
# To handle FLUX.1 tools LoRAs, which change the number of input channels
|
||||
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,)
|
||||
|
||||
|
||||
class NunchakuFluxLoraStack:
|
||||
"""
|
||||
Node for loading and applying multiple LoRAs to a Nunchaku FLUX model with dynamic input.
|
||||
|
||||
This node allows you to configure multiple LoRAs with their respective strengths
|
||||
in a single node, providing the same effect as chaining multiple LoRA nodes.
|
||||
|
||||
Attributes
|
||||
----------
|
||||
RETURN_TYPES : tuple
|
||||
The return type of the node ("MODEL",).
|
||||
OUTPUT_TOOLTIPS : tuple
|
||||
Tooltip for the output.
|
||||
FUNCTION : str
|
||||
The function to call ("load_lora_stack").
|
||||
TITLE : str
|
||||
Node title.
|
||||
CATEGORY : str
|
||||
Node category.
|
||||
DESCRIPTION : str
|
||||
Node description.
|
||||
"""
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
"""
|
||||
Defines the input types for the LoRA stack node.
|
||||
|
||||
Returns
|
||||
-------
|
||||
dict
|
||||
A dictionary specifying the required inputs and optional LoRA inputs.
|
||||
"""
|
||||
# Base inputs
|
||||
inputs = {
|
||||
"required": {
|
||||
"model": (
|
||||
"MODEL",
|
||||
{
|
||||
"tooltip": "The diffusion model the LoRAs will be applied to. "
|
||||
"Make sure the model is loaded by `Nunchaku FLUX DiT Loader`."
|
||||
},
|
||||
),
|
||||
},
|
||||
"optional": {},
|
||||
}
|
||||
|
||||
# Add fixed number of LoRA inputs (15 slots)
|
||||
for i in range(1, 16): # Support up to 15 LoRAs
|
||||
inputs["optional"][f"lora_name_{i}"] = (
|
||||
["None"] + get_filename_list("loras"),
|
||||
{"tooltip": f"The file name of LoRA {i}. Select 'None' to skip this slot."},
|
||||
)
|
||||
inputs["optional"][f"lora_strength_{i}"] = (
|
||||
"FLOAT",
|
||||
{
|
||||
"default": 1.0,
|
||||
"min": -100.0,
|
||||
"max": 100.0,
|
||||
"step": 0.01,
|
||||
"tooltip": f"Strength for LoRA {i}. This value can be negative.",
|
||||
},
|
||||
)
|
||||
|
||||
return inputs
|
||||
|
||||
RETURN_TYPES = ("MODEL",)
|
||||
OUTPUT_TOOLTIPS = ("The modified diffusion model with all LoRAs applied.",)
|
||||
FUNCTION = "load_lora_stack"
|
||||
TITLE = "Nunchaku FLUX LoRA Stack"
|
||||
|
||||
CATEGORY = "Nunchaku"
|
||||
DESCRIPTION = (
|
||||
"Apply multiple LoRAs to a diffusion model in a single node. "
|
||||
"Equivalent to chaining multiple LoRA nodes but more convenient for managing many LoRAs. "
|
||||
"Supports up to 15 LoRAs simultaneously. Set unused slots to 'None' to skip them."
|
||||
)
|
||||
|
||||
def load_lora_stack(self, model, **kwargs):
|
||||
"""
|
||||
Apply multiple LoRAs to a Nunchaku FLUX diffusion model.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
model : object
|
||||
The diffusion model to modify.
|
||||
**kwargs
|
||||
Dynamic LoRA name and strength parameters.
|
||||
|
||||
Returns
|
||||
-------
|
||||
tuple
|
||||
A tuple containing the modified diffusion model.
|
||||
"""
|
||||
# Collect LoRA information to apply
|
||||
loras_to_apply = []
|
||||
|
||||
for i in range(1, 16): # Check all 15 LoRA slots
|
||||
lora_name = kwargs.get(f"lora_name_{i}")
|
||||
lora_strength = kwargs.get(f"lora_strength_{i}", 1.0)
|
||||
|
||||
# Skip unset or None LoRAs
|
||||
if lora_name is None or lora_name == "None" or lora_name == "":
|
||||
continue
|
||||
|
||||
# Skip LoRAs with zero strength
|
||||
if abs(lora_strength) < 1e-5:
|
||||
continue
|
||||
|
||||
loras_to_apply.append((lora_name, lora_strength))
|
||||
|
||||
# If no LoRAs need to be applied, return the original model
|
||||
if not loras_to_apply:
|
||||
return (model,)
|
||||
|
||||
model_wrapper = model.model.diffusion_model
|
||||
assert isinstance(model_wrapper, ComfyFluxWrapper)
|
||||
|
||||
ret_model_wrapper, ret_model = copy_with_ctx(model_wrapper)
|
||||
|
||||
# Clear existing LoRA list
|
||||
ret_model_wrapper.loras = []
|
||||
|
||||
# Track the maximum input channels needed
|
||||
max_in_channels = ret_model.model.model_config.unet_config["in_channels"]
|
||||
|
||||
# Add all LoRAs
|
||||
for lora_name, lora_strength in loras_to_apply:
|
||||
lora_path = get_full_path_or_raise("loras", lora_name)
|
||||
ret_model_wrapper.loras.append((lora_path, lora_strength))
|
||||
|
||||
# Check if input channels need to be updated
|
||||
sd = to_diffusers(lora_path)
|
||||
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
|
||||
max_in_channels = max(max_in_channels, new_in_channels)
|
||||
|
||||
# Update the model's input channels
|
||||
if max_in_channels > ret_model.model.model_config.unet_config["in_channels"]:
|
||||
ret_model.model.model_config.unet_config["in_channels"] = max_in_channels
|
||||
|
||||
return (ret_model,)
|
||||
@@ -1,225 +0,0 @@
|
||||
"""
|
||||
This module provides the :class:`NunchakuZImageDiTLoader` class for loading Nunchaku Z-Image models.
|
||||
"""
|
||||
|
||||
import json
|
||||
import logging
|
||||
|
||||
import comfy.utils
|
||||
import torch
|
||||
from comfy import model_detection, model_management
|
||||
|
||||
from nunchaku.models.transformers.utils import convert_fp16, patch_scale_key
|
||||
from nunchaku.utils import check_hardware_compatibility, get_precision_from_quantization_config, is_turing
|
||||
|
||||
from ...model_configs.zimage import NunchakuZImage
|
||||
from ...model_patcher.zimage import ZImageModelPatcher
|
||||
from ..utils import get_filename_list, get_full_path_or_raise
|
||||
|
||||
|
||||
def _patch_state_dict(state_dict: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]:
|
||||
"""
|
||||
Patch the state dict.
|
||||
|
||||
Convert the keys from diffusers style to Comfy-Org style keys.
|
||||
See https://huggingface.co/Comfy-Org/z_image_turbo/blob/main/split_files/diffusion_models/z_image_turbo_bf16.safetensors
|
||||
|
||||
Parameters
|
||||
----------
|
||||
state_dict (dict[str, torch.Tensor]):
|
||||
the state dict loaded from diffuser style safetensors checkpoint file.
|
||||
|
||||
Returns
|
||||
-------
|
||||
dict[str, torch.Tensor]:
|
||||
the patched state dict which follows Comfy-Org style keys.
|
||||
"""
|
||||
patched_state_dict = {}
|
||||
quant_sub_keys = ["wscales", "wcscales", "wtscale", "smooth_factor_orig", "smooth_factor", "proj_down", "proj_up"]
|
||||
|
||||
for key, value in state_dict.items():
|
||||
# region: attention
|
||||
## case for quantized attention qkv
|
||||
if "attention.to_qkv" in key:
|
||||
patched_state_dict[key.replace("to_qkv", "qkv")] = value
|
||||
## case for non-quantized attetion qkv
|
||||
elif "attention.to_q" in key:
|
||||
q_weight = state_dict[key]
|
||||
k_weight = state_dict[key.replace("to_q", "to_k")]
|
||||
v_weight = state_dict[key.replace("to_q", "to_v")]
|
||||
patched_state_dict[key.replace("to_q", "qkv")] = torch.cat([q_weight, k_weight, v_weight], dim=0)
|
||||
elif "attention.to_k" in key or "attention.to_v" in key:
|
||||
continue
|
||||
## case for attention out
|
||||
elif "attention.to_out" in key:
|
||||
patched_state_dict[key.replace("to_out.0", "out")] = value
|
||||
# end of region
|
||||
# region: feed forward
|
||||
## case for quantized feed forward
|
||||
elif "feed_forward.net.0.proj.qweight" in key:
|
||||
patched_state_dict[key.replace("net.0.proj", "w13")] = value
|
||||
for subkey in quant_sub_keys:
|
||||
quant_param_key = key.replace("qweight", subkey)
|
||||
if quant_param_key in state_dict:
|
||||
patched_state_dict[quant_param_key.replace("net.0.proj", "w13")] = state_dict[quant_param_key]
|
||||
elif any("feed_forward.net.0.proj." + subkey in key for subkey in quant_sub_keys):
|
||||
continue
|
||||
elif "feed_forward.net.2.qweight" in key:
|
||||
patched_state_dict[key.replace("net.2", "w2")] = value
|
||||
for subkey in quant_sub_keys:
|
||||
quant_param_key = key.replace("qweight", subkey)
|
||||
if quant_param_key in state_dict:
|
||||
patched_state_dict[quant_param_key.replace("net.2", "w2")] = state_dict[quant_param_key]
|
||||
elif any("feed_forward.net.2." + subkey in key for subkey in quant_sub_keys):
|
||||
continue
|
||||
## case for non-quantized feed forward
|
||||
elif "feed_forward.net.0.proj.weight" in key:
|
||||
w3, w1 = torch.chunk(value, chunks=2, dim=0)
|
||||
w2 = state_dict[key.replace("0.proj", "2")]
|
||||
patched_state_dict[key.replace("net.0.proj", "w1")] = w1
|
||||
patched_state_dict[key.replace("net.0.proj", "w2")] = w2
|
||||
patched_state_dict[key.replace("net.0.proj", "w3")] = w3
|
||||
elif "feed_forward.net.2.weight" in key:
|
||||
continue
|
||||
# end of region
|
||||
# region: others
|
||||
elif "attention.norm_q" in key:
|
||||
patched_state_dict[key.replace("norm_q", "q_norm")] = value
|
||||
elif "attention.norm_k" in key:
|
||||
patched_state_dict[key.replace("norm_k", "k_norm")] = value
|
||||
elif "all_final_layer.2-1" in key:
|
||||
patched_state_dict[key.replace("all_final_layer.2-1", "final_layer")] = value
|
||||
elif "all_x_embedder.2-1" in key:
|
||||
patched_state_dict[key.replace("all_x_embedder.2-1", "x_embedder")] = value
|
||||
else:
|
||||
patched_state_dict[key] = value
|
||||
# end of region
|
||||
|
||||
return patched_state_dict
|
||||
|
||||
|
||||
def _load(sd: dict[str, torch.Tensor], metadata: dict[str, str] = {}):
|
||||
"""
|
||||
Load a Nunchaku-quantized Z-Image diffusion model.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
sd : dict[str, torch.Tensor]
|
||||
The state dictionary of the model.
|
||||
metadata : dict[str, str], optional
|
||||
Metadata containing quantization configuration (default is empty dict).
|
||||
|
||||
Returns
|
||||
-------
|
||||
comfy.model_patcher.ModelPatcher
|
||||
The patched and loaded Qwen-Image model ready for inference.
|
||||
"""
|
||||
quantization_config = json.loads(metadata.get("quantization_config", "{}"))
|
||||
precision = get_precision_from_quantization_config(quantization_config)
|
||||
rank = quantization_config.get("rank", 32)
|
||||
skip_refiners = quantization_config.get("skip_refiners", False)
|
||||
|
||||
# Allow loading unets from checkpoint files
|
||||
diffusion_model_prefix = model_detection.unet_prefix_from_state_dict(sd)
|
||||
temp_sd = comfy.utils.state_dict_prefix_replace(sd, {diffusion_model_prefix: ""}, filter_keys=True)
|
||||
if len(temp_sd) > 0:
|
||||
sd = temp_sd
|
||||
|
||||
parameters = comfy.utils.calculate_parameters(sd)
|
||||
weight_dtype = comfy.utils.weight_dtype(sd)
|
||||
|
||||
load_device = model_management.get_torch_device()
|
||||
offload_device = model_management.unet_offload_device()
|
||||
check_hardware_compatibility(quantization_config, load_device)
|
||||
|
||||
new_sd = sd
|
||||
|
||||
model_config = NunchakuZImage(rank=rank, precision=precision, skip_refiners=skip_refiners)
|
||||
|
||||
if not is_turing():
|
||||
unet_weight_dtype = list(model_config.supported_inference_dtypes)
|
||||
unet_dtype = model_management.unet_dtype(
|
||||
model_params=parameters, supported_dtypes=unet_weight_dtype, weight_dtype=weight_dtype
|
||||
)
|
||||
manual_cast_dtype = model_management.unet_manual_cast(
|
||||
unet_dtype, load_device, model_config.supported_inference_dtypes
|
||||
)
|
||||
torch_dtype = torch.bfloat16
|
||||
else:
|
||||
unet_dtype = torch.bfloat16
|
||||
manual_cast_dtype = torch.float16
|
||||
torch_dtype = torch.float16
|
||||
logging.info(f"unet_dtype: {unet_dtype}, manual_cast_dtype: {manual_cast_dtype}, svdq_linear_dtype: {torch_dtype}")
|
||||
model_config.set_inference_dtype(unet_dtype, manual_cast_dtype)
|
||||
|
||||
patched_sd = _patch_state_dict(new_sd)
|
||||
model = model_config.get_model(patched_sd, "", torch_dtype=torch_dtype)
|
||||
|
||||
patch_scale_key(model.diffusion_model, patched_sd)
|
||||
if torch_dtype == torch.float16:
|
||||
convert_fp16(model.diffusion_model, patched_sd)
|
||||
|
||||
model.load_model_weights(patched_sd, "")
|
||||
return ZImageModelPatcher(model, load_device=load_device, offload_device=offload_device)
|
||||
|
||||
|
||||
class NunchakuZImageDiTLoader:
|
||||
"""
|
||||
Loader for Nunchaku Z-Image models.
|
||||
|
||||
Attributes
|
||||
----------
|
||||
RETURN_TYPES : tuple
|
||||
Output types for the node ("MODEL",).
|
||||
FUNCTION : str
|
||||
Name of the function to call ("load_model").
|
||||
CATEGORY : str
|
||||
Node category ("Nunchaku").
|
||||
TITLE : str
|
||||
Node title ("Nunchaku Z-Image DiT Loader").
|
||||
"""
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
"""
|
||||
Define the input types and tooltips for the node.
|
||||
|
||||
Returns
|
||||
-------
|
||||
dict
|
||||
A dictionary specifying the required inputs and their descriptions for the node interface.
|
||||
"""
|
||||
return {
|
||||
"required": {
|
||||
"model_name": (
|
||||
get_filename_list("diffusion_models"),
|
||||
{"tooltip": "The Nunchaku Z-Image model."},
|
||||
),
|
||||
},
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("MODEL",)
|
||||
FUNCTION = "load_model"
|
||||
CATEGORY = "Nunchaku"
|
||||
TITLE = "Nunchaku Z-Image DiT Loader"
|
||||
|
||||
def load_model(self, model_name: str, **kwargs):
|
||||
"""
|
||||
Load the Z-Image model from file and return a patched model.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
model_name : str
|
||||
The filename of the Z-Image model to load.
|
||||
|
||||
Returns
|
||||
-------
|
||||
tuple
|
||||
A tuple containing the loaded and patched model.
|
||||
"""
|
||||
model_path = get_full_path_or_raise("diffusion_models", model_name)
|
||||
sd, metadata = comfy.utils.load_torch_file(model_path, return_metadata=True)
|
||||
|
||||
model = _load(sd, metadata=metadata)
|
||||
|
||||
return (model,)
|
||||
@@ -1,426 +0,0 @@
|
||||
import logging
|
||||
import uuid
|
||||
from collections import namedtuple
|
||||
|
||||
import comfy.model_management
|
||||
import comfy.utils
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
from comfy.float import stochastic_rounding
|
||||
from comfy.ldm.lumina.model import NextDiT
|
||||
from comfy.lora import calculate_weight, model_lora_keys_unet
|
||||
from comfy.model_base import BaseModel
|
||||
from comfy.model_management import cast_to_device, lora_compute_dtype
|
||||
from comfy.model_patcher import LowVramPatch, ModelPatcher, get_key_weight, move_weight_functions, string_to_seed
|
||||
from comfy.utils import get_attr, set_attr_param
|
||||
from comfy.weight_adapter.lora import LoRAAdapter
|
||||
|
||||
from nunchaku.lora.flux.nunchaku_converter import pack_lowrank_weight, unpack_lowrank_weight
|
||||
from nunchaku.models.linear import SVDQW4A4Linear
|
||||
|
||||
|
||||
def apply_lora_to_svdq_linear(linear: SVDQW4A4Linear, lora_down: torch.Tensor, lora_up: torch.Tensor):
|
||||
"""
|
||||
Apply LoRA weights into quantized SVDQW4A4Linear module.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
linear : SVDQW4A4Linear
|
||||
The Nunchaku quantized module to apply LoRAs to
|
||||
lora_down : torch.Tensor
|
||||
The down projection part of the LoRA weights.
|
||||
lora_up : torch.Tensor
|
||||
The up projection part of the LoRA weights.
|
||||
"""
|
||||
dtype = linear.proj_down.dtype
|
||||
proj_down = unpack_lowrank_weight(linear.proj_down, down=True)
|
||||
proj_up = unpack_lowrank_weight(linear.proj_up, down=False)
|
||||
concatenated_down = torch.cat([proj_down, lora_down], dim=0).to(dtype=dtype)
|
||||
concatenated_up = torch.cat([proj_up, lora_up], dim=1).to(dtype=dtype)
|
||||
|
||||
linear.proj_down = nn.Parameter(pack_lowrank_weight(concatenated_down, down=True))
|
||||
linear.proj_up = nn.Parameter(pack_lowrank_weight(concatenated_up, down=False))
|
||||
|
||||
|
||||
def concat_lora_weights(
|
||||
base_down: torch.Tensor,
|
||||
base_up: torch.Tensor,
|
||||
new_downs: list[torch.Tensor],
|
||||
new_ups: list[torch.Tensor],
|
||||
strengths: list[float],
|
||||
):
|
||||
"""
|
||||
Concatenate multiple LoRA weights into single down and up weights.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
base_down : torch.Tensor
|
||||
The base LoRA down weight.
|
||||
base_up : torch.Tensor
|
||||
The base LoRA up weight.
|
||||
new_downs : list of torch.Tensor
|
||||
List of new LoRA down weights to concatenate.
|
||||
new_ups : list of torch.Tensor
|
||||
List of new LoRA up weights to concatenate.
|
||||
strengths : list of float
|
||||
List of strength/scale factors for each new LoRA.
|
||||
|
||||
Returns
|
||||
-------
|
||||
tuple of torch.Tensor
|
||||
The concatenated down and up weights.
|
||||
"""
|
||||
assert len(new_downs) == len(new_ups) == len(strengths), "Lengths of new_downs, new_ups, and strengths must match."
|
||||
assert (base_down is None) == (base_up is None), "Both base_down and base_up should be None or not None."
|
||||
|
||||
combined_new_downs = torch.cat([nd * s for nd, s in zip(new_downs, strengths)], dim=0)
|
||||
combined_new_ups = torch.block_diag(*new_ups)
|
||||
|
||||
if base_down is None:
|
||||
return combined_new_downs, combined_new_ups
|
||||
|
||||
assert base_up.shape[1] == base_down.shape[0], "Base up and down weights shapes do not match."
|
||||
assert base_up.shape[0] == sum(
|
||||
new_up.shape[0] for new_up in new_ups
|
||||
), "Total new up weights rows do not match base up weight rows."
|
||||
assert all(
|
||||
new_down.shape[1] == base_down.shape[1] for new_down in new_downs
|
||||
), "New up and down weights shapes do not match."
|
||||
|
||||
concatenated_down = torch.cat([base_down, combined_new_downs], dim=0)
|
||||
concatenated_up = torch.cat([base_up, combined_new_ups], dim=1)
|
||||
return concatenated_down, concatenated_up
|
||||
|
||||
|
||||
class ZImageModelPatcher(ModelPatcher):
|
||||
def __init__(self, model, load_device, offload_device, size=0, weight_inplace_update=False):
|
||||
"""
|
||||
Adapted from comfy.model_patcher.ModelPatcher#clone
|
||||
|
||||
Note
|
||||
----
|
||||
+ Always set `weight_inplace_update` to False
|
||||
+ Add `svdq_backup` dict for loading/unloading lora weights for Nunchaku Z-Image model.
|
||||
"""
|
||||
super().__init__(model, load_device, offload_device, size, weight_inplace_update=False)
|
||||
self.svdq_backup = {}
|
||||
|
||||
def clone(self):
|
||||
"""
|
||||
Adapted from comfy.model_patcher.ModelPatcher#clone
|
||||
|
||||
Note
|
||||
----
|
||||
+ Clone quantized svdq_backup weights for Nunchaku Z-Image model.
|
||||
"""
|
||||
n = ModelPatcher.clone(self)
|
||||
n.svdq_backup = self.svdq_backup
|
||||
logging.debug("ZImageModelPatcher cloned.")
|
||||
return n
|
||||
|
||||
def unpatch_model(self, device_to=None, unpatch_weights=True):
|
||||
"""
|
||||
Adapted from comfy.model_patcher.ModelPatcher#unpatch_model
|
||||
|
||||
Note
|
||||
----
|
||||
+ Unpatch loras and restore original weights from svdq_backup for Nunchaku Z-Image model.
|
||||
"""
|
||||
super().unpatch_model(device_to, unpatch_weights)
|
||||
if unpatch_weights:
|
||||
for k, (proj_down, proj_up) in self.svdq_backup.items():
|
||||
logging.debug(f"<<<< Unpatching svdq_linears, key: {k}, device_to: {device_to}")
|
||||
attr_name = k.replace(".qweight", "")
|
||||
linear = get_attr(self.model, attr_name)
|
||||
assert isinstance(linear, SVDQW4A4Linear)
|
||||
if device_to is not None:
|
||||
proj_down, proj_up = proj_down.to(device=device_to), proj_up.to(device=device_to)
|
||||
linear.proj_down = proj_down
|
||||
linear.proj_up = proj_up
|
||||
self.svdq_backup.clear()
|
||||
|
||||
def do_svdq_linear_backup(self, key: str, linear: SVDQW4A4Linear):
|
||||
assert key.endswith(".qweight"), f">>> Unexpected key {key} to do svdq_linear backup"
|
||||
if key not in self.svdq_backup:
|
||||
self.svdq_backup[key] = (linear.proj_down, linear.proj_up)
|
||||
|
||||
def partially_unload(self, device_to, memory_to_free=0, force_patch_weights=False):
|
||||
"""
|
||||
Adapted from comfy.model_patcher.ModelPatcher#partially_unload
|
||||
|
||||
Note
|
||||
----
|
||||
+ Unload quantized svdq linears for Nunchaku Z-Image model.
|
||||
"""
|
||||
with self.use_ejected():
|
||||
hooks_unpatched = False
|
||||
memory_freed = 0
|
||||
patch_counter = 0
|
||||
unload_list = self._load_list()
|
||||
unload_list.sort()
|
||||
|
||||
offload_buffer = self.model.model_offload_buffer_memory
|
||||
if len(unload_list) > 0:
|
||||
NS = comfy.model_management.NUM_STREAMS
|
||||
offload_weight_factor = [min(offload_buffer / (NS + 1), unload_list[0][1])] * NS
|
||||
|
||||
for unload in unload_list:
|
||||
if memory_to_free + offload_buffer - self.model.model_offload_buffer_memory < memory_freed:
|
||||
break
|
||||
module_offload_mem, module_mem, n, m, params = unload
|
||||
|
||||
potential_offload = module_offload_mem + sum(offload_weight_factor)
|
||||
|
||||
lowvram_possible = hasattr(m, "comfy_cast_weights")
|
||||
if hasattr(m, "comfy_patched_weights") and m.comfy_patched_weights:
|
||||
move_weight = True
|
||||
for param in params:
|
||||
key = "{}.{}".format(n, param)
|
||||
bk = self.backup.get(key, None)
|
||||
svdq_bk = self.svdq_backup.get(key, None)
|
||||
if bk is not None or svdq_bk is not None:
|
||||
if not lowvram_possible:
|
||||
move_weight = False
|
||||
break
|
||||
|
||||
if not hooks_unpatched:
|
||||
self.unpatch_hooks()
|
||||
hooks_unpatched = True
|
||||
if bk is not None:
|
||||
if bk.inplace_update:
|
||||
comfy.utils.copy_to_param(self.model, key, bk.weight)
|
||||
else:
|
||||
comfy.utils.set_attr_param(self.model, key, bk.weight)
|
||||
self.backup.pop(key)
|
||||
# region unload svdq linears
|
||||
if svdq_bk is not None:
|
||||
logging.debug(f">>> partially unload svdq, key: {key}")
|
||||
proj_down, proj_up = svdq_bk
|
||||
attr_name = key.replace(".qweight", "")
|
||||
linear = get_attr(self.model, attr_name)
|
||||
assert isinstance(linear, SVDQW4A4Linear)
|
||||
linear.proj_down = proj_down
|
||||
linear.proj_up = proj_up
|
||||
self.svdq_backup.pop(key)
|
||||
# end of region
|
||||
|
||||
weight_key = "{}.weight".format(n)
|
||||
bias_key = "{}.bias".format(n)
|
||||
if move_weight:
|
||||
cast_weight = self.force_cast_weights
|
||||
m.to(device_to)
|
||||
module_mem += move_weight_functions(m, device_to)
|
||||
if lowvram_possible:
|
||||
if weight_key in self.patches:
|
||||
if force_patch_weights:
|
||||
self.patch_weight_to_device(weight_key)
|
||||
else:
|
||||
_, set_func, convert_func = get_key_weight(self.model, weight_key)
|
||||
m.weight_function.append(
|
||||
LowVramPatch(weight_key, self.patches, convert_func, set_func)
|
||||
)
|
||||
patch_counter += 1
|
||||
if bias_key in self.patches:
|
||||
if force_patch_weights:
|
||||
self.patch_weight_to_device(bias_key)
|
||||
else:
|
||||
_, set_func, convert_func = get_key_weight(self.model, bias_key)
|
||||
m.bias_function.append(LowVramPatch(bias_key, self.patches, convert_func, set_func))
|
||||
patch_counter += 1
|
||||
cast_weight = True
|
||||
|
||||
if cast_weight and hasattr(m, "comfy_cast_weights"):
|
||||
m.prev_comfy_cast_weights = m.comfy_cast_weights
|
||||
m.comfy_cast_weights = True
|
||||
m.comfy_patched_weights = False
|
||||
memory_freed += module_mem
|
||||
offload_buffer = max(offload_buffer, potential_offload)
|
||||
offload_weight_factor.append(module_mem)
|
||||
offload_weight_factor.pop(0)
|
||||
logging.debug("freed {}".format(n))
|
||||
|
||||
for param in params:
|
||||
self.pin_weight_to_device("{}.{}".format(n, param))
|
||||
|
||||
self.model.model_lowvram = True
|
||||
self.model.lowvram_patch_counter += patch_counter
|
||||
self.model.model_loaded_weight_memory -= memory_freed
|
||||
self.model.model_offload_buffer_memory = offload_buffer
|
||||
logging.info(
|
||||
"Unloaded partially: {:.2f} MB freed, {:.2f} MB remains loaded, {:.2f} MB buffer reserved, lowvram patches: {}".format(
|
||||
memory_freed / (1024 * 1024),
|
||||
self.model.model_loaded_weight_memory / (1024 * 1024),
|
||||
offload_buffer / (1024 * 1024),
|
||||
self.model.lowvram_patch_counter,
|
||||
)
|
||||
)
|
||||
return memory_freed
|
||||
|
||||
def add_patches(self, patches, strength_patch=1.0, strength_model=1.0):
|
||||
"""
|
||||
Adapted from comfy.model_patcher.ModelPatcher#add_patches
|
||||
|
||||
Note
|
||||
----
|
||||
+ Modify the valid key_set
|
||||
"""
|
||||
with self.use_ejected():
|
||||
p = set()
|
||||
key_set = set(model_lora_keys_unet(self.model).values())
|
||||
for k in patches:
|
||||
offset = None
|
||||
function = None
|
||||
if isinstance(k, str):
|
||||
key = k
|
||||
else:
|
||||
offset = k[1]
|
||||
key = k[0]
|
||||
if len(k) > 2:
|
||||
function = k[2]
|
||||
if k in key_set:
|
||||
p.add(k)
|
||||
current_patches = self.patches.get(key, [])
|
||||
current_patches.append((strength_patch, patches[k], strength_model, offset, function))
|
||||
self.patches[key] = current_patches
|
||||
|
||||
self.patches_uuid = uuid.uuid4()
|
||||
return list(p)
|
||||
|
||||
def patch_weight_to_device(self, key, device_to=None, inplace_update=False):
|
||||
"""
|
||||
Adapted from comfy.model_patcher.ModelPatcher#patch_weight_to_device
|
||||
|
||||
Note
|
||||
----
|
||||
+ Patch non-quantized linear modules the same way as original ModelPatcher
|
||||
+ Patch quantized svdq_linear modules by concatenate loras into proj_down and proj_up.
|
||||
"""
|
||||
assert isinstance(self.model, BaseModel)
|
||||
assert isinstance(self.model.diffusion_model, NextDiT)
|
||||
assert isinstance(key, str)
|
||||
assert (
|
||||
inplace_update is False
|
||||
), "In-place weight update is not supported for applying LoRA patches to Nunchaku Z-Image models."
|
||||
|
||||
if len(self.patches) == 0:
|
||||
return
|
||||
|
||||
# the `key` argument passed in could be a name of any nn.Parameter in the model
|
||||
if (key not in self.patches) and not (key.endswith(".qweight")):
|
||||
return
|
||||
|
||||
assert key.endswith(".qweight") or key.endswith(".weight"), f"Unexpected key: {key}"
|
||||
attr_name = key.rsplit(".", 1)[0]
|
||||
linear = get_attr(self.model, attr_name)
|
||||
|
||||
if isinstance(linear, nn.Linear):
|
||||
logging.debug(
|
||||
f">>> Applying LoRA to non-quantized linear at {attr_name}, linear weight shape: {linear.weight.shape}, number of patches: {len(self.patches[key])}"
|
||||
)
|
||||
# Apply LoRA to standard linear layer. This is the same way as in comfy.model_patcher.ModelPatcher#patch_weight_to_device
|
||||
weight = linear.weight
|
||||
if key not in self.backup:
|
||||
self.backup[key] = namedtuple("Dimension", ["weight", "inplace_update"])(
|
||||
weight.to(device=self.offload_device, copy=inplace_update), inplace_update
|
||||
)
|
||||
temp_dtype = lora_compute_dtype(device_to)
|
||||
if device_to is not None:
|
||||
temp_weight = cast_to_device(weight, device_to, temp_dtype, copy=True)
|
||||
else:
|
||||
temp_weight = weight.to(temp_dtype, copy=True)
|
||||
out_weight = calculate_weight(self.patches[key], temp_weight, key)
|
||||
out_weight = stochastic_rounding(out_weight, weight.dtype, seed=string_to_seed(key))
|
||||
set_attr_param(self.model, key, out_weight)
|
||||
elif isinstance(linear, SVDQW4A4Linear):
|
||||
|
||||
def _parse_patch(_the_patch):
|
||||
strength = _the_patch[0]
|
||||
lora_adapter: LoRAAdapter = _the_patch[1]
|
||||
up, down, alpha, _, _, _ = lora_adapter.weights
|
||||
rank = down.shape[0]
|
||||
return strength, up, down, alpha, rank
|
||||
|
||||
if ".qkv." in key:
|
||||
patches_list_for_key = self.patches.get(key.replace(".qweight", ".weight"), [])
|
||||
logging.debug(
|
||||
f">>> Applying LoRA to fused qkv at {key}, svdq_linear shape: {(linear.in_features, linear.out_features)}, svdq_linear rank: {linear.rank}, number of patches: {len(patches_list_for_key)}"
|
||||
)
|
||||
if len(patches_list_for_key) > 0:
|
||||
q_patches, k_patches, v_patches = [], [], []
|
||||
for strength_patch, lora_adapter, strength_model, offset, function in patches_list_for_key:
|
||||
assert isinstance(lora_adapter, LoRAAdapter)
|
||||
assert (
|
||||
offset is not None
|
||||
), f"Offset must be provided for qkv LoRA patch, but got None for key {key}"
|
||||
if offset[1] == 0:
|
||||
q_patches.append((strength_patch, lora_adapter, strength_model, offset, function))
|
||||
elif offset[1] == offset[2]:
|
||||
k_patches.append((strength_patch, lora_adapter, strength_model, offset, function))
|
||||
elif offset[1] == offset[2] * 2:
|
||||
v_patches.append((strength_patch, lora_adapter, strength_model, offset, function))
|
||||
else:
|
||||
raise ValueError(f"Invalid offset {offset} for qkv LoRA patch in key {key}")
|
||||
|
||||
assert len(q_patches) == len(k_patches) == len(v_patches)
|
||||
composed_up, composed_down = None, None
|
||||
for q_patch, k_patch, v_patch in zip(q_patches, k_patches, v_patches):
|
||||
q_strength, q_up, q_down, q_alpha, q_rank = _parse_patch(q_patch)
|
||||
k_strength, k_up, k_down, k_alpha, k_rank = _parse_patch(k_patch)
|
||||
v_strength, v_up, v_down, v_alpha, v_rank = _parse_patch(v_patch)
|
||||
composed_down, composed_up = concat_lora_weights(
|
||||
composed_down,
|
||||
composed_up,
|
||||
[q_down, k_down, v_down],
|
||||
[q_up, k_up, v_up],
|
||||
[
|
||||
q_strength * (q_alpha / q_rank) if q_alpha is not None else q_strength,
|
||||
k_strength * (k_alpha / k_rank) if k_alpha is not None else k_strength,
|
||||
v_strength * (v_alpha / v_rank) if v_alpha is not None else v_strength,
|
||||
],
|
||||
)
|
||||
self.do_svdq_linear_backup(key, linear)
|
||||
apply_lora_to_svdq_linear(linear, composed_down, composed_up)
|
||||
elif ".feed_forward.w13." in key:
|
||||
w1_patches = self.patches.get(key.replace("w13.qweight", "w1.weight"), [])
|
||||
w3_patches = self.patches.get(key.replace("w13.qweight", "w3.weight"), [])
|
||||
assert len(w1_patches) == len(w3_patches)
|
||||
logging.debug(
|
||||
f">>> Applying LoRA to fused ff.w13 at {key}, svdq_linear shape: {(linear.in_features, linear.out_features)}, svdq_linear rank: {linear.rank}, number of patches: {len(w1_patches)}"
|
||||
)
|
||||
if len(w1_patches) > 0:
|
||||
composed_up, composed_down = None, None
|
||||
for w1_patch, w3_patch in zip(w1_patches, w3_patches):
|
||||
w1_strength, w1_up, w1_down, w1_alpha, w1_rank = _parse_patch(w1_patch)
|
||||
w3_strength, w3_up, w3_down, w3_alpha, w3_rank = _parse_patch(w3_patch)
|
||||
composed_down, composed_up = concat_lora_weights(
|
||||
composed_down,
|
||||
composed_up,
|
||||
[w3_down, w1_down],
|
||||
[w3_up, w1_up],
|
||||
[
|
||||
w3_strength * (w3_alpha / w3_rank) if w3_alpha is not None else w3_strength,
|
||||
w1_strength * (w1_alpha / w1_rank) if w1_alpha is not None else w1_strength,
|
||||
],
|
||||
)
|
||||
self.do_svdq_linear_backup(key, linear)
|
||||
apply_lora_to_svdq_linear(linear, composed_down, composed_up)
|
||||
else:
|
||||
patches_list_for_key = self.patches.get(key.replace(".qweight", ".weight"), [])
|
||||
logging.debug(
|
||||
f">>> Applying LoRA to non-fused quantized key {key}, svdq_linear shape: {(linear.in_features, linear.out_features)}, svdq_linear rank: {linear.rank}, number of patches: {len(patches_list_for_key)}"
|
||||
)
|
||||
if len(patches_list_for_key) > 0:
|
||||
composed_up, composed_down = None, None
|
||||
for patch in patches_list_for_key:
|
||||
strength, up, down, alpha, rank = _parse_patch(patch)
|
||||
composed_down, composed_up = concat_lora_weights(
|
||||
composed_down,
|
||||
composed_up,
|
||||
[down],
|
||||
[up],
|
||||
[strength * (alpha / rank) if alpha is not None else strength],
|
||||
)
|
||||
self.do_svdq_linear_backup(key, linear)
|
||||
apply_lora_to_svdq_linear(linear, composed_down, composed_up)
|
||||
else:
|
||||
raise ValueError(f"Unsupported module type for LoRA patching: {type(linear)}, key: {key}")
|
||||
Reference in New Issue
Block a user