70 lines
2.4 KiB
Python
70 lines
2.4 KiB
Python
from datetime import datetime
|
|
from pathlib import Path
|
|
|
|
import comfy.utils
|
|
import folder_paths
|
|
import torch
|
|
from comfy_api.latest import io, ui
|
|
|
|
from .nodes_registry import comfy_node
|
|
|
|
|
|
@comfy_node(name="LTXVSaveConditioning")
|
|
class LTXVSaveConditioning(io.ComfyNode):
|
|
@classmethod
|
|
def define_schema(cls) -> io.Schema:
|
|
return io.Schema(
|
|
node_id="LTXVSaveConditioning",
|
|
display_name="🅛🅣🅧 LTXV Save Conditioning",
|
|
category="lightricks/LTXV",
|
|
inputs=[
|
|
io.Conditioning.Input("conditioning"),
|
|
io.String.Input("filename", default="conditioning"),
|
|
io.Combo.Input("dtype", options=["bfloat16", "float16"]),
|
|
],
|
|
is_output_node=True,
|
|
)
|
|
|
|
@classmethod
|
|
def execute(cls, conditioning: list, filename: str, dtype: str) -> io.NodeOutput:
|
|
if not conditioning or len(conditioning) == 0:
|
|
raise ValueError("Conditioning is empty")
|
|
|
|
embeddings_folder = Path(folder_paths.get_folder_paths("embeddings")[0])
|
|
embeddings_folder.mkdir(parents=True, exist_ok=True)
|
|
|
|
sanitized_filename = "".join(
|
|
c for c in filename if c.isalnum() or c in ("_", "-", ".")
|
|
)
|
|
if not sanitized_filename:
|
|
sanitized_filename = "conditioning"
|
|
|
|
output_path = embeddings_folder / f"{sanitized_filename}.safetensors"
|
|
|
|
target_dtype = torch.bfloat16 if dtype == "bfloat16" else torch.float16
|
|
|
|
tensors_to_save: dict[str, torch.Tensor] = {}
|
|
|
|
for idx, (cond_tensor, cond_options) in enumerate(conditioning):
|
|
tensor_converted = cond_tensor.to(dtype=target_dtype).contiguous()
|
|
tensors_to_save[f"conditioning_data_{idx}"] = tensor_converted
|
|
|
|
if "attention_mask" in cond_options:
|
|
mask = cond_options["attention_mask"].contiguous()
|
|
tensors_to_save[f"attention_mask_{idx}"] = mask
|
|
|
|
metadata = {
|
|
"num_conditionings": str(len(conditioning)),
|
|
"dtype": dtype,
|
|
"created_at": str(datetime.now()),
|
|
}
|
|
|
|
comfy.utils.save_torch_file(
|
|
tensors_to_save, str(output_path), metadata=metadata
|
|
)
|
|
|
|
file_size_mb = output_path.stat().st_size / (1024 * 1024)
|
|
return io.NodeOutput(
|
|
ui=ui.PreviewText(f"Saved: {output_path.name} ({file_size_mb:.2f} MB)")
|
|
)
|