136 lines
4.2 KiB
Python
136 lines
4.2 KiB
Python
import comfy
|
|
import comfy.model_detection
|
|
import comfy.model_management
|
|
import comfy.model_patcher
|
|
import comfy.utils
|
|
import folder_paths
|
|
import torch
|
|
|
|
try:
|
|
from q8_kernels.functional.ops import hadamard_transform
|
|
from q8_kernels.integration.patch_transformer import (
|
|
patch_comfyui_native_transformer,
|
|
patch_comfyui_transformer,
|
|
)
|
|
|
|
Q8_AVAILABLE = True
|
|
except ImportError:
|
|
Q8_AVAILABLE = False
|
|
|
|
from .nodes_registry import comfy_node
|
|
|
|
|
|
def check_q8_available():
|
|
if not Q8_AVAILABLE:
|
|
raise ImportError(
|
|
"Q8 kernels are not available. To use this feature install the q8_kernels package from here:."
|
|
"https://github.com/Lightricks/LTX-Video-Q8-Kernels"
|
|
)
|
|
|
|
|
|
@comfy_node(name="LTXQ8Patch")
|
|
class LTXVQ8Patch:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"model": ("MODEL",),
|
|
"use_fp8_attention": (
|
|
"BOOLEAN",
|
|
{"default": False, "tooltip": "Use FP8 attention."},
|
|
),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("MODEL",)
|
|
FUNCTION = "patch"
|
|
CATEGORY = "lightricks/LTXV"
|
|
TITLE = "LTXV Q8 Patcher"
|
|
|
|
def patch(self, model, use_fp8_attention):
|
|
check_q8_available()
|
|
m = model.clone()
|
|
diffusion_key = "diffusion_model"
|
|
diffusion_model = m.get_model_object(diffusion_key)
|
|
if diffusion_model.__class__.__name__ == "LTXVTransformer3D":
|
|
transformer_key = "diffusion_model.transformer"
|
|
patcher = patch_comfyui_transformer
|
|
else:
|
|
transformer_key = "diffusion_model"
|
|
patcher = patch_comfyui_native_transformer
|
|
transformer = m.get_model_object(transformer_key)
|
|
patcher(transformer, use_fp8_attention, True)
|
|
m.add_object_patch(transformer_key, transformer)
|
|
return (m,)
|
|
|
|
|
|
def idendity_quant_fn(x, t):
|
|
return x.to(dtype=t)
|
|
|
|
|
|
@comfy_node(name="LTXVQ8LoraModelLoader")
|
|
class LTXVQ8LoraModelLoader:
|
|
def __init__(self):
|
|
self.loaded_lora = None
|
|
|
|
@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",)
|
|
CATEGORY = "lightricks/LTXV"
|
|
FUNCTION = "load_lora_model_only"
|
|
|
|
def load_lora(self, model, lora_name, strength_model):
|
|
quant_fn = hadamard_transform
|
|
is_patched_transformer = getattr(
|
|
model.get_model_object("diffusion_model"), "is_patched", False
|
|
)
|
|
if not is_patched_transformer or not Q8_AVAILABLE:
|
|
raise ValueError(
|
|
"LTXV Q8 Patcher is not applied to the model. Please use LTXQ8Patch node before loading lora or install q8_kernels."
|
|
)
|
|
|
|
if strength_model == 0:
|
|
return model
|
|
|
|
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)
|
|
new_lora = {}
|
|
for k in lora:
|
|
device = lora[k].device
|
|
if lora[k].ndim == 2:
|
|
if "lora_A" in k:
|
|
new_lora[k] = quant_fn(
|
|
lora[k].to(device="cuda", dtype=torch.bfloat16),
|
|
out_type=torch.bfloat16,
|
|
).to(device)
|
|
else:
|
|
new_lora[k] = lora[k]
|
|
self.loaded_lora = (lora_path, new_lora)
|
|
|
|
model_lora, _ = comfy.sd.load_lora_for_models(
|
|
model, None, new_lora, strength_model, 0
|
|
)
|
|
return model_lora
|
|
|
|
def load_lora_model_only(self, model, lora_name, strength_model):
|
|
return (self.load_lora(model, lora_name, strength_model),)
|