68 lines
2.0 KiB
Python
68 lines
2.0 KiB
Python
from .nodes_registry import comfy_node
|
|
|
|
# Internal keys used to store AV merge metadata inside model options
|
|
_AV_META_KEY = "ltxv_av_meta"
|
|
_AV_META_VIDEO_ONLY_KEYS = "video_only_keys"
|
|
_AV_META_AUDIO_ONLY_KEYS = "audio_only_keys"
|
|
_AV_META_SHARED_KEYS = "shared_keys"
|
|
_AV_META_AUDIO_SHARED_VALUES = "audio_shared_values"
|
|
|
|
# Common key constants
|
|
_KEY_SAMPLES = "samples"
|
|
_KEY_TYPE = "type"
|
|
_OPT_TRANSFORMER = "transformer_options"
|
|
_OPT_AUDIO_LENGTH = "audio_length"
|
|
|
|
|
|
def _build_av_meta(video_latent: dict, audio_latent: dict) -> dict:
|
|
"""Compute metadata that enables lossless split after AV concat.
|
|
|
|
Notes:
|
|
- We keep video values for shared keys in the merged latent, so we must
|
|
persist the original audio values per shared key.
|
|
"""
|
|
video_keys = set(video_latent.keys()) - {_KEY_SAMPLES}
|
|
audio_keys = set(audio_latent.keys()) - {_KEY_SAMPLES, _KEY_TYPE}
|
|
video_only_keys = list(video_keys - audio_keys)
|
|
audio_only_keys = list(audio_keys - video_keys)
|
|
shared_keys = list(video_keys & audio_keys)
|
|
audio_shared_values = {k: audio_latent[k] for k in shared_keys}
|
|
return {
|
|
_AV_META_VIDEO_ONLY_KEYS: video_only_keys,
|
|
_AV_META_AUDIO_ONLY_KEYS: audio_only_keys,
|
|
_AV_META_SHARED_KEYS: shared_keys,
|
|
_AV_META_AUDIO_SHARED_VALUES: audio_shared_values,
|
|
}
|
|
|
|
|
|
@comfy_node(name="LTXFloatToInt", description="Float To Int")
|
|
class FloatToInt:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {"required": {"a": ("FLOAT", {"default": 0.0})}}
|
|
|
|
RETURN_TYPES = ("INT",)
|
|
FUNCTION = "op"
|
|
CATEGORY = "math/conversion"
|
|
|
|
def op(self, a: float) -> tuple[int]:
|
|
return (round(a),)
|
|
|
|
|
|
@comfy_node(description="Image to CPU")
|
|
class ImageToCPU:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"image": ("IMAGE",),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE",)
|
|
FUNCTION = "run"
|
|
CATEGORY = "utility"
|
|
|
|
def run(self, image):
|
|
return (image.cpu(),)
|