Compare commits
45
Commits
pr-2026-01-27
...
master
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
3bf3ca6259 | ||
|
|
5eb48fdcae | ||
|
|
bf2ca0264f | ||
|
|
e6fdfe6686 | ||
|
|
5722b53a08 | ||
|
|
71d8972bd2 | ||
|
|
61ee82b23b | ||
|
|
c75eacae23 | ||
|
|
f8387c893d | ||
|
|
d164816f8a | ||
|
|
cb1f8204dc | ||
|
|
35364168d9 | ||
|
|
dfb2786749 | ||
|
|
63a7f34e64 | ||
|
|
15d09abb5a | ||
|
|
b9446f5972 | ||
|
|
ac4d998390 | ||
|
|
9d1672beb6 | ||
|
|
3b9c5cde47 | ||
|
|
548a39373b | ||
|
|
aceeae9635 | ||
|
|
ec7fe89955 | ||
|
|
4f45fd6c22 | ||
|
|
2e2ac81759 | ||
|
|
229437c6b6 | ||
|
|
75e1ed26ab | ||
|
|
964be9688f | ||
|
|
d50612a64a | ||
|
|
2acf7af899 | ||
|
|
09b90d0f33 | ||
|
|
937831df0e | ||
|
|
b82708a54f | ||
|
|
d43181d572 | ||
|
|
5984491048 | ||
|
|
3686d832c6 | ||
|
|
6df1f3ff90 | ||
|
|
531512f728 | ||
|
|
c5f4724bc2 | ||
|
|
58a5c12c2d | ||
|
|
5c3d20803d | ||
|
|
82bd963cde | ||
|
|
bb70ecba66 | ||
|
|
49add6dddb | ||
|
|
f82614dc6d | ||
|
|
d153ca3f78 |
@@ -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 Dub-It; two stages (with upsampling)`](./example_workflows/2.3/LTX-2.3_ICLoRA_DubIt_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 upgrades SDR video into ACEScct HDR (same working space as pipelines `hdr_ic_lora`), suitable for grading, EXR export, and HLG masters. Use the tiled 2.5 graph [`LTX-2.5_V2V_TiledFusion_SDR_to_HDR.json`](./example_workflows/2.5/LTX-2.5_V2V_TiledFusion_SDR_to_HDR.json), or the 2.3 single-stage [`LTX-2.3_ICLoRA_HDR_Distilled.json`](./example_workflows/2.3/LTX-2.3_ICLoRA_HDR_Distilled.json).
|
||||
|
||||
### 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 via **OpenImageIO** (`pip install openimageio`, also listed in `requirements.txt`). 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)).
|
||||
|
||||
## Dub-It IC-LoRA
|
||||
|
||||
We provide a **Dub-It 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-dubit-0.9.safetensors`](https://huggingface.co/Lightricks/LTX-2.3-22b-IC-LoRA-DubIt)
|
||||
* [`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,10 +161,40 @@ 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
|
||||
|
||||
### Tiled Fusion Sampler
|
||||
|
||||
`LTXVTiledFusionSampler` (🅛🅣🅧 Tiled Fusion Sampler) denoises a canvas larger
|
||||
than the model's trained window by running each denoise step on overlapping
|
||||
spatial tiles that share one latent canvas and one noise field. It is meant for
|
||||
IC-LoRA video-to-video where `LTXAddVideoICLoRAGuide` has already attached the
|
||||
guide. Pair it with `LTXVGetTilingSizes` (🅛🅣🅧 Get Tiling Sizes) for legal
|
||||
canvas / tile / `tile_frames` values. See
|
||||
[tiled_fusion_sampler.md](./tiled_fusion_sampler.md).
|
||||
Do not confuse it with `LTXVTiledSampler`, which samples tiles independently.
|
||||
|
||||
Example graphs (2.5 distilled):
|
||||
|
||||
* [`V2V upscale`](./example_workflows/2.5/LTX-2.5_V2V_TiledFusion_Upscale.json) — detail-refine at FullHD / 4K / 8K
|
||||
* [`V2V native 4K/8K`](./example_workflows/2.5/LTX-2.5_V2V_TiledFusion_Native_4K_8K.json) — FullHD stage 1, then 4K or 8K
|
||||
* [`V2V SDR→HDR`](./example_workflows/2.5/LTX-2.5_V2V_TiledFusion_SDR_to_HDR.json) — same ladder with the HDR IC-LoRA
|
||||
|
||||
Prompt those graphs for look and style only (lighting, palette, sharpness, grade). Do not name specific objects in the scene: after upscale, named details can appear in every spatial tile.
|
||||
|
||||
### Low VRAM
|
||||
* For systems with low VRAM you can use the model loader nodes from [low_vram_loaders.py](./low_vram_loaders.py). Those nodes ensure the correct order of execution and perform the model offloading such that generation fits in 32 GB VRAM.
|
||||
* Use --reserve-vram ComfyUI parameter: `python -m main --reserve-vram 5` (or other number in GB).
|
||||
|
||||
+48
-4
@@ -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 (
|
||||
@@ -9,8 +12,21 @@ from .easy_samplers import (
|
||||
)
|
||||
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_nodes import (
|
||||
LTXVHDRDecodePostprocess,
|
||||
LTXVLoadEXRSequence,
|
||||
LTXVSaveHLG,
|
||||
LTXVSDRToHDRWorkingSpace,
|
||||
LTXVVAEForceFloat32,
|
||||
)
|
||||
from .iclora import (
|
||||
LTXAddVideoICLoRAGuide,
|
||||
LTXAddVideoICLoRAGuideAdvanced,
|
||||
LTXICLoRALoaderModelOnly,
|
||||
LTXVSetAudioRefTokens,
|
||||
)
|
||||
from .latent_norm import (
|
||||
LTXVAdainLatent,
|
||||
LTXVPerStepAdainPatcher,
|
||||
@@ -18,7 +34,6 @@ from .latent_norm import (
|
||||
LTXVStatNormLatent,
|
||||
)
|
||||
from .latents import (
|
||||
LTXVAddLatentGuide,
|
||||
LTXVImgToVideoConditionOnly,
|
||||
LTXVSelectLatents,
|
||||
LTXVSetVideoLatentNoiseMasks,
|
||||
@@ -36,26 +51,31 @@ 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,
|
||||
STGGuiderAdvancedNode,
|
||||
STGGuiderNode,
|
||||
)
|
||||
from .tiled_fusion_sampler import LTXVTiledFusionSampler
|
||||
from .tiled_fusion_sizes import LTXVGetTilingSizes
|
||||
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,
|
||||
"LTXVAddLatentGuide": LTXVAddLatentGuide,
|
||||
"LTXVAddGuideAdvancedAttention": LTXVAddGuideAdvancedAttention,
|
||||
"LTXVAdainLatent": LTXVAdainLatent,
|
||||
"LTXVImgToVideoConditionOnly": LTXVImgToVideoConditionOnly,
|
||||
"LTXVPerStepAdainPatcher": LTXVPerStepAdainPatcher,
|
||||
@@ -73,6 +93,8 @@ NODE_CLASS_MAPPINGS = {
|
||||
"LTXVSelectLatents": LTXVSelectLatents,
|
||||
"LTXVSetVideoLatentNoiseMasks": LTXVSetVideoLatentNoiseMasks,
|
||||
"LTXVTiledSampler": LTXVTiledSampler,
|
||||
"LTXVTiledFusionSampler": LTXVTiledFusionSampler,
|
||||
"LTXVGetTilingSizes": LTXVGetTilingSizes,
|
||||
"LTXVLoopingSampler": LTXVLoopingSampler,
|
||||
"LTXVTiledVAEDecode": LTXVTiledVAEDecode,
|
||||
"MultimodalGuider": MultimodalGuiderNode,
|
||||
@@ -82,6 +104,7 @@ NODE_CLASS_MAPPINGS = {
|
||||
"STGGuiderNode": STGGuiderNode,
|
||||
"LTXVMultiPromptProvider": MultiPromptProvider,
|
||||
"ImageToCPU": ImageToCPU,
|
||||
"LTXFloatToInt": FloatToInt,
|
||||
"LTXVStatNormLatent": LTXVStatNormLatent,
|
||||
"LTXVPerStepStatNormPatcher": LTXVPerStepStatNormPatcher,
|
||||
"LTXVGemmaCLIPModelLoader": LTXVGemmaCLIPModelLoader,
|
||||
@@ -91,6 +114,24 @@ NODE_CLASS_MAPPINGS = {
|
||||
"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,
|
||||
"LTXVSDRToHDRWorkingSpace": LTXVSDRToHDRWorkingSpace,
|
||||
"LTXVLoadEXRSequence": LTXVLoadEXRSequence,
|
||||
"LTXVVAEForceFloat32": LTXVVAEForceFloat32,
|
||||
"LTXVSaveHLG": LTXVSaveHLG,
|
||||
"LTXVAudioOnlyModel": LTXVAudioOnlyModel,
|
||||
"LTXVAudioOnlyEmptyVideoLatent": LTXVAudioOnlyEmptyVideoLatent,
|
||||
}
|
||||
|
||||
# Consistent display names between static and dynamic node mappings in nodes_registry.py,
|
||||
@@ -108,8 +149,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,163 @@
|
||||
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
|
||||
|
||||
# Pipeline HDR IC-LoRA scene embeddings (`hdr_ic_lora._load_video_context`).
|
||||
_PIPELINE_VIDEO_KEYS = ("video_context", "video_prompt_embeds")
|
||||
_PIPELINE_AUDIO_KEYS = ("audio_context", "audio_prompt_embeds")
|
||||
# LTX-AV text encode layout: video 4096 + audio 2048 on the last dim.
|
||||
_VIDEO_CONTEXT_DIM = 4096
|
||||
_AUDIO_CONTEXT_DIM = 2048
|
||||
_AV_CONTEXT_DIM = _VIDEO_CONTEXT_DIM + _AUDIO_CONTEXT_DIM
|
||||
|
||||
|
||||
def _unsqueeze_batch(t: torch.Tensor, *, source: str, kind: str) -> torch.Tensor:
|
||||
"""Ensure conditioning tensors are ``[B, T, D]`` (pipeline files are often ``[T, D]``)."""
|
||||
if t.ndim == 2:
|
||||
return t.unsqueeze(0)
|
||||
if t.ndim == 3:
|
||||
return t
|
||||
raise ValueError(
|
||||
f"{source}: {kind} embedding must be 2D [T, D] or 3D [B, T, D]; "
|
||||
f"got shape {tuple(t.shape)}"
|
||||
)
|
||||
|
||||
|
||||
def _load_pipeline_scene_embedding(f: Any, source: str) -> list[list[Any]]:
|
||||
"""Load ``video_context`` (+ optional ``audio_context``) into Comfy CONDITIONING.
|
||||
|
||||
Matches ``ltx_pipelines.hdr_ic_lora``: the safetensors replaces text encoding.
|
||||
For AV Comfy models, video (4096) and audio (2048) contexts are concatenated on
|
||||
the last dim — same layout as a fully processed LTX-AV text encode.
|
||||
"""
|
||||
keys = list(f.keys())
|
||||
video_key = next((k for k in _PIPELINE_VIDEO_KEYS if k in keys), None)
|
||||
if video_key is None:
|
||||
return []
|
||||
|
||||
video = _unsqueeze_batch(f.get_tensor(video_key), source=source, kind=video_key)
|
||||
video_dim = int(video.shape[-1])
|
||||
audio_key = next((k for k in _PIPELINE_AUDIO_KEYS if k in keys), None)
|
||||
|
||||
if video_dim == _AV_CONTEXT_DIM:
|
||||
# Already video∥audio on the last dim (rare export); use as-is.
|
||||
tensor = video
|
||||
elif video_dim == _VIDEO_CONTEXT_DIM:
|
||||
if audio_key is None:
|
||||
raise ValueError(
|
||||
f"{source}: {video_key} last dim is {_VIDEO_CONTEXT_DIM} but no "
|
||||
f"audio_context / audio_prompt_embeds found. AV conditioning needs "
|
||||
f"{_VIDEO_CONTEXT_DIM}+{_AUDIO_CONTEXT_DIM}="
|
||||
f"{_AV_CONTEXT_DIM} on the last dim."
|
||||
)
|
||||
audio = _unsqueeze_batch(f.get_tensor(audio_key), source=source, kind=audio_key)
|
||||
audio_dim = int(audio.shape[-1])
|
||||
if audio_dim != _AUDIO_CONTEXT_DIM:
|
||||
raise ValueError(
|
||||
f"{source}: {audio_key} last dim must be {_AUDIO_CONTEXT_DIM}; "
|
||||
f"got {audio_dim} with shape {tuple(audio.shape)}"
|
||||
)
|
||||
if audio.shape[:-1] != video.shape[:-1]:
|
||||
raise ValueError(
|
||||
f"{source}: audio/video context shape mismatch "
|
||||
f"(excluding last dim): {tuple(audio.shape)} vs {tuple(video.shape)}"
|
||||
)
|
||||
tensor = torch.cat([video, audio], dim=-1)
|
||||
else:
|
||||
raise ValueError(
|
||||
f"{source}: {video_key} last dim must be {_VIDEO_CONTEXT_DIM} "
|
||||
f"(or pre-concatenated {_AV_CONTEXT_DIM}); got {video_dim} with shape "
|
||||
f"{tuple(video.shape)}. Check that video/audio keys are not swapped."
|
||||
)
|
||||
|
||||
return [[tensor, {}]]
|
||||
|
||||
|
||||
@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",
|
||||
description=(
|
||||
"Load precomputed text conditioning from models/embeddings. "
|
||||
"Supports Comfy conditioning_data_* files and pipeline scene "
|
||||
"embeddings (video_context / audio_context), e.g. HDR IC-LoRA."
|
||||
),
|
||||
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_")]
|
||||
|
||||
if tensor_keys:
|
||||
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])
|
||||
else:
|
||||
conditioning = _load_pipeline_scene_embedding(f, source=file_name)
|
||||
|
||||
if not conditioning:
|
||||
raise ValueError(
|
||||
f"No conditioning data found in file: {file_name}. "
|
||||
"Expected conditioning_data_* or video_context / video_prompt_embeds."
|
||||
)
|
||||
|
||||
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)")
|
||||
)
|
||||
+8
-8
@@ -12,7 +12,7 @@ from comfy_extras.nodes_lt import EmptyLTXVLatentVideo, LTXVAddGuide, LTXVCropGu
|
||||
|
||||
from .guide import blur_internal
|
||||
from .latent_norm import LTXVAdainLatent
|
||||
from .latents import LTXVAddLatentGuide, LTXVSelectLatents
|
||||
from .latents import LTXVSelectLatents, add_latent_guide
|
||||
from .nodes_registry import comfy_node
|
||||
|
||||
|
||||
@@ -248,7 +248,7 @@ class LTXVBaseSampler:
|
||||
positive,
|
||||
negative,
|
||||
latents,
|
||||
) = LTXVAddLatentGuide().generate(
|
||||
) = add_latent_guide(
|
||||
vae=vae,
|
||||
positive=positive,
|
||||
negative=negative,
|
||||
@@ -446,7 +446,7 @@ class LTXVExtendSampler:
|
||||
positive,
|
||||
negative,
|
||||
new_latents,
|
||||
) = LTXVAddLatentGuide().generate(
|
||||
) = add_latent_guide(
|
||||
vae=vae,
|
||||
positive=positive,
|
||||
negative=negative,
|
||||
@@ -504,7 +504,7 @@ class LTXVExtendSampler:
|
||||
positive,
|
||||
negative,
|
||||
new_latents,
|
||||
) = LTXVAddLatentGuide().generate(
|
||||
) = add_latent_guide(
|
||||
vae=vae,
|
||||
positive=positive,
|
||||
negative=negative,
|
||||
@@ -519,7 +519,7 @@ class LTXVExtendSampler:
|
||||
positive,
|
||||
negative,
|
||||
new_latents,
|
||||
) = LTXVAddLatentGuide().generate(
|
||||
) = add_latent_guide(
|
||||
vae=vae,
|
||||
positive=positive,
|
||||
negative=negative,
|
||||
@@ -552,7 +552,7 @@ class LTXVExtendSampler:
|
||||
positive,
|
||||
negative,
|
||||
denoised_output_latents,
|
||||
) = LTXVAddLatentGuide().generate(
|
||||
) = add_latent_guide(
|
||||
vae=vae,
|
||||
positive=positive,
|
||||
negative=negative,
|
||||
@@ -756,7 +756,7 @@ class LTXVInContextSampler:
|
||||
positive,
|
||||
negative,
|
||||
new_latents,
|
||||
) = LTXVAddLatentGuide().generate(
|
||||
) = add_latent_guide(
|
||||
vae=vae,
|
||||
positive=positive,
|
||||
negative=negative,
|
||||
@@ -792,7 +792,7 @@ class LTXVInContextSampler:
|
||||
positive,
|
||||
negative,
|
||||
new_latents,
|
||||
) = LTXVAddLatentGuide().generate(
|
||||
) = add_latent_guide(
|
||||
vae=vae,
|
||||
positive=positive,
|
||||
negative=negative,
|
||||
|
||||
+144
-3
@@ -2,16 +2,29 @@ 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,
|
||||
)
|
||||
from torch import nn
|
||||
|
||||
try:
|
||||
# New core (since ComfyUI PR #15056) replaced interleaved_freqs_cis/split_freqs_cis
|
||||
# with a single freqs_cis_matrix helper.
|
||||
from comfy.ldm.lightricks.model import freqs_cis_matrix
|
||||
|
||||
_USE_FREQS_CIS_MATRIX = True
|
||||
except ImportError:
|
||||
# Legacy core: cos/sin freqs consumed as a (cos, sin, split) tuple.
|
||||
from .rope_utils import interleaved_freqs_cis, split_freqs_cis
|
||||
|
||||
_USE_FREQS_CIS_MATRIX = False
|
||||
|
||||
from .pos_embedding_exp_values import POS_EMBEDDING_EXP_VALUES
|
||||
|
||||
|
||||
@@ -52,6 +65,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 +79,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 +138,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 +163,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,
|
||||
@@ -244,11 +261,23 @@ class Embeddings1DConnector(nn.Module):
|
||||
expected_freqs = dim // 2
|
||||
current_freqs = freqs.shape[-1]
|
||||
pad_size = expected_freqs - current_freqs
|
||||
else:
|
||||
pad_size = dim % n_elem
|
||||
|
||||
if _USE_FREQS_CIS_MATRIX:
|
||||
# New core takes (rotation_matrix, split_pe); let it build the matrix so
|
||||
# this cannot drift from the transformer again.
|
||||
return freqs_cis_matrix(
|
||||
freqs, pad_size, self.split_rope, self.num_attention_heads, self.dtype
|
||||
)
|
||||
|
||||
# Legacy core: returns (cos, sin, split_rope).
|
||||
if self.split_rope:
|
||||
cos_freq, sin_freq = split_freqs_cis(
|
||||
freqs, pad_size, self.num_attention_heads
|
||||
)
|
||||
else:
|
||||
cos_freq, sin_freq = interleaved_freqs_cis(freqs, dim % n_elem)
|
||||
cos_freq, sin_freq = interleaved_freqs_cis(freqs, pad_size)
|
||||
return cos_freq.to(self.dtype), sin_freq.to(self.dtype), self.split_rope
|
||||
|
||||
def forward(
|
||||
@@ -330,3 +359,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,
|
||||
)
|
||||
|
||||
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 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
File diff suppressed because it is too large
Load Diff
+10678
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
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
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,154 @@
|
||||
# LTX-2.5 example workflows
|
||||
|
||||
These graphs generate video (and usually audio) with the **LTX-2.5 distilled** model. They share the same subgraph layout — load models, set inputs, preprocess, sample, decode — so you can switch files without relearning the graph.
|
||||
|
||||
All of them use the distilled transformer (`ltx-2.5-22b-distilled-transformer-bf16.safetensors`): a few sampling steps, meant for fast iteration. Two-stage graphs add a 2× spatial upscale and a short refine pass; single-stage graphs skip that.
|
||||
|
||||
Default values are tuned for lower-VRAM machines. If you have more memory, you can speed things up by increasing the **decode tile size** — see the Decode notes inside each graph (fewer, larger tiles run faster but need more VRAM).
|
||||
|
||||
## Getting started
|
||||
|
||||
This assumes ComfyUI is already installed. If not, see the [ComfyUI download page](https://www.comfy.org/download), the [ComfyUI setup guide](https://docs.ltx.io/open-source-model/integration-tools/comfy-ui), and the [system requirements](https://docs.ltx.io/open-source-model/getting-started/system-requirements).
|
||||
|
||||
1. Open ComfyUI.
|
||||
2. Load a workflow from this folder (**Workflow** → **Open**, or drag the `.json` onto the canvas). The built-in ComfyUI **Templates** search for **LTX-2.5** is the same family of graphs if you prefer that entry point.
|
||||
3. Open the **Workflow Overview** panel. On first use it lists **Missing Models**.
|
||||
4. Click **Download all** to fetch the files into the right `ComfyUI/models/` folders. You only do this once — later runs reuse them.
|
||||
5. Fill in the left-hand **Inputs** (prompt, optional image, duration, and anything the graph asks for such as audio or a reference video).
|
||||
6. Click **Run**.
|
||||
|
||||
Weights also live on the [LTX-2.5 Hugging Face repo](https://huggingface.co/Lightricks/LTX-2.5) if you want to download them by hand. Each graph’s **Model Links** note lists the exact files it needs, including any IC-LoRA.
|
||||
|
||||
Prompting tips: [prompting guide](https://docs.ltx.io/open-source-model/usage-guides/prompting-guide). A beginner walkthrough of text-to-video: [Text-to-Video workflow](https://docs.ltx.io/open-source-model/usage-guides/text-to-video).
|
||||
|
||||
Frame count must be `1 + a multiple of 8`. Duration in the graphs is converted from fps × seconds and may land slightly off the number you typed.
|
||||
|
||||
## Workflows
|
||||
|
||||
### Generation
|
||||
|
||||
| Workflow | What it does |
|
||||
| -------- | ------------ |
|
||||
| [LTX-2.5_T2V_I2V_Two_Stage_Distilled.json](./LTX-2.5_T2V_I2V_Two_Stage_Distilled.json) | **Start here.** Text-to-video, optional first-frame image. Stage 1 at base resolution, 2× spatial upscale, then a 3-step refine. Generates video and audio together. |
|
||||
| [LTX-2.5_T2V_I2V_Single_Stage_Distilled.json](./LTX-2.5_T2V_I2V_Single_Stage_Distilled.json) | Same text / image inputs, one distilled pass, no upscaler. Faster and lighter; lower spatial detail than two-stage. |
|
||||
| [LTX-2.5_A2V_Two_Stage_Distilled.json](./LTX-2.5_A2V_Two_Stage_Distilled.json) | Audio-to-video. Encodes an input clip (with start + duration trim), **freezes** those tokens in both stages, and muxes the original waveform into the output (no audio decode). Optional first-frame image. |
|
||||
| [LTX-2.5_T2A_Single_Stage_Distilled.json](./LTX-2.5_T2A_Single_Stage_Distilled.json) | Text-to-audio only. No video VAE or upscaler. |
|
||||
|
||||
### IC-LoRA (control and edit)
|
||||
|
||||
These keep the distilled 2.5 backbone and add an IC-LoRA so a guide (video, image sheet, tracks, or mask) steers generation.
|
||||
|
||||
| Workflow | What it does |
|
||||
| -------- | ------------ |
|
||||
| [LTX-2.5_ICLoRA_Union_Control_Distilled.json](./LTX-2.5_ICLoRA_Union_Control_Distilled.json) | Video-to-video from a **depth / canny / pose** annotator on a reference clip. Depth is wired by default; switch the annotator inside the graph. |
|
||||
| [LTX-2.5_V2V_ICLoRA_Single_Stage_Distilled.json](./LTX-2.5_V2V_ICLoRA_Single_Stage_Distilled.json) | Video-to-video using the source frames as IC-LoRA guides. Ships with the Instant Shave LoRA as the example; swap the LoRA for other identity/edit LoRAs. Original audio is frozen through. |
|
||||
| [LTX-2.5_ICLoRA_SDR_to_HDR_Distilled.json](./HDR_workflows/LTX-2.5_ICLoRA_SDR_to_HDR_Distilled.json) | **SDR→HDR**: SDR MP4 → ACEScct IDT → HDR IC-LoRA → EXR + HLG. |
|
||||
| [LTX-2.5_ICLoRA_Ingredients_Single_Stage_Distilled.json](./LTX-2.5_ICLoRA_Ingredients_Single_Stage_Distilled.json) | Generate from a **reference sheet** (characters, props, wardrobe, location laid out in one image). Describe each element in the prompt by its place on the sheet. |
|
||||
| [LTX-2.5_ICLoRA_Motion_Track_Distilled.json](./LTX-2.5_ICLoRA_Motion_Track_Distilled.json) | Image-to-video with **sparse motion tracks** you draw on the first frame. Optional still as the opening frame. |
|
||||
| [LTX-2.5_ICLoRA_Inpaint_Two_Stage_Distilled.json](./LTX-2.5_ICLoRA_Inpaint_Two_Stage_Distilled.json) | Fill masked regions of a reference video (two-stage). Source audio can stay frozen. |
|
||||
| [LTX-2.5_ICLoRA_Inpaint_HDR_Two_Stage_Distilled.json](./HDR_workflows/LTX-2.5_ICLoRA_Inpaint_HDR_Two_Stage_Distilled.json) | Same inpaint, but **HDR→HDR**: EXR sequence + mask video → EXR + HLG out. |
|
||||
| [LTX-2.5_ICLoRA_Outpaint_Two_Stage_Distilled.json](./LTX-2.5_ICLoRA_Outpaint_Two_Stage_Distilled.json) | Extend the canvas of a reference video (two-stage). Same in/outpaint LoRA as inpaint. |
|
||||
| [LTX-2.5_I2V_HDR_Two_Stage_Distilled.json](./HDR_workflows/LTX-2.5_I2V_HDR_Two_Stage_Distilled.json) | **HDR still → HDR video**: EXR still → float32 VAE → EXR + HLG out. |
|
||||
|
||||
### Tiled Fusion (large-canvas IC-LoRA)
|
||||
|
||||
These keep one latent canvas and fuse overlapping spatial tiles after every denoise step (`LTXVTiledFusionSampler`). Use them when the canvas is larger than the LoRA window. They share the same subgraph layout as the other 2.5 graphs (Load Models, Inputs, Preprocess, Generate, Decode). **Get Tiling Sizes** and the canvas resize live inside **Preprocess**; the size combos (`tile_size` / `initial_canvas_size` = qHD / HD / FullHD, `output_size` = FullHD / 4K / 8K) are widgets on that subgraph. Source frames come from Load Video, or from Load EXR still on the native HDR→HDR graph. The node emits model-legal canvases (multiples of 32, `tile_frames` = 8n+1); a **Tiling sizes** preview sits below Preprocess. There is no first-frame still / `use image input` / I2V bypass. `use_tiled_encode` on the guide node stays **false**; `use_streaming` is on so each temporal window gets a fresh IC-LoRA encode.
|
||||
|
||||
Upscale prompts should describe **look and style** (lighting, palette, film stock, sharpness, grade), not named objects, people, or props — anything you name can stamp into every spatial tile after upscale. Native and HDR: the stage-1 prompt describes the **change** you want on the input video or HDR plate (plus enhance / API / [prompting guide](https://docs.ltx.io/open-source-model/usage-guides/prompting-guide)); the stage-2 prompt is look/style only, because only that pass uses spatial tiling. The negative prompt is used as written.
|
||||
|
||||
| Workflow | What it does |
|
||||
| -------- | ------------ |
|
||||
| [LTX-2.5_V2V_TiledFusion_Upscale.json](./LTX-2.5_V2V_TiledFusion_Upscale.json) | **V2V upscale.** Preprocess resizes the guide to `output_size` (FullHD / 4K / 8K), then 8 fused steps with the detail refiner. Tiles come from `tile_size` (HD). |
|
||||
| [LTX-2.5_V2V_TiledFusion_Native_4K_8K.json](./LTX-2.5_V2V_TiledFusion_Native_4K_8K.json) | **V2V native / dynamic 4K/8K.** Stage 1 at `initial_canvas_size` FullHD (8 steps, Day-To-Night IC-LoRA). Latent ×2 (×4 when `output_size` is 8K), then 3 fused steps with a fresh prompt and the original clip resized to output. Optional second IC-LoRA on stage 2. |
|
||||
| [LTX-2.5_V2V_TiledFusion_SDR_to_HDR.json](./LTX-2.5_V2V_TiledFusion_SDR_to_HDR.json) | Same recipe as Native, with the HDR IC-LoRA. SDR MP4 → ACEScct, then `LTXVHDRDecodePostprocess` (EXR + HLG). |
|
||||
| [LTX-2.5_V2V_TiledFusion_HDR_to_HDR.json](./LTX-2.5_V2V_TiledFusion_HDR_to_HDR.json) | **Native HDR→HDR refine.** EXR still (ACEScct, float32 VAE, `img_compression` 0) with the detail refiner instead of the HDR IC-LoRA. Same 4K/8K tiled ladder; EXR + HLG out. |
|
||||
|
||||
## Which workflow should I use?
|
||||
|
||||
```text
|
||||
Do you only need audio (no picture)?
|
||||
├─ YES → LTX-2.5_T2A_Single_Stage_Distilled.json
|
||||
│
|
||||
Do you have an audio file the video should follow?
|
||||
├─ YES → LTX-2.5_A2V_Two_Stage_Distilled.json
|
||||
│ (optional still as the first frame)
|
||||
│
|
||||
Do you have an existing video to edit or follow?
|
||||
├─ YES → What kind of control?
|
||||
│ ├─ Depth / edges / pose from the clip
|
||||
│ │ → LTX-2.5_ICLoRA_Union_Control_Distilled.json
|
||||
│ ├─ Keep the footage, change appearance (identity / style LoRA)
|
||||
│ │ → LTX-2.5_V2V_ICLoRA_Single_Stage_Distilled.json
|
||||
│ ├─ Upgrade SDR footage to HDR (ACEScct IC-LoRA)
|
||||
│ │ → HDR_workflows/LTX-2.5_ICLoRA_SDR_to_HDR_Distilled.json
|
||||
│ ├─ Upscale an existing clip at full HD / 4K / 8K (tiled fusion)
|
||||
│ │ → LTX-2.5_V2V_TiledFusion_Upscale.json
|
||||
│ ├─ Viewport-style 4K/8K ladder (full HD, then ×2 / ×4)
|
||||
│ │ → LTX-2.5_V2V_TiledFusion_Native_4K_8K.json
|
||||
│ ├─ Upgrade SDR footage to HDR (ACEScct IC-LoRA)
|
||||
│ │ → LTX-2.5_V2V_TiledFusion_SDR_to_HDR.json (tiled 4K/8K ladder)
|
||||
│ │ → ../2.3/LTX-2.3_ICLoRA_HDR_Distilled.json (single-stage)
|
||||
│ ├─ Refine HDR stills / EXR (native ACEScct I/O, refine LoRA)
|
||||
│ │ → LTX-2.5_V2V_TiledFusion_HDR_to_HDR.json
|
||||
│ ├─ Fill a masked region
|
||||
│ │ → LTX-2.5_ICLoRA_Inpaint_Two_Stage_Distilled.json (SDR)
|
||||
│ │ → HDR_workflows/LTX-2.5_ICLoRA_Inpaint_HDR_Two_Stage_Distilled.json (HDR EXR)
|
||||
│ └─ Grow the frame / canvas
|
||||
│ → LTX-2.5_ICLoRA_Outpaint_Two_Stage_Distilled.json
|
||||
│
|
||||
Do you have stills rather than a video?
|
||||
├─ YES → HDR EXR still → video?
|
||||
│ ├─ YES → HDR_workflows/LTX-2.5_I2V_HDR_Two_Stage_Distilled.json
|
||||
│ └─ NO → A sheet of characters / props / locations?
|
||||
│ ├─ YES → LTX-2.5_ICLoRA_Ingredients_Single_Stage_Distilled.json
|
||||
│ └─ NO → Draw motion paths on a first frame?
|
||||
│ ├─ YES → LTX-2.5_ICLoRA_Motion_Track_Distilled.json
|
||||
│ └─ NO → Text-to-video with an optional opening still
|
||||
│ (see below)
|
||||
│
|
||||
Text-to-video (optional image)?
|
||||
├─ Want 2× spatial upscale and a refine pass?
|
||||
│ └─ YES → LTX-2.5_T2V_I2V_Two_Stage_Distilled.json (recommended)
|
||||
└─ Want the fastest / lightest run?
|
||||
└─ YES → LTX-2.5_T2V_I2V_Single_Stage_Distilled.json
|
||||
```
|
||||
|
||||
Two-stage is the better default when you care about spatial detail. Single-stage is for previews and tighter VRAM.
|
||||
|
||||
## Comparison
|
||||
|
||||
| Workflow | Stages | Upsample | Main conditioning | Best for |
|
||||
| -------- | ------ | -------- | ----------------- | -------- |
|
||||
| T2V / I2V two-stage | 2 | 2× spatial | Text + optional image | Default generation |
|
||||
| T2V / I2V single-stage | 1 | — | Text + optional image | Fast previews |
|
||||
| A2V two-stage | 2 | 2× spatial | Audio (frozen) + optional image | Video that must match a soundtrack |
|
||||
| T2A single-stage | 1 | — | Text | Audio-only clips |
|
||||
| Union Control | 1 | — | Reference video (depth / canny / pose) | Structure-following v2v |
|
||||
| V2V IC-LoRA | 1 | — | Source video (+ frozen audio) | Appearance edits on existing footage |
|
||||
| [SDR→HDR IC-LoRA](./HDR_workflows/LTX-2.5_ICLoRA_SDR_to_HDR_Distilled.json) | 1 | — | SDR video (+ frozen audio) | SDR→ACEScct HDR; EXR + HLG out |
|
||||
| Ingredients | 1 | — | Reference sheet | Cast / props / location from one image |
|
||||
| Motion Track | 1 | — | Image + drawn tracks | Directed motion from a still |
|
||||
| Inpaint two-stage | 2 | 2× spatial | Video + mask (+ frozen audio) | Replacing a region |
|
||||
| Outpaint two-stage | 2 | 2× spatial | Video + target size (+ frozen audio) | Extending the frame |
|
||||
| [I2V HDR two-stage](./HDR_workflows/LTX-2.5_I2V_HDR_Two_Stage_Distilled.json) | 2 | 2× spatial | EXR still (ACEScct) | HDR still → HDR video; EXR + HLG out |
|
||||
| [Inpaint HDR two-stage](./HDR_workflows/LTX-2.5_ICLoRA_Inpaint_HDR_Two_Stage_Distilled.json) | 2 | 2× spatial | EXR sequence + mask video | HDR inpaint; EXR + HLG out |
|
||||
| Tiled Fusion upscale | 1 | — | Source video as IC-LoRA guide | Detail-refine at full HD / 4K / 8K |
|
||||
| Tiled Fusion native 4K/8K | 2 | 2× or 4× latent | Stage 1 Day-To-Night + stage 2 original@x2 | Large-canvas V2V ladder |
|
||||
| Tiled Fusion SDR→HDR | 2 | 2× or 4× latent | Same as native, HDR IC-LoRA | SDR→ACEScct 4K/8K; EXR + HLG |
|
||||
| Tiled Fusion HDR→HDR | 2 | 2× or 4× latent | HDR EXR still + refine IC-LoRA | Native ACEScct refine; EXR + HLG |
|
||||
|
||||
Native HDR I/O: [LTX-2.5_V2V_TiledFusion_HDR_to_HDR.json](./LTX-2.5_V2V_TiledFusion_HDR_to_HDR.json) — `LTXVLoadEXRSequence` → ACEScct → `LTXVVAEForceFloat32` → refine IC-LoRA tiled fusion → `LTXVHDRDecodePostprocess` (`transfer=acescct`, optional EXR) → `LTXVSaveHLG` from `hdr_linear`. `img_compression = 0`. Same idea as pipelines `--hdr` on Distilled / IC-LoRA.
|
||||
|
||||
SDR→HDR IC-LoRA (separate from native EXR HDR): [LTX-2.3_ICLoRA_HDR_Distilled.json](../2.3/LTX-2.3_ICLoRA_HDR_Distilled.json) for a single-stage pass, or [LTX-2.5_V2V_TiledFusion_SDR_to_HDR.json](./LTX-2.5_V2V_TiledFusion_SDR_to_HDR.json) for the tiled 4K/8K ladder. `LoadVideo` → `LTXVSDRToHDRWorkingSpace` (`srgb_gamma`) → HDR IC-LoRA guide → float32 VAE decode → `LTXVHDRDecodePostprocess` (`transfer=acescct`, EXR=`acescct`) → HLG. Sampler **euler**; original audio frozen/muxed. Same idea as pipelines `hdr_ic_lora --text-embeddings`.
|
||||
|
||||
For native HDR, set **`img_compression = 0`** (Python `--hdr` EXR conditioning skips JPEG/CRF; CRF on ACEScct darkens highlights). Turn **prompt enhancer off** when matching goldens. I2V stage-2 distilled sigmas should be `0.909375, 0.725, 0.421875, 0.0` (not `0.85, …`). Inpaint keeps its own 2-step stage-2 schedule.
|
||||
|
||||
HDR graphs live under [`HDR_workflows/`](./HDR_workflows/).
|
||||
|
||||
Native HDR I/O: `LTXVLoadEXRSequence` → float32 VAE → `LTXVHDRDecodePostprocess` (`transfer=acescct`, optional EXR) → `LTXVSaveHLG` from `hdr_linear`. Same idea as pipelines `--hdr` on Distilled / IC-LoRA.
|
||||
|
||||
SDR→HDR IC-LoRA (separate from native EXR HDR): [LTX-2.5_ICLoRA_SDR_to_HDR_Distilled.json](./HDR_workflows/LTX-2.5_ICLoRA_SDR_to_HDR_Distilled.json) — `LoadVideo` → `LTXVSDRToHDRWorkingSpace` (`srgb_gamma`) → `Resize Image/Mask` (scale to multiple 32) → **fixed scene embeddings** (`LTXVLoadConditioning`, not a free-text prompt) → HDR IC-LoRA guide → float32 VAE decode → `LTXVHDRDecodePostprocess` (`transfer=acescct`, EXR=`acescg`) → HLG. Sampler **euler**; original audio frozen/muxed. Same idea as pipelines `hdr_ic_lora --text-embeddings`. Conditioning fps snaps to 60 when source fps is above 30 (playback fps unchanged for decode/export).
|
||||
|
||||
For native HDR, set **`img_compression = 0`** (Python `--hdr` EXR conditioning skips JPEG/CRF; CRF on ACEScct darkens highlights). Turn **prompt enhancer off** when matching goldens. I2V stage-2 distilled sigmas should be `0.909375, 0.725, 0.421875, 0.0` (not `0.85, …`). Inpaint keeps its own 2-step stage-2 schedule.
|
||||
|
||||
Python / pipeline equivalents of these ideas live in [pipeline-selection.md](https://github.com/Lightricks/LTX-2/blob/main/packages/ltx-pipelines/docs/pipeline-selection.md) in the LTX-2 repo.
|
||||
@@ -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: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: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
|
||||
+102
-16
@@ -4,6 +4,7 @@ Replaces the CLIP encoding step entirely using an external API.
|
||||
"""
|
||||
|
||||
import io
|
||||
import json
|
||||
import logging
|
||||
import pickle
|
||||
|
||||
@@ -23,20 +24,85 @@ UPDATE_MESSAGE = (
|
||||
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"
|
||||
MISSING_MODEL_SELECTOR_MESSAGE = (
|
||||
"Cannot identify the text encoder from the provided model file."
|
||||
)
|
||||
MODEL_ID_METADATA_KEY = "encrypted_wandb_properties"
|
||||
GEMMA_SOURCE_CHECKPOINT_METADATA_KEY = "gemma_source_checkpoint"
|
||||
|
||||
|
||||
def extract_model_id(ckpt_name: str) -> str:
|
||||
model_id_key = "encrypted_wandb_properties"
|
||||
MODEL_FOLDERS = ("checkpoints", "diffusion_models")
|
||||
|
||||
|
||||
def model_filename_list() -> list[str]:
|
||||
"""Every model the node can read a text-encoder selector from, de-duplicated.
|
||||
|
||||
A name present in both folders is listed once; resolve_model_path resolves
|
||||
it the same way round, so the entry always refers to the same file the
|
||||
combo offered.
|
||||
"""
|
||||
names = []
|
||||
for folder in MODEL_FOLDERS:
|
||||
for name in folder_paths.get_filename_list(folder):
|
||||
if name not in names:
|
||||
names.append(name)
|
||||
return names
|
||||
|
||||
|
||||
def resolve_model_path(ckpt_name: str) -> str:
|
||||
"""Full path of ckpt_name, looked up in the folder order the combo used."""
|
||||
for folder in MODEL_FOLDERS:
|
||||
path = folder_paths.get_full_path(folder, ckpt_name)
|
||||
if path is not None:
|
||||
return path
|
||||
raise ValueError(
|
||||
f"Model file '{ckpt_name}' was not found in any of: {', '.join(MODEL_FOLDERS)}."
|
||||
)
|
||||
|
||||
|
||||
def parse_gemma_source_checkpoint(raw: str) -> dict[str, str]:
|
||||
"""Parse the ``gemma_source_checkpoint`` safetensors metadata JSON string."""
|
||||
try:
|
||||
obj = json.loads(raw)
|
||||
except json.JSONDecodeError as exc:
|
||||
raise ValueError(MISSING_MODEL_SELECTOR_MESSAGE) from exc
|
||||
if not isinstance(obj, dict):
|
||||
raise ValueError(MISSING_MODEL_SELECTOR_MESSAGE)
|
||||
ltx_version, gemma_version = obj.get("ltx_version"), obj.get("gemma_version")
|
||||
if not ltx_version or not gemma_version:
|
||||
raise ValueError(MISSING_MODEL_SELECTOR_MESSAGE)
|
||||
return {"ltx_version": str(ltx_version), "gemma_version": str(gemma_version)}
|
||||
|
||||
|
||||
def prompt_embedding_selector_from_metadata(
|
||||
metadata: dict[str, str] | None,
|
||||
) -> dict[str, str | dict[str, str]]:
|
||||
"""Build the XOR selector for ``/v1/prompt-embedding``.
|
||||
|
||||
LTX-2.5+ checkpoints identify the text encoder via ``gemma_source_checkpoint``
|
||||
(``ltx_version`` + ``gemma_version``). Pre-2.5 checkpoints use
|
||||
``encrypted_wandb_properties`` as ``model_id``. Prefer the 2.5+ selector when
|
||||
both are present: OS 2.5 files may still carry an unusable ``model_id``.
|
||||
"""
|
||||
if metadata:
|
||||
raw_gemma = metadata.get(GEMMA_SOURCE_CHECKPOINT_METADATA_KEY)
|
||||
if raw_gemma:
|
||||
return {"model": parse_gemma_source_checkpoint(raw_gemma)}
|
||||
model_id = metadata.get(MODEL_ID_METADATA_KEY)
|
||||
if model_id:
|
||||
return {"model_id": model_id}
|
||||
raise ValueError(MISSING_MODEL_SELECTOR_MESSAGE)
|
||||
|
||||
|
||||
def extract_prompt_embedding_selector(
|
||||
ckpt_name: str,
|
||||
) -> dict[str, str | dict[str, str]]:
|
||||
with safe_open(
|
||||
folder_paths.get_full_path_or_raise("checkpoints", ckpt_name),
|
||||
resolve_model_path(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]
|
||||
return prompt_embedding_selector_from_metadata(f.metadata())
|
||||
|
||||
|
||||
@comfy_node(name="GemmaAPITextEncode")
|
||||
@@ -45,8 +111,8 @@ 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.
|
||||
for processing. It requires an API key and selects the text encoder from checkpoint
|
||||
metadata: ``model`` (LTX-2.5+) or ``model_id`` (LTX-2.3).
|
||||
|
||||
Inputs:
|
||||
- api_key: Authentication key for the LTX Video API
|
||||
@@ -78,9 +144,19 @@ class GemmaAPITextEncode:
|
||||
"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."},
|
||||
model_filename_list(),
|
||||
{
|
||||
"tooltip": "The model to read the API text-encoder selector from. "
|
||||
"Either a checkpoint or a diffusion model."
|
||||
},
|
||||
),
|
||||
},
|
||||
}
|
||||
@@ -90,7 +166,9 @@ class GemmaAPITextEncode:
|
||||
FUNCTION = "encode"
|
||||
CATEGORY = "api node/text/Lightricks"
|
||||
|
||||
def encode(self, api_key: str, prompt: str, ckpt_name: str):
|
||||
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")
|
||||
|
||||
@@ -100,10 +178,18 @@ class GemmaAPITextEncode:
|
||||
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}
|
||||
selector = extract_prompt_embedding_selector(ckpt_name)
|
||||
payload = {
|
||||
"prompt": prompt,
|
||||
"enhance_prompt": enhance_prompt,
|
||||
**selector,
|
||||
}
|
||||
if "model_id" in selector:
|
||||
selector_desc = f"model_id: {str(selector['model_id'])[:50]}..."
|
||||
else:
|
||||
selector_desc = f"model: {selector['model']}"
|
||||
logger.info(
|
||||
f"Calling API to encode prompt: {prompt[:50]}... with model_id: {model_id[:50]}..."
|
||||
f"Calling API to encode prompt: {prompt[:50]}... with {selector_desc}"
|
||||
)
|
||||
try:
|
||||
response = requests.post(
|
||||
|
||||
+48
-254
@@ -1,19 +1,12 @@
|
||||
import json
|
||||
import logging
|
||||
from glob import glob
|
||||
from pathlib import Path
|
||||
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,
|
||||
@@ -25,13 +18,11 @@ from transformers import (
|
||||
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."""
|
||||
@@ -56,99 +47,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(
|
||||
@@ -184,22 +82,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",
|
||||
@@ -207,13 +95,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
|
||||
@@ -227,89 +110,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"]
|
||||
@@ -322,31 +136,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)
|
||||
@@ -368,28 +171,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,
|
||||
@@ -398,20 +196,15 @@ def ltxv_gemma_clip(encoder_path, ltxv_path, processor=None, dtype=None):
|
||||
return _LTXVGemmaTextEncoderModel
|
||||
|
||||
|
||||
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 = [
|
||||
Path(p)
|
||||
for p in glob(f"{root_path}/**", recursive=True)
|
||||
if Path(p).match(pattern)
|
||||
]
|
||||
if not matches:
|
||||
def gemma_model_dir(gemma_path: str) -> Path:
|
||||
"""Resolve the Gemma model directory from a selected weights file."""
|
||||
model_dir = Path(folder_paths.get_full_path("text_encoders", gemma_path)).parent
|
||||
if not (model_dir / "config.json").exists():
|
||||
raise FileNotFoundError(
|
||||
f"No files matching pattern '{pattern}' found under {root_path}"
|
||||
f"No config.json found for the selected Gemma model ({model_dir}). "
|
||||
"Ensure the model's config, tokenizer and processor files are present."
|
||||
)
|
||||
return str(matches[0].parent)
|
||||
return model_dir
|
||||
|
||||
|
||||
@comfy_node(name="LTXVGemmaCLIPModelLoader", description="Gemma 3 Model Loader")
|
||||
@@ -443,33 +236,29 @@ class LTXVGemmaCLIPModelLoader:
|
||||
OUTPUT_NODE = False
|
||||
|
||||
def load_model(self, gemma_path: str, ltxv_path: str, max_length: int):
|
||||
path = Path(folder_paths.get_full_path("text_encoders", gemma_path))
|
||||
model_root = path.parents[1]
|
||||
tokenizer_path = Path(find_matching_dir(model_root, "tokenizer.model"))
|
||||
gemma_model_path = Path(find_matching_dir(model_root, "model*.safetensors"))
|
||||
processor_path = Path(find_matching_dir(model_root, "preprocessor_config.json"))
|
||||
tokenizer_class = ltxv_gemma_tokenizer(tokenizer_path, max_length=max_length)
|
||||
model_dir = gemma_model_dir(gemma_path)
|
||||
tokenizer_class = ltxv_gemma_tokenizer(model_dir, max_length=max_length)
|
||||
|
||||
processor = None
|
||||
try:
|
||||
image_processor = AutoImageProcessor.from_pretrained(
|
||||
str(processor_path),
|
||||
str(model_dir),
|
||||
local_files_only=True,
|
||||
)
|
||||
processor = Gemma3Processor(
|
||||
image_processor=image_processor,
|
||||
tokenizer=tokenizer_class().tokenizer,
|
||||
)
|
||||
logger.info(f"Loaded processor from {model_root} - enhancement enabled")
|
||||
logger.info(f"Loaded processor from {model_dir} - enhancement enabled")
|
||||
except Exception as e:
|
||||
logger.warning(f"Could not load processor from {model_root}: {e}")
|
||||
logger.warning(f"Could not load processor from {model_dir}: {e}")
|
||||
|
||||
clip_dtype = torch.bfloat16
|
||||
ltxv_full_path = folder_paths.get_full_path("checkpoints", ltxv_path)
|
||||
clip_target = comfy.supported_models_base.ClipTarget(
|
||||
tokenizer=tokenizer_class,
|
||||
clip=ltxv_gemma_clip(
|
||||
gemma_model_path, ltxv_full_path, processor=processor, dtype=clip_dtype
|
||||
model_dir, ltxv_full_path, processor=processor, dtype=clip_dtype
|
||||
),
|
||||
)
|
||||
|
||||
@@ -702,16 +491,21 @@ def _cat_with_padding(
|
||||
padding_length: int,
|
||||
value: int | float,
|
||||
) -> torch.Tensor:
|
||||
"""Concatenate a tensor with a padding tensor of the given value."""
|
||||
"""Left-pad a tensor (prepend the padding) with the given value.
|
||||
|
||||
Decoder-only LLMs must be LEFT-padded for generation: right-padding makes the
|
||||
trailing pad the sequence end, so generate() reads next-token logits from a
|
||||
masked pad position instead of the real last token, producing degenerate output.
|
||||
"""
|
||||
return torch.cat(
|
||||
[
|
||||
tensor,
|
||||
torch.full(
|
||||
(1, padding_length),
|
||||
value,
|
||||
dtype=tensor.dtype,
|
||||
device=tensor.device,
|
||||
),
|
||||
tensor,
|
||||
],
|
||||
dim=1,
|
||||
)
|
||||
|
||||
@@ -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},
|
||||
)
|
||||
|
||||
@@ -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}"
|
||||
|
||||
+251
@@ -0,0 +1,251 @@
|
||||
"""ACEScct working-space + primaries math for HDR Comfy paths.
|
||||
|
||||
Kept in lockstep with ``ltx_core.hdr`` / ``ltx_core.color.primaries`` on
|
||||
ltx-2-internal: ACEScct curve constants match ``ltx_core.hdr``, and primaries
|
||||
matrices / EXR chromaticities come from ``colour-science`` with Bradford CAT
|
||||
(same source OpenColorIO-Config-ACES uses for its utility CLFs).
|
||||
|
||||
``colour-science`` is required (same as ``ltx_core.color.primaries``).
|
||||
|
||||
Used by HDR Comfy nodes in ``hdr_nodes`` (EXR ↔ VAE, SDR→HDR IC-LoRA ingest
|
||||
via ``SDRColorSpace`` / sRGB EOTF → ACEScct). File I/O helpers live in ``hdr_io``.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import enum
|
||||
import logging
|
||||
|
||||
import colour
|
||||
import torch
|
||||
from torch import Tensor
|
||||
|
||||
logger = logging.getLogger("LTXVideo.hdr_color")
|
||||
|
||||
# ACEScct (AMPAS S-2016-001) — same constants as ``ltx_core.hdr``
|
||||
_ACESCCT_A_LIN = 10.5402377416545
|
||||
_ACESCCT_B_LIN = 0.0729055341958355
|
||||
_ACESCCT_X_BRK = 0.0078125
|
||||
_ACESCCT_Y_BRK = 0.155251141552511
|
||||
_ACESCCT_LOG_M = 17.52
|
||||
_ACESCCT_LOG_B = 9.72
|
||||
|
||||
_CAT = "Bradford"
|
||||
_CS_ACESCG = colour.RGB_COLOURSPACES["ACEScg"]
|
||||
_CS_REC709 = colour.RGB_COLOURSPACES["ITU-R BT.709"]
|
||||
_CS_REC2020 = colour.RGB_COLOURSPACES["ITU-R BT.2020"]
|
||||
|
||||
|
||||
def _rgb_to_rgb(src: colour.RGB_Colourspace, dst: colour.RGB_Colourspace) -> Tensor:
|
||||
"""Linear RGB→RGB matrix from colour-science, as float32."""
|
||||
return torch.as_tensor(
|
||||
colour.matrix_RGB_to_RGB(src, dst, chromatic_adaptation_transform=_CAT),
|
||||
dtype=torch.float32,
|
||||
)
|
||||
|
||||
|
||||
def _exr_chroma(cs: colour.RGB_Colourspace) -> tuple[float, ...]:
|
||||
"""EXR ``chromaticities`` order: Rxy, Gxy, Bxy, Wxy."""
|
||||
r, g, b = cs.primaries
|
||||
w = cs.whitepoint
|
||||
return (
|
||||
float(r[0]),
|
||||
float(r[1]),
|
||||
float(g[0]),
|
||||
float(g[1]),
|
||||
float(b[0]),
|
||||
float(b[1]),
|
||||
float(w[0]),
|
||||
float(w[1]),
|
||||
)
|
||||
|
||||
|
||||
_ACESCG_TO_SRGB = _rgb_to_rgb(_CS_ACESCG, _CS_REC709)
|
||||
_SRGB_TO_ACESCG = torch.linalg.inv(_ACESCG_TO_SRGB.double()).to(torch.float32)
|
||||
_REC709_TO_2020 = _rgb_to_rgb(_CS_REC709, _CS_REC2020)
|
||||
_ACESCG_TO_2020 = _rgb_to_rgb(_CS_ACESCG, _CS_REC2020)
|
||||
_CHROMA_REC709 = _exr_chroma(_CS_REC709)
|
||||
_CHROMA_ACESCG = _exr_chroma(_CS_ACESCG)
|
||||
_MATRICES_SOURCE = "colour-science"
|
||||
|
||||
# Public aliases (parity with ``ltx_core.color.primaries`` exports).
|
||||
ACESCG_TO_SRGB = _ACESCG_TO_SRGB
|
||||
SRGB_TO_ACESCG = _SRGB_TO_ACESCG
|
||||
REC709_TO_2020 = _REC709_TO_2020
|
||||
ACESCG_TO_2020 = _ACESCG_TO_2020
|
||||
|
||||
_CHROMATICITIES = {
|
||||
"rec709": _CHROMA_REC709,
|
||||
"acescg": _CHROMA_ACESCG,
|
||||
}
|
||||
|
||||
|
||||
class Primaries(enum.Enum):
|
||||
"""Authoring / working-space colour primaries (linear light)."""
|
||||
|
||||
REC709 = "rec709"
|
||||
ACESCG = "acescg"
|
||||
|
||||
@property
|
||||
def exr_chromaticities(self) -> tuple[float, ...]:
|
||||
"""EXR ``chromaticities`` header (R/G/B/W x,y)."""
|
||||
return _CHROMATICITIES[self.value]
|
||||
|
||||
@property
|
||||
def matrix_to_rec2020(self) -> Tensor:
|
||||
"""3x3 linear map from this basis to Rec.2020 (for HLG)."""
|
||||
return _REC709_TO_2020 if self is Primaries.REC709 else _ACESCG_TO_2020
|
||||
|
||||
def to(self, target: Primaries, video: Tensor) -> Tensor:
|
||||
"""Convert linear light from this basis into ``target``. Identity when same."""
|
||||
if self is target:
|
||||
return video
|
||||
if (self, target) == (Primaries.REC709, Primaries.ACESCG):
|
||||
return _apply_matrix(video, _SRGB_TO_ACESCG)
|
||||
if (self, target) == (Primaries.ACESCG, Primaries.REC709):
|
||||
return _apply_matrix(video, _ACESCG_TO_SRGB)
|
||||
raise ValueError(f"No primaries conversion from {self!r} to {target!r}.")
|
||||
|
||||
def to_rec2020(self, video: Tensor) -> Tensor:
|
||||
"""Linear light in this basis → Rec.2020 (HLG master space)."""
|
||||
return _apply_matrix(video, self.matrix_to_rec2020)
|
||||
|
||||
|
||||
class HDRColorSpace(enum.Enum):
|
||||
"""Source / EXR authoring colour space (matches pipelines ``--hdr``)."""
|
||||
|
||||
SRGB_LINEAR = "srgb_linear"
|
||||
ACESCG = "acescg"
|
||||
ACESCCT = "acescct"
|
||||
|
||||
@property
|
||||
def is_log_working(self) -> bool:
|
||||
return self is HDRColorSpace.ACESCCT
|
||||
|
||||
@property
|
||||
def source_primaries(self) -> Primaries:
|
||||
if self in (HDRColorSpace.ACESCG, HDRColorSpace.ACESCCT):
|
||||
return Primaries.ACESCG
|
||||
return Primaries.REC709
|
||||
|
||||
def exr_output_tags(self) -> tuple[Primaries, str]:
|
||||
"""``(primaries, colorSpace tag)`` for EXR headers."""
|
||||
if self is HDRColorSpace.ACESCCT:
|
||||
return Primaries.ACESCG, "ACEScct"
|
||||
if self is HDRColorSpace.ACESCG:
|
||||
return Primaries.ACESCG, "ACEScg"
|
||||
return Primaries.REC709, "sRGB"
|
||||
|
||||
|
||||
def _apply_matrix(video: Tensor, mat: Tensor) -> Tensor:
|
||||
"""Apply 3x3 on channel-last ``[..., H, W, 3]`` or channel-first ``[..., 3, H, W]``.
|
||||
|
||||
Comfy IMAGE tensors are channel-last; ``ltx_core`` often uses channel-first.
|
||||
"""
|
||||
mat = mat.to(device=video.device, dtype=video.dtype)
|
||||
if video.shape[-1] == 3:
|
||||
return video @ mat.T
|
||||
if video.shape[-3] == 3:
|
||||
return torch.einsum("...chw,dc->...dhw", video, mat)
|
||||
raise ValueError(f"Expected RGB in last or -3 dim; got shape {tuple(video.shape)}")
|
||||
|
||||
|
||||
# sRGB EOTF (IEC 61966-2-1) — same constants as ``ltx_core.hdr``.
|
||||
_SRGB_A = 0.055
|
||||
_SRGB_LINEAR_THRESHOLD = 0.04045
|
||||
_SRGB_LINEAR_SLOPE = 12.92
|
||||
_SRGB_GAMMA = 2.4
|
||||
|
||||
|
||||
def compress_acescct(hdr: Tensor) -> Tensor:
|
||||
"""Linear HDR [0, ∞) → ACEScct [0, 1]."""
|
||||
x = torch.clamp(hdr, min=0.0)
|
||||
log_part = (torch.log2(torch.clamp(x, min=1e-12)) + _ACESCCT_LOG_B) / _ACESCCT_LOG_M
|
||||
lin_part = _ACESCCT_A_LIN * x + _ACESCCT_B_LIN
|
||||
return torch.clamp(torch.where(x > _ACESCCT_X_BRK, log_part, lin_part), 0.0, 1.0)
|
||||
|
||||
|
||||
def decompress_acescct(ct: Tensor) -> Tensor:
|
||||
"""ACEScct [0, 1] → linear ACEScg [0, ∞)."""
|
||||
ct = torch.clamp(ct, 0.0, 1.0)
|
||||
lin_from_log = torch.pow(2.0, ct * _ACESCCT_LOG_M - _ACESCCT_LOG_B)
|
||||
lin_from_lin = (ct - _ACESCCT_B_LIN) / _ACESCCT_A_LIN
|
||||
return torch.where(ct > _ACESCCT_Y_BRK, lin_from_log, lin_from_lin)
|
||||
|
||||
|
||||
def srgb_eotf_to_linear(srgb_01: Tensor) -> Tensor:
|
||||
"""sRGB-encoded [0, 1] → display-linear Rec.709/sRGB (IEC 61966-2-1 EOTF)."""
|
||||
x = torch.clamp(srgb_01.float(), 0.0, 1.0)
|
||||
return torch.where(
|
||||
x <= _SRGB_LINEAR_THRESHOLD,
|
||||
x / _SRGB_LINEAR_SLOPE,
|
||||
torch.pow((x + _SRGB_A) / (1.0 + _SRGB_A), _SRGB_GAMMA),
|
||||
)
|
||||
|
||||
|
||||
def to_acescct_working_space(linear_rgb: Tensor, source_primaries: Primaries) -> Tensor:
|
||||
"""Scene-linear RGB → ACEScct [0, 1] (channel-last or channel-first)."""
|
||||
native = source_primaries.to(Primaries.ACESCG, linear_rgb.float())
|
||||
return compress_acescct(native.clamp(min=0.0))
|
||||
|
||||
|
||||
def acescct_to_linear(
|
||||
working: Tensor, out_primaries: Primaries = Primaries.REC709
|
||||
) -> Tensor:
|
||||
"""ACEScct [0, 1] → scene-linear in ``out_primaries``."""
|
||||
linear_acescg = decompress_acescct(working.float())
|
||||
return Primaries.ACESCG.to(out_primaries, linear_acescg).clamp(min=0.0)
|
||||
|
||||
|
||||
class SDRColorSpace(enum.Enum):
|
||||
"""SDR ingest colour space for mapping float RGB into ACEScct.
|
||||
|
||||
Matches pipelines ``SDRColorSpace`` / ``hdr_ic_lora`` IDT:
|
||||
|
||||
* ``srgb_gamma`` — display-referred sRGB (apply EOTF, then ACEScct).
|
||||
* ``srgb`` — scene-linear Rec.709 / sRGB primaries (no EOTF).
|
||||
* ``acescg`` — scene-linear ACEScg → ACEScct.
|
||||
* ``acescct`` — already ACEScct log codes; clamp passthrough.
|
||||
"""
|
||||
|
||||
SRGB = "srgb"
|
||||
SRGB_GAMMA = "srgb_gamma"
|
||||
ACESCG = "acescg"
|
||||
ACESCCT = "acescct"
|
||||
|
||||
def to_working_space(self, video: Tensor) -> Tensor:
|
||||
"""Map SDR float RGB into ACEScct ``[0, 1]`` (channel-last or channel-first)."""
|
||||
if not video.is_floating_point():
|
||||
raise ValueError(
|
||||
f"SDRColorSpace.to_working_space expects float RGB; got dtype={video.dtype}. "
|
||||
"Normalize integer pixels before calling."
|
||||
)
|
||||
video = video.float()
|
||||
if self is SDRColorSpace.ACESCCT:
|
||||
return video.clamp(0.0, 1.0)
|
||||
if self is SDRColorSpace.SRGB_GAMMA:
|
||||
return to_acescct_working_space(
|
||||
srgb_eotf_to_linear(video), Primaries.REC709
|
||||
)
|
||||
if self is SDRColorSpace.SRGB:
|
||||
return to_acescct_working_space(video, Primaries.REC709)
|
||||
if self is SDRColorSpace.ACESCG:
|
||||
return to_acescct_working_space(video, Primaries.ACESCG)
|
||||
raise ValueError(f"Unsupported SDR color space: {self!r}")
|
||||
|
||||
|
||||
def images_to_acescct(images: Tensor, color_space: HDRColorSpace) -> Tensor:
|
||||
"""Convert loaded EXR frames ``[F,H,W,3]`` into ACEScct ``[0,1]`` for the VAE."""
|
||||
images = images.float()
|
||||
if color_space.is_log_working:
|
||||
return images.clamp(0.0, 1.0)
|
||||
return to_acescct_working_space(images, color_space.source_primaries)
|
||||
|
||||
|
||||
def sdr_images_to_acescct(images: Tensor, color_space: SDRColorSpace) -> Tensor:
|
||||
"""Convert SDR frames ``[F,H,W,C]`` into ACEScct ``[0,1]`` for HDR IC-LoRA."""
|
||||
rgb = images[..., :3]
|
||||
working = color_space.to_working_space(rgb)
|
||||
if images.shape[-1] > 3:
|
||||
return torch.cat([working, images[..., 3:].float()], dim=-1)
|
||||
return working
|
||||
@@ -0,0 +1,394 @@
|
||||
"""HDR file I/O helpers: EXR read/write, HLG encode, path utilities.
|
||||
|
||||
Shared by Comfy HDR nodes (ingest + delivery). Color math lives in ``hdr_color``;
|
||||
node classes live in ``hdr_nodes``.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
import math
|
||||
import os
|
||||
import re
|
||||
from pathlib import Path
|
||||
|
||||
import folder_paths
|
||||
import numpy as np
|
||||
import torch
|
||||
from torch import Tensor
|
||||
|
||||
from .hdr_color import Primaries
|
||||
|
||||
logger = logging.getLogger("LTXVideo.hdr_io")
|
||||
|
||||
# BT.2020 / HLG (ARIB STD-B67) encode constants — from ltx_core.color.hlg
|
||||
_HLG_A, _HLG_B, _HLG_C = 0.17883277, 0.28466892, 0.55991073
|
||||
_AV_COLOR_PRIMARIES_BT2020 = 9
|
||||
_AV_COLOR_TRC_ARIB_STD_B67 = 18
|
||||
_AV_COLORSPACE_BT2020_NCL = 9
|
||||
_AV_COLOR_RANGE_MPEG = 1
|
||||
_KR_2020, _KG_2020, _KB_2020 = 0.2627, 0.6780, 0.0593
|
||||
_BT2020_RGB_TO_YUV = torch.tensor(
|
||||
[
|
||||
[_KR_2020, _KG_2020, _KB_2020],
|
||||
[-_KR_2020 / 1.8814, -_KG_2020 / 1.8814, 0.5],
|
||||
[0.5, -_KG_2020 / 1.4746, -_KB_2020 / 1.4746],
|
||||
],
|
||||
dtype=torch.float32,
|
||||
)
|
||||
|
||||
|
||||
def _resolve_path(path: str) -> Path:
|
||||
p = Path(path).expanduser()
|
||||
if p.is_absolute() and p.exists():
|
||||
return p
|
||||
# Prefer Comfy input dir for relative paths.
|
||||
candidate = Path(folder_paths.get_input_directory()) / path
|
||||
if candidate.exists():
|
||||
return candidate
|
||||
if p.exists():
|
||||
return p.resolve()
|
||||
raise FileNotFoundError(
|
||||
f"HDR path not found: {path!r} (tried {candidate} and {p.resolve()})"
|
||||
)
|
||||
|
||||
|
||||
def _natural_sort_key(name: str) -> tuple:
|
||||
"""Sort key so ``frame_9.exr`` precedes ``frame_10.exr`` (not lexicographic)."""
|
||||
parts = re.split(r"(\d+)", name)
|
||||
key: list = []
|
||||
for part in parts:
|
||||
if part.isdigit():
|
||||
key.append((0, int(part)))
|
||||
else:
|
||||
key.append((1, part.casefold()))
|
||||
return tuple(key)
|
||||
|
||||
|
||||
def _list_exr_files(path: Path) -> list[Path]:
|
||||
"""Resolve a still ``.exr`` or an EXR folder to a sorted frame list."""
|
||||
if path.is_file():
|
||||
if path.suffix.lower() != ".exr":
|
||||
raise ValueError(f"Expected a .exr file; got {path}")
|
||||
return [path]
|
||||
if path.is_dir():
|
||||
# Case-insensitive match (Linux is case-sensitive; macOS usually is not).
|
||||
files = sorted(
|
||||
{p for p in path.iterdir() if p.is_file() and p.suffix.lower() == ".exr"},
|
||||
key=lambda p: _natural_sort_key(p.name),
|
||||
)
|
||||
if not files:
|
||||
raise RuntimeError(f"No .exr frames in {path}")
|
||||
return files
|
||||
raise ValueError(f"Path must be a .exr file or a directory of EXRs; got {path}")
|
||||
|
||||
|
||||
def _silent_audio(num_frames: int, frame_rate: float, sample_rate: int = 44100) -> dict:
|
||||
"""Stereo silence long enough for the LTX audio VAE mel path.
|
||||
|
||||
EXR plates have no soundtrack; V2V still freezes an audio latent. A 0-length
|
||||
waveform crashes ``torchaudio.functional.resample`` / ``VAEEncodeAudio``.
|
||||
"""
|
||||
if frame_rate <= 0:
|
||||
raise ValueError(f"frame_rate must be > 0; got {frame_rate}")
|
||||
duration_s = max(float(num_frames) / float(frame_rate), 1.0)
|
||||
# Keep a comfortable floor for mel/STFT (typical n_fft is 1024–2048).
|
||||
n_samples = max(int(math.ceil(duration_s * sample_rate)), sample_rate, 8192)
|
||||
return {
|
||||
"waveform": torch.zeros(1, 2, n_samples, dtype=torch.float32),
|
||||
"sample_rate": int(sample_rate),
|
||||
}
|
||||
|
||||
|
||||
def _read_exr(path: Path) -> Tensor:
|
||||
"""Read one EXR as float32 ``[H,W,3]`` (unbounded linear or log codes as stored).
|
||||
|
||||
Requires OpenImageIO (same as ``ltx_pipelines`` media I/O). No OpenCV fallback.
|
||||
"""
|
||||
try:
|
||||
import OpenImageIO as oiio
|
||||
except ImportError as e:
|
||||
raise RuntimeError(
|
||||
"Reading EXR requires OpenImageIO. Install with: pip install openimageio"
|
||||
) from e
|
||||
|
||||
buf = oiio.ImageBuf(str(path))
|
||||
pixels = buf.get_pixels(oiio.FLOAT)
|
||||
if pixels is None or buf.has_error:
|
||||
raise RuntimeError(f"OpenImageIO failed on '{path}': {buf.geterror()}")
|
||||
arr = np.ascontiguousarray(np.asarray(pixels, dtype=np.float32))
|
||||
|
||||
if arr.ndim == 2:
|
||||
arr = np.repeat(arr[..., None], 3, axis=-1)
|
||||
return torch.from_numpy(np.ascontiguousarray(arr[..., :3]))
|
||||
|
||||
|
||||
def _save_exr_frame(
|
||||
frame_hwc: Tensor,
|
||||
path: Path,
|
||||
*,
|
||||
half_precision: bool,
|
||||
primaries: Primaries | None,
|
||||
color_space_tag: str | None,
|
||||
) -> None:
|
||||
"""Write one EXR via OpenImageIO (same backend as ``ltx_pipelines``).
|
||||
|
||||
Optional ``primaries`` / ``color_space_tag`` set EXR header attributes when
|
||||
provided; ``exr_color_space=linear`` may omit them for legacy untagged plates.
|
||||
"""
|
||||
try:
|
||||
import OpenImageIO as oiio
|
||||
except ImportError as e:
|
||||
raise RuntimeError(
|
||||
"Saving EXR requires OpenImageIO. Install with: pip install openimageio"
|
||||
) from e
|
||||
|
||||
rgb = frame_hwc.detach().float().cpu().numpy()
|
||||
path.parent.mkdir(parents=True, exist_ok=True)
|
||||
h, w, _ = rgb.shape
|
||||
fmt = oiio.HALF if half_precision else oiio.FLOAT
|
||||
spec = oiio.ImageSpec(w, h, 3, fmt)
|
||||
spec.channelnames = ("R", "G", "B")
|
||||
spec.attribute("compression", "zip")
|
||||
if primaries is not None:
|
||||
spec.attribute(
|
||||
"chromaticities",
|
||||
"float[8]",
|
||||
primaries.exr_chromaticities,
|
||||
)
|
||||
if color_space_tag is not None:
|
||||
spec.attribute("colorSpace", color_space_tag)
|
||||
out = oiio.ImageOutput.create(str(path))
|
||||
if out is None:
|
||||
raise RuntimeError(
|
||||
f"OpenImageIO cannot create writer for {path}. "
|
||||
"Ensure OpenImageIO is built with OpenEXR support."
|
||||
)
|
||||
try:
|
||||
if not out.open(str(path), spec):
|
||||
raise RuntimeError(
|
||||
f"Failed to access EXR output file at '{path}': {out.geterror()}"
|
||||
)
|
||||
# OIIO converts to HALF from float32 buffer when spec is HALF.
|
||||
if not out.write_image(np.ascontiguousarray(rgb, dtype=np.float32)):
|
||||
raise RuntimeError(f"Failed to write EXR '{path}': {out.geterror()}")
|
||||
finally:
|
||||
out.close()
|
||||
|
||||
|
||||
def _hlg_inverse_oetf_scalar(v: float) -> float:
|
||||
return (
|
||||
(v * v) / 3.0 if v <= 0.5 else (math.exp((v - _HLG_C) / _HLG_A) + _HLG_B) / 12.0
|
||||
)
|
||||
|
||||
|
||||
def _linear_to_hlg_signal(
|
||||
rgb_fhwc: Tensor,
|
||||
primaries: Primaries = Primaries.REC709,
|
||||
white_signal: float = 0.75,
|
||||
) -> Tensor:
|
||||
"""Scene-linear RGB ``[F,H,W,3]`` in ``primaries`` → HLG signal in BT.2020.
|
||||
|
||||
Maps with ``primaries.matrix_to_rec2020`` then clamps (same as
|
||||
``ltx_core.color.hlg.HlgGpuConverter``). Prefer ACEScg→Rec.2020 directly
|
||||
over ACEScg→Rec.709→Rec.2020 so OOG negatives are not crushed early.
|
||||
"""
|
||||
white_x = _hlg_inverse_oetf_scalar(white_signal)
|
||||
roll_k = white_x / (1.0 - white_x)
|
||||
mat = primaries.matrix_to_rec2020.to(device=rgb_fhwc.device, dtype=rgb_fhwc.dtype)
|
||||
lin = torch.nan_to_num((rgb_fhwc @ mat.T).clamp(min=0.0), nan=0.0, neginf=0.0)
|
||||
x = torch.where(
|
||||
lin <= 1.0,
|
||||
lin * white_x,
|
||||
1.0 - (1.0 - white_x) * torch.exp(-roll_k * (lin - 1.0)),
|
||||
)
|
||||
sig = torch.where(
|
||||
x <= 1.0 / 12.0,
|
||||
torch.sqrt((3.0 * x).clamp(min=0.0)),
|
||||
_HLG_A * torch.log((12.0 * x - _HLG_B).clamp(min=1e-12)) + _HLG_C,
|
||||
)
|
||||
return sig.clamp(0.0, 1.0)
|
||||
|
||||
|
||||
def _rgb_hlg_to_yuv420p10(
|
||||
rgb_fhwc: Tensor,
|
||||
) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
|
||||
"""HLG RGB ``[F,H,W,3]`` → planar uint16 Y,U,V (limited / BT.2020)."""
|
||||
if rgb_fhwc.shape[-2] % 2 or rgb_fhwc.shape[-3] % 2:
|
||||
raise ValueError(f"HLG encode needs even H and W; got {tuple(rgb_fhwc.shape)}")
|
||||
rgb = rgb_fhwc.movedim(-1, -3).contiguous() # F,3,H,W
|
||||
mat = _BT2020_RGB_TO_YUV.to(device=rgb.device, dtype=rgb.dtype)
|
||||
pixels = rgb.movedim(-3, -1)
|
||||
yuv = pixels.flatten(-3, -2) @ mat.T
|
||||
yuv = yuv.unflatten(-2, rgb.shape[-2:]).movedim(-1, -3)
|
||||
y = yuv[:, :1]
|
||||
uv = torch.nn.functional.avg_pool2d(
|
||||
yuv[:, 1:3].contiguous(), kernel_size=2, stride=2
|
||||
)
|
||||
# MPEG limited 10-bit
|
||||
y = y.mul(219 * 4).add_(16 * 4)
|
||||
uv = uv.mul(224 * 4).add_(128 * 4)
|
||||
y16 = y[:, 0].round_().clamp_(0, 1023).to(torch.uint16).cpu().numpy()
|
||||
u16 = uv[:, 0].round_().clamp_(0, 1023).to(torch.uint16).cpu().numpy()
|
||||
v16 = uv[:, 1].round_().clamp_(0, 1023).to(torch.uint16).cpu().numpy()
|
||||
return y16, u16, v16
|
||||
|
||||
|
||||
def _encode_hlg_mp4(
|
||||
frames_linear: Tensor,
|
||||
out_path: Path,
|
||||
fps: float,
|
||||
audio: dict | None = None,
|
||||
*,
|
||||
primaries: Primaries = Primaries.REC709,
|
||||
) -> None:
|
||||
"""Encode scene-linear ``[F,H,W,3]`` (in ``primaries``) to BT.2020/HLG/10-bit HEVC.
|
||||
|
||||
Optional Comfy ``AUDIO`` dict ``{waveform, sample_rate}`` is muxed as AAC.
|
||||
The AAC stream is registered *before* video packets (PyAV requires this).
|
||||
"""
|
||||
from fractions import Fraction
|
||||
|
||||
import av
|
||||
|
||||
hlg = _linear_to_hlg_signal(frames_linear.float(), primaries=primaries)
|
||||
y, u, v = _rgb_hlg_to_yuv420p10(hlg)
|
||||
f, height, width = int(y.shape[0]), int(y.shape[1]), int(y.shape[2])
|
||||
out_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
container = av.open(str(out_path), mode="w", options={"movflags": "+faststart"})
|
||||
audio_stream = None
|
||||
try:
|
||||
stream = container.add_stream(
|
||||
"libx265", rate=Fraction(fps).limit_denominator(1000)
|
||||
)
|
||||
stream.width = width
|
||||
stream.height = height
|
||||
stream.pix_fmt = "yuv420p10le"
|
||||
stream.codec_tag = "hvc1"
|
||||
threads = max(1, min(os.cpu_count() or 8, 16))
|
||||
stream.options = {
|
||||
"crf": "12",
|
||||
"preset": "medium",
|
||||
"x265-params": (
|
||||
"colorprim=bt2020:transfer=arib-std-b67:colormatrix=bt2020nc:range=limited:"
|
||||
f"repeat-headers=1:info=0:pools={threads}:frame-threads=4"
|
||||
),
|
||||
}
|
||||
ctx = stream.codec_context
|
||||
ctx.thread_count = threads
|
||||
ctx.color_primaries = _AV_COLOR_PRIMARIES_BT2020
|
||||
ctx.color_trc = _AV_COLOR_TRC_ARIB_STD_B67
|
||||
ctx.colorspace = _AV_COLORSPACE_BT2020_NCL
|
||||
ctx.color_range = _AV_COLOR_RANGE_MPEG
|
||||
|
||||
if audio is not None:
|
||||
audio_stream = _prepare_aac_stream(container, audio)
|
||||
|
||||
for i in range(f):
|
||||
frame = av.VideoFrame(width, height, "yuv420p10le")
|
||||
for plane, src in zip(frame.planes, (y[i], u[i], v[i]), strict=True):
|
||||
dest = np.frombuffer(plane, dtype=np.uint16).reshape(
|
||||
plane.height, plane.line_size // 2
|
||||
)
|
||||
dest[:, : src.shape[1]] = src
|
||||
frame.colorspace = _AV_COLORSPACE_BT2020_NCL
|
||||
frame.color_range = _AV_COLOR_RANGE_MPEG
|
||||
for packet in stream.encode(frame):
|
||||
container.mux(packet)
|
||||
for packet in stream.encode():
|
||||
container.mux(packet)
|
||||
|
||||
if audio is not None and audio_stream is not None:
|
||||
_mux_comfy_audio_aac(container, audio_stream, audio)
|
||||
finally:
|
||||
container.close()
|
||||
|
||||
|
||||
def _comfy_audio_stereo_n2(audio: dict) -> tuple[Tensor, int]:
|
||||
"""Return stereo int16-ready float samples ``[N, 2]`` and sample rate."""
|
||||
waveform = audio.get("waveform")
|
||||
sample_rate = int(audio.get("sample_rate", 0) or 0)
|
||||
if waveform is None or sample_rate <= 0:
|
||||
raise ValueError("AUDIO must include waveform and sample_rate")
|
||||
|
||||
samples = waveform.detach().float().cpu()
|
||||
if samples.ndim == 3:
|
||||
samples = samples[0]
|
||||
if samples.ndim != 2:
|
||||
raise ValueError(
|
||||
f"Expected AUDIO waveform [C,N] or [B,C,N]; got {tuple(waveform.shape)}"
|
||||
)
|
||||
if samples.shape[0] == 2:
|
||||
samples_cn = samples
|
||||
elif samples.shape[1] == 2:
|
||||
samples_cn = samples.T
|
||||
elif samples.shape[0] == 1:
|
||||
samples_cn = samples.repeat(2, 1)
|
||||
else:
|
||||
raise ValueError(
|
||||
f"AUDIO must be mono or stereo; got shape {tuple(samples.shape)}"
|
||||
)
|
||||
return samples_cn.T.contiguous(), sample_rate
|
||||
|
||||
|
||||
def _prepare_aac_stream(container, audio: dict):
|
||||
"""Register a stereo AAC stream before any packets are muxed."""
|
||||
from fractions import Fraction
|
||||
|
||||
_, sample_rate = _comfy_audio_stereo_n2(audio)
|
||||
audio_stream = container.add_stream("aac", rate=sample_rate)
|
||||
audio_stream.codec_context.sample_rate = sample_rate
|
||||
audio_stream.codec_context.layout = "stereo"
|
||||
audio_stream.codec_context.time_base = Fraction(1, sample_rate)
|
||||
return audio_stream
|
||||
|
||||
|
||||
def _mux_comfy_audio_aac(container, audio_stream, audio: dict) -> None:
|
||||
"""Mux Comfy AUDIO onto an already-added AAC stream (after video flush)."""
|
||||
import av
|
||||
|
||||
samples_n2, sample_rate = _comfy_audio_stereo_n2(audio)
|
||||
samples_i16 = (samples_n2.clamp(-1.0, 1.0) * 32767.0).to(torch.int16)
|
||||
|
||||
frame_in = av.AudioFrame.from_ndarray(
|
||||
samples_i16.reshape(1, -1).numpy(),
|
||||
format="s16",
|
||||
layout="stereo",
|
||||
)
|
||||
frame_in.sample_rate = sample_rate
|
||||
|
||||
cc = audio_stream.codec_context
|
||||
resampler = av.audio.resampler.AudioResampler(
|
||||
format=cc.format or "fltp",
|
||||
layout=cc.layout or "stereo",
|
||||
rate=cc.sample_rate or sample_rate,
|
||||
)
|
||||
next_pts = 0
|
||||
for rframe in resampler.resample(frame_in):
|
||||
if rframe.pts is None:
|
||||
rframe.pts = next_pts
|
||||
next_pts += rframe.samples
|
||||
for packet in audio_stream.encode(rframe):
|
||||
container.mux(packet)
|
||||
for rframe in resampler.resample(None):
|
||||
if rframe.pts is None:
|
||||
rframe.pts = next_pts
|
||||
next_pts += rframe.samples
|
||||
for packet in audio_stream.encode(rframe):
|
||||
container.mux(packet)
|
||||
for packet in audio_stream.encode(None):
|
||||
container.mux(packet)
|
||||
|
||||
|
||||
def _trim_to_8k1(images: Tensor) -> Tensor:
|
||||
"""Keep longest legal ``8k+1`` prefix (VAE temporal scale 8)."""
|
||||
n = images.shape[0]
|
||||
legal = ((n - 1) // 8) * 8 + 1
|
||||
if legal < 1:
|
||||
raise ValueError(f"Need at least 1 frame; got {n}")
|
||||
if legal != n:
|
||||
logger.warning(
|
||||
"Trimmed EXR sequence from %d to %d frames (must be 8k+1)", n, legal
|
||||
)
|
||||
return images[:legal]
|
||||
+675
@@ -0,0 +1,675 @@
|
||||
"""ComfyUI HDR nodes: ingest, decode postprocess, and delivery.
|
||||
|
||||
Stage-oriented layout (shared ACEScct working space):
|
||||
|
||||
- Ingest: ``LTXVSDRToHDRWorkingSpace``, ``LTXVLoadEXRSequence``
|
||||
- Runtime: ``LTXVVAEForceFloat32``
|
||||
- Delivery: ``LTXVHDRDecodePostprocess``, ``LTXVSaveHLG``
|
||||
|
||||
Color math: ``hdr_color``. File I/O helpers: ``hdr_io``.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
import os
|
||||
from pathlib import Path
|
||||
|
||||
import folder_paths
|
||||
import torch
|
||||
from torch import Tensor
|
||||
|
||||
from .hdr_color import (
|
||||
HDRColorSpace,
|
||||
Primaries,
|
||||
SDRColorSpace,
|
||||
compress_acescct,
|
||||
images_to_acescct,
|
||||
sdr_images_to_acescct,
|
||||
)
|
||||
from .hdr_io import (
|
||||
_encode_hlg_mp4,
|
||||
_list_exr_files,
|
||||
_read_exr,
|
||||
_resolve_path,
|
||||
_save_exr_frame,
|
||||
_silent_audio,
|
||||
_trim_to_8k1,
|
||||
)
|
||||
from .nodes_registry import comfy_node
|
||||
|
||||
logger = logging.getLogger("LTXVideo.hdr_nodes")
|
||||
|
||||
_EXR_COLOR_SPACE_CHOICES = ["linear", "acescct", "srgb_linear", "acescg"]
|
||||
_SDR_COLOR_SPACE_CHOICES = [c.value for c in SDRColorSpace]
|
||||
_COLOR_SPACE_CHOICES = [c.value for c in HDRColorSpace]
|
||||
|
||||
|
||||
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, transfer: str = "acescct") -> Tensor:
|
||||
"""Decompress VAE-decoded ``[0,1]`` image to linear HDR ``[0, inf)``.
|
||||
|
||||
ComfyUI's VAE decode returns images in ``[0, 1]`` via ``(raw + 1) / 2``.
|
||||
|
||||
* ``acescct`` — ACEScct working codes (SDR→HDR IC-LoRA + native HDR).
|
||||
* ``logc3`` — legacy LogC3 SDR→HDR LoRA only.
|
||||
"""
|
||||
transfer = (transfer or "acescct").lower()
|
||||
if transfer == "acescct":
|
||||
from .hdr_color import decompress_acescct
|
||||
|
||||
# Decoded IMAGE is already ACEScct working-space [0, 1].
|
||||
return decompress_acescct(decoded_01.float())
|
||||
if transfer != "logc3":
|
||||
raise ValueError(
|
||||
f"Unsupported HDR transfer {transfer!r}; use 'logc3' or 'acescct'."
|
||||
)
|
||||
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 Nodes
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@comfy_node(name="LTXVSDRToHDRWorkingSpace")
|
||||
class LTXVSDRToHDRWorkingSpace:
|
||||
"""Map SDR float RGB into ACEScct ``[0, 1]`` for HDR IC-LoRA conditioning.
|
||||
Place **before** resize when the source is display sRGB so
|
||||
the EOTF is not applied to bilinear-filtered codes.
|
||||
|
||||
Default ``srgb_gamma`` matches typical ``LoadVideo`` / MP4 decode.
|
||||
"""
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"image": ("IMAGE",),
|
||||
"color_space": (
|
||||
_SDR_COLOR_SPACE_CHOICES,
|
||||
{
|
||||
"default": "srgb_gamma",
|
||||
"tooltip": (
|
||||
"SDR ingest IDT. srgb_gamma = display sRGB (EOTF then "
|
||||
"ACEScct); srgb = scene-linear Rec.709; acescg / acescct "
|
||||
"for float EXR-style plates."
|
||||
),
|
||||
},
|
||||
),
|
||||
},
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("IMAGE",)
|
||||
RETURN_NAMES = ("acescct",)
|
||||
FUNCTION = "convert"
|
||||
CATEGORY = "Lightricks/HDR"
|
||||
DESCRIPTION = (
|
||||
"SDR → ACEScct working space for HDR IC-LoRA (pipelines SDRColorSpace). "
|
||||
"Wire LoadVideo images through this before resize / IC-LoRA guide."
|
||||
)
|
||||
|
||||
def convert(self, image: torch.Tensor, color_space: str = "srgb_gamma") -> tuple:
|
||||
space = SDRColorSpace(color_space)
|
||||
return (sdr_images_to_acescct(image, space),)
|
||||
|
||||
|
||||
def _exr_payload_for_color_space(
|
||||
image: Tensor,
|
||||
hdr_linear: Tensor,
|
||||
transfer: str,
|
||||
exr_color_space: str,
|
||||
) -> tuple[Tensor, Primaries | None, str | None]:
|
||||
"""Build EXR RGB frames and optional primaries / colorSpace tags.
|
||||
|
||||
``linear`` keeps the legacy 2.3 behaviour (raw decompressed buffer, no tags).
|
||||
Other values match pipelines ``--hdr`` EXR authoring.
|
||||
"""
|
||||
space = (exr_color_space or "linear").lower()
|
||||
transfer = (transfer or "acescct").lower()
|
||||
rgb_in = image[..., :3].float()
|
||||
hdr = hdr_linear[..., :3].float()
|
||||
|
||||
if space == "linear":
|
||||
return hdr, None, None
|
||||
|
||||
if space == "acescct":
|
||||
if transfer == "acescct":
|
||||
frames = rgb_in.clamp(0.0, 1.0)
|
||||
else:
|
||||
# LogC3 linear is treated as Rec.709 scene-linear for ACEScct authoring.
|
||||
frames = compress_acescct(
|
||||
Primaries.REC709.to(Primaries.ACESCG, hdr).clamp(min=0.0)
|
||||
)
|
||||
return frames, Primaries.ACESCG, "ACEScct"
|
||||
|
||||
if space == "acescg":
|
||||
if transfer == "acescct":
|
||||
frames = hdr # decompress_acescct → ACEScg linear
|
||||
else:
|
||||
frames = Primaries.REC709.to(Primaries.ACESCG, hdr).clamp(min=0.0)
|
||||
return frames, Primaries.ACESCG, "ACEScg"
|
||||
|
||||
if space == "srgb_linear":
|
||||
if transfer == "acescct":
|
||||
frames = Primaries.ACESCG.to(Primaries.REC709, hdr).clamp(min=0.0)
|
||||
else:
|
||||
frames = hdr.clamp(min=0.0)
|
||||
return frames, Primaries.REC709, "sRGB"
|
||||
|
||||
raise ValueError(
|
||||
f"Unsupported exr_color_space {exr_color_space!r}; "
|
||||
f"use one of {_EXR_COLOR_SPACE_CHOICES}."
|
||||
)
|
||||
|
||||
|
||||
@comfy_node(name="LTXVHDRDecodePostprocess")
|
||||
class LTXVHDRDecodePostprocess:
|
||||
"""Decompress HDR from VAE output, tonemap for preview, optionally save EXR.
|
||||
|
||||
Place after VAE Decode. Recovers linear HDR from the VAE ``[0,1]`` image and
|
||||
tonemaps to SDR for display. When ``save_exr`` is enabled, writes an EXR
|
||||
sequence (legacy ``linear``, or tagged ``acescct`` / ``srgb_linear`` /
|
||||
``acescg``).
|
||||
|
||||
For BT.2020/HLG masters, wire ``hdr_linear`` into ``LTXVSaveHLG``.
|
||||
|
||||
Outputs:
|
||||
tonemapped: SDR preview [0, 1] after Reinhard tonemap + sRGB gamma.
|
||||
hdr_linear: Raw linear HDR values [0, inf) for further processing.
|
||||
With ``transfer=acescct`` this is ACEScg scene-linear; with
|
||||
``logc3`` it is LogC3/Rec.709-style scene-linear.
|
||||
"""
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"image": ("IMAGE",),
|
||||
},
|
||||
"optional": {
|
||||
"transfer": (
|
||||
["acescct", "logc3"],
|
||||
{
|
||||
"default": "acescct",
|
||||
"tooltip": (
|
||||
"Working-space curve of the VAE decode. "
|
||||
"Use acescct for SDR→HDR IC-LoRA (after "
|
||||
"LTXVSDRToHDRWorkingSpace) and for native HDR EXR; "
|
||||
"logc3 only for legacy LogC3 HDR LoRA checkpoints."
|
||||
),
|
||||
},
|
||||
),
|
||||
"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. "
|
||||
"Affects the tonemapped preview only."
|
||||
),
|
||||
},
|
||||
),
|
||||
"save_exr": (
|
||||
"BOOLEAN",
|
||||
{
|
||||
"default": False,
|
||||
"tooltip": "Save HDR frames as an EXR sequence.",
|
||||
},
|
||||
),
|
||||
"exr_color_space": (
|
||||
_EXR_COLOR_SPACE_CHOICES,
|
||||
{
|
||||
"default": "acescg",
|
||||
"tooltip": (
|
||||
"EXR contents when save_exr is on. "
|
||||
"Default 'acescg' is tagged ACEScg scene-linear. "
|
||||
"'acescct' is ACEScct log; 'srgb_linear' is tagged "
|
||||
"Rec.709 linear; 'linear' is legacy raw decompress "
|
||||
"(untagged). All paths use OpenImageIO."
|
||||
),
|
||||
},
|
||||
),
|
||||
"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 HDR (ACEScct or legacy LogC3), tonemaps for "
|
||||
"preview, and optionally writes EXR (linear / ACEScct / sRGB linear / "
|
||||
"ACEScg). Wire hdr_linear into LTXVSaveHLG for a BT.2020/HLG master."
|
||||
)
|
||||
|
||||
def postprocess(
|
||||
self,
|
||||
image: torch.Tensor,
|
||||
transfer: str = "acescct",
|
||||
exposure: float = 0.0,
|
||||
save_exr: bool = False,
|
||||
exr_color_space: str = "acescg",
|
||||
output_dir: str = "output/hdr_exr",
|
||||
filename_prefix: str = "frame",
|
||||
half_precision: bool = True,
|
||||
) -> tuple:
|
||||
|
||||
hdr = _hdr_decompress(image, transfer=transfer)
|
||||
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:
|
||||
frames, primaries, tag = _exr_payload_for_color_space(
|
||||
image, hdr, transfer, exr_color_space
|
||||
)
|
||||
self._save_exr_frames_tagged(
|
||||
frames,
|
||||
output_dir,
|
||||
filename_prefix,
|
||||
half_precision,
|
||||
primaries=primaries,
|
||||
color_space_tag=tag,
|
||||
)
|
||||
|
||||
return (tonemapped, hdr)
|
||||
|
||||
@staticmethod
|
||||
def _resolve_exr_dir(output_dir: str) -> str:
|
||||
import folder_paths
|
||||
|
||||
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)
|
||||
return output_dir
|
||||
|
||||
@classmethod
|
||||
def _save_exr_frames_tagged(
|
||||
cls,
|
||||
frames: torch.Tensor,
|
||||
output_dir: str,
|
||||
filename_prefix: str,
|
||||
half_precision: bool,
|
||||
*,
|
||||
primaries: Primaries | None,
|
||||
color_space_tag: str | None,
|
||||
) -> None:
|
||||
output_dir = cls._resolve_exr_dir(output_dir)
|
||||
for i, frame in enumerate(frames):
|
||||
_save_exr_frame(
|
||||
frame,
|
||||
Path(output_dir) / f"{filename_prefix}_{i:05d}.exr",
|
||||
half_precision=half_precision,
|
||||
primaries=primaries,
|
||||
color_space_tag=color_space_tag,
|
||||
)
|
||||
logger.info(
|
||||
"Saved %d EXR frame(s) to %s (%s%s)",
|
||||
frames.shape[0],
|
||||
output_dir,
|
||||
"float16" if half_precision else "float32",
|
||||
f", {color_space_tag}" if color_space_tag else ", untagged linear",
|
||||
)
|
||||
|
||||
|
||||
@comfy_node(name="LTXVLoadEXRSequence")
|
||||
class LTXVLoadEXRSequence:
|
||||
"""Load an EXR still or folder and convert to ACEScct ``[0,1]`` IMAGE for the VAE.
|
||||
|
||||
Declares the source colour space the same way as pipelines ``--hdr``:
|
||||
``srgb_linear`` / ``acescg`` (compress to ACEScct) or ``acescct`` (passthrough).
|
||||
"""
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {
|
||||
"path": (
|
||||
"STRING",
|
||||
{
|
||||
"default": "",
|
||||
"tooltip": (
|
||||
"Path to a .exr still or a directory of *.exr frames. "
|
||||
"I2V: single still is fine. V2V: use a folder with 8k+1 frames. "
|
||||
"Relative paths are resolved under ComfyUI's input folder."
|
||||
),
|
||||
},
|
||||
),
|
||||
"color_space": (_COLOR_SPACE_CHOICES, {"default": "srgb_linear"}),
|
||||
"frame_rate": (
|
||||
"FLOAT",
|
||||
{
|
||||
"default": 24.0,
|
||||
"min": 1.0,
|
||||
"max": 120.0,
|
||||
"step": 0.001,
|
||||
"tooltip": "Required for EXR folders (no container fps).",
|
||||
},
|
||||
),
|
||||
},
|
||||
"optional": {
|
||||
"frame_start": ("INT", {"default": 0, "min": 0, "max": 100000}),
|
||||
"frame_cap": (
|
||||
"INT",
|
||||
{
|
||||
"default": 0,
|
||||
"min": 0,
|
||||
"max": 100000,
|
||||
"tooltip": "Max frames to load (0 = all). Still trimmed to 8k+1.",
|
||||
},
|
||||
),
|
||||
"trim_to_8k1": (
|
||||
"BOOLEAN",
|
||||
{
|
||||
"default": True,
|
||||
"tooltip": "Trim to longest 8k+1 prefix required by the video VAE.",
|
||||
},
|
||||
),
|
||||
},
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("IMAGE", "FLOAT", "INT", "INT", "INT", "AUDIO")
|
||||
RETURN_NAMES = ("images", "frame_rate", "width", "height", "frame_count", "audio")
|
||||
FUNCTION = "load"
|
||||
CATEGORY = "Lightricks/HDR"
|
||||
DESCRIPTION = (
|
||||
"Native HDR input: load EXR → ACEScct [0,1] frames for IC-LoRA / VAE encode. "
|
||||
"Also emits silent AUDIO matching duration so V2V graphs that freeze audio still wire. "
|
||||
"Native EXR path only (not SDR→HDR IC-LoRA ingest). "
|
||||
"Downstream: set LTXVPreprocess img_compression=0 — Python --hdr EXR conditioning "
|
||||
"skips JPEG/CRF; compressing ACEScct codes darkens HDR highlights."
|
||||
)
|
||||
|
||||
def load(
|
||||
self,
|
||||
path: str,
|
||||
color_space: str,
|
||||
frame_rate: float,
|
||||
frame_start: int = 0,
|
||||
frame_cap: int = 0,
|
||||
trim_to_8k1: bool = True,
|
||||
):
|
||||
root = _resolve_path(path)
|
||||
files = _list_exr_files(root)
|
||||
files = files[frame_start:]
|
||||
if frame_cap > 0:
|
||||
files = files[:frame_cap]
|
||||
if not files:
|
||||
raise ValueError("No EXR frames left after frame_start / frame_cap.")
|
||||
|
||||
cs = HDRColorSpace(color_space)
|
||||
frames = [_read_exr(f) for f in files]
|
||||
shapes = {tuple(f.shape) for f in frames}
|
||||
if len(shapes) != 1:
|
||||
raise ValueError(
|
||||
"EXR frames must share the same H×W×C; got "
|
||||
+ ", ".join(str(s) for s in sorted(shapes))
|
||||
)
|
||||
images = torch.stack(frames, dim=0)
|
||||
images = images_to_acescct(images, cs)
|
||||
if trim_to_8k1:
|
||||
images = _trim_to_8k1(images)
|
||||
|
||||
h, w = int(images.shape[1]), int(images.shape[2])
|
||||
n = int(images.shape[0])
|
||||
if n < 1:
|
||||
raise ValueError(f"No EXR frames loaded from {root}")
|
||||
audio = _silent_audio(n, float(frame_rate))
|
||||
logger.info(
|
||||
"Loaded %d EXR frame(s) from %s as %s → ACEScct (%dx%d @ %.3f fps); "
|
||||
"silent audio %d samples @ %d Hz",
|
||||
n,
|
||||
root if root.is_dir() else (root.parent if root.is_file() else root),
|
||||
cs.value,
|
||||
w,
|
||||
h,
|
||||
frame_rate,
|
||||
int(audio["waveform"].shape[-1]),
|
||||
int(audio["sample_rate"]),
|
||||
)
|
||||
return (images, float(frame_rate), w, h, n, audio)
|
||||
|
||||
|
||||
@comfy_node(name="LTXVVAEForceFloat32")
|
||||
class LTXVVAEForceFloat32:
|
||||
"""Force the video VAE to encode/decode in float32 (native HDR requirement).
|
||||
|
||||
This node does **not** clone the VAE into a separate float32 copy for the
|
||||
output. It mutates the same object that was passed in (``vae_dtype`` and
|
||||
``first_stage_model`` weights). The output socket is that same instance.
|
||||
"""
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {
|
||||
"vae": ("VAE",),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("VAE",)
|
||||
FUNCTION = "patch"
|
||||
CATEGORY = "Lightricks/HDR"
|
||||
DESCRIPTION = (
|
||||
"Forces float32 on the video VAE (vae_dtype + first_stage_model) for "
|
||||
"native HDR. IMPORTANT: this does NOT clone a float32 VAE for the "
|
||||
"output — it mutates the input object in place and returns the same "
|
||||
"instance. If you also wire that loader's VAE (or this output) into "
|
||||
"another branch, those nodes share the mutated model: they may see "
|
||||
"float16 or float32 depending on execution order (a race with the "
|
||||
"original dtype). Prefer a single chain VAELoader → this node → all "
|
||||
"encode/decode / IC-LoRA guide uses."
|
||||
)
|
||||
|
||||
def patch(self, vae):
|
||||
# In-place: Comfy Module.to(dtype) mutates weights; a wrapper copy would
|
||||
# still share first_stage_model and waste peak memory for no isolation.
|
||||
vae.vae_dtype = torch.float32
|
||||
model = getattr(vae, "first_stage_model", None)
|
||||
if model is not None:
|
||||
try:
|
||||
model.to(dtype=torch.float32)
|
||||
except Exception as exc:
|
||||
raise RuntimeError(
|
||||
"LTXVVAEForceFloat32 could not cast VAE weights to float32 "
|
||||
f"({type(exc).__name__}: {exc}). Native HDR requires float32."
|
||||
) from exc
|
||||
logger.info("VAE dtype forced to float32 for native HDR")
|
||||
return (vae,)
|
||||
|
||||
|
||||
@comfy_node(name="LTXVSaveHLG")
|
||||
class LTXVSaveHLG:
|
||||
"""Encode scene-linear ``hdr_linear`` frames to a BT.2020 / HLG 10-bit HEVC master.
|
||||
|
||||
Wire the ``hdr_linear`` output of ``LTXVHDRDecodePostprocess`` here (not the
|
||||
tonemapped preview). Matches pipelines ``encode_video`` HLG delivery.
|
||||
"""
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {
|
||||
"hdr_linear": (
|
||||
"IMAGE",
|
||||
{
|
||||
"tooltip": (
|
||||
"Scene-linear HDR from LTXVHDRDecodePostprocess. "
|
||||
"With transfer=acescct this is ACEScg linear — leave "
|
||||
"linear_primaries=acescg. For logc3, use rec709."
|
||||
),
|
||||
},
|
||||
),
|
||||
"frame_rate": (
|
||||
"FLOAT",
|
||||
{"default": 24.0, "min": 1.0, "max": 120.0, "step": 0.001},
|
||||
),
|
||||
"filename_prefix": ("STRING", {"default": "hdr/ltxv_hlg"}),
|
||||
},
|
||||
"optional": {
|
||||
"linear_primaries": (
|
||||
["acescg", "rec709"],
|
||||
{
|
||||
"default": "acescg",
|
||||
"tooltip": (
|
||||
"Primaries of hdr_linear. ACEScct decode "
|
||||
"(SDR→HDR IC-LoRA or native HDR) → acescg. "
|
||||
"Legacy LogC3 → rec709."
|
||||
),
|
||||
},
|
||||
),
|
||||
"audio": (
|
||||
"AUDIO",
|
||||
{
|
||||
"tooltip": (
|
||||
"Optional track to mux into the HLG MP4 (e.g. Decode AUDIO). "
|
||||
"Omit for video-only."
|
||||
),
|
||||
},
|
||||
),
|
||||
},
|
||||
}
|
||||
|
||||
RETURN_TYPES = ()
|
||||
OUTPUT_NODE = True
|
||||
FUNCTION = "save"
|
||||
CATEGORY = "Lightricks/HDR"
|
||||
DESCRIPTION = (
|
||||
"Writes a BT.2020/HLG 10-bit HEVC master from scene-linear hdr_linear "
|
||||
"(from LTXVHDRDecodePostprocess). Maps linear_primaries → Rec.2020 "
|
||||
"directly (no Rec.709 hop), then HLG. Optional AUDIO is muxed as AAC."
|
||||
)
|
||||
|
||||
def save(
|
||||
self,
|
||||
hdr_linear: Tensor,
|
||||
frame_rate: float,
|
||||
filename_prefix: str,
|
||||
linear_primaries: str = "acescg",
|
||||
audio: dict | None = None,
|
||||
):
|
||||
if hdr_linear.ndim != 4 or hdr_linear.shape[-1] < 3:
|
||||
raise ValueError(f"Expected IMAGE [F,H,W,C]; got {tuple(hdr_linear.shape)}")
|
||||
|
||||
# ACEScg/Rec.709 → Rec.2020 inside HLG encode (Primaries.to_rec2020 /
|
||||
# matrix_to_rec2020). No Rec.709 intermediate that would crush OOG early.
|
||||
linear = hdr_linear[..., :3].float()
|
||||
prim = (linear_primaries or "acescg").lower()
|
||||
if prim == "acescg":
|
||||
src_primaries = Primaries.ACESCG
|
||||
elif prim == "rec709":
|
||||
src_primaries = Primaries.REC709
|
||||
else:
|
||||
raise ValueError(
|
||||
f"Unsupported linear_primaries {linear_primaries!r}; use acescg or rec709."
|
||||
)
|
||||
|
||||
# Even dims for 4:2:0
|
||||
h, w = linear.shape[1], linear.shape[2]
|
||||
if h % 2 or w % 2:
|
||||
linear = linear[:, : h - (h % 2), : w - (w % 2), :]
|
||||
|
||||
full_output_folder, filename, counter, _subfolder, _ = (
|
||||
folder_paths.get_save_image_path(
|
||||
filename_prefix,
|
||||
folder_paths.get_output_directory(),
|
||||
linear.shape[2],
|
||||
linear.shape[1],
|
||||
)
|
||||
)
|
||||
hlg_path = Path(full_output_folder) / f"{filename}_{counter:05d}_hlg.mp4"
|
||||
_encode_hlg_mp4(
|
||||
linear,
|
||||
hlg_path,
|
||||
fps=float(frame_rate),
|
||||
audio=audio,
|
||||
primaries=src_primaries,
|
||||
)
|
||||
logger.info(
|
||||
"Wrote HLG 10-bit master to %s%s",
|
||||
hlg_path,
|
||||
" (with audio)" if audio is not None else "",
|
||||
)
|
||||
return {}
|
||||
@@ -0,0 +1,906 @@
|
||||
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
|
||||
from .tiled_fusion_plan import snap_frames
|
||||
|
||||
FRAME_IDX_TOOLTIP = (
|
||||
"RoPE position of the reference, in pixel frames.\n"
|
||||
"0 aligns the reference with the generated clip. Use for video references.\n"
|
||||
">0 places the reference at a specific pixel-frame position in the clip.\n"
|
||||
"-1 places the reference outside the generated frame range, while still using it "
|
||||
"for attention conditioning. Use for single-image references.\n"
|
||||
"For multi-frame references, >0 values are rounded down to 1 mod 8."
|
||||
)
|
||||
|
||||
|
||||
@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. "
|
||||
"For LTXVTiledFusionSampler temporal windows, turn on use_streaming and "
|
||||
"set tile_frames to the same 8n+1 value as the sampler (97 for LTX)."
|
||||
),
|
||||
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_IDX_TOOLTIP,
|
||||
),
|
||||
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.",
|
||||
),
|
||||
io.Boolean.Input(
|
||||
"use_streaming",
|
||||
optional=True,
|
||||
default=False,
|
||||
tooltip=(
|
||||
"Off (default): encode the whole clip as one IC-LoRA guide. "
|
||||
"On: re-encode one fresh guide per LTXVTiledFusionSampler "
|
||||
"temporal window and write the plan the sampler reads. "
|
||||
"Set tile_frames on this node and the sampler to the same "
|
||||
"value (97 for LTX)."
|
||||
),
|
||||
),
|
||||
io.Int.Input(
|
||||
"tile_frames",
|
||||
optional=True,
|
||||
default=97,
|
||||
min=0,
|
||||
max=1000,
|
||||
step=1,
|
||||
tooltip=(
|
||||
"Window length in pixel frames, used only when use_streaming "
|
||||
"is on. Snapped to 8n+1. Must match LTXVTiledFusionSampler. "
|
||||
"Ignored when use_streaming is off."
|
||||
),
|
||||
),
|
||||
],
|
||||
outputs=[
|
||||
io.Conditioning.Output("positive"),
|
||||
io.Conditioning.Output("negative"),
|
||||
io.Latent.Output("latent"),
|
||||
],
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def resolve_frame_idx(
|
||||
cls,
|
||||
positive,
|
||||
frame_idx,
|
||||
latent_length,
|
||||
image_length,
|
||||
guide_latent,
|
||||
scale_factors,
|
||||
):
|
||||
"""Map frame_idx to the position append_keyframe should use.
|
||||
|
||||
A negative value is a RoPE position before the clip, not an index counted back
|
||||
from its end, so it bypasses get_latent_index. That is what lets a one-frame
|
||||
reference sit at -1: adjacent to the first generated frame without landing in
|
||||
its slot, which is the only position a one-frame guide can occupy that makes it
|
||||
indistinguishable from the frame being generated. The length assert is
|
||||
meaningless there -- the guide is outside the clip by construction.
|
||||
"""
|
||||
if frame_idx < 0:
|
||||
return frame_idx
|
||||
frame_idx, latent_idx = nodes_lt.LTXVAddGuide.get_latent_index(
|
||||
positive, latent_length, image_length, frame_idx, scale_factors
|
||||
)
|
||||
assert (
|
||||
latent_idx + guide_latent.shape[2] <= latent_length
|
||||
), "Conditioning frames exceed the length of the latent sequence."
|
||||
return frame_idx
|
||||
|
||||
@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 encode_windowed_guides(
|
||||
cls,
|
||||
positive,
|
||||
negative,
|
||||
vae,
|
||||
latent_image,
|
||||
noise_mask,
|
||||
image,
|
||||
latent_length,
|
||||
latent_width,
|
||||
latent_height,
|
||||
scale_factors,
|
||||
strength,
|
||||
crop,
|
||||
tile_frames,
|
||||
attention_strength=1.0,
|
||||
attention_mask=None,
|
||||
frame_idx=0,
|
||||
use_tiled_encode=False,
|
||||
):
|
||||
"""Re-encode each fusion window as its own clip, or return None.
|
||||
|
||||
Mid-clip causal latents cannot be sliced into windows, so each window's
|
||||
pixels are encoded as a clip starting at 0. The shared plan is written
|
||||
onto the conditioning so LTXVTiledFusionSampler pairs window w with
|
||||
guide block w. Returns None when the clip fits one window.
|
||||
"""
|
||||
from .iclora_attention import append_guide_attention_entry
|
||||
from .tiled_fusion_plan import attach_temporal_plan, temporal_plan
|
||||
|
||||
t_scale = scale_factors[0]
|
||||
tf_lat = max(1, (tile_frames - 1) // t_scale + 1)
|
||||
if not (0 < tf_lat < latent_length):
|
||||
return None
|
||||
if int(frame_idx) != 0:
|
||||
raise ValueError(
|
||||
"use_streaming encodes each window as a fresh clip at frame 0; "
|
||||
f"frame_idx={frame_idx} is not supported. Use the whole-clip path "
|
||||
"(use_streaming off) for stills or mid-clip placement."
|
||||
)
|
||||
if use_tiled_encode:
|
||||
raise ValueError(
|
||||
"use_streaming cannot use tiled VAE encode; tiled guide encodes "
|
||||
"imprint a grid into the conditioning. Leave use_tiled_encode off."
|
||||
)
|
||||
keep = ((image.shape[0] - 1) // t_scale) * t_scale + 1
|
||||
image = image[:keep]
|
||||
ts = temporal_plan(latent_length, tf_lat)
|
||||
span_px = (tf_lat - 1) * t_scale + 1
|
||||
need = ts[-1] * t_scale + span_px
|
||||
if image.shape[0] < need:
|
||||
raise ValueError(
|
||||
f"windowed guide needs {need} frames for {latent_length} "
|
||||
f"content latents, found {image.shape[0]}"
|
||||
)
|
||||
for w, t0 in enumerate(ts):
|
||||
span = image[t0 * t_scale : t0 * t_scale + span_px]
|
||||
_, guide_latent = cls.encode(
|
||||
vae,
|
||||
latent_width,
|
||||
latent_height,
|
||||
span,
|
||||
scale_factors,
|
||||
1.0,
|
||||
crop,
|
||||
False,
|
||||
256,
|
||||
64,
|
||||
)
|
||||
shape = list(guide_latent.shape[2:])
|
||||
toks = shape[0] * shape[1] * shape[2]
|
||||
positive, negative, latent_image, noise_mask = (
|
||||
nodes_lt.LTXVAddGuide.append_keyframe(
|
||||
positive,
|
||||
negative,
|
||||
0,
|
||||
latent_image,
|
||||
noise_mask,
|
||||
guide_latent,
|
||||
strength,
|
||||
scale_factors,
|
||||
guide_mask=None,
|
||||
latent_downscale_factor=1.0,
|
||||
causal_fix=True,
|
||||
)
|
||||
)
|
||||
# Pixel mask is (1, 1, F, H, W) after normalize_mask. F==1 broadcasts;
|
||||
# otherwise it must match this window's encoded span, not the full clip.
|
||||
window_mask = attention_mask
|
||||
if window_mask is not None and window_mask.shape[2] != 1:
|
||||
start = t0 * t_scale
|
||||
window_mask = window_mask[:, :, start : start + span_px]
|
||||
positive = append_guide_attention_entry(
|
||||
positive,
|
||||
toks,
|
||||
shape,
|
||||
attention_strength=attention_strength,
|
||||
attention_mask=window_mask,
|
||||
)
|
||||
negative = append_guide_attention_entry(
|
||||
negative,
|
||||
toks,
|
||||
shape,
|
||||
attention_strength=attention_strength,
|
||||
attention_mask=window_mask,
|
||||
)
|
||||
logging.info(
|
||||
"window %d/%d: guide frames %d-%d -> %s",
|
||||
w + 1,
|
||||
len(ts),
|
||||
t0 * t_scale,
|
||||
t0 * t_scale + span_px - 1,
|
||||
tuple(shape),
|
||||
)
|
||||
plan = {
|
||||
"tile_frames": int(tile_frames),
|
||||
"tf_lat": int(tf_lat),
|
||||
"n_content": int(latent_length),
|
||||
"ts": [int(v) for v in ts],
|
||||
}
|
||||
positive = attach_temporal_plan(positive, plan)
|
||||
negative = attach_temporal_plan(negative, plan)
|
||||
logging.info("%d per-window guides (t0 = %s)", len(ts), plan["ts"])
|
||||
return io.NodeOutput(
|
||||
positive, negative, {"samples": latent_image, "noise_mask": noise_mask}
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def maybe_encode_windowed(
|
||||
cls,
|
||||
positive,
|
||||
negative,
|
||||
vae,
|
||||
latent_image,
|
||||
noise_mask,
|
||||
image,
|
||||
latent_length,
|
||||
latent_width,
|
||||
latent_height,
|
||||
scale_factors,
|
||||
strength,
|
||||
crop,
|
||||
use_streaming,
|
||||
tile_frames,
|
||||
latent_downscale_factor=1.0,
|
||||
attention_strength=1.0,
|
||||
attention_mask=None,
|
||||
frame_idx=0,
|
||||
use_tiled_encode=False,
|
||||
):
|
||||
"""Run the windowed encode only when use_streaming is on. Else None."""
|
||||
if not use_streaming:
|
||||
return None
|
||||
if latent_downscale_factor != 1:
|
||||
raise ValueError(
|
||||
"use_streaming requires latent_downscale_factor=1; "
|
||||
"windowed fusion guides cannot be dilated."
|
||||
)
|
||||
tf = snap_frames(int(tile_frames) if tile_frames and tile_frames > 0 else 97)
|
||||
return cls.encode_windowed_guides(
|
||||
positive,
|
||||
negative,
|
||||
vae,
|
||||
latent_image,
|
||||
noise_mask,
|
||||
image,
|
||||
latent_length,
|
||||
latent_width,
|
||||
latent_height,
|
||||
scale_factors,
|
||||
strength,
|
||||
crop,
|
||||
tf,
|
||||
attention_strength=attention_strength,
|
||||
attention_mask=attention_mask,
|
||||
frame_idx=frame_idx,
|
||||
use_tiled_encode=use_tiled_encode,
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def execute(
|
||||
cls,
|
||||
positive,
|
||||
negative,
|
||||
vae,
|
||||
latent,
|
||||
image,
|
||||
frame_idx,
|
||||
strength,
|
||||
latent_downscale_factor,
|
||||
crop,
|
||||
use_tiled_encode,
|
||||
tile_size,
|
||||
tile_overlap,
|
||||
use_streaming=False,
|
||||
tile_frames=0,
|
||||
) -> 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
|
||||
|
||||
windowed = cls.maybe_encode_windowed(
|
||||
positive,
|
||||
negative,
|
||||
vae,
|
||||
latent_image,
|
||||
noise_mask,
|
||||
image,
|
||||
latent_length,
|
||||
latent_width,
|
||||
latent_height,
|
||||
scale_factors,
|
||||
strength,
|
||||
crop,
|
||||
use_streaming,
|
||||
tile_frames,
|
||||
latent_downscale_factor=latent_downscale_factor,
|
||||
frame_idx=frame_idx,
|
||||
use_tiled_encode=use_tiled_encode,
|
||||
)
|
||||
if windowed is not None:
|
||||
return windowed
|
||||
|
||||
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 = cls.resolve_frame_idx(
|
||||
positive, frame_idx, latent_length, len(image), guide_latent, scale_factors
|
||||
)
|
||||
|
||||
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. "
|
||||
"use_streaming + matching tile_frames re-encodes one guide per Tiled Fusion window."
|
||||
),
|
||||
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_IDX_TOOLTIP,
|
||||
),
|
||||
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."
|
||||
),
|
||||
),
|
||||
io.Boolean.Input(
|
||||
"use_streaming",
|
||||
optional=True,
|
||||
default=False,
|
||||
tooltip=(
|
||||
"Off (default): whole-clip guide. On: one fresh encode per "
|
||||
"fusion temporal window. Must match the sampler's tile_frames."
|
||||
),
|
||||
),
|
||||
io.Int.Input(
|
||||
"tile_frames",
|
||||
optional=True,
|
||||
default=97,
|
||||
min=0,
|
||||
max=1000,
|
||||
step=1,
|
||||
tooltip=(
|
||||
"Used only when use_streaming is on. Snapped to 8n+1. "
|
||||
"97 for LTX. Must match the sampler."
|
||||
),
|
||||
),
|
||||
],
|
||||
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,
|
||||
use_streaming=False,
|
||||
tile_frames=0,
|
||||
) -> 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
|
||||
|
||||
windowed = cls.maybe_encode_windowed(
|
||||
positive,
|
||||
negative,
|
||||
vae,
|
||||
latent_image,
|
||||
noise_mask,
|
||||
image,
|
||||
latent_length,
|
||||
latent_width,
|
||||
latent_height,
|
||||
scale_factors,
|
||||
strength,
|
||||
crop,
|
||||
use_streaming,
|
||||
tile_frames,
|
||||
latent_downscale_factor=latent_downscale_factor,
|
||||
attention_strength=attention_strength,
|
||||
attention_mask=normalize_mask(attention_mask),
|
||||
frame_idx=frame_idx,
|
||||
use_tiled_encode=use_tiled_encode,
|
||||
)
|
||||
if windowed is not None:
|
||||
return windowed
|
||||
|
||||
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 = cls.resolve_frame_idx(
|
||||
positive, frame_idx, latent_length, len(image), guide_latent, scale_factors
|
||||
)
|
||||
|
||||
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,80 @@
|
||||
"""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 set_last_guide_attention_strength(conditioning, strength):
|
||||
"""Override the attention strength of the most recently appended guide entry.
|
||||
|
||||
Separate from a guide's denoise strength: one says how clean the reference is
|
||||
kept, the other how strongly the rest of the sequence may attend to it.
|
||||
"""
|
||||
entries = _get_guide_attention_entries(conditioning)
|
||||
if not entries:
|
||||
return conditioning
|
||||
updated = [*entries[:-1], {**entries[-1], "strength": strength}]
|
||||
return _set_guide_attention_entries(conditioning, updated)
|
||||
|
||||
|
||||
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
|
||||
+52
-77
@@ -7,6 +7,7 @@ import torch
|
||||
from comfy.ldm.lightricks.vae.audio_vae import LATENT_DOWNSAMPLE_FACTOR
|
||||
from comfy.nested_tensor import NestedTensor
|
||||
|
||||
from .iclora_attention import set_last_guide_attention_strength
|
||||
from .nodes_registry import comfy_node
|
||||
|
||||
|
||||
@@ -397,85 +398,47 @@ class LTXVDilateLatent:
|
||||
return (latent,)
|
||||
|
||||
|
||||
@comfy_node(name="LTXVAddLatentGuide")
|
||||
class LTXVAddLatentGuide:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"vae": ("VAE",),
|
||||
"positive": ("CONDITIONING",),
|
||||
"negative": ("CONDITIONING",),
|
||||
"latent": ("LATENT",),
|
||||
"guiding_latent": ("LATENT",),
|
||||
"latent_idx": (
|
||||
"INT",
|
||||
{
|
||||
"default": 0,
|
||||
"min": -9999,
|
||||
"max": 9999,
|
||||
"step": 1,
|
||||
"tooltip": "Latent index to start the conditioning at. Can be negative to"
|
||||
"indicate that the conditioning is on the frames before the latent.",
|
||||
},
|
||||
),
|
||||
"strength": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0}),
|
||||
}
|
||||
}
|
||||
def add_latent_guide(
|
||||
*,
|
||||
vae,
|
||||
positive,
|
||||
negative,
|
||||
latent,
|
||||
guiding_latent,
|
||||
latent_idx,
|
||||
strength,
|
||||
attention_mask=None,
|
||||
):
|
||||
"""Pin an already-encoded latent as a guide, via ComfyUI's native node.
|
||||
|
||||
RETURN_TYPES = ("CONDITIONING", "CONDITIONING", "LATENT")
|
||||
RETURN_NAMES = ("positive", "negative", "latent")
|
||||
ComfyUI ships LTXVAddLatentGuide itself (Comfy-Org/ComfyUI#16176), so the
|
||||
implementation lives there rather than being duplicated here; this is only the
|
||||
call-shape adapter the pack's samplers want -- keyword arguments and a plain
|
||||
tuple instead of a NodeOutput.
|
||||
|
||||
CATEGORY = "ltxtricks"
|
||||
FUNCTION = "generate"
|
||||
|
||||
DESCRIPTION = "Adds a keyframe or a video segment at a specific frame index."
|
||||
|
||||
def generate(
|
||||
self, vae, positive, negative, latent, guiding_latent, latent_idx, strength
|
||||
):
|
||||
noise_mask = nodes_lt.get_noise_mask(latent)
|
||||
latent = latent["samples"]
|
||||
guide = guiding_latent["samples"]
|
||||
|
||||
assert (
|
||||
latent.shape[4] % guide.shape[4] == 0
|
||||
and latent.shape[3] % guide.shape[3] == 0
|
||||
), "The ratio of the height and width of the latents and optional_guiding_latents must be an integer"
|
||||
|
||||
guiding_latent = LTXVDilateLatent().dilate_latent(
|
||||
guiding_latent,
|
||||
horizontal_scale=latent.shape[4] // guide.shape[4],
|
||||
vertical_scale=latent.shape[3] // guide.shape[3],
|
||||
)[0]
|
||||
|
||||
guide = guiding_latent["samples"]
|
||||
guide_mask = guiding_latent.get("noise_mask", None)
|
||||
|
||||
scale_factors = vae.downscale_index_formula
|
||||
|
||||
if latent_idx <= 0:
|
||||
frame_idx = latent_idx * scale_factors[0]
|
||||
else:
|
||||
frame_idx = 1 + (latent_idx - 1) * scale_factors[0]
|
||||
|
||||
positive, negative, latent, noise_mask = nodes_lt.LTXVAddGuide.append_keyframe(
|
||||
positive=positive,
|
||||
negative=negative,
|
||||
frame_idx=frame_idx,
|
||||
latent_image=latent,
|
||||
noise_mask=noise_mask,
|
||||
guiding_latent=guide,
|
||||
strength=strength,
|
||||
scale_factors=scale_factors,
|
||||
guide_mask=guide_mask,
|
||||
)
|
||||
|
||||
return (
|
||||
positive,
|
||||
negative,
|
||||
{"samples": latent, "noise_mask": noise_mask},
|
||||
)
|
||||
Deliberately does not forward pixel_frame_offset: that input exists only in
|
||||
Lightricks' ComfyUI fork, and the pack has to keep working against a stock
|
||||
install.
|
||||
"""
|
||||
positive, negative, latent = nodes_lt.LTXVAddLatentGuide.execute(
|
||||
positive=positive,
|
||||
negative=negative,
|
||||
vae=vae,
|
||||
latent=latent,
|
||||
guiding_latent=guiding_latent,
|
||||
latent_idx=latent_idx,
|
||||
strength=strength,
|
||||
attention_mask=attention_mask,
|
||||
)
|
||||
if attention_mask is None:
|
||||
# The core node records `strength` as the guide's attention strength too, so a
|
||||
# weakly pinned guide would also become weakly attended. The implementation this
|
||||
# replaced always left attention at full strength, and the samplers calling this
|
||||
# pass strengths well below 1.0 (0.15 in the tiled sampler, 0.5 when extending),
|
||||
# so keep the two decoupled rather than silently reweighting their attention.
|
||||
positive = set_last_guide_attention_strength(positive, 1.0)
|
||||
negative = set_last_guide_attention_strength(negative, 1.0)
|
||||
return positive, negative, latent
|
||||
|
||||
|
||||
@comfy_node(name="LTXVImgToVideoConditionOnly")
|
||||
@@ -809,6 +772,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,266 @@
|
||||
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,
|
||||
)
|
||||
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 = F.pad(image, (0, pad_right, 0, pad_down), mode="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 = F.pad(images, padding, mode=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,10 @@
|
||||
# Required for HDR primaries (same as ltx-core); no baked-matrix fallback.
|
||||
colour-science>=0.4.4
|
||||
diffusers
|
||||
einops
|
||||
huggingface_hub>=0.25.2
|
||||
kornia
|
||||
ninja~=1.11.1.4
|
||||
# Required for HDR EXR read/write (same as ltx-pipelines); no OpenCV EXR fallback.
|
||||
openimageio
|
||||
transformers[timm]>=4.50.0
|
||||
|
||||
@@ -0,0 +1,46 @@
|
||||
"""RoPE frequency helpers used by the embeddings connector.
|
||||
|
||||
These were previously imported from ``comfy.ldm.lightricks.model``, but upstream
|
||||
ComfyUI does not ship them (it folds the equivalent logic into
|
||||
``freqs_cis_matrix``). They are kept here, local to this node, so the connector
|
||||
keeps its exact original behavior without relying on fork-only additions to core
|
||||
ComfyUI.
|
||||
"""
|
||||
|
||||
import torch
|
||||
|
||||
|
||||
def interleaved_freqs_cis(freqs, pad_size):
|
||||
cos_freq = freqs.cos().repeat_interleave(2, dim=-1)
|
||||
sin_freq = freqs.sin().repeat_interleave(2, dim=-1)
|
||||
if pad_size != 0:
|
||||
cos_padding = torch.ones_like(cos_freq[:, :, :pad_size])
|
||||
sin_padding = torch.zeros_like(cos_freq[:, :, :pad_size])
|
||||
cos_freq = torch.cat([cos_padding, cos_freq], dim=-1)
|
||||
sin_freq = torch.cat([sin_padding, sin_freq], dim=-1)
|
||||
return cos_freq, sin_freq
|
||||
|
||||
|
||||
def split_freqs_cis(freqs, pad_size, num_attention_heads):
|
||||
cos_freq = freqs.cos()
|
||||
sin_freq = freqs.sin()
|
||||
|
||||
if pad_size != 0:
|
||||
cos_padding = torch.ones_like(cos_freq[:, :, :pad_size])
|
||||
sin_padding = torch.zeros_like(sin_freq[:, :, :pad_size])
|
||||
|
||||
cos_freq = torch.concatenate([cos_padding, cos_freq], axis=-1)
|
||||
sin_freq = torch.concatenate([sin_padding, sin_freq], axis=-1)
|
||||
|
||||
B, T, half_HD = cos_freq.shape
|
||||
|
||||
cos_freq = cos_freq.reshape(
|
||||
B, T, num_attention_heads, half_HD // num_attention_heads
|
||||
)
|
||||
sin_freq = sin_freq.reshape(
|
||||
B, T, num_attention_heads, half_HD // num_attention_heads
|
||||
)
|
||||
|
||||
cos_freq = torch.swapaxes(cos_freq, 1, 2)
|
||||
sin_freq = torch.swapaxes(sin_freq, 1, 2)
|
||||
return cos_freq, sin_freq
|
||||
@@ -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)
|
||||
@@ -0,0 +1,80 @@
|
||||
"""Shared temporal-window plan for tiled fusion.
|
||||
|
||||
The guide node encodes one fresh IC-LoRA clip per window and writes the plan
|
||||
onto conditioning. The sampler reads it so window w always pairs with guide
|
||||
block w. Both sides must use the same origins; slicing a whole-clip causal
|
||||
guide cannot be repaired by rewriting coordinates.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
|
||||
# Minimum temporal overlap fraction. Keep in sync with the sampler env knobs.
|
||||
_T_OV = float(os.environ.get("LTXV_TILED_FUSION_T_OVERLAP", "0.5"))
|
||||
|
||||
TEMPORAL_PLAN_KEY = "ltxv_tiled_fusion_temporal_plan"
|
||||
GEAR_TEMPORAL_PLAN_KEY = "gear_temporal_plan"
|
||||
FRAME_STEP = 8
|
||||
|
||||
|
||||
def snap_frames(n: int) -> int:
|
||||
"""Nearest 8n+1 length. Ties keep the shorter clip (no pad)."""
|
||||
n = max(1, int(n))
|
||||
below = ((n - 1) // FRAME_STEP) * FRAME_STEP + 1
|
||||
above = below + FRAME_STEP
|
||||
if above - n < n - below:
|
||||
return above
|
||||
return below
|
||||
|
||||
|
||||
def plan_grid(canvas_lat, tile_lat, min_overlap):
|
||||
"""Latent-aligned tile origins covering [0, canvas_lat) with >= min_overlap."""
|
||||
if canvas_lat <= tile_lat:
|
||||
return [0]
|
||||
tile_lat = max(1, int(tile_lat))
|
||||
canvas_lat = int(canvas_lat)
|
||||
# Overlap cannot reach tile_lat, so an unclamped request would never exit.
|
||||
min_overlap = min(max(0, min_overlap), tile_lat - 1)
|
||||
n = 2
|
||||
while n < canvas_lat and (n * tile_lat - canvas_lat) / (n - 1) < min_overlap:
|
||||
n += 1
|
||||
step = (canvas_lat - tile_lat) / (n - 1)
|
||||
return [round(i * step) for i in range(n)]
|
||||
|
||||
|
||||
def shifted(base, canvas_lat, tile_lat, off):
|
||||
"""Shift interior origins by `off`; endpoints stay pinned so coverage holds."""
|
||||
out = []
|
||||
for i, v in enumerate(base):
|
||||
if i in (0, len(base) - 1):
|
||||
out.append(v)
|
||||
else:
|
||||
out.append(max(1, min(canvas_lat - tile_lat - 1, v + off)))
|
||||
return sorted(set(out))
|
||||
|
||||
|
||||
def temporal_plan(n_lat, tf_lat):
|
||||
"""Shared window plan so a guide node and the sampler always agree."""
|
||||
return plan_grid(n_lat, tf_lat, max(0, int(tf_lat * _T_OV)))
|
||||
|
||||
|
||||
def attach_temporal_plan(cond, plan):
|
||||
"""Write the plan onto every conditioning entry."""
|
||||
out = []
|
||||
for emb, d in cond:
|
||||
d = dict(d)
|
||||
d[TEMPORAL_PLAN_KEY] = plan
|
||||
out.append([emb, d])
|
||||
return out
|
||||
|
||||
|
||||
def read_temporal_plan(positive):
|
||||
"""Return the plan from conditioning, including the Gear-era key."""
|
||||
for _, d in positive:
|
||||
if not isinstance(d, dict):
|
||||
continue
|
||||
plan = d.get(TEMPORAL_PLAN_KEY) or d.get(GEAR_TEMPORAL_PLAN_KEY)
|
||||
if plan:
|
||||
return plan
|
||||
return None
|
||||
@@ -0,0 +1,85 @@
|
||||
# LTXVTiledFusionSampler
|
||||
|
||||
Per-step tiled latent fusion for IC-LoRA (and other LTX conditioning) at
|
||||
resolutions larger than the model's trained window.
|
||||
|
||||
This is **not** `LTXVTiledSampler`. That node independently samples a grid of
|
||||
tiles. This node keeps **one** latent canvas and **one** noise field, runs the
|
||||
model on overlapping crops at **every** denoise step, and Gaussian-blends the
|
||||
stepped tiles back onto the canvas. Tiles share a trajectory, so seams do not
|
||||
form. Peak VRAM tracks **spatial tile size** and the **temporal window**
|
||||
actually stepped (one content window plus one guide block when streaming),
|
||||
not the full canvas. Cached tile latents/noise/masks follow `canvas_device`;
|
||||
processed conds stay on the GPU. Streaming also does one full-res VAE encode
|
||||
per temporal window.
|
||||
|
||||
## Graph shape
|
||||
|
||||
```
|
||||
UNET + IC-LoRA
|
||||
Gemma CLIP → CLIPTextEncode (pos / neg)
|
||||
guide frames → LTX Add Video IC-LoRA Guide (positive, negative, latents)
|
||||
KSamplerSelect (euler / heun / …) → SAMPLER
|
||||
sigmas (Manual Sigmas or any scheduler, descending, ending at 0)
|
||||
│
|
||||
▼
|
||||
LTXV Tiled Fusion Sampler → LATENT → VAE Decode (tiled)
|
||||
```
|
||||
|
||||
The latent **must** come from the IC-LoRA guide node. It carries appended guide
|
||||
frames and a noise mask. This node crops those extra frames itself, so
|
||||
`LTXVCropGuides` is not needed.
|
||||
|
||||
Leave `use_streaming` **off** on the guide (the default) for a whole-clip encode.
|
||||
For clips longer than the trained window, turn `use_streaming` **on** and set the
|
||||
same `tile_frames` (typically 97) on both the guide and this sampler. That
|
||||
re-encodes each temporal window as its own clip and writes
|
||||
`ltxv_tiled_fusion_temporal_plan` onto the conditioning. The sampler disables
|
||||
windowing if the plan is missing. Do not slice a whole-clip guide: mid-clip
|
||||
causal latents are not fresh-clip-start latents.
|
||||
|
||||
## Sampler input
|
||||
|
||||
The `SAMPLER` socket is the same object `SamplerCustomAdvanced` takes. Fusion
|
||||
owns the sigma loop and calls `sampler.sampler_function` once per tile with a
|
||||
two-sigma slice. Use discrete step methods (`euler`, `heun`, `dpm_2`). History
|
||||
methods (`dpmpp_2m`, `lms`) lose their multi-step memory. Adaptive samplers
|
||||
cannot fuse. Ancestral methods add independent noise per tile and can seam.
|
||||
|
||||
## Settings that matter
|
||||
|
||||
| Input | Notes |
|
||||
|---|---|
|
||||
| `tile_width` / `tile_height` | The IC-LoRA's trained spatial window in pixels (multiples of 32). Smaller tiles see content at the wrong scale. |
|
||||
| `overlap_frac` | Stay at **0.5** or above. Below that a periodic grid appears on structured content. |
|
||||
| `blend_var` | Gaussian overlap variance. Default 0.05. |
|
||||
| `grid_cycle` | 1 = fixed grid (cheapest), 4 = shift the grid across steps. Spatial only. |
|
||||
| `tile_frames` | Temporal window in pixel frames (97 for LTX). 0 = one extent. >0 needs `use_streaming` on `LTXAddVideoICLoRAGuide` with the same value. |
|
||||
| `vae` | Optional. Supplies temporal/spatial downscale (otherwise read from the diffusion model). |
|
||||
| `use_tiled_encode` on the guide node | Must stay **false**. Spatially tiled guide encodes imprint a grid into the conditioning. |
|
||||
| `use_streaming` on the guide node | Off by default. On = per-window fresh encode + plan. |
|
||||
|
||||
`overlap_frac`, `blend_var`, and `grid_cycle` are spatial fusion knobs. They do
|
||||
not replace an untiled composition pass when the LoRA is generating a new look
|
||||
at a canvas far above the trained window.
|
||||
|
||||
## Example workflows
|
||||
|
||||
Under [`example_workflows/2.5/`](./example_workflows/2.5/):
|
||||
|
||||
| Graph | What it shows |
|
||||
| ----- | ------------- |
|
||||
| [`LTX-2.5_V2V_TiledFusion_Upscale.json`](./example_workflows/2.5/LTX-2.5_V2V_TiledFusion_Upscale.json) | Detail-refine IC-LoRA; `output_size` FullHD / 4K / 8K, HD tiles |
|
||||
| [`LTX-2.5_V2V_TiledFusion_Native_4K_8K.json`](./example_workflows/2.5/LTX-2.5_V2V_TiledFusion_Native_4K_8K.json) | Stage-1 FullHD composition, then 4K or 8K refine |
|
||||
| [`LTX-2.5_V2V_TiledFusion_SDR_to_HDR.json`](./example_workflows/2.5/LTX-2.5_V2V_TiledFusion_SDR_to_HDR.json) | Same ladder with the HDR IC-LoRA + `LTXVHDRDecodePostprocess` |
|
||||
|
||||
All of them take `SAMPLER` from `KSamplerSelect` (`euler`) and sigmas from
|
||||
`Manual Sigmas` (descending, ending at 0). Size combos live on **Preprocess**
|
||||
(Get Tiling Sizes + canvas resize are inside that subgraph). Keep
|
||||
`use_tiled_encode` **false** on every `LTXAddVideoICLoRAGuide`. Turn
|
||||
`use_streaming` **on** and set the same `tile_frames` on the guide and the
|
||||
sampler when the clip is longer than one window.
|
||||
|
||||
Prompt for look and style only — lighting, palette, sharpness, grade. Do not
|
||||
name specific objects in the scene; after upscale, named details can appear in
|
||||
every spatial tile. The negative prompt is used as written.
|
||||
@@ -0,0 +1,832 @@
|
||||
"""Per-step tiled latent fusion for LTX video models.
|
||||
|
||||
Keeps one latent canvas and one noise field, runs the model once per overlapping
|
||||
spatial tile at every denoise step (each tile conditioned on its own crop of the
|
||||
IC-LoRA guide), and merges the stepped tiles back onto the canvas with Gaussian
|
||||
weights. Tiles follow a single coupled trajectory, so seams do not form. Peak
|
||||
VRAM is set by the spatial tile and by the temporal window actually stepped
|
||||
(content window plus one guide block when streaming), not the full canvas
|
||||
spatial size. Cached tile latents/noise/masks follow canvas_device; processed
|
||||
conds stay on the GPU for the run.
|
||||
|
||||
This is distinct from LTXVTiledSampler, which independently samples a grid of
|
||||
tiles. Fusion has to happen between denoise steps, so this node calls
|
||||
``sampler.sampler_function`` once per tile with a two-sigma slice rather than
|
||||
``sampler.sample()`` on the whole trajectory.
|
||||
|
||||
The latent input must be the one returned by an IC-LoRA guide node: it carries
|
||||
guide frames appended along the frame axis plus the noise mask, and the
|
||||
conditioning carries ``keyframe_idxs`` / ``guide_attention_entries``, all of
|
||||
which this node crops per tile.
|
||||
"""
|
||||
|
||||
import copy
|
||||
import logging
|
||||
import os
|
||||
|
||||
import comfy.model_management
|
||||
import comfy.sampler_helpers
|
||||
import comfy.samplers
|
||||
import comfy.utils
|
||||
import torch
|
||||
|
||||
from .nodes_registry import comfy_node
|
||||
from .tiled_fusion_plan import (
|
||||
plan_grid,
|
||||
read_temporal_plan,
|
||||
shifted,
|
||||
snap_frames,
|
||||
temporal_plan,
|
||||
)
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# Per-step/per-window mean logging for drift forensics (set env before launch).
|
||||
_DEBUG_STATS = os.environ.get("LTXV_TILED_FUSION_DEBUG", "") not in ("", "0")
|
||||
|
||||
# Temporal-window experiment knobs (env, not node inputs, while under study):
|
||||
# LTXV_TILED_FUSION_T_OVERLAP minimum temporal overlap fraction (default 0.5)
|
||||
# LTXV_TILED_FUSION_SKIP_W0_LAT 0 = also merge the first latent of window w>0
|
||||
# (that slot is a fresh-clip-start / causal frame
|
||||
# and poisons a mid-clip canvas position).
|
||||
# Default on: skip that latent.
|
||||
_SKIP_FIRST_LAT = os.environ.get("LTXV_TILED_FUSION_SKIP_W0_LAT", "1") not in (
|
||||
"0",
|
||||
"",
|
||||
)
|
||||
|
||||
_NODE = "LTXV Tiled Fusion Sampler"
|
||||
|
||||
|
||||
def _scale_factors(model, vae=None):
|
||||
"""Return (temporal, height, width) pixel-per-latent downscale factors."""
|
||||
if vae is not None:
|
||||
formula = getattr(vae, "downscale_index_formula", None)
|
||||
if formula is not None:
|
||||
st, sw, sh = formula
|
||||
return int(st), int(sh), int(sw)
|
||||
try:
|
||||
sf = model.model.latent_format
|
||||
sh = int(
|
||||
getattr(sf, "spatial_downscale", None)
|
||||
or getattr(sf, "spacial_downscale_ratio", None)
|
||||
or 32
|
||||
)
|
||||
sw = sh
|
||||
st = int(
|
||||
getattr(sf, "temporal_downscale", None)
|
||||
or getattr(sf, "temporal_downscale_ratio", None)
|
||||
or 8
|
||||
)
|
||||
return st, sh, sw
|
||||
except Exception:
|
||||
return 8, 32, 32
|
||||
|
||||
|
||||
def gaussian_weight(tile_w, tile_h, var, device):
|
||||
"""Separable Gaussian tile weight. Pure Gaussian, deliberately no floor."""
|
||||
|
||||
def axis(n):
|
||||
x = torch.arange(n, dtype=torch.float32, device=device)
|
||||
mid = (n - 1) / 2
|
||||
return torch.exp(-((x - mid) ** 2) / (n * n * 2 * var))
|
||||
|
||||
return axis(tile_h)[:, None] * axis(tile_w)[None, :]
|
||||
|
||||
|
||||
def _crop_guide_metadata(
|
||||
conds, y0, x0, th, tw, canvas_h, canvas_w, sh, sw, entry_idx=None
|
||||
):
|
||||
"""Crop keyframe_idxs + guide_attention_entries to a spatial tile.
|
||||
|
||||
y0/x0/th/tw are in latent units; sh/sw are pixels per latent step. Token
|
||||
order inside each guide's block is the patchifier's row-major (t, h, w),
|
||||
which is what makes the reshape below valid. Spatial positions are
|
||||
normalised to the tile so each tile is the same problem the LoRA was
|
||||
trained on.
|
||||
|
||||
entry_idx selects a SINGLE guide entry: temporal windowing pairs window w
|
||||
with guide block w, each encoded FRESH from that window's pixel span, so
|
||||
its coords already describe a clip starting at 0 and are never rewritten
|
||||
here. None keeps every entry (whole-clip guide).
|
||||
Slicing a whole-clip guide's latents into windows does NOT work and no
|
||||
coordinate rewrite can save it: a mid-clip causal latent (each frame
|
||||
encodes 8 pixels given prior context) is not a fresh-clip latent (frame 0
|
||||
encodes exactly 1 pixel frame).
|
||||
"""
|
||||
out = []
|
||||
py0, px0 = y0 * sh, x0 * sw
|
||||
for emb, d in conds:
|
||||
d = dict(d)
|
||||
entries = d.get("guide_attention_entries")
|
||||
kf = d.get("keyframe_idxs")
|
||||
if kf is not None and entries:
|
||||
pieces = []
|
||||
new_entries = []
|
||||
tok0 = 0
|
||||
for j, e in enumerate(entries):
|
||||
f, gh, gw = (int(v) for v in e["latent_shape"])
|
||||
n_tok = f * gh * gw
|
||||
blk = kf[:, :, tok0 : tok0 + n_tok, :]
|
||||
tok0 += n_tok
|
||||
if entry_idx is not None and j != entry_idx:
|
||||
continue
|
||||
if (gh, gw) != (canvas_h, canvas_w):
|
||||
raise ValueError(
|
||||
f"{_NODE}: guide latent grid {gh}x{gw} does not match "
|
||||
f"the canvas {canvas_h}x{canvas_w}. Downscaled "
|
||||
f"(factor > 1) IC-LoRA guides are not supported; "
|
||||
f"re-attach the guide at factor 1."
|
||||
)
|
||||
b = blk.shape[0]
|
||||
blk = blk.reshape(b, 3, f, gh, gw, 2)
|
||||
blk = blk[:, :, :, y0 : y0 + th, x0 : x0 + tw, :].clone()
|
||||
blk[:, 1, ...] -= py0
|
||||
blk[:, 2, ...] -= px0
|
||||
pieces.append(blk.reshape(b, 3, f * th * tw, 2))
|
||||
ne = dict(e)
|
||||
ne["pre_filter_count"] = f * th * tw
|
||||
ne["latent_shape"] = [f, th, tw]
|
||||
pm = e.get("pixel_mask")
|
||||
if pm is not None:
|
||||
ne["pixel_mask"] = pm[
|
||||
:, :, :, py0 : py0 + th * sh, px0 : px0 + tw * sw
|
||||
]
|
||||
new_entries.append(ne)
|
||||
d["keyframe_idxs"] = torch.cat(pieces, dim=2)
|
||||
d["guide_attention_entries"] = new_entries
|
||||
out.append([emb, d])
|
||||
return out
|
||||
|
||||
|
||||
def _crop_spatial(t, y0, x0, th, tw):
|
||||
"""Crop the last two dims when they are real spatial dims (>1)."""
|
||||
if t is None:
|
||||
return None
|
||||
if t.shape[-2] > 1 or t.shape[-1] > 1:
|
||||
return t[..., y0 : y0 + th, x0 : x0 + tw]
|
||||
return t
|
||||
|
||||
|
||||
class _TileGuider:
|
||||
"""CFGGuider look-alike so KSamplerX0Inpaint / k-diffusion can call us.
|
||||
|
||||
`inner_model` is the object `prepare_sampling` returned (has
|
||||
`scale_latent_inpaint`). `__call__` is the transformer+CFG pass.
|
||||
"""
|
||||
|
||||
def __init__(self, inner_model, conds, cfg):
|
||||
self.inner_model = inner_model
|
||||
self.conds = conds
|
||||
self.cfg = cfg
|
||||
|
||||
def __call__(self, x, timestep, model_options={}, seed=None):
|
||||
return comfy.samplers.sampling_function(
|
||||
self.inner_model,
|
||||
x,
|
||||
timestep,
|
||||
self.conds.get("negative"),
|
||||
self.conds.get("positive"),
|
||||
self.cfg,
|
||||
model_options=model_options,
|
||||
seed=seed,
|
||||
)
|
||||
|
||||
|
||||
def _one_sampler_step(
|
||||
sampler,
|
||||
inner_model,
|
||||
conds,
|
||||
cfg,
|
||||
x,
|
||||
sig,
|
||||
sig_next,
|
||||
noise,
|
||||
latent_image,
|
||||
denoise_mask,
|
||||
model_options,
|
||||
seed,
|
||||
full_sigmas,
|
||||
):
|
||||
"""Run one k-diffusion step (sig -> sig_next) on a tile.
|
||||
|
||||
Uses `sampler.sampler_function` (what KSamplerSelect builds), not
|
||||
`sampler.sample()`, which would re-scale noise and walk the whole
|
||||
schedule — fusion has to happen between steps.
|
||||
"""
|
||||
fn = getattr(sampler, "sampler_function", None)
|
||||
if fn is None:
|
||||
raise ValueError(
|
||||
f"{_NODE} needs a SAMPLER from KSamplerSelect (or another node "
|
||||
"that exposes sampler_function). Adaptive samplers that run the "
|
||||
"whole trajectory in one go cannot fuse tiles between steps."
|
||||
)
|
||||
wrap = _TileGuider(inner_model, conds, cfg)
|
||||
model_k = comfy.samplers.KSamplerX0Inpaint(wrap, full_sigmas)
|
||||
model_k.latent_image = latent_image
|
||||
inpaint_opts = getattr(sampler, "inpaint_options", None) or {}
|
||||
if inpaint_opts.get("random", False):
|
||||
generator = torch.manual_seed(seed + 1)
|
||||
model_k.noise = torch.randn(noise.shape, generator=generator, device="cpu").to(
|
||||
device=noise.device, dtype=noise.dtype
|
||||
)
|
||||
else:
|
||||
model_k.noise = noise
|
||||
extra_args = {
|
||||
"model_options": model_options,
|
||||
"seed": seed,
|
||||
"denoise_mask": denoise_mask,
|
||||
}
|
||||
sigmas_pair = torch.tensor([sig, sig_next], device=x.device, dtype=torch.float32)
|
||||
extra_opts = dict(getattr(sampler, "extra_options", None) or {})
|
||||
return fn(
|
||||
model_k,
|
||||
x,
|
||||
sigmas_pair,
|
||||
extra_args=extra_args,
|
||||
callback=None,
|
||||
disable=True,
|
||||
**extra_opts,
|
||||
)
|
||||
|
||||
|
||||
@comfy_node(
|
||||
name="LTXVTiledFusionSampler",
|
||||
description="Tiled Fusion Sampler",
|
||||
)
|
||||
class LTXVTiledFusionSampler:
|
||||
CATEGORY = "sampling"
|
||||
RETURN_TYPES = ("LATENT",)
|
||||
RETURN_NAMES = ("output",)
|
||||
FUNCTION = "sample"
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"model": (
|
||||
"MODEL",
|
||||
{"tooltip": "The diffusion model to use."},
|
||||
),
|
||||
"positive": (
|
||||
"CONDITIONING",
|
||||
{
|
||||
"tooltip": "Positive conditioning from the IC-LoRA guide node. "
|
||||
"Must carry keyframe_idxs / guide_attention_entries so each "
|
||||
"tile can be cropped to the LoRA's trained window."
|
||||
},
|
||||
),
|
||||
"negative": (
|
||||
"CONDITIONING",
|
||||
{"tooltip": "Negative conditioning from the IC-LoRA guide node."},
|
||||
),
|
||||
"latents": (
|
||||
"LATENT",
|
||||
{
|
||||
"tooltip": "The latent RETURNED BY the IC-LoRA guide node "
|
||||
"(it carries appended guide frames plus the noise mask). "
|
||||
"A bare empty video latent will not work."
|
||||
},
|
||||
),
|
||||
"sigmas": (
|
||||
"SIGMAS",
|
||||
{
|
||||
"tooltip": "Denoise schedule. Must be descending and end at 0. "
|
||||
"Fusion walks this list one step at a time."
|
||||
},
|
||||
),
|
||||
"sampler": (
|
||||
"SAMPLER",
|
||||
{
|
||||
"tooltip": "From KSamplerSelect (same type as SamplerCustomAdvanced). "
|
||||
"Discrete step samplers only: euler, euler_ancestral, heun, dpm_2, … "
|
||||
"History methods (DPM++ 2M, LMS) lose their multi-step memory because "
|
||||
"fusion owns the loop. Adaptive samplers cannot fuse. Ancestral methods "
|
||||
"add independent noise per tile and can seam."
|
||||
},
|
||||
),
|
||||
"seed": (
|
||||
"INT",
|
||||
{
|
||||
"default": 42,
|
||||
"min": 0,
|
||||
"max": 0xFFFFFFFFFFFFFFFF,
|
||||
"control_after_generate": True,
|
||||
"tooltip": "Seed for the single shared noise field across all tiles.",
|
||||
},
|
||||
),
|
||||
"cfg": (
|
||||
"FLOAT",
|
||||
{
|
||||
"default": 1.0,
|
||||
"min": 0.0,
|
||||
"max": 20.0,
|
||||
"step": 0.1,
|
||||
"tooltip": "Classifier-free guidance scale. Distilled models typically use 1.0.",
|
||||
},
|
||||
),
|
||||
"tile_width": (
|
||||
"INT",
|
||||
{
|
||||
"default": 1024,
|
||||
"min": 256,
|
||||
"max": 4096,
|
||||
"step": 32,
|
||||
"tooltip": "Tile width in pixels, multiple of 32. Set this to the "
|
||||
"IC-LoRA's trained spatial window. Tiles smaller than the trained "
|
||||
"window make every tile see content at the wrong scale.",
|
||||
},
|
||||
),
|
||||
"tile_height": (
|
||||
"INT",
|
||||
{
|
||||
"default": 576,
|
||||
"min": 256,
|
||||
"max": 4096,
|
||||
"step": 32,
|
||||
"tooltip": "Tile height in pixels, multiple of 32. Match tile_width to "
|
||||
"the IC-LoRA's trained spatial window.",
|
||||
},
|
||||
),
|
||||
"overlap_frac": (
|
||||
"FLOAT",
|
||||
{
|
||||
"default": 0.5,
|
||||
"min": 0.25,
|
||||
"max": 0.75,
|
||||
"step": 0.05,
|
||||
"tooltip": "Fraction of each tile overlapping its neighbor. Stay at "
|
||||
"0.5 or above: below that a periodic grid appears on structured content.",
|
||||
},
|
||||
),
|
||||
"blend_var": (
|
||||
"FLOAT",
|
||||
{
|
||||
"default": 0.05,
|
||||
"min": 0.01,
|
||||
"max": 0.2,
|
||||
"step": 0.01,
|
||||
"tooltip": "Variance of the Gaussian blend weights in the overlaps. "
|
||||
"Leave at 0.05 unless you are measuring seams.",
|
||||
},
|
||||
),
|
||||
"grid_cycle": (
|
||||
"INT",
|
||||
{
|
||||
"default": 1,
|
||||
"min": 1,
|
||||
"max": 4,
|
||||
"tooltip": "Cycle shifted tile grids across steps. 1 = fixed grid "
|
||||
"(cheapest), 4 = thorough. Spatial only: cycling the temporal grid "
|
||||
"slides content in time.",
|
||||
},
|
||||
),
|
||||
"tile_frames": (
|
||||
"INT",
|
||||
{
|
||||
"default": 0,
|
||||
"min": 0,
|
||||
"max": 1000,
|
||||
"step": 1,
|
||||
"tooltip": "Temporal window in pixel frames (97 for LTX). 0 = one "
|
||||
"extent. Other values are snapped to 8n+1. >0 needs "
|
||||
"LTXAddVideoICLoRAGuide with use_streaming on and the SAME "
|
||||
"tile_frames: that re-encodes a fresh guide per window and "
|
||||
"writes the plan this sampler reads. A whole-clip guide "
|
||||
"cannot be sliced, and windowing is disabled if the plan "
|
||||
"is missing.",
|
||||
},
|
||||
),
|
||||
},
|
||||
"optional": {
|
||||
"vae": (
|
||||
"VAE",
|
||||
{
|
||||
"tooltip": "Used to read the model's temporal and spatial downscale "
|
||||
"factors. If omitted, they are taken from the diffusion model "
|
||||
"(LTX defaults: 8 frames and 32 pixels per latent cell)."
|
||||
},
|
||||
),
|
||||
"canvas_device": (
|
||||
["auto", "gpu", "cpu"],
|
||||
{
|
||||
"default": "auto",
|
||||
"tooltip": "Where to keep the full-canvas tensors and cached tile "
|
||||
"latents/noise/masks. auto = GPU if there is enough free "
|
||||
"memory, otherwise CPU. Processed conds stay on the GPU.",
|
||||
},
|
||||
),
|
||||
},
|
||||
}
|
||||
|
||||
@torch.inference_mode()
|
||||
def sample(
|
||||
self,
|
||||
model,
|
||||
positive,
|
||||
negative,
|
||||
latents,
|
||||
sigmas,
|
||||
sampler,
|
||||
seed,
|
||||
cfg,
|
||||
tile_width,
|
||||
tile_height,
|
||||
overlap_frac,
|
||||
blend_var,
|
||||
grid_cycle,
|
||||
tile_frames=0,
|
||||
vae=None,
|
||||
canvas_device="auto",
|
||||
):
|
||||
samples = latents["samples"]
|
||||
noise_mask = latents.get("noise_mask")
|
||||
if samples.ndim != 5:
|
||||
raise ValueError(f"{_NODE} needs a 5D video latent.")
|
||||
b, C, F_total, H, W = samples.shape
|
||||
if b != 1:
|
||||
raise ValueError(f"{_NODE} supports batch size 1.")
|
||||
if noise_mask is None:
|
||||
raise ValueError(
|
||||
f"{_NODE} requires a noise_mask; the latent must be the one "
|
||||
"returned by the IC-LoRA guide node (it carries noise_mask)."
|
||||
)
|
||||
|
||||
device = comfy.model_management.get_torch_device()
|
||||
st, sh, sw = _scale_factors(model, vae)
|
||||
|
||||
tl_w, tl_h = tile_width // sw, tile_height // sh
|
||||
tl_w, tl_h = min(tl_w, W), min(tl_h, H)
|
||||
|
||||
min_ov_x = max(2, int(tl_w * overlap_frac))
|
||||
min_ov_y = max(2, int(tl_h * overlap_frac))
|
||||
xs = plan_grid(W, tl_w, min_ov_x)
|
||||
ys = plan_grid(H, tl_h, min_ov_y)
|
||||
|
||||
ov_x = tl_w - (xs[1] - xs[0]) if len(xs) > 1 else tl_w
|
||||
ov_y = tl_h - (ys[1] - ys[0]) if len(ys) > 1 else tl_h
|
||||
# Interior shift must leave at least min_ov overlap on cycled grids.
|
||||
max_shift = max(0, min(ov_x - min_ov_x, ov_y - min_ov_y))
|
||||
off = min(4, max_shift)
|
||||
if grid_cycle > 1 and off < 1:
|
||||
logger.warning(
|
||||
"grid_cycle=%d ignored: no interior shift can keep overlap_frac=%.2f",
|
||||
grid_cycle,
|
||||
overlap_frac,
|
||||
)
|
||||
offsets = [(0, 0)]
|
||||
else:
|
||||
offsets = [(0, 0), (off, off), (off, 0), (0, off)][: max(1, grid_cycle)]
|
||||
grids = [
|
||||
[(x, y) for y in shifted(ys, H, tl_h, oy) for x in shifted(xs, W, tl_w, ox)]
|
||||
for ox, oy in offsets
|
||||
]
|
||||
all_origins = sorted({o for g in grids for o in g})
|
||||
logger.info(
|
||||
"tiled fusion: canvas %dx%d lat, tile %dx%d, grid %dx%d, "
|
||||
"cycle=%d, %d unique origins",
|
||||
W,
|
||||
H,
|
||||
tl_w,
|
||||
tl_h,
|
||||
len(xs),
|
||||
len(ys),
|
||||
len(grids),
|
||||
len(all_origins),
|
||||
)
|
||||
|
||||
if len(all_origins) == 1:
|
||||
logger.info("canvas fits one tile; falling back to plain sampling")
|
||||
|
||||
n_content = F_total
|
||||
if noise_mask is not None:
|
||||
fm = (
|
||||
noise_mask.to(dtype=torch.float32)
|
||||
.reshape(1, 1, F_total, -1)
|
||||
.amax(dim=-1)
|
||||
)
|
||||
for f in range(F_total - 1, -1, -1):
|
||||
if float(fm[0, 0, f]) < 0.999:
|
||||
n_content = f
|
||||
else:
|
||||
break
|
||||
n_guide = F_total - n_content
|
||||
|
||||
plan = read_temporal_plan(positive)
|
||||
|
||||
tile_frames = int(tile_frames)
|
||||
if tile_frames > 0:
|
||||
snapped = snap_frames(tile_frames)
|
||||
if snapped != tile_frames:
|
||||
logger.info("tile_frames %d snapped to %d (8n+1)", tile_frames, snapped)
|
||||
tile_frames = snapped
|
||||
|
||||
tf_lat = 0
|
||||
if tile_frames > 0:
|
||||
tf_lat = max(1, (tile_frames - 1) // st + 1)
|
||||
use_temporal = 0 < tf_lat < n_content
|
||||
|
||||
if plan is not None and int(plan["tile_frames"]) != int(tile_frames):
|
||||
raise ValueError(
|
||||
f"{_NODE}: the guide was encoded for tile_frames="
|
||||
f"{plan['tile_frames']} but the sampler got tile_frames="
|
||||
f"{tile_frames}; set both to the same value."
|
||||
)
|
||||
|
||||
per_window_guides = False
|
||||
if use_temporal:
|
||||
if plan is not None:
|
||||
ts = [int(v) for v in plan["ts"]]
|
||||
if int(plan["n_content"]) != n_content or n_guide != len(ts) * tf_lat:
|
||||
raise ValueError(
|
||||
f"{_NODE}: temporal plan does not match the latent "
|
||||
f"(plan: n_content={plan['n_content']}, "
|
||||
f"{len(ts)} windows of {tf_lat}; latent: "
|
||||
f"content={n_content}, guide={n_guide})."
|
||||
)
|
||||
per_window_guides = True
|
||||
elif n_guide == 0:
|
||||
ts = temporal_plan(n_content, tf_lat)
|
||||
else:
|
||||
logger.warning(
|
||||
"tile_frames=%d with a whole-clip guide: temporal windows "
|
||||
"need a fresh IC-LoRA encode per window. Turn on "
|
||||
"use_streaming on LTXAddVideoICLoRAGuide with the same "
|
||||
"tile_frames; windowing DISABLED for this run.",
|
||||
tile_frames,
|
||||
)
|
||||
use_temporal = False
|
||||
|
||||
if use_temporal:
|
||||
logger.info(
|
||||
"temporal windows: %d x %d latent frames (%d content, " "%d guide%s)",
|
||||
len(ts),
|
||||
tf_lat,
|
||||
n_content,
|
||||
n_guide,
|
||||
", per-window guides" if per_window_guides else "",
|
||||
)
|
||||
window_len = min(tf_lat, n_content)
|
||||
if per_window_guides:
|
||||
window_len += tf_lat
|
||||
else:
|
||||
ts = [0]
|
||||
tf_lat = F_total
|
||||
window_len = F_total
|
||||
if tile_frames > 0:
|
||||
logger.info(
|
||||
"temporal windowing not active (clip fits one window, "
|
||||
"or disabled above)"
|
||||
)
|
||||
|
||||
if canvas_device == "cpu":
|
||||
cdev = torch.device("cpu")
|
||||
elif canvas_device == "gpu":
|
||||
cdev = device
|
||||
else:
|
||||
free = comfy.model_management.get_free_memory(device)
|
||||
need = samples.numel() * 4 * 4
|
||||
cdev = device if need < free * 0.25 else torch.device("cpu")
|
||||
logger.info("canvas tensors on %s", cdev)
|
||||
|
||||
g = torch.Generator(device="cpu").manual_seed(seed)
|
||||
noise = torch.randn(samples.shape, generator=g, dtype=torch.float32).to(cdev)
|
||||
|
||||
tile_shape = [1, C, window_len, tl_h, tl_w]
|
||||
conds0 = {
|
||||
"positive": comfy.sampler_helpers.convert_cond(copy.deepcopy(positive)),
|
||||
"negative": comfy.sampler_helpers.convert_cond(copy.deepcopy(negative)),
|
||||
}
|
||||
inner_model, conds0, loaded = comfy.sampler_helpers.prepare_sampling(
|
||||
model, tile_shape, conds0, model.model_options
|
||||
)
|
||||
try:
|
||||
model_sampling = (
|
||||
getattr(inner_model, "model_sampling", None)
|
||||
or model.model.model_sampling
|
||||
)
|
||||
|
||||
model.pre_run()
|
||||
|
||||
latent_in = (
|
||||
model.model.process_latent_in(samples).to(torch.float32).to(cdev)
|
||||
)
|
||||
sigmas = sigmas.to(torch.float32)
|
||||
s0 = sigmas[0].item()
|
||||
|
||||
canvas = model_sampling.noise_scaling(
|
||||
torch.tensor(s0, device=cdev), noise, latent_in, False
|
||||
)
|
||||
|
||||
if noise_mask is not None:
|
||||
mask_full = noise_mask.to(torch.float32).to(cdev)
|
||||
else:
|
||||
mask_full = None
|
||||
|
||||
def _window_slice(w, t0, tensor):
|
||||
"""Take content[t0:t0+tf] plus window w's own guide block."""
|
||||
if not use_temporal:
|
||||
return tensor
|
||||
parts = [tensor[:, :, t0 : t0 + tf_lat]]
|
||||
if per_window_guides:
|
||||
g0 = n_content + w * tf_lat
|
||||
parts.append(tensor[:, :, g0 : g0 + tf_lat])
|
||||
return torch.cat(parts, dim=2) if len(parts) > 1 else parts[0]
|
||||
|
||||
tile_data = {}
|
||||
for w, t0_ in enumerate(ts):
|
||||
eidx = w if (use_temporal and per_window_guides) else None
|
||||
for x0_, y0_ in all_origins:
|
||||
pos_t = _crop_guide_metadata(
|
||||
copy.deepcopy(positive),
|
||||
y0_,
|
||||
x0_,
|
||||
tl_h,
|
||||
tl_w,
|
||||
H,
|
||||
W,
|
||||
sh,
|
||||
sw,
|
||||
entry_idx=eidx,
|
||||
)
|
||||
neg_t = _crop_guide_metadata(
|
||||
copy.deepcopy(negative),
|
||||
y0_,
|
||||
x0_,
|
||||
tl_h,
|
||||
tl_w,
|
||||
H,
|
||||
W,
|
||||
sh,
|
||||
sw,
|
||||
entry_idx=eidx,
|
||||
)
|
||||
lat_t = _window_slice(
|
||||
w,
|
||||
t0_,
|
||||
latent_in[:, :, :, y0_ : y0_ + tl_h, x0_ : x0_ + tl_w],
|
||||
)
|
||||
noise_t = _window_slice(
|
||||
w,
|
||||
t0_,
|
||||
noise[:, :, :, y0_ : y0_ + tl_h, x0_ : x0_ + tl_w],
|
||||
)
|
||||
mask_t = None
|
||||
if mask_full is not None:
|
||||
mask_t = _crop_spatial(mask_full, y0_, x0_, tl_h, tl_w)
|
||||
if mask_t.shape[2] == F_total:
|
||||
mask_t = _window_slice(w, t0_, mask_t)
|
||||
if mask_t.shape[2] == lat_t.shape[2]:
|
||||
mask_t = mask_t.expand(1, 1, lat_t.shape[2], tl_h, tl_w)
|
||||
# process_conds needs GPU tensors; lat/noise/mask stay on cdev
|
||||
# until each step so a CPU canvas does not pin every tile.
|
||||
conds_t = comfy.samplers.process_conds(
|
||||
inner_model,
|
||||
noise_t.to(device),
|
||||
{
|
||||
"positive": comfy.sampler_helpers.convert_cond(pos_t),
|
||||
"negative": comfy.sampler_helpers.convert_cond(neg_t),
|
||||
},
|
||||
device,
|
||||
latent_image=lat_t.to(device),
|
||||
denoise_mask=None if mask_t is None else mask_t.to(device),
|
||||
seed=seed,
|
||||
latent_shapes=[lat_t.shape],
|
||||
)
|
||||
tile_data[(w, x0_, y0_)] = (conds_t, lat_t, noise_t, mask_t)
|
||||
|
||||
w_tile = gaussian_weight(tl_w, tl_h, blend_var, cdev)
|
||||
|
||||
w_len = tf_lat if use_temporal else n_content
|
||||
if use_temporal and len(ts) > 1:
|
||||
xs_t = torch.arange(w_len, dtype=torch.float32, device=cdev)
|
||||
mid_t = (w_len - 1) / 2
|
||||
w_time = torch.exp(
|
||||
-((xs_t - mid_t) ** 2) / (w_len * w_len * 2 * blend_var)
|
||||
)
|
||||
else:
|
||||
w_time = torch.ones(w_len, dtype=torch.float32, device=cdev)
|
||||
|
||||
def _w_skip(w):
|
||||
return int(
|
||||
_SKIP_FIRST_LAT
|
||||
and use_temporal
|
||||
and w > 0
|
||||
and ts[w] < ts[w - 1] + tf_lat
|
||||
)
|
||||
|
||||
wsums = []
|
||||
for grid in grids:
|
||||
ws = torch.zeros((1, 1, n_content, H, W), device=cdev)
|
||||
for w, t0_ in enumerate(ts):
|
||||
sk = _w_skip(w)
|
||||
tw_ = w_time[sk : min(tf_lat, n_content - t0_)].reshape(
|
||||
1, 1, -1, 1, 1
|
||||
)
|
||||
for x0_, y0_ in grid:
|
||||
ws[
|
||||
:,
|
||||
:,
|
||||
t0_ + sk : t0_ + sk + tw_.shape[2],
|
||||
y0_ : y0_ + tl_h,
|
||||
x0_ : x0_ + tl_w,
|
||||
] += (
|
||||
w_tile * tw_
|
||||
)
|
||||
wsums.append(ws.clamp_min(1e-12))
|
||||
|
||||
model_options = model.model_options
|
||||
|
||||
pbar = comfy.utils.ProgressBar(len(sigmas) - 1)
|
||||
for step in range(len(sigmas) - 1):
|
||||
sig = sigmas[step].item()
|
||||
sig_next = sigmas[step + 1].item()
|
||||
grid = grids[step % len(grids)]
|
||||
acc = torch.zeros(
|
||||
(1, C, n_content, H, W), device=cdev, dtype=canvas.dtype
|
||||
)
|
||||
|
||||
for w, t0_ in enumerate(ts):
|
||||
for x0_, y0_ in grid:
|
||||
conds_t, lat_t, noise_t, mask_t = tile_data[(w, x0_, y0_)]
|
||||
x = _window_slice(
|
||||
w,
|
||||
t0_,
|
||||
canvas[:, :, :, y0_ : y0_ + tl_h, x0_ : x0_ + tl_w],
|
||||
).to(device, torch.float32)
|
||||
|
||||
x_next = _one_sampler_step(
|
||||
sampler,
|
||||
inner_model,
|
||||
conds_t,
|
||||
cfg,
|
||||
x,
|
||||
sig,
|
||||
sig_next,
|
||||
noise_t.to(device, torch.float32),
|
||||
lat_t.to(device, torch.float32),
|
||||
None if mask_t is None else mask_t.to(device),
|
||||
model_options,
|
||||
seed,
|
||||
sigmas,
|
||||
)
|
||||
|
||||
nc = min(tf_lat, n_content - t0_) if use_temporal else n_content
|
||||
sk = _w_skip(w)
|
||||
tw_ = w_time[sk:nc].reshape(1, 1, -1, 1, 1)
|
||||
acc[
|
||||
:,
|
||||
:,
|
||||
t0_ + sk : t0_ + nc,
|
||||
y0_ : y0_ + tl_h,
|
||||
x0_ : x0_ + tl_w,
|
||||
] += (
|
||||
x_next[:, :, sk:nc].to(cdev) * w_tile * tw_
|
||||
)
|
||||
|
||||
if _DEBUG_STATS:
|
||||
logger.info(
|
||||
"dbg step %d w%d (t0=%d xy=%d,%d): x %.4f "
|
||||
"x_next %.4f",
|
||||
step,
|
||||
w,
|
||||
t0_,
|
||||
x0_,
|
||||
y0_,
|
||||
float(x[:, :, :nc].mean()),
|
||||
float(x_next[:, :, :nc].mean()),
|
||||
)
|
||||
|
||||
blended = acc / wsums[step % len(grids)]
|
||||
if _DEBUG_STATS:
|
||||
logger.info(
|
||||
"dbg step %d blended: mean %.4f std %.4f",
|
||||
step,
|
||||
float(blended.mean()),
|
||||
float(blended.std()),
|
||||
)
|
||||
if n_guide:
|
||||
canvas = torch.cat([blended, canvas[:, :, n_content:]], dim=2)
|
||||
else:
|
||||
canvas = blended
|
||||
pbar.update(1)
|
||||
|
||||
out = canvas
|
||||
if mask_full is not None:
|
||||
frame_mask = mask_full.reshape(1, 1, F_total, -1).amax(dim=-1)
|
||||
keep = F_total
|
||||
for f in range(F_total - 1, -1, -1):
|
||||
if frame_mask[0, 0, f] < 0.999:
|
||||
keep = f
|
||||
else:
|
||||
break
|
||||
out = canvas[:, :, :keep]
|
||||
if keep < F_total:
|
||||
logger.info("dropped %d appended guide frame(s)", F_total - keep)
|
||||
|
||||
out = model.model.process_latent_out(out.to(torch.float32).cpu())
|
||||
finally:
|
||||
model.cleanup()
|
||||
comfy.sampler_helpers.cleanup_models(conds0, loaded)
|
||||
|
||||
return ({"samples": out},)
|
||||
@@ -0,0 +1,445 @@
|
||||
"""Canvas sizes and temporal window for tiled-fusion graphs.
|
||||
|
||||
Stage 1 aims at full HD in the source aspect without spatial tiling. Stage 2
|
||||
is 2x that (~4K). 8K is another 2x. HD is the 720p step of the same ladder;
|
||||
qHD is the 960 long-edge step.
|
||||
|
||||
Dimensions stay multiples of 32. ``tile_frames`` stays 8n+1, targeting 97,
|
||||
stretching up to 121 only when that lets a slightly longer clip skip chunking.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import math
|
||||
from typing import NamedTuple, Optional
|
||||
|
||||
from .nodes_registry import comfy_node
|
||||
from .tiled_fusion_plan import plan_grid, snap_frames, temporal_plan
|
||||
|
||||
SPATIAL_ALIGN = 32
|
||||
FRAME_STEP = 8
|
||||
SLACK_PX = 200
|
||||
QHD_LONG_EDGE = 960
|
||||
HD_LONG_EDGE = 1280
|
||||
FULLHD_LONG_EDGE = 1920
|
||||
TILE_FRAMES_TARGET = 97
|
||||
TILE_FRAMES_MAX = 121
|
||||
OVERLAP_FRAC = 0.5
|
||||
|
||||
TILE_CHOICES = ("qHD", "HD", "FullHD")
|
||||
OUTPUT_CHOICES = ("FullHD", "4K", "8K")
|
||||
|
||||
_QHD_KEYS = frozenset({"qhd", "quarterhd", "quarter", "960"})
|
||||
_HD_KEYS = frozenset({"hd", "720p", "720"})
|
||||
_FULLHD_KEYS = frozenset({"fullhd", "fhd", "1080p", "1080"})
|
||||
_4K_KEYS = frozenset({"4k", "uhd", "2160p", "2160"})
|
||||
_8K_KEYS = frozenset({"8k", "eightk", "4320p", "4320"})
|
||||
|
||||
|
||||
class TilingSizes(NamedTuple):
|
||||
qhd_width: int
|
||||
qhd_height: int
|
||||
hd_width: int
|
||||
hd_height: int
|
||||
fullhd_width: int
|
||||
fullhd_height: int
|
||||
overlap_frac: float
|
||||
tile_frames: int
|
||||
width_4k: int
|
||||
height_4k: int
|
||||
width_8k: int
|
||||
height_8k: int
|
||||
|
||||
|
||||
class TilingPlan(NamedTuple):
|
||||
initial_canvas_width: int
|
||||
initial_canvas_height: int
|
||||
tile_width: int
|
||||
tile_height: int
|
||||
output_width: int
|
||||
output_height: int
|
||||
overlap_frac: float
|
||||
tile_frames: int
|
||||
spatial_tiles_x: int
|
||||
spatial_tiles_y: int
|
||||
output_tiles_x: int
|
||||
output_tiles_y: int
|
||||
temporal_chunks: int
|
||||
canvas_label: str
|
||||
tile_label: str
|
||||
output_label: str
|
||||
tile_clamped: bool
|
||||
is_8k: bool
|
||||
summary: str
|
||||
|
||||
|
||||
def _norm_label(value: str) -> str:
|
||||
return str(value or "").lower().replace(" ", "").replace("_", "").replace("-", "")
|
||||
|
||||
|
||||
def snap_spatial(value: float, align: int = SPATIAL_ALIGN) -> int:
|
||||
"""Nearest multiple of ``align``, ties round up, never below ``align``."""
|
||||
value = max(float(align), float(value))
|
||||
return max(align, int(math.floor(value / align + 0.5) * align))
|
||||
|
||||
|
||||
def tile_frames_for(frame_count: int) -> int:
|
||||
"""Temporal window: 97 or the clip itself, stretching to 121 to skip chunks."""
|
||||
n = max(1, int(frame_count))
|
||||
if n <= TILE_FRAMES_TARGET:
|
||||
tf = snap_frames(n)
|
||||
return min(TILE_FRAMES_TARGET, tf)
|
||||
if n <= TILE_FRAMES_MAX:
|
||||
covering = ((n - 1 + FRAME_STEP - 1) // FRAME_STEP) * FRAME_STEP + 1
|
||||
return min(TILE_FRAMES_MAX, covering)
|
||||
return TILE_FRAMES_TARGET
|
||||
|
||||
|
||||
def _fit_long_edge(width: float, height: float, long_edge: int) -> tuple[float, float]:
|
||||
long = max(width, height)
|
||||
scale = long_edge / long
|
||||
return width * scale, height * scale
|
||||
|
||||
|
||||
def _keep_source_if_close(
|
||||
target_w: float, target_h: float, src_w: int, src_h: int, slack: int = SLACK_PX
|
||||
) -> tuple[float, float]:
|
||||
if abs(src_w - target_w) <= slack and abs(src_h - target_h) <= slack:
|
||||
return float(src_w), float(src_h)
|
||||
return target_w, target_h
|
||||
|
||||
|
||||
def _aligned_size(width: float, height: float) -> tuple[int, int]:
|
||||
return snap_spatial(width), snap_spatial(height)
|
||||
|
||||
|
||||
def _ladder_step(src_w: int, src_h: int, long_edge: int) -> tuple[int, int]:
|
||||
tw, th = _fit_long_edge(src_w, src_h, long_edge)
|
||||
tw, th = _keep_source_if_close(tw, th, src_w, src_h)
|
||||
return _aligned_size(tw, th)
|
||||
|
||||
|
||||
def compute_tiling_sizes(width: int, height: int, frame_count: int) -> TilingSizes:
|
||||
"""Return the qHD / HD / full-HD / 4K / 8K ladder plus overlap and tile_frames."""
|
||||
src_w = max(1, int(width))
|
||||
src_h = max(1, int(height))
|
||||
|
||||
qhd_w, qhd_h = _ladder_step(src_w, src_h, QHD_LONG_EDGE)
|
||||
hd_w, hd_h = _ladder_step(src_w, src_h, HD_LONG_EDGE)
|
||||
fhd_w, fhd_h = _ladder_step(src_w, src_h, FULLHD_LONG_EDGE)
|
||||
|
||||
k4_w, k4_h = _keep_source_if_close(fhd_w * 2, fhd_h * 2, src_w, src_h)
|
||||
k4_w, k4_h = _aligned_size(k4_w, k4_h)
|
||||
|
||||
k8_w, k8_h = _keep_source_if_close(k4_w * 2, k4_h * 2, src_w, src_h)
|
||||
k8_w, k8_h = _aligned_size(k8_w, k8_h)
|
||||
|
||||
return TilingSizes(
|
||||
qhd_width=qhd_w,
|
||||
qhd_height=qhd_h,
|
||||
hd_width=hd_w,
|
||||
hd_height=hd_h,
|
||||
fullhd_width=fhd_w,
|
||||
fullhd_height=fhd_h,
|
||||
overlap_frac=OVERLAP_FRAC,
|
||||
tile_frames=tile_frames_for(frame_count),
|
||||
width_4k=k4_w,
|
||||
height_4k=k4_h,
|
||||
width_8k=k8_w,
|
||||
height_8k=k8_h,
|
||||
)
|
||||
|
||||
|
||||
def _pick_tile_or_canvas(sizes: TilingSizes, choice: str) -> tuple[int, int, str]:
|
||||
key = _norm_label(choice)
|
||||
if key in _QHD_KEYS:
|
||||
return sizes.qhd_width, sizes.qhd_height, "qHD"
|
||||
if key in _HD_KEYS:
|
||||
return sizes.hd_width, sizes.hd_height, "HD"
|
||||
if key in _FULLHD_KEYS:
|
||||
return sizes.fullhd_width, sizes.fullhd_height, "FullHD"
|
||||
raise ValueError(
|
||||
f"Unknown tile/canvas size {choice!r}; expected qHD, HD, or FullHD."
|
||||
)
|
||||
|
||||
|
||||
def _pick_output(sizes: TilingSizes, choice: str) -> tuple[int, int, str]:
|
||||
key = _norm_label(choice)
|
||||
if key in _8K_KEYS:
|
||||
return sizes.width_8k, sizes.height_8k, "8K"
|
||||
if key in _4K_KEYS:
|
||||
return sizes.width_4k, sizes.height_4k, "4K"
|
||||
if key in _FULLHD_KEYS:
|
||||
return sizes.fullhd_width, sizes.fullhd_height, "FullHD"
|
||||
raise ValueError(f"Unknown output size {choice!r}; expected FullHD, 4K, or 8K.")
|
||||
|
||||
|
||||
def _latent_wh(width: int, height: int) -> tuple[int, int]:
|
||||
return max(1, int(width) // SPATIAL_ALIGN), max(1, int(height) // SPATIAL_ALIGN)
|
||||
|
||||
|
||||
def spatial_tile_grid(
|
||||
canvas_w: int,
|
||||
canvas_h: int,
|
||||
tile_w: int,
|
||||
tile_h: int,
|
||||
overlap_frac: float = OVERLAP_FRAC,
|
||||
) -> tuple[int, int]:
|
||||
"""Number of spatial tiles covering the canvas, matching the sampler grid."""
|
||||
cw, ch = _latent_wh(canvas_w, canvas_h)
|
||||
tw, th = _latent_wh(tile_w, tile_h)
|
||||
tw, th = min(tw, cw), min(th, ch)
|
||||
min_ov_x = max(2, int(tw * overlap_frac))
|
||||
min_ov_y = max(2, int(th * overlap_frac))
|
||||
return (
|
||||
len(plan_grid(cw, tw, min_ov_x)),
|
||||
len(plan_grid(ch, th, min_ov_y)),
|
||||
)
|
||||
|
||||
|
||||
def temporal_chunk_count(frame_count: int, tile_frames: int) -> int:
|
||||
"""Number of temporal windows, matching the sampler / guide handshake."""
|
||||
n = max(1, int(frame_count))
|
||||
tf = max(1, int(tile_frames))
|
||||
n_lat = (n - 1) // FRAME_STEP + 1
|
||||
tf_lat = min(n_lat, (tf - 1) // FRAME_STEP + 1)
|
||||
return len(temporal_plan(n_lat, tf_lat))
|
||||
|
||||
|
||||
def format_summary(
|
||||
*,
|
||||
src_w: int,
|
||||
src_h: int,
|
||||
frame_count: int,
|
||||
plan_fields: dict,
|
||||
) -> str:
|
||||
sx, sy = plan_fields["spatial_tiles_x"], plan_fields["spatial_tiles_y"]
|
||||
ox, oy = plan_fields["output_tiles_x"], plan_fields["output_tiles_y"]
|
||||
chunks = plan_fields["temporal_chunks"]
|
||||
tile_note = " (clamped to canvas)" if plan_fields["tile_clamped"] else ""
|
||||
if (sx, sy) == (ox, oy):
|
||||
spatial = f"spatial {sx}×{sy} tiles"
|
||||
else:
|
||||
spatial = f"spatial {sx}×{sy} on canvas, {ox}×{oy} on output"
|
||||
chunk_word = "chunk" if chunks == 1 else "chunks"
|
||||
return "\n".join(
|
||||
(
|
||||
f"source {src_w}×{src_h} × {frame_count}f",
|
||||
f"canvas {plan_fields['initial_canvas_width']}×{plan_fields['initial_canvas_height']}"
|
||||
f" {plan_fields['canvas_label']}",
|
||||
f"tile {plan_fields['tile_width']}×{plan_fields['tile_height']}"
|
||||
f" {plan_fields['tile_label']}{tile_note}",
|
||||
f"output {plan_fields['output_width']}×{plan_fields['output_height']}"
|
||||
f" {plan_fields['output_label']}",
|
||||
f"overlap {plan_fields['overlap_frac']:g}",
|
||||
f"tile_frames {plan_fields['tile_frames']}",
|
||||
spatial,
|
||||
f"temporal {chunks} {chunk_word}",
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
def resolve_tiling_plan(
|
||||
width: int,
|
||||
height: int,
|
||||
frame_count: int,
|
||||
tile_size: str = "FullHD",
|
||||
initial_canvas_size: str = "FullHD",
|
||||
output_size: str = "FullHD",
|
||||
) -> TilingPlan:
|
||||
"""Pick canvas / tile / output from the ladder and count tiles and chunks."""
|
||||
src_w = max(1, int(width))
|
||||
src_h = max(1, int(height))
|
||||
frames = max(1, int(frame_count))
|
||||
sizes = compute_tiling_sizes(src_w, src_h, frames)
|
||||
|
||||
canvas_w, canvas_h, canvas_label = _pick_tile_or_canvas(sizes, initial_canvas_size)
|
||||
tile_w, tile_h, tile_label = _pick_tile_or_canvas(sizes, tile_size)
|
||||
out_w, out_h, out_label = _pick_output(sizes, output_size)
|
||||
|
||||
clamped_w, clamped_h = min(tile_w, canvas_w), min(tile_h, canvas_h)
|
||||
tile_clamped = (clamped_w, clamped_h) != (tile_w, tile_h)
|
||||
tile_w, tile_h = clamped_w, clamped_h
|
||||
|
||||
sx, sy = spatial_tile_grid(canvas_w, canvas_h, tile_w, tile_h, sizes.overlap_frac)
|
||||
ox, oy = spatial_tile_grid(out_w, out_h, tile_w, tile_h, sizes.overlap_frac)
|
||||
chunks = temporal_chunk_count(frames, sizes.tile_frames)
|
||||
|
||||
fields = dict(
|
||||
initial_canvas_width=canvas_w,
|
||||
initial_canvas_height=canvas_h,
|
||||
tile_width=tile_w,
|
||||
tile_height=tile_h,
|
||||
output_width=out_w,
|
||||
output_height=out_h,
|
||||
overlap_frac=sizes.overlap_frac,
|
||||
tile_frames=sizes.tile_frames,
|
||||
spatial_tiles_x=sx,
|
||||
spatial_tiles_y=sy,
|
||||
output_tiles_x=ox,
|
||||
output_tiles_y=oy,
|
||||
temporal_chunks=chunks,
|
||||
canvas_label=canvas_label,
|
||||
tile_label=tile_label,
|
||||
output_label=out_label,
|
||||
tile_clamped=tile_clamped,
|
||||
is_8k=out_label == "8K",
|
||||
)
|
||||
summary = format_summary(
|
||||
src_w=src_w, src_h=src_h, frame_count=frames, plan_fields=fields
|
||||
)
|
||||
return TilingPlan(summary=summary, **fields)
|
||||
|
||||
|
||||
def _dims_from_image(image) -> tuple[int, int, int]:
|
||||
"""Comfy IMAGE is [frames, height, width, channels]."""
|
||||
frames, height, width = (
|
||||
int(image.shape[0]),
|
||||
int(image.shape[1]),
|
||||
int(image.shape[2]),
|
||||
)
|
||||
return width, height, frames
|
||||
|
||||
|
||||
@comfy_node(
|
||||
name="LTXVGetTilingSizes",
|
||||
description="Get Tiling Sizes",
|
||||
)
|
||||
class LTXVGetTilingSizes:
|
||||
"""Size calculator in the spirit of Comfy's Get Image Size, for tiled fusion."""
|
||||
|
||||
CATEGORY = "ltxvideo"
|
||||
FUNCTION = "get_sizes"
|
||||
OUTPUT_NODE = True
|
||||
RETURN_TYPES = (
|
||||
"INT",
|
||||
"INT",
|
||||
"INT",
|
||||
"INT",
|
||||
"INT",
|
||||
"INT",
|
||||
"FLOAT",
|
||||
"INT",
|
||||
"BOOLEAN",
|
||||
"STRING",
|
||||
)
|
||||
RETURN_NAMES = (
|
||||
"initial_canvas_width",
|
||||
"initial_canvas_height",
|
||||
"tile_width",
|
||||
"tile_height",
|
||||
"output_width",
|
||||
"output_height",
|
||||
"overlap_frac",
|
||||
"tile_frames",
|
||||
"is_8k",
|
||||
"summary",
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {
|
||||
"width": (
|
||||
"INT",
|
||||
{
|
||||
"default": 1920,
|
||||
"min": 1,
|
||||
"max": 16384,
|
||||
"tooltip": "Source width in pixels (from Load Video / Get Image Size).",
|
||||
},
|
||||
),
|
||||
"height": (
|
||||
"INT",
|
||||
{
|
||||
"default": 1080,
|
||||
"min": 1,
|
||||
"max": 16384,
|
||||
"tooltip": "Source height in pixels.",
|
||||
},
|
||||
),
|
||||
"frame_count": (
|
||||
"INT",
|
||||
{
|
||||
"default": 97,
|
||||
"min": 1,
|
||||
"max": 10000,
|
||||
"tooltip": "Source frame count. tile_frames is 8n+1, 97 or less, "
|
||||
"or up to 121 when that avoids temporal chunking.",
|
||||
},
|
||||
),
|
||||
"tile_size": (
|
||||
list(TILE_CHOICES),
|
||||
{
|
||||
"default": "FullHD",
|
||||
"tooltip": "Spatial tile long-edge. qHD=960, HD=1280, FullHD=1920. "
|
||||
"Clamped to the initial canvas so a FullHD tile on a smaller "
|
||||
"canvas does not force tiling.",
|
||||
},
|
||||
),
|
||||
"initial_canvas_size": (
|
||||
list(TILE_CHOICES),
|
||||
{
|
||||
"default": "FullHD",
|
||||
"tooltip": "Stage-1 canvas long-edge (same choices as tile size). "
|
||||
"Source within 200px of FullHD is kept instead of rescaled.",
|
||||
},
|
||||
),
|
||||
"output_size": (
|
||||
list(OUTPUT_CHOICES),
|
||||
{
|
||||
"default": "FullHD",
|
||||
"tooltip": "Final canvas. FullHD is the 1920 long-edge step; "
|
||||
"4K is 2× that; 8K is 4×. Single-stage graphs use this as the "
|
||||
"canvas. Two-stage graphs use initial canvas for stage 1 and "
|
||||
"this for stage 2.",
|
||||
},
|
||||
),
|
||||
},
|
||||
"optional": {
|
||||
"image": (
|
||||
"IMAGE",
|
||||
{
|
||||
"tooltip": "Optional video/image batch. When connected, width, "
|
||||
"height, and frame_count are taken from it (Get Image Size style)."
|
||||
},
|
||||
),
|
||||
},
|
||||
}
|
||||
|
||||
def get_sizes(
|
||||
self,
|
||||
width,
|
||||
height,
|
||||
frame_count,
|
||||
tile_size="FullHD",
|
||||
initial_canvas_size="FullHD",
|
||||
output_size="FullHD",
|
||||
image: Optional[object] = None,
|
||||
**_kwargs,
|
||||
):
|
||||
if image is not None:
|
||||
width, height, frame_count = _dims_from_image(image)
|
||||
plan = resolve_tiling_plan(
|
||||
width,
|
||||
height,
|
||||
frame_count,
|
||||
tile_size=tile_size,
|
||||
initial_canvas_size=initial_canvas_size,
|
||||
output_size=output_size,
|
||||
)
|
||||
return {
|
||||
"ui": {"text": [plan.summary]},
|
||||
"result": (
|
||||
plan.initial_canvas_width,
|
||||
plan.initial_canvas_height,
|
||||
plan.tile_width,
|
||||
plan.tile_height,
|
||||
plan.output_width,
|
||||
plan.output_height,
|
||||
plan.overlap_frac,
|
||||
plan.tile_frames,
|
||||
plan.is_8k,
|
||||
plan.summary,
|
||||
),
|
||||
}
|
||||
+3
-3
@@ -5,7 +5,7 @@ import torch
|
||||
from comfy_extras.nodes_custom_sampler import SamplerCustomAdvanced
|
||||
from comfy_extras.nodes_lt import LTXVAddGuide, LTXVCropGuides
|
||||
|
||||
from .latents import LTXVAddLatentGuide, LTXVSelectLatents
|
||||
from .latents import LTXVSelectLatents, add_latent_guide
|
||||
from .nodes_registry import comfy_node
|
||||
|
||||
|
||||
@@ -233,7 +233,7 @@ class LTXVTiledSampler:
|
||||
tile_positive,
|
||||
tile_negative,
|
||||
tile_latents,
|
||||
) = LTXVAddLatentGuide().generate(
|
||||
) = add_latent_guide(
|
||||
vae=vae,
|
||||
positive=tile_positive,
|
||||
negative=tile_negative,
|
||||
@@ -249,7 +249,7 @@ class LTXVTiledSampler:
|
||||
tile_positive,
|
||||
tile_negative,
|
||||
tile_latents,
|
||||
) = LTXVAddLatentGuide().generate(
|
||||
) = add_latent_guide(
|
||||
vae=vae,
|
||||
positive=tile_positive,
|
||||
negative=tile_negative,
|
||||
|
||||
@@ -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
|
||||
|
||||
+153
@@ -0,0 +1,153 @@
|
||||
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.
|
||||
|
||||
The fill is always SDR display-encoded ``#66FF00`` (``_BG_COLOR_RGB``), even
|
||||
when ``images`` are already ACEScct. Native HDR inpaint IC-LoRA was trained
|
||||
with that pixel-space key; converting the green into ACEScct changes the
|
||||
VAE latent cue and breaks inpainting.
|
||||
"""
|
||||
|
||||
@classmethod
|
||||
def define_schema(cls):
|
||||
return io.Schema(
|
||||
node_id="LTXVInpaintPreprocess",
|
||||
category="Lightricks/image_processing",
|
||||
description=(
|
||||
"Composites images with a green (#66FF00) background where mask "
|
||||
"is active, for inpainting conditioning. Always paints SDR "
|
||||
"#66FF00 (not ACEScct-converted) — that key matches IC-LoRA "
|
||||
"training even on native HDR / ACEScct plates."
|
||||
),
|
||||
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
|
||||
);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,94 @@
|
||||
import { app } from "../../../scripts/app.js";
|
||||
|
||||
const NODE_NAME = "LTXVGetTilingSizes";
|
||||
|
||||
function linesOf(text) {
|
||||
return String(text || "").split("\n").length;
|
||||
}
|
||||
|
||||
function ensurePreview(node, text) {
|
||||
const value = Array.isArray(text) ? text.join("\n") : String(text ?? "");
|
||||
if (node._tilingPreviewEl) {
|
||||
node._tilingPreviewEl.textContent = value;
|
||||
const widget = (node.widgets || []).find((w) => w.name === "tiling_preview");
|
||||
if (widget?.computeSize) {
|
||||
const lines = linesOf(value);
|
||||
widget.computedHeight = Math.max(88, 14 + lines * 16);
|
||||
}
|
||||
return;
|
||||
}
|
||||
|
||||
const el = document.createElement("pre");
|
||||
el.className = "ltx-tiling-sizes-preview";
|
||||
el.textContent = value || "Queue prompt to preview tiling sizes.";
|
||||
Object.assign(el.style, {
|
||||
margin: "0",
|
||||
padding: "6px 8px",
|
||||
fontFamily: "ui-monospace, SFMono-Regular, Menlo, Consolas, monospace",
|
||||
fontSize: "11px",
|
||||
lineHeight: "1.4",
|
||||
whiteSpace: "pre-wrap",
|
||||
color: "var(--fg-color, #ddd)",
|
||||
background: "var(--bg-color, transparent)",
|
||||
border: "1px solid var(--border-color, rgba(255,255,255,0.12))",
|
||||
borderRadius: "6px",
|
||||
width: "100%",
|
||||
boxSizing: "border-box",
|
||||
pointerEvents: "none",
|
||||
});
|
||||
|
||||
const widget = node.addDOMWidget("tiling_preview", "preview", el, {
|
||||
serialize: false,
|
||||
hideOnZoom: false,
|
||||
getValue() {
|
||||
return el.textContent;
|
||||
},
|
||||
setValue(v) {
|
||||
el.textContent = v ?? "";
|
||||
},
|
||||
});
|
||||
widget.computeSize = function (width) {
|
||||
const lines = linesOf(el.textContent);
|
||||
const height = Math.max(88, 14 + lines * 16);
|
||||
this.computedHeight = height;
|
||||
return [width, height];
|
||||
};
|
||||
|
||||
node._tilingPreviewEl = el;
|
||||
}
|
||||
|
||||
app.registerExtension({
|
||||
name: "LTXVideo.GetTilingSizesPreview",
|
||||
async beforeRegisterNodeDef(nodeType, nodeData) {
|
||||
if (nodeData?.name !== NODE_NAME) {
|
||||
return;
|
||||
}
|
||||
|
||||
const onNodeCreated = nodeType.prototype.onNodeCreated;
|
||||
nodeType.prototype.onNodeCreated = function () {
|
||||
onNodeCreated?.apply(this, arguments);
|
||||
ensurePreview(this, "Queue prompt to preview tiling sizes.");
|
||||
};
|
||||
|
||||
const onExecuted = nodeType.prototype.onExecuted;
|
||||
nodeType.prototype.onExecuted = function (message) {
|
||||
onExecuted?.apply(this, arguments);
|
||||
const text = message?.text;
|
||||
if (text == null) {
|
||||
return;
|
||||
}
|
||||
ensurePreview(this, text);
|
||||
requestAnimationFrame(() => {
|
||||
const size = this.computeSize?.() || this.size;
|
||||
if (size && this.onResize) {
|
||||
const next = [
|
||||
Math.max(size[0], this.size?.[0] || 0),
|
||||
Math.max(size[1], this.size?.[1] || 0),
|
||||
];
|
||||
this.onResize(next);
|
||||
}
|
||||
app.graph?.setDirtyCanvas?.(true, false);
|
||||
});
|
||||
};
|
||||
},
|
||||
});
|
||||
Reference in New Issue
Block a user