Compare commits
19
Commits
pr-2026-08-20
...
master
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
44c8a36b7c | ||
|
|
af31f726c0 | ||
|
|
7f22eb16eb | ||
|
|
db476dacc1 | ||
|
|
3bf3ca6259 | ||
|
|
5eb48fdcae | ||
|
|
bf2ca0264f | ||
|
|
e6fdfe6686 | ||
|
|
5722b53a08 | ||
|
|
71d8972bd2 | ||
|
|
61ee82b23b | ||
|
|
c75eacae23 | ||
|
|
f8387c893d | ||
|
|
d164816f8a | ||
|
|
cb1f8204dc | ||
|
|
35364168d9 | ||
|
|
dfb2786749 | ||
|
|
63a7f34e64 | ||
|
|
15d09abb5a |
@@ -0,0 +1,141 @@
|
||||
name: Publish to Comfy Registry
|
||||
|
||||
on:
|
||||
workflow_dispatch:
|
||||
push:
|
||||
branches:
|
||||
- master
|
||||
paths:
|
||||
- pyproject.toml
|
||||
|
||||
concurrency:
|
||||
group: publish-comfy-registry
|
||||
cancel-in-progress: false
|
||||
|
||||
jobs:
|
||||
publish:
|
||||
if: github.repository_owner == 'Lightricks'
|
||||
runs-on: ubuntu-latest
|
||||
timeout-minutes: 30
|
||||
permissions:
|
||||
contents: write
|
||||
steps:
|
||||
- name: Check out the release commit
|
||||
uses: actions/checkout@11d5960a326750d5838078e36cf38b85af677262 # v4
|
||||
with:
|
||||
ref: ${{ github.event_name == 'push' && github.sha || 'master' }}
|
||||
fetch-depth: 0
|
||||
fetch-tags: true
|
||||
lfs: true
|
||||
|
||||
- name: Set up Python
|
||||
uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5
|
||||
with:
|
||||
python-version: "3.12"
|
||||
|
||||
- name: Read the version
|
||||
id: version
|
||||
run: |
|
||||
set -euo pipefail
|
||||
version=$(python3 -c '
|
||||
import sys
|
||||
import tomllib
|
||||
with open("pyproject.toml", "rb") as handle:
|
||||
version = tomllib.load(handle).get("project", {}).get("version")
|
||||
parts = str(version).split(".")
|
||||
if len(parts) != 3 or any(not part.isdigit() for part in parts):
|
||||
print(f"version is not X.Y.Z: {version!r}", file=sys.stderr)
|
||||
raise SystemExit(1)
|
||||
print(version)
|
||||
')
|
||||
echo "version=${version}" >> "$GITHUB_OUTPUT"
|
||||
|
||||
- name: Check the published version
|
||||
id: registry
|
||||
env:
|
||||
VERSION: ${{ steps.version.outputs.version }}
|
||||
run: |
|
||||
set -euo pipefail
|
||||
node=$(curl -fsS https://api.comfy.org/nodes/ComfyUI-LTXVideo)
|
||||
published=$(printf '%s\n' "$node" | python3 -c '
|
||||
import json
|
||||
import sys
|
||||
node = json.load(sys.stdin)
|
||||
latest = node.get("latest_version")
|
||||
if latest is None:
|
||||
raise SystemExit(0)
|
||||
if not isinstance(latest, dict):
|
||||
print("latest_version is not a JSON object", file=sys.stderr)
|
||||
raise SystemExit(1)
|
||||
version = latest.get("version")
|
||||
if version is None:
|
||||
raise SystemExit(0)
|
||||
parts = str(version).split(".")
|
||||
if len(parts) != 3 or any(not part.isdigit() for part in parts):
|
||||
print(f"latest_version.version is not X.Y.Z: {version!r}", file=sys.stderr)
|
||||
raise SystemExit(1)
|
||||
print(version)
|
||||
')
|
||||
if [ "$published" = "$VERSION" ]; then
|
||||
echo "Registry already has ${VERSION}."
|
||||
echo "skip=true" >> "$GITHUB_OUTPUT"
|
||||
else
|
||||
echo "skip=false" >> "$GITHUB_OUTPUT"
|
||||
fi
|
||||
|
||||
- name: Publish to the registry
|
||||
if: steps.registry.outputs.skip != 'true'
|
||||
env:
|
||||
REGISTRY_ACCESS_TOKEN: ${{ secrets.REGISTRY_ACCESS_TOKEN }}
|
||||
run: |
|
||||
set -euo pipefail
|
||||
python3 -m pip install "comfy-cli==1.22.0"
|
||||
comfy --skip-prompt --no-enable-telemetry env
|
||||
comfy node publish --token "$REGISTRY_ACCESS_TOKEN"
|
||||
|
||||
- name: Tag and release
|
||||
env:
|
||||
GH_TOKEN: ${{ github.token }}
|
||||
VERSION: ${{ steps.version.outputs.version }}
|
||||
run: |
|
||||
set -euo pipefail
|
||||
version="$VERSION"
|
||||
tag="v${version}"
|
||||
if git rev-parse -q --verify "refs/tags/${tag}" >/dev/null; then
|
||||
echo "Tag ${tag} already exists."
|
||||
else
|
||||
git tag "${tag}"
|
||||
git push origin "refs/tags/${tag}"
|
||||
fi
|
||||
if gh release view "${tag}" >/dev/null 2>&1; then
|
||||
echo "Release ${tag} already exists."
|
||||
exit 0
|
||||
fi
|
||||
previous=$(
|
||||
git tag --list "v*" | python3 -c '
|
||||
import sys
|
||||
current = tuple(int(part) for part in sys.argv[1].split("."))
|
||||
lower = []
|
||||
for tag in sys.stdin:
|
||||
tag = tag.strip()
|
||||
if not tag.startswith("v"):
|
||||
continue
|
||||
parts = tag[1:].split(".")
|
||||
if len(parts) != 3 or any(not part.isdigit() for part in parts):
|
||||
continue
|
||||
parsed = tuple(int(part) for part in parts)
|
||||
if parsed < current:
|
||||
lower.append((parsed, tag))
|
||||
if lower:
|
||||
print(max(lower)[1])
|
||||
' "$version"
|
||||
)
|
||||
if [ -z "$previous" ]; then
|
||||
notes="Initial release."
|
||||
else
|
||||
notes=$(git log --pretty=format:"%h %s" "${previous}..HEAD")
|
||||
if [ -z "$notes" ]; then
|
||||
notes="No commits since ${previous}."
|
||||
fi
|
||||
fi
|
||||
gh release create "${tag}" --title "${tag}" --notes "${notes}"
|
||||
@@ -86,13 +86,13 @@ The union LoRA is trained to understand and respond to both control signals (dep
|
||||
|
||||
## HDR IC-LoRA
|
||||
|
||||
We provide an **HDR IC-LoRA** that generates linear HDR video encoded in ARRI LogC3, enabling workflows that output high-dynamic-range content suitable for grading and EXR export.
|
||||
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. To enable EXR writing, set `OPENCV_IO_ENABLE_OPENEXR=1` in the environment before starting ComfyUI. The exported EXR sequence is best viewed in [DJV](https://github.com/grizzlypeak3d/DJV) (or [DJV for macOS](https://djv.en.uptodown.com/mac/download)).
|
||||
- **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
|
||||
|
||||
@@ -176,6 +176,25 @@ Download the following models:
|
||||
|
||||
## 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).
|
||||
|
||||
+18
-30
@@ -1,6 +1,5 @@
|
||||
from .audio_only import LTXVAudioOnlyEmptyVideoLatent, LTXVAudioOnlyModel
|
||||
from .conditioning_loader import LTXVLoadConditioning
|
||||
from .conditioning_saver import LTXVSaveConditioning
|
||||
from .decoder_noise import DecoderNoise
|
||||
from .dynamic_conditioning import DynamicConditioning
|
||||
from .easy_samplers import (
|
||||
@@ -12,23 +11,23 @@ from .easy_samplers import (
|
||||
)
|
||||
from .gemma_api_conditioning import GemmaAPITextEncode
|
||||
from .gemma_encoder import LTXVGemmaCLIPModelLoader, LTXVGemmaEnhancePrompt
|
||||
from .guide import LTXVAddGuideAdvanced, LTXVAddGuideAdvancedAttention
|
||||
from .guide import LTXVImgToVideoAdvanced
|
||||
from .guiders import GuiderParametersNode, MultimodalGuiderNode
|
||||
from .hdr import LTXVHDRDecodePostprocess
|
||||
from .hdr_nodes import (
|
||||
LTXVHDRDecodePostprocess,
|
||||
LTXVLoadEXRSequence,
|
||||
LTXVSaveHLG,
|
||||
LTXVSDRToHDRWorkingSpace,
|
||||
LTXVVAEForceFloat32,
|
||||
)
|
||||
from .iclora import (
|
||||
LTXAddVideoICLoRAGuide,
|
||||
LTXAddVideoICLoRAGuideAdvanced,
|
||||
LTXICLoRALoaderModelOnly,
|
||||
LTXVSetAudioRefTokens,
|
||||
)
|
||||
from .latent_norm import (
|
||||
LTXVAdainLatent,
|
||||
LTXVPerStepAdainPatcher,
|
||||
LTXVPerStepStatNormPatcher,
|
||||
LTXVStatNormLatent,
|
||||
)
|
||||
from .latent_norm import LTXVAdainLatent
|
||||
from .latents import (
|
||||
LTXVAddLatentGuide,
|
||||
LTXVImgToVideoConditionOnly,
|
||||
LTXVSelectLatents,
|
||||
LTXVSetVideoLatentNoiseMasks,
|
||||
@@ -39,13 +38,11 @@ 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
|
||||
@@ -55,38 +52,32 @@ from .stg import (
|
||||
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 FloatToInt, ImageToCPU
|
||||
from .vae_patcher import LTXVPatcherVAE
|
||||
from .vanish_nodes import LTXVDilateVideoMask, LTXVInpaintPreprocess
|
||||
|
||||
# Static node mappings, required for ComfyUI-Manager mapping to work
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
"Set VAE Decoder Noise": DecoderNoise,
|
||||
"LTXVLinearOverlapLatentTransition": LinearOverlapLatentTransition,
|
||||
"LTXVAddGuideAdvanced": LTXVAddGuideAdvanced,
|
||||
"LTXVAddGuideAdvancedAttention": LTXVAddGuideAdvancedAttention,
|
||||
"LTXVAddLatentGuide": LTXVAddLatentGuide,
|
||||
"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,
|
||||
@@ -97,8 +88,6 @@ NODE_CLASS_MAPPINGS = {
|
||||
"LTXVMultiPromptProvider": MultiPromptProvider,
|
||||
"ImageToCPU": ImageToCPU,
|
||||
"LTXFloatToInt": FloatToInt,
|
||||
"LTXVStatNormLatent": LTXVStatNormLatent,
|
||||
"LTXVPerStepStatNormPatcher": LTXVPerStepStatNormPatcher,
|
||||
"LTXVGemmaCLIPModelLoader": LTXVGemmaCLIPModelLoader,
|
||||
"LTXVGemmaEnhancePrompt": LTXVGemmaEnhancePrompt,
|
||||
"GemmaAPITextEncode": GemmaAPITextEncode,
|
||||
@@ -111,13 +100,16 @@ NODE_CLASS_MAPPINGS = {
|
||||
"LTXICLoRALoaderModelOnly": LTXICLoRALoaderModelOnly,
|
||||
"LTXVSetAudioRefTokens": LTXVSetAudioRefTokens,
|
||||
"LTXVLoadConditioning": LTXVLoadConditioning,
|
||||
"LTXVSaveConditioning": LTXVSaveConditioning,
|
||||
"LTXVDrawTracks": LTXVDrawTracks,
|
||||
"LTXVSparseTrackEditor": LTXVSparseTrackEditor,
|
||||
"LTXVDilateVideoMask": LTXVDilateVideoMask,
|
||||
"LTXVInpaintPreprocess": LTXVInpaintPreprocess,
|
||||
"LTXVLaplacianPyramidBlend": LTXVLaplacianPyramidBlend,
|
||||
"LTXVHDRDecodePostprocess": LTXVHDRDecodePostprocess,
|
||||
"LTXVSDRToHDRWorkingSpace": LTXVSDRToHDRWorkingSpace,
|
||||
"LTXVLoadEXRSequence": LTXVLoadEXRSequence,
|
||||
"LTXVVAEForceFloat32": LTXVVAEForceFloat32,
|
||||
"LTXVSaveHLG": LTXVSaveHLG,
|
||||
"LTXVAudioOnlyModel": LTXVAudioOnlyModel,
|
||||
"LTXVAudioOnlyEmptyVideoLatent": LTXVAudioOnlyEmptyVideoLatent,
|
||||
}
|
||||
@@ -129,10 +121,6 @@ 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)
|
||||
|
||||
+90
-9
@@ -9,6 +9,76 @@ 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):
|
||||
@@ -21,6 +91,11 @@ class LTXVLoadConditioning(io.ComfyNode):
|
||||
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"]),
|
||||
@@ -47,19 +122,25 @@ class LTXVLoadConditioning(io.ComfyNode):
|
||||
) as f:
|
||||
tensor_keys = [k for k in f.keys() if k.startswith("conditioning_data_")]
|
||||
|
||||
for tensor_key in sorted(tensor_keys):
|
||||
idx = tensor_key.replace("conditioning_data_", "")
|
||||
tensor = f.get_tensor(tensor_key)
|
||||
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)
|
||||
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])
|
||||
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}")
|
||||
raise ValueError(
|
||||
f"No conditioning data found in file: {file_name}. "
|
||||
"Expected conditioning_data_* or video_context / video_prompt_embeds."
|
||||
)
|
||||
|
||||
return io.NodeOutput(conditioning)
|
||||
|
||||
|
||||
@@ -1,69 +0,0 @@
|
||||
from datetime import datetime
|
||||
from pathlib import Path
|
||||
|
||||
import comfy.utils
|
||||
import folder_paths
|
||||
import torch
|
||||
from comfy_api.latest import io, ui
|
||||
|
||||
from .nodes_registry import comfy_node
|
||||
|
||||
|
||||
@comfy_node(name="LTXVSaveConditioning")
|
||||
class LTXVSaveConditioning(io.ComfyNode):
|
||||
@classmethod
|
||||
def define_schema(cls) -> io.Schema:
|
||||
return io.Schema(
|
||||
node_id="LTXVSaveConditioning",
|
||||
display_name="🅛🅣🅧 LTXV Save Conditioning",
|
||||
category="lightricks/LTXV",
|
||||
inputs=[
|
||||
io.Conditioning.Input("conditioning"),
|
||||
io.String.Input("filename", default="conditioning"),
|
||||
io.Combo.Input("dtype", options=["bfloat16", "float16"]),
|
||||
],
|
||||
is_output_node=True,
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def execute(cls, conditioning: list, filename: str, dtype: str) -> io.NodeOutput:
|
||||
if not conditioning or len(conditioning) == 0:
|
||||
raise ValueError("Conditioning is empty")
|
||||
|
||||
embeddings_folder = Path(folder_paths.get_folder_paths("embeddings")[0])
|
||||
embeddings_folder.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
sanitized_filename = "".join(
|
||||
c for c in filename if c.isalnum() or c in ("_", "-", ".")
|
||||
)
|
||||
if not sanitized_filename:
|
||||
sanitized_filename = "conditioning"
|
||||
|
||||
output_path = embeddings_folder / f"{sanitized_filename}.safetensors"
|
||||
|
||||
target_dtype = torch.bfloat16 if dtype == "bfloat16" else torch.float16
|
||||
|
||||
tensors_to_save: dict[str, torch.Tensor] = {}
|
||||
|
||||
for idx, (cond_tensor, cond_options) in enumerate(conditioning):
|
||||
tensor_converted = cond_tensor.to(dtype=target_dtype).contiguous()
|
||||
tensors_to_save[f"conditioning_data_{idx}"] = tensor_converted
|
||||
|
||||
if "attention_mask" in cond_options:
|
||||
mask = cond_options["attention_mask"].contiguous()
|
||||
tensors_to_save[f"attention_mask_{idx}"] = mask
|
||||
|
||||
metadata = {
|
||||
"num_conditionings": str(len(conditioning)),
|
||||
"dtype": dtype,
|
||||
"created_at": str(datetime.now()),
|
||||
}
|
||||
|
||||
comfy.utils.save_torch_file(
|
||||
tensors_to_save, str(output_path), metadata=metadata
|
||||
)
|
||||
|
||||
file_size_mb = output_path.stat().st_size / (1024 * 1024)
|
||||
return io.NodeOutput(
|
||||
ui=ui.PreviewText(f"Saved: {output_path.name} ({file_size_mb:.2f} MB)")
|
||||
)
|
||||
+8
-8
@@ -12,7 +12,7 @@ from comfy_extras.nodes_lt import EmptyLTXVLatentVideo, LTXVAddGuide, LTXVCropGu
|
||||
|
||||
from .guide import blur_internal
|
||||
from .latent_norm import LTXVAdainLatent
|
||||
from .latents import LTXVAddLatentGuide, LTXVSelectLatents
|
||||
from .latents import LTXVSelectLatents, add_latent_guide
|
||||
from .nodes_registry import comfy_node
|
||||
|
||||
|
||||
@@ -248,7 +248,7 @@ class LTXVBaseSampler:
|
||||
positive,
|
||||
negative,
|
||||
latents,
|
||||
) = LTXVAddLatentGuide().generate(
|
||||
) = add_latent_guide(
|
||||
vae=vae,
|
||||
positive=positive,
|
||||
negative=negative,
|
||||
@@ -446,7 +446,7 @@ class LTXVExtendSampler:
|
||||
positive,
|
||||
negative,
|
||||
new_latents,
|
||||
) = LTXVAddLatentGuide().generate(
|
||||
) = add_latent_guide(
|
||||
vae=vae,
|
||||
positive=positive,
|
||||
negative=negative,
|
||||
@@ -504,7 +504,7 @@ class LTXVExtendSampler:
|
||||
positive,
|
||||
negative,
|
||||
new_latents,
|
||||
) = LTXVAddLatentGuide().generate(
|
||||
) = add_latent_guide(
|
||||
vae=vae,
|
||||
positive=positive,
|
||||
negative=negative,
|
||||
@@ -519,7 +519,7 @@ class LTXVExtendSampler:
|
||||
positive,
|
||||
negative,
|
||||
new_latents,
|
||||
) = LTXVAddLatentGuide().generate(
|
||||
) = add_latent_guide(
|
||||
vae=vae,
|
||||
positive=positive,
|
||||
negative=negative,
|
||||
@@ -552,7 +552,7 @@ class LTXVExtendSampler:
|
||||
positive,
|
||||
negative,
|
||||
denoised_output_latents,
|
||||
) = LTXVAddLatentGuide().generate(
|
||||
) = add_latent_guide(
|
||||
vae=vae,
|
||||
positive=positive,
|
||||
negative=negative,
|
||||
@@ -756,7 +756,7 @@ class LTXVInContextSampler:
|
||||
positive,
|
||||
negative,
|
||||
new_latents,
|
||||
) = LTXVAddLatentGuide().generate(
|
||||
) = add_latent_guide(
|
||||
vae=vae,
|
||||
positive=positive,
|
||||
negative=negative,
|
||||
@@ -792,7 +792,7 @@ class LTXVInContextSampler:
|
||||
positive,
|
||||
negative,
|
||||
new_latents,
|
||||
) = LTXVAddLatentGuide().generate(
|
||||
) = add_latent_guide(
|
||||
vae=vae,
|
||||
positive=positive,
|
||||
negative=negative,
|
||||
|
||||
@@ -1119,7 +1119,7 @@
|
||||
"outputs": [],
|
||||
"properties": {},
|
||||
"widgets_values": [
|
||||
"Set OPENCV_IO_ENABLE_OPENEXR=1 in your shell before launching ComfyUI, otherwise save_exr will fail"
|
||||
"EXR export requires OpenImageIO (pip install openimageio). Same backend as ltx-pipelines."
|
||||
],
|
||||
"color": "#432",
|
||||
"bgcolor": "#653"
|
||||
@@ -1407,6 +1407,16 @@
|
||||
"type": "IMAGE",
|
||||
"link": 13664
|
||||
},
|
||||
{
|
||||
"localized_name": "transfer",
|
||||
"name": "transfer",
|
||||
"shape": 7,
|
||||
"type": "COMBO",
|
||||
"widget": {
|
||||
"name": "transfer"
|
||||
},
|
||||
"link": null
|
||||
},
|
||||
{
|
||||
"localized_name": "exposure",
|
||||
"name": "exposure",
|
||||
@@ -1427,6 +1437,16 @@
|
||||
},
|
||||
"link": null
|
||||
},
|
||||
{
|
||||
"localized_name": "exr_color_space",
|
||||
"name": "exr_color_space",
|
||||
"shape": 7,
|
||||
"type": "COMBO",
|
||||
"widget": {
|
||||
"name": "exr_color_space"
|
||||
},
|
||||
"link": null
|
||||
},
|
||||
{
|
||||
"localized_name": "output_dir",
|
||||
"name": "output_dir",
|
||||
@@ -1478,8 +1498,10 @@
|
||||
"Node name for S&R": "LTXVHDRDecodePostprocess"
|
||||
},
|
||||
"widgets_values": [
|
||||
"logc3",
|
||||
7.1,
|
||||
true,
|
||||
"linear",
|
||||
"output/hdr_exr3",
|
||||
"frame",
|
||||
true
|
||||
@@ -3086,8 +3108,10 @@
|
||||
},
|
||||
"5114": {
|
||||
"inputs": {
|
||||
"transfer": "logc3",
|
||||
"exposure": 7.1,
|
||||
"save_exr": true,
|
||||
"exr_color_space": "linear",
|
||||
"output_dir": "output/hdr_exr3",
|
||||
"filename_prefix": "frame",
|
||||
"half_precision": true,
|
||||
|
||||
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
@@ -92,7 +92,7 @@
|
||||
"Node name for S&R": "LoadImage"
|
||||
},
|
||||
"widgets_values": [
|
||||
"",
|
||||
"example.png",
|
||||
"image"
|
||||
],
|
||||
"color": "#232",
|
||||
@@ -7636,7 +7636,7 @@
|
||||
"Node name for S&R": "ManualSigmas"
|
||||
},
|
||||
"widgets_values": [
|
||||
"0.85, 0.7250, 0.4219, 0.0"
|
||||
"0.909375, 0.725, 0.421875, 0.0"
|
||||
]
|
||||
},
|
||||
{
|
||||
|
||||
@@ -263,7 +263,7 @@
|
||||
"ver": "0.3.56"
|
||||
},
|
||||
"widgets_values": [
|
||||
"",
|
||||
"example.png",
|
||||
"image"
|
||||
],
|
||||
"color": "#232",
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
@@ -502,7 +502,7 @@
|
||||
"ver": "0.3.56"
|
||||
},
|
||||
"widgets_values": [
|
||||
"",
|
||||
"example.png",
|
||||
"image"
|
||||
],
|
||||
"color": "#232",
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
@@ -823,7 +823,7 @@
|
||||
"ver": "0.3.56"
|
||||
},
|
||||
"widgets_values": [
|
||||
"",
|
||||
"example.png",
|
||||
"image"
|
||||
],
|
||||
"color": "#232",
|
||||
@@ -5314,7 +5314,7 @@
|
||||
"ver": "0.14.1"
|
||||
},
|
||||
"widgets_values": [
|
||||
"0.85, 0.7250, 0.4219, 0.0"
|
||||
"0.909375, 0.725, 0.421875, 0.0"
|
||||
]
|
||||
},
|
||||
{
|
||||
|
||||
@@ -837,7 +837,7 @@
|
||||
"ver": "0.3.56"
|
||||
},
|
||||
"widgets_values": [
|
||||
"",
|
||||
"example.png",
|
||||
"image"
|
||||
],
|
||||
"color": "#232",
|
||||
|
||||
@@ -536,7 +536,7 @@
|
||||
"ver": "0.3.56"
|
||||
},
|
||||
"widgets_values": [
|
||||
"",
|
||||
"example.png",
|
||||
"image"
|
||||
],
|
||||
"color": "#232",
|
||||
@@ -6019,7 +6019,7 @@
|
||||
"ver": "0.14.1"
|
||||
},
|
||||
"widgets_values": [
|
||||
"0.85, 0.7250, 0.4219, 0.0"
|
||||
"0.909375, 0.725, 0.421875, 0.0"
|
||||
]
|
||||
},
|
||||
{
|
||||
|
||||
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
@@ -42,10 +42,26 @@ These keep the distilled 2.5 backbone and add an IC-LoRA so a guide (video, imag
|
||||
| -------- | ------------ |
|
||||
| [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?
|
||||
|
||||
@@ -63,18 +79,32 @@ Do you have an existing video to edit or follow?
|
||||
│ │ → 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
|
||||
│ │ → 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 → 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)
|
||||
├─ 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?
|
||||
@@ -95,9 +125,30 @@ Two-stage is the better default when you care about spatial detail. Single-stage
|
||||
| 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 + mask (+ frozen audio) | Extending the frame |
|
||||
| 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.
|
||||
|
||||
+56
-16
@@ -4,6 +4,7 @@ Replaces the CLIP encoding step entirely using an external API.
|
||||
"""
|
||||
|
||||
import io
|
||||
import json
|
||||
import logging
|
||||
import pickle
|
||||
|
||||
@@ -23,14 +24,18 @@ 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"
|
||||
|
||||
|
||||
MODEL_FOLDERS = ("checkpoints", "diffusion_models")
|
||||
|
||||
|
||||
def model_filename_list() -> list[str]:
|
||||
"""Every model the node can read a model id from, in folder order and de-duplicated.
|
||||
"""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
|
||||
@@ -55,17 +60,49 @@ def resolve_model_path(ckpt_name: str) -> str:
|
||||
)
|
||||
|
||||
|
||||
def extract_model_id(ckpt_name: str) -> str:
|
||||
model_id_key = "encrypted_wandb_properties"
|
||||
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(
|
||||
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")
|
||||
@@ -74,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
|
||||
@@ -117,9 +154,8 @@ class GemmaAPITextEncode:
|
||||
"ckpt_name": (
|
||||
model_filename_list(),
|
||||
{
|
||||
"tooltip": "The model to read the API model id from. "
|
||||
"Either a checkpoint or a diffusion model; it must "
|
||||
"carry the model id in its safetensors metadata."
|
||||
"tooltip": "The model to read the API text-encoder selector from. "
|
||||
"Either a checkpoint or a diffusion model."
|
||||
},
|
||||
),
|
||||
},
|
||||
@@ -142,14 +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)
|
||||
selector = extract_prompt_embedding_selector(ckpt_name)
|
||||
payload = {
|
||||
"prompt": prompt,
|
||||
"model_id": model_id,
|
||||
"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(
|
||||
|
||||
@@ -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
|
||||
@@ -252,205 +132,3 @@ class LTXVImgToVideoAdvanced:
|
||||
batch_size=batch_size,
|
||||
strength=strength,
|
||||
)
|
||||
|
||||
|
||||
@comfy_node(name="LTXVAddGuideAdvancedAttention")
|
||||
class LTXVAddGuideAdvancedAttention:
|
||||
"""Extended keyframe guide node with per-guide attention strength control.
|
||||
|
||||
Same preprocessing as LTXVAddGuideAdvanced (CRF, blur, interpolation, crop),
|
||||
plus attention_strength and attention_mask inputs to control how strongly
|
||||
this guide's conditioning influences generation via self-attention.
|
||||
"""
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"positive": ("CONDITIONING",),
|
||||
"negative": ("CONDITIONING",),
|
||||
"vae": ("VAE",),
|
||||
"latent": ("LATENT",),
|
||||
"image": ("IMAGE",),
|
||||
"frame_idx": (
|
||||
"INT",
|
||||
{
|
||||
"default": 0,
|
||||
"min": -9999,
|
||||
"max": 9999,
|
||||
"tooltip": (
|
||||
"Frame index to start the conditioning at. "
|
||||
"Negative values are counted from the end of the video."
|
||||
),
|
||||
},
|
||||
),
|
||||
"strength": (
|
||||
"FLOAT",
|
||||
{
|
||||
"default": 1.0,
|
||||
"min": 0.0,
|
||||
"max": 1.0,
|
||||
"tooltip": "Strength of the conditioning. Higher values make it more exact.",
|
||||
},
|
||||
),
|
||||
"crf": (
|
||||
"INT",
|
||||
{
|
||||
"default": 29,
|
||||
"min": 0,
|
||||
"max": 51,
|
||||
"step": 1,
|
||||
"tooltip": "CRF value. Higher = more motion, lower = higher quality.",
|
||||
},
|
||||
),
|
||||
"blur_radius": (
|
||||
"INT",
|
||||
{
|
||||
"default": 0,
|
||||
"min": 0,
|
||||
"max": 7,
|
||||
"step": 1,
|
||||
"tooltip": "Blur kernel radius. Higher = more motion.",
|
||||
},
|
||||
),
|
||||
"interpolation": (
|
||||
[
|
||||
"lanczos",
|
||||
"bislerp",
|
||||
"nearest",
|
||||
"bilinear",
|
||||
"bicubic",
|
||||
"area",
|
||||
"nearest-exact",
|
||||
],
|
||||
{"default": "lanczos"},
|
||||
),
|
||||
"crop": (["center", "disabled"], {"default": "disabled"}),
|
||||
"attention_strength": (
|
||||
"FLOAT",
|
||||
{
|
||||
"default": 1.0,
|
||||
"min": 0.0,
|
||||
"max": 1.0,
|
||||
"step": 0.01,
|
||||
"tooltip": (
|
||||
"Controls how strongly this guide influences generation via "
|
||||
"self-attention. 1.0 = full conditioning, 0.0 = ignore."
|
||||
),
|
||||
},
|
||||
),
|
||||
},
|
||||
"optional": {
|
||||
"attention_mask": (
|
||||
"MASK",
|
||||
{
|
||||
"tooltip": (
|
||||
"Optional pixel-space spatial mask. Shape (F, H, W) or (H, W). "
|
||||
"Values in [0, 1]. Controls per-region conditioning influence. "
|
||||
"Multiplied by attention_strength."
|
||||
),
|
||||
},
|
||||
),
|
||||
},
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("CONDITIONING", "CONDITIONING", "LATENT")
|
||||
RETURN_NAMES = ("positive", "negative", "latent")
|
||||
|
||||
CATEGORY = "conditioning/video_models"
|
||||
FUNCTION = "generate"
|
||||
|
||||
DESCRIPTION = (
|
||||
"Adds a conditioning frame/video at a specific frame index with per-guide "
|
||||
"attention strength control. Same preprocessing as LTXVAddGuideAdvanced, "
|
||||
"plus attention_strength and optional spatial attention_mask."
|
||||
)
|
||||
|
||||
def generate(
|
||||
self,
|
||||
positive,
|
||||
negative,
|
||||
vae,
|
||||
latent,
|
||||
image,
|
||||
frame_idx,
|
||||
strength,
|
||||
crf,
|
||||
blur_radius,
|
||||
interpolation,
|
||||
crop,
|
||||
attention_strength=1.0,
|
||||
attention_mask=None,
|
||||
):
|
||||
from .iclora_attention import append_guide_attention_entry, normalize_mask
|
||||
|
||||
# Preprocessing: resize, CRF, blur (same as LTXVAddGuideAdvanced)
|
||||
scale_factors = vae.downscale_index_formula
|
||||
_, width_scale_factor, height_scale_factor = scale_factors
|
||||
latent_image = latent["samples"]
|
||||
noise_mask = nodes_lt.get_noise_mask(latent)
|
||||
_, _, latent_length, latent_height, latent_width = latent_image.shape
|
||||
|
||||
width = latent_width * width_scale_factor
|
||||
height = latent_height * height_scale_factor
|
||||
image = (
|
||||
comfy.utils.common_upscale(
|
||||
image.movedim(-1, 1), width, height, interpolation, crop=crop
|
||||
)
|
||||
.movedim(1, -1)
|
||||
.clamp(0, 1)
|
||||
)
|
||||
image = nodes_lt.LTXVPreprocess().execute(image, crf)[0]
|
||||
image = blur_internal(image, blur_radius)
|
||||
|
||||
# Encode
|
||||
_, t = nodes_lt.LTXVAddGuide.encode(
|
||||
vae, latent_width, latent_height, image, scale_factors
|
||||
)
|
||||
|
||||
# Compute latent index
|
||||
frame_idx, latent_idx = nodes_lt.LTXVAddGuide.get_latent_index(
|
||||
positive, latent_length, len(image), frame_idx, scale_factors
|
||||
)
|
||||
assert (
|
||||
latent_idx + t.shape[2] <= latent_length
|
||||
), "Conditioning frames exceed the length of the latent sequence."
|
||||
|
||||
# Append keyframe
|
||||
positive, negative, latent_image, noise_mask = (
|
||||
nodes_lt.LTXVAddGuide.append_keyframe(
|
||||
positive,
|
||||
negative,
|
||||
frame_idx,
|
||||
latent_image,
|
||||
noise_mask,
|
||||
t,
|
||||
strength,
|
||||
scale_factors,
|
||||
)
|
||||
)
|
||||
|
||||
# Track with custom attention strength/mask
|
||||
pre_filter_count = t.shape[2] * t.shape[3] * t.shape[4]
|
||||
guide_latent_shape = list(t.shape[2:])
|
||||
norm_mask = normalize_mask(attention_mask)
|
||||
positive = append_guide_attention_entry(
|
||||
positive,
|
||||
pre_filter_count,
|
||||
guide_latent_shape,
|
||||
attention_strength=attention_strength,
|
||||
attention_mask=norm_mask,
|
||||
)
|
||||
negative = append_guide_attention_entry(
|
||||
negative,
|
||||
pre_filter_count,
|
||||
guide_latent_shape,
|
||||
attention_strength=attention_strength,
|
||||
attention_mask=norm_mask,
|
||||
)
|
||||
|
||||
return (
|
||||
positive,
|
||||
negative,
|
||||
{"samples": latent_image, "noise_mask": noise_mask},
|
||||
)
|
||||
|
||||
@@ -1,245 +0,0 @@
|
||||
"""HDR utilities and ComfyUI node for HDR IC-LoRA inference.
|
||||
|
||||
Provides HDR compression/decompression transforms and a single ComfyUI node:
|
||||
|
||||
- LTXVHDRDecodePostprocess: Decompress VAE output, tonemap for preview,
|
||||
and optionally save raw linear HDR frames as EXR.
|
||||
"""
|
||||
|
||||
import logging
|
||||
import os
|
||||
|
||||
import torch
|
||||
from torch import Tensor
|
||||
|
||||
from .nodes_registry import comfy_node
|
||||
|
||||
logger = logging.getLogger("LTXVideo.hdr")
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# HDR Transform Classes (ported from ltx_core/hdr.py)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class LogC3:
|
||||
"""ARRI LogC3 (EI 800) HDR compression.
|
||||
|
||||
Maps linear [0, inf) -> [0, 1] via the camera log curve, then scales to
|
||||
[-1, 1] for VAE input.
|
||||
"""
|
||||
|
||||
A = 5.555556
|
||||
B = 0.052272
|
||||
C = 0.247190
|
||||
D = 0.385537
|
||||
E = 5.367655
|
||||
F = 0.092809
|
||||
CUT = 0.010591
|
||||
|
||||
def compress(self, hdr: Tensor) -> Tensor:
|
||||
x = torch.clamp(hdr, min=0.0)
|
||||
log_part = self.C * torch.log10(self.A * x + self.B) + self.D
|
||||
lin_part = self.E * x + self.F
|
||||
logc = torch.where(x >= self.CUT, log_part, lin_part)
|
||||
logc = torch.clamp(logc, 0.0, 1.0)
|
||||
return logc * 2.0 - 1.0
|
||||
|
||||
def decompress(self, z: Tensor) -> Tensor:
|
||||
logc = torch.clamp((z + 1.0) / 2.0, 0.0, 1.0)
|
||||
cut_log = self.E * self.CUT + self.F
|
||||
lin_from_log = (torch.pow(10.0, (logc - self.D) / self.C) - self.B) / self.A
|
||||
lin_from_lin = (logc - self.F) / self.E
|
||||
return torch.where(logc >= cut_log, lin_from_log, lin_from_lin)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Transform registry — singleton (stateless)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
_LOGC3 = LogC3()
|
||||
|
||||
|
||||
def _hdr_decompress(decoded_01: Tensor) -> Tensor:
|
||||
"""Decompress VAE-decoded image from [0,1] to linear HDR [0, inf).
|
||||
|
||||
ComfyUI's VAE decode returns images in [0, 1] via ``(raw + 1) / 2``.
|
||||
This function reverses that and applies LogC3 HDR decompression.
|
||||
"""
|
||||
raw = decoded_01.float() * 2.0 - 1.0
|
||||
return _LOGC3.decompress(raw)
|
||||
|
||||
|
||||
def _linear_to_srgb(x: Tensor) -> Tensor:
|
||||
"""Convert linear [0, 1] to sRGB [0, 1]."""
|
||||
return torch.where(
|
||||
x <= 0.0031308,
|
||||
12.92 * x,
|
||||
1.055 * torch.pow(x.clamp(min=0.0031308), 1.0 / 2.4) - 0.055,
|
||||
).clamp(0.0, 1.0)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# ComfyUI Node
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@comfy_node(name="LTXVHDRDecodePostprocess")
|
||||
class LTXVHDRDecodePostprocess:
|
||||
"""Decompress HDR from VAE output, tonemap for preview, optionally save EXR.
|
||||
|
||||
Place after VAE Decode in an HDR IC-LoRA workflow. Recovers linear HDR
|
||||
values from the compressed latent space and tonemaps them to SDR for
|
||||
display. When ``save_exr`` is enabled, also writes the raw linear HDR
|
||||
frames as an EXR image sequence.
|
||||
make sure to set OPENCV_IO_ENABLE_OPENEXR=1 environment in the command line # Must be set before cv2 import
|
||||
|
||||
|
||||
Outputs:
|
||||
tonemapped: SDR preview [0, 1] after Reinhard tonemap + sRGB gamma.
|
||||
hdr_linear: Raw linear HDR values [0, inf) for further processing.
|
||||
"""
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"image": ("IMAGE",),
|
||||
},
|
||||
"optional": {
|
||||
"exposure": (
|
||||
"FLOAT",
|
||||
{
|
||||
"default": 0.0,
|
||||
"min": -10.0,
|
||||
"max": 10.0,
|
||||
"step": 0.1,
|
||||
"display": "slider",
|
||||
"tooltip": (
|
||||
"Exposure in stops (EV). 0 = no change, "
|
||||
"+1 = 2x brighter, -1 = half brightness."
|
||||
),
|
||||
},
|
||||
),
|
||||
"save_exr": (
|
||||
"BOOLEAN",
|
||||
{
|
||||
"default": False,
|
||||
"tooltip": "Save raw linear HDR frames as EXR sequence.",
|
||||
},
|
||||
),
|
||||
"output_dir": (
|
||||
"STRING",
|
||||
{
|
||||
"default": "output/hdr_exr",
|
||||
"tooltip": (
|
||||
"Directory for EXR frames (relative to ComfyUI "
|
||||
"output directory, or absolute path)."
|
||||
),
|
||||
},
|
||||
),
|
||||
"filename_prefix": (
|
||||
"STRING",
|
||||
{"default": "frame"},
|
||||
),
|
||||
"half_precision": (
|
||||
"BOOLEAN",
|
||||
{
|
||||
"default": True,
|
||||
"tooltip": (
|
||||
"Save EXR as float16 (half). Smaller files, "
|
||||
"negligible quality loss for most workflows."
|
||||
),
|
||||
},
|
||||
),
|
||||
},
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("IMAGE", "IMAGE")
|
||||
RETURN_NAMES = ("tonemapped", "hdr_linear")
|
||||
OUTPUT_NODE = True
|
||||
FUNCTION = "postprocess"
|
||||
CATEGORY = "Lightricks/HDR"
|
||||
DESCRIPTION = (
|
||||
"Decompresses VAE-decoded output from HDR IC-LoRA (LogC3) and applies "
|
||||
"Reinhard tonemapping. Place after VAE Decode. 'tonemapped' is the "
|
||||
"SDR preview; 'hdr_linear' is raw linear HDR for downstream use. "
|
||||
"Enable 'save_exr' to write an EXR image sequence."
|
||||
"if save_exr is enabled, make sure to set OPENCV_IO_ENABLE_OPENEXR=1 environment in the command line"
|
||||
)
|
||||
|
||||
def postprocess(
|
||||
self,
|
||||
image: torch.Tensor,
|
||||
exposure: float = 0.0,
|
||||
save_exr: bool = False,
|
||||
output_dir: str = "output/hdr_exr",
|
||||
filename_prefix: str = "frame",
|
||||
half_precision: bool = True,
|
||||
) -> tuple:
|
||||
|
||||
hdr = _hdr_decompress(image)
|
||||
hdr = torch.clamp(hdr, min=0.0, max=1e4)
|
||||
|
||||
# Reinhard tonemap with exposure (stops → linear multiplier)
|
||||
exposure_mult = 2.0**exposure
|
||||
hdr_exposed = hdr * exposure_mult
|
||||
tonemapped_linear = (hdr_exposed / (1.0 + hdr_exposed)).clamp(0.0, 1.0)
|
||||
tonemapped = _linear_to_srgb(tonemapped_linear)
|
||||
|
||||
if save_exr:
|
||||
assert os.environ.get("OPENCV_IO_ENABLE_OPENEXR") == "1", (
|
||||
"EXR output is enabled (save_exr = TRUE), but OpenCV does not support EXR by default. "
|
||||
"To enable it, set the environment variable OPENCV_IO_ENABLE_OPENEXR=1 before starting ComfyUI, then restart. "
|
||||
"Alternatively, disable EXR output or switch to PNG/JPG."
|
||||
)
|
||||
self._save_exr_frames(hdr, output_dir, filename_prefix, half_precision)
|
||||
|
||||
return (tonemapped, hdr)
|
||||
|
||||
@staticmethod
|
||||
def _save_exr_frames(
|
||||
hdr_image: torch.Tensor,
|
||||
output_dir: str,
|
||||
filename_prefix: str,
|
||||
half_precision: bool,
|
||||
) -> None:
|
||||
try:
|
||||
import cv2
|
||||
except ImportError:
|
||||
logger.error(
|
||||
"opencv-python is required for EXR export. "
|
||||
"Install with: pip install opencv-python"
|
||||
)
|
||||
return
|
||||
|
||||
import folder_paths
|
||||
import numpy as np
|
||||
|
||||
if not os.path.isabs(output_dir):
|
||||
output_dir = os.path.join(folder_paths.get_output_directory(), output_dir)
|
||||
|
||||
os.makedirs(output_dir, exist_ok=True)
|
||||
|
||||
frames = hdr_image.cpu().numpy()
|
||||
exr_type = (
|
||||
cv2.IMWRITE_EXR_TYPE_HALF if half_precision else cv2.IMWRITE_EXR_TYPE_FLOAT
|
||||
)
|
||||
params = [
|
||||
cv2.IMWRITE_EXR_TYPE,
|
||||
exr_type,
|
||||
cv2.IMWRITE_EXR_COMPRESSION,
|
||||
cv2.IMWRITE_EXR_COMPRESSION_ZIP,
|
||||
]
|
||||
|
||||
for i in range(frames.shape[0]):
|
||||
frame_bgr = frames[i][:, :, ::-1].astype(np.float32).copy()
|
||||
path = os.path.join(output_dir, f"{filename_prefix}_{i:05d}.exr")
|
||||
cv2.imwrite(path, frame_bgr, params)
|
||||
|
||||
logger.info(
|
||||
"Saved %d EXR frame(s) to %s (%s)",
|
||||
frames.shape[0],
|
||||
output_dir,
|
||||
"float16" if half_precision else "float32",
|
||||
)
|
||||
+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 {}
|
||||
@@ -12,6 +12,16 @@ 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")
|
||||
@@ -28,7 +38,9 @@ class LTXAddVideoICLoRAGuide(io.ComfyNode):
|
||||
"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."
|
||||
"(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"),
|
||||
@@ -44,12 +56,7 @@ class LTXAddVideoICLoRAGuide(io.ComfyNode):
|
||||
default=0,
|
||||
min=-9999,
|
||||
max=9999,
|
||||
tooltip=(
|
||||
"Frame index to start the conditioning at. "
|
||||
"For single-frame videos, any frame_idx value is acceptable. "
|
||||
"For videos, frame_idx must be 1 modulo 8, otherwise it will be rounded "
|
||||
"down to the nearest 1 modulo 8. Negative values are counted from the end of the video."
|
||||
),
|
||||
tooltip=FRAME_IDX_TOOLTIP,
|
||||
),
|
||||
io.Float.Input(
|
||||
"strength",
|
||||
@@ -93,6 +100,31 @@ class LTXAddVideoICLoRAGuide(io.ComfyNode):
|
||||
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"),
|
||||
@@ -101,6 +133,35 @@ class LTXAddVideoICLoRAGuide(io.ComfyNode):
|
||||
],
|
||||
)
|
||||
|
||||
@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,
|
||||
@@ -144,6 +205,186 @@ class LTXAddVideoICLoRAGuide(io.ComfyNode):
|
||||
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,
|
||||
@@ -159,6 +400,8 @@ class LTXAddVideoICLoRAGuide(io.ComfyNode):
|
||||
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"]
|
||||
@@ -166,6 +409,28 @@ class LTXAddVideoICLoRAGuide(io.ComfyNode):
|
||||
|
||||
_, _, 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
|
||||
@@ -219,12 +484,9 @@ class LTXAddVideoICLoRAGuide(io.ComfyNode):
|
||||
guide_latent.shape[2] * guide_latent.shape[3] * guide_latent.shape[4]
|
||||
)
|
||||
|
||||
frame_idx, latent_idx = nodes_lt.LTXVAddGuide.get_latent_index(
|
||||
positive, latent_length, len(image), frame_idx, scale_factors
|
||||
frame_idx = cls.resolve_frame_idx(
|
||||
positive, frame_idx, latent_length, len(image), guide_latent, scale_factors
|
||||
)
|
||||
assert (
|
||||
latent_idx + guide_latent.shape[2] <= latent_length
|
||||
), "Conditioning frames exceed the length of the latent sequence."
|
||||
|
||||
positive, negative, latent_image, noise_mask = (
|
||||
nodes_lt.LTXVAddGuide.append_keyframe(
|
||||
@@ -276,7 +538,8 @@ class LTXAddVideoICLoRAGuideAdvanced(LTXAddVideoICLoRAGuide):
|
||||
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."
|
||||
"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"),
|
||||
@@ -292,11 +555,7 @@ class LTXAddVideoICLoRAGuideAdvanced(LTXAddVideoICLoRAGuide):
|
||||
default=0,
|
||||
min=-9999,
|
||||
max=9999,
|
||||
tooltip=(
|
||||
"Frame index to start the conditioning at. The value is rounded to the "
|
||||
"nearest frame and wrapped modulo the number of video frames. Negative "
|
||||
"values are counted from the end of the video before wrapping."
|
||||
),
|
||||
tooltip=FRAME_IDX_TOOLTIP,
|
||||
),
|
||||
io.Float.Input("strength", default=1.0, min=0.0, max=1.0, step=0.01),
|
||||
io.Float.Input(
|
||||
@@ -336,6 +595,27 @@ class LTXAddVideoICLoRAGuideAdvanced(LTXAddVideoICLoRAGuide):
|
||||
"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"),
|
||||
@@ -361,6 +641,8 @@ class LTXAddVideoICLoRAGuideAdvanced(LTXAddVideoICLoRAGuide):
|
||||
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
|
||||
@@ -371,6 +653,30 @@ class LTXAddVideoICLoRAGuideAdvanced(LTXAddVideoICLoRAGuide):
|
||||
|
||||
_, _, 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
|
||||
@@ -420,12 +726,9 @@ class LTXAddVideoICLoRAGuideAdvanced(LTXAddVideoICLoRAGuide):
|
||||
guide_latent.shape[2] * guide_latent.shape[3] * guide_latent.shape[4]
|
||||
)
|
||||
|
||||
frame_idx, latent_idx = nodes_lt.LTXVAddGuide.get_latent_index(
|
||||
positive, latent_length, len(image), frame_idx, scale_factors
|
||||
frame_idx = cls.resolve_frame_idx(
|
||||
positive, frame_idx, latent_length, len(image), guide_latent, scale_factors
|
||||
)
|
||||
assert (
|
||||
latent_idx + guide_latent.shape[2] <= latent_length
|
||||
), "Conditioning frames exceed the length of the latent sequence."
|
||||
|
||||
positive, negative, latent_image, noise_mask = (
|
||||
nodes_lt.LTXVAddGuide.append_keyframe(
|
||||
|
||||
@@ -53,6 +53,19 @@ def append_guide_attention_entry(
|
||||
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.
|
||||
|
||||
|
||||
-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)
|
||||
|
||||
+40
-211
@@ -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,102 +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"]
|
||||
|
||||
# Record original (pre-dilation) guide latent shape for spatial mask downsampling
|
||||
guide_orig_shape = list(guide.shape[2:]) # [F, H_small, W_small]
|
||||
|
||||
assert (
|
||||
latent.shape[4] % guide.shape[4] == 0
|
||||
and latent.shape[3] % guide.shape[3] == 0
|
||||
), "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)
|
||||
|
||||
# Pre-filter token count = product of dilated spatial dims
|
||||
# (before grid_mask filtering removes padding positions)
|
||||
iclora_tokens_added = guide.shape[2] * guide.shape[3] * guide.shape[4]
|
||||
|
||||
scale_factors = vae.downscale_index_formula
|
||||
|
||||
if latent_idx <= 0:
|
||||
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,
|
||||
)
|
||||
|
||||
# Track this guide in guide_attention_entries for per-reference attention control.
|
||||
from .iclora_attention import append_guide_attention_entry
|
||||
|
||||
positive = append_guide_attention_entry(
|
||||
positive, iclora_tokens_added, guide_orig_shape
|
||||
)
|
||||
negative = append_guide_attention_entry(
|
||||
negative, iclora_tokens_added, guide_orig_shape
|
||||
)
|
||||
|
||||
return (
|
||||
positive,
|
||||
negative,
|
||||
{"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")
|
||||
|
||||
@@ -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,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
|
||||
+2
-3
@@ -10,7 +10,6 @@ from kornia.geometry.transform.pyramid import (
|
||||
build_pyramid,
|
||||
find_next_powerof_two,
|
||||
is_powerof_two,
|
||||
pad,
|
||||
)
|
||||
from torch import Tensor
|
||||
|
||||
@@ -28,7 +27,7 @@ def _pad_for_laplacian(image: torch.Tensor) -> tuple[torch.Tensor, tuple[int, in
|
||||
if not (is_powerof_two(h) and is_powerof_two(w)):
|
||||
pad_right = find_next_powerof_two(w) - w
|
||||
pad_down = find_next_powerof_two(h) - h
|
||||
image = pad(image, (0, pad_right, 0, pad_down), "reflect")
|
||||
image = F.pad(image, (0, pad_right, 0, pad_down), mode="reflect")
|
||||
return image, (pad_right, pad_down)
|
||||
|
||||
|
||||
@@ -41,7 +40,7 @@ def _gaussian_pyramid(
|
||||
h, w = images.shape[2], images.shape[3]
|
||||
if not (is_powerof_two(w) and is_powerof_two(h)):
|
||||
padding = (0, find_next_powerof_two(w) - w, 0, find_next_powerof_two(h) - h)
|
||||
images = pad(images, padding, border_type)
|
||||
images = F.pad(images, padding, mode=border_type)
|
||||
return build_pyramid(images, max_level, border_type, align_corners)
|
||||
|
||||
|
||||
|
||||
+5
-1
@@ -1,6 +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
|
||||
|
||||
@@ -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,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
|
||||
@@ -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,)
|
||||
+9
-2
@@ -97,6 +97,11 @@ 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
|
||||
@@ -105,8 +110,10 @@ class LTXVInpaintPreprocess(io.ComfyNode):
|
||||
node_id="LTXVInpaintPreprocess",
|
||||
category="Lightricks/image_processing",
|
||||
description=(
|
||||
"Composites images with a green background where mask is "
|
||||
"active, for inpainting conditioning."
|
||||
"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(
|
||||
|
||||
@@ -0,0 +1,94 @@
|
||||
import { app } from "../../../scripts/app.js";
|
||||
|
||||
const NODE_NAME = "LTXVGetTilingSizes";
|
||||
|
||||
function linesOf(text) {
|
||||
return String(text || "").split("\n").length;
|
||||
}
|
||||
|
||||
function ensurePreview(node, text) {
|
||||
const value = Array.isArray(text) ? text.join("\n") : String(text ?? "");
|
||||
if (node._tilingPreviewEl) {
|
||||
node._tilingPreviewEl.textContent = value;
|
||||
const widget = (node.widgets || []).find((w) => w.name === "tiling_preview");
|
||||
if (widget?.computeSize) {
|
||||
const lines = linesOf(value);
|
||||
widget.computedHeight = Math.max(88, 14 + lines * 16);
|
||||
}
|
||||
return;
|
||||
}
|
||||
|
||||
const el = document.createElement("pre");
|
||||
el.className = "ltx-tiling-sizes-preview";
|
||||
el.textContent = value || "Queue prompt to preview tiling sizes.";
|
||||
Object.assign(el.style, {
|
||||
margin: "0",
|
||||
padding: "6px 8px",
|
||||
fontFamily: "ui-monospace, SFMono-Regular, Menlo, Consolas, monospace",
|
||||
fontSize: "11px",
|
||||
lineHeight: "1.4",
|
||||
whiteSpace: "pre-wrap",
|
||||
color: "var(--fg-color, #ddd)",
|
||||
background: "var(--bg-color, transparent)",
|
||||
border: "1px solid var(--border-color, rgba(255,255,255,0.12))",
|
||||
borderRadius: "6px",
|
||||
width: "100%",
|
||||
boxSizing: "border-box",
|
||||
pointerEvents: "none",
|
||||
});
|
||||
|
||||
const widget = node.addDOMWidget("tiling_preview", "preview", el, {
|
||||
serialize: false,
|
||||
hideOnZoom: false,
|
||||
getValue() {
|
||||
return el.textContent;
|
||||
},
|
||||
setValue(v) {
|
||||
el.textContent = v ?? "";
|
||||
},
|
||||
});
|
||||
widget.computeSize = function (width) {
|
||||
const lines = linesOf(el.textContent);
|
||||
const height = Math.max(88, 14 + lines * 16);
|
||||
this.computedHeight = height;
|
||||
return [width, height];
|
||||
};
|
||||
|
||||
node._tilingPreviewEl = el;
|
||||
}
|
||||
|
||||
app.registerExtension({
|
||||
name: "LTXVideo.GetTilingSizesPreview",
|
||||
async beforeRegisterNodeDef(nodeType, nodeData) {
|
||||
if (nodeData?.name !== NODE_NAME) {
|
||||
return;
|
||||
}
|
||||
|
||||
const onNodeCreated = nodeType.prototype.onNodeCreated;
|
||||
nodeType.prototype.onNodeCreated = function () {
|
||||
onNodeCreated?.apply(this, arguments);
|
||||
ensurePreview(this, "Queue prompt to preview tiling sizes.");
|
||||
};
|
||||
|
||||
const onExecuted = nodeType.prototype.onExecuted;
|
||||
nodeType.prototype.onExecuted = function (message) {
|
||||
onExecuted?.apply(this, arguments);
|
||||
const text = message?.text;
|
||||
if (text == null) {
|
||||
return;
|
||||
}
|
||||
ensurePreview(this, text);
|
||||
requestAnimationFrame(() => {
|
||||
const size = this.computeSize?.() || this.size;
|
||||
if (size && this.onResize) {
|
||||
const next = [
|
||||
Math.max(size[0], this.size?.[0] || 0),
|
||||
Math.max(size[1], this.size?.[1] || 0),
|
||||
];
|
||||
this.onResize(next);
|
||||
}
|
||||
app.graph?.setDirtyCanvas?.(true, false);
|
||||
});
|
||||
};
|
||||
},
|
||||
});
|
||||
Reference in New Issue
Block a user