438 lines
15 KiB
Python
438 lines
15 KiB
Python
import os
|
|
import re
|
|
import json
|
|
import torch
|
|
import folder_paths
|
|
|
|
from .loader import load_flux_mod
|
|
from .sampler import common_ksampler
|
|
import comfy.samplers
|
|
import comfy.cli_args
|
|
from comfy import model_management
|
|
import node_helpers
|
|
|
|
|
|
def using_scaled_fp8(model_patcher):
|
|
return (
|
|
comfy.cli_args.args.fast
|
|
and model_management.supports_fp8_compute(model_patcher.load_device)
|
|
) or model_patcher.model.model_config.scaled_fp8
|
|
|
|
|
|
class FluxModDiffusionLoader:
|
|
NodeId = "FluxModDiffusionLoader"
|
|
NodeName = "FluxMod Unified Model Loader"
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
patches = ["None"] + folder_paths.get_filename_list("diffusion_models")
|
|
checkpoint_paths = folder_paths.get_filename_list("diffusion_models")
|
|
if "unet_gguf" in folder_paths.folder_names_and_paths:
|
|
checkpoint_paths = checkpoint_paths + folder_paths.get_filename_list(
|
|
"unet_gguf"
|
|
)
|
|
return {
|
|
"required": {
|
|
"unet_name": (checkpoint_paths, {"tooltip": "This loads both .safetensors and .GGUF models if ComfyUI-GGUF is installed."}),
|
|
"guidance_name": (patches,),
|
|
"quant_mode": (
|
|
["bf16", "float8_e4m3fn (8 bit)", "float8_e5m2 (also 8 bit)"],
|
|
),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("MODEL",)
|
|
RETURN_NAMES = ("model",)
|
|
FUNCTION = "load_unet"
|
|
CATEGORY = "advanced/FluxMod"
|
|
|
|
def load_unet(
|
|
self, *, unet_name, quant_mode, guidance_name=None, lite_patch_unet_name=None
|
|
):
|
|
dtypes = {
|
|
"bf16": torch.bfloat16,
|
|
"float8_e4m3fn (8 bit)": torch.float8_e4m3fn,
|
|
"float8_e5m2 (also 8 bit)": torch.float8_e5m2,
|
|
}
|
|
|
|
is_gguf = unet_name.lower().endswith(".gguf")
|
|
unet_path = folder_paths.get_full_path(
|
|
"unet_gguf" if is_gguf else "diffusion_models", unet_name
|
|
)
|
|
if guidance_name is not None:
|
|
guidance_path = folder_paths.get_full_path(
|
|
"diffusion_models", guidance_name
|
|
)
|
|
else:
|
|
guidance_path = None
|
|
if lite_patch_unet_name is not None:
|
|
lite_patch_unet_name = folder_paths.get_full_path(
|
|
"diffusion_models", lite_patch_unet_name
|
|
)
|
|
flux_mod = load_flux_mod(
|
|
model_path=unet_path,
|
|
timestep_guidance_path=guidance_path,
|
|
linear_dtypes=dtypes[quant_mode],
|
|
lite_patch_path=lite_patch_unet_name,
|
|
is_gguf=is_gguf,
|
|
)
|
|
return (flux_mod,)
|
|
|
|
|
|
class FluxModDiffusionLoaderMini(FluxModDiffusionLoader):
|
|
NodeId = "FluxModDiffusionLoaderMini"
|
|
NodeName = "FluxMod Mini Unified Model Loader"
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
patches = ["None"] + folder_paths.get_filename_list("diffusion_models")
|
|
result = super().INPUT_TYPES()
|
|
result["required"] |= {
|
|
"lite_patch_unet_name": (patches,),
|
|
}
|
|
return result
|
|
|
|
|
|
class ChromaDiffusionLoader(FluxModDiffusionLoader):
|
|
NodeId = "ChromaDiffusionLoader"
|
|
NodeName = "Chroma Unified Model Loader"
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
result = super().INPUT_TYPES()
|
|
del result["required"]["guidance_name"]
|
|
return result
|
|
|
|
|
|
class ChromaPromptTruncation:
|
|
NodeId = "ChromaPaddingRemoval"
|
|
NodeName = "Padding Removal"
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"conditioning": ("CONDITIONING",),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("CONDITIONING",)
|
|
FUNCTION = "append"
|
|
|
|
CATEGORY = "advanced/conditioning/flux"
|
|
|
|
def append(self, conditioning):
|
|
pruning_idx = conditioning[0][1]["attention_mask"].sum() + 1
|
|
conditioning[0][0] = conditioning[0][0][:, :pruning_idx]
|
|
c = node_helpers.conditioning_set_values(conditioning, {"guidance": 0})
|
|
return (c,)
|
|
|
|
|
|
class ChromaStyleModelApply:
|
|
NodeId = "ChromaStyleModelApply"
|
|
NodeName = "Chroma Style Model"
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {"required": {"conditioning": ("CONDITIONING", ),
|
|
"style_model": ("STYLE_MODEL", ),
|
|
"clip_vision_output": ("CLIP_VISION_OUTPUT", ),
|
|
"strength": ("FLOAT", {"default": 1.0, "min": -10.0, "max": 10.0, "step": 0.001}),
|
|
"truncate_percent": ("FLOAT", {"default": 1.00, "min": -1.0, "max": 1.0, "step": 0.001, "tooltip": "Truncates clipvision conditioning to the first truncate_percent values when > 0. Truncates the last |truncate_percent| values when < 0."}),
|
|
}}
|
|
RETURN_TYPES = ("CONDITIONING",)
|
|
FUNCTION = "apply_stylemodel"
|
|
|
|
CATEGORY = "conditioning/style_model"
|
|
|
|
def apply_stylemodel(self, conditioning, style_model, clip_vision_output, strength, truncate_percent):
|
|
cond = style_model.get_cond(clip_vision_output).flatten(start_dim=0, end_dim=1).unsqueeze(dim=0)
|
|
|
|
if truncate_percent < 1.0 and truncate_percent >= 0.0:
|
|
cond_norm = cond.norm(2) # Take stylemodel cond norm before truncation
|
|
cond = cond[:, :int(cond.shape[1] * truncate_percent), :] # Take first (729 * truncate_percent) blocks
|
|
cond *= cond_norm / cond.norm(2) # Re-normalize
|
|
elif truncate_percent < 0.0 and truncate_percent >= -1.0:
|
|
cond_norm = cond.norm(2)
|
|
cond = cond[:, -int(cond.shape[1] * abs(truncate_percent)):, :] # Take last (729 * truncate_percent) blocks
|
|
cond *= cond_norm / cond.norm(2)
|
|
|
|
cond *= strength # Normal strength
|
|
|
|
c_out = []
|
|
for t in conditioning:
|
|
(txt, keys) = t
|
|
keys = keys.copy()
|
|
|
|
c_out.append([torch.cat((txt, cond), dim=1), keys])
|
|
|
|
return (c_out,)
|
|
|
|
|
|
class ModelMover:
|
|
NodeId = "ModelMover"
|
|
NodeName = "???"
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"samples": ("MODEL",),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("MODEL",)
|
|
RETURN_NAMES = ("model",)
|
|
FUNCTION = "load_unet"
|
|
CATEGORY = "advanced/FluxMod"
|
|
|
|
def load_unet(self, unet_name, guidance_name, quant_mode):
|
|
dtypes = {
|
|
"bf16": torch.bfloat16,
|
|
"float8_e4m3fn (8 bit)": torch.float8_e4m3fn,
|
|
"float8_e5m2 (also 8 bit)": torch.float8_e5m2,
|
|
}
|
|
|
|
unet_path = folder_paths.get_full_path("diffusion_models", unet_name)
|
|
guidance_path = folder_paths.get_full_path("diffusion_models", guidance_name)
|
|
flux_mod = load_flux_mod(
|
|
model_path=unet_path,
|
|
timestep_guidance_path=guidance_path,
|
|
linear_dtypes=dtypes[quant_mode],
|
|
)
|
|
return (flux_mod,)
|
|
|
|
|
|
class SkipLayerForward:
|
|
NodeId = "SkipLayerForward"
|
|
NodeName = "FluxMod Prune Model Layers"
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"model": ("MODEL",),
|
|
"skip_mmdit_layers": ("STRING", {"default": "10", "multiline": False}),
|
|
"skip_dit_layers": ("STRING", {"default": "3, 4", "multiline": False}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("MODEL",)
|
|
RETURN_NAMES = ("model",)
|
|
FUNCTION = "skip_layer"
|
|
CATEGORY = "advanced/FluxMod"
|
|
|
|
DESCRIPTION = "Prune model layers"
|
|
|
|
def skip_layer(self, model, skip_mmdit_layers, skip_dit_layers):
|
|
|
|
skip_mmdit_layers = re.split(r"\s*,\s*", skip_mmdit_layers)
|
|
skip_mmdit_layers = [int(num) for num in skip_mmdit_layers]
|
|
|
|
skip_dit_layers = re.split(r"\s*,\s*", skip_dit_layers)
|
|
skip_dit_layers = [int(num) for num in skip_dit_layers]
|
|
|
|
model.model.diffusion_model.skip_dit = skip_dit_layers
|
|
model.model.diffusion_model.skip_mmdit = skip_mmdit_layers
|
|
return (model,)
|
|
|
|
|
|
class KSamplerMod:
|
|
NodeId = "KSamplerMod"
|
|
NodeName = "FluxMod KSampler"
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"model": (
|
|
"MODEL",
|
|
{"tooltip": "The model used for denoising the input latent."},
|
|
),
|
|
"seed": (
|
|
"INT",
|
|
{
|
|
"default": 0,
|
|
"min": 0,
|
|
"max": 0xFFFFFFFFFFFFFFFF,
|
|
"tooltip": "The random seed used for creating the noise.",
|
|
},
|
|
),
|
|
"steps": (
|
|
"INT",
|
|
{
|
|
"default": 20,
|
|
"min": 1,
|
|
"max": 10000,
|
|
"tooltip": "The number of steps used in the denoising process.",
|
|
},
|
|
),
|
|
"cfg": (
|
|
"FLOAT",
|
|
{
|
|
"default": 8.0,
|
|
"min": 0.0,
|
|
"max": 100.0,
|
|
"step": 0.1,
|
|
"round": 0.01,
|
|
"tooltip": "The Classifier-Free Guidance scale balances creativity and adherence to the prompt. Higher values result in images more closely matching the prompt however too high values will negatively impact quality.",
|
|
},
|
|
),
|
|
"sampler_name": (
|
|
comfy.samplers.KSampler.SAMPLERS,
|
|
{
|
|
"tooltip": "The algorithm used when sampling, this can affect the quality, speed, and style of the generated output."
|
|
},
|
|
),
|
|
"scheduler": (
|
|
comfy.samplers.KSampler.SCHEDULERS,
|
|
{
|
|
"tooltip": "The scheduler controls how noise is gradually removed to form the image."
|
|
},
|
|
),
|
|
"positive": (
|
|
"CONDITIONING",
|
|
{
|
|
"tooltip": "The conditioning describing the attributes you want to include in the image."
|
|
},
|
|
),
|
|
"negative": (
|
|
"CONDITIONING",
|
|
{
|
|
"tooltip": "The conditioning describing the attributes you want to exclude from the image."
|
|
},
|
|
),
|
|
"latent_image": ("LATENT", {"tooltip": "The latent image to denoise."}),
|
|
"denoise": (
|
|
"FLOAT",
|
|
{
|
|
"default": 1.0,
|
|
"min": 0.0,
|
|
"max": 1.0,
|
|
"step": 0.01,
|
|
"tooltip": "The amount of denoising applied, lower values will maintain the structure of the initial image allowing for image to image sampling.",
|
|
},
|
|
),
|
|
"activation_casting": (
|
|
["bf16", "fp16"],
|
|
{
|
|
"default": "bf16",
|
|
"tooltip": "cast model activation to bf16 or fp16, always use bf16 unless your card does not supports it",
|
|
},
|
|
),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("LATENT",)
|
|
OUTPUT_TOOLTIPS = ("The denoised latent.",)
|
|
FUNCTION = "sample"
|
|
CATEGORY = "advanced/FluxMod"
|
|
DESCRIPTION = "Uses the provided model, positive and negative conditioning to denoise the latent image."
|
|
|
|
def sample(
|
|
self,
|
|
model,
|
|
seed,
|
|
steps,
|
|
cfg,
|
|
sampler_name,
|
|
scheduler,
|
|
positive,
|
|
negative,
|
|
latent_image,
|
|
denoise=1.0,
|
|
activation_casting="bf16",
|
|
):
|
|
if using_scaled_fp8(model):
|
|
return common_ksampler(
|
|
model,
|
|
seed,
|
|
steps,
|
|
cfg,
|
|
sampler_name,
|
|
scheduler,
|
|
positive,
|
|
negative,
|
|
latent_image,
|
|
denoise=denoise,
|
|
)
|
|
dtypes = {"bf16": torch.bfloat16, "fp16": torch.float16}
|
|
with torch.autocast(device_type="cuda", dtype=dtypes[activation_casting]):
|
|
return common_ksampler(
|
|
model,
|
|
seed,
|
|
steps,
|
|
cfg,
|
|
sampler_name,
|
|
scheduler,
|
|
positive,
|
|
negative,
|
|
latent_image,
|
|
denoise=denoise,
|
|
)
|
|
|
|
|
|
class FluxModSamplerWrapperNode:
|
|
RETURN_TYPES = ("SAMPLER",)
|
|
FUNCTION = "go"
|
|
CATEGORY = "advanced/FluxMod"
|
|
DESCRIPTION = "Enables FluxMod in float8 quant_mode to be used with advanced sampling nodes by wrapping another SAMPLER. If you are using multiple sampler wrappers, put this node closest to SamplerCustom/SamplerCustomAdvanced/etc."
|
|
NodeId = "FluxModSamplerWrapperNode"
|
|
NodeName = "FluxMod Sampler Wrapper"
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"sampler": ("SAMPLER",),
|
|
"activation_casting": (
|
|
("bf16", "fp16"),
|
|
{
|
|
"default": "bf16",
|
|
"tooltip": "Cast model activation to bf16 or fp16. Always use bf16 unless your card does not support it.",
|
|
},
|
|
),
|
|
},
|
|
}
|
|
|
|
@classmethod
|
|
def go(cls, *, sampler, activation_casting):
|
|
dtype = torch.bfloat16 if activation_casting == "bf16" else torch.float16
|
|
|
|
def wrapper(model, *args, **kwargs):
|
|
if using_scaled_fp8(model.inner_model.model_patcher):
|
|
return sampler.sampler_function(model, *args, **kwargs)
|
|
with torch.autocast(
|
|
device_type=model_management.get_torch_device().type, dtype=dtype
|
|
):
|
|
return sampler.sampler_function(model, *args, **kwargs)
|
|
|
|
sampler_wrapper = comfy.samplers.KSAMPLER(
|
|
wrapper,
|
|
extra_options=sampler.extra_options,
|
|
inpaint_options=sampler.inpaint_options,
|
|
)
|
|
return (sampler_wrapper,)
|
|
|
|
|
|
node_list = [
|
|
# loaders
|
|
FluxModDiffusionLoader,
|
|
FluxModDiffusionLoaderMini,
|
|
ChromaDiffusionLoader,
|
|
# Sampler
|
|
KSamplerMod,
|
|
# Modifiers
|
|
FluxModSamplerWrapperNode,
|
|
SkipLayerForward,
|
|
ChromaPromptTruncation,
|
|
ChromaStyleModelApply,
|
|
]
|
|
|
|
NODE_CLASS_MAPPINGS = {}
|
|
NODE_DISPLAY_NAME_MAPPINGS = {}
|
|
|
|
for node in node_list:
|
|
NODE_CLASS_MAPPINGS[node.NodeId] = node
|
|
NODE_DISPLAY_NAME_MAPPINGS[node.NodeId] = node.NodeName
|