Compare commits
47
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
7f22eb16eb | ||
|
|
db476dacc1 | ||
|
|
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).
|
||||
|
||||
+46
-26
@@ -1,3 +1,5 @@
|
||||
from .audio_only import LTXVAudioOnlyEmptyVideoLatent, LTXVAudioOnlyModel
|
||||
from .conditioning_loader import LTXVLoadConditioning
|
||||
from .decoder_noise import DecoderNoise
|
||||
from .dynamic_conditioning import DynamicConditioning
|
||||
from .easy_samplers import (
|
||||
@@ -9,16 +11,23 @@ from .easy_samplers import (
|
||||
)
|
||||
from .gemma_api_conditioning import GemmaAPITextEncode
|
||||
from .gemma_encoder import LTXVGemmaCLIPModelLoader, LTXVGemmaEnhancePrompt
|
||||
from .guide import LTXVAddGuideAdvanced
|
||||
from .guide import LTXVImgToVideoAdvanced
|
||||
from .guiders import GuiderParametersNode, MultimodalGuiderNode
|
||||
from .latent_norm import (
|
||||
LTXVAdainLatent,
|
||||
LTXVPerStepAdainPatcher,
|
||||
LTXVPerStepStatNormPatcher,
|
||||
LTXVStatNormLatent,
|
||||
from .hdr_nodes import (
|
||||
LTXVHDRDecodePostprocess,
|
||||
LTXVLoadEXRSequence,
|
||||
LTXVSaveHLG,
|
||||
LTXVSDRToHDRWorkingSpace,
|
||||
LTXVVAEForceFloat32,
|
||||
)
|
||||
from .iclora import (
|
||||
LTXAddVideoICLoRAGuide,
|
||||
LTXAddVideoICLoRAGuideAdvanced,
|
||||
LTXICLoRALoaderModelOnly,
|
||||
LTXVSetAudioRefTokens,
|
||||
)
|
||||
from .latent_norm import LTXVAdainLatent
|
||||
from .latents import (
|
||||
LTXVAddLatentGuide,
|
||||
LTXVImgToVideoConditionOnly,
|
||||
LTXVSelectLatents,
|
||||
LTXVSetVideoLatentNoiseMasks,
|
||||
@@ -29,50 +38,46 @@ from .low_vram_loaders import (
|
||||
LowVRAMCheckpointLoader,
|
||||
LowVRAMLatentUpscaleModelLoader,
|
||||
)
|
||||
from .masks import LTXVPreprocessMasks
|
||||
from .nodes_registry import NODE_CLASS_MAPPINGS as RUNTIME_NODE_CLASS_MAPPINGS
|
||||
from .nodes_registry import (
|
||||
NODE_DISPLAY_NAME_MAPPINGS as RUNTIME_NODE_DISPLAY_NAME_MAPPINGS,
|
||||
)
|
||||
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 .vae_patcher import LTXVPatcherVAE
|
||||
from .utiltily_nodes import FloatToInt, ImageToCPU
|
||||
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,
|
||||
"LTXVImgToVideoAdvanced": LTXVImgToVideoAdvanced,
|
||||
"LTXVAdainLatent": LTXVAdainLatent,
|
||||
"LTXVImgToVideoConditionOnly": LTXVImgToVideoConditionOnly,
|
||||
"LTXVPerStepAdainPatcher": LTXVPerStepAdainPatcher,
|
||||
"LTXVApplySTG": LTXVApplySTG,
|
||||
"LTXVBaseSampler": LTXVBaseSampler,
|
||||
"LTXVInContextSampler": LTXVInContextSampler,
|
||||
"LTXVExtendSampler": LTXVExtendSampler,
|
||||
"LTXVNormalizingSampler": LTXVNormalizingSampler,
|
||||
"LTXVPreprocessMasks": LTXVPreprocessMasks,
|
||||
"LTXVPatcherVAE": LTXVPatcherVAE,
|
||||
"LTXVPromptEnhancer": LTXVPromptEnhancer,
|
||||
"LTXVPromptEnhancerLoader": LTXVPromptEnhancerLoader,
|
||||
"LTXQ8Patch": LTXVQ8Patch,
|
||||
"LTXVQ8LoraModelLoader": LTXVQ8LoraModelLoader,
|
||||
"LTXVSelectLatents": LTXVSelectLatents,
|
||||
"LTXVSetVideoLatentNoiseMasks": LTXVSetVideoLatentNoiseMasks,
|
||||
"LTXVTiledSampler": LTXVTiledSampler,
|
||||
"LTXVTiledFusionSampler": LTXVTiledFusionSampler,
|
||||
"LTXVGetTilingSizes": LTXVGetTilingSizes,
|
||||
"LTXVLoopingSampler": LTXVLoopingSampler,
|
||||
"LTXVTiledVAEDecode": LTXVTiledVAEDecode,
|
||||
"MultimodalGuider": MultimodalGuiderNode,
|
||||
@@ -82,8 +87,7 @@ NODE_CLASS_MAPPINGS = {
|
||||
"STGGuiderNode": STGGuiderNode,
|
||||
"LTXVMultiPromptProvider": MultiPromptProvider,
|
||||
"ImageToCPU": ImageToCPU,
|
||||
"LTXVStatNormLatent": LTXVStatNormLatent,
|
||||
"LTXVPerStepStatNormPatcher": LTXVPerStepStatNormPatcher,
|
||||
"LTXFloatToInt": FloatToInt,
|
||||
"LTXVGemmaCLIPModelLoader": LTXVGemmaCLIPModelLoader,
|
||||
"LTXVGemmaEnhancePrompt": LTXVGemmaEnhancePrompt,
|
||||
"GemmaAPITextEncode": GemmaAPITextEncode,
|
||||
@@ -91,6 +95,23 @@ NODE_CLASS_MAPPINGS = {
|
||||
"LowVRAMCheckpointLoader": LowVRAMCheckpointLoader,
|
||||
"LowVRAMAudioVAELoader": LowVRAMAudioVAELoader,
|
||||
"LowVRAMLatentUpscaleModelLoader": LowVRAMLatentUpscaleModelLoader,
|
||||
"LTXAddVideoICLoRAGuide": LTXAddVideoICLoRAGuide,
|
||||
"LTXAddVideoICLoRAGuideAdvanced": LTXAddVideoICLoRAGuideAdvanced,
|
||||
"LTXICLoRALoaderModelOnly": LTXICLoRALoaderModelOnly,
|
||||
"LTXVSetAudioRefTokens": LTXVSetAudioRefTokens,
|
||||
"LTXVLoadConditioning": LTXVLoadConditioning,
|
||||
"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,
|
||||
@@ -100,16 +121,15 @@ NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
for name in NODE_CLASS_MAPPINGS.keys()
|
||||
}
|
||||
|
||||
# Merge the node mappings from tricks into the main mappings
|
||||
NODE_CLASS_MAPPINGS.update(TRICKS_NODE_CLASS_MAPPINGS)
|
||||
NODE_DISPLAY_NAME_MAPPINGS.update(TRICKS_NODE_DISPLAY_NAME_MAPPINGS)
|
||||
|
||||
# Update with runtime mappings (these will override static mappings if there are any differences)
|
||||
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
|
||||
+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,
|
||||
)
|
||||
|
||||
@@ -17,126 +17,6 @@ def blur_internal(image, blur_radius):
|
||||
return image
|
||||
|
||||
|
||||
@comfy_node(name="LTXVAddGuideAdvanced")
|
||||
class LTXVAddGuideAdvanced:
|
||||
@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. For single-frame images or "
|
||||
"videos with 1-8 frames, any frame_idx value is acceptable. For videos with 9+ "
|
||||
"frames, frame_idx must be divisible by 8, otherwise it will be rounded down to "
|
||||
"the nearest multiple of 8. 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 will make the conditioning more exact.",
|
||||
},
|
||||
),
|
||||
"crf": (
|
||||
"INT",
|
||||
{
|
||||
"default": 29,
|
||||
"min": 0,
|
||||
"max": 51,
|
||||
"step": 1,
|
||||
"tooltip": "CRF value for the video. Higher values mean more motion, lower values mean higher quality.",
|
||||
},
|
||||
),
|
||||
"blur_radius": (
|
||||
"INT",
|
||||
{
|
||||
"default": 0,
|
||||
"min": 0,
|
||||
"max": 7,
|
||||
"step": 1,
|
||||
"tooltip": "Blur kernel radius size. Higher values mean more motion, lower values mean higher quality.",
|
||||
},
|
||||
),
|
||||
"interpolation": (
|
||||
[
|
||||
"lanczos",
|
||||
"bislerp",
|
||||
"nearest",
|
||||
"bilinear",
|
||||
"bicubic",
|
||||
"area",
|
||||
"nearest-exact",
|
||||
],
|
||||
{"default": "lanczos"},
|
||||
),
|
||||
"crop": (["center", "disabled"], {"default": "disabled"}),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("CONDITIONING", "CONDITIONING", "LATENT")
|
||||
RETURN_NAMES = ("positive", "negative", "latent")
|
||||
|
||||
CATEGORY = "conditioning/video_models"
|
||||
FUNCTION = "generate"
|
||||
|
||||
DESCRIPTION = (
|
||||
"Adds a conditioning frame or a video at a specific frame index. "
|
||||
"This node is used to add a keyframe or a video segment which should appear in the "
|
||||
"generated video at a specified index. It resizes the image to the correct size and "
|
||||
"applies preprocessing to it."
|
||||
)
|
||||
|
||||
def generate(
|
||||
self,
|
||||
positive,
|
||||
negative,
|
||||
vae,
|
||||
latent,
|
||||
image,
|
||||
frame_idx,
|
||||
strength,
|
||||
crf,
|
||||
blur_radius,
|
||||
interpolation,
|
||||
crop,
|
||||
):
|
||||
_, width_scale_factor, height_scale_factor = vae.downscale_index_formula
|
||||
width, height = (
|
||||
latent["samples"].shape[4] * width_scale_factor,
|
||||
latent["samples"].shape[3] * 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)
|
||||
return nodes_lt.LTXVAddGuide().execute(
|
||||
positive=positive,
|
||||
negative=negative,
|
||||
vae=vae,
|
||||
latent=latent,
|
||||
image=image,
|
||||
frame_idx=frame_idx,
|
||||
strength=strength,
|
||||
)
|
||||
|
||||
|
||||
@comfy_node(name="LTXVImgToVideoAdvanced")
|
||||
class LTXVImgToVideoAdvanced:
|
||||
@classmethod
|
||||
|
||||
@@ -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
|
||||
-259
@@ -1,8 +1,6 @@
|
||||
import copy
|
||||
from typing import Callable
|
||||
|
||||
import torch
|
||||
from comfy.model_patcher import ModelPatcher
|
||||
|
||||
from .nodes_registry import comfy_node
|
||||
|
||||
@@ -72,260 +70,3 @@ class LTXVAdainLatent:
|
||||
|
||||
latents_copy["samples"] = torch.lerp(latents["samples"], t, factor)
|
||||
return (latents_copy,)
|
||||
|
||||
|
||||
@comfy_node(name="LTXVStatNormLatent")
|
||||
class LTXVStatNormLatent:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"latents": ("LATENT",),
|
||||
"target_mean": (
|
||||
"FLOAT",
|
||||
{
|
||||
"default": 0.0,
|
||||
"min": -10.0,
|
||||
"max": 10.0,
|
||||
"step": 0.01,
|
||||
"round": 0.01,
|
||||
},
|
||||
),
|
||||
"target_std": (
|
||||
"FLOAT",
|
||||
{
|
||||
"default": 1.0,
|
||||
"min": 0.01,
|
||||
"max": 10.0,
|
||||
"step": 0.01,
|
||||
"round": 0.01,
|
||||
},
|
||||
),
|
||||
"percentile": (
|
||||
"FLOAT",
|
||||
{
|
||||
"default": 95.0,
|
||||
"min": 50.0,
|
||||
"max": 100.0,
|
||||
"step": 0.1,
|
||||
"round": 0.1,
|
||||
"tooltip": "Percentile of distribution to use for statistics calculation",
|
||||
},
|
||||
),
|
||||
"factor": (
|
||||
"FLOAT",
|
||||
{
|
||||
"default": 1.0,
|
||||
"min": -10.0,
|
||||
"max": 10.0,
|
||||
"step": 0.01,
|
||||
"round": 0.01,
|
||||
},
|
||||
),
|
||||
"clip_outliers": ("BOOLEAN", {"default": False}),
|
||||
},
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("LATENT",)
|
||||
FUNCTION = "statistical_normalize"
|
||||
|
||||
CATEGORY = "Lightricks/latents"
|
||||
|
||||
def statistical_normalize(
|
||||
self, latents, target_mean, target_std, percentile, factor, clip_outliers
|
||||
):
|
||||
latents_copy = copy.deepcopy(latents)
|
||||
t = latents_copy["samples"] # B x C x F x H x W
|
||||
|
||||
# For 95% of distribution, we want to exclude 2.5% from each tail
|
||||
lower_percentile = (100 - percentile) / 2
|
||||
upper_percentile = 100 - lower_percentile
|
||||
|
||||
for i in range(t.size(0)): # batch
|
||||
for c in range(t.size(1)): # channel
|
||||
channel_data = t[i, c]
|
||||
original_shape = channel_data.shape
|
||||
channel_flat = channel_data.flatten()
|
||||
|
||||
# Calculate percentiles
|
||||
lower_bound = torch.quantile(channel_flat, lower_percentile / 100)
|
||||
upper_bound = torch.quantile(channel_flat, upper_percentile / 100)
|
||||
|
||||
# Create mask for values within the percentile range
|
||||
mask = (channel_flat >= lower_bound) & (channel_flat <= upper_bound)
|
||||
|
||||
# Calculate mean and std only on the masked values
|
||||
if mask.sum() > 0:
|
||||
filtered_data = channel_flat[mask]
|
||||
current_mean = filtered_data.mean()
|
||||
current_std = filtered_data.std()
|
||||
|
||||
# Avoid division by zero
|
||||
if current_std > 1e-8:
|
||||
# Normalize all values
|
||||
normalized_flat = (
|
||||
(channel_flat - current_mean) / current_std
|
||||
) * target_std + target_mean
|
||||
|
||||
if clip_outliers:
|
||||
# Calculate the normalized bounds
|
||||
normalized_lower = (
|
||||
(lower_bound - current_mean) / current_std
|
||||
) * target_std + target_mean
|
||||
normalized_upper = (
|
||||
(upper_bound - current_mean) / current_std
|
||||
) * target_std + target_mean
|
||||
|
||||
# Clip outliers to the normalized bounds
|
||||
normalized_flat = torch.where(
|
||||
channel_flat < lower_bound,
|
||||
normalized_lower,
|
||||
normalized_flat,
|
||||
)
|
||||
normalized_flat = torch.where(
|
||||
channel_flat > upper_bound,
|
||||
normalized_upper,
|
||||
normalized_flat,
|
||||
)
|
||||
|
||||
# Reshape back to original shape
|
||||
t[i, c] = normalized_flat.reshape(original_shape)
|
||||
else:
|
||||
# If std is too small, just shift by mean
|
||||
t[i, c] = channel_data - current_mean + target_mean
|
||||
|
||||
latents_copy["samples"] = torch.lerp(latents["samples"], t, factor)
|
||||
return (latents_copy,)
|
||||
|
||||
|
||||
class PerStepNormPatcher:
|
||||
"""
|
||||
Base class for per-step normalization nodes.
|
||||
"""
|
||||
|
||||
@classmethod
|
||||
def required(s):
|
||||
return {
|
||||
"model": ("MODEL",),
|
||||
"factors": (
|
||||
"STRING",
|
||||
{
|
||||
"default": "0.9, 0.75, 0.0",
|
||||
"tooltip": "Comma-separated list of factors, each factor will be used for one step.",
|
||||
},
|
||||
),
|
||||
}
|
||||
|
||||
@classmethod
|
||||
def optional(s):
|
||||
return {}
|
||||
|
||||
@staticmethod
|
||||
def patch_smooth_norm(
|
||||
model: ModelPatcher, factors: str, normalize_fn: Callable
|
||||
) -> ModelPatcher:
|
||||
model = model.clone()
|
||||
factors = [float(x) for x in factors.split(",")]
|
||||
step = 0
|
||||
|
||||
def norm_fn(args):
|
||||
nonlocal step
|
||||
latent = args["denoised"]
|
||||
factor = factors[min(step, len(factors) - 1)]
|
||||
step += 1
|
||||
out_latent = normalize_fn(latent={"samples": latent}, factor=factor)
|
||||
return out_latent[0]["samples"]
|
||||
|
||||
model.set_model_sampler_post_cfg_function(norm_fn)
|
||||
return (model,)
|
||||
|
||||
@classmethod
|
||||
def IS_CHANGED(s, *args, **kwargs):
|
||||
return float("NaN")
|
||||
|
||||
|
||||
@comfy_node(name="LTXVPerStepAdainPatcher")
|
||||
class LTXVPerStepAdainPatcher(PerStepNormPatcher):
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": super().required()
|
||||
| {
|
||||
"reference": ("LATENT",),
|
||||
},
|
||||
"optional": super().optional()
|
||||
| {
|
||||
"per_frame": ("BOOLEAN", {"default": False}),
|
||||
},
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("MODEL",)
|
||||
FUNCTION = "patch_model"
|
||||
|
||||
CATEGORY = "Lightricks/latents"
|
||||
|
||||
def patch_model(self, model, factors, reference, per_frame=False):
|
||||
def cfg_adain(latent, factor):
|
||||
return LTXVAdainLatent().batch_normalize(
|
||||
latent, reference, factor, per_frame
|
||||
)
|
||||
|
||||
return self.patch_smooth_norm(model, factors, cfg_adain)
|
||||
|
||||
|
||||
@comfy_node(name="LTXVPerStepStatNormPatcher")
|
||||
class LTXVPerStepStatNormPatcher(PerStepNormPatcher):
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": super().required()
|
||||
| {
|
||||
"target_mean": (
|
||||
"FLOAT",
|
||||
{
|
||||
"default": 0.0,
|
||||
"min": -10.0,
|
||||
"max": 10.0,
|
||||
"step": 0.01,
|
||||
"round": 0.01,
|
||||
},
|
||||
),
|
||||
"target_std": (
|
||||
"FLOAT",
|
||||
{
|
||||
"default": 1.0,
|
||||
"min": 0.01,
|
||||
"max": 10.0,
|
||||
"step": 0.01,
|
||||
"round": 0.01,
|
||||
},
|
||||
),
|
||||
"percentile": (
|
||||
"FLOAT",
|
||||
{
|
||||
"default": 95.0,
|
||||
"min": 50.0,
|
||||
"max": 100.0,
|
||||
"step": 0.1,
|
||||
"round": 0.1,
|
||||
"tooltip": "Percentile of distribution to use for statistics calculation",
|
||||
},
|
||||
),
|
||||
"clip_outliers": ("BOOLEAN", {"default": False}),
|
||||
},
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("MODEL",)
|
||||
FUNCTION = "patch_model"
|
||||
|
||||
CATEGORY = "Lightricks/latents"
|
||||
|
||||
def patch_model(
|
||||
self, model, factors, target_mean, target_std, percentile, clip_outliers
|
||||
):
|
||||
def cfg_stat_norm(latent, factor):
|
||||
return LTXVStatNormLatent().statistical_normalize(
|
||||
latent, target_mean, target_std, percentile, factor, clip_outliers
|
||||
)
|
||||
|
||||
return self.patch_smooth_norm(model, factors, cfg_stat_norm)
|
||||
|
||||
+52
-194
@@ -1,5 +1,3 @@
|
||||
from typing import Optional
|
||||
|
||||
import comfy.utils
|
||||
import comfy_extras.nodes_lt as nodes_lt
|
||||
import numpy as np
|
||||
@@ -7,6 +5,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
|
||||
|
||||
|
||||
@@ -90,121 +89,6 @@ class LTXVSelectLatents:
|
||||
raise
|
||||
|
||||
|
||||
@comfy_node(name="LTXVAddLatents")
|
||||
class LTXVAddLatents:
|
||||
"""
|
||||
Concatenates two video latents along the frames dimension.
|
||||
|
||||
Features:
|
||||
- Validates dimension compatibility
|
||||
- Handles device placement
|
||||
- Preserves noise masks with proper handling
|
||||
- Supports batch processing
|
||||
"""
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"latents1": ("LATENT",),
|
||||
"latents2": ("LATENT",),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("LATENT",)
|
||||
FUNCTION = "add_latents"
|
||||
CATEGORY = "latent/video"
|
||||
DESCRIPTION = (
|
||||
"Concatenates two video latents along the frames dimension. "
|
||||
"latents1 and latents2 must have the same dimensions except for the frames dimension."
|
||||
)
|
||||
|
||||
def add_latents(
|
||||
self, latents1: torch.Tensor, latents2: torch.Tensor
|
||||
) -> torch.Tensor:
|
||||
"""
|
||||
Concatenates two video latents along the frames dimension.
|
||||
|
||||
Args:
|
||||
latents1 (dict): First video latent dictionary
|
||||
latents2 (dict): Second video latent dictionary
|
||||
|
||||
Returns:
|
||||
tuple: Contains concatenated latent dictionary
|
||||
|
||||
Raises:
|
||||
ValueError: If latent dimensions don't match
|
||||
RuntimeError: If tensor operations fail
|
||||
"""
|
||||
try:
|
||||
s = latents1.copy()
|
||||
video_latent1 = latents1["samples"]
|
||||
video_latent2 = latents2["samples"]
|
||||
|
||||
# Ensure tensors are on the same device
|
||||
target_device = video_latent1.device
|
||||
video_latent2 = video_latent2.to(target_device)
|
||||
|
||||
# Validate dimensions
|
||||
self._validate_dimensions(video_latent1, video_latent2)
|
||||
|
||||
# Concatenate along frames dimension
|
||||
s["samples"] = torch.cat([video_latent1, video_latent2], dim=2)
|
||||
|
||||
# Handle noise masks
|
||||
s["noise_mask"] = self._merge_noise_masks(
|
||||
latents1, latents2, video_latent1.shape[2], video_latent2.shape[2]
|
||||
)
|
||||
|
||||
return (s,)
|
||||
|
||||
except Exception as e:
|
||||
print(f"[LTXVAddLatents] Error: {str(e)}")
|
||||
raise
|
||||
|
||||
def _validate_dimensions(self, latent1: torch.Tensor, latent2: torch.Tensor):
|
||||
"""Validates that latent dimensions match except for frames."""
|
||||
# video latents
|
||||
if latent1.ndim == 5 and latent2.ndim == 5:
|
||||
b1, c1, f1, h1, w1 = latent1.shape
|
||||
b2, c2, f2, h2, w2 = latent2.shape
|
||||
|
||||
if not (b1 == b2 and c1 == c2 and h1 == h2 and w1 == w2):
|
||||
raise ValueError(
|
||||
f"Latent dimensions must match (except frames dimension).\n"
|
||||
f"Got shapes {latent1.shape} and {latent2.shape}"
|
||||
)
|
||||
# audio latents
|
||||
elif latent1.ndim == 4 and latent2.ndim == 4:
|
||||
b1, c1, f1, s1 = latent1.shape
|
||||
b2, c2, f2, s2 = latent2.shape
|
||||
|
||||
if not (b1 == b2 and c1 == c2 and s1 == s2):
|
||||
raise ValueError(
|
||||
f"Latent dimensions must match (except frames dimension).\n"
|
||||
f"Got shapes {latent1.shape} and {latent2.shape}"
|
||||
)
|
||||
else:
|
||||
raise ValueError(
|
||||
f"Latent dimensions must be 4 (audio) or 5 (video) for both inputs.\n"
|
||||
f"Got shapes {latent1.shape} and {latent2.shape}"
|
||||
)
|
||||
|
||||
def _merge_noise_masks(
|
||||
self, latents1: torch.Tensor, latents2: torch.Tensor, frames1: int, frames2: int
|
||||
) -> Optional[torch.Tensor]:
|
||||
"""Merges noise masks from both latents with proper handling."""
|
||||
if "noise_mask" in latents1 and "noise_mask" in latents2:
|
||||
return torch.cat([latents1["noise_mask"], latents2["noise_mask"]], dim=2)
|
||||
elif "noise_mask" in latents1 and latents1["noise_mask"] is not None:
|
||||
zeros = torch.zeros_like(latents1["noise_mask"][:, :, :frames2, :, :])
|
||||
return torch.cat([latents1["noise_mask"], zeros], dim=2)
|
||||
elif "noise_mask" in latents2 and latents2["noise_mask"] is not None:
|
||||
zeros = torch.zeros_like(latents2["noise_mask"][:, :, :frames1, :, :])
|
||||
return torch.cat([zeros, latents2["noise_mask"]], dim=2)
|
||||
return None
|
||||
|
||||
|
||||
@comfy_node(name="LTXVSetVideoLatentNoiseMasks")
|
||||
class LTXVSetVideoLatentNoiseMasks:
|
||||
"""
|
||||
@@ -397,85 +281,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 +655,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,212 +0,0 @@
|
||||
from typing import Any
|
||||
|
||||
import torch
|
||||
from comfy.model_management import is_device_cpu
|
||||
from comfy_extras.nodes_mask import GrowMask
|
||||
|
||||
from .nodes_registry import comfy_node
|
||||
|
||||
|
||||
@comfy_node(name="LTXVPreprocessMasks")
|
||||
class LTXVPreprocessMasks:
|
||||
"""
|
||||
Preprocesses masks for use with the LTXVideo model's latent masking.
|
||||
|
||||
This node handles temporal mask processing by:
|
||||
1. Validating mask dimensions against VAE downscaling factors
|
||||
2. Optionally inverting masks
|
||||
3. Handling the first frame mask separately (can be ignored)
|
||||
4. Combining temporal masks using max pooling
|
||||
5. Growing/shrinking masks with morphological operations
|
||||
6. Clamping mask values to ensure proper opacity
|
||||
|
||||
The output is a set of masks ready for latent-space masking operations.
|
||||
"""
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls: "LTXVPreprocessMasks") -> dict[str, Any]: # noqa: N802
|
||||
return {
|
||||
"required": {
|
||||
"masks": ("MASK",),
|
||||
"vae": ("VAE",),
|
||||
"invert_input_masks": (
|
||||
"BOOLEAN",
|
||||
{
|
||||
"default": False,
|
||||
"tooltip": "Invert the input masks before processing."
|
||||
"Useful for masking vs unmasking operations.",
|
||||
},
|
||||
),
|
||||
"ignore_first_mask": (
|
||||
"BOOLEAN",
|
||||
{
|
||||
"default": True,
|
||||
"tooltip": "Zero out the first mask, typically used when"
|
||||
"it corresponds to the conditioning frame.",
|
||||
},
|
||||
),
|
||||
"pooling_method": (
|
||||
["max", "mean", "min"],
|
||||
{
|
||||
"default": "max",
|
||||
"tooltip": "Method to combine temporal masks. Max preserves strongest values, "
|
||||
"mean averages them, min takes weakest values.",
|
||||
},
|
||||
),
|
||||
"grow_mask": (
|
||||
"INT",
|
||||
{
|
||||
"default": 0,
|
||||
"min": -16384,
|
||||
"max": 16384,
|
||||
"step": 1,
|
||||
"tooltip": "Pixels to grow (positive) or shrink (negative) the mask."
|
||||
"Uses morphological operations.",
|
||||
},
|
||||
),
|
||||
"tapered_corners": (
|
||||
"BOOLEAN",
|
||||
{
|
||||
"default": True,
|
||||
"tooltip": "Use tapered corners in morphological operations for smoother mask edges.",
|
||||
},
|
||||
),
|
||||
"clamp_min": (
|
||||
"FLOAT",
|
||||
{
|
||||
"default": 0.5,
|
||||
"min": 0.0,
|
||||
"max": 1.0,
|
||||
"step": 0.01,
|
||||
"tooltip": "Minimum value to clamp masks to, ensuring minimum opacity.",
|
||||
},
|
||||
),
|
||||
"clamp_max": (
|
||||
"FLOAT",
|
||||
{
|
||||
"default": 1.0,
|
||||
"min": 0.0,
|
||||
"max": 1.0,
|
||||
"step": 0.01,
|
||||
"tooltip": "Maximum value to clamp masks to.",
|
||||
},
|
||||
),
|
||||
},
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("MASK",)
|
||||
FUNCTION = "preprocess_masks"
|
||||
CATEGORY = "Lightricks/mask_operations"
|
||||
DESCRIPTION = (
|
||||
"Preprocess masks to be used for masking latents in the LTXVideo model."
|
||||
)
|
||||
|
||||
def preprocess_masks(
|
||||
self,
|
||||
masks: torch.Tensor,
|
||||
vae, # noqa: ANN001
|
||||
pooling_method: str,
|
||||
grow_mask: int,
|
||||
tapered_corners: bool,
|
||||
clamp_min: float,
|
||||
clamp_max: float,
|
||||
ignore_first_mask: bool,
|
||||
invert_input_masks: bool,
|
||||
) -> tuple[torch.Tensor]:
|
||||
"""
|
||||
Preprocess masks for temporal latent masking.
|
||||
|
||||
Args:
|
||||
masks: Input masks tensor of shape (batch_size, H, W) where batch_size = 1 + (num_frames - 1)
|
||||
The first mask is for the conditioning frame, remaining masks are for generated frames
|
||||
vae: VAE model object containing downscale_index_formula attribute
|
||||
pooling_method: Method to combine temporal masks ('max', 'mean', 'min')
|
||||
grow_mask: Pixels to grow (positive) or shrink (negative) the mask
|
||||
tapered_corners: Whether to use tapered corners in morphological operations
|
||||
clamp_min: Minimum value to clamp masks to (ensures minimum opacity)
|
||||
clamp_max: Maximum value to clamp masks to
|
||||
ignore_first_mask: Whether to zero out the first mask (often used for conditioning)
|
||||
invert_input_masks: Whether to invert all input masks (1 - mask)
|
||||
|
||||
Returns:
|
||||
Tuple containing the processed masks tensor
|
||||
|
||||
Raises:
|
||||
ValueError: If mask dimensions are not compatible with VAE downscaling factors
|
||||
"""
|
||||
|
||||
if masks.ndim != 3:
|
||||
raise ValueError("Masks must be of shape (batch_size, H, W).")
|
||||
|
||||
batch_size, H, W = masks.shape
|
||||
|
||||
# Get VAE downscaling factors for temporal and spatial dimensions
|
||||
time_scale_factor, width_scale_factor, height_scale_factor = (
|
||||
vae.downscale_index_formula
|
||||
)
|
||||
|
||||
# Validate that mask dimensions are compatible with VAE downscaling
|
||||
# The batch size minus 1 (for the first frame) must be divisible by time_scale_factor
|
||||
if (batch_size - 1) % time_scale_factor != 0:
|
||||
raise ValueError(
|
||||
f"Masks batch size must have a multiple of {time_scale_factor} masks + 1."
|
||||
)
|
||||
|
||||
if H % height_scale_factor != 0:
|
||||
raise ValueError(
|
||||
f"Masks height must have a multiple of {height_scale_factor}."
|
||||
)
|
||||
|
||||
if W % width_scale_factor != 0:
|
||||
raise ValueError(
|
||||
f"Masks width must have a multiple of {width_scale_factor}."
|
||||
)
|
||||
|
||||
# Optionally invert all masks (useful for masking vs unmasking operations)
|
||||
if invert_input_masks:
|
||||
masks = 1.0 - masks
|
||||
|
||||
# Handle first mask separately - often used as conditioning frame
|
||||
if ignore_first_mask:
|
||||
first_mask = masks[0:1, :, :] * 0 # Zero out the first mask
|
||||
else:
|
||||
first_mask = masks[0:1, :, :]
|
||||
|
||||
# Process remaining masks by grouping them according to time_scale_factor
|
||||
# and taking the maximum value across each temporal group
|
||||
remaining = masks[1:].view(-1, time_scale_factor, H, W)
|
||||
|
||||
# Apply selected pooling method
|
||||
pooling_methods = {
|
||||
"max": lambda x: x.max(dim=1)[0],
|
||||
"mean": lambda x: x.mean(dim=1),
|
||||
"min": lambda x: x.min(dim=1)[0],
|
||||
}
|
||||
|
||||
combined = pooling_methods[pooling_method](remaining)
|
||||
|
||||
# Apply morphological operations to grow or shrink masks
|
||||
if grow_mask != 0:
|
||||
grow_mask_node = GrowMask()
|
||||
|
||||
# GrowMask node is not supported on GPU, so we need to move the combined tensor to CPU
|
||||
# if it is not already on CPU
|
||||
input_device = combined.device
|
||||
if not is_device_cpu(input_device):
|
||||
combined = combined.cpu()
|
||||
|
||||
(combined,) = grow_mask_node.expand_mask(
|
||||
mask=combined, expand=grow_mask, tapered_corners=tapered_corners
|
||||
)
|
||||
|
||||
# Move the combined tensor back to the original device
|
||||
if not is_device_cpu(input_device):
|
||||
combined = combined.to(input_device)
|
||||
|
||||
# Combine first mask with processed remaining masks and clamp values
|
||||
# Clamping ensures masks have proper opacity values
|
||||
processed_masks = torch.cat(
|
||||
[first_mask, torch.clamp(combined, clamp_min, clamp_max)], dim=0
|
||||
)
|
||||
|
||||
return (processed_masks,)
|
||||
@@ -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
|
||||
|
||||
|
||||
@@ -1,208 +0,0 @@
|
||||
import os
|
||||
import shutil
|
||||
|
||||
import comfy.model_management
|
||||
import comfy.model_patcher
|
||||
import folder_paths
|
||||
import torch
|
||||
from transformers import AutoModelForCausalLM, AutoProcessor, AutoTokenizer
|
||||
|
||||
from .nodes_registry import comfy_node
|
||||
from .prompt_enhancer_utils import generate_cinematic_prompt
|
||||
|
||||
LLM_NAME = ["unsloth/Llama-3.2-3B-Instruct"]
|
||||
|
||||
IMAGE_CAPTIONER = ["MiaoshouAI/Florence-2-large-PromptGen-v2.0"]
|
||||
|
||||
MODELS_PATH_KEY = "LLM"
|
||||
|
||||
|
||||
class PromptEnhancer(torch.nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
image_caption_processor: AutoProcessor,
|
||||
image_caption_model: AutoModelForCausalLM,
|
||||
llm_model: AutoModelForCausalLM,
|
||||
llm_tokenizer: AutoTokenizer,
|
||||
):
|
||||
super().__init__()
|
||||
self.image_caption_processor = image_caption_processor
|
||||
self.image_caption_model = image_caption_model
|
||||
self.llm_model = llm_model
|
||||
self.llm_tokenizer = llm_tokenizer
|
||||
self.device = image_caption_model.device
|
||||
# model parameters and buffer sizes plus some extra 1GB.
|
||||
self.model_size = (
|
||||
self.get_model_size(self.image_caption_model)
|
||||
+ self.get_model_size(self.llm_model)
|
||||
+ 1073741824
|
||||
)
|
||||
|
||||
def forward(self, prompt, image_conditioning, max_resulting_tokens):
|
||||
enhanced_prompt = generate_cinematic_prompt(
|
||||
self.image_caption_model,
|
||||
self.image_caption_processor,
|
||||
self.llm_model,
|
||||
self.llm_tokenizer,
|
||||
prompt,
|
||||
image_conditioning,
|
||||
max_new_tokens=max_resulting_tokens,
|
||||
)
|
||||
|
||||
return enhanced_prompt
|
||||
|
||||
@staticmethod
|
||||
def get_model_size(model):
|
||||
total_size = sum(p.numel() * p.element_size() for p in model.parameters())
|
||||
total_size += sum(b.numel() * b.element_size() for b in model.buffers())
|
||||
return total_size
|
||||
|
||||
def memory_required(self, input_shape):
|
||||
return self.model_size
|
||||
|
||||
|
||||
@comfy_node(name="LTXVPromptEnhancerLoader")
|
||||
class LTXVPromptEnhancerLoader:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"llm_name": (
|
||||
"STRING",
|
||||
{
|
||||
"default": LLM_NAME,
|
||||
"tooltip": "The hugging face name of the llm model to load.",
|
||||
},
|
||||
),
|
||||
"image_captioner_name": (
|
||||
"STRING",
|
||||
{
|
||||
"default": IMAGE_CAPTIONER,
|
||||
"tooltip": "The hugging face name of the image captioning model to load.",
|
||||
},
|
||||
),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("LTXV_PROMPT_ENHANCER",)
|
||||
RETURN_NAMES = ("prompt_enhancer",)
|
||||
FUNCTION = "load"
|
||||
CATEGORY = "lightricks/LTXV"
|
||||
TITLE = "LTXV Prompt Enhancer (Down)Loader"
|
||||
OUTPUT_NODE = False
|
||||
DESCRIPTION = "Downloads and initializes LLM and image captioning models from Hugging Face to enhance text prompts for image generation."
|
||||
|
||||
def model_path_download_if_needed(self, model_name):
|
||||
model_directory = os.path.join(folder_paths.models_dir, MODELS_PATH_KEY)
|
||||
os.makedirs(model_directory, exist_ok=True)
|
||||
|
||||
model_name_ = model_name.rsplit("/", 1)[-1]
|
||||
model_path = os.path.join(model_directory, model_name_)
|
||||
|
||||
if not os.path.exists(model_path):
|
||||
from huggingface_hub import snapshot_download
|
||||
|
||||
try:
|
||||
snapshot_download(
|
||||
repo_id=model_name,
|
||||
local_dir=model_path,
|
||||
local_dir_use_symlinks=False,
|
||||
)
|
||||
except Exception:
|
||||
shutil.rmtree(model_path, ignore_errors=True)
|
||||
raise
|
||||
return model_path
|
||||
|
||||
def down_load_llm_model(self, llm_name, load_device):
|
||||
model_path = self.model_path_download_if_needed(llm_name)
|
||||
llm_model = AutoModelForCausalLM.from_pretrained(
|
||||
model_path,
|
||||
torch_dtype=torch.bfloat16,
|
||||
)
|
||||
|
||||
llm_tokenizer = AutoTokenizer.from_pretrained(
|
||||
model_path,
|
||||
)
|
||||
|
||||
return llm_model, llm_tokenizer
|
||||
|
||||
def down_load_image_captioner(self, image_captioner, load_device):
|
||||
model_path = self.model_path_download_if_needed(image_captioner)
|
||||
image_caption_model = AutoModelForCausalLM.from_pretrained(
|
||||
model_path, trust_remote_code=True
|
||||
)
|
||||
|
||||
image_caption_processor = AutoProcessor.from_pretrained(
|
||||
model_path, trust_remote_code=True
|
||||
)
|
||||
|
||||
return image_caption_model, image_caption_processor
|
||||
|
||||
def load(self, llm_name, image_captioner_name):
|
||||
load_device = comfy.model_management.get_torch_device()
|
||||
offload_device = comfy.model_management.vae_offload_device()
|
||||
llm_model, llm_tokenizer = self.down_load_llm_model(llm_name, load_device)
|
||||
image_caption_model, image_caption_processor = self.down_load_image_captioner(
|
||||
image_captioner_name, load_device
|
||||
)
|
||||
|
||||
enhancer = PromptEnhancer(
|
||||
image_caption_processor, image_caption_model, llm_model, llm_tokenizer
|
||||
)
|
||||
patcher = comfy.model_patcher.ModelPatcher(
|
||||
enhancer,
|
||||
load_device,
|
||||
offload_device,
|
||||
)
|
||||
return (patcher,)
|
||||
|
||||
|
||||
@comfy_node(name="LTXVPromptEnhancer")
|
||||
class LTXVPromptEnhancer:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"prompt": ("STRING",),
|
||||
"prompt_enhancer": ("LTXV_PROMPT_ENHANCER",),
|
||||
"max_resulting_tokens": (
|
||||
"INT",
|
||||
{"default": 256, "min": 32, "max": 512},
|
||||
),
|
||||
},
|
||||
"optional": {
|
||||
"image_prompt": ("IMAGE",),
|
||||
},
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("STRING",)
|
||||
RETURN_NAMES = ("str",)
|
||||
FUNCTION = "enhance"
|
||||
CATEGORY = "lightricks/LTXV"
|
||||
TITLE = "LTXV Prompt Enhancer"
|
||||
OUTPUT_NODE = False
|
||||
DESCRIPTION = (
|
||||
"Enhances text prompts for image generation using LLMs. "
|
||||
"Optionally incorporates reference images to create more contextually relevant descriptions."
|
||||
)
|
||||
|
||||
def enhance(
|
||||
self,
|
||||
prompt,
|
||||
prompt_enhancer: comfy.model_patcher.ModelPatcher,
|
||||
image_prompt: torch.Tensor = None,
|
||||
max_resulting_tokens=256,
|
||||
):
|
||||
comfy.model_management.free_memory(
|
||||
prompt_enhancer.memory_required([]),
|
||||
comfy.model_management.get_torch_device(),
|
||||
)
|
||||
comfy.model_management.load_model_gpu(prompt_enhancer)
|
||||
model = prompt_enhancer.model
|
||||
image_conditioning = None
|
||||
if image_prompt is not None:
|
||||
permuted_image = image_prompt.permute(3, 0, 1, 2)[None, :]
|
||||
image_conditioning = [(permuted_image, 0, 1.0)]
|
||||
|
||||
enhanced_prompt = model(prompt, image_conditioning, max_resulting_tokens)
|
||||
return (enhanced_prompt[0],)
|
||||
@@ -1,258 +0,0 @@
|
||||
import logging
|
||||
import random
|
||||
from typing import List, Optional, Tuple, Union
|
||||
|
||||
import torch
|
||||
from PIL import Image
|
||||
|
||||
logger = logging.getLogger(__name__) # pylint: disable=invalid-name
|
||||
|
||||
T2V_CINEMATIC_PROMPT = """You are an expert cinematic director with many award winning movies, When writing prompts based on the user input, focus on detailed, chronological descriptions of actions and scenes.
|
||||
Include specific movements, appearances, camera angles, and environmental details - all in a single flowing paragraph.
|
||||
Start directly with the action, and keep descriptions literal and precise.
|
||||
Think like a cinematographer describing a shot list.
|
||||
Do not change the user input intent, just enhance it.
|
||||
Keep within 150 words.
|
||||
For best results, build your prompts using this structure:
|
||||
Start with main action in a single sentence
|
||||
Add specific details about movements and gestures
|
||||
Describe character/object appearances precisely
|
||||
Include background and environment details
|
||||
Specify camera angles and movements
|
||||
Describe lighting and colors
|
||||
Note any changes or sudden events
|
||||
Do not exceed the 150 word limit!
|
||||
Output the enhanced prompt only.
|
||||
|
||||
Examples:
|
||||
user prompt: A man drives a toyota car.
|
||||
enhanced prompt: A person is driving a car on a two-lane road, holding the steering wheel with both hands. The person's hands are light-skinned and they are wearing a black long-sleeved shirt. The steering wheel has a Toyota logo in the center and black leather around it. The car's dashboard is visible, showing a speedometer, tachometer, and navigation screen. The road ahead is straight and there are trees and fields visible on either side. The camera is positioned inside the car, providing a view from the driver's perspective. The lighting is natural and overcast, with a slightly cool tone.
|
||||
|
||||
user prompt: A young woman is sitting on a chair.
|
||||
enhanced prompt: A young woman with dark, curly hair and pale skin sits on a chair; she wears a dark, intricately patterned dress with a high collar and long, dark gloves that extend past her elbows; the scene is dimly lit, with light streaming in from a large window behind the characters.
|
||||
|
||||
user prompt: Aerial view of a city skyline.
|
||||
enhanced prompt: The camera pans across a cityscape of tall buildings with a circular building in the center. The camera moves from left to right, showing the tops of the buildings and the circular building in the center. The buildings are various shades of gray and white, and the circular building has a green roof. The camera angle is high, looking down at the city. The lighting is bright, with the sun shining from the upper left, casting shadows from the buildings.
|
||||
"""
|
||||
|
||||
I2V_CINEMATIC_PROMPT = """You are an expert cinematic director with many award winning movies, When writing prompts based on the user input, focus on detailed, chronological descriptions of actions and scenes.
|
||||
Include specific movements, appearances, camera angles, and environmental details - all in a single flowing paragraph.
|
||||
Start directly with the action, and keep descriptions literal and precise.
|
||||
Think like a cinematographer describing a shot list.
|
||||
Keep within 150 words.
|
||||
For best results, build your prompts using this structure:
|
||||
Describe the image first and then add the user input. Image description should be in first priority! Align to the image caption if it contradicts the user text input.
|
||||
Start with main action in a single sentence
|
||||
Add specific details about movements and gestures
|
||||
Describe character/object appearances precisely
|
||||
Include background and environment details
|
||||
Specify camera angles and movements
|
||||
Describe lighting and colors
|
||||
Note any changes or sudden events
|
||||
Align to the image caption if it contradicts the user text input.
|
||||
Do not exceed the 150 word limit!
|
||||
Output the enhanced prompt only.
|
||||
"""
|
||||
|
||||
|
||||
def tensor_to_pil(tensor):
|
||||
# Ensure tensor is in range [-1, 1]
|
||||
assert tensor.min() >= -1 and tensor.max() <= 1
|
||||
|
||||
# Convert from [-1, 1] to [0, 1]
|
||||
tensor = (tensor + 1) / 2
|
||||
|
||||
# Rearrange from [C, H, W] to [H, W, C]
|
||||
tensor = tensor.permute(1, 2, 0)
|
||||
|
||||
# Convert to numpy array and then to uint8 range [0, 255]
|
||||
numpy_image = (tensor.cpu().numpy() * 255).astype("uint8")
|
||||
|
||||
# Convert to PIL Image
|
||||
return Image.fromarray(numpy_image)
|
||||
|
||||
|
||||
def generate_cinematic_prompt(
|
||||
image_caption_model,
|
||||
image_caption_processor,
|
||||
prompt_enhancer_model,
|
||||
prompt_enhancer_tokenizer,
|
||||
prompt: Union[str, List[str]],
|
||||
conditioning_items: Optional[List[Tuple[torch.Tensor, int, float]]] = None,
|
||||
max_new_tokens: int = 256,
|
||||
) -> List[str]:
|
||||
prompts = [prompt] if isinstance(prompt, str) else prompt
|
||||
|
||||
if conditioning_items is None:
|
||||
prompts = _generate_t2v_prompt(
|
||||
prompt_enhancer_model,
|
||||
prompt_enhancer_tokenizer,
|
||||
prompts,
|
||||
max_new_tokens,
|
||||
T2V_CINEMATIC_PROMPT,
|
||||
)
|
||||
else:
|
||||
# if len(conditioning_items) > 1 or conditioning_items[0][1] != 0:
|
||||
# logger.warning(
|
||||
# "prompt enhancement does only support first frame of conditioning items, returning original prompts"
|
||||
# )
|
||||
# return prompts
|
||||
|
||||
first_frame_conditioning_item = conditioning_items[0]
|
||||
first_frames = _get_first_frames_from_conditioning_item(
|
||||
first_frame_conditioning_item
|
||||
)
|
||||
|
||||
assert len(first_frames) == len(
|
||||
prompts
|
||||
), "Number of conditioning frames must match number of prompts"
|
||||
|
||||
prompts = _generate_i2v_prompt(
|
||||
image_caption_model,
|
||||
image_caption_processor,
|
||||
prompt_enhancer_model,
|
||||
prompt_enhancer_tokenizer,
|
||||
prompts,
|
||||
first_frames,
|
||||
max_new_tokens,
|
||||
I2V_CINEMATIC_PROMPT,
|
||||
)
|
||||
|
||||
return prompts
|
||||
|
||||
|
||||
def _get_first_frames_from_conditioning_item(
|
||||
conditioning_item: Tuple[torch.Tensor, int, float]
|
||||
) -> List[Image.Image]:
|
||||
frames_tensor = conditioning_item[0]
|
||||
# tensor shape: [batch_size, 3, num_frames, height, width], take first frame from each sample
|
||||
return [
|
||||
tensor_to_pil(frames_tensor[i, :, 0, :, :])
|
||||
for i in range(frames_tensor.shape[0])
|
||||
]
|
||||
|
||||
|
||||
def _generate_t2v_prompt(
|
||||
prompt_enhancer_model,
|
||||
prompt_enhancer_tokenizer,
|
||||
prompts: List[str],
|
||||
max_new_tokens: int,
|
||||
system_prompt: str,
|
||||
) -> List[str]:
|
||||
messages = [
|
||||
[
|
||||
{"role": "system", "content": system_prompt},
|
||||
{"role": "user", "content": f"user_prompt: {p}"},
|
||||
]
|
||||
for p in prompts
|
||||
]
|
||||
|
||||
texts = [
|
||||
prompt_enhancer_tokenizer.apply_chat_template(
|
||||
m, tokenize=False, add_generation_prompt=True
|
||||
)
|
||||
for m in messages
|
||||
]
|
||||
model_inputs = prompt_enhancer_tokenizer(texts, return_tensors="pt").to(
|
||||
prompt_enhancer_model.device
|
||||
)
|
||||
|
||||
return _generate_and_decode_prompts(
|
||||
prompt_enhancer_model, prompt_enhancer_tokenizer, model_inputs, max_new_tokens
|
||||
)
|
||||
|
||||
|
||||
def _generate_i2v_prompt(
|
||||
image_caption_model,
|
||||
image_caption_processor,
|
||||
prompt_enhancer_model,
|
||||
prompt_enhancer_tokenizer,
|
||||
prompts: List[str],
|
||||
first_frames: List[Image.Image],
|
||||
max_new_tokens: int,
|
||||
system_prompt: str,
|
||||
) -> List[str]:
|
||||
image_captions = _generate_image_captions(
|
||||
image_caption_model, image_caption_processor, first_frames
|
||||
)
|
||||
|
||||
messages = [
|
||||
[
|
||||
{"role": "system", "content": system_prompt},
|
||||
{"role": "user", "content": f"user_prompt: {p}\nimage_caption: {c}"},
|
||||
]
|
||||
for p, c in zip(prompts, image_captions)
|
||||
]
|
||||
|
||||
texts = [
|
||||
prompt_enhancer_tokenizer.apply_chat_template(
|
||||
m, tokenize=False, add_generation_prompt=True
|
||||
)
|
||||
for m in messages
|
||||
]
|
||||
model_inputs = prompt_enhancer_tokenizer(texts, return_tensors="pt").to(
|
||||
prompt_enhancer_model.device
|
||||
)
|
||||
|
||||
return _generate_and_decode_prompts(
|
||||
prompt_enhancer_model, prompt_enhancer_tokenizer, model_inputs, max_new_tokens
|
||||
)
|
||||
|
||||
|
||||
def _generate_image_captions(
|
||||
image_caption_model,
|
||||
image_caption_processor,
|
||||
images: List[Image.Image],
|
||||
system_prompt: str = "<DETAILED_CAPTION>",
|
||||
) -> List[str]:
|
||||
image_caption_prompts = [system_prompt] * len(images)
|
||||
inputs = image_caption_processor(
|
||||
image_caption_prompts, images, return_tensors="pt"
|
||||
).to(image_caption_model.device)
|
||||
|
||||
with torch.inference_mode():
|
||||
generated_ids = image_caption_model.generate(
|
||||
input_ids=inputs["input_ids"],
|
||||
pixel_values=inputs["pixel_values"],
|
||||
max_new_tokens=1024,
|
||||
do_sample=False,
|
||||
num_beams=3,
|
||||
)
|
||||
|
||||
return image_caption_processor.batch_decode(generated_ids, skip_special_tokens=True)
|
||||
|
||||
|
||||
def _get_random_scene_type():
|
||||
"""
|
||||
Randomly select a scene type to add to the prompt.
|
||||
"""
|
||||
types = [
|
||||
"The scene is captured in real-life footage.",
|
||||
"The scene is computer-generated imagery.",
|
||||
"The scene appears to be from a movie.",
|
||||
"The scene appears to be from a TV show.",
|
||||
"The scene is captured in a studio.",
|
||||
]
|
||||
return random.choice(types)
|
||||
|
||||
|
||||
def _generate_and_decode_prompts(
|
||||
prompt_enhancer_model, prompt_enhancer_tokenizer, model_inputs, max_new_tokens: int
|
||||
) -> List[str]:
|
||||
with torch.inference_mode():
|
||||
outputs = prompt_enhancer_model.generate(
|
||||
**model_inputs, max_new_tokens=max_new_tokens
|
||||
)
|
||||
generated_ids = [
|
||||
output_ids[len(input_ids) :]
|
||||
for input_ids, output_ids in zip(model_inputs.input_ids, outputs)
|
||||
]
|
||||
decoded_prompts = prompt_enhancer_tokenizer.batch_decode(
|
||||
generated_ids, skip_special_tokens=True
|
||||
)
|
||||
|
||||
decoded_prompts = [p + f" {_get_random_scene_type()}." for p in decoded_prompts]
|
||||
|
||||
print(decoded_prompts)
|
||||
|
||||
return decoded_prompts
|
||||
@@ -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,)
|
||||
|
||||
|
||||
+6
-1
@@ -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
|
||||
transformers[timm]>=4.50.0
|
||||
# Required for HDR EXR read/write (same as ltx-pipelines); no OpenCV EXR fallback.
|
||||
openimageio
|
||||
transformers>=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]]
|
||||
@@ -33,16 +33,6 @@ def stg(
|
||||
return noise_pred
|
||||
|
||||
|
||||
class MomentumBuffer:
|
||||
def __init__(self, momentum: float):
|
||||
self.momentum = momentum
|
||||
self.running_average = 0
|
||||
|
||||
def update(self, update_value: torch.Tensor):
|
||||
new_average = self.momentum * self.running_average
|
||||
self.running_average = update_value + new_average
|
||||
|
||||
|
||||
def project(v0: torch.Tensor, v1: torch.Tensor):
|
||||
dtype = v0.dtype
|
||||
v0, v1 = v0.double(), v1.double()
|
||||
@@ -792,153 +782,3 @@ class STGAdvancedPresetsNode:
|
||||
|
||||
def get_preset(self, preset=None):
|
||||
return (preset,)
|
||||
|
||||
|
||||
class APGGuider(comfy.samplers.CFGGuider):
|
||||
def __init__(
|
||||
self,
|
||||
model: ModelPatcher,
|
||||
cfg_scale,
|
||||
eta,
|
||||
norm_threshold,
|
||||
momentum_coefficient: float = -0.9,
|
||||
):
|
||||
self.model = model.clone()
|
||||
super().__init__(self.model)
|
||||
self.momentum_coefficient = momentum_coefficient
|
||||
self.momentum_buffer = MomentumBuffer(self.momentum_coefficient)
|
||||
self.eta = eta
|
||||
self.norm_threshold = norm_threshold
|
||||
self.cfg_scale = cfg_scale
|
||||
self.previous_timestep = None
|
||||
|
||||
def set_conds(self, positive, negative):
|
||||
self.inner_set_conds({"positive": positive, "negative": negative})
|
||||
|
||||
def predict_noise(
|
||||
self,
|
||||
x: torch.Tensor,
|
||||
timestep: torch.Tensor,
|
||||
model_options: dict = {},
|
||||
seed=None,
|
||||
):
|
||||
# CFG zero init - skipping steps with timestep bigger than given threshold.
|
||||
|
||||
# in CFGGuider.predict_noise, we call sampling_function(), which uses cfg_function() to compute pos & neg
|
||||
# but we'd rather do a single batch of sampling pos, neg, and perturbed, so we call calc_cond_batch([perturbed,pos,neg]) directly
|
||||
positive_cond = self.conds.get("positive", None)
|
||||
negative_cond = self.conds.get("negative", None)
|
||||
|
||||
noise_pred_pos = comfy.samplers.calc_cond_batch(
|
||||
self.inner_model,
|
||||
[positive_cond],
|
||||
x,
|
||||
timestep,
|
||||
model_options,
|
||||
)[0]
|
||||
|
||||
if (
|
||||
self.previous_timestep is not None
|
||||
and timestep.item() > self.previous_timestep
|
||||
):
|
||||
print("Resetting momentum buffer")
|
||||
self.momentum_buffer = MomentumBuffer(self.momentum_coefficient)
|
||||
|
||||
noise_pred_neg = 0
|
||||
if not math.isclose(self.cfg_scale, 1.0):
|
||||
noise_pred_neg = comfy.samplers.calc_cond_batch(
|
||||
self.inner_model,
|
||||
[negative_cond],
|
||||
x,
|
||||
timestep,
|
||||
model_options,
|
||||
)[0]
|
||||
|
||||
apg_result = apg(
|
||||
noise_pred_pos,
|
||||
noise_pred_neg,
|
||||
self.cfg_scale,
|
||||
self.eta,
|
||||
self.norm_threshold,
|
||||
)
|
||||
|
||||
# normally this would be done in cfg_function, but we skipped
|
||||
# that for efficiency: we can compute the noise predictions in
|
||||
# a single call to calc_cond_batch() (rather than two)
|
||||
# so we replicate the hook here
|
||||
for fn in model_options.get("sampler_post_cfg_function", []):
|
||||
args = {
|
||||
"denoised": apg_result,
|
||||
"cond": positive_cond,
|
||||
"uncond": negative_cond,
|
||||
"model": self.inner_model,
|
||||
"uncond_denoised": noise_pred_neg,
|
||||
"cond_denoised": noise_pred_pos,
|
||||
"sigma": timestep,
|
||||
"model_options": model_options,
|
||||
"input": x,
|
||||
}
|
||||
apg_result = fn(args)
|
||||
|
||||
self.previous_timestep = timestep.item()
|
||||
return apg_result
|
||||
|
||||
|
||||
@comfy_node(name="APGGuider")
|
||||
class APGGuiderNode:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"model": ("MODEL",),
|
||||
"positive": ("CONDITIONING",),
|
||||
"negative": ("CONDITIONING",),
|
||||
"cfg_scale": (
|
||||
"FLOAT",
|
||||
{"default": 1.0, "min": 0.0, "max": 100.0, "step": 0.01},
|
||||
),
|
||||
"eta": (
|
||||
"FLOAT",
|
||||
{"default": 1.0, "min": 0.0, "max": 100.0, "step": 0.01},
|
||||
),
|
||||
"norm_threshold": (
|
||||
"FLOAT",
|
||||
{"default": 0.0, "min": 0.0, "max": 100.0, "step": 0.01},
|
||||
),
|
||||
"momentum_coefficient": (
|
||||
"FLOAT",
|
||||
{"default": -0.9, "min": -3.0, "max": 3.0, "step": 0.01},
|
||||
),
|
||||
},
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("GUIDER",)
|
||||
|
||||
FUNCTION = "get_guider"
|
||||
CATEGORY = "lightricks/LTXV"
|
||||
|
||||
DESCRIPTION = """
|
||||
The APG Guider implements Adaptive Projected Guidance (APG).
|
||||
Reference: https://arxiv.org/abs/2410.02416.
|
||||
"""
|
||||
|
||||
def get_guider(
|
||||
self,
|
||||
model,
|
||||
positive,
|
||||
negative,
|
||||
cfg_scale,
|
||||
eta,
|
||||
norm_threshold,
|
||||
momentum_coefficient,
|
||||
):
|
||||
guider = APGGuider(
|
||||
model,
|
||||
cfg_scale,
|
||||
eta,
|
||||
norm_threshold,
|
||||
momentum_coefficient,
|
||||
)
|
||||
guider.set_conds(positive, negative)
|
||||
guider.raw_conds = (positive, negative)
|
||||
return (guider,)
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -1,54 +0,0 @@
|
||||
from .nodes.attn_bank_nodes import (
|
||||
LTXAttentionBankNode,
|
||||
LTXAttentioOverrideNode,
|
||||
LTXPrepareAttnInjectionsNode,
|
||||
)
|
||||
from .nodes.attn_override_node import LTXAttnOverrideNode
|
||||
from .nodes.latent_guide_node import AddLatentGuideNode
|
||||
from .nodes.ltx_feta_enhance_node import LTXFetaEnhanceNode
|
||||
from .nodes.ltx_flowedit_nodes import LTXFlowEditCFGGuiderNode, LTXFlowEditSamplerNode
|
||||
from .nodes.ltx_inverse_model_pred_nodes import (
|
||||
LTXForwardModelSamplingPredNode,
|
||||
LTXReverseModelSamplingPredNode,
|
||||
)
|
||||
from .nodes.ltx_pag_node import LTXPerturbedAttentionNode
|
||||
from .nodes.modify_ltx_model_node import ModifyLTXModelNode
|
||||
from .nodes.rectified_sampler_nodes import (
|
||||
LTXRFForwardODESamplerNode,
|
||||
LTXRFReverseODESamplerNode,
|
||||
)
|
||||
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
"ModifyLTXModel": ModifyLTXModelNode,
|
||||
"AddLatentGuide": AddLatentGuideNode,
|
||||
"LTXForwardModelSamplingPred": LTXForwardModelSamplingPredNode,
|
||||
"LTXReverseModelSamplingPred": LTXReverseModelSamplingPredNode,
|
||||
"LTXRFForwardODESampler": LTXRFForwardODESamplerNode,
|
||||
"LTXRFReverseODESampler": LTXRFReverseODESamplerNode,
|
||||
"LTXAttentionBank": LTXAttentionBankNode,
|
||||
"LTXPrepareAttnInjections": LTXPrepareAttnInjectionsNode,
|
||||
"LTXAttentioOverride": LTXAttentioOverrideNode,
|
||||
"LTXPerturbedAttention": LTXPerturbedAttentionNode,
|
||||
"LTXAttnOverride": LTXAttnOverrideNode,
|
||||
"LTXFlowEditCFGGuider": LTXFlowEditCFGGuiderNode,
|
||||
"LTXFlowEditSampler": LTXFlowEditSamplerNode,
|
||||
"LTXFetaEnhance": LTXFetaEnhanceNode,
|
||||
}
|
||||
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"ModifyLTXModel": "Modify LTX Model",
|
||||
"AddLatentGuide": "Add LTX Latent Guide",
|
||||
"LTXAddImageGuide": "Add LTX Image Guide",
|
||||
"LTXForwardModelSamplingPred": "LTX Forward Model Pred",
|
||||
"LTXReverseModelSamplingPred": "LTX Reverse Model Pred",
|
||||
"LTXRFForwardODESampler": "LTX Rf-Inv Forward Sampler",
|
||||
"LTXRFReverseODESampler": "LTX Rf-Inv Reverse Sampler",
|
||||
"LTXAttentionBank": "LTX Attention Bank",
|
||||
"LTXPrepareAttnInjections": "LTX Prepare Attn Injection",
|
||||
"LTXAttentioOverride": "LTX Attn Block Override",
|
||||
"LTXPerturbedAttention": "LTX Apply Perturbed Attention",
|
||||
"LTXAttnOverride": "LTX Attention Override",
|
||||
"LTXFlowEditCFGGuider": "LTX Flow Edit CFG Guider",
|
||||
"LTXFlowEditSampler": "LTX Flow Edit Sampler",
|
||||
"LTXFetaEnhance": "LTX Feta Enhance",
|
||||
}
|
||||
@@ -1,257 +0,0 @@
|
||||
import math
|
||||
|
||||
import comfy.ldm.common_dit
|
||||
import comfy.ldm.modules.attention
|
||||
import torch
|
||||
from comfy.ldm.lightricks.model import (
|
||||
BasicTransformerBlock,
|
||||
LTXVModel,
|
||||
apply_rotary_emb,
|
||||
)
|
||||
from comfy.ldm.lightricks.symmetric_patchifier import latent_to_pixel_coords
|
||||
from torch import nn
|
||||
|
||||
from ..utils.feta_enhance_utils import get_feta_scores
|
||||
|
||||
|
||||
class LTXModifiedCrossAttention(nn.Module):
|
||||
def forward(self, x, context=None, mask=None, pe=None, transformer_options={}):
|
||||
context = x if context is None else context
|
||||
context_v = x if context is None else context
|
||||
|
||||
step = transformer_options.get("step", -1)
|
||||
total_steps = transformer_options.get("total_steps", 0)
|
||||
attn_bank = transformer_options.get("attn_bank", None)
|
||||
sample_mode = transformer_options.get("sample_mode", None)
|
||||
if attn_bank is not None and self.idx in attn_bank["block_map"]:
|
||||
len_conds = len(transformer_options["cond_or_uncond"])
|
||||
pred_order = transformer_options["pred_order"]
|
||||
if (
|
||||
sample_mode == "forward"
|
||||
and total_steps - step - 1 < attn_bank["save_steps"]
|
||||
):
|
||||
step_idx = f"{pred_order}_{total_steps-step-1}"
|
||||
attn_bank["block_map"][self.idx][step_idx] = x.cpu()
|
||||
elif sample_mode == "reverse" and step < attn_bank["inject_steps"]:
|
||||
step_idx = f"{pred_order}_{step}"
|
||||
inject_settings = attn_bank.get("inject_settings", {})
|
||||
if len(inject_settings) > 0:
|
||||
inj = (
|
||||
attn_bank["block_map"][self.idx][step_idx]
|
||||
.to(x.device)
|
||||
.repeat(len_conds, 1, 1)
|
||||
)
|
||||
if "q" in inject_settings:
|
||||
x = inj
|
||||
if "k" in inject_settings:
|
||||
context = inj
|
||||
if "v" in inject_settings:
|
||||
context_v = inj
|
||||
|
||||
q = self.to_q(x)
|
||||
k = self.to_k(context)
|
||||
v = self.to_v(context_v)
|
||||
|
||||
q = self.q_norm(q)
|
||||
k = self.k_norm(k)
|
||||
|
||||
if pe is not None:
|
||||
q = apply_rotary_emb(q, pe)
|
||||
k = apply_rotary_emb(k, pe)
|
||||
|
||||
feta_score = None
|
||||
if (
|
||||
transformer_options.get("feta_weight", 0) > 0
|
||||
and self.idx in transformer_options["feta_layers"]["layers"]
|
||||
):
|
||||
feta_score = get_feta_scores(q, k, self.heads, transformer_options)
|
||||
|
||||
alt_attn_fn = (
|
||||
transformer_options.get("patches_replace", {})
|
||||
.get("layer", {})
|
||||
.get(("self_attn", self.idx), None)
|
||||
)
|
||||
if alt_attn_fn is not None:
|
||||
out = alt_attn_fn(
|
||||
q,
|
||||
k,
|
||||
v,
|
||||
self.heads,
|
||||
attn_precision=self.attn_precision,
|
||||
transformer_options=transformer_options,
|
||||
)
|
||||
elif mask is None:
|
||||
out = comfy.ldm.modules.attention.optimized_attention(
|
||||
q, k, v, self.heads, attn_precision=self.attn_precision
|
||||
)
|
||||
else:
|
||||
out = comfy.ldm.modules.attention.optimized_attention_masked(
|
||||
q, k, v, self.heads, mask, attn_precision=self.attn_precision
|
||||
)
|
||||
|
||||
if feta_score is not None:
|
||||
out *= feta_score
|
||||
|
||||
return self.to_out(out)
|
||||
|
||||
|
||||
class LTXModifiedBasicTransformerBlock(BasicTransformerBlock):
|
||||
def forward(
|
||||
self,
|
||||
x,
|
||||
context=None,
|
||||
attention_mask=None,
|
||||
timestep=None,
|
||||
pe=None,
|
||||
transformer_options={},
|
||||
):
|
||||
shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = (
|
||||
self.scale_shift_table[None, None].to(device=x.device, dtype=x.dtype)
|
||||
+ timestep.reshape(
|
||||
x.shape[0], timestep.shape[1], self.scale_shift_table.shape[0], -1
|
||||
)
|
||||
).unbind(dim=2)
|
||||
x += (
|
||||
self.attn1(
|
||||
comfy.ldm.common_dit.rms_norm(x) * (1 + scale_msa) + shift_msa,
|
||||
pe=pe,
|
||||
transformer_options=transformer_options,
|
||||
)
|
||||
* gate_msa
|
||||
)
|
||||
|
||||
x += self.attn2(x, context=context, mask=attention_mask)
|
||||
|
||||
y = comfy.ldm.common_dit.rms_norm(x) * (1 + scale_mlp) + shift_mlp
|
||||
x += self.ff(y) * gate_mlp
|
||||
|
||||
return x
|
||||
|
||||
|
||||
class LTXVModelModified(LTXVModel):
|
||||
|
||||
def forward(
|
||||
self,
|
||||
x,
|
||||
timestep,
|
||||
context,
|
||||
attention_mask,
|
||||
frame_rate=25,
|
||||
transformer_options={},
|
||||
keyframe_idxs=None,
|
||||
**kwargs,
|
||||
):
|
||||
patches_replace = transformer_options.get("patches_replace", {})
|
||||
|
||||
orig_shape = list(x.shape)
|
||||
|
||||
x, latent_coords = self.patchifier.patchify(x)
|
||||
pixel_coords = latent_to_pixel_coords(
|
||||
latent_coords=latent_coords,
|
||||
scale_factors=self.vae_scale_factors,
|
||||
causal_fix=self.causal_temporal_positioning,
|
||||
)
|
||||
|
||||
if keyframe_idxs is not None:
|
||||
pixel_coords[:, :, -keyframe_idxs.shape[2] :] = keyframe_idxs
|
||||
|
||||
fractional_coords = pixel_coords.to(torch.float32)
|
||||
fractional_coords[:, 0] = fractional_coords[:, 0] * (1.0 / frame_rate)
|
||||
|
||||
x = self.patchify_proj(x)
|
||||
timestep = timestep * 1000.0
|
||||
|
||||
if attention_mask is not None and not torch.is_floating_point(attention_mask):
|
||||
attention_mask = (attention_mask - 1).to(x.dtype).reshape(
|
||||
(attention_mask.shape[0], 1, -1, attention_mask.shape[-1])
|
||||
) * torch.finfo(x.dtype).max
|
||||
|
||||
pe = self._precompute_freqs_cis(
|
||||
fractional_coords, dim=self.inner_dim, out_dtype=x.dtype
|
||||
)
|
||||
|
||||
batch_size = x.shape[0]
|
||||
timestep, embedded_timestep = self.adaln_single(
|
||||
timestep.flatten(),
|
||||
{"resolution": None, "aspect_ratio": None},
|
||||
batch_size=batch_size,
|
||||
hidden_dtype=x.dtype,
|
||||
)
|
||||
# Second dimension is 1 or number of tokens (if timestep_per_token)
|
||||
timestep = timestep.view(batch_size, -1, timestep.shape[-1])
|
||||
embedded_timestep = embedded_timestep.view(
|
||||
batch_size, -1, embedded_timestep.shape[-1]
|
||||
)
|
||||
|
||||
# 2. Blocks
|
||||
if self.caption_projection is not None:
|
||||
batch_size = x.shape[0]
|
||||
context = self.caption_projection(context)
|
||||
context = context.view(batch_size, -1, x.shape[-1])
|
||||
|
||||
blocks_replace = patches_replace.get("dit", {})
|
||||
for i, block in enumerate(self.transformer_blocks):
|
||||
if ("double_block", i) in blocks_replace:
|
||||
|
||||
def block_wrap(args):
|
||||
out = {}
|
||||
out["img"] = block(
|
||||
args["img"],
|
||||
context=args["txt"],
|
||||
attention_mask=args["attention_mask"],
|
||||
timestep=args["vec"],
|
||||
pe=args["pe"],
|
||||
)
|
||||
return out
|
||||
|
||||
out = blocks_replace[("double_block", i)](
|
||||
{
|
||||
"img": x,
|
||||
"txt": context,
|
||||
"attention_mask": attention_mask,
|
||||
"vec": timestep,
|
||||
"pe": pe,
|
||||
},
|
||||
{"original_block": block_wrap},
|
||||
)
|
||||
x = out["img"]
|
||||
else:
|
||||
x = block(
|
||||
x,
|
||||
context=context,
|
||||
attention_mask=attention_mask,
|
||||
timestep=timestep,
|
||||
pe=pe,
|
||||
transformer_options=transformer_options,
|
||||
)
|
||||
|
||||
# 3. Output
|
||||
scale_shift_values = (
|
||||
self.scale_shift_table[None, None].to(device=x.device, dtype=x.dtype)
|
||||
+ embedded_timestep[:, :, None]
|
||||
)
|
||||
shift, scale = scale_shift_values[:, :, 0], scale_shift_values[:, :, 1]
|
||||
x = self.norm_out(x)
|
||||
# Modulation
|
||||
x = x * (1 + scale) + shift
|
||||
x = self.proj_out(x)
|
||||
|
||||
x = self.patchifier.unpatchify(
|
||||
latents=x,
|
||||
output_height=orig_shape[3],
|
||||
output_width=orig_shape[4],
|
||||
output_num_frames=orig_shape[2],
|
||||
out_channels=orig_shape[1] // math.prod(self.patchifier.patch_size),
|
||||
)
|
||||
|
||||
return x
|
||||
|
||||
|
||||
def inject_model(diffusion_model):
|
||||
diffusion_model.__class__ = LTXVModelModified
|
||||
for idx, transformer_block in enumerate(diffusion_model.transformer_blocks):
|
||||
transformer_block.__class__ = LTXModifiedBasicTransformerBlock
|
||||
transformer_block.idx = idx
|
||||
transformer_block.attn1.__class__ = LTXModifiedCrossAttention
|
||||
transformer_block.attn1.idx = idx
|
||||
return diffusion_model
|
||||
@@ -1,90 +0,0 @@
|
||||
from ..utils.attn_bank import AttentionBank
|
||||
|
||||
|
||||
class LTXAttentionBankNode:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"save_steps": ("INT", {"default": 0, "min": 0, "max": 1000, "step": 1}),
|
||||
"blocks": ("STRING", {"multiline": True}),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("ATTN_BANK",)
|
||||
FUNCTION = "build"
|
||||
|
||||
CATEGORY = "ltxtricks"
|
||||
|
||||
def build(self, save_steps, blocks=""):
|
||||
block_map = {}
|
||||
block_list = blocks.split(",")
|
||||
for block in block_list:
|
||||
block_idx = int(block)
|
||||
block_map[block_idx] = {}
|
||||
|
||||
bank = AttentionBank(save_steps, block_map)
|
||||
return (bank,)
|
||||
|
||||
|
||||
class LTXPrepareAttnInjectionsNode:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"latent": ("LATENT",),
|
||||
"attn_bank": ("ATTN_BANK",),
|
||||
"query": ("BOOLEAN", {"default": False}),
|
||||
"key": ("BOOLEAN", {"default": False}),
|
||||
"value": ("BOOLEAN", {"default": False}),
|
||||
"inject_steps": (
|
||||
"INT",
|
||||
{"default": 0, "min": 0, "max": 1000, "step": 1},
|
||||
),
|
||||
},
|
||||
"optional": {"blocks": ("LTX_BLOCKS",)},
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("LATENT", "ATTN_INJ")
|
||||
FUNCTION = "prepare"
|
||||
|
||||
CATEGORY = "fluxtapoz"
|
||||
|
||||
def prepare(self, latent, attn_bank, query, key, value, inject_steps, blocks=None):
|
||||
if inject_steps > attn_bank["save_steps"]:
|
||||
raise ValueError("Can not inject more steps than were saved.")
|
||||
attn_bank = AttentionBank(
|
||||
attn_bank["save_steps"], attn_bank["block_map"], inject_steps
|
||||
)
|
||||
attn_bank["inject_settings"] = set([])
|
||||
if query:
|
||||
attn_bank["inject_settings"].add("q")
|
||||
if key:
|
||||
attn_bank["inject_settings"].add("k")
|
||||
if value:
|
||||
attn_bank["inject_settings"].add("v")
|
||||
|
||||
if blocks is not None:
|
||||
attn_bank["block_map"] = {**attn_bank["block_map"]}
|
||||
for key in list(attn_bank["block_map"].keys()):
|
||||
if key not in blocks:
|
||||
del attn_bank["block_map"][key]
|
||||
|
||||
# Hack to force order of operations in ComfyUI graph
|
||||
return (latent, attn_bank)
|
||||
|
||||
|
||||
class LTXAttentioOverrideNode:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {"required": {"blocks": ("STRING", {"multiline": True})}}
|
||||
|
||||
RETURN_TYPES = ("LTX_BLOCKS",)
|
||||
FUNCTION = "build"
|
||||
|
||||
CATEGORY = "ltxtricks"
|
||||
|
||||
def build(self, blocks=""):
|
||||
block_set = set(list(int(block) for block in blocks.split(",")))
|
||||
|
||||
return (block_set,)
|
||||
@@ -1,30 +0,0 @@
|
||||
def is_integer(string):
|
||||
try:
|
||||
int(string)
|
||||
return True
|
||||
except ValueError:
|
||||
return False
|
||||
|
||||
|
||||
class LTXAttnOverrideNode:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"layers": ("STRING", {"multiline": True}),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("ATTN_OVERRIDE",)
|
||||
FUNCTION = "build"
|
||||
|
||||
CATEGORY = "ltxtricks/attn"
|
||||
|
||||
def build(self, layers):
|
||||
layers_map = set([])
|
||||
for block in layers.split(","):
|
||||
block = block.strip()
|
||||
if is_integer(block):
|
||||
layers_map.add(int(block))
|
||||
|
||||
return ({"layers": layers_map},)
|
||||
@@ -1,40 +0,0 @@
|
||||
import comfy_extras.nodes_lt as nodes_lt
|
||||
|
||||
|
||||
class AddLatentGuideNode(nodes_lt.LTXVAddGuide):
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"model": ("MODEL",),
|
||||
"latent": ("LATENT",),
|
||||
"image_latent": ("LATENT",),
|
||||
"index": ("INT", {"default": 0, "min": -1, "max": 9999, "step": 1}),
|
||||
"strength": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0}),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("MODEL", "LATENT")
|
||||
RETURN_NAMES = ("model", "latent")
|
||||
|
||||
CATEGORY = "ltxtricks"
|
||||
FUNCTION = "generate"
|
||||
|
||||
def generate(self, model, latent, image_latent, index, strength):
|
||||
noise_mask = nodes_lt.get_noise_mask(latent)
|
||||
latent = latent["samples"]
|
||||
|
||||
image_latent = image_latent["samples"]
|
||||
|
||||
latent, noise_mask = self.replace_latent_frames(
|
||||
latent,
|
||||
noise_mask,
|
||||
image_latent,
|
||||
index,
|
||||
strength,
|
||||
)
|
||||
|
||||
return (
|
||||
model,
|
||||
{"samples": latent, "noise_mask": noise_mask},
|
||||
)
|
||||
@@ -1,36 +0,0 @@
|
||||
DEFAULT_ATTN = {
|
||||
"layers": [i for i in range(0, 100, 1)],
|
||||
}
|
||||
|
||||
|
||||
class LTXFetaEnhanceNode:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"model": ("MODEL",),
|
||||
"feta_weight": (
|
||||
"FLOAT",
|
||||
{"default": 4, "min": -100.0, "max": 100.0, "step": 0.01},
|
||||
),
|
||||
},
|
||||
"optional": {"attn_override": ("ATTN_OVERRIDE",)},
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("MODEL",)
|
||||
|
||||
CATEGORY = "ltxtricks"
|
||||
FUNCTION = "apply"
|
||||
|
||||
def apply(self, model, feta_weight, attn_override=DEFAULT_ATTN):
|
||||
model = model.clone()
|
||||
|
||||
model_options = model.model_options.copy()
|
||||
transformer_options = model_options["transformer_options"].copy()
|
||||
|
||||
transformer_options["feta_weight"] = feta_weight
|
||||
transformer_options["feta_layers"] = attn_override
|
||||
model_options["transformer_options"] = transformer_options
|
||||
|
||||
model.model_options = model_options
|
||||
return (model,)
|
||||
@@ -1,180 +0,0 @@
|
||||
import torch
|
||||
from comfy.samplers import KSAMPLER, CFGGuider, sampling_function
|
||||
from tqdm import trange
|
||||
|
||||
|
||||
class FlowEditGuider(CFGGuider):
|
||||
def __init__(self, model_patcher):
|
||||
super().__init__(model_patcher)
|
||||
self.cfgs = {}
|
||||
|
||||
def set_conds(self, **kwargs):
|
||||
self.inner_set_conds(kwargs)
|
||||
|
||||
def set_cfgs(self, **kwargs):
|
||||
self.cfgs = {**kwargs}
|
||||
|
||||
def predict_noise(self, x, timestep, model_options={}, seed=None):
|
||||
latent_type = model_options["transformer_options"]["latent_type"]
|
||||
positive = self.conds.get(f"{latent_type}_positive", None)
|
||||
negative = self.conds.get(f"{latent_type}_negative", None)
|
||||
cfg = self.cfgs.get(latent_type, self.cfg)
|
||||
return sampling_function(
|
||||
self.inner_model,
|
||||
x,
|
||||
timestep,
|
||||
negative,
|
||||
positive,
|
||||
cfg,
|
||||
model_options=model_options,
|
||||
seed=seed,
|
||||
)
|
||||
|
||||
|
||||
class LTXFlowEditCFGGuiderNode:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"model": ("MODEL",),
|
||||
"source_pos": ("CONDITIONING",),
|
||||
"source_neg": ("CONDITIONING",),
|
||||
"target_pos": ("CONDITIONING",),
|
||||
"target_neg": ("CONDITIONING",),
|
||||
"source_cfg": (
|
||||
"FLOAT",
|
||||
{"default": 2, "min": 0, "max": 0xFFFFFFFFFFFFFFFF, "step": 0.01},
|
||||
),
|
||||
"target_cfg": (
|
||||
"FLOAT",
|
||||
{"default": 4.5, "min": 0, "max": 0xFFFFFFFFFFFFFFFF, "step": 0.01},
|
||||
),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("GUIDER",)
|
||||
|
||||
FUNCTION = "get_guider"
|
||||
CATEGORY = "ltxtricks"
|
||||
|
||||
def get_guider(
|
||||
self,
|
||||
model,
|
||||
source_pos,
|
||||
source_neg,
|
||||
target_pos,
|
||||
target_neg,
|
||||
source_cfg,
|
||||
target_cfg,
|
||||
):
|
||||
guider = FlowEditGuider(model)
|
||||
guider.set_conds(
|
||||
source_positive=source_pos,
|
||||
source_negative=source_neg,
|
||||
target_positive=target_pos,
|
||||
target_negative=target_neg,
|
||||
)
|
||||
guider.set_cfgs(source=source_cfg, target=target_cfg)
|
||||
return (guider,)
|
||||
|
||||
|
||||
def get_flowedit_sample(skip_steps, refine_steps, seed):
|
||||
generator = torch.manual_seed(seed)
|
||||
|
||||
@torch.no_grad()
|
||||
def flowedit_sample(
|
||||
model, x_init, sigmas, extra_args=None, callback=None, disable=None
|
||||
):
|
||||
extra_args = {} if extra_args is None else extra_args
|
||||
|
||||
model_options = extra_args.get("model_options", {})
|
||||
transformer_options = model_options.get("transformer_options", {})
|
||||
transformer_options = {**transformer_options}
|
||||
model_options["transformer_options"] = transformer_options
|
||||
extra_args["model_options"] = model_options
|
||||
denoise_mask = extra_args.get("denoise_mask", 1)
|
||||
if denoise_mask is None:
|
||||
denoise_mask = 1
|
||||
else:
|
||||
extra_args["denoise_mask"] = torch.ones_like(denoise_mask)
|
||||
|
||||
source_extra_args = {
|
||||
**extra_args,
|
||||
"model_options": {
|
||||
"transformer_options": {**transformer_options, "latent_type": "source"}
|
||||
},
|
||||
}
|
||||
|
||||
sigmas = sigmas[skip_steps:]
|
||||
|
||||
x_tgt = x_init.clone()
|
||||
N = len(sigmas) - 1
|
||||
s_in = x_init.new_ones([x_init.shape[0]])
|
||||
|
||||
for i in trange(N, disable=disable):
|
||||
sigma = sigmas[i]
|
||||
noise = torch.randn(x_init.shape, generator=generator).to(x_init.device)
|
||||
|
||||
zt_src = (1 - sigma) * x_init + sigma * noise
|
||||
|
||||
if i < N - refine_steps:
|
||||
zt_tgt = x_tgt + (zt_src - x_init) * denoise_mask
|
||||
transformer_options["latent_type"] = "source"
|
||||
source_extra_args["model_options"]["transformer_options"][
|
||||
"latent_type"
|
||||
] = "source"
|
||||
vt_src = model(zt_src, sigma * s_in, **source_extra_args)
|
||||
else:
|
||||
if i == N - refine_steps:
|
||||
x_tgt = x_tgt + (zt_src - x_init) * denoise_mask
|
||||
zt_tgt = x_tgt
|
||||
vt_src = 0
|
||||
|
||||
transformer_options["latent_type"] = "target"
|
||||
vt_tgt = model(zt_tgt, sigma * s_in, **extra_args)
|
||||
|
||||
v_delta = (vt_tgt - vt_src) * denoise_mask
|
||||
x_tgt += (sigmas[i + 1] - sigmas[i]) * v_delta
|
||||
|
||||
if callback is not None:
|
||||
callback(
|
||||
{
|
||||
"x": x_tgt,
|
||||
"denoised": x_tgt,
|
||||
"i": i + skip_steps,
|
||||
"sigma": sigmas[i],
|
||||
"sigma_hat": sigmas[i],
|
||||
}
|
||||
)
|
||||
|
||||
return x_tgt
|
||||
|
||||
return flowedit_sample
|
||||
|
||||
|
||||
class LTXFlowEditSamplerNode:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"skip_steps": (
|
||||
"INT",
|
||||
{"default": 4, "min": 0, "max": 0xFFFFFFFFFFFFFFFF},
|
||||
),
|
||||
"refine_steps": (
|
||||
"INT",
|
||||
{"default": 0, "min": 0, "max": 0xFFFFFFFFFFFFFFFF},
|
||||
),
|
||||
"seed": ("INT", {"default": 0, "min": 0, "max": 0xFFFFFFFFFFFFFFFF}),
|
||||
},
|
||||
"optional": {},
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("SAMPLER",)
|
||||
FUNCTION = "build"
|
||||
|
||||
CATEGORY = "ltxtricks"
|
||||
|
||||
def build(self, skip_steps, refine_steps, seed):
|
||||
sampler = KSAMPLER(get_flowedit_sample(skip_steps, refine_steps, seed))
|
||||
return (sampler,)
|
||||
@@ -1,91 +0,0 @@
|
||||
import comfy.latent_formats
|
||||
import comfy.model_sampling
|
||||
import comfy.sd
|
||||
|
||||
|
||||
class InverseCONST:
|
||||
def calculate_input(self, sigma, noise):
|
||||
return noise
|
||||
|
||||
def calculate_denoised(self, sigma, model_output, model_input):
|
||||
sigma = sigma.view(sigma.shape[:1] + (1,) * (model_output.ndim - 1))
|
||||
return model_output
|
||||
|
||||
def noise_scaling(self, sigma, noise, latent_image, max_denoise=False):
|
||||
return latent_image
|
||||
|
||||
def inverse_noise_scaling(self, sigma, latent):
|
||||
return latent
|
||||
|
||||
|
||||
class LTXForwardModelSamplingPredNode:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"model": ("MODEL",),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("MODEL",)
|
||||
FUNCTION = "patch"
|
||||
|
||||
CATEGORY = "ltxtricks"
|
||||
|
||||
def patch(self, model):
|
||||
m = model.clone()
|
||||
|
||||
sampling_base = comfy.model_sampling.ModelSamplingFlux
|
||||
sampling_type = InverseCONST
|
||||
|
||||
class ModelSamplingAdvanced(sampling_base, sampling_type):
|
||||
pass
|
||||
|
||||
model_sampling = ModelSamplingAdvanced(model.model.model_config)
|
||||
model_sampling.set_parameters(shift=1.15)
|
||||
m.add_object_patch("model_sampling", model_sampling)
|
||||
return (m,)
|
||||
|
||||
|
||||
class ReverseCONST:
|
||||
def calculate_input(self, sigma, noise):
|
||||
return noise
|
||||
|
||||
def calculate_denoised(self, sigma, model_output, model_input):
|
||||
sigma = sigma.view(sigma.shape[:1] + (1,) * (model_output.ndim - 1))
|
||||
return model_output # model_input - model_output * sigma
|
||||
|
||||
def noise_scaling(self, sigma, noise, latent_image, max_denoise=False):
|
||||
return latent_image
|
||||
|
||||
def inverse_noise_scaling(self, sigma, latent):
|
||||
return latent / (1.0 - sigma)
|
||||
|
||||
|
||||
class LTXReverseModelSamplingPredNode:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"model": ("MODEL",),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("MODEL",)
|
||||
FUNCTION = "patch"
|
||||
|
||||
CATEGORY = "ltxtricks"
|
||||
|
||||
def patch(self, model):
|
||||
m = model.clone()
|
||||
|
||||
sampling_base = comfy.model_sampling.ModelSamplingFlux
|
||||
sampling_type = ReverseCONST
|
||||
|
||||
class ModelSamplingAdvanced(sampling_base, sampling_type):
|
||||
pass
|
||||
|
||||
model_sampling = ModelSamplingAdvanced(model.model.model_config)
|
||||
model_sampling.set_parameters(shift=1.15)
|
||||
m.add_object_patch("model_sampling", model_sampling)
|
||||
return (m,)
|
||||
@@ -1,149 +0,0 @@
|
||||
import math
|
||||
|
||||
import comfy.model_patcher
|
||||
import comfy.samplers
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
from comfy.ldm.modules.attention import optimized_attention
|
||||
from einops import rearrange
|
||||
|
||||
DEFAULT_PAG_LTX = {"layers": set([14])}
|
||||
|
||||
|
||||
def gaussian_blur_2d(img, kernel_size, sigma):
|
||||
height = img.shape[-1]
|
||||
kernel_size = min(kernel_size, height - (height % 2 - 1))
|
||||
ksize_half = (kernel_size - 1) * 0.5
|
||||
|
||||
x = torch.linspace(-ksize_half, ksize_half, steps=kernel_size)
|
||||
|
||||
pdf = torch.exp(-0.5 * (x / sigma).pow(2))
|
||||
|
||||
x_kernel = pdf / pdf.sum()
|
||||
x_kernel = x_kernel.to(device=img.device, dtype=img.dtype)
|
||||
|
||||
kernel2d = torch.mm(x_kernel[:, None], x_kernel[None, :])
|
||||
kernel2d = kernel2d.expand(img.shape[-3], 1, kernel2d.shape[0], kernel2d.shape[1])
|
||||
|
||||
padding = [kernel_size // 2, kernel_size // 2, kernel_size // 2, kernel_size // 2]
|
||||
|
||||
img = F.pad(img, padding, mode="reflect")
|
||||
img = F.conv2d(img, kernel2d, groups=img.shape[-3])
|
||||
|
||||
return img
|
||||
|
||||
|
||||
class LTXPerturbedAttentionNode:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"model": ("MODEL",),
|
||||
"scale": (
|
||||
"FLOAT",
|
||||
{
|
||||
"default": 2.0,
|
||||
"min": 0.0,
|
||||
"max": 100.0,
|
||||
"step": 0.01,
|
||||
"round": 0.01,
|
||||
},
|
||||
),
|
||||
"rescale": (
|
||||
"FLOAT",
|
||||
{
|
||||
"default": 0.5,
|
||||
"min": 0.0,
|
||||
"max": 100.0,
|
||||
"step": 0.01,
|
||||
"round": 0.01,
|
||||
},
|
||||
),
|
||||
"cfg": (
|
||||
"FLOAT",
|
||||
{
|
||||
"default": 3.0,
|
||||
"min": 0.0,
|
||||
"max": 100.0,
|
||||
"step": 0.01,
|
||||
"round": 0.01,
|
||||
},
|
||||
),
|
||||
},
|
||||
"optional": {
|
||||
"attn_override": ("ATTN_OVERRIDE",),
|
||||
# "attn_type": (["PAG", "SEG"],),
|
||||
},
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("MODEL",)
|
||||
FUNCTION = "patch"
|
||||
|
||||
CATEGORY = "ltxtricks/attn"
|
||||
|
||||
def patch(
|
||||
self, model, scale, rescale, cfg, attn_override=DEFAULT_PAG_LTX, attn_type="PAG"
|
||||
):
|
||||
m = model.clone()
|
||||
|
||||
def pag_fn(q, k, v, heads, attn_precision=None, transformer_options=None):
|
||||
return v
|
||||
|
||||
def seg_fn(q, k, v, heads, attn_precision=None, transformer_options=None):
|
||||
_, sequence_length, _ = q.shape
|
||||
b, c, f, h, w = transformer_options["original_shape"]
|
||||
|
||||
q = rearrange(q, "b (f h w) d -> b (f d) w h", h=h, w=w)
|
||||
kernel_size = math.ceil(6 * scale) + 1 - math.ceil(6 * scale) % 2
|
||||
q = gaussian_blur_2d(q, kernel_size, scale)
|
||||
q = rearrange(q, "b (f d) w h -> b (f h w) d", f=f)
|
||||
return optimized_attention(q, k, v, heads, attn_precision=attn_precision)
|
||||
|
||||
def post_cfg_function(args):
|
||||
model = args["model"]
|
||||
|
||||
cond_pred = args["cond_denoised"]
|
||||
uncond_pred = args["uncond_denoised"]
|
||||
|
||||
len_conds = 1 if args.get("uncond", None) is None else 2
|
||||
|
||||
cond = args["cond"]
|
||||
sigma = args["sigma"]
|
||||
model_options = args["model_options"].copy()
|
||||
x = args["input"]
|
||||
|
||||
if scale == 0:
|
||||
if len_conds == 1:
|
||||
return cond_pred
|
||||
return uncond_pred + (cond_pred - uncond_pred)
|
||||
|
||||
attn_fn = pag_fn if attn_type == "PAG" else seg_fn
|
||||
for block_idx in attn_override["layers"]:
|
||||
model_options = comfy.model_patcher.set_model_options_patch_replace(
|
||||
model_options, attn_fn, "layer", "self_attn", int(block_idx)
|
||||
)
|
||||
|
||||
(perturbed,) = comfy.samplers.calc_cond_batch(
|
||||
model, [cond], x, sigma, model_options
|
||||
)
|
||||
|
||||
# if len_conds == 1:
|
||||
# output = cond_pred + scale * (cond_pred - pag)
|
||||
# else:
|
||||
# output = cond_pred + (scale-1.0) * (cond_pred - uncond_pred) + scale * (cond_pred - pag)
|
||||
|
||||
output = (
|
||||
uncond_pred
|
||||
+ cfg * (cond_pred - uncond_pred)
|
||||
+ scale * (cond_pred - perturbed)
|
||||
)
|
||||
if rescale > 0:
|
||||
factor = cond_pred.std() / output.std()
|
||||
factor = rescale * factor + (1 - rescale)
|
||||
output = output * factor
|
||||
|
||||
return output
|
||||
|
||||
m.set_model_sampler_post_cfg_function(post_cfg_function)
|
||||
|
||||
return (m,)
|
||||
@@ -1,20 +0,0 @@
|
||||
from ..modules.ltx_model import inject_model
|
||||
|
||||
|
||||
class ModifyLTXModelNode:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"model": ("MODEL",),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("MODEL",)
|
||||
|
||||
CATEGORY = "ltxtricks"
|
||||
FUNCTION = "modify"
|
||||
|
||||
def modify(self, model):
|
||||
model.model.diffusion_model = inject_model(model.model.diffusion_model)
|
||||
return (model,)
|
||||
@@ -1,275 +0,0 @@
|
||||
import torch
|
||||
from comfy.samplers import KSAMPLER
|
||||
from tqdm import trange
|
||||
|
||||
|
||||
def generate_trend_values(steps, start_time, end_time, eta, eta_trend):
|
||||
eta_values = [0] * steps
|
||||
|
||||
if eta_trend == "constant":
|
||||
for i in range(start_time, end_time):
|
||||
eta_values[i] = eta
|
||||
elif eta_trend == "linear_increase":
|
||||
for i in range(start_time, end_time):
|
||||
progress = (i - start_time) / (end_time - start_time - 1)
|
||||
eta_values[i] = eta * progress
|
||||
elif eta_trend == "linear_decrease":
|
||||
for i in range(start_time, end_time):
|
||||
progress = 1 - (i - start_time) / (end_time - start_time - 1)
|
||||
eta_values[i] = eta * progress
|
||||
|
||||
return eta_values
|
||||
|
||||
|
||||
def get_sample_forward(
|
||||
gamma, start_step, end_step, gamma_trend, seed, attn_bank=None, order="first"
|
||||
):
|
||||
# Controlled Forward ODE (Algorithm 1)
|
||||
generator = torch.Generator()
|
||||
generator.manual_seed(seed)
|
||||
|
||||
@torch.no_grad()
|
||||
def sample_forward(model, y0, sigmas, extra_args=None, callback=None, disable=None):
|
||||
if attn_bank is not None:
|
||||
for block_idx in attn_bank["block_map"]:
|
||||
attn_bank["block_map"][block_idx].clear()
|
||||
|
||||
extra_args = {} if extra_args is None else extra_args
|
||||
model_options = extra_args.get("model_options", {})
|
||||
model_options = {**model_options}
|
||||
transformer_options = model_options.get("transformer_options", {})
|
||||
transformer_options = {
|
||||
**transformer_options,
|
||||
"total_steps": len(sigmas) - 1,
|
||||
"sample_mode": "forward",
|
||||
"attn_bank": attn_bank,
|
||||
}
|
||||
model_options["transformer_options"] = transformer_options
|
||||
extra_args["model_options"] = model_options
|
||||
|
||||
Y = y0.clone()
|
||||
y1 = torch.randn(Y.shape, generator=generator).to(y0.device)
|
||||
N = len(sigmas) - 1
|
||||
s_in = y0.new_ones([y0.shape[0]])
|
||||
gamma_values = generate_trend_values(
|
||||
N, start_step, end_step, gamma, gamma_trend
|
||||
)
|
||||
for i in trange(N, disable=disable):
|
||||
transformer_options["step"] = i
|
||||
sigma = sigmas[i]
|
||||
sigma_next = sigmas[i + 1]
|
||||
t_i = model.inner_model.inner_model.model_sampling.timestep(sigmas[i])
|
||||
|
||||
conditional_vector_field = (y1 - Y) / (1 - t_i)
|
||||
|
||||
transformer_options["pred_order"] = "first"
|
||||
pred = model(
|
||||
Y, s_in * sigmas[i], **extra_args
|
||||
) # this implementation takes sigma instead of timestep
|
||||
|
||||
if order == "second":
|
||||
transformer_options["pred_order"] = "second"
|
||||
img_mid = Y + (sigma_next - sigma) / 2 * pred
|
||||
sigma_mid = sigma + (sigma_next - sigma) / 2
|
||||
pred_mid = model(img_mid, s_in * sigma_mid, **extra_args)
|
||||
|
||||
first_order = (pred_mid - pred) / ((sigma_next - sigma) / 2)
|
||||
pred = pred + gamma_values[i] * (conditional_vector_field - pred)
|
||||
# first_order = first_order + gamma_values[i] * (conditional_vector_field - first_order)
|
||||
Y = (
|
||||
Y
|
||||
+ (sigma_next - sigma) * pred
|
||||
+ 0.5 * (sigma_next - sigma) ** 2 * first_order
|
||||
)
|
||||
else:
|
||||
pred = pred + gamma_values[i] * (conditional_vector_field - pred)
|
||||
Y = Y + pred * (sigma_next - sigma)
|
||||
|
||||
if callback is not None:
|
||||
callback(
|
||||
{"x": Y, "denoised": Y, "i": i, "sigma": sigma, "sigma_hat": sigma}
|
||||
)
|
||||
|
||||
return Y
|
||||
|
||||
return sample_forward
|
||||
|
||||
|
||||
def get_sample_reverse(
|
||||
latent_image, eta, start_time, end_time, eta_trend, attn_bank=None, order="first"
|
||||
):
|
||||
# Controlled Reverse ODE (Algorithm 2)
|
||||
@torch.no_grad()
|
||||
def sample_reverse(model, y1, sigmas, extra_args=None, callback=None, disable=None):
|
||||
extra_args = {} if extra_args is None else extra_args
|
||||
model_options = extra_args.get("model_options", {})
|
||||
model_options = {**model_options}
|
||||
transformer_options = model_options.get("transformer_options", {})
|
||||
transformer_options = {
|
||||
**transformer_options,
|
||||
"total_steps": len(sigmas) - 1,
|
||||
"sample_mode": "reverse",
|
||||
"attn_bank": attn_bank,
|
||||
}
|
||||
model_options["transformer_options"] = transformer_options
|
||||
extra_args["model_options"] = model_options
|
||||
|
||||
X = y1.clone()
|
||||
N = len(sigmas) - 1
|
||||
y0 = latent_image.clone().to(y1.device)
|
||||
s_in = y0.new_ones([y0.shape[0]])
|
||||
eta_values = generate_trend_values(N, start_time, end_time, eta, eta_trend)
|
||||
for i in trange(N, disable=disable):
|
||||
transformer_options["step"] = i
|
||||
t_i = 1 - model.inner_model.inner_model.model_sampling.timestep(sigmas[i])
|
||||
sigma = sigmas[i]
|
||||
sigma_prev = sigmas[i + 1]
|
||||
|
||||
conditional_vector_field = (y0 - X) / (1 - t_i)
|
||||
|
||||
transformer_options["pred_order"] = "first"
|
||||
pred = model(
|
||||
X, sigma * s_in, **extra_args
|
||||
) # this implementation takes sigma instead of timestep
|
||||
|
||||
if order == "second":
|
||||
transformer_options["pred_order"] = "second"
|
||||
img_mid = X + (sigma_prev - sigma) / 2 * pred
|
||||
sigma_mid = sigma + (sigma_prev - sigma) / 2
|
||||
pred_mid = model(img_mid, s_in * sigma_mid, **extra_args)
|
||||
|
||||
first_order = (pred_mid - pred) / ((sigma_prev - sigma) / 2)
|
||||
pred = -pred + eta_values[i] * (conditional_vector_field + pred)
|
||||
|
||||
first_order = -first_order + eta_values[i] * (
|
||||
conditional_vector_field + first_order
|
||||
)
|
||||
X = (
|
||||
X
|
||||
+ (sigma - sigma_prev) * pred
|
||||
+ 0.5 * (sigma - sigma_prev) ** 2 * first_order
|
||||
)
|
||||
else:
|
||||
controlled_vector_field = -pred + eta_values[i] * (
|
||||
conditional_vector_field + pred
|
||||
)
|
||||
X = X + controlled_vector_field * (sigma - sigma_prev)
|
||||
|
||||
if callback is not None:
|
||||
callback(
|
||||
{
|
||||
"x": X,
|
||||
"denoised": X,
|
||||
"i": i,
|
||||
"sigma": sigmas[i],
|
||||
"sigma_hat": sigmas[i],
|
||||
}
|
||||
)
|
||||
|
||||
return X
|
||||
|
||||
return sample_reverse
|
||||
|
||||
|
||||
class LTXRFForwardODESamplerNode:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"gamma": (
|
||||
"FLOAT",
|
||||
{"default": 0.5, "min": 0.0, "max": 100.0, "step": 0.01},
|
||||
),
|
||||
"start_step": ("INT", {"default": 0, "min": 0, "max": 1000, "step": 1}),
|
||||
"end_step": ("INT", {"default": 5, "min": 0, "max": 1000, "step": 1}),
|
||||
"gamma_trend": (["linear_decrease", "linear_increase", "constant"],),
|
||||
},
|
||||
"optional": {
|
||||
"seed": ("INT", {"default": 0, "min": 0, "max": 0xFFFFFFFFFFFFFFFF}),
|
||||
"attn_bank": ("ATTN_BANK",),
|
||||
"order": (["first", "second"],),
|
||||
},
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("SAMPLER",)
|
||||
FUNCTION = "build"
|
||||
|
||||
CATEGORY = "ltxtricks"
|
||||
|
||||
def build(
|
||||
self,
|
||||
gamma,
|
||||
start_step,
|
||||
end_step,
|
||||
gamma_trend,
|
||||
seed=0,
|
||||
attn_bank=None,
|
||||
order="first",
|
||||
):
|
||||
sampler = KSAMPLER(
|
||||
get_sample_forward(
|
||||
gamma,
|
||||
start_step,
|
||||
end_step,
|
||||
gamma_trend,
|
||||
seed,
|
||||
attn_bank=attn_bank,
|
||||
order=order,
|
||||
)
|
||||
)
|
||||
|
||||
return (sampler,)
|
||||
|
||||
|
||||
class LTXRFReverseODESamplerNode:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"model": ("MODEL",),
|
||||
"latent_image": ("LATENT",),
|
||||
"eta": (
|
||||
"FLOAT",
|
||||
{"default": 0.8, "min": 0.0, "max": 100.0, "step": 0.01},
|
||||
),
|
||||
"start_step": ("INT", {"default": 0, "min": 0, "max": 1000, "step": 1}),
|
||||
"end_step": ("INT", {"default": 15, "min": 0, "max": 1000, "step": 1}),
|
||||
},
|
||||
"optional": {
|
||||
"eta_trend": (["linear_decrease", "linear_increase", "constant"],),
|
||||
"attn_inj": ("ATTN_INJ",),
|
||||
"order": (["first", "second"],),
|
||||
},
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("SAMPLER",)
|
||||
FUNCTION = "build"
|
||||
|
||||
CATEGORY = "ltxtricks"
|
||||
|
||||
def build(
|
||||
self,
|
||||
model,
|
||||
latent_image,
|
||||
eta,
|
||||
start_step,
|
||||
end_step,
|
||||
eta_trend="constant",
|
||||
attn_inj=None,
|
||||
order="first",
|
||||
):
|
||||
process_latent_in = model.get_model_object("process_latent_in")
|
||||
latent_image = process_latent_in(latent_image["samples"])
|
||||
sampler = KSAMPLER(
|
||||
get_sample_reverse(
|
||||
latent_image,
|
||||
eta,
|
||||
start_step,
|
||||
end_step,
|
||||
eta_trend,
|
||||
attn_bank=attn_inj,
|
||||
order=order,
|
||||
)
|
||||
)
|
||||
|
||||
return (sampler,)
|
||||
@@ -1,376 +0,0 @@
|
||||
import torch
|
||||
from comfy.samplers import KSAMPLER
|
||||
from tqdm import trange
|
||||
|
||||
|
||||
def get_sample_forward(attn_bank, save_steps, single_layers, double_layers):
|
||||
@torch.no_grad()
|
||||
def sample_forward(model, x, sigmas, extra_args=None, callback=None, disable=None):
|
||||
attn_bank.clear()
|
||||
attn_bank["save_steps"] = save_steps
|
||||
|
||||
extra_args = {} if extra_args is None else extra_args
|
||||
|
||||
model_options = extra_args.get("model_options", {})
|
||||
model_options = {**model_options}
|
||||
transformer_options = model_options.get("transformer_options", {})
|
||||
transformer_options = {**transformer_options}
|
||||
model_options["transformer_options"] = transformer_options
|
||||
extra_args["model_options"] = model_options
|
||||
|
||||
N = len(sigmas) - 1
|
||||
s_in = x.new_ones([x.shape[0]])
|
||||
for i in trange(N, disable=disable):
|
||||
sigma = sigmas[i]
|
||||
sigma_next = sigmas[i + 1]
|
||||
|
||||
if N - i - 1 < save_steps:
|
||||
attn_bank[N - i - 1] = {"first": {}, "mid": {}}
|
||||
|
||||
transformer_options["rfedit"] = {
|
||||
"step": N - i - 1,
|
||||
"process": "forward" if N - i - 1 < save_steps else None,
|
||||
"pred": "first",
|
||||
"bank": attn_bank,
|
||||
"single_layers": single_layers,
|
||||
"double_layers": double_layers,
|
||||
}
|
||||
|
||||
pred = model(x, s_in * sigma, **extra_args)
|
||||
|
||||
transformer_options["rfedit"] = {
|
||||
"step": N - i - 1,
|
||||
"process": "forward" if N - i - 1 < save_steps else None,
|
||||
"pred": "mid",
|
||||
"bank": attn_bank,
|
||||
"single_layers": single_layers,
|
||||
"double_layers": double_layers,
|
||||
}
|
||||
|
||||
img_mid = x + (sigma_next - sigma) / 2 * pred
|
||||
sigma_mid = sigma + (sigma_next - sigma) / 2
|
||||
pred_mid = model(img_mid, s_in * sigma_mid, **extra_args)
|
||||
|
||||
first_order = (pred_mid - pred) / ((sigma_next - sigma) / 2)
|
||||
x = (
|
||||
x
|
||||
+ (sigma_next - sigma) * pred
|
||||
+ 0.5 * (sigma_next - sigma) ** 2 * first_order
|
||||
)
|
||||
|
||||
if callback is not None:
|
||||
callback(
|
||||
{
|
||||
"x": x,
|
||||
"denoised": x,
|
||||
"i": i,
|
||||
"sigma": sigmas[i],
|
||||
"sigma_hat": sigmas[i],
|
||||
}
|
||||
)
|
||||
|
||||
return x
|
||||
|
||||
return sample_forward
|
||||
|
||||
|
||||
def get_sample_reverse(attn_bank, inject_steps, single_layers, double_layers):
|
||||
@torch.no_grad()
|
||||
def sample_reverse(model, x, sigmas, extra_args=None, callback=None, disable=None):
|
||||
if inject_steps > attn_bank["save_steps"]:
|
||||
raise ValueError(
|
||||
f'You must save at least as many steps as you want to inject. save_steps: {attn_bank["save_steps"]}, inject_steps: {inject_steps}'
|
||||
)
|
||||
|
||||
extra_args = {} if extra_args is None else extra_args
|
||||
|
||||
model_options = extra_args.get("model_options", {})
|
||||
model_options = {**model_options}
|
||||
transformer_options = model_options.get("transformer_options", {})
|
||||
transformer_options = {**transformer_options}
|
||||
model_options["transformer_options"] = transformer_options
|
||||
extra_args["model_options"] = model_options
|
||||
|
||||
N = len(sigmas) - 1
|
||||
s_in = x.new_ones([x.shape[0]])
|
||||
for i in trange(N, disable=disable):
|
||||
sigma = sigmas[i]
|
||||
sigma_prev = sigmas[i + 1]
|
||||
|
||||
transformer_options["rfedit"] = {
|
||||
"step": i,
|
||||
"process": "reverse" if i < inject_steps else None,
|
||||
"pred": "first",
|
||||
"bank": attn_bank,
|
||||
"single_layers": single_layers,
|
||||
"double_layers": double_layers,
|
||||
}
|
||||
|
||||
pred = model(x, s_in * sigma, **extra_args)
|
||||
|
||||
transformer_options["rfedit"] = {
|
||||
"step": i,
|
||||
"process": "reverse" if i < inject_steps else None,
|
||||
"pred": "mid",
|
||||
"bank": attn_bank,
|
||||
"single_layers": single_layers,
|
||||
"double_layers": double_layers,
|
||||
}
|
||||
|
||||
img_mid = x + (sigma_prev - sigma) / 2 * pred
|
||||
sigma_mid = sigma + (sigma_prev - sigma) / 2
|
||||
pred_mid = model(img_mid, s_in * sigma_mid, **extra_args)
|
||||
|
||||
first_order = (pred_mid - pred) / ((sigma_prev - sigma) / 2)
|
||||
x = (
|
||||
x
|
||||
+ (sigma_prev - sigma) * pred
|
||||
+ 0.5 * (sigma_prev - sigma) ** 2 * first_order
|
||||
)
|
||||
|
||||
if callback is not None:
|
||||
callback(
|
||||
{
|
||||
"x": x,
|
||||
"denoised": x,
|
||||
"i": i,
|
||||
"sigma": sigmas[i],
|
||||
"sigma_hat": sigmas[i],
|
||||
}
|
||||
)
|
||||
|
||||
return x
|
||||
|
||||
return sample_reverse
|
||||
|
||||
|
||||
DEFAULT_SINGLE_LAYERS = {}
|
||||
for i in range(38):
|
||||
DEFAULT_SINGLE_LAYERS[f"{i}"] = i > 19
|
||||
|
||||
DEFAULT_DOUBLE_LAYERS = {}
|
||||
for i in range(19):
|
||||
DEFAULT_DOUBLE_LAYERS[f"{i}"] = False
|
||||
|
||||
|
||||
class FlowEditForwardSamplerNode:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"save_steps": (
|
||||
"INT",
|
||||
{"default": 0, "min": 0, "max": 0xFFFFFFFFFFFFFFFF},
|
||||
),
|
||||
},
|
||||
"optional": {
|
||||
"single_layers": ("SINGLE_LAYERS",),
|
||||
"double_layers": ("DOUBLE_LAYERS",),
|
||||
},
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("SAMPLER", "ATTN_INJ")
|
||||
FUNCTION = "build"
|
||||
|
||||
CATEGORY = "fluxtapoz"
|
||||
|
||||
def build(
|
||||
self,
|
||||
save_steps,
|
||||
single_layers=DEFAULT_SINGLE_LAYERS,
|
||||
double_layers=DEFAULT_DOUBLE_LAYERS,
|
||||
):
|
||||
attn_bank = {}
|
||||
sampler = KSAMPLER(
|
||||
get_sample_forward(attn_bank, save_steps, single_layers, double_layers)
|
||||
)
|
||||
|
||||
return (sampler, attn_bank)
|
||||
|
||||
|
||||
# class FlowEditReverseSamplerNode:
|
||||
# @classmethod
|
||||
# def INPUT_TYPES(s):
|
||||
# return {
|
||||
# "required": {
|
||||
# "attn_inj": ("ATTN_INJ",),
|
||||
# "latent_image": ("LATENT",),
|
||||
# "eta": (
|
||||
# "FLOAT",
|
||||
# {"default": 0.8, "min": 0.0, "max": 100.0, "step": 0.01},
|
||||
# ),
|
||||
# "start_step": ("INT", {"default": 0, "min": 0, "max": 1000, "step": 1}),
|
||||
# "end_step": ("INT", {"default": 5, "min": 0, "max": 1000, "step": 1}),
|
||||
# },
|
||||
# "optional": {},
|
||||
# }
|
||||
#
|
||||
# RETURN_TYPES = ("SAMPLER",)
|
||||
# FUNCTION = "build"
|
||||
#
|
||||
# CATEGORY = "fluxtapoz"
|
||||
#
|
||||
# def build(self, latent_image, eta, start_step, end_step):
|
||||
# sampler = KSAMPLER(
|
||||
# get_sample_reverse(attn_inj, inject_steps, single_layers, double_layers)
|
||||
# )
|
||||
# return (sampler,)
|
||||
|
||||
|
||||
def get_sample_reverse2(attn_bank, inject_steps, single_layers, double_layers):
|
||||
@torch.no_grad()
|
||||
def sample_reverse(model, x, sigmas, extra_args=None, callback=None, disable=None):
|
||||
if inject_steps > attn_bank["save_steps"]:
|
||||
raise ValueError(
|
||||
f'You must save at least as many steps as you want to inject. save_steps: {attn_bank["save_steps"]}, inject_steps: {inject_steps}'
|
||||
)
|
||||
|
||||
extra_args = {} if extra_args is None else extra_args
|
||||
|
||||
model_options = extra_args.get("model_options", {})
|
||||
model_options = {**model_options}
|
||||
transformer_options = model_options.get("transformer_options", {})
|
||||
transformer_options = {**transformer_options}
|
||||
model_options["transformer_options"] = transformer_options
|
||||
extra_args["model_options"] = model_options
|
||||
|
||||
N = len(sigmas) - 1
|
||||
s_in = x.new_ones([x.shape[0]])
|
||||
for i in trange(N, disable=disable):
|
||||
sigma = sigmas[i]
|
||||
sigma_prev = sigmas[i + 1]
|
||||
|
||||
transformer_options["rfedit"] = {
|
||||
"step": i,
|
||||
"process": "reverse" if i < inject_steps else None,
|
||||
"pred": "first",
|
||||
"bank": attn_bank,
|
||||
"single_layers": single_layers,
|
||||
"double_layers": double_layers,
|
||||
}
|
||||
|
||||
pred = model(x, s_in * sigma, **extra_args)
|
||||
|
||||
transformer_options["rfedit"] = {
|
||||
"step": i,
|
||||
"process": "reverse" if i < inject_steps else None,
|
||||
"pred": "mid",
|
||||
"bank": attn_bank,
|
||||
"single_layers": single_layers,
|
||||
"double_layers": double_layers,
|
||||
}
|
||||
|
||||
img_mid = x + (sigma_prev - sigma) / 2 * pred
|
||||
sigma_mid = sigma + (sigma_prev - sigma) / 2
|
||||
pred_mid = model(img_mid, s_in * sigma_mid, **extra_args)
|
||||
|
||||
first_order = (pred_mid - pred) / ((sigma_prev - sigma) / 2)
|
||||
x = (
|
||||
x
|
||||
+ (sigma_prev - sigma) * pred
|
||||
+ 0.5 * (sigma_prev - sigma) ** 2 * first_order
|
||||
)
|
||||
|
||||
if callback is not None:
|
||||
callback(
|
||||
{
|
||||
"x": x,
|
||||
"denoised": x,
|
||||
"i": i,
|
||||
"sigma": sigmas[i],
|
||||
"sigma_hat": sigmas[i],
|
||||
}
|
||||
)
|
||||
|
||||
return x
|
||||
|
||||
return sample_reverse
|
||||
|
||||
|
||||
class FlowEdit2ReverseSamplerNode:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"attn_inj": ("ATTN_INJ",),
|
||||
"inject_steps": (
|
||||
"INT",
|
||||
{"default": 0, "min": 0, "max": 1000, "step": 1},
|
||||
),
|
||||
},
|
||||
"optional": {
|
||||
"single_layers": ("SINGLE_LAYERS",),
|
||||
"double_layers": ("DOUBLE_LAYERS",),
|
||||
},
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("SAMPLER",)
|
||||
FUNCTION = "build"
|
||||
|
||||
CATEGORY = "ltxtricks"
|
||||
|
||||
def build(
|
||||
self,
|
||||
attn_inj,
|
||||
inject_steps,
|
||||
single_layers=DEFAULT_SINGLE_LAYERS,
|
||||
double_layers=DEFAULT_DOUBLE_LAYERS,
|
||||
):
|
||||
sampler = KSAMPLER(
|
||||
get_sample_reverse(attn_inj, inject_steps, single_layers, double_layers)
|
||||
)
|
||||
return (sampler,)
|
||||
|
||||
|
||||
class PrepareAttnBankNode:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"latent": ("LATENT",),
|
||||
"attn_inj": ("ATTN_INJ",),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("LATENT", "ATTN_INJ")
|
||||
FUNCTION = "prepare"
|
||||
|
||||
CATEGORY = "ltxtricks"
|
||||
|
||||
def prepare(self, latent, attn_inj):
|
||||
# Hack to force order of operations in ComfyUI graph
|
||||
return (latent, attn_inj)
|
||||
|
||||
|
||||
class RFSingleBlocksOverrideNode:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
layers = {}
|
||||
for i in range(38):
|
||||
layers[f"{i}"] = ("BOOLEAN", {"default": i > 19})
|
||||
return {"required": layers}
|
||||
|
||||
RETURN_TYPES = ("SINGLE_LAYERS",)
|
||||
FUNCTION = "build"
|
||||
|
||||
CATEGORY = "ltxtricks"
|
||||
|
||||
def build(self, *args, **kwargs):
|
||||
return (kwargs,)
|
||||
|
||||
|
||||
class RFDoubleBlocksOverrideNode:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
layers = {}
|
||||
for i in range(19):
|
||||
layers[f"{i}"] = ("BOOLEAN", {"default": False})
|
||||
return {"required": layers}
|
||||
|
||||
RETURN_TYPES = ("DOUBLE_LAYERS",)
|
||||
FUNCTION = "build"
|
||||
|
||||
CATEGORY = "ltxtricks"
|
||||
|
||||
def build(self, *args, **kwargs):
|
||||
return (kwargs,)
|
||||
@@ -1,16 +0,0 @@
|
||||
class AttentionBank:
|
||||
def __init__(self, save_steps, block_map, inject_steps=None):
|
||||
self._data = {
|
||||
"save_steps": save_steps,
|
||||
"block_map": block_map,
|
||||
"inject_steps": inject_steps,
|
||||
}
|
||||
|
||||
def __getitem__(self, key):
|
||||
return self._data[key]
|
||||
|
||||
def __setitem__(self, key, value):
|
||||
self._data[key] = value
|
||||
|
||||
def get(self, key, default=None):
|
||||
return self._data.get(key, default)
|
||||
@@ -1,55 +0,0 @@
|
||||
import torch
|
||||
from einops import rearrange
|
||||
|
||||
|
||||
def _feta_score(query_image, key_image, head_dim, num_frames, enhance_weight):
|
||||
scale = head_dim**-0.5
|
||||
query_image = query_image * scale
|
||||
attn_temp = query_image @ key_image.transpose(-2, -1) # translate attn to float32
|
||||
attn_temp = attn_temp.to(torch.float32)
|
||||
attn_temp = attn_temp.softmax(dim=-1)
|
||||
|
||||
# Reshape to [batch_size * num_tokens, num_frames, num_frames]
|
||||
attn_temp = attn_temp.reshape(-1, num_frames, num_frames)
|
||||
|
||||
# Create a mask for diagonal elements
|
||||
diag_mask = torch.eye(num_frames, device=attn_temp.device).bool()
|
||||
diag_mask = diag_mask.unsqueeze(0).expand(attn_temp.shape[0], -1, -1)
|
||||
|
||||
# Zero out diagonal elements
|
||||
attn_wo_diag = attn_temp.masked_fill(diag_mask, 0)
|
||||
|
||||
# Calculate mean for each token's attention matrix
|
||||
# Number of off-diagonal elements per matrix is n*n - n
|
||||
num_off_diag = num_frames * num_frames - num_frames
|
||||
mean_scores = attn_wo_diag.sum(dim=(1, 2)) / num_off_diag
|
||||
|
||||
enhance_scores = mean_scores.mean() * (num_frames + enhance_weight)
|
||||
enhance_scores = enhance_scores.clamp(min=1)
|
||||
return enhance_scores
|
||||
|
||||
|
||||
def get_feta_scores(img_q, img_k, num_heads, transformer_options):
|
||||
num_frames = transformer_options["original_shape"][2]
|
||||
_, ST, dim = img_q.shape
|
||||
head_dim = dim // num_heads
|
||||
spatial_dim = ST // num_frames
|
||||
|
||||
query_image = rearrange(
|
||||
img_q,
|
||||
"B (T S) (N C) -> (B S) N T C",
|
||||
T=num_frames,
|
||||
S=spatial_dim,
|
||||
N=num_heads,
|
||||
C=head_dim,
|
||||
)
|
||||
key_image = rearrange(
|
||||
img_k,
|
||||
"B (T S) (N C) -> (B S) N T C",
|
||||
T=num_frames,
|
||||
S=spatial_dim,
|
||||
N=num_heads,
|
||||
C=head_dim,
|
||||
)
|
||||
weight = transformer_options.get("feta_weight", 0)
|
||||
return _feta_score(query_image, key_image, head_dim, num_frames, weight)
|
||||
@@ -1,8 +0,0 @@
|
||||
import torch
|
||||
|
||||
|
||||
class LatentGuide(torch.nn.Module):
|
||||
def __init__(self, latent: torch.Tensor, index) -> None:
|
||||
super().__init__()
|
||||
self.index = index
|
||||
self.register_buffer("latent", latent)
|
||||
@@ -1,6 +0,0 @@
|
||||
def isinstance_str(x: object, cls_name: str):
|
||||
for _cls in x.__class__.__mro__:
|
||||
if _cls.__name__ == cls_name:
|
||||
return True
|
||||
|
||||
return False
|
||||
@@ -1,23 +0,0 @@
|
||||
def get_alphacumprod(sigma):
|
||||
return 1 / ((sigma * sigma) + 1)
|
||||
|
||||
|
||||
def add_noise(src_latent, noise, sigma):
|
||||
alphas_cumprod = get_alphacumprod(sigma)
|
||||
|
||||
sqrt_alpha_prod = alphas_cumprod**0.5
|
||||
sqrt_alpha_prod = sqrt_alpha_prod.flatten()
|
||||
while len(sqrt_alpha_prod.shape) < len(src_latent.shape):
|
||||
sqrt_alpha_prod = sqrt_alpha_prod.unsqueeze(-1)
|
||||
|
||||
sqrt_one_minus_alpha_prod = (1 - alphas_cumprod) ** 0.5
|
||||
sqrt_one_minus_alpha_prod = sqrt_one_minus_alpha_prod.flatten()
|
||||
while len(sqrt_one_minus_alpha_prod.shape) < len(src_latent.shape):
|
||||
sqrt_one_minus_alpha_prod = sqrt_one_minus_alpha_prod.unsqueeze(-1)
|
||||
|
||||
noisy_samples = sqrt_alpha_prod * src_latent + sqrt_one_minus_alpha_prod * noise
|
||||
return noisy_samples
|
||||
|
||||
|
||||
def add_noise_flux(src_latent, noise, sigma):
|
||||
return sigma * noise + (1.0 - sigma) * src_latent
|
||||
@@ -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
|
||||
|
||||
@@ -1,38 +0,0 @@
|
||||
from comfy.model_management import dtype_size
|
||||
|
||||
from .nodes_registry import comfy_node
|
||||
|
||||
|
||||
@comfy_node(name="LTXVPatcherVAE")
|
||||
class LTXVPatcherVAE:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"vae": ("VAE",),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("VAE",)
|
||||
FUNCTION = "patch"
|
||||
CATEGORY = "lightricks/LTXV"
|
||||
TITLE = "LTXV VAE Patcher"
|
||||
|
||||
def patch(self, vae):
|
||||
from q8_kernels.integration.patch_vae import patch_vae
|
||||
|
||||
vae_model = vae.first_stage_model
|
||||
vae.memory_used_decode = (
|
||||
lambda shape, dtype: shape[1]
|
||||
* shape[2]
|
||||
* shape[3]
|
||||
* shape[4]
|
||||
* 8
|
||||
* 8
|
||||
* 8
|
||||
* 3
|
||||
* dtype_size(dtype)
|
||||
)
|
||||
|
||||
patch_vae(vae_model, patch_block=4)
|
||||
return (vae,)
|
||||
+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)
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user