Compare commits
28
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
ec7fe89955 | ||
|
|
4f45fd6c22 | ||
|
|
2e2ac81759 | ||
|
|
229437c6b6 | ||
|
|
75e1ed26ab | ||
|
|
964be9688f | ||
|
|
d50612a64a | ||
|
|
2acf7af899 | ||
|
|
09b90d0f33 | ||
|
|
937831df0e | ||
|
|
b82708a54f | ||
|
|
d43181d572 | ||
|
|
5984491048 | ||
|
|
3686d832c6 | ||
|
|
6df1f3ff90 | ||
|
|
531512f728 | ||
|
|
c5f4724bc2 | ||
|
|
58a5c12c2d | ||
|
|
5c3d20803d | ||
|
|
82bd963cde | ||
|
|
bb70ecba66 | ||
|
|
49add6dddb | ||
|
|
f82614dc6d | ||
|
|
d153ca3f78 | ||
|
|
d008795913 | ||
|
|
cd5d371518 | ||
|
|
683d46186a | ||
|
|
6790484426 |
@@ -90,20 +90,20 @@ to be bound by this Agreement.
|
||||
subject to the terms and provisions of a different license (the
|
||||
"Commercial Use Agreement"), as will be provided by the Licensor.
|
||||
Commercial Entities interested in such a commercial license are
|
||||
required to contact Licensor. Any commercial use of LTX-2 or
|
||||
Derivatives of LTX-2 by the Commercial Entities not in accordance
|
||||
with this Agreement and/or the Commercial Use Agreement is strictly
|
||||
prohibited and shall be deemed a material breach of this Agreement.
|
||||
Such material breach will be subject, in addition to any license
|
||||
fees owed to Licensor for the period such Commercial Entity used
|
||||
LTX-2 (as will be determined by Licensor), to liquidated damages,
|
||||
which will be paid to Licensor immediately upon demand, in an
|
||||
amount equal to double the amount that would otherwise have been
|
||||
paid by you for the relevant period of time. Such amount reflects
|
||||
a reasonable estimation of the losses and administrative costs
|
||||
incurred due to such breach. You agree and understand that this
|
||||
remedy does not limit the Licensor's right to pursue other remedies
|
||||
available at law or equity.
|
||||
required to [contact Licensor](https://ltx.io/model/licensing).
|
||||
Any commercial use of LTX-2 or Derivatives of LTX-2 by the
|
||||
Commercial Entities not in accordance with this Agreement and/or
|
||||
the Commercial Use Agreement is strictly prohibited and shall be
|
||||
deemed a material breach of this Agreement. Such material breach
|
||||
will be subject, in addition to any license fees owed to Licensor
|
||||
for the period such Commercial Entity used LTX-2 (as will be
|
||||
determined by Licensor), to liquidated damages, which will be paid
|
||||
to Licensor immediately upon demand, in an amount equal to double
|
||||
the amount that would otherwise have been paid by you for the
|
||||
relevant period of time. Such amount reflects a reasonable estimation
|
||||
of the losses and administrative costs incurred due to such breach.
|
||||
You agree and understand that this remedy does not limit the Licensor's
|
||||
right to pursue other remedies available at law or equity.
|
||||
|
||||
3. Distribution and Redistribution. You may host for third parties
|
||||
remote access purposes (e.g. software-as-a-service), reproduce
|
||||
|
||||
@@ -2,7 +2,7 @@
|
||||
|
||||
[](https://github.com/Lightricks/LTX-2)
|
||||
[](https://ltx.io/model)
|
||||
[](https://huggingface.co/Lightricks/LTX-2)
|
||||
[](https://huggingface.co/Lightricks/LTX-2.3)
|
||||
[](https://github.com/Lightricks/LTX-2/tree/main/packages/ltx-trainer)
|
||||
[](https://app.ltx.studio/ltx-2-playground/i2v)
|
||||
[](https://videos.ltx.io/LTX-2/grants/LTX_2_Technical_Report_compressed.pdf)
|
||||
@@ -43,42 +43,115 @@ The nodes will appear in your node menu under the “LTXVideo” category. Requi
|
||||
|
||||
The ComfyUI-LTXVideo installation includes several example workflows.
|
||||
You can see them all at:
|
||||
'''
|
||||
```
|
||||
ComfyUI/custom_nodes/ComfyUI-LTXVideo/example_workflows/
|
||||
'''
|
||||
```
|
||||
|
||||
* [`Text to video full model`](./example_workflows/LTX-2_T2V_Full_wLora.json)
|
||||
* [`Text to video distilled model (Fast)`](./example_workflows/LTX-2_T2V_Distilled_wLora.json)
|
||||
* [`Image to video full model`](./example_workflows/LTX-2_I2V_Full_wLora.json)
|
||||
* [`Image to video distilled model (Fast)`](./example_workflows/LTX-2_I2V_Distilled_wLora.json)
|
||||
* [`Video to video detailer`](./example_workflows/LTX-2_V2V_Detailer.json)
|
||||
* [`IC-LoRA distilled model (depth + human pose + edges)`](./example_workflows/LTX-2_ICLoRA_All_Distilled.json)
|
||||
LTX-2.3 Workflows:
|
||||
|
||||
* [`Text/image to video full/distilled model; single stage`](./example_workflows/2.3/LTX-2.3_T2V_I2V_Single_Stage_Distilled_Full.json)
|
||||
* [`Text/image to video distilled model; two stages (with upsampling)`](./example_workflows/2.3/LTX-2.3_T2V_I2V_Two_Stage_Distilled.json)
|
||||
* [`IC-LoRA distilled model depth + human pose + edges`](./example_workflows/2.3/LTX-2.3_ICLoRA_Union_Control_Distilled.json)
|
||||
* [`IC-LoRA distilled model I2V motion tracking`](./example_workflows/2.3/LTX-2.3_ICLoRA_Motion_Track_Distilled.json)
|
||||
* [`IC-LoRA distilled model HDR`](./example_workflows/2.3/LTX-2.3_ICLoRA_HDR_Distilled.json)
|
||||
* [`IC-LoRA distilled model Lipdub; two stages (with upsampling)`](./example_workflows/2.3/LTX-2.3_ICLoRA_Lipdub_Two_Stage_Distilled.json)
|
||||
* [`IC-LoRA distilled model pixel spatial upscaling`](./example_workflows/2.3/LTX-2.3_ICLoRA_Pixel_Spatial_Upscaler_Distilled.json)
|
||||
* [`Text to audio distilled model; single stage`](./example_workflows/2.3/LTX-2.3_T2A_Single_Stage_Distilled.json)
|
||||
|
||||
Older Workflows (LTX-2.0):
|
||||
|
||||
* [`Text to video full model`](./example_workflows/2.0/LTX-2_T2V_Full_wLora.json)
|
||||
* [`Text to video distilled model (Fast)`](./example_workflows/2.0/LTX-2_T2V_Distilled_wLora.json)
|
||||
* [`Image to video full model`](./example_workflows/2.0/LTX-2_I2V_Full_wLora.json)
|
||||
* [`Image to video distilled model (Fast)`](./example_workflows/2.0/LTX-2_I2V_Distilled_wLora.json)
|
||||
* [`Video to video detailer`](./example_workflows/2.0/LTX-2_V2V_Detailer.json)
|
||||
* [`IC-LoRA distilled model (depth + human pose + edges)`](./example_workflows/2.0/LTX-2_ICLoRA_All_Distilled.json)
|
||||
* [`IC-LoRA distilled model with downscaled reference latents`](./example_workflows/2.0/LTX-2_ICLoRA_All_Distilled_ref0.5.json)
|
||||
|
||||
## Union IC-LoRA Model
|
||||
|
||||
We introduce a new **Union IC-LoRA** model that combines depth and edge (canny) control conditions into a single unified LoRA.
|
||||
|
||||
### Key Features
|
||||
|
||||
- **Unified Control**: A single LoRA that supports multiple control conditions (depth or edges).
|
||||
- **Downsampled Latent Processing**: The union LoRA operates on a downsampled latent size, which reduces memory usage and significantly speeds up inference while maintaining quality.
|
||||
|
||||
### How It Works
|
||||
|
||||
The union LoRA is trained to understand and respond to both control signals (depth maps and edge maps) within a single model. The model learns to:
|
||||
|
||||
1. **Parse multiple conditions**: Identify which control signals are present in the input
|
||||
2. **Process at reduced resolution**: Work on downsampled latents to improve efficiency
|
||||
|
||||
## HDR IC-LoRA
|
||||
|
||||
We provide an **HDR IC-LoRA** that generates linear HDR video encoded in ARRI LogC3, enabling workflows that output high-dynamic-range content suitable for grading and EXR export.
|
||||
|
||||
### Key Features
|
||||
|
||||
- **Linear HDR output**: The LoRA produces frames in LogC3-compressed space; the `LTXVHDRDecodePostprocess` node decodes these back to linear HDR values.
|
||||
- **SDR preview + raw HDR**: The node outputs both a Reinhard-tonemapped SDR preview and the raw linear HDR tensor for downstream use.
|
||||
- **EXR export**: Optionally writes the linear HDR frames as a 16/32-bit EXR image sequence. To enable EXR writing, set `OPENCV_IO_ENABLE_OPENEXR=1` in the environment before starting ComfyUI. The exported EXR sequence is best viewed in [DJV](https://github.com/grizzlypeak3d/DJV) (or [DJV for macOS](https://djv.en.uptodown.com/mac/download)).
|
||||
|
||||
## Lipdub IC-LoRA
|
||||
|
||||
We provide a **Lipdub IC-LoRA** that dubs or rephrases speech in video. Given a source video and a text prompt containing the desired dialogue, it generates new lip movements and audio that match the target text while preserving the speaker's identity.
|
||||
|
||||
### Key Features
|
||||
|
||||
- **Multilingual dubbing**: Translate speech into another language - the model regenerates lips and audio to match.
|
||||
- **Same-language rephrasing**: Change what the speaker says while keeping the original language.
|
||||
- **Two-stage pipeline**: Stage 1 generates the video and audio at base resolution; Stage 2 upscales while freezing the audio.
|
||||
- **Speaker identity preservation**: Reference audio tokens provide speaker context so the generated voice stays consistent.
|
||||
|
||||
## Pixel Spatial Upscaler IC-LoRA
|
||||
|
||||
We provide **Pixel Spatial Upscaler** IC-LoRAs that creatively upscale low-resolution video by synthesizing fine detail rather than simply interpolating pixels. Given a low-resolution reference clip, the model re-renders it at 2× or 4× resolution with generative spatial detail — making it a creative upsampler, not a pixel-accurate refiner.
|
||||
|
||||
### Key Features
|
||||
|
||||
- **2× and 4× variants**: Choose the 2× upscaler for moderate upscaling or the 4× upscaler for larger resolution jumps.
|
||||
- **Generative detail synthesis**: The model synthesizes texture and structure from the reference rather than faithfully preserving every pixel.
|
||||
- **Draft-then-upscale workflow**: Generate at a low base resolution (e.g. ~280p) to lock in composition and motion, then run the upscaler for the final high-resolution output.
|
||||
- **Tunable fidelity**: LoRA strength, guidance, and step count control how closely the output follows the reference — lower values stay closer to the source; higher values allow more creative detail.
|
||||
|
||||
## Text-to-Audio (T2A)
|
||||
|
||||
LTX-2 is a single joint audio/video transformer, but it can generate audio on its own. The `LTXVAudioOnlyModel` node puts the model into audio-only mode for text-to-audio, with no video output.
|
||||
|
||||
### Key Features
|
||||
|
||||
- **Audio-only sampling**: The node sets the model's `run_vx`, `a2v_cross_attn` and `v2a_cross_attn` flags off, so the audio is denoised with no dependence on the video latent and the video stream is skipped. This matches the reference single-stage T2A pipeline's `video=None` behavior.
|
||||
- **Minimal dummy video latent**: The model splits its input positionally into `[video, audio]`, so the sampler still needs a video latent at index 0. Use the `LTXVAudioOnlyEmptyVideoLatent` node (a fixed 64x64 single-frame placeholder, no params to tweak) joined with the audio latent via `LTXVConcatAVLatent`; with `LTXVAudioOnlyModel` active it is never attended to and adds negligible cost.
|
||||
- **Audio decode**: `LTXVAudioVAEDecode` extracts the audio directly from the joint latent, then save it with a standard built-in audio node (for example `Save Audio (FLAC)`).
|
||||
|
||||
## Required Models
|
||||
|
||||
Download the following models:
|
||||
|
||||
**LTX-2 Model Checkpoint** - Choose and download one of the models to `COMFYUI_ROOT_FOLDER/models/checkpoints` folder.
|
||||
* [`ltx-2-19b-dev-fp8.safetensors`](https://huggingface.co/Lightricks/LTX-2/blob/main/ltx-2-19b-dev-fp8.safetensors)
|
||||
* [`ltx-2-19b-distilled-fp8.safetensors`](https://huggingface.co/Lightricks/LTX-2/blob/main/ltx-2-19b-distilled-fp8.safetensors)
|
||||
* [`ltx-2-19b-dev.safetensors`](https://huggingface.co/Lightricks/LTX-2/blob/main/ltx-2-19b-dev.safetensors)
|
||||
* [`ltx-2-19b-distilled.safetensors`](https://huggingface.co/Lightricks/LTX-2/blob/main/ltx-2-19b-distilled.safetensors)
|
||||
**LTX-2.3 Model Checkpoint** - Choose and download one of the models to `COMFYUI_ROOT_FOLDER/models/checkpoints` folder.
|
||||
* [`ltx-2.3-22b-dev.safetensors`](https://huggingface.co/Lightricks/LTX-2.3/blob/main/ltx-2.3-22b-dev.safetensors)
|
||||
* [`ltx-2.3-22b-distilled-1.1.safetensors`](https://huggingface.co/Lightricks/LTX-2.3/blob/main/ltx-2.3-22b-distilled-1.1.safetensors)
|
||||
|
||||
**Spatial Upscaler** - Required for current two-stage pipeline implementations in this repository. Download to `COMFYUI_ROOT_FOLDER/models/latent_upscale_models` folder.
|
||||
* [`ltx-2-spatial-upscaler-x2-1.0.safetensors`](https://huggingface.co/Lightricks/LTX-2/blob/main/ltx-2-spatial-upscaler-x2-1.0.safetensors)
|
||||
* [`ltx-2.3-spatial-upscaler-x2-1.1.safetensors`](https://huggingface.co/Lightricks/LTX-2.3/blob/main/ltx-2.3-spatial-upscaler-x2-1.1.safetensors)
|
||||
* [`ltx-2.3-spatial-upscaler-x1.5-1.0.safetensors`](https://huggingface.co/Lightricks/LTX-2.3/blob/main/ltx-2.3-spatial-upscaler-x1.5-1.0.safetensors)
|
||||
|
||||
**Temporal Upscaler** - Required for current two-stage pipeline implementations in this repository. Download to `COMFYUI_ROOT_FOLDER/models/latent_upscale_models` folder.
|
||||
* [`ltx-2-temporal-upscaler-x2-1.0.safetensors`](https://huggingface.co/Lightricks/LTX-2/blob/main/ltx-2-temporal-upscaler-x2-1.0.safetensors)
|
||||
* [`ltx-2.3-temporal-upscaler-x2-1.0.safetensors`](https://huggingface.co/Lightricks/LTX-2.3/blob/main/ltx-2.3-temporal-upscaler-x2-1.0.safetensors)
|
||||
|
||||
**Distilled LoRA** - Required for current two-stage pipeline implementations in this repository (except DistilledPipeline and ICLoraPipeline). Download to `COMFYUI_ROOT_FOLDER/models/loras` folder.
|
||||
* [`ltx-2-19b-distilled-lora-384.safetensors`](https://huggingface.co/Lightricks/LTX-2/blob/main/ltx-2-19b-distilled-lora-384.safetensors)
|
||||
* [`ltx-2.3-22b-distilled-lora-384-1.1.safetensors`](https://huggingface.co/Lightricks/LTX-2.3/blob/main/ltx-2.3-22b-distilled-lora-384-1.1.safetensors)
|
||||
|
||||
**Gemma Text Encoder** Download all files from the repository to `COMFYUI_ROOT_FOLDER/models/text_encoders/gemma-3-12b-it-qat-q4_0-unquantized`.
|
||||
* [`Gemma 3`](https://huggingface.co/google/gemma-3-12b-it-qat-q4_0-unquantized)
|
||||
|
||||
**LoRAs** Choose and download to `COMFYUI_ROOT_FOLDER/models/loras` folder.
|
||||
* [`ltx-2-19b-ic-lora-canny-control.safetensors`](https://huggingface.co/Lightricks/LTX-2-19b-IC-LoRA-Canny-Control/blob/main/ltx-2-19b-ic-lora-canny-control.safetensors)
|
||||
* [`ltx-2-19b-ic-lora-depth-control.safetensors`](https://huggingface.co/Lightricks/LTX-2-19b-IC-LoRA-Depth-Control/blob/main/ltx-2-19b-ic-lora-depth-control.safetensors)
|
||||
* [`ltx-2.3-22b-ic-lora-union-control-ref0.5.safetensors`](https://huggingface.co/Lightricks/LTX-2.3-22b-IC-LoRA-Union-Control/blob/main/ltx-2.3-22b-ic-lora-union-control-ref0.5.safetensors)
|
||||
* [`ltx-2.3-22b-ic-lora-motion-track-control-ref0.5.safetensors`](https://huggingface.co/Lightricks/LTX-2.3-22b-IC-LoRA-Motion-Track-Control/blob/main/ltx-2.3-22b-ic-lora-motion-track-control-ref0.5.safetensors)
|
||||
* [`ltx-2.3-22b-ic-lora-hdr-0.9.safetensors`](https://huggingface.co/Lightricks/LTX-2.3-22b-IC-LoRA-HDR)
|
||||
* [`ltx-2.3-22b-ic-lora-lipdub-0.9.safetensors`](https://huggingface.co/Lightricks/LTX-2.3-22b-IC-LoRA-LipDub)
|
||||
* [`ltx-2-19b-ic-lora-detailer.safetensors`](https://huggingface.co/Lightricks/LTX-2-19b-IC-LoRA-Detailer/blob/main/ltx-2-19b-ic-lora-detailer.safetensors)
|
||||
* [`ltx-2-19b-ic-lora-pose-control.safetensors`](https://huggingface.co/Lightricks/LTX-2-19b-IC-LoRA-Pose-Control/blob/main/ltx-2-19b-ic-lora-pose-control.safetensors)
|
||||
* [`ltx-2-19b-lora-camera-control-dolly-in.safetensors`](https://huggingface.co/Lightricks/LTX-2-19b-LoRA-Camera-Control-Dolly-In/blob/main/ltx-2-19b-lora-camera-control-dolly-in.safetensors)
|
||||
@@ -88,6 +161,17 @@ Download the following models:
|
||||
* [`ltx-2-19b-lora-camera-control-jib-down.safetensors`](https://huggingface.co/Lightricks/LTX-2-19b-LoRA-Camera-Control-Jib-Down/blob/main/ltx-2-19b-lora-camera-control-jib-down.safetensors)
|
||||
* [`ltx-2-19b-lora-camera-control-jib-up.safetensors`](https://huggingface.co/Lightricks/LTX-2-19b-LoRA-Camera-Control-Jib-Up/blob/main/ltx-2-19b-lora-camera-control-jib-up.safetensors)
|
||||
* [`ltx-2-19b-lora-camera-control-static.safetensors`](https://huggingface.co/Lightricks/LTX-2-19b-LoRA-Camera-Control-Static/blob/main/ltx-2-19b-lora-camera-control-static.safetensors)
|
||||
* [`ltx-2.3-22b-ic-lora-instant-shave-0.9.safetensors`](https://huggingface.co/Lightricks/LTX-2.3-22b-IC-LoRA-Instant-Shave/blob/main/ltx-2.3-22b-ic-lora-instant-shave-0.9.safetensors)
|
||||
* [`ltx-2.3-22b-ic-lora-colorization-0.9.safetensors`](https://huggingface.co/Lightricks/LTX-2.3-22b-IC-LoRA-Colorization/blob/main/ltx-2.3-22b-ic-lora-colorization-0.9.safetensors)
|
||||
* [`ltx-2.3-22b-ic-lora-cross-eyed-0.9.safetensors`](https://huggingface.co/Lightricks/LTX-2.3-22b-IC-LoRA-Cross-Eyed/blob/main/ltx-2.3-22b-ic-lora-cross-eyed-0.9.safetensors)
|
||||
* [`ltx-2.3-22b-ic-lora-day-to-night-0.9.safetensors`](https://huggingface.co/Lightricks/LTX-2.3-22b-IC-LoRA-Day-To-Night/blob/main/ltx-2.3-22b-ic-lora-day-to-night-0.9.safetensors)
|
||||
* [`ltx-2.3-22b-ic-lora-deblur-0.9.safetensors`](https://huggingface.co/Lightricks/LTX-2.3-22b-IC-LoRA-Deblur/blob/main/ltx-2.3-22b-ic-lora-deblur-0.9.safetensors)
|
||||
* [`ltx-2.3-22b-ic-lora-decompression-0.9.safetensors`](https://huggingface.co/Lightricks/LTX-2.3-22b-IC-LoRA-Decompression/blob/main/ltx-2.3-22b-ic-lora-decompression-0.9.safetensors)
|
||||
* [`ltx-2.3-22b-ic-lora-in-outpainting-0.9.safetensors`](https://huggingface.co/Lightricks/LTX-2.3-22b-IC-LoRA-In-Outpainting/blob/main/ltx-2.3-22b-ic-lora-in-outpainting-0.9.safetensors)
|
||||
* [`ltx-2.3-22b-ic-lora-water-simulation-0.9.safetensors`](https://huggingface.co/Lightricks/LTX-2.3-22b-IC-LoRA-Water-Simulation/blob/main/ltx-2.3-22b-ic-lora-water-simulation-0.9.safetensors)
|
||||
* [`ltx-2.3-22b-ic-lora-ingredients-0.9.safetensors`](https://huggingface.co/Lightricks/LTX-2.3-22b-IC-LoRA-Ingredients/blob/main/ltx-2.3-22b-ic-lora-ingredients-0.9.safetensors)
|
||||
* [`ltx-2.3-22b-ic-lora-pixel-spatial-upscaler-x4-0.9.safetensors`](https://huggingface.co/Lightricks/LTX-2.3-22b-IC-LoRA-Pixel-Spatial-Upscaler/blob/main/ltx-2.3-22b-ic-lora-pixel-spatial-upscaler-x4-0.9.safetensors)
|
||||
* [`ltx-2.3-22b-ic-lora-pixel-spatial-upscaler-x2-0.9.safetensors`](https://huggingface.co/Lightricks/LTX-2.3-22b-IC-LoRA-Pixel-Spatial-Upscaler/blob/main/ltx-2.3-22b-ic-lora-pixel-spatial-upscaler-x2-0.9.safetensors)
|
||||
|
||||
|
||||
## Advanced Techniques
|
||||
|
||||
+38
-2
@@ -1,3 +1,6 @@
|
||||
from .audio_only import LTXVAudioOnlyEmptyVideoLatent, LTXVAudioOnlyModel
|
||||
from .conditioning_loader import LTXVLoadConditioning
|
||||
from .conditioning_saver import LTXVSaveConditioning
|
||||
from .decoder_noise import DecoderNoise
|
||||
from .dynamic_conditioning import DynamicConditioning
|
||||
from .easy_samplers import (
|
||||
@@ -5,10 +8,19 @@ from .easy_samplers import (
|
||||
LTXVBaseSampler,
|
||||
LTXVExtendSampler,
|
||||
LTXVInContextSampler,
|
||||
LTXVNormalizingSampler,
|
||||
)
|
||||
from .gemma_api_conditioning import GemmaAPITextEncode
|
||||
from .gemma_encoder import LTXVGemmaCLIPModelLoader, LTXVGemmaEnhancePrompt
|
||||
from .guide import LTXVAddGuideAdvanced
|
||||
from .guide import LTXVAddGuideAdvanced, LTXVAddGuideAdvancedAttention
|
||||
from .guiders import GuiderParametersNode, MultimodalGuiderNode
|
||||
from .hdr import LTXVHDRDecodePostprocess
|
||||
from .iclora import (
|
||||
LTXAddVideoICLoRAGuide,
|
||||
LTXAddVideoICLoRAGuideAdvanced,
|
||||
LTXICLoRALoaderModelOnly,
|
||||
LTXVSetAudioRefTokens,
|
||||
)
|
||||
from .latent_norm import (
|
||||
LTXVAdainLatent,
|
||||
LTXVPerStepAdainPatcher,
|
||||
@@ -34,7 +46,9 @@ from .nodes_registry import (
|
||||
)
|
||||
from .nodes_registry import NODES_DISPLAY_NAME_PREFIX, camel_case_to_spaces
|
||||
from .prompt_enhancer_nodes import LTXVPromptEnhancer, LTXVPromptEnhancerLoader
|
||||
from .pyramid_blending import LTXVLaplacianPyramidBlend
|
||||
from .q8_nodes import LTXVQ8LoraModelLoader, LTXVQ8Patch
|
||||
from .sparse_tracks import LTXVDrawTracks, LTXVSparseTrackEditor
|
||||
from .stg import (
|
||||
LTXVApplySTG,
|
||||
STGAdvancedPresetsNode,
|
||||
@@ -45,14 +59,16 @@ from .tiled_sampler import LTXVTiledSampler
|
||||
from .tiled_vae_decode import LTXVTiledVAEDecode
|
||||
from .tricks import NODE_CLASS_MAPPINGS as TRICKS_NODE_CLASS_MAPPINGS
|
||||
from .tricks import NODE_DISPLAY_NAME_MAPPINGS as TRICKS_NODE_DISPLAY_NAME_MAPPINGS
|
||||
from .utiltily_nodes import ImageToCPU
|
||||
from .utiltily_nodes import FloatToInt, ImageToCPU
|
||||
from .vae_patcher import LTXVPatcherVAE
|
||||
from .vanish_nodes import LTXVDilateVideoMask, LTXVInpaintPreprocess
|
||||
|
||||
# Static node mappings, required for ComfyUI-Manager mapping to work
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
"Set VAE Decoder Noise": DecoderNoise,
|
||||
"LTXVLinearOverlapLatentTransition": LinearOverlapLatentTransition,
|
||||
"LTXVAddGuideAdvanced": LTXVAddGuideAdvanced,
|
||||
"LTXVAddGuideAdvancedAttention": LTXVAddGuideAdvancedAttention,
|
||||
"LTXVAddLatentGuide": LTXVAddLatentGuide,
|
||||
"LTXVAdainLatent": LTXVAdainLatent,
|
||||
"LTXVImgToVideoConditionOnly": LTXVImgToVideoConditionOnly,
|
||||
@@ -61,6 +77,7 @@ NODE_CLASS_MAPPINGS = {
|
||||
"LTXVBaseSampler": LTXVBaseSampler,
|
||||
"LTXVInContextSampler": LTXVInContextSampler,
|
||||
"LTXVExtendSampler": LTXVExtendSampler,
|
||||
"LTXVNormalizingSampler": LTXVNormalizingSampler,
|
||||
"LTXVPreprocessMasks": LTXVPreprocessMasks,
|
||||
"LTXVPatcherVAE": LTXVPatcherVAE,
|
||||
"LTXVPromptEnhancer": LTXVPromptEnhancer,
|
||||
@@ -79,14 +96,30 @@ NODE_CLASS_MAPPINGS = {
|
||||
"STGGuiderNode": STGGuiderNode,
|
||||
"LTXVMultiPromptProvider": MultiPromptProvider,
|
||||
"ImageToCPU": ImageToCPU,
|
||||
"LTXFloatToInt": FloatToInt,
|
||||
"LTXVStatNormLatent": LTXVStatNormLatent,
|
||||
"LTXVPerStepStatNormPatcher": LTXVPerStepStatNormPatcher,
|
||||
"LTXVGemmaCLIPModelLoader": LTXVGemmaCLIPModelLoader,
|
||||
"LTXVGemmaEnhancePrompt": LTXVGemmaEnhancePrompt,
|
||||
"GemmaAPITextEncode": GemmaAPITextEncode,
|
||||
"DynamicConditioning": DynamicConditioning,
|
||||
"LowVRAMCheckpointLoader": LowVRAMCheckpointLoader,
|
||||
"LowVRAMAudioVAELoader": LowVRAMAudioVAELoader,
|
||||
"LowVRAMLatentUpscaleModelLoader": LowVRAMLatentUpscaleModelLoader,
|
||||
"LTXAddVideoICLoRAGuide": LTXAddVideoICLoRAGuide,
|
||||
"LTXAddVideoICLoRAGuideAdvanced": LTXAddVideoICLoRAGuideAdvanced,
|
||||
"LTXICLoRALoaderModelOnly": LTXICLoRALoaderModelOnly,
|
||||
"LTXVSetAudioRefTokens": LTXVSetAudioRefTokens,
|
||||
"LTXVLoadConditioning": LTXVLoadConditioning,
|
||||
"LTXVSaveConditioning": LTXVSaveConditioning,
|
||||
"LTXVDrawTracks": LTXVDrawTracks,
|
||||
"LTXVSparseTrackEditor": LTXVSparseTrackEditor,
|
||||
"LTXVDilateVideoMask": LTXVDilateVideoMask,
|
||||
"LTXVInpaintPreprocess": LTXVInpaintPreprocess,
|
||||
"LTXVLaplacianPyramidBlend": LTXVLaplacianPyramidBlend,
|
||||
"LTXVHDRDecodePostprocess": LTXVHDRDecodePostprocess,
|
||||
"LTXVAudioOnlyModel": LTXVAudioOnlyModel,
|
||||
"LTXVAudioOnlyEmptyVideoLatent": LTXVAudioOnlyEmptyVideoLatent,
|
||||
}
|
||||
|
||||
# Consistent display names between static and dynamic node mappings in nodes_registry.py,
|
||||
@@ -104,8 +137,11 @@ NODE_DISPLAY_NAME_MAPPINGS.update(TRICKS_NODE_DISPLAY_NAME_MAPPINGS)
|
||||
NODE_CLASS_MAPPINGS.update(RUNTIME_NODE_CLASS_MAPPINGS)
|
||||
NODE_DISPLAY_NAME_MAPPINGS.update(RUNTIME_NODE_DISPLAY_NAME_MAPPINGS)
|
||||
|
||||
WEB_DIRECTORY = "./web"
|
||||
|
||||
# Export so that ComfyUI can pick them up.
|
||||
__all__ = [
|
||||
"NODE_CLASS_MAPPINGS",
|
||||
"NODE_DISPLAY_NAME_MAPPINGS",
|
||||
"WEB_DIRECTORY",
|
||||
]
|
||||
|
||||
@@ -0,0 +1,89 @@
|
||||
"""Audio-only (text-to-audio) mode for the LTX-2 audio/video model.
|
||||
|
||||
The LTX-2 transformer in ComfyUI core (``comfy/ldm/lightricks/av_model.py``)
|
||||
already supports an audio-only path: every block reads ``run_vx``,
|
||||
``a2v_cross_attn`` and ``v2a_cross_attn`` from ``transformer_options`` and skips
|
||||
the video self-attention/FFN and the audio<->video cross-attention when they are
|
||||
``False``. These flags default to ``True``, so the model normally runs jointly.
|
||||
|
||||
``LTXVAudioOnlyModel`` flips those flags, which is the equivalent of the
|
||||
reference pipeline's ``video=None``: the audio is denoised with no dependence on
|
||||
the video latent, and the video stream is skipped for speed.
|
||||
|
||||
Note: the model splits its input positionally into ``[video, audio]``, so the
|
||||
sampler still needs a (tiny) dummy video latent at index 0. Use
|
||||
``LTXVAudioOnlyEmptyVideoLatent`` (a fixed 64x64 single-frame placeholder) joined
|
||||
with the audio latent via ``LTXVConcatAVLatent``. With ``LTXVAudioOnlyModel``
|
||||
active that dummy video is never attended to.
|
||||
"""
|
||||
|
||||
import comfy.model_management
|
||||
import torch
|
||||
from comfy.model_patcher import ModelPatcher
|
||||
|
||||
from .nodes_registry import comfy_node
|
||||
|
||||
|
||||
@comfy_node(name="LTXVAudioOnlyModel")
|
||||
class LTXVAudioOnlyModel:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"model": (
|
||||
"MODEL",
|
||||
{
|
||||
"tooltip": "The LTX-2 audio/video model to run in audio-only mode."
|
||||
},
|
||||
),
|
||||
},
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("MODEL",)
|
||||
RETURN_NAMES = ("model",)
|
||||
FUNCTION = "patch"
|
||||
CATEGORY = "Lightricks/audio"
|
||||
DESCRIPTION = (
|
||||
"Run the LTX-2 audio/video model in audio-only mode for text-to-audio. "
|
||||
"Skips the video stream and disables the audio<->video cross-attention, so "
|
||||
"the generated audio does not depend on the video latent. The sampler still "
|
||||
"needs a minimal dummy video latent (EmptyLTXVLatentVideo at 64x64 length 1, "
|
||||
"joined with LTXVConcatAVLatent); it is never attended to. Remove or bypass "
|
||||
"this node to restore normal joint audio+video generation."
|
||||
)
|
||||
|
||||
def patch(self, model: ModelPatcher):
|
||||
new_model = model.clone()
|
||||
transformer_options = new_model.model_options["transformer_options"]
|
||||
transformer_options["run_vx"] = False
|
||||
transformer_options["a2v_cross_attn"] = False
|
||||
transformer_options["v2a_cross_attn"] = False
|
||||
return (new_model,)
|
||||
|
||||
|
||||
@comfy_node(name="LTXVAudioOnlyEmptyVideoLatent")
|
||||
class LTXVAudioOnlyEmptyVideoLatent:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {"required": {}}
|
||||
|
||||
RETURN_TYPES = ("LATENT",)
|
||||
RETURN_NAMES = ("latent",)
|
||||
FUNCTION = "generate"
|
||||
CATEGORY = "Lightricks/audio"
|
||||
DESCRIPTION = (
|
||||
"Fixed minimal placeholder video latent for audio-only (text-to-audio) "
|
||||
"generation - a single 64x64 frame. LTX-2 needs a video latent at input "
|
||||
"index 0, but with LTXV Audio Only Model it is never attended to, so the "
|
||||
"dimensions are fixed (and hidden) to avoid confusion. Join it with the "
|
||||
"audio latent via LTXVConcatAVLatent."
|
||||
)
|
||||
|
||||
def generate(self):
|
||||
# LTX-2 video latent shape: (batch, 128 channels, frames, H/32, W/32).
|
||||
# Fixed to batch 1, a single frame, 64x64 -> (1, 128, 1, 2, 2).
|
||||
latent = torch.zeros(
|
||||
[1, 128, 1, 2, 2],
|
||||
device=comfy.model_management.intermediate_device(),
|
||||
)
|
||||
return ({"samples": latent},)
|
||||
@@ -0,0 +1,82 @@
|
||||
import hashlib
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
import folder_paths
|
||||
import safetensors
|
||||
import torch
|
||||
from comfy_api.latest import io
|
||||
|
||||
from .nodes_registry import comfy_node
|
||||
|
||||
|
||||
@comfy_node(name="LTXVLoadConditioning")
|
||||
class LTXVLoadConditioning(io.ComfyNode):
|
||||
@classmethod
|
||||
def define_schema(cls) -> io.Schema:
|
||||
files = folder_paths.get_filename_list("embeddings")
|
||||
if not files:
|
||||
files = [""]
|
||||
return io.Schema(
|
||||
node_id="LTXVLoadConditioning",
|
||||
display_name="🅛🅣🅧 LTXV Load Conditioning",
|
||||
category="lightricks/LTXV",
|
||||
inputs=[
|
||||
io.Combo.Input("file_name", options=sorted(files)),
|
||||
io.Combo.Input("device", options=["cpu", "gpu"]),
|
||||
],
|
||||
outputs=[
|
||||
io.Conditioning.Output(),
|
||||
],
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def execute(cls, file_name: str, device: str) -> io.NodeOutput:
|
||||
file_path = folder_paths.get_full_path("embeddings", file_name)
|
||||
if not Path(file_path).exists():
|
||||
raise FileNotFoundError(f"Conditioning file not found: {file_path}")
|
||||
|
||||
target_device = "cpu"
|
||||
if device == "gpu":
|
||||
target_device = "cuda" if torch.cuda.is_available() else "cpu"
|
||||
|
||||
conditioning: list[list[Any]] = []
|
||||
|
||||
with safetensors.safe_open(
|
||||
file_path, framework="pt", device=target_device
|
||||
) as f:
|
||||
tensor_keys = [k for k in f.keys() if k.startswith("conditioning_data_")]
|
||||
|
||||
for tensor_key in sorted(tensor_keys):
|
||||
idx = tensor_key.replace("conditioning_data_", "")
|
||||
tensor = f.get_tensor(tensor_key)
|
||||
|
||||
options: dict[str, Any] = {}
|
||||
mask_key = f"attention_mask_{idx}"
|
||||
if mask_key in f.keys():
|
||||
options["attention_mask"] = f.get_tensor(mask_key)
|
||||
|
||||
conditioning.append([tensor, options])
|
||||
|
||||
if not conditioning:
|
||||
raise ValueError(f"No conditioning data found in file: {file_name}")
|
||||
|
||||
return io.NodeOutput(conditioning)
|
||||
|
||||
@classmethod
|
||||
def fingerprint_inputs(cls, file_name: str, device: str) -> str:
|
||||
file_path = folder_paths.get_full_path("embeddings", file_name)
|
||||
with open(file_path, "rb") as f:
|
||||
return hashlib.sha256(f.read()).hexdigest()
|
||||
|
||||
@classmethod
|
||||
def validate_inputs(cls, file_name: str, device: str) -> bool | str:
|
||||
if not file_name:
|
||||
return "No files found. Please save a conditioning first."
|
||||
try:
|
||||
file_path = folder_paths.get_full_path("embeddings", file_name)
|
||||
if not Path(file_path).exists():
|
||||
return f"File not found: {file_name}"
|
||||
except Exception:
|
||||
return f"Invalid file: {file_name}"
|
||||
return True
|
||||
@@ -0,0 +1,69 @@
|
||||
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)")
|
||||
)
|
||||
@@ -4,6 +4,9 @@ import comfy
|
||||
import comfy_extras
|
||||
import nodes
|
||||
import torch
|
||||
from comfy.ldm.lightricks.av_model import LTXAVModel
|
||||
from comfy.nested_tensor import NestedTensor
|
||||
from comfy_api.latest import io
|
||||
from comfy_extras.nodes_custom_sampler import SamplerCustomAdvanced, SplitSigmas
|
||||
from comfy_extras.nodes_lt import EmptyLTXVLatentVideo, LTXVAddGuide, LTXVCropGuides
|
||||
|
||||
@@ -907,3 +910,142 @@ class LinearOverlapLatentTransition:
|
||||
"batch_index": combined_batch_index,
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
@comfy_node(description="LTXV Normalizing Sampler")
|
||||
class LTXVNormalizingSampler(io.ComfyNode):
|
||||
@classmethod
|
||||
def define_schema(cls):
|
||||
return io.Schema(
|
||||
node_id="LTXVNormalizingSampler",
|
||||
category="utility",
|
||||
inputs=[
|
||||
io.Noise.Input("noise"),
|
||||
io.Guider.Input("guider"),
|
||||
io.Sampler.Input("sampler"),
|
||||
io.Sigmas.Input("sigmas"),
|
||||
io.Latent.Input("latent_image"),
|
||||
io.String.Input(
|
||||
"video_normalization_factors", default="1,1,1,1,1,1,1,1"
|
||||
),
|
||||
io.String.Input(
|
||||
"audio_normalization_factors", default="1,1,0.25,1,1,0.25,1,1"
|
||||
),
|
||||
],
|
||||
outputs=[
|
||||
io.Latent.Output(display_name="denoised_output"),
|
||||
],
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def execute(
|
||||
cls,
|
||||
noise,
|
||||
guider,
|
||||
sampler,
|
||||
sigmas,
|
||||
latent_image,
|
||||
video_normalization_factors,
|
||||
audio_normalization_factors,
|
||||
) -> io.NodeOutput:
|
||||
|
||||
if (
|
||||
guider.model_patcher.model.diffusion_model.__class__.__name__
|
||||
!= "LTXAVModel"
|
||||
):
|
||||
raise ValueError()
|
||||
|
||||
ltxav: LTXAVModel = guider.model_patcher.model.diffusion_model
|
||||
|
||||
video_normalization_factors = video_normalization_factors.split(",")
|
||||
audio_normalization_factors = audio_normalization_factors.split(",")
|
||||
video_normalization_factors = [
|
||||
float(factor) for factor in video_normalization_factors
|
||||
]
|
||||
audio_normalization_factors = [
|
||||
float(factor) for factor in audio_normalization_factors
|
||||
]
|
||||
|
||||
# Extend normalization factors to match the length of sigmas
|
||||
sigmas_len = len(sigmas) - 1
|
||||
if len(video_normalization_factors) < sigmas_len:
|
||||
if len(video_normalization_factors) > 0:
|
||||
video_normalization_factors.extend(
|
||||
[video_normalization_factors[-1]]
|
||||
* (sigmas_len - len(video_normalization_factors))
|
||||
)
|
||||
if len(audio_normalization_factors) < sigmas_len:
|
||||
if len(audio_normalization_factors) > 0:
|
||||
audio_normalization_factors.extend(
|
||||
[audio_normalization_factors[-1]]
|
||||
* (sigmas_len - len(audio_normalization_factors))
|
||||
)
|
||||
|
||||
# Calculate indices where both normalization factors are not 1.0
|
||||
sampling_split_indices = [
|
||||
i + 1
|
||||
for i, (v, a) in enumerate(
|
||||
zip(video_normalization_factors, audio_normalization_factors)
|
||||
)
|
||||
if v != 1.0 or a != 1.0
|
||||
]
|
||||
print("Sampling split indices: %s" % sampling_split_indices, flush=True)
|
||||
|
||||
# Split sigmas according to sampling_split_indices
|
||||
def split_by_indices(arr, indices):
|
||||
"""
|
||||
Splits arr into chunks according to indices (split points).
|
||||
Indices are treated as starting a new chunk at each index in the list.
|
||||
"""
|
||||
if not indices:
|
||||
return [arr]
|
||||
split_points = sorted(set(indices))
|
||||
chunks = []
|
||||
prev = 0
|
||||
for idx in split_points:
|
||||
if prev < idx:
|
||||
chunks.append(arr[prev : idx + 1])
|
||||
prev = idx
|
||||
if prev < len(arr):
|
||||
chunks.append(arr[prev:])
|
||||
return chunks
|
||||
|
||||
sigmas_chunks = split_by_indices(sigmas, sampling_split_indices)
|
||||
print("Sigmas chunks: %s" % sigmas_chunks, flush=True)
|
||||
|
||||
i = 0
|
||||
for sigmas_chunk in sigmas_chunks:
|
||||
i += len(sigmas_chunk) - 1
|
||||
print("Sampling with sigmas %s" % (sigmas_chunk), flush=True)
|
||||
(_, latent_image) = SamplerCustomAdvanced().execute(
|
||||
noise=noise,
|
||||
guider=guider,
|
||||
sampler=sampler,
|
||||
sigmas=sigmas_chunk,
|
||||
latent_image=latent_image,
|
||||
)
|
||||
video_samples, audio_samples = ltxav.separate_audio_and_video_latents(
|
||||
latent_image["samples"].tensors,
|
||||
None,
|
||||
)
|
||||
if i - 1 < len(video_normalization_factors) and i - 1 < len(
|
||||
audio_normalization_factors
|
||||
):
|
||||
video_samples = video_samples * video_normalization_factors[i - 1]
|
||||
audio_samples = audio_samples * audio_normalization_factors[i - 1]
|
||||
latent_image["samples"] = NestedTensor(
|
||||
ltxav.recombine_audio_and_video_latents(
|
||||
video_samples, audio_samples
|
||||
)
|
||||
)
|
||||
print(
|
||||
"After %d steps, the latent image was normalized by %f and %f"
|
||||
% (
|
||||
i,
|
||||
video_normalization_factors[i - 1],
|
||||
audio_normalization_factors[i - 1],
|
||||
),
|
||||
flush=True,
|
||||
)
|
||||
|
||||
return io.NodeOutput(latent_image)
|
||||
|
||||
@@ -2,10 +2,13 @@ import math
|
||||
from typing import Optional
|
||||
|
||||
import comfy.ldm.common_dit
|
||||
import comfy.ops
|
||||
import torch
|
||||
from comfy.ldm.lightricks.model import (
|
||||
CrossAttention,
|
||||
FeedForward,
|
||||
LTXFrequenciesPrecision,
|
||||
LTXRopeType,
|
||||
generate_freq_grid_np,
|
||||
interleaved_freqs_cis,
|
||||
split_freqs_cis,
|
||||
@@ -52,6 +55,7 @@ class BasicTransformerBlock1D(nn.Module):
|
||||
d_head,
|
||||
context_dim=None,
|
||||
attn_precision=None,
|
||||
apply_gated_attention=False,
|
||||
dtype=None,
|
||||
device=None,
|
||||
operations=None,
|
||||
@@ -65,6 +69,7 @@ class BasicTransformerBlock1D(nn.Module):
|
||||
heads=n_heads,
|
||||
dim_head=d_head,
|
||||
context_dim=None,
|
||||
apply_gated_attention=apply_gated_attention,
|
||||
dtype=dtype,
|
||||
device=device,
|
||||
operations=operations,
|
||||
@@ -123,6 +128,7 @@ class Embeddings1DConnector(nn.Module):
|
||||
positional_embedding_max_pos=[1],
|
||||
causal_temporal_positioning=False,
|
||||
num_learnable_registers: Optional[int] = 128,
|
||||
apply_gated_attention=False,
|
||||
dtype=None,
|
||||
device=None,
|
||||
operations=None,
|
||||
@@ -147,6 +153,7 @@ class Embeddings1DConnector(nn.Module):
|
||||
num_attention_heads,
|
||||
attention_head_dim,
|
||||
context_dim=cross_attention_dim,
|
||||
apply_gated_attention=apply_gated_attention,
|
||||
dtype=dtype,
|
||||
device=device,
|
||||
operations=operations,
|
||||
@@ -330,3 +337,115 @@ class Embeddings1DConnector(nn.Module):
|
||||
hidden_states = comfy.ldm.common_dit.rms_norm(hidden_states)
|
||||
|
||||
return hidden_states, attention_mask
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Loading helpers
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
_PREFIX_BASE = "model.diffusion_model."
|
||||
|
||||
|
||||
def load_embeddings_connector(
|
||||
sd,
|
||||
connector_prefix,
|
||||
connector_config,
|
||||
dtype=torch.bfloat16,
|
||||
rope_type=LTXRopeType.INTERLEAVED,
|
||||
frequencies_precision=LTXFrequenciesPrecision.FLOAT32,
|
||||
pe_max_pos=None,
|
||||
):
|
||||
sd_connector = {
|
||||
k[len(connector_prefix) :]: v
|
||||
for k, v in sd.items()
|
||||
if k.startswith(connector_prefix)
|
||||
}
|
||||
|
||||
if len(sd_connector) == 0:
|
||||
return None
|
||||
|
||||
operations = comfy.ops.pick_operations(dtype, dtype, disable_fast_fp8=True)
|
||||
connector = Embeddings1DConnector(
|
||||
num_attention_heads=connector_config["num_attention_heads"],
|
||||
attention_head_dim=connector_config["attention_head_dim"],
|
||||
num_layers=connector_config["num_layers"],
|
||||
apply_gated_attention=connector_config.get("apply_gated_attention", False),
|
||||
dtype=dtype,
|
||||
operations=operations,
|
||||
positional_embedding_max_pos=pe_max_pos if pe_max_pos is not None else [1],
|
||||
split_rope=rope_type == LTXRopeType.SPLIT,
|
||||
double_precision_rope=frequencies_precision == LTXFrequenciesPrecision.FLOAT64,
|
||||
)
|
||||
connector.load_state_dict(sd_connector)
|
||||
return connector
|
||||
|
||||
|
||||
def load_video_embeddings_connector(sd, transformer_config, dtype=torch.bfloat16):
|
||||
rope_type = LTXRopeType.from_dict(transformer_config)
|
||||
frequencies_precision = LTXFrequenciesPrecision.from_dict(transformer_config)
|
||||
pe_max_pos = transformer_config.get("connector_positional_embedding_max_pos", [1])
|
||||
|
||||
video_only_connector_prefix = f"{_PREFIX_BASE}embeddings_connector."
|
||||
av_connector_prefix = f"{_PREFIX_BASE}video_embeddings_connector."
|
||||
prefix = (
|
||||
av_connector_prefix
|
||||
if f"{_PREFIX_BASE}audio_adaln_single.linear.weight" in sd
|
||||
else video_only_connector_prefix
|
||||
)
|
||||
|
||||
connector_config = {
|
||||
"num_attention_heads": transformer_config.get(
|
||||
"connector_num_attention_heads", 30
|
||||
),
|
||||
"attention_head_dim": transformer_config.get(
|
||||
"connector_attention_head_dim", 128
|
||||
),
|
||||
"num_layers": transformer_config.get("connector_num_layers", 2),
|
||||
"apply_gated_attention": transformer_config.get(
|
||||
"connector_apply_gated_attention", False
|
||||
),
|
||||
}
|
||||
|
||||
return load_embeddings_connector(
|
||||
sd,
|
||||
prefix,
|
||||
connector_config,
|
||||
dtype,
|
||||
rope_type,
|
||||
frequencies_precision,
|
||||
pe_max_pos,
|
||||
)
|
||||
|
||||
|
||||
def load_audio_embeddings_connector(sd, transformer_config, dtype=torch.bfloat16):
|
||||
rope_type = LTXRopeType.from_dict(transformer_config)
|
||||
frequencies_precision = LTXFrequenciesPrecision.from_dict(transformer_config)
|
||||
pe_max_pos = transformer_config.get("connector_positional_embedding_max_pos", [1])
|
||||
|
||||
connector_config = {
|
||||
"num_attention_heads": transformer_config.get(
|
||||
"audio_connector_num_attention_heads",
|
||||
transformer_config.get("connector_num_attention_heads", 30),
|
||||
),
|
||||
"attention_head_dim": transformer_config.get(
|
||||
"audio_connector_attention_head_dim",
|
||||
transformer_config.get("connector_attention_head_dim", 128),
|
||||
),
|
||||
"num_layers": transformer_config.get(
|
||||
"audio_connector_num_layers",
|
||||
transformer_config.get("connector_num_layers", 2),
|
||||
),
|
||||
"apply_gated_attention": transformer_config.get(
|
||||
"connector_apply_gated_attention", False
|
||||
),
|
||||
}
|
||||
|
||||
return load_embeddings_connector(
|
||||
sd,
|
||||
f"{_PREFIX_BASE}audio_embeddings_connector.",
|
||||
connector_config,
|
||||
dtype,
|
||||
rope_type,
|
||||
frequencies_precision,
|
||||
pe_max_pos,
|
||||
)
|
||||
|
||||
+1
-1
@@ -565,7 +565,7 @@
|
||||
"title": "Video Size",
|
||||
"properties": {},
|
||||
"widgets_values": [
|
||||
"Width & height settings must be divisible by 32 + 1. \nFrame count must be divisible by 8 + 1. \n\n\nRunning with invalid parameters **will not cause errors**. Instead, the flow will silently choose the closest valid parameters. "
|
||||
"Width & height settings must be divisible by 64. \nFrame count must be divisible by 8 + 1. \n\n\nRunning with invalid parameters **will not cause errors**. Instead, the flow will silently choose the closest valid parameters. "
|
||||
],
|
||||
"color": "#432",
|
||||
"bgcolor": "#653"
|
||||
+1
-1
@@ -743,7 +743,7 @@
|
||||
"title": "Video Size",
|
||||
"properties": {},
|
||||
"widgets_values": [
|
||||
"Width & height settings must be divisible by 32 + 1. \nFrame count must be divisible by 8 + 1. \n\n\nRunning with invalid parameters **will not cause errors**. Instead, the flow will silently choose the closest valid parameters. "
|
||||
"Width & height settings must be divisible by 64. \nFrame count must be divisible by 8 + 1. \n\n\nRunning with invalid parameters **will not cause errors**. Instead, the flow will silently choose the closest valid parameters. "
|
||||
],
|
||||
"color": "#432",
|
||||
"bgcolor": "#653"
|
||||
+1
-1
@@ -57,7 +57,7 @@
|
||||
"title": "Video Size",
|
||||
"properties": {},
|
||||
"widgets_values": [
|
||||
"Width & height settings must be divisible by 32 + 1. \n\nRunning with invalid parameters **will not cause errors**. Instead, the flow will silently choose the closest valid parameters. "
|
||||
"Width & height settings must be divisible by 64. \n\nRunning with invalid parameters **will not cause errors**. Instead, the flow will silently choose the closest valid parameters. "
|
||||
],
|
||||
"color": "#432",
|
||||
"bgcolor": "#653"
|
||||
File diff suppressed because it is too large
Load Diff
+1
-1
@@ -567,7 +567,7 @@
|
||||
"title": "Video Size",
|
||||
"properties": {},
|
||||
"widgets_values": [
|
||||
"Width & height settings must be divisible by 32 + 1. \nFrame count must be divisible by 8 + 1. \n\n\nRunning with invalid parameters **will not cause errors**. Instead, the flow will silently choose the closest valid parameters. "
|
||||
"Width & height settings must be divisible by 64. \nFrame count must be divisible by 8 + 1. \n\n\nRunning with invalid parameters **will not cause errors**. Instead, the flow will silently choose the closest valid parameters. "
|
||||
],
|
||||
"color": "#432",
|
||||
"bgcolor": "#653"
|
||||
+1
-1
@@ -520,7 +520,7 @@
|
||||
"title": "Video Size",
|
||||
"properties": {},
|
||||
"widgets_values": [
|
||||
"Width & height settings must be divisible by 32 + 1. \nFrame count must be divisible by 8 + 1. \n\n\nRunning with invalid parameters **will not cause errors**. Instead, the flow will silently choose the closest valid parameters. "
|
||||
"Width & height settings must be divisible by 64. \nFrame count must be divisible by 8 + 1. \n\n\nRunning with invalid parameters **will not cause errors**. Instead, the flow will silently choose the closest valid parameters. "
|
||||
],
|
||||
"color": "#432",
|
||||
"bgcolor": "#653"
|
||||
File diff suppressed because it is too large
Load Diff
File diff suppressed because it is too large
Load Diff
File diff suppressed because it is too large
Load Diff
File diff suppressed because it is too large
Load Diff
File diff suppressed because one or more lines are too long
File diff suppressed because it is too large
Load Diff
File diff suppressed because it is too large
Load Diff
File diff suppressed because it is too large
Load Diff
File diff suppressed because it is too large
Load Diff
File diff suppressed because it is too large
Load Diff
File diff suppressed because it is too large
Load Diff
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,3 @@
|
||||
version https://git-lfs.github.com/spec/v1
|
||||
oid sha256:96c2a8f737e1e49583e56b1d0f17e487876b1e424b049434efe8132914eb3e64
|
||||
size 1968386
|
||||
@@ -0,0 +1,3 @@
|
||||
version https://git-lfs.github.com/spec/v1
|
||||
oid sha256:6666a36c4acf4ac1b1f0e85df38551ab51616593b4b93d08502f08223d18f307
|
||||
size 250432
|
||||
@@ -0,0 +1,3 @@
|
||||
version https://git-lfs.github.com/spec/v1
|
||||
oid sha256:a04e0a13b03d31c119e3d988bd840dec04b9d15c737e1e8951dc293add136a2a
|
||||
size 3334888
|
||||
@@ -0,0 +1,3 @@
|
||||
version https://git-lfs.github.com/spec/v1
|
||||
oid sha256:a89b03aa9d353dd28441cacd22258deac24ca6c8ff689471e2a968cadd2d8a26
|
||||
size 35624
|
||||
@@ -0,0 +1,3 @@
|
||||
version https://git-lfs.github.com/spec/v1
|
||||
oid sha256:7d64dea45a90dc0a0cecd0317b84bc84979bb3b2074733d5157c53c792a665db
|
||||
size 3653150
|
||||
@@ -0,0 +1,3 @@
|
||||
version https://git-lfs.github.com/spec/v1
|
||||
oid sha256:fa96b9f0b71a388b30b24b64734779b35562d6ab2b8cd377781a5d5a4a388b10
|
||||
size 5226932
|
||||
@@ -0,0 +1,3 @@
|
||||
version https://git-lfs.github.com/spec/v1
|
||||
oid sha256:0c24a1501f188175ac3a3b42fa43814188a3ee38ae264a6859e796a5e2851274
|
||||
size 237417
|
||||
@@ -0,0 +1,3 @@
|
||||
version https://git-lfs.github.com/spec/v1
|
||||
oid sha256:a5ced045ce8bd87b6b6e9d5e9eedb7ffc89d271509780765134aeb0555b427d3
|
||||
size 2738806
|
||||
@@ -0,0 +1,3 @@
|
||||
version https://git-lfs.github.com/spec/v1
|
||||
oid sha256:166913c9ea536e9b48a292312b9c380957d7836ea0facb984caacf3fd69208bd
|
||||
size 461785
|
||||
@@ -0,0 +1,147 @@
|
||||
"""
|
||||
API-based text encoding that returns CONDITIONING for LTX-2.
|
||||
Replaces the CLIP encoding step entirely using an external API.
|
||||
"""
|
||||
|
||||
import io
|
||||
import logging
|
||||
import pickle
|
||||
|
||||
import folder_paths
|
||||
import requests
|
||||
from safetensors import safe_open
|
||||
|
||||
from .nodes_registry import comfy_node
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
LTXV_API_BASE_URL = "https://api.ltx.video"
|
||||
UPDATE_MESSAGE = (
|
||||
"Note: If this error persists, the node might be outdated. "
|
||||
"Please update ComfyUI-LTXVideo to the latest version."
|
||||
)
|
||||
INVALID_API_KEY_MESSAGE = (
|
||||
"Invalid API key. Please generate a new API key at: https://console.ltx.video/"
|
||||
)
|
||||
MISSING_MODEL_ID_MESSAGE = "Model ID cannot be identified from the provided model file"
|
||||
|
||||
|
||||
def extract_model_id(ckpt_name: str) -> str:
|
||||
model_id_key = "encrypted_wandb_properties"
|
||||
with safe_open(
|
||||
folder_paths.get_full_path_or_raise("checkpoints", ckpt_name),
|
||||
framework="pt",
|
||||
device="cpu",
|
||||
) as f:
|
||||
metadata = f.metadata()
|
||||
if not metadata or model_id_key not in metadata:
|
||||
raise ValueError(MISSING_MODEL_ID_MESSAGE)
|
||||
return metadata[model_id_key]
|
||||
|
||||
|
||||
@comfy_node(name="GemmaAPITextEncode")
|
||||
class GemmaAPITextEncode:
|
||||
"""
|
||||
Encodes text prompts using the LTX Video API, returning CONDITIONING for LTX-2 models.
|
||||
|
||||
This node replaces the local CLIP encoding step by sending the prompt to an external API
|
||||
for processing. It requires an API key and automatically extracts the model ID from the
|
||||
checkpoint file metadata.
|
||||
|
||||
Inputs:
|
||||
- api_key: Authentication key for the LTX Video API
|
||||
- prompt: Text prompt to encode
|
||||
- ckpt_name: Checkpoint file containing model metadata
|
||||
|
||||
Returns:
|
||||
- CONDITIONING: Encoded prompt conditioning ready for LTX-2 video generation
|
||||
"""
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {
|
||||
"api_key": (
|
||||
"STRING",
|
||||
{
|
||||
"default": "",
|
||||
"placeholder": "API_KEY",
|
||||
"multiline": False,
|
||||
"tooltip": "API key for authentication",
|
||||
},
|
||||
),
|
||||
"prompt": (
|
||||
"STRING",
|
||||
{
|
||||
"multiline": True,
|
||||
"default": "",
|
||||
"tooltip": "Text prompt to encode",
|
||||
},
|
||||
),
|
||||
"enhance_prompt": (
|
||||
"BOOLEAN",
|
||||
{
|
||||
"default": True,
|
||||
"tooltip": "When enabled, the prompt is enhanced using Gemma 3 before encoding",
|
||||
},
|
||||
),
|
||||
"ckpt_name": (
|
||||
folder_paths.get_filename_list("checkpoints"),
|
||||
{"tooltip": "The name of the checkpoint (model) to load."},
|
||||
),
|
||||
},
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("CONDITIONING",)
|
||||
RETURN_NAMES = ("conditioning",)
|
||||
FUNCTION = "encode"
|
||||
CATEGORY = "api node/text/Lightricks"
|
||||
|
||||
def encode(
|
||||
self, api_key: str, prompt: str, ckpt_name: str, enhance_prompt: bool = False
|
||||
):
|
||||
if not api_key:
|
||||
raise ValueError("API key is required")
|
||||
|
||||
if not prompt.strip():
|
||||
raise ValueError("Text prompt cannot be empty")
|
||||
|
||||
if not ckpt_name or not ckpt_name.strip():
|
||||
raise ValueError("Model path is required")
|
||||
|
||||
model_id = extract_model_id(ckpt_name)
|
||||
payload = {
|
||||
"prompt": prompt,
|
||||
"model_id": model_id,
|
||||
"enhance_prompt": enhance_prompt,
|
||||
}
|
||||
logger.info(
|
||||
f"Calling API to encode prompt: {prompt[:50]}... with model_id: {model_id[:50]}..."
|
||||
)
|
||||
try:
|
||||
response = requests.post(
|
||||
f"{LTXV_API_BASE_URL}/v1/prompt-embedding",
|
||||
json=payload,
|
||||
headers={
|
||||
"Authorization": f"Bearer {api_key}",
|
||||
"Content-Type": "application/json",
|
||||
},
|
||||
timeout=60,
|
||||
)
|
||||
|
||||
if response.status_code == 401:
|
||||
raise RuntimeError(INVALID_API_KEY_MESSAGE)
|
||||
|
||||
if response.status_code != 200:
|
||||
raise RuntimeError(
|
||||
f"API request failed with status {response.status_code}: {response.text}\n"
|
||||
f"{UPDATE_MESSAGE}"
|
||||
)
|
||||
|
||||
conditioning = pickle.load(io.BytesIO(response.content))
|
||||
logger.info("Successfully received conditioning from API")
|
||||
|
||||
return (conditioning,)
|
||||
|
||||
except Exception as e:
|
||||
raise RuntimeError(f"API request failed: {str(e)}\n {UPDATE_MESSAGE}")
|
||||
@@ -0,0 +1,62 @@
|
||||
{
|
||||
"architectures": [
|
||||
"Gemma3ForConditionalGeneration"
|
||||
],
|
||||
"boi_token_index": 255999,
|
||||
"eoi_token_index": 256000,
|
||||
"eos_token_id": [
|
||||
1,
|
||||
106
|
||||
],
|
||||
"image_token_index": 262144,
|
||||
"initializer_range": 0.02,
|
||||
"mm_tokens_per_image": 256,
|
||||
"model_type": "gemma3",
|
||||
"text_config": {
|
||||
"attention_bias": false,
|
||||
"attention_dropout": 0.0,
|
||||
"attn_logit_softcapping": null,
|
||||
"cache_implementation": "hybrid",
|
||||
"final_logit_softcapping": null,
|
||||
"head_dim": 256,
|
||||
"hidden_activation": "gelu_pytorch_tanh",
|
||||
"hidden_size": 3840,
|
||||
"initializer_range": 0.02,
|
||||
"intermediate_size": 15360,
|
||||
"max_position_embeddings": 131072,
|
||||
"model_type": "gemma3_text",
|
||||
"num_attention_heads": 16,
|
||||
"num_hidden_layers": 48,
|
||||
"num_key_value_heads": 8,
|
||||
"query_pre_attn_scalar": 256,
|
||||
"rms_norm_eps": 1e-06,
|
||||
"rope_local_base_freq": 10000,
|
||||
"rope_scaling": {
|
||||
"factor": 8.0,
|
||||
"rope_type": "linear"
|
||||
},
|
||||
"rope_theta": 1000000,
|
||||
"sliding_window": 1024,
|
||||
"sliding_window_pattern": 6,
|
||||
"torch_dtype": "bfloat16",
|
||||
"use_cache": true,
|
||||
"vocab_size": 262208
|
||||
},
|
||||
"torch_dtype": "bfloat16",
|
||||
"transformers_version": "4.51.0",
|
||||
"vision_config": {
|
||||
"attention_dropout": 0.0,
|
||||
"hidden_act": "gelu_pytorch_tanh",
|
||||
"hidden_size": 1152,
|
||||
"image_size": 896,
|
||||
"intermediate_size": 4304,
|
||||
"layer_norm_eps": 1e-06,
|
||||
"model_type": "siglip_vision_model",
|
||||
"num_attention_heads": 16,
|
||||
"num_channels": 3,
|
||||
"num_hidden_layers": 27,
|
||||
"patch_size": 14,
|
||||
"torch_dtype": "bfloat16",
|
||||
"vision_use_head": false
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,29 @@
|
||||
{
|
||||
"do_convert_rgb": null,
|
||||
"do_normalize": true,
|
||||
"do_pan_and_scan": null,
|
||||
"do_rescale": true,
|
||||
"do_resize": true,
|
||||
"image_mean": [
|
||||
0.5,
|
||||
0.5,
|
||||
0.5
|
||||
],
|
||||
"image_processor_type": "Gemma3ImageProcessor",
|
||||
"image_seq_length": 256,
|
||||
"image_std": [
|
||||
0.5,
|
||||
0.5,
|
||||
0.5
|
||||
],
|
||||
"pan_and_scan_max_num_crops": null,
|
||||
"pan_and_scan_min_crop_size": null,
|
||||
"pan_and_scan_min_ratio_to_activate": null,
|
||||
"processor_class": "Gemma3Processor",
|
||||
"resample": 2,
|
||||
"rescale_factor": 0.00392156862745098,
|
||||
"size": {
|
||||
"height": 896,
|
||||
"width": 896
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,4 @@
|
||||
{
|
||||
"image_seq_length": 256,
|
||||
"processor_class": "Gemma3Processor"
|
||||
}
|
||||
+2379611
File diff suppressed because it is too large
Load Diff
Binary file not shown.
File diff suppressed because it is too large
Load Diff
+328
-371
@@ -1,33 +1,29 @@
|
||||
import json
|
||||
import logging
|
||||
from glob import glob
|
||||
from pathlib import Path
|
||||
from typing import Optional
|
||||
from typing import List, Optional, Tuple
|
||||
|
||||
import comfy.model_management
|
||||
import comfy.ops
|
||||
import comfy.sd
|
||||
import comfy.supported_models_base
|
||||
import folder_paths
|
||||
import safetensors
|
||||
import torch
|
||||
from comfy.ldm.lightricks.model import LTXFrequenciesPrecision, LTXRopeType
|
||||
from comfy.utils import load_torch_file
|
||||
from einops import rearrange
|
||||
from PIL import Image
|
||||
from transformers import (
|
||||
AutoImageProcessor,
|
||||
AutoTokenizer,
|
||||
Gemma3Config,
|
||||
Gemma3ForConditionalGeneration,
|
||||
Gemma3Processor,
|
||||
)
|
||||
from transformers.modeling_rope_utils import ROPE_INIT_FUNCTIONS
|
||||
from transformers.models.auto.modeling_auto import MODEL_MAPPING_NAMES
|
||||
|
||||
from .embeddings_connector import Embeddings1DConnector
|
||||
from .nodes_registry import comfy_node
|
||||
from .text_embeddings_connectors import load_text_embeddings_pipeline
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
PREFIX_BASE = "model.diffusion_model."
|
||||
|
||||
|
||||
def _load_system_prompt(filename: str) -> str:
|
||||
"""Load system prompt from file at module level."""
|
||||
@@ -52,99 +48,6 @@ def tensor_to_pil(tensor: torch.Tensor) -> Image.Image:
|
||||
return Image.fromarray(numpy_image)
|
||||
|
||||
|
||||
def load_video_embeddings_connector(ltxv_path, dtype=torch.bfloat16):
|
||||
sd, metadata = load_torch_file(str(ltxv_path), return_metadata=True)
|
||||
config = json.loads(metadata.get("config", "{}"))
|
||||
transformer_config = config.get("transformer", {})
|
||||
rope_type = LTXRopeType.from_dict(transformer_config)
|
||||
frequencies_precision = LTXFrequenciesPrecision.from_dict(transformer_config)
|
||||
pe_max_pos = transformer_config.get("connector_positional_embedding_max_pos", [1])
|
||||
sd_keys = list(sd.keys())
|
||||
|
||||
video_only_connector_prefix = f"{PREFIX_BASE}embeddings_connector."
|
||||
av_connector_prefix = f"{PREFIX_BASE}video_embeddings_connector."
|
||||
prefix = (
|
||||
av_connector_prefix
|
||||
if f"{PREFIX_BASE}audio_adaln_single.linear.weight" in sd_keys
|
||||
else video_only_connector_prefix
|
||||
)
|
||||
return load_embeddings_connector(
|
||||
sd, prefix, dtype, rope_type, frequencies_precision, pe_max_pos
|
||||
)
|
||||
|
||||
|
||||
def load_audio_embeddings_connector(ltxv_path, dtype=torch.bfloat16):
|
||||
sd, metadata = load_torch_file(str(ltxv_path), return_metadata=True)
|
||||
config = json.loads(metadata.get("config", "{}"))
|
||||
transformer_config = config.get("transformer", {})
|
||||
rope_type = LTXRopeType.from_dict(transformer_config)
|
||||
frequencies_precision = LTXFrequenciesPrecision.from_dict(transformer_config)
|
||||
pe_max_pos = transformer_config.get("connector_positional_embedding_max_pos", [1])
|
||||
return load_embeddings_connector(
|
||||
sd,
|
||||
f"{PREFIX_BASE}audio_embeddings_connector.",
|
||||
dtype,
|
||||
rope_type,
|
||||
frequencies_precision,
|
||||
pe_max_pos,
|
||||
)
|
||||
|
||||
|
||||
def load_embeddings_connector(
|
||||
sd,
|
||||
connector_prefix,
|
||||
dtype=torch.bfloat16,
|
||||
rope_type=LTXRopeType.INTERLEAVED,
|
||||
frequencies_precision=LTXFrequenciesPrecision.FLOAT32,
|
||||
pe_max_pos=None,
|
||||
):
|
||||
sd_connector = {
|
||||
k[len(connector_prefix) :]: v
|
||||
for k, v in sd.items()
|
||||
if k.startswith(connector_prefix)
|
||||
}
|
||||
|
||||
if len(sd_connector) == 0:
|
||||
return None
|
||||
|
||||
operations = comfy.ops.pick_operations(dtype, dtype, disable_fast_fp8=True)
|
||||
connector = Embeddings1DConnector(
|
||||
dtype=dtype,
|
||||
operations=operations,
|
||||
positional_embedding_max_pos=pe_max_pos if pe_max_pos is not None else [1],
|
||||
split_rope=rope_type == LTXRopeType.SPLIT,
|
||||
double_precision_rope=frequencies_precision == LTXFrequenciesPrecision.FLOAT64,
|
||||
)
|
||||
connector.load_state_dict(sd_connector)
|
||||
return connector
|
||||
|
||||
|
||||
def load_proj_matrix_from_ltxv(ltxv_path: Path, prefix=""):
|
||||
with safetensors.safe_open(ltxv_path, framework="pt", device="cpu") as f:
|
||||
keys = filter(lambda key: key.startswith(prefix), f.keys())
|
||||
sd = {k.removeprefix(prefix): f.get_tensor(k) for k in keys if k in f.keys()}
|
||||
if not sd:
|
||||
return None
|
||||
model = GemmaFeaturesExtractorProjLinear()
|
||||
model.load_state_dict(sd)
|
||||
return model
|
||||
|
||||
|
||||
def load_proj_matrix_from_checkpoint(checkpoint_path: Path):
|
||||
"""
|
||||
Load model weights from a checkpoint file.
|
||||
:param checkpoint_path: Path to the checkpoint file.
|
||||
"""
|
||||
model = GemmaFeaturesExtractorProjLinear()
|
||||
loaded_state_dict = load_torch_file(str(checkpoint_path), return_metadata=False)
|
||||
if "aggregate_embed.weight" not in loaded_state_dict:
|
||||
raise ValueError(
|
||||
f"Checkpoint {checkpoint_path} does not contain 'aggregate_embed.weight'."
|
||||
)
|
||||
model.load_state_dict(loaded_state_dict)
|
||||
return model
|
||||
|
||||
|
||||
class LTXVGemmaTokenizer:
|
||||
def __init__(self, tokenizer_path: str, max_length: int = 1024):
|
||||
self.tokenizer = AutoTokenizer.from_pretrained(
|
||||
@@ -180,22 +83,12 @@ class LTXVGemmaTokenizer:
|
||||
return out
|
||||
|
||||
|
||||
class GemmaFeaturesExtractorProjLinear(torch.nn.Module):
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
self.aggregate_embed = torch.nn.Linear(3840 * 49, 3840, bias=False)
|
||||
|
||||
def forward(self, x):
|
||||
return self.aggregate_embed(x)
|
||||
|
||||
|
||||
class LTXVGemmaTextEncoderModel(torch.nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
model: Gemma3ForConditionalGeneration,
|
||||
feature_extractor_linear: GemmaFeaturesExtractorProjLinear,
|
||||
embeddings_connector: Embeddings1DConnector,
|
||||
audio_embeddings_connector: Embeddings1DConnector | None = None,
|
||||
feature_extractor, # FeatureExtractorV1/V2
|
||||
embeddings_processor, # VideoEmbeddingsProcessor or AVEmbeddingsProcessor
|
||||
processor: Gemma3Processor | None = None,
|
||||
dtype=torch.bfloat16,
|
||||
device="cpu",
|
||||
@@ -203,13 +96,8 @@ class LTXVGemmaTextEncoderModel(torch.nn.Module):
|
||||
super().__init__()
|
||||
self.model = model
|
||||
self.processor = processor
|
||||
self.feature_extractor_linear = feature_extractor_linear.to(dtype=dtype)
|
||||
self.embeddings_connector = embeddings_connector.to(dtype=dtype)
|
||||
self.audio_embeddings_connector = (
|
||||
audio_embeddings_connector.to(dtype=dtype)
|
||||
if audio_embeddings_connector is not None
|
||||
else None
|
||||
)
|
||||
self.feature_extractor = feature_extractor.to(dtype=dtype)
|
||||
self.embeddings_processor = embeddings_processor.to(dtype=dtype)
|
||||
self.dtypes = set([dtype])
|
||||
# Cache an estimate of memory required to load/keep the model on device
|
||||
# weights size + small overhead
|
||||
@@ -223,89 +111,20 @@ class LTXVGemmaTextEncoderModel(torch.nn.Module):
|
||||
def reset_clip_options(self):
|
||||
pass
|
||||
|
||||
@staticmethod
|
||||
def norm_and_concat_padded_batch(
|
||||
encoded_text: torch.Tensor,
|
||||
sequence_lengths: torch.Tensor,
|
||||
padding_side: str = "right",
|
||||
) -> torch.Tensor:
|
||||
"""
|
||||
Normalize a 4D tensor [B, T, D, L] per sample and per layer, using sequence_lengths to mask.
|
||||
Returns [B, T, D * L] tensor with original padding preserved.
|
||||
|
||||
Args:
|
||||
encoded_text: 4D tensor [B, T, D, L]
|
||||
sequence_lengths: 1D tensor [B] with actual sequence lengths
|
||||
padding_side: "left" or "right" to indicate which side has padding
|
||||
"""
|
||||
B, T, D, L = encoded_text.shape
|
||||
device = encoded_text.device
|
||||
|
||||
# Build mask: [B, T, 1, 1]
|
||||
token_indices = torch.arange(T, device=device)[None, :] # [1, T]
|
||||
|
||||
if padding_side == "right":
|
||||
# For right padding, valid tokens are from 0 to sequence_length-1
|
||||
mask = token_indices < sequence_lengths[:, None] # [B, T]
|
||||
elif padding_side == "left":
|
||||
# For left padding, valid tokens are from (T - sequence_length) to T-1
|
||||
start_indices = T - sequence_lengths[:, None] # [B, 1]
|
||||
mask = token_indices >= start_indices # [B, T]
|
||||
else:
|
||||
raise ValueError(
|
||||
f"padding_side must be 'left' or 'right', got {padding_side}"
|
||||
)
|
||||
|
||||
mask = rearrange(mask, "b t -> b t 1 1")
|
||||
|
||||
# Compute masked mean: [B, 1, 1, L]
|
||||
masked = encoded_text.masked_fill(~mask, 0.0)
|
||||
# denom = mask.sum(dim=(1, 2), keepdim=True) # [B, 1, 1, 1]
|
||||
denom = (sequence_lengths * D).view(B, 1, 1, 1)
|
||||
mean = masked.sum(dim=(1, 2), keepdim=True) / (denom + 1e-6)
|
||||
|
||||
# Compute masked min/max: [B, 1, 1, L]
|
||||
x_min = encoded_text.masked_fill(~mask, float("inf")).amin(
|
||||
dim=(1, 2), keepdim=True
|
||||
)
|
||||
x_max = encoded_text.masked_fill(~mask, float("-inf")).amax(
|
||||
dim=(1, 2), keepdim=True
|
||||
)
|
||||
range_ = x_max - x_min
|
||||
|
||||
# Normalize only the valid tokens
|
||||
normed = 8 * (encoded_text - mean) / (range_ + 1e-6)
|
||||
|
||||
# concat to be [Batch, T, D * L] - this preserves the original structure
|
||||
normed = normed.reshape(B, T, -1) # [B, T, D * L]
|
||||
|
||||
# Apply mask to preserve original padding (set padded positions to 0)
|
||||
mask_flattened = rearrange(mask, "b t 1 1 -> b t 1").expand(-1, -1, D * L)
|
||||
normed = normed.masked_fill(~mask_flattened, 0.0)
|
||||
|
||||
return normed
|
||||
|
||||
def forward(self, input_ids, attention_mask, padding_side="right"):
|
||||
# Block 1: Run Gemma model
|
||||
outputs = self.model(
|
||||
input_ids=input_ids,
|
||||
attention_mask=attention_mask,
|
||||
output_hidden_states=True,
|
||||
)
|
||||
encoded_text_features = torch.stack(outputs.hidden_states, dim=-1)
|
||||
encoded_text_features_dtype = encoded_text_features.dtype
|
||||
all_layer_hiddens = torch.stack(outputs.hidden_states, dim=-1) # [B, T, D, L]
|
||||
|
||||
sequence_lengths = attention_mask.sum(dim=-1)
|
||||
normed_concated_encoded_text_features = (
|
||||
LTXVGemmaTextEncoderModel.norm_and_concat_padded_batch(
|
||||
encoded_text_features, sequence_lengths, padding_side=padding_side
|
||||
)
|
||||
# Block 2: Feature extraction
|
||||
features = self.feature_extractor(
|
||||
all_layer_hiddens, attention_mask, padding_side
|
||||
)
|
||||
|
||||
projected = self.feature_extractor_linear(
|
||||
normed_concated_encoded_text_features.to(encoded_text_features_dtype)
|
||||
)
|
||||
|
||||
return projected
|
||||
return features # dict with "video" and optionally "audio"
|
||||
|
||||
def encode_token_weights(self, token_weight_pairs):
|
||||
token_pairs = token_weight_pairs["gemma"]
|
||||
@@ -318,31 +137,20 @@ class LTXVGemmaTextEncoderModel(torch.nn.Module):
|
||||
|
||||
self.to(self.model.device)
|
||||
|
||||
encoded_input = self(input_ids, attention_mask, padding_side="left")
|
||||
# convert attention mask to format embeddings connector expects.
|
||||
connector_attention_mask = (attention_mask - 1).to(encoded_input.dtype).reshape(
|
||||
features = self(input_ids, attention_mask, padding_side="left")
|
||||
|
||||
# Convert binary mask -> additive mask for processor
|
||||
encoded_input_dtype = next(iter(features.values())).dtype
|
||||
connector_attention_mask = (attention_mask - 1).to(encoded_input_dtype).reshape(
|
||||
(attention_mask.shape[0], 1, -1, attention_mask.shape[-1])
|
||||
) * torch.finfo(encoded_input.dtype).max
|
||||
encoded, encoded_connector_attention_mask = self.embeddings_connector(
|
||||
encoded_input,
|
||||
connector_attention_mask,
|
||||
) * torch.finfo(encoded_input_dtype).max
|
||||
|
||||
# Block 3: Embeddings processor
|
||||
encoded, mask = self.embeddings_processor.create_embeddings(
|
||||
features, connector_attention_mask
|
||||
)
|
||||
|
||||
# restore the mask values to int64
|
||||
attention_mask = (encoded_connector_attention_mask < 0.000001).to(torch.int64)
|
||||
|
||||
attention_mask = attention_mask.reshape([encoded.shape[0], encoded.shape[1], 1])
|
||||
encoded = encoded * attention_mask
|
||||
|
||||
if self.audio_embeddings_connector is not None:
|
||||
encoded_for_audio, _ = self.audio_embeddings_connector(
|
||||
encoded_input, connector_attention_mask
|
||||
)
|
||||
encoded = torch.concat(
|
||||
[encoded, encoded_for_audio], dim=len(encoded.shape) - 1
|
||||
)
|
||||
|
||||
return encoded, None, {"attention_mask": attention_mask.squeeze(-1)}
|
||||
return encoded, None, {"attention_mask": mask}
|
||||
|
||||
def load_sd(self, sd):
|
||||
return self.model.load_state_dict(sd, strict=False)
|
||||
@@ -351,136 +159,6 @@ class LTXVGemmaTextEncoderModel(torch.nn.Module):
|
||||
# Return a conservative estimate in bytesed(input_shape)
|
||||
return self._model_memory_required
|
||||
|
||||
@staticmethod
|
||||
def _cat_with_padding(
|
||||
tensor: torch.Tensor,
|
||||
padding_length: int,
|
||||
value: int | float,
|
||||
) -> torch.Tensor:
|
||||
"""Concatenate a tensor with a padding tensor of the given value."""
|
||||
return torch.cat(
|
||||
[
|
||||
tensor,
|
||||
torch.full(
|
||||
(1, padding_length),
|
||||
value,
|
||||
dtype=tensor.dtype,
|
||||
device=tensor.device,
|
||||
),
|
||||
],
|
||||
dim=1,
|
||||
)
|
||||
|
||||
def _pad_inputs_for_attention_alignment(self, model_inputs, alignment: int = 8):
|
||||
"""Pad sequence length to multiple of alignment for Flash Attention compatibility.
|
||||
|
||||
Flash Attention within SDPA requires sequence lengths aligned to 8 bytes.
|
||||
This pads input_ids, attention_mask, and token_type_ids (if present) to prevent
|
||||
'p.attn_bias_ptr is not correctly aligned' errors.
|
||||
"""
|
||||
seq_len = model_inputs.input_ids.shape[1]
|
||||
padded_len = ((seq_len + alignment - 1) // alignment) * alignment
|
||||
padding_length = padded_len - seq_len
|
||||
|
||||
if padding_length > 0:
|
||||
pad_token_id = (
|
||||
self.processor.tokenizer.pad_token_id
|
||||
if self.processor.tokenizer.pad_token_id is not None
|
||||
else 0
|
||||
)
|
||||
|
||||
model_inputs["input_ids"] = self._cat_with_padding(
|
||||
model_inputs.input_ids, padding_length, pad_token_id
|
||||
)
|
||||
|
||||
model_inputs["attention_mask"] = self._cat_with_padding(
|
||||
model_inputs.attention_mask, padding_length, 0
|
||||
)
|
||||
|
||||
if (
|
||||
"token_type_ids" in model_inputs
|
||||
and model_inputs["token_type_ids"] is not None
|
||||
):
|
||||
model_inputs["token_type_ids"] = self._cat_with_padding(
|
||||
model_inputs["token_type_ids"], padding_length, 0
|
||||
)
|
||||
|
||||
return model_inputs
|
||||
|
||||
def _enhance(
|
||||
self,
|
||||
messages: list,
|
||||
image: Optional[Image.Image] = None,
|
||||
max_new_tokens: int = 512,
|
||||
seed: int = 42,
|
||||
) -> str:
|
||||
"""Common enhancement logic for both T2V and I2V modes."""
|
||||
if self.processor is None:
|
||||
raise ValueError("Processor not loaded - enhancement not available")
|
||||
|
||||
text = self.processor.tokenizer.apply_chat_template(
|
||||
messages, tokenize=False, add_generation_prompt=True
|
||||
)
|
||||
|
||||
model_inputs = self.processor(
|
||||
text=text,
|
||||
images=image,
|
||||
return_tensors="pt",
|
||||
).to(self.model.device)
|
||||
model_inputs = self._pad_inputs_for_attention_alignment(model_inputs)
|
||||
|
||||
with torch.inference_mode(), torch.random.fork_rng(devices=[self.model.device]):
|
||||
torch.manual_seed(seed)
|
||||
outputs = self.model.generate(
|
||||
**model_inputs,
|
||||
max_new_tokens=max_new_tokens,
|
||||
do_sample=True,
|
||||
temperature=0.7,
|
||||
)
|
||||
generated_ids = outputs[0][len(model_inputs.input_ids[0]) :]
|
||||
enhanced_prompt = self.processor.tokenizer.decode(
|
||||
generated_ids, skip_special_tokens=True
|
||||
)
|
||||
|
||||
return enhanced_prompt
|
||||
|
||||
def enhance_t2v(
|
||||
self,
|
||||
prompt: str,
|
||||
system_prompt: str,
|
||||
max_new_tokens: int = 512,
|
||||
seed: int = 42,
|
||||
) -> str:
|
||||
"""Enhance a text prompt for T2V generation."""
|
||||
messages = [
|
||||
{"role": "system", "content": system_prompt},
|
||||
{"role": "user", "content": f"User Raw Input Prompt: {prompt}."},
|
||||
]
|
||||
return self._enhance(messages, max_new_tokens=max_new_tokens, seed=seed)
|
||||
|
||||
def enhance_i2v(
|
||||
self,
|
||||
prompt: str,
|
||||
image: Image.Image,
|
||||
system_prompt: str,
|
||||
max_new_tokens: int = 512,
|
||||
seed: int = 42,
|
||||
) -> str:
|
||||
"""Enhance a text prompt for I2V generation using a reference image."""
|
||||
messages = [
|
||||
{"role": "system", "content": system_prompt},
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{"type": "image"},
|
||||
{"type": "text", "text": f"User Raw Input Prompt: {prompt}."},
|
||||
],
|
||||
},
|
||||
]
|
||||
return self._enhance(
|
||||
messages, image=image, max_new_tokens=max_new_tokens, seed=seed
|
||||
)
|
||||
|
||||
|
||||
def ltxv_gemma_tokenizer(tokenizer_path, max_length=256):
|
||||
class _LTXVGemmaTokenizer(LTXVGemmaTokenizer):
|
||||
@@ -494,28 +172,23 @@ def ltxv_gemma_clip(encoder_path, ltxv_path, processor=None, dtype=None):
|
||||
class _LTXVGemmaTextEncoderModel(LTXVGemmaTextEncoderModel):
|
||||
def __init__(self, device="cpu", dtype=dtype, model_options={}):
|
||||
dtype = torch.bfloat16 # TODO: make this configurable
|
||||
|
||||
gemma_model = Gemma3ForConditionalGeneration.from_pretrained(
|
||||
encoder_path,
|
||||
local_files_only=True,
|
||||
torch_dtype=dtype,
|
||||
)
|
||||
|
||||
feature_extractor_linear = load_proj_matrix_from_ltxv(
|
||||
feature_extractor, embeddings_processor = load_text_embeddings_pipeline(
|
||||
ltxv_path,
|
||||
"text_embedding_projection.",
|
||||
dtype=dtype,
|
||||
fallback_proj_path=encoder_path / "proj_linear.safetensors",
|
||||
)
|
||||
if feature_extractor_linear is None:
|
||||
feature_extractor_linear = load_proj_matrix_from_checkpoint(
|
||||
encoder_path / "proj_linear.safetensors"
|
||||
)
|
||||
|
||||
embeddings_connector = load_video_embeddings_connector(ltxv_path)
|
||||
audio_embeddings_connector = load_audio_embeddings_connector(ltxv_path)
|
||||
super().__init__(
|
||||
model=gemma_model,
|
||||
feature_extractor_linear=feature_extractor_linear,
|
||||
embeddings_connector=embeddings_connector,
|
||||
audio_embeddings_connector=audio_embeddings_connector,
|
||||
feature_extractor=feature_extractor,
|
||||
embeddings_processor=embeddings_processor,
|
||||
processor=processor,
|
||||
dtype=dtype,
|
||||
device=device,
|
||||
@@ -528,8 +201,11 @@ def find_matching_dir(root_path: str, pattern: str) -> str:
|
||||
"""
|
||||
Recursively search for files matching a glob pattern and return the parent directory of the first match.
|
||||
"""
|
||||
|
||||
matches = list(Path(root_path).rglob(pattern))
|
||||
matches = [
|
||||
Path(p)
|
||||
for p in glob(f"{root_path}/**", recursive=True)
|
||||
if Path(p).match(pattern)
|
||||
]
|
||||
if not matches:
|
||||
raise FileNotFoundError(
|
||||
f"No files matching pattern '{pattern}' found under {root_path}"
|
||||
@@ -672,13 +348,18 @@ class LTXVGemmaEnhancePrompt:
|
||||
clip.load_model()
|
||||
encoder = clip.cond_stage_model
|
||||
|
||||
if encoder.processor is None:
|
||||
raise ValueError(
|
||||
"Processor not loaded - enhancement not available. "
|
||||
"Ensure your model directory has chat_template.json, processor_config.json, "
|
||||
"and preprocessor_config.json files."
|
||||
)
|
||||
|
||||
if not hasattr(encoder, "processor") or encoder.processor is None:
|
||||
if hasattr(encoder, "gemma3_12b"):
|
||||
model, processor = transformers_gemma3_from_encoder(encoder)
|
||||
else:
|
||||
raise ValueError(
|
||||
"Processor not loaded - enhancement not available. "
|
||||
"Ensure your model directory has chat_template.json, processor_config.json, "
|
||||
"and preprocessor_config.json files."
|
||||
)
|
||||
elif isinstance(encoder, LTXVGemmaTextEncoderModel):
|
||||
model = encoder.model
|
||||
processor = encoder.processor
|
||||
# Determine mode: use I2V if image is provided and not bypassed
|
||||
use_i2v = image is not None and not bypass_i2v
|
||||
|
||||
@@ -694,7 +375,9 @@ class LTXVGemmaEnhancePrompt:
|
||||
|
||||
if use_i2v:
|
||||
pil_image = tensor_to_pil(image)
|
||||
enhanced_prompt = encoder.enhance_i2v(
|
||||
enhanced_prompt = enhance_i2v(
|
||||
processor=processor,
|
||||
model=model,
|
||||
prompt=prompt,
|
||||
image=pil_image,
|
||||
system_prompt=system_prompt,
|
||||
@@ -702,7 +385,9 @@ class LTXVGemmaEnhancePrompt:
|
||||
seed=seed,
|
||||
)
|
||||
else:
|
||||
enhanced_prompt = encoder.enhance_t2v(
|
||||
enhanced_prompt = enhance_t2v(
|
||||
processor=processor,
|
||||
model=model,
|
||||
prompt=prompt,
|
||||
system_prompt=system_prompt,
|
||||
max_new_tokens=max_tokens,
|
||||
@@ -712,3 +397,275 @@ class LTXVGemmaEnhancePrompt:
|
||||
enhanced_prompt = clean_response(enhanced_prompt)
|
||||
|
||||
return (enhanced_prompt,)
|
||||
|
||||
|
||||
def _enhance(
|
||||
processor: Gemma3Processor,
|
||||
model: Gemma3ForConditionalGeneration,
|
||||
messages: list,
|
||||
image: Optional[Image.Image] = None,
|
||||
max_new_tokens: int = 512,
|
||||
seed: int = 42,
|
||||
) -> str:
|
||||
"""Common enhancement logic for both T2V and I2V modes."""
|
||||
if processor is None:
|
||||
raise ValueError("Processor not loaded - enhancement not available")
|
||||
|
||||
text = processor.tokenizer.apply_chat_template(
|
||||
messages, tokenize=False, add_generation_prompt=True
|
||||
)
|
||||
|
||||
model_inputs = processor(
|
||||
text=text,
|
||||
images=image,
|
||||
return_tensors="pt",
|
||||
).to(model.device)
|
||||
|
||||
pad_token_id = (
|
||||
processor.tokenizer.pad_token_id
|
||||
if processor.tokenizer.pad_token_id is not None
|
||||
else 0
|
||||
)
|
||||
model_inputs = _pad_inputs_for_attention_alignment(model_inputs, pad_token_id)
|
||||
|
||||
with (
|
||||
torch.inference_mode(),
|
||||
torch.random.fork_rng(devices=[model.device]),
|
||||
torch.autocast(device_type=model.device.type, dtype=model.dtype),
|
||||
):
|
||||
torch.manual_seed(seed)
|
||||
outputs = model.generate(
|
||||
**model_inputs,
|
||||
max_new_tokens=max_new_tokens,
|
||||
do_sample=True,
|
||||
temperature=0.7,
|
||||
)
|
||||
generated_ids = outputs[0][len(model_inputs.input_ids[0]) :]
|
||||
enhanced_prompt = processor.tokenizer.decode(
|
||||
generated_ids, skip_special_tokens=True
|
||||
)
|
||||
|
||||
return enhanced_prompt
|
||||
|
||||
|
||||
def enhance_t2v(
|
||||
processor: Gemma3Processor,
|
||||
model: Gemma3ForConditionalGeneration,
|
||||
prompt: str,
|
||||
system_prompt: str,
|
||||
max_new_tokens: int = 512,
|
||||
seed: int = 42,
|
||||
) -> str:
|
||||
"""Enhance a text prompt for T2V generation."""
|
||||
messages = [
|
||||
{"role": "system", "content": system_prompt},
|
||||
{"role": "user", "content": f"User Raw Input Prompt: {prompt}."},
|
||||
]
|
||||
return _enhance(
|
||||
processor, model, messages, max_new_tokens=max_new_tokens, seed=seed
|
||||
)
|
||||
|
||||
|
||||
def enhance_i2v(
|
||||
processor: Gemma3Processor,
|
||||
model: Gemma3ForConditionalGeneration,
|
||||
prompt: str,
|
||||
image: Image.Image,
|
||||
system_prompt: str,
|
||||
max_new_tokens: int = 512,
|
||||
seed: int = 42,
|
||||
) -> str:
|
||||
"""Enhance a text prompt for I2V generation using a reference image."""
|
||||
messages = [
|
||||
{"role": "system", "content": system_prompt},
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{"type": "image"},
|
||||
{"type": "text", "text": f"User Raw Input Prompt: {prompt}."},
|
||||
],
|
||||
},
|
||||
]
|
||||
return _enhance(
|
||||
processor,
|
||||
model,
|
||||
messages,
|
||||
image=image,
|
||||
max_new_tokens=max_new_tokens,
|
||||
seed=seed,
|
||||
)
|
||||
|
||||
|
||||
def _cat_with_padding(
|
||||
tensor: torch.Tensor,
|
||||
padding_length: int,
|
||||
value: int | float,
|
||||
) -> torch.Tensor:
|
||||
"""Concatenate a tensor with a padding tensor of the given value."""
|
||||
return torch.cat(
|
||||
[
|
||||
tensor,
|
||||
torch.full(
|
||||
(1, padding_length),
|
||||
value,
|
||||
dtype=tensor.dtype,
|
||||
device=tensor.device,
|
||||
),
|
||||
],
|
||||
dim=1,
|
||||
)
|
||||
|
||||
|
||||
def _pad_inputs_for_attention_alignment(model_inputs, pad_token_id, alignment: int = 8):
|
||||
"""Pad sequence length to multiple of alignment for Flash Attention compatibility.
|
||||
|
||||
Flash Attention within SDPA requires sequence lengths aligned to 8 bytes.
|
||||
This pads input_ids, attention_mask, and token_type_ids (if present) to prevent
|
||||
'p.attn_bias_ptr is not correctly aligned' errors.
|
||||
"""
|
||||
seq_len = model_inputs.input_ids.shape[1]
|
||||
padded_len = ((seq_len + alignment - 1) // alignment) * alignment
|
||||
padding_length = padded_len - seq_len
|
||||
|
||||
if padding_length > 0:
|
||||
model_inputs["input_ids"] = _cat_with_padding(
|
||||
model_inputs.input_ids, padding_length, pad_token_id
|
||||
)
|
||||
|
||||
model_inputs["attention_mask"] = _cat_with_padding(
|
||||
model_inputs.attention_mask, padding_length, 0
|
||||
)
|
||||
|
||||
if (
|
||||
"token_type_ids" in model_inputs
|
||||
and model_inputs["token_type_ids"] is not None
|
||||
):
|
||||
model_inputs["token_type_ids"] = _cat_with_padding(
|
||||
model_inputs["token_type_ids"], padding_length, 0
|
||||
)
|
||||
|
||||
return model_inputs
|
||||
|
||||
|
||||
def _locate_model_within_model(super_model, model_name):
|
||||
class_name = MODEL_MAPPING_NAMES.get(model_name, None)
|
||||
if class_name is None:
|
||||
return None
|
||||
for module in super_model.modules():
|
||||
if module.__class__.__name__ == class_name:
|
||||
return module
|
||||
return None
|
||||
|
||||
|
||||
def _locate_unique_parameter_owner_by_leaf(
|
||||
root: torch.nn.Module,
|
||||
leaf_param_name: str,
|
||||
must_have_in_path: Optional[str] = None,
|
||||
) -> Optional[Tuple[torch.nn.Module, str, torch.nn.Parameter, str]]:
|
||||
|
||||
modules = dict(root.named_modules())
|
||||
|
||||
candidates: List[Tuple[torch.nn.Module, str, torch.nn.Parameter, str]] = []
|
||||
for full_name, p in root.named_parameters(recurse=True):
|
||||
parts = full_name.split(".")
|
||||
leaf = parts[-1]
|
||||
if leaf != leaf_param_name:
|
||||
continue
|
||||
if must_have_in_path is not None and must_have_in_path not in parts:
|
||||
continue
|
||||
|
||||
owner_path = ".".join(parts[:-1])
|
||||
owner = modules.get(owner_path, root if owner_path == "" else None)
|
||||
if owner is None:
|
||||
continue
|
||||
candidates.append((owner, leaf, p, full_name))
|
||||
|
||||
if not candidates:
|
||||
return None
|
||||
return candidates[0]
|
||||
|
||||
|
||||
def transformers_gemma3_from_encoder(encoder):
|
||||
jsons_path = Path(__file__).parent / "gemma_configs"
|
||||
config = Gemma3Config.from_json_file(jsons_path / "gemma3cfg.json")
|
||||
with torch.device("meta"):
|
||||
metamodel = Gemma3ForConditionalGeneration(config)
|
||||
t_model_name = config.text_config.model_type
|
||||
t_model = _locate_model_within_model(metamodel, t_model_name)
|
||||
if t_model is None:
|
||||
raise ValueError(
|
||||
"Can't locate text model while converting text encoder to Gemma3ForConditionalGeneration"
|
||||
)
|
||||
t_model.load_state_dict(
|
||||
encoder.gemma3_12b.transformer.model.state_dict(), assign=True, strict=False
|
||||
)
|
||||
v_tower_name = config.vision_config.model_type
|
||||
v_tower = _locate_model_within_model(metamodel, v_tower_name)
|
||||
if v_tower is None:
|
||||
raise ValueError(
|
||||
"Can't locate vision model while converting text encoder to Gemma3ForConditionalGeneration"
|
||||
)
|
||||
v_model = v_tower.vision_model
|
||||
v_model.load_state_dict(
|
||||
encoder.gemma3_12b.transformer.vision_model.state_dict(),
|
||||
assign=True,
|
||||
strict=False,
|
||||
)
|
||||
metamodel.multi_modal_projector.load_state_dict(
|
||||
encoder.gemma3_12b.transformer.multi_modal_projector.state_dict(),
|
||||
assign=True,
|
||||
strict=False,
|
||||
)
|
||||
config = config.text_config
|
||||
dim = getattr(config, "head_dim", config.hidden_size // config.num_attention_heads)
|
||||
base = config.rope_local_base_freq
|
||||
|
||||
device = encoder.device
|
||||
positions_length = len(v_model.embeddings.position_ids[0])
|
||||
position_ids = torch.arange(
|
||||
positions_length, dtype=torch.long, device="cpu"
|
||||
).unsqueeze(0)
|
||||
v_model.embeddings.register_buffer("position_ids", position_ids)
|
||||
embed_scale = torch.tensor(config.hidden_size**0.5, device=device)
|
||||
t_model.embed_tokens.register_buffer("embed_scale", embed_scale)
|
||||
local_rope_freqs = 1.0 / (
|
||||
base
|
||||
** (
|
||||
torch.arange(0, dim, 2, dtype=torch.int64).to(
|
||||
device=device, dtype=torch.float
|
||||
)
|
||||
/ dim
|
||||
)
|
||||
)
|
||||
t_model.rotary_emb_local.register_buffer("inv_freq", local_rope_freqs)
|
||||
rope_freqs, _ = ROPE_INIT_FUNCTIONS[config.rope_scaling["rope_type"]](
|
||||
config, device
|
||||
)
|
||||
t_model.rotary_emb.register_buffer("inv_freq", rope_freqs)
|
||||
lm_head_requires_grad = False
|
||||
loc_result = _locate_unique_parameter_owner_by_leaf(
|
||||
metamodel, leaf_param_name="weight", must_have_in_path="lm_head"
|
||||
)
|
||||
if loc_result is None:
|
||||
raise ValueError(
|
||||
"Can't locate lm_head while converting text encoder to Gemma3ForConditionalGeneration"
|
||||
)
|
||||
lm_head_owner, lm_head_attr, _, _ = loc_result
|
||||
real_w = t_model.embed_tokens.weight
|
||||
setattr(
|
||||
lm_head_owner,
|
||||
lm_head_attr,
|
||||
torch.nn.Parameter(real_w, requires_grad=lm_head_requires_grad),
|
||||
)
|
||||
metamodel.to(device)
|
||||
|
||||
tokenizer_class = ltxv_gemma_tokenizer(jsons_path, max_length=1024)
|
||||
image_processor = AutoImageProcessor.from_pretrained(
|
||||
str(jsons_path),
|
||||
local_files_only=True,
|
||||
)
|
||||
processor = Gemma3Processor(
|
||||
image_processor=image_processor,
|
||||
tokenizer=tokenizer_class().tokenizer,
|
||||
)
|
||||
return metamodel, processor
|
||||
|
||||
@@ -252,3 +252,205 @@ class LTXVImgToVideoAdvanced:
|
||||
batch_size=batch_size,
|
||||
strength=strength,
|
||||
)
|
||||
|
||||
|
||||
@comfy_node(name="LTXVAddGuideAdvancedAttention")
|
||||
class LTXVAddGuideAdvancedAttention:
|
||||
"""Extended keyframe guide node with per-guide attention strength control.
|
||||
|
||||
Same preprocessing as LTXVAddGuideAdvanced (CRF, blur, interpolation, crop),
|
||||
plus attention_strength and attention_mask inputs to control how strongly
|
||||
this guide's conditioning influences generation via self-attention.
|
||||
"""
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"positive": ("CONDITIONING",),
|
||||
"negative": ("CONDITIONING",),
|
||||
"vae": ("VAE",),
|
||||
"latent": ("LATENT",),
|
||||
"image": ("IMAGE",),
|
||||
"frame_idx": (
|
||||
"INT",
|
||||
{
|
||||
"default": 0,
|
||||
"min": -9999,
|
||||
"max": 9999,
|
||||
"tooltip": (
|
||||
"Frame index to start the conditioning at. "
|
||||
"Negative values are counted from the end of the video."
|
||||
),
|
||||
},
|
||||
),
|
||||
"strength": (
|
||||
"FLOAT",
|
||||
{
|
||||
"default": 1.0,
|
||||
"min": 0.0,
|
||||
"max": 1.0,
|
||||
"tooltip": "Strength of the conditioning. Higher values make it more exact.",
|
||||
},
|
||||
),
|
||||
"crf": (
|
||||
"INT",
|
||||
{
|
||||
"default": 29,
|
||||
"min": 0,
|
||||
"max": 51,
|
||||
"step": 1,
|
||||
"tooltip": "CRF value. Higher = more motion, lower = higher quality.",
|
||||
},
|
||||
),
|
||||
"blur_radius": (
|
||||
"INT",
|
||||
{
|
||||
"default": 0,
|
||||
"min": 0,
|
||||
"max": 7,
|
||||
"step": 1,
|
||||
"tooltip": "Blur kernel radius. Higher = more motion.",
|
||||
},
|
||||
),
|
||||
"interpolation": (
|
||||
[
|
||||
"lanczos",
|
||||
"bislerp",
|
||||
"nearest",
|
||||
"bilinear",
|
||||
"bicubic",
|
||||
"area",
|
||||
"nearest-exact",
|
||||
],
|
||||
{"default": "lanczos"},
|
||||
),
|
||||
"crop": (["center", "disabled"], {"default": "disabled"}),
|
||||
"attention_strength": (
|
||||
"FLOAT",
|
||||
{
|
||||
"default": 1.0,
|
||||
"min": 0.0,
|
||||
"max": 1.0,
|
||||
"step": 0.01,
|
||||
"tooltip": (
|
||||
"Controls how strongly this guide influences generation via "
|
||||
"self-attention. 1.0 = full conditioning, 0.0 = ignore."
|
||||
),
|
||||
},
|
||||
),
|
||||
},
|
||||
"optional": {
|
||||
"attention_mask": (
|
||||
"MASK",
|
||||
{
|
||||
"tooltip": (
|
||||
"Optional pixel-space spatial mask. Shape (F, H, W) or (H, W). "
|
||||
"Values in [0, 1]. Controls per-region conditioning influence. "
|
||||
"Multiplied by attention_strength."
|
||||
),
|
||||
},
|
||||
),
|
||||
},
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("CONDITIONING", "CONDITIONING", "LATENT")
|
||||
RETURN_NAMES = ("positive", "negative", "latent")
|
||||
|
||||
CATEGORY = "conditioning/video_models"
|
||||
FUNCTION = "generate"
|
||||
|
||||
DESCRIPTION = (
|
||||
"Adds a conditioning frame/video at a specific frame index with per-guide "
|
||||
"attention strength control. Same preprocessing as LTXVAddGuideAdvanced, "
|
||||
"plus attention_strength and optional spatial attention_mask."
|
||||
)
|
||||
|
||||
def generate(
|
||||
self,
|
||||
positive,
|
||||
negative,
|
||||
vae,
|
||||
latent,
|
||||
image,
|
||||
frame_idx,
|
||||
strength,
|
||||
crf,
|
||||
blur_radius,
|
||||
interpolation,
|
||||
crop,
|
||||
attention_strength=1.0,
|
||||
attention_mask=None,
|
||||
):
|
||||
from .iclora_attention import append_guide_attention_entry, normalize_mask
|
||||
|
||||
# Preprocessing: resize, CRF, blur (same as LTXVAddGuideAdvanced)
|
||||
scale_factors = vae.downscale_index_formula
|
||||
_, width_scale_factor, height_scale_factor = scale_factors
|
||||
latent_image = latent["samples"]
|
||||
noise_mask = nodes_lt.get_noise_mask(latent)
|
||||
_, _, latent_length, latent_height, latent_width = latent_image.shape
|
||||
|
||||
width = latent_width * width_scale_factor
|
||||
height = latent_height * height_scale_factor
|
||||
image = (
|
||||
comfy.utils.common_upscale(
|
||||
image.movedim(-1, 1), width, height, interpolation, crop=crop
|
||||
)
|
||||
.movedim(1, -1)
|
||||
.clamp(0, 1)
|
||||
)
|
||||
image = nodes_lt.LTXVPreprocess().execute(image, crf)[0]
|
||||
image = blur_internal(image, blur_radius)
|
||||
|
||||
# Encode
|
||||
_, t = nodes_lt.LTXVAddGuide.encode(
|
||||
vae, latent_width, latent_height, image, scale_factors
|
||||
)
|
||||
|
||||
# Compute latent index
|
||||
frame_idx, latent_idx = nodes_lt.LTXVAddGuide.get_latent_index(
|
||||
positive, latent_length, len(image), frame_idx, scale_factors
|
||||
)
|
||||
assert (
|
||||
latent_idx + t.shape[2] <= latent_length
|
||||
), "Conditioning frames exceed the length of the latent sequence."
|
||||
|
||||
# Append keyframe
|
||||
positive, negative, latent_image, noise_mask = (
|
||||
nodes_lt.LTXVAddGuide.append_keyframe(
|
||||
positive,
|
||||
negative,
|
||||
frame_idx,
|
||||
latent_image,
|
||||
noise_mask,
|
||||
t,
|
||||
strength,
|
||||
scale_factors,
|
||||
)
|
||||
)
|
||||
|
||||
# Track with custom attention strength/mask
|
||||
pre_filter_count = t.shape[2] * t.shape[3] * t.shape[4]
|
||||
guide_latent_shape = list(t.shape[2:])
|
||||
norm_mask = normalize_mask(attention_mask)
|
||||
positive = append_guide_attention_entry(
|
||||
positive,
|
||||
pre_filter_count,
|
||||
guide_latent_shape,
|
||||
attention_strength=attention_strength,
|
||||
attention_mask=norm_mask,
|
||||
)
|
||||
negative = append_guide_attention_entry(
|
||||
negative,
|
||||
pre_filter_count,
|
||||
guide_latent_shape,
|
||||
attention_strength=attention_strength,
|
||||
attention_mask=norm_mask,
|
||||
)
|
||||
|
||||
return (
|
||||
positive,
|
||||
negative,
|
||||
{"samples": latent_image, "noise_mask": noise_mask},
|
||||
)
|
||||
|
||||
@@ -64,7 +64,6 @@ class MultimodalGuider(comfy.samplers.CFGGuider):
|
||||
self,
|
||||
run_vx: bool,
|
||||
run_ax: bool,
|
||||
audio_length: int,
|
||||
audio_ptb: bool,
|
||||
video_ptb: bool,
|
||||
):
|
||||
@@ -72,7 +71,7 @@ class MultimodalGuider(comfy.samplers.CFGGuider):
|
||||
num_self_attns = 0
|
||||
if run_vx:
|
||||
num_self_attns += 1
|
||||
if run_ax and audio_length > 0:
|
||||
if run_ax:
|
||||
num_self_attns += 1
|
||||
|
||||
video_attn_idx = 0
|
||||
@@ -85,6 +84,20 @@ class MultimodalGuider(comfy.samplers.CFGGuider):
|
||||
|
||||
return list(stg_indexes)
|
||||
|
||||
def unpack_latents(self, x: torch.Tensor):
|
||||
latent_shapes = (
|
||||
self.conds.get("positive", {})[0]
|
||||
.get("model_conds", {})
|
||||
.get("latent_shapes", None)
|
||||
.cond
|
||||
)
|
||||
vx, ax = comfy.utils.unpack_latents(x, latent_shapes)
|
||||
return vx, ax
|
||||
|
||||
def pack_latents(self, vx: torch.Tensor, ax: torch.Tensor):
|
||||
x, latent_shapes = comfy.utils.pack_latents([vx, ax])
|
||||
return x, latent_shapes
|
||||
|
||||
def predict_noise(
|
||||
self,
|
||||
x: torch.Tensor,
|
||||
@@ -92,7 +105,6 @@ class MultimodalGuider(comfy.samplers.CFGGuider):
|
||||
model_options: dict = {},
|
||||
seed=None,
|
||||
):
|
||||
x_shape = x.shape
|
||||
# in CFGGuider.predict_noise, we call sampling_function(), which uses cfg_function() to compute pos & neg
|
||||
# but we'd rather do a single batch of sampling pos, neg, and perturbed, so we call calc_cond_batch([perturbed,pos,neg]) directly
|
||||
current_step = self.current_step
|
||||
@@ -108,23 +120,18 @@ class MultimodalGuider(comfy.samplers.CFGGuider):
|
||||
Modality.VIDEO.value, GuiderParameters()
|
||||
)
|
||||
|
||||
audio_length = model_options["transformer_options"].get("audio_length", 0)
|
||||
|
||||
run_vx = not video_params.do_skip(current_step)
|
||||
run_ax = not audio_params.do_skip(current_step)
|
||||
run_a2v = video_params.do_cross_attn(current_step)
|
||||
run_v2a = audio_params.do_cross_attn(current_step)
|
||||
|
||||
vx, ax = self.inner_model.diffusion_model.separate_audio_and_video_latents(
|
||||
x, audio_length
|
||||
)
|
||||
vx, ax = self.unpack_latents(x)
|
||||
if not run_vx:
|
||||
vx = self.last_denoised_v
|
||||
if not run_ax:
|
||||
ax = self.last_denoised_a
|
||||
x = self.inner_model.diffusion_model.recombine_audio_and_video_latents(
|
||||
vx, ax, x_shape
|
||||
)
|
||||
x, _ = self.pack_latents(vx, ax)
|
||||
|
||||
if not run_vx and not run_ax:
|
||||
return x
|
||||
|
||||
@@ -145,11 +152,7 @@ class MultimodalGuider(comfy.samplers.CFGGuider):
|
||||
del model_options["transformer_options"]["run_ax"]
|
||||
del model_options["transformer_options"]["a2v_cross_attn"]
|
||||
del model_options["transformer_options"]["v2a_cross_attn"]
|
||||
v_noise_pred_pos, a_noise_pred_pos = (
|
||||
self.inner_model.diffusion_model.separate_audio_and_video_latents(
|
||||
noise_pred_pos, audio_length
|
||||
)
|
||||
)
|
||||
v_noise_pred_pos, a_noise_pred_pos = self.unpack_latents(noise_pred_pos)
|
||||
|
||||
a_noise_pred_neg, v_noise_pred_neg = 0, 0
|
||||
a_noise_pred_perturbed, v_noise_pred_perturbed = 0, 0
|
||||
@@ -168,11 +171,7 @@ class MultimodalGuider(comfy.samplers.CFGGuider):
|
||||
timestep,
|
||||
model_options,
|
||||
)[0]
|
||||
v_noise_pred_neg, a_noise_pred_neg = (
|
||||
self.inner_model.diffusion_model.separate_audio_and_video_latents(
|
||||
noise_pred_neg, audio_length
|
||||
)
|
||||
)
|
||||
v_noise_pred_neg, a_noise_pred_neg = self.unpack_latents(noise_pred_neg)
|
||||
finally:
|
||||
del model_options["transformer_options"]["run_vx"]
|
||||
del model_options["transformer_options"]["run_ax"]
|
||||
@@ -183,8 +182,7 @@ class MultimodalGuider(comfy.samplers.CFGGuider):
|
||||
try:
|
||||
stg_indexes = self.calc_stg_indexes(
|
||||
run_vx,
|
||||
run_ax,
|
||||
audio_length,
|
||||
run_ax and ax.numel() > 0,
|
||||
audio_params.perturb_attn,
|
||||
video_params.perturb_attn,
|
||||
)
|
||||
@@ -202,10 +200,8 @@ class MultimodalGuider(comfy.samplers.CFGGuider):
|
||||
timestep,
|
||||
model_options,
|
||||
)[0]
|
||||
v_noise_pred_perturbed, a_noise_pred_perturbed = (
|
||||
self.inner_model.diffusion_model.separate_audio_and_video_latents(
|
||||
noise_pred_perturbed, audio_length
|
||||
)
|
||||
v_noise_pred_perturbed, a_noise_pred_perturbed = self.unpack_latents(
|
||||
noise_pred_perturbed
|
||||
)
|
||||
finally:
|
||||
self.stg_flag.do_skip = False
|
||||
@@ -229,10 +225,8 @@ class MultimodalGuider(comfy.samplers.CFGGuider):
|
||||
timestep,
|
||||
model_options,
|
||||
)[0]
|
||||
v_noise_pred_modality, a_noise_pred_modality = (
|
||||
self.inner_model.diffusion_model.separate_audio_and_video_latents(
|
||||
noise_pred_modality, audio_length
|
||||
)
|
||||
v_noise_pred_modality, a_noise_pred_modality = self.unpack_latents(
|
||||
noise_pred_modality
|
||||
)
|
||||
finally:
|
||||
del model_options["transformer_options"]["a2v_cross_attn"]
|
||||
@@ -260,9 +254,7 @@ class MultimodalGuider(comfy.samplers.CFGGuider):
|
||||
else:
|
||||
ax = self.last_denoised_a
|
||||
|
||||
x = self.inner_model.diffusion_model.recombine_audio_and_video_latents(
|
||||
vx, ax, x_shape
|
||||
)
|
||||
x, _ = self.pack_latents(vx, ax)
|
||||
|
||||
# normally this would be done in cfg_function, but we skipped
|
||||
# that for efficiency: we can compute the noise predictions in
|
||||
@@ -285,11 +277,7 @@ class MultimodalGuider(comfy.samplers.CFGGuider):
|
||||
}
|
||||
x = fn(args)
|
||||
|
||||
self.last_denoised_v, self.last_denoised_a = (
|
||||
self.inner_model.diffusion_model.separate_audio_and_video_latents(
|
||||
x, audio_length
|
||||
)
|
||||
)
|
||||
self.last_denoised_v, self.last_denoised_a = self.unpack_latents(x)
|
||||
return x
|
||||
|
||||
|
||||
|
||||
@@ -19,6 +19,8 @@ class GuiderParameters:
|
||||
modality_scale: float = 1.0,
|
||||
skip_step: int = 0,
|
||||
cross_attn: bool = True,
|
||||
cfg_zero_star: bool = False,
|
||||
zero_init_sigma: float = 1.0,
|
||||
):
|
||||
self.cfg_scale = cfg_scale
|
||||
self.stg_scale = stg_scale
|
||||
@@ -27,6 +29,8 @@ class GuiderParameters:
|
||||
self.modality_scale = modality_scale
|
||||
self.skip_step = skip_step
|
||||
self.cross_attn = cross_attn
|
||||
self.cfg_zero_star = cfg_zero_star
|
||||
self.zero_init_sigma = zero_init_sigma
|
||||
|
||||
def __str__(self):
|
||||
return f"cfg_scale: {self.cfg_scale}, stg_scale: {self.stg_scale}, rescale_scale: {self.rescale_scale}, modality_scale: {self.modality_scale}"
|
||||
|
||||
@@ -0,0 +1,245 @@
|
||||
"""HDR utilities and ComfyUI node for HDR IC-LoRA inference.
|
||||
|
||||
Provides HDR compression/decompression transforms and a single ComfyUI node:
|
||||
|
||||
- LTXVHDRDecodePostprocess: Decompress VAE output, tonemap for preview,
|
||||
and optionally save raw linear HDR frames as EXR.
|
||||
"""
|
||||
|
||||
import logging
|
||||
import os
|
||||
|
||||
import torch
|
||||
from torch import Tensor
|
||||
|
||||
from .nodes_registry import comfy_node
|
||||
|
||||
logger = logging.getLogger("LTXVideo.hdr")
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# HDR Transform Classes (ported from ltx_core/hdr.py)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class LogC3:
|
||||
"""ARRI LogC3 (EI 800) HDR compression.
|
||||
|
||||
Maps linear [0, inf) -> [0, 1] via the camera log curve, then scales to
|
||||
[-1, 1] for VAE input.
|
||||
"""
|
||||
|
||||
A = 5.555556
|
||||
B = 0.052272
|
||||
C = 0.247190
|
||||
D = 0.385537
|
||||
E = 5.367655
|
||||
F = 0.092809
|
||||
CUT = 0.010591
|
||||
|
||||
def compress(self, hdr: Tensor) -> Tensor:
|
||||
x = torch.clamp(hdr, min=0.0)
|
||||
log_part = self.C * torch.log10(self.A * x + self.B) + self.D
|
||||
lin_part = self.E * x + self.F
|
||||
logc = torch.where(x >= self.CUT, log_part, lin_part)
|
||||
logc = torch.clamp(logc, 0.0, 1.0)
|
||||
return logc * 2.0 - 1.0
|
||||
|
||||
def decompress(self, z: Tensor) -> Tensor:
|
||||
logc = torch.clamp((z + 1.0) / 2.0, 0.0, 1.0)
|
||||
cut_log = self.E * self.CUT + self.F
|
||||
lin_from_log = (torch.pow(10.0, (logc - self.D) / self.C) - self.B) / self.A
|
||||
lin_from_lin = (logc - self.F) / self.E
|
||||
return torch.where(logc >= cut_log, lin_from_log, lin_from_lin)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Transform registry — singleton (stateless)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
_LOGC3 = LogC3()
|
||||
|
||||
|
||||
def _hdr_decompress(decoded_01: Tensor) -> Tensor:
|
||||
"""Decompress VAE-decoded image from [0,1] to linear HDR [0, inf).
|
||||
|
||||
ComfyUI's VAE decode returns images in [0, 1] via ``(raw + 1) / 2``.
|
||||
This function reverses that and applies LogC3 HDR decompression.
|
||||
"""
|
||||
raw = decoded_01.float() * 2.0 - 1.0
|
||||
return _LOGC3.decompress(raw)
|
||||
|
||||
|
||||
def _linear_to_srgb(x: Tensor) -> Tensor:
|
||||
"""Convert linear [0, 1] to sRGB [0, 1]."""
|
||||
return torch.where(
|
||||
x <= 0.0031308,
|
||||
12.92 * x,
|
||||
1.055 * torch.pow(x.clamp(min=0.0031308), 1.0 / 2.4) - 0.055,
|
||||
).clamp(0.0, 1.0)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# ComfyUI Node
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@comfy_node(name="LTXVHDRDecodePostprocess")
|
||||
class LTXVHDRDecodePostprocess:
|
||||
"""Decompress HDR from VAE output, tonemap for preview, optionally save EXR.
|
||||
|
||||
Place after VAE Decode in an HDR IC-LoRA workflow. Recovers linear HDR
|
||||
values from the compressed latent space and tonemaps them to SDR for
|
||||
display. When ``save_exr`` is enabled, also writes the raw linear HDR
|
||||
frames as an EXR image sequence.
|
||||
make sure to set OPENCV_IO_ENABLE_OPENEXR=1 environment in the command line # Must be set before cv2 import
|
||||
|
||||
|
||||
Outputs:
|
||||
tonemapped: SDR preview [0, 1] after Reinhard tonemap + sRGB gamma.
|
||||
hdr_linear: Raw linear HDR values [0, inf) for further processing.
|
||||
"""
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"image": ("IMAGE",),
|
||||
},
|
||||
"optional": {
|
||||
"exposure": (
|
||||
"FLOAT",
|
||||
{
|
||||
"default": 0.0,
|
||||
"min": -10.0,
|
||||
"max": 10.0,
|
||||
"step": 0.1,
|
||||
"display": "slider",
|
||||
"tooltip": (
|
||||
"Exposure in stops (EV). 0 = no change, "
|
||||
"+1 = 2x brighter, -1 = half brightness."
|
||||
),
|
||||
},
|
||||
),
|
||||
"save_exr": (
|
||||
"BOOLEAN",
|
||||
{
|
||||
"default": False,
|
||||
"tooltip": "Save raw linear HDR frames as EXR sequence.",
|
||||
},
|
||||
),
|
||||
"output_dir": (
|
||||
"STRING",
|
||||
{
|
||||
"default": "output/hdr_exr",
|
||||
"tooltip": (
|
||||
"Directory for EXR frames (relative to ComfyUI "
|
||||
"output directory, or absolute path)."
|
||||
),
|
||||
},
|
||||
),
|
||||
"filename_prefix": (
|
||||
"STRING",
|
||||
{"default": "frame"},
|
||||
),
|
||||
"half_precision": (
|
||||
"BOOLEAN",
|
||||
{
|
||||
"default": True,
|
||||
"tooltip": (
|
||||
"Save EXR as float16 (half). Smaller files, "
|
||||
"negligible quality loss for most workflows."
|
||||
),
|
||||
},
|
||||
),
|
||||
},
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("IMAGE", "IMAGE")
|
||||
RETURN_NAMES = ("tonemapped", "hdr_linear")
|
||||
OUTPUT_NODE = True
|
||||
FUNCTION = "postprocess"
|
||||
CATEGORY = "Lightricks/HDR"
|
||||
DESCRIPTION = (
|
||||
"Decompresses VAE-decoded output from HDR IC-LoRA (LogC3) and applies "
|
||||
"Reinhard tonemapping. Place after VAE Decode. 'tonemapped' is the "
|
||||
"SDR preview; 'hdr_linear' is raw linear HDR for downstream use. "
|
||||
"Enable 'save_exr' to write an EXR image sequence."
|
||||
"if save_exr is enabled, make sure to set OPENCV_IO_ENABLE_OPENEXR=1 environment in the command line"
|
||||
)
|
||||
|
||||
def postprocess(
|
||||
self,
|
||||
image: torch.Tensor,
|
||||
exposure: float = 0.0,
|
||||
save_exr: bool = False,
|
||||
output_dir: str = "output/hdr_exr",
|
||||
filename_prefix: str = "frame",
|
||||
half_precision: bool = True,
|
||||
) -> tuple:
|
||||
|
||||
hdr = _hdr_decompress(image)
|
||||
hdr = torch.clamp(hdr, min=0.0, max=1e4)
|
||||
|
||||
# Reinhard tonemap with exposure (stops → linear multiplier)
|
||||
exposure_mult = 2.0**exposure
|
||||
hdr_exposed = hdr * exposure_mult
|
||||
tonemapped_linear = (hdr_exposed / (1.0 + hdr_exposed)).clamp(0.0, 1.0)
|
||||
tonemapped = _linear_to_srgb(tonemapped_linear)
|
||||
|
||||
if save_exr:
|
||||
assert os.environ.get("OPENCV_IO_ENABLE_OPENEXR") == "1", (
|
||||
"EXR output is enabled (save_exr = TRUE), but OpenCV does not support EXR by default. "
|
||||
"To enable it, set the environment variable OPENCV_IO_ENABLE_OPENEXR=1 before starting ComfyUI, then restart. "
|
||||
"Alternatively, disable EXR output or switch to PNG/JPG."
|
||||
)
|
||||
self._save_exr_frames(hdr, output_dir, filename_prefix, half_precision)
|
||||
|
||||
return (tonemapped, hdr)
|
||||
|
||||
@staticmethod
|
||||
def _save_exr_frames(
|
||||
hdr_image: torch.Tensor,
|
||||
output_dir: str,
|
||||
filename_prefix: str,
|
||||
half_precision: bool,
|
||||
) -> None:
|
||||
try:
|
||||
import cv2
|
||||
except ImportError:
|
||||
logger.error(
|
||||
"opencv-python is required for EXR export. "
|
||||
"Install with: pip install opencv-python"
|
||||
)
|
||||
return
|
||||
|
||||
import folder_paths
|
||||
import numpy as np
|
||||
|
||||
if not os.path.isabs(output_dir):
|
||||
output_dir = os.path.join(folder_paths.get_output_directory(), output_dir)
|
||||
|
||||
os.makedirs(output_dir, exist_ok=True)
|
||||
|
||||
frames = hdr_image.cpu().numpy()
|
||||
exr_type = (
|
||||
cv2.IMWRITE_EXR_TYPE_HALF if half_precision else cv2.IMWRITE_EXR_TYPE_FLOAT
|
||||
)
|
||||
params = [
|
||||
cv2.IMWRITE_EXR_TYPE,
|
||||
exr_type,
|
||||
cv2.IMWRITE_EXR_COMPRESSION,
|
||||
cv2.IMWRITE_EXR_COMPRESSION_ZIP,
|
||||
]
|
||||
|
||||
for i in range(frames.shape[0]):
|
||||
frame_bgr = frames[i][:, :, ::-1].astype(np.float32).copy()
|
||||
path = os.path.join(output_dir, f"{filename_prefix}_{i:05d}.exr")
|
||||
cv2.imwrite(path, frame_bgr, params)
|
||||
|
||||
logger.info(
|
||||
"Saved %d EXR frame(s) to %s (%s)",
|
||||
frames.shape[0],
|
||||
output_dir,
|
||||
"float16" if half_precision else "float32",
|
||||
)
|
||||
@@ -0,0 +1,603 @@
|
||||
import logging
|
||||
|
||||
import comfy
|
||||
import comfy.samplers
|
||||
import comfy.sd
|
||||
import comfy.utils
|
||||
import comfy_extras.nodes_lt as nodes_lt
|
||||
import folder_paths
|
||||
import node_helpers
|
||||
import torch
|
||||
from comfy_api.latest import io
|
||||
|
||||
from .latents import LTXVDilateLatent
|
||||
from .nodes_registry import NODES_DISPLAY_NAME_PREFIX, comfy_node
|
||||
|
||||
|
||||
@comfy_node(name="LTXAddVideoICLoRAGuide")
|
||||
class LTXAddVideoICLoRAGuide(io.ComfyNode):
|
||||
PATCHIFIER = nodes_lt.SymmetricPatchifier(1, start_end=True)
|
||||
|
||||
@classmethod
|
||||
def define_schema(cls):
|
||||
return io.Schema(
|
||||
node_id="LTXAddVideoICLoRAGuide",
|
||||
display_name=NODES_DISPLAY_NAME_PREFIX + " Add Video IC-LoRA Guide",
|
||||
category="Lightricks/IC-LoRA",
|
||||
description=(
|
||||
"Adds one or more conditioning frames starting at the specified frame index. "
|
||||
"Supports both single images and multi-frame videos. "
|
||||
"The latent_downscale_factor resizes input to a fraction of the target size "
|
||||
"(1 = original, 2 = half, 3 = third, etc.) for IC-LoRA on small grids."
|
||||
),
|
||||
inputs=[
|
||||
io.Conditioning.Input("positive"),
|
||||
io.Conditioning.Input("negative"),
|
||||
io.Vae.Input("vae"),
|
||||
io.Latent.Input(
|
||||
"latent",
|
||||
tooltip="Video-only latent to condition. Must be a 5D video latent (batch, channels, frames, height, width).",
|
||||
),
|
||||
io.Image.Input("image"),
|
||||
io.Int.Input(
|
||||
"frame_idx",
|
||||
default=0,
|
||||
min=-9999,
|
||||
max=9999,
|
||||
tooltip=(
|
||||
"Frame index to start the conditioning at. "
|
||||
"For single-frame videos, any frame_idx value is acceptable. "
|
||||
"For videos, frame_idx must be 1 modulo 8, otherwise it will be rounded "
|
||||
"down to the nearest 1 modulo 8. Negative values are counted from the end of the video."
|
||||
),
|
||||
),
|
||||
io.Float.Input(
|
||||
"strength",
|
||||
default=1.0,
|
||||
min=0.0,
|
||||
max=1.0,
|
||||
step=0.01,
|
||||
),
|
||||
io.Float.Input(
|
||||
"latent_downscale_factor",
|
||||
default=1.0,
|
||||
min=1.0,
|
||||
max=10.0,
|
||||
step=1.0,
|
||||
tooltip="For IC-LoRA on small grid. 1 means original size, 2 means half size, 3 means third, etc.",
|
||||
),
|
||||
io.Combo.Input(
|
||||
"crop",
|
||||
options=["disabled", "center"],
|
||||
default="disabled",
|
||||
tooltip="Crop mode when resizing. 'center' crops to fit, 'disabled' stretches to fit.",
|
||||
),
|
||||
io.Boolean.Input(
|
||||
"use_tiled_encode",
|
||||
default=False,
|
||||
tooltip="Enable tiled VAE encoding for large resolutions/long videos to reduce memory usage.",
|
||||
),
|
||||
io.Int.Input(
|
||||
"tile_size",
|
||||
default=256,
|
||||
min=64,
|
||||
max=512,
|
||||
step=32,
|
||||
tooltip="Spatial tile size for tiled encoding. Only used when use_tiled_encode is enabled.",
|
||||
),
|
||||
io.Int.Input(
|
||||
"tile_overlap",
|
||||
default=64,
|
||||
min=16,
|
||||
max=256,
|
||||
step=16,
|
||||
tooltip="Overlap between tiles for tiled encoding. Only used when use_tiled_encode is enabled.",
|
||||
),
|
||||
],
|
||||
outputs=[
|
||||
io.Conditioning.Output("positive"),
|
||||
io.Conditioning.Output("negative"),
|
||||
io.Latent.Output("latent"),
|
||||
],
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def encode(
|
||||
cls,
|
||||
vae,
|
||||
latent_width,
|
||||
latent_height,
|
||||
images,
|
||||
scale_factors,
|
||||
latent_downscale_factor,
|
||||
crop,
|
||||
use_tiled_encode,
|
||||
tile_size,
|
||||
tile_overlap,
|
||||
):
|
||||
time_scale_factor, width_scale_factor, height_scale_factor = scale_factors
|
||||
num_frames_to_keep = (
|
||||
(images.shape[0] - 1) // time_scale_factor
|
||||
) * time_scale_factor + 1
|
||||
images = images[:num_frames_to_keep]
|
||||
# Divide target size by latent_downscale_factor
|
||||
target_width = int(latent_width * width_scale_factor / latent_downscale_factor)
|
||||
target_height = int(
|
||||
latent_height * height_scale_factor / latent_downscale_factor
|
||||
)
|
||||
pixels = comfy.utils.common_upscale(
|
||||
images.movedim(-1, 1),
|
||||
target_width,
|
||||
target_height,
|
||||
"bilinear",
|
||||
crop=crop,
|
||||
).movedim(1, -1)
|
||||
encode_pixels = pixels[:, :, :, :3]
|
||||
if use_tiled_encode:
|
||||
guide_latent = vae.encode_tiled(
|
||||
encode_pixels,
|
||||
tile_x=tile_size,
|
||||
tile_y=tile_size,
|
||||
overlap=tile_overlap,
|
||||
)
|
||||
else:
|
||||
guide_latent = vae.encode(encode_pixels)
|
||||
return encode_pixels, guide_latent
|
||||
|
||||
@classmethod
|
||||
def execute(
|
||||
cls,
|
||||
positive,
|
||||
negative,
|
||||
vae,
|
||||
latent,
|
||||
image,
|
||||
frame_idx,
|
||||
strength,
|
||||
latent_downscale_factor,
|
||||
crop,
|
||||
use_tiled_encode,
|
||||
tile_size,
|
||||
tile_overlap,
|
||||
) -> io.NodeOutput:
|
||||
scale_factors = vae.downscale_index_formula
|
||||
latent_image = latent["samples"]
|
||||
noise_mask = nodes_lt.get_noise_mask(latent)
|
||||
|
||||
_, _, latent_length, latent_height, latent_width = latent_image.shape
|
||||
|
||||
time_scale_factor = scale_factors[0]
|
||||
num_frames_to_keep = (
|
||||
(image.shape[0] - 1) // time_scale_factor
|
||||
) * time_scale_factor + 1
|
||||
causal_fix = frame_idx == 0 or num_frames_to_keep == 1
|
||||
if not causal_fix:
|
||||
image = torch.cat([image[:1], image], dim=0)
|
||||
|
||||
image, guide_latent = cls.encode(
|
||||
vae,
|
||||
latent_width,
|
||||
latent_height,
|
||||
image,
|
||||
scale_factors,
|
||||
latent_downscale_factor,
|
||||
crop,
|
||||
use_tiled_encode,
|
||||
tile_size,
|
||||
tile_overlap,
|
||||
)
|
||||
|
||||
if not causal_fix:
|
||||
guide_latent = guide_latent[:, :, 1:, :, :]
|
||||
image = image[1:]
|
||||
|
||||
# Record original (pre-dilation) guide latent shape for spatial mask downsampling
|
||||
guide_orig_shape = list(guide_latent.shape[2:]) # [F, H_small, W_small]
|
||||
|
||||
guide_mask = None
|
||||
|
||||
# Dilate the latent if latent_downscale_factor > 1
|
||||
if latent_downscale_factor > 1:
|
||||
if (
|
||||
latent_width % latent_downscale_factor != 0
|
||||
or latent_height % latent_downscale_factor != 0
|
||||
):
|
||||
raise ValueError(
|
||||
f"Latent spatial size {latent_width}x{latent_height} must be divisible by latent_downscale_factor {latent_downscale_factor}"
|
||||
)
|
||||
|
||||
dilated = LTXVDilateLatent().dilate_latent(
|
||||
{"samples": guide_latent},
|
||||
horizontal_scale=int(latent_downscale_factor),
|
||||
vertical_scale=int(latent_downscale_factor),
|
||||
)[0]
|
||||
guide_mask = dilated["noise_mask"]
|
||||
guide_latent = dilated["samples"]
|
||||
|
||||
# Pre-filter token count = product of (potentially dilated) spatial dims
|
||||
iclora_tokens_added = (
|
||||
guide_latent.shape[2] * guide_latent.shape[3] * guide_latent.shape[4]
|
||||
)
|
||||
|
||||
frame_idx, latent_idx = nodes_lt.LTXVAddGuide.get_latent_index(
|
||||
positive, latent_length, len(image), frame_idx, scale_factors
|
||||
)
|
||||
assert (
|
||||
latent_idx + guide_latent.shape[2] <= latent_length
|
||||
), "Conditioning frames exceed the length of the latent sequence."
|
||||
|
||||
positive, negative, latent_image, noise_mask = (
|
||||
nodes_lt.LTXVAddGuide.append_keyframe(
|
||||
positive,
|
||||
negative,
|
||||
frame_idx,
|
||||
latent_image,
|
||||
noise_mask,
|
||||
guide_latent,
|
||||
strength,
|
||||
scale_factors,
|
||||
guide_mask=guide_mask,
|
||||
latent_downscale_factor=latent_downscale_factor,
|
||||
causal_fix=causal_fix,
|
||||
)
|
||||
)
|
||||
|
||||
# Track this guide in guide_attention_entries for per-reference attention control.
|
||||
from .iclora_attention import append_guide_attention_entry
|
||||
|
||||
positive = append_guide_attention_entry(
|
||||
positive, iclora_tokens_added, guide_orig_shape
|
||||
)
|
||||
negative = append_guide_attention_entry(
|
||||
negative, iclora_tokens_added, guide_orig_shape
|
||||
)
|
||||
|
||||
return io.NodeOutput(
|
||||
positive, negative, {"samples": latent_image, "noise_mask": noise_mask}
|
||||
)
|
||||
|
||||
|
||||
@comfy_node(name="LTXAddVideoICLoRAGuideAdvanced")
|
||||
class LTXAddVideoICLoRAGuideAdvanced(LTXAddVideoICLoRAGuide):
|
||||
"""Extended IC-LoRA guide node with per-guide attention strength control.
|
||||
|
||||
Same as LTXAddVideoICLoRAGuide, but adds attention_strength and
|
||||
attention_mask inputs to control how strongly this guide's conditioning
|
||||
influences generation via self-attention.
|
||||
"""
|
||||
|
||||
@classmethod
|
||||
def define_schema(cls):
|
||||
return io.Schema(
|
||||
node_id="LTXAddVideoICLoRAGuideAdvanced",
|
||||
display_name=NODES_DISPLAY_NAME_PREFIX
|
||||
+ " Add Video IC-LoRA Guide Advanced",
|
||||
category="Lightricks/IC-LoRA",
|
||||
description=(
|
||||
"Adds IC-LoRA guide conditioning with per-guide attention strength control. "
|
||||
"Same as LTXAddVideoICLoRAGuide, but allows controlling how strongly this "
|
||||
"guide influences generation via self-attention, optionally with a spatial mask."
|
||||
),
|
||||
inputs=[
|
||||
io.Conditioning.Input("positive"),
|
||||
io.Conditioning.Input("negative"),
|
||||
io.Vae.Input("vae"),
|
||||
io.Latent.Input(
|
||||
"latent",
|
||||
tooltip="Video-only latent to condition. Must be a 5D video latent.",
|
||||
),
|
||||
io.Image.Input("image"),
|
||||
io.Int.Input(
|
||||
"frame_idx",
|
||||
default=0,
|
||||
min=-9999,
|
||||
max=9999,
|
||||
tooltip=(
|
||||
"Frame index to start the conditioning at. The value is rounded to the "
|
||||
"nearest frame and wrapped modulo the number of video frames. Negative "
|
||||
"values are counted from the end of the video before wrapping."
|
||||
),
|
||||
),
|
||||
io.Float.Input("strength", default=1.0, min=0.0, max=1.0, step=0.01),
|
||||
io.Float.Input(
|
||||
"latent_downscale_factor",
|
||||
default=1.0,
|
||||
min=1.0,
|
||||
max=10.0,
|
||||
step=1.0,
|
||||
tooltip="For IC-LoRA on small grid. 1 = original size, 2 = half, etc.",
|
||||
),
|
||||
io.Combo.Input(
|
||||
"crop",
|
||||
options=["disabled", "center"],
|
||||
default="disabled",
|
||||
),
|
||||
io.Boolean.Input("use_tiled_encode", default=False),
|
||||
io.Int.Input("tile_size", default=256, min=64, max=512, step=32),
|
||||
io.Int.Input("tile_overlap", default=64, min=16, max=256, step=16),
|
||||
io.Float.Input(
|
||||
"attention_strength",
|
||||
default=1.0,
|
||||
min=0.0,
|
||||
max=1.0,
|
||||
step=0.01,
|
||||
tooltip=(
|
||||
"Controls how strongly this guide influences generation via "
|
||||
"self-attention. 1.0 = full conditioning (default), 0.0 = ignore. "
|
||||
"When an attention_mask is also provided, this multiplies the mask values."
|
||||
),
|
||||
),
|
||||
io.Mask.Input(
|
||||
"attention_mask",
|
||||
optional=True,
|
||||
tooltip=(
|
||||
"Optional pixel-space spatial mask. Shape (F, H, W) or (H, W). "
|
||||
"Values in [0, 1]. Controls per-region conditioning influence. "
|
||||
"Multiplied by attention_strength."
|
||||
),
|
||||
),
|
||||
],
|
||||
outputs=[
|
||||
io.Conditioning.Output("positive"),
|
||||
io.Conditioning.Output("negative"),
|
||||
io.Latent.Output("latent"),
|
||||
],
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def execute(
|
||||
cls,
|
||||
positive,
|
||||
negative,
|
||||
vae,
|
||||
latent,
|
||||
image,
|
||||
frame_idx,
|
||||
strength,
|
||||
latent_downscale_factor,
|
||||
crop,
|
||||
use_tiled_encode,
|
||||
tile_size,
|
||||
tile_overlap,
|
||||
attention_strength=1.0,
|
||||
attention_mask=None,
|
||||
) -> io.NodeOutput:
|
||||
from .iclora_attention import normalize_mask
|
||||
from .latents import LTXVDilateLatent
|
||||
|
||||
scale_factors = vae.downscale_index_formula
|
||||
latent_image = latent["samples"]
|
||||
noise_mask = nodes_lt.get_noise_mask(latent)
|
||||
|
||||
_, _, latent_length, latent_height, latent_width = latent_image.shape
|
||||
|
||||
time_scale_factor = scale_factors[0]
|
||||
num_frames_to_keep = (
|
||||
(image.shape[0] - 1) // time_scale_factor
|
||||
) * time_scale_factor + 1
|
||||
causal_fix = frame_idx == 0 or num_frames_to_keep == 1
|
||||
if not causal_fix:
|
||||
image = torch.cat([image[:1], image], dim=0)
|
||||
|
||||
image, guide_latent = cls.encode(
|
||||
vae,
|
||||
latent_width,
|
||||
latent_height,
|
||||
image,
|
||||
scale_factors,
|
||||
latent_downscale_factor,
|
||||
crop,
|
||||
use_tiled_encode,
|
||||
tile_size,
|
||||
tile_overlap,
|
||||
)
|
||||
|
||||
if not causal_fix:
|
||||
guide_latent = guide_latent[:, :, 1:, :, :]
|
||||
image = image[1:]
|
||||
|
||||
guide_orig_shape = list(guide_latent.shape[2:])
|
||||
guide_mask = None
|
||||
|
||||
if latent_downscale_factor > 1:
|
||||
if (
|
||||
latent_width % latent_downscale_factor != 0
|
||||
or latent_height % latent_downscale_factor != 0
|
||||
):
|
||||
raise ValueError(
|
||||
f"Latent spatial size {latent_width}x{latent_height} must be "
|
||||
f"divisible by latent_downscale_factor {latent_downscale_factor}"
|
||||
)
|
||||
dilated = LTXVDilateLatent().dilate_latent(
|
||||
{"samples": guide_latent},
|
||||
horizontal_scale=int(latent_downscale_factor),
|
||||
vertical_scale=int(latent_downscale_factor),
|
||||
)[0]
|
||||
guide_mask = dilated["noise_mask"]
|
||||
guide_latent = dilated["samples"]
|
||||
|
||||
iclora_tokens_added = (
|
||||
guide_latent.shape[2] * guide_latent.shape[3] * guide_latent.shape[4]
|
||||
)
|
||||
|
||||
frame_idx, latent_idx = nodes_lt.LTXVAddGuide.get_latent_index(
|
||||
positive, latent_length, len(image), frame_idx, scale_factors
|
||||
)
|
||||
assert (
|
||||
latent_idx + guide_latent.shape[2] <= latent_length
|
||||
), "Conditioning frames exceed the length of the latent sequence."
|
||||
|
||||
positive, negative, latent_image, noise_mask = (
|
||||
nodes_lt.LTXVAddGuide.append_keyframe(
|
||||
positive,
|
||||
negative,
|
||||
frame_idx,
|
||||
latent_image,
|
||||
noise_mask,
|
||||
guide_latent,
|
||||
strength,
|
||||
scale_factors,
|
||||
guide_mask=guide_mask,
|
||||
latent_downscale_factor=latent_downscale_factor,
|
||||
causal_fix=causal_fix,
|
||||
)
|
||||
)
|
||||
|
||||
from .iclora_attention import append_guide_attention_entry
|
||||
|
||||
norm_mask = normalize_mask(attention_mask)
|
||||
positive = append_guide_attention_entry(
|
||||
positive,
|
||||
iclora_tokens_added,
|
||||
guide_orig_shape,
|
||||
attention_strength=attention_strength,
|
||||
attention_mask=norm_mask,
|
||||
)
|
||||
negative = append_guide_attention_entry(
|
||||
negative,
|
||||
iclora_tokens_added,
|
||||
guide_orig_shape,
|
||||
attention_strength=attention_strength,
|
||||
attention_mask=norm_mask,
|
||||
)
|
||||
|
||||
return io.NodeOutput(
|
||||
positive, negative, {"samples": latent_image, "noise_mask": noise_mask}
|
||||
)
|
||||
|
||||
|
||||
@comfy_node(name="LTXICLoRALoaderModelOnly")
|
||||
class LTXICLoRALoaderModelOnly(io.ComfyNode):
|
||||
@classmethod
|
||||
def define_schema(cls):
|
||||
return io.Schema(
|
||||
node_id="LTXICLoRALoaderModelOnly",
|
||||
display_name=NODES_DISPLAY_NAME_PREFIX + " IC-LoRA Loader Model Only",
|
||||
category="Lightricks/IC-LoRA",
|
||||
description="Loads a LoRA model and extracts the latent_downscale_factor from the safetensors metadata.",
|
||||
inputs=[
|
||||
io.Model.Input("model"),
|
||||
io.Combo.Input(
|
||||
"lora_name",
|
||||
options=folder_paths.get_filename_list("loras"),
|
||||
),
|
||||
io.Float.Input(
|
||||
"strength_model",
|
||||
default=1.0,
|
||||
min=-100.0,
|
||||
max=100.0,
|
||||
step=0.01,
|
||||
),
|
||||
],
|
||||
outputs=[
|
||||
io.Model.Output("model"),
|
||||
io.Float.Output("latent_downscale_factor"),
|
||||
],
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def execute(cls, model, lora_name, strength_model) -> io.NodeOutput:
|
||||
lora_path = folder_paths.get_full_path_or_raise("loras", lora_name)
|
||||
|
||||
# Load lora and extract metadata
|
||||
lora, metadata = comfy.utils.load_torch_file(
|
||||
lora_path, safe_load=True, return_metadata=True
|
||||
)
|
||||
|
||||
# Extract latent_downscale_factor from metadata
|
||||
try:
|
||||
latent_downscale_factor = float(metadata["reference_downscale_factor"])
|
||||
except (KeyError, ValueError, TypeError):
|
||||
latent_downscale_factor = 1.0
|
||||
logging.warning(
|
||||
"Failed to extract reference_downscale_factor from metadata for %s, using 1.0",
|
||||
lora_path,
|
||||
)
|
||||
|
||||
if strength_model == 0:
|
||||
return io.NodeOutput(model, latent_downscale_factor)
|
||||
|
||||
model_lora, _ = comfy.sd.load_lora_for_models(
|
||||
model, None, lora, strength_model, 0
|
||||
)
|
||||
return io.NodeOutput(model_lora, latent_downscale_factor)
|
||||
|
||||
|
||||
def _patchify_audio_latent(latent):
|
||||
"""Reshape audio latent (b, c, t, f) → ref_audio token dict (b, t, c*f)."""
|
||||
b, c, t, f = latent.shape
|
||||
ref_tokens = latent.permute(0, 2, 1, 3).reshape(b, t, c * f)
|
||||
return {"tokens": ref_tokens}
|
||||
|
||||
|
||||
@comfy_node(name="LTXVSetAudioRefTokens")
|
||||
class LTXVSetAudioRefTokens(io.ComfyNode):
|
||||
"""Provide speaker identity context for audio generation.
|
||||
|
||||
Patchifies the audio latent and attaches it as reference tokens on both
|
||||
positive and negative conditioning. The model prepends these tokens with
|
||||
negative temporal positions so they serve as identity context without
|
||||
being part of the generated output.
|
||||
|
||||
Also outputs a frozen copy of the audio latent (noise_mask=0) for
|
||||
direct use in stage 2 without re-encoding.
|
||||
"""
|
||||
|
||||
@classmethod
|
||||
def define_schema(cls):
|
||||
return io.Schema(
|
||||
node_id="LTXVSetAudioRefTokens",
|
||||
display_name=NODES_DISPLAY_NAME_PREFIX + " Set Audio Ref Tokens",
|
||||
category="Lightricks/IC-LoRA",
|
||||
description=(
|
||||
"Provides speaker identity context for audio generation by attaching "
|
||||
"reference audio tokens to the conditioning. The tokens are prepended "
|
||||
"with negative temporal positions so the model treats them as context "
|
||||
"rather than generation targets."
|
||||
),
|
||||
inputs=[
|
||||
io.Conditioning.Input(
|
||||
"positive",
|
||||
tooltip="Positive conditioning to attach the reference audio tokens to.",
|
||||
),
|
||||
io.Conditioning.Input(
|
||||
"negative",
|
||||
tooltip="Negative conditioning to attach the reference audio tokens to.",
|
||||
),
|
||||
io.Latent.Input(
|
||||
"audio_latent",
|
||||
tooltip="Encoded audio latent from LTXV Audio VAE Encode.",
|
||||
),
|
||||
],
|
||||
outputs=[
|
||||
io.Conditioning.Output(
|
||||
display_name="positive",
|
||||
tooltip="Positive conditioning with reference audio tokens attached.",
|
||||
),
|
||||
io.Conditioning.Output(
|
||||
display_name="negative",
|
||||
tooltip="Negative conditioning with reference audio tokens attached.",
|
||||
),
|
||||
io.Latent.Output(
|
||||
display_name="frozen_audio",
|
||||
tooltip="Audio latent with noise_mask=0, fully frozen during denoising.",
|
||||
),
|
||||
],
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def execute(cls, positive, negative, audio_latent) -> io.NodeOutput:
|
||||
latent = audio_latent["samples"]
|
||||
ref_audio = _patchify_audio_latent(latent)
|
||||
positive = node_helpers.conditioning_set_values(
|
||||
positive, {"ref_audio": ref_audio}
|
||||
)
|
||||
negative = node_helpers.conditioning_set_values(
|
||||
negative, {"ref_audio": ref_audio}
|
||||
)
|
||||
|
||||
frozen = audio_latent.copy()
|
||||
b, c, t, f = latent.shape
|
||||
frozen["noise_mask"] = torch.zeros((b, 1, t, 1), dtype=torch.float32)
|
||||
|
||||
return io.NodeOutput(positive, negative, frozen)
|
||||
@@ -0,0 +1,67 @@
|
||||
"""Helpers for per-guide attention strength tracking in conditioning metadata."""
|
||||
|
||||
|
||||
def _get_guide_attention_entries(conditioning):
|
||||
"""Read the current guide_attention_entries list from conditioning."""
|
||||
for t in conditioning:
|
||||
entries = t[1].get("guide_attention_entries", None)
|
||||
if entries is not None:
|
||||
return entries
|
||||
return []
|
||||
|
||||
|
||||
def _set_guide_attention_entries(conditioning, entries):
|
||||
"""Write guide_attention_entries into conditioning (immutable update)."""
|
||||
import node_helpers
|
||||
|
||||
return node_helpers.conditioning_set_values(
|
||||
conditioning, {"guide_attention_entries": entries}
|
||||
)
|
||||
|
||||
|
||||
def append_guide_attention_entry(
|
||||
conditioning,
|
||||
pre_filter_count,
|
||||
latent_shape,
|
||||
attention_strength=1.0,
|
||||
attention_mask=None,
|
||||
):
|
||||
"""Append a new guide attention entry to conditioning metadata.
|
||||
|
||||
Called by guide-adding nodes after appending tokens.
|
||||
|
||||
Args:
|
||||
conditioning: ComfyUI conditioning list.
|
||||
pre_filter_count: Token count for this guide (before grid_mask filtering).
|
||||
latent_shape: [F, H, W] of the pre-dilation guide latent.
|
||||
attention_strength: Scalar in [0, 1]. 1.0 = full attention (default).
|
||||
attention_mask: Optional pixel-space mask tensor, shape (1, 1, F, H, W).
|
||||
|
||||
Returns:
|
||||
Updated conditioning.
|
||||
"""
|
||||
existing_entries = _get_guide_attention_entries(conditioning)
|
||||
entries = [*existing_entries]
|
||||
entries.append(
|
||||
{
|
||||
"pre_filter_count": pre_filter_count,
|
||||
"strength": attention_strength,
|
||||
"pixel_mask": attention_mask,
|
||||
"latent_shape": latent_shape,
|
||||
}
|
||||
)
|
||||
return _set_guide_attention_entries(conditioning, entries)
|
||||
|
||||
|
||||
def normalize_mask(mask):
|
||||
"""Normalize a ComfyUI MASK to (1, 1, F, H, W) for downstream processing.
|
||||
|
||||
ComfyUI MASK type is typically (B, H, W) or (H, W).
|
||||
"""
|
||||
if mask is None:
|
||||
return None
|
||||
if mask.dim() == 2: # (H, W) → single frame
|
||||
return mask.unsqueeze(0).unsqueeze(0).unsqueeze(0)
|
||||
elif mask.dim() == 3: # (F, H, W) → video mask
|
||||
return mask.unsqueeze(0).unsqueeze(0)
|
||||
return mask
|
||||
+29
@@ -438,6 +438,9 @@ class LTXVAddLatentGuide:
|
||||
latent = latent["samples"]
|
||||
guide = guiding_latent["samples"]
|
||||
|
||||
# Record original (pre-dilation) guide latent shape for spatial mask downsampling
|
||||
guide_orig_shape = list(guide.shape[2:]) # [F, H_small, W_small]
|
||||
|
||||
assert (
|
||||
latent.shape[4] % guide.shape[4] == 0
|
||||
and latent.shape[3] % guide.shape[3] == 0
|
||||
@@ -452,6 +455,10 @@ class LTXVAddLatentGuide:
|
||||
guide = guiding_latent["samples"]
|
||||
guide_mask = guiding_latent.get("noise_mask", None)
|
||||
|
||||
# Pre-filter token count = product of dilated spatial dims
|
||||
# (before grid_mask filtering removes padding positions)
|
||||
iclora_tokens_added = guide.shape[2] * guide.shape[3] * guide.shape[4]
|
||||
|
||||
scale_factors = vae.downscale_index_formula
|
||||
|
||||
if latent_idx <= 0:
|
||||
@@ -471,6 +478,16 @@ class LTXVAddLatentGuide:
|
||||
guide_mask=guide_mask,
|
||||
)
|
||||
|
||||
# Track this guide in guide_attention_entries for per-reference attention control.
|
||||
from .iclora_attention import append_guide_attention_entry
|
||||
|
||||
positive = append_guide_attention_entry(
|
||||
positive, iclora_tokens_added, guide_orig_shape
|
||||
)
|
||||
negative = append_guide_attention_entry(
|
||||
negative, iclora_tokens_added, guide_orig_shape
|
||||
)
|
||||
|
||||
return (
|
||||
positive,
|
||||
negative,
|
||||
@@ -809,6 +826,18 @@ class LTXVSetAudioVideoMaskByTime:
|
||||
:, :, audio_latent_frame_index_start:audio_latent_frame_index_end
|
||||
] = 1.0
|
||||
|
||||
if "noise_mask" in av_latent:
|
||||
base_mask = av_latent["noise_mask"].tensors[0].clone()
|
||||
if (
|
||||
base_mask.shape[0]
|
||||
== base_mask.shape[1]
|
||||
== 1
|
||||
== base_mask.shape[3]
|
||||
== base_mask.shape[4]
|
||||
):
|
||||
for frame in range(base_mask.shape[2]):
|
||||
video_mask[:, :, frame, :, :] *= base_mask[0, 0, frame, 0, 0]
|
||||
|
||||
av_latent["noise_mask"] = NestedTensor(
|
||||
ltxav.recombine_audio_and_video_latents(video_mask, audio_mask)
|
||||
)
|
||||
|
||||
@@ -1,4 +1,5 @@
|
||||
import re
|
||||
from functools import wraps
|
||||
from typing import Callable, Optional, Type
|
||||
|
||||
NODE_CLASS_MAPPINGS = {}
|
||||
@@ -36,6 +37,29 @@ def register_node(node_class: Type, name: str, description: str) -> None:
|
||||
NODE_DISPLAY_NAME_MAPPINGS[name] = description
|
||||
|
||||
|
||||
def _is_v3_node(node_class: Type) -> bool:
|
||||
"""Check if the node class is a v3 node (has define_schema method)."""
|
||||
return hasattr(node_class, "define_schema") and callable(
|
||||
getattr(node_class, "define_schema")
|
||||
)
|
||||
|
||||
|
||||
def _wrap_define_schema(node_class: Type, display_name: str) -> None:
|
||||
"""Wrap the define_schema method to inject the display_name into the schema."""
|
||||
original_define_schema = node_class.define_schema
|
||||
|
||||
@classmethod
|
||||
@wraps(original_define_schema.__func__)
|
||||
def wrapped_define_schema(cls):
|
||||
schema = original_define_schema.__func__(cls)
|
||||
# Only set display_name if not already set in the schema
|
||||
if schema.display_name is None:
|
||||
schema.display_name = display_name
|
||||
return schema
|
||||
|
||||
node_class.define_schema = wrapped_define_schema
|
||||
|
||||
|
||||
def comfy_node(
|
||||
node_class: Optional[Type] = None,
|
||||
name: Optional[str] = None,
|
||||
@@ -82,6 +106,10 @@ def comfy_node(
|
||||
|
||||
description = _format_description(description, name, experimental, deprecated)
|
||||
|
||||
# For v3 nodes, wrap define_schema to inject the display_name
|
||||
if _is_v3_node(node_class):
|
||||
_wrap_define_schema(node_class, description)
|
||||
|
||||
register_node(node_class, name, description)
|
||||
return node_class
|
||||
|
||||
|
||||
@@ -0,0 +1,267 @@
|
||||
import math
|
||||
|
||||
import comfy.model_management
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
from comfy_api.latest import io
|
||||
from kornia.geometry.transform.pyramid import (
|
||||
PyrUp,
|
||||
build_laplacian_pyramid,
|
||||
build_pyramid,
|
||||
find_next_powerof_two,
|
||||
is_powerof_two,
|
||||
pad,
|
||||
)
|
||||
from torch import Tensor
|
||||
|
||||
from .nodes_registry import comfy_node
|
||||
|
||||
_CHUNK_SIZE = 8
|
||||
_MASK_LOW_RES_LONG_SIDE = 64
|
||||
|
||||
|
||||
def _pad_for_laplacian(image: torch.Tensor) -> tuple[torch.Tensor, tuple[int, int]]:
|
||||
"""Pad an image tensor so H and W are powers of two (kornia requirement)."""
|
||||
h, w = image.shape[2], image.shape[3]
|
||||
pad_right = 0
|
||||
pad_down = 0
|
||||
if not (is_powerof_two(h) and is_powerof_two(w)):
|
||||
pad_right = find_next_powerof_two(w) - w
|
||||
pad_down = find_next_powerof_two(h) - h
|
||||
image = pad(image, (0, pad_right, 0, pad_down), "reflect")
|
||||
return image, (pad_right, pad_down)
|
||||
|
||||
|
||||
def _gaussian_pyramid(
|
||||
images: torch.Tensor,
|
||||
max_level: int,
|
||||
border_type: str = "reflect",
|
||||
align_corners: bool = False,
|
||||
) -> list[Tensor]:
|
||||
h, w = images.shape[2], images.shape[3]
|
||||
if not (is_powerof_two(w) and is_powerof_two(h)):
|
||||
padding = (0, find_next_powerof_two(w) - w, 0, find_next_powerof_two(h) - h)
|
||||
images = pad(images, padding, border_type)
|
||||
return build_pyramid(images, max_level, border_type, align_corners)
|
||||
|
||||
|
||||
def _resize_preserving_aspect_ratio(
|
||||
images: torch.Tensor, long_side: int, mode: str
|
||||
) -> torch.Tensor:
|
||||
h, w = images.shape[-2:]
|
||||
current_long_side = max(h, w)
|
||||
if current_long_side == long_side:
|
||||
return images
|
||||
|
||||
scale = long_side / current_long_side
|
||||
resized_h = max(1, int(round(h * scale)))
|
||||
resized_w = max(1, int(round(w * scale)))
|
||||
|
||||
if mode == "nearest":
|
||||
return F.interpolate(images, size=(resized_h, resized_w), mode=mode)
|
||||
|
||||
return F.interpolate(
|
||||
images,
|
||||
size=(resized_h, resized_w),
|
||||
mode=mode,
|
||||
align_corners=False,
|
||||
)
|
||||
|
||||
|
||||
def _apply_low_res_mask_dilation(
|
||||
mask: torch.Tensor,
|
||||
spatial_radius: int,
|
||||
long_side: int = _MASK_LOW_RES_LONG_SIDE,
|
||||
) -> torch.Tensor:
|
||||
if spatial_radius <= 0:
|
||||
return mask
|
||||
|
||||
original_size = mask.shape[-2:]
|
||||
mask_low_res = _resize_preserving_aspect_ratio(
|
||||
mask.float(), long_side, mode="bilinear"
|
||||
)
|
||||
mask_low_res = F.max_pool2d(
|
||||
mask_low_res,
|
||||
kernel_size=spatial_radius * 2 + 1,
|
||||
stride=1,
|
||||
padding=spatial_radius,
|
||||
)
|
||||
return F.interpolate(
|
||||
mask_low_res,
|
||||
size=original_size,
|
||||
mode="bilinear",
|
||||
align_corners=False,
|
||||
)
|
||||
|
||||
|
||||
def _pyramid_blend_chunk(
|
||||
image1: torch.Tensor,
|
||||
image2: torch.Tensor,
|
||||
mask: torch.Tensor,
|
||||
max_level: int = 7,
|
||||
) -> torch.Tensor:
|
||||
"""Blend a single chunk (already padded, already on device)."""
|
||||
pyramid1 = build_laplacian_pyramid(image1, max_level=max_level)
|
||||
pyramid2 = build_laplacian_pyramid(image2, max_level=max_level)
|
||||
pyramid_mask = _gaussian_pyramid(mask, max_level=max_level)
|
||||
pyr_up = PyrUp()
|
||||
|
||||
output = pyramid1[-1] * pyramid_mask[-1] + pyramid2[-1] * (1 - pyramid_mask[-1])
|
||||
for i in range(len(pyramid1) - 2, -1, -1):
|
||||
residual = pyramid1[i] * pyramid_mask[i] + pyramid2[i] * (1 - pyramid_mask[i])
|
||||
output = pyr_up(output) + residual
|
||||
|
||||
return output
|
||||
|
||||
|
||||
def _pyramid_blend(
|
||||
image1: torch.Tensor,
|
||||
image2: torch.Tensor,
|
||||
mask: torch.Tensor,
|
||||
max_level: int = 7,
|
||||
device: torch.device | None = None,
|
||||
output_device: torch.device | None = None,
|
||||
) -> torch.Tensor:
|
||||
if image1.shape != image2.shape:
|
||||
raise ValueError(
|
||||
f"input images must have the same size, {image1.shape} != {image2.shape}"
|
||||
)
|
||||
if image1.shape[0] != mask.shape[0]:
|
||||
raise ValueError(
|
||||
"image_a, image_b, and mask must have the same frame count for blending"
|
||||
)
|
||||
if image1.shape[-2:] != mask.shape[-2:]:
|
||||
raise ValueError(
|
||||
"image_a, image_b, and mask must have the same spatial resolution for blending"
|
||||
)
|
||||
|
||||
_, padding = _pad_for_laplacian(image1[:1])
|
||||
orig_h, orig_w = image1.shape[-2], image1.shape[-1]
|
||||
padded_min = min(orig_h + padding[1], orig_w + padding[0])
|
||||
max_level = min(max_level, int(math.log2(padded_min)))
|
||||
|
||||
if any(padding):
|
||||
mask = torch.nn.functional.pad(
|
||||
mask, (0, padding[0], 0, padding[1]), mode="reflect"
|
||||
)
|
||||
|
||||
B = image1.shape[0]
|
||||
results = []
|
||||
|
||||
for start in range(0, B, _CHUNK_SIZE):
|
||||
end = min(start + _CHUNK_SIZE, B)
|
||||
|
||||
img1_chunk, _ = _pad_for_laplacian(image1[start:end])
|
||||
img2_chunk, _ = _pad_for_laplacian(image2[start:end])
|
||||
mask_chunk = mask[start:end]
|
||||
|
||||
if device is not None:
|
||||
img1_chunk = img1_chunk.to(device)
|
||||
img2_chunk = img2_chunk.to(device)
|
||||
mask_chunk = mask_chunk.to(device)
|
||||
|
||||
out_chunk = _pyramid_blend_chunk(
|
||||
img1_chunk, img2_chunk, mask_chunk, max_level=max_level
|
||||
)
|
||||
cropped = out_chunk[..., :orig_h, :orig_w].clamp(0, 1)
|
||||
results.append(
|
||||
cropped.to(output_device) if output_device is not None else cropped
|
||||
)
|
||||
|
||||
return torch.cat(results, dim=0)
|
||||
|
||||
|
||||
@comfy_node(
|
||||
name="LTXVLaplacianPyramidBlend",
|
||||
description="LTX Laplacian Pyramid Blend",
|
||||
)
|
||||
class LTXVLaplacianPyramidBlend(io.ComfyNode):
|
||||
"""Blend two images seamlessly using Laplacian pyramid blending with a mask."""
|
||||
|
||||
@classmethod
|
||||
def define_schema(cls):
|
||||
return io.Schema(
|
||||
node_id="LTXVLaplacianPyramidBlend",
|
||||
category="Lightricks/utility",
|
||||
description="Blend two images seamlessly using Laplacian pyramid blending.",
|
||||
inputs=[
|
||||
io.Image.Input(
|
||||
"image_a",
|
||||
tooltip="First source image.",
|
||||
),
|
||||
io.Image.Input(
|
||||
"image_b",
|
||||
tooltip="Second source image.",
|
||||
),
|
||||
io.Mask.Input(
|
||||
"mask",
|
||||
tooltip="Blend mask (white = image_a, black = image_b).",
|
||||
),
|
||||
io.Boolean.Input(
|
||||
"trim_to_shortest",
|
||||
default=True,
|
||||
tooltip="Trim image_a, image_b, and mask to the shortest sequence length before blending.",
|
||||
),
|
||||
io.Int.Input(
|
||||
"mask_low_res_dilation",
|
||||
default=5,
|
||||
min=0,
|
||||
max=15,
|
||||
tooltip="Downscale the mask to long side 64, dilate it spatially, then resize it back before blending.",
|
||||
),
|
||||
],
|
||||
outputs=[
|
||||
io.Image.Output("image"),
|
||||
],
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def execute(
|
||||
cls,
|
||||
image_a: torch.Tensor,
|
||||
image_b: torch.Tensor,
|
||||
mask: torch.Tensor,
|
||||
trim_to_shortest: bool,
|
||||
mask_low_res_dilation: int,
|
||||
) -> io.NodeOutput:
|
||||
device = comfy.model_management.get_torch_device()
|
||||
|
||||
if mask.ndim == 4:
|
||||
mask = mask[:, :, :, 0]
|
||||
if mask.ndim == 2:
|
||||
mask = mask.unsqueeze(0)
|
||||
|
||||
if image_a.shape[1:3] != image_b.shape[1:3]:
|
||||
raise ValueError(
|
||||
"image_a and image_b must have the same spatial resolution"
|
||||
)
|
||||
if image_a.shape[1:3] != mask.shape[1:3]:
|
||||
raise ValueError(
|
||||
"image_a, image_b, and mask must have the same spatial resolution"
|
||||
)
|
||||
|
||||
if trim_to_shortest:
|
||||
shortest = min(image_a.shape[0], image_b.shape[0], mask.shape[0])
|
||||
image_a = image_a[:shortest]
|
||||
image_b = image_b[:shortest]
|
||||
mask = mask[:shortest]
|
||||
elif not (image_a.shape[0] == image_b.shape[0] == mask.shape[0]):
|
||||
raise ValueError(
|
||||
"image_a, image_b, and mask must have the same frame count unless trim_to_shortest is enabled"
|
||||
)
|
||||
|
||||
original = image_a.permute(0, 3, 1, 2)
|
||||
target = image_b.permute(0, 3, 1, 2)
|
||||
mask_4d = mask.unsqueeze(1).float()
|
||||
mask_4d = _apply_low_res_mask_dilation(mask_4d, mask_low_res_dilation)
|
||||
|
||||
result = _pyramid_blend(
|
||||
original,
|
||||
target,
|
||||
mask_4d,
|
||||
max_level=7,
|
||||
device=device,
|
||||
output_device=comfy.model_management.intermediate_device(),
|
||||
)
|
||||
|
||||
return io.NodeOutput(result.permute(0, 2, 3, 1))
|
||||
+49
-22
@@ -1,3 +1,6 @@
|
||||
import importlib.metadata
|
||||
import logging
|
||||
|
||||
import comfy
|
||||
import comfy.model_detection
|
||||
import comfy.model_management
|
||||
@@ -9,6 +12,7 @@ import torch
|
||||
try:
|
||||
from q8_kernels.functional.ops import hadamard_transform
|
||||
from q8_kernels.integration.patch_transformer import (
|
||||
patch_comfyui_native_avtransformer,
|
||||
patch_comfyui_native_transformer,
|
||||
)
|
||||
|
||||
@@ -31,6 +35,13 @@ def check_q8_available():
|
||||
)
|
||||
|
||||
|
||||
def check_deprecated():
|
||||
q8_version = tuple(
|
||||
int(x) for x in importlib.metadata.version("q8_kernels").split(".")[:3]
|
||||
)
|
||||
return q8_version >= (0, 1, 5)
|
||||
|
||||
|
||||
@comfy_node(name="LTXQ8Patch")
|
||||
class LTXVQ8Patch:
|
||||
@classmethod
|
||||
@@ -84,30 +95,46 @@ class LTXVQ8Patch:
|
||||
m = model.clone()
|
||||
transformer_key = "diffusion_model"
|
||||
transformer = m.get_model_object(transformer_key)
|
||||
quantize_self_attn, quantize_cross_attn, quantize_ffn = LTXVQ8Patch.PRESETS.get(
|
||||
quantization_preset, (quantize_self_attn, quantize_cross_attn, quantize_ffn)
|
||||
)
|
||||
if getattr(transformer, "quantization_config", None) is not None:
|
||||
if (quantize_self_attn, quantize_cross_attn, quantize_ffn) != getattr(
|
||||
transformer, "quantization_config"
|
||||
):
|
||||
quantize_self_attn, quantize_cross_attn, quantize_ffn = getattr(
|
||||
transformer, "quantization_config"
|
||||
|
||||
is_av = transformer.__class__.__name__ == "LTXAVModel"
|
||||
if is_av:
|
||||
patcher = patch_comfyui_native_avtransformer
|
||||
else:
|
||||
patcher = patch_comfyui_native_transformer
|
||||
|
||||
if check_deprecated():
|
||||
logging.info("This node is deprecated soon. Use new one.")
|
||||
quant_audio_path = False
|
||||
quant_type = "blockwise-fp8"
|
||||
patcher(transformer, use_fp8_attention, True, quant_type, quant_audio_path)
|
||||
else:
|
||||
quantize_self_attn, quantize_cross_attn, quantize_ffn = (
|
||||
LTXVQ8Patch.PRESETS.get(
|
||||
quantization_preset,
|
||||
(quantize_self_attn, quantize_cross_attn, quantize_ffn),
|
||||
)
|
||||
patch_comfyui_native_transformer(
|
||||
transformer,
|
||||
use_fp8_attention,
|
||||
True,
|
||||
quantize_self_attn,
|
||||
quantize_cross_attn,
|
||||
quantize_ffn,
|
||||
)
|
||||
)
|
||||
if getattr(transformer, "quantization_config", None) is not None:
|
||||
if (quantize_self_attn, quantize_cross_attn, quantize_ffn) != getattr(
|
||||
transformer, "quantization_config"
|
||||
):
|
||||
quantize_self_attn, quantize_cross_attn, quantize_ffn = getattr(
|
||||
transformer, "quantization_config"
|
||||
)
|
||||
patcher(
|
||||
transformer,
|
||||
use_fp8_attention,
|
||||
True,
|
||||
quantize_self_attn,
|
||||
quantize_cross_attn,
|
||||
quantize_ffn,
|
||||
)
|
||||
setattr(
|
||||
transformer,
|
||||
"quantization_config",
|
||||
(quantize_self_attn, quantize_cross_attn, quantize_ffn),
|
||||
)
|
||||
setattr(transformer, "is_q8_patched", True)
|
||||
setattr(
|
||||
transformer,
|
||||
"quantization_config",
|
||||
(quantize_self_attn, quantize_cross_attn, quantize_ffn),
|
||||
)
|
||||
setattr(transformer, "use_fp8_attention", use_fp8_attention)
|
||||
return (m,)
|
||||
|
||||
|
||||
@@ -1,5 +1,6 @@
|
||||
diffusers
|
||||
einops
|
||||
huggingface_hub>=0.25.2
|
||||
kornia
|
||||
ninja~=1.11.1.4
|
||||
transformers[timm]>=4.50.0
|
||||
|
||||
@@ -0,0 +1,381 @@
|
||||
import base64
|
||||
import json
|
||||
from io import BytesIO
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
from comfy_api.latest import io
|
||||
from PIL import Image
|
||||
|
||||
from .nodes_registry import comfy_node
|
||||
|
||||
|
||||
def _catmull_rom(p0: dict, p1: dict, p2: dict, p3: dict, t: float) -> dict[str, float]:
|
||||
t2 = t * t
|
||||
t3 = t2 * t
|
||||
return {
|
||||
"x": 0.5
|
||||
* (
|
||||
2 * p1["x"]
|
||||
+ (-p0["x"] + p2["x"]) * t
|
||||
+ (2 * p0["x"] - 5 * p1["x"] + 4 * p2["x"] - p3["x"]) * t2
|
||||
+ (-p0["x"] + 3 * p1["x"] - 3 * p2["x"] + p3["x"]) * t3
|
||||
),
|
||||
"y": 0.5
|
||||
* (
|
||||
2 * p1["y"]
|
||||
+ (-p0["y"] + p2["y"]) * t
|
||||
+ (2 * p0["y"] - 5 * p1["y"] + 4 * p2["y"] - p3["y"]) * t2
|
||||
+ (-p0["y"] + 3 * p1["y"] - 3 * p2["y"] + p3["y"]) * t3
|
||||
),
|
||||
}
|
||||
|
||||
|
||||
def _interpolate_spline(
|
||||
control_points: list[dict], num_samples: int
|
||||
) -> list[dict[str, int]]:
|
||||
"""Catmull-Rom spline interpolation matching the JS frontend logic."""
|
||||
if len(control_points) == 0:
|
||||
return []
|
||||
if len(control_points) == 1:
|
||||
p = control_points[0]
|
||||
return [{"x": round(p["x"]), "y": round(p["y"])} for _ in range(num_samples)]
|
||||
if len(control_points) == 2:
|
||||
a, b = control_points
|
||||
return [
|
||||
{
|
||||
"x": round(a["x"] + (b["x"] - a["x"]) * i / (num_samples - 1)),
|
||||
"y": round(a["y"] + (b["y"] - a["y"]) * i / (num_samples - 1)),
|
||||
}
|
||||
for i in range(num_samples)
|
||||
]
|
||||
pts = [control_points[0], *control_points, control_points[-1]]
|
||||
n_seg = len(pts) - 3
|
||||
result = []
|
||||
for i in range(num_samples):
|
||||
g_t = (i / (num_samples - 1)) * n_seg
|
||||
seg = min(int(g_t), n_seg - 1)
|
||||
l_t = g_t - seg
|
||||
p = _catmull_rom(pts[seg], pts[seg + 1], pts[seg + 2], pts[seg + 3], l_t)
|
||||
result.append({"x": round(p["x"]), "y": round(p["y"])})
|
||||
return result
|
||||
|
||||
|
||||
@comfy_node(name="LTXVSparseTrackEditor", description="LTX Sparse Track Editor")
|
||||
class LTXVSparseTrackEditor(io.ComfyNode):
|
||||
"""Interactive spline editor for drawing sparse motion tracks.
|
||||
|
||||
Provides a canvas widget where users can draw and edit spline control
|
||||
points on top of a reference image. Outputs interpolated track
|
||||
coordinates compatible with LTXVDrawTracks.
|
||||
"""
|
||||
|
||||
@classmethod
|
||||
def define_schema(cls):
|
||||
return io.Schema(
|
||||
node_id="LTXVSparseTrackEditor",
|
||||
category="Lightricks/motion_tracking",
|
||||
description=(
|
||||
"Interactive spline editor for drawing sparse motion tracks "
|
||||
"on a reference image."
|
||||
),
|
||||
inputs=[
|
||||
io.Image.Input(
|
||||
"image",
|
||||
tooltip="Reference image displayed as the editor canvas background.",
|
||||
),
|
||||
io.String.Input(
|
||||
"points_store",
|
||||
default="[]",
|
||||
tooltip="JSON array of spline control points managed by the editor widget.",
|
||||
),
|
||||
io.String.Input(
|
||||
"coordinates",
|
||||
default="[]",
|
||||
tooltip="JSON array of interpolated track coordinates produced by the editor.",
|
||||
),
|
||||
io.Int.Input(
|
||||
"points_to_sample",
|
||||
default=121,
|
||||
min=2,
|
||||
max=10000,
|
||||
tooltip="Number of points sampled along each spline curve.",
|
||||
),
|
||||
],
|
||||
outputs=[
|
||||
io.String.Output("tracks"),
|
||||
],
|
||||
is_output_node=True,
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def execute(
|
||||
cls,
|
||||
image,
|
||||
points_store: str,
|
||||
coordinates: str,
|
||||
points_to_sample: int,
|
||||
) -> io.NodeOutput:
|
||||
# Re-interpolate from control points so that changes to
|
||||
# points_to_sample are always respected, regardless of JS sync.
|
||||
try:
|
||||
splines = json.loads(points_store) if points_store else []
|
||||
except (json.JSONDecodeError, TypeError):
|
||||
splines = []
|
||||
|
||||
if splines and isinstance(splines, list) and isinstance(splines[0], list):
|
||||
interpolated = [_interpolate_spline(sp, points_to_sample) for sp in splines]
|
||||
tracks = json.dumps(interpolated)
|
||||
elif coordinates and coordinates != "[]":
|
||||
tracks = coordinates
|
||||
else:
|
||||
tracks = "[]"
|
||||
|
||||
img_array = (image[0].cpu().numpy() * 255).astype(np.uint8)
|
||||
img = Image.fromarray(img_array)
|
||||
buf = BytesIO()
|
||||
img.save(buf, format="JPEG", quality=75)
|
||||
img_b64 = base64.b64encode(buf.getvalue()).decode("utf-8")
|
||||
return io.NodeOutput(tracks, ui={"bg_image": [img_b64]})
|
||||
|
||||
|
||||
def _parse_tracks(raw: str) -> list[list[dict]]:
|
||||
"""Parse tracks from a JSON string, handling nested/wrapped formats."""
|
||||
parsed = json.loads(raw) if isinstance(raw, str) else raw
|
||||
|
||||
if isinstance(parsed, list):
|
||||
unwrapped = []
|
||||
for item in parsed:
|
||||
unwrapped.append(json.loads(item) if isinstance(item, str) else item)
|
||||
parsed = unwrapped
|
||||
|
||||
tracks: list[list[dict]] = []
|
||||
stack = [parsed]
|
||||
while stack:
|
||||
obj = stack.pop()
|
||||
if isinstance(obj, list) and len(obj) > 0:
|
||||
if isinstance(obj[0], dict) and "x" in obj[0] and "y" in obj[0]:
|
||||
tracks.append(obj)
|
||||
else:
|
||||
stack.extend(obj)
|
||||
return tracks
|
||||
|
||||
|
||||
def _age_color_batch(ratios: torch.Tensor, device: torch.device) -> torch.Tensor:
|
||||
"""Vectorised age-ratio -> RGB [0..1] mapping on GPU.
|
||||
|
||||
Gradient: blue -> green -> yellow -> red.
|
||||
"""
|
||||
colors = torch.zeros(ratios.shape[0], 3, device=device)
|
||||
|
||||
m1 = ratios <= 1 / 3
|
||||
tr1 = ratios[m1] * 3
|
||||
colors[m1, 1] = tr1
|
||||
colors[m1, 2] = 1 - tr1
|
||||
|
||||
m2 = (ratios > 1 / 3) & (ratios <= 2 / 3)
|
||||
tr2 = (ratios[m2] - 1 / 3) * 3
|
||||
colors[m2, 0] = tr2
|
||||
colors[m2, 1] = 1
|
||||
|
||||
m3 = ratios > 2 / 3
|
||||
tr3 = (ratios[m3] - 2 / 3) * 3
|
||||
colors[m3, 0] = 1
|
||||
colors[m3, 1] = 1 - tr3
|
||||
|
||||
return colors
|
||||
|
||||
|
||||
def _render_resolution(width: int, height: int, reference_short_side: int):
|
||||
"""Compute the higher render resolution that preserves aspect ratio."""
|
||||
if height <= width:
|
||||
rw = int(width * reference_short_side / height)
|
||||
rh = reference_short_side
|
||||
else:
|
||||
rw = reference_short_side
|
||||
rh = int(height * reference_short_side / width)
|
||||
scale_x = rw / width
|
||||
scale_y = rh / height
|
||||
return rw, rh, scale_x, scale_y
|
||||
|
||||
|
||||
_MIN_RADIUS = 2
|
||||
_MAX_RADIUS = 8
|
||||
_MAX_TRAIL = 50
|
||||
_REF_SHORT_SIDE = 1080
|
||||
|
||||
|
||||
@comfy_node(name="LTXVDrawTracks", description="LTX Draw Sparse Tracks")
|
||||
class LTXVDrawTracks(io.ComfyNode):
|
||||
"""GPU-accelerated sparse track renderer.
|
||||
|
||||
Renders circles at a high reference resolution and downscales with
|
||||
bilinear interpolation so circle sizes match the CPU version.
|
||||
All work — rasterisation, compositing and resize — stays on GPU.
|
||||
"""
|
||||
|
||||
@classmethod
|
||||
def define_schema(cls):
|
||||
return io.Schema(
|
||||
node_id="LTXVDrawTracks",
|
||||
category="Lightricks/motion_tracking",
|
||||
description=(
|
||||
"GPU-accelerated sparse track renderer. Rasterises circles at "
|
||||
"high resolution and downscales with bilinear interpolation."
|
||||
),
|
||||
inputs=[
|
||||
io.String.Input(
|
||||
"tracks",
|
||||
multiline=True,
|
||||
tooltip="JSON string of track coordinates (list of point lists with x/y keys).",
|
||||
),
|
||||
io.Int.Input(
|
||||
"width",
|
||||
default=512,
|
||||
min=8,
|
||||
max=8192,
|
||||
step=8,
|
||||
tooltip="Output image width in pixels.",
|
||||
),
|
||||
io.Int.Input(
|
||||
"height",
|
||||
default=512,
|
||||
min=8,
|
||||
max=8192,
|
||||
step=8,
|
||||
tooltip="Output image height in pixels.",
|
||||
),
|
||||
],
|
||||
outputs=[
|
||||
io.Image.Output(),
|
||||
],
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def execute(cls, tracks: str, width: int, height: int) -> io.NodeOutput:
|
||||
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
||||
parsed = _parse_tracks(tracks)
|
||||
if not parsed:
|
||||
blank = torch.zeros(1, height, width, 3, device=device, dtype=torch.half)
|
||||
return io.NodeOutput(blank)
|
||||
num_tracks = len(parsed)
|
||||
num_frames = max(len(t) for t in parsed)
|
||||
rw, rh, sx, sy = _render_resolution(width, height, _REF_SHORT_SIDE)
|
||||
|
||||
point_xy = torch.zeros(num_tracks, num_frames, 2, device=device)
|
||||
vis = torch.zeros(num_tracks, num_frames, dtype=torch.bool, device=device)
|
||||
for i, trk in enumerate(parsed):
|
||||
coords = torch.tensor(
|
||||
[[p["x"] * sx, p["y"] * sy] for p in trk],
|
||||
dtype=torch.float32,
|
||||
device=device,
|
||||
)
|
||||
point_xy[i, : len(trk)] = coords
|
||||
vis[i, : len(trk)] = True
|
||||
|
||||
max_d = 2 * _MAX_RADIUS + 3
|
||||
half_d = max_d // 2
|
||||
offsets = torch.arange(max_d, device=device) - half_d
|
||||
oy, ox = torch.meshgrid(offsets, offsets, indexing="ij")
|
||||
template_dist_sq = oy.float().square() + ox.float().square()
|
||||
|
||||
render_frames = torch.zeros(num_frames, rh, rw, 3, device=device)
|
||||
|
||||
for t in range(num_frames):
|
||||
tau_min = max(0, t - _MAX_TRAIL)
|
||||
window = t - tau_min + 1
|
||||
|
||||
active_xy = point_xy[:, tau_min : t + 1]
|
||||
active_vis = vis[:, tau_min : t + 1]
|
||||
|
||||
ages = torch.arange(window - 1, -1, -1, device=device, dtype=torch.float32)
|
||||
ratios = 1.0 - ages / _MAX_TRAIL
|
||||
radii = _MIN_RADIUS + (_MAX_RADIUS - _MIN_RADIUS) * ratios
|
||||
colors = _age_color_batch(ratios, device)
|
||||
|
||||
flat_xy = active_xy.reshape(-1, 2)
|
||||
flat_vis = active_vis.reshape(-1)
|
||||
flat_radii = radii.unsqueeze(0).expand(num_tracks, -1).reshape(-1)
|
||||
flat_colors = colors.unsqueeze(0).expand(num_tracks, -1, -1).reshape(-1, 3)
|
||||
|
||||
idx = flat_vis.nonzero(as_tuple=True)[0]
|
||||
if idx.shape[0] == 0:
|
||||
continue
|
||||
|
||||
pts = flat_xy[idx]
|
||||
r = flat_radii[idx]
|
||||
c = flat_colors[idx]
|
||||
|
||||
flat_ages = ages.unsqueeze(0).expand(num_tracks, -1).reshape(-1)
|
||||
sort_order = flat_ages[idx].argsort(descending=True)
|
||||
pts = pts[sort_order]
|
||||
r = r[sort_order]
|
||||
c = c[sort_order]
|
||||
|
||||
_rasterise_circles(
|
||||
render_frames[t], pts, r, c, template_dist_sq, half_d, max_d, rh, rw
|
||||
)
|
||||
|
||||
out = torch.nn.functional.interpolate(
|
||||
render_frames.permute(0, 3, 1, 2),
|
||||
size=(height, width),
|
||||
mode="bilinear",
|
||||
align_corners=False,
|
||||
).permute(0, 2, 3, 1)
|
||||
|
||||
out = out[..., [2, 1, 0]] # RGB -> BGR to match IC-LoRA training data format
|
||||
return io.NodeOutput(out.half())
|
||||
|
||||
|
||||
def _rasterise_circles(
|
||||
frame: torch.Tensor,
|
||||
pts: torch.Tensor,
|
||||
radii: torch.Tensor,
|
||||
colors: torch.Tensor,
|
||||
template_dist_sq: torch.Tensor,
|
||||
half_d: int,
|
||||
max_d: int,
|
||||
H: int,
|
||||
W: int,
|
||||
) -> None:
|
||||
"""Stamp filled circles onto *frame* fully on-device.
|
||||
|
||||
Uses ``scatter_reduce_`` with ``'amax'`` to resolve overlaps in
|
||||
painter's order (circles are expected oldest-first so the highest
|
||||
index = newest wins).
|
||||
"""
|
||||
M = pts.shape[0]
|
||||
if M == 0:
|
||||
return
|
||||
device = pts.device
|
||||
|
||||
# per-circle masks [M, D, D]
|
||||
radii_sq = (radii * radii).view(M, 1, 1)
|
||||
circle_masks = template_dist_sq.unsqueeze(0) <= radii_sq
|
||||
|
||||
# frame-space indices [M, D, D]
|
||||
cx = pts[:, 0].round().long().view(M, 1, 1)
|
||||
cy = pts[:, 1].round().long().view(M, 1, 1)
|
||||
offsets_y = torch.arange(max_d, device=device).sub(half_d).view(1, max_d, 1)
|
||||
offsets_x = torch.arange(max_d, device=device).sub(half_d).view(1, 1, max_d)
|
||||
fy = (cy + offsets_y).expand(M, max_d, max_d) # [M, D, D]
|
||||
fx = (cx + offsets_x).expand(M, max_d, max_d) # [M, D, D]
|
||||
|
||||
valid = circle_masks & (fy >= 0) & (fy < H) & (fx >= 0) & (fx < W)
|
||||
|
||||
flat_fy = fy[valid]
|
||||
flat_fx = fx[valid]
|
||||
flat_lin = (flat_fy * W + flat_fx).long()
|
||||
|
||||
# circle index per valid pixel (oldest=0 … newest=M-1)
|
||||
j_map = torch.arange(M, device=device, dtype=torch.float32).view(M, 1, 1)
|
||||
j_map = j_map.expand_as(valid)
|
||||
flat_j = j_map[valid]
|
||||
|
||||
# priority map — highest index (newest) wins via 'amax' reduce
|
||||
priority = torch.full((H * W,), -1.0, device=device)
|
||||
priority.scatter_reduce_(0, flat_lin, flat_j, reduce="amax", include_self=False)
|
||||
priority = priority.view(H, W).long()
|
||||
|
||||
has_circle = priority >= 0
|
||||
frame[has_circle] = colors[priority[has_circle]]
|
||||
@@ -0,0 +1,376 @@
|
||||
"""Text Embeddings Pipeline: Feature Extractors + Embeddings Processors.
|
||||
|
||||
Provides the 3-block Gemma text encoder pipeline:
|
||||
1. Gemma Model (external) -- runs the LLM, gets hidden states
|
||||
2. Feature Extractor (V1/V2) -- normalization + projection
|
||||
3. Embeddings Processor (Video / AV) -- wraps Embeddings1DConnector(s)
|
||||
"""
|
||||
|
||||
import json
|
||||
import math
|
||||
from pathlib import Path
|
||||
|
||||
import torch
|
||||
from comfy.utils import load_torch_file
|
||||
from einops import rearrange
|
||||
from torch import nn
|
||||
|
||||
from .embeddings_connector import (
|
||||
Embeddings1DConnector,
|
||||
load_audio_embeddings_connector,
|
||||
load_video_embeddings_connector,
|
||||
)
|
||||
|
||||
_PREFIX_BASE = "model.diffusion_model."
|
||||
_PREFIX_TEXT_PROJ = "text_embedding_projection."
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Helpers
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _to_binary_mask(encoded, encoded_mask):
|
||||
"""Convert additive mask to binary [B, seq, 1] and zero out padding."""
|
||||
binary_mask = (encoded_mask < 0.000001).to(torch.int64)
|
||||
binary_mask = binary_mask.reshape(encoded.shape[0], encoded.shape[1], 1)
|
||||
return encoded * binary_mask, binary_mask
|
||||
|
||||
|
||||
def _rescale_norm(x: torch.Tensor, target_dim: int, source_dim: int) -> torch.Tensor:
|
||||
"""Rescale normalization: x * sqrt(target_dim / source_dim)."""
|
||||
return x * math.sqrt(target_dim / source_dim)
|
||||
|
||||
|
||||
def _filter_sd(sd: dict, prefix: str) -> dict:
|
||||
"""Extract keys with *prefix* and strip the prefix."""
|
||||
return {k[len(prefix) :]: v for k, v in sd.items() if k.startswith(prefix)}
|
||||
|
||||
|
||||
def _load_aggregate_embed(sd: dict, modality: str, dtype) -> nn.Linear:
|
||||
"""Load an aggregate_embed Linear from the state dict.
|
||||
|
||||
Args:
|
||||
sd: Full checkpoint state dict.
|
||||
modality: ``"video"`` or ``"audio"``.
|
||||
dtype: Target dtype.
|
||||
|
||||
Returns ``linear`` module.
|
||||
"""
|
||||
name = f"{modality}_aggregate_embed"
|
||||
weight_key = f"{_PREFIX_TEXT_PROJ}{name}.weight"
|
||||
bias_key = f"{_PREFIX_TEXT_PROJ}{name}.bias"
|
||||
weight = sd[weight_key]
|
||||
out_features, embedding_dim = weight.shape
|
||||
linear = nn.Linear(embedding_dim, out_features, bias=bias_key in sd)
|
||||
linear.load_state_dict(_filter_sd(sd, f"{_PREFIX_TEXT_PROJ}{name}."))
|
||||
return linear.to(dtype=dtype)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Normalization functions (moved from gemma_encoder.py)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def norm_and_concat_padded_batch(
|
||||
encoded_text: torch.Tensor,
|
||||
sequence_lengths: torch.Tensor,
|
||||
padding_side: str = "right",
|
||||
) -> torch.Tensor:
|
||||
"""
|
||||
Normalize a 4D tensor [B, T, D, L] per sample and per layer, using sequence_lengths to mask.
|
||||
Returns [B, T, D * L] tensor with original padding preserved.
|
||||
|
||||
Args:
|
||||
encoded_text: 4D tensor [B, T, D, L]
|
||||
sequence_lengths: 1D tensor [B] with actual sequence lengths
|
||||
padding_side: "left" or "right" to indicate which side has padding
|
||||
"""
|
||||
B, T, D, L = encoded_text.shape
|
||||
device = encoded_text.device
|
||||
|
||||
# Build mask: [B, T, 1, 1]
|
||||
token_indices = torch.arange(T, device=device)[None, :] # [1, T]
|
||||
|
||||
if padding_side == "right":
|
||||
# For right padding, valid tokens are from 0 to sequence_length-1
|
||||
mask = token_indices < sequence_lengths[:, None] # [B, T]
|
||||
elif padding_side == "left":
|
||||
# For left padding, valid tokens are from (T - sequence_length) to T-1
|
||||
start_indices = T - sequence_lengths[:, None] # [B, 1]
|
||||
mask = token_indices >= start_indices # [B, T]
|
||||
else:
|
||||
raise ValueError(f"padding_side must be 'left' or 'right', got {padding_side}")
|
||||
|
||||
mask = rearrange(mask, "b t -> b t 1 1")
|
||||
|
||||
# Compute masked mean: [B, 1, 1, L]
|
||||
masked = encoded_text.masked_fill(~mask, 0.0)
|
||||
# denom = mask.sum(dim=(1, 2), keepdim=True) # [B, 1, 1, 1]
|
||||
denom = (sequence_lengths * D).view(B, 1, 1, 1)
|
||||
mean = masked.sum(dim=(1, 2), keepdim=True) / (denom + 1e-6)
|
||||
|
||||
# Compute masked min/max: [B, 1, 1, L]
|
||||
x_min = encoded_text.masked_fill(~mask, float("inf")).amin(dim=(1, 2), keepdim=True)
|
||||
x_max = encoded_text.masked_fill(~mask, float("-inf")).amax(
|
||||
dim=(1, 2), keepdim=True
|
||||
)
|
||||
range_ = x_max - x_min
|
||||
|
||||
# Normalize only the valid tokens
|
||||
normed = 8 * (encoded_text - mean) / (range_ + 1e-6)
|
||||
|
||||
# concat to be [Batch, T, D * L] - this preserves the original structure
|
||||
normed = normed.reshape(B, T, -1) # [B, T, D * L]
|
||||
|
||||
# Apply mask to preserve original padding (set padded positions to 0)
|
||||
mask_flattened = rearrange(mask, "b t 1 1 -> b t 1").expand(-1, -1, D * L)
|
||||
normed = normed.masked_fill(~mask_flattened, 0.0)
|
||||
|
||||
return normed
|
||||
|
||||
|
||||
def norm_and_concat_per_token_rms(
|
||||
encoded_text: torch.Tensor,
|
||||
attention_mask: torch.Tensor,
|
||||
) -> torch.Tensor:
|
||||
"""Per-token RMSNorm normalization.
|
||||
|
||||
Each token is normalized independently using RMSNorm over its D
|
||||
dimension. This is naturally causal and supports packing since each
|
||||
token's normalization is self-contained.
|
||||
|
||||
Args:
|
||||
encoded_text: 4D tensor [B, T, D, L]
|
||||
attention_mask: [B, T] binary mask (1=real, 0=pad)
|
||||
|
||||
Returns:
|
||||
[B, T, D*L] normalized and flattened tensor with padding zeroed out.
|
||||
"""
|
||||
B, T, D, L = encoded_text.shape
|
||||
variance = torch.mean(encoded_text**2, dim=2, keepdim=True) # [B,T,1,L]
|
||||
normed = encoded_text * torch.rsqrt(variance + 1e-6) # [B,T,D,L]
|
||||
normed = normed.reshape(B, T, D * L)
|
||||
# Zero out padding
|
||||
mask_3d = attention_mask.bool().unsqueeze(-1) # [B, T, 1]
|
||||
normed = torch.where(mask_3d, normed, torch.zeros_like(normed))
|
||||
return normed
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Feature Extractors (Block 2)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class FeatureExtractorV1(nn.Module):
|
||||
"""19B: per-segment norm -> aggregate_embed -> 3840"""
|
||||
|
||||
def __init__(self, aggregate_embed: nn.Module, is_av: bool = False):
|
||||
super().__init__()
|
||||
self.aggregate_embed = aggregate_embed
|
||||
self.is_av = is_av
|
||||
|
||||
def forward(self, all_layer_hiddens, attention_mask, padding_side="left"):
|
||||
sequence_lengths = attention_mask.sum(dim=-1)
|
||||
normed = norm_and_concat_padded_batch(
|
||||
all_layer_hiddens, sequence_lengths, padding_side
|
||||
)
|
||||
normed = normed.to(all_layer_hiddens.dtype)
|
||||
features = self.aggregate_embed(normed)
|
||||
if self.is_av:
|
||||
return {"video": features, "audio": features}
|
||||
return {"video": features}
|
||||
|
||||
|
||||
class FeatureExtractorV2(nn.Module):
|
||||
"""22B: per-token RMS norm -> rescale -> aggregate_embed(s)"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
video_aggregate_embed: nn.Linear,
|
||||
embedding_dim: int,
|
||||
audio_aggregate_embed: nn.Linear | None = None,
|
||||
):
|
||||
super().__init__()
|
||||
self.video_aggregate_embed = video_aggregate_embed
|
||||
self.audio_aggregate_embed = audio_aggregate_embed
|
||||
self.embedding_dim = embedding_dim
|
||||
|
||||
def forward(self, all_layer_hiddens, attention_mask, padding_side="left"):
|
||||
normed = norm_and_concat_per_token_rms(all_layer_hiddens, attention_mask)
|
||||
normed = normed.to(all_layer_hiddens.dtype)
|
||||
v_dim = self.video_aggregate_embed.out_features
|
||||
result = {
|
||||
"video": self.video_aggregate_embed(
|
||||
_rescale_norm(normed, v_dim, self.embedding_dim)
|
||||
)
|
||||
}
|
||||
if self.audio_aggregate_embed is not None:
|
||||
a_dim = self.audio_aggregate_embed.out_features
|
||||
result["audio"] = self.audio_aggregate_embed(
|
||||
_rescale_norm(normed, a_dim, self.embedding_dim)
|
||||
)
|
||||
return result
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Embeddings Processors (Block 3)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class VideoEmbeddingsProcessor(nn.Module):
|
||||
"""Video-only embeddings processor: single connector."""
|
||||
|
||||
def __init__(self, video_connector: Embeddings1DConnector):
|
||||
super().__init__()
|
||||
self.video_connector = video_connector
|
||||
|
||||
def create_embeddings(self, features, attention_mask):
|
||||
encoded, mask = self.video_connector(features["video"], attention_mask)
|
||||
encoded, binary_mask = _to_binary_mask(encoded, mask)
|
||||
return encoded, binary_mask.squeeze(-1)
|
||||
|
||||
|
||||
class AVEmbeddingsProcessor(nn.Module):
|
||||
"""Audio-video embeddings processor: dual connectors + concat."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
video_connector: Embeddings1DConnector,
|
||||
audio_connector: Embeddings1DConnector,
|
||||
):
|
||||
super().__init__()
|
||||
self.video_connector = video_connector
|
||||
self.audio_connector = audio_connector
|
||||
|
||||
def create_embeddings(self, features, attention_mask):
|
||||
encoded, mask = self.video_connector(features["video"], attention_mask)
|
||||
encoded, binary_mask = _to_binary_mask(encoded, mask)
|
||||
audio_encoded, _ = self.audio_connector(features["audio"], attention_mask)
|
||||
encoded = torch.cat([encoded, audio_encoded], dim=-1)
|
||||
return encoded, binary_mask.squeeze(-1)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Loading helpers
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _load_single_aggregate_embed(sd, dtype):
|
||||
"""Load the single aggregate_embed (V1 models).
|
||||
|
||||
Returns:
|
||||
nn.Linear or None if key not found.
|
||||
"""
|
||||
key = f"{_PREFIX_TEXT_PROJ}aggregate_embed.weight"
|
||||
if key not in sd:
|
||||
return None
|
||||
weight = sd[key]
|
||||
in_features, out_features = weight.shape[1], weight.shape[0]
|
||||
# Check for bias
|
||||
bias_key = f"{_PREFIX_TEXT_PROJ}aggregate_embed.bias"
|
||||
has_bias = bias_key in sd
|
||||
linear = nn.Linear(in_features, out_features, bias=has_bias)
|
||||
linear.load_state_dict(_filter_sd(sd, f"{_PREFIX_TEXT_PROJ}aggregate_embed."))
|
||||
return linear.to(dtype=dtype)
|
||||
|
||||
|
||||
def _load_single_aggregate_embed_from_file(path, dtype):
|
||||
"""Load aggregate_embed from a standalone checkpoint file (legacy fallback).
|
||||
|
||||
Args:
|
||||
path: Path to safetensors file containing aggregate_embed weights.
|
||||
dtype: Target dtype.
|
||||
|
||||
Returns:
|
||||
nn.Linear or None if file does not exist.
|
||||
"""
|
||||
path = Path(path)
|
||||
if not path.exists():
|
||||
return None
|
||||
loaded_sd = load_torch_file(str(path), return_metadata=False)
|
||||
if "aggregate_embed.weight" not in loaded_sd:
|
||||
raise ValueError(
|
||||
f"Checkpoint {path} does not contain 'aggregate_embed.weight'."
|
||||
)
|
||||
weight = loaded_sd["aggregate_embed.weight"]
|
||||
in_features, out_features = weight.shape[1], weight.shape[0]
|
||||
bias_key = "aggregate_embed.bias"
|
||||
has_bias = bias_key in loaded_sd
|
||||
linear = nn.Linear(in_features, out_features, bias=has_bias)
|
||||
linear.load_state_dict(
|
||||
{
|
||||
k.removeprefix("aggregate_embed."): v
|
||||
for k, v in loaded_sd.items()
|
||||
if k.startswith("aggregate_embed.")
|
||||
}
|
||||
)
|
||||
return linear.to(dtype=dtype)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Pipeline loader
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def load_text_embeddings_pipeline(
|
||||
ltxv_path, dtype=torch.bfloat16, fallback_proj_path=None
|
||||
):
|
||||
"""Load feature extractor + embeddings processor from LTX-V checkpoint.
|
||||
|
||||
Auto-detects model variant (19B/22B, video-only/AV).
|
||||
|
||||
Args:
|
||||
ltxv_path: Path to the LTX-V checkpoint.
|
||||
dtype: Target dtype for loaded modules.
|
||||
fallback_proj_path: Optional path to a standalone ``proj_linear.safetensors``
|
||||
file. Used as a legacy fallback when the aggregate_embed is not
|
||||
stored inside the LTX-V checkpoint (V1 models).
|
||||
|
||||
Returns:
|
||||
(feature_extractor, embeddings_processor)
|
||||
"""
|
||||
sd, metadata = load_torch_file(str(ltxv_path), return_metadata=True)
|
||||
config = json.loads(metadata.get("config", "{}"))
|
||||
transformer_config = config.get("transformer", {})
|
||||
|
||||
is_av = f"{_PREFIX_BASE}audio_adaln_single.linear.weight" in sd
|
||||
has_dual_aggregate = f"{_PREFIX_TEXT_PROJ}video_aggregate_embed.weight" in sd
|
||||
|
||||
# Load connectors (always needed)
|
||||
video_connector = load_video_embeddings_connector(sd, transformer_config, dtype)
|
||||
|
||||
# Build embeddings processor (Block 3)
|
||||
if is_av:
|
||||
audio_connector = load_audio_embeddings_connector(sd, transformer_config, dtype)
|
||||
processor = AVEmbeddingsProcessor(video_connector, audio_connector)
|
||||
else:
|
||||
processor = VideoEmbeddingsProcessor(video_connector)
|
||||
|
||||
# Build feature extractor (Block 2)
|
||||
if has_dual_aggregate:
|
||||
# V2 (22B): validate config matches expected settings
|
||||
_expected = {
|
||||
"caption_projection_first_linear": False,
|
||||
"caption_proj_input_norm": False,
|
||||
"caption_projection_second_linear": False,
|
||||
"caption_proj_before_connector": True,
|
||||
"text_encoder_norm_type": "per_token_rms",
|
||||
}
|
||||
for key, expected_val in _expected.items():
|
||||
actual = transformer_config.get(key)
|
||||
assert actual == expected_val, (
|
||||
f"Unexpected config for dual-aggregate model: "
|
||||
f"{key}={actual!r}, expected {expected_val!r}"
|
||||
)
|
||||
video_agg = _load_aggregate_embed(sd, "video", dtype)
|
||||
audio_agg = _load_aggregate_embed(sd, "audio", dtype) if is_av else None
|
||||
embedding_dim = transformer_config.get("prompt_embedding_dim", 3840)
|
||||
return (FeatureExtractorV2(video_agg, embedding_dim, audio_agg), processor)
|
||||
# V1 (19B)
|
||||
aggregate_embed = _load_single_aggregate_embed(sd, dtype)
|
||||
if aggregate_embed is None and fallback_proj_path is not None:
|
||||
aggregate_embed = _load_single_aggregate_embed_from_file(
|
||||
fallback_proj_path, dtype
|
||||
)
|
||||
return (FeatureExtractorV1(aggregate_embed, is_av=is_av), processor)
|
||||
@@ -35,6 +35,20 @@ def _build_av_meta(video_latent: dict, audio_latent: dict) -> dict:
|
||||
}
|
||||
|
||||
|
||||
@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
|
||||
|
||||
+146
@@ -0,0 +1,146 @@
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
from comfy_api.latest import io
|
||||
|
||||
from .nodes_registry import comfy_node
|
||||
|
||||
|
||||
@comfy_node(name="LTXVDilateVideoMask")
|
||||
class LTXVDilateVideoMask(io.ComfyNode):
|
||||
"""Dilates a video mask spatially and/or temporally using max-pooling and thresholds the result."""
|
||||
|
||||
@classmethod
|
||||
def define_schema(cls):
|
||||
return io.Schema(
|
||||
node_id="LTXVDilateVideoMask",
|
||||
category="Lightricks/mask_operations",
|
||||
description=(
|
||||
"Dilates a video mask spatially and/or temporally using "
|
||||
"separable max-pooling and thresholds the result."
|
||||
),
|
||||
inputs=[
|
||||
io.Int.Input(
|
||||
"spatial_radius",
|
||||
default=1,
|
||||
min=0,
|
||||
max=300,
|
||||
tooltip="Half-size of the spatial dilation kernel. Kernel = 2*radius+1.",
|
||||
),
|
||||
io.Int.Input(
|
||||
"temporal_radius",
|
||||
default=0,
|
||||
min=0,
|
||||
max=10,
|
||||
tooltip="Half-size of the temporal dilation kernel. Kernel = 2*radius+1.",
|
||||
),
|
||||
io.Mask.Input(
|
||||
"mask",
|
||||
optional=True,
|
||||
tooltip="Video mask to dilate. Either this or image_as_mask must be provided.",
|
||||
),
|
||||
io.Image.Input(
|
||||
"image_as_mask",
|
||||
optional=True,
|
||||
tooltip="Image to use as mask (channel-averaged). Either this or mask must be provided.",
|
||||
),
|
||||
],
|
||||
outputs=[
|
||||
io.Mask.Output("mask"),
|
||||
],
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def execute(
|
||||
cls,
|
||||
spatial_radius: int,
|
||||
temporal_radius: int,
|
||||
mask: torch.Tensor | None = None,
|
||||
image_as_mask: torch.Tensor | None = None,
|
||||
) -> io.NodeOutput:
|
||||
if mask is None and image_as_mask is None:
|
||||
raise ValueError("Either 'mask' or 'image_as_mask' must be provided.")
|
||||
|
||||
if mask is None:
|
||||
mask = image_as_mask.mean(dim=-1)
|
||||
|
||||
if mask.ndim == 4:
|
||||
mask = mask[:, :, :, 0]
|
||||
|
||||
s_kernel = spatial_radius * 2 + 1
|
||||
t_kernel = temporal_radius * 2 + 1
|
||||
|
||||
# Separable dilation: 2D spatial + 1D temporal (much faster than 3D pooling)
|
||||
if s_kernel > 1:
|
||||
mask = mask.unsqueeze(1) # (B, 1, H, W)
|
||||
mask = F.max_pool2d(
|
||||
mask, kernel_size=s_kernel, stride=1, padding=spatial_radius
|
||||
)
|
||||
mask = mask.squeeze(1)
|
||||
|
||||
if t_kernel > 1:
|
||||
B, H, W = mask.shape
|
||||
mask = mask.permute(1, 2, 0).reshape(H * W, 1, B)
|
||||
mask = F.max_pool1d(
|
||||
mask, kernel_size=t_kernel, stride=1, padding=temporal_radius
|
||||
)
|
||||
mask = mask.reshape(H, W, B).permute(2, 0, 1)
|
||||
|
||||
mask = (mask > 0.5).float()
|
||||
return io.NodeOutput(mask)
|
||||
|
||||
|
||||
_BG_COLOR_RGB = (102, 255, 0)
|
||||
|
||||
|
||||
@comfy_node(name="LTXVInpaintPreprocess")
|
||||
class LTXVInpaintPreprocess(io.ComfyNode):
|
||||
"""Composites images with a green (#66FF00) background where mask is active.
|
||||
|
||||
If the mask has a single frame it is broadcast to match the video length.
|
||||
"""
|
||||
|
||||
@classmethod
|
||||
def define_schema(cls):
|
||||
return io.Schema(
|
||||
node_id="LTXVInpaintPreprocess",
|
||||
category="Lightricks/image_processing",
|
||||
description=(
|
||||
"Composites images with a green background where mask is "
|
||||
"active, for inpainting conditioning."
|
||||
),
|
||||
inputs=[
|
||||
io.Image.Input(
|
||||
"images",
|
||||
tooltip="Video frames to composite onto the green background.",
|
||||
),
|
||||
io.Mask.Input(
|
||||
"mask",
|
||||
tooltip="Mask indicating regions to replace with green. Single-frame masks are broadcast.",
|
||||
),
|
||||
],
|
||||
outputs=[
|
||||
io.Image.Output("image"),
|
||||
],
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def execute(
|
||||
cls,
|
||||
images: torch.Tensor,
|
||||
mask: torch.Tensor,
|
||||
) -> io.NodeOutput:
|
||||
if mask.ndim == 4:
|
||||
mask = mask[:, :, :, 0]
|
||||
|
||||
if mask.shape[0] == 1 and images.shape[0] > 1:
|
||||
mask = mask.expand(images.shape[0], -1, -1)
|
||||
|
||||
min_frames = min(mask.shape[0], images.shape[0])
|
||||
mask = mask[:min_frames]
|
||||
images = images[:min_frames]
|
||||
|
||||
mask_4d = mask.unsqueeze(-1) # (B, H, W, 1) for broadcasting
|
||||
bg_color = torch.tensor(_BG_COLOR_RGB).float().to(images.device) / 255
|
||||
|
||||
result = images * (1 - mask_4d) + bg_color.view(1, 1, 1, 3) * mask_4d
|
||||
return io.NodeOutput(result)
|
||||
@@ -0,0 +1,346 @@
|
||||
import { app } from "../../../scripts/app.js";
|
||||
|
||||
const PREFIXES = ["video", "audio"];
|
||||
const STRING_FIELDS = [
|
||||
"sigmas", "cfg", "stg", "rescale", "modality_scale",
|
||||
"perturb_attn", "skip_step", "cross_attn",
|
||||
"cfg_zero_star",
|
||||
];
|
||||
const FIELD_LABELS = {
|
||||
sigmas: "Sigmas",
|
||||
cfg: "CFG",
|
||||
stg: "STG",
|
||||
rescale: "Rescale",
|
||||
modality_scale: "Modality Scale",
|
||||
perturb_attn: "Perturb Attn",
|
||||
skip_step: "Skip Step",
|
||||
cross_attn: "Cross Attn",
|
||||
cfg_zero_star: "CFG Zero*",
|
||||
};
|
||||
|
||||
const _origProps = {};
|
||||
|
||||
function toggleWidget(node, widget, show) {
|
||||
if (!widget) return;
|
||||
if (!_origProps[widget.name]) {
|
||||
_origProps[widget.name] = {
|
||||
origComputeSize: widget.computeSize,
|
||||
};
|
||||
}
|
||||
const origSize = node.size;
|
||||
|
||||
widget.hidden = !show;
|
||||
widget.computeSize = show
|
||||
? _origProps[widget.name].origComputeSize
|
||||
: () => [0, -3.3];
|
||||
|
||||
widget.linkedWidgets?.forEach((w) => toggleWidget(node, w, show));
|
||||
|
||||
const height = show
|
||||
? Math.max(node.computeSize()[1], origSize[1])
|
||||
: node.computeSize()[1];
|
||||
node.setSize([node.size[0], height]);
|
||||
|
||||
if (show) delete widget.computedHeight;
|
||||
else widget.computedHeight = 0;
|
||||
}
|
||||
|
||||
function toggleModalityWidgets(node, prefix, enabled) {
|
||||
for (const w of node.widgets) {
|
||||
if (
|
||||
w.name.startsWith(prefix + "_") &&
|
||||
w.name !== prefix + "_enabled"
|
||||
) {
|
||||
toggleWidget(node, w, enabled);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
function injectMultilineLabels(node) {
|
||||
const labelMap = {};
|
||||
for (const prefix of PREFIXES) {
|
||||
for (const field of STRING_FIELDS) {
|
||||
const key = prefix + "_" + field;
|
||||
labelMap[key] = prefix.toUpperCase() + " " + FIELD_LABELS[field];
|
||||
}
|
||||
}
|
||||
|
||||
function applyLabels() {
|
||||
let allDone = true;
|
||||
for (const w of node.widgets) {
|
||||
if (!(w.name in labelMap)) continue;
|
||||
// w.element is the <textarea> itself, not the wrapper div
|
||||
const textarea = w.element;
|
||||
if (!textarea || textarea.tagName !== "TEXTAREA") {
|
||||
allDone = false;
|
||||
continue;
|
||||
}
|
||||
const container = textarea.parentElement;
|
||||
if (!container) {
|
||||
allDone = false;
|
||||
continue;
|
||||
}
|
||||
if (container.querySelector(".guider-param-label")) continue;
|
||||
|
||||
textarea.placeholder = labelMap[w.name];
|
||||
|
||||
const label = document.createElement("div");
|
||||
label.className = "guider-param-label";
|
||||
label.textContent = labelMap[w.name];
|
||||
label.style.cssText =
|
||||
"font-size:10px;font-weight:600;color:#bbb;" +
|
||||
"padding:0 4px;pointer-events:none;" +
|
||||
"position:absolute;top:-14px;left:0;z-index:1;";
|
||||
container.style.overflow = "visible";
|
||||
container.insertBefore(label, textarea);
|
||||
}
|
||||
return allDone;
|
||||
}
|
||||
|
||||
let retries = 0;
|
||||
function tryApply() {
|
||||
if (applyLabels() || retries > 30) return;
|
||||
retries++;
|
||||
requestAnimationFrame(tryApply);
|
||||
}
|
||||
requestAnimationFrame(tryApply);
|
||||
}
|
||||
|
||||
const MOD_PREFIX_MAP = { VIDEO: "video", AUDIO: "audio" };
|
||||
|
||||
function populateWidgets(node, data) {
|
||||
for (const [modName, fields] of Object.entries(data)) {
|
||||
const prefix = MOD_PREFIX_MAP[modName];
|
||||
if (!prefix) continue;
|
||||
for (const [field, value] of Object.entries(fields)) {
|
||||
const w = node.widgets.find(
|
||||
(w) => w.name === prefix + "_" + field
|
||||
);
|
||||
if (w) w.value = value;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
async function importXlsxValues(node, filename) {
|
||||
try {
|
||||
const resp = await fetch("/guider_params/import", {
|
||||
method: "POST",
|
||||
headers: { "Content-Type": "application/json" },
|
||||
body: JSON.stringify({ filename }),
|
||||
});
|
||||
if (!resp.ok) return;
|
||||
const data = await resp.json();
|
||||
populateWidgets(node, data);
|
||||
} catch (err) {
|
||||
console.error("XLSX import error:", err);
|
||||
}
|
||||
}
|
||||
|
||||
function setupUploadButton(node) {
|
||||
const fileInput = document.createElement("input");
|
||||
fileInput.type = "file";
|
||||
fileInput.accept = ".xlsx";
|
||||
fileInput.style.display = "none";
|
||||
document.body.appendChild(fileInput);
|
||||
|
||||
fileInput.addEventListener("change", async () => {
|
||||
if (!fileInput.files || fileInput.files.length === 0) return;
|
||||
const file = fileInput.files[0];
|
||||
const formData = new FormData();
|
||||
formData.append("image", file, file.name);
|
||||
formData.append("type", "input");
|
||||
formData.append("overwrite", "true");
|
||||
|
||||
try {
|
||||
const resp = await fetch("/upload/image", {
|
||||
method: "POST",
|
||||
body: formData,
|
||||
});
|
||||
if (!resp.ok) {
|
||||
console.error("XLSX upload failed:", resp.statusText);
|
||||
return;
|
||||
}
|
||||
const result = await resp.json();
|
||||
const uploadedName = result.name;
|
||||
|
||||
await importXlsxValues(node, uploadedName);
|
||||
} catch (err) {
|
||||
console.error("XLSX upload error:", err);
|
||||
}
|
||||
fileInput.value = "";
|
||||
});
|
||||
|
||||
node.addWidget("button", "Upload XLSX", "Upload XLSX", () => {
|
||||
fileInput.click();
|
||||
});
|
||||
}
|
||||
|
||||
function hasConnectedInput(node, widgetName) {
|
||||
if (!node.inputs) return false;
|
||||
for (const inp of node.inputs) {
|
||||
if (inp.widget && inp.widget.name === widgetName && inp.link != null) {
|
||||
return true;
|
||||
}
|
||||
}
|
||||
return false;
|
||||
}
|
||||
|
||||
function collectModalityParams(node) {
|
||||
const resolved = node._resolvedValues || {};
|
||||
const modalities = {};
|
||||
for (const prefix of PREFIXES) {
|
||||
const enabledWidget = node.widgets.find(
|
||||
(w) => w.name === prefix + "_enabled"
|
||||
);
|
||||
if (!enabledWidget || !enabledWidget.value) continue;
|
||||
const fields = {};
|
||||
const resolvedMod = resolved[prefix] || {};
|
||||
for (const field of STRING_FIELDS) {
|
||||
const widgetName = prefix + "_" + field;
|
||||
if (hasConnectedInput(node, widgetName) && field in resolvedMod) {
|
||||
fields[field] = resolvedMod[field];
|
||||
} else {
|
||||
const w = node.widgets.find(
|
||||
(w) => w.name === widgetName
|
||||
);
|
||||
if (w) fields[field] = w.value;
|
||||
}
|
||||
}
|
||||
const zisName = prefix + "_zero_init_sigma";
|
||||
if (hasConnectedInput(node, zisName) && "zero_init_sigma" in resolvedMod) {
|
||||
fields["zero_init_sigma"] = resolvedMod["zero_init_sigma"];
|
||||
} else {
|
||||
const zisWidget = node.widgets.find(
|
||||
(w) => w.name === zisName
|
||||
);
|
||||
if (zisWidget) fields["zero_init_sigma"] = String(zisWidget.value);
|
||||
}
|
||||
modalities[prefix === "video" ? "VIDEO" : "AUDIO"] = fields;
|
||||
}
|
||||
return modalities;
|
||||
}
|
||||
|
||||
function setupExportButton(node) {
|
||||
node.addWidget("button", "Export XLSX", "Export XLSX", async () => {
|
||||
const filename = prompt("Enter filename for XLSX export:", "preset");
|
||||
if (!filename) return;
|
||||
|
||||
const modalities = collectModalityParams(node);
|
||||
if (Object.keys(modalities).length === 0) {
|
||||
alert("No modalities enabled. Enable at least one.");
|
||||
return;
|
||||
}
|
||||
|
||||
try {
|
||||
const resp = await fetch("/guider_params/export", {
|
||||
method: "POST",
|
||||
headers: { "Content-Type": "application/json" },
|
||||
body: JSON.stringify({ filename, modalities }),
|
||||
});
|
||||
const result = await resp.json();
|
||||
if (!resp.ok) {
|
||||
alert("Export failed: " + (result.error || resp.statusText));
|
||||
return;
|
||||
}
|
||||
|
||||
const url =
|
||||
"/view?filename=" +
|
||||
encodeURIComponent(result.filename) +
|
||||
"&type=" +
|
||||
encodeURIComponent(result.type);
|
||||
window.open(url, "_blank");
|
||||
} catch (err) {
|
||||
console.error("Export error:", err);
|
||||
alert("Export failed: " + err.message);
|
||||
}
|
||||
});
|
||||
}
|
||||
|
||||
app.registerExtension({
|
||||
name: "ltx.guider_parameters_advanced",
|
||||
|
||||
async nodeCreated(node) {
|
||||
if (node.comfyClass !== "GuiderParametersAdvanced") {
|
||||
return;
|
||||
}
|
||||
|
||||
// Hide internal widgets that buttons manage
|
||||
const saveWidget = node.widgets.find(
|
||||
(w) => w.name === "save_filename"
|
||||
);
|
||||
if (saveWidget) toggleWidget(node, saveWidget, false);
|
||||
|
||||
setupUploadButton(node);
|
||||
setupExportButton(node);
|
||||
injectMultilineLabels(node);
|
||||
|
||||
for (const prefix of PREFIXES) {
|
||||
const enabledWidget = node.widgets.find(
|
||||
(w) => w.name === prefix + "_enabled"
|
||||
);
|
||||
if (!enabledWidget) continue;
|
||||
|
||||
const origCallback = enabledWidget.callback;
|
||||
enabledWidget.callback = function (...args) {
|
||||
if (origCallback) origCallback.apply(this, args);
|
||||
toggleModalityWidgets(node, prefix, enabledWidget.value);
|
||||
};
|
||||
|
||||
toggleModalityWidgets(node, prefix, enabledWidget.value);
|
||||
}
|
||||
},
|
||||
|
||||
async beforeRegisterNodeDef(nodeType, nodeData) {
|
||||
if (nodeData.name !== "GuiderParametersAdvanced") {
|
||||
return;
|
||||
}
|
||||
|
||||
const origOnExecuted = nodeType.prototype.onExecuted;
|
||||
|
||||
nodeType.prototype.onExecuted = function (message) {
|
||||
if (origOnExecuted) {
|
||||
origOnExecuted.apply(this, arguments);
|
||||
}
|
||||
|
||||
if (
|
||||
Array.isArray(message.resolved_values) &&
|
||||
message.resolved_values.length > 0
|
||||
) {
|
||||
this._resolvedValues = message.resolved_values[0];
|
||||
}
|
||||
|
||||
if (
|
||||
Array.isArray(message.xlsx_download) &&
|
||||
message.xlsx_download.length > 0
|
||||
) {
|
||||
const info = message.xlsx_download[0];
|
||||
const url =
|
||||
"/view?filename=" +
|
||||
encodeURIComponent(info.filename) +
|
||||
"&type=" +
|
||||
encodeURIComponent(info.type);
|
||||
|
||||
let linkWidget = this.widgets.find(
|
||||
(w) => w.name === "_download_link"
|
||||
);
|
||||
if (!linkWidget) {
|
||||
linkWidget = this.addWidget(
|
||||
"button",
|
||||
"_download_link",
|
||||
"Download XLSX",
|
||||
() => {
|
||||
window.open(url, "_blank");
|
||||
}
|
||||
);
|
||||
}
|
||||
linkWidget.label = "Download: " + info.filename;
|
||||
linkWidget.callback = () => {
|
||||
window.open(url, "_blank");
|
||||
};
|
||||
|
||||
this.setSize(this.computeSize());
|
||||
}
|
||||
|
||||
};
|
||||
},
|
||||
});
|
||||
@@ -0,0 +1,715 @@
|
||||
import { app } from "../../../scripts/app.js";
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Catmull-Rom spline interpolation (written from scratch)
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
function catmullRom(p0, p1, p2, p3, t) {
|
||||
const t2 = t * t;
|
||||
const t3 = t2 * t;
|
||||
return {
|
||||
x:
|
||||
0.5 *
|
||||
(2 * p1.x +
|
||||
(-p0.x + p2.x) * t +
|
||||
(2 * p0.x - 5 * p1.x + 4 * p2.x - p3.x) * t2 +
|
||||
(-p0.x + 3 * p1.x - 3 * p2.x + p3.x) * t3),
|
||||
y:
|
||||
0.5 *
|
||||
(2 * p1.y +
|
||||
(-p0.y + p2.y) * t +
|
||||
(2 * p0.y - 5 * p1.y + 4 * p2.y - p3.y) * t2 +
|
||||
(-p0.y + 3 * p1.y - 3 * p2.y + p3.y) * t3),
|
||||
};
|
||||
}
|
||||
|
||||
function interpolateSpline(controlPoints, numSamples) {
|
||||
if (controlPoints.length === 0) return [];
|
||||
if (controlPoints.length === 1) {
|
||||
return Array.from({ length: numSamples }, () => ({ ...controlPoints[0] }));
|
||||
}
|
||||
if (controlPoints.length === 2) {
|
||||
const [a, b] = controlPoints;
|
||||
return Array.from({ length: numSamples }, (_, i) => {
|
||||
const t = i / (numSamples - 1);
|
||||
return { x: a.x + (b.x - a.x) * t, y: a.y + (b.y - a.y) * t };
|
||||
});
|
||||
}
|
||||
// Pad with phantom endpoints for Catmull-Rom
|
||||
const pts = [
|
||||
controlPoints[0],
|
||||
...controlPoints,
|
||||
controlPoints[controlPoints.length - 1],
|
||||
];
|
||||
const nSeg = pts.length - 3;
|
||||
const result = [];
|
||||
for (let i = 0; i < numSamples; i++) {
|
||||
const gT = (i / (numSamples - 1)) * nSeg;
|
||||
const seg = Math.min(Math.floor(gT), nSeg - 1);
|
||||
const lT = gT - seg;
|
||||
result.push(catmullRom(pts[seg], pts[seg + 1], pts[seg + 2], pts[seg + 3], lT));
|
||||
}
|
||||
return result;
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Visual constants
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
const SPLINE_COLORS = [
|
||||
"#ef4444", "#22c55e", "#3b82f6", "#f59e0b",
|
||||
"#a855f7", "#06b6d4", "#f97316", "#84cc16",
|
||||
];
|
||||
const POINT_RADIUS = 5;
|
||||
const ACTIVE_POINT_RADIUS = 7;
|
||||
const HIT_TOLERANCE = 14;
|
||||
const CURVE_LINE_WIDTH = 2.5;
|
||||
const CANVAS_MIN_HEIGHT = 128;
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Helpers
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
function dist(a, b) {
|
||||
return Math.hypot(a.x - b.x, a.y - b.y);
|
||||
}
|
||||
|
||||
function hideWidget(w) {
|
||||
if (!w) return;
|
||||
w.hidden = true;
|
||||
w._origComputeSize = w.computeSize;
|
||||
w.computeSize = () => [0, -3.3];
|
||||
w.computedHeight = 0;
|
||||
}
|
||||
|
||||
function closestSegment(pts, p) {
|
||||
let best = Infinity;
|
||||
let bestIdx = 0;
|
||||
for (let i = 0; i < pts.length - 1; i++) {
|
||||
const a = pts[i];
|
||||
const b = pts[i + 1];
|
||||
const dx = b.x - a.x;
|
||||
const dy = b.y - a.y;
|
||||
const len2 = dx * dx + dy * dy;
|
||||
let t = len2 === 0 ? 0 : ((p.x - a.x) * dx + (p.y - a.y) * dy) / len2;
|
||||
t = Math.max(0, Math.min(1, t));
|
||||
const proj = { x: a.x + t * dx, y: a.y + t * dy };
|
||||
const d = dist(proj, p);
|
||||
if (d < best) {
|
||||
best = d;
|
||||
bestIdx = i;
|
||||
}
|
||||
}
|
||||
return { dist: best, segIndex: bestIdx };
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Extension registration
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
app.registerExtension({
|
||||
name: "LTXVideo.SparseTrackEditor",
|
||||
|
||||
async nodeCreated(node) {
|
||||
if (node.comfyClass !== "LTXVSparseTrackEditor") return;
|
||||
initEditor(node);
|
||||
},
|
||||
|
||||
async beforeRegisterNodeDef(nodeType, nodeData) {
|
||||
if (nodeData?.name !== "LTXVSparseTrackEditor") return;
|
||||
|
||||
const origExecuted = nodeType.prototype.onExecuted;
|
||||
nodeType.prototype.onExecuted = function (data) {
|
||||
origExecuted?.apply(this, arguments);
|
||||
if (data?.bg_image?.[0]) {
|
||||
loadBgImage(this, data.bg_image[0]);
|
||||
}
|
||||
};
|
||||
|
||||
const origConfigure = nodeType.prototype.onConfigure;
|
||||
nodeType.prototype.onConfigure = function (info) {
|
||||
origConfigure?.apply(this, arguments);
|
||||
if (this._ed) {
|
||||
reloadState(this);
|
||||
}
|
||||
};
|
||||
|
||||
const origRemoved = nodeType.prototype.onRemoved;
|
||||
nodeType.prototype.onRemoved = function () {
|
||||
origRemoved?.apply(this, arguments);
|
||||
if (this._ed) {
|
||||
cancelAnimationFrame(this._ed.rafId);
|
||||
if (this._ed._docClickHandler) {
|
||||
document.removeEventListener("click", this._ed._docClickHandler);
|
||||
}
|
||||
this._ed = null;
|
||||
}
|
||||
};
|
||||
},
|
||||
});
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Editor initialisation
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
function initEditor(node) {
|
||||
const pointsW = node.widgets.find((w) => w.name === "points_store");
|
||||
const coordsW = node.widgets.find((w) => w.name === "coordinates");
|
||||
hideWidget(pointsW);
|
||||
hideWidget(coordsW);
|
||||
|
||||
const ed = {
|
||||
splines: [],
|
||||
active: 0,
|
||||
imgW: 1024,
|
||||
imgH: 576,
|
||||
bgImg: null,
|
||||
drag: null,
|
||||
hover: null,
|
||||
canvas: null,
|
||||
ctx: null,
|
||||
menu: null,
|
||||
dirty: true,
|
||||
rafId: null,
|
||||
};
|
||||
node._ed = ed;
|
||||
|
||||
// Restore saved state
|
||||
try {
|
||||
const saved = JSON.parse(pointsW.value);
|
||||
if (Array.isArray(saved) && saved.length > 0) {
|
||||
ed.splines = saved;
|
||||
}
|
||||
} catch (_) {
|
||||
/* first run */
|
||||
}
|
||||
if (ed.splines.length === 0) {
|
||||
ed.splines = [
|
||||
[
|
||||
{ x: ed.imgW * 0.3, y: ed.imgH * 0.5 },
|
||||
{ x: ed.imgW * 0.7, y: ed.imgH * 0.5 },
|
||||
],
|
||||
];
|
||||
}
|
||||
|
||||
// DOM
|
||||
const container = document.createElement("div");
|
||||
container.style.position = "relative";
|
||||
container.style.minHeight = `${CANVAS_MIN_HEIGHT}px`;
|
||||
|
||||
const canvas = document.createElement("canvas");
|
||||
canvas.width = 1024;
|
||||
canvas.height = 576;
|
||||
canvas.style.width = "100%";
|
||||
canvas.style.display = "block";
|
||||
canvas.style.background = "#1a1a2e";
|
||||
canvas.style.borderRadius = "4px";
|
||||
canvas.style.cursor = "default";
|
||||
canvas.style.pointerEvents = "auto";
|
||||
canvas.style.touchAction = "none";
|
||||
container.appendChild(canvas);
|
||||
ed.canvas = canvas;
|
||||
ed.ctx = canvas.getContext("2d");
|
||||
|
||||
const menu = buildContextMenu();
|
||||
container.appendChild(menu);
|
||||
ed.menu = menu;
|
||||
|
||||
const widget = node.addDOMWidget("track_editor", "TrackEditorWidget", container, {
|
||||
serialize: false,
|
||||
hideOnZoom: false,
|
||||
});
|
||||
widget.computeSize = () => {
|
||||
const w = node.size[0];
|
||||
const aspect = node._ed.imgH / node._ed.imgW;
|
||||
const h = Math.max(CANVAS_MIN_HEIGHT, Math.round(w * aspect));
|
||||
return [w, h + 10];
|
||||
};
|
||||
|
||||
setupEvents(node);
|
||||
syncWidgets(node);
|
||||
startRenderLoop(node);
|
||||
|
||||
requestAnimationFrame(() => {
|
||||
node.setSize?.(node.computeSize());
|
||||
app.graph?.setDirtyCanvas(true, true);
|
||||
});
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// State reload (fires on onConfigure, i.e. workflow load)
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
function reloadState(node) {
|
||||
const ed = node._ed;
|
||||
const pointsW = node.widgets.find((w) => w.name === "points_store");
|
||||
try {
|
||||
const saved = JSON.parse(pointsW?.value);
|
||||
if (Array.isArray(saved) && saved.length > 0) {
|
||||
ed.splines = saved;
|
||||
}
|
||||
} catch (_) {
|
||||
/* keep current splines */
|
||||
}
|
||||
ed.dirty = true;
|
||||
syncWidgets(node);
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Background image loading
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
function loadBgImage(node, b64) {
|
||||
const ed = node._ed;
|
||||
const img = new window.Image();
|
||||
img.onload = () => {
|
||||
const firstLoad = ed.bgImg === null;
|
||||
ed.bgImg = img;
|
||||
ed.imgW = img.width;
|
||||
ed.imgH = img.height;
|
||||
if (firstLoad && ed.splines.length === 1 && ed.splines[0].length === 2) {
|
||||
ed.splines[0][0] = { x: ed.imgW * 0.3, y: ed.imgH * 0.5 };
|
||||
ed.splines[0][1] = { x: ed.imgW * 0.7, y: ed.imgH * 0.5 };
|
||||
}
|
||||
ed.dirty = true;
|
||||
syncWidgets(node);
|
||||
requestAnimationFrame(() => {
|
||||
node.setSize?.(node.computeSize());
|
||||
app.graph?.setDirtyCanvas(true, true);
|
||||
});
|
||||
};
|
||||
img.src = `data:image/jpeg;base64,${b64}`;
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Coordinate transforms (image space <-> canvas pixel space)
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
function getScale(node) {
|
||||
const ed = node._ed;
|
||||
const c = ed.canvas;
|
||||
const scaleX = c.width / ed.imgW;
|
||||
const scaleY = c.height / ed.imgH;
|
||||
return Math.min(scaleX, scaleY);
|
||||
}
|
||||
|
||||
function imgToCanvas(node, p) {
|
||||
const s = getScale(node);
|
||||
return { x: p.x * s, y: p.y * s };
|
||||
}
|
||||
|
||||
function canvasToImg(node, p) {
|
||||
const s = getScale(node);
|
||||
return { x: p.x / s, y: p.y / s };
|
||||
}
|
||||
|
||||
function mouseToCanvas(node, e) {
|
||||
const rect = node._ed.canvas.getBoundingClientRect();
|
||||
return {
|
||||
x: (e.clientX - rect.left) * (node._ed.canvas.width / rect.width),
|
||||
y: (e.clientY - rect.top) * (node._ed.canvas.height / rect.height),
|
||||
};
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Hit-testing
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
function hitTestPoint(node, canvasPos) {
|
||||
const ed = node._ed;
|
||||
const dpr = window.devicePixelRatio || 1;
|
||||
const tol = HIT_TOLERANCE * dpr;
|
||||
for (let si = 0; si < ed.splines.length; si++) {
|
||||
for (let pi = 0; pi < ed.splines[si].length; pi++) {
|
||||
const cp = imgToCanvas(node, ed.splines[si][pi]);
|
||||
if (dist(cp, canvasPos) < tol) {
|
||||
return { si, pi };
|
||||
}
|
||||
}
|
||||
}
|
||||
return null;
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Context menu
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
function buildContextMenu() {
|
||||
const menu = document.createElement("div");
|
||||
Object.assign(menu.style, {
|
||||
display: "none",
|
||||
position: "absolute",
|
||||
background: "#252530",
|
||||
border: "1px solid #555",
|
||||
borderRadius: "6px",
|
||||
padding: "4px 0",
|
||||
zIndex: "1000",
|
||||
boxShadow: "0 4px 12px rgba(0,0,0,0.4)",
|
||||
minWidth: "160px",
|
||||
fontFamily: "system-ui, sans-serif",
|
||||
fontSize: "13px",
|
||||
});
|
||||
return menu;
|
||||
}
|
||||
|
||||
function showMenu(node, canvasEvt, imgPos) {
|
||||
const ed = node._ed;
|
||||
const menu = ed.menu;
|
||||
menu.innerHTML = "";
|
||||
|
||||
const hit = hitTestPoint(node, mouseToCanvas(node, canvasEvt));
|
||||
|
||||
const items = [];
|
||||
items.push({
|
||||
label: "Add Point at Cursor",
|
||||
action: () => {
|
||||
ed.splines[ed.active].push({ x: imgPos.x, y: imgPos.y });
|
||||
ed.dirty = true;
|
||||
syncWidgets(node);
|
||||
},
|
||||
});
|
||||
items.push({
|
||||
label: "Subdivide Nearest Segment",
|
||||
action: () => {
|
||||
const spline = ed.splines[ed.active];
|
||||
if (spline.length < 2) return;
|
||||
const { segIndex } = closestSegment(spline, imgPos);
|
||||
const a = spline[segIndex];
|
||||
const b = spline[segIndex + 1];
|
||||
const mid = { x: (a.x + b.x) / 2, y: (a.y + b.y) / 2 };
|
||||
spline.splice(segIndex + 1, 0, mid);
|
||||
ed.dirty = true;
|
||||
syncWidgets(node);
|
||||
},
|
||||
});
|
||||
items.push({ separator: true });
|
||||
items.push({
|
||||
label: "New Spline",
|
||||
action: () => {
|
||||
ed.splines.push([
|
||||
{ x: imgPos.x - 40, y: imgPos.y },
|
||||
{ x: imgPos.x + 40, y: imgPos.y },
|
||||
]);
|
||||
ed.active = ed.splines.length - 1;
|
||||
ed.dirty = true;
|
||||
syncWidgets(node);
|
||||
},
|
||||
});
|
||||
items.push({
|
||||
label: "New Static Spline",
|
||||
action: () => {
|
||||
ed.splines.push([{ x: imgPos.x, y: imgPos.y }]);
|
||||
ed.active = ed.splines.length - 1;
|
||||
ed.dirty = true;
|
||||
syncWidgets(node);
|
||||
},
|
||||
});
|
||||
items.push({
|
||||
label: "Delete Spline",
|
||||
action: () => {
|
||||
if (ed.splines.length <= 1) return;
|
||||
ed.splines.splice(ed.active, 1);
|
||||
ed.active = Math.min(ed.active, ed.splines.length - 1);
|
||||
ed.dirty = true;
|
||||
syncWidgets(node);
|
||||
},
|
||||
disabled: ed.splines.length <= 1,
|
||||
});
|
||||
if (hit) {
|
||||
const spline = ed.splines[hit.si];
|
||||
if (spline.length > 1) {
|
||||
items.push({ separator: true });
|
||||
items.push({
|
||||
label: "Delete Point",
|
||||
action: () => {
|
||||
spline.splice(hit.pi, 1);
|
||||
ed.dirty = true;
|
||||
syncWidgets(node);
|
||||
},
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
for (const item of items) {
|
||||
if (item.separator) {
|
||||
const hr = document.createElement("div");
|
||||
hr.style.borderTop = "1px solid #444";
|
||||
hr.style.margin = "4px 8px";
|
||||
menu.appendChild(hr);
|
||||
continue;
|
||||
}
|
||||
const el = document.createElement("div");
|
||||
el.textContent = item.label;
|
||||
const isDisabled = item.disabled;
|
||||
Object.assign(el.style, {
|
||||
padding: "6px 14px",
|
||||
cursor: isDisabled ? "default" : "pointer",
|
||||
color: isDisabled ? "#666" : "#ddd",
|
||||
whiteSpace: "nowrap",
|
||||
});
|
||||
if (!isDisabled) {
|
||||
el.addEventListener("mouseenter", () => (el.style.background = "#3a3a4a"));
|
||||
el.addEventListener("mouseleave", () => (el.style.background = "none"));
|
||||
el.addEventListener("click", (e) => {
|
||||
e.stopPropagation();
|
||||
menu.style.display = "none";
|
||||
item.action();
|
||||
});
|
||||
}
|
||||
menu.appendChild(el);
|
||||
}
|
||||
|
||||
const rect = node._ed.canvas.getBoundingClientRect();
|
||||
const containerRect = node._ed.canvas.parentElement.getBoundingClientRect();
|
||||
menu.style.left = `${canvasEvt.clientX - containerRect.left}px`;
|
||||
menu.style.top = `${canvasEvt.clientY - containerRect.top}px`;
|
||||
menu.style.display = "block";
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Event handling
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
function setupEvents(node) {
|
||||
const ed = node._ed;
|
||||
const canvas = ed.canvas;
|
||||
|
||||
canvas.addEventListener("pointerdown", (e) => {
|
||||
ed.menu.style.display = "none";
|
||||
if (e.button === 2) return; // right-click handled by contextmenu
|
||||
if (e.button !== 0) return;
|
||||
e.preventDefault();
|
||||
e.stopPropagation();
|
||||
canvas.setPointerCapture(e.pointerId);
|
||||
|
||||
const cp = mouseToCanvas(node, e);
|
||||
const hit = hitTestPoint(node, cp);
|
||||
if (hit) {
|
||||
ed.active = hit.si;
|
||||
ed.drag = { ...hit, pointerId: e.pointerId };
|
||||
canvas.style.cursor = "grabbing";
|
||||
ed.dirty = true;
|
||||
}
|
||||
});
|
||||
|
||||
canvas.addEventListener("pointermove", (e) => {
|
||||
const cp = mouseToCanvas(node, e);
|
||||
|
||||
if (ed.drag) {
|
||||
e.preventDefault();
|
||||
e.stopPropagation();
|
||||
const imgP = canvasToImg(node, cp);
|
||||
imgP.x = Math.max(0, Math.min(ed.imgW, imgP.x));
|
||||
imgP.y = Math.max(0, Math.min(ed.imgH, imgP.y));
|
||||
ed.splines[ed.drag.si][ed.drag.pi] = imgP;
|
||||
ed.dirty = true;
|
||||
syncWidgets(node);
|
||||
return;
|
||||
}
|
||||
|
||||
const hit = hitTestPoint(node, cp);
|
||||
if (hit) {
|
||||
canvas.style.cursor = "grab";
|
||||
ed.hover = hit;
|
||||
} else {
|
||||
canvas.style.cursor = "default";
|
||||
ed.hover = null;
|
||||
}
|
||||
ed.dirty = true;
|
||||
});
|
||||
|
||||
canvas.addEventListener("pointerup", (e) => {
|
||||
if (ed.drag) {
|
||||
e.preventDefault();
|
||||
e.stopPropagation();
|
||||
canvas.releasePointerCapture(e.pointerId);
|
||||
ed.drag = null;
|
||||
canvas.style.cursor = "default";
|
||||
ed.dirty = true;
|
||||
syncWidgets(node);
|
||||
}
|
||||
});
|
||||
|
||||
canvas.addEventListener("pointerleave", () => {
|
||||
ed.hover = null;
|
||||
ed.dirty = true;
|
||||
});
|
||||
|
||||
canvas.addEventListener("lostpointercapture", () => {
|
||||
if (ed.drag) {
|
||||
ed.drag = null;
|
||||
canvas.style.cursor = "default";
|
||||
ed.dirty = true;
|
||||
syncWidgets(node);
|
||||
}
|
||||
});
|
||||
|
||||
canvas.addEventListener("contextmenu", (e) => {
|
||||
e.preventDefault();
|
||||
e.stopPropagation();
|
||||
const cp = mouseToCanvas(node, e);
|
||||
const imgP = canvasToImg(node, cp);
|
||||
showMenu(node, e, imgP);
|
||||
});
|
||||
|
||||
const docClickHandler = (e) => {
|
||||
if (!ed.menu.contains(e.target)) {
|
||||
ed.menu.style.display = "none";
|
||||
}
|
||||
};
|
||||
document.addEventListener("click", docClickHandler);
|
||||
ed._docClickHandler = docClickHandler;
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Widget synchronisation
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
function syncWidgets(node) {
|
||||
const ed = node._ed;
|
||||
const pointsW = node.widgets.find((w) => w.name === "points_store");
|
||||
const coordsW = node.widgets.find((w) => w.name === "coordinates");
|
||||
const samplesW = node.widgets.find((w) => w.name === "points_to_sample");
|
||||
|
||||
if (pointsW) pointsW.value = JSON.stringify(ed.splines);
|
||||
|
||||
const numSamples = samplesW ? samplesW.value : 121;
|
||||
const interpolated = ed.splines.map((sp) => interpolateSpline(sp, numSamples));
|
||||
// Round to integers for pixel-level coordinates
|
||||
const rounded = interpolated.map((track) =>
|
||||
track.map((p) => ({ x: Math.round(p.x), y: Math.round(p.y) }))
|
||||
);
|
||||
if (coordsW) coordsW.value = JSON.stringify(rounded);
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Rendering
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
function startRenderLoop(node) {
|
||||
const ed = node._ed;
|
||||
|
||||
function loop() {
|
||||
ed.rafId = requestAnimationFrame(loop);
|
||||
resizeCanvas(node);
|
||||
if (!ed.dirty) return;
|
||||
ed.dirty = false;
|
||||
render(node);
|
||||
}
|
||||
|
||||
loop();
|
||||
}
|
||||
|
||||
function resizeCanvas(node) {
|
||||
const ed = node._ed;
|
||||
const canvas = ed.canvas;
|
||||
const dispW = canvas.clientWidth || node.size?.[0] || 400;
|
||||
|
||||
const aspect = ed.imgH / ed.imgW;
|
||||
const dispH = Math.max(CANVAS_MIN_HEIGHT, Math.round(dispW * aspect));
|
||||
|
||||
canvas.style.height = `${dispH}px`;
|
||||
|
||||
const dpr = window.devicePixelRatio || 1;
|
||||
const bufW = Math.round(dispW * dpr);
|
||||
const bufH = Math.round(dispH * dpr);
|
||||
|
||||
if (canvas.width !== bufW || canvas.height !== bufH) {
|
||||
canvas.width = bufW;
|
||||
canvas.height = bufH;
|
||||
ed.dirty = true;
|
||||
}
|
||||
}
|
||||
|
||||
function render(node) {
|
||||
const ed = node._ed;
|
||||
const ctx = ed.ctx;
|
||||
const c = ed.canvas;
|
||||
|
||||
ctx.clearRect(0, 0, c.width, c.height);
|
||||
|
||||
// Background image
|
||||
if (ed.bgImg) {
|
||||
const s = getScale(node);
|
||||
ctx.globalAlpha = 0.85;
|
||||
ctx.drawImage(ed.bgImg, 0, 0, ed.imgW * s, ed.imgH * s);
|
||||
ctx.globalAlpha = 1.0;
|
||||
}
|
||||
|
||||
const samplesW = node.widgets.find((w) => w.name === "points_to_sample");
|
||||
const numSamples = samplesW ? samplesW.value : 121;
|
||||
const dpr = window.devicePixelRatio || 1;
|
||||
|
||||
// Draw each spline
|
||||
for (let si = 0; si < ed.splines.length; si++) {
|
||||
const isActive = si === ed.active;
|
||||
const color = SPLINE_COLORS[si % SPLINE_COLORS.length];
|
||||
const spline = ed.splines[si];
|
||||
|
||||
// Interpolated curve
|
||||
if (spline.length >= 2) {
|
||||
const curve = interpolateSpline(spline, Math.max(numSamples, 60));
|
||||
ctx.beginPath();
|
||||
const p0 = imgToCanvas(node, curve[0]);
|
||||
ctx.moveTo(p0.x, p0.y);
|
||||
for (let i = 1; i < curve.length; i++) {
|
||||
const p = imgToCanvas(node, curve[i]);
|
||||
ctx.lineTo(p.x, p.y);
|
||||
}
|
||||
ctx.strokeStyle = color;
|
||||
ctx.lineWidth = (isActive ? CURVE_LINE_WIDTH * 1.4 : CURVE_LINE_WIDTH) * dpr;
|
||||
ctx.globalAlpha = isActive ? 1.0 : 0.5;
|
||||
ctx.stroke();
|
||||
ctx.globalAlpha = 1.0;
|
||||
}
|
||||
|
||||
// Control points
|
||||
for (let pi = 0; pi < spline.length; pi++) {
|
||||
const cp = imgToCanvas(node, spline[pi]);
|
||||
const isHov = ed.hover && ed.hover.si === si && ed.hover.pi === pi;
|
||||
const isDrag = ed.drag && ed.drag.si === si && ed.drag.pi === pi;
|
||||
const baseR =
|
||||
isHov || isDrag
|
||||
? ACTIVE_POINT_RADIUS
|
||||
: isActive
|
||||
? POINT_RADIUS
|
||||
: POINT_RADIUS * 0.8;
|
||||
const r = baseR * dpr;
|
||||
|
||||
// Outer ring
|
||||
ctx.beginPath();
|
||||
ctx.arc(cp.x, cp.y, r + 2 * dpr, 0, Math.PI * 2);
|
||||
ctx.fillStyle = isActive ? "#fff" : "rgba(255,255,255,0.5)";
|
||||
ctx.fill();
|
||||
|
||||
// Inner fill
|
||||
ctx.beginPath();
|
||||
ctx.arc(cp.x, cp.y, r, 0, Math.PI * 2);
|
||||
ctx.fillStyle = color;
|
||||
ctx.fill();
|
||||
|
||||
// Point number label
|
||||
if (isActive) {
|
||||
ctx.fillStyle = "#fff";
|
||||
ctx.font = `bold ${Math.round(11 * dpr)}px system-ui`;
|
||||
ctx.textAlign = "center";
|
||||
ctx.fillText(String(pi), cp.x, cp.y - r - 4 * dpr);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Active spline indicator
|
||||
if (ed.splines.length > 1) {
|
||||
const dpr = window.devicePixelRatio || 1;
|
||||
ctx.fillStyle = "rgba(0,0,0,0.5)";
|
||||
ctx.fillRect(4 * dpr, 4 * dpr, 130 * dpr, 20 * dpr);
|
||||
ctx.fillStyle = SPLINE_COLORS[ed.active % SPLINE_COLORS.length];
|
||||
ctx.font = `bold ${11 * dpr}px system-ui`;
|
||||
ctx.textAlign = "left";
|
||||
ctx.fillText(
|
||||
`Spline ${ed.active + 1} / ${ed.splines.length}`,
|
||||
10 * dpr,
|
||||
18 * dpr
|
||||
);
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user