From 6bd282b14209546615f89adfe316f306c7267d2e Mon Sep 17 00:00:00 2001 From: Kruk2 Date: Mon, 28 Sep 2026 13:12:00 +0200 Subject: [PATCH 01/48] build(amd): pin ROCm 10.0 / torch 2.12.0 / torchvision 0.27.0 on both OSes ROCm 10.0.0 (stable, 2026-08-26) publishes torch 2.12.0 and torchvision 0.27.0 cp312 wheels for Windows and Linux on https://stable.repo.amd.com/rocm/whl-next/ with per-GPU kernel packs for gfx1100-gfx1103, gfx1150/gfx1151 and gfx1200/gfx1201, so the Radeon 780M (gfx1103) gets kernels alongside every target the ROCm 7.2.1 build had. Both OSes now share one torch/torchvision pin, the Windows setup installs the same kernel packs as the Linux image keeps, and the README names the Adrenalin 26.8.1 driver ROCm 10 needs on Windows. The protection bump moves the Linux image to ROCm 10 and bundles the ROCm payload on Linux. Refs #408 --- README.ja.md | 2 +- README.md | 2 +- README.zh.md | 2 +- docs/en/development.md | 51 ++++++++++++++++++++---------------------- jasna/protection | 2 +- pyproject.toml | 11 ++++----- 6 files changed, 32 insertions(+), 38 deletions(-) diff --git a/README.ja.md b/README.ja.md index 5e853679..0ec717c6 100644 --- a/README.ja.md +++ b/README.ja.md @@ -45,7 +45,7 @@ Jasna は無料です。支援者には、このプロジェクト用に訓練 - NVIDIA の **GTX 16 シリーズ / RTX 20 シリーズ以降**の GPU。GTX 10 シリーズ以前のカード（GTX 1050/1060/1070/1080）では動作しません。自分の GPU が対応か分からない場合は、NVIDIA の [GPU 一覧](https://developer.nvidia.com/cuda/gpus)で確認してください — コンピュート能力 7.5 以上が必要です。 - Nvidia ドライバーは Windows で **610 以上**、Linux で **580 以上**。 -- AMD 対応は実験的で、ROCm 対応 GPU が必要です。 +- AMD 対応は実験的で、ROCm 対応 GPU（Radeon RX 7000/9000、Radeon 780M/760M/740M、Ryzen AI の内蔵 GPU）が必要です。Windows では AMD Adrenalin ドライバーを 26.8.1 以降に更新してください。 - Jasna は、パスに英語の文字と数字のみを含むフォルダにインストールしてください。 Jasna は VRAM を自動管理します。VRAM が不足すると、待機中のフレームを一時的にシステム RAM へ移します。設定は不要です。 diff --git a/README.md b/README.md index f271b992..a17fe30d 100644 --- a/README.md +++ b/README.md @@ -47,7 +47,7 @@ Join the [SLS Discord](https://discord.gg/uNwQ4mHqgv) for examples, support, and - An NVIDIA **GTX 16-series / RTX 20-series or newer** GPU. GTX 10-series and older cards (GTX 1050/1060/1070/1080) won't work. Not sure about yours? Check NVIDIA's [GPU table](https://developer.nvidia.com/cuda/gpus) — compute capability 7.5+ is required. - Nvidia driver **610 or newer** on Windows, **580 or newer** on Linux. -- AMD support is experimental and needs a ROCm-supported GPU. +- AMD support is experimental and needs a ROCm-supported GPU (Radeon RX 7000/9000, Radeon 780M/760M/740M, Ryzen AI iGPUs). On Windows, update the AMD Adrenalin driver to 26.8.1 or newer. - Install Jasna into a folder whose path contains only English letters and numbers. Jasna manages VRAM automatically: when it runs low, waiting frames are temporarily moved to system RAM. No configuration needed. diff --git a/README.zh.md b/README.zh.md index 04d1539d..0c9e759e 100644 --- a/README.zh.md +++ b/README.zh.md @@ -45,7 +45,7 @@ Jasna 是免费的。支持者会获得一个密钥，用于解锁为本项目 - NVIDIA **GTX 16 系列 / RTX 20 系列或更新**的 GPU。GTX 10 系列及更旧的显卡（GTX 1050/1060/1070/1080）无法使用。不确定自己的显卡？请查看 NVIDIA 的 [GPU 表格](https://developer.nvidia.com/cuda/gpus) — 需要计算能力 7.5 或更高。 - Nvidia 驱动: Windows 需 **610 或更新**，Linux 需 **580 或更新**。 -- AMD 支持是实验性的，需要 ROCm 支持的 GPU。 +- AMD 支持是实验性的，需要 ROCm 支持的 GPU（Radeon RX 7000/9000、Radeon 780M/760M/740M、Ryzen AI 核显）。Windows 上请把 AMD Adrenalin 驱动更新到 26.8.1 或更高版本。 - 请把 Jasna 安装到路径只包含英文字母和数字的文件夹中。 Jasna 会自动管理 VRAM: 显存不足时，等待中的帧会临时移动到系统内存。无需任何配置。 diff --git a/docs/en/development.md b/docs/en/development.md index 94e7ed3a..423b9c0b 100644 --- a/docs/en/development.md +++ b/docs/en/development.md @@ -31,35 +31,35 @@ Install runtime dependencies for the active vendor: # NVIDIA (CUDA 13 wheels) uv pip install ".[nvidia]" --extra-index-url https://download.pytorch.org/whl/cu130 -# AMD Linux (inside a ROCm 7.2 environment; torch/torchvision+rocm come from the -# rocm/pytorch base image, the find-links only backfills any missing rocm wheel) -uv pip install ".[amd]" \ - --find-links https://repo.radeon.com/rocm/manylinux/rocm-rel-7.2.1/ +# AMD Linux (inside the rocm/pytorch ROCm 10.0 image; torch/torchvision+rocm come +# from the image, see jasna/protection/keytool/Dockerfile.amd) +uv pip install ".[amd]" ``` -**AMD Windows** (ROCm 7.2.1 — Python 3.12, Adrenalin ≥ 26.2.2). Install the ROCm -SDK + torch/torchvision ROCm wheels FIRST, with `--no-deps` so pip cannot silently -replace them with the CPU `torch==2.9.1` from PyPI when a later dependency -(torchvision, rfdetr, …) pulls torch — that swap is the usual cause of a -"non-ROCm torch/torchvision" env: +**AMD Windows** (ROCm 10.0.0 — Python 3.12, Adrenalin ≥ 26.8.1). Install torch and +torchvision from AMD's ROCm index FIRST, with the GPU kernel packs the release +ships, so pip cannot silently replace them with the CPU `torch==2.12.0` from PyPI +when a later dependency (rfdetr, …) pulls torch — that swap is the usual cause of +a "non-ROCm torch/torchvision" env: ```powershell -$R = "https://repo.radeon.com/rocm/windows/rocm-rel-7.2.1" -pip install --no-deps ` - "$R/rocm_sdk_core-7.2.1-py3-none-win_amd64.whl" ` - "$R/rocm_sdk_libraries_custom-7.2.1-py3-none-win_amd64.whl" ` - "$R/torch-2.9.1+rocm7.2.1-cp312-cp312-win_amd64.whl" ` - "$R/torchvision-0.24.1+rocm7.2.1-cp312-cp312-win_amd64.whl" +$I = "https://stable.repo.amd.com/rocm/whl-next/" +$D = "device-gfx1100,device-gfx1101,device-gfx1102,device-gfx1103,device-gfx1150,device-gfx1151,device-gfx1200,device-gfx1201" +pip install --index-url $I "torch[$D]==2.12.0+rocm10.0.0" "torchvision[$D]==0.27.0+rocm10.0.0" # then the remaining AMD deps (torch/torchvision above already satisfy the pins) uv pip install ".[amd]" ``` -The `rocm_sdk_core` / `rocm_sdk_libraries_*` wheels are the ROCm runtime itself. -`import torch` reaches them through `rocm_sdk`, which resolves the package name at -call time, so Nuitka sees no import and bundles nothing — the release build copies -both wheel trees into the dist verbatim and then asserts that every library torch -preloads resolves there. Without that, the frozen app dies at startup with -`UnboundLocalError: cannot access local variable 'py_module'` (upstream +Adding a GPU target means adding its `device-gfx*` extra here and keeping it in +the `Dockerfile.amd` keep-list, so both OSes ship the same targets. + +The `rocm_sdk_core` / `rocm_sdk_libraries` / `rocm_sdk_device_gfx*` wheels are the +ROCm runtime itself, and torch's own GPU kernels sit in `torch/.kpack`. +`import torch` reaches the runtime through `rocm_sdk`, which resolves the package +name at call time, so Nuitka sees no import and bundles nothing — both AMD release +builds copy the wheel trees into the dist verbatim and then assert that every +library torch preloads resolves there. Without that, the frozen app dies at startup +with `UnboundLocalError: cannot access local variable 'py_module'` (upstream `rocm_sdk.find_libraries` reports an absent payload package that way). Verify ROCm actually stuck on either OS (the Linux Docker build asserts the same): @@ -68,12 +68,9 @@ Verify ROCm actually stuck on either OS (the Linux Docker build asserts the same python -c "import torch, torchvision; assert torch.version.hip and '+rocm' in torchvision.__version__; print('ROCm OK', torch.__version__, torchvision.__version__)" ``` -`jasna[amd]` pins `torch==2.9.1` as a plain version on purpose: the ROCm wheels -(`2.9.1+rocm7.2.1`) satisfy it, and a `+rocm7.2.1` *local-version* pin can't be -shared across OSes. torchvision is pinned per OS (`0.24.0` on Linux, `0.24.1` -on Windows) because the ROCm channels ship different versions — the Linux -manylinux channel and the `rocm/pytorch` base image only have `0.24.0`, so a -shared `0.24.1` pin would make pip/uv silently install the PyPI CUDA build, +`jasna[amd]` pins `torch==2.12.0` and `torchvision==0.27.0` as plain versions, which +the ROCm wheels (`2.12.0+rocm10.0.0`, `0.27.0+rocm10.0.0`) satisfy on both OSes. A +pin no ROCm wheel satisfies makes pip/uv silently install the PyPI CUDA build, which breaks `torchvision.ops.nms` (and with it RF-DETR) at runtime. The ROCm-build assertion above is the fail-loud guard on top. diff --git a/jasna/protection b/jasna/protection index b8d2bf3c..6152fc76 160000 --- a/jasna/protection +++ b/jasna/protection @@ -1 +1 @@ -Subproject commit b8d2bf3cfdd5dba49432e89414ad6197eed277f0 +Subproject commit 6152fc76ba302074935acb63dc97dfc054dd5daf diff --git a/pyproject.toml b/pyproject.toml index 487a4493..f163ec77 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -46,13 +46,10 @@ nvidia = [ "python_vali==4.8.8", ] amd = [ - "torch==2.9.1", - # The ROCm channels ship different torchvision versions per OS (Linux - # manylinux/base image: 0.24.0, Windows: 0.24.1). A shared ==0.24.1 pin is - # unsatisfiable by any Linux ROCm wheel, so uv silently replaced torchvision - # with the PyPI CUDA build and torchvision::nms failed to register. - "torchvision==0.24.0; sys_platform == 'linux'", - "torchvision==0.24.1; sys_platform == 'win32'", + # Must match the ROCm 10.0 wheels (stable.repo.amd.com/rocm/whl-next) on both + # OSes; an unsatisfiable pin makes uv silently swap in the PyPI CUDA build. + "torch==2.12.0", + "torchvision==0.27.0", # AMD runs RF-DETR through the rfdetr torch model (ROCm), not ONNX/MIGraphX, # which sidesteps the Windows-only MIGraphX wheels (see issue #240). "rfdetr==1.8.3", From 82ecbbadf596b8adac93be404d807a7f0a691fa1 Mon Sep 17 00:00:00 2001 From: Kruk2 Date: Mon, 28 Sep 2026 13:16:03 +0200 Subject: [PATCH 02/48] fix(licensing): ship release notices Bundle application and dependency license texts, record exact source and model correspondence, and include the legal payload in frozen builds.\n\nDocument the private protection source as the sole reproducibility exception. --- LICENSING.md | 77 ++-- NOTICE | 19 + README.ja.md | 8 + README.md | 8 + README.zh.md | 8 + RELEASE_SOURCES.md | 79 ++++ assets/THIRD_PARTY_LICENSES.md | 59 ++- assets/THIRD_PARTY_MODELS.md | 77 +++- assets/licenses/Apache-2.0.txt | 145 +++++++ assets/licenses/BSD-3-Clause.txt | 23 ++ assets/licenses/GPL-3.0.txt | 674 +++++++++++++++++++++++++++++++ assets/licenses/LGPL-2.1.txt | 502 +++++++++++++++++++++++ assets/licenses/LGPL-3.0.txt | 165 ++++++++ assets/licenses/Python-3.13.txt | 277 +++++++++++++ assets/licenses/README.md | 12 + docs/en/development.md | 19 + jasna/__init__.py | 2 +- jasna/protection | 2 +- pyproject.toml | 2 +- 19 files changed, 2111 insertions(+), 47 deletions(-) create mode 100644 NOTICE create mode 100644 RELEASE_SOURCES.md create mode 100644 assets/licenses/Apache-2.0.txt create mode 100644 assets/licenses/BSD-3-Clause.txt create mode 100644 assets/licenses/GPL-3.0.txt create mode 100644 assets/licenses/LGPL-2.1.txt create mode 100644 assets/licenses/LGPL-3.0.txt create mode 100644 assets/licenses/Python-3.13.txt create mode 100644 assets/licenses/README.md diff --git a/LICENSING.md b/LICENSING.md index 696ca5a7..0c545cef 100644 --- a/LICENSING.md +++ b/LICENSING.md @@ -1,45 +1,62 @@ # Licensing and source availability -This document describes how Jasna's public source and optional supporter -components are distributed. It does not replace the license notices in -individual files or grant additional permissions. +This document describes how Jasna's source, release packages, models, and +optional supporter components are distributed. It does not replace the +license notices in individual files or grant additional permissions. -## Public source +## Public application source -Unless an individual file states otherwise, the source code in this repository -is licensed under the GNU Affero General Public License version 3.0 -(`AGPL-3.0`). Files derived from other projects retain their copyright and -license notices. +Unless an individual file states otherwise, the source code in this +repository is licensed under the GNU Affero General Public License version +3.0 (`AGPL-3.0-only`). Files derived from other projects retain their +copyright and license notices. Jasna includes code and model weights from [Lada](https://codeberg.org/ladaapp/lada). Those components remain covered by their applicable AGPL-3.0 notices. The vendored MMagic subset retains its -OpenMMLab copyright and license notices. +OpenMMLab copyright and Apache-2.0 notices. -A clean checkout of this repository can be installed and used for Jasna's free -restoration workflows without the optional `jasna.protection` package. +The complete Jasna license is in [LICENSE](LICENSE). Copyright and attribution +notices are in [NOTICE](NOTICE), +[assets/THIRD_PARTY_LICENSES.md](assets/THIRD_PARTY_LICENSES.md), and +[assets/THIRD_PARTY_MODELS.md](assets/THIRD_PARTY_MODELS.md). -## Optional supporter components - -Official release packages may additionally include a separately developed, -source-unavailable protection component. It performs local supporter-key -validation and in-memory decryption for the supporter-only `unet-4x` and -SD 1.5 restoration models. When that component is unavailable, supporter-only -models are unavailable, while the free restoration workflows remain usable. +## Official release packages -The protection component is maintained as a separate project, can be reused by -other applications, and is distributed under separate proprietary terms. It -is not included in the public Jasna source repository. +Official packages include the files above, the full license texts in +`licenses/`, and the release-specific source correspondence in +[RELEASE_SOURCES.md](RELEASE_SOURCES.md). The release page also provides a +SHA-256 file for every split archive. -The protection component contains debugger and core-dump resistance intended -to make extraction of model-decryption material more difficult. It does not -implement telemetry, persistence, or network communication. +The source tag and release package must have the same version. The tag is the +preferred form for modifying Jasna's public application code. Exact custom +dependency revisions, FFmpeg builds, patches, model hashes, and source links +are recorded in `RELEASE_SOURCES.md` and the two third-party notice files. -## Release reproducibility +## Optional supporter components -The public repository contains the application source and development setup -for the free restoration workflows. It does not contain the proprietary -protection component or the maintainer's standalone release tooling. -Consequently, a public checkout does not reproduce a supporter-enabled -official release package bit for bit. +Official release packages may include a separately maintained, +source-unavailable protection component. It performs local supporter-key +validation and in-memory decryption for supporter-only models. It contains no +telemetry, persistence, or network communication. + +The protection component is compiled into the official executable. There is +no supported way to remove it from a prebuilt package. In the free CLI path, +when no license arguments or supporter models are selected, Jasna does not +import or call it. A public-source installation can run the free restoration +workflows without the component. + +The protection implementation is proprietary and its source is not included +in the public repository. Consequently, the public source does not reproduce +the official supporter-enabled executable bit for bit. The public repository +must not be represented as complete corresponding source for that private +component. + +## Models + +Every bundled or downloadable model has separate artifact-specific terms and +a SHA-256 entry in `assets/THIRD_PARTY_MODELS.md`. Jasna-trained RF-DETR +detectors are licensed under Apache-2.0. Lada models retain AGPL-3.0. Encrypted +supporter models are proprietary and may be used only under their stated +supporter-model terms. diff --git a/NOTICE b/NOTICE new file mode 100644 index 00000000..f6bb073e --- /dev/null +++ b/NOTICE @@ -0,0 +1,19 @@ +Jasna +Copyright 2026 Kruk2 and contributors + +Except where an individual file or bundled component states otherwise, Jasna +is licensed under the GNU Affero General Public License version 3.0. The full +license is in LICENSE. + +Jasna incorporates or distributes work from Lada, OpenMMLab MMagic, +Roboflow RF-DETR, PyAV, VALI, FFmpeg, VideoLAN, Python, and other projects. +Their copyright, attribution, source, and license information is retained in: + +- LICENSING.md +- RELEASE_SOURCES.md +- assets/THIRD_PARTY_LICENSES.md +- assets/THIRD_PARTY_MODELS.md +- assets/licenses/ + +Official binary packages also contain license files collected from the exact +installed Python distributions under licenses/python-packages/. diff --git a/README.ja.md b/README.ja.md index 5e853679..04875f1b 100644 --- a/README.ja.md +++ b/README.ja.md @@ -20,6 +20,7 @@ Jasna は無料です。支援者には、このプロジェクト用に訓練 - [初回実行](#初回実行) - [さらに詳しく](#さらに詳しく) - [ベンチマーク](#ベンチマーク) +- [ライセンス](#ライセンス) - [プロジェクトを支援する](#プロジェクトを支援する) - [謝辞](#謝辞) - [TODO](#todo) @@ -143,6 +144,13 @@ RAM の中央値/ピーク、計測方法、生データ、パフォーマンス 長時間動画の以前の結果は [従来のベンチマークレポート](benchmarks/2026-07-26_legacy_videos.md)（英語）にあります。 +## ライセンス + +Jasna の公開アプリケーションソースは AGPL-3.0 です。モデルと同梱依存関係には、 +それぞれのライセンスが適用されます。[ライセンスとソース公開範囲](LICENSING.md)、 +[モデルの通知](assets/THIRD_PARTY_MODELS.md)、 +[リリースソース対応表](RELEASE_SOURCES.md)をご覧ください。 + ## プロジェクトを支援する 支援は追加モデルの訓練に使われます。主に、GPU のレンタル代と、より大きなデータセットで訓練するための計算時間です。支援者には以下を解除するキーが提供されます: diff --git a/README.md b/README.md index f271b992..44857eb6 100644 --- a/README.md +++ b/README.md @@ -22,6 +22,7 @@ Jasna is free. Supporters get a key that unlocks the extra models trained for th - [First Run](#first-run) - [Learn More](#learn-more) - [Benchmarks](#benchmarks) +- [Licensing](#licensing) - [Supporting the Project](#supporting-the-project) - [Acknowledgments](#acknowledgments) - [TODO](#todo) @@ -146,6 +147,13 @@ the [full benchmark report](benchmarks/2026-07-26_release_matrix.md). Older full-length video results are in the [legacy benchmark report](benchmarks/2026-07-26_legacy_videos.md). +## Licensing + +Jasna's public application source is AGPL-3.0. Models and bundled dependencies +retain their own terms. See [licensing and source availability](LICENSING.md), +[model notices](assets/THIRD_PARTY_MODELS.md), and the +[release source correspondence](RELEASE_SOURCES.md). + ## Supporting the Project Support pays for training extra models, mainly GPU rental and compute time for larger datasets. Supporters get a key that unlocks: diff --git a/README.zh.md b/README.zh.md index 04d1539d..d28fb278 100644 --- a/README.zh.md +++ b/README.zh.md @@ -20,6 +20,7 @@ Jasna 是免费的。支持者会获得一个密钥，用于解锁为本项目 - [首次运行](#首次运行) - [了解更多](#了解更多) - [基准测试](#基准测试) +- [许可证](#许可证) - [支持本项目](#支持本项目) - [致谢](#致谢) - [TODO](#todo) @@ -137,6 +138,13 @@ RAM 的中位数/峰值、测试方法、原始数据及性能提交对比，请 较早的完整长视频结果见 [旧版基准测试报告](benchmarks/2026-07-26_legacy_videos.md)（英文）。 +## 许可证 + +Jasna 的公开应用程序源代码采用 AGPL-3.0。模型和捆绑依赖项保留各自的许可条款。 +请参阅[许可与源代码范围](LICENSING.md)、 +[模型声明](assets/THIRD_PARTY_MODELS.md)和 +[发行版源代码对应表](RELEASE_SOURCES.md)。 + ## 支持本项目 支持用于训练额外模型，主要是租用 GPU，以及在更大数据集上训练所需的计算时间。支持者会获得一个密钥，用于解锁: diff --git a/RELEASE_SOURCES.md b/RELEASE_SOURCES.md new file mode 100644 index 00000000..9b59056c --- /dev/null +++ b/RELEASE_SOURCES.md @@ -0,0 +1,79 @@ +# v0.11.0 release source correspondence + +This file records the custom and patched inputs used by official v0.11.0 +packages. The release tag is the source revision for Jasna itself. Release +packages also include `LICENSE`, `LICENSING.md`, `NOTICE`, the third-party +notices, and full license texts. + +## Source revisions + +| Component | Revision or source | +| --- | --- | +| Jasna | Git tag `v0.11.0`; the tag and package version must match | +| Protection component | The private gitlink revision recorded by the `v0.11.0` tag; source is not published | +| PyAV | [`f6f0a5e3d975aab851e12ca881bdadfdd8f9ec74`](https://github.com/PyAV-Org/PyAV/commit/f6f0a5e3d975aab851e12ca881bdadfdd8f9ec74) | +| VALI | [`0e5c01ee57222a8ef4f9c7591284c4324bb3fba5`](https://codeberg.org/Kruk2/vali/commit/0e5c01ee57222a8ef4f9c7591284c4324bb3fba5) in the public Kruk2 fork | +| RF-DETR | [`1.8.3` / `3bd6bffbcb13cac3a5b1c37da5a0fd5453b50c86`](https://github.com/roboflow/rf-detr/commit/3bd6bffbcb13cac3a5b1c37da5a0fd5453b50c86) | +| MMagic compatibility patch | [`patches/fix_loading_mmengine_weights_on_torch26_and_higher.diff`](patches/fix_loading_mmengine_weights_on_torch26_and_higher.diff) | + +The VALI fork contains the complete modified source, including Jasna's decode +status, corrupt-packet handling, timestamp, and decoder-construction changes. +It is licensed under Apache-2.0. + +## FFmpeg in the PyAV and VALI wheels + +The Linux wheels link against the +[`pyav-ffmpeg` 8.1.2-1](https://github.com/PyAV-Org/pyav-ffmpeg/releases/tag/8.1.2-1) +build of FFmpeg 8.1.2. The source is FFmpeg tag `n8.1.2`, commit +[`38b88335f99e76ed89ff3c93f877fdefce736c13`](https://github.com/FFmpeg/FFmpeg/commit/38b88335f99e76ed89ff3c93f877fdefce736c13). +The exact build scripts are in the pyav-ffmpeg tag. They request x264/x265 and +`--enable-version3`; because the wheel vendors and links the GPL-licensed x264 +library, Jasna distributes the combined wheel payload under GPLv3. FFmpeg's +own runtime metadata identifies its libraries as LGPLv3-or-later. + +Linux x86-64 vendor archive: + +```text +https://github.com/PyAV-Org/pyav-ffmpeg/releases/download/8.1.2-1/ffmpeg-manylinux-x86_64.tar.gz +SHA-256 665b4903251dc753725e9cf64615c1af7708b932fb26bf16356acc77d76e95f8 +``` + +## Bundled FFmpeg command-line tools + +Linux releases use the BtbN `linux64-gpl-8.1` build +`ffmpeg-n8.1.2-34-g9b6c8969e0`. Its FFmpeg source is commit +[`9b6c8969e05b4f0b29f0f85cd501be6b3e582e6b`](https://github.com/FFmpeg/FFmpeg/commit/9b6c8969e05b4f0b29f0f85cd501be6b3e582e6b). +The BtbN GPL variant uses `--enable-gpl --enable-version3 --disable-debug`; +the complete build scripts are in +[BtbN/FFmpeg-Builds](https://github.com/BtbN/FFmpeg-Builds). + +```text +https://github.com/BtbN/FFmpeg-Builds/releases/download/autobuild-2026-07-31-14-10/ffmpeg-n8.1.2-34-g9b6c8969e0-linux64-gpl-8.1.tar.xz +SHA-256 09fc77be269c7053e438b7e96548e4af97604faf96a42c4a3c56a1ad74c22c0a +``` + +Windows packages use the GPLv3 FFmpeg 8 installation selected by the release +builder. Before publishing, its `ffmpeg -version` output, source revision, +configure line, and archive hash must be added to the release notes if they +differ from the Linux build above. + +## Python and build environment + +Linux releases build CPython 3.13.14 from the official source archive: + +```text +https://www.python.org/ftp/python/3.13.14/Python-3.13.14.tar.xz +SHA-256 639e43243c620a308f968213df9e00f2f8f62332f7adbaa7a7eeb9783057c690 +``` + +The application dependencies and exact direct version constraints are in +`pyproject.toml`. Binary packages copy the license files from every installed +Python distribution into `licenses/python-packages/`. + +## Reproducibility limit + +The public tag contains Jasna's public application source, public patches, and +the source correspondence above. The final whole-application build driver and +the protection implementation remain in the private protection project. +Therefore a public checkout can reproduce the free workflows but cannot +reproduce the supporter-enabled official executable bit for bit. diff --git a/assets/THIRD_PARTY_LICENSES.md b/assets/THIRD_PARTY_LICENSES.md index 4d9c34e6..731dae7b 100644 --- a/assets/THIRD_PARTY_LICENSES.md +++ b/assets/THIRD_PARTY_LICENSES.md @@ -1,17 +1,56 @@ # Third-party runtime notices -## libVLC +Official packages include the full texts named below in `licenses/`. +License files supplied by installed Python distributions are copied into +`licenses/python-packages/` during the release build. Exact custom source +revisions and build configurations are in `RELEASE_SOURCES.md`. -Jasna release packages include unmodified libVLC 3 runtime libraries and -plugins from the VideoLAN project. +## Bundled native components -- Project and source: https://www.videolan.org/vlc/ -- License: GNU Lesser General Public License, version 2.1 or later -- License text: https://www.gnu.org/licenses/old-licenses/lgpl-2.1.html +| Component | License | Source | +| --- | --- | --- | +| FFmpeg libraries in the PyAV and VALI wheels, including vendored x264 | GPLv3 for the combined payload; FFmpeg reports LGPLv3-or-later | [FFmpeg 8.1.2](https://github.com/FFmpeg/FFmpeg/tree/n8.1.2), built by [pyav-ffmpeg 8.1.2-1](https://github.com/PyAV-Org/pyav-ffmpeg/releases/tag/8.1.2-1) | +| Bundled `ffmpeg` and `ffprobe` tools | GPLv3 | [FFmpeg commit `9b6c8969e0`](https://github.com/FFmpeg/FFmpeg/commit/9b6c8969e05b4f0b29f0f85cd501be6b3e582e6b), built by [BtbN/FFmpeg-Builds](https://github.com/BtbN/FFmpeg-Builds) | +| VALI / python_vali | Apache-2.0 | [Kruk2/vali](https://codeberg.org/Kruk2/vali) | +| libVLC 3 and plugins | LGPL-2.1-or-later | [VideoLAN VLC](https://www.videolan.org/vlc/) | +| CPython 3.13 | Python-2.0 | [CPython](https://github.com/python/cpython) | -## python-vlc +The command-line build explicitly enables GPLv3. The shared wheel payload +vendors GPL-licensed x264 alongside FFmpeg libraries that report +LGPLv3-or-later, so Jasna applies GPLv3 to the combined payload. Exact source +commits, dependency versions, archive hashes, and build flags are in +`RELEASE_SOURCES.md` and the linked pyav-ffmpeg tag. -Jasna uses python-vlc 3.0.21203, the Python bindings for libVLC. +## Key Python components -- Project and source: https://github.com/oaubert/python-vlc -- License: GNU Lesser General Public License, version 2.1 or later +| Component | License | Source | +| --- | --- | --- | +| PyAV | BSD-3-Clause | [PyAV](https://github.com/PyAV-Org/PyAV) | +| python-vlc | LGPL-2.1-or-later | [python-vlc](https://github.com/oaubert/python-vlc) | +| PyTorch and torchvision | BSD-3-Clause | [PyTorch](https://github.com/pytorch/pytorch), [torchvision](https://github.com/pytorch/vision) | +| RF-DETR | Apache-2.0 | [roboflow/rf-detr](https://github.com/roboflow/rf-detr) | +| MMagic and MMEngine | Apache-2.0 | [OpenMMLab MMagic](https://github.com/open-mmlab/mmagic), [MMEngine](https://github.com/open-mmlab/mmengine) | +| diffusers, accelerate, transformers, huggingface-hub | Apache-2.0 | [Hugging Face](https://github.com/huggingface) | +| OpenCV | Apache-2.0 | [OpenCV](https://github.com/opencv/opencv) | +| Ultralytics | AGPL-3.0 | [ultralytics](https://github.com/ultralytics/ultralytics) | + +This table highlights the components most directly used or patched by Jasna; +the package-specific license directory in each binary release is the complete +notice set generated from that release environment. + +## Vendored source + +Jasna's `jasna/models/basicvsrpp/mmagic/` directory is an inference-only +subset derived from OpenMMLab MMagic. Its existing copyright headers and +Apache-2.0 terms are retained. + +Jasna also contains code derived from +[Lada](https://codeberg.org/ladaapp/lada), licensed under AGPL-3.0. Jasna's +root `LICENSE` contains the complete AGPL-3.0 text. + +## Proprietary vendor components + +NVIDIA TensorRT, CUDA runtime libraries, RTX Video Effects, and AMD ROCm +runtime libraries retain their vendor terms. Their wheel/package license files +are copied into `licenses/python-packages/` when present. These components +are not relicensed by Jasna. diff --git a/assets/THIRD_PARTY_MODELS.md b/assets/THIRD_PARTY_MODELS.md index 32d9a501..6cc1326e 100644 --- a/assets/THIRD_PARTY_MODELS.md +++ b/assets/THIRD_PARTY_MODELS.md @@ -1,6 +1,63 @@ -# Third-party models +# Model licenses and provenance -## ZeLeFans VR Mosaic Detection v2 (accurate) +Hashes below identify the v0.11.0 candidate artifacts. Release preparation +must fail or update this file when a bundled model changes. + +## Lada BasicVSR++ mosaic restoration + +- Jasna model name: `basicvsrpp` +- File: `lada_mosaic_restoration_model_generic_v1.2.pth` +- Upstream: +- Upstream revision: `3bfd69ffc21518bde80ba6b61696d51efd0a398b` +- SHA-256: `d404152576ce64fb5b2f315c03062709dac4f5f8548934866cd01c823c8104ee` +- License: AGPL-3.0 +- Copyright: ladaapp and contributors + +This checkpoint is redistributed unmodified. + +## Lada YOLO v4 fast mosaic detection + +- Jasna model name: `lada-yolo-v4` +- File: `lada_mosaic_detection_model_v4_fast.pt` +- Upstream: +- Upstream revision: `404620fe2f6b72657b92f76e62af914c8b3ee686` +- SHA-256: `9a6b660d1d3e3797d39515e08b0e72fcc59815f38279faa7a4ab374ab2c1e3b4` +- License: AGPL-3.0 +- Copyright: ladaapp and contributors + +This checkpoint is redistributed unmodified. + +## Jasna RF-DETR v6 + +- Jasna model name: `rfdetr-v6` +- NVIDIA file: `rfdetr-v6.onnx` +- NVIDIA SHA-256: `b6555cfce325d1d8bc413422cd46f3a453511a246c6f29ce652382998049d825` +- AMD file: `rfdetr-v6.pt` +- AMD SHA-256: `f10bedc4d105c2721e4259b8680203d51f344f73e55e85710d915619f5731b55` +- Architecture/source: [RF-DETR 1.8.3](https://github.com/roboflow/rf-detr/tree/3bd6bffbcb13cac3a5b1c37da5a0fd5453b50c86) +- License: Apache-2.0 +- Copyright: 2026 Kruk2 + +This project-trained RF-DETR Seg Medium checkpoint uses classes +`Background` and `mosaic`. The ONNX and PyTorch files are deployment +formats of the same trained model. + +## Jasna RF-DETR VR v1 + +- Jasna model name: `rfdetr-vr-v1` +- NVIDIA file: `rfdetr-vr-v1.onnx` +- NVIDIA SHA-256: `6e2ed2043851dccb97f21deda38dc20ea2b8e265e682359752e815c600030a40` +- AMD file: `rfdetr-vr-v1.pt` +- AMD SHA-256: `55543c83911921ef79cd8cae8540bd25e34c7daf488e77f79d233d6926973a2e` +- Architecture/source: [RF-DETR 1.8.3](https://github.com/roboflow/rf-detr/tree/3bd6bffbcb13cac3a5b1c37da5a0fd5453b50c86) +- License: Apache-2.0 +- Copyright: 2026 Kruk2 + +This project-trained RF-DETR Seg Large checkpoint is trained for side-by-side +VR material. The ONNX and PyTorch files are deployment formats of the same +trained model. + +## ZeLeFans VR Mosaic Detection v2 accurate - Jasna model name: `zelefans-vr-yolo-v2` - Upstream: @@ -8,6 +65,18 @@ - Pinned revision: `0f65a21133335f9a4ec6fc5d7da8d3385bfdb8b1` - Upstream file: `lada_vr_mosaic_detection_model_v2_accurate.pt` - SHA-256: `91fe7a48b0e9edf51361918c8a30f752c64511005e643343a7382d951f3fe0f8` -- License: [Apache License 2.0](https://www.apache.org/licenses/LICENSE-2.0) +- License: Apache-2.0 + +This optional checkpoint is not bundled in the standard v0.11.0 package. + +## Supporter models + +`unet-4x.onnx.enc` and the encrypted SD 1.5 Jasna checkpoint are +project-trained, proprietary supporter models. Copyright 2026 Kruk2. They are +provided only for use with Jasna by a holder of a valid supporter key. +Redistribution, extraction, modification, and use outside Jasna are not +granted. -This checkpoint is redistributed unmodified in Jasna release packages. +The base v0.11.0 release may include `unet-4x.onnx.enc`; SD 1.5 is downloaded +separately when requested. The protection implementation and supporter-model +terms are separate from Jasna's AGPL-covered public application source. diff --git a/assets/licenses/Apache-2.0.txt b/assets/licenses/Apache-2.0.txt new file mode 100644 index 00000000..e6d1e128 --- /dev/null +++ b/assets/licenses/Apache-2.0.txt @@ -0,0 +1,145 @@ + Apache License + Version 2.0, January 2004 + http://www.apache.org/licenses/ + + TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION + + 1. Definitions. + + "License" shall mean the terms and conditions for use, reproduction, + and distribution as defined by Sections 1 through 9 of this document. + + "Licensor" shall mean the copyright owner or entity authorized by + the copyright owner that is granting the License. + + "Legal Entity" shall mean the union of the acting entity and all + other entities that control, are controlled by, or are under common + control with that entity. For the purposes of this definition, + "control" means (i) the power, direct or indirect, to cause the + direction or management of such entity, whether by contract or + otherwise, or (ii) ownership of fifty percent (50%) or more of the + outstanding shares, or (iii) beneficial ownership of such entity. + + "You" (or "Your") shall mean an individual or Legal Entity + exercising permissions granted by this License. + + "Source" form shall mean the preferred form for making modifications, + including but not limited to software source code, documentation + source, and configuration files. + + "Object" form shall mean any form resulting from mechanical + transformation or translation of a Source form, including but + not limited to compiled object code, generated documentation, + and conversions to other media types. + + "Work" shall mean the work of authorship, whether in Source or + Object form, made available under the License, as indicated by a + copyright notice that is included in or attached to the work + (an example is provided in the Appendix below). + + "Derivative Works" shall mean any work, whether in Source or Object + form, that is based on (or derived from) the Work and for which the + editorial revisions, annotations, elaborations, or other modifications + represent, as a whole, an original work of authorship. For the purposes + of this License, Derivative Works shall not include works that remain + separable from, or merely link (or bind by name) to the interfaces of, + the Work and Derivative Works thereof. + + "Contribution" shall mean any work of authorship, including + the original version of the Work and any modifications or additions + to that Work or Derivative Works thereof, that is intentionally + submitted to Licensor for inclusion in the Work by the copyright owner + or by an individual or Legal Entity authorized to submit on behalf of + the copyright owner. For the purposes of this definition, "submitted" + means any form of electronic, verbal, or written communication sent + to the Licensor or its representatives, including but not limited to + communication on electronic mailing lists, source code control systems, + and issue tracking systems that are managed by, or on behalf of, the + Licensor for the purpose of discussing and improving the Work, but + excluding communication that is conspicuously marked or otherwise + designated in writing by the copyright owner as "Not a Contribution." + + "Contributor" shall mean Licensor and any individual or Legal Entity + on behalf of whom a Contribution has been received by Licensor and + subsequently incorporated within the Work. + + 2. Grant of Copyright License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + copyright license to reproduce, prepare Derivative Works of, + publicly display, publicly perform, sublicense, and distribute the + Work and such Derivative Works in Source or Object form. + + 3. Grant of Patent License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + (except as stated in this section) patent license to make, have made, + use, offer to sell, sell, import, and otherwise transfer the Work, + where such license applies only to those patent claims licensable + by such Contributor that are necessarily infringed by their + Contribution(s) alone or by combination of their Contribution(s) + with the Work to which such Contribution(s) was submitted. If You + institute patent litigation against any entity (including a + cross-claim or counterclaim in a lawsuit) alleging that the Work + or a Contribution incorporated within the Work constitutes direct + or contributory patent infringement, then any patent licenses + granted to You under this License for that Work shall terminate + as of the date such litigation is filed. + + 4. Redistribution. You may reproduce and distribute copies of the + Work or Derivative Works thereof in any medium, with or without + modifications, and in Source or Object form, provided that You + meet the following conditions: + + (a) You must give any other recipients of the Work or + Derivative Works a copy of this License; and + + (b) You must cause any modified files to carry prominent notices + stating that You changed the files; and + + (c) You must retain, in the Source form of any Derivative Works + that You distribute, all copyright, patent, trademark, and + attribution notices from the Source form of the Work, + excluding those notices that do not pertain to any part of + the Derivative Works; and + + (d) If the Work includes a "NOTICE" text file as part of its + distribution, then any Derivative Works that You distribute must + include a readable copy of the attribution notices contained + within such NOTICE file, excluding those notices that do not + pertain to any part of the Derivative Works, in at least one + of the following places: within a NOTICE text file distributed + as part of the Derivative Works; within the Source form or + documentation, if provided along with the Derivative Works; or, + within a display generated by the Derivative Works, if and + wherever such third-party notices normally appear. The contents + of the NOTICE file are for informational purposes only and + do not modify the License. You may add Your own attribution + notices within Derivative Works that You distribute, alongside + or as an addendum to the NOTICE text from the Work, provided + that such additional attribution notices cannot be construed + as modifying the License. + + You may add Your own copyright statement to Your modifications and + may provide additional or different license terms and conditions + for use, reproduction, or distribution of Your modifications, or + for any such Derivative Works as a whole, provided Your use, + reproduction, and distribution of the Work otherwise complies with + the conditions stated in this License. + + 5. Submission of Contributions. Unless You explicitly state otherwise, + any Contribution intentionally submitted for inclusion in the Work + by You to the Licensor shall be under the terms and conditions of + this License, without any additional terms or conditions. + Notwithstanding the above, nothing herein shall supersede or modify + the terms of any separate license agreement you may have executed + with Licensor regarding such Contributions. + + 6. Trademarks. This License does not grant permission to use the trade + names, trademarks, service marks, or product names of the Licensor, + except as required for reasonable and customary use in describing the + origin of the Work and reproducing the content of the NOTICE file. + + 7. Disclaimer of Warranty. Unless required by applicable law or + agreed to in writing, Licensor provides the Work (and each +... (+57 lines) [see remaining: rtk proxy git '-C' '/home/dominik/tools/vali' show '0e5c01ee57222a8ef4f9c7591284c4324bb3fba5:LICENSE' | tail -n +145] diff --git a/assets/licenses/BSD-3-Clause.txt b/assets/licenses/BSD-3-Clause.txt new file mode 100644 index 00000000..9db1661e --- /dev/null +++ b/assets/licenses/BSD-3-Clause.txt @@ -0,0 +1,23 @@ +Copyright retained by original committers. All rights reserved. + +Redistribution and use in source and binary forms, with or without +modification, are permitted provided that the following conditions are met: + * Redistributions of source code must retain the above copyright + notice, this list of conditions and the following disclaimer. + * Redistributions in binary form must reproduce the above copyright + notice, this list of conditions and the following disclaimer in the + documentation and/or other materials provided with the distribution. + * Neither the name of the project nor the names of its contributors may be + used to endorse or promote products derived from this software without + specific prior written permission. + +THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDERS BE LIABLE FOR ANY DIRECT, +INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, +BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, +DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY +OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING +NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, +EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. diff --git a/assets/licenses/GPL-3.0.txt b/assets/licenses/GPL-3.0.txt new file mode 100644 index 00000000..94a9ed02 --- /dev/null +++ b/assets/licenses/GPL-3.0.txt @@ -0,0 +1,674 @@ + GNU GENERAL PUBLIC LICENSE + Version 3, 29 June 2007 + + Copyright (C) 2007 Free Software Foundation, Inc. + Everyone is permitted to copy and distribute verbatim copies + of this license document, but changing it is not allowed. + + Preamble + + The GNU General Public License is a free, copyleft license for +software and other kinds of works. + + The licenses for most software and other practical works are designed +to take away your freedom to share and change the works. By contrast, +the GNU General Public License is intended to guarantee your freedom to +share and change all versions of a program--to make sure it remains free +software for all its users. We, the Free Software Foundation, use the +GNU General Public License for most of our software; it applies also to +any other work released this way by its authors. You can apply it to +your programs, too. + + When we speak of free software, we are referring to freedom, not +price. Our General Public Licenses are designed to make sure that you +have the freedom to distribute copies of free software (and charge for +them if you wish), that you receive source code or can get it if you +want it, that you can change the software or use pieces of it in new +free programs, and that you know you can do these things. + + To protect your rights, we need to prevent others from denying you +these rights or asking you to surrender the rights. Therefore, you have +certain responsibilities if you distribute copies of the software, or if +you modify it: responsibilities to respect the freedom of others. + + For example, if you distribute copies of such a program, whether +gratis or for a fee, you must pass on to the recipients the same +freedoms that you received. You must make sure that they, too, receive +or can get the source code. And you must show them these terms so they +know their rights. + + Developers that use the GNU GPL protect your rights with two steps: +(1) assert copyright on the software, and (2) offer you this License +giving you legal permission to copy, distribute and/or modify it. + + For the developers' and authors' protection, the GPL clearly explains +that there is no warranty for this free software. For both users' and +authors' sake, the GPL requires that modified versions be marked as +changed, so that their problems will not be attributed erroneously to +authors of previous versions. + + Some devices are designed to deny users access to install or run +modified versions of the software inside them, although the manufacturer +can do so. This is fundamentally incompatible with the aim of +protecting users' freedom to change the software. The systematic +pattern of such abuse occurs in the area of products for individuals to +use, which is precisely where it is most unacceptable. Therefore, we +have designed this version of the GPL to prohibit the practice for those +products. If such problems arise substantially in other domains, we +stand ready to extend this provision to those domains in future versions +of the GPL, as needed to protect the freedom of users. + + Finally, every program is threatened constantly by software patents. +States should not allow patents to restrict development and use of +software on general-purpose computers, but in those that do, we wish to +avoid the special danger that patents applied to a free program could +make it effectively proprietary. To prevent this, the GPL assures that +patents cannot be used to render the program non-free. + + The precise terms and conditions for copying, distribution and +modification follow. + + TERMS AND CONDITIONS + + 0. Definitions. + + "This License" refers to version 3 of the GNU General Public License. + + "Copyright" also means copyright-like laws that apply to other kinds of +works, such as semiconductor masks. + + "The Program" refers to any copyrightable work licensed under this +License. Each licensee is addressed as "you". "Licensees" and +"recipients" may be individuals or organizations. + + To "modify" a work means to copy from or adapt all or part of the work +in a fashion requiring copyright permission, other than the making of an +exact copy. The resulting work is called a "modified version" of the +earlier work or a work "based on" the earlier work. + + A "covered work" means either the unmodified Program or a work based +on the Program. + + To "propagate" a work means to do anything with it that, without +permission, would make you directly or secondarily liable for +infringement under applicable copyright law, except executing it on a +computer or modifying a private copy. Propagation includes copying, +distribution (with or without modification), making available to the +public, and in some countries other activities as well. + + To "convey" a work means any kind of propagation that enables other +parties to make or receive copies. Mere interaction with a user through +a computer network, with no transfer of a copy, is not conveying. + + An interactive user interface displays "Appropriate Legal Notices" +to the extent that it includes a convenient and prominently visible +feature that (1) displays an appropriate copyright notice, and (2) +tells the user that there is no warranty for the work (except to the +extent that warranties are provided), that licensees may convey the +work under this License, and how to view a copy of this License. If +the interface presents a list of user commands or options, such as a +menu, a prominent item in the list meets this criterion. + + 1. Source Code. + + The "source code" for a work means the preferred form of the work +for making modifications to it. "Object code" means any non-source +form of a work. + + A "Standard Interface" means an interface that either is an official +standard defined by a recognized standards body, or, in the case of +interfaces specified for a particular programming language, one that +is widely used among developers working in that language. + + The "System Libraries" of an executable work include anything, other +than the work as a whole, that (a) is included in the normal form of +packaging a Major Component, but which is not part of that Major +Component, and (b) serves only to enable use of the work with that +Major Component, or to implement a Standard Interface for which an +implementation is available to the public in source code form. A +"Major Component", in this context, means a major essential component +(kernel, window system, and so on) of the specific operating system +(if any) on which the executable work runs, or a compiler used to +produce the work, or an object code interpreter used to run it. + + The "Corresponding Source" for a work in object code form means all +the source code needed to generate, install, and (for an executable +work) run the object code and to modify the work, including scripts to +control those activities. However, it does not include the work's +System Libraries, or general-purpose tools or generally available free +programs which are used unmodified in performing those activities but +which are not part of the work. For example, Corresponding Source +includes interface definition files associated with source files for +the work, and the source code for shared libraries and dynamically +linked subprograms that the work is specifically designed to require, +such as by intimate data communication or control flow between those +subprograms and other parts of the work. + + The Corresponding Source need not include anything that users +can regenerate automatically from other parts of the Corresponding +Source. + + The Corresponding Source for a work in source code form is that +same work. + + 2. Basic Permissions. + + All rights granted under this License are granted for the term of +copyright on the Program, and are irrevocable provided the stated +conditions are met. This License explicitly affirms your unlimited +permission to run the unmodified Program. The output from running a +covered work is covered by this License only if the output, given its +content, constitutes a covered work. This License acknowledges your +rights of fair use or other equivalent, as provided by copyright law. + + You may make, run and propagate covered works that you do not +convey, without conditions so long as your license otherwise remains +in force. You may convey covered works to others for the sole purpose +of having them make modifications exclusively for you, or provide you +with facilities for running those works, provided that you comply with +the terms of this License in conveying all material for which you do +not control copyright. Those thus making or running the covered works +for you must do so exclusively on your behalf, under your direction +and control, on terms that prohibit them from making any copies of +your copyrighted material outside their relationship with you. + + Conveying under any other circumstances is permitted solely under +the conditions stated below. Sublicensing is not allowed; section 10 +makes it unnecessary. + + 3. Protecting Users' Legal Rights From Anti-Circumvention Law. + + No covered work shall be deemed part of an effective technological +measure under any applicable law fulfilling obligations under article +11 of the WIPO copyright treaty adopted on 20 December 1996, or +similar laws prohibiting or restricting circumvention of such +measures. + + When you convey a covered work, you waive any legal power to forbid +circumvention of technological measures to the extent such circumvention +is effected by exercising rights under this License with respect to +the covered work, and you disclaim any intention to limit operation or +modification of the work as a means of enforcing, against the work's +users, your or third parties' legal rights to forbid circumvention of +technological measures. + + 4. Conveying Verbatim Copies. + + You may convey verbatim copies of the Program's source code as you +receive it, in any medium, provided that you conspicuously and +appropriately publish on each copy an appropriate copyright notice; +keep intact all notices stating that this License and any +non-permissive terms added in accord with section 7 apply to the code; +keep intact all notices of the absence of any warranty; and give all +recipients a copy of this License along with the Program. + + You may charge any price or no price for each copy that you convey, +and you may offer support or warranty protection for a fee. + + 5. Conveying Modified Source Versions. + + You may convey a work based on the Program, or the modifications to +produce it from the Program, in the form of source code under the +terms of section 4, provided that you also meet all of these conditions: + + a) The work must carry prominent notices stating that you modified + it, and giving a relevant date. + + b) The work must carry prominent notices stating that it is + released under this License and any conditions added under section + 7. This requirement modifies the requirement in section 4 to + "keep intact all notices". + + c) You must license the entire work, as a whole, under this + License to anyone who comes into possession of a copy. This + License will therefore apply, along with any applicable section 7 + additional terms, to the whole of the work, and all its parts, + regardless of how they are packaged. This License gives no + permission to license the work in any other way, but it does not + invalidate such permission if you have separately received it. + + d) If the work has interactive user interfaces, each must display + Appropriate Legal Notices; however, if the Program has interactive + interfaces that do not display Appropriate Legal Notices, your + work need not make them do so. + + A compilation of a covered work with other separate and independent +works, which are not by their nature extensions of the covered work, +and which are not combined with it such as to form a larger program, +in or on a volume of a storage or distribution medium, is called an +"aggregate" if the compilation and its resulting copyright are not +used to limit the access or legal rights of the compilation's users +beyond what the individual works permit. Inclusion of a covered work +in an aggregate does not cause this License to apply to the other +parts of the aggregate. + + 6. Conveying Non-Source Forms. + + You may convey a covered work in object code form under the terms +of sections 4 and 5, provided that you also convey the +machine-readable Corresponding Source under the terms of this License, +in one of these ways: + + a) Convey the object code in, or embodied in, a physical product + (including a physical distribution medium), accompanied by the + Corresponding Source fixed on a durable physical medium + customarily used for software interchange. + + b) Convey the object code in, or embodied in, a physical product + (including a physical distribution medium), accompanied by a + written offer, valid for at least three years and valid for as + long as you offer spare parts or customer support for that product + model, to give anyone who possesses the object code either (1) a + copy of the Corresponding Source for all the software in the + product that is covered by this License, on a durable physical + medium customarily used for software interchange, for a price no + more than your reasonable cost of physically performing this + conveying of source, or (2) access to copy the + Corresponding Source from a network server at no charge. + + c) Convey individual copies of the object code with a copy of the + written offer to provide the Corresponding Source. This + alternative is allowed only occasionally and noncommercially, and + only if you received the object code with such an offer, in accord + with subsection 6b. + + d) Convey the object code by offering access from a designated + place (gratis or for a charge), and offer equivalent access to the + Corresponding Source in the same way through the same place at no + further charge. You need not require recipients to copy the + Corresponding Source along with the object code. If the place to + copy the object code is a network server, the Corresponding Source + may be on a different server (operated by you or a third party) + that supports equivalent copying facilities, provided you maintain + clear directions next to the object code saying where to find the + Corresponding Source. Regardless of what server hosts the + Corresponding Source, you remain obligated to ensure that it is + available for as long as needed to satisfy these requirements. + + e) Convey the object code using peer-to-peer transmission, provided + you inform other peers where the object code and Corresponding + Source of the work are being offered to the general public at no + charge under subsection 6d. + + A separable portion of the object code, whose source code is excluded +from the Corresponding Source as a System Library, need not be +included in conveying the object code work. + + A "User Product" is either (1) a "consumer product", which means any +tangible personal property which is normally used for personal, family, +or household purposes, or (2) anything designed or sold for incorporation +into a dwelling. In determining whether a product is a consumer product, +doubtful cases shall be resolved in favor of coverage. For a particular +product received by a particular user, "normally used" refers to a +typical or common use of that class of product, regardless of the status +of the particular user or of the way in which the particular user +actually uses, or expects or is expected to use, the product. A product +is a consumer product regardless of whether the product has substantial +commercial, industrial or non-consumer uses, unless such uses represent +the only significant mode of use of the product. + + "Installation Information" for a User Product means any methods, +procedures, authorization keys, or other information required to install +and execute modified versions of a covered work in that User Product from +a modified version of its Corresponding Source. The information must +suffice to ensure that the continued functioning of the modified object +code is in no case prevented or interfered with solely because +modification has been made. + + If you convey an object code work under this section in, or with, or +specifically for use in, a User Product, and the conveying occurs as +part of a transaction in which the right of possession and use of the +User Product is transferred to the recipient in perpetuity or for a +fixed term (regardless of how the transaction is characterized), the +Corresponding Source conveyed under this section must be accompanied +by the Installation Information. But this requirement does not apply +if neither you nor any third party retains the ability to install +modified object code on the User Product (for example, the work has +been installed in ROM). + + The requirement to provide Installation Information does not include a +requirement to continue to provide support service, warranty, or updates +for a work that has been modified or installed by the recipient, or for +the User Product in which it has been modified or installed. Access to a +network may be denied when the modification itself materially and +adversely affects the operation of the network or violates the rules and +protocols for communication across the network. + + Corresponding Source conveyed, and Installation Information provided, +in accord with this section must be in a format that is publicly +documented (and with an implementation available to the public in +source code form), and must require no special password or key for +unpacking, reading or copying. + + 7. Additional Terms. + + "Additional permissions" are terms that supplement the terms of this +License by making exceptions from one or more of its conditions. +Additional permissions that are applicable to the entire Program shall +be treated as though they were included in this License, to the extent +that they are valid under applicable law. If additional permissions +apply only to part of the Program, that part may be used separately +under those permissions, but the entire Program remains governed by +this License without regard to the additional permissions. + + When you convey a copy of a covered work, you may at your option +remove any additional permissions from that copy, or from any part of +it. (Additional permissions may be written to require their own +removal in certain cases when you modify the work.) You may place +additional permissions on material, added by you to a covered work, +for which you have or can give appropriate copyright permission. + + Notwithstanding any other provision of this License, for material you +add to a covered work, you may (if authorized by the copyright holders of +that material) supplement the terms of this License with terms: + + a) Disclaiming warranty or limiting liability differently from the + terms of sections 15 and 16 of this License; or + + b) Requiring preservation of specified reasonable legal notices or + author attributions in that material or in the Appropriate Legal + Notices displayed by works containing it; or + + c) Prohibiting misrepresentation of the origin of that material, or + requiring that modified versions of such material be marked in + reasonable ways as different from the original version; or + + d) Limiting the use for publicity purposes of names of licensors or + authors of the material; or + + e) Declining to grant rights under trademark law for use of some + trade names, trademarks, or service marks; or + + f) Requiring indemnification of licensors and authors of that + material by anyone who conveys the material (or modified versions of + it) with contractual assumptions of liability to the recipient, for + any liability that these contractual assumptions directly impose on + those licensors and authors. + + All other non-permissive additional terms are considered "further +restrictions" within the meaning of section 10. If the Program as you +received it, or any part of it, contains a notice stating that it is +governed by this License along with a term that is a further +restriction, you may remove that term. If a license document contains +a further restriction but permits relicensing or conveying under this +License, you may add to a covered work material governed by the terms +of that license document, provided that the further restriction does +not survive such relicensing or conveying. + + If you add terms to a covered work in accord with this section, you +must place, in the relevant source files, a statement of the +additional terms that apply to those files, or a notice indicating +where to find the applicable terms. + + Additional terms, permissive or non-permissive, may be stated in the +form of a separately written license, or stated as exceptions; +the above requirements apply either way. + + 8. Termination. + + You may not propagate or modify a covered work except as expressly +provided under this License. Any attempt otherwise to propagate or +modify it is void, and will automatically terminate your rights under +this License (including any patent licenses granted under the third +paragraph of section 11). + + However, if you cease all violation of this License, then your +license from a particular copyright holder is reinstated (a) +provisionally, unless and until the copyright holder explicitly and +finally terminates your license, and (b) permanently, if the copyright +holder fails to notify you of the violation by some reasonable means +prior to 60 days after the cessation. + + Moreover, your license from a particular copyright holder is +reinstated permanently if the copyright holder notifies you of the +violation by some reasonable means, this is the first time you have +received notice of violation of this License (for any work) from that +copyright holder, and you cure the violation prior to 30 days after +your receipt of the notice. + + Termination of your rights under this section does not terminate the +licenses of parties who have received copies or rights from you under +this License. If your rights have been terminated and not permanently +reinstated, you do not qualify to receive new licenses for the same +material under section 10. + + 9. Acceptance Not Required for Having Copies. + + You are not required to accept this License in order to receive or +run a copy of the Program. Ancillary propagation of a covered work +occurring solely as a consequence of using peer-to-peer transmission +to receive a copy likewise does not require acceptance. However, +nothing other than this License grants you permission to propagate or +modify any covered work. These actions infringe copyright if you do +not accept this License. Therefore, by modifying or propagating a +covered work, you indicate your acceptance of this License to do so. + + 10. Automatic Licensing of Downstream Recipients. + + Each time you convey a covered work, the recipient automatically +receives a license from the original licensors, to run, modify and +propagate that work, subject to this License. You are not responsible +for enforcing compliance by third parties with this License. + + An "entity transaction" is a transaction transferring control of an +organization, or substantially all assets of one, or subdividing an +organization, or merging organizations. If propagation of a covered +work results from an entity transaction, each party to that +transaction who receives a copy of the work also receives whatever +licenses to the work the party's predecessor in interest had or could +give under the previous paragraph, plus a right to possession of the +Corresponding Source of the work from the predecessor in interest, if +the predecessor has it or can get it with reasonable efforts. + + You may not impose any further restrictions on the exercise of the +rights granted or affirmed under this License. For example, you may +not impose a license fee, royalty, or other charge for exercise of +rights granted under this License, and you may not initiate litigation +(including a cross-claim or counterclaim in a lawsuit) alleging that +any patent claim is infringed by making, using, selling, offering for +sale, or importing the Program or any portion of it. + + 11. Patents. + + A "contributor" is a copyright holder who authorizes use under this +License of the Program or a work on which the Program is based. The +work thus licensed is called the contributor's "contributor version". + + A contributor's "essential patent claims" are all patent claims +owned or controlled by the contributor, whether already acquired or +hereafter acquired, that would be infringed by some manner, permitted +by this License, of making, using, or selling its contributor version, +but do not include claims that would be infringed only as a +consequence of further modification of the contributor version. For +purposes of this definition, "control" includes the right to grant +patent sublicenses in a manner consistent with the requirements of +this License. + + Each contributor grants you a non-exclusive, worldwide, royalty-free +patent license under the contributor's essential patent claims, to +make, use, sell, offer for sale, import and otherwise run, modify and +propagate the contents of its contributor version. + + In the following three paragraphs, a "patent license" is any express +agreement or commitment, however denominated, not to enforce a patent +(such as an express permission to practice a patent or covenant not to +sue for patent infringement). To "grant" such a patent license to a +party means to make such an agreement or commitment not to enforce a +patent against the party. + + If you convey a covered work, knowingly relying on a patent license, +and the Corresponding Source of the work is not available for anyone +to copy, free of charge and under the terms of this License, through a +publicly available network server or other readily accessible means, +then you must either (1) cause the Corresponding Source to be so +available, or (2) arrange to deprive yourself of the benefit of the +patent license for this particular work, or (3) arrange, in a manner +consistent with the requirements of this License, to extend the patent +license to downstream recipients. "Knowingly relying" means you have +actual knowledge that, but for the patent license, your conveying the +covered work in a country, or your recipient's use of the covered work +in a country, would infringe one or more identifiable patents in that +country that you have reason to believe are valid. + + If, pursuant to or in connection with a single transaction or +arrangement, you convey, or propagate by procuring conveyance of, a +covered work, and grant a patent license to some of the parties +receiving the covered work authorizing them to use, propagate, modify +or convey a specific copy of the covered work, then the patent license +you grant is automatically extended to all recipients of the covered +work and works based on it. + + A patent license is "discriminatory" if it does not include within +the scope of its coverage, prohibits the exercise of, or is +conditioned on the non-exercise of one or more of the rights that are +specifically granted under this License. You may not convey a covered +work if you are a party to an arrangement with a third party that is +in the business of distributing software, under which you make payment +to the third party based on the extent of your activity of conveying +the work, and under which the third party grants, to any of the +parties who would receive the covered work from you, a discriminatory +patent license (a) in connection with copies of the covered work +conveyed by you (or copies made from those copies), or (b) primarily +for and in connection with specific products or compilations that +contain the covered work, unless you entered into that arrangement, +or that patent license was granted, prior to 28 March 2007. + + Nothing in this License shall be construed as excluding or limiting +any implied license or other defenses to infringement that may +otherwise be available to you under applicable patent law. + + 12. No Surrender of Others' Freedom. + + If conditions are imposed on you (whether by court order, agreement or +otherwise) that contradict the conditions of this License, they do not +excuse you from the conditions of this License. If you cannot convey a +covered work so as to satisfy simultaneously your obligations under this +License and any other pertinent obligations, then as a consequence you may +not convey it at all. For example, if you agree to terms that obligate you +to collect a royalty for further conveying from those to whom you convey +the Program, the only way you could satisfy both those terms and this +License would be to refrain entirely from conveying the Program. + + 13. Use with the GNU Affero General Public License. + + Notwithstanding any other provision of this License, you have +permission to link or combine any covered work with a work licensed +under version 3 of the GNU Affero General Public License into a single +combined work, and to convey the resulting work. The terms of this +License will continue to apply to the part which is the covered work, +but the special requirements of the GNU Affero General Public License, +section 13, concerning interaction through a network will apply to the +combination as such. + + 14. Revised Versions of this License. + + The Free Software Foundation may publish revised and/or new versions of +the GNU General Public License from time to time. Such new versions will +be similar in spirit to the present version, but may differ in detail to +address new problems or concerns. + + Each version is given a distinguishing version number. If the +Program specifies that a certain numbered version of the GNU General +Public License "or any later version" applies to it, you have the +option of following the terms and conditions either of that numbered +version or of any later version published by the Free Software +Foundation. If the Program does not specify a version number of the +GNU General Public License, you may choose any version ever published +by the Free Software Foundation. + + If the Program specifies that a proxy can decide which future +versions of the GNU General Public License can be used, that proxy's +public statement of acceptance of a version permanently authorizes you +to choose that version for the Program. + + Later license versions may give you additional or different +permissions. However, no additional obligations are imposed on any +author or copyright holder as a result of your choosing to follow a +later version. + + 15. Disclaimer of Warranty. + + THERE IS NO WARRANTY FOR THE PROGRAM, TO THE EXTENT PERMITTED BY +APPLICABLE LAW. EXCEPT WHEN OTHERWISE STATED IN WRITING THE COPYRIGHT +HOLDERS AND/OR OTHER PARTIES PROVIDE THE PROGRAM "AS IS" WITHOUT WARRANTY +OF ANY KIND, EITHER EXPRESSED OR IMPLIED, INCLUDING, BUT NOT LIMITED TO, +THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR +PURPOSE. THE ENTIRE RISK AS TO THE QUALITY AND PERFORMANCE OF THE PROGRAM +IS WITH YOU. SHOULD THE PROGRAM PROVE DEFECTIVE, YOU ASSUME THE COST OF +ALL NECESSARY SERVICING, REPAIR OR CORRECTION. + + 16. Limitation of Liability. + + IN NO EVENT UNLESS REQUIRED BY APPLICABLE LAW OR AGREED TO IN WRITING +WILL ANY COPYRIGHT HOLDER, OR ANY OTHER PARTY WHO MODIFIES AND/OR CONVEYS +THE PROGRAM AS PERMITTED ABOVE, BE LIABLE TO YOU FOR DAMAGES, INCLUDING ANY +GENERAL, SPECIAL, INCIDENTAL OR CONSEQUENTIAL DAMAGES ARISING OUT OF THE +USE OR INABILITY TO USE THE PROGRAM (INCLUDING BUT NOT LIMITED TO LOSS OF +DATA OR DATA BEING RENDERED INACCURATE OR LOSSES SUSTAINED BY YOU OR THIRD +PARTIES OR A FAILURE OF THE PROGRAM TO OPERATE WITH ANY OTHER PROGRAMS), +EVEN IF SUCH HOLDER OR OTHER PARTY HAS BEEN ADVISED OF THE POSSIBILITY OF +SUCH DAMAGES. + + 17. Interpretation of Sections 15 and 16. + + If the disclaimer of warranty and limitation of liability provided +above cannot be given local legal effect according to their terms, +reviewing courts shall apply local law that most closely approximates +an absolute waiver of all civil liability in connection with the +Program, unless a warranty or assumption of liability accompanies a +copy of the Program in return for a fee. + + END OF TERMS AND CONDITIONS + + How to Apply These Terms to Your New Programs + + If you develop a new program, and you want it to be of the greatest +possible use to the public, the best way to achieve this is to make it +free software which everyone can redistribute and change under these terms. + + To do so, attach the following notices to the program. It is safest +to attach them to the start of each source file to most effectively +state the exclusion of warranty; and each file should have at least +the "copyright" line and a pointer to where the full notice is found. + + + Copyright (C) + + This program is free software: you can redistribute it and/or modify + it under the terms of the GNU General Public License as published by + the Free Software Foundation, either version 3 of the License, or + (at your option) any later version. + + This program is distributed in the hope that it will be useful, + but WITHOUT ANY WARRANTY; without even the implied warranty of + MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the + GNU General Public License for more details. + + You should have received a copy of the GNU General Public License + along with this program. If not, see . + +Also add information on how to contact you by electronic and paper mail. + + If the program does terminal interaction, make it output a short +notice like this when it starts in an interactive mode: + + Copyright (C) + This program comes with ABSOLUTELY NO WARRANTY; for details type `show w'. + This is free software, and you are welcome to redistribute it + under certain conditions; type `show c' for details. + +The hypothetical commands `show w' and `show c' should show the appropriate +parts of the General Public License. Of course, your program's commands +might be different; for a GUI interface, you would use an "about box". + + You should also get your employer (if you work as a programmer) or school, +if any, to sign a "copyright disclaimer" for the program, if necessary. +For more information on this, and how to apply and follow the GNU GPL, see +. + + The GNU General Public License does not permit incorporating your program +into proprietary programs. If your program is a subroutine library, you +may consider it more useful to permit linking proprietary applications with +the library. If this is what you want to do, use the GNU Lesser General +Public License instead of this License. But first, please read +. diff --git a/assets/licenses/LGPL-2.1.txt b/assets/licenses/LGPL-2.1.txt new file mode 100644 index 00000000..4362b491 --- /dev/null +++ b/assets/licenses/LGPL-2.1.txt @@ -0,0 +1,502 @@ + GNU LESSER GENERAL PUBLIC LICENSE + Version 2.1, February 1999 + + Copyright (C) 1991, 1999 Free Software Foundation, Inc. + 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA + Everyone is permitted to copy and distribute verbatim copies + of this license document, but changing it is not allowed. + +[This is the first released version of the Lesser GPL. It also counts + as the successor of the GNU Library Public License, version 2, hence + the version number 2.1.] + + Preamble + + The licenses for most software are designed to take away your +freedom to share and change it. By contrast, the GNU General Public +Licenses are intended to guarantee your freedom to share and change +free software--to make sure the software is free for all its users. + + This license, the Lesser General Public License, applies to some +specially designated software packages--typically libraries--of the +Free Software Foundation and other authors who decide to use it. You +can use it too, but we suggest you first think carefully about whether +this license or the ordinary General Public License is the better +strategy to use in any particular case, based on the explanations below. + + When we speak of free software, we are referring to freedom of use, +not price. Our General Public Licenses are designed to make sure that +you have the freedom to distribute copies of free software (and charge +for this service if you wish); that you receive source code or can get +it if you want it; that you can change the software and use pieces of +it in new free programs; and that you are informed that you can do +these things. + + To protect your rights, we need to make restrictions that forbid +distributors to deny you these rights or to ask you to surrender these +rights. These restrictions translate to certain responsibilities for +you if you distribute copies of the library or if you modify it. + + For example, if you distribute copies of the library, whether gratis +or for a fee, you must give the recipients all the rights that we gave +you. You must make sure that they, too, receive or can get the source +code. If you link other code with the library, you must provide +complete object files to the recipients, so that they can relink them +with the library after making changes to the library and recompiling +it. And you must show them these terms so they know their rights. + + We protect your rights with a two-step method: (1) we copyright the +library, and (2) we offer you this license, which gives you legal +permission to copy, distribute and/or modify the library. + + To protect each distributor, we want to make it very clear that +there is no warranty for the free library. Also, if the library is +modified by someone else and passed on, the recipients should know +that what they have is not the original version, so that the original +author's reputation will not be affected by problems that might be +introduced by others. + + Finally, software patents pose a constant threat to the existence of +any free program. We wish to make sure that a company cannot +effectively restrict the users of a free program by obtaining a +restrictive license from a patent holder. Therefore, we insist that +any patent license obtained for a version of the library must be +consistent with the full freedom of use specified in this license. + + Most GNU software, including some libraries, is covered by the +ordinary GNU General Public License. This license, the GNU Lesser +General Public License, applies to certain designated libraries, and +is quite different from the ordinary General Public License. We use +this license for certain libraries in order to permit linking those +libraries into non-free programs. + + When a program is linked with a library, whether statically or using +a shared library, the combination of the two is legally speaking a +combined work, a derivative of the original library. The ordinary +General Public License therefore permits such linking only if the +entire combination fits its criteria of freedom. The Lesser General +Public License permits more lax criteria for linking other code with +the library. + + We call this license the "Lesser" General Public License because it +does Less to protect the user's freedom than the ordinary General +Public License. It also provides other free software developers Less +of an advantage over competing non-free programs. These disadvantages +are the reason we use the ordinary General Public License for many +libraries. However, the Lesser license provides advantages in certain +special circumstances. + + For example, on rare occasions, there may be a special need to +encourage the widest possible use of a certain library, so that it becomes +a de-facto standard. To achieve this, non-free programs must be +allowed to use the library. A more frequent case is that a free +library does the same job as widely used non-free libraries. In this +case, there is little to gain by limiting the free library to free +software only, so we use the Lesser General Public License. + + In other cases, permission to use a particular library in non-free +programs enables a greater number of people to use a large body of +free software. For example, permission to use the GNU C Library in +non-free programs enables many more people to use the whole GNU +operating system, as well as its variant, the GNU/Linux operating +system. + + Although the Lesser General Public License is Less protective of the +users' freedom, it does ensure that the user of a program that is +linked with the Library has the freedom and the wherewithal to run +that program using a modified version of the Library. + + The precise terms and conditions for copying, distribution and +modification follow. Pay close attention to the difference between a +"work based on the library" and a "work that uses the library". The +former contains code derived from the library, whereas the latter must +be combined with the library in order to run. + + GNU LESSER GENERAL PUBLIC LICENSE + TERMS AND CONDITIONS FOR COPYING, DISTRIBUTION AND MODIFICATION + + 0. This License Agreement applies to any software library or other +program which contains a notice placed by the copyright holder or +other authorized party saying it may be distributed under the terms of +this Lesser General Public License (also called "this License"). +Each licensee is addressed as "you". + + A "library" means a collection of software functions and/or data +prepared so as to be conveniently linked with application programs +(which use some of those functions and data) to form executables. + + The "Library", below, refers to any such software library or work +which has been distributed under these terms. A "work based on the +Library" means either the Library or any derivative work under +copyright law: that is to say, a work containing the Library or a +portion of it, either verbatim or with modifications and/or translated +straightforwardly into another language. (Hereinafter, translation is +included without limitation in the term "modification".) + + "Source code" for a work means the preferred form of the work for +making modifications to it. For a library, complete source code means +all the source code for all modules it contains, plus any associated +interface definition files, plus the scripts used to control compilation +and installation of the library. + + Activities other than copying, distribution and modification are not +covered by this License; they are outside its scope. The act of +running a program using the Library is not restricted, and output from +such a program is covered only if its contents constitute a work based +on the Library (independent of the use of the Library in a tool for +writing it). Whether that is true depends on what the Library does +and what the program that uses the Library does. + + 1. You may copy and distribute verbatim copies of the Library's +complete source code as you receive it, in any medium, provided that +you conspicuously and appropriately publish on each copy an +appropriate copyright notice and disclaimer of warranty; keep intact +all the notices that refer to this License and to the absence of any +warranty; and distribute a copy of this License along with the +Library. + + You may charge a fee for the physical act of transferring a copy, +and you may at your option offer warranty protection in exchange for a +fee. + + 2. You may modify your copy or copies of the Library or any portion +of it, thus forming a work based on the Library, and copy and +distribute such modifications or work under the terms of Section 1 +above, provided that you also meet all of these conditions: + + a) The modified work must itself be a software library. + + b) You must cause the files modified to carry prominent notices + stating that you changed the files and the date of any change. + + c) You must cause the whole of the work to be licensed at no + charge to all third parties under the terms of this License. + + d) If a facility in the modified Library refers to a function or a + table of data to be supplied by an application program that uses + the facility, other than as an argument passed when the facility + is invoked, then you must make a good faith effort to ensure that, + in the event an application does not supply such function or + table, the facility still operates, and performs whatever part of + its purpose remains meaningful. + + (For example, a function in a library to compute square roots has + a purpose that is entirely well-defined independent of the + application. Therefore, Subsection 2d requires that any + application-supplied function or table used by this function must + be optional: if the application does not supply it, the square + root function must still compute square roots.) + +These requirements apply to the modified work as a whole. If +identifiable sections of that work are not derived from the Library, +and can be reasonably considered independent and separate works in +themselves, then this License, and its terms, do not apply to those +sections when you distribute them as separate works. But when you +distribute the same sections as part of a whole which is a work based +on the Library, the distribution of the whole must be on the terms of +this License, whose permissions for other licensees extend to the +entire whole, and thus to each and every part regardless of who wrote +it. + +Thus, it is not the intent of this section to claim rights or contest +your rights to work written entirely by you; rather, the intent is to +exercise the right to control the distribution of derivative or +collective works based on the Library. + +In addition, mere aggregation of another work not based on the Library +with the Library (or with a work based on the Library) on a volume of +a storage or distribution medium does not bring the other work under +the scope of this License. + + 3. You may opt to apply the terms of the ordinary GNU General Public +License instead of this License to a given copy of the Library. To do +this, you must alter all the notices that refer to this License, so +that they refer to the ordinary GNU General Public License, version 2, +instead of to this License. (If a newer version than version 2 of the +ordinary GNU General Public License has appeared, then you can specify +that version instead if you wish.) Do not make any other change in +these notices. + + Once this change is made in a given copy, it is irreversible for +that copy, so the ordinary GNU General Public License applies to all +subsequent copies and derivative works made from that copy. + + This option is useful when you wish to copy part of the code of +the Library into a program that is not a library. + + 4. You may copy and distribute the Library (or a portion or +derivative of it, under Section 2) in object code or executable form +under the terms of Sections 1 and 2 above provided that you accompany +it with the complete corresponding machine-readable source code, which +must be distributed under the terms of Sections 1 and 2 above on a +medium customarily used for software interchange. + + If distribution of object code is made by offering access to copy +from a designated place, then offering equivalent access to copy the +source code from the same place satisfies the requirement to +distribute the source code, even though third parties are not +compelled to copy the source along with the object code. + + 5. A program that contains no derivative of any portion of the +Library, but is designed to work with the Library by being compiled or +linked with it, is called a "work that uses the Library". Such a +work, in isolation, is not a derivative work of the Library, and +therefore falls outside the scope of this License. + + However, linking a "work that uses the Library" with the Library +creates an executable that is a derivative of the Library (because it +contains portions of the Library), rather than a "work that uses the +library". The executable is therefore covered by this License. +Section 6 states terms for distribution of such executables. + + When a "work that uses the Library" uses material from a header file +that is part of the Library, the object code for the work may be a +derivative work of the Library even though the source code is not. +Whether this is true is especially significant if the work can be +linked without the Library, or if the work is itself a library. The +threshold for this to be true is not precisely defined by law. + + If such an object file uses only numerical parameters, data +structure layouts and accessors, and small macros and small inline +functions (ten lines or less in length), then the use of the object +file is unrestricted, regardless of whether it is legally a derivative +work. (Executables containing this object code plus portions of the +Library will still fall under Section 6.) + + Otherwise, if the work is a derivative of the Library, you may +distribute the object code for the work under the terms of Section 6. +Any executables containing that work also fall under Section 6, +whether or not they are linked directly with the Library itself. + + 6. As an exception to the Sections above, you may also combine or +link a "work that uses the Library" with the Library to produce a +work containing portions of the Library, and distribute that work +under terms of your choice, provided that the terms permit +modification of the work for the customer's own use and reverse +engineering for debugging such modifications. + + You must give prominent notice with each copy of the work that the +Library is used in it and that the Library and its use are covered by +this License. You must supply a copy of this License. If the work +during execution displays copyright notices, you must include the +copyright notice for the Library among them, as well as a reference +directing the user to the copy of this License. Also, you must do one +of these things: + + a) Accompany the work with the complete corresponding + machine-readable source code for the Library including whatever + changes were used in the work (which must be distributed under + Sections 1 and 2 above); and, if the work is an executable linked + with the Library, with the complete machine-readable "work that + uses the Library", as object code and/or source code, so that the + user can modify the Library and then relink to produce a modified + executable containing the modified Library. (It is understood + that the user who changes the contents of definitions files in the + Library will not necessarily be able to recompile the application + to use the modified definitions.) + + b) Use a suitable shared library mechanism for linking with the + Library. A suitable mechanism is one that (1) uses at run time a + copy of the library already present on the user's computer system, + rather than copying library functions into the executable, and (2) + will operate properly with a modified version of the library, if + the user installs one, as long as the modified version is + interface-compatible with the version that the work was made with. + + c) Accompany the work with a written offer, valid for at + least three years, to give the same user the materials + specified in Subsection 6a, above, for a charge no more + than the cost of performing this distribution. + + d) If distribution of the work is made by offering access to copy + from a designated place, offer equivalent access to copy the above + specified materials from the same place. + + e) Verify that the user has already received a copy of these + materials or that you have already sent this user a copy. + + For an executable, the required form of the "work that uses the +Library" must include any data and utility programs needed for +reproducing the executable from it. However, as a special exception, +the materials to be distributed need not include anything that is +normally distributed (in either source or binary form) with the major +components (compiler, kernel, and so on) of the operating system on +which the executable runs, unless that component itself accompanies +the executable. + + It may happen that this requirement contradicts the license +restrictions of other proprietary libraries that do not normally +accompany the operating system. Such a contradiction means you cannot +use both them and the Library together in an executable that you +distribute. + + 7. You may place library facilities that are a work based on the +Library side-by-side in a single library together with other library +facilities not covered by this License, and distribute such a combined +library, provided that the separate distribution of the work based on +the Library and of the other library facilities is otherwise +permitted, and provided that you do these two things: + + a) Accompany the combined library with a copy of the same work + based on the Library, uncombined with any other library + facilities. This must be distributed under the terms of the + Sections above. + + b) Give prominent notice with the combined library of the fact + that part of it is a work based on the Library, and explaining + where to find the accompanying uncombined form of the same work. + + 8. You may not copy, modify, sublicense, link with, or distribute +the Library except as expressly provided under this License. Any +attempt otherwise to copy, modify, sublicense, link with, or +distribute the Library is void, and will automatically terminate your +rights under this License. However, parties who have received copies, +or rights, from you under this License will not have their licenses +terminated so long as such parties remain in full compliance. + + 9. You are not required to accept this License, since you have not +signed it. However, nothing else grants you permission to modify or +distribute the Library or its derivative works. These actions are +prohibited by law if you do not accept this License. Therefore, by +modifying or distributing the Library (or any work based on the +Library), you indicate your acceptance of this License to do so, and +all its terms and conditions for copying, distributing or modifying +the Library or works based on it. + + 10. Each time you redistribute the Library (or any work based on the +Library), the recipient automatically receives a license from the +original licensor to copy, distribute, link with or modify the Library +subject to these terms and conditions. You may not impose any further +restrictions on the recipients' exercise of the rights granted herein. +You are not responsible for enforcing compliance by third parties with +this License. + + 11. If, as a consequence of a court judgment or allegation of patent +infringement or for any other reason (not limited to patent issues), +conditions are imposed on you (whether by court order, agreement or +otherwise) that contradict the conditions of this License, they do not +excuse you from the conditions of this License. If you cannot +distribute so as to satisfy simultaneously your obligations under this +License and any other pertinent obligations, then as a consequence you +may not distribute the Library at all. For example, if a patent +license would not permit royalty-free redistribution of the Library by +all those who receive copies directly or indirectly through you, then +the only way you could satisfy both it and this License would be to +refrain entirely from distribution of the Library. + +If any portion of this section is held invalid or unenforceable under any +particular circumstance, the balance of the section is intended to apply, +and the section as a whole is intended to apply in other circumstances. + +It is not the purpose of this section to induce you to infringe any +patents or other property right claims or to contest validity of any +such claims; this section has the sole purpose of protecting the +integrity of the free software distribution system which is +implemented by public license practices. Many people have made +generous contributions to the wide range of software distributed +through that system in reliance on consistent application of that +system; it is up to the author/donor to decide if he or she is willing +to distribute software through any other system and a licensee cannot +impose that choice. + +This section is intended to make thoroughly clear what is believed to +be a consequence of the rest of this License. + + 12. If the distribution and/or use of the Library is restricted in +certain countries either by patents or by copyrighted interfaces, the +original copyright holder who places the Library under this License may add +an explicit geographical distribution limitation excluding those countries, +so that distribution is permitted only in or among countries not thus +excluded. In such case, this License incorporates the limitation as if +written in the body of this License. + + 13. The Free Software Foundation may publish revised and/or new +versions of the Lesser General Public License from time to time. +Such new versions will be similar in spirit to the present version, +but may differ in detail to address new problems or concerns. + +Each version is given a distinguishing version number. If the Library +specifies a version number of this License which applies to it and +"any later version", you have the option of following the terms and +conditions either of that version or of any later version published by +the Free Software Foundation. If the Library does not specify a +license version number, you may choose any version ever published by +the Free Software Foundation. + + 14. If you wish to incorporate parts of the Library into other free +programs whose distribution conditions are incompatible with these, +write to the author to ask for permission. For software which is +copyrighted by the Free Software Foundation, write to the Free +Software Foundation; we sometimes make exceptions for this. Our +decision will be guided by the two goals of preserving the free status +of all derivatives of our free software and of promoting the sharing +and reuse of software generally. + + NO WARRANTY + + 15. BECAUSE THE LIBRARY IS LICENSED FREE OF CHARGE, THERE IS NO +WARRANTY FOR THE LIBRARY, TO THE EXTENT PERMITTED BY APPLICABLE LAW. +EXCEPT WHEN OTHERWISE STATED IN WRITING THE COPYRIGHT HOLDERS AND/OR +OTHER PARTIES PROVIDE THE LIBRARY "AS IS" WITHOUT WARRANTY OF ANY +KIND, EITHER EXPRESSED OR IMPLIED, INCLUDING, BUT NOT LIMITED TO, THE +IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR +PURPOSE. THE ENTIRE RISK AS TO THE QUALITY AND PERFORMANCE OF THE +LIBRARY IS WITH YOU. SHOULD THE LIBRARY PROVE DEFECTIVE, YOU ASSUME +THE COST OF ALL NECESSARY SERVICING, REPAIR OR CORRECTION. + + 16. IN NO EVENT UNLESS REQUIRED BY APPLICABLE LAW OR AGREED TO IN +WRITING WILL ANY COPYRIGHT HOLDER, OR ANY OTHER PARTY WHO MAY MODIFY +AND/OR REDISTRIBUTE THE LIBRARY AS PERMITTED ABOVE, BE LIABLE TO YOU +FOR DAMAGES, INCLUDING ANY GENERAL, SPECIAL, INCIDENTAL OR +CONSEQUENTIAL DAMAGES ARISING OUT OF THE USE OR INABILITY TO USE THE +LIBRARY (INCLUDING BUT NOT LIMITED TO LOSS OF DATA OR DATA BEING +RENDERED INACCURATE OR LOSSES SUSTAINED BY YOU OR THIRD PARTIES OR A +FAILURE OF THE LIBRARY TO OPERATE WITH ANY OTHER SOFTWARE), EVEN IF +SUCH HOLDER OR OTHER PARTY HAS BEEN ADVISED OF THE POSSIBILITY OF SUCH +DAMAGES. + + END OF TERMS AND CONDITIONS + + How to Apply These Terms to Your New Libraries + + If you develop a new library, and you want it to be of the greatest +possible use to the public, we recommend making it free software that +everyone can redistribute and change. You can do so by permitting +redistribution under these terms (or, alternatively, under the terms of the +ordinary General Public License). + + To apply these terms, attach the following notices to the library. It is +safest to attach them to the start of each source file to most effectively +convey the exclusion of warranty; and each file should have at least the +"copyright" line and a pointer to where the full notice is found. + + + Copyright (C) + + This library is free software; you can redistribute it and/or + modify it under the terms of the GNU Lesser General Public + License as published by the Free Software Foundation; either + version 2.1 of the License, or (at your option) any later version. + + This library is distributed in the hope that it will be useful, + but WITHOUT ANY WARRANTY; without even the implied warranty of + MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU + Lesser General Public License for more details. + + You should have received a copy of the GNU Lesser General Public + License along with this library; if not, write to the Free Software + Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA + +Also add information on how to contact you by electronic and paper mail. + +You should also get your employer (if you work as a programmer) or your +school, if any, to sign a "copyright disclaimer" for the library, if +necessary. Here is a sample; alter the names: + + Yoyodyne, Inc., hereby disclaims all copyright interest in the + library `Frob' (a library for tweaking knobs) written by James Random Hacker. + + , 1 April 1990 + Ty Coon, President of Vice + +That's all there is to it! diff --git a/assets/licenses/LGPL-3.0.txt b/assets/licenses/LGPL-3.0.txt new file mode 100644 index 00000000..65c5ca88 --- /dev/null +++ b/assets/licenses/LGPL-3.0.txt @@ -0,0 +1,165 @@ + GNU LESSER GENERAL PUBLIC LICENSE + Version 3, 29 June 2007 + + Copyright (C) 2007 Free Software Foundation, Inc. + Everyone is permitted to copy and distribute verbatim copies + of this license document, but changing it is not allowed. + + + This version of the GNU Lesser General Public License incorporates +the terms and conditions of version 3 of the GNU General Public +License, supplemented by the additional permissions listed below. + + 0. Additional Definitions. + + As used herein, "this License" refers to version 3 of the GNU Lesser +General Public License, and the "GNU GPL" refers to version 3 of the GNU +General Public License. + + "The Library" refers to a covered work governed by this License, +other than an Application or a Combined Work as defined below. + + An "Application" is any work that makes use of an interface provided +by the Library, but which is not otherwise based on the Library. +Defining a subclass of a class defined by the Library is deemed a mode +of using an interface provided by the Library. + + A "Combined Work" is a work produced by combining or linking an +Application with the Library. The particular version of the Library +with which the Combined Work was made is also called the "Linked +Version". + + The "Minimal Corresponding Source" for a Combined Work means the +Corresponding Source for the Combined Work, excluding any source code +for portions of the Combined Work that, considered in isolation, are +based on the Application, and not on the Linked Version. + + The "Corresponding Application Code" for a Combined Work means the +object code and/or source code for the Application, including any data +and utility programs needed for reproducing the Combined Work from the +Application, but excluding the System Libraries of the Combined Work. + + 1. Exception to Section 3 of the GNU GPL. + + You may convey a covered work under sections 3 and 4 of this License +without being bound by section 3 of the GNU GPL. + + 2. Conveying Modified Versions. + + If you modify a copy of the Library, and, in your modifications, a +facility refers to a function or data to be supplied by an Application +that uses the facility (other than as an argument passed when the +facility is invoked), then you may convey a copy of the modified +version: + + a) under this License, provided that you make a good faith effort to + ensure that, in the event an Application does not supply the + function or data, the facility still operates, and performs + whatever part of its purpose remains meaningful, or + + b) under the GNU GPL, with none of the additional permissions of + this License applicable to that copy. + + 3. Object Code Incorporating Material from Library Header Files. + + The object code form of an Application may incorporate material from +a header file that is part of the Library. You may convey such object +code under terms of your choice, provided that, if the incorporated +material is not limited to numerical parameters, data structure +layouts and accessors, or small macros, inline functions and templates +(ten or fewer lines in length), you do both of the following: + + a) Give prominent notice with each copy of the object code that the + Library is used in it and that the Library and its use are + covered by this License. + + b) Accompany the object code with a copy of the GNU GPL and this license + document. + + 4. Combined Works. + + You may convey a Combined Work under terms of your choice that, +taken together, effectively do not restrict modification of the +portions of the Library contained in the Combined Work and reverse +engineering for debugging such modifications, if you also do each of +the following: + + a) Give prominent notice with each copy of the Combined Work that + the Library is used in it and that the Library and its use are + covered by this License. + + b) Accompany the Combined Work with a copy of the GNU GPL and this license + document. + + c) For a Combined Work that displays copyright notices during + execution, include the copyright notice for the Library among + these notices, as well as a reference directing the user to the + copies of the GNU GPL and this license document. + + d) Do one of the following: + + 0) Convey the Minimal Corresponding Source under the terms of this + License, and the Corresponding Application Code in a form + suitable for, and under terms that permit, the user to + recombine or relink the Application with a modified version of + the Linked Version to produce a modified Combined Work, in the + manner specified by section 6 of the GNU GPL for conveying + Corresponding Source. + + 1) Use a suitable shared library mechanism for linking with the + Library. A suitable mechanism is one that (a) uses at run time + a copy of the Library already present on the user's computer + system, and (b) will operate properly with a modified version + of the Library that is interface-compatible with the Linked + Version. + + e) Provide Installation Information, but only if you would otherwise + be required to provide such information under section 6 of the + GNU GPL, and only to the extent that such information is + necessary to install and execute a modified version of the + Combined Work produced by recombining or relinking the + Application with a modified version of the Linked Version. (If + you use option 4d0, the Installation Information must accompany + the Minimal Corresponding Source and Corresponding Application + Code. If you use option 4d1, you must provide the Installation + Information in the manner specified by section 6 of the GNU GPL + for conveying Corresponding Source.) + + 5. Combined Libraries. + + You may place library facilities that are a work based on the +Library side by side in a single library together with other library +facilities that are not Applications and are not covered by this +License, and convey such a combined library under terms of your +choice, if you do both of the following: + + a) Accompany the combined library with a copy of the same work based + on the Library, uncombined with any other library facilities, + conveyed under the terms of this License. + + b) Give prominent notice with the combined library that part of it + is a work based on the Library, and explaining where to find the + accompanying uncombined form of the same work. + + 6. Revised Versions of the GNU Lesser General Public License. + + The Free Software Foundation may publish revised and/or new versions +of the GNU Lesser General Public License from time to time. Such new +versions will be similar in spirit to the present version, but may +differ in detail to address new problems or concerns. + + Each version is given a distinguishing version number. If the +Library as you received it specifies that a certain numbered version +of the GNU Lesser General Public License "or any later version" +applies to it, you have the option of following the terms and +conditions either of that published version or of any later version +published by the Free Software Foundation. If the Library as you +received it does not specify a version number of the GNU Lesser +General Public License, you may choose any version of the GNU Lesser +General Public License ever published by the Free Software Foundation. + + If the Library as you received it specifies that a proxy can decide +whether future versions of the GNU Lesser General Public License shall +apply, that proxy's public statement of acceptance of any version is +permanent authorization for you to choose that version for the +Library. diff --git a/assets/licenses/Python-3.13.txt b/assets/licenses/Python-3.13.txt new file mode 100644 index 00000000..14603b95 --- /dev/null +++ b/assets/licenses/Python-3.13.txt @@ -0,0 +1,277 @@ +A. HISTORY OF THE SOFTWARE +========================== + +Python was created in the early 1990s by Guido van Rossum at Stichting +Mathematisch Centrum (CWI, see https://www.cwi.nl) in the Netherlands +as a successor of a language called ABC. Guido remains Python's +principal author, although it includes many contributions from others. + +In 1995, Guido continued his work on Python at the Corporation for +National Research Initiatives (CNRI, see https://www.cnri.reston.va.us) +in Reston, Virginia where he released several versions of the +software. + +In May 2000, Guido and the Python core development team moved to +BeOpen.com to form the BeOpen PythonLabs team. In October of the same +year, the PythonLabs team moved to Digital Creations, which became +Zope Corporation. In 2001, the Python Software Foundation (PSF, see +https://www.python.org/psf/) was formed, a non-profit organization +created specifically to own Python-related Intellectual Property. +Zope Corporation was a sponsoring member of the PSF. + +All Python releases are Open Source (see https://opensource.org for +the Open Source Definition). Historically, most, but not all, Python +releases have also been GPL-compatible; the table below summarizes +the various releases. + + Release Derived Year Owner GPL- + from compatible? (1) + + 0.9.0 thru 1.2 1991-1995 CWI yes + 1.3 thru 1.5.2 1.2 1995-1999 CNRI yes + 1.6 1.5.2 2000 CNRI no + 2.0 1.6 2000 BeOpen.com no + 1.6.1 1.6 2001 CNRI yes (2) + 2.1 2.0+1.6.1 2001 PSF no + 2.0.1 2.0+1.6.1 2001 PSF yes + 2.1.1 2.1+2.0.1 2001 PSF yes + 2.1.2 2.1.1 2002 PSF yes + 2.1.3 2.1.2 2002 PSF yes + 2.2 and above 2.1.1 2001-now PSF yes + +Footnotes: + +(1) GPL-compatible doesn't mean that we're distributing Python under + the GPL. All Python licenses, unlike the GPL, let you distribute + a modified version without making your changes open source. The + GPL-compatible licenses make it possible to combine Python with + other software that is released under the GPL; the others don't. + +(2) According to Richard Stallman, 1.6.1 is not GPL-compatible, + because its license has a choice of law clause. According to + CNRI, however, Stallman's lawyer has told CNRI's lawyer that 1.6.1 + is "not incompatible" with the GPL. + +Thanks to the many outside volunteers who have worked under Guido's +direction to make these releases possible. + + +B. TERMS AND CONDITIONS FOR ACCESSING OR OTHERWISE USING PYTHON +=============================================================== + +Python software and documentation are licensed under the +Python Software Foundation License Version 2. + +Starting with Python 3.8.6, examples, recipes, and other code in +the documentation are dual licensed under the PSF License Version 2 +and the Zero-Clause BSD license. + +Some software incorporated into Python is under different licenses. +The licenses are listed with code falling under that license. + + +PYTHON SOFTWARE FOUNDATION LICENSE VERSION 2 +-------------------------------------------- + +1. This LICENSE AGREEMENT is between the Python Software Foundation +("PSF"), and the Individual or Organization ("Licensee") accessing and +otherwise using this software ("Python") in source or binary form and +its associated documentation. + +2. Subject to the terms and conditions of this License Agreement, PSF hereby +grants Licensee a nonexclusive, royalty-free, world-wide license to reproduce, +analyze, test, perform and/or display publicly, prepare derivative works, +distribute, and otherwise use Python alone or in any derivative version, +provided, however, that PSF's License Agreement and PSF's notice of copyright, +i.e., "Copyright (c) 2001-2024 Python Software Foundation; All Rights Reserved" +are retained in Python alone or in any derivative version prepared by Licensee. + +3. In the event Licensee prepares a derivative work that is based on +or incorporates Python or any part thereof, and wants to make +the derivative work available to others as provided herein, then +Licensee hereby agrees to include in any such work a brief summary of +the changes made to Python. + +4. PSF is making Python available to Licensee on an "AS IS" +basis. PSF MAKES NO REPRESENTATIONS OR WARRANTIES, EXPRESS OR +IMPLIED. BY WAY OF EXAMPLE, BUT NOT LIMITATION, PSF MAKES NO AND +DISCLAIMS ANY REPRESENTATION OR WARRANTY OF MERCHANTABILITY OR FITNESS +FOR ANY PARTICULAR PURPOSE OR THAT THE USE OF PYTHON WILL NOT +INFRINGE ANY THIRD PARTY RIGHTS. + +5. PSF SHALL NOT BE LIABLE TO LICENSEE OR ANY OTHER USERS OF PYTHON +FOR ANY INCIDENTAL, SPECIAL, OR CONSEQUENTIAL DAMAGES OR LOSS AS +A RESULT OF MODIFYING, DISTRIBUTING, OR OTHERWISE USING PYTHON, +OR ANY DERIVATIVE THEREOF, EVEN IF ADVISED OF THE POSSIBILITY THEREOF. + +6. This License Agreement will automatically terminate upon a material +breach of its terms and conditions. + +7. Nothing in this License Agreement shall be deemed to create any +relationship of agency, partnership, or joint venture between PSF and +Licensee. This License Agreement does not grant permission to use PSF +trademarks or trade name in a trademark sense to endorse or promote +products or services of Licensee, or any third party. + +8. By copying, installing or otherwise using Python, Licensee +agrees to be bound by the terms and conditions of this License +Agreement. + + +BEOPEN.COM LICENSE AGREEMENT FOR PYTHON 2.0 +------------------------------------------- + +BEOPEN PYTHON OPEN SOURCE LICENSE AGREEMENT VERSION 1 + +1. This LICENSE AGREEMENT is between BeOpen.com ("BeOpen"), having an +office at 160 Saratoga Avenue, Santa Clara, CA 95051, and the +Individual or Organization ("Licensee") accessing and otherwise using +this software in source or binary form and its associated +documentation ("the Software"). + +2. Subject to the terms and conditions of this BeOpen Python License +Agreement, BeOpen hereby grants Licensee a non-exclusive, +royalty-free, world-wide license to reproduce, analyze, test, perform +and/or display publicly, prepare derivative works, distribute, and +otherwise use the Software alone or in any derivative version, +provided, however, that the BeOpen Python License is retained in the +Software, alone or in any derivative version prepared by Licensee. + +3. BeOpen is making the Software available to Licensee on an "AS IS" +basis. BEOPEN MAKES NO REPRESENTATIONS OR WARRANTIES, EXPRESS OR +IMPLIED. BY WAY OF EXAMPLE, BUT NOT LIMITATION, BEOPEN MAKES NO AND +DISCLAIMS ANY REPRESENTATION OR WARRANTY OF MERCHANTABILITY OR FITNESS +FOR ANY PARTICULAR PURPOSE OR THAT THE USE OF THE SOFTWARE WILL NOT +INFRINGE ANY THIRD PARTY RIGHTS. + +4. BEOPEN SHALL NOT BE LIABLE TO LICENSEE OR ANY OTHER USERS OF THE +SOFTWARE FOR ANY INCIDENTAL, SPECIAL, OR CONSEQUENTIAL DAMAGES OR LOSS +AS A RESULT OF USING, MODIFYING OR DISTRIBUTING THE SOFTWARE, OR ANY +DERIVATIVE THEREOF, EVEN IF ADVISED OF THE POSSIBILITY THEREOF. + +5. This License Agreement will automatically terminate upon a material +breach of its terms and conditions. + +6. This License Agreement shall be governed by and interpreted in all +respects by the law of the State of California, excluding conflict of +law provisions. Nothing in this License Agreement shall be deemed to +create any relationship of agency, partnership, or joint venture +between BeOpen and Licensee. This License Agreement does not grant +permission to use BeOpen trademarks or trade names in a trademark +sense to endorse or promote products or services of Licensee, or any +third party. As an exception, the "BeOpen Python" logos available at +http://www.pythonlabs.com/logos.html may be used according to the +permissions granted on that web page. + +7. By copying, installing or otherwise using the software, Licensee +agrees to be bound by the terms and conditions of this License +Agreement. + + +CNRI LICENSE AGREEMENT FOR PYTHON 1.6.1 +--------------------------------------- + +1. This LICENSE AGREEMENT is between the Corporation for National +Research Initiatives, having an office at 1895 Preston White Drive, +Reston, VA 20191 ("CNRI"), and the Individual or Organization +("Licensee") accessing and otherwise using Python 1.6.1 software in +source or binary form and its associated documentation. + +2. Subject to the terms and conditions of this License Agreement, CNRI +hereby grants Licensee a nonexclusive, royalty-free, world-wide +license to reproduce, analyze, test, perform and/or display publicly, +prepare derivative works, distribute, and otherwise use Python 1.6.1 +alone or in any derivative version, provided, however, that CNRI's +License Agreement and CNRI's notice of copyright, i.e., "Copyright (c) +1995-2001 Corporation for National Research Initiatives; All Rights +Reserved" are retained in Python 1.6.1 alone or in any derivative +version prepared by Licensee. Alternately, in lieu of CNRI's License +Agreement, Licensee may substitute the following text (omitting the +quotes): "Python 1.6.1 is made available subject to the terms and +conditions in CNRI's License Agreement. This Agreement together with +Python 1.6.1 may be located on the internet using the following +unique, persistent identifier (known as a handle): 1895.22/1013. This +Agreement may also be obtained from a proxy server on the internet +using the following URL: http://hdl.handle.net/1895.22/1013". + +3. In the event Licensee prepares a derivative work that is based on +or incorporates Python 1.6.1 or any part thereof, and wants to make +the derivative work available to others as provided herein, then +Licensee hereby agrees to include in any such work a brief summary of +the changes made to Python 1.6.1. + +4. CNRI is making Python 1.6.1 available to Licensee on an "AS IS" +basis. CNRI MAKES NO REPRESENTATIONS OR WARRANTIES, EXPRESS OR +IMPLIED. BY WAY OF EXAMPLE, BUT NOT LIMITATION, CNRI MAKES NO AND +DISCLAIMS ANY REPRESENTATION OR WARRANTY OF MERCHANTABILITY OR FITNESS +FOR ANY PARTICULAR PURPOSE OR THAT THE USE OF PYTHON 1.6.1 WILL NOT +INFRINGE ANY THIRD PARTY RIGHTS. + +5. CNRI SHALL NOT BE LIABLE TO LICENSEE OR ANY OTHER USERS OF PYTHON +1.6.1 FOR ANY INCIDENTAL, SPECIAL, OR CONSEQUENTIAL DAMAGES OR LOSS AS +A RESULT OF MODIFYING, DISTRIBUTING, OR OTHERWISE USING PYTHON 1.6.1, +OR ANY DERIVATIVE THEREOF, EVEN IF ADVISED OF THE POSSIBILITY THEREOF. + +6. This License Agreement will automatically terminate upon a material +breach of its terms and conditions. + +7. This License Agreement shall be governed by the federal +intellectual property law of the United States, including without +limitation the federal copyright law, and, to the extent such +U.S. federal law does not apply, by the law of the Commonwealth of +Virginia, excluding Virginia's conflict of law provisions. +Notwithstanding the foregoing, with regard to derivative works based +on Python 1.6.1 that incorporate non-separable material that was +previously distributed under the GNU General Public License (GPL), the +law of the Commonwealth of Virginia shall govern this License +Agreement only as to issues arising under or with respect to +Paragraphs 4, 5, and 7 of this License Agreement. Nothing in this +License Agreement shall be deemed to create any relationship of +agency, partnership, or joint venture between CNRI and Licensee. This +License Agreement does not grant permission to use CNRI trademarks or +trade name in a trademark sense to endorse or promote products or +services of Licensee, or any third party. + +8. By clicking on the "ACCEPT" button where indicated, or by copying, +installing or otherwise using Python 1.6.1, Licensee agrees to be +bound by the terms and conditions of this License Agreement. + + ACCEPT + + +CWI LICENSE AGREEMENT FOR PYTHON 0.9.0 THROUGH 1.2 +-------------------------------------------------- + +Copyright (c) 1991 - 1995, Stichting Mathematisch Centrum Amsterdam, +The Netherlands. All rights reserved. + +Permission to use, copy, modify, and distribute this software and its +documentation for any purpose and without fee is hereby granted, +provided that the above copyright notice appear in all copies and that +both that copyright notice and this permission notice appear in +supporting documentation, and that the name of Stichting Mathematisch +Centrum or CWI not be used in advertising or publicity pertaining to +distribution of the software without specific, written prior +permission. + +STICHTING MATHEMATISCH CENTRUM DISCLAIMS ALL WARRANTIES WITH REGARD TO +THIS SOFTWARE, INCLUDING ALL IMPLIED WARRANTIES OF MERCHANTABILITY AND +FITNESS, IN NO EVENT SHALL STICHTING MATHEMATISCH CENTRUM BE LIABLE +FOR ANY SPECIAL, INDIRECT OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES +WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, WHETHER IN AN +ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT +OF OR IN CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE. + +ZERO-CLAUSE BSD LICENSE FOR CODE IN THE PYTHON DOCUMENTATION +---------------------------------------------------------------------- + +Permission to use, copy, modify, and/or distribute this software for any +purpose with or without fee is hereby granted. + +THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES WITH +REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF MERCHANTABILITY +AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR ANY SPECIAL, DIRECT, +INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES WHATSOEVER RESULTING FROM +LOSS OF USE, DATA OR PROFITS, WHETHER IN AN ACTION OF CONTRACT, NEGLIGENCE OR +OTHER TORTIOUS ACTION, ARISING OUT OF OR IN CONNECTION WITH THE USE OR +PERFORMANCE OF THIS SOFTWARE. diff --git a/assets/licenses/README.md b/assets/licenses/README.md new file mode 100644 index 00000000..e37b61c5 --- /dev/null +++ b/assets/licenses/README.md @@ -0,0 +1,12 @@ +# Included license texts + +- `Apache-2.0.txt`: VALI, RF-DETR, OpenMMLab MMagic, and other Apache-2.0 components. +- `BSD-3-Clause.txt`: PyAV. +- `GPL-3.0.txt`: the GPLv3 FFmpeg libraries and command-line tools. +- `LGPL-2.1.txt`: libVLC and python-vlc under LGPL-2.1-or-later. +- `LGPL-3.0.txt`: FFmpeg library code under LGPL-3.0-or-later. +- `Python-3.13.txt`: CPython 3.13. + +Jasna's AGPL-3.0 license is the repository and release-root `LICENSE` file. +Official binary builds also collect license files from the installed Python +distributions into `licenses/python-packages/`. diff --git a/docs/en/development.md b/docs/en/development.md index 94e7ed3a..a72c1192 100644 --- a/docs/en/development.md +++ b/docs/en/development.md @@ -171,6 +171,25 @@ CSVs are the only full record. `scripts/benchmark_releases.py` compares frozen release archives instead of the working tree, and `scripts/benchmark_lada_flatpak.py` refreshes the Lada baseline column. +## Release licensing checklist + +Before publishing any platform archive: + +1. Set the same version in `pyproject.toml`, `jasna/__init__.py`, the git tag, + and `RELEASE_SOURCES.md`. +2. Confirm `RELEASE_SOURCES.md` matches the PyAV, VALI, FFmpeg, Python, and + protection revisions used by the build. +3. Verify the unpacked package contains `LICENSE`, `LICENSING.md`, `NOTICE`, + `RELEASE_SOURCES.md`, `assets/THIRD_PARTY_*.md`, and `licenses/`. +4. Recalculate every bundled model hash and compare it with + `assets/THIRD_PARTY_MODELS.md`. +5. Upload every archive part **and** its generated `.sha256` file to the same + GitHub release. + +The release builder copies maintained license texts plus the license files +from all installed Python distributions. The public VALI fork is the source +for Jasna's modified `python_vali` wheel. + ## AMD release builds These scripts live in the private protection submodule and are for the diff --git a/jasna/__init__.py b/jasna/__init__.py index 132d0cc1..01ec93ae 100644 --- a/jasna/__init__.py +++ b/jasna/__init__.py @@ -1,3 +1,3 @@ __all__ = ["__version__"] -__version__ = "0.10.0" +__version__ = "0.11.0" diff --git a/jasna/protection b/jasna/protection index b8d2bf3c..8d6af872 160000 --- a/jasna/protection +++ b/jasna/protection @@ -1 +1 @@ -Subproject commit b8d2bf3cfdd5dba49432e89414ad6197eed277f0 +Subproject commit 8d6af872b9434f399161ca9729cc0ff820803248 diff --git a/pyproject.toml b/pyproject.toml index 487a4493..08b039e5 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta" [project] name = "jasna" -version = "0.10.0" +version = "0.11.0" description = "Jasna" requires-python = ">=3.12" dependencies = [ From 7860ae31ad4f0eb38ab0b09a45152a1a3f41910c Mon Sep 17 00:00:00 2001 From: Kruk2 Date: Mon, 28 Sep 2026 13:35:18 +0200 Subject: [PATCH 03/48] build: depend on PyPI PyAV 18.1.0 instead of the fork build 18.1.0 contains the current-CUDA-context and explicit NVENC stream changes the release builds compiled from a pinned PyAV commit. Pin it exactly because the VALI release wheel links the same pyav-ffmpeg 8.1.2-1 libraries. Drop cython/delvewheel from the dev extra; only the PyAV rebuild used them. --- RELEASE_SOURCES.md | 4 ++-- docs/en/development.md | 1 - jasna/protection | 2 +- pyproject.toml | 9 +++------ 4 files changed, 6 insertions(+), 10 deletions(-) diff --git a/RELEASE_SOURCES.md b/RELEASE_SOURCES.md index 9b59056c..75e60b10 100644 --- a/RELEASE_SOURCES.md +++ b/RELEASE_SOURCES.md @@ -11,7 +11,7 @@ notices, and full license texts. | --- | --- | | Jasna | Git tag `v0.11.0`; the tag and package version must match | | Protection component | The private gitlink revision recorded by the `v0.11.0` tag; source is not published | -| PyAV | [`f6f0a5e3d975aab851e12ca881bdadfdd8f9ec74`](https://github.com/PyAV-Org/PyAV/commit/f6f0a5e3d975aab851e12ca881bdadfdd8f9ec74) | +| PyAV | Unmodified PyPI `av` 18.1.0, tag [`v18.1.0`](https://github.com/PyAV-Org/PyAV/commit/7e3d950a8b72062502c1a60d672f8ca565313af5) | | VALI | [`0e5c01ee57222a8ef4f9c7591284c4324bb3fba5`](https://codeberg.org/Kruk2/vali/commit/0e5c01ee57222a8ef4f9c7591284c4324bb3fba5) in the public Kruk2 fork | | RF-DETR | [`1.8.3` / `3bd6bffbcb13cac3a5b1c37da5a0fd5453b50c86`](https://github.com/roboflow/rf-detr/commit/3bd6bffbcb13cac3a5b1c37da5a0fd5453b50c86) | | MMagic compatibility patch | [`patches/fix_loading_mmengine_weights_on_torch26_and_higher.diff`](patches/fix_loading_mmengine_weights_on_torch26_and_higher.diff) | @@ -22,7 +22,7 @@ It is licensed under Apache-2.0. ## FFmpeg in the PyAV and VALI wheels -The Linux wheels link against the +The PyAV wheels and the Linux VALI wheel link against the [`pyav-ffmpeg` 8.1.2-1](https://github.com/PyAV-Org/pyav-ffmpeg/releases/tag/8.1.2-1) build of FFmpeg 8.1.2. The source is FFmpeg tag `n8.1.2`, commit [`38b88335f99e76ed89ff3c93f877fdefce736c13`](https://github.com/FFmpeg/FFmpeg/commit/38b88335f99e76ed89ff3c93f877fdefce736c13). diff --git a/docs/en/development.md b/docs/en/development.md index a406d1a4..bed3480a 100644 --- a/docs/en/development.md +++ b/docs/en/development.md @@ -89,7 +89,6 @@ Developer setup also requires: - `ffmpeg` and `ffprobe` on `PATH`; `ffmpeg` major version must be **8**. - libVLC 3 installed for GUI player audio. The Python binding is installed from `pyproject.toml`; on Ubuntu install `libvlc5` and `vlc-plugin-base`. -- PyAV built from main branch. Currently released 18.0.0 doesn't ship some CUDA context improvements. - Optional: a `python_vali` wheel built from . Only that fork has `DecodeSingleSurfaceAsyncDetailed` and its corrupt-packet tolerance, which the VALI decode backend needs; with the stock PyPI wheel the reader falls back to PyAV for diff --git a/jasna/protection b/jasna/protection index e8e0b3d0..356f3e14 160000 --- a/jasna/protection +++ b/jasna/protection @@ -1 +1 @@ -Subproject commit e8e0b3d0eae1b1ca3449ae0ceff5422b323921b9 +Subproject commit 356f3e1438ba1ae7330d854f6f5a03827741fee6 diff --git a/pyproject.toml b/pyproject.toml index d9510382..0059ac30 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -8,10 +8,9 @@ version = "0.11.0" description = "Jasna" requires-python = ">=3.12" dependencies = [ - # current_ctx support is merged upstream for PyAV 18.1.0 but is not published yet. - # Linux releases override this with the pinned upstream-main wheel in build_linux.sh; - # remove that override and require av>=18.1 once the release reaches PyPI. - "av>=18,<19", + # Exact pin: the release VALI wheel is built against the same pyav-ffmpeg libraries + # this wheel bundles, so bump both together. + "av==18.1.0", "numpy>=1.26", "psutil", "diffusers", @@ -65,8 +64,6 @@ dev = [ "scikit-build", "cmake", "ninja", - "cython>=3.1,<4", - "delvewheel; sys_platform == 'win32'", ] [tool.setuptools] From fb9153f4fa0ac311d06258aa6ea5405ad6950c64 Mon Sep 17 00:00:00 2001 From: Kruk2 Date: Mon, 28 Sep 2026 13:35:18 +0200 Subject: [PATCH 04/48] build(nvidia): pin torch 2.12.1 and torchvision 0.27.1 Patch releases within the 2.12 line that torch-tensorrt 2.12.0 accepts; the dev environment already runs them. --- pyproject.toml | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 0059ac30..b1fdd6e9 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -33,8 +33,8 @@ jasna = "jasna.main:main" [project.optional-dependencies] nvidia = [ - "torch==2.12.0+cu130", - "torchvision==0.27.0+cu130", + "torch==2.12.1+cu130", + "torchvision==0.27.1+cu130", "tensorrt==10.16.1.11", "torch-tensorrt==2.12.0", "nvidia-vfx", From 26183cca4787e07215182737a549124c130923cc Mon Sep 17 00:00:00 2001 From: Kruk2 Date: Mon, 28 Sep 2026 13:38:09 +0200 Subject: [PATCH 05/48] build: mark shell scripts executable Also records the protection commit that fixes the release scripts' modes. --- jasna/protection | 2 +- scripts/build_fatbins.sh | 0 2 files changed, 1 insertion(+), 1 deletion(-) mode change 100644 => 100755 scripts/build_fatbins.sh diff --git a/jasna/protection b/jasna/protection index 356f3e14..c8beb4e8 160000 --- a/jasna/protection +++ b/jasna/protection @@ -1 +1 @@ -Subproject commit 356f3e1438ba1ae7330d854f6f5a03827741fee6 +Subproject commit c8beb4e871b884f13270e7414fddca530d220e41 diff --git a/scripts/build_fatbins.sh b/scripts/build_fatbins.sh old mode 100644 new mode 100755 From dbfe9b9249149f1ea72c2c12f983481be51482c6 Mon Sep 17 00:00:00 2001 From: Kruk2 Date: Mon, 28 Sep 2026 14:59:26 +0200 Subject: [PATCH 06/48] feat(gui): show the official supporter pages and explain rejected keys Show where supporter keys are sold in the system-check wizard and in the license dialog, and give translated messages for keys that are not accepted. --- README.ja.md | 2 ++ README.md | 2 ++ README.zh.md | 2 ++ docs/en/gui.md | 2 ++ docs/ja/gui.md | 2 ++ docs/zh/gui.md | 1 + jasna/gui/components.py | 17 +++++++-- jasna/gui/locales/en.py | 4 +++ jasna/gui/locales/ja.py | 4 +++ jasna/gui/locales/ko.py | 4 +++ jasna/gui/locales/th.py | 4 +++ jasna/gui/locales/zh.py | 4 +++ jasna/gui/wizard.py | 9 +++++ tests/test_gui_official_sellers.py | 57 ++++++++++++++++++++++++++++++ 14 files changed, 112 insertions(+), 2 deletions(-) create mode 100644 tests/test_gui_official_sellers.py diff --git a/README.ja.md b/README.ja.md index 43e09816..d9257382 100644 --- a/README.ja.md +++ b/README.ja.md @@ -170,6 +170,8 @@ Jasna の公開アプリケーションソースは AGPL-3.0 です。モデル - **[Unifans](https://app.unifans.io/c/kruk2)**: プラットフォームメッセージで送信されます。少し遅れる場合があります。 - **[Buy Me a Coffee](https://buymeacoffee.com/kruk2)**、**暗号通貨**を含む: 支援時に使ったメールアドレスまたはハンドルへ送信されます。キーはそのメールアドレスまたはハンドルに紐付きます。 +キーは上記の2つのページで**のみ**販売しています。それ以外で Jasna、キー、改変ファイルを販売している人は詐欺です。そのキーは公式版では使えず、そのファイルはPC上で隠れたコードを実行する可能性があります。Jasna は GitHub からのみダウンロードしてください。 + ## 謝辞 - **[Lada](https://codeberg.org/ladaapp/lada)（Codeberg）** — Jasna は Lada に diff --git a/README.md b/README.md index 44d1dff6..467d2b9f 100644 --- a/README.md +++ b/README.md @@ -173,6 +173,8 @@ How to get a key: - **[Unifans](https://app.unifans.io/c/kruk2)**: sent by platform message. There might be a slight delay. - **[Buy Me a Coffee](https://buymeacoffee.com/kruk2)**, including **crypto**: sent to the email or handle used for the contribution. The key is tied to that email or handle. +Keys are sold **only** on these two pages. Anyone else selling Jasna, keys or modified files is a scam: their keys do not work with the official app, and their files can run hidden code on your PC. Download Jasna only from GitHub. + ## Acknowledgments - **[Lada](https://codeberg.org/ladaapp/lada) (Codeberg)** — Jasna is inspired diff --git a/README.zh.md b/README.zh.md index 4e99e27c..e377a762 100644 --- a/README.zh.md +++ b/README.zh.md @@ -164,6 +164,8 @@ Jasna 的公开应用程序源代码采用 AGPL-3.0。模型和捆绑依赖项 - **[Unifans](https://app.unifans.io/c/kruk2)**: 通过平台消息发送，可能会有轻微延迟。 - **[Buy Me a Coffee](https://buymeacoffee.com/kruk2)**，包括**加密货币**: 发送到贡献时使用的邮箱或账号。密钥与该邮箱或账号绑定。 +密钥**仅**在以上两个页面出售。其他任何出售 Jasna、密钥或修改版文件的人都是骗子：他们的密钥无法用于官方版本，他们的文件可能会在你的电脑上运行隐藏代码。请只从 GitHub 下载 Jasna。 + ## 致谢 - **[Lada](https://codeberg.org/ladaapp/lada)（Codeberg）** — Jasna 受 Lada 启发， diff --git a/docs/en/gui.md b/docs/en/gui.md index a497ea21..4df6824c 100644 --- a/docs/en/gui.md +++ b/docs/en/gui.md @@ -99,3 +99,5 @@ after a driver update or when something misbehaves. bar while processing. - **License**: enter your supporter key from the header; the chip shows whether it's active. + Keys are sold only on Buy Me a Coffee and Unifans (kruk2); keys from other + sellers are fake and do not work. diff --git a/docs/ja/gui.md b/docs/ja/gui.md index 36f8b981..b6e56e55 100644 --- a/docs/ja/gui.md +++ b/docs/ja/gui.md @@ -96,3 +96,5 @@ Space キーで再生と一時停止を切り替え、左右矢印キーで 30 下部バーに表示されます。 - **ライセンス**: 支援者キーはヘッダーから入力します。チップの表示で 有効かどうかが分かります。 + キーは Buy Me a Coffee と Unifans (kruk2) でのみ販売しています。他の販売者の + キーは偽物で、使えません。 diff --git a/docs/zh/gui.md b/docs/zh/gui.md index 66ecba0f..5033deba 100644 --- a/docs/zh/gui.md +++ b/docs/zh/gui.md @@ -79,3 +79,4 @@ 反馈问题时请附上导出的日志。 - **系统状态**: 处理时底部工具条会显示 GPU、VRAM、内存和 CPU 占用。 - **许可证**: 在头部输入你的支持者密钥；状态标签会显示是否已激活。 + 密钥仅在 Buy Me a Coffee 和 Unifans (kruk2) 出售；其他卖家的密钥都是假的，无法使用。 diff --git a/jasna/gui/components.py b/jasna/gui/components.py index f7c709ac..2266cd26 100644 --- a/jasna/gui/components.py +++ b/jasna/gui/components.py @@ -214,6 +214,10 @@ def __init__(self, master, on_activated): ctk.CTkLabel( outer, text=t("license_crypto_info"), text_color=Colors.STATUS_PENDING, font=(Fonts.FAMILY, Fonts.SIZE_SMALL), wraplength=340, justify="left", + ).pack(anchor="w", pady=(0, 4)) + ctk.CTkLabel( + outer, text=t("license_official_sellers"), text_color=Colors.STATUS_PENDING, + font=(Fonts.FAMILY, Fonts.SIZE_SMALL), wraplength=340, justify="left", ).pack(anchor="w", pady=(0, 10)) self._email = ctk.CTkEntry( @@ -257,17 +261,26 @@ def __init__(self, master, on_activated): ) def _activate(self): - from jasna.protection import ProtectionError, license_store + from jasna.protection import ForgedLicenseError, ProtectionError, RetiredLicenseError, license_store email = self._email.get().strip() key = self._key.get().strip() try: license_store.set_license(email, key) + except ForgedLicenseError: + self._show_error(t("license_forged")) + return + except RetiredLicenseError: + self._show_error(t("license_retired")) + return except ProtectionError as exc: - self._status.configure(text=str(exc), text_color=Colors.STATUS_ERROR) + self._show_error(str(exc)) return self._status.configure(text=t("license_active"), text_color=Colors.STATUS_COMPLETED) self._on_activated() + def _show_error(self, text: str) -> None: + self._status.configure(text=text, text_color=Colors.STATUS_ERROR, wraplength=220, justify="left") + class StatusPill(ctk.CTkFrame): """Status indicator pill shown in header.""" diff --git a/jasna/gui/locales/en.py b/jasna/gui/locales/en.py index 84e3257d..7143de92 100644 --- a/jasna/gui/locales/en.py +++ b/jasna/gui/locales/en.py @@ -298,6 +298,10 @@ "ctx_paste": "Paste", "ctx_select_all": "Select all", "license_crypto_info": "Crypto payments are also supported.", + "license_official_sellers": "Keys are sold only on Buy Me a Coffee and Unifans (kruk2). Keys bought anywhere else are fake.", + "license_forged": "This key was not issued by Jasna. Keys are sold only on Buy Me a Coffee and Unifans (kruk2).", + "license_retired": "This key is from an older version. Use the new key sent to you on Unifans or by email.", + "official_sellers_notice": "Jasna is free on GitHub. Keys are sold only through the two buttons below. Anyone else selling Jasna, keys or modified files is a scam, and their files can run hidden code on your PC.", "license_chip_inactive": "Activate Jasna", "license_chip_active": "Active", "secondary_none": "None", diff --git a/jasna/gui/locales/ja.py b/jasna/gui/locales/ja.py index ea4097ad..336f40e8 100644 --- a/jasna/gui/locales/ja.py +++ b/jasna/gui/locales/ja.py @@ -298,6 +298,10 @@ "ctx_paste": "貼り付け", "ctx_select_all": "すべて選択", "license_crypto_info": "暗号資産での支払いにも対応しています。", + "license_official_sellers": "キーは Buy Me a Coffee と Unifans (kruk2) でのみ販売しています。それ以外で購入したキーは偽物です。", + "license_forged": "このキーは Jasna が発行したものではありません。キーは Buy Me a Coffee と Unifans (kruk2) でのみ販売しています。", + "license_retired": "このキーは旧バージョン用です。Unifans のメッセージまたはメールで届いた新しいキーを使ってください。", + "official_sellers_notice": "Jasna は GitHub で無料で公開しています。キーは下の2つのボタンからのみ販売しています。それ以外で Jasna、キー、改変ファイルを販売している人は詐欺です。そのファイルはPC上で隠れたコードを実行する可能性があります。", "license_chip_inactive": "Jasna を有効化", "license_chip_active": "有効", "secondary_none": "なし", diff --git a/jasna/gui/locales/ko.py b/jasna/gui/locales/ko.py index b4d6a621..a6bf93bd 100644 --- a/jasna/gui/locales/ko.py +++ b/jasna/gui/locales/ko.py @@ -279,6 +279,10 @@ "ctx_paste": "붙여넣기", "ctx_select_all": "모두 선택", "license_crypto_info": "암호화폐 결제도 지원합니다.", + "license_official_sellers": "키는 Buy Me a Coffee와 Unifans (kruk2)에서만 판매합니다. 다른 곳에서 산 키는 가짜입니다.", + "license_forged": "이 키는 Jasna가 발급한 키가 아닙니다. 키는 Buy Me a Coffee와 Unifans (kruk2)에서만 판매합니다.", + "license_retired": "이 키는 이전 버전용입니다. Unifans 메시지나 이메일로 받은 새 키를 사용하세요.", + "official_sellers_notice": "Jasna는 GitHub에서 무료입니다. 키는 아래 두 버튼을 통해서만 판매합니다. 그 외에 Jasna, 키 또는 수정된 파일을 파는 사람은 사기이며, 그 파일은 PC에서 숨겨진 코드를 실행할 수 있습니다.", "license_chip_inactive": "Jasna 활성화", "license_chip_active": "활성", "secondary_none": "없음", diff --git a/jasna/gui/locales/th.py b/jasna/gui/locales/th.py index a0d8a834..e7a2cd39 100644 --- a/jasna/gui/locales/th.py +++ b/jasna/gui/locales/th.py @@ -280,6 +280,10 @@ "ctx_paste": "วาง", "ctx_select_all": "เลือกทั้งหมด", "license_crypto_info": "รองรับการชำระเงินด้วยคริปโทเคอร์เรนซีด้วย", + "license_official_sellers": "คีย์ขายเฉพาะที่ Buy Me a Coffee และ Unifans (kruk2) เท่านั้น คีย์ที่ซื้อจากที่อื่นเป็นของปลอม", + "license_forged": "คีย์นี้ไม่ได้ออกโดย Jasna คีย์ขายเฉพาะที่ Buy Me a Coffee และ Unifans (kruk2) เท่านั้น", + "license_retired": "คีย์นี้เป็นของเวอร์ชันเก่า โปรดใช้คีย์ใหม่ที่ส่งถึงคุณทาง Unifans หรืออีเมล", + "official_sellers_notice": "Jasna ใช้ได้ฟรีบน GitHub คีย์ขายผ่านสองปุ่มด้านล่างเท่านั้น ใครก็ตามที่ขาย Jasna คีย์ หรือไฟล์ที่ถูกดัดแปลงคือมิจฉาชีพ และไฟล์ของพวกเขาอาจรันโค้ดที่ซ่อนอยู่บนเครื่องของคุณ", "license_chip_inactive": "เปิดใช้งาน Jasna", "license_chip_active": "เปิดใช้แล้ว", "secondary_none": "ไม่มี", diff --git a/jasna/gui/locales/zh.py b/jasna/gui/locales/zh.py index 7cee4931..f47a362b 100644 --- a/jasna/gui/locales/zh.py +++ b/jasna/gui/locales/zh.py @@ -298,6 +298,10 @@ "ctx_paste": "粘贴", "ctx_select_all": "全选", "license_crypto_info": "也支持加密货币支付。", + "license_official_sellers": "密钥仅在 Buy Me a Coffee 和 Unifans (kruk2) 出售。其他渠道购买的密钥都是假的。", + "license_forged": "此密钥并非 Jasna 官方签发。密钥仅在 Buy Me a Coffee 和 Unifans (kruk2) 出售。", + "license_retired": "此密钥属于旧版本。请使用通过 Unifans 私信或邮件发送给你的新密钥。", + "official_sellers_notice": "Jasna 在 GitHub 上免费提供。密钥仅通过下方两个按钮出售。其他任何出售 Jasna、密钥或修改版文件的人都是骗子，他们的文件可能会在你的电脑上运行隐藏代码。", "license_chip_inactive": "激活 Jasna", "license_chip_active": "已启用", "secondary_none": "无", diff --git a/jasna/gui/wizard.py b/jasna/gui/wizard.py index e47abdfa..25d9242d 100644 --- a/jasna/gui/wizard.py +++ b/jasna/gui/wizard.py @@ -159,6 +159,15 @@ def _build_footer_loading(self): self._footer = ctk.CTkFrame(self, fg_color="transparent") self._footer.pack(fill="x", side="bottom", padx=40, pady=(20, 40)) + ctk.CTkLabel( + self._footer, + text=t("official_sellers_notice"), + font=(Fonts.FAMILY, Fonts.SIZE_SMALL), + text_color=Colors.STATUS_PENDING, + wraplength=_WINDOW_WIDTH - 80, + justify="center", + ).pack(pady=(0, 12)) + # Button container for centering both buttons btn_container = ctk.CTkFrame(self._footer, fg_color="transparent") btn_container.pack() diff --git a/tests/test_gui_official_sellers.py b/tests/test_gui_official_sellers.py new file mode 100644 index 00000000..1208cf4b --- /dev/null +++ b/tests/test_gui_official_sellers.py @@ -0,0 +1,57 @@ +from __future__ import annotations + +from tkinter import TclError + +import customtkinter as ctk +import pytest + +from jasna.gui import wizard +from jasna.gui.locales import t + + +def _root() -> ctk.CTk: + try: + return ctk.CTk() + except TclError as exc: + pytest.skip(f"Tk display unavailable: {exc}") + + +def _label_texts(widget) -> list[str]: + texts = [widget.cget("text")] if isinstance(widget, ctk.CTkLabel) else [] + for child in widget.winfo_children(): + texts += _label_texts(child) + return texts + + +def test_wizard_names_the_official_sellers(monkeypatch): + monkeypatch.setattr(wizard, "run_system_checks", lambda results: None) + root = _root() + try: + dialog = wizard.FirstRunWizard(root) + assert t("official_sellers_notice") in _label_texts(dialog) + finally: + root.destroy() + + +@pytest.mark.parametrize("error_name, message_key", [ + ("ForgedLicenseError", "license_forged"), + ("RetiredLicenseError", "license_retired"), +]) +def test_license_dialog_explains_rejected_keys(monkeypatch, error_name, message_key): + protection = pytest.importorskip("jasna.protection") + from jasna.gui.components import LicenseDialog + from jasna.protection import license_store + + def reject(email, key): + raise getattr(protection, error_name)("english text") + + monkeypatch.setattr(license_store, "load_license", lambda: None) + monkeypatch.setattr(license_store, "set_license", reject) + root = _root() + try: + dialog = LicenseDialog(root, on_activated=lambda: None) + assert t("license_official_sellers") in _label_texts(dialog) + dialog._activate() + assert dialog._status.cget("text") == t(message_key) + finally: + root.destroy() From b376a50e683d43363882b24432db33c4ab709b97 Mon Sep 17 00:00:00 2001 From: Kruk2 Date: Mon, 28 Sep 2026 15:01:18 +0200 Subject: [PATCH 07/48] feat(sd15): use the updated SD15 checkpoint file Newer versions load a new checkpoint file name, so older versions keep downloading the file they expect. The download skips the old file so users fetch only the checkpoint this version needs. --- jasna/engine_paths.py | 2 +- jasna/restorer/sd15_download.py | 10 +++++++++- jasna/restorer/sd15_inpaint_restorer.py | 2 +- tests/test_sd15_download.py | 16 ++++++++++++++++ 4 files changed, 27 insertions(+), 3 deletions(-) diff --git a/jasna/engine_paths.py b/jasna/engine_paths.py index 8e39ed0e..bb8b0852 100644 --- a/jasna/engine_paths.py +++ b/jasna/engine_paths.py @@ -62,7 +62,7 @@ def get_yolo_tensorrt_engine_path(model_path: str | Path, *, fp16: bool) -> Path SD15_HF_REPO = "Kruk2/sd-15-jav" SD15_DIR = model_weights_dir() / "sd-15-jav" SD15_CKPT_PATH = SD15_DIR / "sd15-200000.ckpt" -SD15_CKPT_ENC_PATH = SD15_DIR / "sd15-200000.ckpt.enc" +SD15_CKPT_ENC_PATH = SD15_DIR / "sd15-200000.v2.ckpt.enc" def get_unet4x_engine_path(onnx_path: str | Path | None = None, fp16: bool = True) -> Path: diff --git a/jasna/restorer/sd15_download.py b/jasna/restorer/sd15_download.py index 32a64fec..0f37a9ab 100644 --- a/jasna/restorer/sd15_download.py +++ b/jasna/restorer/sd15_download.py @@ -10,6 +10,9 @@ logger = logging.getLogger(__name__) +# Checkpoints encrypted under retired keys; they stay on the HF repo for older app versions. +RETIRED_CKPT_ENC_NAMES = ["sd15-200000.ckpt.enc"] + def bundle_present(model_dir: Path) -> bool: """True when a usable SD15 bundle already exists at ``model_dir``. @@ -106,7 +109,12 @@ def download_sd15_bundle( from huggingface_hub import snapshot_download logger.info("Downloading SD15 bundle %s -> %s", repo_id, model_dir) - kwargs: dict[str, Any] = {"repo_id": repo_id, "repo_type": "model", "local_dir": str(model_dir)} + kwargs: dict[str, Any] = { + "repo_id": repo_id, + "repo_type": "model", + "local_dir": str(model_dir), + "ignore_patterns": RETIRED_CKPT_ENC_NAMES, + } if progress_callback is not None: kwargs["tqdm_class"] = _progress_tqdm_class(progress_callback) snapshot_download(**kwargs) diff --git a/jasna/restorer/sd15_inpaint_restorer.py b/jasna/restorer/sd15_inpaint_restorer.py index b6f1d505..0d9d2313 100644 --- a/jasna/restorer/sd15_inpaint_restorer.py +++ b/jasna/restorer/sd15_inpaint_restorer.py @@ -47,7 +47,7 @@ def _read_checkpoint(model_dir: Path) -> dict: class Sd15InpaintRestorer: # AES-GCM subkey id — must match the id used by keytool encrypt-model that - # produced sd15-200000.ckpt.enc (verified against the shipped file: "sd-15-jav"). + # produced the SD15 .enc checkpoint (verified against the shipped file: "sd-15-jav"). MODEL_ID = "sd-15-jav" LATENT_SCALE = 0.18215 diff --git a/tests/test_sd15_download.py b/tests/test_sd15_download.py index 970be986..d10d396c 100644 --- a/tests/test_sd15_download.py +++ b/tests/test_sd15_download.py @@ -25,6 +25,12 @@ def test_absent_without_ckpt(self, tmp_path: Path): (tmp_path / "unet" / "config.json").write_text("{}") assert dl.bundle_present(tmp_path) is False + def test_absent_with_only_retired_checkpoint(self, tmp_path: Path): + (tmp_path / "unet").mkdir() + (tmp_path / "unet" / "config.json").write_text("{}") + (tmp_path / "sd15-200000.ckpt.enc").write_bytes(b"x") + assert dl.bundle_present(tmp_path) is False + def test_absent_without_config(self, tmp_path: Path): (tmp_path / SD15_CKPT_ENC_PATH.name).write_bytes(b"x") assert dl.bundle_present(tmp_path) is False @@ -65,6 +71,16 @@ def test_without_callback_uses_default_hf_progress(self, tmp_path: Path): assert kwargs["local_dir"] == str(tmp_path) assert "tqdm_class" not in kwargs + def test_skips_retired_checkpoint_but_not_the_current_one(self, tmp_path: Path): + from fnmatch import fnmatch + + with patch("huggingface_hub.snapshot_download") as mock_snapshot: + dl.download_sd15_bundle(tmp_path) + + ignored = mock_snapshot.call_args.kwargs["ignore_patterns"] + assert any(fnmatch("sd15-200000.ckpt.enc", pattern) for pattern in ignored) + assert not any(fnmatch(SD15_CKPT_ENC_PATH.name, pattern) for pattern in ignored) + def test_callback_uses_silent_byte_progress_class(self, tmp_path: Path): events: list[tuple[int, int | None]] = [] From 9ac5c8afd01d5f9a7f8200f75446076b71017655 Mon Sep 17 00:00:00 2001 From: Kruk2 Date: Mon, 28 Sep 2026 15:03:07 +0200 Subject: [PATCH 08/48] chore: bump protection submodule --- jasna/protection | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/jasna/protection b/jasna/protection index c8beb4e8..1b2fdafa 160000 --- a/jasna/protection +++ b/jasna/protection @@ -1 +1 @@ -Subproject commit c8beb4e871b884f13270e7414fddca530d220e41 +Subproject commit 1b2fdafa7ca70b9144fdee2e71c16a06a41ec603 From c11a37ed2101d28ce65518b149219ca009606415 Mon Sep 17 00:00:00 2001 From: Kruk2 Date: Mon, 28 Sep 2026 15:09:53 +0200 Subject: [PATCH 09/48] chore: bump protection submodule --- jasna/protection | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/jasna/protection b/jasna/protection index 1b2fdafa..896de2ce 160000 --- a/jasna/protection +++ b/jasna/protection @@ -1 +1 @@ -Subproject commit 1b2fdafa7ca70b9144fdee2e71c16a06a41ec603 +Subproject commit 896de2cef21c72a62b0bc51472662c8a002d8ad8 From 4de27c19d9bbfe81c92ace276a5cb9dfadc5dcd8 Mon Sep 17 00:00:00 2001 From: Kruk2 Date: Mon, 28 Sep 2026 17:04:54 +0200 Subject: [PATCH 10/48] chore: bump protection submodule --- jasna/protection | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/jasna/protection b/jasna/protection index 896de2ce..eb08059b 160000 --- a/jasna/protection +++ b/jasna/protection @@ -1 +1 @@ -Subproject commit 896de2cef21c72a62b0bc51472662c8a002d8ad8 +Subproject commit eb08059bf448808eb08fe82feb33f1a89b50cc54 From 3e33c970c86cd02c49516bf062b7bc851da7c4e9 Mon Sep 17 00:00:00 2001 From: Kruk2 Date: Mon, 28 Sep 2026 18:27:45 +0200 Subject: [PATCH 11/48] fix: preserve existing supporter access Keep downloaded models usable after upgrading. --- jasna/engine_paths.py | 2 +- jasna/protection | 2 +- jasna/restorer/sd15_download.py | 2 +- tests/test_sd15_download.py | 10 +++++++--- 4 files changed, 10 insertions(+), 6 deletions(-) diff --git a/jasna/engine_paths.py b/jasna/engine_paths.py index bb8b0852..8e39ed0e 100644 --- a/jasna/engine_paths.py +++ b/jasna/engine_paths.py @@ -62,7 +62,7 @@ def get_yolo_tensorrt_engine_path(model_path: str | Path, *, fp16: bool) -> Path SD15_HF_REPO = "Kruk2/sd-15-jav" SD15_DIR = model_weights_dir() / "sd-15-jav" SD15_CKPT_PATH = SD15_DIR / "sd15-200000.ckpt" -SD15_CKPT_ENC_PATH = SD15_DIR / "sd15-200000.v2.ckpt.enc" +SD15_CKPT_ENC_PATH = SD15_DIR / "sd15-200000.ckpt.enc" def get_unet4x_engine_path(onnx_path: str | Path | None = None, fp16: bool = True) -> Path: diff --git a/jasna/protection b/jasna/protection index eb08059b..05a5a12f 160000 --- a/jasna/protection +++ b/jasna/protection @@ -1 +1 @@ -Subproject commit eb08059bf448808eb08fe82feb33f1a89b50cc54 +Subproject commit 05a5a12f6991e746e47518fc8f1935e67e9b4b74 diff --git a/jasna/restorer/sd15_download.py b/jasna/restorer/sd15_download.py index 0f37a9ab..5596cf98 100644 --- a/jasna/restorer/sd15_download.py +++ b/jasna/restorer/sd15_download.py @@ -11,7 +11,7 @@ logger = logging.getLogger(__name__) # Checkpoints encrypted under retired keys; they stay on the HF repo for older app versions. -RETIRED_CKPT_ENC_NAMES = ["sd15-200000.ckpt.enc"] +RETIRED_CKPT_ENC_NAMES = ["sd15-200000.v2.ckpt.enc", "sd15-200000.v2.ckpt.enc.manifest.json"] def bundle_present(model_dir: Path) -> bool: diff --git a/tests/test_sd15_download.py b/tests/test_sd15_download.py index d10d396c..9c90bd0e 100644 --- a/tests/test_sd15_download.py +++ b/tests/test_sd15_download.py @@ -25,11 +25,15 @@ def test_absent_without_ckpt(self, tmp_path: Path): (tmp_path / "unet" / "config.json").write_text("{}") assert dl.bundle_present(tmp_path) is False - def test_absent_with_only_retired_checkpoint(self, tmp_path: Path): + def test_original_download_is_reused_without_network_or_prompt(self, tmp_path: Path): (tmp_path / "unet").mkdir() (tmp_path / "unet" / "config.json").write_text("{}") (tmp_path / "sd15-200000.ckpt.enc").write_bytes(b"x") - assert dl.bundle_present(tmp_path) is False + assert dl.bundle_present(tmp_path) is True + with patch.object(dl, "download_sd15_bundle") as download, patch("builtins.input") as prompt: + dl.ensure_sd15_bundle(tmp_path) + download.assert_not_called() + prompt.assert_not_called() def test_absent_without_config(self, tmp_path: Path): (tmp_path / SD15_CKPT_ENC_PATH.name).write_bytes(b"x") @@ -78,7 +82,7 @@ def test_skips_retired_checkpoint_but_not_the_current_one(self, tmp_path: Path): dl.download_sd15_bundle(tmp_path) ignored = mock_snapshot.call_args.kwargs["ignore_patterns"] - assert any(fnmatch("sd15-200000.ckpt.enc", pattern) for pattern in ignored) + assert any(fnmatch("sd15-200000.v2.ckpt.enc", pattern) for pattern in ignored) assert not any(fnmatch(SD15_CKPT_ENC_PATH.name, pattern) for pattern in ignored) def test_callback_uses_silent_byte_progress_class(self, tmp_path: Path): From 4727bdc9e8dea774850148db59a3ae06833d20d7 Mon Sep 17 00:00:00 2001 From: Kruk2 Date: Mon, 28 Sep 2026 19:14:07 +0200 Subject: [PATCH 12/48] fix: improve supporter activation messages Explain activation failures and update supporter documentation. --- README.ja.md | 6 ++++++ README.md | 6 ++++++ README.zh.md | 6 ++++++ docs/en/gui.md | 3 +++ docs/ja/gui.md | 3 +++ docs/zh/gui.md | 2 ++ jasna/gui/components.py | 11 ++++++++++- jasna/gui/locales/en.py | 6 ++++-- jasna/gui/locales/ja.py | 6 ++++-- jasna/gui/locales/ko.py | 6 ++++-- jasna/gui/locales/th.py | 6 ++++-- jasna/gui/locales/zh.py | 6 ++++-- jasna/protection | 2 +- tests/test_gui_official_sellers.py | 2 ++ 14 files changed, 59 insertions(+), 12 deletions(-) diff --git a/README.ja.md b/README.ja.md index d9257382..cc80b92e 100644 --- a/README.ja.md +++ b/README.ja.md @@ -158,6 +158,9 @@ Jasna の公開アプリケーションソースは AGPL-3.0 です。モデル - **unet-4x** セカンダリアップスケーラー。よりシャープな 256->1024 復元用です。 - **SD 1.5 画像復元**。実験的な静止画モデルです。 +既存のキーは v0.11 でもこれらのモデルに使えます。LTX には対象バージョン用の +アップグレードキーが必要です。対象条件は公式販売者にお問い合わせください。 + 結果例: - [unet-4x / セカンダリ復元の例（SLS Discord）](https://discord.com/channels/1196376491815092265/1199059436199759943/1516497879684874260) @@ -172,6 +175,9 @@ Jasna の公開アプリケーションソースは AGPL-3.0 です。モデル キーは上記の2つのページで**のみ**販売しています。それ以外で Jasna、キー、改変ファイルを販売している人は詐欺です。そのキーは公式版では使えず、そのファイルはPC上で隠れたコードを実行する可能性があります。Jasna は GitHub からのみダウンロードしてください。 +有効化できない場合は、購入時のメールアドレスまたはユーザー名を確認し、キー全体をコピーし直してください。 +キーが無効というだけでは不正の証拠にはなりません。確認済みの非公式販売者のキーは別のメッセージで通知します。 + ## 謝辞 - **[Lada](https://codeberg.org/ladaapp/lada)（Codeberg）** — Jasna は Lada に diff --git a/README.md b/README.md index 467d2b9f..d5b29f9f 100644 --- a/README.md +++ b/README.md @@ -161,6 +161,9 @@ Support pays for training extra models, mainly GPU rental and compute time for l - **unet-4x** secondary upscaler for sharper 256->1024 restoration. - **SD 1.5 image restoration**, the experimental still-image model. +Existing keys keep unlocking these models in v0.11. LTX versions require an upgraded key +for the requested version; ask your official seller about eligibility. + Example results: - [unet-4x / secondary restoration examples on SLS Discord](https://discord.com/channels/1196376491815092265/1199059436199759943/1516497879684874260) @@ -175,6 +178,9 @@ How to get a key: Keys are sold **only** on these two pages. Anyone else selling Jasna, keys or modified files is a scam: their keys do not work with the official app, and their files can run hidden code on your PC. Download Jasna only from GitHub. +If activation fails, check the exact purchase email or handle and copy the full key again. +An invalid key alone does not prove fraud; Jasna identifies keys from known unofficial sellers separately. + ## Acknowledgments - **[Lada](https://codeberg.org/ladaapp/lada) (Codeberg)** — Jasna is inspired diff --git a/README.zh.md b/README.zh.md index e377a762..887b67e9 100644 --- a/README.zh.md +++ b/README.zh.md @@ -152,6 +152,9 @@ Jasna 的公开应用程序源代码采用 AGPL-3.0。模型和捆绑依赖项 - **unet-4x** 二级放大模型，用于更清晰的 256->1024 修复。 - **SD 1.5 图像修复**，实验性静态图像模型。 +现有密钥在 v0.11 中仍可解锁这些模型。LTX 需要包含目标版本的升级密钥； +请向官方卖家咨询升级资格。 + 结果示例: - [unet-4x / 二级修复示例（SLS Discord）](https://discord.com/channels/1196376491815092265/1199059436199759943/1516497879684874260) @@ -166,6 +169,9 @@ Jasna 的公开应用程序源代码采用 AGPL-3.0。模型和捆绑依赖项 密钥**仅**在以上两个页面出售。其他任何出售 Jasna、密钥或修改版文件的人都是骗子：他们的密钥无法用于官方版本，他们的文件可能会在你的电脑上运行隐藏代码。请只从 GitHub 下载 Jasna。 +若激活失败，请核对购买时使用的邮箱或账号，并重新复制完整密钥。 +密钥无效本身并不能证明欺诈；Jasna 会单独提示来自已知非官方卖家的密钥。 + ## 致谢 - **[Lada](https://codeberg.org/ladaapp/lada)（Codeberg）** — Jasna 受 Lada 启发， diff --git a/docs/en/gui.md b/docs/en/gui.md index 4df6824c..f61a6eac 100644 --- a/docs/en/gui.md +++ b/docs/en/gui.md @@ -101,3 +101,6 @@ after a driver update or when something misbehaves. whether it's active. Keys are sold only on Buy Me a Coffee and Unifans (kruk2); keys from other sellers are fake and do not work. + If verification fails, check the exact purchase email or handle and paste the complete + key again. Known unofficial keys get a separate warning. Existing keys still unlock + UNet and SD; a newer LTX version needs an upgraded key for that version. diff --git a/docs/ja/gui.md b/docs/ja/gui.md index b6e56e55..07d4514c 100644 --- a/docs/ja/gui.md +++ b/docs/ja/gui.md @@ -98,3 +98,6 @@ Space キーで再生と一時停止を切り替え、左右矢印キーで 30 有効かどうかが分かります。 キーは Buy Me a Coffee と Unifans (kruk2) でのみ販売しています。他の販売者の キーは偽物で、使えません。 + 確認に失敗した場合は、購入時のメールアドレスまたはユーザー名を確認し、キー全体を貼り付け直してください。 + 確認済みの非公式キーは別の警告で通知します。既存のキーは引き続き UNet と SD に使えます。 + 新しい LTX バージョンには、そのバージョン用のアップグレードキーが必要です。 diff --git a/docs/zh/gui.md b/docs/zh/gui.md index 5033deba..5f93bebb 100644 --- a/docs/zh/gui.md +++ b/docs/zh/gui.md @@ -80,3 +80,5 @@ - **系统状态**: 处理时底部工具条会显示 GPU、VRAM、内存和 CPU 占用。 - **许可证**: 在头部输入你的支持者密钥；状态标签会显示是否已激活。 密钥仅在 Buy Me a Coffee 和 Unifans (kruk2) 出售；其他卖家的密钥都是假的，无法使用。 + 若验证失败，请核对购买时使用的邮箱或账号，并重新粘贴完整密钥。已知的非官方密钥会单独提示。 + 现有密钥仍可解锁 UNet 和 SD；新的 LTX 版本需要包含该版本的升级密钥。 diff --git a/jasna/gui/components.py b/jasna/gui/components.py index 2266cd26..0acc763a 100644 --- a/jasna/gui/components.py +++ b/jasna/gui/components.py @@ -261,7 +261,10 @@ def __init__(self, master, on_activated): ) def _activate(self): - from jasna.protection import ForgedLicenseError, ProtectionError, RetiredLicenseError, license_store + from jasna.protection import ( + ForgedLicenseError, LicenseError, MalformedLicenseError, + ProtectionError, RetiredLicenseError, license_store, + ) email = self._email.get().strip() key = self._key.get().strip() try: @@ -272,6 +275,12 @@ def _activate(self): except RetiredLicenseError: self._show_error(t("license_retired")) return + except MalformedLicenseError: + self._show_error(t("license_malformed")) + return + except LicenseError: + self._show_error(t("license_invalid")) + return except ProtectionError as exc: self._show_error(str(exc)) return diff --git a/jasna/gui/locales/en.py b/jasna/gui/locales/en.py index 7143de92..dee5fa6e 100644 --- a/jasna/gui/locales/en.py +++ b/jasna/gui/locales/en.py @@ -299,8 +299,10 @@ "ctx_select_all": "Select all", "license_crypto_info": "Crypto payments are also supported.", "license_official_sellers": "Keys are sold only on Buy Me a Coffee and Unifans (kruk2). Keys bought anywhere else are fake.", - "license_forged": "This key was not issued by Jasna. Keys are sold only on Buy Me a Coffee and Unifans (kruk2).", - "license_retired": "This key is from an older version. Use the new key sent to you on Unifans or by email.", + "license_forged": "This key comes from a known unofficial seller, not Jasna. It cannot activate the official app. Get an official key through Buy Me a Coffee or Unifans (kruk2).", + "license_retired": "This key was retired. Contact your official seller for a replacement. Original v0.10 keys are still accepted.", + "license_invalid": "Could not verify this key. Check the exact purchase email or handle and paste the complete key. Keys for modified or unofficial builds do not work here. Still failing? Contact your official seller.", + "license_malformed": "The key is incomplete or has an invalid format. Copy the complete key from your official purchase message.", "official_sellers_notice": "Jasna is free on GitHub. Keys are sold only through the two buttons below. Anyone else selling Jasna, keys or modified files is a scam, and their files can run hidden code on your PC.", "license_chip_inactive": "Activate Jasna", "license_chip_active": "Active", diff --git a/jasna/gui/locales/ja.py b/jasna/gui/locales/ja.py index 336f40e8..c9db9e4b 100644 --- a/jasna/gui/locales/ja.py +++ b/jasna/gui/locales/ja.py @@ -299,8 +299,10 @@ "ctx_select_all": "すべて選択", "license_crypto_info": "暗号資産での支払いにも対応しています。", "license_official_sellers": "キーは Buy Me a Coffee と Unifans (kruk2) でのみ販売しています。それ以外で購入したキーは偽物です。", - "license_forged": "このキーは Jasna が発行したものではありません。キーは Buy Me a Coffee と Unifans (kruk2) でのみ販売しています。", - "license_retired": "このキーは旧バージョン用です。Unifans のメッセージまたはメールで届いた新しいキーを使ってください。", + "license_forged": "このキーは確認済みの非公式販売者によるもので、Jasna が発行したものではありません。公式アプリでは使えません。Buy Me a Coffee または Unifans (kruk2) で公式キーを入手してください。", + "license_retired": "このキーは無効になりました。公式販売者に交換を依頼してください。元の v0.10 キーは引き続き使えます。", + "license_invalid": "キーを確認できませんでした。購入時のメールアドレスまたはユーザー名を確認し、キー全体を貼り付けてください。改変版や非公式版のキーは使えません。解決しない場合は公式販売者にお問い合わせください。", + "license_malformed": "キーが不完全か、形式が正しくありません。公式の購入メッセージからキー全体をコピーしてください。", "official_sellers_notice": "Jasna は GitHub で無料で公開しています。キーは下の2つのボタンからのみ販売しています。それ以外で Jasna、キー、改変ファイルを販売している人は詐欺です。そのファイルはPC上で隠れたコードを実行する可能性があります。", "license_chip_inactive": "Jasna を有効化", "license_chip_active": "有効", diff --git a/jasna/gui/locales/ko.py b/jasna/gui/locales/ko.py index a6bf93bd..d806203f 100644 --- a/jasna/gui/locales/ko.py +++ b/jasna/gui/locales/ko.py @@ -280,8 +280,10 @@ "ctx_select_all": "모두 선택", "license_crypto_info": "암호화폐 결제도 지원합니다.", "license_official_sellers": "키는 Buy Me a Coffee와 Unifans (kruk2)에서만 판매합니다. 다른 곳에서 산 키는 가짜입니다.", - "license_forged": "이 키는 Jasna가 발급한 키가 아닙니다. 키는 Buy Me a Coffee와 Unifans (kruk2)에서만 판매합니다.", - "license_retired": "이 키는 이전 버전용입니다. Unifans 메시지나 이메일로 받은 새 키를 사용하세요.", + "license_forged": "이 키는 확인된 비공식 판매자가 발급했으며 Jasna의 공식 키가 아닙니다. 공식 앱에서 사용할 수 없습니다. Buy Me a Coffee 또는 Unifans (kruk2)에서 공식 키를 받으세요.", + "license_retired": "이 키는 사용 중지되었습니다. 공식 판매자에게 교체를 요청하세요. 기존 v0.10 키는 계속 사용할 수 있습니다.", + "license_invalid": "키를 확인할 수 없습니다. 구매 시 사용한 이메일이나 사용자 이름을 확인하고 전체 키를 붙여넣으세요. 수정본 또는 비공식 버전의 키는 사용할 수 없습니다. 계속 실패하면 공식 판매자에게 문의하세요.", + "license_malformed": "키가 불완전하거나 형식이 잘못되었습니다. 공식 구매 메시지에서 전체 키를 복사하세요.", "official_sellers_notice": "Jasna는 GitHub에서 무료입니다. 키는 아래 두 버튼을 통해서만 판매합니다. 그 외에 Jasna, 키 또는 수정된 파일을 파는 사람은 사기이며, 그 파일은 PC에서 숨겨진 코드를 실행할 수 있습니다.", "license_chip_inactive": "Jasna 활성화", "license_chip_active": "활성", diff --git a/jasna/gui/locales/th.py b/jasna/gui/locales/th.py index e7a2cd39..c21c2810 100644 --- a/jasna/gui/locales/th.py +++ b/jasna/gui/locales/th.py @@ -281,8 +281,10 @@ "ctx_select_all": "เลือกทั้งหมด", "license_crypto_info": "รองรับการชำระเงินด้วยคริปโทเคอร์เรนซีด้วย", "license_official_sellers": "คีย์ขายเฉพาะที่ Buy Me a Coffee และ Unifans (kruk2) เท่านั้น คีย์ที่ซื้อจากที่อื่นเป็นของปลอม", - "license_forged": "คีย์นี้ไม่ได้ออกโดย Jasna คีย์ขายเฉพาะที่ Buy Me a Coffee และ Unifans (kruk2) เท่านั้น", - "license_retired": "คีย์นี้เป็นของเวอร์ชันเก่า โปรดใช้คีย์ใหม่ที่ส่งถึงคุณทาง Unifans หรืออีเมล", + "license_forged": "คีย์นี้มาจากผู้ขายที่ตรวจพบว่าไม่ใช่ผู้ขายทางการของ Jasna จึงเปิดใช้งานแอปทางการไม่ได้ โปรดรับคีย์ทางการผ่าน Buy Me a Coffee หรือ Unifans (kruk2)", + "license_retired": "คีย์นี้ถูกยกเลิกแล้ว โปรดติดต่อผู้ขายทางการเพื่อขอคีย์ทดแทน คีย์เดิมของ v0.10 ยังใช้ได้", + "license_invalid": "ตรวจสอบคีย์นี้ไม่ได้ โปรดตรวจสอบอีเมลหรือชื่อผู้ใช้ที่ใช้ซื้อและวางคีย์ให้ครบ คีย์สำหรับแอปดัดแปลงหรือแอปที่ไม่เป็นทางการใช้ที่นี่ไม่ได้ หากยังไม่สำเร็จ โปรดติดต่อผู้ขายทางการ", + "license_malformed": "คีย์ไม่ครบหรือรูปแบบไม่ถูกต้อง โปรดคัดลอกคีย์ทั้งหมดจากข้อความยืนยันการซื้ออย่างเป็นทางการ", "official_sellers_notice": "Jasna ใช้ได้ฟรีบน GitHub คีย์ขายผ่านสองปุ่มด้านล่างเท่านั้น ใครก็ตามที่ขาย Jasna คีย์ หรือไฟล์ที่ถูกดัดแปลงคือมิจฉาชีพ และไฟล์ของพวกเขาอาจรันโค้ดที่ซ่อนอยู่บนเครื่องของคุณ", "license_chip_inactive": "เปิดใช้งาน Jasna", "license_chip_active": "เปิดใช้แล้ว", diff --git a/jasna/gui/locales/zh.py b/jasna/gui/locales/zh.py index f47a362b..2ecc9168 100644 --- a/jasna/gui/locales/zh.py +++ b/jasna/gui/locales/zh.py @@ -299,8 +299,10 @@ "ctx_select_all": "全选", "license_crypto_info": "也支持加密货币支付。", "license_official_sellers": "密钥仅在 Buy Me a Coffee 和 Unifans (kruk2) 出售。其他渠道购买的密钥都是假的。", - "license_forged": "此密钥并非 Jasna 官方签发。密钥仅在 Buy Me a Coffee 和 Unifans (kruk2) 出售。", - "license_retired": "此密钥属于旧版本。请使用通过 Unifans 私信或邮件发送给你的新密钥。", + "license_forged": "此密钥来自已知的非官方卖家，并非 Jasna 签发，无法激活官方应用。请通过 Buy Me a Coffee 或 Unifans (kruk2) 获取官方密钥。", + "license_retired": "此密钥已停用，请联系官方卖家获取替换密钥。原版 v0.10 密钥仍然有效。", + "license_invalid": "无法验证此密钥。请核对购买时使用的邮箱或用户名，并粘贴完整密钥。修改版或非官方版本的密钥在此无效。若仍失败，请联系官方卖家。", + "license_malformed": "密钥不完整或格式无效。请从官方购买消息中复制完整密钥。", "official_sellers_notice": "Jasna 在 GitHub 上免费提供。密钥仅通过下方两个按钮出售。其他任何出售 Jasna、密钥或修改版文件的人都是骗子，他们的文件可能会在你的电脑上运行隐藏代码。", "license_chip_inactive": "激活 Jasna", "license_chip_active": "已启用", diff --git a/jasna/protection b/jasna/protection index 05a5a12f..8cdc9890 160000 --- a/jasna/protection +++ b/jasna/protection @@ -1 +1 @@ -Subproject commit 05a5a12f6991e746e47518fc8f1935e67e9b4b74 +Subproject commit 8cdc9890c691f780c761a8671e11df83e0ea5423 diff --git a/tests/test_gui_official_sellers.py b/tests/test_gui_official_sellers.py index 1208cf4b..d82f3aa8 100644 --- a/tests/test_gui_official_sellers.py +++ b/tests/test_gui_official_sellers.py @@ -36,6 +36,8 @@ def test_wizard_names_the_official_sellers(monkeypatch): @pytest.mark.parametrize("error_name, message_key", [ ("ForgedLicenseError", "license_forged"), ("RetiredLicenseError", "license_retired"), + ("LicenseError", "license_invalid"), + ("MalformedLicenseError", "license_malformed"), ]) def test_license_dialog_explains_rejected_keys(monkeypatch, error_name, message_key): protection = pytest.importorskip("jasna.protection") From ebfd78527f5728bad5fd2cbe4f4fcb44ca784b9f Mon Sep 17 00:00:00 2001 From: Kruk2 Date: Mon, 28 Sep 2026 19:17:13 +0200 Subject: [PATCH 13/48] chore: update private component --- jasna/protection | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/jasna/protection b/jasna/protection index 8cdc9890..37a800ff 160000 --- a/jasna/protection +++ b/jasna/protection @@ -1 +1 @@ -Subproject commit 8cdc9890c691f780c761a8671e11df83e0ea5423 +Subproject commit 37a800fffb365be0cc54c33176141b360ad59a1b From 83f37d187e6e62d78418e2b5ae558ce5a7c6bcc0 Mon Sep 17 00:00:00 2001 From: Kruk2 Date: Mon, 28 Sep 2026 19:17:37 +0200 Subject: [PATCH 14/48] feat(ltx): add LTX-2.5 diffusion video restoration (--restoration-model-name ltx) Inference-only restorer for the fine-tuned LTX-2.5 model, run in phased passes with one model on the GPU at a time: scan (detect, cuts, plan) -> VAE encode -> denoise -> decode + composite + encode. - jasna/ltx: track planning with one zooming follow camera per track and 121-frame windows at hop 88; video-only transformer port with prequantized INT8 ConvRot linears; lockstep sampler with overlap fusion and shared-prefix STG; block streaming from pinned RAM for low-VRAM GPUs; GPU crop/composite. - jasna/models/ltx_vae: vendored pure-PyTorch subset of the ltx-core video VAE (causal encoder, single-step diffusion decoder, tiled decode). - CLI: --restoration-model-name ltx, --ltx-large-canvas, --ltx-seed; --restoration-model-path now defaults per model. - Licensing: LTX-2.x Community License text and notices. --- LICENSING.md | 5 +- NOTICE | 2 +- assets/THIRD_PARTY_LICENSES.md | 8 + assets/THIRD_PARTY_MODELS.md | 15 + assets/licenses/LTX-2.x-Community-License.txt | 580 ++++++++++++++++++ assets/licenses/README.md | 2 + jasna/benchmark/__init__.py | 3 +- jasna/cli_help.py | 13 + jasna/engine_paths.py | 10 + jasna/gui/engine_preflight.py | 4 +- jasna/gui/video_session.py | 9 +- jasna/ltx/blocks.py | 130 ++++ jasna/ltx/camera.py | 182 ++++++ jasna/ltx/compose.py | 179 ++++++ jasna/ltx/plan.py | 333 ++++++++++ jasna/ltx/regions.py | 51 ++ jasna/ltx/restore.py | 301 +++++++++ jasna/ltx/sampler.py | 180 ++++++ jasna/ltx/transformer.py | 306 +++++++++ jasna/main.py | 37 +- jasna/models/ltx_vae/__init__.py | 15 + jasna/models/ltx_vae/diffusion_tiling.py | 488 +++++++++++++++ .../models/ltx_vae/diffusion_video_decoder.py | 380 ++++++++++++ jasna/models/ltx_vae/eager_na.py | 148 +++++ jasna/models/ltx_vae/loader.py | 174 ++++++ jasna/models/ltx_vae/ops.py | 50 ++ jasna/models/ltx_vae/tiling.py | 230 +++++++ jasna/models/ltx_vae/transformer.py | 418 +++++++++++++ jasna/models/ltx_vae/video_encoder.py | 165 +++++ jasna/pipeline.py | 73 ++- jasna/session_config.py | 6 + jasna/session_factory.py | 43 +- tests/factories.py | 4 + tests/test_ltx_blocks.py | 31 + tests/test_ltx_compose.py | 71 +++ tests/test_ltx_plan.py | 117 ++++ tests/test_ltx_transformer.py | 81 +++ tests/test_ltx_vae.py | 227 +++++++ tests/test_session_config.py | 19 + 39 files changed, 5059 insertions(+), 31 deletions(-) create mode 100644 assets/licenses/LTX-2.x-Community-License.txt create mode 100644 jasna/ltx/blocks.py create mode 100644 jasna/ltx/camera.py create mode 100644 jasna/ltx/compose.py create mode 100644 jasna/ltx/plan.py create mode 100644 jasna/ltx/regions.py create mode 100644 jasna/ltx/restore.py create mode 100644 jasna/ltx/sampler.py create mode 100644 jasna/ltx/transformer.py create mode 100644 jasna/models/ltx_vae/__init__.py create mode 100644 jasna/models/ltx_vae/diffusion_tiling.py create mode 100644 jasna/models/ltx_vae/diffusion_video_decoder.py create mode 100644 jasna/models/ltx_vae/eager_na.py create mode 100644 jasna/models/ltx_vae/loader.py create mode 100644 jasna/models/ltx_vae/ops.py create mode 100644 jasna/models/ltx_vae/tiling.py create mode 100644 jasna/models/ltx_vae/transformer.py create mode 100644 jasna/models/ltx_vae/video_encoder.py create mode 100644 tests/test_ltx_blocks.py create mode 100644 tests/test_ltx_compose.py create mode 100644 tests/test_ltx_plan.py create mode 100644 tests/test_ltx_transformer.py create mode 100644 tests/test_ltx_vae.py diff --git a/LICENSING.md b/LICENSING.md index 0c545cef..0c9349a4 100644 --- a/LICENSING.md +++ b/LICENSING.md @@ -14,7 +14,10 @@ copyright and license notices. Jasna includes code and model weights from [Lada](https://codeberg.org/ladaapp/lada). Those components remain covered by their applicable AGPL-3.0 notices. The vendored MMagic subset retains its -OpenMMLab copyright and Apache-2.0 notices. +OpenMMLab copyright and Apache-2.0 notices. The LTX restoration code ported +from Lightricks LTX-2.5 (`jasna/ltx/transformer.py`, `jasna/models/ltx_vae/`) +and the LTX restoration model weights are covered by the LTX-2.x Community +License. The complete Jasna license is in [LICENSE](LICENSE). Copyright and attribution notices are in [NOTICE](NOTICE), diff --git a/NOTICE b/NOTICE index f6bb073e..1d612358 100644 --- a/NOTICE +++ b/NOTICE @@ -6,7 +6,7 @@ is licensed under the GNU Affero General Public License version 3.0. The full license is in LICENSE. Jasna incorporates or distributes work from Lada, OpenMMLab MMagic, -Roboflow RF-DETR, PyAV, VALI, FFmpeg, VideoLAN, Python, and other projects. +Lightricks LTX-2.5, Roboflow RF-DETR, PyAV, VALI, FFmpeg, VideoLAN, Python, and other projects. Their copyright, attribution, source, and license information is retained in: - LICENSING.md diff --git a/assets/THIRD_PARTY_LICENSES.md b/assets/THIRD_PARTY_LICENSES.md index 731dae7b..9e6aa138 100644 --- a/assets/THIRD_PARTY_LICENSES.md +++ b/assets/THIRD_PARTY_LICENSES.md @@ -44,6 +44,14 @@ Jasna's `jasna/models/basicvsrpp/mmagic/` directory is an inference-only subset derived from OpenMMLab MMagic. Its existing copyright headers and Apache-2.0 terms are retained. +Jasna's `jasna/models/ltx_vae/` directory is an inference-only subset of the +video VAE in [Lightricks LTX-2.5](https://github.com/Lightricks/LTX-2) `ltx-core` +(commit `6ea1527869a5ce57452e215595eae189a7cf65cc`), and +`jasna/ltx/transformer.py` is a port of its video transformer. Both are covered +by the LTX-2.x Community License (`assets/licenses/LTX-2.x-Community-License.txt`). +`jasna/models/ltx_vae/eager_na.py` originates from Comfy Org comfy-kitchen and +keeps its Apache-2.0 header. + Jasna also contains code derived from [Lada](https://codeberg.org/ladaapp/lada), licensed under AGPL-3.0. Jasna's root `LICENSE` contains the complete AGPL-3.0 text. diff --git a/assets/THIRD_PARTY_MODELS.md b/assets/THIRD_PARTY_MODELS.md index 6cc1326e..644492bc 100644 --- a/assets/THIRD_PARTY_MODELS.md +++ b/assets/THIRD_PARTY_MODELS.md @@ -27,6 +27,21 @@ This checkpoint is redistributed unmodified. This checkpoint is redistributed unmodified. +## LTX restoration (LTX-2.5 video VAE and fine-tuned transformer) + +- Jasna model name: `ltx` +- Folder: `model_weights/ltx-restore/` +- `vae.safetensors`: Lightricks `ltx-2.5-video-vae-bf16.safetensors`, redistributed unmodified +- `vae.safetensors` SHA-256: `847e14ca7f3355debca0cea4eaa24ac0fbcdf0061da054ac89ca638a869ddba3` +- `vae-decoder.safetensors`: Jasna fine-tune of the LTX-2.5 video VAE decoder +- `transformer.safetensors`: Jasna fine-tune of the LTX-2.5 22B transformer (video path only, INT8) +- Upstream: [Lightricks LTX-2.5](https://huggingface.co/Lightricks/LTX-2.5) +- License: LTX-2.x Community License (`assets/licenses/LTX-2.x-Community-License.txt`) +- Copyright: Lightricks Ltd.; fine-tunes 2026 Kruk2 + +These files are not part of release packages yet; hashes for the fine-tuned +files are added when they are. + ## Jasna RF-DETR v6 - Jasna model name: `rfdetr-v6` diff --git a/assets/licenses/LTX-2.x-Community-License.txt b/assets/licenses/LTX-2.x-Community-License.txt new file mode 100644 index 00000000..aa1e32a7 --- /dev/null +++ b/assets/licenses/LTX-2.x-Community-License.txt @@ -0,0 +1,580 @@ + LTX-2.x Community License Agreement + License date: August 11, 2026 + + + By downloading, using, accessing or distributing any portion or + element of LTX-2.x, you agree that you have read and accepted to be + bound by this Agreement. + + 1. Definitions + + 1.1 "Agreement" means the terms and conditions for this LTX-2.x + Community License Agreement and the exhibits, attachments, and + Complementary Materials, as specified in this document. + + 1.2 "Complementary Materials" means the accompanying + documentation, tutorials, examples, configuration files and + other materials made available by Licensor together with the + LTX-2.x model weights and parameters, in each case as + distributed by Licensor. + + 1.3 "Control" means the direct or indirect ownership of more than + fifty percent (50%) of the voting securities or other + ownership interests, or the power to direct the management and + policies of such Entity through voting rights, contract, or + otherwise. + + 1.4 "Data" means a collection of information and/or content + extracted from the dataset used with LTX-2.x, including to + train, pretrain, or otherwise evaluate LTX-2.x. The Data is + not licensed under this Agreement. + + 1.5 "Derivatives of LTX-2.x" means all modifications to LTX-2.x, + works based on LTX-2.x, or any other model which is created or + initialized by transfer of patterns of the weights, + parameters, activations or output of LTX-2.x, to the other + model, in order to cause the other model to perform similarly + to LTX-2.x, including - but not limited to - distillation + methods entailing the use of intermediate data representations + or methods based on the generation of synthetic data by + LTX-2.x for training the other model. For clarity, Derivatives + of LTX-2.x include: (i) any fine-tuned or adapted weights, + parameters, or checkpoints derived from LTX-2.x; (ii) + derivative model architectures that incorporate or are based + upon LTX-2.x's architecture; and (iii) any modified or + extended versions of the Complementary Materials. + + 1.6 "Entity" means any individual, corporation, partnership, + limited liability company, or other legal entity. For purposes + of this Agreement, an Entity shall be deemed to include, on an + aggregative basis, all subsidiaries, affiliates, and other + companies under common Control with such Entity. When + determining whether an Entity meets any threshold under this + Agreement (including revenue thresholds in Section 2.1), all + subsidiaries, affiliates, and companies under common Control + shall be considered collectively. + + 1.7 "Harm" includes but is not limited to physical, mental, + psychological, financial and reputational damage, pain, or + loss. + + 1.8 "Licensor" or "LTX" means the owner that is granting the + license under this Agreement. For the purposes of this + Agreement, the Licensor is Lightricks Ltd. + + 1.9 "LTX-2.x" means the large generative models, + text/image/video/audio/3D generation models, and multimodal + large language models and their software and algorithms, + including trained model weights, parameters (including + optimizer states), machine-learning model code, + inference-enabling code, training-enabling code, fine-tuning + enabling code, accompanying source code, scripts, and all + other elements of the foregoing distributed and made publicly + available by LTX (including, for example, at + https://github.com/Lightricks/LTX-2). This license is + applicable to all LTX-2.5 versions released since August 11, + 2026, and all future releases of LTX-2.x under this license. + + 1.10 "Output" means the results of operating LTX-2.x as embodied + in informational content resulting therefrom. + + 1.11 "you" (or "your") means an individual or legal Entity + licensing LTX-2.x in accordance with this Agreement and/or + otherwise downloading, accessing, distributing or using + LTX-2.x for whichever purpose and in any field of use, + including usage of LTX-2.x in an end-use application - e.g. + chatbot, translator, image generator. + + 2. Grant of License. + + 2.1 Subject to your compliance with the terms and conditions of + this Agreement, you are granted a non-exclusive, worldwide, + non-transferable and royalty-free limited license under + Licensor's intellectual property or other rights owned by + Licensor embodied in LTX-2.x to use, reproduce, prepare, + distribute, publicly display, publicly perform, sublicense, + copy, create derivative works of, and make modifications to + LTX-2.x, for any purpose, subject to the restrictions set + forth in Attachment A; provided however, that Entities with + annual revenues of at least $10,000,000 (the "Commercial + Entities") are required to obtain a paid license for any use + (excluding use solely for a Non-Commercial Purpose as set + forth in Section 2.2) of LTX-2.x and Derivatives of LTX-2.x + (such paid license referred to herein as a "Commercial Use + Agreement"), as will be provided by the Licensor. Commercial + Entities interested in such a commercial license are required + to contact Licensor (ltxv-licensing@lightricks.com). Any use + of LTX-2.x or Derivatives of LTX-2.x by Commercial Entities + not in accordance with this Agreement and/or the Commercial + Use Agreement is strictly prohibited and shall be deemed a + material breach of this Agreement. In the event of such + material breach, and without limiting Licensor's right to + terminate the Agreement or to pursue any other remedies + available at law or in equity, you shall pay Licensor the + license fees owed for the period such Commercial Entity used + LTX-2.x (calculated at Licensor's standard commercial license + fees, in effect during the relevant period or, absent + published standard fees, a reasonable market rate for a + comparable license), within thirty (30) days of Licensor's + written demand. + + 2.2 Notwithstanding the foregoing or anything to the contrary in + this Agreement, a Commercial Entity may download and use + LTX-2.x and Derivatives of LTX-2.x without obtaining the + Commercial Use Agreement solely for a Non-Commercial Purpose. + "Non-Commercial Purpose" means any of the following uses, but + only so far as such Commercial Entity does not receive any + direct or indirect payment arising from the use of LTX-2.x or + Derivatives of LTX-2.x: (i) use by an individual acting in a + personal capacity for research, experimentation, learning, + private study, hobby or recreational projects, or personal + entertainment, in each case where such use is not connected, + directly or indirectly, to any commercial activity, business + operation, or the performance of duties for an employer or any + other Entity; and (ii) use by a Commercial Entity for testing, + evaluation, or non-commercial research and development in a + non-production or development environment. For clarity, use + (a) for revenue-generating activity in any manner, whether + direct or indirect, (b) in direct interactions with or that + has impact on end users, or (c) to train, fine-tune, or + distill any model (including any Derivative of LTX-2.x) for + commercial use, in each case, is not a Non-Commercial Purpose + and requires all Commercial Entities to obtain a paid license + under the Commercial Use Agreement prior to such use. For the + avoidance of doubt, the permission granted under this Section + 2.2 is a limited right of use only and does not convey or + transfer any ownership right, title, or interest in or to + LTX-2.x or any Derivatives of LTX-2.x, and all Derivatives of + LTX-2.x created or used pursuant to this Section 2.2 remain + subject to the terms of this Agreement, including Section 1.5. + + 3. Distribution and Redistribution. You may host for third parties + remote access purposes (e.g. software-as-a-service), reproduce and + distribute copies of LTX-2.x or Derivatives of LTX-2.x thereof in + any medium, with or without modifications, provided that you meet + the following conditions: + + 3.1 Use-based restrictions as referenced in Section 4 and all + provisions of Attachment A MUST be included as an enforceable + provision by you in any type of legal agreement (e.g. a + license) governing the use and/or distribution of LTX-2.x or + Derivatives of LTX-2.x, and you shall give notice to + subsequent users you distribute to, that LTX-2.x or + Derivatives of LTX-2.x are subject to Section 4 and Attachment + A in their entirety, including all use restrictions and + acceptable use policies; + + 3.2 You must provide any third-party recipients of LTX-2.x or + Derivatives of LTX-2.x a copy of this Agreement, including all + attachments and use policies. Any Derivative of LTX-2.x (as + defined in Section 1.5, including but not limited to + fine-tuned weights, modified training code, models trained on + Outputs, or any other derivative) must be distributed + exclusively under the terms of this Agreement, subject to + Section 3.6, with a complete copy of this Agreement included; + + 3.3 You must cause any modified files to carry prominent notices + stating that you changed the files; + + 3.4 You must retain all copyright, patent, trademark, and + attribution notices excluding those notices that do not + pertain to any part of LTX-2.x, Derivatives of LTX-2.x. + + 3.5 Transfer of Derivatives. No transfer of any Derivative of + LTX-2.x (including any fine-tuned weights, LoRA adapters, or + similar adaptations) to a third party shall grant such third + party any right, title, license, or authorization to access, + use, reproduce, distribute, or exploit LTX-2.x, or any + Derivative of LTX-2.x beyond the rights granted under this + Agreement. If the transferee is a Commercial Entity (as + defined in Section 2), it must obtain a paid license from + Licensor prior to any use of any Derivative of LTX-2.x, + regardless of who created such Derivative. Prior to or at the + time of any such transfer, you shall notify the transferee in + writing that (i) use of such Derivative of LTX-2.x is subject + to the terms of this Agreement, and (ii) if the transferee is + a Commercial Entity, it must obtain a separate paid license to + LTX-2.x from Licensor. You shall not transfer any Derivative + of LTX-2.x to a Commercial Entity unless such Commercial + Entity has obtained the required paid license from Licensor + prior to any use, and unless the proposed transferee has been + so informed. You and the transferee shall each be responsible + for ensuring the transferee obtains the required license from + Licensor prior to any use of LTX-2.x or Derivative of LTX-2.x. + Nothing in this Section 3.5 shall require a Commercial Entity + to obtain a paid license for use solely for a Non-Commercial + Purpose as permitted under Section 2.2. + + 3.6 You may add your own copyright statement to your modifications + and may provide additional license terms and conditions - + respecting Section 3.1 - for use, reproduction, or + distribution of your modifications, or for any such + Derivatives of LTX-2.x as a whole, provided your use, + reproduction, and distribution of LTX-2.x otherwise complies + with the conditions stated in this Agreement, and you provide + a complete copy of this Agreement with any such use, + reproduction and distribution of LTX-2.x and any Derivatives + thereof; provided that any such additional terms shall be + additive only and shall not derogate from, conflict with, + waive, or purport to modify any term of this Agreement, and + this Agreement shall govern in the event of any conflict. + + 4. Use-based restrictions. The restrictions set forth in Attachment A + are considered Use-based restrictions. Therefore, you cannot use + LTX-2.x and the Derivatives of LTX-2.x in violation of the + specified restricted uses. You may use LTX-2.x subject to this + Agreement, only for lawful purposes and in accordance with the + Agreement. "Use" may include creating any content with, + fine-tuning, updating, running, training, evaluating and/or + re-parametrizing LTX-2.x. You shall require all of your users who + use LTX-2.x or a Derivative of LTX-2.x to comply with the terms of + this Section 4. + + 5. The Output You Generate. Except as set forth herein, Licensor + claims no rights in the Output you generate using LTX-2.x. You are + accountable for input you insert into LTX-2.x, the Output you + generate and its subsequent uses. No use of the Output can + contravene any provision as stated in the Agreement. + + 6. Updates and Runtime Restrictions; AI Regulations. To the maximum + extent permitted by law, Licensor reserves the right to restrict + (remotely or otherwise) usage of LTX-2.x in violation of this + Agreement, update LTX-2.x through electronic means, or modify the + Output of LTX-2.x based on updates. You shall undertake reasonable + efforts to use the latest version of LTX-2.x. Any use of the + non-current version of LTX-2.x is done solely at your risk. To the + extent applicable to you, you shall comply with all laws and + regulations governing artificial intelligence that apply to your + use, deployment, or distribution of LTX-2.x, Derivatives of + LTX-2.x, or Outputs, including Regulation (EU) 2024/1689 (the "EU + AI Act") and the California AI Transparency Act (Cal. Bus. & Prof. + Code § 22757 et seq.), each as amended from time to time and any + other applicable laws, regulations, or binding guidance relating + to artificial intelligence, transparency, content provenance, or + synthetic media, together with any documentation made available by + Licensor regarding compliance with the same (collectively, "AI + Regulations"). You shall maintain (including within any + application or service through which LTX-2.x, any Derivative of + LTX-2.x, or any Output is made available), and shall not remove, + disable, alter, or circumvent, any safety or security measures, + disclosures, metadata, watermarking, content provenance, latent + disclosure, or other transparency features or functionalities + included or embedded within LTX-2.x or any Derivative of LTX-2.x, + or applied to any Output, in furtherance of AI Regulations, + including any capability of LTX-2.x to include latent disclosures + in Outputs, and you shall include equivalent obligations in any + agreement governing your distribution of LTX-2.x or any Derivative + of LTX-2.x. You are solely responsible for any transparency, + disclosure, marking, or labeling obligations applicable to you + under AI Regulations as a provider or deployer of LTX-2.x, any + Derivative of LTX-2.x, or any system incorporating any of the + foregoing, including any obligation to disclose that content is + artificially generated or manipulated. If Licensor knows or + reasonably believes that you have modified LTX-2.x or any + Derivative of LTX-2.x such that it is no longer capable of + including any disclosure required by AI Regulations in Outputs, or + that you have otherwise removed, disabled, or circumvented any + feature or functionality described in this Section, Licensor may + in its sole discretion revoke the license granted under this + Agreement effective immediately upon notice to you, and upon such + revocation you shall immediately cease all use of LTX-2.x and + Derivatives of LTX-2.x. Licensor makes no representation or + warranty that LTX-2.x, any Derivative of LTX-2.x, or any Output + complies with any AI Regulations applicable to your specific use + case or deployment, and you are solely responsible for determining + the applicability of, and ensuring your compliance with, all AI + Regulations. You shall indemnify, defend, and hold harmless + Licensor and its affiliates from and against any and all claims, + liabilities, losses, damages, costs, and expenses (including + reasonable attorneys' fees) arising out of or relating to your + use, deployment, distribution, or modification of LTX-2.x, any + Derivative of LTX-2.x, or any Output in violation of, or your + other failure to comply with, any AI Regulations. + + For purposes of the EU AI Act, Licensor makes LTX-2.x openly + available under this community license and intends that LTX-2.x be + treated as a free and open-source general purpose AI model within + the meaning of Article 53(2) of the EU AI Act. You acknowledge and + agree that (a) to the extent the free and open source derogations + under Article 53(2) of the EU AI Act apply, Licensor's obligations + under the EU AI Act with respect to LTX-2.x are limited to those + applicable to providers of free and open source general purpose AI + models (it being acknowledged that such derogations do not extend + to the obligations under Article 53(1)(c) and (d)), (b) you + acknowledge that LTX-2.x is not intended to be integrated into a + high risk AI system, and shall be fully and solely responsible for + any obligation resulting from such integration, (c) if you + integrate LTX-2.x or any Derivative of LTX-2.x into a high-risk AI + system you shall be solely responsible for all provider + obligations that would otherwise apply to Licensor under the EU AI + Act, and (d) you shall not take any action, or omit to take any + action, that would cause Licensor to lose the benefit of the free + and open source derogations under the EU AI Act, and you shall + indemnify and hold Licensor harmless from any liability, costs, or + expenses arising from your breach of this Section. + + 7. Export Controls and Sanctions Compliance. You acknowledge that + LTX-2.x, Derivatives of LTX-2.x may be subject to export control + laws and regulations, including but not limited to the U.S. Export + Administration Regulations and sanctions programs administered by + the Office of Foreign Assets Control (OFAC). You represent and + warrant that you and any users of LTX-2.x are not (i) located in, + organized under the laws of, or ordinarily resident in any country + or territory subject to comprehensive sanctions; (ii) identified + on any U.S. government restricted party list, including the + Specially Designated Nationals and Blocked Persons List; or (iii) + otherwise prohibited from receiving LTX-2.x under applicable law. + You shall not export, re-export, or transfer LTX-2.x, directly or + indirectly, in violation of any applicable export control or + sanctions laws or regulations. You agree to comply with all + applicable trade control laws and shall indemnify and hold + Licensor harmless from any claims arising from your failure to + comply with such laws. + + 8. Trademarks; Reservation of Rights. Nothing in this Agreement + permits you to make use of Licensor's trademarks, trade names, + logos or to otherwise suggest endorsement or misrepresent the + relationship between the parties; and any rights not expressly + granted herein are reserved by the Licensor. Except as expressly + set forth in this Agreement, Licensor does not grant, directly or + by implication, estoppel, statute or otherwise, any right or + license in its, or its affiliates', intellectual property rights + or other proprietary rights. For avoidance of doubt, all + intellectual property rights in Derivatives of LTX-2.x shall be + subject to the terms of this Agreement, and you acquire no right, + title, or interest in or to LTX-2.x itself, which is and remains + the exclusive property of Licensor. You shall not assert any + ownership or other right in LTX-2.x or any Derivative of LTX-2.x + in any manner that restricts, encumbers, or is inconsistent with + the rights retained by Licensor or granted to other licensees + under this Agreement. + + 9. Disclaimer of Warranty. Unless required by applicable law or + agreed to in writing, Licensor provides LTX-2.x on an "AS IS" + BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either + express or implied, including, without limitation, any warranties + or conditions of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or + FITNESS FOR A PARTICULAR PURPOSE. You are solely responsible for + determining the appropriateness of using or redistributing LTX-2.x + and Derivatives of LTX-2.x and assume any risks associated with + your exercise of permissions under this Agreement. + + 10. Limitation of Liability. To the fullest extent permitted by + applicable law, in no event and under no legal theory, whether in + tort (including negligence), contract, or otherwise, unless + required by applicable law (such as deliberate and grossly + negligent acts) or agreed to in writing, shall Licensor be liable + to you or any other individual or Entity for damages, including + any direct, indirect, special, incidental, or consequential + damages of any character arising as a result of this Agreement or + out of the use of, or inability to use LTX-2.x or any Derivative + of LTX-2.x (including but not limited to damages for loss of + goodwill, work stoppage, computer failure or malfunction, or any + and all other commercial damages or losses), even if Licensor has + been advised of the possibility of such damages. + + 11. Accepting Warranty or Additional Liability. While redistributing + LTX-2.x and Derivatives of LTX-2.x, you may, provided you do not + violate the terms of this Agreement, choose to offer and charge a + fee for, acceptance of support, warranty, indemnity, or other + liability obligations. However, in accepting such obligations, + you may act only on your own behalf and on your sole + responsibility, not on behalf of Licensor, and only if you agree + to indemnify, defend, and hold Licensor harmless for any + liability incurred by, or claims asserted against Licensor, by + reason of your accepting any such warranty or additional + liability. + + 12. Governing Law. This Agreement and all relations, disputes, claims + and other matters arising hereunder (including non-contractual + disputes or claims) will be governed exclusively by, and + construed exclusively in accordance with, the laws of the State + of New York and applicable U.S. federal law. To the extent + permitted by law, choice of laws rules and the United Nations + Convention on Contracts for the International Sale of Goods will + not apply. The prevailing party in any claim or dispute between + the parties under this Agreement will be entitled to + reimbursement of its reasonable attorneys' fees and costs. + + 13. Term and Termination. This Agreement is effective upon your + acceptance and continues until terminated. Licensor may terminate + this Agreement immediately upon written notice to you if you + breach any provision of this Agreement, including but not limited + to violations of the use restrictions in Attachment A or + unauthorized commercial use. This Agreement also terminates + immediately and automatically, without notice, upon any material + breach of this Agreement, including any use in violation of + applicable AI Regulations or any unauthorized commercial use of + LTX-2.x or Derivatives of LTX-2.x by a Commercial Entity. Upon + termination: (a) all rights granted to you under this Agreement + will immediately cease; (b) you must immediately cease all use of + LTX-2.x and Derivatives of LTX-2.x; (c) you must delete or + destroy all copies of LTX-2.x and Derivatives of LTX-2.x in your + possession or control; and (d) you must notify any third parties + to whom you distributed LTX-2.x or Derivatives of LTX-2.x of the + termination. Sections 2, 3, 4, 6-16 and Attachment A shall + survive termination of this Agreement. Termination does not + relieve you of any obligations incurred prior to termination, + including payment obligations under Section 2 and adhering to the + restrictions under Section 3. In addition, if You commence a + lawsuit or other proceedings (including a cross-claim or + counterclaim in a lawsuit) against Licensor or any person or + entity alleging that LTX-2.x or any Output, or any portion of any + of the foregoing, infringe any intellectual property or other + right owned or licensable by you, then all licenses granted to + you under this Agreement shall terminate as of the date such + lawsuit or other proceeding is filed. + + 14. Disputes and Arbitration; Waiver of Jury Trial; Class Action + Waiver. IF YOU ARE NOT ACTING AS A CONSUMER UNDER APPLICABLE LAW, + YOU HEREBY WAIVE THE RIGHT TO A TRIAL BY JURY, TO PARTICIPATE IN + A CLASS OR REPRESENTATIVE ACTION (INCLUDING IN ARBITRATION), OR + TO COMBINE INDIVIDUAL PROCEEDINGS IN COURT OR IN ARBITRATION + WITHOUT THE CONSENT OF ALL PARTIES. All disputes arising in + connection with this Agreement shall be finally settled by + arbitration under the Rules of Arbitration of the International + Chamber of Commerce ("ICC Rules"), by one (1) arbitrator + appointed in accordance with the ICC Rules. The seat of + arbitration shall be New York, NY, USA, and the proceedings shall + be conducted in English. The arbitrator shall be empowered to + grant any relief that a court could grant. Judgment on the + arbitration award may be entered by any court having jurisdiction + thereof. Notwithstanding the foregoing, either party may seek + injunctive or other equitable relief in respect of any actual or + threatened breach of the license restrictions under this + Agreement (including Attachment A and the Acceptable Use Policy) + or any actual or threatened infringement, misappropriation, or + violation of Licensor's intellectual property rights, in the + state or federal courts located in the County of New York, State + of New York, and each party irrevocably consents to the + jurisdiction of, and venue in, such courts for that limited + purpose. The foregoing waivers do not apply to, and are not + enforceable against, any licensee acting as a consumer under the + mandatory consumer-protection laws of its jurisdiction of + residence (including, without limitation, the European Union, the + United Kingdom, and the State of California), and nothing in this + Agreement limits any rights under such laws that cannot be waived + or limited by contract. If any waiver in this Section is held + invalid or unenforceable as to a particular licensee or dispute, + such waiver shall be severed to that extent only and shall not + affect the validity or enforceability of the remainder of this + Section. + + 15. In the event of any exception to the application of binding + arbitration, all disputes, claims, and other matters arising + hereunder shall be brought exclusively in the state or federal + courts located in the County of New York, State of New York. You + waive all defenses of lack of personal jurisdiction and forum non + conveniens with respect to venue and jurisdiction in such courts, + and consent to their exclusive jurisdiction and venue. + + 16. Severability. If any provision of this Agreement is held to be + invalid, illegal or unenforceable, the remaining provisions shall + be unaffected thereby and remain valid as if such provision had + not been set forth herein. + + END OF TERMS AND CONDITIONS + + Attachment A + Use Restrictions + + When using the Outputs, LTX-2.x and any Derivatives thereof, you + agree to comply with the Acceptable Use Policy + (https://static.lightricks.com/legal/ltx-acceptable-use-policy.pdf) + which is hereby incorporated into and made part of this Agreement by + reference. Licensor may update it from time to time, and the version + in effect at the time of your use governs; continued use after an + update constitutes acceptance. Licensor shall post each version of + the Acceptable Use Policy with its effective date, and no update + shall apply retroactively to use occurring before that effective + date. In addition, you agree not to use the Outputs, LTX-2.x or its + Derivatives in any of the following ways: + + 1) In any way that violates any applicable national, federal, + state, local or international law or regulation; + + 2) For the purpose of exploiting, Harming or attempting to exploit + or Harm minors in any way; + + 3) Knowingly generate or disseminate verifiably false information + and/or content with the intent to deceive, defraud, or + otherwise unlawfully Harm others; + + 4) To generate or disseminate personal identifiable information + that can be used to Harm an individual; + + 5) To generate or disseminate information and/or content (e.g. + images, code, posts, articles), and place the information + and/or content in any context (e.g. bot generating tweets) + without expressly and intelligibly disclaiming that the + information and/or content is machine generated; + + 6) To defame others, or to engage in the unlawful harassment of + others; + + 7) To impersonate or attempt to impersonate (e.g. deepfakes) + others without their consent; + + 8) For fully automated decision making that adversely impacts an + individual's legal rights or otherwise creates or modifies a + binding, enforceable obligation; + + 9) For any use intended to or which has the effect of + discriminating against or Harming individuals or groups based + on online or offline social behavior or known or predicted + personal or personality characteristics; + + 10) To exploit any of the vulnerabilities of a specific group of + persons based on their age, social, physical or mental + characteristics, in order to materially distort the behavior + of a person pertaining to that group in a manner that causes + or is likely to cause that person or another person physical + or psychological Harm; + + 11) For any use intended to or which has the effect of + discriminating against individuals or groups based on legally + protected characteristics or categories; + + 12) To provide medical advice and medical results interpretation; + + 13) To generate or disseminate information for the purpose to be + used for administration of justice, law enforcement, + immigration or asylum processes, such as predicting an + individual will commit fraud/crime commitment (e.g. by text + profiling, drawing causal relationships between assertions + made in documents, indiscriminate and arbitrarily-targeted + use); + + 14) To generate and/or disseminate malware (including - but not + limited to - ransomware) or any other content to be used for + the purpose of harming electronic systems; + + 15) To engage in, promote, incite, or facilitate discrimination or + other unlawful or harmful conduct in the provision of + employment, employment benefits, credit, housing, or other + essential goods and services; + + 16) To engage in, promote, incite, or facilitate the harassment, + abuse, threatening, or bullying of individuals or groups of + individuals; + + 17) For military, warfare, nuclear industries or applications, + weapons development, or any use in connection with activities + that may cause death, personal injury, or severe physical or + environmental damage; + + 18) For commercial use only: To train, improve, or fine-tune any + other machine learning model, artificial intelligence system, + or competing model, except for Derivatives of LTX-2.x as + expressly permitted under this Agreement; + + 19) To circumvent, disable, or interfere with any technical + limitations, safety features, content filters, watermarking, + content provenance or latent disclosure functionalities, or + use restrictions implemented in LTX-2.x by Licensor; + + 20) To use LTX-2.x or Derivatives of LTX-2.x in any product, + service, or application that directly competes with Licensor's + commercial products or services, or is designed to replace or + substitute Licensor's offerings in the market, without + obtaining a separate commercial license from Licensor. diff --git a/assets/licenses/README.md b/assets/licenses/README.md index e37b61c5..ba1c454a 100644 --- a/assets/licenses/README.md +++ b/assets/licenses/README.md @@ -3,6 +3,8 @@ - `Apache-2.0.txt`: VALI, RF-DETR, OpenMMLab MMagic, and other Apache-2.0 components. - `BSD-3-Clause.txt`: PyAV. - `GPL-3.0.txt`: the GPLv3 FFmpeg libraries and command-line tools. +- `LTX-2.x-Community-License.txt`: Lightricks LTX-2.5 code ported into + `jasna/ltx/` and `jasna/models/ltx_vae/`, and the LTX restoration model weights. - `LGPL-2.1.txt`: libVLC and python-vlc under LGPL-2.1-or-later. - `LGPL-3.0.txt`: FFmpeg library code under LGPL-3.0-or-later. - `Python-3.13.txt`: CPython 3.13. diff --git a/jasna/benchmark/__init__.py b/jasna/benchmark/__init__.py index ac07a597..d8b95a87 100644 --- a/jasna/benchmark/__init__.py +++ b/jasna/benchmark/__init__.py @@ -6,6 +6,7 @@ import torch +from jasna.engine_paths import default_restoration_model_path from jasna.benchmark.basicvsrpp_restoration import benchmark_basicvsrpp_restoration from jasna.benchmark.detection_speed import benchmark_lada_yolo_detection_speed, benchmark_rfdetr_detection_speed from jasna.os_utils import check_required_executables, check_supported_gpu, gpu_check_error @@ -114,7 +115,7 @@ def run_benchmark_cli(args: Namespace) -> None: if args.detection_score_threshold is None else float(args.detection_score_threshold) ), - restoration_model_path=Path(args.restoration_model_path), + restoration_model_path=Path(args.restoration_model_path or default_restoration_model_path("basicvsrpp")), compile_basicvsrpp=bool(args.compile_basicvsrpp), benchmark_filter=getattr(args, "benchmark_filter", None), ) diff --git a/jasna/cli_help.py b/jasna/cli_help.py index ffed4901..876de5d9 100644 --- a/jasna/cli_help.py +++ b/jasna/cli_help.py @@ -7,6 +7,19 @@ """ CLI_HELP: dict[str, str] = { + "restoration_model_name": ( + "Restoration model for video input: basicvsrpp (fast) or ltx (diffusion model, much " + "slower, best quality; no streaming, segments, secondary restoration or denoise). " + "(default: %(default)s)" + ), + "restoration_model_path": ( + "Restoration model file (basicvsrpp) or folder (ltx). Default: the model in model_weights/." + ), + "ltx_seed": "ltx: noise seed. Another seed gives another take on the same restoration. (default: %(default)s)", + "ltx_large_canvas": ( + "ltx: restore large mosaics at 768 px instead of 512 px. Sharper, about 3x slower for " + "those mosaics and needs more VRAM. (default: %(default)s)" + ), "fp16": "Use FP16 where supported (restoration + TensorRT). Reduces VRAM usage and might improve performance.", "compile_basicvsrpp": "Compile BasicVSR++ for big performance boost (at cost of VRAM usage). Not recommended to use big clip sizes. (default: %(default)s)", "max_clip_size": "Maximum clip size for tracking (default: %(default)s)", diff --git a/jasna/engine_paths.py b/jasna/engine_paths.py index 8e39ed0e..8b8a90c8 100644 --- a/jasna/engine_paths.py +++ b/jasna/engine_paths.py @@ -20,6 +20,16 @@ def model_weights_dir() -> Path: return Path("model_weights") +_DEFAULT_RESTORATION_MODELS = { + "basicvsrpp": "lada_mosaic_restoration_model_generic_v1.2.pth", + "ltx": "ltx-restore", +} + + +def default_restoration_model_path(model_name: str) -> Path: + return model_weights_dir() / _DEFAULT_RESTORATION_MODELS[model_name] + + def engine_system_suffix() -> str: return ".win" if os.name == "nt" else ".linux" diff --git a/jasna/gui/engine_preflight.py b/jasna/gui/engine_preflight.py index 4fc94bb8..a91b5f7d 100644 --- a/jasna/gui/engine_preflight.py +++ b/jasna/gui/engine_preflight.py @@ -36,11 +36,11 @@ def run_engine_preflight(settings: AppSettings) -> EnginePreflightResult: from jasna.accelerator import is_amd_device from jasna.engine_paths import ( + default_restoration_model_path, expected_unet4x_engine_path, get_basicvsrpp_sub_engine_paths, get_onnx_tensorrt_engine_path, get_yolo_tensorrt_engine_path, - model_weights_dir, ) from jasna.mosaic.detection_registry import ( coerce_detection_model_name, @@ -93,7 +93,7 @@ def run_engine_preflight(settings: AppSettings) -> EnginePreflightResult: ) ) - restoration_model_path = model_weights_dir() / "lada_mosaic_restoration_model_generic_v1.2.pth" + restoration_model_path = default_restoration_model_path("basicvsrpp") if bool(settings.compile_basicvsrpp) and not amd: sub_paths = get_basicvsrpp_sub_engine_paths(str(restoration_model_path), bool(settings.fp16_mode)) all_engine_paths = tuple(Path(p) for p in sub_paths.values()) diff --git a/jasna/gui/video_session.py b/jasna/gui/video_session.py index a5fb4a76..546dc590 100644 --- a/jasna/gui/video_session.py +++ b/jasna/gui/video_session.py @@ -10,7 +10,7 @@ from typing import TYPE_CHECKING, Callable, Mapping from jasna.gui.models import AppSettings -from jasna.session_config import SessionConfig +from jasna.session_config import LTX_DEFAULT_SEED, SessionConfig from jasna.session_factory import RestorationSession if TYPE_CHECKING: @@ -54,7 +54,7 @@ def video_session_config( codec: str, encoder_settings: Mapping[str, object], ) -> SessionConfig: - from jasna.engine_paths import model_weights_dir + from jasna.engine_paths import default_restoration_model_path from jasna.mosaic.detection_registry import coerce_detection_model_name, require_detection_model_weights det_name = coerce_detection_model_name(str(settings.detection_model)) @@ -68,7 +68,10 @@ def video_session_config( max_detection_gap=int(settings.max_detection_gap), min_detection_duration=int(settings.min_detection_duration), scene_detection=bool(settings.scene_detection), - restoration_model_path=model_weights_dir() / "lada_mosaic_restoration_model_generic_v1.2.pth", + restoration_model_name="basicvsrpp", + restoration_model_path=default_restoration_model_path("basicvsrpp"), + ltx_large_canvas=True, + ltx_seed=LTX_DEFAULT_SEED, compile_basicvsrpp=bool(settings.compile_basicvsrpp), max_clip_size=int(settings.max_clip_size), temporal_overlap=int(settings.temporal_overlap), diff --git a/jasna/ltx/blocks.py b/jasna/ltx/blocks.py new file mode 100644 index 00000000..7d3997cb --- /dev/null +++ b/jasna/ltx/blocks.py @@ -0,0 +1,130 @@ +"""Transformer block weights, resident on the GPU or streamed from pinned host memory. + +Each block is packed into one flat byte buffer so it moves in a single copy. As many +leading blocks as the VRAM budget allows stay resident; the rest live in pinned host +memory and are copied on a side stream into one of two GPU slots, one block ahead of +the block being computed. +""" + +from __future__ import annotations + +import logging +from dataclasses import dataclass + +import torch + +logger = logging.getLogger(__name__) + +_ALIGN = 256 + + +@dataclass(frozen=True) +class _Entry: + offset: int + dtype: torch.dtype + shape: tuple[int, ...] + + @property + def nbytes(self) -> int: + count = 1 + for dim in self.shape: + count *= dim + return count * torch.empty((), dtype=self.dtype).element_size() + + +def _layout(tensors: dict[str, torch.Tensor]) -> tuple[dict[str, _Entry], int]: + layout: dict[str, _Entry] = {} + offset = 0 + for name, tensor in tensors.items(): + entry = _Entry(offset, tensor.dtype, tuple(tensor.shape)) + layout[name] = entry + offset += -(-entry.nbytes // _ALIGN) * _ALIGN + return layout, offset + + +def _pack(tensors: dict[str, torch.Tensor], layout: dict[str, _Entry], buffer: torch.Tensor) -> None: + for name, tensor in tensors.items(): + entry = layout[name] + buffer[entry.offset : entry.offset + entry.nbytes].copy_(tensor.contiguous().view(-1).view(torch.uint8)) + + +def _views(buffer: torch.Tensor, layout: dict[str, _Entry]) -> dict[str, torch.Tensor]: + return { + name: buffer[entry.offset : entry.offset + entry.nbytes].view(entry.dtype).view(entry.shape) + for name, entry in layout.items() + } + + +class BlockStore: + """Hands out each block's weights on ``device`` in a fixed cyclic block order.""" + + def __init__(self, blocks: list[dict[str, torch.Tensor]], device: torch.device, *, resident: int) -> None: + self.device = device + self.count = len(blocks) + self.resident = min(resident, self.count) + self._layouts: list[dict[str, _Entry]] = [] + self._resident: list[dict[str, torch.Tensor]] = [] + self._host: list[torch.Tensor] = [] + slot_bytes = 0 + for index, tensors in enumerate(blocks): + layout, nbytes = _layout(tensors) + self._layouts.append(layout) + if index < self.resident: + buffer = torch.empty(nbytes, dtype=torch.uint8, device=device) + _pack(tensors, layout, buffer) + self._resident.append(_views(buffer, layout)) + else: + buffer = torch.empty(nbytes, dtype=torch.uint8, pin_memory=True) + _pack(tensors, layout, buffer) + self._host.append(buffer) + slot_bytes = max(slot_bytes, nbytes) + streamed = self.count - self.resident + self._slots = [torch.empty(slot_bytes, dtype=torch.uint8, device=device) for _ in range(min(2, streamed))] + self._slot_block = [-1] * len(self._slots) + self._copy_done = [torch.cuda.Event() for _ in self._slots] + self._compute_done = [torch.cuda.Event() for _ in self._slots] + self._copy_stream = torch.cuda.Stream(device) if self._slots else None + logger.info("LTX blocks: %d resident, %d streamed", self.resident, streamed) + + def _slot_of(self, index: int) -> int | None: + return self._slot_block.index(index) if index in self._slot_block else None + + def _prefetch(self, index: int, keep: int | None) -> None: + if index < self.resident or self._slot_of(index) is not None: + return + busy = self._slot_of(keep) if keep is not None else None + slot = next(s for s in range(len(self._slots)) if s != busy) + source = self._host[index - self.resident] + with torch.cuda.stream(self._copy_stream): + self._copy_stream.wait_event(self._compute_done[slot]) + self._slots[slot][: source.numel()].copy_(source, non_blocking=True) + self._copy_done[slot].record(self._copy_stream) + self._slot_block[slot] = index + + def acquire(self, index: int) -> dict[str, torch.Tensor]: + """Weights of block ``index``; also starts copying the next streamed block.""" + if index < self.resident: + weights = self._resident[index] + else: + self._prefetch(index, None) + slot = self._slot_of(index) + torch.cuda.current_stream(self.device).wait_event(self._copy_done[slot]) + weights = _views(self._slots[slot], self._layouts[index]) + following = (index + 1) % self.count + if following < self.resident and self.resident < self.count: + following = self.resident + if following != index: + self._prefetch(following, index) + return weights + + def release(self, index: int) -> None: + """Mark block ``index``'s slot reusable once the work queued so far is done. + Call before acquiring the next block, so its prefetch waits for this work.""" + slot = self._slot_of(index) if index >= self.resident else None + if slot is not None: + self._compute_done[slot].record(torch.cuda.current_stream(self.device)) + + def close(self) -> None: + self._resident.clear() + self._host.clear() + self._slots.clear() diff --git a/jasna/ltx/camera.py b/jasna/ltx/camera.py new file mode 100644 index 00000000..4fed90ee --- /dev/null +++ b/jasna/ltx/camera.py @@ -0,0 +1,182 @@ +"""The LTX restoration crop camera: a smoothed follow crop with a slowly zooming side. + +A track gets one camera; its centre follows the Lada crop, Gaussian-smoothed, and is +clamped so the mosaic support always stays inside. The crop side follows the raw Lada +side through a rate-limited envelope (at most ``ZOOM_PER_WINDOW`` per window). +""" + +from __future__ import annotations + +import math +from collections.abc import Sequence +from dataclasses import dataclass + +import numpy as np + +Box = tuple[float, float, float, float] + +PLAN_CANVAS = 512 +SMOOTH_SIGMA_FRAMES = 12.0 +WINDOW_FRAMES = 121 +ZOOM_PER_WINDOW = 1.25 +ZOOM_LOG_STEP = math.log(ZOOM_PER_WINDOW) / (WINDOW_FRAMES - 1) +LADA_BORDER_RATIO = 0.06 +LADA_MIN_BORDER = 20 + + +@dataclass(frozen=True) +class CropFrame: + """One source-frame crop letterboxed into a square model canvas.""" + + crop_box: tuple[int, int, int, int] # xyxy half-open, source px + resize_hw: tuple[int, int] + pad: tuple[int, int, int, int] # top, bottom, left, right + + @property + def canvas(self) -> int: + return self.resize_hw[0] + self.pad[0] + self.pad[1] + + @property + def side(self) -> int: + x1, y1, x2, y2 = self.crop_box + return max(x2 - x1, y2 - y1) + + +def letterbox(crop_box: tuple[int, int, int, int], canvas: int) -> CropFrame: + x1, y1, x2, y2 = crop_box + width, height = x2 - x1, y2 - y1 + if width <= 0 or height <= 0: + raise ValueError(f"crop must be a positive rectangle: {crop_box}") + scale = min(canvas / width, canvas / height) + resize_w = min(canvas, max(1, int(round(width * scale)))) + resize_h = min(canvas, max(1, int(round(height * scale)))) + pad_left = (canvas - resize_w) // 2 + pad_top = (canvas - resize_h) // 2 + return CropFrame( + crop_box=crop_box, + resize_hw=(resize_h, resize_w), + pad=(pad_top, canvas - resize_h - pad_top, pad_left, canvas - resize_w - pad_left), + ) + + +def _lada_crop_box(box: Box, frame_w: int, frame_h: int) -> tuple[int, int, int, int]: + x1 = max(0, min(int(math.floor(box[0])), frame_w)) + y1 = max(0, min(int(math.floor(box[1])), frame_h)) + x2 = max(0, min(int(math.ceil(box[2])), frame_w)) + y2 = max(0, min(int(math.ceil(box[3])), frame_h)) + if x2 <= x1 or y2 <= y1: + raise ValueError(f"degenerate mosaic box {box}") + + border = max(LADA_MIN_BORDER, int(max(x2 - x1, y2 - y1) * LADA_BORDER_RATIO)) + x1, y1 = max(0, x1 - border), max(0, y1 - border) + x2, y2 = min(frame_w, x2 + border), min(frame_h, y2 + border) + + width, height = x2 - x1, y2 - y1 + down_scale = min(PLAN_CANVAS / width, PLAN_CANVAS / height, 1.0) + missing_w = int((PLAN_CANVAS - width * down_scale) / down_scale) + missing_h = int((PLAN_CANVAS - height * down_scale) / down_scale) + + def expand(available_low: int, available_high: int, missing: int, budget: int) -> tuple[int, int, int]: + both = min(available_low, available_high, missing // 2, budget) + low = min(available_low - both, missing - 2 * both, budget - both) + high = min(available_high - both, missing - 2 * both - low, budget - both - low) + return both, low, high + + lr, left, right = expand(x1, frame_w - x2, missing_w, width) + tb, top, bottom = expand(y1, frame_h - y2, missing_h, height) + return ( + x1 - (math.floor(lr / 2) + left), + y1 - (math.floor(tb / 2) + top), + x2 + math.ceil(lr / 2) + right, + y2 + math.ceil(tb / 2) + bottom, + ) + + +def lada_crop_boxes(boxes: Sequence[Box], frame_w: int, frame_h: int) -> np.ndarray: + return np.asarray([_lada_crop_box(box, frame_w, frame_h) for box in boxes], dtype=np.int64) + + +def gaussian_smooth(values: np.ndarray, sigma: float) -> np.ndarray: + """Zero-phase Gaussian smoothing with edge replication.""" + if sigma <= 0 or len(values) < 2: + return values.astype(np.float64, copy=True) + radius = max(1, int(math.ceil(3.0 * sigma))) + offsets = np.arange(-radius, radius + 1, dtype=np.float64) + kernel = np.exp(-0.5 * np.square(offsets / sigma)) + kernel /= kernel.sum() + padded = np.pad(values.astype(np.float64), (radius, radius), mode="edge") + return np.convolve(padded, kernel, mode="valid") + + +def zoom_crop_sides(raw_sides: np.ndarray, max_log_step: float = ZOOM_LOG_STEP) -> np.ndarray: + """Smooth per-frame sides, never below any frame's raw side, zooming at most + ``exp(max_log_step)`` per frame (log-space running max, rate limit, Gaussian).""" + values = np.log(np.asarray(raw_sides, dtype=np.float64)) + radius = max(1, int(math.ceil(3.0 * SMOOTH_SIGMA_FRAMES))) + padded = np.pad(values, (radius, radius), mode="edge") + envelope = np.lib.stride_tricks.sliding_window_view(padded, 2 * radius + 1).max(axis=1) + for index in range(1, len(envelope)): + envelope[index] = max(envelope[index], envelope[index - 1] - max_log_step) + for index in range(len(envelope) - 2, -1, -1): + envelope[index] = max(envelope[index], envelope[index + 1] - max_log_step) + return np.ceil(np.exp(gaussian_smooth(envelope, SMOOTH_SIGMA_FRAMES)) - 1e-6).astype(np.int64) + + +def _feasible_interval(lower: float, upper: float, side: int, extent: int, margin: float) -> tuple[int, int]: + minimum = max(0, min(extent, int(math.ceil(upper + margin))) - side) + maximum = min(extent - side, max(0, int(math.floor(lower - margin)))) + if minimum <= maximum: + return minimum, maximum + minimum = max(0, int(math.ceil(upper)) - side) + maximum = min(extent - side, int(math.floor(lower))) + if minimum > maximum: + raise ValueError("mosaic support has no valid crop placement") + return minimum, maximum + + +def follow_camera( + boxes: Sequence[Box], sides: np.ndarray, *, frame_w: int, frame_h: int, margin: int +) -> list[CropFrame]: + """Place a crop of ``sides[i]`` on every frame, centred on the smoothed Lada crop and + clamped so the mosaic box (plus ``margin`` where it fits) stays inside.""" + raw = lada_crop_boxes(boxes, frame_w, frame_h) + support = np.asarray(boxes, dtype=np.float64) + widths = np.minimum(sides, frame_w).astype(int) + heights = np.minimum(sides, frame_h).astype(int) + if ((support[:, 2] - support[:, 0]) > widths).any() or ((support[:, 3] - support[:, 1]) > heights).any(): + raise ValueError("mosaic support does not fit inside the crop") + target_x = gaussian_smooth((raw[:, 0] + raw[:, 2]) / 2.0, SMOOTH_SIGMA_FRAMES) + target_y = gaussian_smooth((raw[:, 1] + raw[:, 3]) / 2.0, SMOOTH_SIGMA_FRAMES) + frames = [] + for index, (x1, y1, x2, y2) in enumerate(support): + width, height = int(widths[index]), int(heights[index]) + min_left, max_left = _feasible_interval(x1, x2, width, frame_w, margin) + min_top, max_top = _feasible_interval(y1, y2, height, frame_h, margin) + left = min(max(int(round(target_x[index] - width / 2)), min_left), max_left) + top = min(max(int(round(target_y[index] - height / 2)), min_top), max_top) + frames.append(letterbox((left, top, left + width, top + height), PLAN_CANVAS)) + return frames + + +def _edge_gap(crop: tuple[int, int, int, int], box: Box, frame_w: int, frame_h: int) -> float: + """Distance from the mosaic box to the nearest crop edge that is not the frame border.""" + x1, y1, x2, y2 = crop + return float( + min( + box[0] - x1 if x1 > 0 else np.inf, + box[1] - y1 if y1 > 0 else np.inf, + x2 - box[2] if x2 < frame_w else np.inf, + y2 - box[3] if y2 < frame_h else np.inf, + ) + ) + + +def track_camera(boxes: Sequence[Box], *, frame_w: int, frame_h: int, feather: int) -> list[CropFrame]: + """The zooming track camera, widened by a feather per side (with a feather margin) + when the tight one brings the mask within a feather of an inner crop edge.""" + raw = lada_crop_boxes(boxes, frame_w, frame_h) + sides = zoom_crop_sides(np.maximum(raw[:, 2] - raw[:, 0], raw[:, 3] - raw[:, 1])) + tight = follow_camera(boxes, sides, frame_w=frame_w, frame_h=frame_h, margin=0) + if all(_edge_gap(crop.crop_box, box, frame_w, frame_h) >= feather for crop, box in zip(tight, boxes)): + return tight + return follow_camera(boxes, sides + 2 * feather, frame_w=frame_w, frame_h=frame_h, margin=feather) diff --git a/jasna/ltx/compose.py b/jasna/ltx/compose.py new file mode 100644 index 00000000..ee7f8e9d --- /dev/null +++ b/jasna/ltx/compose.py @@ -0,0 +1,179 @@ +"""Crops into the model canvas and the composite of restored windows back into frames. + +Frames are ``[3, H, W]`` uint8 RGB tensors. A window paints only its own track's +feathered mask and never a pixel inside another track's mask; outside every mask the +source stays bit-exact. +""" + +from __future__ import annotations + +import functools +from collections.abc import Sequence +from dataclasses import dataclass + +import cv2 +import numpy as np +import torch +import torch.nn.functional as F + +from jasna.ltx.camera import CropFrame +from jasna.ltx.plan import Polygon + +COLOUR_MATCH_STRENGTH = 0.75 + + +def _reflect101_index(size: int, before: int, after: int, device: torch.device) -> torch.Tensor: + index = torch.arange(-before, size + after, device=device) + period = 2 * (size - 1) if size > 1 else 1 + index = index.remainder(period) + return torch.where(index >= size, period - index, index) + + +def crop_to_canvas(frame: torch.Tensor, crop: CropFrame) -> torch.Tensor: + """Crop, bilinear-resize and reflect-101 pad one frame onto the model canvas.""" + x1, y1, x2, y2 = crop.crop_box + content = frame[:, y1:y2, x1:x2] + if content.shape[1:] != crop.resize_hw: + resized = F.interpolate(content[None].float(), size=crop.resize_hw, mode="bilinear", align_corners=False) + content = resized[0].round_().clamp_(0, 255).to(torch.uint8) + top, bottom, left, right = crop.pad + if top or bottom or left or right: + rows = _reflect101_index(content.shape[1], top, bottom, frame.device) + cols = _reflect101_index(content.shape[2], left, right, frame.device) + content = content[:, rows][:, :, cols] + return content + + +def _area_weights(out_size: int, in_size: int, device: torch.device) -> torch.Tensor: + """``[out, in]`` pixel-area resampling weights (OpenCV INTER_AREA for downscaling).""" + scale = in_size / out_size + starts = torch.arange(out_size, dtype=torch.float64) * scale + edges = torch.arange(in_size + 1, dtype=torch.float64) + lo = torch.maximum(starts[:, None], edges[None, :-1]) + hi = torch.minimum(starts[:, None] + scale, edges[None, 1:]) + return ((hi - lo).clamp(min=0) / scale).to(device=device, dtype=torch.float32) + + +def canvas_to_source(canvas: torch.Tensor, crop: CropFrame) -> torch.Tensor: + """The restored canvas mapped back to its source crop box, float ``[3, h, w]`` in 0..255 + rounded to whole grey levels (bicubic when enlarging, pixel area when shrinking).""" + top, _, left, _ = crop.pad + resize_h, resize_w = crop.resize_hw + content = canvas[:, top : top + resize_h, left : left + resize_w].float() + x1, y1, x2, y2 = crop.crop_box + height, width = y2 - y1, x2 - x1 + if max(height, width) > max(resize_h, resize_w): + out = F.interpolate(content[None], size=(height, width), mode="bicubic", align_corners=False)[0] + else: + rows = _area_weights(height, resize_h, canvas.device) + cols = _area_weights(width, resize_w, canvas.device) + out = torch.einsum("hy,cyx,wx->chw", rows, content, cols) + return out.round_().clamp_(0, 255) + + +@functools.cache +def _ellipse(radius: int, device: torch.device) -> torch.Tensor: + kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (2 * radius + 1,) * 2) + return torch.from_numpy(kernel.astype(np.float32)).to(device)[None, None] + + +@functools.cache +def _gaussian(radius: int, device: torch.device) -> torch.Tensor: + kernel = cv2.getGaussianKernel(2 * radius + 1, 0.0, cv2.CV_32F)[:, 0] + return torch.from_numpy(kernel).to(device) + + +def feather_alpha(mask: torch.Tensor, feather: int) -> torch.Tensor: + """Dilate a bool ``[H, W]`` mask by an ellipse of ``feather`` px and Gaussian-blur it + (kernel ``2*feather+1``, reflect-101 border) into a [0, 1] alpha.""" + dilated = F.conv2d(F.pad(mask[None, None].float(), (feather,) * 4), _ellipse(feather, mask.device)) > 0 + kernel = _gaussian(feather, mask.device) + x = F.pad(dilated.float(), (feather, feather, 0, 0), mode="reflect") + x = F.conv2d(x, kernel.view(1, 1, 1, -1)) + x = F.pad(x, (0, 0, feather, feather), mode="reflect") + return F.conv2d(x, kernel.view(1, 1, -1, 1))[0, 0] + + +def crop_edge_ramp(box: tuple[int, int, int, int], shape: tuple[int, int], feather: int, device: torch.device) -> torch.Tensor: + """1 inside the crop, falling to 0 over ``feather`` px at every crop edge that is not the + region border, 0 outside.""" + height, width = shape + x1, y1, x2, y2 = box + + def axis(low: int, high: int, size: int) -> torch.Tensor: + centre = torch.arange(size, dtype=torch.float32, device=device) + 0.5 + ramp = ((centre > low) & (centre < high)).float() + if low > 0: + ramp = torch.minimum(ramp, (centre - low) / max(feather, 1)) + if high < size: + ramp = torch.minimum(ramp, (high - centre) / max(feather, 1)) + return ramp.clamp(0.0, 1.0) + + return torch.minimum(axis(y1, y2, height)[:, None], axis(x1, x2, width)[None, :]) + + +def polygon_mask(polygons: Sequence[Polygon], shape: tuple[int, int], offset: tuple[int, int], device: torch.device) -> torch.Tensor: + """Bool ``[h, w]`` fill of ``polygons`` inside the region at ``offset`` (x, y).""" + mask = np.zeros(shape, dtype=np.uint8) + contours = [ + (np.asarray(polygon, dtype=np.float32).round().astype(np.int32) - np.asarray(offset, dtype=np.int32)) + for polygon in polygons + if len(polygon) >= 3 + ] + if contours: + cv2.fillPoly(mask, contours, 255) + return torch.from_numpy(mask).to(device) > 0 + + +@dataclass(frozen=True) +class Candidate: + """One window's restored canvas for a frame, with its crop, crossfade weight and the + polygons of the track it restores.""" + + canvas: torch.Tensor + crop: CropFrame + weight: float + polygons: Sequence[Polygon] + + +def composite_frame(source: torch.Tensor, candidates: Sequence[Candidate], *, feather: int) -> torch.Tensor: + live = [c for c in candidates if c.weight > 0.0] + if not live: + return source + _, height, width = source.shape + reach = 2 * feather + 1 + x0 = max(min(c.crop.crop_box[0] for c in live) - reach, 0) + y0 = max(min(c.crop.crop_box[1] for c in live) - reach, 0) + x3 = min(max(c.crop.crop_box[2] for c in live) + reach, width) + y3 = min(max(c.crop.crop_box[3] for c in live) + reach, height) + shape = (y3 - y0, x3 - x0) + device = source.device + rois = [polygon_mask(c.polygons, shape, (x0, y0), device) for c in live] + union = torch.stack(rois).any(dim=0) + outside = feather_alpha(union, feather) <= 0.0 + region = source[:, y0:y3, x0:x3].float() + accumulator = torch.zeros_like(region) + total = torch.zeros(shape, dtype=torch.float32, device=device) + for candidate, roi in zip(live, rois): + x1, y1, x2, y2 = candidate.crop.crop_box + x1, y1, x2, y2 = x1 - x0, y1 - y0, x2 - x0, y2 - y0 + ramp = crop_edge_ramp((x1, y1, x2, y2), shape, feather, device) + covered = ramp > 0.0 + restored = region.clone() + restored[:, y1:y2, x1:x2] = canvas_to_source(candidate.canvas, candidate.crop) + edge = torch.where(roi, covered.float(), ramp) + alpha = feather_alpha(roi, feather) * edge * candidate.weight + alpha = alpha * (roi | ~union).float() + ring = covered & outside + if ring.any(): + offset = torch.quantile((region - restored)[:, ring], 0.5, dim=1) * COLOUR_MATCH_STRENGTH + restored = restored + offset[:, None, None] + accumulator += restored * alpha + total += alpha + over = total > 1.0 + accumulator = torch.where(over, accumulator / total, accumulator) + total = torch.where(over, torch.ones_like(total), total) + blended = accumulator + region * (1.0 - total) + out = source.clone() + out[:, y0:y3, x0:x3] = blended.round_().clamp_(0, 255).to(torch.uint8) + return out diff --git a/jasna/ltx/plan.py b/jasna/ltx/plan.py new file mode 100644 index 00000000..023caf1c --- /dev/null +++ b/jasna/ltx/plan.py @@ -0,0 +1,333 @@ +"""Mosaic tracks and their 121-frame model windows for LTX restoration. + +Detections are grouped per frame, linked into cut-safe tracks, continuations are folded +into their track, and each track gets one zooming camera sliced into windows at hop 88. +Windows of a track are denoised together and share their overlapping latent frames. +""" + +from __future__ import annotations + +import bisect +from collections.abc import Sequence +from dataclasses import dataclass, field + +import numpy as np + +from jasna.ltx.camera import PLAN_CANVAS, WINDOW_FRAMES, Box, CropFrame, letterbox, track_camera +from jasna.tracking.clip_tracker import compute_iou_matrix + +Polygon = list[list[float]] + +WINDOW_HOP = 88 +TEMPORAL_SCALE = 8 +GROUP_IOU = 0.05 +LINK_IOU = 0.1 +MAX_GAP = 3 +CONTINUATION_COVER = 0.5 +LARGE_CANVAS = 768 +LARGE_CANVAS_MIN_SIDE = 640 +FEATHER_FRACTION = 0.028 + + +@dataclass(frozen=True) +class Region: + """One detected mosaic: model box and the traced outline of its dilated mask.""" + + box: Box + polygon: Polygon + + +@dataclass(frozen=True) +class RegionGroup: + box: Box + polygons: tuple[Polygon, ...] + + +@dataclass +class Track: + track_id: int + frames: dict[int, RegionGroup] = field(default_factory=dict) + + @property + def first(self) -> int: + return min(self.frames) + + @property + def last(self) -> int: + return max(self.frames) + + +@dataclass(frozen=True) +class Window: + index: int # global across tracks; also the noise seed offset + track_id: int + start: int + real_frames: int # the model always sees WINDOW_FRAMES; the tail is padded + crops: tuple[CropFrame, ...] # one per real frame + + @property + def stop(self) -> int: + return self.start + self.real_frames + + @property + def canvas(self) -> int: + return self.crops[0].canvas + + def model_crops(self) -> list[CropFrame]: + return list(self.crops) + [self.crops[-1]] * (WINDOW_FRAMES - self.real_frames) + + +@dataclass(frozen=True) +class TrackPlan: + track_id: int + start: int + polygons: tuple[tuple[Polygon, ...], ...] # per frame from ``start`` + windows: tuple[Window, ...] + + @property + def stop(self) -> int: + return self.start + len(self.polygons) + + +def feather_pixels(frame_h: int) -> int: + return max(1, int(round(FEATHER_FRACTION * frame_h))) + + +def _union_box(boxes: Sequence[Box]) -> Box: + return ( + min(box[0] for box in boxes), + min(box[1] for box in boxes), + max(box[2] for box in boxes), + max(box[3] for box in boxes), + ) + + +def group_regions(regions: Sequence[Region]) -> list[RegionGroup]: + """Union-find merge of regions whose boxes overlap by IoU >= ``GROUP_IOU``.""" + count = len(regions) + if count == 0: + return [] + parent = list(range(count)) + + def find(item: int) -> int: + while parent[item] != item: + parent[item] = parent[parent[item]] + item = parent[item] + return item + + boxes = np.asarray([r.box for r in regions], dtype=np.float64) + iou = compute_iou_matrix(boxes, boxes) + for i in range(count): + for j in range(i + 1, count): + if iou[i, j] >= GROUP_IOU: + root_i, root_j = find(i), find(j) + if root_i != root_j: + parent[root_i] = root_j + buckets: dict[int, list[Region]] = {} + for i in range(count): + buckets.setdefault(find(i), []).append(regions[i]) + return [ + RegionGroup(box=_union_box([r.box for r in members]), polygons=tuple(r.polygon for r in members)) + for members in buckets.values() + ] + + +def build_tracks(groups_per_frame: Sequence[Sequence[RegionGroup]], cuts: set[int]) -> list[Track]: + """Greedy best-IoU linking; a track tolerates ``MAX_GAP`` missing frames and never + crosses a scene cut (``cuts`` holds first-frame-after-cut indices).""" + tracks: list[Track] = [] + active: list[Track] = [] + for frame_idx, groups in enumerate(groups_per_frame): + active = [ + track + for track in active + if frame_idx - track.last - 1 <= MAX_GAP + and not any(track.last < cut <= frame_idx for cut in cuts) + ] + candidates: list[tuple[float, int, int]] = [] + if active and groups: + iou = compute_iou_matrix( + np.asarray([track.frames[track.last].box for track in active], dtype=np.float64), + np.asarray([group.box for group in groups], dtype=np.float64), + ) + candidates = [ + (float(iou[ti, gi]), ti, gi) + for ti in range(len(active)) + for gi in range(len(groups)) + if iou[ti, gi] >= LINK_IOU + ] + candidates.sort(reverse=True) + used_tracks: set[int] = set() + used_groups: set[int] = set() + for _, ti, gi in candidates: + if ti in used_tracks or gi in used_groups: + continue + active[ti].frames[frame_idx] = groups[gi] + used_tracks.add(ti) + used_groups.add(gi) + for gi, group in enumerate(groups): + if gi not in used_groups: + track = Track(track_id=len(tracks), frames={frame_idx: group}) + tracks.append(track) + active.append(track) + return tracks + + +def _cover(outer: Box, inner: Box) -> float: + """Intersection of two boxes over the area of ``inner``.""" + width = min(outer[2], inner[2]) - max(outer[0], inner[0]) + height = min(outer[3], inner[3]) - max(outer[1], inner[1]) + area = (inner[2] - inner[0]) * (inner[3] - inner[1]) + return max(0.0, width) * max(0.0, height) / area if area > 0 else 0.0 + + +def merge_continuations(tracks: Sequence[Track], cuts: set[int]) -> list[Track]: + """Fold a track whose first box lies >= ``CONTINUATION_COVER`` inside an earlier + track's box from <= ``MAX_GAP`` frames before (no cut between) into that track: + a mosaic the detector split in two, or a link lost for a few frames.""" + merged: list[Track] = [] + for track in sorted(tracks, key=lambda item: (item.first, item.track_id)): + head_idx = track.first + head = track.frames[head_idx] + parent, best = None, CONTINUATION_COVER + for candidate in merged: + before = [idx for idx in candidate.frames if idx < head_idx] + if not before: + continue + last = max(before) + if head_idx - last - 1 > MAX_GAP or any(last < cut <= head_idx for cut in cuts): + continue + cover = _cover(candidate.frames[last].box, head.box) + if cover >= best: + parent, best = candidate, cover + if parent is None: + merged.append(Track(track_id=track.track_id, frames=dict(track.frames))) + continue + for idx, group in track.frames.items(): + other = parent.frames.get(idx) + parent.frames[idx] = ( + group + if other is None + else RegionGroup(_union_box([other.box, group.box]), other.polygons + group.polygons) + ) + parent.frames = dict(sorted(parent.frames.items())) + return merged + + +def _affine_polygons(polygons: Sequence[Polygon], source: Box, target: Box) -> tuple[Polygon, ...]: + scale_x = (target[2] - target[0]) / max(source[2] - source[0], 1e-6) + scale_y = (target[3] - target[1]) / max(source[3] - source[1], 1e-6) + return tuple( + [[target[0] + (x - source[0]) * scale_x, target[1] + (y - source[1]) * scale_y] for x, y in polygon] + for polygon in polygons + ) + + +def track_polygons(track: Track) -> list[tuple[Polygon, ...]]: + """Per-frame polygons over the whole track span; a detection gap gets the box + interpolated between its neighbours and the nearer one's polygons mapped onto it.""" + anchors = sorted(track.frames) + out = [] + for frame_idx in range(anchors[0], anchors[-1] + 1): + group = track.frames.get(frame_idx) + if group is not None: + out.append(group.polygons) + continue + position = bisect.bisect_left(anchors, frame_idx) + left, right = track.frames[anchors[position - 1]], track.frames[anchors[position]] + alpha = (frame_idx - anchors[position - 1]) / (anchors[position] - anchors[position - 1]) + box = tuple(a + (b - a) * alpha for a, b in zip(left.box, right.box)) + near = left if alpha <= 0.5 else right + out.append(_affine_polygons(near.polygons, near.box, box)) + return out + + +def polygons_box(polygons: Sequence[Polygon], frame_w: int, frame_h: int) -> Box: + points = np.asarray([point for polygon in polygons for point in polygon], dtype=np.float64) + if not len(points): + raise ValueError("a tracked frame carries no polygon") + x1, y1 = points.min(axis=0) + x2, y2 = points.max(axis=0) + return ( + float(max(0.0, min(x1, frame_w - 1.0))), + float(max(0.0, min(y1, frame_h - 1.0))), + float(min(float(frame_w), max(x2, x1 + 1.0))), + float(min(float(frame_h), max(y2, y1 + 1.0))), + ) + + +def window_starts(frames: int) -> list[int]: + starts = [0] + while starts[-1] + WINDOW_FRAMES < frames: + starts.append(starts[-1] + WINDOW_HOP) + return starts + + +def track_canvas(camera: Sequence[CropFrame], *, large_canvas: bool) -> int: + """768 for a track whose largest crop is over 640 px (when allowed), else 512.""" + if large_canvas and max(crop.side for crop in camera) > LARGE_CANVAS_MIN_SIDE: + return LARGE_CANVAS + return PLAN_CANVAS + + +def plan_track(track: Track, *, first_index: int, frame_w: int, frame_h: int, large_canvas: bool) -> TrackPlan: + polygons = track_polygons(track) + boxes = [polygons_box(item, frame_w, frame_h) for item in polygons] + camera = track_camera(boxes, frame_w=frame_w, frame_h=frame_h, feather=feather_pixels(frame_h)) + canvas = track_canvas(camera, large_canvas=large_canvas) + if canvas != PLAN_CANVAS: + camera = [letterbox(crop.crop_box, canvas) for crop in camera] + windows = tuple( + Window( + index=first_index + offset, + track_id=track.track_id, + start=track.first + local, + real_frames=min(WINDOW_FRAMES, len(polygons) - local), + crops=tuple(camera[local : local + WINDOW_FRAMES]), + ) + for offset, local in enumerate(window_starts(len(polygons))) + ) + return TrackPlan(track_id=track.track_id, start=track.first, polygons=tuple(polygons), windows=windows) + + +def plan_video( + regions_per_frame: Sequence[Sequence[Region]], + cuts: set[int], + *, + frame_w: int, + frame_h: int, + large_canvas: bool, +) -> list[TrackPlan]: + """Every mosaic track of a video with its windows, in window-index order.""" + groups = [group_regions(regions) for regions in regions_per_frame] + tracks = merge_continuations(build_tracks(groups, cuts), cuts) + plans = [] + for track in tracks: + plan = plan_track( + track, + first_index=sum(len(p.windows) for p in plans), + frame_w=frame_w, + frame_h=frame_h, + large_canvas=large_canvas, + ) + plans.append(plan) + return plans + + +def crossfade_weights(windows: Sequence[Window]) -> list[np.ndarray]: + """Per-frame compose weights of one track's windows: linear ramps across each overlap.""" + weights = [np.ones(window.real_frames, dtype=np.float32) for window in windows] + for index, (left, right) in enumerate(zip(windows, windows[1:])): + overlap = left.stop - right.start + ramp = np.linspace(0.0, 1.0, overlap, dtype=np.float32) + weights[index][len(weights[index]) - overlap :] = 1.0 - ramp + weights[index + 1][:overlap] = ramp + return weights + + +LATENT_FRAMES = 1 + (WINDOW_FRAMES - 1) // TEMPORAL_SCALE +SHARED_LATENT_SHIFT = WINDOW_HOP // TEMPORAL_SCALE +# Window k's latents 12..15 encode the same pixel frames as window k+1's latents 1..4. +LEFT_SHARED_LATENTS = range(SHARED_LATENT_SHIFT + 1, LATENT_FRAMES) +RIGHT_SHARED_LATENTS = range(1, LATENT_FRAMES - SHARED_LATENT_SHIFT) +FUSION_RAMP = tuple((i + 1) / (len(LEFT_SHARED_LATENTS) + 1) for i in range(len(LEFT_SHARED_LATENTS))) diff --git a/jasna/ltx/regions.py b/jasna/ltx/regions.py new file mode 100644 index 00000000..a204317b --- /dev/null +++ b/jasna/ltx/regions.py @@ -0,0 +1,51 @@ +"""Mosaic detections as LTX planning regions: the model box plus the outline of its mask. + +Each low-res mask is bilinearly upsampled to the frame and thresholded at 0.5, dilated +by an ellipse of ``EXPAND_PIXELS``, and its largest external contour becomes the region +polygon. The box is the model box truncated to whole pixels. +""" + +from __future__ import annotations + +import cv2 +import numpy as np +import torch +import torch.nn.functional as F + +from jasna.ltx.plan import Region +from jasna.mosaic.detections import Detections + +EXPAND_PIXELS = 20 +_KERNEL = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (2 * EXPAND_PIXELS + 1,) * 2) + + +def trace_region(box: np.ndarray, mask: torch.Tensor, frame_h: int, frame_w: int) -> Region | None: + full = F.interpolate(mask[None, None].float(), size=(frame_h, frame_w), mode="bilinear", align_corners=False)[0, 0] > 0.5 + rows = torch.nonzero(full.any(dim=1)).flatten() + if not len(rows): + return None + cols = torch.nonzero(full.any(dim=0)).flatten() + y0 = max(0, int(rows[0]) - EXPAND_PIXELS - 1) + y1 = min(frame_h, int(rows[-1]) + EXPAND_PIXELS + 2) + x0 = max(0, int(cols[0]) - EXPAND_PIXELS - 1) + x1 = min(frame_w, int(cols[-1]) + EXPAND_PIXELS + 2) + crop = full[y0:y1, x0:x1].to(torch.uint8).mul_(255).cpu().numpy() + contours, _ = cv2.findContours(cv2.dilate(crop, _KERNEL), cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE) + if not contours: + return None + largest = max(contours, key=cv2.contourArea) + if len(largest) < 3: + return None + x_min, y_min, x_max, y_max = (float(int(v)) for v in box.tolist()) + return Region( + box=(x_min, y_min, x_max, y_max), + polygon=[[float(x + x0), float(y + y0)] for x, y in largest[:, 0, :]], + ) + + +def regions_per_frame(detections: Detections, frame_h: int, frame_w: int) -> list[list[Region]]: + out = [] + for boxes, masks in zip(detections.boxes_xyxy, detections.masks): + traced = (trace_region(box, mask, frame_h, frame_w) for box, mask in zip(boxes, masks)) + out.append([region for region in traced if region is not None]) + return out diff --git a/jasna/ltx/restore.py b/jasna/ltx/restore.py new file mode 100644 index 00000000..383ac066 --- /dev/null +++ b/jasna/ltx/restore.py @@ -0,0 +1,301 @@ +"""Restore a whole video with the LTX model in four passes, one model on the GPU at a time. + +1. scan: decode, detect mosaics and scene cuts, plan tracks, cameras and windows; +2. encode: crop every window's 121 frames onto its canvas and VAE-encode them; +3. denoise: run each track's windows through the transformer in lockstep; +4. compose: decode each window's latent when the video reaches it, composite, write. + +The camera and fusion need a whole track before its first window can be denoised, so +the video is decoded three times (scan, encode, compose). Per-window latents (1-2 MB) +are kept in a temporary directory between passes, so memory does not grow with length. +""" + +from __future__ import annotations + +import logging +import tempfile +import threading +from collections.abc import Callable, Iterator +from dataclasses import dataclass +from pathlib import Path + +import torch +from tqdm import tqdm + +from jasna.ltx.camera import WINDOW_FRAMES +from jasna.ltx.compose import Candidate, composite_frame, crop_to_canvas +from jasna.ltx.plan import Region, TrackPlan, Window, crossfade_weights, feather_pixels, plan_video +from jasna.ltx.regions import regions_per_frame +from jasna.ltx.sampler import LtxTransformer +from jasna.models.ltx_vae import decode_latent, load_video_decoder, load_video_encoder +from jasna.tracking.scene_detector import SceneCutDetector + +logger = logging.getLogger(__name__) + +FrameBatches = Iterator[tuple[torch.Tensor, list[int]]] +FrameSource = Callable[[], FrameBatches] + + +@dataclass(frozen=True) +class LtxModelFiles: + transformer: Path + vae: Path + tuned_decoder: Path + + @classmethod + def from_dir(cls, directory: Path) -> LtxModelFiles: + files = cls( + transformer=directory / "transformer.safetensors", + vae=directory / "vae.safetensors", + tuned_decoder=directory / "vae-decoder.safetensors", + ) + missing = [str(path) for path in (files.transformer, files.vae, files.tuned_decoder) if not path.is_file()] + if missing: + raise FileNotFoundError(f"LTX model files missing: {', '.join(missing)}") + return files + + +class Cancelled(Exception): + pass + + +@dataclass(frozen=True) +class Progress: + """One console bar per pass; ``frames`` sizes the frame-based passes.""" + + frames: int + disable: bool + + def bar(self, name: str, total: int, unit: str = "frame") -> tqdm: + return tqdm(total=total, desc=f"LTX {name}", unit=unit, dynamic_ncols=True, disable=self.disable) + + +class _LatentStore: + def __init__(self, directory: Path, kind: str) -> None: + self._directory = directory + self._kind = kind + + def _path(self, index: int) -> Path: + return self._directory / f"{index:06d}.{self._kind}.pt" + + def put(self, index: int, latent: torch.Tensor) -> None: + torch.save(latent.to("cpu", torch.bfloat16).contiguous(), self._path(index)) + + def get(self, index: int) -> torch.Tensor: + return torch.load(self._path(index), weights_only=True) + + def delete(self, index: int) -> None: + self._path(index).unlink() + + +def _check(cancel: threading.Event) -> None: + if cancel.is_set(): + raise Cancelled() + + +def scan( + frames: FrameSource, + detector, + *, + frame_h: int, + frame_w: int, + large_canvas: bool, + progress: Progress, + cancel: threading.Event, +) -> list[TrackPlan]: + regions: list[list[Region]] = [] + cuts: set[int] = set() + scene = SceneCutDetector() + with progress.bar("scan", progress.frames) as bar: + for batch, _pts in frames(): + _check(cancel) + cuts.update(len(regions) + offset for offset in scene.find_cuts(batch)) + regions.extend(regions_per_frame(detector(batch, target_hw=(frame_h, frame_w)), frame_h, frame_w)) + bar.update(len(batch)) + plans = plan_video(regions, cuts, frame_w=frame_w, frame_h=frame_h, large_canvas=large_canvas) + logger.info( + "LTX plan: %d frames, %d scene cuts, %d tracks, %d windows", + len(regions), + len(cuts), + len(plans), + sum(len(plan.windows) for plan in plans), + ) + return plans + + +def _vae_pixels(canvases: list[torch.Tensor], device: torch.device) -> torch.Tensor: + stacked = torch.stack(canvases + [canvases[-1]] * (WINDOW_FRAMES - len(canvases)), dim=1) + return stacked.to(device).float().div_(127.5).sub_(1.0).to(torch.bfloat16).unsqueeze(0) + + +def encode_references( + frames: FrameSource, + plans: list[TrackPlan], + encoder, + store: _LatentStore, + *, + device: torch.device, + progress: Progress, + cancel: threading.Event, +) -> None: + windows = sorted((w for plan in plans for w in plan.windows), key=lambda w: w.start) + pending = iter(windows) + upcoming = next(pending, None) + active: dict[int, tuple[Window, list[torch.Tensor]]] = {} + frame_idx = 0 + bar = progress.bar("encode", progress.frames) + for batch, _pts in frames(): + _check(cancel) + bar.update(len(batch)) + for frame in batch: + while upcoming is not None and upcoming.start == frame_idx: + active[upcoming.index] = (upcoming, []) + upcoming = next(pending, None) + for index, (window, canvases) in list(active.items()): + canvases.append(crop_to_canvas(frame, window.crops[frame_idx - window.start]).cpu()) + if len(canvases) == window.real_frames: + with torch.inference_mode(): + store.put(index, encoder(_vae_pixels(canvases, device))) + del active[index] + frame_idx += 1 + bar.close() + if active or upcoming is not None: + raise RuntimeError("the video ended before every planned window was filled") + + +def denoise( + plans: list[TrackPlan], + transformer: LtxTransformer, + references: _LatentStore, + finals: _LatentStore, + *, + seed: int, + progress: Progress, + cancel: threading.Event, +) -> None: + bar = progress.bar("denoise", sum(len(plan.windows) for plan in plans), unit="window") + for plan in plans: + _check(cancel) + latents = transformer.denoise_chain( + [references.get(w.index) for w in plan.windows], [seed + w.index for w in plan.windows] + ) + for window, latent in zip(plan.windows, latents): + finals.put(window.index, latent) + references.delete(window.index) + bar.update(len(plan.windows)) + bar.close() + + +def _decode_window(decoder, latent: torch.Tensor, seed: int, device: torch.device) -> torch.Tensor: + """Decoded window as uint8 ``[F, 3, S, S]`` on ``device``.""" + torch.cuda.empty_cache() + free, _ = torch.cuda.mem_get_info(device) + generator = torch.Generator(device=device).manual_seed(seed) + with torch.inference_mode(): + video = decode_latent(decoder, latent.to(device), free_bytes=free, generator=generator) + video = ((video + 1.0) / 2.0).clamp(0.0, 1.0)[0].float() + return video.mul_(255.0).round_().to(torch.uint8).permute(1, 0, 2, 3).contiguous() + + +def compose( + frames: FrameSource, + plans: list[TrackPlan], + decoder, + finals: _LatentStore, + write: Callable[[torch.Tensor, int], None], + *, + frame_h: int, + seed: int, + device: torch.device, + progress: Progress, + cancel: threading.Event, +) -> None: + feather = feather_pixels(frame_h) + weights: dict[int, torch.Tensor] = {} + tracks: dict[int, TrackPlan] = {} + for plan in plans: + for window, weight in zip(plan.windows, crossfade_weights(plan.windows)): + weights[window.index] = weight + tracks[window.index] = plan + pending = iter(sorted((w for plan in plans for w in plan.windows), key=lambda w: w.start)) + upcoming = next(pending, None) + covering: list[Window] = [] + decoded: dict[int, torch.Tensor] = {} + frame_idx = 0 + bar = progress.bar("compose", progress.frames) + for batch, pts in frames(): + _check(cancel) + bar.update(len(batch)) + for frame, frame_pts in zip(batch, pts): + while upcoming is not None and upcoming.start == frame_idx: + covering.append(upcoming) + upcoming = next(pending, None) + for window in [w for w in covering if w.stop <= frame_idx]: + covering.remove(window) + decoded.pop(window.index, None) + candidates = [] + for window in covering: + if window.index not in decoded: + decoded[window.index] = _decode_window(decoder, finals.get(window.index), seed + window.index, device) + local = frame_idx - window.start + track = tracks[window.index] + candidates.append( + Candidate( + canvas=decoded[window.index][local], + crop=window.crops[local], + weight=float(weights[window.index][local]), + polygons=track.polygons[frame_idx - track.start], + ) + ) + write(composite_frame(frame, candidates, feather=feather), frame_pts) + frame_idx += 1 + bar.close() + + +def restore_video( + frames: FrameSource, + write: Callable[[torch.Tensor, int], None], + *, + detector, + files: LtxModelFiles, + frame_h: int, + frame_w: int, + large_canvas: bool, + seed: int, + device: torch.device, + work_dir: Path, + progress: Progress, + cancel: threading.Event, +) -> None: + """Restore every mosaic of the video ``frames`` yields and ``write`` each output frame + with its pts. ``detector`` is only used by the scan pass; ``large_canvas`` allows 768 px + windows for large mosaics; window ``i`` draws its noise from ``seed + i``.""" + plans = scan(frames, detector, frame_h=frame_h, frame_w=frame_w, large_canvas=large_canvas, progress=progress, cancel=cancel) + with tempfile.TemporaryDirectory(dir=work_dir, prefix=".ltx-") as temp: + references = _LatentStore(Path(temp), "reference") + finals = _LatentStore(Path(temp), "final") + if plans: + encoder = load_video_encoder(files.vae, device) + encode_references(frames, plans, encoder, references, device=device, progress=progress, cancel=cancel) + del encoder + torch.cuda.empty_cache() + transformer = LtxTransformer(files.transformer, device) + try: + denoise(plans, transformer, references, finals, seed=seed, progress=progress, cancel=cancel) + finally: + transformer.close() + del transformer + torch.cuda.empty_cache() + decoder = load_video_decoder(files.vae, files.tuned_decoder, device) if plans else None + compose( + frames, + plans, + decoder, + finals, + write, + frame_h=frame_h, + seed=seed, + device=device, + progress=progress, + cancel=cancel, + ) diff --git a/jasna/ltx/sampler.py b/jasna/ltx/sampler.py new file mode 100644 index 00000000..eb8ec9ae --- /dev/null +++ b/jasna/ltx/sampler.py @@ -0,0 +1,180 @@ +"""The LTX restoration transformer and its fused-chain sampler. + +A track's windows are denoised in lockstep: every step runs all windows through each +block before the next block is fetched, so a streamed block crosses PCIe once per step +per window group. After each step the windows' shared latent frames are replaced by a +ramp-weighted mean (MultiDiffusion-style overlap fusion). Sampler: Euler over a fixed +sigma schedule, CFG 1, STG on one block whose perturbed pass reuses the conditional +pass's hidden state up to that block. +""" + +from __future__ import annotations + +import logging +from collections.abc import Sequence +from pathlib import Path + +import torch +from safetensors import safe_open + +from jasna.ltx import transformer as T +from jasna.ltx.blocks import BlockStore +from jasna.ltx.plan import FUSION_RAMP, LEFT_SHARED_LATENTS, RIGHT_SHARED_LATENTS + +logger = logging.getLogger(__name__) + +LOWTAIL15_SIGMAS = ( + 1.0, 0.982, 0.959, 0.93, 0.892, 0.839, 0.762, 0.637, 0.4, + 0.32, 0.184, 0.106, 0.061, 0.035, 0.02, 0.0, +) # fmt: skip +STG_SCALE = 1.0 +STG_BLOCK = 28 +ACTIVATION_RESERVE_BYTES = 3 << 30 +_ADALN_PREFIXES = ("adaln_single.", "prompt_adaln_single.") + + +def patchify(latent: torch.Tensor) -> torch.Tensor: + """``[1, C, F, H, W]`` -> ``[1, F*H*W, C]``.""" + return latent.flatten(2).transpose(1, 2) + + +def unpatchify(tokens: torch.Tensor, frames: int, height: int, width: int) -> torch.Tensor: + return tokens.transpose(1, 2).reshape(1, -1, frames, height, width) + + +def fuse_overlaps(latents: Sequence[torch.Tensor], tokens_per_frame: int, right_weights: Sequence[float]) -> None: + """Write the weighted mean of each consecutive pair's shared latent frames into both, + in place. All-zero weights copy the earlier window into the later one.""" + for left, right in zip(latents, latents[1:]): + a = left[:, LEFT_SHARED_LATENTS.start * tokens_per_frame : LEFT_SHARED_LATENTS.stop * tokens_per_frame] + b = right[:, RIGHT_SHARED_LATENTS.start * tokens_per_frame : RIGHT_SHARED_LATENTS.stop * tokens_per_frame] + weight = torch.tensor(right_weights, dtype=torch.float32, device=a.device) + weight = weight.repeat_interleave(tokens_per_frame).view(1, -1, 1) + fused = (a.float() * (1.0 - weight) + b.float() * weight).to(a.dtype) + a.copy_(fused) + b.copy_(fused) + + +def _window_state_bytes(tokens: int) -> int: + """Hidden states one window keeps between blocks (conditional + perturbed).""" + return 2 * 2 * tokens * 4096 * 2 + + +def _window_workspace_bytes(tokens: int) -> int: + """Transient activations of one window inside one block.""" + return 2 * tokens * 4096 * 2 * 10 + T.CHUNK_ROWS * 16384 * 4 * 3 + + +class LtxTransformer: + """The restoration transformer loaded from a jasna LTX model file.""" + + def __init__(self, path: Path, device: torch.device, *, sigmas: Sequence[float] = LOWTAIL15_SIGMAS) -> None: + self.device = device + self.sigmas = torch.tensor(sigmas, dtype=torch.float32) + with safe_open(str(path), framework="pt", device="cpu") as handle: + keys = list(handle.keys()) + top_keys = [k for k in keys if not k.startswith("blocks.") and k != "prompt_context"] + top = {k: handle.get_tensor(k).to(device) for k in top_keys} + self.context = handle.get_tensor("prompt_context").to(device) + self.conditions = [T.step_conditioning(top, float(s), device) for s in self.sigmas[:-1]] + self.top = {k: v for k, v in top.items() if not k.startswith(_ADALN_PREFIXES)} + del top + blocks = [ + {k.split(".", 2)[2]: handle.get_tensor(k) for k in keys if k.startswith(f"blocks.{i}.")} + for i in range(T.NUM_BLOCKS) + ] + self.blocks = BlockStore(blocks, device, resident=self._resident_blocks(blocks)) + self._ropes: dict[tuple[int, int, int], T.Rope] = {} + + def _resident_blocks(self, blocks: list[dict[str, torch.Tensor]]) -> int: + sizes = [sum(t.numel() * t.element_size() for t in b.values()) for b in blocks] + free, _ = torch.cuda.mem_get_info(self.device) + budget = free - ACTIVATION_RESERVE_BYTES - 2 * max(sizes) + resident, used = 0, 0 + for size in sizes: + if used + size > budget: + break + used += size + resident += 1 + return resident + + def close(self) -> None: + self.blocks.close() + self.top.clear() + + def _rope(self, frames: int, height: int, width: int) -> T.Rope: + key = (frames, height, width) + if key not in self._ropes: + dim = self.top["patchify_proj.weight"].shape[0] + self._ropes[key] = T.build_rope(frames, height, width, dim=dim, heads=T.HEADS, device=self.device) + return self._ropes[key] + + def _group_size(self, windows: int, tokens: int) -> int: + free, _ = torch.cuda.mem_get_info(self.device) + free += torch.cuda.memory_reserved(self.device) - torch.cuda.memory_allocated(self.device) + fits = (free - _window_workspace_bytes(tokens)) // _window_state_bytes(tokens) + return max(1, min(windows, int(fits))) + + @torch.inference_mode() + def denoise_chain(self, references: Sequence[torch.Tensor], seeds: Sequence[int]) -> list[torch.Tensor]: + """Final latents ``[1, 128, F, H, W]`` of one track's windows, given each window's + reference latent (the encoded mosaic crop) and noise seed.""" + frames, height, width = references[0].shape[2:] + tokens_per_frame = height * width + tokens = frames * tokens_per_frame + rope = self._rope(frames, height, width) + refs = [patchify(r.to(self.device, torch.bfloat16)) for r in references] + latents = [] + for ref, seed in zip(refs, seeds): + generator = torch.Generator(device=self.device).manual_seed(int(seed)) + noise = torch.randn(1, 2 * tokens, ref.shape[-1], device=self.device, dtype=torch.bfloat16, generator=generator) + latents.append(noise[:, :tokens].clone()) + fuse_overlaps(latents, tokens_per_frame, [0.0] * len(FUSION_RAMP)) + group = self._group_size(len(latents), 2 * tokens) + sigmas = self.sigmas.to(self.device) + for step, cond in enumerate(self.conditions): + for start in range(0, len(latents), group): + members = range(start, min(start + group, len(latents))) + denoised = self._denoise( + [latents[i] for i in members], [refs[i] for i in members], cond, rope, tokens_per_frame, sigmas[step] + ) + for i, value in zip(members, denoised): + latents[i] = _euler(latents[i], value, sigmas, step) + fuse_overlaps(latents, tokens_per_frame, FUSION_RAMP) + return [unpatchify(latent, frames, height, width) for latent in latents] + + def _denoise( + self, + latents: list[torch.Tensor], + refs: list[torch.Tensor], + cond: T.StepConditioning, + rope: T.Rope, + tokens_per_frame: int, + sigma: torch.Tensor, + ) -> list[torch.Tensor]: + xs = [T.embed_tokens(self.top, torch.cat([lat, ref], dim=1), tokens_per_frame) for lat, ref in zip(latents, refs)] + perturbed: list[torch.Tensor | None] = [None] * len(xs) + for index in range(T.NUM_BLOCKS): + weights = self.blocks.acquire(index) + for j in range(len(xs)): + if STG_SCALE and index >= STG_BLOCK: + source = xs[j] if index == STG_BLOCK else perturbed[j] + perturbed[j] = T.block_forward( + weights, source, cond, self.context, rope, skip_self_attention=index == STG_BLOCK + ) + xs[j] = T.block_forward(weights, xs[j], cond, self.context, rope) + self.blocks.release(index) + out = [] + for lat, x, xp in zip(latents, xs, perturbed): + denoised = (lat.float() - T.velocity(self.top, x, cond).float() * sigma).to(lat.dtype) + if STG_SCALE: + weak = (lat.float() - T.velocity(self.top, xp, cond).float() * sigma).to(lat.dtype) + denoised = (denoised.float() + STG_SCALE * (denoised.float() - weak.float())).to(lat.dtype) + out.append(denoised) + return out + + +def _euler(sample: torch.Tensor, denoised: torch.Tensor, sigmas: torch.Tensor, step: int) -> torch.Tensor: + sigma = sigmas[step] + velocity = ((sample.float() - denoised.float()) / sigma.item()).to(sample.dtype) + return (sample.float() + velocity.float() * (sigmas[step + 1] - sigma)).to(sample.dtype) diff --git a/jasna/ltx/transformer.py b/jasna/ltx/transformer.py new file mode 100644 index 00000000..de5219c4 --- /dev/null +++ b/jasna/ltx/transformer.py @@ -0,0 +1,306 @@ +"""Video-only LTX-2.5 transformer as plain functions over per-block weight dicts. + +A port of the Lightricks LTX-2.5 ``ltx-core`` 22B forward (``model/transformer``, commit +6ea1527869a5ce57452e215595eae189a7cf65cc; LTX-2.x Community License) restricted to what +restoration runs: video tokens, cross-attention AdaLN, gated attention, split RoPE, and a +per-window token layout of ``n`` target tokens followed by ``n`` reference tokens sharing +the same positions. Target tokens carry timestep ``sigma`` and reference tokens 0, so +every AdaLN modulation is one row per half instead of one per token. + +Blocks take their weights as a dict so the executor can hand in resident tensors or views +into a streaming slot. A Linear is stored either as ``.weight`` (bf16) or as ConvRot +W8A8 (``.qdata`` int8 rotated codes, ``.scales`` fp32 per output channel). +""" + +from __future__ import annotations + +import functools +import math +from collections.abc import Mapping +from dataclasses import dataclass + +import numpy as np +import torch +import torch.nn.functional as F +from torch.nn.attention import SDPBackend, sdpa_kernel + +Weights = Mapping[str, torch.Tensor] + +NUM_BLOCKS = 48 +HEADS = 32 +EPS = 1e-6 +TIMESTEP_SCALE = 1000.0 +ROPE_THETA = 10000.0 +ROPE_MAX_POS = (20, 2048, 2048) +POSITION_FPS = 24.0 +LATENT_TEMPORAL_SCALE = 8 +LATENT_SPATIAL_SCALE = 32 +ROT_SIZE = 256 +ACT_QMAX = 127 +CHUNK_ROWS = 4096 +SDPA_PRIORITY = [ + SDPBackend.CUDNN_ATTENTION, + SDPBackend.FLASH_ATTENTION, + SDPBackend.EFFICIENT_ATTENTION, + SDPBackend.MATH, +] + +LINEAR_NAMES = tuple( + f"{attn}.{proj}" + for attn in ("attn1", "attn2") + for proj in ("to_q", "to_k", "to_v", "to_out.0", "to_gate_logits") +) + ("ff.net.0.proj", "ff.net.2") + + +@functools.cache +def _regular_hadamard(size: int, device: torch.device, dtype: torch.dtype) -> torch.Tensor: + r4 = torch.tensor([[1.0, 1, 1, -1], [1, 1, -1, 1], [1, -1, 1, 1], [-1, 1, 1, 1]]) + h = r4.clone() + while h.shape[0] < size: + h = torch.kron(h, r4) + return (h / size**0.5).to(device=device, dtype=dtype) + + +def rotate(x: torch.Tensor) -> torch.Tensor: + """ConvRot block regular-Hadamard rotation along the last dim (self-inverse).""" + h = _regular_hadamard(ROT_SIZE, x.device, x.dtype) + shape = x.shape + return torch.matmul(x.reshape(-1, shape[-1] // ROT_SIZE, ROT_SIZE), h).reshape(shape) + + +def quantize_rows(x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + """Symmetric per-row int8 (round half to even), rows padded to a multiple of 32 + for ``torch._int_mm``. Returns (int8 codes, fp32 scales).""" + rows = x.shape[0] + padded = -(-rows // 32) * 32 + xf = x.float() + scales = xf.abs().amax(dim=1) / ACT_QMAX + scales = torch.where(scales > 0, scales, torch.ones_like(scales)) + codes = torch.round(xf / scales.unsqueeze(1)).clamp_(-ACT_QMAX, ACT_QMAX).to(torch.int8) + if padded != rows: + codes = F.pad(codes, (0, 0, 0, padded - rows)) + scales = F.pad(scales, (0, padded - rows), value=1.0) + return codes, scales + + +def int8_linear( + x: torch.Tensor, qdata: torch.Tensor, scales: torch.Tensor, bias: torch.Tensor | None +) -> torch.Tensor: + """ConvRot W8A8: rotate, quantize per token, INT8 tensor-core GEMM, rescale. + Rows run in chunks of ``CHUNK_ROWS``; every step is per row, so chunking only caps + the fp32/int32 temporaries.""" + out_features, in_features = qdata.shape + flat = x.reshape(-1, in_features) + out = torch.empty(flat.shape[0], out_features, dtype=x.dtype, device=x.device) + for start in range(0, flat.shape[0], CHUNK_ROWS): + chunk = flat[start : start + CHUNK_ROWS] + rows = chunk.shape[0] + codes, act_scales = quantize_rows(rotate(chunk)) + acc = torch._int_mm(codes, qdata.t())[:rows] + result = acc.float() * (act_scales[:rows].unsqueeze(1) * scales) + if bias is not None: + result = result + bias.float() + out[start : start + rows] = result + return out.reshape(*x.shape[:-1], out_features) + + +def linear(w: Weights, name: str, x: torch.Tensor) -> torch.Tensor: + bias = w.get(f"{name}.bias") + weight = w.get(f"{name}.weight") + if weight is not None: + return F.linear(x, weight, bias) + return int8_linear(x, w[f"{name}.qdata"], w[f"{name}.scales"], bias) + + +def rms_norm(x: torch.Tensor, weight: torch.Tensor | None = None) -> torch.Tensor: + return F.rms_norm(x, (x.shape[-1],), weight=weight, eps=EPS) + + +# --------------------------------------------------------------------------- # +# Positions +# --------------------------------------------------------------------------- # +@dataclass(frozen=True) +class Rope: + """cos/sin ``[1, heads, n, head_dim/2]`` for the ``n`` target tokens; the reference + half reuses them (identical positions).""" + + cos: torch.Tensor + sin: torch.Tensor + + +def _freq_grid(dim: int) -> torch.Tensor: + count = dim // (2 * len(ROPE_MAX_POS)) + grid = np.power( + ROPE_THETA, + np.linspace(0.0, np.log(ROPE_THETA) / np.log(ROPE_THETA), count, dtype=np.float64), + ) + return torch.as_tensor(grid * math.pi / 2, dtype=torch.float32) + + +def latent_positions(frames: int, height: int, width: int) -> torch.Tensor: + """``[3, n]`` token-centre positions (seconds, pixel y, pixel x), causal first frame.""" + f, h, w = torch.meshgrid(torch.arange(frames), torch.arange(height), torch.arange(width), indexing="ij") + starts = torch.stack([f, h, w]).reshape(3, -1) + bounds = torch.stack([starts, starts + 1], dim=-1) + scale = torch.tensor([LATENT_TEMPORAL_SCALE, LATENT_SPATIAL_SCALE, LATENT_SPATIAL_SCALE]).view(3, 1, 1) + pixel = bounds * scale + pixel[0] = (pixel[0] + 1 - LATENT_TEMPORAL_SCALE).clamp(min=0) + pixel = pixel.float() + pixel[0] = pixel[0] / POSITION_FPS + return (pixel[..., 0] + pixel[..., 1]) / 2.0 + + +def build_rope(frames: int, height: int, width: int, *, dim: int, heads: int, device: torch.device) -> Rope: + positions = latent_positions(frames, height, width) + fractional = torch.stack([positions[i] / ROPE_MAX_POS[i] for i in range(3)], dim=-1)[None] + indices = _freq_grid(dim) + freqs = (indices * (fractional.unsqueeze(-1) * 2 - 1)).transpose(-1, -2).flatten(2) + pad = dim // 2 - freqs.shape[-1] + cos, sin = freqs.cos(), freqs.sin() + cos = torch.cat([torch.ones_like(cos[:, :, :pad]), cos], dim=-1) + sin = torch.cat([torch.zeros_like(sin[:, :, :pad]), sin], dim=-1) + tokens = cos.shape[1] + cos = cos.reshape(1, tokens, heads, -1).swapaxes(1, 2) + sin = sin.reshape(1, tokens, heads, -1).swapaxes(1, 2) + return Rope(cos=cos.to(device=device, dtype=torch.bfloat16), sin=sin.to(device=device, dtype=torch.bfloat16)) + + +def _apply_rope(x: torch.Tensor, rope: Rope, heads: int) -> torch.Tensor: + """``x`` ``[1, 2n, heads*d]`` -> ``[1, heads, 2n, d]`` rotated (split halves).""" + n = rope.cos.shape[2] + halves = x.unflatten(-1, (heads, 2, -1)).transpose(1, 2).unflatten(2, (2, n)) # [1,H,2,n,2,d/2] + cos = rope.cos[:, :, None, :, None, :] + sin = rope.sin[:, :, None, :, None, :] + out = halves * cos + out[..., :1, :].addcmul_(-sin, halves[..., 1:, :]) + out[..., 1:, :].addcmul_(sin, halves[..., :1, :]) + return out.flatten(4).flatten(2, 3) + + +# --------------------------------------------------------------------------- # +# Timestep conditioning +# --------------------------------------------------------------------------- # +def _sinusoid(t: torch.Tensor) -> torch.Tensor: + half = 128 + exponent = -math.log(10000) * torch.arange(half, dtype=torch.float32, device=t.device) / half + emb = t[:, None].float() * torch.exp(exponent)[None, :] + return torch.cat([torch.cos(emb), torch.sin(emb)], dim=-1) + + +def adaln(w: Weights, prefix: str, t: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + """AdaLayerNormSingle on timesteps ``t`` (already scaled): (modulation, embedding).""" + proj = _sinusoid(t).to(torch.bfloat16) + emb = linear(w, f"{prefix}.emb.timestep_embedder.linear_2", F.silu(linear(w, f"{prefix}.emb.timestep_embedder.linear_1", proj))) + return linear(w, f"{prefix}.linear", F.silu(emb)), emb + + +@dataclass(frozen=True) +class StepConditioning: + """Everything one denoising step needs besides the tokens, for (target, reference).""" + + block_rows: torch.Tensor # [1, 2, 9, D]: AdaLN rows per half, before each block's table + output_rows: torch.Tensor # [1, 2, D]: output-norm embedding per half + prompt_rows: torch.Tensor # [1, 1, 2, D]: cross-attention K/V modulation before each table + + +def step_conditioning(w: Weights, sigma: float, device: torch.device) -> StepConditioning: + sigma_t = torch.tensor([sigma], dtype=torch.float32, device=device) + t = torch.cat([sigma_t, torch.zeros_like(sigma_t)]) * TIMESTEP_SCALE + rows, embedded = adaln(w, "adaln_single", t) + prompt, _ = adaln(w, "prompt_adaln_single", sigma_t * TIMESTEP_SCALE) + dim = embedded.shape[-1] + return StepConditioning( + block_rows=rows.view(1, 2, -1, dim), + output_rows=embedded.view(1, 2, dim), + prompt_rows=prompt.view(1, 1, 2, dim), + ) + + +# --------------------------------------------------------------------------- # +# Block +# --------------------------------------------------------------------------- # +def _modulate(x: torch.Tensor, shift: torch.Tensor, scale: torch.Tensor) -> torch.Tensor: + """``x * (1 + scale) + shift`` with one row per half: x ``[1, 2n, D]``, rows ``[1, 2, D]``.""" + halves = x.unflatten(1, (2, -1)) + return (halves * (1 + scale[:, :, None]) + shift[:, :, None]).flatten(1, 2) + + +def _gate(x: torch.Tensor, gate: torch.Tensor) -> torch.Tensor: + return (x.unflatten(1, (2, -1)) * gate[:, :, None]).flatten(1, 2) + + +def _attention( + w: Weights, + prefix: str, + x: torch.Tensor, + context: torch.Tensor, + rope: Rope | None, + heads: int, + *, + skip: bool, +) -> torch.Tensor: + v = linear(w, f"{prefix}.to_v", context) + if skip: + out = v + else: + q = rms_norm(linear(w, f"{prefix}.to_q", x), w[f"{prefix}.q_norm.weight"]) + k = rms_norm(linear(w, f"{prefix}.to_k", context), w[f"{prefix}.k_norm.weight"]) + if rope is not None: + q, k = _apply_rope(q, rope, heads), _apply_rope(k, rope, heads) + else: + q = q.unflatten(-1, (heads, -1)).transpose(1, 2) + k = k.unflatten(-1, (heads, -1)).transpose(1, 2) + v_heads = v.unflatten(-1, (heads, -1)).transpose(1, 2) + with sdpa_kernel(SDPA_PRIORITY, set_priority=True): + out = F.scaled_dot_product_attention(q, k, v_heads) + out = out.transpose(1, 2).flatten(2) + gates = 2.0 * torch.sigmoid(linear(w, f"{prefix}.to_gate_logits", x)) + out = (out.unflatten(-1, (heads, -1)) * gates.unsqueeze(-1)).flatten(2) + return linear(w, f"{prefix}.to_out.0", out) + + +def block_forward( + w: Weights, + x: torch.Tensor, + cond: StepConditioning, + context: torch.Tensor, + rope: Rope, + *, + heads: int = HEADS, + skip_self_attention: bool = False, +) -> torch.Tensor: + """One transformer block on ``x`` ``[1, 2n, D]`` (target half, then reference half).""" + ada = w["scale_shift_table"][None, None] + cond.block_rows # [1, 2, 9, D] + attn_in = _modulate(rms_norm(x), ada[:, :, 0], ada[:, :, 1]) + attn_out = _attention(w, "attn1", attn_in, attn_in, rope, heads, skip=skip_self_attention) + x = x + _gate(attn_out, ada[:, :, 2]) + kv = w["prompt_scale_shift_table"][None, None] + cond.prompt_rows + query = _modulate(rms_norm(x), ada[:, :, 6], ada[:, :, 7]) + prompt = context * (1 + kv[:, :, 1]) + kv[:, :, 0] + x = x + _gate(_attention(w, "attn2", query, prompt, None, heads, skip=False), ada[:, :, 8]) + ff_in = _modulate(rms_norm(x), ada[:, :, 3], ada[:, :, 4]) + ff_out = torch.empty_like(ff_in) + for start in range(0, ff_in.shape[1], CHUNK_ROWS): + hidden = F.gelu(linear(w, "ff.net.0.proj", ff_in[:, start : start + CHUNK_ROWS]), approximate="tanh") + ff_out[:, start : start + CHUNK_ROWS] = linear(w, "ff.net.2", hidden) + return x + _gate(ff_out, ada[:, :, 5]) + + +# --------------------------------------------------------------------------- # +# Input / output +# --------------------------------------------------------------------------- # +def embed_tokens(w: Weights, tokens: torch.Tensor, tokens_per_frame: int) -> torch.Tensor: + """Project ``[1, 2n, 128]`` latent tokens; the target's first (single-pixel-frame) + latent frame gets the keyframe marker embedding.""" + x = linear(w, "patchify_proj", tokens) + x[:, :tokens_per_frame] = x[:, :tokens_per_frame] + w["keyframes_abs_pos_embedding"] + return x + + +def velocity(w: Weights, x: torch.Tensor, cond: StepConditioning) -> torch.Tensor: + """Predicted velocity of the target half.""" + n = x.shape[1] // 2 + mod = w["scale_shift_table"][None, None] + cond.output_rows[:, :1, None] # [1, 1, 2, D] + target = F.layer_norm(x[:, :n], (x.shape[-1],), eps=EPS) + target = target * (1 + mod[:, :, 1]) + mod[:, :, 0] + return linear(w, "proj_out", target) diff --git a/jasna/main.py b/jasna/main.py index d187b287..0dbc3897 100644 --- a/jasna/main.py +++ b/jasna/main.py @@ -5,7 +5,7 @@ from jasna import __version__ from jasna.cli_help import CLI_HELP -from jasna.engine_paths import model_weights_dir +from jasna.engine_paths import default_restoration_model_path from jasna.media.probe import UnsupportedColorspaceError from jasna.os_utils import ( MIN_DRIVER_VERSION, @@ -16,7 +16,8 @@ check_windows_nvidia_sysmem_fallback_policy, gpu_check_error, ) -from jasna.session_config import SessionConfig +from jasna.session_config import LTX_DEFAULT_SEED, SessionConfig + def _session_config_from_args( @@ -40,7 +41,10 @@ def _session_config_from_args( max_detection_gap=int(args.max_detection_gap), min_detection_duration=int(args.min_detection_duration), scene_detection=bool(args.scene_detection), + restoration_model_name=str(args.restoration_model_name), restoration_model_path=restoration_model_path, + ltx_large_canvas=bool(args.ltx_large_canvas), + ltx_seed=int(args.ltx_seed), compile_basicvsrpp=bool(args.compile_basicvsrpp), max_clip_size=int(args.max_clip_size), temporal_overlap=int(args.temporal_overlap), @@ -190,14 +194,26 @@ def build_parser() -> argparse.ArgumentParser: "--restoration-model-name", type=str, default="basicvsrpp", - choices=["basicvsrpp"], - help='Restoration model for video input (only "basicvsrpp" supported for now).', + choices=["basicvsrpp", "ltx"], + help=CLI_HELP["restoration_model_name"], ) restoration.add_argument( "--restoration-model-path", type=str, - default=str(model_weights_dir() / "lada_mosaic_restoration_model_generic_v1.2.pth"), - help="Path to restoration model (default: %(default)s)", + default=None, + help=CLI_HELP["restoration_model_path"], + ) + restoration.add_argument( + "--ltx-large-canvas", + default=True, + action=argparse.BooleanOptionalAction, + help=CLI_HELP["ltx_large_canvas"], + ) + restoration.add_argument( + "--ltx-seed", + type=int, + default=LTX_DEFAULT_SEED, + help=CLI_HELP["ltx_seed"], ) restoration.add_argument( "--compile-basicvsrpp", @@ -838,7 +854,14 @@ def _run_post_export_action() -> None: str(args.detection_model), str(args.detection_model_path), args.detection_score_threshold ) - restoration_model_path = Path(args.restoration_model_path) + if args.restoration_model_name == "ltx": + if is_streaming or segments_spec: + parser.error("--restoration-model-name ltx does not support --stream or --segments") + if args.secondary_restoration != "none" or args.denoise != "none": + parser.error("--restoration-model-name ltx does not support --secondary-restoration or --denoise") + restoration_model_path = Path( + args.restoration_model_path or default_restoration_model_path(args.restoration_model_name) + ) if not restoration_model_path.exists(): raise FileNotFoundError(str(restoration_model_path)) diff --git a/jasna/models/ltx_vae/__init__.py b/jasna/models/ltx_vae/__init__.py new file mode 100644 index 00000000..cb2b1629 --- /dev/null +++ b/jasna/models/ltx_vae/__init__.py @@ -0,0 +1,15 @@ +"""LTX-2.5 video VAE (causal conv encoder + single-step diffusion decoder), inference only. + +Pure-PyTorch subset ported from the Lightricks LTX-2.5 monorepo, package ``ltx-core`` +(``ltx_core/model/video_vae``, ``tiling.py``, ``types.py``, timestep embedding, normalization, +state-dict ops), commit 6ea1527869a5ce57452e215595eae189a7cf65cc. The decoder reproduces the +``chunked_eager`` mode with the eager tiled-SDPA neighborhood-attention fallback +(``eager_na.py``, originally from comfy-kitchen, Apache-2.0). + +This code and the model weights are covered by the LTX-2.x Community License +(see the ``LICENSE`` file of the LTX-2.5 repository). +""" + +from jasna.models.ltx_vae.loader import decode_latent, load_video_decoder, load_video_encoder + +__all__ = ["decode_latent", "load_video_decoder", "load_video_encoder"] diff --git a/jasna/models/ltx_vae/diffusion_tiling.py b/jasna/models/ltx_vae/diffusion_tiling.py new file mode 100644 index 00000000..dbeca25a --- /dev/null +++ b/jasna/models/ltx_vae/diffusion_tiling.py @@ -0,0 +1,488 @@ +"""DiffVAE tiling helpers: schedule, pad/crop/size-floor, blend utilities. + +Subset of ``ltx_core/model/video_vae/diffusion_tiling.py`` for the ``chunked_eager`` mode. +""" + +from __future__ import annotations + +import itertools +import math +from dataclasses import dataclass +from typing import List, Literal, Sequence, Tuple + +import torch + +from jasna.models.ltx_vae.tiling import ( + VIDEO_SCALE_FACTORS, + DimensionInterval, + DimensionSizeConfig, + SpatioTemporalScaleFactors, + Tile, + TileSizeConfig, + compute_trapezoidal_mask_1d, + split_by_size, + untiled_mask_1d, +) + +ResizeAxisMode = Literal["repeat_last", "symmetric"] + +# Peak-activation heuristic (bytes): +# stage-4 input feature (resident for the whole decode): s4_t x s4_h x s4_w x stage4_channels x bf16 +# accumulator: H x W x (2 x tile_t) x out_channels x fp16 +# stage-5: stage5_tokens x stage5_channels x bf16 x coef (NA working-set multiplicity) +_STAGE5_MEM_COEF: float = 5 +_ELEMENT_SIZE: int = 2 +_MIN_MODEL_BYTES_FLOOR: int = 1 << 30 +_BUDGET_SAFETY_BYTES: int = 1 << 30 + + +def stage4_feature_bytes( + *, + height: int, + width: int, + num_frames: int, + upsample_strides: Sequence[Tuple[int, int, int]], + stage4_channels: int, + natten_trailing_pad_latent_frames: int, +) -> int: + """Resident stages-1-3 output size (full volume, including the trailing latent pad).""" + latent_frames = (num_frames - 1) // VIDEO_SCALE_FACTORS.time + 1 + s4_t, s4_h, s4_w = stage4_thw_from_latent( + upsample_strides[:3], + latent_frames + natten_trailing_pad_latent_frames, + height // VIDEO_SCALE_FACTORS.height, + width // VIDEO_SCALE_FACTORS.width, + drop_leading_frame=True, + ) + return s4_t * s4_h * s4_w * stage4_channels * _ELEMENT_SIZE + + +def recommended_decode_tiling_config( + *, + tile_halos: Tuple[Tuple[int, int, int], Tuple[int, int, int]], + pixel_scale: SpatioTemporalScaleFactors, + min_tile_size_s4: Tuple[int, int, int], + patch_size: int, + height: int, + width: int, + num_frames: int, + free_bytes: int, + stage5_channels: int, + stage4_channels: int, + upsample_strides: Sequence[Tuple[int, int, int]], + model_bytes: int, + natten_trailing_pad_latent_frames: int, + out_channels: int, +) -> TileSizeConfig: + """Pick decode tiling from stage-4/5 halos and the activation budget. + + Enumerates legal tile sizes, drops those whose peak-bytes estimate exceeds + ``free - max(model, 1 GiB) - 1 GiB - stage4_feature``, and picks the one with the + least overlap recompute (then the largest volume). + """ + overlap_t, overlap_hw = recommended_pixel_overlaps(tile_halos, pixel_scale) + + ft, fh, fw = pixel_scale.time, pixel_scale.height, pixel_scale.width + step_t = math.lcm(ft, VIDEO_SCALE_FACTORS.time) + step_h = math.lcm(fh, VIDEO_SCALE_FACTORS.height) + step_w = math.lcm(fw, VIDEO_SCALE_FACTORS.width) + # ``2 * overlap`` so left+right ramps fit (else masks are not complementary). + min_t_px = _round_up(max(2 * ft, 2 * overlap_t, _round_up(min_tile_size_s4[0] * ft, ft), 16), step_t) + min_h_px = _round_up(max(2 * fh, 2 * overlap_hw, _round_up(min_tile_size_s4[1] * fh, fh), 64), step_h) + min_w_px = _round_up(max(2 * fw, 2 * overlap_hw, _round_up(min_tile_size_s4[2] * fw, fw), 64), step_w) + + model_cost = max(model_bytes, _MIN_MODEL_BYTES_FLOOR) + s4_feat_bytes = stage4_feature_bytes( + height=height, + width=width, + num_frames=num_frames, + upsample_strides=upsample_strides, + stage4_channels=stage4_channels, + natten_trailing_pad_latent_frames=natten_trailing_pad_latent_frames, + ) + usable = max(0, free_bytes - model_cost - _BUDGET_SAFETY_BYTES - s4_feat_bytes) + s5_bytes_per_token = max(1.0, float(stage5_channels) * float(_ELEMENT_SIZE) * _STAGE5_MEM_COEF) + acc_bytes_per_pixel = out_channels * _ELEMENT_SIZE + + t_cands = _axis_candidates(num_frames, overlap_t, min_t_px, step_t) + h_cands = _axis_candidates(height, overlap_hw, min_h_px, step_h) + w_cands = _axis_candidates(width, overlap_hw, min_w_px, step_w) + + scored: list[tuple[float, int, int, int, int, int]] = [] + for tile_t, n_t in t_cands: + # Current group buffer + still-live emit/stub during temporal handoff. + acc_bytes = 2 * tile_t * height * width * acc_bytes_per_pixel + if acc_bytes >= usable: + continue + max_s5_tokens = int((usable - acc_bytes) // s5_bytes_per_token) + for tile_h, n_h in h_cands: + for tile_w, n_w in w_cands: + if stage5_tokens_for_pixel_tile(tile_t, tile_h, tile_w, patch_size=patch_size) > max_s5_tokens: + continue + waste = (n_t * n_h * n_w * tile_t * tile_h * tile_w) / max(1, num_frames * height * width) + scored.append((waste, -tile_t * tile_h * tile_w, n_t * n_h * n_w, tile_t, tile_h, tile_w)) + + if not scored: + raise ValueError( + "Cannot fit a DiffVAE decode tile under the memory budget: " + f"min tile ~{min_t_px}f x {min_h_px}x{min_w_px}px " + f"(overlaps T={overlap_t}, HW={overlap_hw}), " + f"coef={_STAGE5_MEM_COEF}, stage5_channels={stage5_channels}, " + f"stage4_feature_bytes={s4_feat_bytes}, usable_bytes={usable}. " + "Reduce resolution or free GPU memory." + ) + + scored.sort() + _waste, _vol, _ntiles, tile_t, tile_h, tile_w = scored[0] + return TileSizeConfig( + frames=DimensionSizeConfig(tile_size=tile_t, overlap=overlap_t), + height=DimensionSizeConfig(tile_size=tile_h, overlap=overlap_hw), + width=DimensionSizeConfig(tile_size=tile_w, overlap=overlap_hw), + ) + + +def prepare_tile_schedule( + stage4_shape_bcthw: torch.Size, + tiling_config: TileSizeConfig, + *, + upsample3_stride: Tuple[int, int, int], + patch_size: int, + min_tile_size: Tuple[int, int, int], +) -> List[Tile]: + """Build pixel-blend tiles whose ``in_coords`` land on the stage-4 input grid. + + Temporal intervals propagate through the pixel-shuffle hop (``drop_leading_frame`` + only on the origin tile); ramps are symmetric so masks stay complementary. + """ + pixel_scale = stage4_to_pixel_scale_factors(upsample3_stride, patch_size) + split_sizes = tiling_config.split_sizes(pixel_scale) + + def axis_specs(axis: int, *, propagate_causal: bool, apply_patch: bool) -> list[tuple[slice, slice, torch.Tensor]]: + tile, overlap = split_sizes[axis] + stride_component = upsample3_stride[axis] + specs = [] + for iv in split_by_size(stage4_shape_bcthw[2 + axis], tile, overlap, min_tile_size[axis]): + stage5 = _propagate_interval_through_upsample_hops(iv, [stride_component], propagate_causal) + pixel = _propagate_interval_through_upsample_hops(stage5, [patch_size], False) if apply_patch else stage5 + mask_pixel = compute_trapezoidal_mask_1d( + pixel.end - pixel.start, pixel.left_ramp, pixel.right_ramp, left_starts_from_0=False + ) + specs.append((slice(iv.start, iv.end), slice(pixel.start, pixel.end), mask_pixel)) + return specs + + t_specs = axis_specs(0, propagate_causal=True, apply_patch=False) + h_specs = axis_specs(1, propagate_causal=False, apply_patch=True) + w_specs = axis_specs(2, propagate_causal=False, apply_patch=True) + + tiles: List[Tile] = [] + for t_spec, h_spec, w_spec in itertools.product(t_specs, h_specs, w_specs): + t_s4, t_px, t_mask = t_spec + h_s4, h_px, h_mask = h_spec + w_s4, w_px, w_mask = w_spec + tiles.append( + Tile( + in_coords=(slice(None), t_s4, h_s4, w_s4, slice(None)), + out_coords=(slice(None), slice(None), t_px, h_px, w_px), + masks_1d=(untiled_mask_1d(), untiled_mask_1d(), t_mask, h_mask, w_mask), + ) + ) + return tiles + + +def slice_stage4_tile( + feat_s4: torch.Tensor, + tile: Tile, + *, + content_frames: int, +) -> tuple[torch.Tensor, bool, bool, tuple[int, int, int]]: + """Slice a stage-4 feature tile, extending trailing tiles to include ghost frames.""" + is_origin = tile.in_coords[1].start in (0, None) + _, stop, _ = tile.in_coords[1].indices(content_frames) + pad_trailing = stop == content_frames + _b, t_coord, h_coord, w_coord, _c = tile.in_coords + t0, t1, _ = t_coord.indices(content_frames) + h0, h1, _ = h_coord.indices(feat_s4.shape[2]) + w0, w1, _ = w_coord.indices(feat_s4.shape[3]) + content_thw = (t1 - t0, h1 - h0, w1 - w0) + if pad_trailing: + t1 = feat_s4.shape[1] + feat_tile = feat_s4[:, t0:t1, h_coord, w_coord, :] + return feat_tile, is_origin, pad_trailing, content_thw + + +@dataclass(frozen=True) +class AxisPad: + """How many elements were added (pad) or removed (crop) on each side of one axis.""" + + before: int + after: int + + +def resize_axis(x: torch.Tensor, dim: int, size: int, *, mode: ResizeAxisMode) -> tuple[torch.Tensor, AxisPad]: + """Pad or crop axis ``dim`` to ``size``. + + Pad: ``repeat_last`` appends copies of the last slice; ``symmetric`` edge-replicates + both ends (leftover goes to the end). Crop: ``repeat_last`` drops from the end; + ``symmetric`` drops from both ends with the same split. + """ + length = x.shape[dim] + if length == size: + return x, AxisPad(0, 0) + + if length < size: + need = size - length + if mode == "repeat_last": + last = x.narrow(dim, length - 1, 1) + expand_shape = list(x.shape) + expand_shape[dim] = need + return torch.cat([x, last.expand(expand_shape)], dim=dim), AxisPad(0, need) + + before = need // 2 + after = need - before + first = x.narrow(dim, 0, 1) + last = x.narrow(dim, length - 1, 1) + parts: list[torch.Tensor] = [] + if before: + expand_shape = list(x.shape) + expand_shape[dim] = before + parts.append(first.expand(expand_shape)) + parts.append(x) + if after: + expand_shape = list(x.shape) + expand_shape[dim] = after + parts.append(last.expand(expand_shape)) + return torch.cat(parts, dim=dim), AxisPad(before, after) + + need = length - size + if mode == "repeat_last": + return x.narrow(dim, 0, size).contiguous(), AxisPad(0, need) + + before = need // 2 + after = need - before + return x.narrow(dim, before, size).contiguous(), AxisPad(before, after) + + +def ensure_min_latent_shape( + latent: torch.Tensor, + min_tile_sizes: Tuple[int, int, int], +) -> tuple[torch.Tensor, tuple[AxisPad, AxisPad, AxisPad]]: + """Pad latent ``(B, C, T, H, W)`` up to ``min_tile_sizes`` if needed.""" + min_t, min_h, min_w = min_tile_sizes + t_pad = AxisPad(0, 0) + h_pad = AxisPad(0, 0) + w_pad = AxisPad(0, 0) + x = latent + if x.shape[2] < min_t: + x, t_pad = resize_axis(x, 2, min_t, mode="repeat_last") + if x.shape[3] < min_h: + x, h_pad = resize_axis(x, 3, min_h, mode="symmetric") + if x.shape[4] < min_w: + x, w_pad = resize_axis(x, 4, min_w, mode="symmetric") + return x, (t_pad, h_pad, w_pad) + + +def scale_axis_pad(pad: AxisPad, scale: int) -> AxisPad: + return AxisPad(pad.before * scale, pad.after * scale) + + +def crop_pixels_to_content( + pixels: torch.Tensor, + frames: int, + height: int, + width: int, + *, + h_pad: AxisPad | None = None, + w_pad: AxisPad | None = None, + spatial_scale: Tuple[int, int] = (1, 1), +) -> torch.Tensor: + """Crop padded decode output ``(B, C, F, H, W)`` back to the content shape. + + Temporal pad is always trailing. Spatial size-floor pads must pass the recorded + ``h_pad`` / ``w_pad`` (latent units) plus ``spatial_scale``; otherwise the crop is centered. + """ + x, _ = resize_axis(pixels, 2, frames, mode="repeat_last") + scale_h, scale_w = spatial_scale + if h_pad is not None: + x = x.narrow(3, scale_axis_pad(h_pad, scale_h).before, height).contiguous() + else: + x, _ = resize_axis(x, 3, height, mode="symmetric") + if w_pad is not None: + x = x.narrow(4, scale_axis_pad(w_pad, scale_w).before, width).contiguous() + else: + x, _ = resize_axis(x, 4, width, mode="symmetric") + return x + + +def stage5_pixel_shape_from_stage4( + stage4_t: int, + stage4_h: int, + stage4_w: int, + *, + upsample_stride: Tuple[int, int, int], + patch_size: int, + stage5_kernel_t: int, + drop_leading_frame: bool, + pad_trailing: bool, +) -> tuple[int, int, int]: + """Pixel ``(F, H, W)`` for a stage-4-input extent (one remaining NA hop + patch).""" + st, sh, sw = upsample_stride + frames = stage4_t * st - 1 if drop_leading_frame and st == 2 else stage4_t * st + if pad_trailing: + frames = max(frames, stage5_kernel_t) + return frames, stage4_h * sh * patch_size, stage4_w * sw * patch_size + + +def pad_trailing_latent_for_natten_border(latent: torch.Tensor, n_frames: int) -> torch.Tensor: + """Replicate the last latent frame ``n_frames`` times for the NA last-frame border.""" + if n_frames <= 0: + return latent + padded, _ = resize_axis(latent, 2, latent.shape[2] + n_frames, mode="repeat_last") + return padded + + +def crop_trailing_context_natten_pad( + context: torch.Tensor, + *, + n_latent_frames: int, + time_scale: int, + stage5_kernel_t: int, +) -> torch.Tensor: + """Crop the ghosting appendix before stage 5, leaving at least ``stage5_kernel_t``.""" + if n_latent_frames <= 0: + return context + ghost = n_latent_frames * time_scale + content_t = max(context.shape[1] - ghost, 1) + keep = min(context.shape[1], max(content_t, stage5_kernel_t)) + cropped, _ = resize_axis(context, 1, keep, mode="repeat_last") + return cropped + + +def weight_floor(dtype: torch.dtype) -> float: + """Smallest divisor that safely guards ``buffer / weights`` in ``dtype``.""" + return max(1e-8, torch.finfo(dtype).tiny) + + +def stage4_thw_from_latent( + upsample_strides: Sequence[Tuple[int, int, int]], + latent_t: int, + latent_h: int, + latent_w: int, + *, + drop_leading_frame: bool, +) -> Tuple[int, int, int]: + """Stage-4 input ``(T, H, W)`` after the first three upsample hops.""" + t, h, w = latent_t, latent_h, latent_w + for st, sh, sw in upsample_strides[:3]: + t, h, w = t * st, h * sh, w * sw + if st == 2 and drop_leading_frame: + t -= 1 + return t, h, w + + +def stage4_to_pixel_scale_factors(upsample_stride: Tuple[int, int, int], patch_size: int) -> SpatioTemporalScaleFactors: + """Pixel/frame units per stage-4-input cell (last NA hop + unpatchify).""" + st, sh, sw = upsample_stride + return SpatioTemporalScaleFactors(time=st, height=sh * patch_size, width=sw * patch_size) + + +def compute_tile_min_size( + stage4_kernel: Tuple[int, int, int], + stage5_kernel: Tuple[int, int, int], + upsample3_stride: Tuple[int, int, int], +) -> Tuple[int, int, int]: + """Min stage-4-input ``(T, H, W)`` so stages 4 and 5 each see ``>= kernel``.""" + return tuple(max(stage4_kernel[a], -(-stage5_kernel[a] // upsample3_stride[a])) for a in range(3)) + + +def compute_tile_halos( + stage4_kernel: Tuple[int, int, int], + stage4_depth: int, + stage5_kernel: Tuple[int, int, int], + stage5_depth: int, + upsample3_stride: Tuple[int, int, int], +) -> Tuple[Tuple[int, int, int], Tuple[int, int, int]]: + """One-sided halos in stage-4-input units for stages 4 and 5.""" + halo4 = tuple(stage4_depth * (stage4_kernel[a] // 2) for a in range(3)) + halo5 = tuple(-(-(stage5_depth * (stage5_kernel[a] // 2)) // upsample3_stride[a]) for a in range(3)) + return halo4, halo5 # type: ignore[return-value] + + +def all_stages_min_tile_size( + stage_kernels: Sequence[Tuple[int, int, int]], + upsamples: Sequence[Tuple[Tuple[int, int, int], int]], + stage5_kernel: Tuple[int, int, int], +) -> Tuple[int, int, int]: + """Per-axis latent-grid floor so every stage's NA sees dims ``>= kernel_size``.""" + cumulative = [(1, 1, 1)] + t, h, w = 1, 1, 1 + for stride, _ in upsamples: + t, h, w = t * stride[0], h * stride[1], w * stride[2] + cumulative.append((t, h, w)) + mins = [1, 1, 1] + for stage_i in range(len(upsamples)): + for axis in range(3): + mins[axis] = max(mins[axis], -(-stage_kernels[stage_i][axis] // cumulative[stage_i][axis])) + for axis in range(3): + mins[axis] = max(mins[axis], -(-stage5_kernel[axis] // cumulative[len(upsamples)][axis])) + return (mins[0], mins[1], mins[2]) + + +def pixel_tile_shape(full_shape: tuple[int, ...], out_coords: tuple[slice, ...]) -> tuple[int, ...]: + return tuple(len(range(*coord.indices(size))) for size, coord in zip(full_shape, out_coords, strict=True)) + + +def _round_up(value: int, multiple: int) -> int: + return -(-value // multiple) * multiple + + +def recommended_pixel_overlaps( + tile_halos: Tuple[Tuple[int, int, int], Tuple[int, int, int]], + pixel_scale: SpatioTemporalScaleFactors, +) -> Tuple[int, int]: + """Stage-4/5-safe ``(temporal_overlap_frames, spatial_overlap_pixels)``.""" + + def dominant(axis: int) -> int: + return max(tile_halos[i][axis] for i in range(len(tile_halos))) + + overlap_t = _round_up(dominant(0) * pixel_scale.time, 8) + halo_hw = max(dominant(1), dominant(2)) + overlap_hw = _round_up(halo_hw * pixel_scale.height, 32) + return overlap_t, overlap_hw + + +def stage5_tokens_for_pixel_tile(tile_frames: int, tile_height: int, tile_width: int, *, patch_size: int) -> int: + """Pre-unpatchify stage-5 token count for a pixel-space tile.""" + return tile_frames * max(1, tile_height // patch_size) * max(1, tile_width // patch_size) + + +def _axis_candidates(length: int, overlap: int, min_size: int, multiple: int) -> list[tuple[int, int]]: + """``(tile_size, num_tiles)`` for every legal size on ``multiple``'s grid.""" + out: list[tuple[int, int]] = [] + max_size = max(_round_up(length, multiple), min_size) + for size in range(min_size, max_size + multiple, multiple): + if size <= overlap: + continue + out.append((size, len(split_by_size(length, size, overlap, None)))) + return out + + +def _propagate_interval_through_upsample_hops( + interval: DimensionInterval, + strides: Sequence[int], + causal: bool, +) -> DimensionInterval: + """Forward-propagate one interval through upsample hops on one axis. + + Multiply by ``stride``; on the causal temporal axis with ``stride == 2`` drop the + duplicate frame (``end -= 1``; non-origin tiles also ``start -= 1``). + """ + x = interval + for stride in strides: + start = x.start * stride + end = x.end * stride + if causal and stride == 2: + end -= 1 + if x.start != 0: + start -= 1 + x = DimensionInterval(start=start, end=end, left_ramp=x.left_ramp * stride, right_ramp=x.right_ramp * stride) + return x diff --git a/jasna/models/ltx_vae/diffusion_video_decoder.py b/jasna/models/ltx_vae/diffusion_video_decoder.py new file mode 100644 index 00000000..7c3abf3b --- /dev/null +++ b/jasna/models/ltx_vae/diffusion_video_decoder.py @@ -0,0 +1,380 @@ +"""Single-step x0 diffusion video VAE decoder (``NADiffusionDecoder``), ``chunked_eager`` mode. + +Ported from ``ltx_core/model/video_vae/diffusion_video_decoder.py`` and +``ltx_core/model/transformer/timestep_embedding.py``. +""" + +from __future__ import annotations + +import math +from typing import Iterator, List, Tuple + +import torch +from torch import nn + +from jasna.models.ltx_vae import diffusion_tiling +from jasna.models.ltx_vae.ops import PerChannelStatistics, patchify, unpatchify +from jasna.models.ltx_vae.tiling import ( + SpatioTemporalScaleFactors, + Tile, + TileSizeConfig, + VideoLatentShape, + group_tiles_by_temporal_slice, + masks_are_complementary, + scale_by_masks_1d, +) +from jasna.models.ltx_vae.transformer import ( + AdaLNZero, + ChunkedDiffusionNABlock, + LinearPixelShuffleUpsample, + NABlock, +) + +_TIMESTEP_PROJ_DIM = 256 + + +def get_timestep_embedding(timesteps: torch.Tensor) -> torch.Tensor: + """Sinusoidal ``[cos, sin]`` embedding of 1-D ``timesteps`` (256 dims, max period 10000).""" + half_dim = _TIMESTEP_PROJ_DIM // 2 + exponent = -math.log(10000) * torch.arange(start=0, end=half_dim, dtype=torch.float32, device=timesteps.device) + exponent = exponent / half_dim + emb = timesteps[:, None].float() * torch.exp(exponent)[None, :] + return torch.cat([torch.cos(emb), torch.sin(emb)], dim=-1) + + +class TimestepEmbedding(nn.Module): + def __init__(self, in_channels: int, time_embed_dim: int) -> None: + super().__init__() + self.linear_1 = nn.Linear(in_channels, time_embed_dim, bias=True) + self.act = nn.SiLU() + self.linear_2 = nn.Linear(time_embed_dim, time_embed_dim, bias=True) + + def forward(self, sample: torch.Tensor) -> torch.Tensor: + return self.linear_2(self.act(self.linear_1(sample))) + + +class PixArtAlphaCombinedTimestepSizeEmbeddings(nn.Module): + def __init__(self, embedding_dim: int) -> None: + super().__init__() + self.timestep_embedder = TimestepEmbedding(in_channels=_TIMESTEP_PROJ_DIM, time_embed_dim=embedding_dim) + + def forward(self, timestep: torch.Tensor, hidden_dtype: torch.dtype) -> torch.Tensor: + return self.timestep_embedder(get_timestep_embedding(timestep).to(dtype=hidden_dtype)) + + +class DiffusionVideoDecoder(nn.Module): + """Diffusion video VAE decoder with a neighborhood-attention backbone. + + Stages 1-4 deterministically upsample the latent into a context volume; stage 5 runs + ``ChunkedDiffusionNABlock``s that predict clean patchified pixels (x0) from per-tile + noise in one step, guided by that context. The last latent frame is replicated + ``(stage1_K_t // 2) * 2`` times through stages 1-4 against the NA last-frame border and + cropped off before stage 5. Latents below ``stage_min_tile_sizes`` are edge-padded first; + leftover pad is cropped from the final pixels. + """ + + def __init__( + self, + *, + in_channels: int, + out_channels: int, + patch_size: int, + head_dim: int, + stage_channels: Tuple[int, ...], + stage_depths: Tuple[int, ...], + stage_kernels: Tuple[Tuple[int, int, int], ...], + upsamples: Tuple[Tuple[Tuple[int, int, int], int], ...], + stage5_kernel: Tuple[int, int, int], + stage5_channels: int, + t_emb_dim: int, + timestep_scale_multiplier: float, + ) -> None: + super().__init__() + self.patch_size = patch_size + self.out_channels = out_channels + self.stage_channels = stage_channels + self.video_downscale_factors = SpatioTemporalScaleFactors.default() + self.stage5_kernel = tuple(stage5_kernel) + self._natten_trailing_pad_latent_frames = (stage_kernels[0][0] // 2) * 2 + + self.per_channel_statistics = PerChannelStatistics(latent_channels=in_channels) + self.conv_in = nn.Linear(in_channels, stage_channels[0], bias=True) + + self.det_stages = nn.ModuleList() + self.upsamples = nn.ModuleList() + for stage_i in range(len(stage_channels) - 1): + c = stage_channels[stage_i] + self.det_stages.append( + nn.ModuleList( + [ + NABlock(dim=c, kernel_size=stage_kernels[stage_i], head_dim=head_dim) + for _ in range(stage_depths[stage_i]) + ] + ) + ) + stride, reduction = upsamples[stage_i] + self.upsamples.append( + LinearPixelShuffleUpsample(in_channels=c, stride=stride, out_channels_reduction_factor=reduction) + ) + + self.t_embedder = PixArtAlphaCombinedTimestepSizeEmbeddings(embedding_dim=t_emb_dim) + + self.context_channels = stage_channels[-1] + noised_pixel_channels = out_channels * (patch_size**2) + + self.stage_min_tile_sizes = diffusion_tiling.all_stages_min_tile_size(stage_kernels, upsamples, stage5_kernel) + up3_stride = upsamples[3][0] + self.tile_min_sizes = diffusion_tiling.compute_tile_min_size(stage_kernels[3], stage5_kernel, up3_stride) + self.tile_halos = diffusion_tiling.compute_tile_halos( + stage_kernels[3], stage_depths[3], stage5_kernel, stage_depths[-1], up3_stride + ) + self.conv_in_x_t = nn.Linear(noised_pixel_channels, stage5_channels, bias=True) + self.shared_adaln = AdaLNZero(dim=stage5_channels, t_emb_dim=t_emb_dim) + self.diff_blocks = nn.ModuleList( + [ + ChunkedDiffusionNABlock( + dim=stage5_channels, + kernel_size=stage5_kernel, + context_channels=self.context_channels, + head_dim=head_dim, + ) + for _ in range(stage_depths[-1]) + ] + ) + self.norm_out = nn.RMSNorm(stage5_channels, eps=1e-6) + self.conv_out = nn.Linear(stage5_channels, noised_pixel_channels, bias=True) + self.timestep_scale_multiplier = timestep_scale_multiplier + + def recommended_tiling_config(self, *, height: int, width: int, num_frames: int, free_bytes: int) -> TileSizeConfig: + """Stage-4/5 overlaps + tile sizes that fit ``free_bytes`` of activations.""" + return diffusion_tiling.recommended_decode_tiling_config( + tile_halos=self.tile_halos, + pixel_scale=diffusion_tiling.stage4_to_pixel_scale_factors(tuple(self.upsamples[3].stride), self.patch_size), + min_tile_size_s4=self.tile_min_sizes, + patch_size=self.patch_size, + height=height, + width=width, + num_frames=num_frames, + free_bytes=free_bytes, + stage5_channels=self.conv_in_x_t.out_features, + stage4_channels=self.stage_channels[3], + upsample_strides=tuple(tuple(u.stride) for u in self.upsamples), + model_bytes=sum(p.numel() * p.element_size() for p in self.parameters()), + natten_trailing_pad_latent_frames=self._natten_trailing_pad_latent_frames, + out_channels=self.out_channels, + ) + + def _run_det_stage(self, x: torch.Tensor, stage_i: int, drop_leading_frame: bool) -> torch.Tensor: + for block in self.det_stages[stage_i]: + x = block(x) + return self.upsamples[stage_i](x, drop_leading_frame=drop_leading_frame) + + def forward_stages_1_to_3(self, z_noisy: torch.Tensor) -> torch.Tensor: + """Stages 1-3 on the full (ghost-padded) latent -> channels-last stage-4 input feature.""" + z_noisy = self.per_channel_statistics.un_normalize(z_noisy) + x = self.conv_in(z_noisy.permute(0, 2, 3, 4, 1)) + for stage_i in range(3): + x = self._run_det_stage(x, stage_i, drop_leading_frame=True) + return x + + def forward_stage_4(self, x: torch.Tensor, pad_trailing: bool) -> torch.Tensor: + """Stage-4 NA blocks on one tile (the upsample is deferred into the diffusion blocks).""" + for block in self.det_stages[3]: + x = block(x) + if pad_trailing: + up_t = int(self.upsamples[3].stride[0]) + x = diffusion_tiling.crop_trailing_context_natten_pad( + x, + n_latent_frames=self._natten_trailing_pad_latent_frames, + time_scale=self.video_downscale_factors.time // up_t, + stage5_kernel_t=max(1, -(-self.stage5_kernel[0] // up_t)), + ) + return x + + def forward_diff_step( + self, x_t: torch.Tensor, stage4_feat: torch.Tensor, t: torch.Tensor, *, drop_leading_frame: bool + ) -> torch.Tensor: + """The stage-5 x0 prediction in pixel space from noise ``x_t`` and stage-4 context.""" + x = self.conv_in_x_t(patchify(x_t, patch_size_hw=self.patch_size).permute(0, 2, 3, 4, 1)) + t_emb = self.t_embedder(self.timestep_scale_multiplier * t, hidden_dtype=x.dtype) + modulation = self.shared_adaln(t_emb) + for block in self.diff_blocks: + x = block(x, stage4_feat, modulation, self.upsamples[3], drop_leading_frame=drop_leading_frame) + x = self.norm_out(x) + x = self.conv_out(x) + x = x.permute(0, 4, 1, 2, 3).contiguous() + return unpatchify(x, patch_size_hw=self.patch_size) + + def _decode_temporal_group_isolated( + self, + tiles: List[Tile], + feat_s4: torch.Tensor, + content_s4_frames: int, + timestep: torch.Tensor, + full_video_shape: VideoLatentShape, + curr_temporal_slice: slice, + generator: torch.Generator | None, + *, + complementary: bool, + ) -> Tuple[torch.Tensor, torch.Tensor | None]: + """Decode every tile of one temporal group in isolation and blend into an fp16 accumulator.""" + group_temporal_len = curr_temporal_slice.stop - curr_temporal_slice.start + group_shape = full_video_shape._replace(frames=group_temporal_len) + full_torch_shape = full_video_shape.to_torch_shape() + accum_dtype = torch.float16 if feat_s4.dtype == torch.bfloat16 else feat_s4.dtype + buffer = torch.zeros(group_shape.to_torch_shape(), device=feat_s4.device, dtype=accum_dtype) + weights: torch.Tensor | None = None if complementary else torch.zeros_like(buffer) + local_temporal_slice = slice(0, group_temporal_len) + + randn_device = generator.device if generator is not None else feat_s4.device + up3_stride = tuple(self.upsamples[3].stride) + + for tile in tiles: + feat_tile, is_origin, pad_trailing, content_thw = diffusion_tiling.slice_stage4_tile( + feat_s4, tile, content_frames=content_s4_frames + ) + content_pixel_shape = diffusion_tiling.pixel_tile_shape(full_torch_shape, tile.out_coords) + stage5_f, stage5_h, stage5_w = diffusion_tiling.stage5_pixel_shape_from_stage4( + *content_thw, + upsample_stride=up3_stride, + patch_size=self.patch_size, + stage5_kernel_t=self.stage5_kernel[0], + drop_leading_frame=is_origin, + pad_trailing=pad_trailing, + ) + x_t = torch.randn( + (content_pixel_shape[0], content_pixel_shape[1], stage5_f, stage5_h, stage5_w), + dtype=feat_s4.dtype, + generator=generator, + device=randn_device, + ).to(feat_s4.device) + + context_tile = self.forward_stage_4(feat_tile, pad_trailing=pad_trailing) + pixel_tile = self.forward_diff_step(x_t, context_tile, timestep, drop_leading_frame=is_origin) + pixel_tile = diffusion_tiling.crop_pixels_to_content( + pixel_tile, content_pixel_shape[2], content_pixel_shape[3], content_pixel_shape[4] + ).to(buffer.dtype) + + masks = tuple(m.to(device=buffer.device, dtype=torch.float32) for m in tile.masks_1d) + local_coords = ( + tile.out_coords[0], + tile.out_coords[1], + local_temporal_slice, + tile.out_coords[3], + tile.out_coords[4], + ) + buffer[local_coords] += scale_by_masks_1d(pixel_tile, masks) + if weights is not None: + strength = torch.ones(pixel_tile.shape, device=buffer.device, dtype=buffer.dtype) + weights[local_coords] += scale_by_masks_1d(strength, masks) + + return buffer, weights + + def tiled_decode( + self, + latent: torch.Tensor, + tiling_config: TileSizeConfig, + generator: torch.Generator | None, + ) -> Iterator[torch.Tensor]: + """Decode latent to ``(B, C, F, H, W)`` pixels in ~[-1, 1], yielding temporal chunks. + + Stages 1-3 run once on the full volume; stages 4-5 run per tile with pixel blend. + Across temporal groups only the trailing overlap is kept between iterations. + """ + content_shape = VideoLatentShape.from_torch_shape(latent.shape) + content_pixel = content_shape.upscale(self.video_downscale_factors)._replace(channels=self.out_channels) + + latent, (_t_pad, h_pad, w_pad) = diffusion_tiling.ensure_min_latent_shape(latent, self.stage_min_tile_sizes) + spatial_scale = (self.video_downscale_factors.height, self.video_downscale_factors.width) + work_shape = VideoLatentShape.from_torch_shape(latent.shape) + full_video_shape = work_shape.upscale(self.video_downscale_factors)._replace(channels=self.out_channels) + target_shape = full_video_shape.to_torch_shape() + + strides = [tuple(u.stride) for u in self.upsamples] + s4_t, s4_h, s4_w = diffusion_tiling.stage4_thw_from_latent( + strides, latent.shape[2], latent.shape[3], latent.shape[4], drop_leading_frame=True + ) + tiles = diffusion_tiling.prepare_tile_schedule( + torch.Size([latent.shape[0], latent.shape[1], s4_t, s4_h, s4_w]), + tiling_config, + upsample3_stride=tuple(self.upsamples[3].stride), + patch_size=self.patch_size, + min_tile_size=self.tile_min_sizes, + ) + + latent_padded = diffusion_tiling.pad_trailing_latent_for_natten_border( + latent, self._natten_trailing_pad_latent_frames + ) + feat_s4 = self.forward_stages_1_to_3(latent_padded) + timestep = torch.ones(latent.shape[0], dtype=torch.float32, device=latent.device) + + complementary = masks_are_complementary(tiles, target_shape) + groups = group_tiles_by_temporal_slice(tiles) + group_slices = [slice(*group[0].out_coords[2].indices(target_shape[2])[:2]) for group in groups] + + overlap_stub: torch.Tensor | None = None + overlap_stub_weights: torch.Tensor | None = None + + def _finalize(buf: torch.Tensor, wts: torch.Tensor | None) -> torch.Tensor: + if complementary: + return buf.to(latent.dtype) + wts = wts.clamp(min=diffusion_tiling.weight_floor(wts.dtype)) + return (buf / wts).to(latent.dtype) + + def _crop_emit(buf: torch.Tensor, wts: torch.Tensor | None, global_start: int) -> torch.Tensor | None: + if global_start >= content_pixel.frames or buf.shape[2] < 1: + return None + frames_keep = min(buf.shape[2], content_pixel.frames - global_start) + chunk = _finalize(buf[:, :, :frames_keep], None if wts is None else wts[:, :, :frames_keep]) + return diffusion_tiling.crop_pixels_to_content( + chunk, + frames_keep, + content_pixel.height, + content_pixel.width, + h_pad=h_pad, + w_pad=w_pad, + spatial_scale=spatial_scale, + ) + + for gi, group in enumerate(groups): + curr_temporal_slice = group_slices[gi] + buffer, weights = self._decode_temporal_group_isolated( + group, + feat_s4, + s4_t, + timestep, + full_video_shape, + curr_temporal_slice, + generator, + complementary=complementary, + ) + + if overlap_stub is not None: + overlap_len = int(overlap_stub.shape[2]) + if overlap_len > 0: + overlap_stub += buffer[:, :, :overlap_len] + buffer[:, :, :overlap_len] = overlap_stub + if not complementary: + overlap_stub_weights += weights[:, :, :overlap_len] + weights[:, :, :overlap_len] = overlap_stub_weights + overlap_stub = None + overlap_stub_weights = None + + if gi + 1 < len(groups): + next_start = group_slices[gi + 1].start + exclusive_len = min(max(0, next_start - curr_temporal_slice.start), buffer.shape[2]) + emitted = _crop_emit( + buffer[:, :, :exclusive_len], + None if weights is None else weights[:, :, :exclusive_len], + curr_temporal_slice.start, + ) + if emitted is not None: + yield emitted + overlap_stub = buffer[:, :, exclusive_len:].clone() + if not complementary: + overlap_stub_weights = weights[:, :, exclusive_len:].clone() + del buffer, weights + else: + emitted = _crop_emit(buffer, weights, curr_temporal_slice.start) + if emitted is not None: + yield emitted diff --git a/jasna/models/ltx_vae/eager_na.py b/jasna/models/ltx_vae/eager_na.py new file mode 100644 index 00000000..05b4783f --- /dev/null +++ b/jasna/models/ltx_vae/eager_na.py @@ -0,0 +1,148 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025 Comfy Org. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Limited-workspace 3D neighborhood attention (NATTEN ``na3d`` semantics) in pure torch. + +Vendored via ltx-core ``fallback_na/eager.py`` from comfy-kitchen ``backends/eager/na.py`` +(non-causal, pre-scaled Q only). Queries are tiled; tiles that share window geometry +stack into batched ``scaled_dot_product_attention`` calls with one additive mask per group. +""" + +from __future__ import annotations + +import math + +import torch +from torch.nn import functional + +# Element budget for one tile's [Nq, Nk] attention mask (bounds the mask +# allocation and, on CPU, the math-backend score materialization). +NA_SCORE_BUDGET = 2**25 +# Element budget for the stacked K/V copies of one batched SDPA call on CUDA. +NA_KV_STACK_BUDGET = 2**28 + + +def _window_bounds(length: int, kernel: int) -> tuple[list[int], list[int]]: + """Per-index (start, end) of the attended window along one axis (NATTEN shifts it inward at borders).""" + starts: list[int] = [] + ends: list[int] = [] + kernel = min(kernel, length) + lo = length - kernel + half = kernel // 2 + for i in range(length): + start = min(max(i - half, 0), lo) + starts.append(start) + ends.append(start + kernel) + return starts, ends + + +def _pick_tiles(dims: tuple[int, int, int], kernels: list[int]) -> list[int]: + """Per-axis query-tile lengths keeping one tile's [Nq, Nk] under budget.""" + tiles = list(dims) + + def cost(ts: list[int]) -> int: + nq = math.prod(ts) + nk = math.prod(min(d, t + k - 1) for t, k, d in zip(ts, kernels, dims, strict=True)) + return nq * nk + + while cost(tiles) > NA_SCORE_BUDGET and max(tiles) > 1: + i = max(range(3), key=lambda a: tiles[a] / kernels[a]) + if tiles[i] <= 1: + break + tiles[i] = max(1, (tiles[i] + 1) // 2) + return tiles + + +def _group_mask( + rel_bounds: tuple[tuple[tuple[int, ...], tuple[int, ...]], ...], + dtype: torch.dtype, + device: torch.device, +) -> torch.Tensor: + """Additive ``[1, 1, Nq, Nk]`` mask for one tile-geometry group.""" + bools = [] + for starts, ends in rel_bounds: + st = torch.tensor(starts, device=device) + en = torch.tensor(ends, device=device) + kj = torch.arange(int(en.max()), device=device) + bools.append((kj[None, :] >= st[:, None]) & (kj[None, :] < en[:, None])) + visible = ( + bools[0][:, None, None, :, None, None] + & bools[1][None, :, None, None, :, None] + & bools[2][None, None, :, None, None, :] + ) + nq = visible.shape[0] * visible.shape[1] * visible.shape[2] + nk = visible.shape[3] * visible.shape[4] * visible.shape[5] + mask = torch.zeros((nq, nk), dtype=dtype, device=device) + mask.masked_fill_(~visible.reshape(nq, nk), torch.finfo(dtype).min) + return mask.reshape(1, 1, nq, nk) + + +def na3d(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, kernel_size: tuple[int, int, int]) -> torch.Tensor: + """3D neighborhood attention over ``(B, T, H, W, NH, HD)`` tensors; Q must already be scaled.""" + batch, t, h, w, nh, hd = q.shape + dims = (t, h, w) + kernels = [min(k_, d) for k_, d in zip(kernel_size, dims, strict=True)] + device = q.device + + bounds = [_window_bounds(d, k_) for d, k_ in zip(dims, kernels, strict=True)] + tile_t, tile_h, tile_w = _pick_tiles(dims, kernels) + + groups: dict[ + tuple[ + tuple[tuple[int, ...], tuple[int, ...]], + tuple[tuple[int, ...], tuple[int, ...]], + tuple[tuple[int, ...], tuple[int, ...]], + ], + list[tuple[tuple[slice, slice, slice], tuple[slice, slice, slice]]], + ] = {} + for t0 in range(0, t, tile_t): + t1 = min(t0 + tile_t, t) + rt0, rt1 = bounds[0][0][t0], bounds[0][1][t1 - 1] + rel_t = ( + tuple(s - rt0 for s in bounds[0][0][t0:t1]), + tuple(e - rt0 for e in bounds[0][1][t0:t1]), + ) + for h0 in range(0, h, tile_h): + h1 = min(h0 + tile_h, h) + rh0, rh1 = bounds[1][0][h0], bounds[1][1][h1 - 1] + rel_h = ( + tuple(s - rh0 for s in bounds[1][0][h0:h1]), + tuple(e - rh0 for e in bounds[1][1][h0:h1]), + ) + for w0 in range(0, w, tile_w): + w1 = min(w0 + tile_w, w) + rw0, rw1 = bounds[2][0][w0], bounds[2][1][w1 - 1] + rel_w = ( + tuple(s - rw0 for s in bounds[2][0][w0:w1]), + tuple(e - rw0 for e in bounds[2][1][w0:w1]), + ) + groups.setdefault((rel_t, rel_h, rel_w), []).append( + ( + (slice(t0, t1), slice(h0, h1), slice(w0, w1)), + (slice(rt0, rt1), slice(rh0, rh1), slice(rw0, rw1)), + ) + ) + + out = torch.empty((batch, t, h, w, nh, hd), device=device, dtype=v.dtype) + for rel, tiles in groups.items(): + mask = _group_mask(rel, q.dtype, device) + nq, nk = mask.shape[2], mask.shape[3] + g_max = max(1, NA_KV_STACK_BUDGET // max(1, batch * nh * nk * hd * 2)) if device.type == "cuda" else 1 + qs0, _ = tiles[0] + tq = qs0[0].stop - qs0[0].start + th = qs0[1].stop - qs0[1].start + tw = qs0[2].stop - qs0[2].start + for c0 in range(0, len(tiles), g_max): + chunk = tiles[c0 : c0 + g_max] + g = len(chunk) + q_s = torch.stack([q[:, qs[0], qs[1], qs[2]] for qs, _ in chunk]) + k_s = torch.stack([k[:, rs[0], rs[1], rs[2]] for _, rs in chunk]) + v_s = torch.stack([v[:, rs[0], rs[1], rs[2]] for _, rs in chunk]) + q_s = q_s.permute(0, 1, 5, 2, 3, 4, 6).reshape(g * batch, nh, nq, hd) + k_s = k_s.permute(0, 1, 5, 2, 3, 4, 6).reshape(g * batch, nh, nk, hd) + v_s = v_s.permute(0, 1, 5, 2, 3, 4, 6).reshape(g * batch, nh, nk, hd) + o = functional.scaled_dot_product_attention(q_s, k_s, v_s, attn_mask=mask, scale=1.0) + o = o.view(g, batch, nh, tq, th, tw, hd).permute(0, 1, 3, 4, 5, 2, 6) + for i, (qs, _) in enumerate(chunk): + out[:, qs[0], qs[1], qs[2]] = o[i] + + return out diff --git a/jasna/models/ltx_vae/loader.py b/jasna/models/ltx_vae/loader.py new file mode 100644 index 00000000..eeab3d8c --- /dev/null +++ b/jasna/models/ltx_vae/loader.py @@ -0,0 +1,174 @@ +"""Build the encoder / decoder from a ``CausalDiffusionVAE`` checkpoint (its metadata config + weights).""" + +from __future__ import annotations + +import json +from collections.abc import Callable +from pathlib import Path +from typing import TypeVar + +import safetensors +import torch +from safetensors.torch import load_file +from torch import nn + +from jasna.models.ltx_vae.diffusion_video_decoder import DiffusionVideoDecoder +from jasna.models.ltx_vae.video_encoder import VideoEncoder + +_VAE_CLASS_NAME = "CausalDiffusionVAE" +_DECODER_CLASS_NAME = "NADiffusionDecoder" +_GATE_SUFFIXES = (".gate_msa", ".gate_mlp", ".gate_ctx") +_QKV_LEAVES = ("weight", "bias") +_UNUSED_DECODER_KEYS = frozenset({"type_emb"}) + +_Model = TypeVar("_Model", bound=nn.Module) + + +def read_vae_config(vae_path: Path) -> dict: + with safetensors.safe_open(str(vae_path), framework="pt") as handle: + config = json.loads(handle.metadata()["config"])["vae"] + if config.get("_class_name") != _VAE_CLASS_NAME: + raise ValueError(f"{vae_path}: VAE class {config.get('_class_name')!r} is not supported") + return config + + +def _require(config: dict, key: str, expected: object, section: str) -> None: + if config.get(key) != expected: + raise ValueError(f"VAE {section} config {key}={config.get(key)!r} is not supported (expected {expected!r})") + + +def encoder_kwargs(config: dict) -> dict: + """``VideoEncoder`` kwargs from a ``CausalDiffusionVAE`` config; rejects unported variants.""" + encoder = config["encoder"] + _require(encoder, "dims", 3, "encoder") + _require(encoder, "norm_layer", "pixel_norm", "encoder") + _require(encoder, "latent_log_var", "constant", "encoder") + _require(encoder, "spatial_padding_mode", "zeros", "encoder") + return { + "in_channels": encoder["in_channels"], + "out_channels": encoder["out_channels"], + "encoder_blocks": encoder["blocks"], + "patch_size": encoder["patch_size"], + } + + +def decoder_kwargs(config: dict) -> dict: + """``DiffusionVideoDecoder`` kwargs from a ``CausalDiffusionVAE`` config; rejects unported variants.""" + decoder = config["decoder"] + _require(decoder, "_class_name", _DECODER_CLASS_NAME, "decoder") + _require(config, "model_output_type", "x0", "top-level") + _require(decoder, "default_num_inference_steps", 1, "decoder") + _require(decoder, "resampler_kind", "linear", "decoder") + _require(decoder, "spatial_padding_mode", "zeros", "decoder") + stage_channels = tuple(decoder["stage_channels"]) + return { + "in_channels": decoder["in_channels"], + "out_channels": decoder["out_channels"], + "patch_size": decoder["patch_size"], + "head_dim": decoder["head_dim"], + "stage_channels": stage_channels, + "stage_depths": tuple(decoder["stage_depths"]), + "stage_kernels": tuple(tuple(kernel) for kernel in decoder["stage_kernels"]), + "upsamples": tuple((tuple(stride), reduction) for stride, reduction in decoder["upsamples"]), + "stage5_kernel": tuple(decoder["stage5_kernel"]), + "stage5_channels": decoder.get("stage5_channels", stage_channels[-1]), + "t_emb_dim": decoder.get("t_emb_dim", 384), + "timestep_scale_multiplier": decoder["timestep_scale_multiplier"], + } + + +def _read_tensors(path: Path, prefixes: tuple[str, ...]) -> dict[str, torch.Tensor]: + with safetensors.safe_open(str(path), framework="pt", device="cpu") as handle: + return {key: handle.get_tensor(key) for key in handle.keys() if key.startswith(prefixes)} + + +def encoder_state_dict(checkpoint: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]: + """Checkpoint ``encoder.*`` / ``per_channel_statistics.*`` tensors -> ``VideoEncoder`` keys.""" + return { + key.removeprefix("encoder."): value + for key, value in checkpoint.items() + if key.startswith(("encoder.", "per_channel_statistics.")) + } + + +def decoder_state_dict(checkpoint: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]: + """Checkpoint ``decoder.*`` / ``per_channel_statistics.*`` tensors -> ``DiffusionVideoDecoder`` keys. + + Strips ``decoder.``, renames ``t_embedder.mlp.{0,2}`` to ``t_embedder.timestep_embedder.linear_{1,2}``, + splits fused ``qkv`` into ``to_q`` / ``to_k`` / ``to_v`` and drops the ``coarse_*`` preview head + and the unused ``type_emb``. + """ + out: dict[str, torch.Tensor] = {} + for key, value in checkpoint.items(): + if key.startswith("per_channel_statistics."): + out[key] = value + continue + if not key.startswith("decoder."): + continue + key = key.removeprefix("decoder.") + if key.startswith("coarse_") or key in _UNUSED_DECODER_KEYS: + continue + if key.endswith(_GATE_SUFFIXES): + raise ValueError(f"gated decoder checkpoints are not supported (found {key!r})") + key = key.replace("t_embedder.mlp.0.", "t_embedder.timestep_embedder.linear_1.") + key = key.replace("t_embedder.mlp.2.", "t_embedder.timestep_embedder.linear_2.") + leaf = key.rsplit(".", 1)[-1] + if key.endswith(".qkv." + leaf) and leaf in _QKV_LEAVES: + prefix = key[: -len(leaf)] + for name, part in zip(("to_q", "to_k", "to_v"), value.chunk(3), strict=True): + out[f"{prefix}{name}.{leaf}"] = part.clone() + continue + out[key] = value + return out + + +def _build_bf16(factory: Callable[[], _Model], state: dict[str, torch.Tensor], device: torch.device) -> _Model: + with torch.device("meta"): + model = factory() + model.load_state_dict({k: v.to(torch.bfloat16) for k, v in state.items()}, strict=True, assign=True) + return model.to(device).eval().requires_grad_(False) + + +def load_video_encoder(vae_path: Path, device: torch.device) -> VideoEncoder: + """bf16 eval ``VideoEncoder``; call it on ``(B, 3, 1+8k, H, W)`` pixels in [-1, 1].""" + kwargs = encoder_kwargs(read_vae_config(vae_path)) + state = encoder_state_dict(_read_tensors(vae_path, ("encoder.", "per_channel_statistics."))) + return _build_bf16(lambda: VideoEncoder(**kwargs), state, device) + + +def load_video_decoder(vae_path: Path, tuned_decoder_path: Path, device: torch.device) -> DiffusionVideoDecoder: + """bf16 eval ``DiffusionVideoDecoder`` with the fine-tuned decoder tensors overlaid.""" + kwargs = decoder_kwargs(read_vae_config(vae_path)) + state = decoder_state_dict(_read_tensors(vae_path, ("decoder.", "per_channel_statistics."))) + decoder = _build_bf16(lambda: DiffusionVideoDecoder(**kwargs), state, device) + + tuned = load_file(str(tuned_decoder_path)) + own = decoder.state_dict() + unknown = sorted(set(tuned) - set(own)) + if unknown: + raise KeyError(f"{tuned_decoder_path}: {len(unknown)} tensors not in the decoder, e.g. {unknown[:3]}") + decoder.load_state_dict({k: v.to(own[k].dtype) for k, v in tuned.items()}, strict=False) + return decoder + + +@torch.inference_mode() +def decode_latent( + decoder: DiffusionVideoDecoder, + latent: torch.Tensor, + free_bytes: int, + generator: torch.Generator | None, +) -> torch.Tensor: + """Tiled decode of a normalized ``(B, 128, T, h, w)`` latent to ``(B, 3, F, H, W)`` pixels in ~[-1, 1]. + + The tile layout is chosen from ``free_bytes`` (activation budget); raises ``ValueError`` + when even the minimum tile does not fit. + """ + _b, _c, frames, height, width = latent.shape + scale = decoder.video_downscale_factors + tiling = decoder.recommended_tiling_config( + height=height * scale.height, + width=width * scale.width, + num_frames=(frames - 1) * scale.time + 1, + free_bytes=free_bytes, + ) + return torch.cat(list(decoder.tiled_decode(latent, tiling, generator)), dim=2) diff --git a/jasna/models/ltx_vae/ops.py b/jasna/models/ltx_vae/ops.py new file mode 100644 index 00000000..bf04a331 --- /dev/null +++ b/jasna/models/ltx_vae/ops.py @@ -0,0 +1,50 @@ +"""Shared VAE ops: patchify, per-channel latent statistics, PixelNorm.""" + +import torch +from torch import nn + + +def patchify(x: torch.Tensor, patch_size_hw: int) -> torch.Tensor: + """Space-to-depth ``(B, C, F, H, W) -> (B, C*p*p, F, H/p, W/p)``; channel order ``(c, w-patch, h-patch)``.""" + b, c, f, h, w = x.shape + p = patch_size_hw + x = x.reshape(b, c, f, h // p, p, w // p, p).permute(0, 1, 6, 4, 2, 3, 5) + return x.reshape(b, c * p * p, f, h // p, w // p) + + +def unpatchify(x: torch.Tensor, patch_size_hw: int) -> torch.Tensor: + """Inverse of :func:`patchify`.""" + b, cpp, f, h, w = x.shape + p = patch_size_hw + x = x.reshape(b, cpp // (p * p), p, p, f, h, w).permute(0, 1, 4, 5, 3, 6, 2) + return x.reshape(b, cpp // (p * p), f, h * p, w * p) + + +class PerChannelStatistics(nn.Module): + """Dataset per-channel latent statistics used to (un)normalize latents.""" + + def __init__(self, latent_channels: int): + super().__init__() + self.register_buffer("std-of-means", torch.ones(latent_channels)) + self.register_buffer("mean-of-means", torch.zeros(latent_channels)) + + def un_normalize(self, x: torch.Tensor) -> torch.Tensor: + return (x * self.get_buffer("std-of-means").view(1, -1, 1, 1, 1).to(x)) + self.get_buffer("mean-of-means").view( + 1, -1, 1, 1, 1 + ).to(x) + + def normalize(self, x: torch.Tensor) -> torch.Tensor: + return (x - self.get_buffer("mean-of-means").view(1, -1, 1, 1, 1).to(x)) / self.get_buffer("std-of-means").view( + 1, -1, 1, 1, 1 + ).to(x) + + +class PixelNorm(nn.Module): + """Per-location RMS normalization over the channel dimension.""" + + eps = 1e-8 + + def forward(self, x: torch.Tensor) -> torch.Tensor: + mean_sq = torch.mean(x**2, dim=1, keepdim=True) + rms = torch.sqrt(mean_sq + self.eps) + return x / rms diff --git a/jasna/models/ltx_vae/tiling.py b/jasna/models/ltx_vae/tiling.py new file mode 100644 index 00000000..b55f9a0d --- /dev/null +++ b/jasna/models/ltx_vae/tiling.py @@ -0,0 +1,230 @@ +"""Video-grid tiling primitives (subset of ``ltx_core/tiling.py`` and ``ltx_core/types.py``).""" + +from __future__ import annotations + +from dataclasses import dataclass, replace +from typing import NamedTuple, Sequence + +import torch + + +class SpatioTemporalScaleFactors(NamedTuple): + """Spatiotemporal downscaling between decoded video space and the VAE latent grid.""" + + time: int + height: int + width: int + + @classmethod + def default(cls) -> SpatioTemporalScaleFactors: + return cls(time=8, height=32, width=32) + + @classmethod + def from_blocks(cls, blocks: list, patch_size: int) -> SpatioTemporalScaleFactors: + """Each ``compress_*`` block halves its target axes; patchify adds ``patch_size`` spatially.""" + spatial_steps = 0 + temporal_steps = 0 + for block_name, _ in blocks: + if block_name.startswith(("compress_space", "compress_all")): + spatial_steps += 1 + if block_name.startswith(("compress_time", "compress_all")): + temporal_steps += 1 + spatial = patch_size * (2**spatial_steps) + return cls(time=2**temporal_steps, height=spatial, width=spatial) + + +VIDEO_SCALE_FACTORS = SpatioTemporalScaleFactors.default() + + +class VideoLatentShape(NamedTuple): + """``(batch, channels, frames, height, width)`` of a video latent.""" + + batch: int + channels: int + frames: int + height: int + width: int + + def to_torch_shape(self) -> torch.Size: + return torch.Size([self.batch, self.channels, self.frames, self.height, self.width]) + + @staticmethod + def from_torch_shape(shape: torch.Size) -> VideoLatentShape: + return VideoLatentShape(batch=shape[0], channels=shape[1], frames=shape[2], height=shape[3], width=shape[4]) + + def upscale(self, scale_factors: SpatioTemporalScaleFactors) -> VideoLatentShape: + return self._replace( + channels=3, + frames=(self.frames - 1) * scale_factors.time + 1, + height=self.height * scale_factors.height, + width=self.width * scale_factors.width, + ) + + +def compute_trapezoidal_mask_1d( + length: int, + ramp_left: int, + ramp_right: int, + left_starts_from_0: bool, +) -> torch.Tensor: + """1D blending mask with linear fade-in / fade-out ramps, values in [0, 1].""" + if length <= 0: + raise ValueError("Mask length must be positive.") + + ramp_left = max(0, min(ramp_left, length)) + ramp_right = max(0, min(ramp_right, length)) + + mask = torch.ones(length) + + if ramp_left > 0: + interval_length = ramp_left + 1 if left_starts_from_0 else ramp_left + 2 + fade_in = torch.linspace(0.0, 1.0, interval_length)[:-1] + if not left_starts_from_0: + fade_in = fade_in[1:] + mask[:ramp_left] *= fade_in + + if ramp_right > 0: + fade_out = torch.linspace(1.0, 0.0, steps=ramp_right + 2)[1:-1] + mask[-ramp_right:] *= fade_out + + return mask.clamp_(0, 1) + + +@dataclass(frozen=True) +class DimensionInterval: + start: int + end: int + left_ramp: int + right_ramp: int + + +def untiled_mask_1d() -> torch.Tensor: + """Length-1 ones that broadcast over an untiled axis.""" + return torch.ones(1) + + +def _grow_last_tile_to_min(intervals: list[DimensionInterval], min_tile_size: int) -> list[DimensionInterval]: + """Grow a short last tile left to ``min_tile_size``; widen penultimate ``right_ramp``.""" + if len(intervals) <= 1: + return list(intervals) + last = intervals[-1] + if last.end - last.start >= min_tile_size: + return list(intervals) + new_start = last.end - min_tile_size + prev = intervals[-2] + new_overlap = prev.end - new_start + return [ + *intervals[:-2], + replace(prev, right_ramp=new_overlap), + replace(last, start=new_start, left_ramp=new_overlap), + ] + + +def split_by_size(dimension_size: int, size: int, overlap: int, min_tile_size: int | None) -> list[DimensionInterval]: + """Split ``dimension_size`` into tiles of ``size`` sharing ``overlap`` elements. + + The last tile may be shorter; with ``min_tile_size`` it is grown leftward instead. + """ + if (min_tile_size is not None and dimension_size < min_tile_size) or dimension_size <= size: + return [DimensionInterval(start=0, end=dimension_size, left_ramp=0, right_ramp=0)] + amount = (dimension_size + size - 2 * overlap - 1) // (size - overlap) + intervals = [ + DimensionInterval(start=0, end=size, left_ramp=0, right_ramp=overlap), + *( + DimensionInterval( + start=i * (size - overlap), + end=i * (size - overlap) + size, + left_ramp=overlap, + right_ramp=overlap, + ) + for i in range(1, amount - 1) + ), + DimensionInterval(start=(amount - 1) * (size - overlap), end=dimension_size, left_ramp=overlap, right_ramp=0), + ] + if min_tile_size is not None: + intervals = _grow_last_tile_to_min(intervals, min_tile_size) + return intervals + + +class Tile(NamedTuple): + """``in_coords`` cut the tile from the input, ``out_coords`` place its output; ``masks_1d`` blend it.""" + + in_coords: tuple[slice, ...] + out_coords: tuple[slice, ...] + masks_1d: tuple[torch.Tensor, ...] + + +def scale_by_masks_1d(x: torch.Tensor, masks_1d: Sequence[torch.Tensor]) -> torch.Tensor: + """Multiply ``x`` by separable 1d masks with broadcasting (one mask per axis).""" + out = x + for axis, mask in enumerate(masks_1d): + view_shape = [1] * x.ndim + view_shape[axis] = -1 + out = out * mask.reshape(*view_shape) + return out + + +def masks_are_complementary(tiles: Sequence[Tile], full_shape: Sequence[int]) -> bool: + """Whether per-axis 1d blend masks partition unity, so blending needs no weight buffer.""" + for axis, length in enumerate(full_shape): + acc = torch.zeros(length, dtype=torch.float32, device="cpu") + seen: set[tuple[int | None, int | None]] = set() + for tile in tiles: + sl = tile.out_coords[axis] + key = (sl.start, sl.stop) + if key in seen: + continue + seen.add(key) + acc[sl] += tile.masks_1d[axis].detach().float().cpu() + if not torch.allclose(acc, torch.ones(length, dtype=torch.float32), atol=1e-5, rtol=0.0): + return False + return True + + +def group_tiles_by_temporal_slice(tiles: list[Tile]) -> list[list[Tile]]: + """Group consecutive tiles sharing a temporal ``out_coords`` slice (temporal axis varies slowest).""" + groups = [] + current_slice = tiles[0].out_coords[2] + current_group = [] + for tile in tiles: + tile_slice = tile.out_coords[2] + if tile_slice == current_slice: + current_group.append(tile) + else: + groups.append(current_group) + current_slice = tile_slice + current_group = [tile] + groups.append(current_group) + return groups + + +@dataclass(frozen=True) +class DimensionSizeConfig: + """Tile size and overlap for one video axis, in pixel / frame units.""" + + tile_size: int + overlap: int + + +@dataclass(frozen=True) +class TileSizeConfig: + """Size-based tiling layout for a ``(F, H, W)`` video.""" + + frames: DimensionSizeConfig + height: DimensionSizeConfig + width: DimensionSizeConfig + + def split_sizes( + self, scale_factors: SpatioTemporalScaleFactors + ) -> tuple[tuple[int, int], tuple[int, int], tuple[int, int]]: + """``(tile, overlap)`` per ``(T, H, W)`` axis on the grid scaled down by ``scale_factors``.""" + + def axis(cfg: DimensionSizeConfig, factor: int) -> tuple[int, int]: + overlap = cfg.overlap // factor + return max(max(2, overlap + 1), cfg.tile_size // factor), overlap + + return ( + axis(self.frames, scale_factors.time), + axis(self.height, scale_factors.height), + axis(self.width, scale_factors.width), + ) diff --git a/jasna/models/ltx_vae/transformer.py b/jasna/models/ltx_vae/transformer.py new file mode 100644 index 00000000..068234a8 --- /dev/null +++ b/jasna/models/ltx_vae/transformer.py @@ -0,0 +1,418 @@ +"""Neighborhood-attention transformer blocks of the diffusion VAE decoder, ``chunked_eager`` path only. + +Ported from ``ltx_core/model/video_vae/transformer/`` (layers, qkv, rope_math, det_attn_rope, +attention, swiglu, blocks, chunked/*): eager tiled-SDPA neighborhood attention, torch SwiGLU, +W-chunked diffusion residual with halos and deferred stage-4 context inject. +""" + +from __future__ import annotations + +import math + +import numpy as np +import torch +import torch.nn.functional as F +from torch import nn + +from jasna.models.ltx_vae.eager_na import na3d + +DET_ROPE_NUM_TILES = 4 +DIFFUSION_W_CHUNKS = 4 +SWIGLU_TILE_TOKENS = 16_384 +_NORM_EPS = 1e-6 + + +def pixel_shuffle_channels_last(x: torch.Tensor, stride: tuple[int, int, int]) -> torch.Tensor: + """``b t h w (c p1 p2 p3) -> b (t p1) (h p2) (w p3) c``.""" + b, t, h, w, cp = x.shape + p1, p2, p3 = stride + c = cp // (p1 * p2 * p3) + x = x.reshape(b, t, h, w, c, p1, p2, p3).permute(0, 1, 5, 2, 6, 3, 7, 4) + return x.reshape(b, t * p1, h * p2, w * p3, c) + + +class LinearPixelShuffleUpsample(nn.Module): + """Linear channel-expand, then channels-last pixel shuffle. + + With ``stride[0] == 2`` the shuffle yields a duplicate leading frame; ``drop_leading_frame`` + must be True only for the chunk holding the tensor's true temporal origin. + """ + + def __init__(self, in_channels: int, stride: tuple[int, int, int], out_channels_reduction_factor: int) -> None: + super().__init__() + self.stride = stride + self.proj = nn.Linear(in_channels, math.prod(stride) * in_channels // out_channels_reduction_factor, bias=True) + + def forward(self, x: torch.Tensor, drop_leading_frame: bool) -> torch.Tensor: + x = pixel_shuffle_channels_last(self.proj(x), self.stride) + if self.stride[0] == 2 and drop_leading_frame: + x = x[:, 1:, :, :, :] + return x + + +class AdaLNZero(nn.Module): + """``t_emb`` -> 7 modulation chunks (scale/shift/gate slots; gates are unused).""" + + NUM_CHUNKS: int = 7 + + def __init__(self, dim: int, t_emb_dim: int) -> None: + super().__init__() + self.proj = nn.Linear(t_emb_dim, self.NUM_CHUNKS * dim, bias=True) + + def forward(self, t_emb: torch.Tensor) -> tuple[torch.Tensor, ...]: + chunks = self.proj(F.silu(t_emb)).chunk(self.NUM_CHUNKS, dim=-1) + return tuple(c[:, None, None, None, :] for c in chunks) + + +def default_rope_dim_split(head_dim: int) -> tuple[int, int, int]: + """Split of ``head_dim`` across (T, H, W) RoPE chunks.""" + d_t = (head_dim // 4) // 2 * 2 + d_hw = (head_dim - d_t) // 2 + if d_hw % 2 != 0: + d_t -= 2 + d_hw = (head_dim - d_t) // 2 + return (d_t, d_hw, d_hw) + + +def rope_inv_freqs(dim: int) -> torch.Tensor: + exponents = np.arange(0, dim, 2, dtype=np.float64) / dim + return torch.from_numpy(1.0 / np.power(10000.0, exponents)).to(torch.float32) + + +def rot_abs_axis_impl(xc: torch.Tensor, pos: torch.Tensor, inv: torch.Tensor, axis: int) -> torch.Tensor: + """Absolute RoPE (fp32 math) on one axis chunk ``xc[..., D]``.""" + out_dtype = xc.dtype + pairs = xc.reshape(*xc.shape[:-1], xc.shape[-1] // 2, 2) + xe = pairs[..., 0].to(torch.float32) + xo = pairs[..., 1].to(torch.float32) + shape = [1, 1, 1, 1, 1, inv.shape[0]] + shape[axis] = pos.shape[0] + ang = (pos[:, None] * inv[None, :]).reshape(shape) + c = ang.cos() + s = ang.sin() + re = xe * c - xo * s + ro = xe * s + xo * c + out = torch.stack([re, ro], dim=-1).reshape(xc.shape) + return out.to(out_dtype) if out.dtype != out_dtype else out + + +def apply_abs_rope_slab(x: torch.Tensor, attn: NeighborhoodAttention3D, w_pos: torch.Tensor) -> torch.Tensor: + """Rotate one W-extent ``(B, T, H, W, NH, HD)``; T/H positions are local, W positions are ``w_pos``.""" + d_t, d_h, _ = attn.rope_dim_split + t_pos = torch.arange(x.shape[1], dtype=torch.float32, device=x.device) + h_pos = torch.arange(x.shape[2], dtype=torch.float32, device=x.device) + xt = rot_abs_axis_impl(x[..., :d_t], t_pos, attn.rope_inv_t, axis=1) + xh = rot_abs_axis_impl(x[..., d_t : d_t + d_h], h_pos, attn.rope_inv_h, axis=2) + xw = rot_abs_axis_impl(x[..., d_t + d_h :], w_pos, attn.rope_inv_w, axis=3) + return torch.cat([xt, xh, xw], dim=-1) + + +def apply_tiled_abs_rope(x: torch.Tensor, attn: NeighborhoodAttention3D) -> torch.Tensor: + """Full-volume abs-RoPE computed over ``DET_ROPE_NUM_TILES`` W slabs.""" + w_off = 0 + parts: list[torch.Tensor] = [] + for slab in torch.chunk(x, DET_ROPE_NUM_TILES, dim=3): + w_slab = slab.shape[3] + w_pos = torch.arange(w_slab, dtype=torch.float32, device=x.device) + w_off + parts.append(apply_abs_rope_slab(slab, attn, w_pos)) + w_off = w_off + w_slab + return torch.cat(parts, dim=3) + + +class QKVProjections(nn.Module): + """Separate Q/K/V linears (checkpoints ship a fused ``qkv``; the loader splits it).""" + + def __init__(self, dim: int) -> None: + super().__init__() + self.to_q = nn.Linear(dim, dim, bias=True) + self.to_k = nn.Linear(dim, dim, bias=True) + self.to_v = nn.Linear(dim, dim, bias=True) + + def forward(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + return self.to_q(x), self.to_k(x), self.to_v(x) + + +def eager_sdpa_attention( + attn: NeighborhoodAttention3D, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor +) -> torch.Tensor: + if q.dtype != v.dtype or k.dtype != v.dtype: + q = q.to(dtype=v.dtype) + k = k.to(dtype=v.dtype) + return na3d(q, k, v, kernel_size=attn.kernel_size) + + +class NeighborhoodAttention3D(nn.Module): + """3D neighborhood attention with absolute RoPE on channels-last ``(B, T, H, W, C)``.""" + + def __init__(self, dim: int, kernel_size: tuple[int, int, int], head_dim: int) -> None: + super().__init__() + self.dim = dim + self.num_heads = dim // head_dim + self.head_dim = head_dim + self.kernel_size = tuple(kernel_size) + self.scale = head_dim**-0.5 + self.rope_dim_split = default_rope_dim_split(head_dim) + + self.register_buffer("rope_inv_t", rope_inv_freqs(self.rope_dim_split[0]), persistent=False) + self.register_buffer("rope_inv_h", rope_inv_freqs(self.rope_dim_split[1]), persistent=False) + self.register_buffer("rope_inv_w", rope_inv_freqs(self.rope_dim_split[2]), persistent=False) + + self.qkv = QKVProjections(dim) + self.proj = nn.Linear(dim, dim, bias=True) + self.q_norm = nn.RMSNorm(head_dim, eps=_NORM_EPS) + self.k_norm = nn.RMSNorm(head_dim, eps=_NORM_EPS) + + def project_qkv(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + batch, t, h, w, _ = x.shape + q, k, v = self.qkv(x) + shape = (batch, t, h, w, self.num_heads, self.head_dim) + return q.view(shape), k.view(shape), v.view(shape) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + """Deterministic-stage NA over the whole volume.""" + batch, t, h, w, _ = x.shape + kt, kh, kw = self.kernel_size + if t < kt or h < kh or w < kw: + raise ValueError( + f"3D neighborhood attention requires spatial dims >= kernel_size; " + f"got (T,H,W)=({t},{h},{w}) vs kernel={self.kernel_size}" + ) + q, k, v = self.project_qkv(x) + q = self.q_norm(q) + k = self.k_norm(k) + q = q * self.scale + q = apply_tiled_abs_rope(q, self) + k = apply_tiled_abs_rope(k, self) + out = eager_sdpa_attention(self, q.contiguous(), k.contiguous(), v.contiguous()) + out = out.reshape(batch, t, h, w, self.dim) + return self.proj(out) + + +class SwiGLU(nn.Module): + """Gated MLP weights: ``w_down(silu(w_gate(x)) * w_up(x))``.""" + + def __init__(self, dim: int, hidden_dim: int) -> None: + super().__init__() + self.w_up = nn.Linear(dim, hidden_dim, bias=False) + self.w_gate = nn.Linear(dim, hidden_dim, bias=False) + self.w_down = nn.Linear(hidden_dim, dim, bias=False) + + +def _swiglu_hidden_dim(dim: int) -> int: + return (int(dim * 4.0) + 15) // 16 * 16 + + +def _swiglu_into(xc: torch.Tensor, mlp: SwiGLU, workspace: torch.Tensor, out: torch.Tensor) -> None: + """``out = w_down(silu(xc @ W_gateᵀ) * (xc @ W_upᵀ))`` through a reusable ``workspace``.""" + torch.mm(xc, mlp.w_gate.weight.t(), out=workspace) + F.silu(workspace, inplace=True) + workspace.mul_(F.linear(xc, mlp.w_up.weight)) + torch.mm(workspace, mlp.w_down.weight.t(), out=out) + + +def swiglu_token_chunked(x: torch.Tensor, mlp: SwiGLU) -> torch.Tensor: + """SwiGLU over ``SWIGLU_TILE_TOKENS``-token chunks with one ``(chunk, hidden)`` workspace.""" + leading = x.shape[:-1] + dim = x.shape[-1] + x_flat = x.reshape(-1, dim).contiguous() + n_tok = x_flat.shape[0] + out_flat = torch.empty((n_tok, dim), device=x.device, dtype=x.dtype) + workspace = torch.empty( + (min(n_tok, SWIGLU_TILE_TOKENS), mlp.w_gate.weight.shape[0]), device=x.device, dtype=x.dtype + ) + for start in range(0, n_tok, SWIGLU_TILE_TOKENS): + end = min(n_tok, start + SWIGLU_TILE_TOKENS) + _swiglu_into(x_flat[start:end], mlp, workspace[: end - start], out_flat[start:end]) + return out_flat.view(*leading, dim) + + +def residual_modulating_mlp( + x: torch.Tensor, mlp: SwiGLU, norm: nn.RMSNorm, scale: torch.Tensor, shift: torch.Tensor +) -> torch.Tensor: + """In-place token-chunked ``x += swiglu(modulate(rms_norm(x)))``.""" + if not x.is_contiguous(): + x = x.contiguous() + dim = x.shape[-1] + x_flat = x.reshape(-1, dim) + n_tok = x_flat.shape[0] + s = scale.reshape(-1, dim) + sh = shift.reshape(-1, dim) + max_chunk = min(n_tok, SWIGLU_TILE_TOKENS) + workspace = torch.empty((max_chunk, mlp.w_gate.weight.shape[0]), device=x.device, dtype=x.dtype) + y_buf = torch.empty((max_chunk, dim), device=x.device, dtype=x.dtype) + out_buf = torch.empty((max_chunk, dim), device=x.device, dtype=x.dtype) + for start in range(0, n_tok, SWIGLU_TILE_TOKENS): + end = min(n_tok, start + SWIGLU_TILE_TOKENS) + n = end - start + xc = x_flat[start:end] + y = y_buf[:n] + y.copy_(F.rms_norm(xc, (dim,), norm.weight, _NORM_EPS)) + y.mul_(1.0 + s).add_(sh) + _swiglu_into(y, mlp, workspace[:n], out_buf[:n]) + xc.add_(out_buf[:n]) + return x + + +class NABlock(nn.Module): + """Pre-norm transformer block: NA -> SwiGLU MLP with residual adds (channels-last).""" + + def __init__(self, dim: int, kernel_size: tuple[int, int, int], head_dim: int) -> None: + super().__init__() + self.norm1 = nn.RMSNorm(dim, eps=_NORM_EPS) + self.attn = NeighborhoodAttention3D(dim, kernel_size, head_dim=head_dim) + self.norm2 = nn.RMSNorm(dim, eps=_NORM_EPS) + self.mlp = SwiGLU(dim, _swiglu_hidden_dim(dim)) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + x = x + self.attn(self.norm1(x)) + return x + swiglu_token_chunked(self.norm2(x), self.mlp) + + +def _attn_on_w_slab( + attn: NeighborhoodAttention3D, + norm: nn.RMSNorm, + x_chunk: torch.Tensor, + w_pos: torch.Tensor, + scale: torch.Tensor, + shift: torch.Tensor, +) -> torch.Tensor: + """One W slab: rms_norm + modulate + QKV / in-slab RoPE / NA / proj.""" + batch, t, h, ext_w, _ = x_chunk.shape + y = F.rms_norm(x_chunk, (attn.dim,), norm.weight, _NORM_EPS) + y = y * (1.0 + scale) + shift + + head_shape = (batch, t, h, ext_w, attn.num_heads, attn.head_dim) + q = F.linear(y, attn.qkv.to_q.weight, attn.qkv.to_q.bias).view(head_shape) + q = F.rms_norm(q, (attn.head_dim,), attn.q_norm.weight, _NORM_EPS) * float(attn.scale) + q = apply_abs_rope_slab(q, attn, w_pos) + k = F.linear(y, attn.qkv.to_k.weight, attn.qkv.to_k.bias).view(head_shape) + k = F.rms_norm(k, (attn.head_dim,), attn.k_norm.weight, _NORM_EPS) + k = apply_abs_rope_slab(k, attn, w_pos) + v = F.linear(y, attn.qkv.to_v.weight, attn.qkv.to_v.bias).view(head_shape) + del y + + out = eager_sdpa_attention(attn, q.contiguous(), k.contiguous(), v.contiguous()) + del q, k, v + return F.linear(out.reshape(batch, t, h, ext_w, attn.dim), attn.proj.weight, attn.proj.bias) + + +def w_chunked_attention_residual( + x: torch.Tensor, + attn: NeighborhoodAttention3D, + norm: nn.RMSNorm, + scale: torch.Tensor, + shift: torch.Tensor, +) -> torch.Tensor: + """In-place ``x += NA(modulate(norm(x)))`` over ``DIFFUSION_W_CHUNKS`` fixed-extent W slabs. + + Each slab carries ``kernel_w // 2`` halo columns copied from its neighbors; missing + halo slots at the true W boundary are edge-replicated so NA never sees zero padding. + """ + w_chunks = DIFFUSION_W_CHUNKS + halo = attn.kernel_size[2] // 2 + _, _, _, w, c = x.shape + chunk_w = (w + w_chunks - 1) // w_chunks + extent = chunk_w + 2 * halo + left_halo: torch.Tensor | None = None + + for i in range(w_chunks): + core_start = i * chunk_w + core_end = min(w, (i + 1) * chunk_w) + core_len = core_end - core_start + + buf = x.new_zeros(*x.shape[:3], extent, c) + if i > 0: + lh = left_halo.shape[3] + buf[:, :, :, halo - lh : halo, :] = left_halo + buf[:, :, :, halo : halo + core_len, :] = x[:, :, :, core_start:core_end, :] + right_filled = 0 + if i + 1 < w_chunks: + right_end = min(w, core_end + halo) + right = x[:, :, :, core_end:right_end, :] + right_filled = right.shape[3] + buf[:, :, :, halo + core_len : halo + core_len + right_filled, :] = right + + if i == 0 and halo > 0 and core_len > 0: + edge_l = buf[:, :, :, halo : halo + 1, :] + buf[:, :, :, :halo, :] = edge_l.expand(*x.shape[:3], halo, c) + missing_right = extent - (halo + core_len + right_filled) + if missing_right > 0 and core_len > 0: + edge_r = buf[:, :, :, halo + core_len - 1 : halo + core_len, :] + lo_r = halo + core_len + right_filled + buf[:, :, :, lo_r:extent, :] = edge_r.expand(*x.shape[:3], missing_right, c) + + if i + 1 < w_chunks: + take = min(halo, core_len) + left_halo = x[:, :, :, core_end - take : core_end, :].clone() + + w_pos = torch.arange(extent, device=x.device, dtype=torch.float32) + (core_start - halo) + out = _attn_on_w_slab(attn, norm, buf, w_pos, scale, shift) + x[:, :, :, core_start:core_end, :].add_(out[:, :, :, halo : halo + core_len, :]) + + return x + + +def _upsample_then_ctx( + feat: torch.Tensor, upsample: LinearPixelShuffleUpsample, context_proj: nn.Linear, *, drop_leading_frame: bool +) -> torch.Tensor: + """``context_proj(pixel_shuffle(upsample.proj(feat)))``.""" + up = pixel_shuffle_channels_last(F.linear(feat, upsample.proj.weight, upsample.proj.bias), upsample.stride) + if upsample.stride[0] == 2 and drop_leading_frame: + up = up[:, 1:, :, :, :] + return F.linear(up, context_proj.weight, context_proj.bias) + + +def inject_deferred_context( + x: torch.Tensor, + stage4_feat: torch.Tensor, + upsample: LinearPixelShuffleUpsample, + context_proj: nn.Linear, + *, + drop_leading_frame: bool, +) -> torch.Tensor: + """Chunk stage-4 feat along W, upsample + ``context_proj`` each slab, ``add_`` into ``x``.""" + w_hi = x.shape[3] + lo = 0 + for feat_chunk in torch.chunk(stage4_feat, DIFFUSION_W_CHUNKS, dim=3): + hi = min(w_hi, lo + feat_chunk.shape[3] * upsample.stride[2]) + ctx = _upsample_then_ctx(feat_chunk, upsample, context_proj, drop_leading_frame=drop_leading_frame) + x[:, :, :, lo:hi, :].add_(ctx[:, :, :, : hi - lo, :]) + lo = hi + return x + + +class ChunkedDiffusionNABlock(nn.Module): + """Diffusion NA + SwiGLU block with shared AdaLN modulation and deferred stage-4 context.""" + + def __init__(self, dim: int, kernel_size: tuple[int, int, int], context_channels: int, head_dim: int) -> None: + super().__init__() + self.context_proj = nn.Linear(context_channels, dim, bias=True) + self.scale_shift_table = nn.Parameter(torch.zeros(AdaLNZero.NUM_CHUNKS, dim)) + self.norm1 = nn.RMSNorm(dim, eps=_NORM_EPS) + self.attn = NeighborhoodAttention3D(dim, kernel_size, head_dim=head_dim) + self.norm2 = nn.RMSNorm(dim, eps=_NORM_EPS) + self.mlp = SwiGLU(dim, _swiglu_hidden_dim(dim)) + + def _modulation( + self, modulation: tuple[torch.Tensor, ...] + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + scale_msa, shift_msa, _, scale_mlp, shift_mlp, _, _ = [ + modulation[i] + self.scale_shift_table[i].view(1, 1, 1, 1, -1) for i in range(AdaLNZero.NUM_CHUNKS) + ] + return scale_msa, shift_msa, scale_mlp, shift_mlp + + def forward( + self, + x: torch.Tensor, + stage4_feat: torch.Tensor, + modulation: tuple[torch.Tensor, ...], + stage4_upsample: LinearPixelShuffleUpsample, + *, + drop_leading_frame: bool, + ) -> torch.Tensor: + scale_msa, shift_msa, scale_mlp, shift_mlp = self._modulation(modulation) + x = inject_deferred_context( + x, stage4_feat, stage4_upsample, self.context_proj, drop_leading_frame=drop_leading_frame + ) + x = w_chunked_attention_residual(x, self.attn, self.norm1, scale_msa, shift_msa) + return residual_modulating_mlp(x, self.mlp, self.norm2, scale_mlp, shift_mlp) diff --git a/jasna/models/ltx_vae/video_encoder.py b/jasna/models/ltx_vae/video_encoder.py new file mode 100644 index 00000000..cc763282 --- /dev/null +++ b/jasna/models/ltx_vae/video_encoder.py @@ -0,0 +1,165 @@ +"""Causal 3D-conv video VAE encoder (``ltx_core/model/video_vae/{video_vae,convolution,resnet,sampling}.py``).""" + +import logging +import math +from typing import Any + +import torch +from torch import nn + +from jasna.models.ltx_vae.ops import PerChannelStatistics, PixelNorm, patchify +from jasna.models.ltx_vae.tiling import SpatioTemporalScaleFactors + +logger = logging.getLogger(__name__) + + +class CausalConv3d(nn.Module): + """3x3x3 conv; time is padded causally by repeating the first frame.""" + + time_kernel_size = 3 + + def __init__(self, in_channels: int, out_channels: int) -> None: + super().__init__() + self.conv = nn.Conv3d(in_channels, out_channels, 3, padding=(0, 1, 1), padding_mode="zeros") + + def forward(self, x: torch.Tensor) -> torch.Tensor: + first_frame_pad = x[:, :, :1, :, :].repeat((1, 1, self.time_kernel_size - 1, 1, 1)) + x = torch.concatenate((first_frame_pad, x), dim=2) + return self.conv(x) + + +class ResnetBlock3D(nn.Module): + """PixelNorm -> SiLU -> conv, twice, plus identity skip.""" + + def __init__(self, channels: int) -> None: + super().__init__() + self.norm1 = PixelNorm() + self.non_linearity = nn.SiLU() + self.conv1 = CausalConv3d(channels, channels) + self.norm2 = PixelNorm() + self.conv2 = CausalConv3d(channels, channels) + + def forward(self, input_tensor: torch.Tensor) -> torch.Tensor: + hidden_states = self.norm1(input_tensor) + hidden_states = self.non_linearity(hidden_states) + hidden_states = self.conv1(hidden_states) + hidden_states = self.norm2(hidden_states) + hidden_states = self.non_linearity(hidden_states) + hidden_states = self.conv2(hidden_states) + return input_tensor + hidden_states + + +class UNetMidBlock3D(nn.Module): + def __init__(self, channels: int, num_layers: int) -> None: + super().__init__() + self.res_blocks = nn.ModuleList([ResnetBlock3D(channels) for _ in range(num_layers)]) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + for resnet in self.res_blocks: + hidden_states = resnet(hidden_states) + return hidden_states + + +def space_to_depth(x: torch.Tensor, stride: tuple[int, int, int]) -> torch.Tensor: + """``b c (d p1) (h p2) (w p3) -> b (c p1 p2 p3) d h w``.""" + b, c, d, h, w = x.shape + p1, p2, p3 = stride + x = x.reshape(b, c, d // p1, p1, h // p2, p2, w // p3, p3).permute(0, 1, 3, 5, 7, 2, 4, 6) + return x.reshape(b, c * p1 * p2 * p3, d // p1, h // p2, w // p3) + + +class SpaceToDepthDownsample(nn.Module): + """Strided downsample: conv + space-to-depth, plus a group-mean space-to-depth skip.""" + + def __init__(self, in_channels: int, out_channels: int, stride: tuple[int, int, int]) -> None: + super().__init__() + self.stride = stride + self.group_size = in_channels * math.prod(stride) // out_channels + self.conv = CausalConv3d(in_channels, out_channels // math.prod(stride)) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + if self.stride[0] == 2: + x = torch.cat([x[:, :, :1, :, :], x], dim=2) + + x_in = space_to_depth(x, self.stride) + b, c, d, h, w = x_in.shape + x_in = x_in.view(b, c // self.group_size, self.group_size, d, h, w).mean(dim=2) + + x = self.conv(x) + x = space_to_depth(x, self.stride) + return x + x_in + + +_DOWNSAMPLE_STRIDES = { + "compress_space_res": (1, 2, 2), + "compress_time_res": (2, 1, 1), + "compress_all_res": (2, 2, 2), +} + + +def _make_encoder_block(block_name: str, block_config: dict[str, Any], in_channels: int) -> tuple[nn.Module, int]: + if block_name == "res_x": + return UNetMidBlock3D(in_channels, block_config["num_layers"]), in_channels + if block_name in _DOWNSAMPLE_STRIDES: + out_channels = in_channels * block_config.get("multiplier", 2) + return SpaceToDepthDownsample(in_channels, out_channels, _DOWNSAMPLE_STRIDES[block_name]), out_channels + raise ValueError(f"unsupported encoder block: {block_name}") + + +class VideoEncoder(nn.Module): + """Encodes ``(B, 3, F, H, W)`` video in [-1, 1] into normalized latent means. + + ``F' = 1 + (F-1)/8, H' = H/32, W' = W/32``; F must be ``1 + 8k`` (extra frames are cropped). + """ + + def __init__( + self, + in_channels: int, + out_channels: int, + encoder_blocks: list[tuple[str, dict[str, Any]]], + patch_size: int, + ) -> None: + super().__init__() + self.patch_size = patch_size + self.latent_channels = out_channels + self.video_scale_factors = SpatioTemporalScaleFactors.from_blocks(encoder_blocks, patch_size) + self.per_channel_statistics = PerChannelStatistics(latent_channels=out_channels) + + feature_channels = out_channels + self.conv_in = CausalConv3d(in_channels * patch_size**2, feature_channels) + + self.down_blocks = nn.ModuleList([]) + for block_name, block_config in encoder_blocks: + block, feature_channels = _make_encoder_block(block_name, block_config, feature_channels) + self.down_blocks.append(block) + + self.conv_norm_out = PixelNorm() + self.conv_act = nn.SiLU() + # One extra channel: the checkpoint's (unused) constant log-variance. + self.conv_out = CausalConv3d(feature_channels, out_channels + 1) + + def forward(self, sample: torch.Tensor) -> torch.Tensor: + temporal_factor = self.video_scale_factors.time + frames_count = sample.shape[2] + if ((frames_count - 1) % temporal_factor) != 0: + frames_to_crop = (frames_count - 1) % temporal_factor + logger.warning( + "Invalid number of frames %s for encode; cropping last %s frames to satisfy 1 + %s*k.", + frames_count, + frames_to_crop, + temporal_factor, + ) + sample = sample[:, :, :-frames_to_crop, ...] + + sample = patchify(sample, patch_size_hw=self.patch_size) + sample = self.conv_in(sample) + + for down_block in self.down_blocks: + sample = down_block(sample) + + sample = self.conv_norm_out(sample) + sample = self.conv_act(sample) + sample = self.conv_out(sample) + + means = sample[:, :-1, ...] + return self.per_channel_statistics.normalize(means) diff --git a/jasna/pipeline.py b/jasna/pipeline.py index cd290076..29dd09e5 100644 --- a/jasna/pipeline.py +++ b/jasna/pipeline.py @@ -96,6 +96,9 @@ def __init__( self.vr_projection = config.vr_projection self.detection_model = session.detection_model_for(config) self.restoration_pipeline = session.restoration_pipeline + self.ltx_files = session.ltx_files + self.ltx_large_canvas = config.ltx_large_canvas + self.ltx_seed = config.ltx_seed self.disable_progress = config.disable_progress self.progress_callback = progress_callback self.lut_path = config.lut_path @@ -203,9 +206,7 @@ def validate_metadata(self, metadata) -> None: "Only BT.709, BT.601, and BT.2020 non-constant-luminance are supported." ) - def _run_full( - self, metadata, *, effect_ranges: tuple[tuple[int, int], ...] | None = None - ) -> None: + def _resolve_frame_rate(self, metadata): frame_rate = resolve_frame_rate_retarget( metadata.video_fps_exact, enabled=self.retarget_high_fps, @@ -230,19 +231,15 @@ def _run_full( "Frame-rate retargeting requested, but %s fps is not a supported source rate; keeping source rate", frame_rate.source_fps, ) - output_frame_count = frame_rate.output_frame_count(metadata.num_frames) - progress = Progressbar( - total_frames=output_frame_count, - video_fps=float(frame_rate.output_fps), - disable=self.disable_progress, - callback=self.progress_callback, - ) + return frame_rate + + def _video_encoder(self, metadata, frame_rate) -> VideoEncoder: if self.fmp4 and self.output_video.suffix.lower() not in MOV_SUFFIXES: log.info( "Fragmented MP4 has no effect on %s output; it is already playable while it grows", self.output_video.suffix, ) - encoder_ctx = VideoEncoder( + return VideoEncoder( str(self.output_video), device=self.device, metadata=metadata, @@ -253,6 +250,19 @@ def _run_full( output_fps=frame_rate.output_fps, fmp4=self.fmp4, ) + + def _run_full( + self, metadata, *, effect_ranges: tuple[tuple[int, int], ...] | None = None + ) -> None: + frame_rate = self._resolve_frame_rate(metadata) + output_frame_count = frame_rate.output_frame_count(metadata.num_frames) + progress = Progressbar( + total_frames=output_frame_count, + video_fps=float(frame_rate.output_fps), + disable=self.disable_progress, + callback=self.progress_callback, + ) + encoder_ctx = self._video_encoder(metadata, frame_rate) try: self._run_pass( metadata=metadata, @@ -264,6 +274,43 @@ def _run_full( finally: progress.close(ensure_completed_bar=True) + def _run_ltx(self, metadata) -> None: + from jasna.ltx.restore import Cancelled, Progress, restore_video + from jasna.media.video_decoder import VideoReader + + if self.vr_resolution.is_sbs: + raise ValueError("LTX restoration does not support VR180 side-by-side video") + frame_rate = self._resolve_frame_rate(metadata) + writer = _OfflineFrameWriter(self._video_encoder(metadata, frame_rate), [time.monotonic()]) + + def frames(): + with VideoReader( + str(self.input_video), self.batch_size, self.device, metadata, frame_stride=frame_rate.frame_stride + ) as reader: + yield from reader.frames(None) + + try: + restore_video( + frames, + writer.write, + detector=self.job_detection_model, + files=self.ltx_files, + frame_h=int(metadata.video_height), + frame_w=int(metadata.video_width), + large_canvas=self.ltx_large_canvas, + seed=self.ltx_seed, + device=self.device, + work_dir=self.working_dir or self.output_video.parent, + progress=Progress( + frames=frame_rate.output_frame_count(metadata.num_frames), disable=self.disable_progress + ), + cancel=self._cancel_event, + ) + except Cancelled: + log.info("LTX restoration cancelled") + finally: + writer.close() + def _run_smart(self, metadata) -> None: codec = validate_smart_render( metadata, @@ -390,7 +437,9 @@ def run(self) -> None: metadata = get_video_meta_data(str(self.input_video)) self.validate_metadata(metadata) self.configure_vr(metadata) - if self.segments: + if self.ltx_files is not None: + self._run_ltx(metadata) + elif self.segments: if self.fmp4: log.warning( "Fragmented MP4 is not available with segment processing; " diff --git a/jasna/session_config.py b/jasna/session_config.py index 2d47e868..cdc62838 100644 --- a/jasna/session_config.py +++ b/jasna/session_config.py @@ -11,6 +11,9 @@ from pathlib import Path from typing import Literal, Mapping +RestorationModelName = Literal["basicvsrpp", "ltx"] + +LTX_DEFAULT_SEED = 20260923 SecondaryRestorationName = Literal["none", "unet-4x", "tvai", "rtx-super-res"] DenoiseStrengthName = Literal["none", "low", "medium", "high"] DenoiseStepName = Literal["after_primary", "after_secondary"] @@ -32,7 +35,10 @@ class SessionConfig: max_detection_gap: int min_detection_duration: int scene_detection: bool + restoration_model_name: RestorationModelName restoration_model_path: Path + ltx_large_canvas: bool + ltx_seed: int compile_basicvsrpp: bool max_clip_size: int temporal_overlap: int diff --git a/jasna/session_factory.py b/jasna/session_factory.py index 7b3eb9cc..483c1666 100644 --- a/jasna/session_factory.py +++ b/jasna/session_factory.py @@ -19,6 +19,7 @@ if TYPE_CHECKING: import torch + from jasna.ltx.restore import LtxModelFiles from jasna.media.splice import SplicePlan from jasna.mosaic.rfdetr import RfDetrMosaicDetectionModel from jasna.mosaic.yolo import YoloMosaicDetectionModel @@ -32,7 +33,8 @@ @dataclass class RestorationSession: device: "torch.device" - restoration_pipeline: "RestorationPipeline" + restoration_pipeline: "RestorationPipeline | None" + ltx_files: "LtxModelFiles | None" = None _detection_key: tuple | None = None _detection_model: "DetectionModel | None" = None @@ -65,9 +67,10 @@ def close(self) -> None: if self._detection_model is not None: self._detection_model.close() self._detection_model = None - self.restoration_pipeline.restorer.close() - if self.restoration_pipeline.secondary_restorer is not None: - self.restoration_pipeline.secondary_restorer.close() + if self.restoration_pipeline is not None: + self.restoration_pipeline.restorer.close() + if self.restoration_pipeline.secondary_restorer is not None: + self.restoration_pipeline.secondary_restorer.close() def build_compiled_detection_model( @@ -158,6 +161,9 @@ def build_restoration_session( f"Secondary restoration '{config.secondary_restoration}' is not available in the AMD build yet" ) + if config.restoration_model_name == "ltx": + return _build_ltx_session(config, device, log_callback=log_callback) + compile_result = ensure_engines_compiled( EngineCompilationRequest( device=str(device), @@ -190,6 +196,35 @@ def build_restoration_session( return RestorationSession(device=device, restoration_pipeline=restoration_pipeline) +def _build_ltx_session( + config: SessionConfig, device: "torch.device", *, log_callback: Callable[[str], None] | None +) -> RestorationSession: + from jasna.accelerator import is_nvidia_device + from jasna.engine_compiler import EngineCompilationRequest, ensure_engines_compiled + from jasna.ltx.restore import LtxModelFiles + + if not is_nvidia_device(device): + raise ValueError("LTX restoration needs an NVIDIA GPU") + if config.secondary_restoration != "none" or config.denoise_strength != "none": + raise ValueError("LTX restoration does not support secondary restoration or denoise") + ensure_engines_compiled( + EngineCompilationRequest( + device=str(device), + fp16=bool(config.fp16), + detection=True, + detection_model_name=config.detection_model_name, + detection_model_path=str(config.detection_model_path), + detection_batch_size=int(config.batch_size), + ), + log_callback=log_callback, + ) + return RestorationSession( + device=device, + restoration_pipeline=None, + ltx_files=LtxModelFiles.from_dir(config.restoration_model_path), + ) + + def build_pipeline( config: SessionConfig, session: RestorationSession, diff --git a/tests/factories.py b/tests/factories.py index f02f2c94..a61eab2f 100644 --- a/tests/factories.py +++ b/tests/factories.py @@ -25,7 +25,10 @@ def session_config(**overrides) -> SessionConfig: max_detection_gap=2, min_detection_duration=2, scene_detection=True, + restoration_model_name="basicvsrpp", restoration_model_path=Path("restore.pth"), + ltx_large_canvas=True, + ltx_seed=0, compile_basicvsrpp=True, max_clip_size=90, temporal_overlap=8, @@ -63,6 +66,7 @@ def fake_session(*, device: torch.device, restoration_pipeline, detection_model= return SimpleNamespace( device=device, restoration_pipeline=restoration_pipeline, + ltx_files=None, detection_model_for=lambda config: detection_model, ) diff --git a/tests/test_ltx_blocks.py b/tests/test_ltx_blocks.py new file mode 100644 index 00000000..cdbdd9c9 --- /dev/null +++ b/tests/test_ltx_blocks.py @@ -0,0 +1,31 @@ +import pytest +import torch + +from jasna.ltx.blocks import BlockStore + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA") + + +def _blocks(count: int) -> list[dict[str, torch.Tensor]]: + generator = torch.Generator().manual_seed(0) + return [ + { + "a.qdata": torch.randint(-127, 128, (64, 48), dtype=torch.int8, generator=generator), + "a.scales": torch.rand(64, generator=generator), + "norm.weight": torch.randn(33, generator=generator).to(torch.bfloat16), + } + for _ in range(count) + ] + + +@pytest.mark.parametrize("resident", [0, 2, 5]) +def test_streamed_and_resident_blocks_hand_out_the_same_weights(resident): + blocks = _blocks(5) + store = BlockStore(blocks, torch.device("cuda"), resident=resident) + for _ in range(3): + for index, expected in enumerate(blocks): + weights = store.acquire(index) + for name, tensor in expected.items(): + assert torch.equal(weights[name].cpu(), tensor) + store.release(index) + store.close() diff --git a/tests/test_ltx_compose.py b/tests/test_ltx_compose.py new file mode 100644 index 00000000..d5a3fc65 --- /dev/null +++ b/tests/test_ltx_compose.py @@ -0,0 +1,71 @@ +import cv2 +import numpy as np +import torch + +from jasna.ltx.camera import letterbox +from jasna.ltx.compose import Candidate, canvas_to_source, composite_frame, crop_to_canvas, feather_alpha + + +def _frame(seed: int = 0) -> torch.Tensor: + generator = torch.Generator().manual_seed(seed) + return torch.randint(0, 256, (3, 360, 640), dtype=torch.uint8, generator=generator) + + +def test_feather_alpha_matches_opencv(): + mask = np.zeros((200, 300), dtype=np.uint8) + cv2.circle(mask, (150, 100), 40, 255, -1) + feather = 12 + kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (2 * feather + 1,) * 2) + expected = cv2.GaussianBlur(cv2.dilate(mask, kernel).astype(np.float32) / 255.0, (2 * feather + 1,) * 2, 0.0) + actual = feather_alpha(torch.from_numpy(mask > 0), feather).numpy() + assert np.abs(actual - expected).max() < 1e-5 + + +def test_crop_round_trip_at_unit_scale(): + frame = _frame() + crop = letterbox((100, 50, 356, 306), 256) + canvas = crop_to_canvas(frame, crop) + assert canvas.shape == (3, 256, 256) + assert torch.equal(canvas_to_source(canvas, crop).to(torch.uint8), frame[:, 50:306, 100:356]) + + +def test_crop_pads_by_reflection(): + frame = _frame() + crop = letterbox((0, 0, 200, 100), 200) + canvas = crop_to_canvas(frame, crop) + top = crop.pad[0] + assert torch.equal(canvas[:, top - 1], canvas[:, top + 1]) + + +def test_composite_touches_only_the_feathered_mask(): + source = _frame(1) + crop = letterbox((200, 100, 456, 356), 256) + restored = torch.full((3, 256, 256), 128, dtype=torch.uint8) + polygon = [[300.0, 200.0], [360.0, 200.0], [360.0, 260.0], [300.0, 260.0]] + out = composite_frame(source, [Candidate(restored, crop, 1.0, [polygon])], feather=8) + changed = (out != source).any(dim=0) + ys, xs = torch.nonzero(changed, as_tuple=True) + assert changed[230, 330] + assert xs.min() >= 300 - 2 * 8 - 1 and xs.max() <= 360 + 2 * 8 + 1 + assert ys.min() >= 200 - 2 * 8 - 1 and ys.max() <= 260 + 2 * 8 + 1 + + +def test_zero_weight_and_no_candidates_keep_the_source(): + source = _frame(2) + crop = letterbox((200, 100, 456, 356), 256) + restored = torch.zeros((3, 256, 256), dtype=torch.uint8) + polygon = [[300.0, 200.0], [360.0, 200.0], [360.0, 260.0]] + assert torch.equal(composite_frame(source, [], feather=8), source) + assert torch.equal(composite_frame(source, [Candidate(restored, crop, 0.0, [polygon])], feather=8), source) + + +def test_two_agreeing_windows_equal_one(): + source = _frame(3) + crop = letterbox((200, 100, 456, 356), 256) + restored = torch.full((3, 256, 256), 90, dtype=torch.uint8) + polygon = [[300.0, 200.0], [360.0, 200.0], [360.0, 260.0], [300.0, 260.0]] + one = composite_frame(source, [Candidate(restored, crop, 1.0, [polygon])], feather=8) + two = composite_frame( + source, [Candidate(restored, crop, 0.3, [polygon]), Candidate(restored, crop, 0.7, [polygon])], feather=8 + ) + assert (one.int() - two.int()).abs().max() <= 1 diff --git a/tests/test_ltx_plan.py b/tests/test_ltx_plan.py new file mode 100644 index 00000000..246b9ab6 --- /dev/null +++ b/tests/test_ltx_plan.py @@ -0,0 +1,117 @@ +import numpy as np +import pytest + +from jasna.ltx.camera import PLAN_CANVAS, ZOOM_PER_WINDOW, letterbox, track_camera, zoom_crop_sides +from jasna.ltx.plan import ( + LARGE_CANVAS, + LEFT_SHARED_LATENTS, + RIGHT_SHARED_LATENTS, + WINDOW_HOP, + Region, + RegionGroup, + Track, + build_tracks, + crossfade_weights, + group_regions, + merge_continuations, + plan_video, + track_polygons, + window_starts, +) + + +def _square(x: float, y: float, side: float) -> Region: + box = (x, y, x + side, y + side) + return Region(box=box, polygon=[[x, y], [x + side, y], [x + side, y + side], [x, y + side]]) + + +def _group(x: float, y: float, side: float) -> RegionGroup: + region = _square(x, y, side) + return RegionGroup(region.box, (region.polygon,)) + + +def test_group_regions_merges_touching_boxes_only(): + groups = group_regions([_square(0, 0, 100), _square(90, 0, 100), _square(500, 500, 50)]) + assert sorted(len(g.polygons) for g in groups) == [1, 2] + merged = next(g for g in groups if len(g.polygons) == 2) + assert merged.box == (0, 0, 190, 100) + + +def test_build_tracks_links_across_short_gaps_but_not_cuts(): + frames = [[_group(100, 100, 80)] for _ in range(10)] + frames[3] = [] + frames[4] = [] + assert len(build_tracks(frames, cuts=set())) == 1 + assert len(build_tracks(frames, cuts={6})) == 2 + frames[5] = frames[6] = frames[7] = [] + frames[3] = frames[4] = [] + assert len(build_tracks(frames, cuts=set())) == 2 + + +def test_merge_continuations_folds_a_split_mosaic(): + parent = Track(0, {i: _group(100, 100, 200) for i in range(10)}) + child = Track(1, {i: _group(150, 150, 60) for i in range(5, 15)}) + elsewhere = Track(2, {i: _group(900, 900, 60) for i in range(5, 15)}) + merged = merge_continuations([parent, child, elsewhere], cuts=set()) + assert [t.track_id for t in merged] == [0, 2] + assert merged[0].first == 0 and merged[0].last == 14 + assert len(merged[0].frames[6].polygons) == 2 + assert len(merge_continuations([parent, child], cuts={5})) == 2 + + +def test_track_polygons_interpolate_gaps(): + track = Track(0, {0: _group(0, 0, 100), 4: _group(40, 0, 100)}) + polygons = track_polygons(track) + assert len(polygons) == 5 + xs = [p[0][0][0] for p in polygons] + assert xs == pytest.approx([0, 10, 20, 30, 40]) + + +def test_window_starts_cover_with_hop(): + assert window_starts(1) == [0] + assert window_starts(121) == [0] + assert window_starts(122) == [0, WINDOW_HOP] + assert window_starts(385) == [0, 88, 176, 264] + assert window_starts(386) == [0, 88, 176, 264, 352] + + +def test_shared_latents_encode_the_same_frames(): + def first_pixel(latent: int) -> int: + return 0 if latent == 0 else 8 * latent - 7 + + for left, right in zip(LEFT_SHARED_LATENTS, RIGHT_SHARED_LATENTS): + assert first_pixel(left) == WINDOW_HOP + first_pixel(right) + + +def test_crossfade_weights_partition_unity(): + plans = plan_video([[_square(300 + t, 300, 120)] for t in range(300)], set(), frame_w=1920, frame_h=1080, large_canvas=True) + windows = plans[0].windows + total = np.zeros(300) + for window, weight in zip(windows, crossfade_weights(windows)): + total[window.start : window.stop] += weight + assert np.allclose(total, 1.0) + + +def test_zoom_sides_cover_raw_and_respect_rate(): + raw = np.array([200] * 100 + [800] * 100 + [200] * 200) + sides = zoom_crop_sides(raw) + assert (sides >= raw).all() + span = sides[120:] / sides[:-120] + assert max(span.max(), (1 / span).max()) <= ZOOM_PER_WINDOW * (1 + 1 / raw.min()) + + +def test_track_camera_contains_support_and_letterboxes(): + boxes = [(100.0 + 3 * t, 200.0, 400.0 + 3 * t, 500.0) for t in range(150)] + camera = track_camera(boxes, frame_w=1920, frame_h=1080, feather=30) + for crop, box in zip(camera, boxes): + x1, y1, x2, y2 = crop.crop_box + assert x1 <= box[0] and y1 <= box[1] and x2 >= box[2] and y2 >= box[3] + assert crop.canvas == PLAN_CANVAS + + +def test_large_crops_use_the_large_canvas(): + plans = plan_video([[_square(200, 100, 700)] for _ in range(50)], set(), frame_w=1920, frame_h=1080, large_canvas=True) + assert plans[0].windows[0].canvas == LARGE_CANVAS + capped = plan_video([[_square(200, 100, 700)] for _ in range(50)], set(), frame_w=1920, frame_h=1080, large_canvas=False) + assert capped[0].windows[0].canvas == PLAN_CANVAS + assert letterbox((0, 0, 900, 900), LARGE_CANVAS).resize_hw == (LARGE_CANVAS, LARGE_CANVAS) diff --git a/tests/test_ltx_transformer.py b/tests/test_ltx_transformer.py new file mode 100644 index 00000000..2c4e23a5 --- /dev/null +++ b/tests/test_ltx_transformer.py @@ -0,0 +1,81 @@ +import pytest +import torch + +from jasna.ltx import transformer as T +from jasna.ltx.sampler import fuse_overlaps, patchify, unpatchify + + +def _block_weights(dim: int, heads: int, generator: torch.Generator) -> dict[str, torch.Tensor]: + def rand(*shape: int) -> torch.Tensor: + return (torch.randn(*shape, generator=generator) * 0.05).to(torch.bfloat16) + + weights = {} + for attn in ("attn1", "attn2"): + for proj in ("to_q", "to_k", "to_v", "to_out.0"): + weights[f"{attn}.{proj}.weight"] = rand(dim, dim) + weights[f"{attn}.{proj}.bias"] = rand(dim) + weights[f"{attn}.to_gate_logits.weight"] = rand(heads, dim) + weights[f"{attn}.to_gate_logits.bias"] = rand(heads) + weights[f"{attn}.q_norm.weight"] = rand(dim) + 1 + weights[f"{attn}.k_norm.weight"] = rand(dim) + 1 + weights["ff.net.0.proj.weight"] = rand(4 * dim, dim) + weights["ff.net.2.weight"] = rand(dim, 4 * dim) + weights["scale_shift_table"] = rand(9, dim) + weights["prompt_scale_shift_table"] = rand(2, dim) + return weights + + +def test_rotation_is_its_own_inverse(): + x = torch.randn(4, 512, dtype=torch.float32) + assert torch.allclose(T.rotate(T.rotate(x)), x, atol=1e-5) + + +def test_quantize_rows_pads_and_bounds(): + x = torch.randn(40, 64) * 3 + codes, scales = T.quantize_rows(x) + assert codes.shape == (64, 64) and scales.shape == (64,) + assert codes.abs().max() == T.ACT_QMAX + assert torch.allclose(codes[:40].float() * scales[:40, None], x, atol=float(scales.max())) + + +def test_block_matches_across_token_chunking(monkeypatch): + generator = torch.Generator().manual_seed(0) + dim, heads = 64, 4 + weights = _block_weights(dim, heads, generator) + rope = T.build_rope(3, 2, 2, dim=dim, heads=heads, device=torch.device("cpu")) + cond = T.StepConditioning( + block_rows=(torch.randn(1, 2, 9, dim, generator=generator) * 0.1).to(torch.bfloat16), + output_rows=torch.zeros(1, 2, dim, dtype=torch.bfloat16), + prompt_rows=(torch.randn(1, 1, 2, dim, generator=generator) * 0.1).to(torch.bfloat16), + ) + x = torch.randn(1, 24, dim, generator=generator).to(torch.bfloat16) + context = torch.randn(1, 5, dim, generator=generator).to(torch.bfloat16) + whole = T.block_forward(weights, x, cond, context, rope, heads=heads) + monkeypatch.setattr(T, "CHUNK_ROWS", 5) + chunked = T.block_forward(weights, x, cond, context, rope, heads=heads) + assert torch.equal(whole, chunked) + skipped = T.block_forward(weights, x, cond, context, rope, heads=heads, skip_self_attention=True) + assert not torch.equal(whole, skipped) + + +def test_latent_positions_are_causal_seconds(): + positions = T.latent_positions(3, 1, 1) + assert positions[0].tolist() == pytest.approx([0.5 / 24, 5 / 24, 13 / 24]) + assert positions[1].tolist() == [16.0, 16.0, 16.0] + + +def test_patchify_round_trip(): + latent = torch.randn(1, 128, 16, 2, 3) + assert torch.equal(unpatchify(patchify(latent), 16, 2, 3), latent) + + +def test_fuse_overlaps_writes_the_same_mean_to_both_windows(): + tokens_per_frame = 2 + left = torch.randn(1, 16 * tokens_per_frame, 4) + right = torch.randn(1, 16 * tokens_per_frame, 4) + right_before = right.clone() + fuse_overlaps([left, right], tokens_per_frame, [0.0, 0.0, 0.0, 0.0]) + assert torch.equal(right[:, 2:10], left[:, 24:32]) + assert torch.equal(right[:, 10:], right_before[:, 10:]) + fuse_overlaps([left, right], tokens_per_frame, [0.2, 0.4, 0.6, 0.8]) + assert torch.equal(right[:, 2:10], left[:, 24:32]) diff --git a/tests/test_ltx_vae.py b/tests/test_ltx_vae.py new file mode 100644 index 00000000..6d2cc8bb --- /dev/null +++ b/tests/test_ltx_vae.py @@ -0,0 +1,227 @@ +import json + +import pytest +import torch +from safetensors.torch import save_file + +from jasna.models.ltx_vae import diffusion_tiling, load_video_encoder +from jasna.models.ltx_vae.diffusion_video_decoder import DiffusionVideoDecoder +from jasna.models.ltx_vae.loader import decoder_kwargs, decoder_state_dict, encoder_kwargs, encoder_state_dict +from jasna.models.ltx_vae.tiling import DimensionSizeConfig, TileSizeConfig, scale_by_masks_1d +from jasna.models.ltx_vae.video_encoder import VideoEncoder + +TINY_ENCODER_CONFIG = { + "_class_name": "Encoder", + "dims": 3, + "in_channels": 3, + "out_channels": 8, + "blocks": [ + ["res_x", {"num_layers": 1}], + ["compress_space_res", {"multiplier": 2}], + ["compress_time_res", {"multiplier": 2}], + ["compress_all_res", {"multiplier": 2}], + ["compress_all_res", {"multiplier": 1}], + ], + "patch_size": 4, + "latent_log_var": "constant", + "norm_layer": "pixel_norm", + "spatial_padding_mode": "zeros", +} + +LTX25_DECODER_CONFIG = { + "_class_name": "NADiffusionDecoder", + "in_channels": 128, + "out_channels": 3, + "patch_size": 4, + "head_dim": 64, + "stage_channels": [2048, 1024, 512, 512, 256], + "stage_depths": [4, 6, 4, 2, 8], + "stage_kernels": [[3, 7, 7], [3, 7, 7], [3, 5, 5], [3, 5, 5], [11, 11, 11]], + "upsamples": [[[1, 2, 2], 2], [[2, 1, 1], 2], [[2, 2, 2], 1], [[2, 2, 2], 2]], + "spatial_padding_mode": "zeros", + "resampler_kind": "linear", + "stage5_kernel": [11, 11, 11], + "timestep_scale_multiplier": 1000.0, + "default_num_inference_steps": 1, +} + + +def _vae_config(encoder: dict, decoder: dict) -> dict: + return {"_class_name": "CausalDiffusionVAE", "encoder": encoder, "decoder": decoder, "model_output_type": "x0"} + + +def _meta_ltx25_decoder() -> DiffusionVideoDecoder: + with torch.device("meta"): + decoder = DiffusionVideoDecoder(**decoder_kwargs(_vae_config(TINY_ENCODER_CONFIG, LTX25_DECODER_CONFIG))) + return decoder.to(torch.bfloat16) + + +def test_decoder_state_dict_renames_splits_and_drops(): + qkv_weight = torch.arange(12.0).reshape(6, 2) + qkv_bias = torch.arange(6.0) + checkpoint = { + "decoder.t_embedder.mlp.0.weight": torch.zeros(1), + "decoder.t_embedder.mlp.2.bias": torch.zeros(1), + "decoder.det_stages.0.0.attn.qkv.weight": qkv_weight, + "decoder.diff_blocks.1.attn.qkv.bias": qkv_bias, + "decoder.diff_blocks.1.attn.proj.weight": torch.ones(2, 2), + "decoder.coarse_head.weight": torch.zeros(1), + "decoder.type_emb": torch.zeros(1), + "per_channel_statistics.mean-of-means": torch.zeros(1), + "encoder.conv_in.conv.weight": torch.zeros(1), + } + + state = decoder_state_dict(checkpoint) + + assert set(state) == { + "t_embedder.timestep_embedder.linear_1.weight", + "t_embedder.timestep_embedder.linear_2.bias", + "det_stages.0.0.attn.qkv.to_q.weight", + "det_stages.0.0.attn.qkv.to_k.weight", + "det_stages.0.0.attn.qkv.to_v.weight", + "diff_blocks.1.attn.qkv.to_q.bias", + "diff_blocks.1.attn.qkv.to_k.bias", + "diff_blocks.1.attn.qkv.to_v.bias", + "diff_blocks.1.attn.proj.weight", + "per_channel_statistics.mean-of-means", + } + assert torch.equal(state["det_stages.0.0.attn.qkv.to_k.weight"], qkv_weight[2:4]) + assert torch.equal(state["diff_blocks.1.attn.qkv.to_v.bias"], qkv_bias[4:]) + + +def test_decoder_state_dict_rejects_gated_checkpoints(): + with pytest.raises(ValueError, match="gated"): + decoder_state_dict({"decoder.diff_blocks.0.gate_msa": torch.zeros(1)}) + + +def test_ltx25_decoder_keys_match_module(): + decoder = _meta_ltx25_decoder() + expected = set(decoder.state_dict()) + assert "diff_blocks.7.attn.qkv.to_q.weight" in expected + assert "t_embedder.timestep_embedder.linear_2.weight" in expected + assert not any("stage4_upsample" in key for key in expected) + + +def test_encoder_state_dict_strips_prefix(): + state = encoder_state_dict( + { + "encoder.conv_in.conv.weight": torch.zeros(1), + "per_channel_statistics.std-of-means": torch.ones(1), + "decoder.conv_in.weight": torch.zeros(1), + } + ) + assert set(state) == {"conv_in.conv.weight", "per_channel_statistics.std-of-means"} + + +def test_unported_configs_are_rejected(): + with pytest.raises(ValueError, match="norm_layer"): + encoder_kwargs(_vae_config({**TINY_ENCODER_CONFIG, "norm_layer": "group_norm"}, LTX25_DECODER_CONFIG)) + with pytest.raises(ValueError, match="model_output_type"): + decoder_kwargs({**_vae_config(TINY_ENCODER_CONFIG, LTX25_DECODER_CONFIG), "model_output_type": "v"}) + with pytest.raises(ValueError, match="default_num_inference_steps"): + decoder_kwargs( + _vae_config(TINY_ENCODER_CONFIG, {**LTX25_DECODER_CONFIG, "default_num_inference_steps": 2}) + ) + + +def test_tiny_encoder_output_shape_and_frame_crop(): + torch.manual_seed(0) + encoder = VideoEncoder(**encoder_kwargs(_vae_config(TINY_ENCODER_CONFIG, LTX25_DECODER_CONFIG))).eval() + with torch.inference_mode(): + latent = encoder(torch.rand(1, 3, 18, 64, 96) * 2 - 1) + assert latent.shape == (1, 8, 3, 2, 3) + + +def test_load_video_encoder_from_checkpoint(tmp_path): + torch.manual_seed(0) + source = VideoEncoder(**encoder_kwargs(_vae_config(TINY_ENCODER_CONFIG, LTX25_DECODER_CONFIG))) + tensors = {f"encoder.{k}": v for k, v in source.state_dict().items() if not k.startswith("per_channel")} + tensors["per_channel_statistics.std-of-means"] = torch.full((8,), 2.0) + tensors["per_channel_statistics.mean-of-means"] = torch.full((8,), 0.5) + path = tmp_path / "vae.safetensors" + save_file(tensors, str(path), metadata={"config": json.dumps({"vae": _vae_config(TINY_ENCODER_CONFIG, {})})}) + + encoder = load_video_encoder(path, torch.device("cpu")) + + assert all(p.dtype == torch.bfloat16 and not p.requires_grad for p in encoder.parameters()) + assert torch.equal(encoder.conv_in.conv.weight, source.conv_in.conv.weight.bfloat16()) + assert torch.equal(encoder.per_channel_statistics.get_buffer("std-of-means"), torch.full((8,), 2.0).bfloat16()) + + +def test_ltx25_tiling_needs_minimum_budget(): + decoder = _meta_ltx25_decoder() + with pytest.raises(ValueError, match="Cannot fit"): + decoder.recommended_tiling_config(height=512, width=512, num_frames=121, free_bytes=3 << 30) + config = decoder.recommended_tiling_config(height=512, width=512, num_frames=121, free_bytes=4_200_000_000) + assert config.frames == DimensionSizeConfig(tile_size=80, overlap=40) + assert config.height.overlap == config.width.overlap == 160 + assert config.height.tile_size < 512 and config.width.tile_size < 512 + roomy = decoder.recommended_tiling_config(height=512, width=512, num_frames=121, free_bytes=24 << 30) + assert (roomy.frames.tile_size, roomy.height.tile_size, roomy.width.tile_size) == (128, 512, 512) + + +def test_ltx25_tile_schedule_masks_partition_unity(): + decoder = _meta_ltx25_decoder() + frames, height, width = 121, 512, 768 + config = TileSizeConfig( + frames=DimensionSizeConfig(tile_size=80, overlap=40), + height=DimensionSizeConfig(tile_size=320, overlap=160), + width=DimensionSizeConfig(tile_size=384, overlap=160), + ) + strides = [tuple(u.stride) for u in decoder.upsamples] + s4 = diffusion_tiling.stage4_thw_from_latent(strides, 16, 16, 24, drop_leading_frame=True) + tiles = diffusion_tiling.prepare_tile_schedule( + torch.Size([1, 128, *s4]), + config, + upsample3_stride=strides[3], + patch_size=decoder.patch_size, + min_tile_size=decoder.tile_min_sizes, + ) + + assert len({t.out_coords[2] for t in tiles}) > 1 + assert len({t.out_coords[3] for t in tiles}) > 1 + coverage = torch.zeros(1, 1, frames, height, width) + for tile in tiles: + region = torch.ones(1, 1, *(s.stop - s.start for s in tile.out_coords[2:])) + coverage[tile.out_coords] += scale_by_masks_1d(region, tile.masks_1d) + torch.testing.assert_close(coverage, torch.ones_like(coverage)) + + +def _tiny_decoder() -> DiffusionVideoDecoder: + torch.manual_seed(0) + return DiffusionVideoDecoder( + in_channels=8, + out_channels=3, + patch_size=4, + head_dim=16, + stage_channels=(128, 64, 32, 32, 16), + stage_depths=(1, 1, 1, 1, 1), + stage_kernels=((3, 3, 3),) * 5, + upsamples=(((1, 2, 2), 2), ((2, 1, 1), 2), ((2, 2, 2), 1), ((2, 2, 2), 2)), + stage5_kernel=(3, 3, 3), + stage5_channels=16, + t_emb_dim=32, + timestep_scale_multiplier=1000.0, + ).eval() + + +def test_tiny_decoder_multi_tile_decode_is_seeded_and_shaped(): + decoder = _tiny_decoder() + latent = torch.randn(1, 8, 3, 3, 4) + config = TileSizeConfig( + frames=DimensionSizeConfig(tile_size=16, overlap=8), + height=DimensionSizeConfig(tile_size=64, overlap=32), + width=DimensionSizeConfig(tile_size=64, overlap=32), + ) + + def decode(seed: int) -> list[torch.Tensor]: + with torch.inference_mode(): + return list(decoder.tiled_decode(latent, config, torch.Generator().manual_seed(seed))) + + chunks = decode(1) + pixels = torch.cat(chunks, dim=2) + assert len(chunks) > 1 + assert pixels.shape == (1, 3, 17, 96, 128) + assert torch.isfinite(pixels).all() + assert torch.equal(pixels, torch.cat(decode(1), dim=2)) + assert not torch.equal(pixels, torch.cat(decode(2), dim=2)) diff --git a/tests/test_session_config.py b/tests/test_session_config.py index 79061efe..b670f6f1 100644 --- a/tests/test_session_config.py +++ b/tests/test_session_config.py @@ -120,6 +120,25 @@ def test_cli_non_default_args_are_mapped() -> None: assert config.scene_detection is False +def test_cli_ltx_args_are_mapped() -> None: + defaults = _cli_config() + assert defaults.restoration_model_name == "basicvsrpp" + assert defaults.ltx_large_canvas is True + assert defaults.ltx_seed == 20260923 + + config = _cli_config(["--restoration-model-name", "ltx", "--no-ltx-large-canvas", "--ltx-seed", "7"]) + assert config.restoration_model_name == "ltx" + assert config.ltx_large_canvas is False + assert config.ltx_seed == 7 + + +def test_restoration_model_path_defaults_per_model() -> None: + from jasna.engine_paths import default_restoration_model_path + + assert default_restoration_model_path("basicvsrpp").name == "lada_mosaic_restoration_model_generic_v1.2.pth" + assert default_restoration_model_path("ltx").name == "ltx-restore" + + def test_gui_defaults_match_cli_defaults() -> None: cli = _cli_config() gui = _gui_config(AppSettings()) From 4c91a795cc279fba11a793199eb63878e39b0ed8 Mon Sep 17 00:00:00 2001 From: Kruk2 Date: Mon, 28 Sep 2026 19:27:34 +0200 Subject: [PATCH 15/48] fix(ltx): load the supporter model bundle Use shared model loading for all LTX components and report missing access before engine setup. Clarify model access after app updates. --- README.ja.md | 1 + README.md | 1 + README.zh.md | 1 + docs/en/gui.md | 1 + docs/ja/gui.md | 1 + docs/zh/gui.md | 1 + jasna/ltx/model_files.py | 34 ++++++++++++++++++++++++++++++++++ jasna/ltx/restore.py | 20 +------------------- jasna/ltx/sampler.py | 4 ++-- jasna/models/ltx_vae/loader.py | 9 ++++----- jasna/protection | 2 +- jasna/session_factory.py | 7 ++++--- tests/test_ltx_model_files.py | 32 ++++++++++++++++++++++++++++++++ 13 files changed, 84 insertions(+), 30 deletions(-) create mode 100644 jasna/ltx/model_files.py create mode 100644 tests/test_ltx_model_files.py diff --git a/README.ja.md b/README.ja.md index cc80b92e..a93c0708 100644 --- a/README.ja.md +++ b/README.ja.md @@ -160,6 +160,7 @@ Jasna の公開アプリケーションソースは AGPL-3.0 です。モデル 既存のキーは v0.11 でもこれらのモデルに使えます。LTX には対象バージョン用の アップグレードキーが必要です。対象条件は公式販売者にお問い合わせください。 +Jasna を更新しても、キーで使えるモデルは変わりません。 結果例: diff --git a/README.md b/README.md index d5b29f9f..835ef63f 100644 --- a/README.md +++ b/README.md @@ -163,6 +163,7 @@ Support pays for training extra models, mainly GPU rental and compute time for l Existing keys keep unlocking these models in v0.11. LTX versions require an upgraded key for the requested version; ask your official seller about eligibility. +Updating Jasna does not add or remove the models included in your key. Example results: diff --git a/README.zh.md b/README.zh.md index 887b67e9..a46be44e 100644 --- a/README.zh.md +++ b/README.zh.md @@ -154,6 +154,7 @@ Jasna 的公开应用程序源代码采用 AGPL-3.0。模型和捆绑依赖项 现有密钥在 v0.11 中仍可解锁这些模型。LTX 需要包含目标版本的升级密钥； 请向官方卖家咨询升级资格。 +更新 Jasna 不会增加或减少密钥中包含的模型。 结果示例: diff --git a/docs/en/gui.md b/docs/en/gui.md index f61a6eac..3bd75a28 100644 --- a/docs/en/gui.md +++ b/docs/en/gui.md @@ -104,3 +104,4 @@ after a driver update or when something misbehaves. If verification fails, check the exact purchase email or handle and paste the complete key again. Known unofficial keys get a separate warning. Existing keys still unlock UNet and SD; a newer LTX version needs an upgraded key for that version. + Updating Jasna does not change which models your key includes. diff --git a/docs/ja/gui.md b/docs/ja/gui.md index 07d4514c..fdfde9d7 100644 --- a/docs/ja/gui.md +++ b/docs/ja/gui.md @@ -101,3 +101,4 @@ Space キーで再生と一時停止を切り替え、左右矢印キーで 30 確認に失敗した場合は、購入時のメールアドレスまたはユーザー名を確認し、キー全体を貼り付け直してください。 確認済みの非公式キーは別の警告で通知します。既存のキーは引き続き UNet と SD に使えます。 新しい LTX バージョンには、そのバージョン用のアップグレードキーが必要です。 + Jasna を更新しても、キーで使えるモデルは変わりません。 diff --git a/docs/zh/gui.md b/docs/zh/gui.md index 5f93bebb..59634438 100644 --- a/docs/zh/gui.md +++ b/docs/zh/gui.md @@ -82,3 +82,4 @@ 密钥仅在 Buy Me a Coffee 和 Unifans (kruk2) 出售；其他卖家的密钥都是假的，无法使用。 若验证失败，请核对购买时使用的邮箱或账号，并重新粘贴完整密钥。已知的非官方密钥会单独提示。 现有密钥仍可解锁 UNet 和 SD；新的 LTX 版本需要包含该版本的升级密钥。 + 更新 Jasna 不会改变密钥中包含的模型。 diff --git a/jasna/ltx/model_files.py b/jasna/ltx/model_files.py new file mode 100644 index 00000000..b3303753 --- /dev/null +++ b/jasna/ltx/model_files.py @@ -0,0 +1,34 @@ +"""LTX model bundle discovery and tensor access.""" +from __future__ import annotations + +from dataclasses import dataclass +from pathlib import Path + +from jasna._frozen import is_frozen + + +def open_tensors(path: Path): + if not is_frozen() and path.suffix != ".enc" and path.is_file(): + from safetensors import safe_open + return safe_open(str(path), framework="pt", device="cpu") + from jasna.protection.protected_tensors import open_ltx_tensors + return open_ltx_tensors(path) + + +@dataclass(frozen=True) +class LtxModelFiles: + transformer: Path + vae: Path + tuned_decoder: Path + + @classmethod + def from_dir(cls, directory: Path) -> LtxModelFiles: + paths = [directory / f"{name}.safetensors" for name in ("transformer", "vae", "vae-decoder")] + paths = [path.with_name(path.name + ".enc") if is_frozen() or not path.is_file() else path for path in paths] + if any(path.suffix == ".enc" for path in paths): + from jasna.protection.protected_tensors import check_ltx_license + check_ltx_license() + missing = [str(path) for path in paths if not path.is_file()] + if missing: + raise FileNotFoundError(f"LTX model files missing: {', '.join(missing)}") + return cls(*paths) diff --git a/jasna/ltx/restore.py b/jasna/ltx/restore.py index 383ac066..e6d6c263 100644 --- a/jasna/ltx/restore.py +++ b/jasna/ltx/restore.py @@ -24,6 +24,7 @@ from jasna.ltx.camera import WINDOW_FRAMES from jasna.ltx.compose import Candidate, composite_frame, crop_to_canvas +from jasna.ltx.model_files import LtxModelFiles from jasna.ltx.plan import Region, TrackPlan, Window, crossfade_weights, feather_pixels, plan_video from jasna.ltx.regions import regions_per_frame from jasna.ltx.sampler import LtxTransformer @@ -36,25 +37,6 @@ FrameSource = Callable[[], FrameBatches] -@dataclass(frozen=True) -class LtxModelFiles: - transformer: Path - vae: Path - tuned_decoder: Path - - @classmethod - def from_dir(cls, directory: Path) -> LtxModelFiles: - files = cls( - transformer=directory / "transformer.safetensors", - vae=directory / "vae.safetensors", - tuned_decoder=directory / "vae-decoder.safetensors", - ) - missing = [str(path) for path in (files.transformer, files.vae, files.tuned_decoder) if not path.is_file()] - if missing: - raise FileNotFoundError(f"LTX model files missing: {', '.join(missing)}") - return files - - class Cancelled(Exception): pass diff --git a/jasna/ltx/sampler.py b/jasna/ltx/sampler.py index eb8ec9ae..523dede7 100644 --- a/jasna/ltx/sampler.py +++ b/jasna/ltx/sampler.py @@ -15,10 +15,10 @@ from pathlib import Path import torch -from safetensors import safe_open from jasna.ltx import transformer as T from jasna.ltx.blocks import BlockStore +from jasna.ltx.model_files import open_tensors from jasna.ltx.plan import FUSION_RAMP, LEFT_SHARED_LATENTS, RIGHT_SHARED_LATENTS logger = logging.getLogger(__name__) @@ -71,7 +71,7 @@ class LtxTransformer: def __init__(self, path: Path, device: torch.device, *, sigmas: Sequence[float] = LOWTAIL15_SIGMAS) -> None: self.device = device self.sigmas = torch.tensor(sigmas, dtype=torch.float32) - with safe_open(str(path), framework="pt", device="cpu") as handle: + with open_tensors(path) as handle: keys = list(handle.keys()) top_keys = [k for k in keys if not k.startswith("blocks.") and k != "prompt_context"] top = {k: handle.get_tensor(k).to(device) for k in top_keys} diff --git a/jasna/models/ltx_vae/loader.py b/jasna/models/ltx_vae/loader.py index eeab3d8c..d888e880 100644 --- a/jasna/models/ltx_vae/loader.py +++ b/jasna/models/ltx_vae/loader.py @@ -7,11 +7,10 @@ from pathlib import Path from typing import TypeVar -import safetensors import torch -from safetensors.torch import load_file from torch import nn +from jasna.ltx.model_files import open_tensors from jasna.models.ltx_vae.diffusion_video_decoder import DiffusionVideoDecoder from jasna.models.ltx_vae.video_encoder import VideoEncoder @@ -25,7 +24,7 @@ def read_vae_config(vae_path: Path) -> dict: - with safetensors.safe_open(str(vae_path), framework="pt") as handle: + with open_tensors(vae_path) as handle: config = json.loads(handle.metadata()["config"])["vae"] if config.get("_class_name") != _VAE_CLASS_NAME: raise ValueError(f"{vae_path}: VAE class {config.get('_class_name')!r} is not supported") @@ -78,7 +77,7 @@ def decoder_kwargs(config: dict) -> dict: def _read_tensors(path: Path, prefixes: tuple[str, ...]) -> dict[str, torch.Tensor]: - with safetensors.safe_open(str(path), framework="pt", device="cpu") as handle: + with open_tensors(path) as handle: return {key: handle.get_tensor(key) for key in handle.keys() if key.startswith(prefixes)} @@ -142,7 +141,7 @@ def load_video_decoder(vae_path: Path, tuned_decoder_path: Path, device: torch.d state = decoder_state_dict(_read_tensors(vae_path, ("decoder.", "per_channel_statistics."))) decoder = _build_bf16(lambda: DiffusionVideoDecoder(**kwargs), state, device) - tuned = load_file(str(tuned_decoder_path)) + tuned = _read_tensors(tuned_decoder_path, ("",)) own = decoder.state_dict() unknown = sorted(set(tuned) - set(own)) if unknown: diff --git a/jasna/protection b/jasna/protection index 37a800ff..474bcca9 160000 --- a/jasna/protection +++ b/jasna/protection @@ -1 +1 @@ -Subproject commit 37a800fffb365be0cc54c33176141b360ad59a1b +Subproject commit 474bcca944038908f20b82d6a9a4b2985385bae4 diff --git a/jasna/session_factory.py b/jasna/session_factory.py index 483c1666..d4b4a9a8 100644 --- a/jasna/session_factory.py +++ b/jasna/session_factory.py @@ -19,7 +19,7 @@ if TYPE_CHECKING: import torch - from jasna.ltx.restore import LtxModelFiles + from jasna.ltx.model_files import LtxModelFiles from jasna.media.splice import SplicePlan from jasna.mosaic.rfdetr import RfDetrMosaicDetectionModel from jasna.mosaic.yolo import YoloMosaicDetectionModel @@ -201,12 +201,13 @@ def _build_ltx_session( ) -> RestorationSession: from jasna.accelerator import is_nvidia_device from jasna.engine_compiler import EngineCompilationRequest, ensure_engines_compiled - from jasna.ltx.restore import LtxModelFiles + from jasna.ltx.model_files import LtxModelFiles if not is_nvidia_device(device): raise ValueError("LTX restoration needs an NVIDIA GPU") if config.secondary_restoration != "none" or config.denoise_strength != "none": raise ValueError("LTX restoration does not support secondary restoration or denoise") + files = LtxModelFiles.from_dir(config.restoration_model_path) ensure_engines_compiled( EngineCompilationRequest( device=str(device), @@ -221,7 +222,7 @@ def _build_ltx_session( return RestorationSession( device=device, restoration_pipeline=None, - ltx_files=LtxModelFiles.from_dir(config.restoration_model_path), + ltx_files=files, ) diff --git a/tests/test_ltx_model_files.py b/tests/test_ltx_model_files.py new file mode 100644 index 00000000..123c0ff8 --- /dev/null +++ b/tests/test_ltx_model_files.py @@ -0,0 +1,32 @@ +from unittest.mock import Mock + +import pytest +import torch +from safetensors.torch import save_file + +from jasna.ltx.model_files import LtxModelFiles, open_tensors + + +def test_source_model_bundle_and_reader(tmp_path): + value = torch.arange(6).reshape(2, 3) + for name in ("transformer", "vae", "vae-decoder"): + save_file({"weight": value}, str(tmp_path / f"{name}.safetensors"), metadata={"config": "example"}) + files = LtxModelFiles.from_dir(tmp_path) + with open_tensors(files.transformer) as handle: + assert list(handle.keys()) == ["weight"] + assert handle.metadata() == {"config": "example"} + assert torch.equal(handle.get_tensor("weight"), value) + + +def test_model_check_precedes_engine_compilation(tmp_path, monkeypatch): + from factories import session_config + from jasna import accelerator, engine_compiler, session_factory + + compile_engines = Mock() + monkeypatch.setattr(accelerator, "is_nvidia_device", lambda device: True) + monkeypatch.setattr(engine_compiler, "ensure_engines_compiled", compile_engines) + monkeypatch.setattr(LtxModelFiles, "from_dir", Mock(side_effect=ValueError("model unavailable"))) + config = session_config(restoration_model_name="ltx", restoration_model_path=tmp_path) + with pytest.raises(ValueError, match="model unavailable"): + session_factory._build_ltx_session(config, torch.device("cuda:0"), log_callback=None) + compile_engines.assert_not_called() From 946898d54460d7d511b0e0d75580366163915dd2 Mon Sep 17 00:00:00 2001 From: Kruk2 Date: Mon, 28 Sep 2026 19:32:00 +0200 Subject: [PATCH 16/48] docs: remove redundant license update note --- README.ja.md | 1 - README.md | 1 - README.zh.md | 1 - 3 files changed, 3 deletions(-) diff --git a/README.ja.md b/README.ja.md index a93c0708..cc80b92e 100644 --- a/README.ja.md +++ b/README.ja.md @@ -160,7 +160,6 @@ Jasna の公開アプリケーションソースは AGPL-3.0 です。モデル 既存のキーは v0.11 でもこれらのモデルに使えます。LTX には対象バージョン用の アップグレードキーが必要です。対象条件は公式販売者にお問い合わせください。 -Jasna を更新しても、キーで使えるモデルは変わりません。 結果例: diff --git a/README.md b/README.md index 835ef63f..d5b29f9f 100644 --- a/README.md +++ b/README.md @@ -163,7 +163,6 @@ Support pays for training extra models, mainly GPU rental and compute time for l Existing keys keep unlocking these models in v0.11. LTX versions require an upgraded key for the requested version; ask your official seller about eligibility. -Updating Jasna does not add or remove the models included in your key. Example results: diff --git a/README.zh.md b/README.zh.md index a46be44e..887b67e9 100644 --- a/README.zh.md +++ b/README.zh.md @@ -154,7 +154,6 @@ Jasna 的公开应用程序源代码采用 AGPL-3.0。模型和捆绑依赖项 现有密钥在 v0.11 中仍可解锁这些模型。LTX 需要包含目标版本的升级密钥； 请向官方卖家咨询升级资格。 -更新 Jasna 不会增加或减少密钥中包含的模型。 结果示例: From 5405e2b8ed5c67e903499f7c5d2cb3b68bc44ca1 Mon Sep 17 00:00:00 2001 From: Kruk2 Date: Tue, 29 Sep 2026 00:24:19 +0200 Subject: [PATCH 17/48] docs: drop unreleased LTX key note from README --- README.ja.md | 3 --- README.md | 3 --- README.zh.md | 3 --- 3 files changed, 9 deletions(-) diff --git a/README.ja.md b/README.ja.md index cc80b92e..376ed584 100644 --- a/README.ja.md +++ b/README.ja.md @@ -158,9 +158,6 @@ Jasna の公開アプリケーションソースは AGPL-3.0 です。モデル - **unet-4x** セカンダリアップスケーラー。よりシャープな 256->1024 復元用です。 - **SD 1.5 画像復元**。実験的な静止画モデルです。 -既存のキーは v0.11 でもこれらのモデルに使えます。LTX には対象バージョン用の -アップグレードキーが必要です。対象条件は公式販売者にお問い合わせください。 - 結果例: - [unet-4x / セカンダリ復元の例（SLS Discord）](https://discord.com/channels/1196376491815092265/1199059436199759943/1516497879684874260) diff --git a/README.md b/README.md index d5b29f9f..5dba0a8a 100644 --- a/README.md +++ b/README.md @@ -161,9 +161,6 @@ Support pays for training extra models, mainly GPU rental and compute time for l - **unet-4x** secondary upscaler for sharper 256->1024 restoration. - **SD 1.5 image restoration**, the experimental still-image model. -Existing keys keep unlocking these models in v0.11. LTX versions require an upgraded key -for the requested version; ask your official seller about eligibility. - Example results: - [unet-4x / secondary restoration examples on SLS Discord](https://discord.com/channels/1196376491815092265/1199059436199759943/1516497879684874260) diff --git a/README.zh.md b/README.zh.md index 887b67e9..3308548b 100644 --- a/README.zh.md +++ b/README.zh.md @@ -152,9 +152,6 @@ Jasna 的公开应用程序源代码采用 AGPL-3.0。模型和捆绑依赖项 - **unet-4x** 二级放大模型，用于更清晰的 256->1024 修复。 - **SD 1.5 图像修复**，实验性静态图像模型。 -现有密钥在 v0.11 中仍可解锁这些模型。LTX 需要包含目标版本的升级密钥； -请向官方卖家咨询升级资格。 - 结果示例: - [unet-4x / 二级修复示例（SLS Discord）](https://discord.com/channels/1196376491815092265/1199059436199759943/1516497879684874260) From e15d00fe7cbdae9dcdc6a6a66e3b8093ae657cb9 Mon Sep 17 00:00:00 2001 From: Kruk2 Date: Tue, 29 Sep 2026 00:26:44 +0200 Subject: [PATCH 18/48] docs: drop unreleased LTX key note from GUI docs --- docs/en/gui.md | 3 +-- docs/ja/gui.md | 3 +-- docs/zh/gui.md | 1 - 3 files changed, 2 insertions(+), 5 deletions(-) diff --git a/docs/en/gui.md b/docs/en/gui.md index 3bd75a28..8a3d9c4a 100644 --- a/docs/en/gui.md +++ b/docs/en/gui.md @@ -102,6 +102,5 @@ after a driver update or when something misbehaves. Keys are sold only on Buy Me a Coffee and Unifans (kruk2); keys from other sellers are fake and do not work. If verification fails, check the exact purchase email or handle and paste the complete - key again. Known unofficial keys get a separate warning. Existing keys still unlock - UNet and SD; a newer LTX version needs an upgraded key for that version. + key again. Known unofficial keys get a separate warning. Updating Jasna does not change which models your key includes. diff --git a/docs/ja/gui.md b/docs/ja/gui.md index fdfde9d7..dd43026c 100644 --- a/docs/ja/gui.md +++ b/docs/ja/gui.md @@ -99,6 +99,5 @@ Space キーで再生と一時停止を切り替え、左右矢印キーで 30 キーは Buy Me a Coffee と Unifans (kruk2) でのみ販売しています。他の販売者の キーは偽物で、使えません。 確認に失敗した場合は、購入時のメールアドレスまたはユーザー名を確認し、キー全体を貼り付け直してください。 - 確認済みの非公式キーは別の警告で通知します。既存のキーは引き続き UNet と SD に使えます。 - 新しい LTX バージョンには、そのバージョン用のアップグレードキーが必要です。 + 確認済みの非公式キーは別の警告で通知します。 Jasna を更新しても、キーで使えるモデルは変わりません。 diff --git a/docs/zh/gui.md b/docs/zh/gui.md index 59634438..a434679c 100644 --- a/docs/zh/gui.md +++ b/docs/zh/gui.md @@ -81,5 +81,4 @@ - **许可证**: 在头部输入你的支持者密钥；状态标签会显示是否已激活。 密钥仅在 Buy Me a Coffee 和 Unifans (kruk2) 出售；其他卖家的密钥都是假的，无法使用。 若验证失败，请核对购买时使用的邮箱或账号，并重新粘贴完整密钥。已知的非官方密钥会单独提示。 - 现有密钥仍可解锁 UNet 和 SD；新的 LTX 版本需要包含该版本的升级密钥。 更新 Jasna 不会改变密钥中包含的模型。 From ebd33290cb688e7ca9ed1dbb4d025abb98efade1 Mon Sep 17 00:00:00 2001 From: Kruk2 Date: Tue, 29 Sep 2026 11:22:31 +0200 Subject: [PATCH 19/48] fix(ltx): finish LTX videos without a decode crash --- jasna/ltx/restore.py | 2 +- tests/test_ltx_restore.py | 13 +++++++++++++ 2 files changed, 14 insertions(+), 1 deletion(-) create mode 100644 tests/test_ltx_restore.py diff --git a/jasna/ltx/restore.py b/jasna/ltx/restore.py index e6d6c263..a47d007b 100644 --- a/jasna/ltx/restore.py +++ b/jasna/ltx/restore.py @@ -176,7 +176,7 @@ def _decode_window(decoder, latent: torch.Tensor, seed: int, device: torch.devic with torch.inference_mode(): video = decode_latent(decoder, latent.to(device), free_bytes=free, generator=generator) video = ((video + 1.0) / 2.0).clamp(0.0, 1.0)[0].float() - return video.mul_(255.0).round_().to(torch.uint8).permute(1, 0, 2, 3).contiguous() + return video.mul_(255.0).round_().to(torch.uint8).permute(1, 0, 2, 3).contiguous() def compose( diff --git a/tests/test_ltx_restore.py b/tests/test_ltx_restore.py new file mode 100644 index 00000000..d2617a7d --- /dev/null +++ b/tests/test_ltx_restore.py @@ -0,0 +1,13 @@ +import torch + +from jasna.ltx import restore + + +def test_decode_window_converts_a_float32_decode(monkeypatch): + decoded = torch.tensor([-1.0, 0.0, 1.0]).view(1, 1, 3, 1, 1).expand(1, 3, 3, 1, 1).contiguous() + monkeypatch.setattr(restore, "decode_latent", lambda decoder, latent, free_bytes, generator: decoded.clone()) + monkeypatch.setattr(torch.cuda, "mem_get_info", lambda device: (1, 1)) + frames = restore._decode_window(None, torch.zeros(1, 128, 1, 1, 1), 0, torch.device("cpu")) + assert frames.dtype == torch.uint8 and frames.shape == (3, 3, 1, 1) + assert frames[:, 0, 0, 0].tolist() == [0, 128, 255] + From 6eb357807d3128bcfe8885e6c927e74b9cfb6674 Mon Sep 17 00:00:00 2001 From: Kruk2 Date: Tue, 29 Sep 2026 11:22:31 +0200 Subject: [PATCH 20/48] feat(ltx): add fast mode and faster LTX inference --ltx-fast uses an NVFP4 model on RTX 50-series GPUs. SageAttention, torch.compile and NATTEN are used when installed. --- jasna/cli_help.py | 4 + jasna/gui/video_session.py | 1 + jasna/ltx/model_files.py | 6 +- jasna/ltx/sampler.py | 17 +++- jasna/ltx/transformer.py | 124 ++++++++++++++++++++++++++-- jasna/main.py | 7 ++ jasna/models/ltx_vae/transformer.py | 26 +++++- jasna/session_config.py | 1 + jasna/session_factory.py | 6 +- tests/factories.py | 1 + tests/test_ltx_model_files.py | 8 +- tests/test_ltx_transformer.py | 52 ++++++++++++ 12 files changed, 236 insertions(+), 17 deletions(-) diff --git a/jasna/cli_help.py b/jasna/cli_help.py index 876de5d9..b1e76742 100644 --- a/jasna/cli_help.py +++ b/jasna/cli_help.py @@ -16,6 +16,10 @@ "Restoration model file (basicvsrpp) or folder (ltx). Default: the model in model_weights/." ), "ltx_seed": "ltx: noise seed. Another seed gives another take on the same restoration. (default: %(default)s)", + "ltx_fast": ( + "ltx: about 1.4x faster with slightly less detail. Needs an RTX 50-series (Blackwell) GPU " + "and the fast model file. (default: %(default)s)" + ), "ltx_large_canvas": ( "ltx: restore large mosaics at 768 px instead of 512 px. Sharper, about 3x slower for " "those mosaics and needs more VRAM. (default: %(default)s)" diff --git a/jasna/gui/video_session.py b/jasna/gui/video_session.py index 546dc590..f565e2df 100644 --- a/jasna/gui/video_session.py +++ b/jasna/gui/video_session.py @@ -72,6 +72,7 @@ def video_session_config( restoration_model_path=default_restoration_model_path("basicvsrpp"), ltx_large_canvas=True, ltx_seed=LTX_DEFAULT_SEED, + ltx_fast=False, compile_basicvsrpp=bool(settings.compile_basicvsrpp), max_clip_size=int(settings.max_clip_size), temporal_overlap=int(settings.temporal_overlap), diff --git a/jasna/ltx/model_files.py b/jasna/ltx/model_files.py index b3303753..c08284cf 100644 --- a/jasna/ltx/model_files.py +++ b/jasna/ltx/model_files.py @@ -22,8 +22,10 @@ class LtxModelFiles: tuned_decoder: Path @classmethod - def from_dir(cls, directory: Path) -> LtxModelFiles: - paths = [directory / f"{name}.safetensors" for name in ("transformer", "vae", "vae-decoder")] + def from_dir(cls, directory: Path, *, fast: bool) -> LtxModelFiles: + """The model bundle; ``fast`` picks the NVFP4 transformer over the INT8 one.""" + transformer = "transformer-fast" if fast else "transformer" + paths = [directory / f"{name}.safetensors" for name in (transformer, "vae", "vae-decoder")] paths = [path.with_name(path.name + ".enc") if is_frozen() or not path.is_file() else path for path in paths] if any(path.suffix == ".enc" for path in paths): from jasna.protection.protected_tensors import check_ltx_license diff --git a/jasna/ltx/sampler.py b/jasna/ltx/sampler.py index 523dede7..d69819af 100644 --- a/jasna/ltx/sampler.py +++ b/jasna/ltx/sampler.py @@ -10,6 +10,7 @@ from __future__ import annotations +import importlib.util import logging from collections.abc import Sequence from pathlib import Path @@ -85,6 +86,8 @@ def __init__(self, path: Path, device: torch.device, *, sigmas: Sequence[float] ] self.blocks = BlockStore(blocks, device, resident=self._resident_blocks(blocks)) self._ropes: dict[tuple[int, int, int], T.Rope] = {} + logger.info("LTX self-attention: %s", "SageAttention 2" if T.enable_sage_attention() else "SDPA") + self._block_forward = compiled_block_forward() def _resident_blocks(self, blocks: list[dict[str, torch.Tensor]]) -> int: sizes = [sum(t.numel() * t.element_size() for t in b.values()) for b in blocks] @@ -159,10 +162,10 @@ def _denoise( for j in range(len(xs)): if STG_SCALE and index >= STG_BLOCK: source = xs[j] if index == STG_BLOCK else perturbed[j] - perturbed[j] = T.block_forward( + perturbed[j] = self._block_forward( weights, source, cond, self.context, rope, skip_self_attention=index == STG_BLOCK ) - xs[j] = T.block_forward(weights, xs[j], cond, self.context, rope) + xs[j] = self._block_forward(weights, xs[j], cond, self.context, rope) self.blocks.release(index) out = [] for lat, x, xp in zip(latents, xs, perturbed): @@ -174,6 +177,16 @@ def _denoise( return out +def compiled_block_forward(): + """``transformer.block_forward`` compiled with static shapes when Triton is available + (one graph per canvas and per STG variant, shared by all blocks), else eager.""" + if importlib.util.find_spec("triton") is None: + return T.block_forward + torch._dynamo.config.cache_size_limit = max(torch._dynamo.config.cache_size_limit, 64) + torch._dynamo.config.accumulated_cache_size_limit = max(torch._dynamo.config.accumulated_cache_size_limit, 1024) + return torch.compile(T.block_forward, dynamic=False) + + def _euler(sample: torch.Tensor, denoised: torch.Tensor, sigmas: torch.Tensor, step: int) -> torch.Tensor: sigma = sigmas[step] velocity = ((sample.float() - denoised.float()) / sigma.item()).to(sample.dtype) diff --git a/jasna/ltx/transformer.py b/jasna/ltx/transformer.py index de5219c4..1d7165b2 100644 --- a/jasna/ltx/transformer.py +++ b/jasna/ltx/transformer.py @@ -8,13 +8,16 @@ every AdaLN modulation is one row per half instead of one per token. Blocks take their weights as a dict so the executor can hand in resident tensors or views -into a streaming slot. A Linear is stored either as ``.weight`` (bf16) or as ConvRot -W8A8 (``.qdata`` int8 rotated codes, ``.scales`` fp32 per output channel). +into a streaming slot. A Linear is stored as ``.weight`` (bf16), as ConvRot W8A8 +(``.qdata`` int8 rotated codes, ``.scales`` fp32 per output channel), or as NVFP4 +W4A4 (``.weight`` uint8 packed E2M1, high nibble first, ``.weight_scale`` +float8_e4m3fn block scales in the cuBLAS 128x4 tiled layout, ``.weight_scale_2`` fp32). """ from __future__ import annotations import functools +import importlib.util import math from collections.abc import Mapping from dataclasses import dataclass @@ -37,6 +40,13 @@ LATENT_SPATIAL_SCALE = 32 ROT_SIZE = 256 ACT_QMAX = 127 +FP4_BLOCK = 16 +FP4_MAX = 6.0 +FP8_E4M3_MAX = 448.0 +# E2M1 magnitudes 0, .5, 1, 1.5, 2, 3, 4, 6: midpoints rounding to the even code sit below +# the tie, the others above, so a tie lands on the even code (round half to even). +_FP4_TIES_DOWN = (0.25, 1.25, 2.5, 5.0) +_FP4_TIES_UP = (0.75, 1.75, 3.5) CHUNK_ROWS = 4096 SDPA_PRIORITY = [ SDPBackend.CUDNN_ATTENTION, @@ -104,12 +114,109 @@ def int8_linear( return out.reshape(*x.shape[:-1], out_features) +def fp4_codes(x: torch.Tensor) -> torch.Tensor: + """4-bit E2M1 codes (sign bit 3), round half to even, saturating at +-6.""" + magnitude = x.abs() + index = sum((magnitude > tie).to(torch.uint8) for tie in _FP4_TIES_DOWN) + index = index + sum((magnitude >= tie).to(torch.uint8) for tie in _FP4_TIES_UP) + return index | ((x < 0).to(torch.uint8) << 3) + + +def swizzle_block_scales(scales: torch.Tensor) -> torch.Tensor: + """``(rows, k/16)`` float8 block scales -> zero-padded cuBLAS 128x4 tiled layout.""" + rows, cols = scales.shape + padded = F.pad(scales.view(torch.uint8), (0, -cols % 4, 0, -rows % 128)) + tiled_rows, tiled_cols = padded.shape + tiles = padded.view(tiled_rows // 128, 4, 32, tiled_cols // 4, 4).permute(0, 3, 2, 1, 4) + return tiles.reshape(tiled_rows, tiled_cols).view(torch.float8_e4m3fn) + + +def fp4_tensor_scale(x: torch.Tensor) -> torch.Tensor: + """Dynamic per-tensor decode scale ``amax / (6 * 448)`` (1 for an all-zero tensor).""" + scale = x.abs().amax().float() / (FP4_MAX * FP8_E4M3_MAX) + return torch.where(scale > 0, scale, torch.ones_like(scale)) + + +def quantize_fp4(x: torch.Tensor, tensor_scale: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + """``(rows, k)`` -> (packed uint8 ``(rows, k/2)``, tiled block scales), one block scale + per 16 values under the given per-tensor decode scale.""" + xf = x.float() + blocks = xf.view(x.shape[0], -1, FP4_BLOCK) + block_scales = (blocks.abs().amax(-1) / FP4_MAX / tensor_scale).clamp(max=FP8_E4M3_MAX).to(torch.float8_e4m3fn) + divisor = block_scales.float() * tensor_scale + divisor = torch.where(divisor > 0, divisor, torch.ones_like(divisor)) + codes = fp4_codes(blocks / divisor.unsqueeze(-1)).view(x.shape[0], -1, 2) + packed = (codes[..., 0] << 4) | codes[..., 1] + return packed, swizzle_block_scales(block_scales) + + +def nvfp4_linear( + x: torch.Tensor, + weight: torch.Tensor, + block_scales: torch.Tensor, + tensor_scale: torch.Tensor, + bias: torch.Tensor | None, +) -> torch.Tensor: + """NVFP4 W4A4: activations quantized per call (one dynamic tensor scale for the call), then + the block-scaled FP4 tensor-core GEMM. Both operands pack high nibble first; the block scale + covers whole pairs, so the order only permutes each dot product's terms.""" + out_features = weight.shape[0] + flat = x.reshape(-1, weight.shape[1] * 2) + out = torch.empty(flat.shape[0], out_features, dtype=x.dtype, device=x.device) + act_scale = fp4_tensor_scale(flat) + for start in range(0, flat.shape[0], CHUNK_ROWS): + chunk = flat[start : start + CHUNK_ROWS] + packed, act_blocks = quantize_fp4(chunk, act_scale) + out[start : start + chunk.shape[0]] = F.scaled_mm( + packed.view(torch.float4_e2m1fn_x2), + weight.view(torch.float4_e2m1fn_x2).t(), + scale_a=[act_blocks, act_scale.reshape(1)], + scale_recipe_a=[F.ScalingType.BlockWise1x16, F.ScalingType.TensorWise], + scale_b=[block_scales, tensor_scale.reshape(1)], + scale_recipe_b=[F.ScalingType.BlockWise1x16, F.ScalingType.TensorWise], + swizzle_a=[F.SwizzleType.SWIZZLE_32_4_4, F.SwizzleType.NO_SWIZZLE], + swizzle_b=[F.SwizzleType.SWIZZLE_32_4_4, F.SwizzleType.NO_SWIZZLE], + bias=bias, + output_dtype=x.dtype, + ) + return out.reshape(*x.shape[:-1], out_features) + + def linear(w: Weights, name: str, x: torch.Tensor) -> torch.Tensor: bias = w.get(f"{name}.bias") weight = w.get(f"{name}.weight") - if weight is not None: - return F.linear(x, weight, bias) - return int8_linear(x, w[f"{name}.qdata"], w[f"{name}.scales"], bias) + if weight is None: + return int8_linear(x, w[f"{name}.qdata"], w[f"{name}.scales"], bias) + if weight.dtype == torch.uint8: + return nvfp4_linear(x, weight, w[f"{name}.weight_scale"], w[f"{name}.weight_scale_2"], bias) + return F.linear(x, weight, bias) + + +_sage_attention = None + + +def enable_sage_attention() -> bool: + """Route video self-attention through SageAttention 2 (an opaque op, so it compiles as + one call) when it is installed. Returns whether it is in use.""" + global _sage_attention + if _sage_attention is None and importlib.util.find_spec("sageattention") is not None: + from sageattention import sageattn + + @torch.library.custom_op("jasna::sage_attention", mutates_args=()) + def sage_attention(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor) -> torch.Tensor: + return sageattn(q, k, v, tensor_layout="HND", is_causal=False).contiguous() + + sage_attention.register_fake(lambda q, k, v: torch.empty_like(q, memory_format=torch.contiguous_format)) + _sage_attention = sage_attention + return _sage_attention is not None + + +def self_attention(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor) -> torch.Tensor: + """Video self-attention ``[1, heads, tokens, d]``: SageAttention 2 once enabled.""" + if _sage_attention is not None and q.is_cuda: + return _sage_attention(q, k, v) + with sdpa_kernel(SDPA_PRIORITY, set_priority=True): + return F.scaled_dot_product_attention(q, k, v) def rms_norm(x: torch.Tensor, weight: torch.Tensor | None = None) -> torch.Tensor: @@ -251,8 +358,11 @@ def _attention( q = q.unflatten(-1, (heads, -1)).transpose(1, 2) k = k.unflatten(-1, (heads, -1)).transpose(1, 2) v_heads = v.unflatten(-1, (heads, -1)).transpose(1, 2) - with sdpa_kernel(SDPA_PRIORITY, set_priority=True): - out = F.scaled_dot_product_attention(q, k, v_heads) + if rope is not None: + out = self_attention(q, k, v_heads) + else: + with sdpa_kernel(SDPA_PRIORITY, set_priority=True): + out = F.scaled_dot_product_attention(q, k, v_heads) out = out.transpose(1, 2).flatten(2) gates = 2.0 * torch.sigmoid(linear(w, f"{prefix}.to_gate_logits", x)) out = (out.unflatten(-1, (heads, -1)) * gates.unsqueeze(-1)).flatten(2) diff --git a/jasna/main.py b/jasna/main.py index 0dbc3897..1d2a7e18 100644 --- a/jasna/main.py +++ b/jasna/main.py @@ -45,6 +45,7 @@ def _session_config_from_args( restoration_model_path=restoration_model_path, ltx_large_canvas=bool(args.ltx_large_canvas), ltx_seed=int(args.ltx_seed), + ltx_fast=bool(args.ltx_fast), compile_basicvsrpp=bool(args.compile_basicvsrpp), max_clip_size=int(args.max_clip_size), temporal_overlap=int(args.temporal_overlap), @@ -215,6 +216,12 @@ def build_parser() -> argparse.ArgumentParser: default=LTX_DEFAULT_SEED, help=CLI_HELP["ltx_seed"], ) + restoration.add_argument( + "--ltx-fast", + default=False, + action=argparse.BooleanOptionalAction, + help=CLI_HELP["ltx_fast"], + ) restoration.add_argument( "--compile-basicvsrpp", default=True, diff --git a/jasna/models/ltx_vae/transformer.py b/jasna/models/ltx_vae/transformer.py index 068234a8..27f02eb2 100644 --- a/jasna/models/ltx_vae/transformer.py +++ b/jasna/models/ltx_vae/transformer.py @@ -1,12 +1,15 @@ """Neighborhood-attention transformer blocks of the diffusion VAE decoder, ``chunked_eager`` path only. Ported from ``ltx_core/model/video_vae/transformer/`` (layers, qkv, rope_math, det_attn_rope, -attention, swiglu, blocks, chunked/*): eager tiled-SDPA neighborhood attention, torch SwiGLU, +attention, swiglu, blocks, chunked/*): NATTEN neighborhood attention when installed, else eager +tiled SDPA (same window semantics), torch SwiGLU, W-chunked diffusion residual with halos and deferred stage-4 context inject. """ from __future__ import annotations +import functools +import importlib.util import math import numpy as np @@ -132,12 +135,27 @@ def forward(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, torch.Te return self.to_q(x), self.to_k(x), self.to_v(x) -def eager_sdpa_attention( +@functools.cache +def _fused_na3d(): + """NATTEN's fused ``na3d`` when installed, else None.""" + if importlib.util.find_spec("natten") is None: + return None + import natten + + return natten.na3d + + +def neighborhood_attention( attn: NeighborhoodAttention3D, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor ) -> torch.Tensor: + """NATTEN's fused kernel on CUDA when installed, else the pure-torch ``na3d`` + (same window semantics). Q is already scaled.""" if q.dtype != v.dtype or k.dtype != v.dtype: q = q.to(dtype=v.dtype) k = k.to(dtype=v.dtype) + fused = _fused_na3d() + if fused is not None and q.is_cuda: + return fused(q, k, v, kernel_size=attn.kernel_size, scale=1.0) return na3d(q, k, v, kernel_size=attn.kernel_size) @@ -183,7 +201,7 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: q = q * self.scale q = apply_tiled_abs_rope(q, self) k = apply_tiled_abs_rope(k, self) - out = eager_sdpa_attention(self, q.contiguous(), k.contiguous(), v.contiguous()) + out = neighborhood_attention(self, q.contiguous(), k.contiguous(), v.contiguous()) out = out.reshape(batch, t, h, w, self.dim) return self.proj(out) @@ -291,7 +309,7 @@ def _attn_on_w_slab( v = F.linear(y, attn.qkv.to_v.weight, attn.qkv.to_v.bias).view(head_shape) del y - out = eager_sdpa_attention(attn, q.contiguous(), k.contiguous(), v.contiguous()) + out = neighborhood_attention(attn, q.contiguous(), k.contiguous(), v.contiguous()) del q, k, v return F.linear(out.reshape(batch, t, h, ext_w, attn.dim), attn.proj.weight, attn.proj.bias) diff --git a/jasna/session_config.py b/jasna/session_config.py index cdc62838..6ba80f97 100644 --- a/jasna/session_config.py +++ b/jasna/session_config.py @@ -39,6 +39,7 @@ class SessionConfig: restoration_model_path: Path ltx_large_canvas: bool ltx_seed: int + ltx_fast: bool compile_basicvsrpp: bool max_clip_size: int temporal_overlap: int diff --git a/jasna/session_factory.py b/jasna/session_factory.py index d4b4a9a8..e45d26cb 100644 --- a/jasna/session_factory.py +++ b/jasna/session_factory.py @@ -199,6 +199,8 @@ def build_restoration_session( def _build_ltx_session( config: SessionConfig, device: "torch.device", *, log_callback: Callable[[str], None] | None ) -> RestorationSession: + import torch + from jasna.accelerator import is_nvidia_device from jasna.engine_compiler import EngineCompilationRequest, ensure_engines_compiled from jasna.ltx.model_files import LtxModelFiles @@ -207,7 +209,9 @@ def _build_ltx_session( raise ValueError("LTX restoration needs an NVIDIA GPU") if config.secondary_restoration != "none" or config.denoise_strength != "none": raise ValueError("LTX restoration does not support secondary restoration or denoise") - files = LtxModelFiles.from_dir(config.restoration_model_path) + if config.ltx_fast and torch.cuda.get_device_capability(device)[0] < 10: + raise ValueError("The fast LTX model needs an RTX 50-series (Blackwell) GPU") + files = LtxModelFiles.from_dir(config.restoration_model_path, fast=config.ltx_fast) ensure_engines_compiled( EngineCompilationRequest( device=str(device), diff --git a/tests/factories.py b/tests/factories.py index a61eab2f..f0cd70ee 100644 --- a/tests/factories.py +++ b/tests/factories.py @@ -29,6 +29,7 @@ def session_config(**overrides) -> SessionConfig: restoration_model_path=Path("restore.pth"), ltx_large_canvas=True, ltx_seed=0, + ltx_fast=False, compile_basicvsrpp=True, max_clip_size=90, temporal_overlap=8, diff --git a/tests/test_ltx_model_files.py b/tests/test_ltx_model_files.py index 123c0ff8..cf8cdabf 100644 --- a/tests/test_ltx_model_files.py +++ b/tests/test_ltx_model_files.py @@ -11,7 +11,7 @@ def test_source_model_bundle_and_reader(tmp_path): value = torch.arange(6).reshape(2, 3) for name in ("transformer", "vae", "vae-decoder"): save_file({"weight": value}, str(tmp_path / f"{name}.safetensors"), metadata={"config": "example"}) - files = LtxModelFiles.from_dir(tmp_path) + files = LtxModelFiles.from_dir(tmp_path, fast=False) with open_tensors(files.transformer) as handle: assert list(handle.keys()) == ["weight"] assert handle.metadata() == {"config": "example"} @@ -30,3 +30,9 @@ def test_model_check_precedes_engine_compilation(tmp_path, monkeypatch): with pytest.raises(ValueError, match="model unavailable"): session_factory._build_ltx_session(config, torch.device("cuda:0"), log_callback=None) compile_engines.assert_not_called() + + +def test_fast_bundle_uses_the_fast_transformer(tmp_path): + for name in ("transformer-fast", "vae", "vae-decoder"): + save_file({"weight": torch.zeros(1)}, str(tmp_path / f"{name}.safetensors")) + assert LtxModelFiles.from_dir(tmp_path, fast=True).transformer == tmp_path / "transformer-fast.safetensors" diff --git a/tests/test_ltx_transformer.py b/tests/test_ltx_transformer.py index 2c4e23a5..7b45b104 100644 --- a/tests/test_ltx_transformer.py +++ b/tests/test_ltx_transformer.py @@ -79,3 +79,55 @@ def test_fuse_overlaps_writes_the_same_mean_to_both_windows(): assert torch.equal(right[:, 10:], right_before[:, 10:]) fuse_overlaps([left, right], tokens_per_frame, [0.2, 0.4, 0.6, 0.8]) assert torch.equal(right[:, 2:10], left[:, 24:32]) + + +def test_fp4_codes_round_half_to_even_and_saturate(): + values = torch.tensor([0.0, 0.25, 0.3, 0.75, 1.25, 1.75, 2.5, 3.5, 5.0, 5.5, 9.0, -0.5, -6.0]) + expected = [0, 0, 1, 2, 2, 4, 4, 6, 6, 7, 7, 8 | 1, 8 | 7] + assert T.fp4_codes(values).tolist() == expected + + +def test_block_scale_swizzle_matches_cublas_tile_addresses(): + rows, cols = 200, 6 + scales = torch.arange(rows * cols, dtype=torch.float32).remainder(200).view(rows, cols).to(torch.float8_e4m3fn) + tiled = T.swizzle_block_scales(scales).view(torch.uint8).flatten() + padded_cols = 8 + for row in (0, 31, 32, 127, 128, 199): + for col in range(cols): + tile = (row // 128) * (padded_cols // 4) + col // 4 + offset = tile * 512 + (row % 32) * 16 + ((row % 128) // 32) * 4 + col % 4 + assert tiled[offset] == scales[row, col].view(torch.uint8) + + +def test_quantize_fp4_reconstructs_within_fp4_error(): + generator = torch.Generator().manual_seed(0) + x = torch.randn(130, 64, generator=generator) + tensor_scale = T.fp4_tensor_scale(x) + packed, tiled = T.quantize_fp4(x, tensor_scale) + assert packed.shape == (130, 32) and packed.dtype == torch.uint8 + assert tiled.shape == (256, 4) and tiled.dtype == torch.float8_e4m3fn + table = torch.tensor([0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0]) + codes = torch.stack([packed >> 4, packed & 15], dim=-1).flatten(1).long() + values = table[codes & 7] * (1 - 2 * (codes >> 3)).float() + blocks = x.view(130, 4, 16) + scales = (blocks.abs().amax(-1) / T.FP4_MAX / tensor_scale).clamp(max=T.FP8_E4M3_MAX) + scales = scales.to(torch.float8_e4m3fn).float() + rebuilt = values.view(130, 4, 16) * scales.unsqueeze(-1) * tensor_scale + assert (rebuilt - blocks).norm() / blocks.norm() < 0.15 + + +def test_self_attention_falls_back_to_sdpa_off_gpu(): + generator = torch.Generator().manual_seed(0) + q, k, v = torch.randn(3, 1, 2, 8, 16, generator=generator).unbind(0) + expected = torch.nn.functional.scaled_dot_product_attention(q, k, v) + assert torch.allclose(T.self_attention(q, k, v), expected) + + +def test_block_forward_runs_eager_without_triton(monkeypatch): + import importlib.util + + from jasna.ltx import sampler + + real_find_spec = importlib.util.find_spec + monkeypatch.setattr(importlib.util, "find_spec", lambda name: None if name == "triton" else real_find_spec(name)) + assert sampler.compiled_block_forward() is T.block_forward From d55328ced4c54b60433cdfc24c00ab6e43318936 Mon Sep 17 00:00:00 2001 From: Kruk2 Date: Tue, 29 Sep 2026 11:22:31 +0200 Subject: [PATCH 21/48] feat(ltx): use 512 px windows below 10 GB free VRAM --- jasna/cli_help.py | 2 +- jasna/ltx/restore.py | 11 +++++++++++ tests/test_ltx_restore.py | 6 ++++++ 3 files changed, 18 insertions(+), 1 deletion(-) diff --git a/jasna/cli_help.py b/jasna/cli_help.py index b1e76742..3b80c200 100644 --- a/jasna/cli_help.py +++ b/jasna/cli_help.py @@ -22,7 +22,7 @@ ), "ltx_large_canvas": ( "ltx: restore large mosaics at 768 px instead of 512 px. Sharper, about 3x slower for " - "those mosaics and needs more VRAM. (default: %(default)s)" + "those mosaics and needs more VRAM (off automatically below 10 GB free). (default: %(default)s)" ), "fp16": "Use FP16 where supported (restoration + TensorRT). Reduces VRAM usage and might improve performance.", "compile_basicvsrpp": "Compile BasicVSR++ for big performance boost (at cost of VRAM usage). Not recommended to use big clip sizes. (default: %(default)s)", diff --git a/jasna/ltx/restore.py b/jasna/ltx/restore.py index a47d007b..22a70a1e 100644 --- a/jasna/ltx/restore.py +++ b/jasna/ltx/restore.py @@ -35,6 +35,7 @@ FrameBatches = Iterator[tuple[torch.Tensor, list[int]]] FrameSource = Callable[[], FrameBatches] +LARGE_CANVAS_MIN_FREE_BYTES = 10 << 30 class Cancelled(Exception): @@ -234,6 +235,15 @@ def compose( bar.close() +def large_canvas_fits(requested: bool, free_bytes: int) -> bool: + """768 px windows need ~7.5 GiB of VAE-encoder activations alone; below 10 GiB free + every window runs at 512 px instead.""" + if requested and free_bytes < LARGE_CANVAS_MIN_FREE_BYTES: + logger.warning("LTX: %.1f GiB free VRAM, restoring large mosaics at 512 px", free_bytes / 2**30) + return False + return requested + + def restore_video( frames: FrameSource, write: Callable[[torch.Tensor, int], None], @@ -252,6 +262,7 @@ def restore_video( """Restore every mosaic of the video ``frames`` yields and ``write`` each output frame with its pts. ``detector`` is only used by the scan pass; ``large_canvas`` allows 768 px windows for large mosaics; window ``i`` draws its noise from ``seed + i``.""" + large_canvas = large_canvas_fits(large_canvas, torch.cuda.mem_get_info(device)[0]) plans = scan(frames, detector, frame_h=frame_h, frame_w=frame_w, large_canvas=large_canvas, progress=progress, cancel=cancel) with tempfile.TemporaryDirectory(dir=work_dir, prefix=".ltx-") as temp: references = _LatentStore(Path(temp), "reference") diff --git a/tests/test_ltx_restore.py b/tests/test_ltx_restore.py index d2617a7d..14abff29 100644 --- a/tests/test_ltx_restore.py +++ b/tests/test_ltx_restore.py @@ -11,3 +11,9 @@ def test_decode_window_converts_a_float32_decode(monkeypatch): assert frames.dtype == torch.uint8 and frames.shape == (3, 3, 1, 1) assert frames[:, 0, 0, 0].tolist() == [0, 128, 255] + +def test_large_canvas_needs_ten_gib_free(): + gib = 1 << 30 + assert restore.large_canvas_fits(True, 10 * gib) + assert not restore.large_canvas_fits(True, 9 * gib) + assert not restore.large_canvas_fits(False, 30 * gib) From 0e4051ad3ee4e6e8d07898a73ca295d9cb26292f Mon Sep 17 00:00:00 2001 From: Kruk2 Date: Tue, 29 Sep 2026 15:15:01 +0200 Subject: [PATCH 22/48] feat(ltx): read the sampler from the model file Sigma schedule and STG scale come from the file's metadata; STG 0 skips the perturbed pass. --- jasna/ltx/sampler.py | 35 ++++++++++++++++++----------------- tests/test_ltx_transformer.py | 15 +++++++++++++++ 2 files changed, 33 insertions(+), 17 deletions(-) diff --git a/jasna/ltx/sampler.py b/jasna/ltx/sampler.py index d69819af..58b9d8cc 100644 --- a/jasna/ltx/sampler.py +++ b/jasna/ltx/sampler.py @@ -3,14 +3,15 @@ A track's windows are denoised in lockstep: every step runs all windows through each block before the next block is fetched, so a streamed block crosses PCIe once per step per window group. After each step the windows' shared latent frames are replaced by a -ramp-weighted mean (MultiDiffusion-style overlap fusion). Sampler: Euler over a fixed -sigma schedule, CFG 1, STG on one block whose perturbed pass reuses the conditional -pass's hidden state up to that block. +ramp-weighted mean (MultiDiffusion-style overlap fusion). Sampler: Euler over the model +file's sigma schedule, CFG 1, and the file's STG scale on one block, whose perturbed pass +reuses the conditional pass's hidden state up to that block (skipped at scale 0). """ from __future__ import annotations import importlib.util +import json import logging from collections.abc import Sequence from pathlib import Path @@ -24,11 +25,6 @@ logger = logging.getLogger(__name__) -LOWTAIL15_SIGMAS = ( - 1.0, 0.982, 0.959, 0.93, 0.892, 0.839, 0.762, 0.637, 0.4, - 0.32, 0.184, 0.106, 0.061, 0.035, 0.02, 0.0, -) # fmt: skip -STG_SCALE = 1.0 STG_BLOCK = 28 ACTIVATION_RESERVE_BYTES = 3 << 30 _ADALN_PREFIXES = ("adaln_single.", "prompt_adaln_single.") @@ -56,9 +52,14 @@ def fuse_overlaps(latents: Sequence[torch.Tensor], tokens_per_frame: int, right_ b.copy_(fused) -def _window_state_bytes(tokens: int) -> int: - """Hidden states one window keeps between blocks (conditional + perturbed).""" - return 2 * 2 * tokens * 4096 * 2 +def sampler_settings(metadata: dict[str, str]) -> tuple[torch.Tensor, float]: + """The model file's sigma schedule and STG scale (0 for a distilled student).""" + return torch.tensor(json.loads(metadata["sigmas"]), dtype=torch.float32), float(json.loads(metadata["stg_scale"])) + + +def _window_state_bytes(tokens: int, *, stg: bool) -> int: + """Hidden states one window keeps between blocks (conditional, plus perturbed with STG).""" + return (2 if stg else 1) * 2 * tokens * 4096 * 2 def _window_workspace_bytes(tokens: int) -> int: @@ -69,10 +70,10 @@ def _window_workspace_bytes(tokens: int) -> int: class LtxTransformer: """The restoration transformer loaded from a jasna LTX model file.""" - def __init__(self, path: Path, device: torch.device, *, sigmas: Sequence[float] = LOWTAIL15_SIGMAS) -> None: + def __init__(self, path: Path, device: torch.device) -> None: self.device = device - self.sigmas = torch.tensor(sigmas, dtype=torch.float32) with open_tensors(path) as handle: + self.sigmas, self.stg_scale = sampler_settings(handle.metadata()) keys = list(handle.keys()) top_keys = [k for k in keys if not k.startswith("blocks.") and k != "prompt_context"] top = {k: handle.get_tensor(k).to(device) for k in top_keys} @@ -115,7 +116,7 @@ def _rope(self, frames: int, height: int, width: int) -> T.Rope: def _group_size(self, windows: int, tokens: int) -> int: free, _ = torch.cuda.mem_get_info(self.device) free += torch.cuda.memory_reserved(self.device) - torch.cuda.memory_allocated(self.device) - fits = (free - _window_workspace_bytes(tokens)) // _window_state_bytes(tokens) + fits = (free - _window_workspace_bytes(tokens)) // _window_state_bytes(tokens, stg=bool(self.stg_scale)) return max(1, min(windows, int(fits))) @torch.inference_mode() @@ -160,7 +161,7 @@ def _denoise( for index in range(T.NUM_BLOCKS): weights = self.blocks.acquire(index) for j in range(len(xs)): - if STG_SCALE and index >= STG_BLOCK: + if self.stg_scale and index >= STG_BLOCK: source = xs[j] if index == STG_BLOCK else perturbed[j] perturbed[j] = self._block_forward( weights, source, cond, self.context, rope, skip_self_attention=index == STG_BLOCK @@ -170,9 +171,9 @@ def _denoise( out = [] for lat, x, xp in zip(latents, xs, perturbed): denoised = (lat.float() - T.velocity(self.top, x, cond).float() * sigma).to(lat.dtype) - if STG_SCALE: + if self.stg_scale: weak = (lat.float() - T.velocity(self.top, xp, cond).float() * sigma).to(lat.dtype) - denoised = (denoised.float() + STG_SCALE * (denoised.float() - weak.float())).to(lat.dtype) + denoised = (denoised.float() + self.stg_scale * (denoised.float() - weak.float())).to(lat.dtype) out.append(denoised) return out diff --git a/tests/test_ltx_transformer.py b/tests/test_ltx_transformer.py index 7b45b104..fa8cd01b 100644 --- a/tests/test_ltx_transformer.py +++ b/tests/test_ltx_transformer.py @@ -131,3 +131,18 @@ def test_block_forward_runs_eager_without_triton(monkeypatch): real_find_spec = importlib.util.find_spec monkeypatch.setattr(importlib.util, "find_spec", lambda name: None if name == "triton" else real_find_spec(name)) assert sampler.compiled_block_forward() is T.block_forward + + +def test_sampler_settings_come_from_the_model_file(): + from jasna.ltx.sampler import sampler_settings + + sigmas, stg = sampler_settings({"sigmas": "[1.0, 0.4, 0.0]", "stg_scale": "0.0"}) + assert sigmas.tolist() == pytest.approx([1.0, 0.4, 0.0]) and stg == 0.0 + with pytest.raises(KeyError): + sampler_settings({"format": "jasna-ltx-restore"}) + + +def test_window_state_halves_without_stg(): + from jasna.ltx import sampler + + assert sampler._window_state_bytes(100, stg=True) == 2 * sampler._window_state_bytes(100, stg=False) From e403bf7ac44d651bfd3cb2556d468d0aaeac8c98 Mon Sep 17 00:00:00 2001 From: Kruk2 Date: Tue, 29 Sep 2026 15:29:41 +0200 Subject: [PATCH 23/48] feat(ltx): make 768 px windows opt-in (--ltx-large-canvas) --- jasna/cli_help.py | 4 ++-- jasna/gui/video_session.py | 2 +- jasna/main.py | 2 +- tests/test_session_config.py | 6 +++--- 4 files changed, 7 insertions(+), 7 deletions(-) diff --git a/jasna/cli_help.py b/jasna/cli_help.py index 3b80c200..1d7d6527 100644 --- a/jasna/cli_help.py +++ b/jasna/cli_help.py @@ -21,8 +21,8 @@ "and the fast model file. (default: %(default)s)" ), "ltx_large_canvas": ( - "ltx: restore large mosaics at 768 px instead of 512 px. Sharper, about 3x slower for " - "those mosaics and needs more VRAM (off automatically below 10 GB free). (default: %(default)s)" + "ltx: restore large mosaics at 768 px instead of 512 px. Sharper, but about 3x slower for " + "those mosaics; needs 10 GB of free VRAM (ignored below that). (default: %(default)s)" ), "fp16": "Use FP16 where supported (restoration + TensorRT). Reduces VRAM usage and might improve performance.", "compile_basicvsrpp": "Compile BasicVSR++ for big performance boost (at cost of VRAM usage). Not recommended to use big clip sizes. (default: %(default)s)", diff --git a/jasna/gui/video_session.py b/jasna/gui/video_session.py index f565e2df..6ee88a7b 100644 --- a/jasna/gui/video_session.py +++ b/jasna/gui/video_session.py @@ -70,7 +70,7 @@ def video_session_config( scene_detection=bool(settings.scene_detection), restoration_model_name="basicvsrpp", restoration_model_path=default_restoration_model_path("basicvsrpp"), - ltx_large_canvas=True, + ltx_large_canvas=False, ltx_seed=LTX_DEFAULT_SEED, ltx_fast=False, compile_basicvsrpp=bool(settings.compile_basicvsrpp), diff --git a/jasna/main.py b/jasna/main.py index 1d2a7e18..973ea46e 100644 --- a/jasna/main.py +++ b/jasna/main.py @@ -206,7 +206,7 @@ def build_parser() -> argparse.ArgumentParser: ) restoration.add_argument( "--ltx-large-canvas", - default=True, + default=False, action=argparse.BooleanOptionalAction, help=CLI_HELP["ltx_large_canvas"], ) diff --git a/tests/test_session_config.py b/tests/test_session_config.py index b670f6f1..4638bbf0 100644 --- a/tests/test_session_config.py +++ b/tests/test_session_config.py @@ -123,12 +123,12 @@ def test_cli_non_default_args_are_mapped() -> None: def test_cli_ltx_args_are_mapped() -> None: defaults = _cli_config() assert defaults.restoration_model_name == "basicvsrpp" - assert defaults.ltx_large_canvas is True + assert defaults.ltx_large_canvas is False assert defaults.ltx_seed == 20260923 - config = _cli_config(["--restoration-model-name", "ltx", "--no-ltx-large-canvas", "--ltx-seed", "7"]) + config = _cli_config(["--restoration-model-name", "ltx", "--ltx-large-canvas", "--ltx-seed", "7"]) assert config.restoration_model_name == "ltx" - assert config.ltx_large_canvas is False + assert config.ltx_large_canvas is True assert config.ltx_seed == 7 From 311b22946d0fc2ea6ae6fbb6928e41a9c45d38c5 Mon Sep 17 00:00:00 2001 From: Kruk2 Date: Tue, 29 Sep 2026 15:55:55 +0200 Subject: [PATCH 24/48] fix(ltx): show denoise progress per step and silence the compile cache warning Progress advances after every sampler step, and the Hadamard matrix is built directly so Dynamo no longer warns. --- jasna/ltx/restore.py | 12 +++++----- jasna/ltx/sampler.py | 10 +++++--- jasna/ltx/transformer.py | 12 ++++++++-- tests/test_ltx_restore.py | 43 +++++++++++++++++++++++++++++++++++ tests/test_ltx_transformer.py | 12 ++++++++++ 5 files changed, 78 insertions(+), 11 deletions(-) diff --git a/jasna/ltx/restore.py b/jasna/ltx/restore.py index 22a70a1e..779af6fc 100644 --- a/jasna/ltx/restore.py +++ b/jasna/ltx/restore.py @@ -89,7 +89,7 @@ def scan( regions: list[list[Region]] = [] cuts: set[int] = set() scene = SceneCutDetector() - with progress.bar("scan", progress.frames) as bar: + with progress.bar("1/4 scan", progress.frames) as bar: for batch, _pts in frames(): _check(cancel) cuts.update(len(regions) + offset for offset in scene.find_cuts(batch)) @@ -126,7 +126,7 @@ def encode_references( upcoming = next(pending, None) active: dict[int, tuple[Window, list[torch.Tensor]]] = {} frame_idx = 0 - bar = progress.bar("encode", progress.frames) + bar = progress.bar("2/4 encode", progress.frames) for batch, _pts in frames(): _check(cancel) bar.update(len(batch)) @@ -156,16 +156,16 @@ def denoise( progress: Progress, cancel: threading.Event, ) -> None: - bar = progress.bar("denoise", sum(len(plan.windows) for plan in plans), unit="window") + steps = len(transformer.conditions) + bar = progress.bar("3/4 denoise", steps * sum(len(plan.windows) for plan in plans), unit="step") for plan in plans: _check(cancel) latents = transformer.denoise_chain( - [references.get(w.index) for w in plan.windows], [seed + w.index for w in plan.windows] + [references.get(w.index) for w in plan.windows], [seed + w.index for w in plan.windows], bar.update ) for window, latent in zip(plan.windows, latents): finals.put(window.index, latent) references.delete(window.index) - bar.update(len(plan.windows)) bar.close() @@ -205,7 +205,7 @@ def compose( covering: list[Window] = [] decoded: dict[int, torch.Tensor] = {} frame_idx = 0 - bar = progress.bar("compose", progress.frames) + bar = progress.bar("4/4 compose", progress.frames) for batch, pts in frames(): _check(cancel) bar.update(len(batch)) diff --git a/jasna/ltx/sampler.py b/jasna/ltx/sampler.py index 58b9d8cc..30f430fc 100644 --- a/jasna/ltx/sampler.py +++ b/jasna/ltx/sampler.py @@ -13,7 +13,7 @@ import importlib.util import json import logging -from collections.abc import Sequence +from collections.abc import Callable, Sequence from pathlib import Path import torch @@ -120,9 +120,12 @@ def _group_size(self, windows: int, tokens: int) -> int: return max(1, min(windows, int(fits))) @torch.inference_mode() - def denoise_chain(self, references: Sequence[torch.Tensor], seeds: Sequence[int]) -> list[torch.Tensor]: + def denoise_chain( + self, references: Sequence[torch.Tensor], seeds: Sequence[int], advance: Callable[[int], None] + ) -> list[torch.Tensor]: """Final latents ``[1, 128, F, H, W]`` of one track's windows, given each window's - reference latent (the encoded mosaic crop) and noise seed.""" + reference latent (the encoded mosaic crop) and noise seed. ``advance(n)`` reports n + window-steps done.""" frames, height, width = references[0].shape[2:] tokens_per_frame = height * width tokens = frames * tokens_per_frame @@ -144,6 +147,7 @@ def denoise_chain(self, references: Sequence[torch.Tensor], seeds: Sequence[int] ) for i, value in zip(members, denoised): latents[i] = _euler(latents[i], value, sigmas, step) + advance(len(members)) fuse_overlaps(latents, tokens_per_frame, FUSION_RAMP) return [unpatchify(latent, frames, height, width) for latent in latents] diff --git a/jasna/ltx/transformer.py b/jasna/ltx/transformer.py index 1d7165b2..aa8fee7d 100644 --- a/jasna/ltx/transformer.py +++ b/jasna/ltx/transformer.py @@ -62,8 +62,7 @@ ) + ("ff.net.0.proj", "ff.net.2") -@functools.cache -def _regular_hadamard(size: int, device: torch.device, dtype: torch.dtype) -> torch.Tensor: +def _build_regular_hadamard(size: int, device: torch.device, dtype: torch.dtype) -> torch.Tensor: r4 = torch.tensor([[1.0, 1, 1, -1], [1, 1, -1, 1], [1, -1, 1, 1], [-1, 1, 1, 1]]) h = r4.clone() while h.shape[0] < size: @@ -71,6 +70,15 @@ def _regular_hadamard(size: int, device: torch.device, dtype: torch.dtype) -> to return (h / size**0.5).to(device=device, dtype=dtype) +_cached_regular_hadamard = functools.cache(_build_regular_hadamard) + + +def _regular_hadamard(size: int, device: torch.device, dtype: torch.dtype) -> torch.Tensor: + if torch.compiler.is_compiling(): + return _build_regular_hadamard(size, device, dtype) + return _cached_regular_hadamard(size, device, dtype) + + def rotate(x: torch.Tensor) -> torch.Tensor: """ConvRot block regular-Hadamard rotation along the last dim (self-inverse).""" h = _regular_hadamard(ROT_SIZE, x.device, x.dtype) diff --git a/tests/test_ltx_restore.py b/tests/test_ltx_restore.py index 14abff29..e2aa6211 100644 --- a/tests/test_ltx_restore.py +++ b/tests/test_ltx_restore.py @@ -17,3 +17,46 @@ def test_large_canvas_needs_ten_gib_free(): assert restore.large_canvas_fits(True, 10 * gib) assert not restore.large_canvas_fits(True, 9 * gib) assert not restore.large_canvas_fits(False, 30 * gib) + + +def test_denoise_reports_every_window_step(monkeypatch): + class Bar: + def __init__(self, total): + self.total, self.n = total, 0 + + def update(self, n): + self.n += n + + def close(self): + pass + + class Store: + def get(self, index): + return index + + def put(self, index, latent): + pass + + def delete(self, index): + pass + + class Transformer: + conditions = [None] * 8 + + def denoise_chain(self, references, seeds, advance): + for _ in self.conditions: + advance(len(references)) + return list(references) + + bars = [] + monkeypatch.setattr(restore.Progress, "bar", lambda self, name, total, unit="frame": bars.append(Bar(total)) or bars[-1]) + window = lambda i: restore.Window(index=i, track_id=0, start=0, real_frames=1, crops=()) + plans = [ + restore.TrackPlan(track_id=0, start=0, polygons=(), windows=(window(0), window(1), window(2))), + restore.TrackPlan(track_id=1, start=0, polygons=(), windows=(window(3),)), + ] + restore.denoise( + plans, Transformer(), Store(), Store(), seed=0, + progress=restore.Progress(frames=0, disable=True), cancel=restore.threading.Event(), + ) + assert bars[0].total == bars[0].n == 32 diff --git a/tests/test_ltx_transformer.py b/tests/test_ltx_transformer.py index fa8cd01b..cc4014f9 100644 --- a/tests/test_ltx_transformer.py +++ b/tests/test_ltx_transformer.py @@ -30,6 +30,18 @@ def test_rotation_is_its_own_inverse(): assert torch.allclose(T.rotate(T.rotate(x)), x, atol=1e-5) + +def test_compiled_rotation_emits_no_cache_warning(): + import warnings + + torch._dynamo.reset() + x = torch.randn(4, 512, dtype=torch.float32) + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always") + rotated = torch.compile(T.rotate, backend="eager")(x) + assert torch.equal(rotated, T.rotate(x)) + assert not [w for w in caught if "lru_cache" in str(w.message)] + def test_quantize_rows_pads_and_bounds(): x = torch.randn(40, 64) * 3 codes, scales = T.quantize_rows(x) From 8c5a5c97518ef9a7fdce19787d27205da75ad28c Mon Sep 17 00:00:00 2001 From: Kruk2 Date: Tue, 29 Sep 2026 16:01:11 +0200 Subject: [PATCH 25/48] feat(gui): pick the LTX restoration model in settings Standard and LTX cards plus seed, fast mode and sharper large mosaics options. LTX is disabled without an NVIDIA GPU or its model files. --- jasna/gui/app.py | 34 ++- jasna/gui/engine_preflight.py | 5 +- jasna/gui/locales/en.py | 16 ++ jasna/gui/locales/ja.py | 16 ++ jasna/gui/locales/ko.py | 16 ++ jasna/gui/locales/th.py | 16 ++ jasna/gui/locales/zh.py | 16 ++ jasna/gui/models.py | 7 + jasna/gui/restoration_preview.py | 3 + jasna/gui/segment_editor.py | 4 + jasna/gui/settings_panel.py | 30 ++- jasna/gui/settings_sections/advanced.py | 21 +- jasna/gui/settings_sections/basic.py | 23 +- .../settings_sections/restoration_model.py | 206 ++++++++++++++++++ jasna/gui/settings_sections/secondary.py | 8 + jasna/gui/settings_sections/widgets.py | 9 + jasna/gui/validation.py | 10 +- jasna/gui/video_player.py | 1 + jasna/gui/video_session.py | 25 ++- jasna/ltx/model_files.py | 15 +- tests/test_gui_engine_preflight.py | 9 + tests/test_gui_settings_persistence_paths.py | 10 + tests/test_gui_settings_scroll.py | 4 +- tests/test_gui_settings_sections.py | 126 ++++++++++- tests/test_gui_startup_and_wizard.py | 24 +- tests/test_gui_tvai_validation.py | 25 ++- tests/test_gui_video_player.py | 11 + tests/test_ltx_model_files.py | 10 + tests/test_restoration_preview.py | 33 +++ tests/test_video_session.py | 43 ++++ 30 files changed, 712 insertions(+), 64 deletions(-) create mode 100644 jasna/gui/settings_sections/restoration_model.py diff --git a/jasna/gui/app.py b/jasna/gui/app.py index d8f7210d..7fafe61d 100644 --- a/jasna/gui/app.py +++ b/jasna/gui/app.py @@ -50,17 +50,22 @@ _SHUTDOWN_COUNTDOWN_SECONDS = 60 -def _warm_up_cuda() -> None: +def _warm_up_cuda() -> tuple[bool, bool]: """Import torch and create the CUDA context. Run off the UI thread after the window paints, so the window shows first and the first job skips cold torch/CUDA init. - An unsupported GPU gets no kernel launch; the system check reports it instead.""" + An unsupported GPU gets no kernel launch; the system check reports it instead. + Returns (NVIDIA GPU, Blackwell GPU) for the LTX model options.""" from jasna.os_utils import check_supported_gpu gpu_ok, _ = check_supported_gpu() - if gpu_ok: - import torch + if not gpu_ok: + return False, False + import torch + from jasna.accelerator import is_nvidia_device - torch.zeros(1, device="cuda") + torch.zeros(1, device="cuda") + nvidia = is_nvidia_device() + return nvidia, nvidia and torch.cuda.get_device_capability(0)[0] >= 10 class JasnaApp(ctk.CTk, TkinterDnD.DnDWrapper): @@ -324,7 +329,14 @@ def _build_main_body(self): self._queue_panel.configure(cursor="arrow") self._queue_panel.set_on_jobs_changed(self._on_jobs_changed) - self._settings_panel = SettingsPanel(self._workspace, self._preset_manager) + from jasna.engine_paths import default_restoration_model_path + from jasna.ltx.model_files import bundle_present as ltx_bundle_present + + self._settings_panel = SettingsPanel( + self._workspace, + self._preset_manager, + ltx_installed=ltx_bundle_present(default_restoration_model_path("ltx")), + ) self._settings_panel.configure(cursor="arrow") self._settings_panel.set_on_interactive_image_restore(self._open_interactive_image_restore) @@ -393,9 +405,13 @@ def _log_startup_timing(self): def _start_cuda_warmup(self): def _run(): try: - _warm_up_cuda() + nvidia, blackwell = _warm_up_cuda() except Exception: - logger.debug("CUDA warm-up failed", exc_info=True) + logger.warning("CUDA warm-up failed", exc_info=True) + nvidia, blackwell = False, False + self._main_thread.post( + lambda: self._settings_panel.set_gpu_support(nvidia=nvidia, blackwell=blackwell) + ) threading.Thread(target=_run, daemon=True, name="cuda-warmup").start() @@ -578,7 +594,7 @@ def _on_start(self): settings = self._settings_panel.get_settings() from jasna.gui.validation import validate_gui_start - errors = validate_gui_start(settings) + errors = validate_gui_start(settings, jobs) if errors: from tkinter import messagebox diff --git a/jasna/gui/engine_preflight.py b/jasna/gui/engine_preflight.py index a91b5f7d..96f7b989 100644 --- a/jasna/gui/engine_preflight.py +++ b/jasna/gui/engine_preflight.py @@ -93,8 +93,9 @@ def run_engine_preflight(settings: AppSettings) -> EnginePreflightResult: ) ) + standard_model = settings.restoration_model == "basicvsrpp" restoration_model_path = default_restoration_model_path("basicvsrpp") - if bool(settings.compile_basicvsrpp) and not amd: + if standard_model and bool(settings.compile_basicvsrpp) and not amd: sub_paths = get_basicvsrpp_sub_engine_paths(str(restoration_model_path), bool(settings.fp16_mode)) all_engine_paths = tuple(Path(p) for p in sub_paths.values()) missing_paths = tuple(p for p in all_engine_paths if not p.is_file()) @@ -108,7 +109,7 @@ def run_engine_preflight(settings: AppSettings) -> EnginePreflightResult: ) ) - if settings.secondary_restoration == "unet-4x": + if standard_model and settings.secondary_restoration == "unet-4x": unet_engine = expected_unet4x_engine_path(fp16=bool(settings.fp16_mode)) unet_exists = unet_engine.is_file() reqs.append( diff --git a/jasna/gui/locales/en.py b/jasna/gui/locales/en.py index dee5fa6e..cc8cbe55 100644 --- a/jasna/gui/locales/en.py +++ b/jasna/gui/locales/en.py @@ -77,6 +77,7 @@ "segments_preview_reset_view": "Reset view", "segments_preview_reset_view_hint": "Reset zoom and pan to show the full frame", "segments_restore_gpu_busy": "Unavailable while processing is running", + "segments_restore_ltx_unavailable": "Preview is not available with LTX yet", "segments_restore_loading_models": "Loading restoration models…", "segments_restore_restoring": "Restoring…", "segments_restore_failed": "Restoration preview failed: {message}", @@ -196,6 +197,16 @@ # Sections "section_basic": "Basic Processing", + "section_restoration_model": "Restoration Model", + "model_basicvsrpp": "Standard", + "model_basicvsrpp_description": "Fast. Good for most videos.", + "model_ltx": "LTX", + "model_ltx_description": "Best detail. Much slower and needs a strong NVIDIA GPU.", + "model_ltx_needs_nvidia": "Needs an NVIDIA GPU.", + "model_ltx_not_installed": "Model not installed.", + "ltx_seed": "Seed", + "ltx_fast": "Fast mode", + "ltx_large_canvas": "Sharper large mosaics", "section_advanced": "Advanced Processing", "section_secondary": "Secondary Restoration", "section_image_restoration": "Image Restoration (still images)", @@ -397,6 +408,10 @@ "tip_vr_mode": "Controls side-by-side VR180 processing.\n\nAuto enables it for exact 2:1 frames taller than 1080 pixels, trusted studio filename tokens, or compatible spatial metadata, and routes each detected mosaic region's restoration projection — raw, fisheye, or gnomonic — by studio (raw when the studio is unknown). SBS forces per-eye processing with the same routing; SBS + fisheye forces fisheye conditioning for every region. Detection, tracking, masks, and blending always stay in source coordinates, so only the restored delta is reprojected back.\n\nThe segment editor displays the left eye while scans and exports process both eyes.\nDefault: Auto", "tip_fp16_mode": "Uses half-precision math to reduce VRAM usage and often run faster. No visible quality loss on modern GPUs.\n\nRecommended: ON for RTX 20-series and newer.\nDefault: ON", "tip_compile_basicvsrpp": "Compiles the restoration model into TensorRT sub-engines for a big speed boost (~2-3x faster).\nFirst compilation takes 15-60 minutes. Close all other applications (including browsers) and avoid using the PC during compilation.\nEngines are cached and reused on subsequent runs.\n\nEngine VRAM: ~1.9 GB (clip 60), ~5.4 GB (clip 180).\nPeak VRAM during processing: ~7.6 GB (clip 60), ~14.7 GB (clip 180).\nWithout compilation: ~6 GB (clip 60), ~10.4 GB (clip 180).\n\nIf you run out of VRAM, disable this or lower clip size.\n\nRecommended: ON with clip size 60-90.\nDefault: ON", + "tip_ltx_seed": "Another seed gives another take on the same video.\nKeep a seed you like to get the same result again.", + "tip_ltx_new_seed": "Try a new random seed", + "tip_ltx_fast": "About 1.4x faster with slightly less detail.\nRTX 50-series only.", + "tip_ltx_large_canvas": "Restores big mosaics in more detail, but about 3x slower for them.\nNeeds at least 10 GB of free video memory.", "tip_denoise_strength": "Reduces noise and grain in restored areas. Higher = smoother but may lose fine detail.\n\nNone: no denoising. Low/Medium: good starting point. High: heavy smoothing.\nDefault: None", "tip_denoise_step": "When to apply denoising in the pipeline:\n- After Primary: before upscaling (secondary restoration). Denoises at 256x256.\n- After Secondary: after upscaling, right before final output. Denoises at full resolution.\n\nDefault: After Primary", "tip_secondary_restoration": "Optional second pass that upscales restored areas from 256x256 to 1024 pixels. Improves sharpness, especially for close-ups and 4K video.\n\nUNet 4x is a supporter-only model: much faster than TVAI and high quality. Activate Jasna to use it.\nRTX Super Res is a fast alternative with ok quality.\nTopaz TVAI requires a separate purchase and install.", @@ -450,6 +465,7 @@ "error_cannot_start": "Cannot start processing:", "error_invalid_tvai": "Invalid TVAI configuration", "error_post_export_command_required": "Custom post-export command is required", + "error_ltx_segments": "LTX processes whole videos. Clear the segments or pick Standard: {files}", # Settings panel "dialog_select_tvai_ffmpeg": "Select Topaz Video ffmpeg.exe", diff --git a/jasna/gui/locales/ja.py b/jasna/gui/locales/ja.py index c9db9e4b..384443e3 100644 --- a/jasna/gui/locales/ja.py +++ b/jasna/gui/locales/ja.py @@ -77,6 +77,7 @@ "segments_preview_reset_view": "表示をリセット", "segments_preview_reset_view_hint": "移動とズームをリセットしてフレーム全体を表示します", "segments_restore_gpu_busy": "処理の実行中は使用できません", + "segments_restore_ltx_unavailable": "LTX ではまだプレビューできません", "segments_restore_loading_models": "復元モデルを読み込み中…", "segments_restore_restoring": "復元中…", "segments_restore_failed": "復元プレビューに失敗しました：{message}", @@ -196,6 +197,16 @@ # Sections "section_basic": "基本設定", + "section_restoration_model": "修復モデル", + "model_basicvsrpp": "標準", + "model_basicvsrpp_description": "高速。ほとんどの動画に向いています。", + "model_ltx": "LTX", + "model_ltx_description": "最も精細。かなり遅く、高性能な NVIDIA GPU が必要です。", + "model_ltx_needs_nvidia": "NVIDIA GPU が必要です。", + "model_ltx_not_installed": "モデルがインストールされていません。", + "ltx_seed": "シード", + "ltx_fast": "高速モード", + "ltx_large_canvas": "大きなモザイクをより鮮明に", "section_advanced": "詳細設定", "section_secondary": "二次修復", "section_image_restoration": "画像修復（静止画）", @@ -397,6 +408,10 @@ "tip_vr_mode": "サイドバイサイド VR180 の処理方法を設定します。\n\n自動は、高さが 1080 ピクセルを超える正確な 2:1 フレーム、信頼できるスタジオ名トークン、または互換性のある空間メタデータがある場合に有効になり、検出した各モザイク領域の復元用投影（raw／魚眼／グノモニック）をスタジオごとに振り分けます（不明なスタジオは raw）。SBS は同じ振り分けで両目を個別に処理し、SBS + 魚眼はすべての領域を魚眼で条件付けします。検出・トラッキング・マスク・合成は常にソース座標のままで、復元された差分のみを元へ再投影します。\n\n区間エディターは左目のみ表示しますが、スキャンと書き出しは両目を処理します。\nデフォルト：自動", "tip_fp16_mode": "半精度計算で VRAM 使用量を削減し、速度も向上することが多いです。最新の GPU では画質の劣化はほぼありません。\n\n推奨：RTX 20 シリーズ以降で ON。\nデフォルト：ON", "tip_compile_basicvsrpp": "修復モデルを TensorRT サブエンジンにコンパイルし、大幅に高速化します（約 2-3 倍）。\n初回コンパイルには 15～60 分かかります。他のアプリ（ブラウザ含む）をすべて閉じ、コンパイル中は PC を使用しないでください。\nエンジンはキャッシュされ、次回以降自動的に再利用されます。\n\nエンジン VRAM：約 1.9 GB（クリップ 60）、約 5.4 GB（クリップ 180）。\n処理時ピーク VRAM：約 7.6 GB（クリップ 60）、約 14.7 GB（クリップ 180）。\nコンパイルなし：約 6 GB（クリップ 60）、約 10.4 GB（クリップ 180）。\n\nVRAM 不足の場合は、この設定を無効にするかクリップサイズを下げてください。\n\n推奨：クリップサイズ 60-90 で ON。\nデフォルト：ON", + "tip_ltx_seed": "シードを変えると、同じ動画でも別の仕上がりになります。\n気に入ったシードを残せば同じ結果を再現できます。", + "tip_ltx_new_seed": "新しいランダムなシードを試す", + "tip_ltx_fast": "約 1.4 倍速くなりますが、細部はやや減ります。\nRTX 50 シリーズのみ。", + "tip_ltx_large_canvas": "大きなモザイクをより細かく修復しますが、その部分は約 3 倍遅くなります。\n10 GB 以上の空き VRAM が必要です。", "tip_denoise_strength": "修復された部分のノイズや粒子を低減します。高いほど滑らかになりますが、細部が失われることがあります。\n\nなし：ノイズ除去なし。低/中：まずこちらをお試しください。高：強力な平滑化。\nデフォルト：なし", "tip_denoise_step": "ノイズ除去を適用するタイミング：\n- 一次修復後：アップスケール（二次修復）前に適用。256x256 の解像度で処理。\n- 二次修復後：アップスケール後、最終出力の直前に適用。フル解像度で処理。\n\nデフォルト：一次修復後", "tip_secondary_restoration": "修復された領域を 256x256 から 1024 ピクセルに拡大するオプションの追加処理です。アップ時や 4K 動画で鮮明さが向上します。\n\nUNet 4x は支援者限定モデルです。TVAI よりかなり高速で高品質です。Jasna を有効化すると使えます。\nRTX Super Res は高速ですが品質はまずまずです。\nTopaz TVAI は別途購入とインストールが必要です。", @@ -450,6 +465,7 @@ "error_cannot_start": "処理を開始できません:", "error_invalid_tvai": "無効な TVAI 設定", "error_post_export_command_required": "エクスポート後のカスタムコマンドが必要です", + "error_ltx_segments": "LTX は動画全体を処理します。セグメントを消去するか「標準」を選んでください：{files}", # Settings panel "dialog_select_tvai_ffmpeg": "Topaz Video の ffmpeg.exe を選択", diff --git a/jasna/gui/locales/ko.py b/jasna/gui/locales/ko.py index d806203f..aab5e0d6 100644 --- a/jasna/gui/locales/ko.py +++ b/jasna/gui/locales/ko.py @@ -76,6 +76,7 @@ "segments_preview_reset_view": "보기 초기화", "segments_preview_reset_view_hint": "이동과 확대/축소를 초기화하고 전체 프레임을 표시합니다", "segments_restore_gpu_busy": "처리가 진행 중일 때는 사용할 수 없습니다", + "segments_restore_ltx_unavailable": "LTX에서는 아직 미리보기를 사용할 수 없습니다", "segments_restore_loading_models": "복원 모델을 불러오는 중…", "segments_restore_restoring": "복원 중…", "segments_restore_failed": "복원 미리보기 실패: {message}", @@ -195,6 +196,16 @@ # Sections "section_basic": "기본 처리", + "section_restoration_model": "복원 모델", + "model_basicvsrpp": "표준", + "model_basicvsrpp_description": "빠릅니다. 대부분의 영상에 적합합니다.", + "model_ltx": "LTX", + "model_ltx_description": "가장 세밀합니다. 훨씬 느리고 고성능 NVIDIA GPU가 필요합니다.", + "model_ltx_needs_nvidia": "NVIDIA GPU가 필요합니다.", + "model_ltx_not_installed": "모델이 설치되지 않았습니다.", + "ltx_seed": "시드", + "ltx_fast": "빠른 모드", + "ltx_large_canvas": "큰 모자이크 더 선명하게", "section_advanced": "고급 처리", "section_secondary": "2차 복원", "section_image_restoration": "이미지 복원 (정지 이미지)", @@ -377,6 +388,10 @@ "tip_vr_mode": "좌우 분할 VR180 처리 방식을 설정합니다.\n\n자동은 높이가 1080픽셀을 넘는 정확한 2:1 프레임, 신뢰할 수 있는 스튜디오 파일명 토큰 또는 호환 공간 메타데이터가 있을 때 활성화되며, 감지된 각 모자이크 영역의 복원 투영(raw／어안／그노모닉)을 스튜디오별로 선택합니다(알 수 없는 스튜디오는 raw). SBS는 동일한 라우팅으로 양쪽 눈을 각각 처리하고, SBS + 어안은 모든 영역에 어안 조건화를 강제합니다. 감지·추적·마스크·합성은 항상 소스 좌표에 유지되며 복원된 차분만 원래대로 재투영됩니다.\n\n구간 편집기는 왼쪽 눈만 표시하지만 스캔과 내보내기는 양쪽 눈을 처리합니다.\n기본값: 자동", "tip_fp16_mode": "반정밀도 계산으로 VRAM 사용량을 줄이고 속도도 빨라지는 경우가 많습니다. 최신 GPU에서 화질 차이는 거의 없습니다.\n\n권장: RTX 20 시리즈 이상에서 켜기.\n기본값: 켜기", "tip_compile_basicvsrpp": "복원 모델을 TensorRT 서브엔진으로 컴파일하여 큰 속도 향상을 얻습니다 (약 2-3배).\n첫 컴파일에 15-60분이 소요됩니다. 다른 모든 앱(브라우저 포함)을 닫고, 컴파일 중 PC를 사용하지 마세요.\n엔진은 캐시되어 이후 실행 시 자동으로 재사용됩니다.\n\n엔진 VRAM: 약 1.9 GB (클립 60), 약 5.4 GB (클립 180).\n처리 시 피크 VRAM: 약 7.6 GB (클립 60), 약 14.7 GB (클립 180).\n컴파일 없이: 약 6 GB (클립 60), 약 10.4 GB (클립 180).\n\nVRAM 부족 시 이 설정을 끄거나 클립 크기를 줄이세요.\n\n권장: 클립 크기 60-90에서 켜기.\n기본값: 켜기", + "tip_ltx_seed": "시드를 바꾸면 같은 영상도 다른 결과가 나옵니다.\n마음에 드는 시드를 두면 같은 결과를 다시 얻을 수 있습니다.", + "tip_ltx_new_seed": "새 무작위 시드 사용", + "tip_ltx_fast": "약 1.4배 빠르지만 디테일이 조금 줄어듭니다.\nRTX 50 시리즈 전용.", + "tip_ltx_large_canvas": "큰 모자이크를 더 세밀하게 복원하지만 해당 부분은 약 3배 느려집니다.\n여유 VRAM이 10 GB 이상 필요합니다.", "tip_denoise_strength": "복원된 영역의 노이즈와 입자를 줄여줍니다. 높을수록 부드럽지만 세부 디테일이 줄어들 수 있습니다.\n\n없음: 노이즈 제거 안 함. 낮음/중간: 좋은 시작점. 높음: 강력한 평활화.\n기본값: 없음", "tip_denoise_step": "노이즈 제거를 적용하는 시점:\n- 1차 복원 후: 업스케일링(2차 복원) 전에 적용. 256x256 해상도에서 처리.\n- 2차 복원 후: 업스케일링 후, 최종 출력 직전에 적용. 전체 해상도에서 처리.\n\n기본값: 1차 복원 후", "tip_secondary_restoration": "복원된 영역을 256x256에서 1024 픽셀로 확대하는 선택적 추가 처리입니다. 클로즈업이나 4K 영상에서 선명도가 향상됩니다.\n\nUNet 4x는 서포터 전용 모델입니다. TVAI보다 훨씬 빠르고 고품질입니다. Jasna를 활성화하면 사용할 수 있습니다.\nRTX Super Res는 빠르지만 품질은 보통입니다.\nTopaz TVAI는 별도 구매 및 설치가 필요합니다.", @@ -430,6 +445,7 @@ "error_cannot_start": "처리를 시작할 수 없습니다:", "error_invalid_tvai": "잘못된 TVAI 구성", "error_post_export_command_required": "내보내기 후 사용자 정의 명령이 필요합니다", + "error_ltx_segments": "LTX는 영상 전체를 처리합니다. 구간을 지우거나 '표준'을 선택하세요: {files}", # Settings panel "dialog_select_tvai_ffmpeg": "Topaz Video ffmpeg.exe 선택", diff --git a/jasna/gui/locales/th.py b/jasna/gui/locales/th.py index c21c2810..f3dddd01 100644 --- a/jasna/gui/locales/th.py +++ b/jasna/gui/locales/th.py @@ -77,6 +77,7 @@ "segments_preview_reset_view": "รีเซ็ตมุมมอง", "segments_preview_reset_view_hint": "รีเซ็ตการซูมและการเลื่อนเพื่อแสดงทั้งเฟรม", "segments_restore_gpu_busy": "ใช้งานไม่ได้ขณะกำลังประมวลผล", + "segments_restore_ltx_unavailable": "ยังดูตัวอย่างกับ LTX ไม่ได้", "segments_restore_loading_models": "กำลังโหลดโมเดลกู้คืน…", "segments_restore_restoring": "กำลังกู้คืน…", "segments_restore_failed": "ตัวอย่างการกู้คืนล้มเหลว: {message}", @@ -196,6 +197,16 @@ # Sections "section_basic": "การประมวลผลพื้นฐาน", + "section_restoration_model": "โมเดลฟื้นฟู", + "model_basicvsrpp": "มาตรฐาน", + "model_basicvsrpp_description": "เร็ว เหมาะกับวิดีโอส่วนใหญ่", + "model_ltx": "LTX", + "model_ltx_description": "รายละเอียดดีที่สุด ช้ากว่ามาก และต้องใช้ GPU NVIDIA ที่แรง", + "model_ltx_needs_nvidia": "ต้องใช้ GPU NVIDIA", + "model_ltx_not_installed": "ยังไม่ได้ติดตั้งโมเดล", + "ltx_seed": "ซีด", + "ltx_fast": "โหมดเร็ว", + "ltx_large_canvas": "โมเสกขนาดใหญ่คมชัดขึ้น", "section_advanced": "การประมวลผลขั้นสูง", "section_secondary": "การฟื้นฟูขั้นที่สอง", "section_image_restoration": "การซ่อมแซมภาพ (ภาพนิ่ง)", @@ -379,6 +390,10 @@ "tip_vr_mode": "กำหนดวิธีประมวลผล VR180 แบบวางซ้ายขวา\n\nอัตโนมัติจะเปิดเมื่อเฟรมเป็น 2:1 พอดีและสูงกว่า 1080 พิกเซล พบโทเคนชื่อสตูดิโอที่เชื่อถือได้ หรือมีเมทาดาทาเชิงพื้นที่ที่เข้ากันได้ และจะเลือกการฉายสำหรับการฟื้นฟูของแต่ละบริเวณโมเสก (raw／ฟิชอาย／gnomonic) ตามสตูดิโอ (ใช้ raw เมื่อไม่รู้จักสตูดิโอ) SBS บังคับประมวลผลแยกแต่ละตาด้วยการจัดเส้นทางเดียวกัน ส่วน SBS + ฟิชอายบังคับใช้การปรับสภาพแบบฟิชอายกับทุกบริเวณ การตรวจจับ การติดตาม มาสก์ และการผสมจะอยู่ในพิกัดต้นฉบับเสมอ โดยฉายกลับเฉพาะส่วนต่างที่ฟื้นฟูแล้ว\n\nตัวแก้ไขช่วงแสดงเฉพาะตาซ้าย แต่การสแกนและส่งออกจะประมวลผลทั้งสองตา\nค่าเริ่มต้น: อัตโนมัติ", "tip_fp16_mode": "ใช้การคำนวณ half-precision เพื่อลดการใช้ VRAM และมักจะเร็วขึ้น ไม่มีความแตกต่างด้านคุณภาพบน GPU รุ่นใหม่\n\nแนะนำ: เปิดสำหรับ RTX 20-series ขึ้นไป\nค่าเริ่มต้น: เปิด", "tip_compile_basicvsrpp": "คอมไพล์โมเดลฟื้นฟูเป็น TensorRT sub-engine เพื่อเพิ่มความเร็วอย่างมาก (ประมาณ 2-3 เท่า)\nการคอมไพล์ครั้งแรกใช้เวลา 15-60 นาที ปิดแอปอื่นทั้งหมด (รวมเบราว์เซอร์) และอย่าใช้ PC ระหว่างคอมไพล์\nEngine จะถูกแคชและนำกลับมาใช้ในครั้งต่อไปโดยอัตโนมัติ\n\nVRAM ของ Engine: ~1.9 GB (คลิป 60), ~5.4 GB (คลิป 180)\nVRAM สูงสุดขณะประมวลผล: ~7.6 GB (คลิป 60), ~14.7 GB (คลิป 180)\nไม่คอมไพล์: ~6 GB (คลิป 60), ~10.4 GB (คลิป 180)\n\nหาก VRAM ไม่พอ ให้ปิดตัวเลือกนี้หรือลดขนาดคลิป\n\nแนะนำ: เปิดพร้อมขนาดคลิป 60-90\nค่าเริ่มต้น: เปิด", + "tip_ltx_seed": "ซีดอื่นจะให้ผลลัพธ์อีกแบบกับวิดีโอเดียวกัน\nเก็บซีดที่ชอบไว้เพื่อได้ผลเหมือนเดิมอีกครั้ง", + "tip_ltx_new_seed": "ลองซีดสุ่มใหม่", + "tip_ltx_fast": "เร็วขึ้นประมาณ 1.4 เท่า รายละเอียดลดลงเล็กน้อย\nเฉพาะ RTX ซีรีส์ 50", + "tip_ltx_large_canvas": "ฟื้นฟูโมเสกขนาดใหญ่ได้ละเอียดขึ้น แต่ส่วนนั้นช้าลงประมาณ 3 เท่า\nต้องมี VRAM ว่างอย่างน้อย 10 GB", "tip_denoise_strength": "ลดสัญญาณรบกวนและเม็ดในพื้นที่ที่ฟื้นฟูแล้ว ค่ามาก = ราบรื่นขึ้นแต่อาจสูญเสียรายละเอียด\n\nไม่มี: ไม่ลดสัญญาณรบกวน ต่ำ/ปานกลาง: จุดเริ่มต้นที่ดี สูง: ปรับเรียบอย่างหนัก\nค่าเริ่มต้น: ไม่มี", "tip_denoise_step": "เวลาที่จะลดสัญญาณรบกวนในขั้นตอนการประมวลผล:\n- หลังขั้นแรก: ก่อนการขยาย (การฟื้นฟูขั้นที่สอง) ลดสัญญาณรบกวนที่ 256x256\n- หลังขั้นที่สอง: หลังการขยาย ก่อนเอาต์พุตสุดท้าย ลดสัญญาณรบกวนที่ความละเอียดเต็ม\n\nค่าเริ่มต้น: หลังขั้นแรก", "tip_secondary_restoration": "ขั้นตอนเสริมที่ขยายพื้นที่ฟื้นฟูจาก 256x256 เป็น 1024 พิกเซล ช่วยเพิ่มความคมชัดโดยเฉพาะสำหรับภาพใกล้และวิดีโอ 4K\n\nUNet 4x เป็นโมเดลเฉพาะผู้สนับสนุน เร็วกว่า TVAI มากและคุณภาพสูง เปิดใช้งาน Jasna เพื่อใช้โมเดลนี้\nRTX Super Res เร็วแต่คุณภาพพอใช้\nTopaz TVAI ต้องซื้อและติดตั้งแยกต่างหาก", @@ -432,6 +447,7 @@ "error_cannot_start": "ไม่สามารถเริ่มประมวลผล:", "error_invalid_tvai": "การกำหนดค่า TVAI ไม่ถูกต้อง", "error_post_export_command_required": "ต้องระบุคำสั่งกำหนดเองหลังส่งออก", + "error_ltx_segments": "LTX ประมวลผลทั้งวิดีโอ ล้างช่วงที่เลือกหรือเลือก “มาตรฐาน”: {files}", # Settings panel "dialog_select_tvai_ffmpeg": "เลือก Topaz Video ffmpeg.exe", diff --git a/jasna/gui/locales/zh.py b/jasna/gui/locales/zh.py index 2ecc9168..4dcad8da 100644 --- a/jasna/gui/locales/zh.py +++ b/jasna/gui/locales/zh.py @@ -77,6 +77,7 @@ "segments_preview_reset_view": "重置视图", "segments_preview_reset_view_hint": "重置缩放和平移并显示完整画面", "segments_restore_gpu_busy": "处理进行中，暂不可用", + "segments_restore_ltx_unavailable": "LTX 暂不支持预览", "segments_restore_loading_models": "正在加载修复模型…", "segments_restore_restoring": "正在修复…", "segments_restore_failed": "修复预览失败：{message}", @@ -196,6 +197,16 @@ # Sections "section_basic": "基本处理", + "section_restoration_model": "修复模型", + "model_basicvsrpp": "标准", + "model_basicvsrpp_description": "速度快，适合大多数视频。", + "model_ltx": "LTX", + "model_ltx_description": "细节最好。速度慢很多，需要高性能 NVIDIA 显卡。", + "model_ltx_needs_nvidia": "需要 NVIDIA 显卡。", + "model_ltx_not_installed": "模型未安装。", + "ltx_seed": "种子", + "ltx_fast": "快速模式", + "ltx_large_canvas": "大马赛克更清晰", "section_advanced": "高级处理", "section_secondary": "二次修复", "section_image_restoration": "图像修复（静态图片）", @@ -397,6 +408,10 @@ "tip_vr_mode": "控制并排 VR180 视频的处理方式。\n\n自动模式会在画面为严格 2:1 且高度超过 1080 像素、文件名包含可信的片商标记或具有兼容空间元数据时启用，并按片商为每个检测到的马赛克区域选择修复投影（raw／鱼眼／等距投影，未知片商时使用 raw）。SBS 采用相同路由分别处理双眼；SBS + 鱼眼则对所有区域强制使用鱼眼条件化。检测、跟踪、蒙版和混合始终保持在源坐标中，仅将修复后的差分重新投影回去。\n\n区间编辑器只显示左眼，但扫描和导出会处理双眼。\n默认值：自动", "tip_fp16_mode": "使用半精度计算来减少显存占用，通常还能提升速度。在现代显卡上几乎无画质损失。\n\n建议：RTX 20 系列及以上显卡开启。\n默认值：开启", "tip_compile_basicvsrpp": "将修复模型编译为 TensorRT 子引擎，大幅提升速度（约 2-3 倍）。\n首次编译需要 15-60 分钟。请关闭所有其他应用程序（包括浏览器），编译期间请勿使用电脑。\n引擎会被缓存，后续运行自动复用。\n\n引擎显存：约 1.9 GB（片段 60）、约 5.4 GB（片段 180）。\n处理时峰值显存：约 7.6 GB（片段 60）、约 14.7 GB（片段 180）。\n不编译时：约 6 GB（片段 60）、约 10.4 GB（片段 180）。\n\n如果显存不足，请关闭此选项或降低片段大小。\n\n建议：配合片段大小 60-90 开启。\n默认值：开启", + "tip_ltx_seed": "换一个种子，同一视频会得到另一种效果。\n保留喜欢的种子即可再次得到相同结果。", + "tip_ltx_new_seed": "换一个随机种子", + "tip_ltx_fast": "约快 1.4 倍，细节略少。\n仅限 RTX 50 系列。", + "tip_ltx_large_canvas": "大马赛克修复得更细致，但这些部分约慢 3 倍。\n需要至少 10 GB 可用显存。", "tip_denoise_strength": "降低修复区域的噪点和颗粒感。强度越高画面越平滑，但可能丢失细节。\n\n无：不降噪。低/中：推荐起步值。高：强力平滑。\n默认值：无", "tip_denoise_step": "降噪在处理流程中的应用时机：\n- 主修复后：在放大（二次修复）之前降噪，在 256x256 分辨率下处理。\n- 二次修复后：在放大之后、最终输出之前降噪，在完整分辨率下处理。\n\n默认值：主修复后", "tip_secondary_restoration": "可选的第二步处理，将修复区域从 256x256 放大到 1024 像素。可提升清晰度，特别是近景和 4K 视频。\n\nUNet 4x 是支持者专属模型，比 TVAI 快很多，质量高。激活 Jasna 后可用。\nRTX Super Res 是速度快但质量一般的替代方案。\nTopaz TVAI 需要单独购买和安装。", @@ -450,6 +465,7 @@ "error_cannot_start": "无法开始处理：", "error_invalid_tvai": "无效的 TVAI 配置", "error_post_export_command_required": "需要填写导出后的自定义命令", + "error_ltx_segments": "LTX 会处理整个视频。请清除片段或选择“标准”：{files}", # Settings panel "dialog_select_tvai_ffmpeg": "选择 Topaz Video ffmpeg.exe", diff --git a/jasna/gui/models.py b/jasna/gui/models.py index 860978f5..110b06a1 100644 --- a/jasna/gui/models.py +++ b/jasna/gui/models.py @@ -10,6 +10,7 @@ from jasna.gui.paths import get_settings_path from jasna.segments import SegmentRange +from jasna.session_config import LTX_DEFAULT_SEED logger = logging.getLogger(__name__) @@ -103,6 +104,12 @@ def begin_processing(self) -> JobProcessingSnapshot | None: @dataclass class AppSettings: + # Restoration model + restoration_model: str = "basicvsrpp" # basicvsrpp, ltx + ltx_seed: int = LTX_DEFAULT_SEED + ltx_fast: bool = False + ltx_large_canvas: bool = False + # Basic processing batch_size: int = 4 max_clip_size: int = 90 diff --git a/jasna/gui/restoration_preview.py b/jasna/gui/restoration_preview.py index a8350739..b0f3d37a 100644 --- a/jasna/gui/restoration_preview.py +++ b/jasna/gui/restoration_preview.py @@ -15,6 +15,7 @@ from PIL import Image +from jasna.gui.locales import t from jasna.gui.queues import replace_pending from jasna.gui.models import AppSettings from jasna.gui.video_session import build_video_session, release_session_memory, video_session_key @@ -325,6 +326,8 @@ def _run(self) -> None: if isinstance(command, _Cancel): continue try: + if command.settings.restoration_model == "ltx": + raise ValueError(t("segments_restore_ltx_unavailable")) key = video_session_key(command.settings) if session is None or key != session_key: pipeline = None diff --git a/jasna/gui/segment_editor.py b/jasna/gui/segment_editor.py index c93db112..969ff229 100644 --- a/jasna/gui/segment_editor.py +++ b/jasna/gui/segment_editor.py @@ -994,6 +994,10 @@ def _toggle_restoration_preview(self) -> None: if self._is_gpu_busy(): self._restore_toggle.deselect() return + if self._get_settings().restoration_model == "ltx": + self._restore_toggle.deselect() + self._show_preview_message(t("segments_restore_ltx_unavailable"), Colors.STATUS_WARNING) + return self._restore_active = True self._set_playing(False) self._restore_toggle.select() diff --git a/jasna/gui/settings_panel.py b/jasna/gui/settings_panel.py index 546c81df..b56ddce3 100644 --- a/jasna/gui/settings_panel.py +++ b/jasna/gui/settings_panel.py @@ -24,6 +24,7 @@ from jasna.gui.settings_sections.encoding import EncodingSection from jasna.gui.settings_sections.image_restoration import ImageRestorationSection from jasna.gui.settings_sections.post_export import PostExportSection +from jasna.gui.settings_sections.restoration_model import RestorationModelSection from jasna.gui.settings_sections.secondary import SecondarySection logger = logging.getLogger(__name__) @@ -34,7 +35,7 @@ class SettingsPanel(ctk.CTkFrame): """Right panel composing the settings sections; widgets live in the sections.""" - def __init__(self, master, preset_manager: PresetManager, **kwargs): + def __init__(self, master, preset_manager: PresetManager, *, ltx_installed: bool, **kwargs): super().__init__( master, fg_color=Colors.BG_PANEL, @@ -52,7 +53,7 @@ def __init__(self, master, preset_manager: PresetManager, **kwargs): self._build_preset_bar() self._build_scrollable() - self._build_sections() + self._build_sections(ltx_installed) self._apply_preset(self._current_preset) def _build_preset_bar(self): @@ -211,8 +212,15 @@ def _build_scrollable(self): ) self._scroll.pack(fill="both", expand=True, padx=Sizing.PADDING_MEDIUM, pady=(0, Sizing.PADDING_MEDIUM)) - def _build_sections(self): - self._sections = [ + def _build_sections(self, ltx_installed: bool): + self._model_section = RestorationModelSection( + self._scroll, + self._widgets, + self._mark_modified, + self._on_restoration_model_changed, + ltx_installed=ltx_installed, + ) + self._model_aware_sections = [ BasicSection( self._scroll, self._widgets, @@ -221,6 +229,10 @@ def _build_sections(self): ), AdvancedSection(self._scroll, self._widgets, self._mark_modified), SecondarySection(self._scroll, self._widgets), + ] + self._sections = [ + self._model_section, + *self._model_aware_sections, ImageRestorationSection( self._scroll, self._widgets, @@ -232,6 +244,14 @@ def _build_sections(self): PostExportSection(self._scroll, self._widgets, self._mark_modified), ] + def _on_restoration_model_changed(self, model: str): + for section in self._model_aware_sections: + section.set_model(model) + + def set_gpu_support(self, *, nvidia: bool, blackwell: bool): + self._model_section.set_gpu_support(nvidia=nvidia, blackwell=blackwell) + self._update_modified_indicator() + def set_on_interactive_image_restore(self, callback: callable): self._on_interactive_image_restore = callback @@ -355,6 +375,8 @@ def set_enabled(self, enabled: bool): """Enable or disable the settings; the queue-wide post-export action stays editable while processing.""" state = "normal" if enabled else "disabled" + self._model_section.set_enabled(enabled) + # Preset bar buttons self._preset_dropdown.configure(state=state) self._create_btn.configure(state=state) diff --git a/jasna/gui/settings_sections/advanced.py b/jasna/gui/settings_sections/advanced.py index a1084a0b..4879b91f 100644 --- a/jasna/gui/settings_sections/advanced.py +++ b/jasna/gui/settings_sections/advanced.py @@ -10,6 +10,7 @@ ValueOptionMenu, create_slider_value_label, get_tooltip, + pack_rows, ) from jasna.gui.theme import Colors, Fonts, Sizing @@ -31,7 +32,6 @@ def __init__(self, parent, widgets: dict, on_modified): # Temporal Overlap row row1 = ctk.CTkFrame(inner, fg_color="transparent") - row1.pack(fill="x", pady=(0, Sizing.PADDING_SMALL)) add_setting_label(row1, "temporal_overlap") @@ -49,7 +49,6 @@ def __init__(self, parent, widgets: dict, on_modified): # Max Detection Gap row gap_row = ctk.CTkFrame(inner, fg_color="transparent") - gap_row.pack(fill="x", pady=(0, Sizing.PADDING_SMALL)) add_setting_label(gap_row, "max_detection_gap") @@ -68,7 +67,6 @@ def __init__(self, parent, widgets: dict, on_modified): # Min Detection Duration row mindur_row = ctk.CTkFrame(inner, fg_color="transparent") - mindur_row.pack(fill="x", pady=(0, Sizing.PADDING_SMALL)) add_setting_label(mindur_row, "min_detection_duration") @@ -87,7 +85,6 @@ def __init__(self, parent, widgets: dict, on_modified): # Scene cut detection toggle scene_row = ctk.CTkFrame(inner, fg_color="transparent") - scene_row.pack(fill="x", pady=(0, Sizing.PADDING_SMALL)) scene_frame = ctk.CTkFrame(scene_row, fg_color=Colors.BG_CARD, corner_radius=6) scene_frame.pack(fill="x") @@ -102,7 +99,6 @@ def __init__(self, parent, widgets: dict, on_modified): # Crossfade toggle row2 = ctk.CTkFrame(inner, fg_color="transparent") - row2.pack(fill="x", pady=(0, Sizing.PADDING_SMALL)) crossfade_frame = ctk.CTkFrame(row2, fg_color=Colors.BG_CARD, corner_radius=6) crossfade_frame.pack(fill="x") @@ -116,7 +112,6 @@ def __init__(self, parent, widgets: dict, on_modified): self._widgets["enable_crossfade"].select() row_vr = ctk.CTkFrame(inner, fg_color="transparent") - row_vr.pack(fill="x", pady=(0, Sizing.PADDING_SMALL)) vr_label = ctk.CTkLabel( row_vr, text=t("vr_mode"), @@ -155,7 +150,6 @@ def __init__(self, parent, widgets: dict, on_modified): # Denoising Strength row3 = ctk.CTkFrame(inner, fg_color="transparent") - row3.pack(fill="x", pady=(0, Sizing.PADDING_SMALL)) add_setting_label(row3, "denoise_strength") @@ -178,7 +172,6 @@ def __init__(self, parent, widgets: dict, on_modified): # Denoise Step row4 = ctk.CTkFrame(inner, fg_color="transparent") - row4.pack(fill="x") add_setting_label(row4, "denoise_step") @@ -197,6 +190,18 @@ def __init__(self, parent, widgets: dict, on_modified): self._widgets["denoise_step"].pack(side="right") self._widgets["denoise_step"].set_value("after_primary") + row_gap = dict(fill="x", pady=(0, Sizing.PADDING_SMALL)) + self._rows = [ + (row1, row_gap), (gap_row, row_gap), (mindur_row, row_gap), (scene_row, row_gap), + (row2, row_gap), (row_vr, row_gap), (row3, row_gap), (row4, dict(fill="x")), + ] + self._vr_row = row_vr + self.set_model("basicvsrpp") + + def set_model(self, model: str) -> None: + hidden = tuple(row for row, _options in self._rows if row is not self._vr_row) if model == "ltx" else () + pack_rows(self._rows, hidden) + def _on_slider_change(self, key: str, value: int): self._widgets[f"{key}_val"].configure(text=str(value)) self._on_modified() diff --git a/jasna/gui/settings_sections/basic.py b/jasna/gui/settings_sections/basic.py index eca74b05..ed3a2fc5 100644 --- a/jasna/gui/settings_sections/basic.py +++ b/jasna/gui/settings_sections/basic.py @@ -9,6 +9,7 @@ add_setting_label, ValueOptionMenu, create_slider_value_label, + pack_rows, ) from jasna.gui.theme import Colors, Fonts, Sizing @@ -46,7 +47,6 @@ def __init__(self, parent, widgets: dict, on_modified, on_max_clip_size_change): # Max Clip Size slider row1 = ctk.CTkFrame(inner, fg_color="transparent") - row1.pack(fill="x", pady=(0, Sizing.PADDING_SMALL)) add_setting_label(row1, "max_clip_size") @@ -64,7 +64,6 @@ def __init__(self, parent, widgets: dict, on_modified, on_max_clip_size_change): # Detection Model row2 = ctk.CTkFrame(inner, fg_color="transparent") - row2.pack(fill="x", pady=(0, Sizing.PADDING_SMALL)) add_setting_label(row2, "detection_model") @@ -83,7 +82,6 @@ def __init__(self, parent, widgets: dict, on_modified, on_max_clip_size_change): # Detection Threshold row3 = ctk.CTkFrame(inner, fg_color="transparent") - row3.pack(fill="x", pady=(0, Sizing.PADDING_SMALL)) add_setting_label(row3, "detection_threshold", "detection_score_threshold") @@ -101,10 +99,8 @@ def __init__(self, parent, widgets: dict, on_modified, on_max_clip_size_change): # Toggles row - FP16 Mode and Compile BasicVSR++ row4 = ctk.CTkFrame(inner, fg_color="transparent") - row4.pack(fill="x", pady=(Sizing.PADDING_SMALL, 0)) fp16_frame = ctk.CTkFrame(row4, fg_color=Colors.BG_CARD, corner_radius=6) - fp16_frame.pack(side="left", fill="x", expand=True, padx=(0, 4)) add_setting_label(fp16_frame, "fp16_mode", in_card=True) self._widgets["fp16_mode"] = CompactSwitch( fp16_frame, @@ -115,7 +111,6 @@ def __init__(self, parent, widgets: dict, on_modified, on_max_clip_size_change): self._widgets["fp16_mode"].select() compile_frame = ctk.CTkFrame(row4, fg_color=Colors.BG_CARD, corner_radius=6) - compile_frame.pack(side="right", fill="x", expand=True, padx=(4, 0)) add_setting_label(compile_frame, "compile_basicvsrpp", in_card=True) self._widgets["compile_basicvsrpp"] = CompactSwitch( compile_frame, @@ -127,7 +122,6 @@ def __init__(self, parent, widgets: dict, on_modified, on_max_clip_size_change): # File Conflict dropdown row5 = ctk.CTkFrame(inner, fg_color="transparent") - row5.pack(fill="x", pady=(Sizing.PADDING_SMALL, 0)) add_setting_label(row5, "file_conflict") @@ -152,6 +146,21 @@ def __init__(self, parent, widgets: dict, on_modified, on_max_clip_size_change): self._widgets["file_conflict"].pack(side="right") self._widgets["file_conflict"].set_value("auto_rename") + row_gap = dict(fill="x", pady=(0, Sizing.PADDING_SMALL)) + toggle_gap = dict(fill="x", pady=(Sizing.PADDING_SMALL, 0)) + self._rows = [(row1, row_gap), (row2, row_gap), (row3, row_gap), (row4, toggle_gap), (row5, toggle_gap)] + self._toggles = [ + (fp16_frame, dict(side="left", fill="x", expand=True, padx=(0, 4))), + (compile_frame, dict(side="right", fill="x", expand=True, padx=(4, 0))), + ] + self._standard_only = (row1, compile_frame) + self.set_model("basicvsrpp") + + def set_model(self, model: str) -> None: + hidden = self._standard_only if model == "ltx" else () + pack_rows(self._rows, hidden) + pack_rows(self._toggles, hidden) + def _on_max_clip_size_slider(self, value: float): max_clip_size = int(value) self._widgets["max_clip_size_val"].configure(text=str(max_clip_size)) diff --git a/jasna/gui/settings_sections/restoration_model.py b/jasna/gui/settings_sections/restoration_model.py new file mode 100644 index 00000000..60f8665e --- /dev/null +++ b/jasna/gui/settings_sections/restoration_model.py @@ -0,0 +1,206 @@ +"""Restoration model choice (Standard BasicVSR++ or LTX) and the LTX options.""" + +import random +from typing import Callable + +import customtkinter as ctk + +from jasna.gui.components import CollapsibleSection, Tooltip +from jasna.gui.icons import CompactSwitch, NativeIconButton +from jasna.gui.locales import t +from jasna.gui.settings_sections.widgets import add_setting_label, get_tooltip +from jasna.gui.theme import Colors, Fonts, Sizing +from jasna.session_config import LTX_DEFAULT_SEED + +_MODELS = ("basicvsrpp", "ltx") + + +def ltx_unavailable_reason(*, installed: bool, nvidia: bool | None) -> str | None: + """Locale key saying why LTX cannot be picked; ``nvidia`` is None until the GPU is known.""" + if nvidia is False: + return "model_ltx_needs_nvidia" + if not installed: + return "model_ltx_not_installed" + return None + + +def parse_seed(text: str) -> int: + try: + return int(text.strip()) + except ValueError: + return LTX_DEFAULT_SEED + + +class _ModelCard(ctk.CTkFrame): + def __init__(self, master, variable: ctk.StringVar, value: str, on_select: Callable[[str], None]): + super().__init__(master, fg_color=Colors.BG_CARD, corner_radius=8, border_width=2, border_color=Colors.BG_CARD) + self._value = value + self._on_select = on_select + self._enabled = True + self.radio = ctk.CTkRadioButton( + self, + text=t(f"model_{value}"), + variable=variable, + value=value, + command=lambda: on_select(value), + font=(Fonts.FAMILY, Fonts.SIZE_NORMAL, "bold"), + fg_color=Colors.PRIMARY, + hover_color=Colors.PRIMARY_HOVER, + text_color=Colors.TEXT_PRIMARY, + ) + self.radio.pack(anchor="w", padx=12, pady=(10, 2)) + self.description = ctk.CTkLabel( + self, + text=t(f"model_{value}_description"), + text_color=Colors.STATUS_PENDING, + font=(Fonts.FAMILY, Fonts.SIZE_SMALL), + anchor="w", + justify="left", + wraplength=260, + ) + self.description.pack(fill="x", padx=(40, 12), pady=(0, 10)) + for widget in (self, self.description): + widget.bind("", self._clicked) + + def _clicked(self, _event) -> None: + if self._enabled: + self.radio.invoke() + + def set_enabled(self, enabled: bool) -> None: + self._enabled = enabled + self.radio.configure(state="normal" if enabled else "disabled") + + def set_selected(self, selected: bool) -> None: + self.configure(border_color=Colors.PRIMARY if selected else Colors.BG_CARD) + + +class RestorationModelSection: + def __init__(self, parent, widgets: dict, on_modified, on_model_changed, *, ltx_installed: bool): + self._widgets = widgets + self._on_modified = on_modified + self._on_model_changed = on_model_changed + self._ltx_installed = ltx_installed + self._nvidia: bool | None = None + self._enabled = True + + section = CollapsibleSection(parent, t("section_restoration_model"), expanded=True) + section.pack(fill="x", pady=(0, Sizing.PADDING_SMALL)) + content = section.content + content.configure(corner_radius=Sizing.BORDER_RADIUS) + + inner = ctk.CTkFrame(content, fg_color="transparent") + inner.pack(fill="x", padx=Sizing.PADDING_MEDIUM, pady=Sizing.PADDING_MEDIUM) + + self._model = self._widgets["restoration_model"] = ctk.StringVar(value="basicvsrpp") + cards = ctk.CTkFrame(inner, fg_color="transparent") + cards.pack(fill="x") + cards.grid_columnconfigure((0, 1), weight=1, uniform="model") + self._cards = {} + for column, value in enumerate(_MODELS): + card = _ModelCard(cards, self._model, value, self._on_card_selected) + card.grid(row=0, column=column, sticky="nsew", padx=(0, 4) if column == 0 else (4, 0)) + self._cards[value] = card + + self._ltx_options = ctk.CTkFrame(inner, fg_color=Colors.BG_CARD, corner_radius=6) + seed_row = ctk.CTkFrame(self._ltx_options, fg_color="transparent") + seed_row.pack(fill="x") + add_setting_label(seed_row, "ltx_seed", in_card=True) + self._new_seed_btn = NativeIconButton( + seed_row, + "reset", + 16, + Colors.TEXT_PRIMARY, + Colors.BG_CARD, + Colors.BORDER_LIGHT, + Colors.BORDER_LIGHT, + self._on_new_seed, + 28, + 28, + ) + self._new_seed_btn.pack(side="right", padx=(4, 12), pady=8) + Tooltip(self._new_seed_btn, get_tooltip("ltx_new_seed")) + self._widgets["ltx_seed"] = ctk.CTkEntry( + seed_row, width=120, fg_color=Colors.BG_PANEL, text_color=Colors.TEXT_PRIMARY, border_color=Colors.BORDER_LIGHT, + ) + self._widgets["ltx_seed"].pack(side="right", pady=8) + self._widgets["ltx_seed"].bind("", lambda _event: self._on_modified()) + + self._fast_row = ctk.CTkFrame(self._ltx_options, fg_color="transparent") + add_setting_label(self._fast_row, "ltx_fast", in_card=True) + self._widgets["ltx_fast"] = CompactSwitch(self._fast_row, self._on_modified, Colors.BG_CARD) + self._widgets["ltx_fast"].pack(side="right", padx=12, pady=8) + + large_row = ctk.CTkFrame(self._ltx_options, fg_color="transparent") + large_row.pack(fill="x") + add_setting_label(large_row, "ltx_large_canvas", in_card=True) + self._widgets["ltx_large_canvas"] = CompactSwitch(large_row, self._on_modified, Colors.BG_CARD) + self._widgets["ltx_large_canvas"].pack(side="right", padx=12, pady=8) + self._large_row = large_row + + self._refresh_ltx_availability() + + def _on_card_selected(self, value: str) -> None: + self._show_model(value) + self._on_modified() + + def _show_model(self, value: str) -> None: + self._model.set(value) + for name, card in self._cards.items(): + card.set_selected(name == value) + if value == "ltx": + self._ltx_options.pack(fill="x", pady=(Sizing.PADDING_SMALL, 0)) + else: + self._ltx_options.pack_forget() + self._on_model_changed(value) + + def _on_new_seed(self) -> None: + self._set_seed(random.randrange(1, 2**31)) + self._on_modified() + + def _set_seed(self, seed: int) -> None: + self._widgets["ltx_seed"].delete(0, "end") + self._widgets["ltx_seed"].insert(0, str(seed)) + + def _refresh_ltx_availability(self) -> None: + reason = ltx_unavailable_reason(installed=self._ltx_installed, nvidia=self._nvidia) + card = self._cards["ltx"] + card.set_enabled(self._enabled and reason is None) + card.description.configure( + text=t(reason) if reason else t("model_ltx_description"), + text_color=Colors.STATUS_WARNING if reason else Colors.STATUS_PENDING, + ) + if reason is not None and self._model.get() == "ltx": + self._show_model("basicvsrpp") + + def set_gpu_support(self, *, nvidia: bool, blackwell: bool) -> None: + self._nvidia = nvidia + if blackwell: + self._fast_row.pack(fill="x", before=self._large_row) + else: + self._widgets["ltx_fast"].deselect() + self._fast_row.pack_forget() + self._refresh_ltx_availability() + + def set_enabled(self, enabled: bool) -> None: + self._enabled = enabled + self._cards["basicvsrpp"].set_enabled(enabled) + self._new_seed_btn.configure(state="normal" if enabled else "disabled") + self._refresh_ltx_availability() + + def apply(self, preset) -> None: + self._set_seed(preset.ltx_seed) + for key, selected in (("ltx_fast", preset.ltx_fast), ("ltx_large_canvas", preset.ltx_large_canvas)): + if selected: + self._widgets[key].select() + else: + self._widgets[key].deselect() + self._show_model(preset.restoration_model) + self._refresh_ltx_availability() + + def collect(self) -> dict: + return { + "restoration_model": self._widgets["restoration_model"].get(), + "ltx_seed": parse_seed(self._widgets["ltx_seed"].get()), + "ltx_fast": self._widgets["ltx_fast"].get() == 1, + "ltx_large_canvas": self._widgets["ltx_large_canvas"].get() == 1, + } diff --git a/jasna/gui/settings_sections/secondary.py b/jasna/gui/settings_sections/secondary.py index 9eb34acc..12c48052 100644 --- a/jasna/gui/settings_sections/secondary.py +++ b/jasna/gui/settings_sections/secondary.py @@ -14,8 +14,10 @@ class SecondarySection: def __init__(self, parent, widgets: dict): self._widgets = widgets + self._previous_section = parent.pack_slaves()[-1] section = CollapsibleSection(parent, t("section_secondary"), expanded=False) section.pack(fill="x", pady=(0, Sizing.PADDING_SMALL)) + self._section = section content = section.content content.configure(corner_radius=Sizing.BORDER_RADIUS) @@ -256,6 +258,12 @@ def _on_secondary_changed(self): elif secondary == "rtx-super-res": self._rtx_frame.pack(fill="x", pady=(Sizing.PADDING_SMALL, 0)) + def set_model(self, model: str) -> None: + if model == "ltx": + self._section.pack_forget() + else: + self._section.pack(fill="x", pady=(0, Sizing.PADDING_SMALL), after=self._previous_section) + def apply(self, preset): self._widgets["secondary_var"].set(preset.secondary_restoration) diff --git a/jasna/gui/settings_sections/widgets.py b/jasna/gui/settings_sections/widgets.py index b6a16815..6f28a668 100644 --- a/jasna/gui/settings_sections/widgets.py +++ b/jasna/gui/settings_sections/widgets.py @@ -76,3 +76,12 @@ def set_value(self, value: str): if label is None: label = next(iter(self._value_to_label.values())) self.set(label) + + +def pack_rows(rows: list[tuple[tk.Misc, dict]], hidden: tuple[tk.Misc, ...]) -> None: + """Pack ``rows`` (widget, pack options) in order, leaving out ``hidden``.""" + for row, _options in rows: + row.pack_forget() + for row, options in rows: + if row not in hidden: + row.pack(**options) diff --git a/jasna/gui/validation.py b/jasna/gui/validation.py index 91483526..24aad60e 100644 --- a/jasna/gui/validation.py +++ b/jasna/gui/validation.py @@ -1,11 +1,11 @@ import os from pathlib import Path -from jasna.gui.models import AppSettings +from jasna.gui.models import AppSettings, JobItem from jasna.gui.locales import t -def validate_gui_start(settings: AppSettings) -> list[str]: +def validate_gui_start(settings: AppSettings, jobs: list[JobItem]) -> list[str]: errors: list[str] = [] from jasna.post_export_action import validate_post_export_action @@ -14,6 +14,12 @@ def validate_gui_start(settings: AppSettings) -> list[str]: except ValueError: errors.append(t("error_post_export_command_required")) + if settings.restoration_model == "ltx": + segmented = [job.filename for job in jobs if job.snapshot_segments()] + if segmented: + errors.append(t("error_ltx_segments", files=", ".join(segmented))) + return errors + if settings.secondary_restoration != "tvai": return errors diff --git a/jasna/gui/video_player.py b/jasna/gui/video_player.py index 7caebb58..4c6a3048 100644 --- a/jasna/gui/video_player.py +++ b/jasna/gui/video_player.py @@ -686,6 +686,7 @@ def _start_playback( def _playback_settings(self) -> AppSettings: return replace( self._base_settings, + restoration_model="basicvsrpp", detection_model=self._model.get(), detection_score_threshold=float(self._threshold.get()), secondary_restoration=self._secondary.get_value(), diff --git a/jasna/gui/video_session.py b/jasna/gui/video_session.py index 6ee88a7b..c53f8f90 100644 --- a/jasna/gui/video_session.py +++ b/jasna/gui/video_session.py @@ -10,7 +10,7 @@ from typing import TYPE_CHECKING, Callable, Mapping from jasna.gui.models import AppSettings -from jasna.session_config import LTX_DEFAULT_SEED, SessionConfig +from jasna.session_config import SessionConfig from jasna.session_factory import RestorationSession if TYPE_CHECKING: @@ -28,7 +28,10 @@ def video_session_key(settings: AppSettings) -> tuple: settings.denoise_strength, settings.denoise_step, settings.secondary_restoration, + settings.restoration_model, ) + if settings.restoration_model == "ltx": + key += (settings.ltx_fast,) if settings.secondary_restoration == "tvai": key += ( settings.tvai_ffmpeg_path, @@ -58,6 +61,8 @@ def video_session_config( from jasna.mosaic.detection_registry import coerce_detection_model_name, require_detection_model_weights det_name = coerce_detection_model_name(str(settings.detection_model)) + standard_model = settings.restoration_model == "basicvsrpp" + secondary_restoration = settings.secondary_restoration if standard_model else "none" return SessionConfig( device="cuda:0", fp16=bool(settings.fp16_mode), @@ -68,26 +73,24 @@ def video_session_config( max_detection_gap=int(settings.max_detection_gap), min_detection_duration=int(settings.min_detection_duration), scene_detection=bool(settings.scene_detection), - restoration_model_name="basicvsrpp", - restoration_model_path=default_restoration_model_path("basicvsrpp"), - ltx_large_canvas=False, - ltx_seed=LTX_DEFAULT_SEED, - ltx_fast=False, + restoration_model_name=settings.restoration_model, + restoration_model_path=default_restoration_model_path(settings.restoration_model), + ltx_large_canvas=bool(settings.ltx_large_canvas), + ltx_seed=int(settings.ltx_seed), + ltx_fast=bool(settings.ltx_fast), compile_basicvsrpp=bool(settings.compile_basicvsrpp), max_clip_size=int(settings.max_clip_size), temporal_overlap=int(settings.temporal_overlap), enable_crossfade=bool(settings.enable_crossfade), - denoise_strength=settings.denoise_strength, + denoise_strength=settings.denoise_strength if standard_model else "none", denoise_step=settings.denoise_step, - secondary_restoration=settings.secondary_restoration, + secondary_restoration=secondary_restoration, tvai_ffmpeg_path=settings.tvai_ffmpeg_path, tvai_model=settings.tvai_model, tvai_scale=int(settings.tvai_scale), tvai_args=settings.tvai_args, tvai_workers=int(settings.tvai_workers), - tvai_denoise=bool( - settings.tvai_denoise and settings.secondary_restoration == "tvai" - ), + tvai_denoise=bool(settings.tvai_denoise and secondary_restoration == "tvai"), rtx_scale=int(settings.rtx_scale), rtx_quality=settings.rtx_quality.lower(), rtx_denoise=settings.rtx_denoise.lower(), diff --git a/jasna/ltx/model_files.py b/jasna/ltx/model_files.py index c08284cf..ec8b932c 100644 --- a/jasna/ltx/model_files.py +++ b/jasna/ltx/model_files.py @@ -24,9 +24,7 @@ class LtxModelFiles: @classmethod def from_dir(cls, directory: Path, *, fast: bool) -> LtxModelFiles: """The model bundle; ``fast`` picks the NVFP4 transformer over the INT8 one.""" - transformer = "transformer-fast" if fast else "transformer" - paths = [directory / f"{name}.safetensors" for name in (transformer, "vae", "vae-decoder")] - paths = [path.with_name(path.name + ".enc") if is_frozen() or not path.is_file() else path for path in paths] + paths = _bundle_paths(directory, fast=fast) if any(path.suffix == ".enc" for path in paths): from jasna.protection.protected_tensors import check_ltx_license check_ltx_license() @@ -34,3 +32,14 @@ def from_dir(cls, directory: Path, *, fast: bool) -> LtxModelFiles: if missing: raise FileNotFoundError(f"LTX model files missing: {', '.join(missing)}") return cls(*paths) + + +def _bundle_paths(directory: Path, *, fast: bool) -> list[Path]: + transformer = "transformer-fast" if fast else "transformer" + paths = [directory / f"{name}.safetensors" for name in (transformer, "vae", "vae-decoder")] + return [path.with_name(path.name + ".enc") if is_frozen() or not path.is_file() else path for path in paths] + + +def bundle_present(directory: Path) -> bool: + """Whether the quality model files are installed; checks no license.""" + return all(path.is_file() for path in _bundle_paths(directory, fast=False)) diff --git a/tests/test_gui_engine_preflight.py b/tests/test_gui_engine_preflight.py index 30bc8f9d..dd3b7f3c 100644 --- a/tests/test_gui_engine_preflight.py +++ b/tests/test_gui_engine_preflight.py @@ -173,3 +173,12 @@ def test_amd_preflight_has_no_engine_requirements( assert result.requirements == () assert result.missing == () assert result.should_warn_first_run_slow is False + + +def test_ltx_preflight_skips_standard_model_engines(monkeypatch, tmp_path: Path) -> None: + monkeypatch.chdir(tmp_path) + (tmp_path / "model_weights").mkdir(parents=True, exist_ok=True) + + settings = AppSettings(restoration_model="ltx", secondary_restoration="unet-4x") + + assert {r.key for r in run_engine_preflight(settings).requirements} == {"rfdetr"} diff --git a/tests/test_gui_settings_persistence_paths.py b/tests/test_gui_settings_persistence_paths.py index 45e5fe48..955e5c73 100644 --- a/tests/test_gui_settings_persistence_paths.py +++ b/tests/test_gui_settings_persistence_paths.py @@ -181,3 +181,13 @@ def test_preset_manager_saves_do_not_overwrite_keys_changed_elsewhere(monkeypatc assert reloaded.get_preset("MyPreset") is not None assert reloaded.get_system_check_passed_version() == "9.9.9" assert reloaded.get_last_selected() == "MyPreset" + + +def test_preset_manager_persists_ltx_settings(monkeypatch, tmp_path: Path) -> None: + monkeypatch.setattr(os_utils.sys, "platform", "win32", raising=False) + monkeypatch.setenv("APPDATA", str(tmp_path / "Roaming")) + ltx = AppSettings(restoration_model="ltx", ltx_seed=42, ltx_fast=True, ltx_large_canvas=True) + + assert PresetManager().create_preset("Ltx", ltx) + + assert PresetManager().get_preset("Ltx") == ltx diff --git a/tests/test_gui_settings_scroll.py b/tests/test_gui_settings_scroll.py index 336efd1b..d2a9d724 100644 --- a/tests/test_gui_settings_scroll.py +++ b/tests/test_gui_settings_scroll.py @@ -12,7 +12,7 @@ def _settings_panel(root: ctk.CTk) -> SettingsPanel: - panel = SettingsPanel(root, PresetManager()) + panel = SettingsPanel(root, PresetManager(), ltx_installed=True) panel.pack(fill="both", expand=True) root.update() return panel @@ -46,7 +46,7 @@ def test_settings_scrollbar_only_appears_when_sections_overflow() -> None: pytest.skip(f"Tk display unavailable: {exc}") try: - root.geometry("420x420") + root.geometry("420x480") panel = _settings_panel(root) assert panel._scroll._scrollbar.winfo_ismapped() diff --git a/tests/test_gui_settings_sections.py b/tests/test_gui_settings_sections.py index 197cf7e2..79b2860b 100644 --- a/tests/test_gui_settings_sections.py +++ b/tests/test_gui_settings_sections.py @@ -16,6 +16,7 @@ from jasna.gui.settings_sections.encoding import EncodingSection from jasna.gui.settings_sections.image_restoration import ImageRestorationSection from jasna.gui.settings_sections.post_export import PostExportSection +from jasna.gui.settings_sections.restoration_model import RestorationModelSection from jasna.gui.settings_sections.secondary import SecondarySection from jasna.gui.settings_sections.widgets import ValueOptionMenu @@ -63,6 +64,10 @@ def test_value_option_menu_falls_back_to_first_option_for_unknown_value() -> Non def _fake_section_widgets() -> dict: return { + "restoration_model": _FakeWidget("ltx"), + "ltx_seed": _FakeWidget(" 123 "), + "ltx_fast": _FakeWidget(1), + "ltx_large_canvas": _FakeWidget(0), "max_clip_size": _FakeWidget(90), "fp16_mode": _FakeWidget(1), "detection_model": _FakeWidget("rfdetr-v6"), @@ -110,6 +115,7 @@ def _collect_all(widgets: dict) -> dict: fake = SimpleNamespace(_widgets=widgets) values: dict = {} for section in ( + RestorationModelSection, BasicSection, AdvancedSection, SecondarySection, @@ -143,6 +149,10 @@ def test_sections_collect_internal_values_without_translation_lookups() -> None: assert values["retarget_high_fps"] is True assert values["fmp4"] is True assert values["sharpen_strength"] == 0.35 + assert values["restoration_model"] == "ltx" + assert values["ltx_seed"] == 123 + assert values["ltx_fast"] is True + assert values["ltx_large_canvas"] is False def test_sections_collect_covers_all_widget_backed_appsettings_fields() -> None: @@ -182,9 +192,9 @@ def _basic_section_panel(monkeypatch, tmp_path): from jasna.gui.settings_panel import SettingsPanel - panel = SettingsPanel(root, PresetManager()) + panel = SettingsPanel(root, PresetManager(), ltx_installed=True) try: - yield panel, panel._sections[0] + yield panel, next(section for section in panel._sections if isinstance(section, BasicSection)) finally: root.destroy() @@ -353,7 +363,7 @@ def test_settings_panel_get_settings_is_locale_independent(monkeypatch, tmp_path try: from jasna.gui.settings_panel import SettingsPanel - panel = SettingsPanel(root, PresetManager()) + panel = SettingsPanel(root, PresetManager(), ltx_installed=True) assert panel.get_settings() == replace(AppSettings(), encoder_cq=28) assert panel._saved_preset_settings == panel.get_settings() finally: @@ -394,3 +404,113 @@ def test_only_queue_wide_post_export_action_stays_editable_while_processing(_bas assert panel._widgets["post_export_command"].cget("state") == "normal" assert panel._widgets["post_export_video_command"].cget("state") == "disabled" assert panel._widgets["max_clip_size"].cget("state") == "disabled" + + +def test_ltx_unavailable_reason() -> None: + from jasna.gui.settings_sections.restoration_model import ltx_unavailable_reason + + assert ltx_unavailable_reason(installed=True, nvidia=None) is None + assert ltx_unavailable_reason(installed=True, nvidia=True) is None + assert ltx_unavailable_reason(installed=True, nvidia=False) == "model_ltx_needs_nvidia" + assert ltx_unavailable_reason(installed=False, nvidia=True) == "model_ltx_not_installed" + + +def test_parse_seed_falls_back_to_default() -> None: + from jasna.gui.settings_sections.restoration_model import parse_seed + from jasna.session_config import LTX_DEFAULT_SEED + + assert parse_seed(" 42 ") == 42 + assert parse_seed("abc") == LTX_DEFAULT_SEED + + +@pytest.fixture +def _panel_factory(monkeypatch, tmp_path): + monkeypatch.setattr(os_utils.sys, "platform", "linux", raising=False) + monkeypatch.setenv("XDG_CONFIG_HOME", str(tmp_path / "xdg")) + try: + root = ctk.CTk() + except TclError as exc: + pytest.skip(f"Tk display unavailable: {exc}") + + from jasna.gui.settings_panel import SettingsPanel + + try: + yield lambda installed: SettingsPanel(root, PresetManager(), ltx_installed=installed) + finally: + root.destroy() + + +def _select_ltx(panel) -> None: + panel._model_section._cards["ltx"].radio.invoke() + + +def test_ltx_hides_standard_only_settings_and_restores_them(_panel_factory) -> None: + panel = _panel_factory(True) + secondary = panel._model_aware_sections[2] + panel.update_idletasks() + clip_row = panel._widgets["max_clip_size"].master + overlap_row = panel._widgets["temporal_overlap"].master + + _select_ltx(panel) + + assert panel.get_settings().restoration_model == "ltx" + assert not clip_row.winfo_manager() + assert not panel._widgets["compile_basicvsrpp"].master.winfo_manager() + assert not overlap_row.winfo_manager() + assert panel._widgets["vr_mode"].master.winfo_manager() + assert not secondary._section.winfo_manager() + assert panel._model_section._ltx_options.winfo_manager() + + panel._model_section._cards["basicvsrpp"].radio.invoke() + + assert clip_row.winfo_manager() and overlap_row.winfo_manager() + assert secondary._section.winfo_manager() + assert not panel._model_section._ltx_options.winfo_manager() + order = panel._scroll.pack_slaves() + assert order.index(secondary._section) == order.index(secondary._previous_section) + 1 + + +def test_ltx_card_disabled_when_model_missing(_panel_factory) -> None: + panel = _panel_factory(False) + + assert panel._model_section._cards["ltx"].radio.cget("state") == "disabled" + panel._model_section.apply(AppSettings(restoration_model="ltx")) + assert panel.get_settings().restoration_model == "basicvsrpp" + + +def test_gpu_support_disables_ltx_and_shows_fast_mode_on_blackwell(_panel_factory) -> None: + panel = _panel_factory(True) + section = panel._model_section + section.apply(AppSettings(restoration_model="ltx", ltx_fast=True)) + + section.set_gpu_support(nvidia=True, blackwell=True) + assert section._fast_row.winfo_manager() + assert panel.get_settings().ltx_fast is True + + section.set_gpu_support(nvidia=True, blackwell=False) + assert not section._fast_row.winfo_manager() + assert panel.get_settings().ltx_fast is False + + section.set_gpu_support(nvidia=False, blackwell=False) + assert panel.get_settings().restoration_model == "basicvsrpp" + assert section._cards["ltx"].radio.cget("state") == "disabled" + + +def test_ltx_settings_round_trip_through_panel(_panel_factory) -> None: + panel = _panel_factory(True) + ltx = replace(AppSettings(), restoration_model="ltx", ltx_seed=99, ltx_large_canvas=True, encoder_cq=28) + + for section in panel._sections: + section.apply(ltx) + + assert panel.get_settings() == ltx + + +def test_disabling_settings_keeps_unavailable_ltx_disabled(_panel_factory) -> None: + panel = _panel_factory(False) + + panel.set_enabled(False) + panel.set_enabled(True) + + assert panel._model_section._cards["ltx"].radio.cget("state") == "disabled" + assert panel._model_section._cards["basicvsrpp"].radio.cget("state") == "normal" diff --git a/tests/test_gui_startup_and_wizard.py b/tests/test_gui_startup_and_wizard.py index 40a88adf..e7c962eb 100644 --- a/tests/test_gui_startup_and_wizard.py +++ b/tests/test_gui_startup_and_wizard.py @@ -38,22 +38,31 @@ def test_evaluate_sysmem_only_failure_is_warning_not_required(): assert required_failure is False # but sysmem is warning-only, not blocking -def _warm_up_with_gpu_check(monkeypatch, gpu_check_result) -> dict: - from jasna import os_utils +def _warm_up_with_gpu_check(monkeypatch, gpu_check_result, *, nvidia=True, major=8) -> tuple[dict, tuple]: + from jasna import accelerator, os_utils calls = {} - fake_torch = types.SimpleNamespace(zeros=lambda *a, **k: calls.setdefault("device", k.get("device"))) + fake_torch = types.SimpleNamespace( + zeros=lambda *a, **k: calls.setdefault("device", k.get("device")), + cuda=types.SimpleNamespace(get_device_capability=lambda _index: (major, 0)), + ) monkeypatch.setitem(sys.modules, "torch", fake_torch) monkeypatch.setattr(os_utils, "check_supported_gpu", lambda: gpu_check_result) + monkeypatch.setattr(accelerator, "is_nvidia_device", lambda: nvidia) from jasna.gui.app import _warm_up_cuda - _warm_up_cuda() - return calls + return calls, _warm_up_cuda() def test_warm_up_cuda_inits_context_when_gpu_supported(monkeypatch): - calls = _warm_up_with_gpu_check(monkeypatch, (True, "RTX 4090")) + calls, support = _warm_up_with_gpu_check(monkeypatch, (True, "RTX 4090")) assert calls["device"] == "cuda" + assert support == (True, False) + + +@pytest.mark.parametrize(("nvidia", "major", "support"), [(True, 12, (True, True)), (False, 12, (False, False))]) +def test_warm_up_cuda_reports_ltx_gpu_support(monkeypatch, nvidia, major, support): + assert _warm_up_with_gpu_check(monkeypatch, (True, "GPU"), nvidia=nvidia, major=major)[1] == support @pytest.mark.parametrize( @@ -61,5 +70,6 @@ def test_warm_up_cuda_inits_context_when_gpu_supported(monkeypatch): [(False, "no_cuda"), (False, ("arch_unsupported", "gfx1103"))], ) def test_warm_up_cuda_launches_no_kernel_on_unsupported_gpu(monkeypatch, gpu_check_result): - calls = _warm_up_with_gpu_check(monkeypatch, gpu_check_result) + calls, support = _warm_up_with_gpu_check(monkeypatch, gpu_check_result) assert calls == {} + assert support == (False, False) diff --git a/tests/test_gui_tvai_validation.py b/tests/test_gui_tvai_validation.py index 56fb2032..219a5adb 100644 --- a/tests/test_gui_tvai_validation.py +++ b/tests/test_gui_tvai_validation.py @@ -9,12 +9,12 @@ def test_validate_gui_start_non_tvai_returns_empty() -> None: settings = AppSettings(secondary_restoration="none") - assert validate_gui_start(settings) == [] + assert validate_gui_start(settings, []) == [] def test_validate_gui_start_custom_post_export_requires_command() -> None: settings = AppSettings(secondary_restoration="none", post_export_action="command", post_export_command="") - assert t("error_post_export_command_required") in validate_gui_start(settings) + assert t("error_post_export_command_required") in validate_gui_start(settings, []) def test_validate_gui_start_tvai_missing_env_vars(tmp_path: Path, monkeypatch) -> None: @@ -25,7 +25,7 @@ def test_validate_gui_start_tvai_missing_env_vars(tmp_path: Path, monkeypatch) - ffmpeg.write_bytes(b"") settings = AppSettings(secondary_restoration="tvai", tvai_ffmpeg_path=str(ffmpeg)) - errors = validate_gui_start(settings) + errors = validate_gui_start(settings, []) assert t("error_tvai_data_dir_not_set") in errors assert t("error_tvai_model_dir_not_set") in errors @@ -38,7 +38,7 @@ def test_validate_gui_start_tvai_env_dirs_must_exist(tmp_path: Path, monkeypatch ffmpeg.write_bytes(b"") settings = AppSettings(secondary_restoration="tvai", tvai_ffmpeg_path=str(ffmpeg)) - errors = validate_gui_start(settings) + errors = validate_gui_start(settings, []) assert len(errors) == 2 assert str(tmp_path / "missing_data") in errors[0] assert str(tmp_path / "missing_model") in errors[1] @@ -53,7 +53,7 @@ def test_validate_gui_start_tvai_ffmpeg_path_must_exist(tmp_path: Path, monkeypa monkeypatch.setenv("TVAI_MODEL_DIR", str(model_dir)) settings = AppSettings(secondary_restoration="tvai", tvai_ffmpeg_path=str(tmp_path / "missing_ffmpeg.exe")) - errors = validate_gui_start(settings) + errors = validate_gui_start(settings, []) assert len(errors) == 1 assert str(tmp_path / "missing_ffmpeg.exe") in errors[0] @@ -70,4 +70,17 @@ def test_validate_gui_start_tvai_ok(tmp_path: Path, monkeypatch) -> None: ffmpeg.write_bytes(b"") settings = AppSettings(secondary_restoration="tvai", tvai_ffmpeg_path=str(ffmpeg)) - assert validate_gui_start(settings) == [] + assert validate_gui_start(settings, []) == [] + + +def test_validate_gui_start_blocks_ltx_with_segmented_jobs() -> None: + from jasna.gui.models import JobItem + from jasna.segments import SegmentRange + + whole = JobItem(path=Path("whole.mp4")) + cut = JobItem(path=Path("cut.mp4"), segments=(SegmentRange(1.0, 2.0),)) + ltx = AppSettings(restoration_model="ltx", secondary_restoration="tvai") + + assert validate_gui_start(ltx, [whole]) == [] + assert validate_gui_start(ltx, [whole, cut]) == [t("error_ltx_segments", files="cut.mp4")] + assert validate_gui_start(AppSettings(), [cut]) == [] diff --git a/tests/test_gui_video_player.py b/tests/test_gui_video_player.py index 72c7b333..41e763ac 100644 --- a/tests/test_gui_video_player.py +++ b/tests/test_gui_video_player.py @@ -780,3 +780,14 @@ def test_player_setting_change_reloads_at_current_position() -> None: assert not dialog._playing assert dialog._buffering assert not dialog._eof + + +def test_player_always_plays_with_the_standard_model() -> None: + dialog = SimpleNamespace( + _base_settings=AppSettings(restoration_model="ltx"), + _model=SimpleNamespace(get=lambda: "rfdetr-v6"), + _threshold=SimpleNamespace(get=lambda: 0.4), + _secondary=SimpleNamespace(get_value=lambda: "none"), + ) + + assert VideoPlayerDialog._playback_settings(dialog).restoration_model == "basicvsrpp" diff --git a/tests/test_ltx_model_files.py b/tests/test_ltx_model_files.py index cf8cdabf..bfe46d06 100644 --- a/tests/test_ltx_model_files.py +++ b/tests/test_ltx_model_files.py @@ -36,3 +36,13 @@ def test_fast_bundle_uses_the_fast_transformer(tmp_path): for name in ("transformer-fast", "vae", "vae-decoder"): save_file({"weight": torch.zeros(1)}, str(tmp_path / f"{name}.safetensors")) assert LtxModelFiles.from_dir(tmp_path, fast=True).transformer == tmp_path / "transformer-fast.safetensors" + + +def test_bundle_present_needs_every_quality_file(tmp_path): + from jasna.ltx.model_files import bundle_present + + for name in ("transformer-fast", "vae", "vae-decoder"): + (tmp_path / f"{name}.safetensors").touch() + assert not bundle_present(tmp_path) + (tmp_path / "transformer.safetensors.enc").touch() + assert bundle_present(tmp_path) diff --git a/tests/test_restoration_preview.py b/tests/test_restoration_preview.py index 2b00edd7..7f735183 100644 --- a/tests/test_restoration_preview.py +++ b/tests/test_restoration_preview.py @@ -475,3 +475,36 @@ def test_app_prepares_entire_queue_before_starting_processor() -> None: app._on_start() assert calls == ["reset", "running", "start"] + + +def test_restoration_worker_refuses_ltx_without_loading_models(monkeypatch) -> None: + build = MagicMock() + monkeypatch.setattr(restoration_preview, "build_video_session", build) + worker = RestorationPreviewWorker("unused.mp4", _metadata()) + worker.start() + generation = worker.request(1.0, AppSettings(restoration_model="ltx"), projection="auto") + + event = worker.events.get(timeout=2) + + assert event == restoration_preview.RestorationFailed( + restoration_preview.t("segments_restore_ltx_unavailable"), generation + ) + build.assert_not_called() + worker.close() + worker.join(timeout=2) + + +def test_segment_editor_explains_that_ltx_has_no_preview() -> None: + editor = SegmentEditor.__new__(SegmentEditor) + editor._state = SimpleNamespace() + editor._restore_active = False + editor._is_gpu_busy = lambda: False + editor._get_settings = lambda: AppSettings(restoration_model="ltx") + editor._restore_toggle = MagicMock() + editor._show_preview_message = MagicMock() + + editor._toggle_restoration_preview() + + editor._restore_toggle.deselect.assert_called_once_with() + editor._show_preview_message.assert_called_once() + assert editor._restore_active is False diff --git a/tests/test_video_session.py b/tests/test_video_session.py index 0e262d1a..5f09da74 100644 --- a/tests/test_video_session.py +++ b/tests/test_video_session.py @@ -110,3 +110,46 @@ def test_video_session_close_closes_restorers() -> None: session.restoration_pipeline.restorer.close.assert_called_once_with() session.restoration_pipeline.secondary_restorer.close.assert_called_once_with() + + +def _config(settings: AppSettings): + with ( + patch("jasna.engine_paths.model_weights_dir", return_value=Path("weights")), + patch("jasna.mosaic.detection_registry.coerce_detection_model_name", side_effect=lambda name: name), + patch("jasna.mosaic.detection_registry.require_detection_model_weights", return_value=Path("det.engine")), + ): + return video_session_config(settings, codec="hevc", encoder_settings={}) + + +def test_video_session_config_maps_ltx_settings() -> None: + config = _config( + replace(AppSettings(), restoration_model="ltx", ltx_seed=7, ltx_fast=True, ltx_large_canvas=True) + ) + + assert config.restoration_model_name == "ltx" + assert config.restoration_model_path == Path("weights") / "ltx-restore" + assert (config.ltx_seed, config.ltx_fast, config.ltx_large_canvas) == (7, True, True) + + +def test_video_session_config_ignores_standard_only_extras_for_ltx() -> None: + settings = replace( + AppSettings(), + restoration_model="ltx", + secondary_restoration="tvai", + tvai_denoise=True, + denoise_strength="high", + ) + + config = _config(settings) + + assert (config.secondary_restoration, config.tvai_denoise, config.denoise_strength) == ("none", False, "none") + standard = _config(replace(settings, restoration_model="basicvsrpp")) + assert (standard.secondary_restoration, standard.tvai_denoise, standard.denoise_strength) == ("tvai", True, "high") + + +def test_video_session_key_tracks_model_and_ltx_variant() -> None: + ltx = replace(AppSettings(), restoration_model="ltx") + + assert video_session_key(ltx) != video_session_key(AppSettings()) + assert video_session_key(replace(ltx, ltx_fast=True)) != video_session_key(ltx) + assert video_session_key(replace(ltx, ltx_seed=1, ltx_large_canvas=True)) == video_session_key(ltx) From 9b1058e752da06184d374e054e4426595056c183 Mon Sep 17 00:00:00 2001 From: Kruk2 Date: Tue, 29 Sep 2026 16:02:02 +0200 Subject: [PATCH 26/48] fix(gui): title start errors "Can't start" instead of a TVAI error --- jasna/gui/app.py | 2 +- jasna/gui/locales/en.py | 2 +- jasna/gui/locales/ja.py | 2 +- jasna/gui/locales/ko.py | 2 +- jasna/gui/locales/th.py | 2 +- jasna/gui/locales/zh.py | 2 +- 6 files changed, 6 insertions(+), 6 deletions(-) diff --git a/jasna/gui/app.py b/jasna/gui/app.py index 7fafe61d..e2f6a397 100644 --- a/jasna/gui/app.py +++ b/jasna/gui/app.py @@ -600,7 +600,7 @@ def _on_start(self): msg = t("error_cannot_start") + "\n\n" + "\n".join(f"- {e}" for e in errors) self._log_panel.error(msg) - messagebox.showerror(t("error_invalid_tvai"), msg) + messagebox.showerror(t("error_start_title"), msg) return try: diff --git a/jasna/gui/locales/en.py b/jasna/gui/locales/en.py index cc8cbe55..bfd795b3 100644 --- a/jasna/gui/locales/en.py +++ b/jasna/gui/locales/en.py @@ -463,7 +463,7 @@ # App messages "error_cannot_start": "Cannot start processing:", - "error_invalid_tvai": "Invalid TVAI configuration", + "error_start_title": "Can't start", "error_post_export_command_required": "Custom post-export command is required", "error_ltx_segments": "LTX processes whole videos. Clear the segments or pick Standard: {files}", diff --git a/jasna/gui/locales/ja.py b/jasna/gui/locales/ja.py index 384443e3..327715b3 100644 --- a/jasna/gui/locales/ja.py +++ b/jasna/gui/locales/ja.py @@ -463,7 +463,7 @@ # App messages "error_cannot_start": "処理を開始できません:", - "error_invalid_tvai": "無効な TVAI 設定", + "error_start_title": "開始できません", "error_post_export_command_required": "エクスポート後のカスタムコマンドが必要です", "error_ltx_segments": "LTX は動画全体を処理します。セグメントを消去するか「標準」を選んでください：{files}", diff --git a/jasna/gui/locales/ko.py b/jasna/gui/locales/ko.py index aab5e0d6..3a796973 100644 --- a/jasna/gui/locales/ko.py +++ b/jasna/gui/locales/ko.py @@ -443,7 +443,7 @@ # App messages "error_cannot_start": "처리를 시작할 수 없습니다:", - "error_invalid_tvai": "잘못된 TVAI 구성", + "error_start_title": "시작할 수 없음", "error_post_export_command_required": "내보내기 후 사용자 정의 명령이 필요합니다", "error_ltx_segments": "LTX는 영상 전체를 처리합니다. 구간을 지우거나 '표준'을 선택하세요: {files}", diff --git a/jasna/gui/locales/th.py b/jasna/gui/locales/th.py index f3dddd01..65454afd 100644 --- a/jasna/gui/locales/th.py +++ b/jasna/gui/locales/th.py @@ -445,7 +445,7 @@ # App messages "error_cannot_start": "ไม่สามารถเริ่มประมวลผล:", - "error_invalid_tvai": "การกำหนดค่า TVAI ไม่ถูกต้อง", + "error_start_title": "เริ่มไม่ได้", "error_post_export_command_required": "ต้องระบุคำสั่งกำหนดเองหลังส่งออก", "error_ltx_segments": "LTX ประมวลผลทั้งวิดีโอ ล้างช่วงที่เลือกหรือเลือก “มาตรฐาน”: {files}", diff --git a/jasna/gui/locales/zh.py b/jasna/gui/locales/zh.py index 4dcad8da..319d4f68 100644 --- a/jasna/gui/locales/zh.py +++ b/jasna/gui/locales/zh.py @@ -463,7 +463,7 @@ # App messages "error_cannot_start": "无法开始处理：", - "error_invalid_tvai": "无效的 TVAI 配置", + "error_start_title": "无法开始", "error_post_export_command_required": "需要填写导出后的自定义命令", "error_ltx_segments": "LTX 会处理整个视频。请清除片段或选择“标准”：{files}", From cee90b0e1ac2b568f216af9dc46416d6cedfb9a0 Mon Sep 17 00:00:00 2001 From: Kruk2 Date: Tue, 29 Sep 2026 16:20:00 +0200 Subject: [PATCH 27/48] feat(segments): restore each segment with its own model and LTX seed Segments without one follow the job's model, and overlapping ranges are split by the later one. Each LTX segment is planned from its own frames and all LTX spans run in one batch. --- jasna/cli_help.py | 2 +- jasna/ltx/restore.py | 530 +++++++++++++++++++++++--------- jasna/main.py | 4 +- jasna/media/splice.py | 31 +- jasna/pipeline.py | 171 +++++++++-- jasna/segments.py | 63 +++- jasna/session_factory.py | 88 ++++-- tests/test_ltx_model_files.py | 2 +- tests/test_ltx_restore.py | 293 +++++++++++++++--- tests/test_main_validation.py | 22 ++ tests/test_pipeline_segments.py | 123 +++++++- tests/test_segments.py | 40 +++ tests/test_session_factory.py | 31 +- tests/test_splice.py | 26 +- 14 files changed, 1151 insertions(+), 275 deletions(-) diff --git a/jasna/cli_help.py b/jasna/cli_help.py index 1d7d6527..a170a301 100644 --- a/jasna/cli_help.py +++ b/jasna/cli_help.py @@ -9,7 +9,7 @@ CLI_HELP: dict[str, str] = { "restoration_model_name": ( "Restoration model for video input: basicvsrpp (fast) or ltx (diffusion model, much " - "slower, best quality; no streaming, segments, secondary restoration or denoise). " + "slower, best quality; no streaming, secondary restoration or denoise). " "(default: %(default)s)" ), "restoration_model_path": ( diff --git a/jasna/ltx/restore.py b/jasna/ltx/restore.py index 779af6fc..9bdcacbb 100644 --- a/jasna/ltx/restore.py +++ b/jasna/ltx/restore.py @@ -1,13 +1,16 @@ -"""Restore a whole video with the LTX model in four passes, one model on the GPU at a time. +"""Restore video with the LTX model in four passes, one model on the GPU at a time. 1. scan: decode, detect mosaics and scene cuts, plan tracks, cameras and windows; 2. encode: crop every window's 121 frames onto its canvas and VAE-encode them; 3. denoise: run each track's windows through the transformer in lockstep; 4. compose: decode each window's latent when the video reaches it, composite, write. -The camera and fusion need a whole track before its first window can be denoised, so -the video is decoded three times (scan, encode, compose). Per-window latents (1-2 MB) -are kept in a temporary directory between passes, so memory does not grow with length. +The unit of work is a segment: the frames of one pts range, planned on their own with +frame indices from the segment's first frame, so a segment restores the same way whatever +surrounds it. A whole video is one segment. Spans are the decoded stretches holding the +segments; every pass decodes each span again (the camera and fusion need a whole track +before its first window can be denoised). Per-window latents (1-2 MB) are kept in a +temporary directory between passes, so memory does not grow with length. """ from __future__ import annotations @@ -15,17 +18,18 @@ import logging import tempfile import threading -from collections.abc import Callable, Iterator +from collections.abc import Callable, Iterator, Sequence from dataclasses import dataclass from pathlib import Path +from typing import Protocol import torch from tqdm import tqdm -from jasna.ltx.camera import WINDOW_FRAMES +from jasna.ltx.camera import PLAN_CANVAS, WINDOW_FRAMES from jasna.ltx.compose import Candidate, composite_frame, crop_to_canvas from jasna.ltx.model_files import LtxModelFiles -from jasna.ltx.plan import Region, TrackPlan, Window, crossfade_weights, feather_pixels, plan_video +from jasna.ltx.plan import LARGE_CANVAS, Region, TrackPlan, Window, crossfade_weights, feather_pixels, plan_video from jasna.ltx.regions import regions_per_frame from jasna.ltx.sampler import LtxTransformer from jasna.models.ltx_vae import decode_latent, load_video_decoder, load_video_encoder @@ -35,13 +39,46 @@ FrameBatches = Iterator[tuple[torch.Tensor, list[int]]] FrameSource = Callable[[], FrameBatches] +DecodeBudget = Callable[[int], int] LARGE_CANVAS_MIN_FREE_BYTES = 10 << 30 +LARGE_CANVAS_MIN_TOTAL_BYTES = 15 << 30 +SEGMENT_DECODE_BUDGET_BYTES = {PLAN_CANVAS: 6 << 30, LARGE_CANVAS: 8 << 30} +WHOLE_VIDEO_PTS = (-(1 << 63), (1 << 63) - 1) class Cancelled(Exception): pass +class FrameWriter(Protocol): + def write(self, frame: torch.Tensor, pts: int, *, apply_lut: bool = True) -> None: ... + + def close(self) -> None: ... + + +@dataclass(frozen=True) +class LtxSegment: + """The frames with ``start_pts <= pts < end_pts``; window ``i`` draws noise from ``seed + i``.""" + + start_pts: int + end_pts: int + seed: int + + def contains(self, pts: int) -> bool: + return self.start_pts <= pts < self.end_pts + + +@dataclass(frozen=True) +class LtxSpan: + """``frames`` decodes a stretch of video holding ``segments``; every frame goes to the + writer ``open_writer`` makes when the span is composed, a frame outside the segments + unchanged and without the LUT.""" + + frames: FrameSource + segments: tuple[LtxSegment, ...] + open_writer: Callable[[], FrameWriter] + + @dataclass(frozen=True) class Progress: """One console bar per pass; ``frames`` sizes the frame-based passes.""" @@ -53,7 +90,9 @@ def bar(self, name: str, total: int, unit: str = "frame") -> tqdm: return tqdm(total=total, desc=f"LTX {name}", unit=unit, dynamic_ncols=True, disable=self.disable) -class _LatentStore: +class LatentStore: + """Window latents of one segment on disk, keyed by window index.""" + def __init__(self, directory: Path, kind: str) -> None: self._directory = directory self._kind = kind @@ -76,33 +115,84 @@ def _check(cancel: threading.Event) -> None: raise Cancelled() -def scan( +def segment_batches( + batches: FrameBatches, segments: Sequence[LtxSegment], batch_size: int +) -> Iterator[tuple[int | None, torch.Tensor, list[int]]]: + """``(segment index, frames, pts)`` in decode order. A segment's frames come in batches + of ``batch_size`` counted from its first frame (the last may be short), so detection + sees the same batches wherever decoding started; frames outside every segment come as + decoded with index None.""" + owner: int | None = None + pending: list[tuple[torch.Tensor, list[int]]] = [] + + def full_batches(flush: bool) -> Iterator[tuple[int | None, torch.Tensor, list[int]]]: + count = sum(len(pts) for _, pts in pending) + if count == 0 or (count < batch_size and not flush): + return + if len(pending) == 1: + frames, pts = pending[0] + else: + frames = torch.cat([chunk for chunk, _ in pending]) + pts = [p for _, chunk_pts in pending for p in chunk_pts] + pending.clear() + whole = len(pts) if flush else len(pts) - len(pts) % batch_size + for start in range(0, whole, batch_size): + yield owner, frames[start : start + batch_size], pts[start : start + batch_size] + if whole < len(pts): + pending.append((frames[whole:], pts[whole:])) + + for frames, pts in batches: + owners = [next((i for i, s in enumerate(segments) if s.contains(int(p))), None) for p in pts] + start = 0 + while start < len(pts): + stop = start + 1 + while stop < len(pts) and owners[stop] == owners[start]: + stop += 1 + if owners[start] != owner: + yield from full_batches(flush=True) + owner = owners[start] + if owner is None: + yield None, frames[start:stop], list(pts[start:stop]) + else: + pending.append((frames[start:stop], list(pts[start:stop]))) + yield from full_batches(flush=False) + start = stop + yield from full_batches(flush=True) + + +def scan_span( frames: FrameSource, + segments: Sequence[LtxSegment], detector, *, + batch_size: int, frame_h: int, frame_w: int, large_canvas: bool, - progress: Progress, + bar: tqdm, cancel: threading.Event, -) -> list[TrackPlan]: - regions: list[list[Region]] = [] - cuts: set[int] = set() - scene = SceneCutDetector() - with progress.bar("1/4 scan", progress.frames) as bar: - for batch, _pts in frames(): - _check(cancel) - cuts.update(len(regions) + offset for offset in scene.find_cuts(batch)) - regions.extend(regions_per_frame(detector(batch, target_hw=(frame_h, frame_w)), frame_h, frame_w)) - bar.update(len(batch)) - plans = plan_video(regions, cuts, frame_w=frame_w, frame_h=frame_h, large_canvas=large_canvas) - logger.info( - "LTX plan: %d frames, %d scene cuts, %d tracks, %d windows", - len(regions), - len(cuts), - len(plans), - sum(len(plan.windows) for plan in plans), - ) +) -> list[list[TrackPlan]]: + """The track plans of each segment, frame indices counted from the segment's first frame.""" + regions: list[list[list[Region]]] = [[] for _ in segments] + cuts: list[set[int]] = [set() for _ in segments] + scenes = [SceneCutDetector() for _ in segments] + for owner, batch, _pts in segment_batches(frames(), segments, batch_size): + _check(cancel) + if owner is not None: + cuts[owner].update(len(regions[owner]) + offset for offset in scenes[owner].find_cuts(batch)) + regions[owner].extend(regions_per_frame(detector(batch, target_hw=(frame_h, frame_w)), frame_h, frame_w)) + bar.update(len(batch)) + plans = [] + for segment_regions, segment_cuts in zip(regions, cuts): + segment_plans = plan_video(segment_regions, segment_cuts, frame_w=frame_w, frame_h=frame_h, large_canvas=large_canvas) + logger.info( + "LTX plan: %d frames, %d scene cuts, %d tracks, %d windows", + len(segment_regions), + len(segment_cuts), + len(segment_plans), + sum(len(plan.windows) for plan in segment_plans), + ) + plans.append(segment_plans) return plans @@ -111,68 +201,82 @@ def _vae_pixels(canvases: list[torch.Tensor], device: torch.device) -> torch.Ten return stacked.to(device).float().div_(127.5).sub_(1.0).to(torch.bfloat16).unsqueeze(0) -def encode_references( +class _SegmentEncoder: + def __init__(self, plans: list[TrackPlan], encoder, store: LatentStore, device: torch.device) -> None: + self._pending = iter(sorted((w for plan in plans for w in plan.windows), key=lambda w: w.start)) + self._upcoming = next(self._pending, None) + self._active: dict[int, tuple[Window, list[torch.Tensor]]] = {} + self._encoder = encoder + self._store = store + self._device = device + self._frame_idx = 0 + + def frame(self, frame: torch.Tensor) -> None: + while self._upcoming is not None and self._upcoming.start == self._frame_idx: + self._active[self._upcoming.index] = (self._upcoming, []) + self._upcoming = next(self._pending, None) + for index, (window, canvases) in list(self._active.items()): + canvases.append(crop_to_canvas(frame, window.crops[self._frame_idx - window.start]).cpu()) + if len(canvases) == window.real_frames: + with torch.inference_mode(): + self._store.put(index, self._encoder(_vae_pixels(canvases, self._device))) + del self._active[index] + self._frame_idx += 1 + + def finish(self) -> None: + if self._active or self._upcoming is not None: + raise RuntimeError("the video ended before every planned window was filled") + + +def encode_span( frames: FrameSource, - plans: list[TrackPlan], + segments: Sequence[LtxSegment], + plans: Sequence[list[TrackPlan]], encoder, - store: _LatentStore, + stores: Sequence[LatentStore], *, + batch_size: int, device: torch.device, - progress: Progress, + bar: tqdm, cancel: threading.Event, ) -> None: - windows = sorted((w for plan in plans for w in plan.windows), key=lambda w: w.start) - pending = iter(windows) - upcoming = next(pending, None) - active: dict[int, tuple[Window, list[torch.Tensor]]] = {} - frame_idx = 0 - bar = progress.bar("2/4 encode", progress.frames) - for batch, _pts in frames(): + """VAE-encode every planned window's reference crops into its segment's store.""" + encoders = [_SegmentEncoder(p, encoder, store, device) for p, store in zip(plans, stores)] + for owner, batch, _pts in segment_batches(frames(), segments, batch_size): _check(cancel) bar.update(len(batch)) - for frame in batch: - while upcoming is not None and upcoming.start == frame_idx: - active[upcoming.index] = (upcoming, []) - upcoming = next(pending, None) - for index, (window, canvases) in list(active.items()): - canvases.append(crop_to_canvas(frame, window.crops[frame_idx - window.start]).cpu()) - if len(canvases) == window.real_frames: - with torch.inference_mode(): - store.put(index, encoder(_vae_pixels(canvases, device))) - del active[index] - frame_idx += 1 - bar.close() - if active or upcoming is not None: - raise RuntimeError("the video ended before every planned window was filled") - - -def denoise( - plans: list[TrackPlan], - transformer: LtxTransformer, - references: _LatentStore, - finals: _LatentStore, - *, - seed: int, - progress: Progress, - cancel: threading.Event, -) -> None: - steps = len(transformer.conditions) - bar = progress.bar("3/4 denoise", steps * sum(len(plan.windows) for plan in plans), unit="step") - for plan in plans: - _check(cancel) - latents = transformer.denoise_chain( - [references.get(w.index) for w in plan.windows], [seed + w.index for w in plan.windows], bar.update - ) - for window, latent in zip(plan.windows, latents): - finals.put(window.index, latent) - references.delete(window.index) - bar.close() + if owner is not None: + for frame in batch: + encoders[owner].frame(frame) + for segment_encoder in encoders: + segment_encoder.finish() + + +def denoise_track( + plan: TrackPlan, transformer: LtxTransformer, references: LatentStore, *, seed: int, advance: Callable[[int], None] +) -> list[torch.Tensor]: + """The final latents of one track's windows, in window order.""" + return transformer.denoise_chain( + [references.get(w.index) for w in plan.windows], [seed + w.index for w in plan.windows], advance + ) + + +def free_vram_budget(device: torch.device) -> DecodeBudget: + return lambda _canvas: torch.cuda.mem_get_info(device)[0] + + +def segment_decode_budget(canvas: int) -> int: + """A fixed decode budget, so a segment's tile layout (and pixels) never depend on the + VRAM free at that moment.""" + return SEGMENT_DECODE_BUDGET_BYTES[canvas] -def _decode_window(decoder, latent: torch.Tensor, seed: int, device: torch.device) -> torch.Tensor: +def _decode_window( + decoder, latent: torch.Tensor, seed: int, canvas: int, budget: DecodeBudget, device: torch.device +) -> torch.Tensor: """Decoded window as uint8 ``[F, 3, S, S]`` on ``device``.""" torch.cuda.empty_cache() - free, _ = torch.cuda.mem_get_info(device) + free = budget(canvas) generator = torch.Generator(device=device).manual_seed(seed) with torch.inference_mode(): video = decode_latent(decoder, latent.to(device), free_bytes=free, generator=generator) @@ -180,59 +284,100 @@ def _decode_window(decoder, latent: torch.Tensor, seed: int, device: torch.devic return video.mul_(255.0).round_().to(torch.uint8).permute(1, 0, 2, 3).contiguous() -def compose( - frames: FrameSource, - plans: list[TrackPlan], +class _SegmentComposer: + def __init__( + self, + plans: list[TrackPlan], + decoder, + finals: LatentStore, + *, + seed: int, + feather: int, + budget: DecodeBudget, + device: torch.device, + ) -> None: + self._weights: dict[int, torch.Tensor] = {} + self._tracks: dict[int, TrackPlan] = {} + for plan in plans: + for window, weight in zip(plan.windows, crossfade_weights(plan.windows)): + self._weights[window.index] = weight + self._tracks[window.index] = plan + self._pending = iter(sorted((w for plan in plans for w in plan.windows), key=lambda w: w.start)) + self._upcoming = next(self._pending, None) + self._covering: list[Window] = [] + self._decoded: dict[int, torch.Tensor] = {} + self._decoder = decoder + self._finals = finals + self._seed = seed + self._feather = feather + self._budget = budget + self._device = device + self._frame_idx = 0 + + def frame(self, frame: torch.Tensor) -> torch.Tensor: + frame_idx = self._frame_idx + while self._upcoming is not None and self._upcoming.start == frame_idx: + self._covering.append(self._upcoming) + self._upcoming = next(self._pending, None) + for window in [w for w in self._covering if w.stop <= frame_idx]: + self._covering.remove(window) + self._decoded.pop(window.index, None) + candidates = [] + for window in self._covering: + if window.index not in self._decoded: + self._decoded[window.index] = _decode_window( + self._decoder, + self._finals.get(window.index), + self._seed + window.index, + window.canvas, + self._budget, + self._device, + ) + local = frame_idx - window.start + track = self._tracks[window.index] + candidates.append( + Candidate( + canvas=self._decoded[window.index][local], + crop=window.crops[local], + weight=float(self._weights[window.index][local]), + polygons=track.polygons[frame_idx - track.start], + ) + ) + self._frame_idx += 1 + return composite_frame(frame, candidates, feather=self._feather) + + +def compose_span( + span: LtxSpan, + plans: Sequence[list[TrackPlan]], decoder, - finals: _LatentStore, - write: Callable[[torch.Tensor, int], None], + finals: Sequence[LatentStore], *, frame_h: int, - seed: int, + batch_size: int, + budget: DecodeBudget, device: torch.device, - progress: Progress, + bar: tqdm, cancel: threading.Event, ) -> None: + """Write every frame of ``span``: segment frames composited, the rest unchanged.""" feather = feather_pixels(frame_h) - weights: dict[int, torch.Tensor] = {} - tracks: dict[int, TrackPlan] = {} - for plan in plans: - for window, weight in zip(plan.windows, crossfade_weights(plan.windows)): - weights[window.index] = weight - tracks[window.index] = plan - pending = iter(sorted((w for plan in plans for w in plan.windows), key=lambda w: w.start)) - upcoming = next(pending, None) - covering: list[Window] = [] - decoded: dict[int, torch.Tensor] = {} - frame_idx = 0 - bar = progress.bar("4/4 compose", progress.frames) - for batch, pts in frames(): - _check(cancel) - bar.update(len(batch)) - for frame, frame_pts in zip(batch, pts): - while upcoming is not None and upcoming.start == frame_idx: - covering.append(upcoming) - upcoming = next(pending, None) - for window in [w for w in covering if w.stop <= frame_idx]: - covering.remove(window) - decoded.pop(window.index, None) - candidates = [] - for window in covering: - if window.index not in decoded: - decoded[window.index] = _decode_window(decoder, finals.get(window.index), seed + window.index, device) - local = frame_idx - window.start - track = tracks[window.index] - candidates.append( - Candidate( - canvas=decoded[window.index][local], - crop=window.crops[local], - weight=float(weights[window.index][local]), - polygons=track.polygons[frame_idx - track.start], - ) - ) - write(composite_frame(frame, candidates, feather=feather), frame_pts) - frame_idx += 1 - bar.close() + composers = [ + _SegmentComposer(p, decoder, store, seed=segment.seed, feather=feather, budget=budget, device=device) + for p, store, segment in zip(plans, finals, span.segments) + ] + writer = span.open_writer() + try: + for owner, batch, pts in segment_batches(span.frames(), span.segments, batch_size): + _check(cancel) + bar.update(len(batch)) + for frame, frame_pts in zip(batch, pts): + if owner is None: + writer.write(frame, frame_pts, apply_lut=False) + else: + writer.write(composers[owner].frame(frame), frame_pts) + finally: + writer.close() def large_canvas_fits(requested: bool, free_bytes: int) -> bool: @@ -244,51 +389,134 @@ def large_canvas_fits(requested: bool, free_bytes: int) -> bool: return requested -def restore_video( - frames: FrameSource, - write: Callable[[torch.Tensor, int], None], +def segment_large_canvas(requested: bool, total_bytes: int) -> bool: + """The 768 px decision for segments, from the GPU's size so it is the same every run.""" + if requested and total_bytes < LARGE_CANVAS_MIN_TOTAL_BYTES: + logger.warning("LTX: %.1f GiB GPU, restoring large mosaics at 512 px", total_bytes / 2**30) + return False + return requested + + +def restore_spans( + spans: Sequence[LtxSpan], *, detector, files: LtxModelFiles, frame_h: int, frame_w: int, + batch_size: int, large_canvas: bool, - seed: int, + budget: DecodeBudget, device: torch.device, work_dir: Path, progress: Progress, cancel: threading.Event, ) -> None: - """Restore every mosaic of the video ``frames`` yields and ``write`` each output frame - with its pts. ``detector`` is only used by the scan pass; ``large_canvas`` allows 768 px - windows for large mosaics; window ``i`` draws its noise from ``seed + i``.""" - large_canvas = large_canvas_fits(large_canvas, torch.cuda.mem_get_info(device)[0]) - plans = scan(frames, detector, frame_h=frame_h, frame_w=frame_w, large_canvas=large_canvas, progress=progress, cancel=cancel) + """Restore every segment of ``spans`` and write each span's frames to a writer of its own + (closed once the span is written). ``large_canvas`` allows 768 px windows for large mosaics; + ``budget`` gives the decode memory budget per canvas size.""" + with progress.bar("1/4 scan", progress.frames) as bar: + plans = [ + scan_span( + span.frames, + span.segments, + detector, + batch_size=batch_size, + frame_h=frame_h, + frame_w=frame_w, + large_canvas=large_canvas, + bar=bar, + cancel=cancel, + ) + for span in spans + ] + windows = sum(len(plan.windows) for span_plans in plans for segment_plans in span_plans for plan in segment_plans) with tempfile.TemporaryDirectory(dir=work_dir, prefix=".ltx-") as temp: - references = _LatentStore(Path(temp), "reference") - finals = _LatentStore(Path(temp), "final") - if plans: + stores = [ + [ + (LatentStore(Path(temp), f"{s}-{g}.reference"), LatentStore(Path(temp), f"{s}-{g}.final")) + for g in range(len(span.segments)) + ] + for s, span in enumerate(spans) + ] + if windows: encoder = load_video_encoder(files.vae, device) - encode_references(frames, plans, encoder, references, device=device, progress=progress, cancel=cancel) + with progress.bar("2/4 encode", progress.frames) as bar: + for span, span_plans, span_stores in zip(spans, plans, stores): + encode_span( + span.frames, + span.segments, + span_plans, + encoder, + [references for references, _ in span_stores], + batch_size=batch_size, + device=device, + bar=bar, + cancel=cancel, + ) del encoder torch.cuda.empty_cache() transformer = LtxTransformer(files.transformer, device) try: - denoise(plans, transformer, references, finals, seed=seed, progress=progress, cancel=cancel) + with progress.bar("3/4 denoise", len(transformer.conditions) * windows, unit="step") as bar: + for span, span_plans, span_stores in zip(spans, plans, stores): + for segment, segment_plans, (references, finals) in zip(span.segments, span_plans, span_stores): + for plan in segment_plans: + _check(cancel) + latents = denoise_track(plan, transformer, references, seed=segment.seed, advance=bar.update) + for window, latent in zip(plan.windows, latents): + finals.put(window.index, latent) + references.delete(window.index) finally: transformer.close() del transformer torch.cuda.empty_cache() - decoder = load_video_decoder(files.vae, files.tuned_decoder, device) if plans else None - compose( - frames, - plans, - decoder, - finals, - write, - frame_h=frame_h, - seed=seed, - device=device, - progress=progress, - cancel=cancel, - ) + decoder = load_video_decoder(files.vae, files.tuned_decoder, device) if windows else None + with progress.bar("4/4 compose", progress.frames) as bar: + for span, span_plans, span_stores in zip(spans, plans, stores): + compose_span( + span, + span_plans, + decoder, + [finals for _, finals in span_stores], + frame_h=frame_h, + batch_size=batch_size, + budget=budget, + device=device, + bar=bar, + cancel=cancel, + ) + + +def restore_video( + frames: FrameSource, + writer: FrameWriter, + *, + detector, + files: LtxModelFiles, + frame_h: int, + frame_w: int, + batch_size: int, + large_canvas: bool, + seed: int, + device: torch.device, + work_dir: Path, + progress: Progress, + cancel: threading.Event, +) -> None: + """Restore every mosaic of the video ``frames`` yields as one segment; the 768 px + decision and the decode tiles follow the VRAM free when they are made.""" + restore_spans( + [LtxSpan(frames, (LtxSegment(*WHOLE_VIDEO_PTS, seed),), lambda: writer)], + detector=detector, + files=files, + frame_h=frame_h, + frame_w=frame_w, + batch_size=batch_size, + large_canvas=large_canvas_fits(large_canvas, torch.cuda.mem_get_info(device)[0]), + budget=free_vram_budget(device), + device=device, + work_dir=work_dir, + progress=progress, + cancel=cancel, + ) diff --git a/jasna/main.py b/jasna/main.py index 973ea46e..87bb75d6 100644 --- a/jasna/main.py +++ b/jasna/main.py @@ -862,8 +862,8 @@ def _run_post_export_action() -> None: ) if args.restoration_model_name == "ltx": - if is_streaming or segments_spec: - parser.error("--restoration-model-name ltx does not support --stream or --segments") + if is_streaming: + parser.error("--restoration-model-name ltx does not support --stream") if args.secondary_restoration != "none" or args.denoise != "none": parser.error("--restoration-model-name ltx does not support --secondary-restoration or --denoise") restoration_model_path = Path( diff --git a/jasna/media/splice.py b/jasna/media/splice.py index 77f51055..e715b779 100644 --- a/jasna/media/splice.py +++ b/jasna/media/splice.py @@ -80,6 +80,21 @@ class SplicePlan: def render_spans(self) -> tuple[SpliceSpan, ...]: return tuple(span for span in self.spans if span.is_render) + def render_span_segments(self) -> tuple[tuple[SegmentRange, ...], ...]: + """Per render span, the segment behind each of its effect ranges.""" + remaining = iter(self.segments) + return tuple(tuple(next(remaining) for _ in span.effect_ranges) for span in self.render_spans) + + +def require_one_model(segments: tuple[SegmentRange, ...]) -> None: + """One render span is restored by one model; segments still on the job's model match any.""" + if len({segment.restoration.model for segment in segments if segment.restoration is not None}) > 1: + raise SmartRenderCompatibilityError( + "Ranges restored by different models are too close together to split the video " + "between them; move them apart or use one model for both", + reason="mixed_models", + ) + def canonical_codec(name: str) -> str: value = str(name).lower() @@ -387,19 +402,21 @@ def build_splice_plan( render_end = index.pts[right_index] if right_index < len(index.pts) else index.end_pts if render_end <= render_start: render_end = index.end_pts - expanded.append((render_start, render_end, [(effect_start, effect_end)])) + expanded.append((render_start, render_end, [(effect_start, effect_end)], [segment])) - merged: list[tuple[int, int, list[tuple[int, int]]]] = [] - for start, end, effects in expanded: + merged: list[tuple[int, int, list[tuple[int, int]], list[SegmentRange]]] = [] + for start, end, effects, owners in expanded: if merged and start <= merged[-1][1]: - old_start, old_end, old_effects = merged[-1] - merged[-1] = (old_start, max(old_end, end), old_effects + effects) + old_start, old_end, old_effects, old_owners = merged[-1] + merged[-1] = (old_start, max(old_end, end), old_effects + effects, old_owners + owners) else: - merged.append((start, end, effects)) + merged.append((start, end, effects, owners)) + for *_, owners in merged: + require_one_model(tuple(owners)) spans: list[SpliceSpan] = [] cursor = index.start_pts - for start, end, effects in merged: + for start, end, effects, _owners in merged: if cursor < start: spans.append(SpliceSpan("copy", cursor, start)) spans.append(SpliceSpan("render", start, end, tuple(effects))) diff --git a/jasna/pipeline.py b/jasna/pipeline.py index 29dd09e5..82178567 100644 --- a/jasna/pipeline.py +++ b/jasna/pipeline.py @@ -18,20 +18,23 @@ from jasna.media.video_encoder import VideoEncoder from jasna.media.frame_rate import resolve_frame_rate_retarget from jasna.media.splice import ( + KeyframeIndex, SplicePlan, + SpliceSpan, build_splice_plan, concatenate_fragments, create_copy_fragment, mux_final_output, normalize_fragment, probe_keyframes, + require_one_model, resolve_smart_encoder_settings, validate_smart_render, ) from jasna.pipeline_threads import run_restoration_pass from jasna.progressbar import Progressbar from jasna.restorer.secondary_restorer import AsyncSecondaryRestorer -from jasna.segments import SegmentRange +from jasna.segments import SegmentRange, job_restoration, resolve_restorations from jasna.session_config import SessionConfig from jasna.vr180 import ( SbsDetectionAdapter, @@ -97,6 +100,7 @@ def __init__( self.detection_model = session.detection_model_for(config) self.restoration_pipeline = session.restoration_pipeline self.ltx_files = session.ltx_files + self.restoration_model_name = config.restoration_model_name self.ltx_large_canvas = config.ltx_large_canvas self.ltx_seed = config.ltx_seed self.disable_progress = config.disable_progress @@ -274,29 +278,44 @@ def _run_full( finally: progress.close(ensure_completed_bar=True) - def _run_ltx(self, metadata) -> None: - from jasna.ltx.restore import Cancelled, Progress, restore_video + def _ltx_frames(self, metadata, *, frame_stride: int, seek_ts: float | None, end_pts: int | None): from jasna.media.video_decoder import VideoReader - if self.vr_resolution.is_sbs: - raise ValueError("LTX restoration does not support VR180 side-by-side video") - frame_rate = self._resolve_frame_rate(metadata) - writer = _OfflineFrameWriter(self._video_encoder(metadata, frame_rate), [time.monotonic()]) - def frames(): with VideoReader( - str(self.input_video), self.batch_size, self.device, metadata, frame_stride=frame_rate.frame_stride + str(self.input_video), self.batch_size, self.device, metadata, frame_stride=frame_stride ) as reader: - yield from reader.frames(None) + for batch, pts in reader.frames(seek_ts): + if end_pts is None: + yield batch, pts + continue + keep = next((i for i, p in enumerate(pts) if int(p) >= end_pts), len(pts)) + if keep: + yield batch[:keep], pts[:keep] + if keep < len(pts): + return + + return frames + + def _require_ltx_video(self) -> None: + if self.vr_resolution.is_sbs: + raise ValueError("LTX restoration does not support VR180 side-by-side video") + def _run_ltx(self, metadata) -> None: + from jasna.ltx.restore import Cancelled, Progress, restore_video + + self._require_ltx_video() + frame_rate = self._resolve_frame_rate(metadata) + writer = _OfflineFrameWriter(self._video_encoder(metadata, frame_rate), [time.monotonic()]) try: restore_video( - frames, - writer.write, + self._ltx_frames(metadata, frame_stride=frame_rate.frame_stride, seek_ts=None, end_pts=None), + writer, detector=self.job_detection_model, files=self.ltx_files, frame_h=int(metadata.video_height), frame_w=int(metadata.video_width), + batch_size=self.batch_size, large_canvas=self.ltx_large_canvas, seed=self.ltx_seed, device=self.device, @@ -311,6 +330,61 @@ def frames(): finally: writer.close() + def _run_ltx_spans( + self, + metadata, + index: KeyframeIndex, + spans: list[tuple[SpliceSpan, tuple[SegmentRange, ...], Callable[[], _OfflineFrameWriter]]], + work_dir: Path, + ) -> None: + """Restore the LTX render spans in one batched run, so every model loads once.""" + from jasna.ltx.restore import ( + Cancelled, + LtxSegment, + LtxSpan, + Progress, + restore_spans, + segment_decode_budget, + segment_large_canvas, + ) + + self._require_ltx_video() + ltx_spans = [ + LtxSpan( + self._ltx_frames( + metadata, frame_stride=1, seek_ts=index.seconds_for_pts(span.start_pts), end_pts=span.end_pts + ), + tuple( + LtxSegment(start, end, segment.restoration.ltx_seed) + for (start, end), segment in zip(span.effect_ranges, segments) + ), + open_writer, + ) + for span, segments, open_writer in spans + ] + frames = sum( + round((span.end_pts - span.start_pts) * index.time_base * metadata.video_fps) for span, _, _ in spans + ) + try: + restore_spans( + ltx_spans, + detector=self.job_detection_model, + files=self.ltx_files, + frame_h=int(metadata.video_height), + frame_w=int(metadata.video_width), + batch_size=self.batch_size, + large_canvas=segment_large_canvas( + self.ltx_large_canvas, torch.cuda.get_device_properties(self.device).total_memory + ), + budget=segment_decode_budget, + device=self.device, + work_dir=work_dir, + progress=Progress(frames=frames, disable=self.disable_progress), + cancel=self._cancel_event, + ) + except Cancelled: + log.info("LTX restoration cancelled") + def _run_smart(self, metadata) -> None: codec = validate_smart_render( metadata, @@ -326,6 +400,14 @@ def _run_smart(self, metadata) -> None: if plan.segments != tuple(self.segments or ()): raise ValueError("Precomputed splice plan does not match pipeline segments") index = plan.index + default = job_restoration(self.restoration_model_name, self.ltx_seed) + segments_of = { + span: resolve_restorations(segments, default) + for span, segments in zip(plan.render_spans, plan.render_span_segments()) + } + for segments in segments_of.values(): + require_one_model(segments) + ltx_spans = {span for span, segments in segments_of.items() if segments[0].restoration.model == "ltx"} # AMF's H.264 encoder caps at 3 consecutive B-frames, so it cannot # match sources using more; re-render segments would not stitch # cleanly against the stream-copied ones. Fall back to a full @@ -362,6 +444,7 @@ def _run_smart(self, metadata) -> None: sum( round((span.end_pts - span.start_pts) * index.time_base * metadata.video_fps) for span in plan.render_spans + if span not in ltx_spans ), ) progress = Progressbar( @@ -382,36 +465,60 @@ def _run_smart(self, metadata) -> None: temp_dir = Path(temp_dir_name) fragments: list[tuple[Path, float]] = [] fragment_suffix = ".ts" if codec in {"h264", "hevc"} else ".mkv" + + def raw_path(span_index: int) -> Path: + return temp_dir / f"{span_index:04d}-raw.nut" + + def fragment_encoder(span_index: int, span: SpliceSpan) -> VideoEncoder: + return VideoEncoder( + str(raw_path(span_index)), + device=self.device, + metadata=metadata, + codec=codec, + encoder_settings=smart_encoder_settings, + lut_path=self.lut_path, + sharpen_strength=self.sharpen_strength, + output_fps=metadata.video_fps_exact, + pts_origin=span.start_pts, + smart_fragment=True, + ) + + def fragment_writer(span_index: int, span: SpliceSpan) -> Callable[[], _OfflineFrameWriter]: + return lambda: _OfflineFrameWriter(fragment_encoder(span_index, span), [time.monotonic()]) + + if ltx_spans: + self._run_ltx_spans( + metadata, + index, + [ + ( + span, + segments_of[span], + fragment_writer(span_index, span), + ) + for span_index, span in enumerate(plan.spans) + if span in ltx_spans + ], + work_root, + ) for span_index, span in enumerate(plan.spans): if self._cancel_event.is_set(): return - raw = temp_dir / f"{span_index:04d}-raw.nut" + raw = raw_path(span_index) normalized = temp_dir / f"{span_index:04d}{fragment_suffix}" duration = float((span.end_pts - span.start_pts) * index.time_base) - if span.is_render: - encoder_ctx = VideoEncoder( - str(raw), - device=self.device, - metadata=metadata, - codec=codec, - encoder_settings=smart_encoder_settings, - lut_path=self.lut_path, - sharpen_strength=self.sharpen_strength, - output_fps=metadata.video_fps_exact, - pts_origin=span.start_pts, - smart_fragment=True, - ) + if not span.is_render: + create_copy_fragment(self.input_video, span, index, raw, codec=codec) + elif span not in ltx_spans: self._run_pass( metadata=metadata, - encoder_ctx=encoder_ctx, + encoder_ctx=fragment_encoder(span_index, span), progress=progress, seek_ts=index.seconds_for_pts(span.start_pts), end_pts=span.end_pts, effect_ranges=span.effect_ranges, output_frame_count=max(1, round(duration * metadata.video_fps)), ) - else: - create_copy_fragment(self.input_video, span, index, raw, codec=codec) normalize_fragment(raw, normalized, codec=codec) fragments.append((normalized, duration)) @@ -437,9 +544,7 @@ def run(self) -> None: metadata = get_video_meta_data(str(self.input_video)) self.validate_metadata(metadata) self.configure_vr(metadata) - if self.ltx_files is not None: - self._run_ltx(metadata) - elif self.segments: + if self.segments: if self.fmp4: log.warning( "Fragmented MP4 is not available with segment processing; " @@ -447,6 +552,8 @@ def run(self) -> None: ) self.fmp4 = False self._run_smart(metadata) + elif self.restoration_model_name == "ltx": + self._run_ltx(metadata) else: self._run_full(metadata) self.completed = not self._cancel_event.is_set() diff --git a/jasna/segments.py b/jasna/segments.py index 72a073e9..1947b627 100644 --- a/jasna/segments.py +++ b/jasna/segments.py @@ -1,14 +1,37 @@ from __future__ import annotations -from dataclasses import dataclass +from dataclasses import dataclass, replace +from jasna.session_config import RestorationModelName -@dataclass(frozen=True, order=True) + +@dataclass(frozen=True) +class SegmentRestoration: + """The model that restores a segment; ``ltx_seed`` is set exactly for ``ltx``.""" + + model: RestorationModelName + ltx_seed: int | None + + def __post_init__(self) -> None: + if (self.model == "ltx") != (self.ltx_seed is not None): + raise ValueError("an LTX segment needs a seed and only an LTX segment has one") + + +def job_restoration(model: RestorationModelName, ltx_seed: int) -> SegmentRestoration: + return SegmentRestoration(model, ltx_seed if model == "ltx" else None) + + +@dataclass(frozen=True) class SegmentRange: - """A user-visible half-open time range, in seconds.""" + """A user-visible half-open time range, in seconds. ``restoration`` None means the + job's model; ``resolve_restorations`` fills it in before processing.""" start: float end: float + restoration: SegmentRestoration | None = None + + def __lt__(self, other: SegmentRange) -> bool: + return (self.start, self.end) < (other.start, other.end) def __post_init__(self) -> None: start = float(self.start) @@ -75,16 +98,42 @@ def normalize_segments( f"{format_timestamp(duration)}" ) + painted: list[SegmentRange] = [] + for segment in segments: + painted = [piece for old in painted for piece in _outside(old, segment)] + [segment] merged: list[SegmentRange] = [] - for segment in ordered: - if merged and segment.start <= merged[-1].end + 1e-9: - previous = merged[-1] - merged[-1] = SegmentRange(previous.start, max(previous.end, segment.end)) + for segment in sorted(painted): + previous = merged[-1] if merged else None + if ( + previous is not None + and segment.restoration == previous.restoration + and segment.start <= previous.end + 1e-9 + ): + merged[-1] = replace(previous, end=max(previous.end, segment.end)) else: merged.append(segment) return tuple(merged) +def _outside(segment: SegmentRange, cut: SegmentRange) -> list[SegmentRange]: + """The parts of ``segment`` that ``cut`` does not cover.""" + pieces = [] + if cut.start - segment.start > 1e-9: + pieces.append(replace(segment, end=min(segment.end, cut.start))) + if segment.end - cut.end > 1e-9: + pieces.append(replace(segment, start=max(segment.start, cut.end))) + return pieces + + +def resolve_restorations( + segments: tuple[SegmentRange, ...], default: SegmentRestoration +) -> tuple[SegmentRange, ...]: + return tuple( + segment if segment.restoration is not None else replace(segment, restoration=default) + for segment in segments + ) + + def parse_segments(spec: str, *, duration: float | None = None) -> tuple[SegmentRange, ...]: text = str(spec or "").strip() if not text: diff --git a/jasna/session_factory.py b/jasna/session_factory.py index e45d26cb..f1ae8f10 100644 --- a/jasna/session_factory.py +++ b/jasna/session_factory.py @@ -14,7 +14,7 @@ from pathlib import Path from typing import TYPE_CHECKING, Callable -from jasna.session_config import SessionConfig +from jasna.session_config import RestorationModelName, SessionConfig if TYPE_CHECKING: import torch @@ -147,29 +147,64 @@ def build_restoration_session( import torch from jasna.accelerator import is_amd_device - from jasna.engine_compiler import EngineCompilationRequest, ensure_engines_compiled - from jasna.restorer.basicvsrpp_mosaic_restorer import BasicvsrppMosaicRestorer - from jasna.restorer.denoise import DenoiseStep, DenoiseStrength - from jasna.restorer.restoration_pipeline import RestorationPipeline device = torch.device(config.device) - amd = is_amd_device(device) if config.tvai_denoise and config.secondary_restoration != "tvai": raise ValueError("TVAI Denoise requires secondary restoration 'tvai'") - if amd and config.secondary_restoration != "none": + if is_amd_device(device) and config.secondary_restoration != "none": raise ValueError( f"Secondary restoration '{config.secondary_restoration}' is not available in the AMD build yet" ) + if config.restoration_model_name == "ltx" and ( + config.secondary_restoration != "none" or config.denoise_strength != "none" + ): + raise ValueError("LTX restoration does not support secondary restoration or denoise") + session = RestorationSession(device=device, restoration_pipeline=None) + provide_restoration_models( + config, session, frozenset({config.restoration_model_name}), log_callback=log_callback + ) + return session + + +def provide_restoration_models( + config: SessionConfig, + session: RestorationSession, + models: frozenset[RestorationModelName], + *, + log_callback: Callable[[str], None] | None, +) -> None: + """Load each of ``models`` the session does not hold yet. ``--restoration-model-path`` + belongs to the job's model; another model loads from its default path.""" + if "basicvsrpp" in models and session.restoration_pipeline is None: + session.restoration_pipeline = _build_basicvsrpp_pipeline(config, session.device, log_callback=log_callback) + if "ltx" in models and session.ltx_files is None: + session.ltx_files = _ltx_model_files(config, session.device, log_callback=log_callback) + + +def _restoration_model_path(config: SessionConfig, name: RestorationModelName) -> Path: + from jasna.engine_paths import default_restoration_model_path - if config.restoration_model_name == "ltx": - return _build_ltx_session(config, device, log_callback=log_callback) + if name == config.restoration_model_name: + return config.restoration_model_path + return default_restoration_model_path(name) + +def _build_basicvsrpp_pipeline( + config: SessionConfig, device: "torch.device", *, log_callback: Callable[[str], None] | None +) -> "RestorationPipeline": + from jasna.accelerator import is_amd_device + from jasna.engine_compiler import EngineCompilationRequest, ensure_engines_compiled + from jasna.restorer.basicvsrpp_mosaic_restorer import BasicvsrppMosaicRestorer + from jasna.restorer.denoise import DenoiseStep, DenoiseStrength + from jasna.restorer.restoration_pipeline import RestorationPipeline + + model_path = _restoration_model_path(config, "basicvsrpp") compile_result = ensure_engines_compiled( EngineCompilationRequest( device=str(device), fp16=bool(config.fp16), - basicvsrpp=config.compile_basicvsrpp and not amd, - basicvsrpp_model_path=str(config.restoration_model_path), + basicvsrpp=config.compile_basicvsrpp and not is_amd_device(device), + basicvsrpp_model_path=str(model_path), detection=True, detection_model_name=config.detection_model_name, detection_model_path=str(config.detection_model_path), @@ -178,27 +213,23 @@ def build_restoration_session( ), log_callback=log_callback, ) - - secondary_restorer = _build_secondary_restorer(config, device) - restoration_pipeline = RestorationPipeline( + return RestorationPipeline( restorer=BasicvsrppMosaicRestorer( - checkpoint_path=str(config.restoration_model_path), + checkpoint_path=str(model_path), device=device, max_clip_size=int(config.max_clip_size), use_tensorrt=compile_result.use_basicvsrpp_tensorrt, fp16=bool(config.fp16), ), - secondary_restorer=secondary_restorer, + secondary_restorer=_build_secondary_restorer(config, device), denoise_strength=DenoiseStrength(config.denoise_strength), denoise_step=DenoiseStep(config.denoise_step), ) - return RestorationSession(device=device, restoration_pipeline=restoration_pipeline) - -def _build_ltx_session( +def _ltx_model_files( config: SessionConfig, device: "torch.device", *, log_callback: Callable[[str], None] | None -) -> RestorationSession: +) -> "LtxModelFiles": import torch from jasna.accelerator import is_nvidia_device @@ -207,11 +238,9 @@ def _build_ltx_session( if not is_nvidia_device(device): raise ValueError("LTX restoration needs an NVIDIA GPU") - if config.secondary_restoration != "none" or config.denoise_strength != "none": - raise ValueError("LTX restoration does not support secondary restoration or denoise") if config.ltx_fast and torch.cuda.get_device_capability(device)[0] < 10: raise ValueError("The fast LTX model needs an RTX 50-series (Blackwell) GPU") - files = LtxModelFiles.from_dir(config.restoration_model_path, fast=config.ltx_fast) + files = LtxModelFiles.from_dir(_restoration_model_path(config, "ltx"), fast=config.ltx_fast) ensure_engines_compiled( EngineCompilationRequest( device=str(device), @@ -223,11 +252,7 @@ def _build_ltx_session( ), log_callback=log_callback, ) - return RestorationSession( - device=device, - restoration_pipeline=None, - ltx_files=files, - ) + return files def build_pipeline( @@ -240,8 +265,15 @@ def build_pipeline( segments: "tuple[SegmentRange, ...] | None" = None, splice_plan: "SplicePlan | None" = None, ) -> "Pipeline": + """A per-video ``Pipeline``; the session first loads any model the segments ask for.""" from jasna.pipeline import Pipeline + from jasna.segments import job_restoration, resolve_restorations + default = job_restoration(config.restoration_model_name, config.ltx_seed) + models = frozenset( + segment.restoration.model for segment in resolve_restorations(tuple(segments or ()), default) + ) or frozenset({config.restoration_model_name}) + provide_restoration_models(config, session, models, log_callback=None) return Pipeline( config=config, session=session, diff --git a/tests/test_ltx_model_files.py b/tests/test_ltx_model_files.py index bfe46d06..0a80c8e9 100644 --- a/tests/test_ltx_model_files.py +++ b/tests/test_ltx_model_files.py @@ -28,7 +28,7 @@ def test_model_check_precedes_engine_compilation(tmp_path, monkeypatch): monkeypatch.setattr(LtxModelFiles, "from_dir", Mock(side_effect=ValueError("model unavailable"))) config = session_config(restoration_model_name="ltx", restoration_model_path=tmp_path) with pytest.raises(ValueError, match="model unavailable"): - session_factory._build_ltx_session(config, torch.device("cuda:0"), log_callback=None) + session_factory._ltx_model_files(config, torch.device("cuda:0"), log_callback=None) compile_engines.assert_not_called() diff --git a/tests/test_ltx_restore.py b/tests/test_ltx_restore.py index e2aa6211..a8c2f5da 100644 --- a/tests/test_ltx_restore.py +++ b/tests/test_ltx_restore.py @@ -1,62 +1,271 @@ +import threading +from pathlib import Path + +import pytest import torch from jasna.ltx import restore +from jasna.ltx.plan import Region +FILES = restore.LtxModelFiles(Path("t"), Path("v"), Path("d")) +MOSAIC = Region(box=(10.0, 10.0, 40.0, 40.0), polygon=[[10.0, 10.0], [40.0, 10.0], [40.0, 40.0], [10.0, 40.0]]) -def test_decode_window_converts_a_float32_decode(monkeypatch): - decoded = torch.tensor([-1.0, 0.0, 1.0]).view(1, 1, 3, 1, 1).expand(1, 3, 3, 1, 1).contiguous() - monkeypatch.setattr(restore, "decode_latent", lambda decoder, latent, free_bytes, generator: decoded.clone()) - monkeypatch.setattr(torch.cuda, "mem_get_info", lambda device: (1, 1)) - frames = restore._decode_window(None, torch.zeros(1, 128, 1, 1, 1), 0, torch.device("cpu")) - assert frames.dtype == torch.uint8 and frames.shape == (3, 3, 1, 1) - assert frames[:, 0, 0, 0].tolist() == [0, 128, 255] +def _video(frames: int, mosaic: range, pts_step: int = 10) -> list[tuple[torch.Tensor, int]]: + out = [] + for i in range(frames): + frame = torch.zeros(3, 64, 64, dtype=torch.uint8) + frame[1, 0, 0] = i + frame[0, 0, 0] = 1 if i in mosaic else 0 + out.append((frame, i * pts_step)) + return out + + +def _source(video, first: int = 0, batch: int = 4): + def frames(): + for start in range(first, len(video), batch): + chunk = video[start : start + batch] + yield torch.stack([frame for frame, _ in chunk]), [pts for _, pts in chunk] + + return frames + + +class _Detector: + def __init__(self): + self.batches: list[list[int]] = [] + + def __call__(self, batch, target_hw): + self.batches.append([int(frame[1, 0, 0]) for frame in batch]) + return batch -def test_large_canvas_needs_ten_gib_free(): - gib = 1 << 30 - assert restore.large_canvas_fits(True, 10 * gib) - assert not restore.large_canvas_fits(True, 9 * gib) - assert not restore.large_canvas_fits(False, 30 * gib) +class _Bar: + def __init__(self, total=0): + self.total, self.n = total, 0 -def test_denoise_reports_every_window_step(monkeypatch): - class Bar: - def __init__(self, total): - self.total, self.n = total, 0 + def update(self, n): + self.n += n - def update(self, n): - self.n += n + def __enter__(self): + return self - def close(self): - pass + def __exit__(self, *exc): + return False - class Store: - def get(self, index): - return index - def put(self, index, latent): - pass +class _Transformer: + seeds: list[list[int]] = [] - def delete(self, index): - pass + def __init__(self, path, device): + self.conditions = [None] * 8 - class Transformer: - conditions = [None] * 8 + def denoise_chain(self, references, seeds, advance): + _Transformer.seeds.append(list(seeds)) + for _ in self.conditions: + advance(len(references)) + return [torch.zeros(1) for _ in references] - def denoise_chain(self, references, seeds, advance): - for _ in self.conditions: - advance(len(references)) - return list(references) + def close(self): + pass + +class _Writer: + def __init__(self): + self.frames: list[tuple[int, int, bool]] = [] + self.closed = False + + def write(self, frame, pts, *, apply_lut=True): + self.frames.append((pts, int(frame[2, 0, 0]), apply_lut)) + + def close(self): + self.closed = True + + +@pytest.fixture +def fakes(monkeypatch): + budgets = [] bars = [] - monkeypatch.setattr(restore.Progress, "bar", lambda self, name, total, unit="frame": bars.append(Bar(total)) or bars[-1]) - window = lambda i: restore.Window(index=i, track_id=0, start=0, real_frames=1, crops=()) - plans = [ - restore.TrackPlan(track_id=0, start=0, polygons=(), windows=(window(0), window(1), window(2))), - restore.TrackPlan(track_id=1, start=0, polygons=(), windows=(window(3),)), + + def decode(decoder, latent, free_bytes, generator): + budgets.append(free_bytes) + return torch.zeros(1, 3, restore.WINDOW_FRAMES, 2, 2) + + def composite(source, candidates, *, feather): + out = source.clone() + out[2, 0, 0] = len(candidates) + return out + + def bar(self, name, total, unit="frame"): + bars.append(_Bar(total)) + return bars[-1] + + monkeypatch.setattr( + restore, "regions_per_frame", lambda frames, h, w: [[MOSAIC] if int(f[0, 0, 0]) else [] for f in frames] + ) + monkeypatch.setattr(restore, "crop_to_canvas", lambda frame, crop: torch.zeros(1)) + monkeypatch.setattr(restore, "_vae_pixels", lambda canvases, device: torch.zeros(1)) + monkeypatch.setattr(restore, "load_video_encoder", lambda path, device: lambda pixels: torch.zeros(1)) + monkeypatch.setattr(restore, "load_video_decoder", lambda vae, tuned, device: object()) + monkeypatch.setattr(restore, "LtxTransformer", _Transformer) + monkeypatch.setattr(restore, "decode_latent", decode) + monkeypatch.setattr(restore, "composite_frame", composite) + monkeypatch.setattr(restore.Progress, "bar", bar) + monkeypatch.setattr(torch.cuda, "mem_get_info", lambda device: (123 << 20, 1 << 40)) + monkeypatch.setattr(torch.cuda, "empty_cache", lambda: None) + _Transformer.seeds = [] + return budgets, bars + + +def _run_spans(tmp_path, spans): + restore.restore_spans( + spans, + detector=_Detector(), + files=FILES, + frame_h=64, + frame_w=64, + batch_size=4, + large_canvas=False, + budget=restore.segment_decode_budget, + device=torch.device("cpu"), + work_dir=tmp_path, + progress=restore.Progress(frames=0, disable=True), + cancel=threading.Event(), + ) + + +def test_segment_batches_count_batches_from_each_segment_start(): + video = _video(12, range(0)) + segment = restore.LtxSegment(20, 90, seed=0) + out = [ + (owner, [int(f[1, 0, 0]) for f in frames], pts) + for owner, frames, pts in restore.segment_batches(_source(video)(), [segment], 4) ] - restore.denoise( - plans, Transformer(), Store(), Store(), seed=0, - progress=restore.Progress(frames=0, disable=True), cancel=restore.threading.Event(), + assert out == [ + (None, [0, 1], [0, 10]), + (0, [2, 3, 4, 5], [20, 30, 40, 50]), + (0, [6, 7, 8], [60, 70, 80]), + (None, [9, 10, 11], [90, 100, 110]), + ] + + +def test_a_whole_video_segment_passes_decoded_batches_through(): + batches = list(_source(_video(10, range(0)))()) + whole = restore.LtxSegment(*restore.WHOLE_VIDEO_PTS, seed=0) + out = list(restore.segment_batches(iter(batches), [whole], 4)) + assert [owner for owner, _, _ in out] == [0, 0, 0] + assert all(torch.equal(frames, batch) and pts == batch_pts for (_, frames, pts), (batch, batch_pts) in zip(out, batches)) + + +def test_a_segment_plans_the_same_wherever_decoding_starts(fakes): + video = _video(30, range(3, 28)) + segment = restore.LtxSegment(50, 150, seed=0) + results = [] + for first in (0, 3): + detector = _Detector() + plans = restore.scan_span( + _source(video, first), + [segment], + detector, + batch_size=4, + frame_h=64, + frame_w=64, + large_canvas=False, + bar=_Bar(), + cancel=threading.Event(), + ) + results.append((plans, detector.batches)) + assert results[0] == results[1] + (plans,), batches = results[0] + assert batches == [[5, 6, 7, 8], [9, 10, 11, 12], [13, 14]] + assert [w.index for plan in plans for w in plan.windows] == [0] + assert plans[0].start == 0 and len(plans[0].polygons) == 10 + + +def test_restore_spans_seeds_segments_and_leaves_other_frames_untouched(fakes, tmp_path): + budgets, bars = fakes + video = _video(30, range(3, 28)) + writer = _Writer() + segments = (restore.LtxSegment(50, 150, seed=100), restore.LtxSegment(200, 260, seed=7)) + _run_spans(tmp_path, [restore.LtxSpan(_source(video), segments, lambda: writer)]) + + assert _Transformer.seeds == [[100], [7]] + assert budgets == [restore.SEGMENT_DECODE_BUDGET_BYTES[512]] * 2 + assert bars[2].total == bars[2].n == 16 + inside = {pts for pts in range(50, 150, 10)} | {pts for pts in range(200, 260, 10)} + assert [pts for pts, _, _ in writer.frames] == list(range(0, 300, 10)) + for pts, marker, apply_lut in writer.frames: + assert (marker, apply_lut) == ((1, True) if pts in inside else (0, False)) + assert writer.closed + + +def test_restore_spans_opens_each_span_writer_when_it_is_composed(fakes, tmp_path): + video = _video(30, range(3, 28)) + opened = [] + + def opener(name): + def open_writer(): + opened.append(name) + return _Writer() + + return open_writer + + spans = [ + restore.LtxSpan(_source(video[:15]), (restore.LtxSegment(50, 150, seed=1),), opener("a")), + restore.LtxSpan(_source(video, 15), (restore.LtxSegment(200, 260, seed=1),), opener("b")), + ] + _run_spans(tmp_path, spans) + assert opened == ["a", "b"] + assert _Transformer.seeds == [[1], [1]] + + +def test_restore_video_decodes_with_the_free_vram(fakes, tmp_path): + budgets, bars = fakes + video = _video(12, range(2, 10)) + writer = _Writer() + restore.restore_video( + _source(video), + writer, + detector=_Detector(), + files=FILES, + frame_h=64, + frame_w=64, + batch_size=4, + large_canvas=False, + seed=5, + device=torch.device("cpu"), + work_dir=tmp_path, + progress=restore.Progress(frames=12, disable=True), + cancel=threading.Event(), ) - assert bars[0].total == bars[0].n == 32 + assert budgets == [123 << 20] + assert _Transformer.seeds == [[5]] + assert bars[2].total == bars[2].n == 8 + assert [apply_lut for _, _, apply_lut in writer.frames] == [True] * 12 + + +def test_decode_window_converts_a_float32_decode(monkeypatch): + decoded = torch.tensor([-1.0, 0.0, 1.0]).view(1, 1, 3, 1, 1).expand(1, 3, 3, 1, 1).contiguous() + monkeypatch.setattr(restore, "decode_latent", lambda decoder, latent, free_bytes, generator: decoded.clone()) + frames = restore._decode_window(None, torch.zeros(1, 128, 1, 1, 1), 0, 512, lambda canvas: 1, torch.device("cpu")) + assert frames.dtype == torch.uint8 and frames.shape == (3, 3, 1, 1) + assert frames[:, 0, 0, 0].tolist() == [0, 128, 255] + + +def test_segment_decode_budget_ignores_free_vram(monkeypatch): + monkeypatch.setattr(torch.cuda, "mem_get_info", lambda device: (1, 1)) + assert restore.segment_decode_budget(512) == 6 << 30 + assert restore.segment_decode_budget(768) == 8 << 30 + + +def test_large_canvas_needs_ten_gib_free(): + gib = 1 << 30 + assert restore.large_canvas_fits(True, 10 * gib) + assert not restore.large_canvas_fits(True, 9 * gib) + assert not restore.large_canvas_fits(False, 30 * gib) + + +def test_segment_large_canvas_follows_the_gpu_size(): + gib = 1 << 30 + assert restore.segment_large_canvas(True, 16 * gib) + assert not restore.segment_large_canvas(True, 12 * gib) + assert not restore.segment_large_canvas(False, 32 * gib) diff --git a/tests/test_main_validation.py b/tests/test_main_validation.py index 807e1fb8..006f932e 100644 --- a/tests/test_main_validation.py +++ b/tests/test_main_validation.py @@ -61,6 +61,28 @@ def test_segments_auto_select_source_codec_and_reach_pipeline(self, tmp_path): assert pipeline_cls.call_args.kwargs["segments"] == (SegmentRange(1, 2),) assert pipeline_cls.call_args.kwargs["splice_plan"] is splice_plan + def test_ltx_accepts_segments(self, tmp_path): + metadata = MagicMock(codec_name="h264", duration=10.0) + with ( + patch("jasna.media.probe.get_video_meta_data", return_value=metadata), + patch("jasna.media.splice.validate_smart_render"), + patch("jasna.media.splice.probe_keyframes", return_value=MagicMock()), + patch("jasna.media.splice.build_splice_plan", return_value=MagicMock()), + patch("jasna.accelerator.is_nvidia_device", return_value=True), + patch("jasna.ltx.model_files.LtxModelFiles.from_dir", return_value=MagicMock()), + ): + pipeline_cls = _run_main_with_args( + tmp_path, ["--segments", "1-2", "--restoration-model-name", "ltx", "--ltx-seed", "3"] + ) + + config = pipeline_cls.call_args.kwargs["config"] + assert (config.restoration_model_name, config.ltx_seed) == ("ltx", 3) + assert pipeline_cls.call_args.kwargs["segments"] == (SegmentRange(1, 2),) + + def test_ltx_rejects_streaming(self, tmp_path): + with pytest.raises(SystemExit): + _run_main_with_args(tmp_path, ["--stream", "--restoration-model-name", "ltx"]) + def test_segments_reject_explicit_codec_mismatch(self, tmp_path): metadata = MagicMock(codec_name="h264", duration=10.0) with patch("jasna.media.probe.get_video_meta_data", return_value=metadata): diff --git a/tests/test_pipeline_segments.py b/tests/test_pipeline_segments.py index 48229af3..42e34b22 100644 --- a/tests/test_pipeline_segments.py +++ b/tests/test_pipeline_segments.py @@ -9,13 +9,15 @@ import torch from jasna.accelerator import AcceleratorVendor -from jasna.media.splice import KeyframeIndex, SplicePlan, SpliceSpan +from jasna.media.splice import KeyframeIndex, SmartRenderCompatibilityError, SplicePlan, SpliceSpan from jasna.pipeline import Pipeline -from jasna.segments import SegmentRange +from jasna.segments import SegmentRange, SegmentRestoration def test_smart_run_processes_only_render_spans_and_assembles_full_output(tmp_path) -> None: pipeline = object.__new__(Pipeline) + pipeline.restoration_model_name = "basicvsrpp" + pipeline.ltx_seed = 0 pipeline.input_video = tmp_path / "input.mp4" pipeline.output_video = tmp_path / "output.mp4" pipeline.codec = "h264" @@ -97,6 +99,8 @@ def test_smart_run_processes_only_render_spans_and_assembles_full_output(tmp_pat def test_smart_run_uses_working_dir_for_temp_files(tmp_path) -> None: pipeline = object.__new__(Pipeline) + pipeline.restoration_model_name = "basicvsrpp" + pipeline.ltx_seed = 0 pipeline.input_video = tmp_path / "input.mp4" pipeline.output_video = tmp_path / "out" / "output.mp4" pipeline.codec = "h264" @@ -143,6 +147,8 @@ def test_smart_run_uses_working_dir_for_temp_files(tmp_path) -> None: def test_amf_h264_full_reencode_preserves_selected_ranges(tmp_path) -> None: pipeline = object.__new__(Pipeline) + pipeline.restoration_model_name = "basicvsrpp" + pipeline.ltx_seed = 0 pipeline.input_video = tmp_path / "input.mp4" pipeline.output_video = tmp_path / "output.mp4" pipeline.codec = "h264" @@ -199,6 +205,8 @@ def test_smart_run_keeps_smart_render_unless_amd_exceeds_b_frame_cap( vendor, max_b_frames ) -> None: pipeline = object.__new__(Pipeline) + pipeline.restoration_model_name = "basicvsrpp" + pipeline.ltx_seed = 0 pipeline.input_video = Path("input.mp4") pipeline.output_video = Path("output.mp4") pipeline.codec = "h264" @@ -229,6 +237,8 @@ def test_smart_run_keeps_smart_render_unless_amd_exceeds_b_frame_cap( def test_smart_run_rejects_precomputed_plan_for_different_segments() -> None: pipeline = object.__new__(Pipeline) + pipeline.restoration_model_name = "basicvsrpp" + pipeline.ltx_seed = 0 pipeline.input_video = Path("input.mp4") pipeline.output_video = Path("output.mp4") pipeline.codec = "h264" @@ -245,3 +255,112 @@ def test_smart_run_rejects_precomputed_plan_for_different_segments() -> None: pytest.raises(ValueError, match="does not match"), ): pipeline._run_smart(MagicMock(duration=4.0, video_fps=30.0)) + + +def _mixed_pipeline(tmp_path, default_model: str, segments: tuple[SegmentRange, ...], spans) -> Pipeline: + pipeline = object.__new__(Pipeline) + pipeline.restoration_model_name = default_model + pipeline.ltx_seed = 5 + pipeline.input_video = tmp_path / "input.mp4" + pipeline.output_video = tmp_path / "output.mp4" + pipeline.codec = "h264" + pipeline.encoder_settings = {} + pipeline.device = torch.device("cuda:0") + pipeline.disable_progress = True + pipeline.progress_callback = None + pipeline.lut_path = None + pipeline.sharpen_strength = 0.0 + pipeline.retarget_high_fps = False + pipeline.working_dir = None + pipeline.segments = segments + pipeline.splice_plan = SplicePlan( + index=KeyframeIndex((0, 60, 120, 180), Fraction(1, 30), 0, 240), spans=spans, segments=segments + ) + pipeline._cancel_event = threading.Event() + pipeline._run_pass = MagicMock() + pipeline._run_ltx_spans = MagicMock() + return pipeline + + +def _run_smart_mocked(pipeline): + with ( + patch("jasna.pipeline.vendor_for_device", return_value=AcceleratorVendor.NVIDIA), + patch("jasna.pipeline.validate_smart_render", return_value="h264"), + patch("jasna.pipeline.resolve_smart_encoder_settings", return_value={}), + patch("jasna.pipeline.VideoEncoder"), + patch("jasna.pipeline.create_copy_fragment"), + patch("jasna.pipeline.normalize_fragment"), + patch("jasna.pipeline.concatenate_fragments") as concatenate, + patch("jasna.pipeline.mux_final_output"), + ): + pipeline._run_smart(MagicMock(duration=8.0, video_fps=30.0, video_fps_exact=Fraction(30, 1))) + return concatenate + + +def test_mixed_job_batches_ltx_spans_and_keeps_fragments_in_span_order(tmp_path) -> None: + ltx = SegmentRestoration("ltx", 42) + segments = (SegmentRange(2.5, 3.0), SegmentRange(6.5, 7.0, ltx)) + spans = ( + SpliceSpan("copy", 0, 60), + SpliceSpan("render", 60, 120, ((75, 90),)), + SpliceSpan("copy", 120, 180), + SpliceSpan("render", 180, 240, ((195, 210),)), + ) + pipeline = _mixed_pipeline(tmp_path, "basicvsrpp", segments, spans) + + concatenate = _run_smart_mocked(pipeline) + + pipeline._run_pass.assert_called_once() + assert pipeline._run_pass.call_args.kwargs["effect_ranges"] == ((75, 90),) + (_, _, ltx_spans, _), _ = pipeline._run_ltx_spans.call_args + assert [(span, segs) for span, segs, _ in ltx_spans] == [(spans[3], (SegmentRange(6.5, 7.0, ltx),))] + fragments = concatenate.call_args.args[0] + assert [path.name for path, _ in fragments] == ["0000.ts", "0001.ts", "0002.ts", "0003.ts"] + + +def test_segments_on_the_job_model_follow_an_ltx_job(tmp_path) -> None: + segments = (SegmentRange(2.5, 3.0),) + spans = (SpliceSpan("copy", 0, 60), SpliceSpan("render", 60, 120, ((75, 90),)), SpliceSpan("copy", 120, 240)) + pipeline = _mixed_pipeline(tmp_path, "ltx", segments, spans) + + _run_smart_mocked(pipeline) + + pipeline._run_pass.assert_not_called() + (_, _, ltx_spans, _), _ = pipeline._run_ltx_spans.call_args + assert ltx_spans[0][1] == (SegmentRange(2.5, 3.0, SegmentRestoration("ltx", 5)),) + + +def test_a_render_span_resolving_to_two_models_is_rejected(tmp_path) -> None: + segments = (SegmentRange(2.2, 2.4), SegmentRange(3.0, 3.2, SegmentRestoration("basicvsrpp", None))) + spans = (SpliceSpan("copy", 0, 60), SpliceSpan("render", 60, 120, ((66, 72), (90, 96))), SpliceSpan("copy", 120, 240)) + pipeline = _mixed_pipeline(tmp_path, "ltx", segments, spans) + + with pytest.raises(SmartRenderCompatibilityError): + _run_smart_mocked(pipeline) + + +@pytest.mark.parametrize( + ("model", "segments", "expected"), + [ + ("basicvsrpp", None, "_run_full"), + ("ltx", None, "_run_ltx"), + ("ltx", (SegmentRange(1, 2),), "_run_smart"), + ], +) +def test_run_routes_by_the_job_model_and_segments(model, segments, expected) -> None: + pipeline = object.__new__(Pipeline) + pipeline.restoration_model_name = model + pipeline.segments = segments + pipeline.input_video = Path("input.mp4") + pipeline.fmp4 = False + pipeline.ltx_files = MagicMock() + pipeline._cancel_event = threading.Event() + pipeline.validate_metadata = MagicMock() + pipeline.configure_vr = MagicMock() + for name in ("_run_full", "_run_ltx", "_run_smart"): + setattr(pipeline, name, MagicMock()) + + with patch("jasna.pipeline.get_video_meta_data"): + pipeline.run() + + assert [name for name in ("_run_full", "_run_ltx", "_run_smart") if getattr(pipeline, name).called] == [expected] diff --git a/tests/test_segments.py b/tests/test_segments.py index c5fbcee4..c72f9156 100644 --- a/tests/test_segments.py +++ b/tests/test_segments.py @@ -4,11 +4,14 @@ from jasna.segments import ( SegmentRange, + SegmentRestoration, format_segments, format_timestamp, + job_restoration, normalize_segments, parse_segments, parse_timestamp, + resolve_restorations, ) @@ -48,3 +51,40 @@ def test_segment_formatting_is_round_trippable() -> None: segments = (SegmentRange(1.25, 62.5), SegmentRange(3600, 3601)) assert parse_segments(format_segments(segments)) == segments assert format_timestamp(62.5) == "00:01:02.500" + + +LTX_A = SegmentRestoration("ltx", 1) +LTX_B = SegmentRestoration("ltx", 2) +STANDARD = SegmentRestoration("basicvsrpp", None) + + +def test_segment_restoration_seed_belongs_to_ltx_only() -> None: + assert job_restoration("basicvsrpp", 7) == STANDARD + assert job_restoration("ltx", 7) == SegmentRestoration("ltx", 7) + for model, seed in (("ltx", None), ("basicvsrpp", 3)): + with pytest.raises(ValueError): + SegmentRestoration(model, seed) + + +def test_normalize_merges_only_neighbours_with_the_same_restoration() -> None: + assert normalize_segments( + (SegmentRange(0, 5, LTX_A), SegmentRange(5, 8, LTX_A), SegmentRange(8, 9, LTX_B)) + ) == (SegmentRange(0, 8, LTX_A), SegmentRange(8, 9, LTX_B)) + + +def test_normalize_lets_a_later_range_cut_an_earlier_one() -> None: + assert normalize_segments((SegmentRange(0, 10, LTX_A), SegmentRange(3, 5, STANDARD))) == ( + SegmentRange(0, 3, LTX_A), + SegmentRange(3, 5, STANDARD), + SegmentRange(5, 10, LTX_A), + ) + assert normalize_segments((SegmentRange(3, 5, STANDARD), SegmentRange(0, 10, LTX_A))) == ( + SegmentRange(0, 10, LTX_A), + ) + + +def test_resolve_restorations_fills_in_only_the_job_default() -> None: + assert resolve_restorations((SegmentRange(0, 1), SegmentRange(2, 3, LTX_B)), LTX_A) == ( + SegmentRange(0, 1, LTX_A), + SegmentRange(2, 3, LTX_B), + ) diff --git a/tests/test_session_factory.py b/tests/test_session_factory.py index f63d418c..270715c9 100644 --- a/tests/test_session_factory.py +++ b/tests/test_session_factory.py @@ -5,7 +5,8 @@ import pytest -from jasna.segments import SegmentRange +from jasna.engine_paths import default_restoration_model_path +from jasna.segments import SegmentRange, SegmentRestoration from factories import session_config from jasna.session_config import SessionConfig from jasna.session_factory import RestorationSession, build_pipeline, build_restoration_session @@ -161,3 +162,31 @@ def test_session_reuses_its_detector_until_detection_settings_change() -> None: session.close() changed.close.assert_called_once_with() + + +def test_build_pipeline_loads_a_segment_model_the_session_lacks() -> None: + session = RestorationSession(device=MagicMock(), restoration_pipeline=MagicMock()) + segments = (SegmentRange(1, 2), SegmentRange(3, 4, SegmentRestoration("ltx", 9))) + + with ( + patch("jasna.accelerator.is_nvidia_device", return_value=True), + patch("jasna.engine_compiler.ensure_engines_compiled"), + patch("jasna.ltx.model_files.LtxModelFiles.from_dir") as from_dir, + patch("jasna.pipeline.Pipeline"), + ): + build_pipeline(session_config(), session, Path("in.mp4"), Path("out.mp4"), segments=segments) + build_pipeline(session_config(), session, Path("in.mp4"), Path("out.mp4"), segments=segments) + + from_dir.assert_called_once_with(default_restoration_model_path("ltx"), fast=False) + assert session.ltx_files is from_dir.return_value + + +def test_build_pipeline_keeps_a_session_that_holds_the_job_model() -> None: + pipeline = MagicMock() + session = RestorationSession(device=MagicMock(), restoration_pipeline=pipeline) + + with patch("jasna.pipeline.Pipeline"): + build_pipeline(session_config(), session, Path("in.mp4"), Path("out.mp4"), segments=(SegmentRange(1, 2),)) + + assert session.restoration_pipeline is pipeline + assert session.ltx_files is None diff --git a/tests/test_splice.py b/tests/test_splice.py index eb3788f1..03ede27d 100644 --- a/tests/test_splice.py +++ b/tests/test_splice.py @@ -20,7 +20,7 @@ resolve_smart_encoder_settings, validate_smart_render, ) -from jasna.segments import SegmentRange +from jasna.segments import SegmentRange, SegmentRestoration def _metadata(codec: str = "h264", **overrides) -> VideoMetadata: @@ -71,6 +71,30 @@ def test_plan_merges_two_selections_in_one_gop_bridge() -> None: assert plan.render_spans[0].effect_ranges == ((66, 72), (90, 96)) + +def test_plan_rejects_two_models_in_one_render_span() -> None: + ltx = SegmentRestoration("ltx", 1) + with pytest.raises(SmartRenderCompatibilityError) as exc: + build_splice_plan( + [SegmentRange(2.2, 2.4, ltx), SegmentRange(3.0, 3.2, SegmentRestoration("basicvsrpp", None))], + _index(), + duration=6, + ) + assert exc.value.reason == "mixed_models" + + +def test_plan_keeps_each_segment_behind_its_render_span_effect_range() -> None: + first, third = SegmentRestoration("ltx", 1), SegmentRestoration("ltx", 3) + plan = build_splice_plan( + [SegmentRange(2.2, 2.4, first), SegmentRange(3.0, 3.2), SegmentRange(5.0, 5.5, third)], + _index(points=(0, 60, 120, 150)), + duration=6, + ) + assert [[segment.restoration for segment in segments] for segments in plan.render_span_segments()] == [ + [first, None], + [third], + ] + def test_plan_offsets_effect_pts_by_stream_start_time() -> None: index = KeyframeIndex((900, 960, 1020), Fraction(1, 30), 900, 1080) plan = build_splice_plan([SegmentRange(2.5, 3)], index, duration=6) From d8f1b3db878be05b58d64b940fd11ef209fb5516 Mon Sep 17 00:00:00 2001 From: Kruk2 Date: Tue, 29 Sep 2026 16:33:06 +0200 Subject: [PATCH 28/48] feat(gui): show LTX progress, stage and time left in the progress bar Each LTX pass reports its share of the run and the time left; the control bar shows the pass name in place of the FPS. --- jasna/gui/app.py | 1 + jasna/gui/control_bar.py | 6 ++- jasna/gui/locales/en.py | 4 ++ jasna/gui/locales/ja.py | 4 ++ jasna/gui/locales/ko.py | 4 ++ jasna/gui/locales/th.py | 4 ++ jasna/gui/locales/zh.py | 4 ++ jasna/gui/processor.py | 6 ++- jasna/ltx/restore.py | 83 ++++++++++++++++++++++++++++++------ jasna/pipeline.py | 22 ++++++---- jasna/progressbar.py | 10 +++-- tests/test_gui_components.py | 8 ++++ tests/test_ltx_restore.py | 29 ++++++++++++- tests/test_pipeline_init.py | 8 ++++ tests/test_progressbar.py | 1 + 15 files changed, 165 insertions(+), 29 deletions(-) diff --git a/jasna/gui/app.py b/jasna/gui/app.py index e2f6a397..118d5358 100644 --- a/jasna/gui/app.py +++ b/jasna/gui/app.py @@ -703,6 +703,7 @@ def _handle_progress(self, update: ProgressUpdate): eta_seconds=update.eta_seconds, queue_current=queue_current, queue_total=len(jobs), + stage=update.stage, ) try: self._queue_panel.set_running(True, processing_job_id=job_id) diff --git a/jasna/gui/control_bar.py b/jasna/gui/control_bar.py index 6e8d1013..afc44349 100644 --- a/jasna/gui/control_bar.py +++ b/jasna/gui/control_bar.py @@ -389,11 +389,15 @@ def update_progress( eta_seconds: float = 0.0, queue_current: int = 0, queue_total: int = 0, + stage: str = "", ): self._filename_label.configure(text=filename or t("no_file_processing")) self._progress_bar.set(percent / 100.0) self._percent_label.configure(text=f"{int(percent)}%") - self._fps_label.configure(text=f"FPS: {fps:.1f}" if fps > 0 else "FPS: --") + if stage: + self._fps_label.configure(text=t(f"ltx_stage_{stage}")) + else: + self._fps_label.configure(text=f"FPS: {fps:.1f}" if fps > 0 else "FPS: --") if eta_seconds > 0: self._eta_label.configure(text=f"ETA: {format_duration(eta_seconds)}") diff --git a/jasna/gui/locales/en.py b/jasna/gui/locales/en.py index bfd795b3..b1cdbd10 100644 --- a/jasna/gui/locales/en.py +++ b/jasna/gui/locales/en.py @@ -207,6 +207,10 @@ "ltx_seed": "Seed", "ltx_fast": "Fast mode", "ltx_large_canvas": "Sharper large mosaics", + "ltx_stage_scan": "Finding mosaics", + "ltx_stage_encode": "Preparing", + "ltx_stage_denoise": "Restoring", + "ltx_stage_compose": "Saving", "section_advanced": "Advanced Processing", "section_secondary": "Secondary Restoration", "section_image_restoration": "Image Restoration (still images)", diff --git a/jasna/gui/locales/ja.py b/jasna/gui/locales/ja.py index 327715b3..2822260a 100644 --- a/jasna/gui/locales/ja.py +++ b/jasna/gui/locales/ja.py @@ -207,6 +207,10 @@ "ltx_seed": "シード", "ltx_fast": "高速モード", "ltx_large_canvas": "大きなモザイクをより鮮明に", + "ltx_stage_scan": "モザイクを検出中", + "ltx_stage_encode": "準備中", + "ltx_stage_denoise": "修復中", + "ltx_stage_compose": "保存中", "section_advanced": "詳細設定", "section_secondary": "二次修復", "section_image_restoration": "画像修復（静止画）", diff --git a/jasna/gui/locales/ko.py b/jasna/gui/locales/ko.py index 3a796973..f56a1ee3 100644 --- a/jasna/gui/locales/ko.py +++ b/jasna/gui/locales/ko.py @@ -206,6 +206,10 @@ "ltx_seed": "시드", "ltx_fast": "빠른 모드", "ltx_large_canvas": "큰 모자이크 더 선명하게", + "ltx_stage_scan": "모자이크 찾는 중", + "ltx_stage_encode": "준비 중", + "ltx_stage_denoise": "복원 중", + "ltx_stage_compose": "저장 중", "section_advanced": "고급 처리", "section_secondary": "2차 복원", "section_image_restoration": "이미지 복원 (정지 이미지)", diff --git a/jasna/gui/locales/th.py b/jasna/gui/locales/th.py index 65454afd..d2e2e45e 100644 --- a/jasna/gui/locales/th.py +++ b/jasna/gui/locales/th.py @@ -207,6 +207,10 @@ "ltx_seed": "ซีด", "ltx_fast": "โหมดเร็ว", "ltx_large_canvas": "โมเสกขนาดใหญ่คมชัดขึ้น", + "ltx_stage_scan": "กำลังค้นหาโมเสก", + "ltx_stage_encode": "กำลังเตรียม", + "ltx_stage_denoise": "กำลังฟื้นฟู", + "ltx_stage_compose": "กำลังบันทึก", "section_advanced": "การประมวลผลขั้นสูง", "section_secondary": "การฟื้นฟูขั้นที่สอง", "section_image_restoration": "การซ่อมแซมภาพ (ภาพนิ่ง)", diff --git a/jasna/gui/locales/zh.py b/jasna/gui/locales/zh.py index 319d4f68..1eab5b95 100644 --- a/jasna/gui/locales/zh.py +++ b/jasna/gui/locales/zh.py @@ -207,6 +207,10 @@ "ltx_seed": "种子", "ltx_fast": "快速模式", "ltx_large_canvas": "大马赛克更清晰", + "ltx_stage_scan": "正在查找马赛克", + "ltx_stage_encode": "正在准备", + "ltx_stage_denoise": "正在修复", + "ltx_stage_compose": "正在保存", "section_advanced": "高级处理", "section_secondary": "二次修复", "section_image_restoration": "图像修复（静态图片）", diff --git a/jasna/gui/processor.py b/jasna/gui/processor.py index 5a6b7e6e..c04cb1a8 100644 --- a/jasna/gui/processor.py +++ b/jasna/gui/processor.py @@ -29,6 +29,7 @@ class ProgressUpdate: frames_processed: int = 0 total_frames: int = 0 message: str = "" + stage: str = "" class ProcessingStopped(Exception): @@ -370,7 +371,9 @@ def _run_video_job( raise ProcessingStopped("Processing stopped") last_update_time = [0.0] - def progress_callback(progress_pct: float, fps: float, eta_seconds: float, frames_done: int, total: int): + def progress_callback( + progress_pct: float, fps: float, eta_seconds: float, frames_done: int, total: int, stage: str + ): current_time = time.time() if current_time - last_update_time[0] < 0.1: return @@ -387,6 +390,7 @@ def progress_callback(progress_pct: float, fps: float, eta_seconds: float, frame eta_seconds=eta_seconds, frames_processed=frames_done, total_frames=total, + stage=stage, )) pipeline = None diff --git a/jasna/ltx/restore.py b/jasna/ltx/restore.py index 9bdcacbb..c7e0b337 100644 --- a/jasna/ltx/restore.py +++ b/jasna/ltx/restore.py @@ -18,10 +18,11 @@ import logging import tempfile import threading +import time from collections.abc import Callable, Iterator, Sequence from dataclasses import dataclass from pathlib import Path -from typing import Protocol +from typing import Literal, Protocol import torch from tqdm import tqdm @@ -79,15 +80,69 @@ class LtxSpan: open_writer: Callable[[], FrameWriter] -@dataclass(frozen=True) +LtxStage = Literal["scan", "encode", "denoise", "compose"] +ProgressReport = Callable[[LtxStage, float, float], None] +_STAGE_NUMBER = {"scan": 1, "encode": 2, "denoise": 3, "compose": 4} +_STAGE_SHARE = {"scan": 0.02, "encode": 0.08, "denoise": 0.8, "compose": 0.1} + + class Progress: - """One console bar per pass; ``frames`` sizes the frame-based passes.""" + """One console bar per pass; ``frames`` sizes the frame-based passes. ``report`` gets + ``(stage, fraction of the whole run, seconds left)``; seconds left is 0 until the + denoise pass has timed a step.""" + + def __init__(self, frames: int, *, disable: bool, report: ProgressReport | None) -> None: + self.frames = frames + self._disable = disable + self._report = report + self._started = time.monotonic() + self._done_share = 0.0 + self._timed = False + + def bar(self, stage: LtxStage, total: int, unit: str = "frame") -> _StageBar: + tqdm_bar = tqdm( + total=total, + desc=f"LTX {_STAGE_NUMBER[stage]}/4 {stage}", + unit=unit, + dynamic_ncols=True, + disable=self._disable, + ) + return _StageBar(self, stage, tqdm_bar) + + def _advanced(self, stage: LtxStage, done: int, total: int) -> None: + if self._report is None: + return + fraction = self._done_share + _STAGE_SHARE[stage] * (min(1.0, done / total) if total else 1.0) + self._timed = self._timed or (stage == "denoise" and done > 0) + elapsed = time.monotonic() - self._started + self._report(stage, fraction, elapsed * (1.0 - fraction) / fraction if self._timed else 0.0) + + def _finished(self, stage: LtxStage) -> None: + self._done_share += _STAGE_SHARE[stage] + + +class _StageBar: + def __init__(self, progress: Progress, stage: LtxStage, bar: tqdm) -> None: + self._progress = progress + self._stage = stage + self._bar = bar + self._done = 0 + progress._advanced(stage, 0, bar.total) + + def update(self, n: int) -> None: + self._bar.update(n) + self._done += n + self._progress._advanced(self._stage, self._done, self._bar.total) + + def close(self) -> None: + self._bar.close() + self._progress._finished(self._stage) - frames: int - disable: bool + def __enter__(self) -> _StageBar: + return self - def bar(self, name: str, total: int, unit: str = "frame") -> tqdm: - return tqdm(total=total, desc=f"LTX {name}", unit=unit, dynamic_ncols=True, disable=self.disable) + def __exit__(self, *_exc) -> None: + self.close() class LatentStore: @@ -169,7 +224,7 @@ def scan_span( frame_h: int, frame_w: int, large_canvas: bool, - bar: tqdm, + bar: _StageBar, cancel: threading.Event, ) -> list[list[TrackPlan]]: """The track plans of each segment, frame indices counted from the segment's first frame.""" @@ -237,7 +292,7 @@ def encode_span( *, batch_size: int, device: torch.device, - bar: tqdm, + bar: _StageBar, cancel: threading.Event, ) -> None: """VAE-encode every planned window's reference crops into its segment's store.""" @@ -357,7 +412,7 @@ def compose_span( batch_size: int, budget: DecodeBudget, device: torch.device, - bar: tqdm, + bar: _StageBar, cancel: threading.Event, ) -> None: """Write every frame of ``span``: segment frames composited, the rest unchanged.""" @@ -415,7 +470,7 @@ def restore_spans( """Restore every segment of ``spans`` and write each span's frames to a writer of its own (closed once the span is written). ``large_canvas`` allows 768 px windows for large mosaics; ``budget`` gives the decode memory budget per canvas size.""" - with progress.bar("1/4 scan", progress.frames) as bar: + with progress.bar("scan", progress.frames) as bar: plans = [ scan_span( span.frames, @@ -441,7 +496,7 @@ def restore_spans( ] if windows: encoder = load_video_encoder(files.vae, device) - with progress.bar("2/4 encode", progress.frames) as bar: + with progress.bar("encode", progress.frames) as bar: for span, span_plans, span_stores in zip(spans, plans, stores): encode_span( span.frames, @@ -458,7 +513,7 @@ def restore_spans( torch.cuda.empty_cache() transformer = LtxTransformer(files.transformer, device) try: - with progress.bar("3/4 denoise", len(transformer.conditions) * windows, unit="step") as bar: + with progress.bar("denoise", len(transformer.conditions) * windows, unit="step") as bar: for span, span_plans, span_stores in zip(spans, plans, stores): for segment, segment_plans, (references, finals) in zip(span.segments, span_plans, span_stores): for plan in segment_plans: @@ -472,7 +527,7 @@ def restore_spans( del transformer torch.cuda.empty_cache() decoder = load_video_decoder(files.vae, files.tuned_decoder, device) if windows else None - with progress.bar("4/4 compose", progress.frames) as bar: + with progress.bar("compose", progress.frames) as bar: for span, span_plans, span_stores in zip(spans, plans, stores): compose_span( span, diff --git a/jasna/pipeline.py b/jasna/pipeline.py index 82178567..64ccfb25 100644 --- a/jasna/pipeline.py +++ b/jasna/pipeline.py @@ -32,7 +32,7 @@ validate_smart_render, ) from jasna.pipeline_threads import run_restoration_pass -from jasna.progressbar import Progressbar +from jasna.progressbar import ProgressCallback, Progressbar from jasna.restorer.secondary_restorer import AsyncSecondaryRestorer from jasna.segments import SegmentRange, job_restoration, resolve_restorations from jasna.session_config import SessionConfig @@ -78,7 +78,7 @@ def __init__( session: RestorationSession, input_video: Path, output_video: Path, - progress_callback: Callable[[dict], None] | None, + progress_callback: ProgressCallback | None, segments: tuple[SegmentRange, ...] | None, splice_plan: SplicePlan | None, ) -> None: @@ -297,12 +297,21 @@ def frames(): return frames + def _ltx_progress(self, frames: int): + from jasna.ltx.restore import Progress + + callback = self.progress_callback + report = None + if callback is not None: + report = lambda stage, fraction, eta: callback(fraction * 100.0, 0.0, eta, 0, 0, stage) + return Progress(frames, disable=self.disable_progress, report=report) + def _require_ltx_video(self) -> None: if self.vr_resolution.is_sbs: raise ValueError("LTX restoration does not support VR180 side-by-side video") def _run_ltx(self, metadata) -> None: - from jasna.ltx.restore import Cancelled, Progress, restore_video + from jasna.ltx.restore import Cancelled, restore_video self._require_ltx_video() frame_rate = self._resolve_frame_rate(metadata) @@ -320,9 +329,7 @@ def _run_ltx(self, metadata) -> None: seed=self.ltx_seed, device=self.device, work_dir=self.working_dir or self.output_video.parent, - progress=Progress( - frames=frame_rate.output_frame_count(metadata.num_frames), disable=self.disable_progress - ), + progress=self._ltx_progress(frame_rate.output_frame_count(metadata.num_frames)), cancel=self._cancel_event, ) except Cancelled: @@ -342,7 +349,6 @@ def _run_ltx_spans( Cancelled, LtxSegment, LtxSpan, - Progress, restore_spans, segment_decode_budget, segment_large_canvas, @@ -379,7 +385,7 @@ def _run_ltx_spans( budget=segment_decode_budget, device=self.device, work_dir=work_dir, - progress=Progress(frames=frames, disable=self.disable_progress), + progress=self._ltx_progress(frames), cancel=self._cancel_event, ) except Cancelled: diff --git a/jasna/progressbar.py b/jasna/progressbar.py index 84256806..e22ab214 100644 --- a/jasna/progressbar.py +++ b/jasna/progressbar.py @@ -5,6 +5,10 @@ from tqdm import tqdm +ProgressCallback = Callable[[float, float, float, int, int, str], None] +"""``(percent, fps, eta_seconds, frames_done, total_frames, stage)``; ``stage`` names the LTX +pass running, or is empty for frame-by-frame restoration.""" + class Progressbar: """Progress bar with time remaining estimation and speed display.""" @@ -14,7 +18,7 @@ def __init__( total_frames: int, video_fps: float, disable: bool = False, - callback: Callable[[float, float, float, int, int], None] | None = None, + callback: ProgressCallback | None = None, ): """Initialize progressbar. @@ -22,7 +26,7 @@ def __init__( total_frames: Total number of frames to process video_fps: Video FPS for buffer sizing disable: If True, suppress tqdm console output - callback: Optional callback(progress_pct, fps, eta_seconds, frames_done, total_frames) + callback: Optional ``ProgressCallback`` """ self.total_frames = total_frames self.callback = callback @@ -104,7 +108,7 @@ def update(self, n: int = 1) -> None: fps = 1.0 / mean_duration if mean_duration > 0 else 0.0 remaining = self.total_frames - self.frames_processed eta_seconds = remaining * mean_duration - self.callback(progress_pct, fps, eta_seconds, self.frames_processed, self.total_frames) + self.callback(progress_pct, fps, eta_seconds, self.frames_processed, self.total_frames, "") def _get_mean_processing_duration(self) -> float: """Calculate mean frame processing duration from the buffer.""" diff --git a/tests/test_gui_components.py b/tests/test_gui_components.py index 6a85d387..a03d2f2f 100644 --- a/tests/test_gui_components.py +++ b/tests/test_gui_components.py @@ -259,3 +259,11 @@ def test_grab_modal_restores_minimized_dialog_on_focus(tk_root) -> None: dialog.event_generate("") tk_root.update() assert dialog.winfo_ismapped() + + +def test_control_bar_shows_the_ltx_stage_instead_of_fps(tk_root) -> None: + control_bar = ControlBar(tk_root) + control_bar.update_progress(filename="a.mp4", percent=42.0, eta_seconds=90.0, stage="denoise") + assert control_bar._fps_label.cget("text") == t("ltx_stage_denoise") + control_bar.update_progress(filename="a.mp4", percent=50.0, fps=30.0) + assert control_bar._fps_label.cget("text") == "FPS: 30.0" diff --git a/tests/test_ltx_restore.py b/tests/test_ltx_restore.py index a8c2f5da..fcd9644d 100644 --- a/tests/test_ltx_restore.py +++ b/tests/test_ltx_restore.py @@ -128,7 +128,7 @@ def _run_spans(tmp_path, spans): budget=restore.segment_decode_budget, device=torch.device("cpu"), work_dir=tmp_path, - progress=restore.Progress(frames=0, disable=True), + progress=restore.Progress(0, disable=True, report=None), cancel=threading.Event(), ) @@ -234,7 +234,7 @@ def test_restore_video_decodes_with_the_free_vram(fakes, tmp_path): seed=5, device=torch.device("cpu"), work_dir=tmp_path, - progress=restore.Progress(frames=12, disable=True), + progress=restore.Progress(12, disable=True, report=None), cancel=threading.Event(), ) assert budgets == [123 << 20] @@ -269,3 +269,28 @@ def test_segment_large_canvas_follows_the_gpu_size(): assert restore.segment_large_canvas(True, 16 * gib) assert not restore.segment_large_canvas(True, 12 * gib) assert not restore.segment_large_canvas(False, 32 * gib) + + +def test_progress_reports_the_stage_and_the_share_of_the_whole_run(monkeypatch): + clock = iter([0.0, 100.0, 100.0, 100.0, 300.0, 300.0]) + monkeypatch.setattr(restore.time, "monotonic", lambda: next(clock)) + reports = [] + progress = restore.Progress(10, disable=True, report=lambda *args: reports.append(args)) + with progress.bar("scan", 10) as bar: + bar.update(10) + with progress.bar("denoise", 4, unit="step") as bar: + bar.update(2) + assert [(stage, round(fraction, 3)) for stage, fraction, _eta in reports] == [ + ("scan", 0.0), + ("scan", 0.02), + ("denoise", 0.02), + ("denoise", 0.42), + ] + assert [eta for *_rest, eta in reports[:3]] == [0.0, 0.0, 0.0] + assert reports[-1][2] == pytest.approx(300.0 * 0.58 / 0.42) + + +def test_progress_without_a_report_only_drives_the_console_bar(): + progress = restore.Progress(3, disable=True, report=None) + with progress.bar("compose", 3) as bar: + bar.update(3) diff --git a/tests/test_pipeline_init.py b/tests/test_pipeline_init.py index 3315e506..78ea97e0 100644 --- a/tests/test_pipeline_init.py +++ b/tests/test_pipeline_init.py @@ -3,6 +3,7 @@ from types import SimpleNamespace from unittest.mock import MagicMock, patch +import pytest import torch from factories import make_pipeline @@ -45,6 +46,13 @@ def test_progress_callback(self): p = _make_pipeline(progress_callback=cb) assert p.progress_callback is cb + def test_ltx_progress_reaches_the_progress_callback_with_its_stage(self): + cb = MagicMock() + p = _make_pipeline(progress_callback=cb, disable_progress=True) + with p._ltx_progress(10).bar("scan", 10) as bar: + bar.update(10) + assert cb.call_args[0] == (pytest.approx(2.0), 0.0, 0.0, 0, 0, "scan") + def test_retarget_high_fps_defaults_off_and_can_be_enabled(self): assert _make_pipeline().retarget_high_fps is False assert _make_pipeline(retarget_high_fps=True).retarget_high_fps is True diff --git a/tests/test_progressbar.py b/tests/test_progressbar.py index 7c6f098b..b88591ba 100644 --- a/tests/test_progressbar.py +++ b/tests/test_progressbar.py @@ -77,6 +77,7 @@ def test_callback_invoked(self): assert last_call[0][0] == pytest.approx(100.0) assert last_call[0][3] == 10 assert last_call[0][4] == 10 + assert last_call[0][5] == "" def test_buffer_respects_max_len(self): pb = Progressbar(total_frames=200, video_fps=1, disable=True) From c35c8eaf7ec02b7a5929d54280cb9108ac58f642 Mon Sep 17 00:00:00 2001 From: Kruk2 Date: Tue, 29 Sep 2026 16:38:21 +0200 Subject: [PATCH 29/48] feat(gui): choose the model and LTX seed per range in the segment editor New ranges take the settings panel's model and seed, the selected range is edited in place, and "Use for all ranges" copies it. LTX ranges are teal on the timeline. --- jasna/gui/app.py | 5 +- jasna/gui/locales/en.py | 13 +- jasna/gui/locales/ja.py | 13 +- jasna/gui/locales/ko.py | 13 +- jasna/gui/locales/th.py | 13 +- jasna/gui/locales/zh.py | 13 +- jasna/gui/queue_panel.py | 3 + jasna/gui/restoration_preview.py | 3 - jasna/gui/segment_editor.py | 200 +++++++++++++++++- jasna/gui/segment_editor_state.py | 52 ++++- jasna/gui/segment_timeline.py | 19 +- jasna/gui/settings_panel.py | 4 + .../settings_sections/restoration_model.py | 5 +- jasna/gui/theme.py | 3 + jasna/gui/validation.py | 19 +- jasna/gui/video_session.py | 3 +- tests/test_gui_tvai_validation.py | 29 +-- tests/test_locales.py | 4 +- tests/test_restoration_preview.py | 52 ++--- tests/test_segment_editor.py | 91 +++++++- tests/test_segment_editor_state.py | 129 +++++++++-- tests/test_segment_timeline.py | 10 +- tests/test_video_session.py | 4 + 23 files changed, 582 insertions(+), 118 deletions(-) diff --git a/jasna/gui/app.py b/jasna/gui/app.py index 118d5358..efa7e6fa 100644 --- a/jasna/gui/app.py +++ b/jasna/gui/app.py @@ -356,6 +356,7 @@ def _build_main_body(self): self._settings_panel.get_settings, lambda: self._processor is not None and self._processor.is_running(), self._set_preview_gpu_busy, + self._settings_panel.ltx_unavailable_reason, ) self._queue_panel.set_initial_output( self._preset_manager.get_last_output_folder(), @@ -594,7 +595,9 @@ def _on_start(self): settings = self._settings_panel.get_settings() from jasna.gui.validation import validate_gui_start - errors = validate_gui_start(settings, jobs) + errors = validate_gui_start( + settings, jobs, ltx_available=self._settings_panel.ltx_unavailable_reason() is None + ) if errors: from tkinter import messagebox diff --git a/jasna/gui/locales/en.py b/jasna/gui/locales/en.py index b1cdbd10..2be662a1 100644 --- a/jasna/gui/locales/en.py +++ b/jasna/gui/locales/en.py @@ -77,7 +77,7 @@ "segments_preview_reset_view": "Reset view", "segments_preview_reset_view_hint": "Reset zoom and pan to show the full frame", "segments_restore_gpu_busy": "Unavailable while processing is running", - "segments_restore_ltx_unavailable": "Preview is not available with LTX yet", + "segments_restore_ltx_unavailable": "LTX ranges can't use this preview. Use Try seed instead.", "segments_restore_loading_models": "Loading restoration models…", "segments_restore_restoring": "Restoring…", "segments_restore_failed": "Restoration preview failed: {message}", @@ -125,9 +125,13 @@ "segments_scan_model_changed": "Model changed — scan again to update detections", "segments_timeline_title": "Timeline", "segments_timeline_hint": "Detailed view · full-video overview below", - "segments_legend_selected": "Selected for restoration", "segments_legend_detected": "Detected by scan", "segments_legend_playhead": "Current frame", + "segments_restore_with": "Restore with", + "segments_use_for_all": "Use for all ranges", + "tip_segments_use_for_all": "Give every range this model and seed.", + "segments_ltx_needs_flat_video": "LTX can't restore VR side-by-side videos.", + "segments_ltx_ranges_unavailable": "LTX can't run on this PC, so the LTX ranges can't be restored.", "segments_suggest_mask": "Suggest better mask", "segments_suggest_mask_hint": "Draw the correct mosaic outline on this frame and share it anonymously to improve future detection models", "mask_editor_title": "Suggest a better mosaic mask", @@ -169,7 +173,7 @@ "segments_time_out_of_bounds": "Times must be inside the video duration.", "segments_merged": "Merged with {count} overlapping or adjacent range(s).", "segments_drag_hint": "No ranges selected — the full video will be restored. Drag across the timeline, or use I and O to mark exact frames.", - "segments_range_row": "{index}. {start} → {end}\n {duration:.1f}s", + "segments_range_row": "{index}. {start} → {end}\n {duration:.1f}s · {model}", "segments_delete_range": "Delete range", "segments_zoom_out": "Zoom out (mouse wheel down)", "segments_zoom_in": "Zoom in (mouse wheel up)", @@ -183,6 +187,7 @@ "segments_smart_render_range_too_short": "A selected range is shorter than one video frame. Make it longer.", "segments_smart_render_before_first_keyframe": "The first range starts before the first frame where this video can be split safely. Move its start slightly later.", "segments_smart_render_whole_video": "These ranges leave no safe places to split the video, so processing them separately would re-encode the whole video. Select a smaller area, or clear all ranges to process the full video.", + "segments_mixed_models_too_close": "A Standard range and an LTX range are too close. Move them apart or use one model for both.", "segments_discard_title": "Discard range changes?", "segments_discard_changes": "Close without applying your restoration-range changes?", "segments_edit_tooltip": "Choose which parts of this video should be restored", @@ -468,8 +473,8 @@ # App messages "error_cannot_start": "Cannot start processing:", "error_start_title": "Can't start", + "error_ltx_ranges_unavailable": "These videos have LTX ranges, but LTX can't run on this PC: {files}", "error_post_export_command_required": "Custom post-export command is required", - "error_ltx_segments": "LTX processes whole videos. Clear the segments or pick Standard: {files}", # Settings panel "dialog_select_tvai_ffmpeg": "Select Topaz Video ffmpeg.exe", diff --git a/jasna/gui/locales/ja.py b/jasna/gui/locales/ja.py index 2822260a..ee84033c 100644 --- a/jasna/gui/locales/ja.py +++ b/jasna/gui/locales/ja.py @@ -77,7 +77,7 @@ "segments_preview_reset_view": "表示をリセット", "segments_preview_reset_view_hint": "移動とズームをリセットしてフレーム全体を表示します", "segments_restore_gpu_busy": "処理の実行中は使用できません", - "segments_restore_ltx_unavailable": "LTX ではまだプレビューできません", + "segments_restore_ltx_unavailable": "LTX の範囲ではこのプレビューは使えません。「シードを試す」を使ってください。", "segments_restore_loading_models": "復元モデルを読み込み中…", "segments_restore_restoring": "復元中…", "segments_restore_failed": "復元プレビューに失敗しました：{message}", @@ -125,9 +125,13 @@ "segments_scan_model_changed": "モデルを変更しました — 検出結果を更新するには再スキャンしてください", "segments_timeline_title": "タイムライン", "segments_timeline_hint": "詳細表示 · 下段は動画全体の概要", - "segments_legend_selected": "復元対象", "segments_legend_detected": "スキャン検出", "segments_legend_playhead": "現在のフレーム", + "segments_restore_with": "復元方法", + "segments_use_for_all": "すべての範囲に使う", + "tip_segments_use_for_all": "すべての範囲にこのモデルとシードを使います。", + "segments_ltx_needs_flat_video": "LTX は VR のサイドバイサイド動画を復元できません。", + "segments_ltx_ranges_unavailable": "この PC では LTX を使えないため、LTX の範囲は復元できません。", "segments_suggest_mask": "より良いマスクを提案", "segments_suggest_mask_hint": "このフレームに正しいモザイクの輪郭を描き、匿名で共有して今後の検出モデルの改善に役立てます", "mask_editor_title": "より良いモザイクマスクを提案", @@ -169,7 +173,7 @@ "segments_time_out_of_bounds": "時刻は動画の長さの範囲内にしてください。", "segments_merged": "重複または隣接する {count} 区間と結合しました。", "segments_drag_hint": "区間未選択なら動画全体を復元します。タイムラインをドラッグするか、I と O で正確なフレームを指定します。", - "segments_range_row": "{index}. {start} → {end}\n {duration:.1f}秒", + "segments_range_row": "{index}. {start} → {end}\n {duration:.1f}秒 · {model}", "segments_delete_range": "区間を削除", "segments_zoom_out": "縮小（マウスホイール下）", "segments_zoom_in": "拡大（マウスホイール上）", @@ -183,6 +187,7 @@ "segments_smart_render_range_too_short": "選択区間が1フレームより短いため、長くしてください。", "segments_smart_render_before_first_keyframe": "最初の区間が安全なカット位置より前から始まっています。開始位置を少し後ろへ移動してください。", "segments_smart_render_whole_video": "使用できる安全なカット位置が残らないため、個別処理すると動画全体の再エンコードが必要です。範囲を小さくするか、全区間を消して動画全体を処理してください。", + "segments_mixed_models_too_close": "「標準」の範囲と LTX の範囲が近すぎます。離すか、両方を同じモデルにしてください。", "segments_discard_title": "区間の変更を破棄しますか？", "segments_discard_changes": "復元区間の変更を適用せずに閉じますか？", "segments_edit_tooltip": "この動画で復元する部分を選択", @@ -468,8 +473,8 @@ # App messages "error_cannot_start": "処理を開始できません:", "error_start_title": "開始できません", + "error_ltx_ranges_unavailable": "次の動画には LTX の範囲がありますが、この PC では LTX を使えません：{files}", "error_post_export_command_required": "エクスポート後のカスタムコマンドが必要です", - "error_ltx_segments": "LTX は動画全体を処理します。セグメントを消去するか「標準」を選んでください：{files}", # Settings panel "dialog_select_tvai_ffmpeg": "Topaz Video の ffmpeg.exe を選択", diff --git a/jasna/gui/locales/ko.py b/jasna/gui/locales/ko.py index f56a1ee3..5c48f9a9 100644 --- a/jasna/gui/locales/ko.py +++ b/jasna/gui/locales/ko.py @@ -76,7 +76,7 @@ "segments_preview_reset_view": "보기 초기화", "segments_preview_reset_view_hint": "이동과 확대/축소를 초기화하고 전체 프레임을 표시합니다", "segments_restore_gpu_busy": "처리가 진행 중일 때는 사용할 수 없습니다", - "segments_restore_ltx_unavailable": "LTX에서는 아직 미리보기를 사용할 수 없습니다", + "segments_restore_ltx_unavailable": "LTX 구간에서는 이 미리보기를 쓸 수 없습니다. 대신 '시드 시험'을 사용하세요.", "segments_restore_loading_models": "복원 모델을 불러오는 중…", "segments_restore_restoring": "복원 중…", "segments_restore_failed": "복원 미리보기 실패: {message}", @@ -124,9 +124,13 @@ "segments_scan_model_changed": "모델이 변경되었습니다 — 감지 결과를 갱신하려면 다시 스캔하세요", "segments_timeline_title": "타임라인", "segments_timeline_hint": "상세 보기 · 아래는 전체 비디오 개요", - "segments_legend_selected": "복원 대상으로 선택", "segments_legend_detected": "스캔 감지 결과", "segments_legend_playhead": "현재 프레임", + "segments_restore_with": "복원 방식", + "segments_use_for_all": "모든 구간에 사용", + "tip_segments_use_for_all": "모든 구간에 이 모델과 시드를 적용합니다.", + "segments_ltx_needs_flat_video": "LTX는 VR 사이드 바이 사이드 영상을 복원할 수 없습니다.", + "segments_ltx_ranges_unavailable": "이 PC에서는 LTX를 실행할 수 없어 LTX 구간을 복원할 수 없습니다.", "segments_suggest_mask": "더 나은 마스크 제안", "segments_suggest_mask_hint": "이 프레임에 올바른 모자이크 윤곽을 그려 익명으로 공유하면 향후 감지 모델 개선에 도움이 됩니다", "mask_editor_title": "더 나은 모자이크 마스크 제안", @@ -168,7 +172,7 @@ "segments_time_out_of_bounds": "시간은 비디오 길이 안에 있어야 합니다.", "segments_merged": "겹치거나 인접한 구간 {count}개와 병합했습니다.", "segments_drag_hint": "구간을 선택하지 않으면 전체 비디오를 복원합니다. 타임라인을 드래그하거나 I와 O로 정확한 프레임을 표시하세요.", - "segments_range_row": "{index}. {start} → {end}\n {duration:.1f}초", + "segments_range_row": "{index}. {start} → {end}\n {duration:.1f}초 · {model}", "segments_delete_range": "구간 삭제", "segments_zoom_out": "축소(마우스 휠 아래)", "segments_zoom_in": "확대(마우스 휠 위)", @@ -182,6 +186,7 @@ "segments_smart_render_range_too_short": "선택 구간이 비디오 한 프레임보다 짧습니다. 구간을 늘리세요.", "segments_smart_render_before_first_keyframe": "첫 구간이 안전한 자르기 지점보다 먼저 시작합니다. 시작 위치를 조금 뒤로 옮기세요.", "segments_smart_render_whole_video": "사용 가능한 안전한 자르기 지점이 남지 않아 별도 처리 시 전체 비디오를 다시 인코딩해야 합니다. 더 작은 영역을 선택하거나 모든 구간을 지워 전체 비디오를 처리하세요.", + "segments_mixed_models_too_close": "'표준' 구간과 LTX 구간이 너무 가깝습니다. 떨어뜨리거나 둘 다 같은 모델을 쓰세요.", "segments_discard_title": "구간 변경 사항을 버릴까요?", "segments_discard_changes": "복원 구간 변경 사항을 적용하지 않고 닫을까요?", "segments_edit_tooltip": "이 비디오에서 복원할 부분 선택", @@ -448,8 +453,8 @@ # App messages "error_cannot_start": "처리를 시작할 수 없습니다:", "error_start_title": "시작할 수 없음", + "error_ltx_ranges_unavailable": "다음 영상에 LTX 구간이 있지만 이 PC에서는 LTX를 실행할 수 없습니다: {files}", "error_post_export_command_required": "내보내기 후 사용자 정의 명령이 필요합니다", - "error_ltx_segments": "LTX는 영상 전체를 처리합니다. 구간을 지우거나 '표준'을 선택하세요: {files}", # Settings panel "dialog_select_tvai_ffmpeg": "Topaz Video ffmpeg.exe 선택", diff --git a/jasna/gui/locales/th.py b/jasna/gui/locales/th.py index d2e2e45e..d604d0b9 100644 --- a/jasna/gui/locales/th.py +++ b/jasna/gui/locales/th.py @@ -77,7 +77,7 @@ "segments_preview_reset_view": "รีเซ็ตมุมมอง", "segments_preview_reset_view_hint": "รีเซ็ตการซูมและการเลื่อนเพื่อแสดงทั้งเฟรม", "segments_restore_gpu_busy": "ใช้งานไม่ได้ขณะกำลังประมวลผล", - "segments_restore_ltx_unavailable": "ยังดูตัวอย่างกับ LTX ไม่ได้", + "segments_restore_ltx_unavailable": "ช่วง LTX ใช้ตัวอย่างนี้ไม่ได้ ให้ใช้ “ลองซีด” แทน", "segments_restore_loading_models": "กำลังโหลดโมเดลกู้คืน…", "segments_restore_restoring": "กำลังกู้คืน…", "segments_restore_failed": "ตัวอย่างการกู้คืนล้มเหลว: {message}", @@ -125,9 +125,13 @@ "segments_scan_model_changed": "เปลี่ยนโมเดลแล้ว — สแกนอีกครั้งเพื่ออัปเดตผลการตรวจจับ", "segments_timeline_title": "ไทม์ไลน์", "segments_timeline_hint": "มุมมองละเอียด · ภาพรวมวิดีโอทั้งหมดด้านล่าง", - "segments_legend_selected": "เลือกเพื่อกู้คืน", "segments_legend_detected": "ตรวจพบจากการสแกน", "segments_legend_playhead": "เฟรมปัจจุบัน", + "segments_restore_with": "กู้คืนด้วย", + "segments_use_for_all": "ใช้กับทุกช่วง", + "tip_segments_use_for_all": "ให้ทุกช่วงใช้โมเดลและซีดนี้", + "segments_ltx_needs_flat_video": "LTX กู้คืนวิดีโอ VR แบบซ้าย-ขวาไม่ได้", + "segments_ltx_ranges_unavailable": "PC เครื่องนี้ใช้ LTX ไม่ได้ จึงกู้คืนช่วง LTX ไม่ได้", "segments_suggest_mask": "เสนอมาสก์ที่ดีกว่า", "segments_suggest_mask_hint": "วาดโครงร่างโมเสคที่ถูกต้องบนเฟรมนี้และแชร์แบบไม่ระบุตัวตนเพื่อช่วยปรับปรุงโมเดลตรวจจับในอนาคต", "mask_editor_title": "เสนอมาสก์โมเสคที่ดีกว่า", @@ -169,7 +173,7 @@ "segments_time_out_of_bounds": "เวลาต้องอยู่ภายในความยาวของวิดีโอ", "segments_merged": "รวมกับช่วงที่ทับซ้อนหรือติดกัน {count} ช่วงแล้ว", "segments_drag_hint": "หากไม่เลือกช่วง ระบบจะกู้คืนวิดีโอทั้งหมด ลากบนไทม์ไลน์หรือใช้ I และ O เพื่อทำเครื่องหมายเฟรม", - "segments_range_row": "{index}. {start} → {end}\n {duration:.1f} วินาที", + "segments_range_row": "{index}. {start} → {end}\n {duration:.1f} วินาที · {model}", "segments_delete_range": "ลบช่วง", "segments_zoom_out": "ซูมออก (เลื่อนล้อเมาส์ลง)", "segments_zoom_in": "ซูมเข้า (เลื่อนล้อเมาส์ขึ้น)", @@ -183,6 +187,7 @@ "segments_smart_render_range_too_short": "ช่วงที่เลือกสั้นกว่าหนึ่งเฟรมวิดีโอ โปรดทำให้ยาวขึ้น", "segments_smart_render_before_first_keyframe": "ช่วงแรกเริ่มก่อนจุดตัดที่ปลอดภัยของวิดีโอ โปรดเลื่อนจุดเริ่มไปข้างหน้าเล็กน้อย", "segments_smart_render_whole_video": "ช่วงเหล่านี้ไม่เหลือจุดตัดที่ปลอดภัยให้ใช้ การประมวลผลแยกจะต้องเข้ารหัสวิดีโอทั้งไฟล์ใหม่ โปรดเลือกพื้นที่เล็กลงหรือล้างทุกช่วงเพื่อประมวลผลวิดีโอทั้งไฟล์", + "segments_mixed_models_too_close": "ช่วง “มาตรฐาน” กับช่วง LTX อยู่ใกล้กันเกินไป ให้แยกออกจากกันหรือใช้โมเดลเดียวกันทั้งสองช่วง", "segments_discard_title": "ละทิ้งการเปลี่ยนแปลงช่วงหรือไม่", "segments_discard_changes": "ปิดโดยไม่ใช้การเปลี่ยนแปลงช่วงการกู้คืนหรือไม่", "segments_edit_tooltip": "เลือกส่วนของวิดีโอนี้ที่จะกู้คืน", @@ -450,8 +455,8 @@ # App messages "error_cannot_start": "ไม่สามารถเริ่มประมวลผล:", "error_start_title": "เริ่มไม่ได้", + "error_ltx_ranges_unavailable": "วิดีโอเหล่านี้มีช่วง LTX แต่ PC เครื่องนี้ใช้ LTX ไม่ได้: {files}", "error_post_export_command_required": "ต้องระบุคำสั่งกำหนดเองหลังส่งออก", - "error_ltx_segments": "LTX ประมวลผลทั้งวิดีโอ ล้างช่วงที่เลือกหรือเลือก “มาตรฐาน”: {files}", # Settings panel "dialog_select_tvai_ffmpeg": "เลือก Topaz Video ffmpeg.exe", diff --git a/jasna/gui/locales/zh.py b/jasna/gui/locales/zh.py index 1eab5b95..158fb3fe 100644 --- a/jasna/gui/locales/zh.py +++ b/jasna/gui/locales/zh.py @@ -77,7 +77,7 @@ "segments_preview_reset_view": "重置视图", "segments_preview_reset_view_hint": "重置缩放和平移并显示完整画面", "segments_restore_gpu_busy": "处理进行中，暂不可用", - "segments_restore_ltx_unavailable": "LTX 暂不支持预览", + "segments_restore_ltx_unavailable": "LTX 片段不能使用此预览。请改用“试试种子”。", "segments_restore_loading_models": "正在加载修复模型…", "segments_restore_restoring": "正在修复…", "segments_restore_failed": "修复预览失败：{message}", @@ -125,9 +125,13 @@ "segments_scan_model_changed": "模型已更改 — 请重新扫描以更新检测结果", "segments_timeline_title": "时间线", "segments_timeline_hint": "详细视图 · 下方为完整视频概览", - "segments_legend_selected": "已选择修复", "segments_legend_detected": "扫描检测结果", "segments_legend_playhead": "当前帧", + "segments_restore_with": "修复方式", + "segments_use_for_all": "用于所有片段", + "tip_segments_use_for_all": "让所有片段都使用这个模型和种子。", + "segments_ltx_needs_flat_video": "LTX 无法修复 VR 左右并排视频。", + "segments_ltx_ranges_unavailable": "这台电脑无法运行 LTX，所以 LTX 片段无法修复。", "segments_suggest_mask": "提交更好的遮罩", "segments_suggest_mask_hint": "在此帧上勾画正确的马赛克轮廓并匿名分享，帮助改进未来的检测模型", "mask_editor_title": "提交更好的马赛克遮罩", @@ -169,7 +173,7 @@ "segments_time_out_of_bounds": "时间必须在视频时长范围内。", "segments_merged": "已与 {count} 个重叠或相邻区间合并。", "segments_drag_hint": "未选择区间时将修复完整视频。在时间线上拖动，或使用 I 和 O 标记精确帧。", - "segments_range_row": "{index}. {start} → {end}\n {duration:.1f} 秒", + "segments_range_row": "{index}. {start} → {end}\n {duration:.1f} 秒 · {model}", "segments_delete_range": "删除区间", "segments_zoom_out": "缩小（鼠标滚轮向下）", "segments_zoom_in": "放大（鼠标滚轮向上）", @@ -183,6 +187,7 @@ "segments_smart_render_range_too_short": "所选区间短于一个视频帧，请将其延长。", "segments_smart_render_before_first_keyframe": "第一个区间开始于视频的安全切点之前，请将开始时间稍微后移。", "segments_smart_render_whole_video": "这些区间没有留下可用的安全切点，单独处理会导致整个视频重新编码。请选择更小的区域，或清除所有区间以处理完整视频。", + "segments_mixed_models_too_close": "“标准”片段和 LTX 片段离得太近。请把它们分开，或两者使用同一模型。", "segments_discard_title": "放弃区间更改？", "segments_discard_changes": "不应用修复区间更改并关闭吗？", "segments_edit_tooltip": "选择此视频中需要修复的部分", @@ -468,8 +473,8 @@ # App messages "error_cannot_start": "无法开始处理：", "error_start_title": "无法开始", + "error_ltx_ranges_unavailable": "这些视频含有 LTX 片段，但这台电脑无法运行 LTX：{files}", "error_post_export_command_required": "需要填写导出后的自定义命令", - "error_ltx_segments": "LTX 会处理整个视频。请清除片段或选择“标准”：{files}", # Settings panel "dialog_select_tvai_ffmpeg": "选择 Topaz Video ffmpeg.exe", diff --git a/jasna/gui/queue_panel.py b/jasna/gui/queue_panel.py index 1785f005..7b2794de 100644 --- a/jasna/gui/queue_panel.py +++ b/jasna/gui/queue_panel.py @@ -506,6 +506,7 @@ def _edit_segments(self, job: JobItem) -> None: self._get_settings, self._is_gpu_busy, self._set_preview_gpu_busy, + self._ltx_unavailable_reason(), lambda segments, j=job: self._segments_saved(j, segments), self._segment_editor_closed, ) @@ -595,10 +596,12 @@ def set_segment_editor_context( get_settings: callable, is_gpu_busy: callable, set_preview_gpu_busy: callable, + ltx_unavailable_reason: callable, ): self._get_settings = get_settings self._is_gpu_busy = is_gpu_busy self._set_preview_gpu_busy = set_preview_gpu_busy + self._ltx_unavailable_reason = ltx_unavailable_reason def set_on_output_changed(self, callback: callable): self._on_output_changed = callback diff --git a/jasna/gui/restoration_preview.py b/jasna/gui/restoration_preview.py index b0f3d37a..a8350739 100644 --- a/jasna/gui/restoration_preview.py +++ b/jasna/gui/restoration_preview.py @@ -15,7 +15,6 @@ from PIL import Image -from jasna.gui.locales import t from jasna.gui.queues import replace_pending from jasna.gui.models import AppSettings from jasna.gui.video_session import build_video_session, release_session_memory, video_session_key @@ -326,8 +325,6 @@ def _run(self) -> None: if isinstance(command, _Cancel): continue try: - if command.settings.restoration_model == "ltx": - raise ValueError(t("segments_restore_ltx_unavailable")) key = video_session_key(command.settings) if session is None or key != session_key: pipeline = None diff --git a/jasna/gui/segment_editor.py b/jasna/gui/segment_editor.py index 969ff229..f44bde6d 100644 --- a/jasna/gui/segment_editor.py +++ b/jasna/gui/segment_editor.py @@ -2,9 +2,11 @@ import logging import queue +import random import threading import tkinter as tk from collections.abc import Callable +from dataclasses import replace import customtkinter as ctk from PIL import Image @@ -14,7 +16,7 @@ from jasna.gui.locales import t from jasna.gui.models import AppSettings, JobItem from jasna.gui.components import Tooltip, grab_modal -from jasna.gui.icons import CompactSwitch +from jasna.gui.icons import CompactSwitch, NativeIconButton from jasna.gui.restoration_preview import ( RestorationClip, RestorationFailed, @@ -50,10 +52,17 @@ from jasna.gui.settings_sections.encoding import CODEC_CANONICAL_TO_LABEL from jasna.gui.settings_sections.widgets import ValueOptionMenu from jasna.gui.segment_timeline import SegmentTimeline +from jasna.gui.settings_sections.restoration_model import parse_seed from jasna.gui.theme import Colors, Fonts, Sizing from jasna.media.probe import VideoMetadata from jasna.media.splice import canonical_codec -from jasna.segments import SegmentRange, format_timestamp, parse_timestamp +from jasna.segments import ( + SegmentRange, + SegmentRestoration, + format_timestamp, + job_restoration, + parse_timestamp, +) logger = logging.getLogger(__name__) @@ -68,6 +77,7 @@ def __init__( get_settings: Callable[[], AppSettings], is_gpu_busy: Callable[[], bool], set_preview_gpu_busy: Callable[[bool], None], + ltx_unavailable_reason: str | None, on_saved: Callable[[tuple[SegmentRange, ...]], None], on_closed: Callable[[], None] | None = None, ) -> None: @@ -76,6 +86,7 @@ def __init__( self._get_settings = get_settings self._is_gpu_busy = is_gpu_busy self._set_preview_gpu_busy = set_preview_gpu_busy + self._ltx_unavailable_reason = ltx_unavailable_reason self._on_saved = on_saved self._on_closed = on_closed self._state: SegmentEditorState | None = None @@ -184,10 +195,13 @@ def _build_editor(self, metadata: VideoMetadata) -> None: self._loading_bar.stop() self._loading.destroy() self._metadata = metadata + settings = self._get_settings() + self._job_ltx_seed = int(settings.ltx_seed) self._state = SegmentEditorState( duration=float(metadata.duration), fps=max(1.0, float(metadata.video_fps)), segments=self._job.snapshot_segments(), + default_restoration=job_restoration(settings.restoration_model, settings.ltx_seed), ) self._job.duration_seconds = self._state.duration @@ -462,7 +476,6 @@ def _build_range_panel(self, body: ctk.CTkFrame) -> None: hover_color=Colors.BORDER_LIGHT, command=self._undo, ) - self._undo_btn.pack(side="right") Tooltip(self._undo_btn, t("segments_undo")) self._redo_btn = ctk.CTkButton( range_header, @@ -473,7 +486,8 @@ def _build_range_panel(self, body: ctk.CTkFrame) -> None: hover_color=Colors.BORDER_LIGHT, command=self._redo, ) - self._redo_btn.pack(side="right", padx=4) + self._redo_btn.pack(side="right") + self._undo_btn.pack(side="right", padx=4) Tooltip(self._redo_btn, t("segments_redo")) self._segment_list = ctk.CTkScrollableFrame( @@ -541,6 +555,147 @@ def _build_range_panel(self, body: ctk.CTkFrame) -> None: command=self._add_or_update, ) self._range_action.grid(row=2, column=0, columnspan=3, sticky="ew", pady=(5, 0)) + self._build_restoration_inspector(range_panel) + + def _build_restoration_inspector(self, range_panel: ctk.CTkFrame) -> None: + self._inspector = ctk.CTkFrame(range_panel, fg_color=Colors.BG_PANEL, corner_radius=Sizing.BORDER_RADIUS) + self._inspector.grid_columnconfigure(1, weight=1) + ctk.CTkLabel( + self._inspector, + text=t("segments_restore_with"), + text_color=Colors.TEXT_PRIMARY, + font=(Fonts.FAMILY, Fonts.SIZE_SMALL), + ).grid(row=0, column=0, sticky="w", padx=(8, 6), pady=(6, 3)) + models = ctk.CTkFrame(self._inspector, fg_color="transparent") + models.grid(row=0, column=1, columnspan=2, sticky="w", pady=(6, 3)) + self._range_model = ctk.StringVar(value="basicvsrpp") + self._model_radios = {} + for model in ("basicvsrpp", "ltx"): + radio = ctk.CTkRadioButton( + models, + text=t(f"model_{model}"), + variable=self._range_model, + value=model, + command=self._on_range_model_changed, + font=(Fonts.FAMILY, Fonts.SIZE_SMALL), + fg_color=Colors.MODEL_LTX if model == "ltx" else Colors.PRIMARY, + text_color=Colors.TEXT_PRIMARY, + radiobutton_width=16, + radiobutton_height=16, + ) + radio.pack(side="left", padx=(0, 10)) + self._model_radios[model] = radio + self._ltx_radio_tooltip = Tooltip(self._model_radios["ltx"], "") + self._seed_label = ctk.CTkLabel( + self._inspector, + text=t("ltx_seed"), + text_color=Colors.TEXT_PRIMARY, + font=(Fonts.FAMILY, Fonts.SIZE_SMALL), + ) + self._seed_entry = ctk.CTkEntry(self._inspector, font=(Fonts.FAMILY_MONO, Fonts.SIZE_SMALL)) + self._seed_entry.bind("", self._commit_seed_entry) + self._seed_entry.bind("", self._commit_seed_entry) + Tooltip(self._seed_entry, t("tip_ltx_seed")) + self._new_seed_btn = NativeIconButton( + self._inspector, + "reset", + 14, + Colors.TEXT_PRIMARY, + Colors.BG_PANEL, + Colors.BORDER_LIGHT, + Colors.BORDER_LIGHT, + self._new_range_seed, + 28, + 28, + ) + Tooltip(self._new_seed_btn, t("tip_ltx_new_seed")) + self._use_for_all_btn = ctk.CTkButton( + self._inspector, + text=t("segments_use_for_all"), + height=26, + fg_color=Colors.BG_CARD, + hover_color=Colors.BORDER_LIGHT, + font=(Fonts.FAMILY, Fonts.SIZE_SMALL), + command=self._use_restoration_for_all, + ) + self._use_for_all_btn.grid(row=2, column=0, columnspan=3, sticky="ew", padx=8, pady=(3, 8)) + Tooltip(self._use_for_all_btn, t("tip_segments_use_for_all")) + + def _ltx_blocked_reason(self) -> str | None: + """Locale key saying why a range cannot use LTX here, or None when it can.""" + if self._ltx_unavailable_reason is not None: + return self._ltx_unavailable_reason + if self.vr_resolution.is_sbs: + return "segments_ltx_needs_flat_video" + return None + + def _selected_restoration(self) -> SegmentRestoration: + model = self._range_model.get() + if model != "ltx": + return SegmentRestoration(model, None) + selected = self._require_state().selected_segment + current = None if selected is None else selected.restoration.ltx_seed + seed = parse_seed(self._seed_entry.get()) if self._seed_entry.get().strip() else current + return SegmentRestoration("ltx", self._job_ltx_seed if seed is None else seed) + + def _apply_selected_restoration(self, restoration: SegmentRestoration) -> None: + if self._require_state().set_selected_restoration(restoration): + self._edit_notice = None + self._refresh_all() + else: + self._refresh_inspector() + + def _on_range_model_changed(self) -> None: + self._seed_entry.delete(0, "end") + self._apply_selected_restoration(self._selected_restoration()) + + def _commit_seed_entry(self, _event=None) -> None: + if self._range_model.get() == "ltx" and self._require_state().selected_segment is not None: + self._apply_selected_restoration(self._selected_restoration()) + + def _new_range_seed(self) -> None: + self._apply_selected_restoration(SegmentRestoration("ltx", random.randrange(1, 2**31))) + + def _use_restoration_for_all(self) -> None: + selected = self._require_state().selected_segment + if selected is not None and self._require_state().set_all_restoration(selected.restoration): + self._edit_notice = None + self._refresh_all() + + def _refresh_inspector(self) -> None: + selected = self._require_state().selected_segment + if selected is None: + self._inspector.grid_forget() + return + self._inspector.grid(row=4, column=0, sticky="ew", padx=8, pady=(0, 8)) + restoration = selected.restoration + self._range_model.set(restoration.model) + blocked = self._ltx_blocked_reason() + locked = self._scanning() + self._model_radios["basicvsrpp"].configure(state="disabled" if locked else "normal") + self._model_radios["ltx"].configure( + state="disabled" if locked or (blocked is not None and restoration.model != "ltx") else "normal" + ) + self._ltx_radio_tooltip.set_text(t(blocked) if blocked else t("model_ltx_description")) + if restoration.model == "ltx": + self._seed_label.grid(row=1, column=0, sticky="w", padx=(8, 6), pady=3) + self._seed_entry.grid(row=1, column=1, sticky="ew", pady=3) + self._new_seed_btn.grid(row=1, column=2, padx=(4, 8), pady=3) + self._seed_entry.delete(0, "end") + self._seed_entry.insert(0, str(restoration.ltx_seed)) + else: + for widget in (self._seed_label, self._seed_entry, self._new_seed_btn): + widget.grid_forget() + state = "disabled" if locked else "normal" + for widget in (self._seed_entry, self._new_seed_btn, self._use_for_all_btn): + widget.configure(state=state) + + @staticmethod + def _model_text(segment: SegmentRange) -> str: + restoration = segment.restoration + if restoration.model == "ltx": + return f"{t('model_ltx')} {restoration.ltx_seed}" + return t(f"model_{restoration.model}") def _build_scan_and_timeline(self, metadata: VideoMetadata) -> None: self._timeline = SegmentTimeline( @@ -617,7 +772,8 @@ def _build_scan_and_timeline(self, metadata: VideoMetadata) -> None: legend = ctk.CTkFrame(self, fg_color="transparent") legend.pack(fill="x", padx=20, pady=(0, 2)) for color, label in ( - (Colors.PRIMARY, t("segments_legend_selected")), + (Colors.PRIMARY, t("model_basicvsrpp")), + (Colors.MODEL_LTX, t("model_ltx")), (Colors.STATUS_WARNING, t("segments_legend_detected")), ("#f8fafc", t("segments_legend_playhead")), ): @@ -994,10 +1150,6 @@ def _toggle_restoration_preview(self) -> None: if self._is_gpu_busy(): self._restore_toggle.deselect() return - if self._get_settings().restoration_model == "ltx": - self._restore_toggle.deselect() - self._show_preview_message(t("segments_restore_ltx_unavailable"), Colors.STATUS_WARNING) - return self._restore_active = True self._set_playing(False) self._restore_toggle.select() @@ -1045,16 +1197,26 @@ def _schedule_restoration_preview(self) -> None: self.after_cancel(self._restore_after) self._restore_after = self.after(400, self._request_restoration_preview) + def _playhead_uses_ltx(self) -> bool: + return self._require_state().restoration_at(self._current).model == "ltx" + def _request_restoration_preview(self) -> None: self._restore_after = None if not self._restore_active or self._closed.is_set(): return + if self._playhead_uses_ltx(): + self._restore_play_pending = False + self._restored_clip = () + self._restored_source = None + self._play.configure(text="▶") + self._show_preview_message(t("segments_restore_ltx_unavailable"), Colors.STATUS_WARNING) + return self._restore_play_pending = False self._restored_clip = () self._play.configure(text="▶") self._restore_generation = self._restoration_worker.request( self._current, - self._current_video_settings(), + self._standard_preview_settings(), projection=self._vr_projection, ) if self._restored_source is None: @@ -1066,6 +1228,9 @@ def _request_restoration_preview(self) -> None: def _request_restoration_playback(self, start_seconds: float) -> None: if not self._restore_active or self._closed.is_set(): return + if self._require_state().restoration_at(start_seconds).model == "ltx": + self._show_preview_message(t("segments_restore_ltx_unavailable"), Colors.STATUS_WARNING) + return if self._restore_after is not None: self.after_cancel(self._restore_after) self._restore_after = None @@ -1080,7 +1245,7 @@ def _request_restoration_playback(self, start_seconds: float) -> None: ) self._restore_generation = self._restoration_worker.request( start_seconds, - self._current_video_settings(), + self._standard_preview_settings(), projection=self._vr_projection, playback=True, ) @@ -1190,6 +1355,10 @@ def _lockable_widgets(self) -> tuple: self._start_entry, self._end_entry, self._suggest_btn, + *self._model_radios.values(), + self._seed_entry, + self._new_seed_btn, + self._use_for_all_btn, *self._timeline_zoom_buttons, *self._segment_action_widgets, ) @@ -1219,6 +1388,9 @@ def _scanning(self) -> bool: def _current_video_settings(self) -> AppSettings: return self._scan_panel.video_settings() + def _standard_preview_settings(self) -> AppSettings: + return replace(self._current_video_settings(), restoration_model="basicvsrpp") + def _suggest_mask(self) -> None: self._require_state() if self._suggest_busy: @@ -1403,6 +1575,7 @@ def _render_segment_list(self) -> None: start=format_timestamp(segment.start), end=format_timestamp(segment.end), duration=segment.duration, + model=self._model_text(segment), ), anchor="w", fg_color="transparent", @@ -1450,6 +1623,7 @@ def _refresh_all(self) -> None: state="normal" if state.can_redo and not self._scanning() else "disabled" ) self._render_segment_list() + self._refresh_inspector() self._refresh_timeline() self._refresh_workload() self._refresh_notice() @@ -1519,6 +1693,10 @@ def _refresh_notice(self) -> None: text=self._edit_notice, text_color=Colors.STATUS_WARNING if self._edit_notice_warning else Colors.STATUS_ERROR, ) + elif self._ltx_blocked_reason() is not None and any( + segment.restoration.model == "ltx" for segment in self._state.segments + ): + self._notice.configure(text=t("segments_ltx_ranges_unavailable"), text_color=Colors.STATUS_WARNING) else: self._notice.configure(text="") diff --git a/jasna/gui/segment_editor_state.py b/jasna/gui/segment_editor_state.py index 00160b54..93afb162 100644 --- a/jasna/gui/segment_editor_state.py +++ b/jasna/gui/segment_editor_state.py @@ -1,9 +1,9 @@ from __future__ import annotations -from dataclasses import dataclass +from dataclasses import dataclass, replace from jasna.media.splice import KeyframeIndex, SmartRenderCompatibilityError, build_splice_plan -from jasna.segments import SegmentRange, normalize_segments +from jasna.segments import SegmentRange, SegmentRestoration, normalize_segments, resolve_restorations @dataclass(frozen=True) @@ -21,14 +21,16 @@ class _Snapshot: class SegmentEditorState: - """UI-independent state and history for the segment editor.""" + """UI-independent state and history for the segment editor. Every range carries its + own restoration; new ranges get ``default_restoration``.""" def __init__( self, *, duration: float, fps: float, - segments: tuple[SegmentRange, ...] = (), + segments: tuple[SegmentRange, ...], + default_restoration: SegmentRestoration, ) -> None: duration = float(duration) fps = float(fps) @@ -37,7 +39,8 @@ def __init__( if fps <= 0: raise ValueError("video frame rate must be greater than zero") - normalized = normalize_segments(segments, duration=duration) + normalized = normalize_segments(resolve_restorations(segments, default_restoration), duration=duration) + self.default_restoration = default_restoration self.duration = duration self.fps = fps self.segments = normalized @@ -105,7 +108,7 @@ def clear_marks(self) -> None: def add(self, start: float, end: float) -> SegmentEditResult: start, end = self._validated_bounds(start, end) - segment = SegmentRange(start, end) + segment = SegmentRange(start, end, self.default_restoration) overlaps = sum( existing.start <= segment.end and segment.start <= existing.end for existing in self.segments @@ -120,8 +123,8 @@ def replace_selected(self, start: float, end: float) -> SegmentEditResult: if self.selected_index is None: return self.add(start, end) start, end = self._validated_bounds(start, end) - replacement = SegmentRange(start, end) original_index = self.selected_index + replacement = replace(self.segments[original_index], start=start, end=end) remaining = [ segment for index, segment in enumerate(self.segments) if index != original_index @@ -145,7 +148,7 @@ def add_many(self, ranges: tuple[SegmentRange, ...]) -> int: """Add several ranges as one undoable step; returns how many were new.""" fresh = [ - candidate + replace(candidate, restoration=self.default_restoration) for candidate in ranges if not any( existing.start <= candidate.start and existing.end >= candidate.end @@ -162,6 +165,38 @@ def add_many(self, ranges: tuple[SegmentRange, ...]) -> int: self.clear_marks() return len(fresh) + def set_selected_restoration(self, restoration: SegmentRestoration) -> bool: + """Give the selected range ``restoration``; returns whether anything changed.""" + selected = self.selected_segment + if selected is None or selected.restoration == restoration: + return False + changed = replace(selected, restoration=restoration) + self._record_change() + others = tuple(segment for index, segment in enumerate(self.segments) if index != self.selected_index) + self.segments = normalize_segments((*others, changed), duration=self.duration) + self.selected_index = self._index_containing(changed) + return True + + def set_all_restoration(self, restoration: SegmentRestoration) -> bool: + """Give every range ``restoration``; returns whether anything changed.""" + if all(segment.restoration == restoration for segment in self.segments): + return False + selected = self.selected_segment + self._record_change() + self.segments = normalize_segments( + tuple(replace(segment, restoration=restoration) for segment in self.segments), + duration=self.duration, + ) + self.selected_index = None if selected is None else self._index_containing(selected) + return True + + def restoration_at(self, seconds: float) -> SegmentRestoration: + """The restoration of the range holding ``seconds``, else the default.""" + for segment in self.segments: + if segment.start <= seconds < segment.end: + return segment.restoration + return self.default_restoration + def delete_selected(self) -> bool: if self.selected_index is None: return False @@ -302,6 +337,7 @@ def crop_box(self, width: int, height: int) -> tuple[int, int, int, int] | None: "range_too_short": "segments_smart_render_range_too_short", "before_first_keyframe": "segments_smart_render_before_first_keyframe", "whole_video_reencode": "segments_smart_render_whole_video", + "mixed_models": "segments_mixed_models_too_close", } diff --git a/jasna/gui/segment_timeline.py b/jasna/gui/segment_timeline.py index ca7bf999..0e7ee4cd 100644 --- a/jasna/gui/segment_timeline.py +++ b/jasna/gui/segment_timeline.py @@ -2,6 +2,7 @@ import tkinter as tk from collections.abc import Callable +from dataclasses import replace import customtkinter as ctk @@ -17,6 +18,13 @@ _HANDLE_HIT_PX = 8 +def segment_fill(segment: SegmentRange) -> str: + """Timeline colour of a range: its restoration model.""" + if segment.restoration is not None and segment.restoration.model == "ltx": + return Colors.MODEL_LTX + return Colors.PRIMARY + + def timeline_seconds_to_x( seconds: float, *, @@ -171,7 +179,7 @@ def _redraw_timeline(self) -> None: width, top=_MAIN_TOP + 7, bottom=_MAIN_BOTTOM - 7, - fill=Colors.PRIMARY, + fill=segment_fill(segment), outline="#a5b4fc" if index == self._selected_index else "", outline_width=2 if index == self._selected_index else 0, ) @@ -257,7 +265,7 @@ def _draw_overview(self, width: int) -> None: _OVERVIEW_TOP + 2, max(x1 + 1, x2), _OVERVIEW_BOTTOM - 2, - fill=Colors.PRIMARY, + fill=segment_fill(segment), outline="", ) for run in self._detections: @@ -359,13 +367,10 @@ def _on_motion(self, event) -> None: elif self._drag_segment is not None: if self._drag_kind == "start": start = min(seconds, self._drag_segment.end - 1 / self.fps) - self._drag_segment = SegmentRange(max(0.0, start), self._drag_segment.end) + self._drag_segment = replace(self._drag_segment, start=max(0.0, start)) else: end = max(seconds, self._drag_segment.start + 1 / self.fps) - self._drag_segment = SegmentRange( - self._drag_segment.start, - min(self.duration, end), - ) + self._drag_segment = replace(self._drag_segment, end=min(self.duration, end)) self._redraw_timeline() def _on_release(self, event) -> None: diff --git a/jasna/gui/settings_panel.py b/jasna/gui/settings_panel.py index b56ddce3..ba94a171 100644 --- a/jasna/gui/settings_panel.py +++ b/jasna/gui/settings_panel.py @@ -248,6 +248,10 @@ def _on_restoration_model_changed(self, model: str): for section in self._model_aware_sections: section.set_model(model) + def ltx_unavailable_reason(self) -> str | None: + """Locale key saying why LTX cannot run on this PC, or None when it can.""" + return self._model_section.unavailable_reason() + def set_gpu_support(self, *, nvidia: bool, blackwell: bool): self._model_section.set_gpu_support(nvidia=nvidia, blackwell=blackwell) self._update_modified_indicator() diff --git a/jasna/gui/settings_sections/restoration_model.py b/jasna/gui/settings_sections/restoration_model.py index 60f8665e..cef77e0a 100644 --- a/jasna/gui/settings_sections/restoration_model.py +++ b/jasna/gui/settings_sections/restoration_model.py @@ -161,8 +161,11 @@ def _set_seed(self, seed: int) -> None: self._widgets["ltx_seed"].delete(0, "end") self._widgets["ltx_seed"].insert(0, str(seed)) + def unavailable_reason(self) -> str | None: + return ltx_unavailable_reason(installed=self._ltx_installed, nvidia=self._nvidia) + def _refresh_ltx_availability(self) -> None: - reason = ltx_unavailable_reason(installed=self._ltx_installed, nvidia=self._nvidia) + reason = self.unavailable_reason() card = self._cards["ltx"] card.set_enabled(self._enabled and reason is None) card.description.configure( diff --git a/jasna/gui/theme.py b/jasna/gui/theme.py index b4db9fac..e4dd5222 100644 --- a/jasna/gui/theme.py +++ b/jasna/gui/theme.py @@ -36,6 +36,9 @@ class Colors: PRIMARY_HOVER = "#4338ca"# Indigo-700 PRIMARY_DARK = "#3730a3" # Indigo-800 + # LTX ranges in the segment editor + MODEL_LTX = "#0d9488" # Teal-600 + # Video player call to action PLAYER = "#0284c7" # Sky-600 PLAYER_HOVER = "#0369a1" # Sky-700 diff --git a/jasna/gui/validation.py b/jasna/gui/validation.py index 24aad60e..0122fd39 100644 --- a/jasna/gui/validation.py +++ b/jasna/gui/validation.py @@ -3,9 +3,18 @@ from jasna.gui.models import AppSettings, JobItem from jasna.gui.locales import t +from jasna.segments import job_restoration, resolve_restorations -def validate_gui_start(settings: AppSettings, jobs: list[JobItem]) -> list[str]: +def _uses_ltx(job: JobItem, settings: AppSettings) -> bool: + segments = job.snapshot_segments() + if not segments: + return settings.restoration_model == "ltx" + default = job_restoration(settings.restoration_model, settings.ltx_seed) + return any(segment.restoration.model == "ltx" for segment in resolve_restorations(segments, default)) + + +def validate_gui_start(settings: AppSettings, jobs: list[JobItem], *, ltx_available: bool) -> list[str]: errors: list[str] = [] from jasna.post_export_action import validate_post_export_action @@ -14,10 +23,12 @@ def validate_gui_start(settings: AppSettings, jobs: list[JobItem]) -> list[str]: except ValueError: errors.append(t("error_post_export_command_required")) + if not ltx_available: + needs_ltx = [job.filename for job in jobs if _uses_ltx(job, settings)] + if needs_ltx: + errors.append(t("error_ltx_ranges_unavailable", files=", ".join(needs_ltx))) + if settings.restoration_model == "ltx": - segmented = [job.filename for job in jobs if job.snapshot_segments()] - if segmented: - errors.append(t("error_ltx_segments", files=", ".join(segmented))) return errors if settings.secondary_restoration != "tvai": diff --git a/jasna/gui/video_session.py b/jasna/gui/video_session.py index c53f8f90..d4b1815d 100644 --- a/jasna/gui/video_session.py +++ b/jasna/gui/video_session.py @@ -29,9 +29,8 @@ def video_session_key(settings: AppSettings) -> tuple: settings.denoise_step, settings.secondary_restoration, settings.restoration_model, + settings.ltx_fast, ) - if settings.restoration_model == "ltx": - key += (settings.ltx_fast,) if settings.secondary_restoration == "tvai": key += ( settings.tvai_ffmpeg_path, diff --git a/tests/test_gui_tvai_validation.py b/tests/test_gui_tvai_validation.py index 219a5adb..97644e23 100644 --- a/tests/test_gui_tvai_validation.py +++ b/tests/test_gui_tvai_validation.py @@ -9,12 +9,12 @@ def test_validate_gui_start_non_tvai_returns_empty() -> None: settings = AppSettings(secondary_restoration="none") - assert validate_gui_start(settings, []) == [] + assert validate_gui_start(settings, [], ltx_available=True) == [] def test_validate_gui_start_custom_post_export_requires_command() -> None: settings = AppSettings(secondary_restoration="none", post_export_action="command", post_export_command="") - assert t("error_post_export_command_required") in validate_gui_start(settings, []) + assert t("error_post_export_command_required") in validate_gui_start(settings, [], ltx_available=True) def test_validate_gui_start_tvai_missing_env_vars(tmp_path: Path, monkeypatch) -> None: @@ -25,7 +25,7 @@ def test_validate_gui_start_tvai_missing_env_vars(tmp_path: Path, monkeypatch) - ffmpeg.write_bytes(b"") settings = AppSettings(secondary_restoration="tvai", tvai_ffmpeg_path=str(ffmpeg)) - errors = validate_gui_start(settings, []) + errors = validate_gui_start(settings, [], ltx_available=True) assert t("error_tvai_data_dir_not_set") in errors assert t("error_tvai_model_dir_not_set") in errors @@ -38,7 +38,7 @@ def test_validate_gui_start_tvai_env_dirs_must_exist(tmp_path: Path, monkeypatch ffmpeg.write_bytes(b"") settings = AppSettings(secondary_restoration="tvai", tvai_ffmpeg_path=str(ffmpeg)) - errors = validate_gui_start(settings, []) + errors = validate_gui_start(settings, [], ltx_available=True) assert len(errors) == 2 assert str(tmp_path / "missing_data") in errors[0] assert str(tmp_path / "missing_model") in errors[1] @@ -53,7 +53,7 @@ def test_validate_gui_start_tvai_ffmpeg_path_must_exist(tmp_path: Path, monkeypa monkeypatch.setenv("TVAI_MODEL_DIR", str(model_dir)) settings = AppSettings(secondary_restoration="tvai", tvai_ffmpeg_path=str(tmp_path / "missing_ffmpeg.exe")) - errors = validate_gui_start(settings, []) + errors = validate_gui_start(settings, [], ltx_available=True) assert len(errors) == 1 assert str(tmp_path / "missing_ffmpeg.exe") in errors[0] @@ -70,17 +70,22 @@ def test_validate_gui_start_tvai_ok(tmp_path: Path, monkeypatch) -> None: ffmpeg.write_bytes(b"") settings = AppSettings(secondary_restoration="tvai", tvai_ffmpeg_path=str(ffmpeg)) - assert validate_gui_start(settings, []) == [] + assert validate_gui_start(settings, [], ltx_available=True) == [] -def test_validate_gui_start_blocks_ltx_with_segmented_jobs() -> None: +def test_validate_gui_start_blocks_ltx_ranges_only_when_ltx_cannot_run() -> None: from jasna.gui.models import JobItem - from jasna.segments import SegmentRange + from jasna.segments import SegmentRange, SegmentRestoration whole = JobItem(path=Path("whole.mp4")) - cut = JobItem(path=Path("cut.mp4"), segments=(SegmentRange(1.0, 2.0),)) + standard = JobItem(path=Path("standard.mp4"), segments=(SegmentRange(1.0, 2.0),)) + ltx_range = JobItem(path=Path("cut.mp4"), segments=(SegmentRange(1.0, 2.0, SegmentRestoration("ltx", 3)),)) ltx = AppSettings(restoration_model="ltx", secondary_restoration="tvai") - assert validate_gui_start(ltx, [whole]) == [] - assert validate_gui_start(ltx, [whole, cut]) == [t("error_ltx_segments", files="cut.mp4")] - assert validate_gui_start(AppSettings(), [cut]) == [] + assert validate_gui_start(ltx, [whole, standard, ltx_range], ltx_available=True) == [] + assert validate_gui_start(AppSettings(), [whole, standard, ltx_range], ltx_available=False) == [ + t("error_ltx_ranges_unavailable", files="cut.mp4") + ] + assert validate_gui_start(ltx, [whole, standard], ltx_available=False) == [ + t("error_ltx_ranges_unavailable", files="whole.mp4, standard.mp4") + ] diff --git a/tests/test_locales.py b/tests/test_locales.py index a4371406..efc22116 100644 --- a/tests/test_locales.py +++ b/tests/test_locales.py @@ -192,7 +192,9 @@ def test_format_placeholders_match(lang: str) -> None: "segments_scan_model_changed", "segments_timeline_title", "segments_timeline_hint", - "segments_legend_selected", + "segments_mixed_models_too_close", + "segments_restore_with", + "segments_use_for_all", "segments_legend_detected", "segments_legend_playhead", } diff --git a/tests/test_restoration_preview.py b/tests/test_restoration_preview.py index 7f735183..e4b180c2 100644 --- a/tests/test_restoration_preview.py +++ b/tests/test_restoration_preview.py @@ -477,34 +477,34 @@ def test_app_prepares_entire_queue_before_starting_processor() -> None: assert calls == ["reset", "running", "start"] -def test_restoration_worker_refuses_ltx_without_loading_models(monkeypatch) -> None: - build = MagicMock() - monkeypatch.setattr(restoration_preview, "build_video_session", build) - worker = RestorationPreviewWorker("unused.mp4", _metadata()) - worker.start() - generation = worker.request(1.0, AppSettings(restoration_model="ltx"), projection="auto") - - event = worker.events.get(timeout=2) - - assert event == restoration_preview.RestorationFailed( - restoration_preview.t("segments_restore_ltx_unavailable"), generation - ) - build.assert_not_called() - worker.close() - worker.join(timeout=2) +def test_restoration_preview_refuses_ltx_ranges_and_previews_others_with_standard() -> None: + from jasna.gui.segment_editor_state import SegmentEditorState + from jasna.segments import SegmentRange, SegmentRestoration - -def test_segment_editor_explains_that_ltx_has_no_preview() -> None: + standard = SegmentRestoration("basicvsrpp", None) editor = SegmentEditor.__new__(SegmentEditor) - editor._state = SimpleNamespace() - editor._restore_active = False - editor._is_gpu_busy = lambda: False - editor._get_settings = lambda: AppSettings(restoration_model="ltx") - editor._restore_toggle = MagicMock() + editor._state = SegmentEditorState( + duration=10, + fps=10, + segments=(SegmentRange(1, 2, SegmentRestoration("ltx", 5)),), + default_restoration=standard, + ) + editor._restore_active = True + editor._closed = threading.Event() + editor._restoration_worker = MagicMock() + editor._restored_source = None + editor._play = MagicMock() + editor._vr_projection = "auto" editor._show_preview_message = MagicMock() + editor._scan_panel = MagicMock() + editor._scan_panel.video_settings.return_value = AppSettings(restoration_model="ltx") - editor._toggle_restoration_preview() - - editor._restore_toggle.deselect.assert_called_once_with() + editor._current = 1.5 + editor._request_restoration_preview() + editor._restoration_worker.request.assert_not_called() editor._show_preview_message.assert_called_once() - assert editor._restore_active is False + + editor._current = 5.0 + editor._request_restoration_preview() + settings = editor._restoration_worker.request.call_args[0][1] + assert settings.restoration_model == "basicvsrpp" diff --git a/tests/test_segment_editor.py b/tests/test_segment_editor.py index 0177e585..6816a997 100644 --- a/tests/test_segment_editor.py +++ b/tests/test_segment_editor.py @@ -20,11 +20,12 @@ from jasna.gui.segment_editor import SegmentEditor from jasna.gui.segment_editor_state import PREVIEW_ZOOM_MAX, PreviewView, SegmentEditorState from jasna.media.splice import KeyframeIndex +from jasna.segments import SegmentRange, SegmentRestoration def test_out_of_bounds_range_uses_specific_message() -> None: editor = object.__new__(SegmentEditor) - editor._state = SegmentEditorState(duration=10, fps=30) + editor._state = SegmentEditorState(duration=10, fps=30, segments=(), default_restoration=SegmentRestoration("basicvsrpp", None)) editor._start_entry = MagicMock() editor._start_entry.get.return_value = "9" editor._end_entry = MagicMock() @@ -59,6 +60,7 @@ def test_segment_editor_maps_before_taking_modal_grab(monkeypatch) -> None: lambda: AppSettings(), lambda: False, MagicMock(), + None, MagicMock(), ) @@ -106,6 +108,9 @@ def _build_editor_with_ui( *, metadata=None, path: Path = Path("video.mp4"), + settings: AppSettings = AppSettings(), + job: JobItem | None = None, + ltx_unavailable_reason: str | None = None, ) -> SegmentEditor: worker = MagicMock() worker.events = queue.Queue() @@ -115,10 +120,11 @@ def _build_editor_with_ui( root.update() editor = SegmentEditor( root, - JobItem(path), - lambda: AppSettings(), + job or JobItem(path), + lambda: settings, lambda: False, MagicMock(), + ltx_unavailable_reason, MagicMock(), ) worker.events.put(segment_editor.PreviewLoaded(metadata or _fake_metadata())) @@ -340,7 +346,7 @@ def test_editor_height_grows_on_tall_screens(monkeypatch) -> None: def test_previous_frame_uses_exact_decoder_predecessor() -> None: editor = object.__new__(SegmentEditor) - editor._state = SegmentEditorState(duration=10.0, fps=30.0) + editor._state = SegmentEditorState(duration=10.0, fps=30.0, segments=(), default_restoration=SegmentRestoration("basicvsrpp", None)) editor._current = 1.0 editor._set_playing = MagicMock() editor._time_label = MagicMock() @@ -775,3 +781,80 @@ def test_fit_to_label_compensates_widget_scaling(monkeypatch, scaling: float) -> ) assert rendered[0] <= 900 and rendered[1] <= 540 assert result._light_image.size == (900, 506) + + +def _tk_root(): + try: + return ctk.CTk() + except TclError as exc: + pytest.skip(f"Tk display unavailable: {exc}") + + +def test_inspector_sets_the_model_and_seed_of_the_selected_range(monkeypatch) -> None: + root = _tk_root() + editor = None + try: + editor = _build_editor_with_ui(root, monkeypatch, settings=AppSettings(ltx_seed=42)) + editor._timeline_create(1.0, 2.0) + editor._timeline_create(4.0, 5.0) + root.update() + assert editor._inspector.winfo_manager() == "grid" + assert editor._range_model.get() == "basicvsrpp" + assert not editor._seed_entry.winfo_manager() + + editor._model_radios["ltx"].invoke() + root.update() + assert editor._state.selected_segment.restoration == SegmentRestoration("ltx", 42) + assert editor._seed_entry.get() == "42" + + editor._seed_entry.delete(0, "end") + editor._seed_entry.insert(0, "777") + editor._commit_seed_entry() + assert editor._state.selected_segment.restoration == SegmentRestoration("ltx", 777) + + editor._use_for_all_btn.invoke() + assert {s.restoration for s in editor._state.segments} == {SegmentRestoration("ltx", 777)} + assert editor._timeline._segments == editor._state.segments + + editor._undo() + assert editor._state.segments[0].restoration == SegmentRestoration("basicvsrpp", None) + finally: + if editor is not None: + editor._finish_close() + root.destroy() + + +def test_new_ranges_follow_the_job_model_and_rows_name_it(monkeypatch) -> None: + root = _tk_root() + editor = None + try: + editor = _build_editor_with_ui(root, monkeypatch, settings=AppSettings(restoration_model="ltx", ltx_seed=9)) + editor._timeline_create(1.0, 2.0) + root.update() + assert editor._state.segments[0].restoration == SegmentRestoration("ltx", 9) + texts = [w.cget("text") for w in editor._segment_action_widgets if isinstance(w, ctk.CTkButton)] + assert any(f"{t('model_ltx')} 9" in text for text in texts) + finally: + if editor is not None: + editor._finish_close() + root.destroy() + + +def test_ltx_is_blocked_with_a_reason_when_it_cannot_run(monkeypatch) -> None: + root = _tk_root() + editor = None + try: + job = JobItem(Path("video.mp4"), segments=(SegmentRange(4.0, 5.0, SegmentRestoration("ltx", 3)),)) + editor = _build_editor_with_ui( + root, monkeypatch, job=job, ltx_unavailable_reason="model_ltx_not_installed" + ) + assert editor._notice.cget("text") == t("segments_ltx_ranges_unavailable") + editor._timeline_create(1.0, 2.0) + root.update() + assert editor._model_radios["ltx"].cget("state") == "disabled" + editor._select_range(1) + assert editor._model_radios["ltx"].cget("state") == "normal" + finally: + if editor is not None: + editor._finish_close() + root.destroy() diff --git a/tests/test_segment_editor_state.py b/tests/test_segment_editor_state.py index f79f59af..c393be14 100644 --- a/tests/test_segment_editor_state.py +++ b/tests/test_segment_editor_state.py @@ -6,23 +6,30 @@ from jasna.gui.segment_editor_state import SegmentEditorState, smart_render_error_key from jasna.media.splice import KeyframeIndex -from jasna.segments import SegmentRange +from jasna.segments import SegmentRange, SegmentRestoration + +STANDARD = SegmentRestoration("basicvsrpp", None) +LTX_7 = SegmentRestoration("ltx", 7) + + +def _r(start: float, end: float, restoration: SegmentRestoration = STANDARD) -> SegmentRange: + return SegmentRange(start, end, restoration) def test_empty_editor_means_full_video() -> None: - state = SegmentEditorState(duration=10, fps=30) + state = SegmentEditorState(duration=10, fps=30, segments=(), default_restoration=STANDARD) assert state.output_segments == () assert not state.dirty def test_add_snaps_to_frames_and_selects_new_range() -> None: - state = SegmentEditorState(duration=10, fps=10) + state = SegmentEditorState(duration=10, fps=10, segments=(), default_restoration=STANDARD) state.add(1.04, 2.06) - assert state.output_segments == (SegmentRange(1.0, 2.1),) - assert state.selected_segment == SegmentRange(1.0, 2.1) + assert state.output_segments == (_r(1.0, 2.1),) + assert state.selected_segment == _r(1.0, 2.1) assert state.selected_duration == pytest.approx(1.1) @@ -30,14 +37,15 @@ def test_overlapping_ranges_merge_without_losing_selection() -> None: state = SegmentEditorState( duration=10, fps=10, - segments=(SegmentRange(1, 2), SegmentRange(4, 5)), + segments=(_r(1, 2), _r(4, 5)), + default_restoration=STANDARD, ) state.select(None) result = state.add(1.5, 4.5) assert result.merged_count == 2 - assert state.segments == (SegmentRange(1, 5),) + assert state.segments == (_r(1, 5),) assert state.selected_index == 0 @@ -45,35 +53,37 @@ def test_replace_selected_can_merge_with_neighbor() -> None: state = SegmentEditorState( duration=10, fps=10, - segments=(SegmentRange(1, 2), SegmentRange(4, 5)), + segments=(_r(1, 2), _r(4, 5)), + default_restoration=STANDARD, ) result = state.replace_selected(1, 4.2) assert result.merged_count == 1 - assert state.segments == (SegmentRange(1, 5),) + assert state.segments == (_r(1, 5),) def test_delete_undo_and_redo_restore_range_state() -> None: state = SegmentEditorState( duration=10, fps=30, - segments=(SegmentRange(1, 2), SegmentRange(3, 4)), + segments=(_r(1, 2), _r(3, 4)), + default_restoration=STANDARD, ) state.select(1) assert state.delete_selected() - assert state.segments == (SegmentRange(1, 2),) + assert state.segments == (_r(1, 2),) assert state.undo() - assert state.segments == (SegmentRange(1, 2), SegmentRange(3, 4)) + assert state.segments == (_r(1, 2), _r(3, 4)) assert state.selected_index == 1 assert state.redo() - assert state.segments == (SegmentRange(1, 2),) + assert state.segments == (_r(1, 2),) def test_clearing_all_ranges_means_full_video() -> None: - original = (SegmentRange(1, 2),) - state = SegmentEditorState(duration=10, fps=30, segments=original) + original = (_r(1, 2),) + state = SegmentEditorState(duration=10, fps=30, segments=original, default_restoration=STANDARD) assert state.clear() @@ -86,7 +96,7 @@ def test_clearing_all_ranges_means_full_video() -> None: def test_invalid_or_subframe_range_is_rejected() -> None: - state = SegmentEditorState(duration=10, fps=30) + state = SegmentEditorState(duration=10, fps=30, segments=(), default_restoration=STANDARD) with pytest.raises(ValueError, match="greater than start"): state.add(1.001, 1.002) @@ -95,4 +105,89 @@ def test_invalid_or_subframe_range_is_rejected() -> None: def test_smart_render_error_key_explains_whole_video_reencode() -> None: index = KeyframeIndex(pts=(0,), time_base=Fraction(1, 1), start_pts=0, end_pts=60) - assert smart_render_error_key((SegmentRange(0.0, 58.0),), index, 60.0) == "segments_smart_render_whole_video" + assert smart_render_error_key((_r(0.0, 58.0),), index, 60.0) == "segments_smart_render_whole_video" + + +def test_new_and_scanned_ranges_get_the_default_restoration() -> None: + state = SegmentEditorState(duration=10, fps=10, segments=(), default_restoration=LTX_7) + + state.add(1, 2) + assert state.add_many((SegmentRange(4, 5), SegmentRange(6, 7))) == 2 + + assert {segment.restoration for segment in state.segments} == {LTX_7} + + +def test_ranges_without_a_restoration_get_the_default_when_the_editor_opens() -> None: + state = SegmentEditorState(duration=10, fps=10, segments=(SegmentRange(1, 2),), default_restoration=LTX_7) + + assert state.segments == (_r(1, 2, LTX_7),) + assert not state.dirty + + +def test_moving_a_range_keeps_its_restoration() -> None: + state = SegmentEditorState( + duration=10, fps=10, segments=(_r(1, 2, LTX_7), _r(5, 6)), default_restoration=STANDARD + ) + state.select(0) + + state.adjust_selected(1.5, 3) + state.select(0) + state.replace_selected(0.5, 3) + + assert state.segments == (_r(0.5, 3, LTX_7), _r(5, 6)) + + +def test_setting_the_selected_restoration_is_one_undo_step() -> None: + state = SegmentEditorState(duration=10, fps=10, segments=(_r(1, 2), _r(5, 6)), default_restoration=STANDARD) + state.select(1) + + assert state.set_selected_restoration(LTX_7) + assert not state.set_selected_restoration(LTX_7) + assert state.segments == (_r(1, 2), _r(5, 6, LTX_7)) + assert state.selected_index == 1 + assert state.dirty + + assert state.undo() + assert state.segments == (_r(1, 2), _r(5, 6)) + assert state.redo() + assert state.segments[1].restoration == LTX_7 + + +def test_touching_ranges_merge_once_they_share_a_restoration() -> None: + state = SegmentEditorState( + duration=10, fps=10, segments=(_r(1, 2), _r(2, 3, LTX_7)), default_restoration=STANDARD + ) + assert len(state.segments) == 2 + state.select(1) + + state.set_selected_restoration(STANDARD) + + assert state.segments == (_r(1, 3),) + assert state.selected_index == 0 + + +def test_set_all_restoration_gives_every_range_one_choice() -> None: + state = SegmentEditorState(duration=10, fps=10, segments=(_r(1, 2), _r(5, 6)), default_restoration=STANDARD) + + assert state.set_all_restoration(LTX_7) + assert not state.set_all_restoration(LTX_7) + + assert state.segments == (_r(1, 2, LTX_7), _r(5, 6, LTX_7)) + assert state.undo() + assert state.segments == (_r(1, 2), _r(5, 6)) + + +def test_restoration_at_follows_the_range_under_the_playhead() -> None: + state = SegmentEditorState(duration=10, fps=10, segments=(_r(1, 2, LTX_7),), default_restoration=STANDARD) + + assert state.restoration_at(1.5) == LTX_7 + assert state.restoration_at(2.0) == STANDARD + assert state.restoration_at(0.5) == STANDARD + + +def test_smart_render_error_key_explains_mixed_models_too_close() -> None: + index = KeyframeIndex(pts=(0, 10_000), time_base=Fraction(1, 1000), start_pts=0, end_pts=20_000) + + segments = (_r(1.0, 2.0), _r(3.0, 4.0, LTX_7)) + + assert smart_render_error_key(segments, index, 20.0) == "segments_mixed_models_too_close" diff --git a/tests/test_segment_timeline.py b/tests/test_segment_timeline.py index 36643330..daf84fe2 100644 --- a/tests/test_segment_timeline.py +++ b/tests/test_segment_timeline.py @@ -2,7 +2,9 @@ import pytest -from jasna.gui.segment_timeline import timeline_seconds_to_x, timeline_x_to_seconds +from jasna.gui.segment_timeline import segment_fill, timeline_seconds_to_x, timeline_x_to_seconds +from jasna.gui.theme import Colors +from jasna.segments import SegmentRange, SegmentRestoration def test_timeline_coordinate_mapping_round_trips() -> None: @@ -37,3 +39,9 @@ def test_timeline_x_mapping_clamps_to_visible_range() -> None: view_end=40, width=300, ) == 40 + + +def test_ranges_are_coloured_by_their_model() -> None: + assert segment_fill(SegmentRange(0, 1, SegmentRestoration("ltx", 3))) == Colors.MODEL_LTX + assert segment_fill(SegmentRange(0, 1, SegmentRestoration("basicvsrpp", None))) == Colors.PRIMARY + assert segment_fill(SegmentRange(0, 1)) == Colors.PRIMARY diff --git a/tests/test_video_session.py b/tests/test_video_session.py index 5f09da74..8bfbb800 100644 --- a/tests/test_video_session.py +++ b/tests/test_video_session.py @@ -153,3 +153,7 @@ def test_video_session_key_tracks_model_and_ltx_variant() -> None: assert video_session_key(ltx) != video_session_key(AppSettings()) assert video_session_key(replace(ltx, ltx_fast=True)) != video_session_key(ltx) assert video_session_key(replace(ltx, ltx_seed=1, ltx_large_canvas=True)) == video_session_key(ltx) + + +def test_video_session_key_tracks_the_ltx_variant_for_ltx_ranges_of_standard_jobs() -> None: + assert video_session_key(replace(AppSettings(), ltx_fast=True)) != video_session_key(AppSettings()) From 02c811500d2bfcecddf2ec27c8479b302b503e46 Mon Sep 17 00:00:00 2001 From: Kruk2 Date: Tue, 29 Sep 2026 16:40:21 +0200 Subject: [PATCH 30/48] refactor(ltx): split segment renders into prepare, denoise and compose restore_spans runs prepare_spans, denoise_spans and compose_spans over one LtxRender, so prepared references can be denoised with another seed. --- jasna/ltx/restore.py | 254 ++++++++++++++++++++++++------------ jasna/media/splice.py | 45 ++++--- jasna/pipeline.py | 118 +++++++++-------- tests/test_ltx_restore.py | 35 ++++- tests/test_pipeline_init.py | 20 +++ tests/test_splice.py | 12 ++ 6 files changed, 328 insertions(+), 156 deletions(-) diff --git a/jasna/ltx/restore.py b/jasna/ltx/restore.py index c7e0b337..612d0072 100644 --- a/jasna/ltx/restore.py +++ b/jasna/ltx/restore.py @@ -165,6 +165,35 @@ def delete(self, index: int) -> None: self._path(index).unlink() +def video_frames( + path: Path, + metadata, + *, + batch_size: int, + device: torch.device, + frame_stride: int, + seek_ts: float | None, + end_pts: int | None, +) -> FrameSource: + """Decoded batches of ``path`` from the keyframe at ``seek_ts`` (the start when None), + stopping before ``end_pts`` (the end when None).""" + from jasna.media.video_decoder import VideoReader + + def frames() -> FrameBatches: + with VideoReader(str(path), batch_size, device, metadata, frame_stride=frame_stride) as reader: + for batch, pts in reader.frames(seek_ts): + if end_pts is None: + yield batch, pts + continue + keep = next((i for i, p in enumerate(pts) if int(p) >= end_pts), len(pts)) + if keep: + yield batch[:keep], pts[:keep] + if keep < len(pts): + return + + return frames + + def _check(cancel: threading.Event) -> None: if cancel.is_set(): raise Cancelled() @@ -452,95 +481,158 @@ def segment_large_canvas(requested: bool, total_bytes: int) -> bool: return requested -def restore_spans( - spans: Sequence[LtxSpan], - *, - detector, - files: LtxModelFiles, - frame_h: int, - frame_w: int, - batch_size: int, - large_canvas: bool, - budget: DecodeBudget, - device: torch.device, - work_dir: Path, - progress: Progress, - cancel: threading.Event, -) -> None: - """Restore every segment of ``spans`` and write each span's frames to a writer of its own - (closed once the span is written). ``large_canvas`` allows 768 px windows for large mosaics; - ``budget`` gives the decode memory budget per canvas size.""" +@dataclass(frozen=True) +class LtxRender: + """What every pass needs besides the spans: the detector, the model files, the frame + size, detection batch size, the 768 px decision, the decode budget per canvas and the + GPU.""" + + detector: Callable + files: LtxModelFiles + frame_h: int + frame_w: int + batch_size: int + large_canvas: bool + budget: DecodeBudget + device: torch.device + + +@dataclass(frozen=True) +class PreparedSpan: + """A span planned and encoded: per segment its track plans and reference latents. + Neither depends on the seed, so one prepared span can be denoised with many seeds.""" + + span: LtxSpan + plans: tuple[list[TrackPlan], ...] + references: tuple[LatentStore, ...] + + @property + def windows(self) -> int: + return sum(len(plan.windows) for segment_plans in self.plans for plan in segment_plans) + + +def prepare_spans( + spans: Sequence[LtxSpan], render: LtxRender, directory: Path, *, progress: Progress, cancel: threading.Event +) -> list[PreparedSpan]: + """Scan and VAE-encode every span; reference latents go to ``directory``.""" with progress.bar("scan", progress.frames) as bar: plans = [ scan_span( span.frames, span.segments, - detector, - batch_size=batch_size, - frame_h=frame_h, - frame_w=frame_w, - large_canvas=large_canvas, + render.detector, + batch_size=render.batch_size, + frame_h=render.frame_h, + frame_w=render.frame_w, + large_canvas=render.large_canvas, bar=bar, cancel=cancel, ) for span in spans ] - windows = sum(len(plan.windows) for span_plans in plans for segment_plans in span_plans for plan in segment_plans) + prepared = [ + PreparedSpan( + span, + tuple(span_plans), + tuple(LatentStore(directory, f"{s}-{g}.reference") for g in range(len(span.segments))), + ) + for s, (span, span_plans) in enumerate(zip(spans, plans)) + ] + if sum(p.windows for p in prepared): + encoder = load_video_encoder(render.files.vae, render.device) + with progress.bar("encode", progress.frames) as bar: + for p in prepared: + encode_span( + p.span.frames, + p.span.segments, + p.plans, + encoder, + p.references, + batch_size=render.batch_size, + device=render.device, + bar=bar, + cancel=cancel, + ) + del encoder + torch.cuda.empty_cache() + return prepared + + +def final_stores(prepared: Sequence[PreparedSpan], directory: Path) -> list[tuple[LatentStore, ...]]: + return [ + tuple(LatentStore(directory, f"{s}-{g}.final") for g in range(len(p.span.segments))) + for s, p in enumerate(prepared) + ] + + +def denoise_spans( + prepared: Sequence[PreparedSpan], + finals: Sequence[tuple[LatentStore, ...]], + transformer: LtxTransformer, + *, + keep_references: bool, + progress: Progress, + cancel: threading.Event, +) -> None: + """Denoise every track with its segment's seed into ``finals``; ``keep_references`` + keeps the reference latents for another seed.""" + windows = sum(p.windows for p in prepared) + with progress.bar("denoise", len(transformer.conditions) * windows, unit="step") as bar: + for p, span_finals in zip(prepared, finals): + for segment, segment_plans, references, store in zip(p.span.segments, p.plans, p.references, span_finals): + for plan in segment_plans: + _check(cancel) + latents = denoise_track(plan, transformer, references, seed=segment.seed, advance=bar.update) + for window, latent in zip(plan.windows, latents): + store.put(window.index, latent) + if not keep_references: + references.delete(window.index) + + +def compose_spans( + prepared: Sequence[PreparedSpan], + finals: Sequence[tuple[LatentStore, ...]], + decoder, + render: LtxRender, + *, + progress: Progress, + cancel: threading.Event, +) -> None: + with progress.bar("compose", progress.frames) as bar: + for p, span_finals in zip(prepared, finals): + compose_span( + p.span, + p.plans, + decoder, + span_finals, + frame_h=render.frame_h, + batch_size=render.batch_size, + budget=render.budget, + device=render.device, + bar=bar, + cancel=cancel, + ) + + +def restore_spans( + spans: Sequence[LtxSpan], render: LtxRender, *, work_dir: Path, progress: Progress, cancel: threading.Event +) -> None: + """Restore every segment of ``spans`` and write each span's frames to a writer of its own + (closed once the span is written).""" with tempfile.TemporaryDirectory(dir=work_dir, prefix=".ltx-") as temp: - stores = [ - [ - (LatentStore(Path(temp), f"{s}-{g}.reference"), LatentStore(Path(temp), f"{s}-{g}.final")) - for g in range(len(span.segments)) - ] - for s, span in enumerate(spans) - ] + prepared = prepare_spans(spans, render, Path(temp), progress=progress, cancel=cancel) + finals = final_stores(prepared, Path(temp)) + windows = sum(p.windows for p in prepared) if windows: - encoder = load_video_encoder(files.vae, device) - with progress.bar("encode", progress.frames) as bar: - for span, span_plans, span_stores in zip(spans, plans, stores): - encode_span( - span.frames, - span.segments, - span_plans, - encoder, - [references for references, _ in span_stores], - batch_size=batch_size, - device=device, - bar=bar, - cancel=cancel, - ) - del encoder - torch.cuda.empty_cache() - transformer = LtxTransformer(files.transformer, device) + transformer = LtxTransformer(render.files.transformer, render.device) try: - with progress.bar("denoise", len(transformer.conditions) * windows, unit="step") as bar: - for span, span_plans, span_stores in zip(spans, plans, stores): - for segment, segment_plans, (references, finals) in zip(span.segments, span_plans, span_stores): - for plan in segment_plans: - _check(cancel) - latents = denoise_track(plan, transformer, references, seed=segment.seed, advance=bar.update) - for window, latent in zip(plan.windows, latents): - finals.put(window.index, latent) - references.delete(window.index) + denoise_spans(prepared, finals, transformer, keep_references=False, progress=progress, cancel=cancel) finally: transformer.close() del transformer torch.cuda.empty_cache() - decoder = load_video_decoder(files.vae, files.tuned_decoder, device) if windows else None - with progress.bar("compose", progress.frames) as bar: - for span, span_plans, span_stores in zip(spans, plans, stores): - compose_span( - span, - span_plans, - decoder, - [finals for _, finals in span_stores], - frame_h=frame_h, - batch_size=batch_size, - budget=budget, - device=device, - bar=bar, - cancel=cancel, - ) + decoder = load_video_decoder(render.files.vae, render.files.tuned_decoder, render.device) if windows else None + compose_spans(prepared, finals, decoder, render, progress=progress, cancel=cancel) def restore_video( @@ -563,14 +655,16 @@ def restore_video( decision and the decode tiles follow the VRAM free when they are made.""" restore_spans( [LtxSpan(frames, (LtxSegment(*WHOLE_VIDEO_PTS, seed),), lambda: writer)], - detector=detector, - files=files, - frame_h=frame_h, - frame_w=frame_w, - batch_size=batch_size, - large_canvas=large_canvas_fits(large_canvas, torch.cuda.mem_get_info(device)[0]), - budget=free_vram_budget(device), - device=device, + LtxRender( + detector=detector, + files=files, + frame_h=frame_h, + frame_w=frame_w, + batch_size=batch_size, + large_canvas=large_canvas_fits(large_canvas, torch.cuda.mem_get_info(device)[0]), + budget=free_vram_budget(device), + device=device, + ), work_dir=work_dir, progress=progress, cancel=cancel, diff --git a/jasna/media/splice.py b/jasna/media/splice.py index e715b779..03881e24 100644 --- a/jasna/media/splice.py +++ b/jasna/media/splice.py @@ -372,6 +372,30 @@ def _is_safe_random_access_packet( return any(16 <= nal_type <= 20 for nal_type in nal_types) +def segment_render_span(segment: SegmentRange, index: KeyframeIndex) -> SpliceSpan: + """The render span of one segment alone: from the keyframe at or before it to the + keyframe at or after it, with the segment as its one effect range.""" + effect_start = index.start_pts + round(segment.start / index.time_base) + effect_end = index.start_pts + round(segment.end / index.time_base) + if effect_end <= effect_start: + raise SmartRenderCompatibilityError( + "A selected range is shorter than one video timestamp interval", + reason="range_too_short", + ) + left_index = bisect.bisect_right(index.pts, effect_start) - 1 + if left_index < 0: + raise SmartRenderCompatibilityError( + "The first selected segment begins before the first random-access keyframe", + reason="before_first_keyframe", + ) + render_start = index.pts[left_index] + right_index = bisect.bisect_left(index.pts, effect_end) + render_end = index.pts[right_index] if right_index < len(index.pts) else index.end_pts + if render_end <= render_start: + render_end = index.end_pts + return SpliceSpan("render", render_start, render_end, ((effect_start, effect_end),)) + + def build_splice_plan( segments: tuple[SegmentRange, ...] | list[SegmentRange], index: KeyframeIndex, @@ -384,25 +408,8 @@ def build_splice_plan( expanded: list[tuple[int, int, list[tuple[int, int]]]] = [] for segment in normalized: - effect_start = index.start_pts + round(segment.start / index.time_base) - effect_end = index.start_pts + round(segment.end / index.time_base) - if effect_end <= effect_start: - raise SmartRenderCompatibilityError( - "A selected range is shorter than one video timestamp interval", - reason="range_too_short", - ) - left_index = bisect.bisect_right(index.pts, effect_start) - 1 - if left_index < 0: - raise SmartRenderCompatibilityError( - "The first selected segment begins before the first random-access keyframe", - reason="before_first_keyframe", - ) - render_start = index.pts[left_index] - right_index = bisect.bisect_left(index.pts, effect_end) - render_end = index.pts[right_index] if right_index < len(index.pts) else index.end_pts - if render_end <= render_start: - render_end = index.end_pts - expanded.append((render_start, render_end, [(effect_start, effect_end)], [segment])) + span = segment_render_span(segment, index) + expanded.append((span.start_pts, span.end_pts, list(span.effect_ranges), [segment])) merged: list[tuple[int, int, list[tuple[int, int]], list[SegmentRange]]] = [] for start, end, effects, owners in expanded: diff --git a/jasna/pipeline.py b/jasna/pipeline.py index 64ccfb25..26ca4d9c 100644 --- a/jasna/pipeline.py +++ b/jasna/pipeline.py @@ -43,6 +43,7 @@ from jasna.vr_projection import build_vr_projector if TYPE_CHECKING: + from jasna.ltx.restore import FrameWriter, LtxRender, LtxSpan from jasna.session_factory import RestorationSession log = logging.getLogger(__name__) @@ -278,25 +279,6 @@ def _run_full( finally: progress.close(ensure_completed_bar=True) - def _ltx_frames(self, metadata, *, frame_stride: int, seek_ts: float | None, end_pts: int | None): - from jasna.media.video_decoder import VideoReader - - def frames(): - with VideoReader( - str(self.input_video), self.batch_size, self.device, metadata, frame_stride=frame_stride - ) as reader: - for batch, pts in reader.frames(seek_ts): - if end_pts is None: - yield batch, pts - continue - keep = next((i for i, p in enumerate(pts) if int(p) >= end_pts), len(pts)) - if keep: - yield batch[:keep], pts[:keep] - if keep < len(pts): - return - - return frames - def _ltx_progress(self, frames: int): from jasna.ltx.restore import Progress @@ -311,14 +293,22 @@ def _require_ltx_video(self) -> None: raise ValueError("LTX restoration does not support VR180 side-by-side video") def _run_ltx(self, metadata) -> None: - from jasna.ltx.restore import Cancelled, restore_video + from jasna.ltx.restore import Cancelled, restore_video, video_frames self._require_ltx_video() frame_rate = self._resolve_frame_rate(metadata) writer = _OfflineFrameWriter(self._video_encoder(metadata, frame_rate), [time.monotonic()]) try: restore_video( - self._ltx_frames(metadata, frame_stride=frame_rate.frame_stride, seek_ts=None, end_pts=None), + video_frames( + self.input_video, + metadata, + batch_size=self.batch_size, + device=self.device, + frame_stride=frame_rate.frame_stride, + seek_ts=None, + end_pts=None, + ), writer, detector=self.job_detection_model, files=self.ltx_files, @@ -337,6 +327,55 @@ def _run_ltx(self, metadata) -> None: finally: writer.close() + def ltx_segment_render(self, metadata) -> "LtxRender": + """How LTX restores segments: the 768 px decision and the decode tiles follow the + GPU's size, never the VRAM free at that moment, so a segment renders the same in a + job and in a seed preview.""" + from jasna.ltx.restore import LtxRender, segment_decode_budget, segment_large_canvas + + self._require_ltx_video() + return LtxRender( + detector=self.job_detection_model, + files=self.ltx_files, + frame_h=int(metadata.video_height), + frame_w=int(metadata.video_width), + batch_size=self.batch_size, + large_canvas=segment_large_canvas( + self.ltx_large_canvas, torch.cuda.get_device_properties(self.device).total_memory + ), + budget=segment_decode_budget, + device=self.device, + ) + + def ltx_span( + self, + metadata, + index: KeyframeIndex, + span: SpliceSpan, + segments: tuple[SegmentRange, ...], + open_writer: Callable[[], "FrameWriter"], + ) -> "LtxSpan": + """The LTX work of one render span: decoded from its keyframe, one LTX segment per + effect range with that range's seed.""" + from jasna.ltx.restore import LtxSegment, LtxSpan, video_frames + + return LtxSpan( + video_frames( + self.input_video, + metadata, + batch_size=self.batch_size, + device=self.device, + frame_stride=1, + seek_ts=index.seconds_for_pts(span.start_pts), + end_pts=span.end_pts, + ), + tuple( + LtxSegment(start, end, segment.restoration.ltx_seed) + for (start, end), segment in zip(span.effect_ranges, segments) + ), + open_writer, + ) + def _run_ltx_spans( self, metadata, @@ -345,45 +384,16 @@ def _run_ltx_spans( work_dir: Path, ) -> None: """Restore the LTX render spans in one batched run, so every model loads once.""" - from jasna.ltx.restore import ( - Cancelled, - LtxSegment, - LtxSpan, - restore_spans, - segment_decode_budget, - segment_large_canvas, - ) + from jasna.ltx.restore import Cancelled, restore_spans - self._require_ltx_video() - ltx_spans = [ - LtxSpan( - self._ltx_frames( - metadata, frame_stride=1, seek_ts=index.seconds_for_pts(span.start_pts), end_pts=span.end_pts - ), - tuple( - LtxSegment(start, end, segment.restoration.ltx_seed) - for (start, end), segment in zip(span.effect_ranges, segments) - ), - open_writer, - ) - for span, segments, open_writer in spans - ] + render = self.ltx_segment_render(metadata) frames = sum( round((span.end_pts - span.start_pts) * index.time_base * metadata.video_fps) for span, _, _ in spans ) try: restore_spans( - ltx_spans, - detector=self.job_detection_model, - files=self.ltx_files, - frame_h=int(metadata.video_height), - frame_w=int(metadata.video_width), - batch_size=self.batch_size, - large_canvas=segment_large_canvas( - self.ltx_large_canvas, torch.cuda.get_device_properties(self.device).total_memory - ), - budget=segment_decode_budget, - device=self.device, + [self.ltx_span(metadata, index, span, segments, open_writer) for span, segments, open_writer in spans], + render, work_dir=work_dir, progress=self._ltx_progress(frames), cancel=self._cancel_event, diff --git a/tests/test_ltx_restore.py b/tests/test_ltx_restore.py index fcd9644d..70180ac3 100644 --- a/tests/test_ltx_restore.py +++ b/tests/test_ltx_restore.py @@ -116,9 +116,8 @@ def bar(self, name, total, unit="frame"): return budgets, bars -def _run_spans(tmp_path, spans): - restore.restore_spans( - spans, +def _render(): + return restore.LtxRender( detector=_Detector(), files=FILES, frame_h=64, @@ -127,6 +126,13 @@ def _run_spans(tmp_path, spans): large_canvas=False, budget=restore.segment_decode_budget, device=torch.device("cpu"), + ) + + +def _run_spans(tmp_path, spans): + restore.restore_spans( + spans, + _render(), work_dir=tmp_path, progress=restore.Progress(0, disable=True, report=None), cancel=threading.Event(), @@ -294,3 +300,26 @@ def test_progress_without_a_report_only_drives_the_console_bar(): progress = restore.Progress(3, disable=True, report=None) with progress.bar("compose", 3) as bar: bar.update(3) + + +def test_a_prepared_span_denoises_with_many_seeds(fakes, tmp_path): + video = _video(30, range(3, 28)) + span = restore.LtxSpan(_source(video), (restore.LtxSegment(50, 150, seed=1),), _Writer) + progress = restore.Progress(30, disable=True, report=None) + cancel = threading.Event() + (prepared,) = restore.prepare_spans([span], _render(), tmp_path, progress=progress, cancel=cancel) + transformer = _Transformer(None, None) + _Transformer.seeds = [] + + for seed in (11, 22): + reseeded = restore.PreparedSpan( + restore.LtxSpan(span.frames, (restore.LtxSegment(50, 150, seed=seed),), span.open_writer), + prepared.plans, + prepared.references, + ) + finals = restore.final_stores([reseeded], tmp_path / str(seed)) + (tmp_path / str(seed)).mkdir() + restore.denoise_spans([reseeded], finals, transformer, keep_references=True, progress=progress, cancel=cancel) + + assert _Transformer.seeds == [[11], [22]] + assert prepared.windows == 1 diff --git a/tests/test_pipeline_init.py b/tests/test_pipeline_init.py index 78ea97e0..d3aae656 100644 --- a/tests/test_pipeline_init.py +++ b/tests/test_pipeline_init.py @@ -141,3 +141,23 @@ def test_configure_vr_honors_per_job_projection_override(self): assert pipeline.vr_resolution.projection == "fisheye" assert isinstance(pipeline.vr_projector, FisheyeProjector) + + +def test_ltx_span_gives_each_effect_range_its_segment_seed() -> None: + from jasna.ltx.restore import LtxSegment + from jasna.media.splice import KeyframeIndex, SpliceSpan + from jasna.segments import SegmentRange, SegmentRestoration + + p = _make_pipeline() + index = KeyframeIndex(pts=(0, 2000), time_base=Fraction(1, 1000), start_pts=0, end_pts=6000) + span = SpliceSpan("render", 2000, 6000, ((2500, 3000), (4000, 5000))) + segments = ( + SegmentRange(2.5, 3.0, SegmentRestoration("ltx", 7)), + SegmentRange(4.0, 5.0, SegmentRestoration("ltx", 9)), + ) + opener = MagicMock() + + ltx = p.ltx_span(SimpleNamespace(), index, span, segments, opener) + + assert ltx.segments == (LtxSegment(2500, 3000, 7), LtxSegment(4000, 5000, 9)) + assert ltx.open_writer is opener diff --git a/tests/test_splice.py b/tests/test_splice.py index 03ede27d..dfb72adf 100644 --- a/tests/test_splice.py +++ b/tests/test_splice.py @@ -358,3 +358,15 @@ def test_copy_fragment_seeks_before_demux(tmp_path: Path) -> None: ] assert packet.pts == 5 assert packet.dts == 3 + + +def test_segment_render_span_widens_one_segment_to_its_keyframes() -> None: + from fractions import Fraction + + from jasna.media.splice import KeyframeIndex, SpliceSpan, segment_render_span + from jasna.segments import SegmentRange + + index = KeyframeIndex(pts=(0, 2000, 4000), time_base=Fraction(1, 1000), start_pts=0, end_pts=6000) + + assert segment_render_span(SegmentRange(2.5, 3.0), index) == SpliceSpan("render", 2000, 4000, ((2500, 3000),)) + assert segment_render_span(SegmentRange(4.5, 6.0), index) == SpliceSpan("render", 4000, 6000, ((4500, 6000),)) From 95edce6cd3068f4cbd5ba8fd84256d79ab539a78 Mon Sep 17 00:00:00 2001 From: Kruk2 Date: Tue, 29 Sep 2026 16:48:40 +0200 Subject: [PATCH 31/48] feat(gui): try LTX seeds for a range in the segment editor Try seed restores the selected range in the background like a job and keeps the prepared latents, so another seed reruns only denoise and decode. Use this seed applies a tried seed. --- jasna/gui/locales/en.py | 11 + jasna/gui/locales/ja.py | 11 + jasna/gui/locales/ko.py | 11 + jasna/gui/locales/th.py | 11 + jasna/gui/locales/zh.py | 11 + jasna/gui/ltx_seed_preview.py | 354 +++++++++++++++++++++++++++++ jasna/gui/restoration_preview.py | 6 +- jasna/gui/segment_editor.py | 294 +++++++++++++++++++++++- jasna/ltx/restore.py | 5 + tests/test_ltx_seed_preview.py | 119 ++++++++++ tests/test_ltx_seed_preview_gpu.py | 101 ++++++++ tests/test_segment_editor.py | 94 ++++++++ 12 files changed, 1022 insertions(+), 6 deletions(-) create mode 100644 jasna/gui/ltx_seed_preview.py create mode 100644 tests/test_ltx_seed_preview.py create mode 100644 tests/test_ltx_seed_preview_gpu.py diff --git a/jasna/gui/locales/en.py b/jasna/gui/locales/en.py index 2be662a1..9b322272 100644 --- a/jasna/gui/locales/en.py +++ b/jasna/gui/locales/en.py @@ -132,6 +132,17 @@ "tip_segments_use_for_all": "Give every range this model and seed.", "segments_ltx_needs_flat_video": "LTX can't restore VR side-by-side videos.", "segments_ltx_ranges_unavailable": "LTX can't run on this PC, so the LTX ranges can't be restored.", + "segments_try_seed": "Try seed", + "tip_segments_try_seed": "Restore this range with its seed so you can watch it before processing. Takes a few minutes.", + "segments_seed_starting": "Seed {seed}: starting…", + "segments_seed_progress": "Seed {seed}: {stage} {percent:.0f}% · about {eta} left", + "segments_seed_progress_no_eta": "Seed {seed}: {stage} {percent:.0f}%", + "segments_seed_tried": "Tried:", + "segments_seed_use": "Use this seed", + "tip_segments_seed_chip": "Show this seed in the preview.", + "segments_seed_show": "Show seed {seed}", + "segments_seed_failed": "Seed preview failed: {message}", + "segments_seed_stop_confirm": "A seed preview is still running. Stop it and close?", "segments_suggest_mask": "Suggest better mask", "segments_suggest_mask_hint": "Draw the correct mosaic outline on this frame and share it anonymously to improve future detection models", "mask_editor_title": "Suggest a better mosaic mask", diff --git a/jasna/gui/locales/ja.py b/jasna/gui/locales/ja.py index ee84033c..6504c0d4 100644 --- a/jasna/gui/locales/ja.py +++ b/jasna/gui/locales/ja.py @@ -132,6 +132,17 @@ "tip_segments_use_for_all": "すべての範囲にこのモデルとシードを使います。", "segments_ltx_needs_flat_video": "LTX は VR のサイドバイサイド動画を復元できません。", "segments_ltx_ranges_unavailable": "この PC では LTX を使えないため、LTX の範囲は復元できません。", + "segments_try_seed": "シードを試す", + "tip_segments_try_seed": "この範囲をそのシードで復元し、処理前に確認できます。数分かかります。", + "segments_seed_starting": "シード {seed}：開始中…", + "segments_seed_progress": "シード {seed}：{stage} {percent:.0f}% · 残り約 {eta}", + "segments_seed_progress_no_eta": "シード {seed}：{stage} {percent:.0f}%", + "segments_seed_tried": "試したシード：", + "segments_seed_use": "このシードを使う", + "tip_segments_seed_chip": "このシードをプレビューに表示します。", + "segments_seed_show": "シード {seed} を表示", + "segments_seed_failed": "シードのプレビューに失敗しました：{message}", + "segments_seed_stop_confirm": "シードのプレビューがまだ実行中です。停止して閉じますか？", "segments_suggest_mask": "より良いマスクを提案", "segments_suggest_mask_hint": "このフレームに正しいモザイクの輪郭を描き、匿名で共有して今後の検出モデルの改善に役立てます", "mask_editor_title": "より良いモザイクマスクを提案", diff --git a/jasna/gui/locales/ko.py b/jasna/gui/locales/ko.py index 5c48f9a9..dfeeb97a 100644 --- a/jasna/gui/locales/ko.py +++ b/jasna/gui/locales/ko.py @@ -131,6 +131,17 @@ "tip_segments_use_for_all": "모든 구간에 이 모델과 시드를 적용합니다.", "segments_ltx_needs_flat_video": "LTX는 VR 사이드 바이 사이드 영상을 복원할 수 없습니다.", "segments_ltx_ranges_unavailable": "이 PC에서는 LTX를 실행할 수 없어 LTX 구간을 복원할 수 없습니다.", + "segments_try_seed": "시드 시험", + "tip_segments_try_seed": "이 구간을 해당 시드로 복원해 처리 전에 미리 볼 수 있습니다. 몇 분 걸립니다.", + "segments_seed_starting": "시드 {seed}: 시작 중…", + "segments_seed_progress": "시드 {seed}: {stage} {percent:.0f}% · 약 {eta} 남음", + "segments_seed_progress_no_eta": "시드 {seed}: {stage} {percent:.0f}%", + "segments_seed_tried": "시험한 시드:", + "segments_seed_use": "이 시드 사용", + "tip_segments_seed_chip": "미리보기에 이 시드를 표시합니다.", + "segments_seed_show": "시드 {seed} 보기", + "segments_seed_failed": "시드 미리보기 실패: {message}", + "segments_seed_stop_confirm": "시드 미리보기가 아직 실행 중입니다. 중지하고 닫을까요?", "segments_suggest_mask": "더 나은 마스크 제안", "segments_suggest_mask_hint": "이 프레임에 올바른 모자이크 윤곽을 그려 익명으로 공유하면 향후 감지 모델 개선에 도움이 됩니다", "mask_editor_title": "더 나은 모자이크 마스크 제안", diff --git a/jasna/gui/locales/th.py b/jasna/gui/locales/th.py index d604d0b9..a8ec09db 100644 --- a/jasna/gui/locales/th.py +++ b/jasna/gui/locales/th.py @@ -132,6 +132,17 @@ "tip_segments_use_for_all": "ให้ทุกช่วงใช้โมเดลและซีดนี้", "segments_ltx_needs_flat_video": "LTX กู้คืนวิดีโอ VR แบบซ้าย-ขวาไม่ได้", "segments_ltx_ranges_unavailable": "PC เครื่องนี้ใช้ LTX ไม่ได้ จึงกู้คืนช่วง LTX ไม่ได้", + "segments_try_seed": "ลองซีด", + "tip_segments_try_seed": "กู้คืนช่วงนี้ด้วยซีดนี้เพื่อดูผลก่อนประมวลผล ใช้เวลาไม่กี่นาที", + "segments_seed_starting": "ซีด {seed}: กำลังเริ่ม…", + "segments_seed_progress": "ซีด {seed}: {stage} {percent:.0f}% · เหลืออีกประมาณ {eta}", + "segments_seed_progress_no_eta": "ซีด {seed}: {stage} {percent:.0f}%", + "segments_seed_tried": "ที่ลองแล้ว:", + "segments_seed_use": "ใช้ซีดนี้", + "tip_segments_seed_chip": "แสดงซีดนี้ในตัวอย่าง", + "segments_seed_show": "แสดงซีด {seed}", + "segments_seed_failed": "ดูตัวอย่างซีดไม่สำเร็จ: {message}", + "segments_seed_stop_confirm": "ยังดูตัวอย่างซีดอยู่ หยุดแล้วปิดเลยไหม", "segments_suggest_mask": "เสนอมาสก์ที่ดีกว่า", "segments_suggest_mask_hint": "วาดโครงร่างโมเสคที่ถูกต้องบนเฟรมนี้และแชร์แบบไม่ระบุตัวตนเพื่อช่วยปรับปรุงโมเดลตรวจจับในอนาคต", "mask_editor_title": "เสนอมาสก์โมเสคที่ดีกว่า", diff --git a/jasna/gui/locales/zh.py b/jasna/gui/locales/zh.py index 158fb3fe..d048aa35 100644 --- a/jasna/gui/locales/zh.py +++ b/jasna/gui/locales/zh.py @@ -132,6 +132,17 @@ "tip_segments_use_for_all": "让所有片段都使用这个模型和种子。", "segments_ltx_needs_flat_video": "LTX 无法修复 VR 左右并排视频。", "segments_ltx_ranges_unavailable": "这台电脑无法运行 LTX，所以 LTX 片段无法修复。", + "segments_try_seed": "试试种子", + "tip_segments_try_seed": "用这个种子修复此片段，处理前就能先看效果。需要几分钟。", + "segments_seed_starting": "种子 {seed}：正在开始…", + "segments_seed_progress": "种子 {seed}：{stage} {percent:.0f}% · 约剩 {eta}", + "segments_seed_progress_no_eta": "种子 {seed}：{stage} {percent:.0f}%", + "segments_seed_tried": "已试：", + "segments_seed_use": "使用这个种子", + "tip_segments_seed_chip": "在预览中显示这个种子。", + "segments_seed_show": "显示种子 {seed}", + "segments_seed_failed": "种子预览失败：{message}", + "segments_seed_stop_confirm": "种子预览仍在运行。要停止并关闭吗？", "segments_suggest_mask": "提交更好的遮罩", "segments_suggest_mask_hint": "在此帧上勾画正确的马赛克轮廓并匿名分享，帮助改进未来的检测模型", "mask_editor_title": "提交更好的马赛克遮罩", diff --git a/jasna/gui/ltx_seed_preview.py b/jasna/gui/ltx_seed_preview.py new file mode 100644 index 00000000..c18dbe90 --- /dev/null +++ b/jasna/gui/ltx_seed_preview.py @@ -0,0 +1,354 @@ +"""LTX seed preview for the segment editor. + +Restores one range with one seed the way a job restores it (the same render span, +detection batches, canvas decision, decode tiles and seed rule, through the same +``Pipeline`` and ``ltx.restore`` code), so what the user picks is what the job makes. +The plan and reference latents of a range are kept, so another seed only reruns +denoise and decode; the transformer stays loaded between seeds while it leaves room +to decode. Restored frames are saved as JPEG files, one per source frame. +""" + +from __future__ import annotations + +import queue +import shutil +import tempfile +import threading +from collections.abc import Callable +from dataclasses import dataclass, replace +from pathlib import Path + +import torch + +from jasna.gui.models import AppSettings +from jasna.gui.queues import replace_pending +from jasna.gui.video_session import build_video_session, release_session_memory, video_session_key +from jasna.media.probe import VideoMetadata +from jasna.media.splice import KeyframeIndex, segment_render_span +from jasna.segments import SegmentRange, SegmentRestoration + +DECODE_HEADROOM_BYTES = 2 << 30 +JPEG_QUALITY = 95 + + +@dataclass(frozen=True) +class SeedFrame: + seconds: float + path: Path + + +@dataclass(frozen=True) +class SeedProgress: + seed: int + stage: str + fraction: float + eta_seconds: float + generation: int + + +@dataclass(frozen=True) +class SeedReady: + seed: int + frames: tuple[SeedFrame, ...] + generation: int + + +@dataclass(frozen=True) +class SeedFailed: + message: str + generation: int + + +SeedEvent = SeedProgress | SeedReady | SeedFailed + + +class SeedFrameWriter: + """FrameWriter that saves the frames of one effect range as JPEG files.""" + + def __init__(self, directory: Path, effect_range: tuple[int, int], metadata: VideoMetadata, lut_applier) -> None: + self._directory = directory + self._start_pts, self._end_pts = effect_range + self._metadata = metadata + self._lut_applier = lut_applier + self.frames: list[SeedFrame] = [] + + def write(self, frame: torch.Tensor, pts: int, *, apply_lut: bool = True) -> None: + if not self._start_pts <= int(pts) < self._end_pts: + return + from jasna.gui.restoration_preview import frame_image + + size = (int(self._metadata.video_width), int(self._metadata.video_height)) + path = self._directory / f"{int(pts)}.jpg" + frame_image(frame, size, self._lut_applier, apply_lut=apply_lut).save( + path, quality=JPEG_QUALITY, subsampling=0 + ) + seconds = (int(pts) - self._metadata.start_pts) * float(self._metadata.time_base) + self.frames.append(SeedFrame(max(0.0, seconds), path)) + + def close(self) -> None: + pass + + +def prepared_key(segment: SegmentRange, settings: AppSettings) -> tuple: + """What a range's plan and reference latents depend on (not the seed).""" + return (segment.start, segment.end, video_session_key(settings), settings.ltx_large_canvas) + + +class SeedRenderer: + """Renders LTX ranges seed by seed on the calling thread, keeping the session, the + transformer and every range's prepared references between calls.""" + + def __init__(self, path: Path, metadata: VideoMetadata, index: KeyframeIndex, work_dir: Path) -> None: + self.path = Path(path) + self.metadata = metadata + self.index = index + self._temp = Path(tempfile.mkdtemp(prefix=".jasna-seeds-", dir=work_dir)) + self._session = None + self._session_key: tuple | None = None + self._prepared: dict[tuple, object] = {} + self._transformer = None + self._transformer_path: Path | None = None + self._runs = 0 + + def render( + self, + segment: SegmentRange, + seed: int, + settings: AppSettings, + *, + writer_for: Callable[[Path, tuple[int, int]], object], + report: Callable[[str, float, float], None], + cancel: threading.Event, + ): + """Restore ``segment`` with ``seed``; returns the writer ``writer_for(directory, + effect range)`` made, after every frame of the range went through it.""" + from jasna.gui.video_session import video_session_config + from jasna.ltx.plan import LARGE_CANVAS + from jasna.ltx.camera import PLAN_CANVAS + from jasna.ltx.restore import ( + PreparedSpan, + Progress, + compose_spans, + denoise_spans, + final_stores, + prepare_spans, + ) + from jasna.ltx.sampler import LtxTransformer + from jasna.models.ltx_vae import load_video_decoder + from jasna.session_factory import build_pipeline + + settings = replace(settings, restoration_model="ltx") + seeded = replace(segment, restoration=SegmentRestoration("ltx", int(seed))) + self._ensure_session(settings) + config = video_session_config(settings, codec=settings.codec, encoder_settings={}) + pipeline = build_pipeline(config, self._session, self.path, self.path, segments=(seeded,)) + pipeline.configure_vr(self.metadata) + render = pipeline.ltx_segment_render(self.metadata) + span = segment_render_span(seeded, self.index) + self._runs += 1 + run_dir = self._temp / f"run-{self._runs}" + run_dir.mkdir() + writer = writer_for(run_dir, span.effect_ranges[0]) + ltx_span = pipeline.ltx_span(self.metadata, self.index, span, (seeded,), lambda: writer) + frames = round((span.end_pts - span.start_pts) * self.index.time_base * self.metadata.video_fps) + progress = Progress(frames, disable=True, report=report) + + key = prepared_key(segment, settings) + prepared = self._prepared.get(key) + if prepared is None: + directory = self._temp / f"prepared-{len(self._prepared)}" + directory.mkdir() + (prepared,) = prepare_spans([ltx_span], render, directory, progress=progress, cancel=cancel) + self._prepared[key] = prepared + else: + progress.skip("scan", "encode") + prepared = PreparedSpan(ltx_span, prepared.plans, prepared.references) + finals = final_stores([prepared], run_dir) + decoder = None + if prepared.windows: + if self._transformer is None or self._transformer_path != render.files.transformer: + self._drop_transformer() + self._transformer = LtxTransformer(render.files.transformer, render.device) + self._transformer_path = render.files.transformer + denoise_spans([prepared], finals, self._transformer, keep_references=True, progress=progress, cancel=cancel) + decode_bytes = render.budget(LARGE_CANVAS if render.large_canvas else PLAN_CANVAS) + DECODE_HEADROOM_BYTES + if torch.cuda.mem_get_info(render.device)[0] < decode_bytes: + self._drop_transformer() + decoder = load_video_decoder(render.files.vae, render.files.tuned_decoder, render.device) + try: + compose_spans([prepared], finals, decoder, render, progress=progress, cancel=cancel) + finally: + del decoder + torch.cuda.empty_cache() + return writer + + def _ensure_session(self, settings: AppSettings) -> None: + key = video_session_key(settings) + if self._session is not None and key == self._session_key: + return + self._drop_session() + self._session = build_video_session(settings, log=lambda _msg: None) + self._session_key = key + + def _drop_transformer(self) -> None: + if self._transformer is not None: + self._transformer.close() + self._transformer = None + self._transformer_path = None + torch.cuda.empty_cache() + + def _drop_session(self) -> None: + self._drop_transformer() + if self._session is not None: + device = self._session.device + self._session.close() + self._session = None + self._session_key = None + release_session_memory(device) + + def release(self) -> None: + """Free the GPU; the prepared references stay on disk for the next seed.""" + self._drop_session() + + def close(self) -> None: + self._drop_session() + self._prepared.clear() + shutil.rmtree(self._temp) + + +@dataclass(frozen=True) +class _TrySeed: + segment: SegmentRange + seed: int + settings: AppSettings + generation: int + + +@dataclass(frozen=True) +class _Release: + pass + + +@dataclass(frozen=True) +class _Stop: + pass + + +class LtxSeedPreviewWorker: + """Background thread around ``SeedRenderer``: one seed at a time, a new request + cancels the running one. Posts ``SeedEvent``s to ``events``; never touches Tk.""" + + def __init__( + self, + path: str | Path, + metadata: VideoMetadata, + index: KeyframeIndex, + work_dir: Path, + *, + on_stopped: Callable[[], None], + ) -> None: + self.path = Path(path) + self.metadata = metadata + self.index = index + self.work_dir = Path(work_dir) + self._on_stopped = on_stopped + self.events: queue.Queue[SeedEvent] = queue.Queue() + self._commands: queue.Queue[_TrySeed | _Release | _Stop] = queue.Queue(maxsize=1) + self._closed = threading.Event() + self._generation = 0 + self._active_cancel: threading.Event | None = None + self._cancel_lock = threading.Lock() + self._thread = threading.Thread(target=self._run, name=f"ltx-seed-preview-{self.path.name}", daemon=True) + + def start(self) -> None: + self._thread.start() + + def try_seed(self, segment: SegmentRange, seed: int, settings: AppSettings) -> int: + self._generation += 1 + self._cancel_active() + replace_pending(self._commands, _TrySeed(segment, int(seed), settings, self._generation)) + return self._generation + + def cancel(self) -> None: + self._generation += 1 + self._cancel_active() + + def release(self) -> None: + """Cancel any run and free the GPU for another preview or a scan.""" + self.cancel() + replace_pending(self._commands, _Release()) + + def close(self) -> None: + if self._closed.is_set(): + return + self._closed.set() + self._cancel_active() + replace_pending(self._commands, _Stop()) + + def join(self, timeout: float | None = None) -> None: + self._thread.join(timeout=timeout) + + def _cancel_active(self) -> None: + with self._cancel_lock: + if self._active_cancel is not None: + self._active_cancel.set() + + def _make_renderer(self) -> SeedRenderer: + return SeedRenderer(self.path, self.metadata, self.index, self.work_dir) + + def _run(self) -> None: + from jasna.ltx.restore import Cancelled + + renderer: SeedRenderer | None = None + try: + while True: + command = self._commands.get() + if isinstance(command, _Stop): + break + if isinstance(command, _Release): + if renderer is not None: + renderer.release() + continue + cancel = threading.Event() + with self._cancel_lock: + self._active_cancel = cancel + generation = command.generation + try: + if renderer is None: + renderer = self._make_renderer() + writer = renderer.render( + command.segment, + command.seed, + command.settings, + writer_for=lambda directory, effect: SeedFrameWriter( + directory, effect, self.metadata, _lut_applier(command.settings) + ), + report=lambda stage, fraction, eta: self.events.put( + SeedProgress(command.seed, stage, fraction, eta, generation) + ), + cancel=cancel, + ) + self.events.put(SeedReady(command.seed, tuple(writer.frames), generation)) + except Cancelled: + pass + except Exception as exc: + if not self._closed.is_set(): + self.events.put(SeedFailed(str(exc), generation)) + finally: + with self._cancel_lock: + self._active_cancel = None + finally: + try: + if renderer is not None: + renderer.close() + finally: + self._on_stopped() + + +def _lut_applier(settings: AppSettings): + lut_path = (settings.lut_path or "").strip() + if not lut_path: + return None + from jasna.media.lut import GpuLutApplier, parse_cube_file + + return GpuLutApplier(parse_cube_file(lut_path), torch.device("cuda:0")) diff --git a/jasna/gui/restoration_preview.py b/jasna/gui/restoration_preview.py index a8350739..01e66dca 100644 --- a/jasna/gui/restoration_preview.py +++ b/jasna/gui/restoration_preview.py @@ -106,7 +106,7 @@ def playback_window(metadata: VideoMetadata, start_seconds: float, max_clip_size return PreviewWindow(seek_ts=start_seconds, end_pts=end_pts, center_pts=start_pts) -def _frame_image( +def frame_image( frame, max_size: tuple[int, int], lut_applier=None, @@ -171,7 +171,7 @@ def has_result(self) -> bool: return self._best_frame is not None def result_image(self, max_size: tuple[int, int]) -> Image.Image: - return _frame_image(self._best_frame, max_size, apply_lut=False) + return frame_image(self._best_frame, max_size, apply_lut=False) class _PlaybackFrameCollector: @@ -197,7 +197,7 @@ def write(self, frame, pts: int, *, apply_lut: bool = True) -> None: self._frames.append( ( pts, - _frame_image( + frame_image( frame, self._max_size, self._lut_applier, diff --git a/jasna/gui/segment_editor.py b/jasna/gui/segment_editor.py index f44bde6d..d7219243 100644 --- a/jasna/gui/segment_editor.py +++ b/jasna/gui/segment_editor.py @@ -3,6 +3,7 @@ import logging import queue import random +from pathlib import Path import threading import tkinter as tk from collections.abc import Callable @@ -15,8 +16,9 @@ from jasna.gui import scaling from jasna.gui.locales import t from jasna.gui.models import AppSettings, JobItem -from jasna.gui.components import Tooltip, grab_modal +from jasna.gui.components import Tooltip, format_duration, grab_modal from jasna.gui.icons import CompactSwitch, NativeIconButton +from jasna.gui.ltx_seed_preview import LtxSeedPreviewWorker, SeedFailed, SeedFrame, SeedProgress, SeedReady from jasna.gui.restoration_preview import ( RestorationClip, RestorationFailed, @@ -124,6 +126,13 @@ def __init__( self._timeline_zoom_buttons: list = [] self._mask_feedback_worker = MaskFeedbackWorker() self._suggest_busy = False + self._seed_worker: LtxSeedPreviewWorker | None = None + self._seed_generation = 0 + self._seed_running: tuple[tuple[float, float], int] | None = None + self._tried_seeds: dict[tuple[float, float], dict[int, tuple[SeedFrame, ...]]] = {} + self._shown_seed: tuple[tuple[float, float], int] | None = None + self._seed_view_on = False + self._seed_image: tuple[SeedFrame, Image.Image] | None = None self.title(t("segments_title")) self.configure(fg_color=Colors.BG_MAIN) @@ -409,6 +418,16 @@ def _build_preview_card(self, body: ctk.CTkFrame) -> None: font=(Fonts.FAMILY, Fonts.SIZE_SMALL), ).pack(side="right", padx=(0, 6)) self._restore_toggle_tooltip = Tooltip(self._restore_toggle, t("segments_restore_preview_hint")) + self._seed_view_control = ctk.CTkFrame(preview_options, fg_color="transparent") + self._seed_view_toggle = CompactSwitch(self._seed_view_control, self._toggle_seed_view, Colors.BG_CARD) + self._seed_view_toggle.pack(side="right") + self._seed_view_label = ctk.CTkLabel( + self._seed_view_control, + text="", + text_color=Colors.TEXT_PRIMARY, + font=(Fonts.FAMILY, Fonts.SIZE_SMALL), + ) + self._seed_view_label.pack(side="right", padx=(0, 6)) projection_names = { "raw": t("segments_vr_projection_raw"), "fisheye": t("segments_vr_projection_fisheye"), @@ -621,6 +640,40 @@ def _build_restoration_inspector(self, range_panel: ctk.CTkFrame) -> None: self._use_for_all_btn.grid(row=2, column=0, columnspan=3, sticky="ew", padx=8, pady=(3, 8)) Tooltip(self._use_for_all_btn, t("tip_segments_use_for_all")) + self._seed_box = ctk.CTkFrame(self._inspector, fg_color="transparent") + self._seed_box.grid_columnconfigure(0, weight=1) + self._try_seed_btn = ctk.CTkButton( + self._seed_box, + text=t("segments_try_seed"), + height=28, + fg_color=Colors.MODEL_LTX, + hover_color=Colors.PRIMARY_HOVER, + command=self._try_seed, + ) + self._try_seed_btn.grid(row=0, column=0, sticky="ew") + self._try_seed_tooltip = Tooltip(self._try_seed_btn, t("tip_segments_try_seed")) + self._seed_cancel_btn = ctk.CTkButton( + self._seed_box, + text=t("btn_cancel"), + width=72, + height=28, + fg_color=Colors.BG_CARD, + hover_color=Colors.BORDER_LIGHT, + command=self._cancel_seed, + ) + self._seed_progress = ctk.CTkProgressBar(self._seed_box, height=8, progress_color=Colors.MODEL_LTX) + self._seed_status = ctk.CTkLabel( + self._seed_box, + text="", + font=(Fonts.FAMILY, Fonts.SIZE_TINY), + text_color=Colors.STATUS_PENDING, + anchor="w", + justify="left", + wraplength=320, + ) + self._seed_chips = ctk.CTkFrame(self._seed_box, fg_color="transparent") + self._seed_chip_widgets: list = [] + def _ltx_blocked_reason(self) -> str | None: """Locale key saying why a range cannot use LTX here, or None when it can.""" if self._ltx_unavailable_reason is not None: @@ -678,17 +731,232 @@ def _refresh_inspector(self) -> None: ) self._ltx_radio_tooltip.set_text(t(blocked) if blocked else t("model_ltx_description")) if restoration.model == "ltx": + self._seed_box.grid(row=3, column=0, columnspan=3, sticky="ew", padx=8, pady=(0, 8)) self._seed_label.grid(row=1, column=0, sticky="w", padx=(8, 6), pady=3) self._seed_entry.grid(row=1, column=1, sticky="ew", pady=3) self._new_seed_btn.grid(row=1, column=2, padx=(4, 8), pady=3) self._seed_entry.delete(0, "end") self._seed_entry.insert(0, str(restoration.ltx_seed)) else: - for widget in (self._seed_label, self._seed_entry, self._new_seed_btn): + for widget in (self._seed_label, self._seed_entry, self._new_seed_btn, self._seed_box): widget.grid_forget() state = "disabled" if locked else "normal" for widget in (self._seed_entry, self._new_seed_btn, self._use_for_all_btn): widget.configure(state=state) + self._refresh_seed_controls() + + @staticmethod + def _range_key(segment: SegmentRange) -> tuple[float, float]: + return (segment.start, segment.end) + + def _try_seed_block_reason(self) -> str | None: + """Locale key saying why Try seed is not available now, or None when it is.""" + blocked = self._ltx_blocked_reason() + if blocked is not None: + return blocked + if self._is_gpu_busy(): + return "segments_restore_gpu_busy" + if self._keyframe_index is None: + return "segments_loading_preview" + return None + + def _refresh_seed_controls(self) -> None: + selected = self._require_state().selected_segment + if selected is None or selected.restoration.model != "ltx": + self._refresh_seed_view_control() + return + running = self._seed_running is not None + reason = self._try_seed_block_reason() + self._try_seed_btn.configure( + state="disabled" if running or reason is not None or self._scanning() else "normal" + ) + self._try_seed_tooltip.set_text(t(reason) if reason else t("tip_segments_try_seed")) + if running: + self._seed_cancel_btn.grid(row=0, column=1, padx=(6, 0)) + self._seed_progress.grid(row=1, column=0, columnspan=2, sticky="ew", pady=(6, 0)) + else: + self._seed_cancel_btn.grid_forget() + self._seed_progress.grid_forget() + if self._seed_status.cget("text"): + self._seed_status.grid(row=2, column=0, columnspan=2, sticky="ew", pady=(4, 0)) + else: + self._seed_status.grid_forget() + self._render_seed_chips(selected) + self._refresh_seed_view_control() + + def _render_seed_chips(self, selected: SegmentRange) -> None: + for widget in self._seed_chip_widgets: + widget.destroy() + self._seed_chip_widgets = [] + tried = self._tried_seeds.get(self._range_key(selected), {}) + if not tried: + self._seed_chips.grid_forget() + return + self._seed_chips.grid(row=3, column=0, columnspan=2, sticky="ew", pady=(6, 0)) + label = ctk.CTkLabel( + self._seed_chips, + text=t("segments_seed_tried"), + font=(Fonts.FAMILY, Fonts.SIZE_TINY), + text_color=Colors.STATUS_PENDING, + ) + label.pack(side="left", padx=(0, 4)) + self._seed_chip_widgets.append(label) + shown = self._shown_seed_for(selected) + for seed in tried: + in_use = seed == selected.restoration.ltx_seed + chip = ctk.CTkButton( + self._seed_chips, + text=f"{seed} ✓" if in_use else str(seed), + width=10, + height=24, + font=(Fonts.FAMILY_MONO, Fonts.SIZE_TINY), + fg_color=Colors.MODEL_LTX if seed == shown else Colors.BG_CARD, + hover_color=Colors.BORDER_LIGHT, + command=lambda value=seed: self._show_tried_seed(value), + ) + chip.pack(side="left", padx=(0, 4)) + Tooltip(chip, t("tip_segments_seed_chip")) + self._seed_chip_widgets.append(chip) + if shown is not None and shown != selected.restoration.ltx_seed: + use = ctk.CTkButton( + self._seed_chips, + text=t("segments_seed_use"), + height=24, + font=(Fonts.FAMILY, Fonts.SIZE_TINY), + command=lambda value=shown: self._use_tried_seed(value), + ) + use.pack(side="left", padx=(4, 0)) + self._seed_chip_widgets.append(use) + + def _shown_seed_for(self, segment: SegmentRange) -> int | None: + if self._shown_seed is None or self._shown_seed[0] != self._range_key(segment): + return None + return self._shown_seed[1] + + def _refresh_seed_view_control(self) -> None: + if self._shown_seed is None: + self._seed_view_control.pack_forget() + return + self._seed_view_label.configure(text=t("segments_seed_show", seed=self._shown_seed[1])) + if self._seed_view_on: + self._seed_view_toggle.select() + else: + self._seed_view_toggle.deselect() + if not self._seed_view_control.winfo_manager(): + self._seed_view_control.pack(side="left") + + def _try_seed(self) -> None: + selected = self._require_state().selected_segment + if selected is None or selected.restoration.model != "ltx" or self._try_seed_block_reason() is not None: + return + self._claim_gpu_for_seed() + if self._seed_worker is None: + settings = self._get_settings() + work_dir = Path(settings.working_directory) if settings.working_directory else Path(self._job.path).parent + self._seed_worker = LtxSeedPreviewWorker( + self._job.path, + self._metadata, + self._keyframe_index, + work_dir, + on_stopped=lambda: self._set_preview_gpu_busy(False), + ) + self._seed_worker.start() + seed = int(selected.restoration.ltx_seed) + self._seed_generation = self._seed_worker.try_seed(selected, seed, self._current_video_settings()) + self._seed_running = (self._range_key(selected), seed) + self._seed_progress.set(0.0) + self._seed_status.configure(text=t("segments_seed_starting", seed=seed), text_color=Colors.STATUS_PENDING) + self._refresh_seed_controls() + + def _cancel_seed(self) -> None: + if self._seed_worker is not None: + self._seed_worker.cancel() + self._seed_running = None + self._seed_status.configure(text="") + self._refresh_seed_controls() + + def _claim_gpu_for_seed(self) -> None: + if self._restore_active: + self._deactivate_restoration_preview() + if self._restoration_worker is not None: + self._restoration_worker.close() + self._restoration_worker = None + self._set_playing(False) + self._set_preview_gpu_busy(True) + + def _release_seed_gpu(self) -> None: + if self._seed_worker is None: + return + self._seed_worker.release() + if self._seed_running is not None: + self._seed_running = None + self._seed_status.configure(text="") + if self._state is not None: + self._refresh_seed_controls() + + def _handle_seed_event(self, event) -> None: + if event.generation != self._seed_generation or self._seed_running is None: + return + if isinstance(event, SeedProgress): + self._seed_progress.set(event.fraction) + values = dict(seed=event.seed, stage=t(f"ltx_stage_{event.stage}"), percent=event.fraction * 100) + text = ( + t("segments_seed_progress", eta=format_duration(event.eta_seconds), **values) + if event.eta_seconds > 0 + else t("segments_seed_progress_no_eta", **values) + ) + self._seed_status.configure(text=text, text_color=Colors.STATUS_PENDING) + elif isinstance(event, SeedReady): + key, seed = self._seed_running + self._seed_running = None + self._tried_seeds.setdefault(key, {})[seed] = event.frames + self._seed_status.configure(text="") + self._shown_seed = (key, seed) + self._seed_view_on = True + if not key[0] <= self._current < key[1]: + self._seek(key[0]) + self._refresh_all() + self._refresh_preview_image() + return + elif isinstance(event, SeedFailed): + self._seed_running = None + self._seed_status.configure( + text=t("segments_seed_failed", message=event.message), text_color=Colors.STATUS_ERROR + ) + self._refresh_seed_controls() + + def _show_tried_seed(self, seed: int) -> None: + selected = self._require_state().selected_segment + if selected is None: + return + key = self._range_key(selected) + self._shown_seed = (key, seed) + self._seed_view_on = True + if not key[0] <= self._current < key[1]: + self._seek(key[0]) + self._refresh_seed_controls() + self._refresh_preview_image() + + def _use_tried_seed(self, seed: int) -> None: + self._apply_selected_restoration(SegmentRestoration("ltx", seed)) + + def _toggle_seed_view(self) -> None: + self._seed_view_on = not self._seed_view_on + self._refresh_seed_view_control() + self._refresh_preview_image() + + def _seed_view_image(self) -> Image.Image | None: + if not self._seed_view_on or self._shown_seed is None: + return None + key, seed = self._shown_seed + frames = self._tried_seeds.get(key, {}).get(seed, ()) + if not frames or not key[0] <= self._current < key[1]: + return None + frame = min(frames, key=lambda item: abs(item.seconds - self._current)) + if self._seed_image is None or self._seed_image[0] != frame: + with Image.open(frame.path) as image: + self._seed_image = (frame, image.convert("RGB")) + return self._seed_image[1] @staticmethod def _model_text(segment: SegmentRange) -> str: @@ -868,6 +1136,12 @@ def _poll_workers(self) -> None: self._handle_restoration_event(self._restoration_worker.events.get_nowait()) except queue.Empty: pass + if self._seed_worker is not None: + try: + while True: + self._handle_seed_event(self._seed_worker.events.get_nowait()) + except queue.Empty: + pass if self._scan_panel is not None: self._scan_panel.poll() try: @@ -901,6 +1175,9 @@ def _preview_resized(self, _event=None) -> None: self._resize_after = self.after(60, self._refresh_preview_image) def _active_preview_source(self) -> Image.Image | None: + seed_image = self._seed_view_image() + if seed_image is not None: + return seed_image return self._restored_source if self._restore_active else self._preview_source def _preview_image_geometry( @@ -1034,7 +1311,7 @@ def _refresh_preview_image(self) -> None: source = self._active_preview_source() if source is None or self._closed.is_set(): return - if not self._restore_active: + if source is self._preview_source: source = self._scan_panel.apply_overlay(source) source = self._preview_crop(source) self._preview_image = self._fit_to_label(self._preview, source) @@ -1150,6 +1427,7 @@ def _toggle_restoration_preview(self) -> None: if self._is_gpu_busy(): self._restore_toggle.deselect() return + self._release_seed_gpu() self._restore_active = True self._set_playing(False) self._restore_toggle.select() @@ -1329,6 +1607,7 @@ def _refresh_restore_toggle(self) -> None: ) def _claim_gpu_for_scan(self) -> None: + self._release_seed_gpu() if self._restore_active: self._deactivate_restoration_preview() if self._restoration_worker is not None: @@ -1359,6 +1638,7 @@ def _lockable_widgets(self) -> tuple: self._seed_entry, self._new_seed_btn, self._use_for_all_btn, + self._try_seed_btn, *self._timeline_zoom_buttons, *self._segment_action_widgets, ) @@ -1726,6 +2006,12 @@ def _request_close(self) -> None: if self._scanning(): self._scan_panel.stop() return + if self._seed_running is not None and not messagebox.askyesno( + t("segments_try_seed"), + t("segments_seed_stop_confirm"), + parent=self, + ): + return if ( not self._saved and self._state is not None @@ -1747,6 +2033,8 @@ def _finish_close(self) -> None: self._preview_worker.close() if self._restoration_worker is not None: self._restoration_worker.close() + if self._seed_worker is not None: + self._seed_worker.close() if self._scan_panel is not None: self._scan_panel.close() self.grab_release() diff --git a/jasna/ltx/restore.py b/jasna/ltx/restore.py index 612d0072..426400e9 100644 --- a/jasna/ltx/restore.py +++ b/jasna/ltx/restore.py @@ -109,6 +109,11 @@ def bar(self, stage: LtxStage, total: int, unit: str = "frame") -> _StageBar: ) return _StageBar(self, stage, tqdm_bar) + def skip(self, *stages: LtxStage) -> None: + """Count ``stages`` as done, for a run that reuses their results.""" + for stage in stages: + self._finished(stage) + def _advanced(self, stage: LtxStage, done: int, total: int) -> None: if self._report is None: return diff --git a/tests/test_ltx_seed_preview.py b/tests/test_ltx_seed_preview.py new file mode 100644 index 00000000..9c05e105 --- /dev/null +++ b/tests/test_ltx_seed_preview.py @@ -0,0 +1,119 @@ +from __future__ import annotations + +import threading +import time +from dataclasses import replace +from fractions import Fraction +from types import SimpleNamespace + +import torch + +from jasna.gui import ltx_seed_preview +from jasna.gui.ltx_seed_preview import ( + LtxSeedPreviewWorker, + SeedFrame, + SeedFrameWriter, + SeedProgress, + SeedReady, + prepared_key, +) +from jasna.gui.models import AppSettings +from jasna.ltx.restore import Cancelled +from jasna.segments import SegmentRange + +METADATA = SimpleNamespace(video_width=8, video_height=4, start_pts=0, time_base=Fraction(1, 100)) + + +def test_seed_frame_writer_keeps_only_the_range_frames(tmp_path) -> None: + writer = SeedFrameWriter(tmp_path, (100, 300), METADATA, None) + + for pts in (0, 100, 200, 300): + writer.write(torch.full((3, 4, 8), 200, dtype=torch.uint8), pts, apply_lut=pts != 0) + + assert writer.frames == [SeedFrame(1.0, tmp_path / "100.jpg"), SeedFrame(2.0, tmp_path / "200.jpg")] + assert all(frame.path.is_file() for frame in writer.frames) + + +def test_prepared_key_follows_detection_and_canvas_but_not_the_look() -> None: + segment = SegmentRange(1, 2) + settings = AppSettings() + + assert prepared_key(segment, replace(settings, lut_path="x.cube", ltx_seed=5)) == prepared_key(segment, settings) + assert prepared_key(segment, replace(settings, detection_score_threshold=0.9)) != prepared_key(segment, settings) + assert prepared_key(segment, replace(settings, ltx_large_canvas=True)) != prepared_key(segment, settings) + assert prepared_key(SegmentRange(1, 3), settings) != prepared_key(segment, settings) + + +class _Renderer: + def __init__(self, tmp_path, *, block: bool = False) -> None: + self.tmp_path = tmp_path + self.block = block + self.started = threading.Event() + self.released = 0 + self.closed = False + + def render(self, segment, seed, settings, *, writer_for, report, cancel): + self.started.set() + if self.block: + cancel.wait(5) + raise Cancelled() + report("denoise", 0.5, 12.0) + writer = writer_for(self.tmp_path, (0, 1000)) + writer.frames.append(SeedFrame(0.5, self.tmp_path / "50.jpg")) + return writer + + def release(self): + self.released += 1 + + def close(self): + self.closed = True + + +def _worker(monkeypatch, renderer) -> tuple[LtxSeedPreviewWorker, threading.Event]: + stopped = threading.Event() + worker = LtxSeedPreviewWorker( + "video.mp4", METADATA, SimpleNamespace(), renderer.tmp_path, on_stopped=stopped.set + ) + monkeypatch.setattr(worker, "_make_renderer", lambda: renderer) + worker.start() + return worker, stopped + + +def test_worker_reports_progress_then_the_seed_frames(monkeypatch, tmp_path) -> None: + renderer = _Renderer(tmp_path) + worker, stopped = _worker(monkeypatch, renderer) + + generation = worker.try_seed(SegmentRange(1, 2), 42, AppSettings()) + + assert worker.events.get(timeout=5) == SeedProgress(42, "denoise", 0.5, 12.0, generation) + assert worker.events.get(timeout=5) == SeedReady(42, (SeedFrame(0.5, tmp_path / "50.jpg"),), generation) + worker.release() + deadline = time.monotonic() + 5 + while not renderer.released and time.monotonic() < deadline: + time.sleep(0.01) + worker.close() + worker.join(timeout=5) + assert renderer.released == 1 and renderer.closed and stopped.is_set() + + +def test_cancel_stops_a_running_seed_without_a_result(monkeypatch, tmp_path) -> None: + renderer = _Renderer(tmp_path, block=True) + worker, stopped = _worker(monkeypatch, renderer) + + worker.try_seed(SegmentRange(1, 2), 42, AppSettings()) + assert renderer.started.wait(5) + worker.cancel() + worker.close() + worker.join(timeout=5) + + assert worker.events.empty() + assert renderer.closed and stopped.is_set() + + +def test_seed_renderer_deletes_its_temp_dir_on_close(tmp_path) -> None: + renderer = ltx_seed_preview.SeedRenderer(tmp_path / "video.mp4", METADATA, SimpleNamespace(), tmp_path) + (temp,) = tmp_path.glob(".jasna-seeds-*") + + renderer.close() + + assert not temp.exists() diff --git a/tests/test_ltx_seed_preview_gpu.py b/tests/test_ltx_seed_preview_gpu.py new file mode 100644 index 00000000..a0e6c5d4 --- /dev/null +++ b/tests/test_ltx_seed_preview_gpu.py @@ -0,0 +1,101 @@ +"""A seed preview restores a range exactly as a segment job does. + +Needs an NVIDIA GPU, the LTX model in model_weights/ltx-restore and a video with a +mosaic: set JASNA_LTX_TEST_VIDEO to its path and JASNA_LTX_TEST_RANGES to two ranges in +seconds, e.g. "9-11,21-24" (the first is previewed; the job restores both). +""" + +from __future__ import annotations + +import os +import threading +from pathlib import Path + +import pytest +import torch + +from jasna.engine_paths import default_restoration_model_path +from jasna.ltx.model_files import bundle_present + +VIDEO = os.environ.get("JASNA_LTX_TEST_VIDEO", "") + +pytestmark = pytest.mark.skipif( + not VIDEO or not torch.cuda.is_available() or not bundle_present(default_restoration_model_path("ltx")), + reason="needs JASNA_LTX_TEST_VIDEO, an NVIDIA GPU and the LTX model", +) + + +class _Capture: + def __init__(self, effect_range: tuple[int, int]) -> None: + self.start, self.end = effect_range + self.frames: dict[int, torch.Tensor] = {} + + def write(self, frame: torch.Tensor, pts: int, *, apply_lut: bool = True) -> None: + if self.start <= int(pts) < self.end: + self.frames[int(pts)] = frame.cpu() + + def close(self) -> None: + pass + + +def _ranges() -> list[tuple[float, float]]: + spec = os.environ.get("JASNA_LTX_TEST_RANGES", "9-11,21-24") + return [tuple(float(value) for value in part.split("-")) for part in spec.split(",")] + + +def test_seed_preview_restores_a_range_exactly_like_a_segment_job(tmp_path) -> None: + from jasna.gui.ltx_seed_preview import SeedRenderer + from jasna.gui.models import AppSettings + from jasna.gui.video_session import build_video_session, video_session_config + from jasna.media.probe import get_video_meta_data + from jasna.media.splice import build_splice_plan, probe_keyframes + from jasna.segments import SegmentRange, SegmentRestoration + from jasna.session_factory import build_pipeline + + video = Path(VIDEO) + (first_start, first_end), (second_start, second_end) = _ranges() + previewed = SegmentRange(first_start, first_end, SegmentRestoration("ltx", 1234)) + other = SegmentRange(second_start, second_end, SegmentRestoration("ltx", 99)) + settings = AppSettings() + metadata = get_video_meta_data(str(video)) + index = probe_keyframes(video, metadata) + + renderer = SeedRenderer(video, metadata, index, tmp_path) + try: + preview = renderer.render( + previewed, + 1234, + settings, + writer_for=lambda _directory, effect: _Capture(effect), + report=lambda *_args: None, + cancel=threading.Event(), + ) + finally: + renderer.close() + + session = build_video_session(settings, log=lambda _msg: None) + try: + config = video_session_config(settings, codec="hevc", encoder_settings={}) + pipeline = build_pipeline(config, session, video, tmp_path / "out.mkv", segments=(previewed, other)) + pipeline.configure_vr(metadata) + plan = build_splice_plan((previewed, other), index, duration=metadata.duration) + captures = [] + + def opener(span): + capture = _Capture(span.effect_ranges[0]) + captures.append(capture) + return lambda: capture + + pipeline._run_ltx_spans( + metadata, + index, + [(span, segments, opener(span)) for span, segments in zip(plan.render_spans, plan.render_span_segments())], + tmp_path, + ) + finally: + session.close() + + job = captures[0] + assert preview.frames and preview.frames.keys() == job.frames.keys() + differing = [pts for pts in job.frames if not torch.equal(job.frames[pts], preview.frames[pts])] + assert not differing, f"{len(differing)} of {len(job.frames)} frames differ" diff --git a/tests/test_segment_editor.py b/tests/test_segment_editor.py index 6816a997..1e2d5901 100644 --- a/tests/test_segment_editor.py +++ b/tests/test_segment_editor.py @@ -858,3 +858,97 @@ def test_ltx_is_blocked_with_a_reason_when_it_cannot_run(monkeypatch) -> None: if editor is not None: editor._finish_close() root.destroy() + + +class _FakeSeedWorker: + instances: list = [] + + def __init__(self, path, metadata, index, work_dir, *, on_stopped): + self.events = queue.Queue() + self.requests = [] + self.released = 0 + self.closed = False + self.work_dir = work_dir + _FakeSeedWorker.instances.append(self) + + def start(self): + pass + + def try_seed(self, segment, seed, settings): + self.requests.append((segment, seed, settings)) + return len(self.requests) + + def cancel(self): + pass + + def release(self): + self.released += 1 + + def close(self): + self.closed = True + + +def _seed_frames(tmp_path, seed, color): + frames = [] + for seconds in (1.0, 1.5): + path = tmp_path / f"{seed}-{seconds}.jpg" + Image.new("RGB", (160, 90), color).save(path) + frames.append(segment_editor.SeedFrame(seconds, path)) + return tuple(frames) + + +def test_try_seed_shows_progress_then_lets_the_user_compare_and_pick_seeds(monkeypatch, tmp_path) -> None: + root = _tk_root() + editor = None + _FakeSeedWorker.instances = [] + monkeypatch.setattr(segment_editor, "LtxSeedPreviewWorker", _FakeSeedWorker) + try: + job = JobItem(tmp_path / "video.mp4", segments=(SegmentRange(1.0, 2.0, SegmentRestoration("ltx", 5)),)) + editor = _build_editor_with_ui(root, monkeypatch, job=job) + editor._keyframe_index = KeyframeIndex(pts=(0,), time_base=Fraction(1, 1000), start_pts=0, end_pts=60_000) + editor._select_range(0) + editor._preview_source = Image.new("RGB", (160, 90), "red") + root.update() + assert editor._try_seed_btn.cget("state") == "normal" + + editor._try_seed_btn.invoke() + worker = _FakeSeedWorker.instances[0] + assert worker.work_dir == tmp_path + assert worker.requests[0][:2] == (job.segments[0], 5) + assert editor._try_seed_btn.cget("state") == "disabled" + worker.events.put(segment_editor.SeedProgress(5, "denoise", 0.5, 70.0, 1)) + editor._poll_workers() + assert editor._seed_status.cget("text") == t( + "segments_seed_progress", seed=5, stage=t("ltx_stage_denoise"), percent=50.0, eta="1m 10s" + ) + + worker.events.put(segment_editor.SeedReady(5, _seed_frames(tmp_path, 5, "blue"), 1)) + editor._poll_workers() + root.update() + assert editor._seed_view_on and 1.0 <= editor._current < 2.0 + assert editor._active_preview_source().getpixel((0, 0))[2] > 200 + assert editor._seed_view_label.cget("text") == t("segments_seed_show", seed=5) + + editor._new_range_seed() + new_seed = editor._state.selected_segment.restoration.ltx_seed + editor._try_seed_btn.invoke() + worker.events.put(segment_editor.SeedReady(new_seed, _seed_frames(tmp_path, new_seed, "green"), 2)) + editor._poll_workers() + chips = [w.cget("text") for w in editor._seed_chip_widgets if isinstance(w, ctk.CTkButton)] + assert chips == ["5", f"{new_seed} ✓"] + + editor._show_tried_seed(5) + assert editor._active_preview_source().getpixel((0, 0))[2] > 200 + editor._use_tried_seed(5) + assert editor._state.selected_segment.restoration == SegmentRestoration("ltx", 5) + + editor._toggle_seed_view() + assert editor._active_preview_source().getpixel((0, 0))[0] > 200 + + editor._claim_gpu_for_scan() + assert worker.released == 1 + finally: + if editor is not None: + editor._finish_close() + root.destroy() + assert _FakeSeedWorker.instances[0].closed From d969d96caec98cc80042232e599ff72888a650d4 Mon Sep 17 00:00:00 2001 From: Kruk2 Date: Tue, 29 Sep 2026 17:00:30 +0200 Subject: [PATCH 32/48] fix(ltx): make a seed give the same pixels on every run Inductor picked kernel configs by timing them, so the output changed between runs. The block forward now compiles in deterministic mode. --- jasna/ltx/sampler.py | 6 ++++-- tests/test_ltx_transformer.py | 15 +++++++++++++++ 2 files changed, 19 insertions(+), 2 deletions(-) diff --git a/jasna/ltx/sampler.py b/jasna/ltx/sampler.py index 30f430fc..09989b4a 100644 --- a/jasna/ltx/sampler.py +++ b/jasna/ltx/sampler.py @@ -184,12 +184,14 @@ def _denoise( def compiled_block_forward(): """``transformer.block_forward`` compiled with static shapes when Triton is available - (one graph per canvas and per STG variant, shared by all blocks), else eager.""" + (one graph per canvas and per STG variant, shared by all blocks), else eager. Kernels + are picked without timing them, so a seed gives the same pixels whatever else the GPU + is doing and wherever the block weights live.""" if importlib.util.find_spec("triton") is None: return T.block_forward torch._dynamo.config.cache_size_limit = max(torch._dynamo.config.cache_size_limit, 64) torch._dynamo.config.accumulated_cache_size_limit = max(torch._dynamo.config.accumulated_cache_size_limit, 1024) - return torch.compile(T.block_forward, dynamic=False) + return torch.compile(T.block_forward, dynamic=False, options={"deterministic": True}) def _euler(sample: torch.Tensor, denoised: torch.Tensor, sigmas: torch.Tensor, step: int) -> torch.Tensor: diff --git a/tests/test_ltx_transformer.py b/tests/test_ltx_transformer.py index cc4014f9..f1a79e98 100644 --- a/tests/test_ltx_transformer.py +++ b/tests/test_ltx_transformer.py @@ -158,3 +158,18 @@ def test_window_state_halves_without_stg(): from jasna.ltx import sampler assert sampler._window_state_bytes(100, stg=True) == 2 * sampler._window_state_bytes(100, stg=False) + + +def test_block_forward_compiles_without_timing_kernels(monkeypatch): + import importlib.util + + from jasna.ltx import sampler + + if importlib.util.find_spec("triton") is None: + pytest.skip("needs Triton") + + calls = [] + monkeypatch.setattr(torch, "compile", lambda fn, **kwargs: calls.append(kwargs) or fn) + + assert sampler.compiled_block_forward() is T.block_forward + assert calls == [{"dynamic": False, "options": {"deterministic": True}}] From f506391c2e9e7fae6d42edffe63bca296eab5b75 Mon Sep 17 00:00:00 2001 From: Kruk2 Date: Tue, 29 Sep 2026 17:04:46 +0200 Subject: [PATCH 33/48] fix(gui): keep the range list visible under the range inspector --- jasna/gui/segment_editor.py | 13 +++++++------ 1 file changed, 7 insertions(+), 6 deletions(-) diff --git a/jasna/gui/segment_editor.py b/jasna/gui/segment_editor.py index d7219243..b40c9701 100644 --- a/jasna/gui/segment_editor.py +++ b/jasna/gui/segment_editor.py @@ -584,9 +584,9 @@ def _build_restoration_inspector(self, range_panel: ctk.CTkFrame) -> None: text=t("segments_restore_with"), text_color=Colors.TEXT_PRIMARY, font=(Fonts.FAMILY, Fonts.SIZE_SMALL), - ).grid(row=0, column=0, sticky="w", padx=(8, 6), pady=(6, 3)) + ).grid(row=0, column=0, sticky="w", padx=(8, 6), pady=6) models = ctk.CTkFrame(self._inspector, fg_color="transparent") - models.grid(row=0, column=1, columnspan=2, sticky="w", pady=(6, 3)) + models.grid(row=0, column=1, sticky="w", pady=6) self._range_model = ctk.StringVar(value="basicvsrpp") self._model_radios = {} for model in ("basicvsrpp", "ltx"): @@ -631,13 +631,14 @@ def _build_restoration_inspector(self, range_panel: ctk.CTkFrame) -> None: self._use_for_all_btn = ctk.CTkButton( self._inspector, text=t("segments_use_for_all"), - height=26, + width=10, + height=24, fg_color=Colors.BG_CARD, hover_color=Colors.BORDER_LIGHT, font=(Fonts.FAMILY, Fonts.SIZE_SMALL), command=self._use_restoration_for_all, ) - self._use_for_all_btn.grid(row=2, column=0, columnspan=3, sticky="ew", padx=8, pady=(3, 8)) + self._use_for_all_btn.grid(row=0, column=2, sticky="e", padx=(4, 8), pady=6) Tooltip(self._use_for_all_btn, t("tip_segments_use_for_all")) self._seed_box = ctk.CTkFrame(self._inspector, fg_color="transparent") @@ -731,10 +732,10 @@ def _refresh_inspector(self) -> None: ) self._ltx_radio_tooltip.set_text(t(blocked) if blocked else t("model_ltx_description")) if restoration.model == "ltx": - self._seed_box.grid(row=3, column=0, columnspan=3, sticky="ew", padx=8, pady=(0, 8)) + self._seed_box.grid(row=2, column=0, columnspan=3, sticky="ew", padx=8, pady=(3, 8)) self._seed_label.grid(row=1, column=0, sticky="w", padx=(8, 6), pady=3) self._seed_entry.grid(row=1, column=1, sticky="ew", pady=3) - self._new_seed_btn.grid(row=1, column=2, padx=(4, 8), pady=3) + self._new_seed_btn.grid(row=1, column=2, sticky="w", padx=(4, 8), pady=3) self._seed_entry.delete(0, "end") self._seed_entry.insert(0, str(restoration.ltx_seed)) else: From d4474009e3cf6d4153c058cbbc621316315bfd38 Mon Sep 17 00:00:00 2001 From: Kruk2 Date: Tue, 29 Sep 2026 17:27:18 +0200 Subject: [PATCH 34/48] test(gui): build SegmentEditorState with a default restoration in scan tests --- tests/test_mosaic_scan.py | 10 ++++++---- 1 file changed, 6 insertions(+), 4 deletions(-) diff --git a/tests/test_mosaic_scan.py b/tests/test_mosaic_scan.py index 79ddd49f..c73abbca 100644 --- a/tests/test_mosaic_scan.py +++ b/tests/test_mosaic_scan.py @@ -9,7 +9,9 @@ segments_from_scores, ) from jasna.gui.segment_editor_state import SegmentEditorState -from jasna.segments import SegmentRange +from jasna.segments import SegmentRange, SegmentRestoration + +STANDARD = SegmentRestoration("basicvsrpp", None) def test_stride_follows_fps(): @@ -122,7 +124,7 @@ class FakeTorch: def test_add_many_is_single_undo_step(): - state = SegmentEditorState(duration=100.0, fps=25.0) + state = SegmentEditorState(duration=100.0, fps=25.0, segments=(), default_restoration=STANDARD) added = state.add_many((SegmentRange(1.0, 2.0), SegmentRange(5.0, 7.0))) assert added == 2 assert len(state.segments) == 2 @@ -131,12 +133,12 @@ def test_add_many_is_single_undo_step(): def test_add_many_skips_already_covered_ranges(): - state = SegmentEditorState(duration=100.0, fps=25.0) + state = SegmentEditorState(duration=100.0, fps=25.0, segments=(), default_restoration=STANDARD) state.add(0.0, 10.0) assert state.add_many((SegmentRange(2.0, 3.0),)) == 0 added = state.add_many((SegmentRange(2.0, 3.0), SegmentRange(20.0, 21.0))) assert added == 1 - assert state.segments == (SegmentRange(0.0, 10.0), SegmentRange(20.0, 21.0)) + assert state.segments == (SegmentRange(0.0, 10.0, STANDARD), SegmentRange(20.0, 21.0, STANDARD)) def test_scan_decoder_count_parallel_only_for_4k_on_nvidia(): From efed032f49b29aabd312bb2ade60d6509c929f53 Mon Sep 17 00:00:00 2001 From: Kruk2 Date: Tue, 29 Sep 2026 17:41:05 +0200 Subject: [PATCH 35/48] fix(gui): keep three range rows visible beside the range settings Start, End and the add button share a row when there is room, and the range settings scroll on their own. --- jasna/gui/segment_editor.py | 163 +++++++++++++++++++++++------------ tests/test_segment_editor.py | 64 ++++++++++++++ 2 files changed, 171 insertions(+), 56 deletions(-) diff --git a/jasna/gui/segment_editor.py b/jasna/gui/segment_editor.py index b40c9701..4b8924ff 100644 --- a/jasna/gui/segment_editor.py +++ b/jasna/gui/segment_editor.py @@ -16,7 +16,7 @@ from jasna.gui import scaling from jasna.gui.locales import t from jasna.gui.models import AppSettings, JobItem -from jasna.gui.components import Tooltip, format_duration, grab_modal +from jasna.gui.components import AutoHidingScrollableFrame, Tooltip, format_duration, grab_modal from jasna.gui.icons import CompactSwitch, NativeIconButton from jasna.gui.ltx_seed_preview import LtxSeedPreviewWorker, SeedFailed, SeedFrame, SeedProgress, SeedReady from jasna.gui.restoration_preview import ( @@ -68,6 +68,15 @@ logger = logging.getLogger(__name__) +RANGE_LIST_MIN_HEIGHT = 3 * 44 + 24 +RANGE_TIMES_ONE_ROW_MIN_WIDTH = 520 + + +def range_controls_height(shared: int, content: int, list_min: int) -> int: + """Height of the range controls under the range list, which share ``shared`` pixels: + all of ``content`` while the list keeps ``list_min``, else the rest (never under half).""" + return max(0, min(content, max(shared - list_min, shared // 2))) + class SegmentEditor(ctk.CTkToplevel): """Modal, frame-aware editor for per-job restoration ranges.""" @@ -478,6 +487,7 @@ def _build_range_panel(self, body: ctk.CTkFrame) -> None: range_panel.grid(row=0, column=1, sticky="nsew", padx=(6, 0)) range_panel.grid_columnconfigure(0, weight=1) range_panel.grid_rowconfigure(1, weight=1) + self._range_panel = range_panel range_header = ctk.CTkFrame(range_panel, fg_color="transparent") range_header.grid(row=0, column=0, sticky="ew", padx=10, pady=(8, 4)) ctk.CTkLabel( @@ -535,58 +545,86 @@ def _build_range_panel(self, body: ctk.CTkFrame) -> None: ) self._clear_btn.pack(side="right") - editor = ctk.CTkFrame(range_panel, fg_color="transparent") - editor.grid(row=3, column=0, sticky="ew", padx=8, pady=(4, 8)) - editor.grid_columnconfigure(1, weight=1) - ctk.CTkLabel( - editor, - text=t("segments_start"), - text_color=Colors.TEXT_PRIMARY, - ).grid(row=0, column=0, sticky="w", padx=(0, 6), pady=3) - self._start_entry = ctk.CTkEntry(editor, font=(Fonts.FAMILY_MONO, Fonts.SIZE_SMALL)) - self._start_entry.grid(row=0, column=1, sticky="ew", pady=3) - self._mark_in_btn = ctk.CTkButton( - editor, - text=t("segments_mark_in_short"), - width=42, - command=self._set_mark_in, - ) - self._mark_in_btn.grid(row=0, column=2, padx=(5, 0), pady=3) - Tooltip(self._mark_in_btn, t("segments_mark_in_hint")) - ctk.CTkLabel( - editor, - text=t("segments_end"), - text_color=Colors.TEXT_PRIMARY, - ).grid(row=1, column=0, sticky="w", padx=(0, 6), pady=3) - self._end_entry = ctk.CTkEntry(editor, font=(Fonts.FAMILY_MONO, Fonts.SIZE_SMALL)) - self._end_entry.grid(row=1, column=1, sticky="ew", pady=3) - self._mark_out_btn = ctk.CTkButton( - editor, - text=t("segments_mark_out_short"), - width=42, - command=self._set_mark_out, - ) - self._mark_out_btn.grid(row=1, column=2, padx=(5, 0), pady=3) - Tooltip(self._mark_out_btn, t("segments_mark_out_hint")) + self._range_controls = AutoHidingScrollableFrame(range_panel, fg_color="transparent", height=1) + self._range_controls.grid(row=3, column=0, sticky="nsew", padx=4, pady=(0, 4)) + self._range_controls.grid_columnconfigure(0, weight=1) + range_panel.bind("", self._fit_range_controls, add="+") + self._range_controls.bind("", self._fit_range_controls, add="+") + + editor = ctk.CTkFrame(self._range_controls, fg_color="transparent") + editor.grid(row=0, column=0, sticky="ew", padx=4, pady=(2, 4)) + editor.bind("", self._layout_range_times, add="+") + self._range_times: list[tuple[ctk.CTkLabel, ctk.CTkEntry, ctk.CTkButton]] = [] + for column, (label, entry_name, mark_name, mark_text, mark_hint, mark_command) in enumerate( + ( + ("segments_start", "_start_entry", "_mark_in_btn", "segments_mark_in_short", + "segments_mark_in_hint", self._set_mark_in), + ("segments_end", "_end_entry", "_mark_out_btn", "segments_mark_out_short", + "segments_mark_out_hint", self._set_mark_out), + ) + ): + text = ctk.CTkLabel(editor, text=t(label), text_color=Colors.TEXT_PRIMARY) + entry = ctk.CTkEntry(editor, width=40, font=(Fonts.FAMILY_MONO, Fonts.SIZE_SMALL)) + setattr(self, entry_name, entry) + mark = ctk.CTkButton(editor, text=t(mark_text), width=30, command=mark_command) + Tooltip(mark, t(mark_hint)) + setattr(self, mark_name, mark) + self._range_times.append((text, entry, mark)) self._range_action = ctk.CTkButton( editor, text=t("segments_add_range"), command=self._add_or_update, ) - self._range_action.grid(row=2, column=0, columnspan=3, sticky="ew", pady=(5, 0)) - self._build_restoration_inspector(range_panel) + self._range_times_one_row: bool | None = None + self._layout_range_times() + self._build_restoration_inspector(self._range_controls) - def _build_restoration_inspector(self, range_panel: ctk.CTkFrame) -> None: - self._inspector = ctk.CTkFrame(range_panel, fg_color=Colors.BG_PANEL, corner_radius=Sizing.BORDER_RADIUS) + def _layout_range_times(self, event=None) -> None: + """Start, End and the add/update button share one row when the panel is wide enough, else stack.""" + width = RANGE_TIMES_ONE_ROW_MIN_WIDTH if event is None else event.width + margin = 24 if self._range_times_one_row is False else 0 + one_row = width >= scaling.raw_tk_size(self, RANGE_TIMES_ONE_ROW_MIN_WIDTH + margin) + if one_row == self._range_times_one_row: + return + self._range_times_one_row = one_row + editor = self._range_action.master + for column in range(7): + editor.grid_columnconfigure(column, weight=0) + editor.grid_columnconfigure((1, 4) if one_row else 1, weight=1) + for index, (text, entry, mark) in enumerate(self._range_times): + row, first = (0, index * 3) if one_row else (index, 0) + text.grid(row=row, column=first, sticky="w", padx=(0 if first == 0 else 8, 4), pady=3) + entry.grid(row=row, column=first + 1, sticky="ew", pady=3) + mark.grid(row=row, column=first + 2, padx=(3, 0), pady=3) + if one_row: + self._range_action.grid(row=0, column=6, columnspan=1, sticky="ew", padx=(8, 0), pady=3) + else: + self._range_action.grid(row=2, column=0, columnspan=3, sticky="ew", padx=0, pady=(5, 0)) + + def _fit_range_controls(self, _event=None) -> None: + """Size the range controls to their content, but keep room for about three list rows.""" + panel = self._range_panel + fixed = panel.grid_bbox(0, 0)[3] + panel.grid_bbox(0, 2)[3] + scaling.raw_tk_size(panel, 4) + height = range_controls_height( + panel.winfo_height() - fixed, + self._range_controls.winfo_reqheight(), + scaling.raw_tk_size(panel, RANGE_LIST_MIN_HEIGHT), + ) + logical = max(1, round(height / scaling.widget_scaling(panel))) + if logical != self._range_controls.cget("height"): + self._range_controls.configure(height=logical) + + def _build_restoration_inspector(self, parent: ctk.CTkFrame) -> None: + self._inspector = ctk.CTkFrame(parent, fg_color=Colors.BG_PANEL, corner_radius=Sizing.BORDER_RADIUS) self._inspector.grid_columnconfigure(1, weight=1) ctk.CTkLabel( self._inspector, text=t("segments_restore_with"), text_color=Colors.TEXT_PRIMARY, font=(Fonts.FAMILY, Fonts.SIZE_SMALL), - ).grid(row=0, column=0, sticky="w", padx=(8, 6), pady=6) + ).grid(row=0, column=0, sticky="w", padx=(8, 6), pady=4) models = ctk.CTkFrame(self._inspector, fg_color="transparent") - models.grid(row=0, column=1, sticky="w", pady=6) + models.grid(row=0, column=1, sticky="w", pady=4) self._range_model = ctk.StringVar(value="basicvsrpp") self._model_radios = {} for model in ("basicvsrpp", "ltx"): @@ -638,25 +676,25 @@ def _build_restoration_inspector(self, range_panel: ctk.CTkFrame) -> None: font=(Fonts.FAMILY, Fonts.SIZE_SMALL), command=self._use_restoration_for_all, ) - self._use_for_all_btn.grid(row=0, column=2, sticky="e", padx=(4, 8), pady=6) + self._use_for_all_btn.grid(row=0, column=2, columnspan=2, sticky="e", padx=(4, 8), pady=4) Tooltip(self._use_for_all_btn, t("tip_segments_use_for_all")) self._seed_box = ctk.CTkFrame(self._inspector, fg_color="transparent") self._seed_box.grid_columnconfigure(0, weight=1) self._try_seed_btn = ctk.CTkButton( - self._seed_box, + self._inspector, text=t("segments_try_seed"), + width=80, height=28, fg_color=Colors.MODEL_LTX, hover_color=Colors.PRIMARY_HOVER, command=self._try_seed, ) - self._try_seed_btn.grid(row=0, column=0, sticky="ew") self._try_seed_tooltip = Tooltip(self._try_seed_btn, t("tip_segments_try_seed")) self._seed_cancel_btn = ctk.CTkButton( - self._seed_box, + self._inspector, text=t("btn_cancel"), - width=72, + width=80, height=28, fg_color=Colors.BG_CARD, hover_color=Colors.BORDER_LIGHT, @@ -721,7 +759,7 @@ def _refresh_inspector(self) -> None: if selected is None: self._inspector.grid_forget() return - self._inspector.grid(row=4, column=0, sticky="ew", padx=8, pady=(0, 8)) + self._inspector.grid(row=1, column=0, sticky="ew", padx=4, pady=(0, 4)) restoration = selected.restoration self._range_model.set(restoration.model) blocked = self._ltx_blocked_reason() @@ -732,14 +770,20 @@ def _refresh_inspector(self) -> None: ) self._ltx_radio_tooltip.set_text(t(blocked) if blocked else t("model_ltx_description")) if restoration.model == "ltx": - self._seed_box.grid(row=2, column=0, columnspan=3, sticky="ew", padx=8, pady=(3, 8)) - self._seed_label.grid(row=1, column=0, sticky="w", padx=(8, 6), pady=3) - self._seed_entry.grid(row=1, column=1, sticky="ew", pady=3) - self._new_seed_btn.grid(row=1, column=2, sticky="w", padx=(4, 8), pady=3) + self._seed_label.grid(row=1, column=0, sticky="w", padx=(8, 6), pady=(0, 6)) + self._seed_entry.grid(row=1, column=1, sticky="ew", pady=(0, 6)) + self._new_seed_btn.grid(row=1, column=2, sticky="w", padx=4, pady=(0, 6)) self._seed_entry.delete(0, "end") self._seed_entry.insert(0, str(restoration.ltx_seed)) else: - for widget in (self._seed_label, self._seed_entry, self._new_seed_btn, self._seed_box): + for widget in ( + self._seed_label, + self._seed_entry, + self._new_seed_btn, + self._try_seed_btn, + self._seed_cancel_btn, + self._seed_box, + ): widget.grid_forget() state = "disabled" if locked else "normal" for widget in (self._seed_entry, self._new_seed_btn, self._use_for_all_btn): @@ -772,17 +816,24 @@ def _refresh_seed_controls(self) -> None: state="disabled" if running or reason is not None or self._scanning() else "normal" ) self._try_seed_tooltip.set_text(t(reason) if reason else t("tip_segments_try_seed")) + seed_button, hidden_button = ( + (self._seed_cancel_btn, self._try_seed_btn) if running else (self._try_seed_btn, self._seed_cancel_btn) + ) + hidden_button.grid_forget() + seed_button.grid(row=1, column=3, sticky="e", padx=(0, 8), pady=(0, 6)) if running: - self._seed_cancel_btn.grid(row=0, column=1, padx=(6, 0)) - self._seed_progress.grid(row=1, column=0, columnspan=2, sticky="ew", pady=(6, 0)) + self._seed_progress.grid(row=0, column=0, sticky="ew", pady=(2, 0)) else: - self._seed_cancel_btn.grid_forget() self._seed_progress.grid_forget() if self._seed_status.cget("text"): - self._seed_status.grid(row=2, column=0, columnspan=2, sticky="ew", pady=(4, 0)) + self._seed_status.grid(row=1, column=0, sticky="ew", pady=(2, 0)) else: self._seed_status.grid_forget() self._render_seed_chips(selected) + if self._seed_box.grid_slaves(): + self._seed_box.grid(row=2, column=0, columnspan=4, sticky="ew", padx=8, pady=(0, 6)) + else: + self._seed_box.grid_forget() self._refresh_seed_view_control() def _render_seed_chips(self, selected: SegmentRange) -> None: @@ -793,7 +844,7 @@ def _render_seed_chips(self, selected: SegmentRange) -> None: if not tried: self._seed_chips.grid_forget() return - self._seed_chips.grid(row=3, column=0, columnspan=2, sticky="ew", pady=(6, 0)) + self._seed_chips.grid(row=2, column=0, sticky="ew", pady=(4, 0)) label = ctk.CTkLabel( self._seed_chips, text=t("segments_seed_tried"), diff --git a/tests/test_segment_editor.py b/tests/test_segment_editor.py index 1e2d5901..cbf2128e 100644 --- a/tests/test_segment_editor.py +++ b/tests/test_segment_editor.py @@ -952,3 +952,67 @@ def test_try_seed_shows_progress_then_lets_the_user_compare_and_pick_seeds(monke editor._finish_close() root.destroy() assert _FakeSeedWorker.instances[0].closed + + +def test_range_controls_fit_their_content_while_the_list_keeps_its_rows() -> None: + assert segment_editor.range_controls_height(500, 200, 156) == 200 + assert segment_editor.range_controls_height(320, 200, 156) == 164 + assert segment_editor.range_controls_height(200, 200, 156) == 100 + + +def _visible_in(widget, viewport) -> bool: + top, bottom = viewport.winfo_rooty(), viewport.winfo_rooty() + viewport.winfo_height() + return widget.winfo_ismapped() and top <= widget.winfo_rooty() and widget.winfo_rooty() + widget.winfo_height() <= bottom + + +def test_range_list_keeps_three_rows_beside_a_busy_seed_inspector(monkeypatch) -> None: + root = _tk_root() + editor = None + try: + editor = _build_editor_with_ui(root, monkeypatch) + editor.geometry("1856x880") + for start in (1.0, 4.0, 7.0, 10.0): + editor._timeline_create(start, start + 1.0) + editor._select_range(0) + editor._model_radios["ltx"].invoke() + key = editor._range_key(editor._state.selected_segment) + editor._tried_seeds[key] = {1: (), 2: ()} + editor._shown_seed = (key, 2) + editor._seed_running = (key, 1) + editor._seed_status.configure(text="Seed 1: Restoring 30% · about 15s left") + editor._refresh_seed_controls() + for _ in range(5): + root.update() + + rows = editor._segment_list.winfo_children() + viewport = editor._segment_list._parent_canvas + assert all(_visible_in(row, viewport) for row in rows[:3]) + assert editor._range_controls._parent_canvas.yview() == (0.0, 1.0) + assert _visible_in(editor._seed_chips, editor._range_controls._parent_canvas) + assert editor._range_action.grid_info()["row"] == 0 + finally: + if editor is not None: + editor._seed_running = None + editor._finish_close() + root.destroy() + + +def test_range_times_stack_and_controls_scroll_at_minimum_size(monkeypatch) -> None: + root = _tk_root() + editor = None + try: + editor = _build_editor_with_ui(root, monkeypatch) + editor.geometry("900x640") + editor._timeline_create(1.0, 2.0) + editor._model_radios["ltx"].invoke() + for _ in range(5): + root.update() + + assert editor._end_entry.grid_info()["row"] == 1 + assert editor._range_action.grid_info()["row"] == 2 + assert editor._range_controls._parent_canvas.yview()[1] < 1.0 + assert editor._segment_list._parent_canvas.winfo_height() > 0 + finally: + if editor is not None: + editor._finish_close() + root.destroy() From 06c2d8b4c5ba661247da0553a0f5242c046bb6c3 Mon Sep 17 00:00:00 2001 From: Kruk2 Date: Tue, 29 Sep 2026 17:43:12 +0200 Subject: [PATCH 36/48] feat(progress): show one bar for jobs that mix LTX and standard ranges An LTX frame counts as about 30 standard frames, and time left comes from the whole job. --- jasna/pipeline.py | 44 ++++++++++++++++++------------ jasna/progressbar.py | 34 +++++++++++++++++++++++ tests/test_ltx_seed_preview_gpu.py | 1 + tests/test_pipeline_init.py | 2 +- tests/test_pipeline_segments.py | 25 +++++++++++++++-- tests/test_progressbar.py | 26 ++++++++++++++++++ 6 files changed, 112 insertions(+), 20 deletions(-) diff --git a/jasna/pipeline.py b/jasna/pipeline.py index 26ca4d9c..6ccd48d0 100644 --- a/jasna/pipeline.py +++ b/jasna/pipeline.py @@ -32,7 +32,7 @@ validate_smart_render, ) from jasna.pipeline_threads import run_restoration_pass -from jasna.progressbar import ProgressCallback, Progressbar +from jasna.progressbar import LTX_WORK_PER_FRAME, JobProgress, ProgressCallback, Progressbar from jasna.restorer.secondary_restorer import AsyncSecondaryRestorer from jasna.segments import SegmentRange, job_restoration, resolve_restorations from jasna.session_config import SessionConfig @@ -49,6 +49,10 @@ log = logging.getLogger(__name__) +def _span_frames(span: SpliceSpan, index: KeyframeIndex, metadata) -> int: + return round((span.end_pts - span.start_pts) * index.time_base * metadata.video_fps) + + class _OfflineFrameWriter: def __init__(self, encoder_ctx: VideoEncoder, encode_heartbeat: list[float]): self._encoder_ctx = encoder_ctx @@ -279,10 +283,9 @@ def _run_full( finally: progress.close(ensure_completed_bar=True) - def _ltx_progress(self, frames: int): + def _ltx_progress(self, frames: int, callback: ProgressCallback | None): from jasna.ltx.restore import Progress - callback = self.progress_callback report = None if callback is not None: report = lambda stage, fraction, eta: callback(fraction * 100.0, 0.0, eta, 0, 0, stage) @@ -319,7 +322,9 @@ def _run_ltx(self, metadata) -> None: seed=self.ltx_seed, device=self.device, work_dir=self.working_dir or self.output_video.parent, - progress=self._ltx_progress(frame_rate.output_frame_count(metadata.num_frames)), + progress=self._ltx_progress( + frame_rate.output_frame_count(metadata.num_frames), self.progress_callback + ), cancel=self._cancel_event, ) except Cancelled: @@ -382,20 +387,19 @@ def _run_ltx_spans( index: KeyframeIndex, spans: list[tuple[SpliceSpan, tuple[SegmentRange, ...], Callable[[], _OfflineFrameWriter]]], work_dir: Path, + progress_callback: ProgressCallback | None, ) -> None: """Restore the LTX render spans in one batched run, so every model loads once.""" from jasna.ltx.restore import Cancelled, restore_spans render = self.ltx_segment_render(metadata) - frames = sum( - round((span.end_pts - span.start_pts) * index.time_base * metadata.video_fps) for span, _, _ in spans - ) + frames = sum(_span_frames(span, index, metadata) for span, _, _ in spans) try: restore_spans( [self.ltx_span(metadata, index, span, segments, open_writer) for span, segments, open_writer in spans], render, work_dir=work_dir, - progress=self._ltx_progress(frames), + progress=self._ltx_progress(frames, progress_callback), cancel=self._cancel_event, ) except Cancelled: @@ -455,19 +459,24 @@ def _run_smart(self, metadata) -> None: self.encoder_settings, vendor=vendor_for_device(self.device), ) - total_frames = max( - 1, - sum( - round((span.end_pts - span.start_pts) * index.time_base * metadata.video_fps) - for span in plan.render_spans - if span not in ltx_spans - ), + standard_frames = sum( + _span_frames(span, index, metadata) for span in plan.render_spans if span not in ltx_spans ) + ltx_callback = standard_callback = self.progress_callback + if self.progress_callback is not None and ltx_spans and standard_frames: + job_progress = JobProgress( + self.progress_callback, + { + "ltx": LTX_WORK_PER_FRAME * sum(_span_frames(span, index, metadata) for span in ltx_spans), + "standard": float(standard_frames), + }, + ) + ltx_callback, standard_callback = job_progress.part("ltx"), job_progress.part("standard") progress = Progressbar( - total_frames=total_frames, + total_frames=max(1, standard_frames), video_fps=metadata.video_fps, disable=self.disable_progress, - callback=self.progress_callback, + callback=standard_callback, ) self.output_video.parent.mkdir(parents=True, exist_ok=True) work_root = self.working_dir or self.output_video.parent @@ -516,6 +525,7 @@ def fragment_writer(span_index: int, span: SpliceSpan) -> Callable[[], _OfflineF if span in ltx_spans ], work_root, + ltx_callback, ) for span_index, span in enumerate(plan.spans): if self._cancel_event.is_set(): diff --git a/jasna/progressbar.py b/jasna/progressbar.py index e22ab214..871a6b44 100644 --- a/jasna/progressbar.py +++ b/jasna/progressbar.py @@ -9,6 +9,40 @@ """``(percent, fps, eta_seconds, frames_done, total_frames, stage)``; ``stage`` names the LTX pass running, or is empty for frame-by-frame restoration.""" +LTX_WORK_PER_FRAME = 30.0 +"""Rough cost of one LTX-restored frame in frames of the standard model.""" + + +class JobProgress: + """One bar for a job made of parts that report on their own (e.g. LTX ranges, then + standard ranges). Each part fills its share of the bar by its estimated ``work``, in + the order given; time left comes from the whole job's elapsed time and fraction, so + it does not jump when the next part starts.""" + + def __init__(self, callback: ProgressCallback, work: dict[str, float]) -> None: + self._callback = callback + total = sum(work.values()) + self._shares: dict[str, tuple[float, float]] = {} + offset = 0.0 + for part, amount in work.items(): + share = amount / total if total > 0 else 0.0 + self._shares[part] = (offset, share) + offset += share + self._started = time.monotonic() + self._timed = False + + def part(self, name: str) -> ProgressCallback: + offset, share = self._shares[name] + + def report(percent: float, fps: float, eta_seconds: float, done: int, total: int, stage: str) -> None: + fraction = min(1.0, offset + share * percent / 100.0) + self._timed = self._timed or eta_seconds > 0 + elapsed = time.monotonic() - self._started + eta = elapsed * (1.0 - fraction) / fraction if self._timed and fraction > 0 else 0.0 + self._callback(fraction * 100.0, fps, eta, done, total, stage) + + return report + class Progressbar: """Progress bar with time remaining estimation and speed display.""" diff --git a/tests/test_ltx_seed_preview_gpu.py b/tests/test_ltx_seed_preview_gpu.py index a0e6c5d4..2c069904 100644 --- a/tests/test_ltx_seed_preview_gpu.py +++ b/tests/test_ltx_seed_preview_gpu.py @@ -91,6 +91,7 @@ def opener(span): index, [(span, segments, opener(span)) for span, segments in zip(plan.render_spans, plan.render_span_segments())], tmp_path, + None, ) finally: session.close() diff --git a/tests/test_pipeline_init.py b/tests/test_pipeline_init.py index d3aae656..7c4c7316 100644 --- a/tests/test_pipeline_init.py +++ b/tests/test_pipeline_init.py @@ -49,7 +49,7 @@ def test_progress_callback(self): def test_ltx_progress_reaches_the_progress_callback_with_its_stage(self): cb = MagicMock() p = _make_pipeline(progress_callback=cb, disable_progress=True) - with p._ltx_progress(10).bar("scan", 10) as bar: + with p._ltx_progress(10, cb).bar("scan", 10) as bar: bar.update(10) assert cb.call_args[0] == (pytest.approx(2.0), 0.0, 0.0, 0, 0, "scan") diff --git a/tests/test_pipeline_segments.py b/tests/test_pipeline_segments.py index 42e34b22..e7908f5b 100644 --- a/tests/test_pipeline_segments.py +++ b/tests/test_pipeline_segments.py @@ -312,12 +312,33 @@ def test_mixed_job_batches_ltx_spans_and_keeps_fragments_in_span_order(tmp_path) pipeline._run_pass.assert_called_once() assert pipeline._run_pass.call_args.kwargs["effect_ranges"] == ((75, 90),) - (_, _, ltx_spans, _), _ = pipeline._run_ltx_spans.call_args + (_, _, ltx_spans, _, _), _ = pipeline._run_ltx_spans.call_args assert [(span, segs) for span, segs, _ in ltx_spans] == [(spans[3], (SegmentRange(6.5, 7.0, ltx),))] fragments = concatenate.call_args.args[0] assert [path.name for path, _ in fragments] == ["0000.ts", "0001.ts", "0002.ts", "0003.ts"] + +def test_mixed_job_reports_one_bar_weighted_by_estimated_work(tmp_path) -> None: + ltx = SegmentRestoration("ltx", 42) + segments = (SegmentRange(2.5, 3.0), SegmentRange(6.5, 7.0, ltx)) + spans = ( + SpliceSpan("copy", 0, 60), + SpliceSpan("render", 60, 120, ((75, 90),)), + SpliceSpan("copy", 120, 180), + SpliceSpan("render", 180, 240, ((195, 210),)), + ) + pipeline = _mixed_pipeline(tmp_path, "basicvsrpp", segments, spans) + pipeline.progress_callback = MagicMock() + + _run_smart_mocked(pipeline) + + (*_, ltx_callback), _ = pipeline._run_ltx_spans.call_args + ltx_callback(100.0, 0.0, 0.0, 0, 0, "compose") + ltx_share = 30.0 * 60 / (30.0 * 60 + 60) + assert pipeline.progress_callback.call_args.args[0] == pytest.approx(100.0 * ltx_share) + assert pipeline._run_pass.call_args.kwargs["progress"].callback is not pipeline.progress_callback + def test_segments_on_the_job_model_follow_an_ltx_job(tmp_path) -> None: segments = (SegmentRange(2.5, 3.0),) spans = (SpliceSpan("copy", 0, 60), SpliceSpan("render", 60, 120, ((75, 90),)), SpliceSpan("copy", 120, 240)) @@ -326,7 +347,7 @@ def test_segments_on_the_job_model_follow_an_ltx_job(tmp_path) -> None: _run_smart_mocked(pipeline) pipeline._run_pass.assert_not_called() - (_, _, ltx_spans, _), _ = pipeline._run_ltx_spans.call_args + (_, _, ltx_spans, _, _), _ = pipeline._run_ltx_spans.call_args assert ltx_spans[0][1] == (SegmentRange(2.5, 3.0, SegmentRestoration("ltx", 5)),) diff --git a/tests/test_progressbar.py b/tests/test_progressbar.py index b88591ba..fd541d2a 100644 --- a/tests/test_progressbar.py +++ b/tests/test_progressbar.py @@ -92,3 +92,29 @@ def test_speed_display_after_enough_data(self): for _ in range(5): pb.update(1) assert "Speed:" in pb.tqdm.desc or "?" in pb.tqdm.desc + + +class TestJobProgress: + def test_parts_fill_their_work_share_of_one_bar(self, monkeypatch): + from jasna import progressbar + + clock = [100.0] + monkeypatch.setattr(progressbar.time, "monotonic", lambda: clock[0]) + calls = [] + job = progressbar.JobProgress(lambda *args: calls.append(args), {"ltx": 300.0, "standard": 100.0}) + ltx, standard = job.part("ltx"), job.part("standard") + + clock[0] = 110.0 + ltx(50.0, 0.0, 0.0, 0, 0, "scan") + assert calls[-1] == (37.5, 0.0, 0.0, 0, 0, "scan") + + clock[0] = 130.0 + ltx(100.0, 0.0, 5.0, 0, 0, "compose") + assert calls[-1][0] == 75.0 + assert calls[-1][2] == pytest.approx(10.0) + + clock[0] = 134.0 + standard(50.0, 30.0, 0.0, 60, 120, "") + assert calls[-1][:2] == (87.5, 30.0) + assert calls[-1][2] == pytest.approx(34.0 / 7.0) + assert calls[-1][3:] == (60, 120, "") From 86d6eb807bda946d009be457bedfa2c66247c1e7 Mon Sep 17 00:00:00 2001 From: Kruk2 Date: Tue, 29 Sep 2026 17:47:19 +0200 Subject: [PATCH 37/48] feat(gui): name the default restoration model BasicVSR++ instead of Standard --- jasna/gui/locales/en.py | 4 ++-- jasna/gui/locales/ja.py | 4 ++-- jasna/gui/locales/ko.py | 4 ++-- jasna/gui/locales/th.py | 4 ++-- jasna/gui/locales/zh.py | 4 ++-- jasna/gui/settings_sections/restoration_model.py | 2 +- 6 files changed, 11 insertions(+), 11 deletions(-) diff --git a/jasna/gui/locales/en.py b/jasna/gui/locales/en.py index 9b322272..239d5a42 100644 --- a/jasna/gui/locales/en.py +++ b/jasna/gui/locales/en.py @@ -198,7 +198,7 @@ "segments_smart_render_range_too_short": "A selected range is shorter than one video frame. Make it longer.", "segments_smart_render_before_first_keyframe": "The first range starts before the first frame where this video can be split safely. Move its start slightly later.", "segments_smart_render_whole_video": "These ranges leave no safe places to split the video, so processing them separately would re-encode the whole video. Select a smaller area, or clear all ranges to process the full video.", - "segments_mixed_models_too_close": "A Standard range and an LTX range are too close. Move them apart or use one model for both.", + "segments_mixed_models_too_close": "A BasicVSR++ range and an LTX range are too close. Move them apart or use one model for both.", "segments_discard_title": "Discard range changes?", "segments_discard_changes": "Close without applying your restoration-range changes?", "segments_edit_tooltip": "Choose which parts of this video should be restored", @@ -214,7 +214,7 @@ # Sections "section_basic": "Basic Processing", "section_restoration_model": "Restoration Model", - "model_basicvsrpp": "Standard", + "model_basicvsrpp": "BasicVSR++", "model_basicvsrpp_description": "Fast. Good for most videos.", "model_ltx": "LTX", "model_ltx_description": "Best detail. Much slower and needs a strong NVIDIA GPU.", diff --git a/jasna/gui/locales/ja.py b/jasna/gui/locales/ja.py index 6504c0d4..83a1eca3 100644 --- a/jasna/gui/locales/ja.py +++ b/jasna/gui/locales/ja.py @@ -198,7 +198,7 @@ "segments_smart_render_range_too_short": "選択区間が1フレームより短いため、長くしてください。", "segments_smart_render_before_first_keyframe": "最初の区間が安全なカット位置より前から始まっています。開始位置を少し後ろへ移動してください。", "segments_smart_render_whole_video": "使用できる安全なカット位置が残らないため、個別処理すると動画全体の再エンコードが必要です。範囲を小さくするか、全区間を消して動画全体を処理してください。", - "segments_mixed_models_too_close": "「標準」の範囲と LTX の範囲が近すぎます。離すか、両方を同じモデルにしてください。", + "segments_mixed_models_too_close": "BasicVSR++ の範囲と LTX の範囲が近すぎます。離すか、両方を同じモデルにしてください。", "segments_discard_title": "区間の変更を破棄しますか？", "segments_discard_changes": "復元区間の変更を適用せずに閉じますか？", "segments_edit_tooltip": "この動画で復元する部分を選択", @@ -214,7 +214,7 @@ # Sections "section_basic": "基本設定", "section_restoration_model": "修復モデル", - "model_basicvsrpp": "標準", + "model_basicvsrpp": "BasicVSR++", "model_basicvsrpp_description": "高速。ほとんどの動画に向いています。", "model_ltx": "LTX", "model_ltx_description": "最も精細。かなり遅く、高性能な NVIDIA GPU が必要です。", diff --git a/jasna/gui/locales/ko.py b/jasna/gui/locales/ko.py index dfeeb97a..d0767dd0 100644 --- a/jasna/gui/locales/ko.py +++ b/jasna/gui/locales/ko.py @@ -197,7 +197,7 @@ "segments_smart_render_range_too_short": "선택 구간이 비디오 한 프레임보다 짧습니다. 구간을 늘리세요.", "segments_smart_render_before_first_keyframe": "첫 구간이 안전한 자르기 지점보다 먼저 시작합니다. 시작 위치를 조금 뒤로 옮기세요.", "segments_smart_render_whole_video": "사용 가능한 안전한 자르기 지점이 남지 않아 별도 처리 시 전체 비디오를 다시 인코딩해야 합니다. 더 작은 영역을 선택하거나 모든 구간을 지워 전체 비디오를 처리하세요.", - "segments_mixed_models_too_close": "'표준' 구간과 LTX 구간이 너무 가깝습니다. 떨어뜨리거나 둘 다 같은 모델을 쓰세요.", + "segments_mixed_models_too_close": "BasicVSR++ 구간과 LTX 구간이 너무 가깝습니다. 떨어뜨리거나 둘 다 같은 모델을 쓰세요.", "segments_discard_title": "구간 변경 사항을 버릴까요?", "segments_discard_changes": "복원 구간 변경 사항을 적용하지 않고 닫을까요?", "segments_edit_tooltip": "이 비디오에서 복원할 부분 선택", @@ -213,7 +213,7 @@ # Sections "section_basic": "기본 처리", "section_restoration_model": "복원 모델", - "model_basicvsrpp": "표준", + "model_basicvsrpp": "BasicVSR++", "model_basicvsrpp_description": "빠릅니다. 대부분의 영상에 적합합니다.", "model_ltx": "LTX", "model_ltx_description": "가장 세밀합니다. 훨씬 느리고 고성능 NVIDIA GPU가 필요합니다.", diff --git a/jasna/gui/locales/th.py b/jasna/gui/locales/th.py index a8ec09db..01def76e 100644 --- a/jasna/gui/locales/th.py +++ b/jasna/gui/locales/th.py @@ -198,7 +198,7 @@ "segments_smart_render_range_too_short": "ช่วงที่เลือกสั้นกว่าหนึ่งเฟรมวิดีโอ โปรดทำให้ยาวขึ้น", "segments_smart_render_before_first_keyframe": "ช่วงแรกเริ่มก่อนจุดตัดที่ปลอดภัยของวิดีโอ โปรดเลื่อนจุดเริ่มไปข้างหน้าเล็กน้อย", "segments_smart_render_whole_video": "ช่วงเหล่านี้ไม่เหลือจุดตัดที่ปลอดภัยให้ใช้ การประมวลผลแยกจะต้องเข้ารหัสวิดีโอทั้งไฟล์ใหม่ โปรดเลือกพื้นที่เล็กลงหรือล้างทุกช่วงเพื่อประมวลผลวิดีโอทั้งไฟล์", - "segments_mixed_models_too_close": "ช่วง “มาตรฐาน” กับช่วง LTX อยู่ใกล้กันเกินไป ให้แยกออกจากกันหรือใช้โมเดลเดียวกันทั้งสองช่วง", + "segments_mixed_models_too_close": "ช่วง BasicVSR++ กับช่วง LTX อยู่ใกล้กันเกินไป ให้แยกออกจากกันหรือใช้โมเดลเดียวกันทั้งสองช่วง", "segments_discard_title": "ละทิ้งการเปลี่ยนแปลงช่วงหรือไม่", "segments_discard_changes": "ปิดโดยไม่ใช้การเปลี่ยนแปลงช่วงการกู้คืนหรือไม่", "segments_edit_tooltip": "เลือกส่วนของวิดีโอนี้ที่จะกู้คืน", @@ -214,7 +214,7 @@ # Sections "section_basic": "การประมวลผลพื้นฐาน", "section_restoration_model": "โมเดลฟื้นฟู", - "model_basicvsrpp": "มาตรฐาน", + "model_basicvsrpp": "BasicVSR++", "model_basicvsrpp_description": "เร็ว เหมาะกับวิดีโอส่วนใหญ่", "model_ltx": "LTX", "model_ltx_description": "รายละเอียดดีที่สุด ช้ากว่ามาก และต้องใช้ GPU NVIDIA ที่แรง", diff --git a/jasna/gui/locales/zh.py b/jasna/gui/locales/zh.py index d048aa35..183aee5b 100644 --- a/jasna/gui/locales/zh.py +++ b/jasna/gui/locales/zh.py @@ -198,7 +198,7 @@ "segments_smart_render_range_too_short": "所选区间短于一个视频帧，请将其延长。", "segments_smart_render_before_first_keyframe": "第一个区间开始于视频的安全切点之前，请将开始时间稍微后移。", "segments_smart_render_whole_video": "这些区间没有留下可用的安全切点，单独处理会导致整个视频重新编码。请选择更小的区域，或清除所有区间以处理完整视频。", - "segments_mixed_models_too_close": "“标准”片段和 LTX 片段离得太近。请把它们分开，或两者使用同一模型。", + "segments_mixed_models_too_close": "BasicVSR++ 片段和 LTX 片段离得太近。请把它们分开，或两者使用同一模型。", "segments_discard_title": "放弃区间更改？", "segments_discard_changes": "不应用修复区间更改并关闭吗？", "segments_edit_tooltip": "选择此视频中需要修复的部分", @@ -214,7 +214,7 @@ # Sections "section_basic": "基本处理", "section_restoration_model": "修复模型", - "model_basicvsrpp": "标准", + "model_basicvsrpp": "BasicVSR++", "model_basicvsrpp_description": "速度快，适合大多数视频。", "model_ltx": "LTX", "model_ltx_description": "细节最好。速度慢很多，需要高性能 NVIDIA 显卡。", diff --git a/jasna/gui/settings_sections/restoration_model.py b/jasna/gui/settings_sections/restoration_model.py index cef77e0a..3ec6b9b1 100644 --- a/jasna/gui/settings_sections/restoration_model.py +++ b/jasna/gui/settings_sections/restoration_model.py @@ -1,4 +1,4 @@ -"""Restoration model choice (Standard BasicVSR++ or LTX) and the LTX options.""" +"""Restoration model choice (BasicVSR++ or LTX) and the LTX options.""" import random from typing import Callable From 1a287beef5c45e122a9995f140023ada61dcab4b Mon Sep 17 00:00:00 2001 From: Kruk2 Date: Tue, 29 Sep 2026 18:21:41 +0200 Subject: [PATCH 38/48] fix(ltx): restore overlapping mosaic regions in full cv2.fillPoly with several polygons uses the even-odd rule, so overlapping regions were left out of the blend mask. Fill each polygon separately. --- jasna/ltx/compose.py | 14 ++++++-------- tests/test_ltx_compose.py | 18 +++++++++++++++++- 2 files changed, 23 insertions(+), 9 deletions(-) diff --git a/jasna/ltx/compose.py b/jasna/ltx/compose.py index ee7f8e9d..9f94c2b1 100644 --- a/jasna/ltx/compose.py +++ b/jasna/ltx/compose.py @@ -113,15 +113,13 @@ def axis(low: int, high: int, size: int) -> torch.Tensor: def polygon_mask(polygons: Sequence[Polygon], shape: tuple[int, int], offset: tuple[int, int], device: torch.device) -> torch.Tensor: - """Bool ``[h, w]`` fill of ``polygons`` inside the region at ``offset`` (x, y).""" + """Bool ``[h, w]`` union of ``polygons`` inside the region at ``offset`` (x, y).""" mask = np.zeros(shape, dtype=np.uint8) - contours = [ - (np.asarray(polygon, dtype=np.float32).round().astype(np.int32) - np.asarray(offset, dtype=np.int32)) - for polygon in polygons - if len(polygon) >= 3 - ] - if contours: - cv2.fillPoly(mask, contours, 255) + for polygon in polygons: + if len(polygon) >= 3: + contour = np.asarray(polygon, dtype=np.float32).round().astype(np.int32) - np.asarray(offset, dtype=np.int32) + # One fillPoly call per polygon: a single call uses the even-odd rule and punches holes where polygons overlap. + cv2.fillPoly(mask, [contour], 255) return torch.from_numpy(mask).to(device) > 0 diff --git a/tests/test_ltx_compose.py b/tests/test_ltx_compose.py index d5a3fc65..9a94748b 100644 --- a/tests/test_ltx_compose.py +++ b/tests/test_ltx_compose.py @@ -3,7 +3,14 @@ import torch from jasna.ltx.camera import letterbox -from jasna.ltx.compose import Candidate, canvas_to_source, composite_frame, crop_to_canvas, feather_alpha +from jasna.ltx.compose import ( + Candidate, + canvas_to_source, + composite_frame, + crop_to_canvas, + feather_alpha, + polygon_mask, +) def _frame(seed: int = 0) -> torch.Tensor: @@ -50,6 +57,15 @@ def test_composite_touches_only_the_feathered_mask(): assert ys.min() >= 200 - 2 * 8 - 1 and ys.max() <= 260 + 2 * 8 + 1 +def test_polygon_mask_is_the_union_of_overlapping_polygons(): + outer = [[0.0, 0.0], [99.0, 0.0], [99.0, 99.0], [0.0, 99.0]] + nested = [[20.0, 20.0], [60.0, 20.0], [60.0, 60.0], [20.0, 60.0]] + crossing = [[50.0, 50.0], [150.0, 50.0], [150.0, 150.0], [50.0, 150.0]] + mask = polygon_mask([outer, nested, crossing], (160, 160), (0, 0), torch.device("cpu")) + assert mask[40, 40] and mask[55, 55] and mask[80, 80] and mask[140, 140] + assert mask.sum() == 100 * 100 + 101 * 101 - 50 * 50 + + def test_zero_weight_and_no_candidates_keep_the_source(): source = _frame(2) crop = letterbox((200, 100, 456, 356), 256) From a0e4291faeb2555dd82cdc1cde5f2696c20a85ef Mon Sep 17 00:00:00 2001 From: Kruk2 Date: Tue, 29 Sep 2026 18:21:41 +0200 Subject: [PATCH 39/48] perf(masks): share one exact GPU ellipse dilation blending.dilate_ellipse matches cv2.dilate with MORPH_ELLIPSE pixel for pixel and replaces the LTX tracing, LTX feather and SD 1.5 copies. LTX tracing runs about 2x faster. --- jasna/image_restore.py | 13 +++++----- jasna/ltx/compose.py | 9 ++----- jasna/ltx/regions.py | 6 ++--- jasna/tracking/blending.py | 51 ++++++++++++++++++++++++++++++++++++++ tests/test_blending.py | 40 +++++++++++++++++++++++++++++- 5 files changed, 101 insertions(+), 18 deletions(-) diff --git a/jasna/image_restore.py b/jasna/image_restore.py index eb664242..857ab0c3 100644 --- a/jasna/image_restore.py +++ b/jasna/image_restore.py @@ -5,11 +5,12 @@ from dataclasses import dataclass from pathlib import Path -import cv2 import numpy as np import torch import torch.nn.functional as F +from jasna.tracking.blending import dilate_ellipse + logger = logging.getLogger(__name__) SD15_RESTORATION_SIZE = 512 @@ -97,15 +98,13 @@ def _union_bbox(boxes: np.ndarray, indices: list[int]) -> np.ndarray: def _build_group_mask(masks: torch.Tensor, indices: list[int], h: int, w: int, dilate_px: int) -> np.ndarray: - acc = np.zeros((h, w), dtype=np.uint8) + union = torch.zeros((h, w), dtype=torch.bool, device=masks.device) for idx in indices: m = masks[idx].unsqueeze(0).unsqueeze(0).float() - up = F.interpolate(m, size=(h, w), mode="bilinear", align_corners=False)[0, 0] > 0.5 - acc = np.maximum(acc, (up.cpu().numpy() * 255).astype(np.uint8)) + union |= F.interpolate(m, size=(h, w), mode="bilinear", align_corners=False)[0, 0] > 0.5 if dilate_px > 0: - kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (dilate_px * 2 + 1, dilate_px * 2 + 1)) - acc = cv2.dilate(acc, kernel, iterations=1) - return acc + union = dilate_ellipse(union, dilate_px) + return union.to(torch.uint8).mul_(255).cpu().numpy() def prepare_image_restore( diff --git a/jasna/ltx/compose.py b/jasna/ltx/compose.py index 9f94c2b1..5fc24cdc 100644 --- a/jasna/ltx/compose.py +++ b/jasna/ltx/compose.py @@ -18,6 +18,7 @@ from jasna.ltx.camera import CropFrame from jasna.ltx.plan import Polygon +from jasna.tracking.blending import dilate_ellipse COLOUR_MATCH_STRENGTH = 0.75 @@ -71,12 +72,6 @@ def canvas_to_source(canvas: torch.Tensor, crop: CropFrame) -> torch.Tensor: return out.round_().clamp_(0, 255) -@functools.cache -def _ellipse(radius: int, device: torch.device) -> torch.Tensor: - kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (2 * radius + 1,) * 2) - return torch.from_numpy(kernel.astype(np.float32)).to(device)[None, None] - - @functools.cache def _gaussian(radius: int, device: torch.device) -> torch.Tensor: kernel = cv2.getGaussianKernel(2 * radius + 1, 0.0, cv2.CV_32F)[:, 0] @@ -86,7 +81,7 @@ def _gaussian(radius: int, device: torch.device) -> torch.Tensor: def feather_alpha(mask: torch.Tensor, feather: int) -> torch.Tensor: """Dilate a bool ``[H, W]`` mask by an ellipse of ``feather`` px and Gaussian-blur it (kernel ``2*feather+1``, reflect-101 border) into a [0, 1] alpha.""" - dilated = F.conv2d(F.pad(mask[None, None].float(), (feather,) * 4), _ellipse(feather, mask.device)) > 0 + dilated = dilate_ellipse(mask, feather)[None, None] kernel = _gaussian(feather, mask.device) x = F.pad(dilated.float(), (feather, feather, 0, 0), mode="reflect") x = F.conv2d(x, kernel.view(1, 1, 1, -1)) diff --git a/jasna/ltx/regions.py b/jasna/ltx/regions.py index a204317b..014683ab 100644 --- a/jasna/ltx/regions.py +++ b/jasna/ltx/regions.py @@ -14,9 +14,9 @@ from jasna.ltx.plan import Region from jasna.mosaic.detections import Detections +from jasna.tracking.blending import dilate_ellipse EXPAND_PIXELS = 20 -_KERNEL = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (2 * EXPAND_PIXELS + 1,) * 2) def trace_region(box: np.ndarray, mask: torch.Tensor, frame_h: int, frame_w: int) -> Region | None: @@ -29,8 +29,8 @@ def trace_region(box: np.ndarray, mask: torch.Tensor, frame_h: int, frame_w: int y1 = min(frame_h, int(rows[-1]) + EXPAND_PIXELS + 2) x0 = max(0, int(cols[0]) - EXPAND_PIXELS - 1) x1 = min(frame_w, int(cols[-1]) + EXPAND_PIXELS + 2) - crop = full[y0:y1, x0:x1].to(torch.uint8).mul_(255).cpu().numpy() - contours, _ = cv2.findContours(cv2.dilate(crop, _KERNEL), cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE) + dilated = dilate_ellipse(full[y0:y1, x0:x1], EXPAND_PIXELS).to(torch.uint8).mul_(255).cpu().numpy() + contours, _ = cv2.findContours(dilated, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE) if not contours: return None largest = max(contours, key=cv2.contourArea) diff --git a/jasna/tracking/blending.py b/jasna/tracking/blending.py index 980b38ac..e0e72be6 100644 --- a/jasna/tracking/blending.py +++ b/jasna/tracking/blending.py @@ -1,5 +1,9 @@ from __future__ import annotations +import functools + +import cv2 +import numpy as np import torch import torch.nn.functional as F @@ -72,6 +76,53 @@ def _box_blur(x: torch.Tensor, kernel_h: int, kernel_w: int) -> torch.Tensor: return _prefix_box_blur(x, kernel_h, kernel_w) +def _ellipse_structuring_element(radius: int) -> np.ndarray: + return cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (2 * radius + 1,) * 2) + + +@functools.cache +def _ellipse_kernel(radius: int, device: torch.device) -> torch.Tensor: + return torch.from_numpy(_ellipse_structuring_element(radius).astype(np.float32)).to(device)[None, None] + + +@functools.cache +def _ellipse_levels(radius: int) -> tuple[tuple[int, int], ...]: + """``(half width, vertical reach)`` per distinct row width of the ellipse: its rows are + centred runs that narrow away from the middle, so the dilation is the union of one + horizontal-then-vertical window pass per width.""" + widths = [radius - int(np.flatnonzero(row)[0]) for row in _ellipse_structuring_element(radius)] + return tuple( + (width, max(abs(dy - radius) for dy, row_width in enumerate(widths) if row_width >= width)) + for width in sorted(set(widths)) + ) + + +def _conv_dilate_ellipse(mask: torch.Tensor, radius: int) -> torch.Tensor: + padded = F.pad(mask[None, None].float(), (radius,) * 4) + return F.conv2d(padded, _ellipse_kernel(radius, mask.device))[0, 0] > 0.5 + + +def _any_in_window(x: torch.Tensor, half: int, dim: int) -> torch.Tensor: + padding = (half, half) if dim == 1 else (0, 0, half, half) + return _moving_average(F.pad(x, padding), 2 * half + 1, dim) > 0 + + +def _prefix_dilate_ellipse(mask: torch.Tensor, radius: int) -> torch.Tensor: + x = mask.float() + out = torch.zeros_like(mask) + for half_width, reach in _ellipse_levels(radius): + out |= _any_in_window(_any_in_window(x, half_width, 1).float(), reach, 0) + return out + + +def dilate_ellipse(mask: torch.Tensor, radius: int) -> torch.Tensor: + """Bool ``[H, W]`` mask dilated on its own device; same pixels as ``cv2.dilate`` with a + ``MORPH_ELLIPSE`` kernel of ``2*radius+1``.""" + if is_nvidia_device(mask.device): + return _conv_dilate_ellipse(mask, radius) + return _prefix_dilate_ellipse(mask, radius) + + def create_blend_mask( crop_mask: torch.Tensor, frame_height: int, diff --git a/tests/test_blending.py b/tests/test_blending.py index 1cb9c891..aa1ce941 100644 --- a/tests/test_blending.py +++ b/tests/test_blending.py @@ -1,10 +1,18 @@ from unittest.mock import MagicMock +import cv2 +import numpy as np import pytest import torch import torch.nn.functional as F -from jasna.tracking.blending import _box_blur, create_bbox_blend_mask, create_blend_mask +from jasna.tracking.blending import ( + _box_blur, + _conv_dilate_ellipse, + _prefix_dilate_ellipse, + create_bbox_blend_mask, + create_blend_mask, +) @pytest.mark.parametrize("kernel_size", [5, 61, 121]) @@ -144,3 +152,33 @@ def test_bbox_blend_mask_same_res_mask_matches_reference_exactly() -> None: ref = _reference_fullres_blend_mask(mask, bbox, frame_shape) out = create_bbox_blend_mask(mask, bbox, frame_shape) assert torch.allclose(out, ref, atol=1e-5) + + +@pytest.mark.parametrize("dilate", [_conv_dilate_ellipse, _prefix_dilate_ellipse]) +@pytest.mark.parametrize("radius", [1, 2, 3, 12, 20, 30]) +def test_dilate_ellipse_matches_opencv(dilate, radius: int) -> None: + rng = np.random.default_rng(radius) + mask = (rng.random((150, 230)) < 0.01).astype(np.uint8) * 255 + mask[0, 5] = mask[149, 229] = 255 + kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (2 * radius + 1,) * 2) + + actual = dilate(torch.from_numpy(mask > 0), radius).numpy() + + assert np.array_equal(actual, cv2.dilate(mask, kernel) > 0) + + +def test_dilate_ellipse_uses_convolution_only_for_nvidia(monkeypatch) -> None: + import jasna.tracking.blending as module + + mask = torch.ones((8, 8), dtype=torch.bool) + conv = MagicMock(return_value=mask) + prefix = MagicMock(return_value=mask) + monkeypatch.setattr(module, "_conv_dilate_ellipse", conv) + monkeypatch.setattr(module, "_prefix_dilate_ellipse", prefix) + + monkeypatch.setattr(module, "is_nvidia_device", lambda _device: False) + module.dilate_ellipse(mask, 3) + monkeypatch.setattr(module, "is_nvidia_device", lambda _device: True) + module.dilate_ellipse(mask, 3) + + assert prefix.call_count == 1 and conv.call_count == 1 From 4f4539be4ed9a9104e80193f643fb70e35ca868f Mon Sep 17 00:00:00 2001 From: Kruk2 Date: Tue, 29 Sep 2026 18:56:02 +0200 Subject: [PATCH 40/48] feat(ltx): choose distilled or undistilled model and download missing models --ltx-model picks the model; each variant has its own file name. A missing model is offered for download with its size (y/N in the CLI); the download location is not set yet. --- jasna/cli_help.py | 4 ++ jasna/gui/models.py | 3 +- jasna/gui/video_session.py | 2 + jasna/ltx/model_files.py | 123 +++++++++++++++++++++++++++++++--- jasna/main.py | 39 ++++++++++- jasna/session_config.py | 3 + jasna/session_factory.py | 2 +- tests/factories.py | 1 + tests/test_ltx_model_files.py | 79 +++++++++++++++++----- 9 files changed, 228 insertions(+), 28 deletions(-) diff --git a/jasna/cli_help.py b/jasna/cli_help.py index a170a301..017a3ed0 100644 --- a/jasna/cli_help.py +++ b/jasna/cli_help.py @@ -16,6 +16,10 @@ "Restoration model file (basicvsrpp) or folder (ltx). Default: the model in model_weights/." ), "ltx_seed": "ltx: noise seed. Another seed gives another take on the same restoration. (default: %(default)s)", + "ltx_model": ( + "ltx: distilled (8 steps, about 3x faster) or undistilled (15 steps with STG). Both look " + "about the same; a missing model is offered for download. (default: %(default)s)" + ), "ltx_fast": ( "ltx: about 1.4x faster with slightly less detail. Needs an RTX 50-series (Blackwell) GPU " "and the fast model file. (default: %(default)s)" diff --git a/jasna/gui/models.py b/jasna/gui/models.py index 110b06a1..69a2b344 100644 --- a/jasna/gui/models.py +++ b/jasna/gui/models.py @@ -10,7 +10,7 @@ from jasna.gui.paths import get_settings_path from jasna.segments import SegmentRange -from jasna.session_config import LTX_DEFAULT_SEED +from jasna.session_config import LTX_DEFAULT_MODEL, LTX_DEFAULT_SEED logger = logging.getLogger(__name__) @@ -108,6 +108,7 @@ class AppSettings: restoration_model: str = "basicvsrpp" # basicvsrpp, ltx ltx_seed: int = LTX_DEFAULT_SEED ltx_fast: bool = False + ltx_model: str = LTX_DEFAULT_MODEL ltx_large_canvas: bool = False # Basic processing diff --git a/jasna/gui/video_session.py b/jasna/gui/video_session.py index d4b1815d..35a176eb 100644 --- a/jasna/gui/video_session.py +++ b/jasna/gui/video_session.py @@ -30,6 +30,7 @@ def video_session_key(settings: AppSettings) -> tuple: settings.secondary_restoration, settings.restoration_model, settings.ltx_fast, + settings.ltx_model, ) if settings.secondary_restoration == "tvai": key += ( @@ -77,6 +78,7 @@ def video_session_config( ltx_large_canvas=bool(settings.ltx_large_canvas), ltx_seed=int(settings.ltx_seed), ltx_fast=bool(settings.ltx_fast), + ltx_model=settings.ltx_model, compile_basicvsrpp=bool(settings.compile_basicvsrpp), max_clip_size=int(settings.max_clip_size), temporal_overlap=int(settings.temporal_overlap), diff --git a/jasna/ltx/model_files.py b/jasna/ltx/model_files.py index ec8b932c..00572a95 100644 --- a/jasna/ltx/model_files.py +++ b/jasna/ltx/model_files.py @@ -1,10 +1,59 @@ -"""LTX model bundle discovery and tensor access.""" +"""LTX model bundle discovery, tensor access and download. + +A bundle is one transformer file (the model choice x the precision) plus the shared VAE +and tuned decoder. Every file name is unique, so all variants can sit in one folder. +""" from __future__ import annotations +import hashlib +import logging +import urllib.request +from collections.abc import Callable from dataclasses import dataclass from pathlib import Path +from typing import get_args from jasna._frozen import is_frozen +from jasna.session_config import LtxModelName + +logger = logging.getLogger(__name__) + +LTX_MODELS: tuple[LtxModelName, ...] = get_args(LtxModelName) + +# Where the model files are published; each file is fetched from f"{LTX_DOWNLOAD_URL}/{name}". +LTX_DOWNLOAD_URL: str | None = None + +DownloadProgressCallback = Callable[[int, int], None] + + +@dataclass(frozen=True) +class DownloadableFile: + name: str + size_bytes: int + sha256: str | None + + +_TRANSFORMER_NAMES: dict[tuple[LtxModelName, bool], str] = { + ("distilled", False): "ltx-restore-alpha1-distill8-int8.safetensors", + ("distilled", True): "ltx-restore-alpha1-distill8-nvfp4.safetensors", + ("undistilled", False): "ltx-restore-alpha1-int8.safetensors", + ("undistilled", True): "ltx-restore-alpha1-nvfp4.safetensors", +} +_SHARED_NAMES = ("vae.safetensors", "vae-decoder.safetensors") + +# The published files (``.enc`` in releases): size for the download prompt, and a sha256 +# that is filled in with the release that publishes the file. +LTX_DOWNLOADS: dict[str, DownloadableFile] = { + name: DownloadableFile(name, size, None) + for name, size in ( + ("ltx-restore-alpha1-distill8-int8.safetensors", 13_368_044_744), + ("ltx-restore-alpha1-distill8-nvfp4.safetensors", 7_733_292_672), + ("ltx-restore-alpha1-int8.safetensors", 13_368_044_792), + ("ltx-restore-alpha1-nvfp4.safetensors", 7_733_292_712), + ("vae.safetensors", 1_472_223_346), + ("vae-decoder.safetensors", 38_718_560), + ) +} def open_tensors(path: Path): @@ -15,6 +64,16 @@ def open_tensors(path: Path): return open_ltx_tensors(path) +def bundle_names(model: LtxModelName, *, fast: bool) -> tuple[str, str, str]: + """(transformer, vae, tuned decoder) file names of one model choice.""" + return (_TRANSFORMER_NAMES[(model, fast)], *_SHARED_NAMES) + + +def _bundle_paths(directory: Path, model: LtxModelName, *, fast: bool) -> list[Path]: + paths = [directory / name for name in bundle_names(model, fast=fast)] + return [path.with_name(path.name + ".enc") if is_frozen() or not path.is_file() else path for path in paths] + + @dataclass(frozen=True) class LtxModelFiles: transformer: Path @@ -22,9 +81,9 @@ class LtxModelFiles: tuned_decoder: Path @classmethod - def from_dir(cls, directory: Path, *, fast: bool) -> LtxModelFiles: - """The model bundle; ``fast`` picks the NVFP4 transformer over the INT8 one.""" - paths = _bundle_paths(directory, fast=fast) + def from_dir(cls, directory: Path, model: LtxModelName, *, fast: bool) -> LtxModelFiles: + """The bundle of ``model``; ``fast`` picks its NVFP4 transformer over the INT8 one.""" + paths = _bundle_paths(directory, model, fast=fast) if any(path.suffix == ".enc" for path in paths): from jasna.protection.protected_tensors import check_ltx_license check_ltx_license() @@ -34,12 +93,56 @@ def from_dir(cls, directory: Path, *, fast: bool) -> LtxModelFiles: return cls(*paths) -def _bundle_paths(directory: Path, *, fast: bool) -> list[Path]: - transformer = "transformer-fast" if fast else "transformer" - paths = [directory / f"{name}.safetensors" for name in (transformer, "vae", "vae-decoder")] - return [path.with_name(path.name + ".enc") if is_frozen() or not path.is_file() else path for path in paths] +def model_installed(directory: Path, model: LtxModelName, *, fast: bool) -> bool: + """Whether every file of this model choice is on disk; checks no license.""" + return all(path.is_file() for path in _bundle_paths(directory, model, fast=fast)) def bundle_present(directory: Path) -> bool: - """Whether the quality model files are installed; checks no license.""" - return all(path.is_file() for path in _bundle_paths(directory, fast=False)) + """Whether any LTX model can run from ``directory``; checks no license.""" + return any(model_installed(directory, model, fast=fast) for model in LTX_MODELS for fast in (False, True)) + + +def missing_downloads(directory: Path, model: LtxModelName, *, fast: bool) -> list[DownloadableFile]: + """The files this model choice still needs, as published (``.enc`` in releases).""" + missing = [] + for path in _bundle_paths(directory, model, fast=fast): + if not path.is_file(): + entry = LTX_DOWNLOADS[path.name.removesuffix(".enc")] + missing.append(DownloadableFile(path.name, entry.size_bytes, entry.sha256)) + return missing + + +def download_size_text(files: list[DownloadableFile]) -> str: + return f"{sum(f.size_bytes for f in files) / 1e9:.1f} GB" + + +def download_files( + directory: Path, files: list[DownloadableFile], progress_callback: DownloadProgressCallback | None = None +) -> None: + """Download ``files`` into ``directory``, verifying each sha256 before it replaces + anything. ``progress_callback(done_bytes, total_bytes)`` covers all files together.""" + if LTX_DOWNLOAD_URL is None: + raise RuntimeError("The LTX model download location is not published yet; get the model files manually.") + unverifiable = [f.name for f in files if f.sha256 is None] + if unverifiable: + raise RuntimeError(f"No published checksum for {', '.join(unverifiable)}; refusing to download.") + directory.mkdir(parents=True, exist_ok=True) + total = sum(f.size_bytes for f in files) + done = 0 + for file in files: + target = directory / file.name + partial = target.with_name(target.name + ".part") + digest = hashlib.sha256() + logger.info("Downloading %s", file.name) + with urllib.request.urlopen(f"{LTX_DOWNLOAD_URL}/{file.name}") as response, partial.open("wb") as out: + while chunk := response.read(1 << 22): + out.write(chunk) + digest.update(chunk) + done += len(chunk) + if progress_callback is not None: + progress_callback(done, total) + if digest.hexdigest() != file.sha256: + partial.unlink() + raise RuntimeError(f"{file.name} failed its checksum; the download was discarded.") + partial.replace(target) diff --git a/jasna/main.py b/jasna/main.py index 87bb75d6..3ddb3872 100644 --- a/jasna/main.py +++ b/jasna/main.py @@ -2,6 +2,7 @@ import logging import sys from pathlib import Path +from typing import get_args from jasna import __version__ from jasna.cli_help import CLI_HELP @@ -16,7 +17,7 @@ check_windows_nvidia_sysmem_fallback_policy, gpu_check_error, ) -from jasna.session_config import LTX_DEFAULT_SEED, SessionConfig +from jasna.session_config import LTX_DEFAULT_MODEL, LTX_DEFAULT_SEED, LtxModelName, SessionConfig @@ -46,6 +47,7 @@ def _session_config_from_args( ltx_large_canvas=bool(args.ltx_large_canvas), ltx_seed=int(args.ltx_seed), ltx_fast=bool(args.ltx_fast), + ltx_model=str(args.ltx_model), compile_basicvsrpp=bool(args.compile_basicvsrpp), max_clip_size=int(args.max_clip_size), temporal_overlap=int(args.temporal_overlap), @@ -216,6 +218,12 @@ def build_parser() -> argparse.ArgumentParser: default=LTX_DEFAULT_SEED, help=CLI_HELP["ltx_seed"], ) + restoration.add_argument( + "--ltx-model", + choices=list(get_args(LtxModelName)), + default=LTX_DEFAULT_MODEL, + help=CLI_HELP["ltx_model"], + ) restoration.add_argument( "--ltx-fast", default=False, @@ -763,6 +771,33 @@ def _run_videos( return all_commands_ok +def _ensure_ltx_model(config: SessionConfig, segments) -> None: + """Offer to download the selected LTX model when the job needs it and it is missing. + Without a terminal to ask in, fail with what to download instead of waiting.""" + from jasna.ltx.model_files import download_files, download_size_text, missing_downloads + + uses_ltx = config.restoration_model_name == "ltx" or any( + segment.restoration is not None and segment.restoration.model == "ltx" for segment in segments or () + ) + if not uses_ltx: + return + directory = config.restoration_model_path if config.restoration_model_name == "ltx" else default_restoration_model_path("ltx") + missing = missing_downloads(directory, config.ltx_model, fast=config.ltx_fast) + if not missing: + return + names = ", ".join(file.name for file in missing) + size = download_size_text(missing) + if not sys.stdin.isatty(): + raise FileNotFoundError(f"LTX model files missing in {directory}: {names}. Run jasna in a terminal to download them ({size}).") + answer = input(f"The LTX {config.ltx_model} model is missing in {directory}:\n {names}\nDownload {size} now? [y/N]: ") + if answer.strip().lower() not in ("y", "yes"): + raise FileNotFoundError(f"LTX model files missing in {directory}: {names}") + from tqdm import tqdm + + with tqdm(total=sum(file.size_bytes for file in missing), unit="B", unit_scale=True, desc="LTX download") as bar: + download_files(directory, missing, lambda done, total: bar.update(done - bar.n)) + + def main() -> None: parser = build_parser() args = parser.parse_args() @@ -908,6 +943,8 @@ def _run_post_export_action() -> None: lut_path=lut_arg or None, ) + _ensure_ltx_model(config, segments) + from jasna.session_factory import build_pipeline, build_restoration_session with device_context(config.device): diff --git a/jasna/session_config.py b/jasna/session_config.py index 6ba80f97..17cdafd8 100644 --- a/jasna/session_config.py +++ b/jasna/session_config.py @@ -12,6 +12,8 @@ from typing import Literal, Mapping RestorationModelName = Literal["basicvsrpp", "ltx"] +LtxModelName = Literal["distilled", "undistilled"] +LTX_DEFAULT_MODEL: LtxModelName = "distilled" LTX_DEFAULT_SEED = 20260923 SecondaryRestorationName = Literal["none", "unet-4x", "tvai", "rtx-super-res"] @@ -40,6 +42,7 @@ class SessionConfig: ltx_large_canvas: bool ltx_seed: int ltx_fast: bool + ltx_model: LtxModelName compile_basicvsrpp: bool max_clip_size: int temporal_overlap: int diff --git a/jasna/session_factory.py b/jasna/session_factory.py index f1ae8f10..e76dd09b 100644 --- a/jasna/session_factory.py +++ b/jasna/session_factory.py @@ -240,7 +240,7 @@ def _ltx_model_files( raise ValueError("LTX restoration needs an NVIDIA GPU") if config.ltx_fast and torch.cuda.get_device_capability(device)[0] < 10: raise ValueError("The fast LTX model needs an RTX 50-series (Blackwell) GPU") - files = LtxModelFiles.from_dir(_restoration_model_path(config, "ltx"), fast=config.ltx_fast) + files = LtxModelFiles.from_dir(_restoration_model_path(config, "ltx"), config.ltx_model, fast=config.ltx_fast) ensure_engines_compiled( EngineCompilationRequest( device=str(device), diff --git a/tests/factories.py b/tests/factories.py index f0cd70ee..6e4d631c 100644 --- a/tests/factories.py +++ b/tests/factories.py @@ -30,6 +30,7 @@ def session_config(**overrides) -> SessionConfig: ltx_large_canvas=True, ltx_seed=0, ltx_fast=False, + ltx_model="distilled", compile_basicvsrpp=True, max_clip_size=90, temporal_overlap=8, diff --git a/tests/test_ltx_model_files.py b/tests/test_ltx_model_files.py index 0a80c8e9..d0797802 100644 --- a/tests/test_ltx_model_files.py +++ b/tests/test_ltx_model_files.py @@ -1,17 +1,19 @@ +import hashlib from unittest.mock import Mock import pytest import torch from safetensors.torch import save_file -from jasna.ltx.model_files import LtxModelFiles, open_tensors +from jasna.ltx import model_files +from jasna.ltx.model_files import LtxModelFiles, bundle_names, open_tensors def test_source_model_bundle_and_reader(tmp_path): value = torch.arange(6).reshape(2, 3) - for name in ("transformer", "vae", "vae-decoder"): - save_file({"weight": value}, str(tmp_path / f"{name}.safetensors"), metadata={"config": "example"}) - files = LtxModelFiles.from_dir(tmp_path, fast=False) + for name in bundle_names("distilled", fast=False): + save_file({"weight": value}, str(tmp_path / name), metadata={"config": "example"}) + files = LtxModelFiles.from_dir(tmp_path, "distilled", fast=False) with open_tensors(files.transformer) as handle: assert list(handle.keys()) == ["weight"] assert handle.metadata() == {"config": "example"} @@ -32,17 +34,64 @@ def test_model_check_precedes_engine_compilation(tmp_path, monkeypatch): compile_engines.assert_not_called() -def test_fast_bundle_uses_the_fast_transformer(tmp_path): - for name in ("transformer-fast", "vae", "vae-decoder"): - save_file({"weight": torch.zeros(1)}, str(tmp_path / f"{name}.safetensors")) - assert LtxModelFiles.from_dir(tmp_path, fast=True).transformer == tmp_path / "transformer-fast.safetensors" +def test_every_model_choice_has_its_own_transformer_file(): + transformers = {bundle_names(model, fast=fast)[0] for model in model_files.LTX_MODELS for fast in (False, True)} + assert len(transformers) == 4 + assert all(name in model_files.LTX_DOWNLOADS for name in transformers) + assert bundle_names("undistilled", fast=True)[1:] == bundle_names("distilled", fast=False)[1:] -def test_bundle_present_needs_every_quality_file(tmp_path): - from jasna.ltx.model_files import bundle_present +def test_installed_checks_the_chosen_model_only(tmp_path): + for name in bundle_names("undistilled", fast=True): + (tmp_path / name).touch() + assert model_files.model_installed(tmp_path, "undistilled", fast=True) + assert not model_files.model_installed(tmp_path, "distilled", fast=True) + assert model_files.bundle_present(tmp_path) + assert [f.name for f in model_files.missing_downloads(tmp_path, "distilled", fast=True)] == [ + "ltx-restore-alpha1-distill8-nvfp4.safetensors.enc" + ] - for name in ("transformer-fast", "vae", "vae-decoder"): - (tmp_path / f"{name}.safetensors").touch() - assert not bundle_present(tmp_path) - (tmp_path / "transformer.safetensors.enc").touch() - assert bundle_present(tmp_path) + +def test_download_refuses_without_a_published_location(tmp_path, monkeypatch): + monkeypatch.setattr(model_files, "LTX_DOWNLOAD_URL", None) + with pytest.raises(RuntimeError, match="not published"): + model_files.download_files(tmp_path, model_files.missing_downloads(tmp_path, "distilled", fast=False)) + + +def test_download_verifies_and_replaces_atomically(tmp_path, monkeypatch): + source = tmp_path / "published" + source.mkdir() + good, bad = b"model bytes", b"tampered" + (source / "a.safetensors.enc").write_bytes(good) + (source / "b.safetensors.enc").write_bytes(bad) + monkeypatch.setattr(model_files, "LTX_DOWNLOAD_URL", source.as_uri()) + target = tmp_path / "models" + progress = [] + ok = model_files.DownloadableFile("a.safetensors.enc", len(good), hashlib.sha256(good).hexdigest()) + model_files.download_files(target, [ok], lambda done, total: progress.append((done, total))) + assert (target / "a.safetensors.enc").read_bytes() == good and progress[-1] == (len(good), len(good)) + wrong = model_files.DownloadableFile("b.safetensors.enc", len(bad), hashlib.sha256(good).hexdigest()) + with pytest.raises(RuntimeError, match="checksum"): + model_files.download_files(target, [wrong]) + assert sorted(p.name for p in target.iterdir()) == ["a.safetensors.enc"] + + +def test_cli_offers_the_missing_model_only_when_ltx_runs(tmp_path, monkeypatch): + import sys + + from factories import session_config + from jasna import main + + asked = [] + monkeypatch.setattr("builtins.input", lambda prompt: asked.append(prompt) or "n") + main._ensure_ltx_model(session_config(restoration_model_name="basicvsrpp"), None) + assert asked == [] + + ltx = session_config(restoration_model_name="ltx", restoration_model_path=tmp_path) + monkeypatch.setattr(sys.stdin, "isatty", lambda: False) + with pytest.raises(FileNotFoundError, match="Run jasna in a terminal"): + main._ensure_ltx_model(ltx, None) + monkeypatch.setattr(sys.stdin, "isatty", lambda: True) + with pytest.raises(FileNotFoundError, match="missing"): + main._ensure_ltx_model(ltx, None) + assert "GB" in asked[0] and "ltx-restore-alpha1-distill8-int8" in asked[0] From 145035ca11d8052ddaf48900c08e1d646ec44b56 Mon Sep 17 00:00:00 2001 From: Kruk2 Date: Tue, 29 Sep 2026 18:58:48 +0200 Subject: [PATCH 41/48] fix(gui): use the default LTX model until the GUI can choose one --- jasna/gui/models.py | 3 +-- jasna/gui/video_session.py | 5 ++--- tests/test_main_validation.py | 1 + 3 files changed, 4 insertions(+), 5 deletions(-) diff --git a/jasna/gui/models.py b/jasna/gui/models.py index 69a2b344..110b06a1 100644 --- a/jasna/gui/models.py +++ b/jasna/gui/models.py @@ -10,7 +10,7 @@ from jasna.gui.paths import get_settings_path from jasna.segments import SegmentRange -from jasna.session_config import LTX_DEFAULT_MODEL, LTX_DEFAULT_SEED +from jasna.session_config import LTX_DEFAULT_SEED logger = logging.getLogger(__name__) @@ -108,7 +108,6 @@ class AppSettings: restoration_model: str = "basicvsrpp" # basicvsrpp, ltx ltx_seed: int = LTX_DEFAULT_SEED ltx_fast: bool = False - ltx_model: str = LTX_DEFAULT_MODEL ltx_large_canvas: bool = False # Basic processing diff --git a/jasna/gui/video_session.py b/jasna/gui/video_session.py index 35a176eb..cbdb9575 100644 --- a/jasna/gui/video_session.py +++ b/jasna/gui/video_session.py @@ -10,7 +10,7 @@ from typing import TYPE_CHECKING, Callable, Mapping from jasna.gui.models import AppSettings -from jasna.session_config import SessionConfig +from jasna.session_config import LTX_DEFAULT_MODEL, SessionConfig from jasna.session_factory import RestorationSession if TYPE_CHECKING: @@ -30,7 +30,6 @@ def video_session_key(settings: AppSettings) -> tuple: settings.secondary_restoration, settings.restoration_model, settings.ltx_fast, - settings.ltx_model, ) if settings.secondary_restoration == "tvai": key += ( @@ -78,7 +77,7 @@ def video_session_config( ltx_large_canvas=bool(settings.ltx_large_canvas), ltx_seed=int(settings.ltx_seed), ltx_fast=bool(settings.ltx_fast), - ltx_model=settings.ltx_model, + ltx_model=LTX_DEFAULT_MODEL, compile_basicvsrpp=bool(settings.compile_basicvsrpp), max_clip_size=int(settings.max_clip_size), temporal_overlap=int(settings.temporal_overlap), diff --git a/tests/test_main_validation.py b/tests/test_main_validation.py index 006f932e..533841cd 100644 --- a/tests/test_main_validation.py +++ b/tests/test_main_validation.py @@ -70,6 +70,7 @@ def test_ltx_accepts_segments(self, tmp_path): patch("jasna.media.splice.build_splice_plan", return_value=MagicMock()), patch("jasna.accelerator.is_nvidia_device", return_value=True), patch("jasna.ltx.model_files.LtxModelFiles.from_dir", return_value=MagicMock()), + patch("jasna.ltx.model_files.missing_downloads", return_value=[]), ): pipeline_cls = _run_main_with_args( tmp_path, ["--segments", "1-2", "--restoration-model-name", "ltx", "--ltx-seed", "3"] From 4c8c9d8c549f7dfdb84b38f2ab8815ed6835be45 Mon Sep 17 00:00:00 2001 From: Kruk2 Date: Tue, 29 Sep 2026 19:01:05 +0200 Subject: [PATCH 42/48] feat(ltx): pick the undistilled model as --restoration-model-name ltx-undistilled --- jasna/cli_help.py | 11 ++++------- jasna/main.py | 27 ++++++++++++++------------- tests/test_session_config.py | 4 ++++ 3 files changed, 22 insertions(+), 20 deletions(-) diff --git a/jasna/cli_help.py b/jasna/cli_help.py index 017a3ed0..c8b793bf 100644 --- a/jasna/cli_help.py +++ b/jasna/cli_help.py @@ -8,18 +8,15 @@ CLI_HELP: dict[str, str] = { "restoration_model_name": ( - "Restoration model for video input: basicvsrpp (fast) or ltx (diffusion model, much " - "slower, best quality; no streaming, secondary restoration or denoise). " - "(default: %(default)s)" + "Restoration model for video input: basicvsrpp (fast), ltx (diffusion model, much " + "slower, best quality; no streaming, secondary restoration or denoise) or " + "ltx-undistilled (the same model without the 8-step speed-up, about 3x slower, looks " + "about the same). A missing ltx model is offered for download. (default: %(default)s)" ), "restoration_model_path": ( "Restoration model file (basicvsrpp) or folder (ltx). Default: the model in model_weights/." ), "ltx_seed": "ltx: noise seed. Another seed gives another take on the same restoration. (default: %(default)s)", - "ltx_model": ( - "ltx: distilled (8 steps, about 3x faster) or undistilled (15 steps with STG). Both look " - "about the same; a missing model is offered for download. (default: %(default)s)" - ), "ltx_fast": ( "ltx: about 1.4x faster with slightly less detail. Needs an RTX 50-series (Blackwell) GPU " "and the fast model file. (default: %(default)s)" diff --git a/jasna/main.py b/jasna/main.py index 3ddb3872..4d5ffef0 100644 --- a/jasna/main.py +++ b/jasna/main.py @@ -2,7 +2,6 @@ import logging import sys from pathlib import Path -from typing import get_args from jasna import __version__ from jasna.cli_help import CLI_HELP @@ -17,7 +16,14 @@ check_windows_nvidia_sysmem_fallback_policy, gpu_check_error, ) -from jasna.session_config import LTX_DEFAULT_MODEL, LTX_DEFAULT_SEED, LtxModelName, SessionConfig +from jasna.session_config import LTX_DEFAULT_MODEL, LTX_DEFAULT_SEED, LtxModelName, RestorationModelName, SessionConfig + +# CLI restoration model name -> (restoration backend, LTX model variant). +CLI_RESTORATION_MODELS: dict[str, tuple[RestorationModelName, LtxModelName]] = { + "basicvsrpp": ("basicvsrpp", LTX_DEFAULT_MODEL), + "ltx": ("ltx", "distilled"), + "ltx-undistilled": ("ltx", "undistilled"), +} @@ -42,12 +48,12 @@ def _session_config_from_args( max_detection_gap=int(args.max_detection_gap), min_detection_duration=int(args.min_detection_duration), scene_detection=bool(args.scene_detection), - restoration_model_name=str(args.restoration_model_name), + restoration_model_name=CLI_RESTORATION_MODELS[args.restoration_model_name][0], restoration_model_path=restoration_model_path, ltx_large_canvas=bool(args.ltx_large_canvas), ltx_seed=int(args.ltx_seed), ltx_fast=bool(args.ltx_fast), - ltx_model=str(args.ltx_model), + ltx_model=CLI_RESTORATION_MODELS[args.restoration_model_name][1], compile_basicvsrpp=bool(args.compile_basicvsrpp), max_clip_size=int(args.max_clip_size), temporal_overlap=int(args.temporal_overlap), @@ -197,7 +203,7 @@ def build_parser() -> argparse.ArgumentParser: "--restoration-model-name", type=str, default="basicvsrpp", - choices=["basicvsrpp", "ltx"], + choices=list(CLI_RESTORATION_MODELS), help=CLI_HELP["restoration_model_name"], ) restoration.add_argument( @@ -218,12 +224,6 @@ def build_parser() -> argparse.ArgumentParser: default=LTX_DEFAULT_SEED, help=CLI_HELP["ltx_seed"], ) - restoration.add_argument( - "--ltx-model", - choices=list(get_args(LtxModelName)), - default=LTX_DEFAULT_MODEL, - help=CLI_HELP["ltx_model"], - ) restoration.add_argument( "--ltx-fast", default=False, @@ -896,13 +896,14 @@ def _run_post_export_action() -> None: str(args.detection_model), str(args.detection_model_path), args.detection_score_threshold ) - if args.restoration_model_name == "ltx": + restoration_backend = CLI_RESTORATION_MODELS[args.restoration_model_name][0] + if restoration_backend == "ltx": if is_streaming: parser.error("--restoration-model-name ltx does not support --stream") if args.secondary_restoration != "none" or args.denoise != "none": parser.error("--restoration-model-name ltx does not support --secondary-restoration or --denoise") restoration_model_path = Path( - args.restoration_model_path or default_restoration_model_path(args.restoration_model_name) + args.restoration_model_path or default_restoration_model_path(restoration_backend) ) if not restoration_model_path.exists(): raise FileNotFoundError(str(restoration_model_path)) diff --git a/tests/test_session_config.py b/tests/test_session_config.py index 4638bbf0..c8c8ac0a 100644 --- a/tests/test_session_config.py +++ b/tests/test_session_config.py @@ -126,6 +126,10 @@ def test_cli_ltx_args_are_mapped() -> None: assert defaults.ltx_large_canvas is False assert defaults.ltx_seed == 20260923 + assert _cli_config(["--restoration-model-name", "ltx"]).ltx_model == "distilled" + teacher = _cli_config(["--restoration-model-name", "ltx-undistilled"]) + assert (teacher.restoration_model_name, teacher.ltx_model) == ("ltx", "undistilled") + config = _cli_config(["--restoration-model-name", "ltx", "--ltx-large-canvas", "--ltx-seed", "7"]) assert config.restoration_model_name == "ltx" assert config.ltx_large_canvas is True From b43b2cc2f8958a928fde3b06be3e202138be25bd Mon Sep 17 00:00:00 2001 From: Kruk2 Date: Tue, 29 Sep 2026 19:01:41 +0200 Subject: [PATCH 43/48] test(session): expect the LTX model name in the segment-model load --- tests/test_session_factory.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/test_session_factory.py b/tests/test_session_factory.py index 270715c9..61361008 100644 --- a/tests/test_session_factory.py +++ b/tests/test_session_factory.py @@ -177,7 +177,7 @@ def test_build_pipeline_loads_a_segment_model_the_session_lacks() -> None: build_pipeline(session_config(), session, Path("in.mp4"), Path("out.mp4"), segments=segments) build_pipeline(session_config(), session, Path("in.mp4"), Path("out.mp4"), segments=segments) - from_dir.assert_called_once_with(default_restoration_model_path("ltx"), fast=False) + from_dir.assert_called_once_with(default_restoration_model_path("ltx"), "distilled", fast=False) assert session.ltx_files is from_dir.return_value From 11e18f6d64ab55b96018c696c8e2df138dba6ee1 Mon Sep 17 00:00:00 2001 From: Kruk2 Date: Tue, 29 Sep 2026 19:59:54 +0200 Subject: [PATCH 44/48] feat(gui): choose the LTX model and download missing LTX models A Model menu picks distilled or undistilled. A missing model is offered for download with its size and blocks Start until done. One download worker now serves SD 1.5 and LTX. --- jasna/gui/app.py | 32 +- ..._download_worker.py => download_worker.py} | 19 +- jasna/gui/locales/en.py | 24 ++ jasna/gui/locales/ja.py | 24 ++ jasna/gui/locales/ko.py | 24 ++ jasna/gui/locales/th.py | 24 ++ jasna/gui/locales/zh.py | 24 ++ jasna/gui/ltx_models.py | 124 ++++++++ jasna/gui/models.py | 3 +- jasna/gui/segment_editor.py | 28 +- jasna/gui/settings_panel.py | 15 +- .../settings_sections/image_restoration.py | 10 +- .../settings_sections/restoration_model.py | 220 ++++++++++--- jasna/gui/validation.py | 4 + jasna/gui/video_session.py | 5 +- tests/factories.py | 15 + tests/test_gui_download_worker.py | 29 ++ tests/test_gui_hidpi_scaling.py | 3 + tests/test_gui_ltx_models.py | 289 ++++++++++++++++++ tests/test_gui_queue_layout.py | 5 + tests/test_gui_sd15_download_worker.py | 41 --- tests/test_gui_settings_scroll.py | 13 +- tests/test_gui_settings_sections.py | 20 +- tests/test_ltx_seed_preview.py | 7 + tests/test_restoration_preview.py | 1 + tests/test_segment_editor.py | 53 ++++ tests/test_video_session.py | 10 + 27 files changed, 929 insertions(+), 137 deletions(-) rename jasna/gui/{sd15_download_worker.py => download_worker.py} (70%) create mode 100644 jasna/gui/ltx_models.py create mode 100644 tests/test_gui_download_worker.py create mode 100644 tests/test_gui_ltx_models.py delete mode 100644 tests/test_gui_sd15_download_worker.py diff --git a/jasna/gui/app.py b/jasna/gui/app.py index efa7e6fa..6a319c66 100644 --- a/jasna/gui/app.py +++ b/jasna/gui/app.py @@ -330,13 +330,12 @@ def _build_main_body(self): self._queue_panel.set_on_jobs_changed(self._on_jobs_changed) from jasna.engine_paths import default_restoration_model_path - from jasna.ltx.model_files import bundle_present as ltx_bundle_present + from jasna.gui.ltx_models import LtxModels - self._settings_panel = SettingsPanel( - self._workspace, - self._preset_manager, - ltx_installed=ltx_bundle_present(default_restoration_model_path("ltx")), + self._ltx_models = LtxModels( + default_restoration_model_path("ltx"), self._main_thread, self._on_ltx_models_changed ) + self._settings_panel = SettingsPanel(self._workspace, self._preset_manager, ltx_models=self._ltx_models) self._settings_panel.configure(cursor="arrow") self._settings_panel.set_on_interactive_image_restore(self._open_interactive_image_restore) @@ -438,7 +437,16 @@ def _stop_system_stats_poller(self): self._system_stats_thread.join(timeout=1.0) self._system_stats_thread = None + def _on_ltx_models_changed(self): + self._settings_panel.refresh_ltx_models() + self._update_start_button_state() + def _on_close(self): + if self._ltx_models.downloading: + from tkinter import messagebox + + if not messagebox.askyesno(t("ltx_download_title"), t("ltx_download_close_confirm")): + return if self._video_player_dialog is not None: self._closing_after_player = True self._video_player_dialog.request_close() @@ -550,9 +558,12 @@ def _video_player_closed(self): def _update_start_button_state(self): jobs = self._queue_panel.get_jobs() - can_start = bool(jobs) and not self._preview_gpu_busy + downloading = self._ltx_models.downloading + can_start = bool(jobs) and not self._preview_gpu_busy and not downloading if can_start: self._control_bar.set_start_enabled(True) + elif downloading: + self._control_bar.set_start_enabled(False, t("ltx_downloading_short")) elif self._preview_gpu_busy: self._control_bar.set_start_enabled(False, t("segments_restore_restoring")) else: @@ -594,7 +605,14 @@ def _on_start(self): output_pattern = self._queue_panel.get_output_pattern() settings = self._settings_panel.get_settings() - from jasna.gui.validation import validate_gui_start + from jasna.gui.validation import any_job_uses_ltx, validate_gui_start + + if ( + self._settings_panel.ltx_unavailable_reason() == "model_ltx_not_downloaded" + and any_job_uses_ltx(jobs, settings) + and self._ltx_models.ensure(settings.ltx_model, settings.ltx_fast, on_ready=self._on_start) + ): + return errors = validate_gui_start( settings, jobs, ltx_available=self._settings_panel.ltx_unavailable_reason() is None ) diff --git a/jasna/gui/sd15_download_worker.py b/jasna/gui/download_worker.py similarity index 70% rename from jasna/gui/sd15_download_worker.py rename to jasna/gui/download_worker.py index 0db1f932..418c58e5 100644 --- a/jasna/gui/sd15_download_worker.py +++ b/jasna/gui/download_worker.py @@ -1,18 +1,17 @@ -"""Background SD15 bundle download, kept off the widget layer.""" +"""Background model downloads, kept off the widget layer.""" import threading -from pathlib import Path +from typing import Callable -from jasna.restorer.sd15_download import download_sd15_bundle +DownloadProgress = Callable[[int, int | None], None] -def start_sd15_download( - model_dir: Path, - repo_id: str, - on_percent, - on_done, +def start_download( + download: Callable[[DownloadProgress], None], + on_percent: Callable[[int], None], + on_done: Callable[[str | None], None], ) -> threading.Thread: - """Download the SD15 bundle on a daemon thread. + """Run ``download(progress_callback)`` on a daemon thread. ``on_percent(percent)`` fires on every whole-percent change and ``on_done(error)`` fires once with ``None`` on success or the error text on @@ -35,7 +34,7 @@ def progress(downloaded: int, total: int | None): def worker(): error = None try: - download_sd15_bundle(model_dir, repo_id, progress_callback=progress) + download(progress) except Exception as exc: # surface failure to the user error = str(exc) on_done(error) diff --git a/jasna/gui/locales/en.py b/jasna/gui/locales/en.py index 239d5a42..a6475fca 100644 --- a/jasna/gui/locales/en.py +++ b/jasna/gui/locales/en.py @@ -220,6 +220,22 @@ "model_ltx_description": "Best detail. Much slower and needs a strong NVIDIA GPU.", "model_ltx_needs_nvidia": "Needs an NVIDIA GPU.", "model_ltx_not_installed": "Model not installed.", + "ltx_model": "Model", + "ltx_model_distilled": "Distilled (faster)", + "ltx_model_undistilled": "Undistilled (slightly better, slower)", + "model_ltx_variant_not_installed": "That model is not installed, so the installed one is used.", + "tip_ltx_model": "Undistilled gives slightly better detail but takes longer.\nDistilled is faster.", + "model_ltx_undistilled": "LTX (undistilled)", + "model_ltx_download_needed": "Not installed · {size} download", + "model_ltx_not_downloaded": "Not downloaded yet. Pick it in the settings to download it.", + "model_ltx_fast_not_installed": "Fast mode stays off: its model file is not installed.", + "model_ltx_quality_not_installed": "Fast mode stays on: only the fast model file is installed.", + "ltx_download_title": "LTX model download", + "ltx_download_confirm": "Download {model} ({size})?\nIt can take a while.", + "ltx_downloading": "Downloading the LTX model… {percent}%", + "ltx_downloading_short": "Downloading a model…", + "ltx_download_failed": "The download failed: {message}", + "ltx_download_close_confirm": "A model download is still running. Stop it and close?", "ltx_seed": "Seed", "ltx_fast": "Fast mode", "ltx_large_canvas": "Sharper large mosaics", @@ -432,6 +448,14 @@ "tip_ltx_new_seed": "Try a new random seed", "tip_ltx_fast": "About 1.4x faster with slightly less detail.\nRTX 50-series only.", "tip_ltx_large_canvas": "Restores big mosaics in more detail, but about 3x slower for them.\nNeeds at least 10 GB of free video memory.", + "tip_model_basicvsrpp": "The usual model. Fast and good for most videos.\nPick it unless you want the most detail.", + "tip_model_ltx": "Restores the most detail, but is much slower.\nNeeds a strong NVIDIA GPU. Good for short clips or favorite scenes.", + "tip_model_ltx_undistilled": "LTX with the undistilled model: slightly better detail, but slower.", + "tip_ltx_download": "The model is downloaded once and then kept on this PC.", + "tip_segments_restore_with": "Pick which model restores this range.", + "tip_segments_seed_cancel": "Stop this seed preview.", + "tip_segments_seed_use": "Keep this seed for the range.", + "tip_segments_seed_show": "Switch the preview between the original video and this seed.", "tip_denoise_strength": "Reduces noise and grain in restored areas. Higher = smoother but may lose fine detail.\n\nNone: no denoising. Low/Medium: good starting point. High: heavy smoothing.\nDefault: None", "tip_denoise_step": "When to apply denoising in the pipeline:\n- After Primary: before upscaling (secondary restoration). Denoises at 256x256.\n- After Secondary: after upscaling, right before final output. Denoises at full resolution.\n\nDefault: After Primary", "tip_secondary_restoration": "Optional second pass that upscales restored areas from 256x256 to 1024 pixels. Improves sharpness, especially for close-ups and 4K video.\n\nUNet 4x is a supporter-only model: much faster than TVAI and high quality. Activate Jasna to use it.\nRTX Super Res is a fast alternative with ok quality.\nTopaz TVAI requires a separate purchase and install.", diff --git a/jasna/gui/locales/ja.py b/jasna/gui/locales/ja.py index 83a1eca3..7345fcac 100644 --- a/jasna/gui/locales/ja.py +++ b/jasna/gui/locales/ja.py @@ -220,6 +220,22 @@ "model_ltx_description": "最も精細。かなり遅く、高性能な NVIDIA GPU が必要です。", "model_ltx_needs_nvidia": "NVIDIA GPU が必要です。", "model_ltx_not_installed": "モデルがインストールされていません。", + "ltx_model": "モデル", + "ltx_model_distilled": "蒸留（速い）", + "ltx_model_undistilled": "非蒸留（少し高品質・遅い）", + "model_ltx_variant_not_installed": "そのモデルはインストールされていないため、インストール済みのモデルを使います。", + "tip_ltx_model": "非蒸留は少し精細になりますが、時間がかかります。\n蒸留は速いです。", + "model_ltx_undistilled": "LTX（非蒸留）", + "model_ltx_download_needed": "未インストール · {size} のダウンロード", + "model_ltx_not_downloaded": "まだダウンロードされていません。設定で選ぶとダウンロードできます。", + "model_ltx_fast_not_installed": "高速モードはオフのままです：そのモデルファイルがありません。", + "model_ltx_quality_not_installed": "高速モードはオンのままです：高速モデルのファイルしかありません。", + "ltx_download_title": "LTX モデルのダウンロード", + "ltx_download_confirm": "{model}（{size}）をダウンロードしますか？\n時間がかかることがあります。", + "ltx_downloading": "LTX モデルをダウンロード中… {percent}%", + "ltx_downloading_short": "モデルをダウンロード中…", + "ltx_download_failed": "ダウンロードに失敗しました：{message}", + "ltx_download_close_confirm": "モデルのダウンロード中です。中止して閉じますか？", "ltx_seed": "シード", "ltx_fast": "高速モード", "ltx_large_canvas": "大きなモザイクをより鮮明に", @@ -432,6 +448,14 @@ "tip_ltx_new_seed": "新しいランダムなシードを試す", "tip_ltx_fast": "約 1.4 倍速くなりますが、細部はやや減ります。\nRTX 50 シリーズのみ。", "tip_ltx_large_canvas": "大きなモザイクをより細かく修復しますが、その部分は約 3 倍遅くなります。\n10 GB 以上の空き VRAM が必要です。", + "tip_model_basicvsrpp": "標準のモデル。速く、ほとんどの動画に向いています。\n最高の精細さが必要なければこれを選んでください。", + "tip_model_ltx": "最も精細に修復しますが、かなり遅くなります。\n高性能な NVIDIA GPU が必要です。短いクリップやお気に入りのシーン向けです。", + "tip_model_ltx_undistilled": "非蒸留モデルの LTX：少し精細ですが、遅くなります。", + "tip_ltx_download": "モデルは一度だけダウンロードされ、この PC に保存されます。", + "tip_segments_restore_with": "この範囲を修復するモデルを選びます。", + "tip_segments_seed_cancel": "このシードのプレビューを止めます。", + "tip_segments_seed_use": "この範囲にこのシードを使います。", + "tip_segments_seed_show": "プレビューを元の動画とこのシードで切り替えます。", "tip_denoise_strength": "修復された部分のノイズや粒子を低減します。高いほど滑らかになりますが、細部が失われることがあります。\n\nなし：ノイズ除去なし。低/中：まずこちらをお試しください。高：強力な平滑化。\nデフォルト：なし", "tip_denoise_step": "ノイズ除去を適用するタイミング：\n- 一次修復後：アップスケール（二次修復）前に適用。256x256 の解像度で処理。\n- 二次修復後：アップスケール後、最終出力の直前に適用。フル解像度で処理。\n\nデフォルト：一次修復後", "tip_secondary_restoration": "修復された領域を 256x256 から 1024 ピクセルに拡大するオプションの追加処理です。アップ時や 4K 動画で鮮明さが向上します。\n\nUNet 4x は支援者限定モデルです。TVAI よりかなり高速で高品質です。Jasna を有効化すると使えます。\nRTX Super Res は高速ですが品質はまずまずです。\nTopaz TVAI は別途購入とインストールが必要です。", diff --git a/jasna/gui/locales/ko.py b/jasna/gui/locales/ko.py index d0767dd0..345e2ab9 100644 --- a/jasna/gui/locales/ko.py +++ b/jasna/gui/locales/ko.py @@ -219,6 +219,22 @@ "model_ltx_description": "가장 세밀합니다. 훨씬 느리고 고성능 NVIDIA GPU가 필요합니다.", "model_ltx_needs_nvidia": "NVIDIA GPU가 필요합니다.", "model_ltx_not_installed": "모델이 설치되지 않았습니다.", + "ltx_model": "모델", + "ltx_model_distilled": "증류 (더 빠름)", + "ltx_model_undistilled": "비증류 (조금 더 좋음, 느림)", + "model_ltx_variant_not_installed": "그 모델이 설치되지 않아 설치된 모델을 사용합니다.", + "tip_ltx_model": "비증류는 디테일이 조금 더 좋지만 오래 걸립니다.\n증류는 더 빠릅니다.", + "model_ltx_undistilled": "LTX (비증류)", + "model_ltx_download_needed": "설치 안 됨 · {size} 다운로드", + "model_ltx_not_downloaded": "아직 다운로드하지 않았습니다. 설정에서 선택하면 다운로드됩니다.", + "model_ltx_fast_not_installed": "빠른 모드는 꺼진 상태로 둡니다: 해당 모델 파일이 없습니다.", + "model_ltx_quality_not_installed": "빠른 모드는 켜진 상태로 둡니다: 빠른 모델 파일만 있습니다.", + "ltx_download_title": "LTX 모델 다운로드", + "ltx_download_confirm": "{model}({size})을(를) 다운로드할까요?\n시간이 걸릴 수 있습니다.", + "ltx_downloading": "LTX 모델 다운로드 중… {percent}%", + "ltx_downloading_short": "모델 다운로드 중…", + "ltx_download_failed": "다운로드 실패: {message}", + "ltx_download_close_confirm": "모델 다운로드가 아직 진행 중입니다. 중지하고 닫을까요?", "ltx_seed": "시드", "ltx_fast": "빠른 모드", "ltx_large_canvas": "큰 모자이크 더 선명하게", @@ -412,6 +428,14 @@ "tip_ltx_new_seed": "새 무작위 시드 사용", "tip_ltx_fast": "약 1.4배 빠르지만 디테일이 조금 줄어듭니다.\nRTX 50 시리즈 전용.", "tip_ltx_large_canvas": "큰 모자이크를 더 세밀하게 복원하지만 해당 부분은 약 3배 느려집니다.\n여유 VRAM이 10 GB 이상 필요합니다.", + "tip_model_basicvsrpp": "기본 모델입니다. 빠르고 대부분의 영상에 알맞습니다.\n최고 디테일이 필요 없다면 이것을 고르세요.", + "tip_model_ltx": "디테일을 가장 많이 살리지만 훨씬 느립니다.\n강력한 NVIDIA GPU가 필요합니다. 짧은 클립이나 좋아하는 장면에 좋습니다.", + "tip_model_ltx_undistilled": "비증류 모델을 쓰는 LTX: 디테일이 조금 더 좋지만 느립니다.", + "tip_ltx_download": "모델은 한 번만 다운로드되어 이 PC에 보관됩니다.", + "tip_segments_restore_with": "이 구간을 복원할 모델을 고릅니다.", + "tip_segments_seed_cancel": "이 시드 미리보기를 멈춥니다.", + "tip_segments_seed_use": "이 구간에 이 시드를 사용합니다.", + "tip_segments_seed_show": "미리보기를 원본 영상과 이 시드 사이에서 전환합니다.", "tip_denoise_strength": "복원된 영역의 노이즈와 입자를 줄여줍니다. 높을수록 부드럽지만 세부 디테일이 줄어들 수 있습니다.\n\n없음: 노이즈 제거 안 함. 낮음/중간: 좋은 시작점. 높음: 강력한 평활화.\n기본값: 없음", "tip_denoise_step": "노이즈 제거를 적용하는 시점:\n- 1차 복원 후: 업스케일링(2차 복원) 전에 적용. 256x256 해상도에서 처리.\n- 2차 복원 후: 업스케일링 후, 최종 출력 직전에 적용. 전체 해상도에서 처리.\n\n기본값: 1차 복원 후", "tip_secondary_restoration": "복원된 영역을 256x256에서 1024 픽셀로 확대하는 선택적 추가 처리입니다. 클로즈업이나 4K 영상에서 선명도가 향상됩니다.\n\nUNet 4x는 서포터 전용 모델입니다. TVAI보다 훨씬 빠르고 고품질입니다. Jasna를 활성화하면 사용할 수 있습니다.\nRTX Super Res는 빠르지만 품질은 보통입니다.\nTopaz TVAI는 별도 구매 및 설치가 필요합니다.", diff --git a/jasna/gui/locales/th.py b/jasna/gui/locales/th.py index 01def76e..e5badfce 100644 --- a/jasna/gui/locales/th.py +++ b/jasna/gui/locales/th.py @@ -220,6 +220,22 @@ "model_ltx_description": "รายละเอียดดีที่สุด ช้ากว่ามาก และต้องใช้ GPU NVIDIA ที่แรง", "model_ltx_needs_nvidia": "ต้องใช้ GPU NVIDIA", "model_ltx_not_installed": "ยังไม่ได้ติดตั้งโมเดล", + "ltx_model": "โมเดล", + "ltx_model_distilled": "แบบกลั่น (เร็วกว่า)", + "ltx_model_undistilled": "แบบไม่กลั่น (ดีขึ้นเล็กน้อย ช้ากว่า)", + "model_ltx_variant_not_installed": "ไม่ได้ติดตั้งโมเดลนั้น จึงใช้โมเดลที่ติดตั้งไว้", + "tip_ltx_model": "แบบไม่กลั่นให้รายละเอียดดีขึ้นเล็กน้อย แต่ใช้เวลานานกว่า\nแบบกลั่นเร็วกว่า", + "model_ltx_undistilled": "LTX (ไม่กลั่น)", + "model_ltx_download_needed": "ยังไม่ได้ติดตั้ง · ดาวน์โหลด {size}", + "model_ltx_not_downloaded": "ยังไม่ได้ดาวน์โหลด เลือกในการตั้งค่าเพื่อดาวน์โหลด", + "model_ltx_fast_not_installed": "โหมดเร็วยังปิดอยู่: ไม่มีไฟล์โมเดลของโหมดนี้", + "model_ltx_quality_not_installed": "โหมดเร็วยังเปิดอยู่: มีเฉพาะไฟล์โมเดลแบบเร็ว", + "ltx_download_title": "ดาวน์โหลดโมเดล LTX", + "ltx_download_confirm": "ดาวน์โหลด {model} ({size}) ไหม?\nอาจใช้เวลาสักพัก", + "ltx_downloading": "กำลังดาวน์โหลดโมเดล LTX… {percent}%", + "ltx_downloading_short": "กำลังดาวน์โหลดโมเดล…", + "ltx_download_failed": "ดาวน์โหลดไม่สำเร็จ: {message}", + "ltx_download_close_confirm": "ยังดาวน์โหลดโมเดลอยู่ หยุดแล้วปิดไหม?", "ltx_seed": "ซีด", "ltx_fast": "โหมดเร็ว", "ltx_large_canvas": "โมเสกขนาดใหญ่คมชัดขึ้น", @@ -414,6 +430,14 @@ "tip_ltx_new_seed": "ลองซีดสุ่มใหม่", "tip_ltx_fast": "เร็วขึ้นประมาณ 1.4 เท่า รายละเอียดลดลงเล็กน้อย\nเฉพาะ RTX ซีรีส์ 50", "tip_ltx_large_canvas": "ฟื้นฟูโมเสกขนาดใหญ่ได้ละเอียดขึ้น แต่ส่วนนั้นช้าลงประมาณ 3 เท่า\nต้องมี VRAM ว่างอย่างน้อย 10 GB", + "tip_model_basicvsrpp": "โมเดลปกติ เร็วและเหมาะกับวิดีโอส่วนใหญ่\nเลือกอันนี้ถ้าไม่ต้องการรายละเอียดสูงสุด", + "tip_model_ltx": "กู้รายละเอียดได้มากที่สุด แต่ช้ากว่ามาก\nต้องใช้ GPU NVIDIA แรงๆ เหมาะกับคลิปสั้นหรือฉากโปรด", + "tip_model_ltx_undistilled": "LTX แบบไม่กลั่น: รายละเอียดดีขึ้นเล็กน้อย แต่ช้ากว่า", + "tip_ltx_download": "ดาวน์โหลดโมเดลครั้งเดียว แล้วเก็บไว้ในเครื่องนี้", + "tip_segments_restore_with": "เลือกโมเดลที่ใช้กู้ช่วงนี้", + "tip_segments_seed_cancel": "หยุดพรีวิวซีดนี้", + "tip_segments_seed_use": "ใช้ซีดนี้กับช่วงนี้", + "tip_segments_seed_show": "สลับพรีวิวระหว่างวิดีโอต้นฉบับกับซีดนี้", "tip_denoise_strength": "ลดสัญญาณรบกวนและเม็ดในพื้นที่ที่ฟื้นฟูแล้ว ค่ามาก = ราบรื่นขึ้นแต่อาจสูญเสียรายละเอียด\n\nไม่มี: ไม่ลดสัญญาณรบกวน ต่ำ/ปานกลาง: จุดเริ่มต้นที่ดี สูง: ปรับเรียบอย่างหนัก\nค่าเริ่มต้น: ไม่มี", "tip_denoise_step": "เวลาที่จะลดสัญญาณรบกวนในขั้นตอนการประมวลผล:\n- หลังขั้นแรก: ก่อนการขยาย (การฟื้นฟูขั้นที่สอง) ลดสัญญาณรบกวนที่ 256x256\n- หลังขั้นที่สอง: หลังการขยาย ก่อนเอาต์พุตสุดท้าย ลดสัญญาณรบกวนที่ความละเอียดเต็ม\n\nค่าเริ่มต้น: หลังขั้นแรก", "tip_secondary_restoration": "ขั้นตอนเสริมที่ขยายพื้นที่ฟื้นฟูจาก 256x256 เป็น 1024 พิกเซล ช่วยเพิ่มความคมชัดโดยเฉพาะสำหรับภาพใกล้และวิดีโอ 4K\n\nUNet 4x เป็นโมเดลเฉพาะผู้สนับสนุน เร็วกว่า TVAI มากและคุณภาพสูง เปิดใช้งาน Jasna เพื่อใช้โมเดลนี้\nRTX Super Res เร็วแต่คุณภาพพอใช้\nTopaz TVAI ต้องซื้อและติดตั้งแยกต่างหาก", diff --git a/jasna/gui/locales/zh.py b/jasna/gui/locales/zh.py index 183aee5b..ad15dd63 100644 --- a/jasna/gui/locales/zh.py +++ b/jasna/gui/locales/zh.py @@ -220,6 +220,22 @@ "model_ltx_description": "细节最好。速度慢很多，需要高性能 NVIDIA 显卡。", "model_ltx_needs_nvidia": "需要 NVIDIA 显卡。", "model_ltx_not_installed": "模型未安装。", + "ltx_model": "模型", + "ltx_model_distilled": "蒸馏（更快）", + "ltx_model_undistilled": "未蒸馏（稍好，更慢）", + "model_ltx_variant_not_installed": "该模型未安装，改用已安装的模型。", + "tip_ltx_model": "未蒸馏的细节稍好，但耗时更长。\n蒸馏版更快。", + "model_ltx_undistilled": "LTX（未蒸馏）", + "model_ltx_download_needed": "未安装 · 需下载 {size}", + "model_ltx_not_downloaded": "尚未下载。在设置中选择它即可下载。", + "model_ltx_fast_not_installed": "快速模式保持关闭：未安装其模型文件。", + "model_ltx_quality_not_installed": "快速模式保持开启：只安装了快速模型文件。", + "ltx_download_title": "下载 LTX 模型", + "ltx_download_confirm": "下载 {model}（{size}）？\n可能需要一段时间。", + "ltx_downloading": "正在下载 LTX 模型… {percent}%", + "ltx_downloading_short": "正在下载模型…", + "ltx_download_failed": "下载失败：{message}", + "ltx_download_close_confirm": "模型仍在下载。停止并关闭？", "ltx_seed": "种子", "ltx_fast": "快速模式", "ltx_large_canvas": "大马赛克更清晰", @@ -432,6 +448,14 @@ "tip_ltx_new_seed": "换一个随机种子", "tip_ltx_fast": "约快 1.4 倍，细节略少。\n仅限 RTX 50 系列。", "tip_ltx_large_canvas": "大马赛克修复得更细致，但这些部分约慢 3 倍。\n需要至少 10 GB 可用显存。", + "tip_model_basicvsrpp": "常用模型。速度快，适合大多数视频。\n除非需要最高细节，否则选它。", + "tip_model_ltx": "修复细节最多，但慢得多。\n需要强劲的 NVIDIA GPU。适合短片或喜欢的场景。", + "tip_model_ltx_undistilled": "使用未蒸馏模型的 LTX：细节稍好，但更慢。", + "tip_ltx_download": "模型只需下载一次，之后保存在这台电脑上。", + "tip_segments_restore_with": "选择修复此片段的模型。", + "tip_segments_seed_cancel": "停止此种子预览。", + "tip_segments_seed_use": "此片段保留这个种子。", + "tip_segments_seed_show": "在原始视频和此种子之间切换预览。", "tip_denoise_strength": "降低修复区域的噪点和颗粒感。强度越高画面越平滑，但可能丢失细节。\n\n无：不降噪。低/中：推荐起步值。高：强力平滑。\n默认值：无", "tip_denoise_step": "降噪在处理流程中的应用时机：\n- 主修复后：在放大（二次修复）之前降噪，在 256x256 分辨率下处理。\n- 二次修复后：在放大之后、最终输出之前降噪，在完整分辨率下处理。\n\n默认值：主修复后", "tip_secondary_restoration": "可选的第二步处理，将修复区域从 256x256 放大到 1024 像素。可提升清晰度，特别是近景和 4K 视频。\n\nUNet 4x 是支持者专属模型，比 TVAI 快很多，质量高。激活 Jasna 后可用。\nRTX Super Res 是速度快但质量一般的替代方案。\nTopaz TVAI 需要单独购买和安装。", diff --git a/jasna/gui/ltx_models.py b/jasna/gui/ltx_models.py new file mode 100644 index 00000000..826dd537 --- /dev/null +++ b/jasna/gui/ltx_models.py @@ -0,0 +1,124 @@ +"""Which LTX models the GUI can run, and the one download that fetches a missing model.""" + +from collections.abc import Callable, Mapping +from dataclasses import dataclass +from pathlib import Path + +from jasna.gui.download_worker import start_download +from jasna.gui.locales import t +from jasna.gui.queues import MainThreadCalls +from jasna.ltx import model_files +from jasna.ltx.model_files import LTX_MODELS, DownloadableFile, download_size_text +from jasna.session_config import LtxModelName + +Variant = tuple[LtxModelName, bool] + + +def model_key(restoration_model: str, ltx_model: str) -> str: + """Locale stem naming a restoration choice: ``model_`` and ``tip_model_``.""" + if restoration_model != "ltx": + return "basicvsrpp" + return "ltx" if ltx_model == "distilled" else f"ltx_{ltx_model}" + + +@dataclass(frozen=True) +class LtxInstallState: + """The files each (model, fast) variant still needs, and whether they can be downloaded.""" + + missing: Mapping[Variant, tuple[DownloadableFile, ...]] + downloadable: bool + + def installed(self, model: LtxModelName, fast: bool) -> bool: + return not self.missing[(model, fast)] + + def usable(self, model: LtxModelName, fast: bool) -> bool: + return self.downloadable or self.installed(model, fast) + + def download_size(self, model: LtxModelName, fast: bool) -> str: + return download_size_text(list(self.missing[(model, fast)])) + + +def read_install_state(directory: Path) -> LtxInstallState: + return LtxInstallState( + { + (model, fast): tuple(model_files.missing_downloads(directory, model, fast=fast)) + for model in LTX_MODELS + for fast in (False, True) + }, + model_files.LTX_DOWNLOAD_URL is not None, + ) + + +def card_unavailable_reason(state: LtxInstallState, *, nvidia: bool | None) -> str | None: + """Locale key saying why the LTX card cannot be picked; ``nvidia`` is None until the GPU is known.""" + if nvidia is False: + return "model_ltx_needs_nvidia" + if not any(state.usable(model, fast) for model in LTX_MODELS for fast in (False, True)): + return "model_ltx_not_installed" + return None + + +def run_unavailable_reason( + state: LtxInstallState, model: LtxModelName, fast: bool, *, nvidia: bool | None +) -> str | None: + """Locale key saying why LTX cannot restore right now with this variant.""" + if nvidia is False: + return "model_ltx_needs_nvidia" + if state.installed(model, fast): + return None + return "model_ltx_not_downloaded" if state.downloadable else "model_ltx_not_installed" + + +class LtxModels: + """The install state plus the single running download. UI glue: it asks with a message + box and runs the download on a worker thread, reporting back on the main thread.""" + + def __init__(self, directory: Path, main_thread: MainThreadCalls, on_change: Callable[[], None]) -> None: + self._directory = directory + self._main_thread = main_thread + self._on_change = on_change + self.state = read_install_state(directory) + self.percent: int | None = None + + @property + def downloading(self) -> bool: + return self.percent is not None + + def ensure(self, model: LtxModelName, fast: bool, *, on_ready: Callable[[], None]) -> bool: + """True when the variant is installed, or the user agreed to download it and the + download started (``on_ready`` runs after it succeeds); False otherwise.""" + from tkinter import messagebox + + files = self.state.missing[(model, fast)] + if not files: + return True + if self.downloading or not self.state.downloadable: + return False + name = t(f"model_{model_key('ltx', model)}") + if not messagebox.askyesno( + t("ltx_download_title"), t("ltx_download_confirm", model=name, size=download_size_text(list(files))) + ): + return False + self.percent = 0 + self._on_change() + start_download( + lambda progress: model_files.download_files(self._directory, list(files), progress), + lambda percent: self._main_thread.post(lambda: self._set_percent(percent)), + lambda error: self._main_thread.post(lambda: self._finish(error, on_ready)), + ) + return True + + def _set_percent(self, percent: int) -> None: + self.percent = percent + self._on_change() + + def _finish(self, error: str | None, on_ready: Callable[[], None]) -> None: + from tkinter import messagebox + + self.percent = None + self.state = read_install_state(self._directory) + self._on_change() + if error: + messagebox.showerror(t("ltx_download_title"), t("ltx_download_failed", message=error)) + else: + on_ready() diff --git a/jasna/gui/models.py b/jasna/gui/models.py index 110b06a1..4b9ed0f7 100644 --- a/jasna/gui/models.py +++ b/jasna/gui/models.py @@ -10,7 +10,7 @@ from jasna.gui.paths import get_settings_path from jasna.segments import SegmentRange -from jasna.session_config import LTX_DEFAULT_SEED +from jasna.session_config import LTX_DEFAULT_MODEL, LTX_DEFAULT_SEED logger = logging.getLogger(__name__) @@ -106,6 +106,7 @@ def begin_processing(self) -> JobProcessingSnapshot | None: class AppSettings: # Restoration model restoration_model: str = "basicvsrpp" # basicvsrpp, ltx + ltx_model: str = LTX_DEFAULT_MODEL # distilled, undistilled ltx_seed: int = LTX_DEFAULT_SEED ltx_fast: bool = False ltx_large_canvas: bool = False diff --git a/jasna/gui/segment_editor.py b/jasna/gui/segment_editor.py index 4b8924ff..54d32493 100644 --- a/jasna/gui/segment_editor.py +++ b/jasna/gui/segment_editor.py @@ -18,6 +18,7 @@ from jasna.gui.models import AppSettings, JobItem from jasna.gui.components import AutoHidingScrollableFrame, Tooltip, format_duration, grab_modal from jasna.gui.icons import CompactSwitch, NativeIconButton +from jasna.gui.ltx_models import model_key from jasna.gui.ltx_seed_preview import LtxSeedPreviewWorker, SeedFailed, SeedFrame, SeedProgress, SeedReady from jasna.gui.restoration_preview import ( RestorationClip, @@ -98,6 +99,8 @@ def __init__( self._is_gpu_busy = is_gpu_busy self._set_preview_gpu_busy = set_preview_gpu_busy self._ltx_unavailable_reason = ltx_unavailable_reason + self._ltx_key = model_key("ltx", get_settings().ltx_model) + self._ltx_label = t(f"model_{self._ltx_key}") self._on_saved = on_saved self._on_closed = on_closed self._state: SegmentEditorState | None = None @@ -437,6 +440,8 @@ def _build_preview_card(self, body: ctk.CTkFrame) -> None: font=(Fonts.FAMILY, Fonts.SIZE_SMALL), ) self._seed_view_label.pack(side="right", padx=(0, 6)) + for widget in (self._seed_view_label, self._seed_view_toggle): + Tooltip(widget, t("tip_segments_seed_show")) projection_names = { "raw": t("segments_vr_projection_raw"), "fisheye": t("segments_vr_projection_fisheye"), @@ -617,12 +622,14 @@ def _fit_range_controls(self, _event=None) -> None: def _build_restoration_inspector(self, parent: ctk.CTkFrame) -> None: self._inspector = ctk.CTkFrame(parent, fg_color=Colors.BG_PANEL, corner_radius=Sizing.BORDER_RADIUS) self._inspector.grid_columnconfigure(1, weight=1) - ctk.CTkLabel( + self._restore_with_label = ctk.CTkLabel( self._inspector, text=t("segments_restore_with"), text_color=Colors.TEXT_PRIMARY, font=(Fonts.FAMILY, Fonts.SIZE_SMALL), - ).grid(row=0, column=0, sticky="w", padx=(8, 6), pady=4) + ) + self._restore_with_label.grid(row=0, column=0, sticky="w", padx=(8, 6), pady=4) + Tooltip(self._restore_with_label, t("tip_segments_restore_with")) models = ctk.CTkFrame(self._inspector, fg_color="transparent") models.grid(row=0, column=1, sticky="w", pady=4) self._range_model = ctk.StringVar(value="basicvsrpp") @@ -630,7 +637,7 @@ def _build_restoration_inspector(self, parent: ctk.CTkFrame) -> None: for model in ("basicvsrpp", "ltx"): radio = ctk.CTkRadioButton( models, - text=t(f"model_{model}"), + text=self._ltx_label if model == "ltx" else t("model_basicvsrpp"), variable=self._range_model, value=model, command=self._on_range_model_changed, @@ -642,6 +649,7 @@ def _build_restoration_inspector(self, parent: ctk.CTkFrame) -> None: ) radio.pack(side="left", padx=(0, 10)) self._model_radios[model] = radio + Tooltip(self._model_radios["basicvsrpp"], t("tip_model_basicvsrpp")) self._ltx_radio_tooltip = Tooltip(self._model_radios["ltx"], "") self._seed_label = ctk.CTkLabel( self._inspector, @@ -652,7 +660,8 @@ def _build_restoration_inspector(self, parent: ctk.CTkFrame) -> None: self._seed_entry = ctk.CTkEntry(self._inspector, font=(Fonts.FAMILY_MONO, Fonts.SIZE_SMALL)) self._seed_entry.bind("", self._commit_seed_entry) self._seed_entry.bind("", self._commit_seed_entry) - Tooltip(self._seed_entry, t("tip_ltx_seed")) + for widget in (self._seed_label, self._seed_entry): + Tooltip(widget, t("tip_ltx_seed")) self._new_seed_btn = NativeIconButton( self._inspector, "reset", @@ -700,6 +709,7 @@ def _build_restoration_inspector(self, parent: ctk.CTkFrame) -> None: hover_color=Colors.BORDER_LIGHT, command=self._cancel_seed, ) + Tooltip(self._seed_cancel_btn, t("tip_segments_seed_cancel")) self._seed_progress = ctk.CTkProgressBar(self._seed_box, height=8, progress_color=Colors.MODEL_LTX) self._seed_status = ctk.CTkLabel( self._seed_box, @@ -768,7 +778,7 @@ def _refresh_inspector(self) -> None: self._model_radios["ltx"].configure( state="disabled" if locked or (blocked is not None and restoration.model != "ltx") else "normal" ) - self._ltx_radio_tooltip.set_text(t(blocked) if blocked else t("model_ltx_description")) + self._ltx_radio_tooltip.set_text(t(blocked) if blocked else t(f"tip_model_{self._ltx_key}")) if restoration.model == "ltx": self._seed_label.grid(row=1, column=0, sticky="w", padx=(8, 6), pady=(0, 6)) self._seed_entry.grid(row=1, column=1, sticky="ew", pady=(0, 6)) @@ -878,6 +888,7 @@ def _render_seed_chips(self, selected: SegmentRange) -> None: command=lambda value=shown: self._use_tried_seed(value), ) use.pack(side="left", padx=(4, 0)) + Tooltip(use, t("tip_segments_seed_use")) self._seed_chip_widgets.append(use) def _shown_seed_for(self, segment: SegmentRange) -> int | None: @@ -1010,11 +1021,10 @@ def _seed_view_image(self) -> Image.Image | None: self._seed_image = (frame, image.convert("RGB")) return self._seed_image[1] - @staticmethod - def _model_text(segment: SegmentRange) -> str: + def _model_text(self, segment: SegmentRange) -> str: restoration = segment.restoration if restoration.model == "ltx": - return f"{t('model_ltx')} {restoration.ltx_seed}" + return f"{self._ltx_label} {restoration.ltx_seed}" return t(f"model_{restoration.model}") def _build_scan_and_timeline(self, metadata: VideoMetadata) -> None: @@ -1093,7 +1103,7 @@ def _build_scan_and_timeline(self, metadata: VideoMetadata) -> None: legend.pack(fill="x", padx=20, pady=(0, 2)) for color, label in ( (Colors.PRIMARY, t("model_basicvsrpp")), - (Colors.MODEL_LTX, t("model_ltx")), + (Colors.MODEL_LTX, self._ltx_label), (Colors.STATUS_WARNING, t("segments_legend_detected")), ("#f8fafc", t("segments_legend_playhead")), ): diff --git a/jasna/gui/settings_panel.py b/jasna/gui/settings_panel.py index ba94a171..15b2835d 100644 --- a/jasna/gui/settings_panel.py +++ b/jasna/gui/settings_panel.py @@ -16,6 +16,7 @@ ) from jasna.gui.icons import NativeIconButton from jasna.gui.locales import t +from jasna.gui.ltx_models import LtxModels from jasna.gui.settings_sections.advanced import ( TEMPORAL_FILTER_SLIDER_MAX, AdvancedSection, @@ -35,7 +36,7 @@ class SettingsPanel(ctk.CTkFrame): """Right panel composing the settings sections; widgets live in the sections.""" - def __init__(self, master, preset_manager: PresetManager, *, ltx_installed: bool, **kwargs): + def __init__(self, master, preset_manager: PresetManager, *, ltx_models: LtxModels, **kwargs): super().__init__( master, fg_color=Colors.BG_PANEL, @@ -53,7 +54,7 @@ def __init__(self, master, preset_manager: PresetManager, *, ltx_installed: bool self._build_preset_bar() self._build_scrollable() - self._build_sections(ltx_installed) + self._build_sections(ltx_models) self._apply_preset(self._current_preset) def _build_preset_bar(self): @@ -212,13 +213,13 @@ def _build_scrollable(self): ) self._scroll.pack(fill="both", expand=True, padx=Sizing.PADDING_MEDIUM, pady=(0, Sizing.PADDING_MEDIUM)) - def _build_sections(self, ltx_installed: bool): + def _build_sections(self, ltx_models: LtxModels): self._model_section = RestorationModelSection( self._scroll, self._widgets, self._mark_modified, self._on_restoration_model_changed, - ltx_installed=ltx_installed, + ltx_models=ltx_models, ) self._model_aware_sections = [ BasicSection( @@ -249,9 +250,13 @@ def _on_restoration_model_changed(self, model: str): section.set_model(model) def ltx_unavailable_reason(self) -> str | None: - """Locale key saying why LTX cannot run on this PC, or None when it can.""" + """Locale key saying why the chosen LTX model cannot run now, or None when it can.""" return self._model_section.unavailable_reason() + def refresh_ltx_models(self): + self._model_section.refresh() + self._update_modified_indicator() + def set_gpu_support(self, *, nvidia: bool, blackwell: bool): self._model_section.set_gpu_support(nvidia=nvidia, blackwell=blackwell) self._update_modified_indicator() diff --git a/jasna/gui/settings_sections/image_restoration.py b/jasna/gui/settings_sections/image_restoration.py index 7105dd18..93c82283 100644 --- a/jasna/gui/settings_sections/image_restoration.py +++ b/jasna/gui/settings_sections/image_restoration.py @@ -6,7 +6,7 @@ from jasna.gui.icons import CompactSwitch from jasna.gui.locales import t from jasna.gui.queues import MainThreadCalls -from jasna.gui.sd15_download_worker import start_sd15_download +from jasna.gui.download_worker import start_download from jasna.gui.settings_sections.widgets import add_setting_label, create_slider_value_label from jasna.gui.theme import Colors, Fonts, Sizing @@ -174,7 +174,13 @@ def finish(): main_thread.post(finish) - start_sd15_download(SD15_DIR, SD15_HF_REPO, on_percent, on_done) + from jasna.restorer.sd15_download import download_sd15_bundle + + start_download( + lambda progress: download_sd15_bundle(SD15_DIR, SD15_HF_REPO, progress_callback=progress), + on_percent, + on_done, + ) def _finish_download(self, error: str | None): btn = self._widgets["image_restore_download_btn"] diff --git a/jasna/gui/settings_sections/restoration_model.py b/jasna/gui/settings_sections/restoration_model.py index 3ec6b9b1..3d9e9b1e 100644 --- a/jasna/gui/settings_sections/restoration_model.py +++ b/jasna/gui/settings_sections/restoration_model.py @@ -8,22 +8,15 @@ from jasna.gui.components import CollapsibleSection, Tooltip from jasna.gui.icons import CompactSwitch, NativeIconButton from jasna.gui.locales import t -from jasna.gui.settings_sections.widgets import add_setting_label, get_tooltip +from jasna.gui.ltx_models import LtxModels, card_unavailable_reason, run_unavailable_reason +from jasna.gui.settings_sections.widgets import ValueOptionMenu, add_setting_label, get_tooltip from jasna.gui.theme import Colors, Fonts, Sizing -from jasna.session_config import LTX_DEFAULT_SEED +from jasna.ltx.model_files import LTX_MODELS +from jasna.session_config import LTX_DEFAULT_MODEL, LTX_DEFAULT_SEED _MODELS = ("basicvsrpp", "ltx") -def ltx_unavailable_reason(*, installed: bool, nvidia: bool | None) -> str | None: - """Locale key saying why LTX cannot be picked; ``nvidia`` is None until the GPU is known.""" - if nvidia is False: - return "model_ltx_needs_nvidia" - if not installed: - return "model_ltx_not_installed" - return None - - def parse_seed(text: str) -> int: try: return int(text.strip()) @@ -34,8 +27,6 @@ def parse_seed(text: str) -> int: class _ModelCard(ctk.CTkFrame): def __init__(self, master, variable: ctk.StringVar, value: str, on_select: Callable[[str], None]): super().__init__(master, fg_color=Colors.BG_CARD, corner_radius=8, border_width=2, border_color=Colors.BG_CARD) - self._value = value - self._on_select = on_select self._enabled = True self.radio = ctk.CTkRadioButton( self, @@ -56,11 +47,13 @@ def __init__(self, master, variable: ctk.StringVar, value: str, on_select: Calla font=(Fonts.FAMILY, Fonts.SIZE_SMALL), anchor="w", justify="left", - wraplength=260, + wraplength=190, ) self.description.pack(fill="x", padx=(40, 12), pady=(0, 10)) for widget in (self, self.description): widget.bind("", self._clicked) + for widget in (self.radio, self.description): + Tooltip(widget, get_tooltip(f"model_{value}")) def _clicked(self, _event) -> None: if self._enabled: @@ -75,11 +68,11 @@ def set_selected(self, selected: bool) -> None: class RestorationModelSection: - def __init__(self, parent, widgets: dict, on_modified, on_model_changed, *, ltx_installed: bool): + def __init__(self, parent, widgets: dict, on_modified, on_model_changed, *, ltx_models: LtxModels): self._widgets = widgets self._on_modified = on_modified self._on_model_changed = on_model_changed - self._ltx_installed = ltx_installed + self._ltx_models = ltx_models self._nvidia: bool | None = None self._enabled = True @@ -92,16 +85,50 @@ def __init__(self, parent, widgets: dict, on_modified, on_model_changed, *, ltx_ inner.pack(fill="x", padx=Sizing.PADDING_MEDIUM, pady=Sizing.PADDING_MEDIUM) self._model = self._widgets["restoration_model"] = ctk.StringVar(value="basicvsrpp") - cards = ctk.CTkFrame(inner, fg_color="transparent") - cards.pack(fill="x") - cards.grid_columnconfigure((0, 1), weight=1, uniform="model") + self._cards_frame = ctk.CTkFrame(inner, fg_color="transparent") + self._cards_frame.pack(fill="x") + self._cards_frame.grid_columnconfigure((0, 1), weight=1, uniform="model") self._cards = {} for column, value in enumerate(_MODELS): - card = _ModelCard(cards, self._model, value, self._on_card_selected) + card = _ModelCard(self._cards_frame, self._model, value, self._on_card_selected) card.grid(row=0, column=column, sticky="nsew", padx=(0, 4) if column == 0 else (4, 0)) self._cards[value] = card + self._download_row = ctk.CTkFrame(inner, fg_color="transparent") + self._download_label = ctk.CTkLabel( + self._download_row, text="", text_color=Colors.TEXT_PRIMARY, font=(Fonts.FAMILY, Fonts.SIZE_SMALL), anchor="w" + ) + self._download_label.pack(fill="x") + self._download_bar = ctk.CTkProgressBar(self._download_row, progress_color=Colors.PRIMARY, fg_color=Colors.BG_CARD) + self._download_bar.pack(fill="x", pady=(2, 0)) + for widget in (self._download_label, self._download_bar): + Tooltip(widget, get_tooltip("ltx_download")) + self._ltx_options = ctk.CTkFrame(inner, fg_color=Colors.BG_CARD, corner_radius=6) + model_row = ctk.CTkFrame(self._ltx_options, fg_color="transparent") + model_row.pack(fill="x") + add_setting_label(model_row, "ltx_model", in_card=True) + self._ltx_model = self._widgets["ltx_model"] = ValueOptionMenu( + model_row, + options={model: t(f"ltx_model_{model}") for model in LTX_MODELS}, + command=self._on_ltx_model_selected, + fg_color=Colors.BG_PANEL, + button_color=Colors.BG_PANEL, + button_hover_color=Colors.BORDER_LIGHT, + dropdown_fg_color=Colors.BG_CARD, + dropdown_hover_color=Colors.PRIMARY, + text_color=Colors.TEXT_PRIMARY, + width=260, + ) + self._ltx_model.pack(side="right", padx=12, pady=8) + Tooltip(self._ltx_model, get_tooltip("ltx_model")) + self._set_ltx_model(LTX_DEFAULT_MODEL) + self._model_row = model_row + self._install_status = ctk.CTkLabel( + self._ltx_options, text="", text_color=Colors.STATUS_WARNING, font=(Fonts.FAMILY, Fonts.SIZE_SMALL), anchor="e" + ) + Tooltip(self._install_status, get_tooltip("ltx_download")) + seed_row = ctk.CTkFrame(self._ltx_options, fg_color="transparent") seed_row.pack(fill="x") add_setting_label(seed_row, "ltx_seed", in_card=True) @@ -124,25 +151,90 @@ def __init__(self, parent, widgets: dict, on_modified, on_model_changed, *, ltx_ ) self._widgets["ltx_seed"].pack(side="right", pady=8) self._widgets["ltx_seed"].bind("", lambda _event: self._on_modified()) + Tooltip(self._widgets["ltx_seed"], get_tooltip("ltx_seed")) self._fast_row = ctk.CTkFrame(self._ltx_options, fg_color="transparent") add_setting_label(self._fast_row, "ltx_fast", in_card=True) - self._widgets["ltx_fast"] = CompactSwitch(self._fast_row, self._on_modified, Colors.BG_CARD) + self._widgets["ltx_fast"] = CompactSwitch(self._fast_row, self._on_fast_toggled, Colors.BG_CARD) self._widgets["ltx_fast"].pack(side="right", padx=12, pady=8) + Tooltip(self._widgets["ltx_fast"], get_tooltip("ltx_fast")) large_row = ctk.CTkFrame(self._ltx_options, fg_color="transparent") large_row.pack(fill="x") add_setting_label(large_row, "ltx_large_canvas", in_card=True) self._widgets["ltx_large_canvas"] = CompactSwitch(large_row, self._on_modified, Colors.BG_CARD) self._widgets["ltx_large_canvas"].pack(side="right", padx=12, pady=8) + Tooltip(self._widgets["ltx_large_canvas"], get_tooltip("ltx_large_canvas")) self._large_row = large_row - self._refresh_ltx_availability() + self._notice = ctk.CTkLabel( + self._ltx_options, text="", text_color=Colors.STATUS_WARNING, font=(Fonts.FAMILY, Fonts.SIZE_SMALL), anchor="w" + ) + + self._refresh_widgets() + + def _fast(self) -> bool: + return self._widgets["ltx_fast"].get() == 1 + + def _set_fast(self, fast: bool) -> None: + if fast: + self._widgets["ltx_fast"].select() + else: + self._widgets["ltx_fast"].deselect() + + def _set_ltx_model(self, model: str) -> None: + self._ltx_model.set_value(model) + self._accepted_ltx_model = self._ltx_model.get_value() + + def _show_notice(self, key: str | None) -> None: + if key: + self._notice.configure(text=t(key)) + self._notice.pack(fill="x", padx=12, pady=(0, 8)) + else: + self._notice.pack_forget() + + def _ensure_variant(self, model: str, fast: bool) -> bool: + """Make ``(model, fast)`` usable: download it when the user agrees, else switch to an + installed variant with a notice. False when the user declined or nothing is installed.""" + self._show_notice(None) + state = self._ltx_models.state + if state.installed(model, fast): + return True + if state.downloadable: + return self._ltx_models.ensure(model, fast, on_ready=lambda: None) + others = [(model, not fast)] + [(m, f) for m in LTX_MODELS if m != model for f in (fast, not fast)] + for other_model, other_fast in others: + if state.installed(other_model, other_fast): + self._set_ltx_model(other_model) + self._set_fast(other_fast) + if other_model != model: + self._show_notice("model_ltx_variant_not_installed") + else: + self._show_notice("model_ltx_fast_not_installed" if fast else "model_ltx_quality_not_installed") + return True + return False def _on_card_selected(self, value: str) -> None: + if value == "ltx" and not self._ensure_variant(self._ltx_model.get_value(), self._fast()): + self._show_model("basicvsrpp") + return self._show_model(value) self._on_modified() + def _on_ltx_model_selected(self, model: str) -> None: + if self._ensure_variant(model, self._fast()): + self._accepted_ltx_model = self._ltx_model.get_value() + else: + self._ltx_model.set_value(self._accepted_ltx_model) + self._refresh_widgets() + self._on_modified() + + def _on_fast_toggled(self) -> None: + if not self._ensure_variant(self._ltx_model.get_value(), self._fast()): + self._set_fast(not self._fast()) + self._refresh_widgets() + self._on_modified() + def _show_model(self, value: str) -> None: self._model.set(value) for name, card in self._cards.items(): @@ -152,6 +244,7 @@ def _show_model(self, value: str) -> None: else: self._ltx_options.pack_forget() self._on_model_changed(value) + self._refresh_widgets() def _on_new_seed(self) -> None: self._set_seed(random.randrange(1, 2**31)) @@ -162,47 +255,90 @@ def _set_seed(self, seed: int) -> None: self._widgets["ltx_seed"].insert(0, str(seed)) def unavailable_reason(self) -> str | None: - return ltx_unavailable_reason(installed=self._ltx_installed, nvidia=self._nvidia) + """Why the chosen LTX model cannot restore right now (None when it can).""" + return run_unavailable_reason( + self._ltx_models.state, self._ltx_model.get_value(), self._fast(), nvidia=self._nvidia + ) - def _refresh_ltx_availability(self) -> None: - reason = self.unavailable_reason() - card = self._cards["ltx"] - card.set_enabled(self._enabled and reason is None) - card.description.configure( + def _refresh_widgets(self) -> None: + state = self._ltx_models.state + editable = self._enabled and not self._ltx_models.downloading + reason = card_unavailable_reason(state, nvidia=self._nvidia) + self._cards["basicvsrpp"].set_enabled(editable) + ltx_card = self._cards["ltx"] + ltx_card.set_enabled(editable and reason is None) + ltx_card.description.configure( text=t(reason) if reason else t("model_ltx_description"), text_color=Colors.STATUS_WARNING if reason else Colors.STATUS_PENDING, ) - if reason is not None and self._model.get() == "ltx": - self._show_model("basicvsrpp") + self._ltx_model.configure(state="normal" if editable else "disabled") + self._new_seed_btn.configure(state="normal" if self._enabled else "disabled") + model, fast = self._ltx_model.get_value(), self._fast() + if not state.installed(model, fast) and state.downloadable: + self._install_status.configure(text=t("model_ltx_download_needed", size=state.download_size(model, fast))) + self._install_status.pack(fill="x", padx=12, pady=(0, 4), after=self._model_row) + else: + self._install_status.pack_forget() + + def _usable_model(self) -> str: + """The chosen model, or BasicVSR++ when LTX cannot run; switches to the installed + precision when only that one is there.""" + state = self._ltx_models.state + if self._model.get() != "ltx" or card_unavailable_reason(state, nvidia=self._nvidia) is not None: + return "basicvsrpp" + model, fast = self._ltx_model.get_value(), self._fast() + if not state.usable(model, fast) and state.installed(model, not fast): + self._set_fast(not fast) + return "ltx" + + def _refresh_download(self) -> None: + percent = self._ltx_models.percent + if percent is None: + self._download_row.pack_forget() + else: + self._download_label.configure(text=t("ltx_downloading", percent=percent)) + self._download_bar.set(percent / 100) + self._download_row.pack(fill="x", pady=(Sizing.PADDING_SMALL, 0), after=self._cards_frame) + + def refresh(self) -> None: + """Show the current install state and download, falling back to a usable choice.""" + model = self._usable_model() + if model != self._model.get(): + self._show_model(model) + else: + self._refresh_widgets() + self._refresh_download() def set_gpu_support(self, *, nvidia: bool, blackwell: bool) -> None: self._nvidia = nvidia if blackwell: self._fast_row.pack(fill="x", before=self._large_row) else: - self._widgets["ltx_fast"].deselect() + self._set_fast(False) self._fast_row.pack_forget() - self._refresh_ltx_availability() + self.refresh() def set_enabled(self, enabled: bool) -> None: self._enabled = enabled - self._cards["basicvsrpp"].set_enabled(enabled) - self._new_seed_btn.configure(state="normal" if enabled else "disabled") - self._refresh_ltx_availability() + self._refresh_widgets() def apply(self, preset) -> None: self._set_seed(preset.ltx_seed) - for key, selected in (("ltx_fast", preset.ltx_fast), ("ltx_large_canvas", preset.ltx_large_canvas)): - if selected: - self._widgets[key].select() - else: - self._widgets[key].deselect() - self._show_model(preset.restoration_model) - self._refresh_ltx_availability() + self._set_fast(preset.ltx_fast) + if preset.ltx_large_canvas: + self._widgets["ltx_large_canvas"].select() + else: + self._widgets["ltx_large_canvas"].deselect() + self._set_ltx_model(preset.ltx_model) + self._model.set(preset.restoration_model if preset.restoration_model in _MODELS else "basicvsrpp") + self._show_notice(None) + self._show_model(self._usable_model()) + self._refresh_download() def collect(self) -> dict: return { "restoration_model": self._widgets["restoration_model"].get(), + "ltx_model": self._widgets["ltx_model"].get_value(), "ltx_seed": parse_seed(self._widgets["ltx_seed"].get()), "ltx_fast": self._widgets["ltx_fast"].get() == 1, "ltx_large_canvas": self._widgets["ltx_large_canvas"].get() == 1, diff --git a/jasna/gui/validation.py b/jasna/gui/validation.py index 0122fd39..ac27747e 100644 --- a/jasna/gui/validation.py +++ b/jasna/gui/validation.py @@ -14,6 +14,10 @@ def _uses_ltx(job: JobItem, settings: AppSettings) -> bool: return any(segment.restoration.model == "ltx" for segment in resolve_restorations(segments, default)) +def any_job_uses_ltx(jobs: list[JobItem], settings: AppSettings) -> bool: + return any(_uses_ltx(job, settings) for job in jobs) + + def validate_gui_start(settings: AppSettings, jobs: list[JobItem], *, ltx_available: bool) -> list[str]: errors: list[str] = [] diff --git a/jasna/gui/video_session.py b/jasna/gui/video_session.py index cbdb9575..e3e459e7 100644 --- a/jasna/gui/video_session.py +++ b/jasna/gui/video_session.py @@ -10,7 +10,7 @@ from typing import TYPE_CHECKING, Callable, Mapping from jasna.gui.models import AppSettings -from jasna.session_config import LTX_DEFAULT_MODEL, SessionConfig +from jasna.session_config import SessionConfig from jasna.session_factory import RestorationSession if TYPE_CHECKING: @@ -29,6 +29,7 @@ def video_session_key(settings: AppSettings) -> tuple: settings.denoise_step, settings.secondary_restoration, settings.restoration_model, + settings.ltx_model, settings.ltx_fast, ) if settings.secondary_restoration == "tvai": @@ -77,7 +78,7 @@ def video_session_config( ltx_large_canvas=bool(settings.ltx_large_canvas), ltx_seed=int(settings.ltx_seed), ltx_fast=bool(settings.ltx_fast), - ltx_model=LTX_DEFAULT_MODEL, + ltx_model=settings.ltx_model, compile_basicvsrpp=bool(settings.compile_basicvsrpp), max_clip_size=int(settings.max_clip_size), temporal_overlap=int(settings.temporal_overlap), diff --git a/tests/factories.py b/tests/factories.py index 6e4d631c..9c3d209e 100644 --- a/tests/factories.py +++ b/tests/factories.py @@ -149,3 +149,18 @@ def write_double_adts_aac_source(tmp_path: Path) -> Path: packet.stream = video_out dst.mux(packet) return source + + +def ltx_models(root, directory: Path, *, installed: bool, on_change=lambda: None): + """A real ``LtxModels`` over ``directory``; ``installed`` puts every model file there.""" + from jasna.gui.ltx_models import LtxModels + from jasna.gui.queues import MainThreadCalls + from jasna.ltx.model_files import LTX_MODELS, bundle_names + + directory.mkdir(parents=True, exist_ok=True) + if installed: + for model in LTX_MODELS: + for fast in (False, True): + for name in bundle_names(model, fast=fast): + (directory / name).touch() + return LtxModels(directory, MainThreadCalls(root, 20), on_change) diff --git a/tests/test_gui_download_worker.py b/tests/test_gui_download_worker.py new file mode 100644 index 00000000..d220b350 --- /dev/null +++ b/tests/test_gui_download_worker.py @@ -0,0 +1,29 @@ +from jasna.gui import download_worker + + +def test_worker_reports_deduplicated_percent_and_success() -> None: + def fake_download(progress_callback): + progress_callback(0, None) + progress_callback(50, 100) + progress_callback(50, 100) + progress_callback(100, 100) + + percents: list[int] = [] + done: list[str | None] = [] + thread = download_worker.start_download(fake_download, percents.append, done.append) + thread.join(timeout=5) + + assert not thread.is_alive() + assert percents == [50, 100] + assert done == [None] + + +def test_worker_reports_error_text_on_failure() -> None: + def fake_download(progress_callback): + raise RuntimeError("disk full") + + done: list[str | None] = [] + thread = download_worker.start_download(fake_download, lambda _percent: None, done.append) + thread.join(timeout=5) + + assert done == ["disk full"] diff --git a/tests/test_gui_hidpi_scaling.py b/tests/test_gui_hidpi_scaling.py index ba64dfc3..94371e7d 100644 --- a/tests/test_gui_hidpi_scaling.py +++ b/tests/test_gui_hidpi_scaling.py @@ -16,6 +16,7 @@ from jasna.gui import scaling from jasna.gui.app import JasnaApp +from jasna.gui.queues import MainThreadCalls from jasna.gui.icons import CompactSwitch, NativeIconButton from jasna.gui.models import PresetManager from jasna.gui.settings_sections.widgets import create_slider_value_label @@ -53,6 +54,8 @@ def _stub_main_body_callbacks(root) -> None: root._on_output_changed = lambda *_args: None root._preset_manager = PresetManager() root.TkdndVersion = None + root._main_thread = MainThreadCalls(root, 50) + root._on_ltx_models_changed = lambda: None def _stub_footer_callbacks(root) -> None: diff --git a/tests/test_gui_ltx_models.py b/tests/test_gui_ltx_models.py new file mode 100644 index 00000000..7631887c --- /dev/null +++ b/tests/test_gui_ltx_models.py @@ -0,0 +1,289 @@ +from __future__ import annotations + +from dataclasses import replace +from pathlib import Path +from tkinter import TclError, messagebox +from types import SimpleNamespace + +import customtkinter as ctk +import pytest + +from factories import ltx_models +from jasna import os_utils +from jasna.gui import ltx_models as ltx_models_module +from jasna.gui.locales import t +from jasna.gui.ltx_models import ( + LtxModels, + model_key, + card_unavailable_reason, + read_install_state, + run_unavailable_reason, +) +from jasna.gui.models import AppSettings, JobItem, PresetManager +from jasna.ltx import model_files +from jasna.ltx.model_files import bundle_names +from jasna.segments import SegmentRange, SegmentRestoration + + +class _NowCalls: + def post(self, callback) -> None: + callback() + + +def _install(directory: Path, model: str, *, fast: bool) -> None: + directory.mkdir(parents=True, exist_ok=True) + for name in bundle_names(model, fast=fast): + (directory / name).touch() + + +@pytest.fixture +def downloadable(monkeypatch): + monkeypatch.setattr(model_files, "LTX_DOWNLOAD_URL", "https://example.invalid/ltx") + + +def test_model_key_names_each_model_choice() -> None: + assert model_key("basicvsrpp", "undistilled") == "basicvsrpp" + assert model_key("ltx", "distilled") == "ltx" + assert model_key("ltx", "undistilled") == "ltx_undistilled" + + +def test_availability_without_a_download_location(tmp_path) -> None: + _install(tmp_path, "distilled", fast=False) + state = read_install_state(tmp_path) + + assert not state.downloadable + assert state.installed("distilled", False) and not state.installed("distilled", True) + assert card_unavailable_reason(state, nvidia=True) is None + assert card_unavailable_reason(state, nvidia=False) == "model_ltx_needs_nvidia" + assert card_unavailable_reason(read_install_state(tmp_path / "empty"), nvidia=True) == "model_ltx_not_installed" + assert run_unavailable_reason(state, "undistilled", False, nvidia=True) == "model_ltx_not_installed" + assert run_unavailable_reason(state, "distilled", False, nvidia=None) is None + assert run_unavailable_reason(state, "distilled", True, nvidia=True) == "model_ltx_not_installed" + + +def test_availability_with_a_download_location(tmp_path, downloadable) -> None: + state = read_install_state(tmp_path) + + assert state.downloadable + assert card_unavailable_reason(state, nvidia=True) is None + assert run_unavailable_reason(state, "undistilled", False, nvidia=True) == "model_ltx_not_downloaded" + assert state.download_size("distilled", False) == "14.9 GB" + + +def _answer(monkeypatch, yes: bool) -> list[str]: + asked: list[str] = [] + monkeypatch.setattr(messagebox, "askyesno", lambda _title, message: asked.append(message) or yes) + return asked + + +def _fake_download(monkeypatch, error: str | None) -> list: + downloads = [] + + def start(download, on_percent, on_done): + downloads.append(download) + on_percent(40) + on_done(error) + + monkeypatch.setattr(ltx_models_module, "start_download", start) + return downloads + + +def test_ensure_asks_and_downloads_the_missing_files(tmp_path, monkeypatch, downloadable) -> None: + asked = _answer(monkeypatch, yes=True) + downloads = _fake_download(monkeypatch, error=None) + changes: list[int | None] = [] + models = LtxModels(tmp_path, _NowCalls(), lambda: changes.append(models.percent)) + ready: list[bool] = [] + + assert models.ensure("undistilled", False, on_ready=lambda: ready.append(True)) + + assert asked == [t("ltx_download_confirm", model=t("model_ltx_undistilled"), size="14.9 GB")] + assert len(downloads) == 1 + assert changes == [0, 40, None] + assert ready == [True] + + +def test_ensure_does_nothing_when_the_user_declines(tmp_path, monkeypatch, downloadable) -> None: + _answer(monkeypatch, yes=False) + downloads = _fake_download(monkeypatch, error=None) + models = LtxModels(tmp_path, _NowCalls(), lambda: None) + + assert not models.ensure("distilled", False, on_ready=lambda: pytest.fail("not downloaded")) + assert downloads == [] and not models.downloading + + +def test_ensure_shows_the_download_error(tmp_path, monkeypatch, downloadable) -> None: + _answer(monkeypatch, yes=True) + _fake_download(monkeypatch, error="checksum failed") + errors: list[str] = [] + monkeypatch.setattr(messagebox, "showerror", lambda _title, message: errors.append(message)) + models = LtxModels(tmp_path, _NowCalls(), lambda: None) + + assert models.ensure("distilled", False, on_ready=lambda: pytest.fail("download failed")) + assert errors == [t("ltx_download_failed", message="checksum failed")] + assert not models.downloading + + +def test_ensure_skips_installed_models_and_unpublished_downloads(tmp_path, monkeypatch) -> None: + _answer(monkeypatch, yes=True) + _install(tmp_path, "distilled", fast=False) + models = LtxModels(tmp_path, _NowCalls(), lambda: None) + + assert models.ensure("distilled", False, on_ready=lambda: None) + assert not models.ensure("undistilled", False, on_ready=lambda: None) + + +@pytest.fixture +def tk_root(monkeypatch, tmp_path): + monkeypatch.setattr(os_utils.sys, "platform", "linux", raising=False) + monkeypatch.setenv("XDG_CONFIG_HOME", str(tmp_path / "xdg")) + try: + root = ctk.CTk() + except TclError as exc: + pytest.skip(f"Tk display unavailable: {exc}") + yield root + root.destroy() + + +def _panel(root, models): + from jasna.gui.settings_panel import SettingsPanel + + return SettingsPanel(root, PresetManager(), ltx_models=models) + + +def test_the_model_menu_round_trips_the_undistilled_model(tk_root, tmp_path) -> None: + panel = _panel(tk_root, ltx_models(tk_root, tmp_path / "ltx", installed=True)) + undistilled = replace(AppSettings(), restoration_model="ltx", ltx_model="undistilled", ltx_seed=3, encoder_cq=28) + + for section in panel._sections: + section.apply(undistilled) + + assert panel.get_settings() == undistilled + assert set(panel._model_section._cards) == {"basicvsrpp", "ltx"} + _pick_ltx_model(panel._model_section, "distilled") + assert panel.get_settings().ltx_model == "distilled" + panel._model_section._cards["basicvsrpp"].radio.invoke() + assert (panel.get_settings().restoration_model, panel.get_settings().ltx_model) == ("basicvsrpp", "distilled") + + +def _pick_ltx_model(section, model: str) -> None: + section._ltx_model.set_value(model) + section._on_ltx_model_selected(model) + + +def test_picking_a_missing_model_asks_and_reverts_on_no(tk_root, tmp_path, monkeypatch, downloadable) -> None: + directory = tmp_path / "ltx" + _install(directory, "distilled", fast=False) + asked = _answer(monkeypatch, yes=False) + panel = _panel(tk_root, ltx_models(tk_root, directory, installed=False)) + section = panel._model_section + section._cards["ltx"].radio.invoke() + assert asked == [] and panel.get_settings().restoration_model == "ltx" + + _pick_ltx_model(section, "undistilled") + + assert len(asked) == 1 + assert panel.get_settings().ltx_model == "distilled" + assert not section._install_status.winfo_manager() + + +def test_a_missing_model_shows_its_download_size(tk_root, tmp_path, downloadable) -> None: + panel = _panel(tk_root, ltx_models(tk_root, tmp_path / "ltx", installed=False)) + section = panel._model_section + + section.apply(replace(AppSettings(), restoration_model="ltx", ltx_model="undistilled")) + + assert section._cards["ltx"].radio.cget("state") == "normal" + assert section._install_status.cget("text") == t("model_ltx_download_needed", size="14.9 GB") + assert section.unavailable_reason() == "model_ltx_not_downloaded" + + +def test_picking_ltx_without_its_files_downloads_and_locks_the_choice(tk_root, tmp_path, monkeypatch, downloadable) -> None: + _answer(monkeypatch, yes=True) + started = [] + monkeypatch.setattr(ltx_models_module, "start_download", lambda *args: started.append(args)) + models = ltx_models(tk_root, tmp_path / "ltx", installed=False) + panel = _panel(tk_root, models) + models._on_change = panel.refresh_ltx_models + + panel._model_section._cards["ltx"].radio.invoke() + + assert len(started) == 1 + assert panel.get_settings().restoration_model == "ltx" + assert panel._model_section._download_row.winfo_manager() + assert all(card.radio.cget("state") == "disabled" for card in panel._model_section._cards.values()) + assert panel._model_section._ltx_model.cget("state") == "disabled" + + +def test_ltx_card_only_needs_some_model(tk_root, tmp_path) -> None: + directory = tmp_path / "ltx" + _install(directory, "undistilled", fast=False) + panel = _panel(tk_root, ltx_models(tk_root, directory, installed=False)) + section = panel._model_section + + section._cards["ltx"].radio.invoke() + + assert panel.get_settings().restoration_model == "ltx" + assert panel.get_settings().ltx_model == "undistilled" + assert section._notice.cget("text") == t("model_ltx_variant_not_installed") + + +def test_fast_mode_without_its_model_falls_back_with_a_notice(tk_root, tmp_path) -> None: + directory = tmp_path / "ltx" + _install(directory, "distilled", fast=False) + panel = _panel(tk_root, ltx_models(tk_root, directory, installed=False)) + section = panel._model_section + section.set_gpu_support(nvidia=True, blackwell=True) + section._cards["ltx"].radio.invoke() + + section._widgets["ltx_fast"]._toggle() + + assert panel.get_settings().ltx_fast is False + assert section._notice.cget("text") == t("model_ltx_fast_not_installed") + + +def test_every_model_option_has_a_tooltip(tk_root, tmp_path, monkeypatch) -> None: + from jasna.gui.settings_sections import restoration_model, widgets + + tips: dict[str, list] = {} + for tooltip_class in {restoration_model.Tooltip, widgets.Tooltip}: + original = tooltip_class.__init__ + + def record(self, widget, text, original=original): + tips.setdefault(text, []).append(widget) + original(self, widget, text) + + monkeypatch.setattr(tooltip_class, "__init__", record) + panel = _panel(tk_root, ltx_models(tk_root, tmp_path / "ltx", installed=True)) + section = panel._model_section + + for key, card in section._cards.items(): + assert card.radio in tips[t(f"tip_model_{key}")] + assert section._ltx_model in tips[t("tip_ltx_model")] + assert section._install_status in tips[t("tip_ltx_download")] + assert section._new_seed_btn in tips[t("tip_ltx_new_seed")] + assert section._download_bar in tips[t("tip_ltx_download")] + for key in ("ltx_seed", "ltx_fast", "ltx_large_canvas"): + assert section._widgets[key] in tips[t(f"tip_{key}")] + + +def test_start_downloads_a_missing_ltx_model_first(monkeypatch) -> None: + from jasna.gui.app import JasnaApp + + ensured: list[tuple] = [] + settings = replace(AppSettings(), ltx_model="undistilled") + job = JobItem(Path("a.mp4"), segments=(SegmentRange(1, 2, SegmentRestoration("ltx", 5)),)) + app = SimpleNamespace( + _preview_gpu_busy=False, + _processor=None, + _queue_panel=SimpleNamespace(get_jobs=lambda: [job], get_output_folder=lambda: "", get_output_pattern=lambda: ""), + _settings_panel=SimpleNamespace( + get_settings=lambda: settings, ltx_unavailable_reason=lambda: "model_ltx_not_downloaded" + ), + _ltx_models=SimpleNamespace(ensure=lambda model, fast, on_ready: ensured.append((model, fast)) or True), + ) + app._on_start = lambda: None + + JasnaApp._on_start(app) + + assert ensured == [("undistilled", False)] diff --git a/tests/test_gui_queue_layout.py b/tests/test_gui_queue_layout.py index 1d4dca60..29bc4f1a 100644 --- a/tests/test_gui_queue_layout.py +++ b/tests/test_gui_queue_layout.py @@ -10,6 +10,7 @@ import pytest from jasna.gui.app import JasnaApp +from jasna.gui.queues import MainThreadCalls from jasna.gui.models import JobItem, JobStatus, PresetManager from jasna.gui import queue_panel as queue_panel_module from jasna.gui.queue_panel import QueuePanel @@ -507,6 +508,8 @@ def test_workspace_sash_cursor_stays_on_the_sash() -> None: root._on_output_changed = lambda *_args: None root._preset_manager = PresetManager() root.TkdndVersion = None + root._main_thread = MainThreadCalls(root, 50) + root._on_ltx_models_changed = lambda: None JasnaApp._build_main_body(root) root.update_idletasks() @@ -583,6 +586,8 @@ def test_main_workspace_starts_wider_and_can_resize_queue_panel(hidpi) -> None: root._on_output_changed = lambda *_args: None root._preset_manager = PresetManager() root.TkdndVersion = None + root._main_thread = MainThreadCalls(root, 50) + root._on_ltx_models_changed = lambda: None JasnaApp._build_main_body(root) root.update_idletasks() diff --git a/tests/test_gui_sd15_download_worker.py b/tests/test_gui_sd15_download_worker.py deleted file mode 100644 index f41c3e2d..00000000 --- a/tests/test_gui_sd15_download_worker.py +++ /dev/null @@ -1,41 +0,0 @@ -from pathlib import Path - -from jasna.gui import sd15_download_worker - - -def test_worker_reports_deduplicated_percent_and_success(monkeypatch, tmp_path: Path) -> None: - def fake_download(model_dir, repo_id, progress_callback): - assert model_dir == tmp_path - assert repo_id == "some/repo" - progress_callback(0, None) - progress_callback(50, 100) - progress_callback(50, 100) - progress_callback(100, 100) - - monkeypatch.setattr(sd15_download_worker, "download_sd15_bundle", fake_download) - - percents: list[int] = [] - done: list[str | None] = [] - thread = sd15_download_worker.start_sd15_download( - tmp_path, "some/repo", percents.append, done.append - ) - thread.join(timeout=5) - - assert not thread.is_alive() - assert percents == [50, 100] - assert done == [None] - - -def test_worker_reports_error_text_on_failure(monkeypatch, tmp_path: Path) -> None: - def fake_download(model_dir, repo_id, progress_callback): - raise RuntimeError("disk full") - - monkeypatch.setattr(sd15_download_worker, "download_sd15_bundle", fake_download) - - done: list[str | None] = [] - thread = sd15_download_worker.start_sd15_download( - tmp_path, "some/repo", lambda _percent: None, done.append - ) - thread.join(timeout=5) - - assert done == ["disk full"] diff --git a/tests/test_gui_settings_scroll.py b/tests/test_gui_settings_scroll.py index d2a9d724..53572928 100644 --- a/tests/test_gui_settings_scroll.py +++ b/tests/test_gui_settings_scroll.py @@ -9,10 +9,11 @@ from jasna.gui.models import PresetManager from jasna.gui.settings_panel import SettingsPanel from jasna.gui.theme import Colors +from factories import ltx_models -def _settings_panel(root: ctk.CTk) -> SettingsPanel: - panel = SettingsPanel(root, PresetManager(), ltx_installed=True) +def _settings_panel(root: ctk.CTk, tmp_path) -> SettingsPanel: + panel = SettingsPanel(root, PresetManager(), ltx_models=ltx_models(root, tmp_path, installed=True)) panel.pack(fill="both", expand=True) root.update() return panel @@ -24,14 +25,14 @@ def _collapse_all_sections(panel: SettingsPanel) -> None: section._toggle() -def test_settings_scrollbar_thumb_contrasts_with_panel_without_hover() -> None: +def test_settings_scrollbar_thumb_contrasts_with_panel_without_hover(tmp_path) -> None: try: root = ctk.CTk() except TclError as exc: pytest.skip(f"Tk display unavailable: {exc}") try: - panel = _settings_panel(root) + panel = _settings_panel(root, tmp_path) assert panel._scroll.cget("scrollbar_button_color") == Colors.BORDER_LIGHT assert panel._scroll.cget("scrollbar_button_color") != Colors.BG_PANEL @@ -39,7 +40,7 @@ def test_settings_scrollbar_thumb_contrasts_with_panel_without_hover() -> None: root.destroy() -def test_settings_scrollbar_only_appears_when_sections_overflow() -> None: +def test_settings_scrollbar_only_appears_when_sections_overflow(tmp_path) -> None: try: root = ctk.CTk() except TclError as exc: @@ -47,7 +48,7 @@ def test_settings_scrollbar_only_appears_when_sections_overflow() -> None: try: root.geometry("420x480") - panel = _settings_panel(root) + panel = _settings_panel(root, tmp_path) assert panel._scroll._scrollbar.winfo_ismapped() diff --git a/tests/test_gui_settings_sections.py b/tests/test_gui_settings_sections.py index 79b2860b..acab2d4b 100644 --- a/tests/test_gui_settings_sections.py +++ b/tests/test_gui_settings_sections.py @@ -19,6 +19,7 @@ from jasna.gui.settings_sections.restoration_model import RestorationModelSection from jasna.gui.settings_sections.secondary import SecondarySection from jasna.gui.settings_sections.widgets import ValueOptionMenu +from factories import ltx_models class _FakeValueMenu: @@ -65,6 +66,7 @@ def test_value_option_menu_falls_back_to_first_option_for_unknown_value() -> Non def _fake_section_widgets() -> dict: return { "restoration_model": _FakeWidget("ltx"), + "ltx_model": _FakeValueMenu({"distilled": "Distilled", "undistilled": "Undistilled"}, "undistilled"), "ltx_seed": _FakeWidget(" 123 "), "ltx_fast": _FakeWidget(1), "ltx_large_canvas": _FakeWidget(0), @@ -150,6 +152,7 @@ def test_sections_collect_internal_values_without_translation_lookups() -> None: assert values["fmp4"] is True assert values["sharpen_strength"] == 0.35 assert values["restoration_model"] == "ltx" + assert values["ltx_model"] == "undistilled" assert values["ltx_seed"] == 123 assert values["ltx_fast"] is True assert values["ltx_large_canvas"] is False @@ -192,7 +195,7 @@ def _basic_section_panel(monkeypatch, tmp_path): from jasna.gui.settings_panel import SettingsPanel - panel = SettingsPanel(root, PresetManager(), ltx_installed=True) + panel = SettingsPanel(root, PresetManager(), ltx_models=ltx_models(root, tmp_path / "ltx", installed=True)) try: yield panel, next(section for section in panel._sections if isinstance(section, BasicSection)) finally: @@ -363,7 +366,7 @@ def test_settings_panel_get_settings_is_locale_independent(monkeypatch, tmp_path try: from jasna.gui.settings_panel import SettingsPanel - panel = SettingsPanel(root, PresetManager(), ltx_installed=True) + panel = SettingsPanel(root, PresetManager(), ltx_models=ltx_models(root, tmp_path / "ltx", installed=True)) assert panel.get_settings() == replace(AppSettings(), encoder_cq=28) assert panel._saved_preset_settings == panel.get_settings() finally: @@ -406,15 +409,6 @@ def test_only_queue_wide_post_export_action_stays_editable_while_processing(_bas assert panel._widgets["max_clip_size"].cget("state") == "disabled" -def test_ltx_unavailable_reason() -> None: - from jasna.gui.settings_sections.restoration_model import ltx_unavailable_reason - - assert ltx_unavailable_reason(installed=True, nvidia=None) is None - assert ltx_unavailable_reason(installed=True, nvidia=True) is None - assert ltx_unavailable_reason(installed=True, nvidia=False) == "model_ltx_needs_nvidia" - assert ltx_unavailable_reason(installed=False, nvidia=True) == "model_ltx_not_installed" - - def test_parse_seed_falls_back_to_default() -> None: from jasna.gui.settings_sections.restoration_model import parse_seed from jasna.session_config import LTX_DEFAULT_SEED @@ -435,7 +429,9 @@ def _panel_factory(monkeypatch, tmp_path): from jasna.gui.settings_panel import SettingsPanel try: - yield lambda installed: SettingsPanel(root, PresetManager(), ltx_installed=installed) + yield lambda installed: SettingsPanel( + root, PresetManager(), ltx_models=ltx_models(root, tmp_path / "ltx", installed=installed) + ) finally: root.destroy() diff --git a/tests/test_ltx_seed_preview.py b/tests/test_ltx_seed_preview.py index 9c05e105..4c15d466 100644 --- a/tests/test_ltx_seed_preview.py +++ b/tests/test_ltx_seed_preview.py @@ -117,3 +117,10 @@ def test_seed_renderer_deletes_its_temp_dir_on_close(tmp_path) -> None: renderer.close() assert not temp.exists() + + +def test_prepared_key_never_mixes_ltx_models() -> None: + segment = SegmentRange(1, 2) + settings = replace(AppSettings(), restoration_model="ltx") + + assert prepared_key(segment, replace(settings, ltx_model="undistilled")) != prepared_key(segment, settings) diff --git a/tests/test_restoration_preview.py b/tests/test_restoration_preview.py index e4b180c2..5a14d41f 100644 --- a/tests/test_restoration_preview.py +++ b/tests/test_restoration_preview.py @@ -416,6 +416,7 @@ def test_app_disables_queue_start_while_preview_gpu_is_tearing_down() -> None: app._queue_panel = MagicMock() app._queue_panel.get_jobs.return_value = [object()] app._control_bar = MagicMock() + app._ltx_models = SimpleNamespace(downloading=False) app._update_start_button_state() diff --git a/tests/test_segment_editor.py b/tests/test_segment_editor.py index cbf2128e..5f8df89b 100644 --- a/tests/test_segment_editor.py +++ b/tests/test_segment_editor.py @@ -1,6 +1,7 @@ from __future__ import annotations import queue +from dataclasses import replace import threading from fractions import Fraction from pathlib import Path @@ -954,6 +955,58 @@ def test_try_seed_shows_progress_then_lets_the_user_compare_and_pick_seeds(monke assert _FakeSeedWorker.instances[0].closed + +def test_seed_controls_have_tooltips_and_name_the_job_ltx_model(monkeypatch, tmp_path) -> None: + tips: dict[str, list] = {} + original = segment_editor.Tooltip.__init__ + + def record(self, widget, text): + tips.setdefault(text, []).append(widget) + original(self, widget, text) + + monkeypatch.setattr(segment_editor.Tooltip, "__init__", record) + root = _tk_root() + editor = None + _FakeSeedWorker.instances = [] + monkeypatch.setattr(segment_editor, "LtxSeedPreviewWorker", _FakeSeedWorker) + try: + job = JobItem(tmp_path / "video.mp4", segments=(SegmentRange(1.0, 2.0, SegmentRestoration("ltx", 5)),)) + settings = replace(AppSettings(), ltx_model="undistilled") + editor = _build_editor_with_ui(root, monkeypatch, job=job, settings=settings) + editor._keyframe_index = KeyframeIndex(pts=(0,), time_base=Fraction(1, 1000), start_pts=0, end_pts=60_000) + editor._select_range(0) + editor._preview_source = Image.new("RGB", (160, 90), "red") + root.update() + editor._try_seed_btn.invoke() + worker = _FakeSeedWorker.instances[0] + worker.events.put(segment_editor.SeedReady(5, _seed_frames(tmp_path, 5, "blue"), 1)) + editor._poll_workers() + editor._new_range_seed() + root.update() + + assert editor._model_radios["ltx"].cget("text") == t("model_ltx_undistilled") + expected = { + "tip_segments_restore_with": editor._restore_with_label, + "tip_model_basicvsrpp": editor._model_radios["basicvsrpp"], + "tip_segments_use_for_all": editor._use_for_all_btn, + "tip_ltx_seed": editor._seed_entry, + "tip_ltx_new_seed": editor._new_seed_btn, + "tip_segments_try_seed": editor._try_seed_btn, + "tip_segments_seed_cancel": editor._seed_cancel_btn, + "tip_segments_seed_show": editor._seed_view_toggle, + } + for key, widget in expected.items(): + assert widget in tips[t(key)], key + chips = [w for w in editor._seed_chip_widgets if w.cget("text") == "5"] + use = [w for w in editor._seed_chip_widgets if w.cget("text") == t("segments_seed_use")] + assert chips and chips[0] in tips[t("tip_segments_seed_chip")] + assert use and use[0] in tips[t("tip_segments_seed_use")] + assert editor._ltx_radio_tooltip._text == t("tip_model_ltx_undistilled") + finally: + if editor is not None: + editor._finish_close() + root.destroy() + def test_range_controls_fit_their_content_while_the_list_keeps_its_rows() -> None: assert segment_editor.range_controls_height(500, 200, 156) == 200 assert segment_editor.range_controls_height(320, 200, 156) == 164 diff --git a/tests/test_video_session.py b/tests/test_video_session.py index 8bfbb800..b9d7c293 100644 --- a/tests/test_video_session.py +++ b/tests/test_video_session.py @@ -157,3 +157,13 @@ def test_video_session_key_tracks_model_and_ltx_variant() -> None: def test_video_session_key_tracks_the_ltx_variant_for_ltx_ranges_of_standard_jobs() -> None: assert video_session_key(replace(AppSettings(), ltx_fast=True)) != video_session_key(AppSettings()) + + +def test_video_session_config_and_key_follow_the_ltx_model() -> None: + undistilled = replace(AppSettings(), restoration_model="ltx", ltx_model="undistilled") + + assert _config(undistilled).ltx_model == "undistilled" + assert _config(replace(undistilled, ltx_model="distilled")).ltx_model == "distilled" + assert video_session_key(undistilled) != video_session_key(replace(undistilled, ltx_model="distilled")) + standard = replace(AppSettings(), ltx_model="undistilled") + assert video_session_key(standard) != video_session_key(AppSettings()) From 3798dc566f6c32cd7cbc5c954c973b1012455ec7 Mon Sep 17 00:00:00 2001 From: Kruk2 Date: Tue, 29 Sep 2026 19:35:25 +0200 Subject: [PATCH 45/48] feat(ltx): add --ltx-trial, a speed test without model files or a license LtxModelFiles.placeholder() serves seeded random tensors with the real names, shapes and dtypes (jasna/ltx/trial_shapes.py), so every pass runs as in a real job. The output is not restored. --- jasna/cli_help.py | 4 + jasna/gui/video_session.py | 1 + jasna/ltx/model_files.py | 97 ++++- jasna/ltx/trial_shapes.py | 739 ++++++++++++++++++++++++++++++++++ jasna/main.py | 13 +- jasna/session_config.py | 1 + jasna/session_factory.py | 13 +- scripts/ltx_trial_shapes.py | 93 +++++ tests/factories.py | 1 + tests/test_ltx_trial.py | 96 +++++ tests/test_main_validation.py | 17 + 11 files changed, 1068 insertions(+), 7 deletions(-) create mode 100644 jasna/ltx/trial_shapes.py create mode 100644 scripts/ltx_trial_shapes.py create mode 100644 tests/test_ltx_trial.py diff --git a/jasna/cli_help.py b/jasna/cli_help.py index c8b793bf..d3e04f36 100644 --- a/jasna/cli_help.py +++ b/jasna/cli_help.py @@ -21,6 +21,10 @@ "ltx: about 1.4x faster with slightly less detail. Needs an RTX 50-series (Blackwell) GPU " "and the fast model file. (default: %(default)s)" ), + "ltx_trial": ( + "Speed test without a license: runs LTX with placeholder weights, so the picture will look " + "wrong. Use it to see if your GPU can run LTX and how fast. (default: %(default)s)" + ), "ltx_large_canvas": ( "ltx: restore large mosaics at 768 px instead of 512 px. Sharper, but about 3x slower for " "those mosaics; needs 10 GB of free VRAM (ignored below that). (default: %(default)s)" diff --git a/jasna/gui/video_session.py b/jasna/gui/video_session.py index e3e459e7..5bfcf0ed 100644 --- a/jasna/gui/video_session.py +++ b/jasna/gui/video_session.py @@ -79,6 +79,7 @@ def video_session_config( ltx_seed=int(settings.ltx_seed), ltx_fast=bool(settings.ltx_fast), ltx_model=settings.ltx_model, + ltx_trial=False, compile_basicvsrpp=bool(settings.compile_basicvsrpp), max_clip_size=int(settings.max_clip_size), temporal_overlap=int(settings.temporal_overlap), diff --git a/jasna/ltx/model_files.py b/jasna/ltx/model_files.py index 00572a95..ebb2807e 100644 --- a/jasna/ltx/model_files.py +++ b/jasna/ltx/model_files.py @@ -6,7 +6,9 @@ from __future__ import annotations import hashlib +import json import logging +import math import urllib.request from collections.abc import Callable from dataclasses import dataclass @@ -56,7 +58,84 @@ class DownloadableFile: } -def open_tensors(path: Path): +@dataclass(frozen=True) +class PlaceholderTensors: + """The trial's stand-in for one model file: the real tensor names, shapes and dtypes, + filled with random values. ``part`` is ``int8``/``nvfp4`` (the transformer of + ``model``), ``vae`` or ``vae-decoder``.""" + + part: str + model: LtxModelName + + +# Chosen so a random INT8 / FP4 weight comes out near 1/sqrt(4096), like a trained one. +_PLACEHOLDER_SCALES = {".scales": 2e-4, ".weight_scale_2": 5e-3, ".weight_scale": 1.0} +_TORCH_DTYPES = {"BF16": "bfloat16", "F32": "float32", "I8": "int8", "U8": "uint8", "F8_E4M3": "float8_e4m3fn"} + + +class _PlaceholderHandle: + """``safe_open``-like reader of placeholder tensors. Every tensor of one dtype is a + fresh copy of the same seeded random pool, so it costs what reading the file costs in + RAM, not what generating 13 GB of random numbers would.""" + + def __init__(self, source: PlaceholderTensors) -> None: + from jasna.ltx import trial_shapes as S + + if source.part in S.TRANSFORMER: + parts = S.TRANSFORMER[source.part] + self._specs = dict(parts["top"]) + for index in range(S.BLOCKS): + self._specs.update({f"blocks.{index}.{name}": spec for name, spec in parts["block"].items()}) + steps, stg = S.STEPS[source.model] + self._metadata = { + "sigmas": json.dumps([round(1 - step / steps, 6) for step in range(steps + 1)]), + "stg_scale": "1.0" if stg else "0.0", + } + elif source.part == "vae": + self._specs, self._metadata = dict(S.VAE), {"config": json.dumps(S.VAE_CONFIG)} + else: + self._specs, self._metadata = dict(S.VAE_DECODER), {} + self._pools: dict = {} + + def __enter__(self) -> _PlaceholderHandle: + return self + + def __exit__(self, *exc) -> None: + self._pools.clear() + + def metadata(self) -> dict[str, str]: + return self._metadata + + def keys(self) -> list[str]: + return list(self._specs) + + def get_tensor(self, name: str): + import torch + + dtype, shape = self._specs[name] + scale = next((value for suffix, value in _PLACEHOLDER_SCALES.items() if name.endswith(suffix)), None) + if scale is not None: + return torch.full(shape, scale).to(getattr(torch, _TORCH_DTYPES[dtype])) + return self._pool(dtype)[: math.prod(shape)].view(shape).clone() + + def _pool(self, dtype: str): + import torch + + if dtype not in self._pools: + numel = max(math.prod(shape) for kind, shape in self._specs.values() if kind == dtype) + generator = torch.Generator().manual_seed(0) + torch_dtype = getattr(torch, _TORCH_DTYPES[dtype]) + if dtype in ("I8", "U8"): + low, high = (-127, 128) if dtype == "I8" else (0, 256) + self._pools[dtype] = torch.randint(low, high, (numel,), generator=generator, dtype=torch_dtype) + else: + self._pools[dtype] = (torch.randn(numel, generator=generator) * 0.02).to(torch_dtype) + return self._pools[dtype] + + +def open_tensors(path: Path | PlaceholderTensors): + if isinstance(path, PlaceholderTensors): + return _PlaceholderHandle(path) if not is_frozen() and path.suffix != ".enc" and path.is_file(): from safetensors import safe_open return safe_open(str(path), framework="pt", device="cpu") @@ -76,9 +155,19 @@ def _bundle_paths(directory: Path, model: LtxModelName, *, fast: bool) -> list[P @dataclass(frozen=True) class LtxModelFiles: - transformer: Path - vae: Path - tuned_decoder: Path + transformer: Path | PlaceholderTensors + vae: Path | PlaceholderTensors + tuned_decoder: Path | PlaceholderTensors + + @classmethod + def placeholder(cls, model: LtxModelName, *, fast: bool) -> LtxModelFiles: + """Random weights shaped like ``model``'s files, for the trial: runs every pass at + full cost with no model files and no license, and restores nothing.""" + return cls( + PlaceholderTensors("nvfp4" if fast else "int8", model), + PlaceholderTensors("vae", model), + PlaceholderTensors("vae-decoder", model), + ) @classmethod def from_dir(cls, directory: Path, model: LtxModelName, *, fast: bool) -> LtxModelFiles: diff --git a/jasna/ltx/trial_shapes.py b/jasna/ltx/trial_shapes.py new file mode 100644 index 00000000..64ba5e75 --- /dev/null +++ b/jasna/ltx/trial_shapes.py @@ -0,0 +1,739 @@ +"""Tensor names, shapes and dtypes of the LTX model files, for the trial's placeholder +weights. Generated by scripts/ltx_trial_shapes.py; architecture only, no weights.""" + +BLOCKS = 48 + +# model -> (sampler steps, uses STG) +STEPS = {'distilled': (8, False), 'undistilled': (15, True)} + +VAE_CONFIG = {'vae': {'_class_name': 'CausalDiffusionVAE', 'dims': 3, 'encoder': {'_class_name': 'Encoder', 'dims': 3, 'in_channels': 3, 'out_channels': 128, 'blocks': [['res_x', {'num_layers': 4}], ['compress_space_res', {'multiplier': 2}], ['res_x', {'num_layers': 6}], ['compress_time_res', {'multiplier': 2}], ['res_x', {'num_layers': 4}], ['compress_all_res', {'multiplier': 2}], ['res_x', {'num_layers': 2}], ['compress_all_res', {'multiplier': 1}], ['res_x', {'num_layers': 2}]], 'patch_size': 4, 'latent_log_var': 'constant', 'latent_log_var_value': -7.824046010856292, 'norm_layer': 'pixel_norm', 'base_channels': 128, 'spatial_padding_mode': 'zeros'}, 'decoder': {'_class_name': 'NADiffusionDecoder', 'in_channels': 128, 'out_channels': 3, 'patch_size': 4, 'head_dim': 64, 'stage_channels': [2048, 1024, 512, 512, 256], 'stage_depths': [4, 6, 4, 2, 8], 'stage_kernels': [[3, 7, 7], [3, 7, 7], [3, 5, 5], [3, 5, 5], [11, 11, 11]], 'upsamples': [[[1, 2, 2], 2], [[2, 1, 1], 2], [[2, 2, 2], 1], [[2, 2, 2], 2]], 'spatial_padding_mode': 'zeros', 'resampler_kind': 'linear', 'stage5_kernel': [11, 11, 11], 'timestep_scale_multiplier': 1000.0, 'default_num_inference_steps': 1}, 'model_output_type': 'x0', 'spatial_padding_mode': 'zeros'}} + +TRANSFORMER = { + 'int8': { + 'top': { + 'adaln_single.emb.timestep_embedder.linear_1.bias': ('BF16', (4096,)), + 'adaln_single.emb.timestep_embedder.linear_1.weight': ('BF16', (4096, 256)), + 'adaln_single.emb.timestep_embedder.linear_2.bias': ('BF16', (4096,)), + 'adaln_single.emb.timestep_embedder.linear_2.weight': ('BF16', (4096, 4096)), + 'adaln_single.linear.bias': ('BF16', (36864,)), + 'adaln_single.linear.weight': ('BF16', (36864, 4096)), + 'keyframes_abs_pos_embedding': ('BF16', (1, 4096)), + 'patchify_proj.bias': ('BF16', (4096,)), + 'patchify_proj.weight': ('BF16', (4096, 128)), + 'proj_out.bias': ('BF16', (128,)), + 'proj_out.weight': ('BF16', (128, 4096)), + 'prompt_adaln_single.emb.timestep_embedder.linear_1.bias': ('BF16', (4096,)), + 'prompt_adaln_single.emb.timestep_embedder.linear_1.weight': ('BF16', (4096, 256)), + 'prompt_adaln_single.emb.timestep_embedder.linear_2.bias': ('BF16', (4096,)), + 'prompt_adaln_single.emb.timestep_embedder.linear_2.weight': ('BF16', (4096, 4096)), + 'prompt_adaln_single.linear.bias': ('BF16', (8192,)), + 'prompt_adaln_single.linear.weight': ('BF16', (8192, 4096)), + 'prompt_context': ('BF16', (1, 1024, 4096)), + 'scale_shift_table': ('BF16', (2, 4096)), + }, + 'block': { + 'attn1.k_norm.weight': ('BF16', (4096,)), + 'attn1.q_norm.weight': ('BF16', (4096,)), + 'attn1.to_gate_logits.bias': ('BF16', (32,)), + 'attn1.to_gate_logits.qdata': ('I8', (32, 4096)), + 'attn1.to_gate_logits.scales': ('F32', (32,)), + 'attn1.to_k.bias': ('BF16', (4096,)), + 'attn1.to_k.qdata': ('I8', (4096, 4096)), + 'attn1.to_k.scales': ('F32', (4096,)), + 'attn1.to_out.0.bias': ('BF16', (4096,)), + 'attn1.to_out.0.qdata': ('I8', (4096, 4096)), + 'attn1.to_out.0.scales': ('F32', (4096,)), + 'attn1.to_q.bias': ('BF16', (4096,)), + 'attn1.to_q.qdata': ('I8', (4096, 4096)), + 'attn1.to_q.scales': ('F32', (4096,)), + 'attn1.to_v.bias': ('BF16', (4096,)), + 'attn1.to_v.qdata': ('I8', (4096, 4096)), + 'attn1.to_v.scales': ('F32', (4096,)), + 'attn2.k_norm.weight': ('BF16', (4096,)), + 'attn2.q_norm.weight': ('BF16', (4096,)), + 'attn2.to_gate_logits.bias': ('BF16', (32,)), + 'attn2.to_gate_logits.qdata': ('I8', (32, 4096)), + 'attn2.to_gate_logits.scales': ('F32', (32,)), + 'attn2.to_k.bias': ('BF16', (4096,)), + 'attn2.to_k.qdata': ('I8', (4096, 4096)), + 'attn2.to_k.scales': ('F32', (4096,)), + 'attn2.to_out.0.bias': ('BF16', (4096,)), + 'attn2.to_out.0.qdata': ('I8', (4096, 4096)), + 'attn2.to_out.0.scales': ('F32', (4096,)), + 'attn2.to_q.bias': ('BF16', (4096,)), + 'attn2.to_q.qdata': ('I8', (4096, 4096)), + 'attn2.to_q.scales': ('F32', (4096,)), + 'attn2.to_v.bias': ('BF16', (4096,)), + 'attn2.to_v.qdata': ('I8', (4096, 4096)), + 'attn2.to_v.scales': ('F32', (4096,)), + 'ff.net.0.proj.qdata': ('I8', (16384, 4096)), + 'ff.net.0.proj.scales': ('F32', (16384,)), + 'ff.net.2.qdata': ('I8', (4096, 16384)), + 'ff.net.2.scales': ('F32', (4096,)), + 'prompt_scale_shift_table': ('BF16', (2, 4096)), + 'scale_shift_table': ('BF16', (9, 4096)), + }, + }, + 'nvfp4': { + 'top': { + 'adaln_single.emb.timestep_embedder.linear_1.bias': ('BF16', (4096,)), + 'adaln_single.emb.timestep_embedder.linear_1.weight': ('BF16', (4096, 256)), + 'adaln_single.emb.timestep_embedder.linear_2.bias': ('BF16', (4096,)), + 'adaln_single.emb.timestep_embedder.linear_2.weight': ('BF16', (4096, 4096)), + 'adaln_single.linear.bias': ('BF16', (36864,)), + 'adaln_single.linear.weight': ('BF16', (36864, 4096)), + 'keyframes_abs_pos_embedding': ('BF16', (1, 4096)), + 'patchify_proj.bias': ('BF16', (4096,)), + 'patchify_proj.weight': ('BF16', (4096, 128)), + 'proj_out.bias': ('BF16', (128,)), + 'proj_out.weight': ('BF16', (128, 4096)), + 'prompt_adaln_single.emb.timestep_embedder.linear_1.bias': ('BF16', (4096,)), + 'prompt_adaln_single.emb.timestep_embedder.linear_1.weight': ('BF16', (4096, 256)), + 'prompt_adaln_single.emb.timestep_embedder.linear_2.bias': ('BF16', (4096,)), + 'prompt_adaln_single.emb.timestep_embedder.linear_2.weight': ('BF16', (4096, 4096)), + 'prompt_adaln_single.linear.bias': ('BF16', (8192,)), + 'prompt_adaln_single.linear.weight': ('BF16', (8192, 4096)), + 'prompt_context': ('BF16', (1, 1024, 4096)), + 'scale_shift_table': ('BF16', (2, 4096)), + }, + 'block': { + 'attn1.k_norm.weight': ('BF16', (4096,)), + 'attn1.q_norm.weight': ('BF16', (4096,)), + 'attn1.to_gate_logits.bias': ('BF16', (32,)), + 'attn1.to_gate_logits.weight': ('BF16', (32, 4096)), + 'attn1.to_k.bias': ('BF16', (4096,)), + 'attn1.to_k.weight': ('U8', (4096, 2048)), + 'attn1.to_k.weight_scale': ('F8_E4M3', (4096, 256)), + 'attn1.to_k.weight_scale_2': ('F32', ()), + 'attn1.to_out.0.bias': ('BF16', (4096,)), + 'attn1.to_out.0.weight': ('U8', (4096, 2048)), + 'attn1.to_out.0.weight_scale': ('F8_E4M3', (4096, 256)), + 'attn1.to_out.0.weight_scale_2': ('F32', ()), + 'attn1.to_q.bias': ('BF16', (4096,)), + 'attn1.to_q.weight': ('U8', (4096, 2048)), + 'attn1.to_q.weight_scale': ('F8_E4M3', (4096, 256)), + 'attn1.to_q.weight_scale_2': ('F32', ()), + 'attn1.to_v.bias': ('BF16', (4096,)), + 'attn1.to_v.weight': ('U8', (4096, 2048)), + 'attn1.to_v.weight_scale': ('F8_E4M3', (4096, 256)), + 'attn1.to_v.weight_scale_2': ('F32', ()), + 'attn2.k_norm.weight': ('BF16', (4096,)), + 'attn2.q_norm.weight': ('BF16', (4096,)), + 'attn2.to_gate_logits.bias': ('BF16', (32,)), + 'attn2.to_gate_logits.weight': ('BF16', (32, 4096)), + 'attn2.to_k.bias': ('BF16', (4096,)), + 'attn2.to_k.weight': ('U8', (4096, 2048)), + 'attn2.to_k.weight_scale': ('F8_E4M3', (4096, 256)), + 'attn2.to_k.weight_scale_2': ('F32', ()), + 'attn2.to_out.0.bias': ('BF16', (4096,)), + 'attn2.to_out.0.weight': ('U8', (4096, 2048)), + 'attn2.to_out.0.weight_scale': ('F8_E4M3', (4096, 256)), + 'attn2.to_out.0.weight_scale_2': ('F32', ()), + 'attn2.to_q.bias': ('BF16', (4096,)), + 'attn2.to_q.weight': ('U8', (4096, 2048)), + 'attn2.to_q.weight_scale': ('F8_E4M3', (4096, 256)), + 'attn2.to_q.weight_scale_2': ('F32', ()), + 'attn2.to_v.bias': ('BF16', (4096,)), + 'attn2.to_v.weight': ('U8', (4096, 2048)), + 'attn2.to_v.weight_scale': ('F8_E4M3', (4096, 256)), + 'attn2.to_v.weight_scale_2': ('F32', ()), + 'ff.net.0.proj.weight': ('U8', (16384, 2048)), + 'ff.net.0.proj.weight_scale': ('F8_E4M3', (16384, 256)), + 'ff.net.0.proj.weight_scale_2': ('F32', ()), + 'ff.net.2.weight': ('U8', (4096, 8192)), + 'ff.net.2.weight_scale': ('F8_E4M3', (4096, 1024)), + 'ff.net.2.weight_scale_2': ('F32', ()), + 'prompt_scale_shift_table': ('BF16', (2, 4096)), + 'scale_shift_table': ('BF16', (9, 4096)), + }, + }, +} + +VAE = { + 'decoder.conv_in.bias': ('BF16', (2048,)), + 'decoder.conv_in.weight': ('BF16', (2048, 128)), + 'decoder.conv_in_x_t.bias': ('BF16', (256,)), + 'decoder.conv_in_x_t.weight': ('BF16', (256, 48)), + 'decoder.conv_out.bias': ('BF16', (48,)), + 'decoder.conv_out.weight': ('BF16', (48, 256)), + 'decoder.det_stages.0.0.attn.k_norm.weight': ('BF16', (64,)), + 'decoder.det_stages.0.0.attn.proj.bias': ('BF16', (2048,)), + 'decoder.det_stages.0.0.attn.proj.weight': ('BF16', (2048, 2048)), + 'decoder.det_stages.0.0.attn.q_norm.weight': ('BF16', (64,)), + 'decoder.det_stages.0.0.attn.qkv.bias': ('BF16', (6144,)), + 'decoder.det_stages.0.0.attn.qkv.weight': ('BF16', (6144, 2048)), + 'decoder.det_stages.0.0.mlp.w_down.weight': ('BF16', (2048, 8192)), + 'decoder.det_stages.0.0.mlp.w_gate.weight': ('BF16', (8192, 2048)), + 'decoder.det_stages.0.0.mlp.w_up.weight': ('BF16', (8192, 2048)), + 'decoder.det_stages.0.0.norm1.weight': ('BF16', (2048,)), + 'decoder.det_stages.0.0.norm2.weight': ('BF16', (2048,)), + 'decoder.det_stages.0.1.attn.k_norm.weight': ('BF16', (64,)), + 'decoder.det_stages.0.1.attn.proj.bias': ('BF16', (2048,)), + 'decoder.det_stages.0.1.attn.proj.weight': ('BF16', (2048, 2048)), + 'decoder.det_stages.0.1.attn.q_norm.weight': ('BF16', (64,)), + 'decoder.det_stages.0.1.attn.qkv.bias': ('BF16', (6144,)), + 'decoder.det_stages.0.1.attn.qkv.weight': ('BF16', (6144, 2048)), + 'decoder.det_stages.0.1.mlp.w_down.weight': ('BF16', (2048, 8192)), + 'decoder.det_stages.0.1.mlp.w_gate.weight': ('BF16', (8192, 2048)), + 'decoder.det_stages.0.1.mlp.w_up.weight': ('BF16', (8192, 2048)), + 'decoder.det_stages.0.1.norm1.weight': ('BF16', (2048,)), + 'decoder.det_stages.0.1.norm2.weight': ('BF16', (2048,)), + 'decoder.det_stages.0.2.attn.k_norm.weight': ('BF16', (64,)), + 'decoder.det_stages.0.2.attn.proj.bias': ('BF16', (2048,)), + 'decoder.det_stages.0.2.attn.proj.weight': ('BF16', (2048, 2048)), + 'decoder.det_stages.0.2.attn.q_norm.weight': ('BF16', (64,)), + 'decoder.det_stages.0.2.attn.qkv.bias': ('BF16', (6144,)), + 'decoder.det_stages.0.2.attn.qkv.weight': ('BF16', (6144, 2048)), + 'decoder.det_stages.0.2.mlp.w_down.weight': ('BF16', (2048, 8192)), + 'decoder.det_stages.0.2.mlp.w_gate.weight': ('BF16', (8192, 2048)), + 'decoder.det_stages.0.2.mlp.w_up.weight': ('BF16', (8192, 2048)), + 'decoder.det_stages.0.2.norm1.weight': ('BF16', (2048,)), + 'decoder.det_stages.0.2.norm2.weight': ('BF16', (2048,)), + 'decoder.det_stages.0.3.attn.k_norm.weight': ('BF16', (64,)), + 'decoder.det_stages.0.3.attn.proj.bias': ('BF16', (2048,)), + 'decoder.det_stages.0.3.attn.proj.weight': ('BF16', (2048, 2048)), + 'decoder.det_stages.0.3.attn.q_norm.weight': ('BF16', (64,)), + 'decoder.det_stages.0.3.attn.qkv.bias': ('BF16', (6144,)), + 'decoder.det_stages.0.3.attn.qkv.weight': ('BF16', (6144, 2048)), + 'decoder.det_stages.0.3.mlp.w_down.weight': ('BF16', (2048, 8192)), + 'decoder.det_stages.0.3.mlp.w_gate.weight': ('BF16', (8192, 2048)), + 'decoder.det_stages.0.3.mlp.w_up.weight': ('BF16', (8192, 2048)), + 'decoder.det_stages.0.3.norm1.weight': ('BF16', (2048,)), + 'decoder.det_stages.0.3.norm2.weight': ('BF16', (2048,)), + 'decoder.det_stages.1.0.attn.k_norm.weight': ('BF16', (64,)), + 'decoder.det_stages.1.0.attn.proj.bias': ('BF16', (1024,)), + 'decoder.det_stages.1.0.attn.proj.weight': ('BF16', (1024, 1024)), + 'decoder.det_stages.1.0.attn.q_norm.weight': ('BF16', (64,)), + 'decoder.det_stages.1.0.attn.qkv.bias': ('BF16', (3072,)), + 'decoder.det_stages.1.0.attn.qkv.weight': ('BF16', (3072, 1024)), + 'decoder.det_stages.1.0.mlp.w_down.weight': ('BF16', (1024, 4096)), + 'decoder.det_stages.1.0.mlp.w_gate.weight': ('BF16', (4096, 1024)), + 'decoder.det_stages.1.0.mlp.w_up.weight': ('BF16', (4096, 1024)), + 'decoder.det_stages.1.0.norm1.weight': ('BF16', (1024,)), + 'decoder.det_stages.1.0.norm2.weight': ('BF16', (1024,)), + 'decoder.det_stages.1.1.attn.k_norm.weight': ('BF16', (64,)), + 'decoder.det_stages.1.1.attn.proj.bias': ('BF16', (1024,)), + 'decoder.det_stages.1.1.attn.proj.weight': ('BF16', (1024, 1024)), + 'decoder.det_stages.1.1.attn.q_norm.weight': ('BF16', (64,)), + 'decoder.det_stages.1.1.attn.qkv.bias': ('BF16', (3072,)), + 'decoder.det_stages.1.1.attn.qkv.weight': ('BF16', (3072, 1024)), + 'decoder.det_stages.1.1.mlp.w_down.weight': ('BF16', (1024, 4096)), + 'decoder.det_stages.1.1.mlp.w_gate.weight': ('BF16', (4096, 1024)), + 'decoder.det_stages.1.1.mlp.w_up.weight': ('BF16', (4096, 1024)), + 'decoder.det_stages.1.1.norm1.weight': ('BF16', (1024,)), + 'decoder.det_stages.1.1.norm2.weight': ('BF16', (1024,)), + 'decoder.det_stages.1.2.attn.k_norm.weight': ('BF16', (64,)), + 'decoder.det_stages.1.2.attn.proj.bias': ('BF16', (1024,)), + 'decoder.det_stages.1.2.attn.proj.weight': ('BF16', (1024, 1024)), + 'decoder.det_stages.1.2.attn.q_norm.weight': ('BF16', (64,)), + 'decoder.det_stages.1.2.attn.qkv.bias': ('BF16', (3072,)), + 'decoder.det_stages.1.2.attn.qkv.weight': ('BF16', (3072, 1024)), + 'decoder.det_stages.1.2.mlp.w_down.weight': ('BF16', (1024, 4096)), + 'decoder.det_stages.1.2.mlp.w_gate.weight': ('BF16', (4096, 1024)), + 'decoder.det_stages.1.2.mlp.w_up.weight': ('BF16', (4096, 1024)), + 'decoder.det_stages.1.2.norm1.weight': ('BF16', (1024,)), + 'decoder.det_stages.1.2.norm2.weight': ('BF16', (1024,)), + 'decoder.det_stages.1.3.attn.k_norm.weight': ('BF16', (64,)), + 'decoder.det_stages.1.3.attn.proj.bias': ('BF16', (1024,)), + 'decoder.det_stages.1.3.attn.proj.weight': ('BF16', (1024, 1024)), + 'decoder.det_stages.1.3.attn.q_norm.weight': ('BF16', (64,)), + 'decoder.det_stages.1.3.attn.qkv.bias': ('BF16', (3072,)), + 'decoder.det_stages.1.3.attn.qkv.weight': ('BF16', (3072, 1024)), + 'decoder.det_stages.1.3.mlp.w_down.weight': ('BF16', (1024, 4096)), + 'decoder.det_stages.1.3.mlp.w_gate.weight': ('BF16', (4096, 1024)), + 'decoder.det_stages.1.3.mlp.w_up.weight': ('BF16', (4096, 1024)), + 'decoder.det_stages.1.3.norm1.weight': ('BF16', (1024,)), + 'decoder.det_stages.1.3.norm2.weight': ('BF16', (1024,)), + 'decoder.det_stages.1.4.attn.k_norm.weight': ('BF16', (64,)), + 'decoder.det_stages.1.4.attn.proj.bias': ('BF16', (1024,)), + 'decoder.det_stages.1.4.attn.proj.weight': ('BF16', (1024, 1024)), + 'decoder.det_stages.1.4.attn.q_norm.weight': ('BF16', (64,)), + 'decoder.det_stages.1.4.attn.qkv.bias': ('BF16', (3072,)), + 'decoder.det_stages.1.4.attn.qkv.weight': ('BF16', (3072, 1024)), + 'decoder.det_stages.1.4.mlp.w_down.weight': ('BF16', (1024, 4096)), + 'decoder.det_stages.1.4.mlp.w_gate.weight': ('BF16', (4096, 1024)), + 'decoder.det_stages.1.4.mlp.w_up.weight': ('BF16', (4096, 1024)), + 'decoder.det_stages.1.4.norm1.weight': ('BF16', (1024,)), + 'decoder.det_stages.1.4.norm2.weight': ('BF16', (1024,)), + 'decoder.det_stages.1.5.attn.k_norm.weight': ('BF16', (64,)), + 'decoder.det_stages.1.5.attn.proj.bias': ('BF16', (1024,)), + 'decoder.det_stages.1.5.attn.proj.weight': ('BF16', (1024, 1024)), + 'decoder.det_stages.1.5.attn.q_norm.weight': ('BF16', (64,)), + 'decoder.det_stages.1.5.attn.qkv.bias': ('BF16', (3072,)), + 'decoder.det_stages.1.5.attn.qkv.weight': ('BF16', (3072, 1024)), + 'decoder.det_stages.1.5.mlp.w_down.weight': ('BF16', (1024, 4096)), + 'decoder.det_stages.1.5.mlp.w_gate.weight': ('BF16', (4096, 1024)), + 'decoder.det_stages.1.5.mlp.w_up.weight': ('BF16', (4096, 1024)), + 'decoder.det_stages.1.5.norm1.weight': ('BF16', (1024,)), + 'decoder.det_stages.1.5.norm2.weight': ('BF16', (1024,)), + 'decoder.det_stages.2.0.attn.k_norm.weight': ('BF16', (64,)), + 'decoder.det_stages.2.0.attn.proj.bias': ('BF16', (512,)), + 'decoder.det_stages.2.0.attn.proj.weight': ('BF16', (512, 512)), + 'decoder.det_stages.2.0.attn.q_norm.weight': ('BF16', (64,)), + 'decoder.det_stages.2.0.attn.qkv.bias': ('BF16', (1536,)), + 'decoder.det_stages.2.0.attn.qkv.weight': ('BF16', (1536, 512)), + 'decoder.det_stages.2.0.mlp.w_down.weight': ('BF16', (512, 2048)), + 'decoder.det_stages.2.0.mlp.w_gate.weight': ('BF16', (2048, 512)), + 'decoder.det_stages.2.0.mlp.w_up.weight': ('BF16', (2048, 512)), + 'decoder.det_stages.2.0.norm1.weight': ('BF16', (512,)), + 'decoder.det_stages.2.0.norm2.weight': ('BF16', (512,)), + 'decoder.det_stages.2.1.attn.k_norm.weight': ('BF16', (64,)), + 'decoder.det_stages.2.1.attn.proj.bias': ('BF16', (512,)), + 'decoder.det_stages.2.1.attn.proj.weight': ('BF16', (512, 512)), + 'decoder.det_stages.2.1.attn.q_norm.weight': ('BF16', (64,)), + 'decoder.det_stages.2.1.attn.qkv.bias': ('BF16', (1536,)), + 'decoder.det_stages.2.1.attn.qkv.weight': ('BF16', (1536, 512)), + 'decoder.det_stages.2.1.mlp.w_down.weight': ('BF16', (512, 2048)), + 'decoder.det_stages.2.1.mlp.w_gate.weight': ('BF16', (2048, 512)), + 'decoder.det_stages.2.1.mlp.w_up.weight': ('BF16', (2048, 512)), + 'decoder.det_stages.2.1.norm1.weight': ('BF16', (512,)), + 'decoder.det_stages.2.1.norm2.weight': ('BF16', (512,)), + 'decoder.det_stages.2.2.attn.k_norm.weight': ('BF16', (64,)), + 'decoder.det_stages.2.2.attn.proj.bias': ('BF16', (512,)), + 'decoder.det_stages.2.2.attn.proj.weight': ('BF16', (512, 512)), + 'decoder.det_stages.2.2.attn.q_norm.weight': ('BF16', (64,)), + 'decoder.det_stages.2.2.attn.qkv.bias': ('BF16', (1536,)), + 'decoder.det_stages.2.2.attn.qkv.weight': ('BF16', (1536, 512)), + 'decoder.det_stages.2.2.mlp.w_down.weight': ('BF16', (512, 2048)), + 'decoder.det_stages.2.2.mlp.w_gate.weight': ('BF16', (2048, 512)), + 'decoder.det_stages.2.2.mlp.w_up.weight': ('BF16', (2048, 512)), + 'decoder.det_stages.2.2.norm1.weight': ('BF16', (512,)), + 'decoder.det_stages.2.2.norm2.weight': ('BF16', (512,)), + 'decoder.det_stages.2.3.attn.k_norm.weight': ('BF16', (64,)), + 'decoder.det_stages.2.3.attn.proj.bias': ('BF16', (512,)), + 'decoder.det_stages.2.3.attn.proj.weight': ('BF16', (512, 512)), + 'decoder.det_stages.2.3.attn.q_norm.weight': ('BF16', (64,)), + 'decoder.det_stages.2.3.attn.qkv.bias': ('BF16', (1536,)), + 'decoder.det_stages.2.3.attn.qkv.weight': ('BF16', (1536, 512)), + 'decoder.det_stages.2.3.mlp.w_down.weight': ('BF16', (512, 2048)), + 'decoder.det_stages.2.3.mlp.w_gate.weight': ('BF16', (2048, 512)), + 'decoder.det_stages.2.3.mlp.w_up.weight': ('BF16', (2048, 512)), + 'decoder.det_stages.2.3.norm1.weight': ('BF16', (512,)), + 'decoder.det_stages.2.3.norm2.weight': ('BF16', (512,)), + 'decoder.det_stages.3.0.attn.k_norm.weight': ('BF16', (64,)), + 'decoder.det_stages.3.0.attn.proj.bias': ('BF16', (512,)), + 'decoder.det_stages.3.0.attn.proj.weight': ('BF16', (512, 512)), + 'decoder.det_stages.3.0.attn.q_norm.weight': ('BF16', (64,)), + 'decoder.det_stages.3.0.attn.qkv.bias': ('BF16', (1536,)), + 'decoder.det_stages.3.0.attn.qkv.weight': ('BF16', (1536, 512)), + 'decoder.det_stages.3.0.mlp.w_down.weight': ('BF16', (512, 2048)), + 'decoder.det_stages.3.0.mlp.w_gate.weight': ('BF16', (2048, 512)), + 'decoder.det_stages.3.0.mlp.w_up.weight': ('BF16', (2048, 512)), + 'decoder.det_stages.3.0.norm1.weight': ('BF16', (512,)), + 'decoder.det_stages.3.0.norm2.weight': ('BF16', (512,)), + 'decoder.det_stages.3.1.attn.k_norm.weight': ('BF16', (64,)), + 'decoder.det_stages.3.1.attn.proj.bias': ('BF16', (512,)), + 'decoder.det_stages.3.1.attn.proj.weight': ('BF16', (512, 512)), + 'decoder.det_stages.3.1.attn.q_norm.weight': ('BF16', (64,)), + 'decoder.det_stages.3.1.attn.qkv.bias': ('BF16', (1536,)), + 'decoder.det_stages.3.1.attn.qkv.weight': ('BF16', (1536, 512)), + 'decoder.det_stages.3.1.mlp.w_down.weight': ('BF16', (512, 2048)), + 'decoder.det_stages.3.1.mlp.w_gate.weight': ('BF16', (2048, 512)), + 'decoder.det_stages.3.1.mlp.w_up.weight': ('BF16', (2048, 512)), + 'decoder.det_stages.3.1.norm1.weight': ('BF16', (512,)), + 'decoder.det_stages.3.1.norm2.weight': ('BF16', (512,)), + 'decoder.diff_blocks.0.attn.k_norm.weight': ('BF16', (64,)), + 'decoder.diff_blocks.0.attn.proj.bias': ('BF16', (256,)), + 'decoder.diff_blocks.0.attn.proj.weight': ('BF16', (256, 256)), + 'decoder.diff_blocks.0.attn.q_norm.weight': ('BF16', (64,)), + 'decoder.diff_blocks.0.attn.qkv.bias': ('BF16', (768,)), + 'decoder.diff_blocks.0.attn.qkv.weight': ('BF16', (768, 256)), + 'decoder.diff_blocks.0.context_proj.bias': ('BF16', (256,)), + 'decoder.diff_blocks.0.context_proj.weight': ('BF16', (256, 256)), + 'decoder.diff_blocks.0.mlp.w_down.weight': ('BF16', (256, 1024)), + 'decoder.diff_blocks.0.mlp.w_gate.weight': ('BF16', (1024, 256)), + 'decoder.diff_blocks.0.mlp.w_up.weight': ('BF16', (1024, 256)), + 'decoder.diff_blocks.0.norm1.weight': ('BF16', (256,)), + 'decoder.diff_blocks.0.norm2.weight': ('BF16', (256,)), + 'decoder.diff_blocks.0.scale_shift_table': ('BF16', (7, 256)), + 'decoder.diff_blocks.1.attn.k_norm.weight': ('BF16', (64,)), + 'decoder.diff_blocks.1.attn.proj.bias': ('BF16', (256,)), + 'decoder.diff_blocks.1.attn.proj.weight': ('BF16', (256, 256)), + 'decoder.diff_blocks.1.attn.q_norm.weight': ('BF16', (64,)), + 'decoder.diff_blocks.1.attn.qkv.bias': ('BF16', (768,)), + 'decoder.diff_blocks.1.attn.qkv.weight': ('BF16', (768, 256)), + 'decoder.diff_blocks.1.context_proj.bias': ('BF16', (256,)), + 'decoder.diff_blocks.1.context_proj.weight': ('BF16', (256, 256)), + 'decoder.diff_blocks.1.mlp.w_down.weight': ('BF16', (256, 1024)), + 'decoder.diff_blocks.1.mlp.w_gate.weight': ('BF16', (1024, 256)), + 'decoder.diff_blocks.1.mlp.w_up.weight': ('BF16', (1024, 256)), + 'decoder.diff_blocks.1.norm1.weight': ('BF16', (256,)), + 'decoder.diff_blocks.1.norm2.weight': ('BF16', (256,)), + 'decoder.diff_blocks.1.scale_shift_table': ('BF16', (7, 256)), + 'decoder.diff_blocks.2.attn.k_norm.weight': ('BF16', (64,)), + 'decoder.diff_blocks.2.attn.proj.bias': ('BF16', (256,)), + 'decoder.diff_blocks.2.attn.proj.weight': ('BF16', (256, 256)), + 'decoder.diff_blocks.2.attn.q_norm.weight': ('BF16', (64,)), + 'decoder.diff_blocks.2.attn.qkv.bias': ('BF16', (768,)), + 'decoder.diff_blocks.2.attn.qkv.weight': ('BF16', (768, 256)), + 'decoder.diff_blocks.2.context_proj.bias': ('BF16', (256,)), + 'decoder.diff_blocks.2.context_proj.weight': ('BF16', (256, 256)), + 'decoder.diff_blocks.2.mlp.w_down.weight': ('BF16', (256, 1024)), + 'decoder.diff_blocks.2.mlp.w_gate.weight': ('BF16', (1024, 256)), + 'decoder.diff_blocks.2.mlp.w_up.weight': ('BF16', (1024, 256)), + 'decoder.diff_blocks.2.norm1.weight': ('BF16', (256,)), + 'decoder.diff_blocks.2.norm2.weight': ('BF16', (256,)), + 'decoder.diff_blocks.2.scale_shift_table': ('BF16', (7, 256)), + 'decoder.diff_blocks.3.attn.k_norm.weight': ('BF16', (64,)), + 'decoder.diff_blocks.3.attn.proj.bias': ('BF16', (256,)), + 'decoder.diff_blocks.3.attn.proj.weight': ('BF16', (256, 256)), + 'decoder.diff_blocks.3.attn.q_norm.weight': ('BF16', (64,)), + 'decoder.diff_blocks.3.attn.qkv.bias': ('BF16', (768,)), + 'decoder.diff_blocks.3.attn.qkv.weight': ('BF16', (768, 256)), + 'decoder.diff_blocks.3.context_proj.bias': ('BF16', (256,)), + 'decoder.diff_blocks.3.context_proj.weight': ('BF16', (256, 256)), + 'decoder.diff_blocks.3.mlp.w_down.weight': ('BF16', (256, 1024)), + 'decoder.diff_blocks.3.mlp.w_gate.weight': ('BF16', (1024, 256)), + 'decoder.diff_blocks.3.mlp.w_up.weight': ('BF16', (1024, 256)), + 'decoder.diff_blocks.3.norm1.weight': ('BF16', (256,)), + 'decoder.diff_blocks.3.norm2.weight': ('BF16', (256,)), + 'decoder.diff_blocks.3.scale_shift_table': ('BF16', (7, 256)), + 'decoder.diff_blocks.4.attn.k_norm.weight': ('BF16', (64,)), + 'decoder.diff_blocks.4.attn.proj.bias': ('BF16', (256,)), + 'decoder.diff_blocks.4.attn.proj.weight': ('BF16', (256, 256)), + 'decoder.diff_blocks.4.attn.q_norm.weight': ('BF16', (64,)), + 'decoder.diff_blocks.4.attn.qkv.bias': ('BF16', (768,)), + 'decoder.diff_blocks.4.attn.qkv.weight': ('BF16', (768, 256)), + 'decoder.diff_blocks.4.context_proj.bias': ('BF16', (256,)), + 'decoder.diff_blocks.4.context_proj.weight': ('BF16', (256, 256)), + 'decoder.diff_blocks.4.mlp.w_down.weight': ('BF16', (256, 1024)), + 'decoder.diff_blocks.4.mlp.w_gate.weight': ('BF16', (1024, 256)), + 'decoder.diff_blocks.4.mlp.w_up.weight': ('BF16', (1024, 256)), + 'decoder.diff_blocks.4.norm1.weight': ('BF16', (256,)), + 'decoder.diff_blocks.4.norm2.weight': ('BF16', (256,)), + 'decoder.diff_blocks.4.scale_shift_table': ('BF16', (7, 256)), + 'decoder.diff_blocks.5.attn.k_norm.weight': ('BF16', (64,)), + 'decoder.diff_blocks.5.attn.proj.bias': ('BF16', (256,)), + 'decoder.diff_blocks.5.attn.proj.weight': ('BF16', (256, 256)), + 'decoder.diff_blocks.5.attn.q_norm.weight': ('BF16', (64,)), + 'decoder.diff_blocks.5.attn.qkv.bias': ('BF16', (768,)), + 'decoder.diff_blocks.5.attn.qkv.weight': ('BF16', (768, 256)), + 'decoder.diff_blocks.5.context_proj.bias': ('BF16', (256,)), + 'decoder.diff_blocks.5.context_proj.weight': ('BF16', (256, 256)), + 'decoder.diff_blocks.5.mlp.w_down.weight': ('BF16', (256, 1024)), + 'decoder.diff_blocks.5.mlp.w_gate.weight': ('BF16', (1024, 256)), + 'decoder.diff_blocks.5.mlp.w_up.weight': ('BF16', (1024, 256)), + 'decoder.diff_blocks.5.norm1.weight': ('BF16', (256,)), + 'decoder.diff_blocks.5.norm2.weight': ('BF16', (256,)), + 'decoder.diff_blocks.5.scale_shift_table': ('BF16', (7, 256)), + 'decoder.diff_blocks.6.attn.k_norm.weight': ('BF16', (64,)), + 'decoder.diff_blocks.6.attn.proj.bias': ('BF16', (256,)), + 'decoder.diff_blocks.6.attn.proj.weight': ('BF16', (256, 256)), + 'decoder.diff_blocks.6.attn.q_norm.weight': ('BF16', (64,)), + 'decoder.diff_blocks.6.attn.qkv.bias': ('BF16', (768,)), + 'decoder.diff_blocks.6.attn.qkv.weight': ('BF16', (768, 256)), + 'decoder.diff_blocks.6.context_proj.bias': ('BF16', (256,)), + 'decoder.diff_blocks.6.context_proj.weight': ('BF16', (256, 256)), + 'decoder.diff_blocks.6.mlp.w_down.weight': ('BF16', (256, 1024)), + 'decoder.diff_blocks.6.mlp.w_gate.weight': ('BF16', (1024, 256)), + 'decoder.diff_blocks.6.mlp.w_up.weight': ('BF16', (1024, 256)), + 'decoder.diff_blocks.6.norm1.weight': ('BF16', (256,)), + 'decoder.diff_blocks.6.norm2.weight': ('BF16', (256,)), + 'decoder.diff_blocks.6.scale_shift_table': ('BF16', (7, 256)), + 'decoder.diff_blocks.7.attn.k_norm.weight': ('BF16', (64,)), + 'decoder.diff_blocks.7.attn.proj.bias': ('BF16', (256,)), + 'decoder.diff_blocks.7.attn.proj.weight': ('BF16', (256, 256)), + 'decoder.diff_blocks.7.attn.q_norm.weight': ('BF16', (64,)), + 'decoder.diff_blocks.7.attn.qkv.bias': ('BF16', (768,)), + 'decoder.diff_blocks.7.attn.qkv.weight': ('BF16', (768, 256)), + 'decoder.diff_blocks.7.context_proj.bias': ('BF16', (256,)), + 'decoder.diff_blocks.7.context_proj.weight': ('BF16', (256, 256)), + 'decoder.diff_blocks.7.mlp.w_down.weight': ('BF16', (256, 1024)), + 'decoder.diff_blocks.7.mlp.w_gate.weight': ('BF16', (1024, 256)), + 'decoder.diff_blocks.7.mlp.w_up.weight': ('BF16', (1024, 256)), + 'decoder.diff_blocks.7.norm1.weight': ('BF16', (256,)), + 'decoder.diff_blocks.7.norm2.weight': ('BF16', (256,)), + 'decoder.diff_blocks.7.scale_shift_table': ('BF16', (7, 256)), + 'decoder.norm_out.weight': ('BF16', (256,)), + 'decoder.shared_adaln.proj.bias': ('BF16', (1792,)), + 'decoder.shared_adaln.proj.weight': ('BF16', (1792, 384)), + 'decoder.t_embedder.mlp.0.bias': ('BF16', (384,)), + 'decoder.t_embedder.mlp.0.weight': ('BF16', (384, 256)), + 'decoder.t_embedder.mlp.2.bias': ('BF16', (384,)), + 'decoder.t_embedder.mlp.2.weight': ('BF16', (384, 384)), + 'decoder.type_emb': ('BF16', (128,)), + 'decoder.upsamples.0.proj.bias': ('BF16', (4096,)), + 'decoder.upsamples.0.proj.weight': ('BF16', (4096, 2048)), + 'decoder.upsamples.1.proj.bias': ('BF16', (1024,)), + 'decoder.upsamples.1.proj.weight': ('BF16', (1024, 1024)), + 'decoder.upsamples.2.proj.bias': ('BF16', (4096,)), + 'decoder.upsamples.2.proj.weight': ('BF16', (4096, 512)), + 'decoder.upsamples.3.proj.bias': ('BF16', (2048,)), + 'decoder.upsamples.3.proj.weight': ('BF16', (2048, 512)), + 'encoder.conv_in.conv.bias': ('BF16', (128,)), + 'encoder.conv_in.conv.weight': ('BF16', (128, 48, 3, 3, 3)), + 'encoder.conv_out.conv.bias': ('BF16', (129,)), + 'encoder.conv_out.conv.weight': ('BF16', (129, 1024, 3, 3, 3)), + 'encoder.down_blocks.0.res_blocks.0.conv1.conv.bias': ('BF16', (128,)), + 'encoder.down_blocks.0.res_blocks.0.conv1.conv.weight': ('BF16', (128, 128, 3, 3, 3)), + 'encoder.down_blocks.0.res_blocks.0.conv2.conv.bias': ('BF16', (128,)), + 'encoder.down_blocks.0.res_blocks.0.conv2.conv.weight': ('BF16', (128, 128, 3, 3, 3)), + 'encoder.down_blocks.0.res_blocks.1.conv1.conv.bias': ('BF16', (128,)), + 'encoder.down_blocks.0.res_blocks.1.conv1.conv.weight': ('BF16', (128, 128, 3, 3, 3)), + 'encoder.down_blocks.0.res_blocks.1.conv2.conv.bias': ('BF16', (128,)), + 'encoder.down_blocks.0.res_blocks.1.conv2.conv.weight': ('BF16', (128, 128, 3, 3, 3)), + 'encoder.down_blocks.0.res_blocks.2.conv1.conv.bias': ('BF16', (128,)), + 'encoder.down_blocks.0.res_blocks.2.conv1.conv.weight': ('BF16', (128, 128, 3, 3, 3)), + 'encoder.down_blocks.0.res_blocks.2.conv2.conv.bias': ('BF16', (128,)), + 'encoder.down_blocks.0.res_blocks.2.conv2.conv.weight': ('BF16', (128, 128, 3, 3, 3)), + 'encoder.down_blocks.0.res_blocks.3.conv1.conv.bias': ('BF16', (128,)), + 'encoder.down_blocks.0.res_blocks.3.conv1.conv.weight': ('BF16', (128, 128, 3, 3, 3)), + 'encoder.down_blocks.0.res_blocks.3.conv2.conv.bias': ('BF16', (128,)), + 'encoder.down_blocks.0.res_blocks.3.conv2.conv.weight': ('BF16', (128, 128, 3, 3, 3)), + 'encoder.down_blocks.1.conv.conv.bias': ('BF16', (64,)), + 'encoder.down_blocks.1.conv.conv.weight': ('BF16', (64, 128, 3, 3, 3)), + 'encoder.down_blocks.2.res_blocks.0.conv1.conv.bias': ('BF16', (256,)), + 'encoder.down_blocks.2.res_blocks.0.conv1.conv.weight': ('BF16', (256, 256, 3, 3, 3)), + 'encoder.down_blocks.2.res_blocks.0.conv2.conv.bias': ('BF16', (256,)), + 'encoder.down_blocks.2.res_blocks.0.conv2.conv.weight': ('BF16', (256, 256, 3, 3, 3)), + 'encoder.down_blocks.2.res_blocks.1.conv1.conv.bias': ('BF16', (256,)), + 'encoder.down_blocks.2.res_blocks.1.conv1.conv.weight': ('BF16', (256, 256, 3, 3, 3)), + 'encoder.down_blocks.2.res_blocks.1.conv2.conv.bias': ('BF16', (256,)), + 'encoder.down_blocks.2.res_blocks.1.conv2.conv.weight': ('BF16', (256, 256, 3, 3, 3)), + 'encoder.down_blocks.2.res_blocks.2.conv1.conv.bias': ('BF16', (256,)), + 'encoder.down_blocks.2.res_blocks.2.conv1.conv.weight': ('BF16', (256, 256, 3, 3, 3)), + 'encoder.down_blocks.2.res_blocks.2.conv2.conv.bias': ('BF16', (256,)), + 'encoder.down_blocks.2.res_blocks.2.conv2.conv.weight': ('BF16', (256, 256, 3, 3, 3)), + 'encoder.down_blocks.2.res_blocks.3.conv1.conv.bias': ('BF16', (256,)), + 'encoder.down_blocks.2.res_blocks.3.conv1.conv.weight': ('BF16', (256, 256, 3, 3, 3)), + 'encoder.down_blocks.2.res_blocks.3.conv2.conv.bias': ('BF16', (256,)), + 'encoder.down_blocks.2.res_blocks.3.conv2.conv.weight': ('BF16', (256, 256, 3, 3, 3)), + 'encoder.down_blocks.2.res_blocks.4.conv1.conv.bias': ('BF16', (256,)), + 'encoder.down_blocks.2.res_blocks.4.conv1.conv.weight': ('BF16', (256, 256, 3, 3, 3)), + 'encoder.down_blocks.2.res_blocks.4.conv2.conv.bias': ('BF16', (256,)), + 'encoder.down_blocks.2.res_blocks.4.conv2.conv.weight': ('BF16', (256, 256, 3, 3, 3)), + 'encoder.down_blocks.2.res_blocks.5.conv1.conv.bias': ('BF16', (256,)), + 'encoder.down_blocks.2.res_blocks.5.conv1.conv.weight': ('BF16', (256, 256, 3, 3, 3)), + 'encoder.down_blocks.2.res_blocks.5.conv2.conv.bias': ('BF16', (256,)), + 'encoder.down_blocks.2.res_blocks.5.conv2.conv.weight': ('BF16', (256, 256, 3, 3, 3)), + 'encoder.down_blocks.3.conv.conv.bias': ('BF16', (256,)), + 'encoder.down_blocks.3.conv.conv.weight': ('BF16', (256, 256, 3, 3, 3)), + 'encoder.down_blocks.4.res_blocks.0.conv1.conv.bias': ('BF16', (512,)), + 'encoder.down_blocks.4.res_blocks.0.conv1.conv.weight': ('BF16', (512, 512, 3, 3, 3)), + 'encoder.down_blocks.4.res_blocks.0.conv2.conv.bias': ('BF16', (512,)), + 'encoder.down_blocks.4.res_blocks.0.conv2.conv.weight': ('BF16', (512, 512, 3, 3, 3)), + 'encoder.down_blocks.4.res_blocks.1.conv1.conv.bias': ('BF16', (512,)), + 'encoder.down_blocks.4.res_blocks.1.conv1.conv.weight': ('BF16', (512, 512, 3, 3, 3)), + 'encoder.down_blocks.4.res_blocks.1.conv2.conv.bias': ('BF16', (512,)), + 'encoder.down_blocks.4.res_blocks.1.conv2.conv.weight': ('BF16', (512, 512, 3, 3, 3)), + 'encoder.down_blocks.4.res_blocks.2.conv1.conv.bias': ('BF16', (512,)), + 'encoder.down_blocks.4.res_blocks.2.conv1.conv.weight': ('BF16', (512, 512, 3, 3, 3)), + 'encoder.down_blocks.4.res_blocks.2.conv2.conv.bias': ('BF16', (512,)), + 'encoder.down_blocks.4.res_blocks.2.conv2.conv.weight': ('BF16', (512, 512, 3, 3, 3)), + 'encoder.down_blocks.4.res_blocks.3.conv1.conv.bias': ('BF16', (512,)), + 'encoder.down_blocks.4.res_blocks.3.conv1.conv.weight': ('BF16', (512, 512, 3, 3, 3)), + 'encoder.down_blocks.4.res_blocks.3.conv2.conv.bias': ('BF16', (512,)), + 'encoder.down_blocks.4.res_blocks.3.conv2.conv.weight': ('BF16', (512, 512, 3, 3, 3)), + 'encoder.down_blocks.5.conv.conv.bias': ('BF16', (128,)), + 'encoder.down_blocks.5.conv.conv.weight': ('BF16', (128, 512, 3, 3, 3)), + 'encoder.down_blocks.6.res_blocks.0.conv1.conv.bias': ('BF16', (1024,)), + 'encoder.down_blocks.6.res_blocks.0.conv1.conv.weight': ('BF16', (1024, 1024, 3, 3, 3)), + 'encoder.down_blocks.6.res_blocks.0.conv2.conv.bias': ('BF16', (1024,)), + 'encoder.down_blocks.6.res_blocks.0.conv2.conv.weight': ('BF16', (1024, 1024, 3, 3, 3)), + 'encoder.down_blocks.6.res_blocks.1.conv1.conv.bias': ('BF16', (1024,)), + 'encoder.down_blocks.6.res_blocks.1.conv1.conv.weight': ('BF16', (1024, 1024, 3, 3, 3)), + 'encoder.down_blocks.6.res_blocks.1.conv2.conv.bias': ('BF16', (1024,)), + 'encoder.down_blocks.6.res_blocks.1.conv2.conv.weight': ('BF16', (1024, 1024, 3, 3, 3)), + 'encoder.down_blocks.7.conv.conv.bias': ('BF16', (128,)), + 'encoder.down_blocks.7.conv.conv.weight': ('BF16', (128, 1024, 3, 3, 3)), + 'encoder.down_blocks.8.res_blocks.0.conv1.conv.bias': ('BF16', (1024,)), + 'encoder.down_blocks.8.res_blocks.0.conv1.conv.weight': ('BF16', (1024, 1024, 3, 3, 3)), + 'encoder.down_blocks.8.res_blocks.0.conv2.conv.bias': ('BF16', (1024,)), + 'encoder.down_blocks.8.res_blocks.0.conv2.conv.weight': ('BF16', (1024, 1024, 3, 3, 3)), + 'encoder.down_blocks.8.res_blocks.1.conv1.conv.bias': ('BF16', (1024,)), + 'encoder.down_blocks.8.res_blocks.1.conv1.conv.weight': ('BF16', (1024, 1024, 3, 3, 3)), + 'encoder.down_blocks.8.res_blocks.1.conv2.conv.bias': ('BF16', (1024,)), + 'encoder.down_blocks.8.res_blocks.1.conv2.conv.weight': ('BF16', (1024, 1024, 3, 3, 3)), + 'per_channel_statistics.mean-of-means': ('BF16', (128,)), + 'per_channel_statistics.std-of-means': ('BF16', (128,)), +} + +VAE_DECODER = { + 'conv_in_x_t.bias': ('BF16', (256,)), + 'conv_in_x_t.weight': ('BF16', (256, 48)), + 'conv_out.bias': ('BF16', (48,)), + 'conv_out.weight': ('BF16', (48, 256)), + 'det_stages.3.0.attn.k_norm.weight': ('BF16', (64,)), + 'det_stages.3.0.attn.proj.bias': ('BF16', (512,)), + 'det_stages.3.0.attn.proj.weight': ('BF16', (512, 512)), + 'det_stages.3.0.attn.q_norm.weight': ('BF16', (64,)), + 'det_stages.3.0.attn.qkv.to_k.bias': ('BF16', (512,)), + 'det_stages.3.0.attn.qkv.to_k.weight': ('BF16', (512, 512)), + 'det_stages.3.0.attn.qkv.to_q.bias': ('BF16', (512,)), + 'det_stages.3.0.attn.qkv.to_q.weight': ('BF16', (512, 512)), + 'det_stages.3.0.attn.qkv.to_v.bias': ('BF16', (512,)), + 'det_stages.3.0.attn.qkv.to_v.weight': ('BF16', (512, 512)), + 'det_stages.3.0.mlp.w_down.weight': ('BF16', (512, 2048)), + 'det_stages.3.0.mlp.w_gate.weight': ('BF16', (2048, 512)), + 'det_stages.3.0.mlp.w_up.weight': ('BF16', (2048, 512)), + 'det_stages.3.0.norm1.weight': ('BF16', (512,)), + 'det_stages.3.0.norm2.weight': ('BF16', (512,)), + 'det_stages.3.1.attn.k_norm.weight': ('BF16', (64,)), + 'det_stages.3.1.attn.proj.bias': ('BF16', (512,)), + 'det_stages.3.1.attn.proj.weight': ('BF16', (512, 512)), + 'det_stages.3.1.attn.q_norm.weight': ('BF16', (64,)), + 'det_stages.3.1.attn.qkv.to_k.bias': ('BF16', (512,)), + 'det_stages.3.1.attn.qkv.to_k.weight': ('BF16', (512, 512)), + 'det_stages.3.1.attn.qkv.to_q.bias': ('BF16', (512,)), + 'det_stages.3.1.attn.qkv.to_q.weight': ('BF16', (512, 512)), + 'det_stages.3.1.attn.qkv.to_v.bias': ('BF16', (512,)), + 'det_stages.3.1.attn.qkv.to_v.weight': ('BF16', (512, 512)), + 'det_stages.3.1.mlp.w_down.weight': ('BF16', (512, 2048)), + 'det_stages.3.1.mlp.w_gate.weight': ('BF16', (2048, 512)), + 'det_stages.3.1.mlp.w_up.weight': ('BF16', (2048, 512)), + 'det_stages.3.1.norm1.weight': ('BF16', (512,)), + 'det_stages.3.1.norm2.weight': ('BF16', (512,)), + 'diff_blocks.0.attn.k_norm.weight': ('BF16', (64,)), + 'diff_blocks.0.attn.proj.bias': ('BF16', (256,)), + 'diff_blocks.0.attn.proj.weight': ('BF16', (256, 256)), + 'diff_blocks.0.attn.q_norm.weight': ('BF16', (64,)), + 'diff_blocks.0.attn.qkv.to_k.bias': ('BF16', (256,)), + 'diff_blocks.0.attn.qkv.to_k.weight': ('BF16', (256, 256)), + 'diff_blocks.0.attn.qkv.to_q.bias': ('BF16', (256,)), + 'diff_blocks.0.attn.qkv.to_q.weight': ('BF16', (256, 256)), + 'diff_blocks.0.attn.qkv.to_v.bias': ('BF16', (256,)), + 'diff_blocks.0.attn.qkv.to_v.weight': ('BF16', (256, 256)), + 'diff_blocks.0.context_proj.bias': ('BF16', (256,)), + 'diff_blocks.0.context_proj.weight': ('BF16', (256, 256)), + 'diff_blocks.0.mlp.w_down.weight': ('BF16', (256, 1024)), + 'diff_blocks.0.mlp.w_gate.weight': ('BF16', (1024, 256)), + 'diff_blocks.0.mlp.w_up.weight': ('BF16', (1024, 256)), + 'diff_blocks.0.norm1.weight': ('BF16', (256,)), + 'diff_blocks.0.norm2.weight': ('BF16', (256,)), + 'diff_blocks.0.scale_shift_table': ('BF16', (7, 256)), + 'diff_blocks.1.attn.k_norm.weight': ('BF16', (64,)), + 'diff_blocks.1.attn.proj.bias': ('BF16', (256,)), + 'diff_blocks.1.attn.proj.weight': ('BF16', (256, 256)), + 'diff_blocks.1.attn.q_norm.weight': ('BF16', (64,)), + 'diff_blocks.1.attn.qkv.to_k.bias': ('BF16', (256,)), + 'diff_blocks.1.attn.qkv.to_k.weight': ('BF16', (256, 256)), + 'diff_blocks.1.attn.qkv.to_q.bias': ('BF16', (256,)), + 'diff_blocks.1.attn.qkv.to_q.weight': ('BF16', (256, 256)), + 'diff_blocks.1.attn.qkv.to_v.bias': ('BF16', (256,)), + 'diff_blocks.1.attn.qkv.to_v.weight': ('BF16', (256, 256)), + 'diff_blocks.1.context_proj.bias': ('BF16', (256,)), + 'diff_blocks.1.context_proj.weight': ('BF16', (256, 256)), + 'diff_blocks.1.mlp.w_down.weight': ('BF16', (256, 1024)), + 'diff_blocks.1.mlp.w_gate.weight': ('BF16', (1024, 256)), + 'diff_blocks.1.mlp.w_up.weight': ('BF16', (1024, 256)), + 'diff_blocks.1.norm1.weight': ('BF16', (256,)), + 'diff_blocks.1.norm2.weight': ('BF16', (256,)), + 'diff_blocks.1.scale_shift_table': ('BF16', (7, 256)), + 'diff_blocks.2.attn.k_norm.weight': ('BF16', (64,)), + 'diff_blocks.2.attn.proj.bias': ('BF16', (256,)), + 'diff_blocks.2.attn.proj.weight': ('BF16', (256, 256)), + 'diff_blocks.2.attn.q_norm.weight': ('BF16', (64,)), + 'diff_blocks.2.attn.qkv.to_k.bias': ('BF16', (256,)), + 'diff_blocks.2.attn.qkv.to_k.weight': ('BF16', (256, 256)), + 'diff_blocks.2.attn.qkv.to_q.bias': ('BF16', (256,)), + 'diff_blocks.2.attn.qkv.to_q.weight': ('BF16', (256, 256)), + 'diff_blocks.2.attn.qkv.to_v.bias': ('BF16', (256,)), + 'diff_blocks.2.attn.qkv.to_v.weight': ('BF16', (256, 256)), + 'diff_blocks.2.context_proj.bias': ('BF16', (256,)), + 'diff_blocks.2.context_proj.weight': ('BF16', (256, 256)), + 'diff_blocks.2.mlp.w_down.weight': ('BF16', (256, 1024)), + 'diff_blocks.2.mlp.w_gate.weight': ('BF16', (1024, 256)), + 'diff_blocks.2.mlp.w_up.weight': ('BF16', (1024, 256)), + 'diff_blocks.2.norm1.weight': ('BF16', (256,)), + 'diff_blocks.2.norm2.weight': ('BF16', (256,)), + 'diff_blocks.2.scale_shift_table': ('BF16', (7, 256)), + 'diff_blocks.3.attn.k_norm.weight': ('BF16', (64,)), + 'diff_blocks.3.attn.proj.bias': ('BF16', (256,)), + 'diff_blocks.3.attn.proj.weight': ('BF16', (256, 256)), + 'diff_blocks.3.attn.q_norm.weight': ('BF16', (64,)), + 'diff_blocks.3.attn.qkv.to_k.bias': ('BF16', (256,)), + 'diff_blocks.3.attn.qkv.to_k.weight': ('BF16', (256, 256)), + 'diff_blocks.3.attn.qkv.to_q.bias': ('BF16', (256,)), + 'diff_blocks.3.attn.qkv.to_q.weight': ('BF16', (256, 256)), + 'diff_blocks.3.attn.qkv.to_v.bias': ('BF16', (256,)), + 'diff_blocks.3.attn.qkv.to_v.weight': ('BF16', (256, 256)), + 'diff_blocks.3.context_proj.bias': ('BF16', (256,)), + 'diff_blocks.3.context_proj.weight': ('BF16', (256, 256)), + 'diff_blocks.3.mlp.w_down.weight': ('BF16', (256, 1024)), + 'diff_blocks.3.mlp.w_gate.weight': ('BF16', (1024, 256)), + 'diff_blocks.3.mlp.w_up.weight': ('BF16', (1024, 256)), + 'diff_blocks.3.norm1.weight': ('BF16', (256,)), + 'diff_blocks.3.norm2.weight': ('BF16', (256,)), + 'diff_blocks.3.scale_shift_table': ('BF16', (7, 256)), + 'diff_blocks.4.attn.k_norm.weight': ('BF16', (64,)), + 'diff_blocks.4.attn.proj.bias': ('BF16', (256,)), + 'diff_blocks.4.attn.proj.weight': ('BF16', (256, 256)), + 'diff_blocks.4.attn.q_norm.weight': ('BF16', (64,)), + 'diff_blocks.4.attn.qkv.to_k.bias': ('BF16', (256,)), + 'diff_blocks.4.attn.qkv.to_k.weight': ('BF16', (256, 256)), + 'diff_blocks.4.attn.qkv.to_q.bias': ('BF16', (256,)), + 'diff_blocks.4.attn.qkv.to_q.weight': ('BF16', (256, 256)), + 'diff_blocks.4.attn.qkv.to_v.bias': ('BF16', (256,)), + 'diff_blocks.4.attn.qkv.to_v.weight': ('BF16', (256, 256)), + 'diff_blocks.4.context_proj.bias': ('BF16', (256,)), + 'diff_blocks.4.context_proj.weight': ('BF16', (256, 256)), + 'diff_blocks.4.mlp.w_down.weight': ('BF16', (256, 1024)), + 'diff_blocks.4.mlp.w_gate.weight': ('BF16', (1024, 256)), + 'diff_blocks.4.mlp.w_up.weight': ('BF16', (1024, 256)), + 'diff_blocks.4.norm1.weight': ('BF16', (256,)), + 'diff_blocks.4.norm2.weight': ('BF16', (256,)), + 'diff_blocks.4.scale_shift_table': ('BF16', (7, 256)), + 'diff_blocks.5.attn.k_norm.weight': ('BF16', (64,)), + 'diff_blocks.5.attn.proj.bias': ('BF16', (256,)), + 'diff_blocks.5.attn.proj.weight': ('BF16', (256, 256)), + 'diff_blocks.5.attn.q_norm.weight': ('BF16', (64,)), + 'diff_blocks.5.attn.qkv.to_k.bias': ('BF16', (256,)), + 'diff_blocks.5.attn.qkv.to_k.weight': ('BF16', (256, 256)), + 'diff_blocks.5.attn.qkv.to_q.bias': ('BF16', (256,)), + 'diff_blocks.5.attn.qkv.to_q.weight': ('BF16', (256, 256)), + 'diff_blocks.5.attn.qkv.to_v.bias': ('BF16', (256,)), + 'diff_blocks.5.attn.qkv.to_v.weight': ('BF16', (256, 256)), + 'diff_blocks.5.context_proj.bias': ('BF16', (256,)), + 'diff_blocks.5.context_proj.weight': ('BF16', (256, 256)), + 'diff_blocks.5.mlp.w_down.weight': ('BF16', (256, 1024)), + 'diff_blocks.5.mlp.w_gate.weight': ('BF16', (1024, 256)), + 'diff_blocks.5.mlp.w_up.weight': ('BF16', (1024, 256)), + 'diff_blocks.5.norm1.weight': ('BF16', (256,)), + 'diff_blocks.5.norm2.weight': ('BF16', (256,)), + 'diff_blocks.5.scale_shift_table': ('BF16', (7, 256)), + 'diff_blocks.6.attn.k_norm.weight': ('BF16', (64,)), + 'diff_blocks.6.attn.proj.bias': ('BF16', (256,)), + 'diff_blocks.6.attn.proj.weight': ('BF16', (256, 256)), + 'diff_blocks.6.attn.q_norm.weight': ('BF16', (64,)), + 'diff_blocks.6.attn.qkv.to_k.bias': ('BF16', (256,)), + 'diff_blocks.6.attn.qkv.to_k.weight': ('BF16', (256, 256)), + 'diff_blocks.6.attn.qkv.to_q.bias': ('BF16', (256,)), + 'diff_blocks.6.attn.qkv.to_q.weight': ('BF16', (256, 256)), + 'diff_blocks.6.attn.qkv.to_v.bias': ('BF16', (256,)), + 'diff_blocks.6.attn.qkv.to_v.weight': ('BF16', (256, 256)), + 'diff_blocks.6.context_proj.bias': ('BF16', (256,)), + 'diff_blocks.6.context_proj.weight': ('BF16', (256, 256)), + 'diff_blocks.6.mlp.w_down.weight': ('BF16', (256, 1024)), + 'diff_blocks.6.mlp.w_gate.weight': ('BF16', (1024, 256)), + 'diff_blocks.6.mlp.w_up.weight': ('BF16', (1024, 256)), + 'diff_blocks.6.norm1.weight': ('BF16', (256,)), + 'diff_blocks.6.norm2.weight': ('BF16', (256,)), + 'diff_blocks.6.scale_shift_table': ('BF16', (7, 256)), + 'diff_blocks.7.attn.k_norm.weight': ('BF16', (64,)), + 'diff_blocks.7.attn.proj.bias': ('BF16', (256,)), + 'diff_blocks.7.attn.proj.weight': ('BF16', (256, 256)), + 'diff_blocks.7.attn.q_norm.weight': ('BF16', (64,)), + 'diff_blocks.7.attn.qkv.to_k.bias': ('BF16', (256,)), + 'diff_blocks.7.attn.qkv.to_k.weight': ('BF16', (256, 256)), + 'diff_blocks.7.attn.qkv.to_q.bias': ('BF16', (256,)), + 'diff_blocks.7.attn.qkv.to_q.weight': ('BF16', (256, 256)), + 'diff_blocks.7.attn.qkv.to_v.bias': ('BF16', (256,)), + 'diff_blocks.7.attn.qkv.to_v.weight': ('BF16', (256, 256)), + 'diff_blocks.7.context_proj.bias': ('BF16', (256,)), + 'diff_blocks.7.context_proj.weight': ('BF16', (256, 256)), + 'diff_blocks.7.mlp.w_down.weight': ('BF16', (256, 1024)), + 'diff_blocks.7.mlp.w_gate.weight': ('BF16', (1024, 256)), + 'diff_blocks.7.mlp.w_up.weight': ('BF16', (1024, 256)), + 'diff_blocks.7.norm1.weight': ('BF16', (256,)), + 'diff_blocks.7.norm2.weight': ('BF16', (256,)), + 'diff_blocks.7.scale_shift_table': ('BF16', (7, 256)), + 'norm_out.weight': ('BF16', (256,)), + 'shared_adaln.proj.bias': ('BF16', (1792,)), + 'shared_adaln.proj.weight': ('BF16', (1792, 384)), + 't_embedder.timestep_embedder.linear_1.bias': ('BF16', (384,)), + 't_embedder.timestep_embedder.linear_1.weight': ('BF16', (384, 256)), + 't_embedder.timestep_embedder.linear_2.bias': ('BF16', (384,)), + 't_embedder.timestep_embedder.linear_2.weight': ('BF16', (384, 384)), + 'upsamples.3.proj.bias': ('BF16', (2048,)), + 'upsamples.3.proj.weight': ('BF16', (2048, 512)), +} diff --git a/jasna/main.py b/jasna/main.py index 4d5ffef0..b4ddcc8b 100644 --- a/jasna/main.py +++ b/jasna/main.py @@ -54,6 +54,7 @@ def _session_config_from_args( ltx_seed=int(args.ltx_seed), ltx_fast=bool(args.ltx_fast), ltx_model=CLI_RESTORATION_MODELS[args.restoration_model_name][1], + ltx_trial=bool(args.ltx_trial), compile_basicvsrpp=bool(args.compile_basicvsrpp), max_clip_size=int(args.max_clip_size), temporal_overlap=int(args.temporal_overlap), @@ -230,6 +231,12 @@ def build_parser() -> argparse.ArgumentParser: action=argparse.BooleanOptionalAction, help=CLI_HELP["ltx_fast"], ) + restoration.add_argument( + "--ltx-trial", + default=False, + action=argparse.BooleanOptionalAction, + help=CLI_HELP["ltx_trial"], + ) restoration.add_argument( "--compile-basicvsrpp", default=True, @@ -779,7 +786,7 @@ def _ensure_ltx_model(config: SessionConfig, segments) -> None: uses_ltx = config.restoration_model_name == "ltx" or any( segment.restoration is not None and segment.restoration.model == "ltx" for segment in segments or () ) - if not uses_ltx: + if not uses_ltx or config.ltx_trial: return directory = config.restoration_model_path if config.restoration_model_name == "ltx" else default_restoration_model_path("ltx") missing = missing_downloads(directory, config.ltx_model, fast=config.ltx_fast) @@ -902,10 +909,12 @@ def _run_post_export_action() -> None: parser.error("--restoration-model-name ltx does not support --stream") if args.secondary_restoration != "none" or args.denoise != "none": parser.error("--restoration-model-name ltx does not support --secondary-restoration or --denoise") + elif args.ltx_trial: + parser.error("--ltx-trial needs --restoration-model-name ltx or ltx-undistilled") restoration_model_path = Path( args.restoration_model_path or default_restoration_model_path(restoration_backend) ) - if not restoration_model_path.exists(): + if not args.ltx_trial and not restoration_model_path.exists(): raise FileNotFoundError(str(restoration_model_path)) segments = None diff --git a/jasna/session_config.py b/jasna/session_config.py index 17cdafd8..02277638 100644 --- a/jasna/session_config.py +++ b/jasna/session_config.py @@ -43,6 +43,7 @@ class SessionConfig: ltx_seed: int ltx_fast: bool ltx_model: LtxModelName + ltx_trial: bool compile_basicvsrpp: bool max_clip_size: int temporal_overlap: int diff --git a/jasna/session_factory.py b/jasna/session_factory.py index e76dd09b..ffba2d00 100644 --- a/jasna/session_factory.py +++ b/jasna/session_factory.py @@ -10,6 +10,7 @@ from __future__ import annotations +import logging from dataclasses import dataclass from pathlib import Path from typing import TYPE_CHECKING, Callable @@ -30,6 +31,9 @@ DetectionModel = RfDetrMosaicDetectionModel | YoloMosaicDetectionModel +logger = logging.getLogger(__name__) + + @dataclass class RestorationSession: device: "torch.device" @@ -240,7 +244,14 @@ def _ltx_model_files( raise ValueError("LTX restoration needs an NVIDIA GPU") if config.ltx_fast and torch.cuda.get_device_capability(device)[0] < 10: raise ValueError("The fast LTX model needs an RTX 50-series (Blackwell) GPU") - files = LtxModelFiles.from_dir(_restoration_model_path(config, "ltx"), config.ltx_model, fast=config.ltx_fast) + if config.ltx_trial: + logger.warning( + "LTX trial: placeholder weights, so the output is not restored. Timings match a real run " + "except model loading, which skips reading and decrypting the model files." + ) + files = LtxModelFiles.placeholder(config.ltx_model, fast=config.ltx_fast) + else: + files = LtxModelFiles.from_dir(_restoration_model_path(config, "ltx"), config.ltx_model, fast=config.ltx_fast) ensure_engines_compiled( EngineCompilationRequest( device=str(device), diff --git a/scripts/ltx_trial_shapes.py b/scripts/ltx_trial_shapes.py new file mode 100644 index 00000000..350338e2 --- /dev/null +++ b/scripts/ltx_trial_shapes.py @@ -0,0 +1,93 @@ +"""Regenerate jasna/ltx/trial_shapes.py from local LTX model files (dev only). + +Reads only the safetensors headers: tensor names, shapes and dtypes, the VAE architecture +config, and each transformer's step count and whether it uses STG. Weights, the sigma +schedule and the license text are not copied. + + python scripts/ltx_trial_shapes.py [model_weights/ltx-restore] +""" + +import json +import struct +import sys +from pathlib import Path + +REPO = Path(__file__).resolve().parents[1] +sys.path.insert(0, str(REPO)) + +from jasna.ltx.model_files import LTX_MODELS, bundle_names # noqa: E402 + +OUTPUT = REPO / "jasna" / "ltx" / "trial_shapes.py" + + +def read_header(path: Path) -> tuple[dict[str, tuple[str, tuple[int, ...]]], dict[str, str]]: + with path.open("rb") as handle: + size = struct.unpack(" tuple[int, dict, dict]: + top = {name: spec for name, spec in tensors.items() if not name.startswith("blocks.")} + blocks: dict[int, dict] = {} + for name, spec in tensors.items(): + if name.startswith("blocks."): + _, index, rest = name.split(".", 2) + blocks.setdefault(int(index), {})[rest] = spec + first = blocks[0] + if sorted(blocks) != list(range(len(blocks))) or any(block != first for block in blocks.values()): + raise ValueError("transformer blocks differ; the trial assumes identical blocks") + return len(blocks), top, first + + +def main() -> None: + directory = Path(sys.argv[1]) if len(sys.argv) > 1 else REPO / "model_weights" / "ltx-restore" + transformer: dict[str, tuple[dict, dict]] = {} + steps: dict[str, tuple[int, bool]] = {} + block_count = None + for model in LTX_MODELS: + for fast, precision in ((False, "int8"), (True, "nvfp4")): + tensors, metadata = read_header(directory / bundle_names(model, fast=fast)[0]) + count, top, block = split_blocks(tensors) + if transformer.setdefault(precision, (top, block)) != (top, block) or block_count not in (None, count): + raise ValueError(f"{model} {precision}: shapes differ from the other model") + block_count = count + sampler = (len(json.loads(metadata["sigmas"])) - 1, float(metadata["stg_scale"]) > 0) + if steps.setdefault(model, sampler) != sampler: + raise ValueError(f"{model}: int8 and nvfp4 sample differently") + _, vae_name, decoder_name = bundle_names(LTX_MODELS[0], fast=False) + vae, vae_metadata = read_header(directory / vae_name) + decoder, _ = read_header(directory / decoder_name) + + lines = [ + '"""Tensor names, shapes and dtypes of the LTX model files, for the trial\'s placeholder', + "weights. Generated by scripts/ltx_trial_shapes.py; architecture only, no weights.\"\"\"", + "", + f"BLOCKS = {block_count}", + "", + "# model -> (sampler steps, uses STG)", + f"STEPS = {steps!r}", + "", + f"VAE_CONFIG = {json.loads(vae_metadata['config'])!r}", + "", + "TRANSFORMER = {", + ] + for precision, (top, block) in transformer.items(): + lines.append(f" {precision!r}: {{") + for part, specs in (("top", top), ("block", block)): + lines.append(f" {part!r}: {{") + lines += [f" {name!r}: {spec!r}," for name, spec in specs.items()] + lines.append(" },") + lines.append(" },") + lines.append("}") + for name, specs in (("VAE", vae), ("VAE_DECODER", decoder)): + lines += ["", f"{name} = {{"] + lines += [f" {tensor!r}: {spec!r}," for tensor, spec in specs.items()] + lines.append("}") + OUTPUT.write_text("\n".join(lines) + "\n") + print(f"wrote {OUTPUT}") + + +if __name__ == "__main__": + main() diff --git a/tests/factories.py b/tests/factories.py index 9c3d209e..6a868f37 100644 --- a/tests/factories.py +++ b/tests/factories.py @@ -31,6 +31,7 @@ def session_config(**overrides) -> SessionConfig: ltx_seed=0, ltx_fast=False, ltx_model="distilled", + ltx_trial=False, compile_basicvsrpp=True, max_clip_size=90, temporal_overlap=8, diff --git a/tests/test_ltx_trial.py b/tests/test_ltx_trial.py new file mode 100644 index 00000000..eb0dfed1 --- /dev/null +++ b/tests/test_ltx_trial.py @@ -0,0 +1,96 @@ +import json +import struct +from pathlib import Path +from unittest.mock import MagicMock, patch + +import pytest +import torch + +from factories import session_config +from jasna.ltx import transformer as T +from jasna.ltx.model_files import LTX_MODELS, LtxModelFiles, bundle_names, open_tensors +from jasna.ltx.sampler import sampler_settings +from jasna.models.ltx_vae.loader import load_video_decoder, load_video_encoder + +DEV_FILES = Path(__file__).resolve().parents[1] / "model_weights" / "ltx-restore" +CASES = [(model, fast) for model in LTX_MODELS for fast in (False, True)] + + +def _header(path: Path) -> dict[str, tuple[str, tuple[int, ...]]]: + with path.open("rb") as handle: + header = json.loads(handle.read(struct.unpack(" dict[str, tuple[str, tuple[int, ...]]]: + names = {torch.bfloat16: "BF16", torch.float32: "F32", torch.int8: "I8", torch.uint8: "U8", + torch.float8_e4m3fn: "F8_E4M3"} + with open_tensors(source) as handle: + return {name: (names[t.dtype], tuple(t.shape)) for name in handle.keys() for t in [handle.get_tensor(name)]} + + +@pytest.mark.parametrize("model,fast", CASES) +def test_placeholders_match_the_real_files(model, fast): + names = bundle_names(model, fast=fast) + if not all((DEV_FILES / name).is_file() for name in names): + pytest.skip("raw LTX model files not present") + files = LtxModelFiles.placeholder(model, fast=fast) + for source, name in zip((files.transformer, files.vae, files.tuned_decoder), names): + assert _placeholder_specs(source) == _header(DEV_FILES / name) + + +@pytest.mark.parametrize("model,steps,stg", [("distilled", 8, 0.0), ("undistilled", 15, 1.0)]) +def test_placeholder_sampler_has_the_model_step_count(model, steps, stg): + with open_tensors(LtxModelFiles.placeholder(model, fast=False).transformer) as handle: + sigmas, stg_scale = sampler_settings(handle.metadata()) + assert len(sigmas) == steps + 1 and stg_scale == stg + assert sigmas[0] == 1 and sigmas[-1] == 0 and bool((sigmas[1:] < sigmas[:-1]).all()) + + +def test_placeholders_are_deterministic_and_finite(): + source = LtxModelFiles.placeholder("distilled", fast=True).transformer + with open_tensors(source) as first, open_tensors(source) as second: + for name in ("blocks.3.attn1.to_q.weight", "blocks.3.attn1.to_q.weight_scale", "blocks.3.attn1.to_q.weight_scale_2", + "blocks.3.attn1.to_q.bias", "prompt_context"): + tensor = first.get_tensor(name) + assert torch.equal(tensor.reshape(-1).view(torch.uint8), second.get_tensor(name).reshape(-1).view(torch.uint8)) + if tensor.dtype != torch.uint8: + assert torch.isfinite(tensor.float()).all() + assert (first.get_tensor("blocks.0.ff.net.2.weight_scale").float() > 0).all() + + +def test_placeholder_vae_loads_into_the_encoder_and_decoder(): + files = LtxModelFiles.placeholder("distilled", fast=False) + load_video_encoder(files.vae, torch.device("cpu")) + load_video_decoder(files.vae, files.tuned_decoder, torch.device("cpu")) + + +def test_trial_session_uses_placeholders_without_files_or_license(): + from jasna.session_factory import _ltx_model_files + + config = session_config(restoration_model_name="ltx", ltx_trial=True, ltx_model="undistilled") + with ( + patch("jasna.accelerator.is_nvidia_device", return_value=True), + patch("jasna.engine_compiler.ensure_engines_compiled"), + patch("jasna.ltx.model_files.LtxModelFiles.from_dir", side_effect=AssertionError("reads model files")), + ): + files = _ltx_model_files(config, MagicMock(), log_callback=None) + assert files == LtxModelFiles.placeholder("undistilled", fast=False) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA") +@pytest.mark.parametrize("fast", [False, True]) +def test_placeholder_block_runs_on_the_gpu(fast): + device = torch.device("cuda") + if fast and torch.cuda.get_device_capability(device)[0] < 10: + pytest.skip("NVFP4 needs Blackwell") + with open_tensors(LtxModelFiles.placeholder("distilled", fast=fast).transformer) as handle: + keys = handle.keys() + top = {k: handle.get_tensor(k).to(device) for k in keys if not k.startswith("blocks.")} + block = {k.split(".", 2)[2]: handle.get_tensor(k).to(device) for k in keys if k.startswith("blocks.0.")} + cond = T.step_conditioning(top, 0.5, device) + rope = T.build_rope(2, 4, 4, dim=4096, heads=T.HEADS, device=device) + tokens = torch.randn(1, 64, 128, device=device, dtype=torch.bfloat16) + x = T.block_forward(block, T.embed_tokens(top, tokens, 16), cond, top["prompt_context"], rope) + assert torch.isfinite(T.velocity(top, x, cond).float()).all() diff --git a/tests/test_main_validation.py b/tests/test_main_validation.py index 533841cd..dc2da31b 100644 --- a/tests/test_main_validation.py +++ b/tests/test_main_validation.py @@ -80,6 +80,23 @@ def test_ltx_accepts_segments(self, tmp_path): assert (config.restoration_model_name, config.ltx_seed) == ("ltx", 3) assert pipeline_cls.call_args.kwargs["segments"] == (SegmentRange(1, 2),) + def test_ltx_trial_needs_no_model_files_or_download(self, tmp_path): + with ( + patch("jasna.accelerator.is_nvidia_device", return_value=True), + patch("jasna.ltx.model_files.LtxModelFiles.from_dir", side_effect=AssertionError("reads model files")), + patch("jasna.ltx.model_files.missing_downloads", side_effect=AssertionError("offers a download")), + ): + pipeline_cls = _run_main_with_args( + tmp_path, ["--restoration-model-name", "ltx-undistilled", "--ltx-trial"], create_restoration=False + ) + + config = pipeline_cls.call_args.kwargs["config"] + assert (config.restoration_model_name, config.ltx_model, config.ltx_trial) == ("ltx", "undistilled", True) + + def test_ltx_trial_rejects_basicvsrpp(self, tmp_path): + with pytest.raises(SystemExit): + _run_main_with_args(tmp_path, ["--ltx-trial"]) + def test_ltx_rejects_streaming(self, tmp_path): with pytest.raises(SystemExit): _run_main_with_args(tmp_path, ["--stream", "--restoration-model-name", "ltx"]) From f3615be9f355f6f4fb72a26cfbed56958dc2957a Mon Sep 17 00:00:00 2001 From: Kruk2 Date: Tue, 29 Sep 2026 19:36:37 +0200 Subject: [PATCH 46/48] fix(ltx): print the trial notice in the CLI, where warnings are hidden by default --- jasna/ltx/model_files.py | 5 +++++ jasna/main.py | 7 +++++-- jasna/session_factory.py | 7 +++---- tests/test_main_validation.py | 3 ++- 4 files changed, 15 insertions(+), 7 deletions(-) diff --git a/jasna/ltx/model_files.py b/jasna/ltx/model_files.py index ebb2807e..458b6fc7 100644 --- a/jasna/ltx/model_files.py +++ b/jasna/ltx/model_files.py @@ -68,6 +68,11 @@ class PlaceholderTensors: model: LtxModelName +LTX_TRIAL_NOTICE = ( + "LTX trial: placeholder weights, so the output is not restored. Timings match a real run " + "except model loading, which skips reading and decrypting the model files." +) + # Chosen so a random INT8 / FP4 weight comes out near 1/sqrt(4096), like a trained one. _PLACEHOLDER_SCALES = {".scales": 2e-4, ".weight_scale_2": 5e-3, ".weight_scale": 1.0} _TORCH_DTYPES = {"BF16": "bfloat16", "F32": "float32", "I8": "int8", "U8": "uint8", "F8_E4M3": "float8_e4m3fn"} diff --git a/jasna/main.py b/jasna/main.py index b4ddcc8b..1ce2a130 100644 --- a/jasna/main.py +++ b/jasna/main.py @@ -781,12 +781,15 @@ def _run_videos( def _ensure_ltx_model(config: SessionConfig, segments) -> None: """Offer to download the selected LTX model when the job needs it and it is missing. Without a terminal to ask in, fail with what to download instead of waiting.""" - from jasna.ltx.model_files import download_files, download_size_text, missing_downloads + from jasna.ltx.model_files import LTX_TRIAL_NOTICE, download_files, download_size_text, missing_downloads uses_ltx = config.restoration_model_name == "ltx" or any( segment.restoration is not None and segment.restoration.model == "ltx" for segment in segments or () ) - if not uses_ltx or config.ltx_trial: + if not uses_ltx: + return + if config.ltx_trial: + print(LTX_TRIAL_NOTICE) return directory = config.restoration_model_path if config.restoration_model_name == "ltx" else default_restoration_model_path("ltx") missing = missing_downloads(directory, config.ltx_model, fast=config.ltx_fast) diff --git a/jasna/session_factory.py b/jasna/session_factory.py index ffba2d00..08fd3e43 100644 --- a/jasna/session_factory.py +++ b/jasna/session_factory.py @@ -245,10 +245,9 @@ def _ltx_model_files( if config.ltx_fast and torch.cuda.get_device_capability(device)[0] < 10: raise ValueError("The fast LTX model needs an RTX 50-series (Blackwell) GPU") if config.ltx_trial: - logger.warning( - "LTX trial: placeholder weights, so the output is not restored. Timings match a real run " - "except model loading, which skips reading and decrypting the model files." - ) + from jasna.ltx.model_files import LTX_TRIAL_NOTICE + + logger.warning(LTX_TRIAL_NOTICE) files = LtxModelFiles.placeholder(config.ltx_model, fast=config.ltx_fast) else: files = LtxModelFiles.from_dir(_restoration_model_path(config, "ltx"), config.ltx_model, fast=config.ltx_fast) diff --git a/tests/test_main_validation.py b/tests/test_main_validation.py index dc2da31b..f810960d 100644 --- a/tests/test_main_validation.py +++ b/tests/test_main_validation.py @@ -80,7 +80,7 @@ def test_ltx_accepts_segments(self, tmp_path): assert (config.restoration_model_name, config.ltx_seed) == ("ltx", 3) assert pipeline_cls.call_args.kwargs["segments"] == (SegmentRange(1, 2),) - def test_ltx_trial_needs_no_model_files_or_download(self, tmp_path): + def test_ltx_trial_needs_no_model_files_or_download(self, tmp_path, capsys): with ( patch("jasna.accelerator.is_nvidia_device", return_value=True), patch("jasna.ltx.model_files.LtxModelFiles.from_dir", side_effect=AssertionError("reads model files")), @@ -92,6 +92,7 @@ def test_ltx_trial_needs_no_model_files_or_download(self, tmp_path): config = pipeline_cls.call_args.kwargs["config"] assert (config.restoration_model_name, config.ltx_model, config.ltx_trial) == ("ltx", "undistilled", True) + assert "LTX trial: placeholder weights" in capsys.readouterr().out def test_ltx_trial_rejects_basicvsrpp(self, tmp_path): with pytest.raises(SystemExit): From 81dc8b053fb317c063390daab1dab8289c2094df Mon Sep 17 00:00:00 2001 From: Kruk2 Date: Tue, 29 Sep 2026 20:15:33 +0200 Subject: [PATCH 47/48] feat(gui): add an LTX trial run that needs no license or model files A "Trial run" switch runs LTX with placeholder weights to test speed. It is forced on when no LTX model is installed and offered at Start when a license is missing. --- jasna/gui/app.py | 19 ++ jasna/gui/control_bar.py | 8 +- jasna/gui/locales/en.py | 9 + jasna/gui/locales/ja.py | 9 + jasna/gui/locales/ko.py | 9 + jasna/gui/locales/th.py | 9 + jasna/gui/locales/zh.py | 9 + jasna/gui/ltx_models.py | 33 +++- jasna/gui/models.py | 1 + jasna/gui/segment_editor.py | 7 +- jasna/gui/settings_panel.py | 7 +- .../settings_sections/restoration_model.py | 57 +++++- jasna/gui/video_session.py | 3 +- tests/test_gui_ltx_models.py | 170 +++++++++++++++++- tests/test_gui_settings_sections.py | 12 +- 15 files changed, 327 insertions(+), 35 deletions(-) diff --git a/jasna/gui/app.py b/jasna/gui/app.py index 6a319c66..a7d6b760 100644 --- a/jasna/gui/app.py +++ b/jasna/gui/app.py @@ -613,6 +613,20 @@ def _on_start(self): and self._ltx_models.ensure(settings.ltx_model, settings.ltx_fast, on_ready=self._on_start) ): return + if ( + not settings.ltx_trial + and self._settings_panel.ltx_unavailable_reason() is None + and any_job_uses_ltx(jobs, settings) + ): + from jasna.gui.ltx_models import license_missing + + if license_missing(self._ltx_models.directory, settings.ltx_model, settings.ltx_fast): + from tkinter import messagebox + + if messagebox.askyesno(t("ltx_license_trial_title"), t("ltx_license_trial_confirm")): + self._settings_panel.set_ltx_trial(True) + self._on_start() + return errors = validate_gui_start( settings, jobs, ltx_available=self._settings_panel.ltx_unavailable_reason() is None ) @@ -651,6 +665,11 @@ def _engine_name(key: str) -> str: except Exception as e: self._log_panel.warning(f"Engine preflight warning failed: {e}") + trial = settings.ltx_trial and any_job_uses_ltx(jobs, settings) + if trial: + self._show_toast(t("toast_ltx_trial"), "warning") + self._log_panel.warning(t("toast_ltx_trial")) + self._control_bar.set_trial(trial) self._queue_panel.reset_jobs_for_run() self._status_pill.set_status("PROCESSING", Colors.STATUS_PROCESSING) diff --git a/jasna/gui/control_bar.py b/jasna/gui/control_bar.py index afc44349..1241f47b 100644 --- a/jasna/gui/control_bar.py +++ b/jasna/gui/control_bar.py @@ -197,6 +197,7 @@ def __init__(self, master, **kwargs): self._on_stop: callable = None self._on_toggle_logs: callable = None self._start_disabled_tooltip = None + self._trial = False self._build_controls() @@ -372,6 +373,10 @@ def set_start_enabled(self, enabled: bool, disabled_tooltip: str = ""): if disabled_tooltip: self._start_disabled_tooltip = Tooltip(self._start_btn, disabled_tooltip) + def set_trial(self, trial: bool) -> None: + """Mark the LTX stages of this run as a trial.""" + self._trial = trial + def set_running(self, running: bool): if running: @@ -395,7 +400,8 @@ def update_progress( self._progress_bar.set(percent / 100.0) self._percent_label.configure(text=f"{int(percent)}%") if stage: - self._fps_label.configure(text=t(f"ltx_stage_{stage}")) + label = t(f"ltx_stage_{stage}") + self._fps_label.configure(text=t("ltx_trial_stage", stage=label) if self._trial else label) else: self._fps_label.configure(text=f"FPS: {fps:.1f}" if fps > 0 else "FPS: --") diff --git a/jasna/gui/locales/en.py b/jasna/gui/locales/en.py index a6475fca..005be2d1 100644 --- a/jasna/gui/locales/en.py +++ b/jasna/gui/locales/en.py @@ -239,6 +239,14 @@ "ltx_seed": "Seed", "ltx_fast": "Fast mode", "ltx_large_canvas": "Sharper large mosaics", + "ltx_trial": "Trial run (no license)", + "ltx_trial_notice": "Trial: videos will not be restored.", + "model_ltx_trial_only": "No LTX model installed, so only a trial run is possible.", + "ltx_license_trial_title": "LTX license", + "ltx_license_trial_confirm": "LTX needs a license. Run a trial instead? The picture will look wrong.", + "toast_ltx_trial": "Trial run: output will not be restored", + "ltx_trial_stage": "{stage} (trial)", + "segments_seed_tried_trial": "Tried (trial):", "ltx_stage_scan": "Finding mosaics", "ltx_stage_encode": "Preparing", "ltx_stage_denoise": "Restoring", @@ -448,6 +456,7 @@ "tip_ltx_new_seed": "Try a new random seed", "tip_ltx_fast": "About 1.4x faster with slightly less detail.\nRTX 50-series only.", "tip_ltx_large_canvas": "Restores big mosaics in more detail, but about 3x slower for them.\nNeeds at least 10 GB of free video memory.", + "tip_ltx_trial": "Speed test without a license or model files. The picture will look wrong.\nUse it to see if your PC can run LTX and how fast.", "tip_model_basicvsrpp": "The usual model. Fast and good for most videos.\nPick it unless you want the most detail.", "tip_model_ltx": "Restores the most detail, but is much slower.\nNeeds a strong NVIDIA GPU. Good for short clips or favorite scenes.", "tip_model_ltx_undistilled": "LTX with the undistilled model: slightly better detail, but slower.", diff --git a/jasna/gui/locales/ja.py b/jasna/gui/locales/ja.py index 7345fcac..f6ea3fc8 100644 --- a/jasna/gui/locales/ja.py +++ b/jasna/gui/locales/ja.py @@ -239,6 +239,14 @@ "ltx_seed": "シード", "ltx_fast": "高速モード", "ltx_large_canvas": "大きなモザイクをより鮮明に", + "ltx_trial": "お試し実行（ライセンス不要）", + "ltx_trial_notice": "お試し：動画は修復されません。", + "model_ltx_trial_only": "LTX モデルが入っていないため、お試し実行のみできます。", + "ltx_license_trial_title": "LTX ライセンス", + "ltx_license_trial_confirm": "LTX にはライセンスが必要です。代わりにお試し実行しますか？映像は正しく修復されません。", + "toast_ltx_trial": "お試し実行：出力は修復されません", + "ltx_trial_stage": "{stage}（お試し）", + "segments_seed_tried_trial": "試したシード（お試し）：", "ltx_stage_scan": "モザイクを検出中", "ltx_stage_encode": "準備中", "ltx_stage_denoise": "修復中", @@ -448,6 +456,7 @@ "tip_ltx_new_seed": "新しいランダムなシードを試す", "tip_ltx_fast": "約 1.4 倍速くなりますが、細部はやや減ります。\nRTX 50 シリーズのみ。", "tip_ltx_large_canvas": "大きなモザイクをより細かく修復しますが、その部分は約 3 倍遅くなります。\n10 GB 以上の空き VRAM が必要です。", + "tip_ltx_trial": "ライセンスもモデルファイルも不要の速度テストです。映像は正しく修復されません。\nこの PC で LTX が動くか、どのくらい速いかを確かめるときに使います。", "tip_model_basicvsrpp": "標準のモデル。速く、ほとんどの動画に向いています。\n最高の精細さが必要なければこれを選んでください。", "tip_model_ltx": "最も精細に修復しますが、かなり遅くなります。\n高性能な NVIDIA GPU が必要です。短いクリップやお気に入りのシーン向けです。", "tip_model_ltx_undistilled": "非蒸留モデルの LTX：少し精細ですが、遅くなります。", diff --git a/jasna/gui/locales/ko.py b/jasna/gui/locales/ko.py index 345e2ab9..d5d78cd5 100644 --- a/jasna/gui/locales/ko.py +++ b/jasna/gui/locales/ko.py @@ -238,6 +238,14 @@ "ltx_seed": "시드", "ltx_fast": "빠른 모드", "ltx_large_canvas": "큰 모자이크 더 선명하게", + "ltx_trial": "체험 실행(라이선스 불필요)", + "ltx_trial_notice": "체험: 영상이 복원되지 않습니다.", + "model_ltx_trial_only": "LTX 모델이 설치되지 않아 체험 실행만 할 수 있습니다.", + "ltx_license_trial_title": "LTX 라이선스", + "ltx_license_trial_confirm": "LTX에는 라이선스가 필요합니다. 대신 체험 실행할까요? 화면은 제대로 복원되지 않습니다.", + "toast_ltx_trial": "체험 실행: 출력이 복원되지 않습니다", + "ltx_trial_stage": "{stage} (체험)", + "segments_seed_tried_trial": "시험한 시드(체험):", "ltx_stage_scan": "모자이크 찾는 중", "ltx_stage_encode": "준비 중", "ltx_stage_denoise": "복원 중", @@ -428,6 +436,7 @@ "tip_ltx_new_seed": "새 무작위 시드 사용", "tip_ltx_fast": "약 1.4배 빠르지만 디테일이 조금 줄어듭니다.\nRTX 50 시리즈 전용.", "tip_ltx_large_canvas": "큰 모자이크를 더 세밀하게 복원하지만 해당 부분은 약 3배 느려집니다.\n여유 VRAM이 10 GB 이상 필요합니다.", + "tip_ltx_trial": "라이선스나 모델 파일 없이 하는 속도 테스트입니다. 화면은 제대로 복원되지 않습니다.\n이 PC에서 LTX가 돌아가는지, 얼마나 빠른지 확인할 때 쓰세요.", "tip_model_basicvsrpp": "기본 모델입니다. 빠르고 대부분의 영상에 알맞습니다.\n최고 디테일이 필요 없다면 이것을 고르세요.", "tip_model_ltx": "디테일을 가장 많이 살리지만 훨씬 느립니다.\n강력한 NVIDIA GPU가 필요합니다. 짧은 클립이나 좋아하는 장면에 좋습니다.", "tip_model_ltx_undistilled": "비증류 모델을 쓰는 LTX: 디테일이 조금 더 좋지만 느립니다.", diff --git a/jasna/gui/locales/th.py b/jasna/gui/locales/th.py index e5badfce..e825d5d2 100644 --- a/jasna/gui/locales/th.py +++ b/jasna/gui/locales/th.py @@ -239,6 +239,14 @@ "ltx_seed": "ซีด", "ltx_fast": "โหมดเร็ว", "ltx_large_canvas": "โมเสกขนาดใหญ่คมชัดขึ้น", + "ltx_trial": "ทดลองรัน (ไม่ต้องใช้ไลเซนส์)", + "ltx_trial_notice": "ทดลอง: วิดีโอจะไม่ถูกฟื้นฟู", + "model_ltx_trial_only": "ไม่ได้ติดตั้งโมเดล LTX จึงทำได้แค่ทดลองรัน", + "ltx_license_trial_title": "ไลเซนส์ LTX", + "ltx_license_trial_confirm": "LTX ต้องใช้ไลเซนส์ จะทดลองรันแทนไหม? ภาพที่ได้จะไม่ถูกต้อง", + "toast_ltx_trial": "ทดลองรัน: ผลลัพธ์จะไม่ถูกฟื้นฟู", + "ltx_trial_stage": "{stage} (ทดลอง)", + "segments_seed_tried_trial": "ที่ลองแล้ว (ทดลอง):", "ltx_stage_scan": "กำลังค้นหาโมเสก", "ltx_stage_encode": "กำลังเตรียม", "ltx_stage_denoise": "กำลังฟื้นฟู", @@ -430,6 +438,7 @@ "tip_ltx_new_seed": "ลองซีดสุ่มใหม่", "tip_ltx_fast": "เร็วขึ้นประมาณ 1.4 เท่า รายละเอียดลดลงเล็กน้อย\nเฉพาะ RTX ซีรีส์ 50", "tip_ltx_large_canvas": "ฟื้นฟูโมเสกขนาดใหญ่ได้ละเอียดขึ้น แต่ส่วนนั้นช้าลงประมาณ 3 เท่า\nต้องมี VRAM ว่างอย่างน้อย 10 GB", + "tip_ltx_trial": "ทดสอบความเร็วโดยไม่ต้องมีไลเซนส์หรือไฟล์โมเดล ภาพที่ได้จะไม่ถูกต้อง\nใช้เพื่อดูว่าเครื่องของคุณรัน LTX ได้ไหมและเร็วแค่ไหน", "tip_model_basicvsrpp": "โมเดลปกติ เร็วและเหมาะกับวิดีโอส่วนใหญ่\nเลือกอันนี้ถ้าไม่ต้องการรายละเอียดสูงสุด", "tip_model_ltx": "กู้รายละเอียดได้มากที่สุด แต่ช้ากว่ามาก\nต้องใช้ GPU NVIDIA แรงๆ เหมาะกับคลิปสั้นหรือฉากโปรด", "tip_model_ltx_undistilled": "LTX แบบไม่กลั่น: รายละเอียดดีขึ้นเล็กน้อย แต่ช้ากว่า", diff --git a/jasna/gui/locales/zh.py b/jasna/gui/locales/zh.py index ad15dd63..a6cb29f7 100644 --- a/jasna/gui/locales/zh.py +++ b/jasna/gui/locales/zh.py @@ -239,6 +239,14 @@ "ltx_seed": "种子", "ltx_fast": "快速模式", "ltx_large_canvas": "大马赛克更清晰", + "ltx_trial": "试运行（无需许可证）", + "ltx_trial_notice": "试运行：视频不会被修复。", + "model_ltx_trial_only": "未安装 LTX 模型，只能试运行。", + "ltx_license_trial_title": "LTX 许可证", + "ltx_license_trial_confirm": "LTX 需要许可证。改为试运行吗？画面会不正确。", + "toast_ltx_trial": "试运行：输出不会被修复", + "ltx_trial_stage": "{stage}（试运行）", + "segments_seed_tried_trial": "已试（试运行）：", "ltx_stage_scan": "正在查找马赛克", "ltx_stage_encode": "正在准备", "ltx_stage_denoise": "正在修复", @@ -448,6 +456,7 @@ "tip_ltx_new_seed": "换一个随机种子", "tip_ltx_fast": "约快 1.4 倍，细节略少。\n仅限 RTX 50 系列。", "tip_ltx_large_canvas": "大马赛克修复得更细致，但这些部分约慢 3 倍。\n需要至少 10 GB 可用显存。", + "tip_ltx_trial": "无需许可证或模型文件的速度测试。画面会不正确。\n用它来看你的电脑能否运行 LTX 以及速度如何。", "tip_model_basicvsrpp": "常用模型。速度快，适合大多数视频。\n除非需要最高细节，否则选它。", "tip_model_ltx": "修复细节最多，但慢得多。\n需要强劲的 NVIDIA GPU。适合短片或喜欢的场景。", "tip_model_ltx_undistilled": "使用未蒸馏模型的 LTX：细节稍好，但更慢。", diff --git a/jasna/gui/ltx_models.py b/jasna/gui/ltx_models.py index 826dd537..5c5dcb9e 100644 --- a/jasna/gui/ltx_models.py +++ b/jasna/gui/ltx_models.py @@ -49,26 +49,40 @@ def read_install_state(directory: Path) -> LtxInstallState: ) -def card_unavailable_reason(state: LtxInstallState, *, nvidia: bool | None) -> str | None: - """Locale key saying why the LTX card cannot be picked; ``nvidia`` is None until the GPU is known.""" - if nvidia is False: - return "model_ltx_needs_nvidia" - if not any(state.usable(model, fast) for model in LTX_MODELS for fast in (False, True)): - return "model_ltx_not_installed" - return None +def card_unavailable_reason(*, nvidia: bool | None) -> str | None: + """Locale key saying why the LTX card cannot be picked; ``nvidia`` is None until the GPU is known. + Without model files LTX still runs as a trial.""" + return "model_ltx_needs_nvidia" if nvidia is False else None + + +def trial_only(state: LtxInstallState) -> bool: + """No LTX model is installed or downloadable, so only a trial run is possible.""" + return not any(state.usable(model, fast) for model in LTX_MODELS for fast in (False, True)) def run_unavailable_reason( - state: LtxInstallState, model: LtxModelName, fast: bool, *, nvidia: bool | None + state: LtxInstallState, model: LtxModelName, fast: bool, *, nvidia: bool | None, trial: bool ) -> str | None: """Locale key saying why LTX cannot restore right now with this variant.""" if nvidia is False: return "model_ltx_needs_nvidia" - if state.installed(model, fast): + if trial or state.installed(model, fast): return None return "model_ltx_not_downloaded" if state.downloadable else "model_ltx_not_installed" +def license_missing(directory: Path, model: LtxModelName, fast: bool) -> bool: + """Whether the installed ``model`` needs a license this PC lacks. Verifies the stored + license only; reads no model data.""" + from jasna.protection import ProtectionError + + try: + model_files.LtxModelFiles.from_dir(directory, model, fast=fast) + except ProtectionError: + return True + return False + + class LtxModels: """The install state plus the single running download. UI glue: it asks with a message box and runs the download on a worker thread, reporting back on the main thread.""" @@ -77,6 +91,7 @@ def __init__(self, directory: Path, main_thread: MainThreadCalls, on_change: Cal self._directory = directory self._main_thread = main_thread self._on_change = on_change + self.directory = directory self.state = read_install_state(directory) self.percent: int | None = None diff --git a/jasna/gui/models.py b/jasna/gui/models.py index 4b9ed0f7..22ed1c5b 100644 --- a/jasna/gui/models.py +++ b/jasna/gui/models.py @@ -110,6 +110,7 @@ class AppSettings: ltx_seed: int = LTX_DEFAULT_SEED ltx_fast: bool = False ltx_large_canvas: bool = False + ltx_trial: bool = False # Basic processing batch_size: int = 4 diff --git a/jasna/gui/segment_editor.py b/jasna/gui/segment_editor.py index 54d32493..f5d6a919 100644 --- a/jasna/gui/segment_editor.py +++ b/jasna/gui/segment_editor.py @@ -857,7 +857,7 @@ def _render_seed_chips(self, selected: SegmentRange) -> None: self._seed_chips.grid(row=2, column=0, sticky="ew", pady=(4, 0)) label = ctk.CTkLabel( self._seed_chips, - text=t("segments_seed_tried"), + text=t("segments_seed_tried_trial" if self._current_video_settings().ltx_trial else "segments_seed_tried"), font=(Fonts.FAMILY, Fonts.SIZE_TINY), text_color=Colors.STATUS_PENDING, ) @@ -962,7 +962,10 @@ def _handle_seed_event(self, event) -> None: return if isinstance(event, SeedProgress): self._seed_progress.set(event.fraction) - values = dict(seed=event.seed, stage=t(f"ltx_stage_{event.stage}"), percent=event.fraction * 100) + stage = t(f"ltx_stage_{event.stage}") + if self._current_video_settings().ltx_trial: + stage = t("ltx_trial_stage", stage=stage) + values = dict(seed=event.seed, stage=stage, percent=event.fraction * 100) text = ( t("segments_seed_progress", eta=format_duration(event.eta_seconds), **values) if event.eta_seconds > 0 diff --git a/jasna/gui/settings_panel.py b/jasna/gui/settings_panel.py index 15b2835d..f78d5e41 100644 --- a/jasna/gui/settings_panel.py +++ b/jasna/gui/settings_panel.py @@ -253,6 +253,10 @@ def ltx_unavailable_reason(self) -> str | None: """Locale key saying why the chosen LTX model cannot run now, or None when it can.""" return self._model_section.unavailable_reason() + def set_ltx_trial(self, trial: bool): + self._model_section.set_trial(trial) + self._update_modified_indicator() + def refresh_ltx_models(self): self._model_section.refresh() self._update_modified_indicator() @@ -384,8 +388,6 @@ def set_enabled(self, enabled: bool): """Enable or disable the settings; the queue-wide post-export action stays editable while processing.""" state = "normal" if enabled else "disabled" - self._model_section.set_enabled(enabled) - # Preset bar buttons self._preset_dropdown.configure(state=state) self._create_btn.configure(state=state) @@ -401,3 +403,4 @@ def set_enabled(self, enabled: bool): widget.configure(state=state) except Exception: logger.debug("Widget %r does not support state=%s", key, state, exc_info=True) + self._model_section.set_enabled(enabled) diff --git a/jasna/gui/settings_sections/restoration_model.py b/jasna/gui/settings_sections/restoration_model.py index 3d9e9b1e..ea1301e1 100644 --- a/jasna/gui/settings_sections/restoration_model.py +++ b/jasna/gui/settings_sections/restoration_model.py @@ -8,7 +8,7 @@ from jasna.gui.components import CollapsibleSection, Tooltip from jasna.gui.icons import CompactSwitch, NativeIconButton from jasna.gui.locales import t -from jasna.gui.ltx_models import LtxModels, card_unavailable_reason, run_unavailable_reason +from jasna.gui.ltx_models import LtxModels, card_unavailable_reason, run_unavailable_reason, trial_only from jasna.gui.settings_sections.widgets import ValueOptionMenu, add_setting_label, get_tooltip from jasna.gui.theme import Colors, Fonts, Sizing from jasna.ltx.model_files import LTX_MODELS @@ -167,6 +167,17 @@ def __init__(self, parent, widgets: dict, on_modified, on_model_changed, *, ltx_ Tooltip(self._widgets["ltx_large_canvas"], get_tooltip("ltx_large_canvas")) self._large_row = large_row + trial_row = ctk.CTkFrame(self._ltx_options, fg_color="transparent") + trial_row.pack(fill="x") + add_setting_label(trial_row, "ltx_trial", in_card=True) + self._widgets["ltx_trial"] = CompactSwitch(trial_row, self._on_trial_toggled, Colors.BG_CARD) + self._widgets["ltx_trial"].pack(side="right", padx=12, pady=8) + Tooltip(self._widgets["ltx_trial"], get_tooltip("ltx_trial")) + self._trial_status = ctk.CTkLabel( + self._ltx_options, text="", text_color=Colors.STATUS_WARNING, font=(Fonts.FAMILY, Fonts.SIZE_SMALL), anchor="w" + ) + self._trial_row = trial_row + self._notice = ctk.CTkLabel( self._ltx_options, text="", text_color=Colors.STATUS_WARNING, font=(Fonts.FAMILY, Fonts.SIZE_SMALL), anchor="w" ) @@ -182,6 +193,16 @@ def _set_fast(self, fast: bool) -> None: else: self._widgets["ltx_fast"].deselect() + def _trial(self) -> bool: + return self._widgets["ltx_trial"].get() == 1 + + def set_trial(self, trial: bool) -> None: + if trial: + self._widgets["ltx_trial"].select() + else: + self._widgets["ltx_trial"].deselect() + self._refresh_widgets() + def _set_ltx_model(self, model: str) -> None: self._ltx_model.set_value(model) self._accepted_ltx_model = self._ltx_model.get_value() @@ -198,7 +219,9 @@ def _ensure_variant(self, model: str, fast: bool) -> bool: installed variant with a notice. False when the user declined or nothing is installed.""" self._show_notice(None) state = self._ltx_models.state - if state.installed(model, fast): + if trial_only(state): + self.set_trial(True) + if self._trial() or state.installed(model, fast): return True if state.downloadable: return self._ltx_models.ensure(model, fast, on_ready=lambda: None) @@ -235,6 +258,12 @@ def _on_fast_toggled(self) -> None: self._refresh_widgets() self._on_modified() + def _on_trial_toggled(self) -> None: + if not self._trial() and not self._ensure_variant(self._ltx_model.get_value(), self._fast()): + self.set_trial(True) + self._refresh_widgets() + self._on_modified() + def _show_model(self, value: str) -> None: self._model.set(value) for name, card in self._cards.items(): @@ -257,13 +286,13 @@ def _set_seed(self, seed: int) -> None: def unavailable_reason(self) -> str | None: """Why the chosen LTX model cannot restore right now (None when it can).""" return run_unavailable_reason( - self._ltx_models.state, self._ltx_model.get_value(), self._fast(), nvidia=self._nvidia + self._ltx_models.state, self._ltx_model.get_value(), self._fast(), nvidia=self._nvidia, trial=self._trial() ) def _refresh_widgets(self) -> None: state = self._ltx_models.state editable = self._enabled and not self._ltx_models.downloading - reason = card_unavailable_reason(state, nvidia=self._nvidia) + reason = card_unavailable_reason(nvidia=self._nvidia) self._cards["basicvsrpp"].set_enabled(editable) ltx_card = self._cards["ltx"] ltx_card.set_enabled(editable and reason is None) @@ -273,8 +302,17 @@ def _refresh_widgets(self) -> None: ) self._ltx_model.configure(state="normal" if editable else "disabled") self._new_seed_btn.configure(state="normal" if self._enabled else "disabled") + locked = trial_only(state) + if locked: + self._widgets["ltx_trial"].select() + self._widgets["ltx_trial"].configure(state="normal" if editable and not locked else "disabled") + if self._trial(): + self._trial_status.configure(text=t("model_ltx_trial_only" if locked else "ltx_trial_notice")) + self._trial_status.pack(fill="x", padx=12, pady=(0, 8), after=self._trial_row) + else: + self._trial_status.pack_forget() model, fast = self._ltx_model.get_value(), self._fast() - if not state.installed(model, fast) and state.downloadable: + if not self._trial() and not state.installed(model, fast) and state.downloadable: self._install_status.configure(text=t("model_ltx_download_needed", size=state.download_size(model, fast))) self._install_status.pack(fill="x", padx=12, pady=(0, 4), after=self._model_row) else: @@ -284,10 +322,10 @@ def _usable_model(self) -> str: """The chosen model, or BasicVSR++ when LTX cannot run; switches to the installed precision when only that one is there.""" state = self._ltx_models.state - if self._model.get() != "ltx" or card_unavailable_reason(state, nvidia=self._nvidia) is not None: + if self._model.get() != "ltx" or card_unavailable_reason(nvidia=self._nvidia) is not None: return "basicvsrpp" model, fast = self._ltx_model.get_value(), self._fast() - if not state.usable(model, fast) and state.installed(model, not fast): + if not self._trial() and not state.usable(model, fast) and state.installed(model, not fast): self._set_fast(not fast) return "ltx" @@ -329,6 +367,10 @@ def apply(self, preset) -> None: self._widgets["ltx_large_canvas"].select() else: self._widgets["ltx_large_canvas"].deselect() + if preset.ltx_trial: + self._widgets["ltx_trial"].select() + else: + self._widgets["ltx_trial"].deselect() self._set_ltx_model(preset.ltx_model) self._model.set(preset.restoration_model if preset.restoration_model in _MODELS else "basicvsrpp") self._show_notice(None) @@ -342,4 +384,5 @@ def collect(self) -> dict: "ltx_seed": parse_seed(self._widgets["ltx_seed"].get()), "ltx_fast": self._widgets["ltx_fast"].get() == 1, "ltx_large_canvas": self._widgets["ltx_large_canvas"].get() == 1, + "ltx_trial": self._widgets["ltx_trial"].get() == 1, } diff --git a/jasna/gui/video_session.py b/jasna/gui/video_session.py index 5bfcf0ed..38ef7d89 100644 --- a/jasna/gui/video_session.py +++ b/jasna/gui/video_session.py @@ -31,6 +31,7 @@ def video_session_key(settings: AppSettings) -> tuple: settings.restoration_model, settings.ltx_model, settings.ltx_fast, + settings.ltx_trial, ) if settings.secondary_restoration == "tvai": key += ( @@ -79,7 +80,7 @@ def video_session_config( ltx_seed=int(settings.ltx_seed), ltx_fast=bool(settings.ltx_fast), ltx_model=settings.ltx_model, - ltx_trial=False, + ltx_trial=bool(settings.ltx_trial), compile_basicvsrpp=bool(settings.compile_basicvsrpp), max_clip_size=int(settings.max_clip_size), temporal_overlap=int(settings.temporal_overlap), diff --git a/tests/test_gui_ltx_models.py b/tests/test_gui_ltx_models.py index 7631887c..84a77f3c 100644 --- a/tests/test_gui_ltx_models.py +++ b/tests/test_gui_ltx_models.py @@ -18,6 +18,7 @@ card_unavailable_reason, read_install_state, run_unavailable_reason, + trial_only, ) from jasna.gui.models import AppSettings, JobItem, PresetManager from jasna.ltx import model_files @@ -53,20 +54,22 @@ def test_availability_without_a_download_location(tmp_path) -> None: assert not state.downloadable assert state.installed("distilled", False) and not state.installed("distilled", True) - assert card_unavailable_reason(state, nvidia=True) is None - assert card_unavailable_reason(state, nvidia=False) == "model_ltx_needs_nvidia" - assert card_unavailable_reason(read_install_state(tmp_path / "empty"), nvidia=True) == "model_ltx_not_installed" - assert run_unavailable_reason(state, "undistilled", False, nvidia=True) == "model_ltx_not_installed" - assert run_unavailable_reason(state, "distilled", False, nvidia=None) is None - assert run_unavailable_reason(state, "distilled", True, nvidia=True) == "model_ltx_not_installed" + assert card_unavailable_reason(nvidia=True) is None + assert card_unavailable_reason(nvidia=False) == "model_ltx_needs_nvidia" + assert not trial_only(state) and trial_only(read_install_state(tmp_path / "empty")) + assert run_unavailable_reason(state, "undistilled", False, nvidia=True, trial=False) == "model_ltx_not_installed" + assert run_unavailable_reason(state, "undistilled", False, nvidia=True, trial=True) is None + assert run_unavailable_reason(state, "distilled", False, nvidia=None, trial=False) is None + assert run_unavailable_reason(state, "distilled", True, nvidia=True, trial=False) == "model_ltx_not_installed" + assert run_unavailable_reason(state, "distilled", True, nvidia=False, trial=True) == "model_ltx_needs_nvidia" def test_availability_with_a_download_location(tmp_path, downloadable) -> None: state = read_install_state(tmp_path) assert state.downloadable - assert card_unavailable_reason(state, nvidia=True) is None - assert run_unavailable_reason(state, "undistilled", False, nvidia=True) == "model_ltx_not_downloaded" + assert not trial_only(state) + assert run_unavailable_reason(state, "undistilled", False, nvidia=True, trial=False) == "model_ltx_not_downloaded" assert state.download_size("distilled", False) == "14.9 GB" @@ -263,7 +266,7 @@ def record(self, widget, text, original=original): assert section._install_status in tips[t("tip_ltx_download")] assert section._new_seed_btn in tips[t("tip_ltx_new_seed")] assert section._download_bar in tips[t("tip_ltx_download")] - for key in ("ltx_seed", "ltx_fast", "ltx_large_canvas"): + for key in ("ltx_seed", "ltx_fast", "ltx_large_canvas", "ltx_trial"): assert section._widgets[key] in tips[t(f"tip_{key}")] @@ -287,3 +290,152 @@ def test_start_downloads_a_missing_ltx_model_first(monkeypatch) -> None: JasnaApp._on_start(app) assert ensured == [("undistilled", False)] + + +def test_the_trial_switch_round_trips_and_skips_the_download(tk_root, tmp_path, monkeypatch, downloadable) -> None: + asked = _answer(monkeypatch, yes=True) + panel = _panel(tk_root, ltx_models(tk_root, tmp_path / "ltx", installed=False)) + section = panel._model_section + trial = replace(AppSettings(), restoration_model="ltx", ltx_model="undistilled", ltx_trial=True) + + section.apply(trial) + + assert panel.get_settings().ltx_trial is True + assert section.unavailable_reason() is None + assert not section._install_status.winfo_manager() + assert section._trial_status.cget("text") == t("ltx_trial_notice") + _pick_ltx_model(section, "distilled") + assert asked == [] and panel.get_settings().ltx_model == "distilled" + + +def test_turning_the_trial_off_asks_for_the_missing_model(tk_root, tmp_path, monkeypatch, downloadable) -> None: + asked = _answer(monkeypatch, yes=False) + panel = _panel(tk_root, ltx_models(tk_root, tmp_path / "ltx", installed=False)) + section = panel._model_section + section.apply(replace(AppSettings(), restoration_model="ltx", ltx_trial=True)) + + section._widgets["ltx_trial"]._toggle() + + assert len(asked) == 1 + assert panel.get_settings().ltx_trial is True + + +def test_without_any_ltx_model_only_a_locked_trial_is_possible(tk_root, tmp_path) -> None: + panel = _panel(tk_root, ltx_models(tk_root, tmp_path / "ltx", installed=False)) + section = panel._model_section + + section._cards["ltx"].radio.invoke() + + settings = panel.get_settings() + assert (settings.restoration_model, settings.ltx_trial) == ("ltx", True) + assert section._widgets["ltx_trial"].cget("cursor") == "" + section._widgets["ltx_trial"]._toggle() + assert panel.get_settings().ltx_trial is True + assert section._trial_status.cget("text") == t("model_ltx_trial_only") + panel.set_enabled(False) + panel.set_enabled(True) + assert section._widgets["ltx_trial"].cget("cursor") == "" + + +def _start_app(settings: AppSettings, *, license_missing: bool, yes: bool, monkeypatch): + from jasna.gui.app import JasnaApp + + monkeypatch.setattr(ltx_models_module, "license_missing", lambda directory, model, fast: license_missing) + asked = _answer(monkeypatch, yes=yes) + trials: list[bool] = [] + restarted: list[bool] = [] + app = SimpleNamespace( + _preview_gpu_busy=False, + _processor=None, + _queue_panel=SimpleNamespace(get_jobs=lambda: [JobItem(Path("a.mp4"))], get_output_folder=lambda: "", get_output_pattern=lambda: ""), + _settings_panel=SimpleNamespace( + get_settings=lambda: settings, ltx_unavailable_reason=lambda: None, set_ltx_trial=trials.append + ), + _ltx_models=SimpleNamespace(directory=Path("ltx")), + _on_start=lambda: restarted.append(True), + ) + return app, asked, trials, restarted, JasnaApp._on_start + + +def test_start_offers_a_trial_when_the_license_is_missing(monkeypatch) -> None: + settings = replace(AppSettings(), restoration_model="ltx") + app, asked, trials, restarted, on_start = _start_app(settings, license_missing=True, yes=True, monkeypatch=monkeypatch) + + on_start(app) + + assert asked == [t("ltx_license_trial_confirm")] + assert trials == [True] and restarted == [True] + + +def test_declining_the_trial_offer_does_not_start(monkeypatch) -> None: + settings = replace(AppSettings(), restoration_model="ltx") + app, asked, trials, restarted, on_start = _start_app(settings, license_missing=True, yes=False, monkeypatch=monkeypatch) + + on_start(app) + + assert len(asked) == 1 and trials == [] and restarted == [] + + +def test_trial_session_config_and_key() -> None: + from jasna.gui.video_session import video_session_config, video_session_key + + real = replace(AppSettings(), restoration_model="ltx") + trial = replace(real, ltx_trial=True) + + assert video_session_config(trial, codec="hevc", encoder_settings={}).ltx_trial is True + assert video_session_config(real, codec="hevc", encoder_settings={}).ltx_trial is False + assert video_session_key(trial) != video_session_key(real) + + +def test_a_trial_seed_preview_never_reuses_a_real_prepared_range() -> None: + from jasna.gui.ltx_seed_preview import prepared_key + + segment = SegmentRange(1, 2, SegmentRestoration("ltx", 5)) + real = replace(AppSettings(), restoration_model="ltx") + + assert prepared_key(segment, real) != prepared_key(segment, replace(real, ltx_trial=True)) + + +def test_license_missing_only_for_encrypted_models(tmp_path, monkeypatch) -> None: + from jasna.protection import LicenseError + + directory = tmp_path / "ltx" + _install(directory, "distilled", fast=False) + assert not ltx_models_module.license_missing(directory, "distilled", False) + + def refuse(*_args, **_kwargs): + raise LicenseError("No license found.") + + monkeypatch.setattr(model_files.LtxModelFiles, "from_dir", refuse) + assert ltx_models_module.license_missing(directory, "distilled", False) + + +def test_start_with_a_trial_skips_the_license_offer(monkeypatch) -> None: + from jasna.gui import validation + + class Validated(Exception): + pass + + def validated(*_args, **_kwargs): + raise Validated + + settings = replace(AppSettings(), restoration_model="ltx", ltx_trial=True) + app, asked, trials, restarted, on_start = _start_app(settings, license_missing=True, yes=True, monkeypatch=monkeypatch) + monkeypatch.setattr(validation, "validate_gui_start", validated) + + with pytest.raises(Validated): + on_start(app) + + assert asked == [] and trials == [] and restarted == [] + + +def test_the_progress_stage_is_marked_as_a_trial(tk_root) -> None: + from jasna.gui.control_bar import ControlBar + + bar = ControlBar(tk_root) + bar.set_trial(True) + bar.update_progress(percent=10, stage="denoise") + assert bar._fps_label.cget("text") == t("ltx_trial_stage", stage=t("ltx_stage_denoise")) + bar.set_trial(False) + bar.update_progress(percent=10, stage="denoise") + assert bar._fps_label.cget("text") == t("ltx_stage_denoise") diff --git a/tests/test_gui_settings_sections.py b/tests/test_gui_settings_sections.py index acab2d4b..02193e53 100644 --- a/tests/test_gui_settings_sections.py +++ b/tests/test_gui_settings_sections.py @@ -70,6 +70,7 @@ def _fake_section_widgets() -> dict: "ltx_seed": _FakeWidget(" 123 "), "ltx_fast": _FakeWidget(1), "ltx_large_canvas": _FakeWidget(0), + "ltx_trial": _FakeWidget(1), "max_clip_size": _FakeWidget(90), "fp16_mode": _FakeWidget(1), "detection_model": _FakeWidget("rfdetr-v6"), @@ -156,6 +157,7 @@ def test_sections_collect_internal_values_without_translation_lookups() -> None: assert values["ltx_seed"] == 123 assert values["ltx_fast"] is True assert values["ltx_large_canvas"] is False + assert values["ltx_trial"] is True def test_sections_collect_covers_all_widget_backed_appsettings_fields() -> None: @@ -466,12 +468,13 @@ def test_ltx_hides_standard_only_settings_and_restores_them(_panel_factory) -> N assert order.index(secondary._section) == order.index(secondary._previous_section) + 1 -def test_ltx_card_disabled_when_model_missing(_panel_factory) -> None: +def test_ltx_without_model_files_runs_as_a_trial(_panel_factory) -> None: panel = _panel_factory(False) - assert panel._model_section._cards["ltx"].radio.cget("state") == "disabled" + assert panel._model_section._cards["ltx"].radio.cget("state") == "normal" panel._model_section.apply(AppSettings(restoration_model="ltx")) - assert panel.get_settings().restoration_model == "basicvsrpp" + settings = panel.get_settings() + assert (settings.restoration_model, settings.ltx_trial) == ("ltx", True) def test_gpu_support_disables_ltx_and_shows_fast_mode_on_blackwell(_panel_factory) -> None: @@ -503,7 +506,8 @@ def test_ltx_settings_round_trip_through_panel(_panel_factory) -> None: def test_disabling_settings_keeps_unavailable_ltx_disabled(_panel_factory) -> None: - panel = _panel_factory(False) + panel = _panel_factory(True) + panel._model_section.set_gpu_support(nvidia=False, blackwell=False) panel.set_enabled(False) panel.set_enabled(True) From 17b69b015b2cf8201897ae17c79ee9799fdedfff Mon Sep 17 00:00:00 2001 From: Codex Date: Sat, 3 Oct 2026 23:16:18 +0800 Subject: [PATCH 48/48] fix(windows): support ROCm imports and vendor-aware system checks Independent upstream-main feature with bilingual guides. --- docs/DETECTION_TEST_VENDOR_ISOLATION_CN.md | 18 +++ docs/MEDIA_TEST_ENVIRONMENT_ISOLATION_CN.md | 21 ++++ docs/WINDOWS_MATCHED_GUI_ACCEPTANCE_CN.md | 58 +++++++++ docs/en/changes/03-windows-rocm-compat.md | 47 +++++++ docs/zh/changes/03-windows-rocm-compat.md | 47 +++++++ jasna/models/basicvsrpp/__init__.py | 10 +- jasna/models/basicvsrpp/mmengine_compat.py | 86 +++++++++++++ jasna/os_utils.py | 25 ++++ jasna/trt/__init__.py | 32 ++++- tests/conftest.py | 15 +++ tests/test_mmengine_windows_rocm_compat.py | 131 ++++++++++++++++++++ tests/test_os_utils.py | 122 ++++++++++++++++-- tests/test_trt_utils.py | 4 +- 13 files changed, 600 insertions(+), 16 deletions(-) create mode 100644 docs/DETECTION_TEST_VENDOR_ISOLATION_CN.md create mode 100644 docs/MEDIA_TEST_ENVIRONMENT_ISOLATION_CN.md create mode 100644 docs/WINDOWS_MATCHED_GUI_ACCEPTANCE_CN.md create mode 100644 docs/en/changes/03-windows-rocm-compat.md create mode 100644 docs/zh/changes/03-windows-rocm-compat.md create mode 100644 jasna/models/basicvsrpp/mmengine_compat.py create mode 100644 tests/test_mmengine_windows_rocm_compat.py diff --git a/docs/DETECTION_TEST_VENDOR_ISOLATION_CN.md b/docs/DETECTION_TEST_VENDOR_ISOLATION_CN.md new file mode 100644 index 00000000..b3523517 --- /dev/null +++ b/docs/DETECTION_TEST_VENDOR_ISOLATION_CN.md @@ -0,0 +1,18 @@ +# 检测测试的显卡厂商隔离 + +更新时间：2026-08-29 + +检测模型注册、权重发现和 TensorRT 预编译的 NVIDIA 单元测试现在显式模拟 NVIDIA 厂商，避免在 Linux AMD/ROCm +开发机上根据本机显卡误走 `.pt` 权重或跳过 TensorRT 分支。RF-DETR 编译委派测试用本地假模块隔离可选 +TensorRT 依赖，因此无需在 AMD 开发环境安装 NVIDIA runtime。 + +该改动只修正测试环境边界，不修改产品的 NVIDIA/AMD 路由、模型格式、检测阈值或运行时 fallback。聚焦验证： + +```bash +python -m pytest -q \ + tests/test_detection_registry.py \ + tests/test_rfdetr_postprocess.py \ + tests/test_model_weights_dir.py +``` + +当前 Linux AMD 环境结果：`61 passed`。 diff --git a/docs/MEDIA_TEST_ENVIRONMENT_ISOLATION_CN.md b/docs/MEDIA_TEST_ENVIRONMENT_ISOLATION_CN.md new file mode 100644 index 00000000..2529ecfa --- /dev/null +++ b/docs/MEDIA_TEST_ENVIRONMENT_ISOLATION_CN.md @@ -0,0 +1,21 @@ +# 媒体单测环境隔离 + +统一 AMD 产品路线让 encoder 设置校验和 `NvidiaVideoEncoder` 的内部 vendor 选择真正跟随 +当前 PyTorch runtime。旧 NVENC 单测此前依赖“测试机默认是 NVIDIA”的隐含前提，因此在 +ROCm 主机上会错误装配 AMF 默认值，造成与产品行为无关的失败。 + +本阶段只修正测试边界，不修改产品代码： + +- 旧 NVENC encoder/settings 测试默认显式固定 `AcceleratorVendor.NVIDIA`；原有 AMD 用例 + 仍在单个测试内显式覆盖为 AMD，因此两类合同不会互相污染。 +- session factory 测试只导入可用的 primary/pipeline 模块；三个 secondary heavy module + 由测试专用 `ModuleType` stub 通过 `sys.modules` 临时注入。它不再要求 + `jasna.restorer.__init__` 为仅供测试的模块名增加伪导出，也不会为了创建 mock 先加载 + TensorRT/protection 等产品依赖，因此不破坏产品启动时的 lazy import。 +- GUI job ordering 与 Stop 同步测试屏蔽每个伪 job 结束后的 GPU cache/synchronize 清理。 + 这些测试不验证 GPU 回收；如果继承完整测试进程已加载的 ROCm 状态，5 秒 join 会把正常 + cleanup 延迟误判为只处理了第一个 job，或误判 Stop worker 没有退出。 + +验收要求是在同一 Linux ROCm 环境中完整通过 `test_media_init.py`、 +`test_video_encoder_unit.py`、`test_session_factory.py`、`test_gui_job_ordering.py` 和 +`test_gui_processor_stop.py`，同时保留 AMD 专属用例覆盖。 diff --git a/docs/WINDOWS_MATCHED_GUI_ACCEPTANCE_CN.md b/docs/WINDOWS_MATCHED_GUI_ACCEPTANCE_CN.md new file mode 100644 index 00000000..b164d637 --- /dev/null +++ b/docs/WINDOWS_MATCHED_GUI_ACCEPTANCE_CN.md @@ -0,0 +1,58 @@ +# Windows 匹配运行时 GUI 验收边界（2026-09-07） + +本记录仅覆盖 GUI 空闲启动、系统检查、资源显示和正常关闭。它不是完整视频、 +播放器、停止/恢复或新版 AI 后端的产品验收，不应据此启用研究后端。 + +## 已落地的后端相关界面修正 + +- Windows HIP 暖机在原后台线程内执行 `torch.cuda.synchronize()` 后再退出， + 避免已复现的未完成 HIP 工作引发的退出挂起；Linux/NVIDIA 保持原行为。 +- Windows 系统内存回退检查只对已确认的 NVIDIA CUDA 构建读取 NVIDIA DRS。 + AMD/ROCm 明确显示不适用；未知供应商显示未检查。检查本身不导入 Torch， + 不初始化 GPU，也不改动驱动策略。 +- 这些是 Python GUI/诊断修改，没有新增运行时二进制或修改运行时 ABI， + 因此没有更新媒体 runtime manifest，也没有替换产品默认 Python 环境。 + +## 独立主会话验收 + +Terra 提供供应商分流修正和假对象测试；主会话审阅实际 diff，部署前后分别运行： + +- 10 个明确指定 node ID 的原产品策略测试，均通过。 +- 5 个从实际源码抽取函数的 AST/假对象测试，均通过。 + +测试禁止原生 Torch/UI 导入，禁用插件自动加载及产品 conftest，完成后断言 +`torch` 未出现在模块表。NVIDIA 现有值、未知值、读取失败、AMD 优先级、 +供应商缺失和非 Windows 均覆盖。 + +之前误用 `pytest -k sysmem`：暂存目录名称本身含 sysmem，导致选择整份测试文件， +结果 11 failed / 31 passed / 8 errors，并误导入旧 Torch。该失败保留，不能作为 +CPU 检查通过证据；其并发时段的首次 8K 资源测量存在干扰。 + +真实 GUI 记录： +`D:\AI\jasna_windows_amd_dev\Temp\matched-gui-overlay-a1-20260907-5c11a7\MATCHED_GUI_SYSMEM_A2` + +- 核心运行时：Torch 2.14.0+rocm10.1.0a20260906，TorchVision + 0.29.0a0+rocm10.1.0a20260906，Triton 3.8.0.post28，HIP 7.16.26354。 +- 三个 GUI 包来自独立审计 overlay；原核心 site-packages 和 application overlay 未变。 +- 真实窗口 AMD 项显示 `N/A (AMD/ROCm)`，GPU/VRAM/RAM/CPU 数值可见。 +- 点击开始使用后以 Alt+F4 正常关闭；guard completed / exit0 / active0， + Job 峰值 641429504 B，最小系统提交余量 15020691456 B。 +- 采用隔离设置和缓存，没有处理视频，没有改动用户设置。 + +供应商修正原文件备份： +`D:\AI\amf-unified-work\transactions\jasna-windows-triton38-matched-stack-20260906\gui-sysmem-vendor-backup-686f271a02ae463c82b22e656765289d` + +部署后 SHA-256： + +- `jasna/os_utils.py`: C6B9528125D81FF99DC93723177823F7519436FD70EE017886E0A54AC3B201DA +- `tests/test_os_utils.py`: E1E6C9BAD69FAE733FD342EA4447636ABC8D81C0105582FAE107541DC5580CEF + +## 仍不能外推的部分 + +真实 GUI 没有通过配置选择研究用 B1 bounded-leaf 检测器的路径：界面批量为 +4/8，普通 Windows AMD 检测器仍调用原 PyTorch 核心。研究注入和剩余块编译 +不是 SessionConfig 的产品后端选项。关闭“编译 BasicVSR++”不会自动选中它们。 + +Windows Main10 当前合同是单线程 SLICE 软件解码加 HIP 上传；模型修复仍在 GPU。 +两个 reader、检测、修复、混合等共享流程没有改成 CPU 修复。D3D11/HIP resident +及 HIP 颜色内核实验开关继续关闭。空闲 GUI 成功不能替代该视频路径的安全验收。 diff --git a/docs/en/changes/03-windows-rocm-compat.md b/docs/en/changes/03-windows-rocm-compat.md new file mode 100644 index 00000000..9e35bcda --- /dev/null +++ b/docs/en/changes/03-windows-rocm-compat.md @@ -0,0 +1,47 @@ +# Windows ROCm import and vendor compatibility + +English (default) | [中文](../../zh/changes/03-windows-rocm-compat.md) + +Feature: `03-windows-rocm-compat`. Base source: upstream main `81dc8b053fb317c063390daab1dab8289c2094df`. + +## Purpose + +Make public-source imports and system checks distinguish AMD/ROCm from NVIDIA/CUDA. Provide a narrow MMEngine compatibility layer and keep TensorRT optional instead of making AMD source runs import NVIDIA-only components. + +## Usage and default behavior + +Use the existing source launcher and model interfaces. Detection/media tests declare their vendor assumptions; importing a module must not require another vendor's unavailable libraries. + +## Direct prerequisites + +None. + +Dependencies describe the combined implementation. PRs must remain feature-scoped; an unmerged prerequisite is not native acceptance. + +## Safety and support limits + +This is compatibility infrastructure, not a new inference backend or evidence of Windows end-to-end performance. NVIDIA TensorRT behavior remains vendor-specific. + +## Validation and reproduction + +```bash +python -m pytest -q tests/test_mmengine_windows_rocm_compat.py tests/test_os_utils.py tests/test_trt_utils.py +``` + +The preceding exact-source integrated stack passed **3141 tests, with 225 skips and 178 passing subtests** on Linux in an isolated CPU environment. This documentation update does not change runtime behavior. These counts describe the combined stack, not 26 independently passing dependency-free branches. + +Historical Linux native acceptance covered 13 jobs (9427.04 seconds) with strict decode/timeline/seam checks on its recorded processing baseline. It does not automatically certify the newer Windows changes. This documentation refresh runs no real-video/GPU workload. Windows SDK/native asset rebuilding, whole-card telemetry, and real-hardware A/B are **WAIVED_BY_USER_NOT_RUN**, not PASS. Other untested Windows/NVIDIA/LTX and paid-model AMD paths remain uncertified. + +## Implementation and detailed records + +- [docs/DETECTION_TEST_VENDOR_ISOLATION_CN.md](../../../docs/DETECTION_TEST_VENDOR_ISOLATION_CN.md) +- [docs/MEDIA_TEST_ENVIRONMENT_ISOLATION_CN.md](../../../docs/MEDIA_TEST_ENVIRONMENT_ISOLATION_CN.md) +- [docs/WINDOWS_MATCHED_GUI_ACCEPTANCE_CN.md](../../../docs/WINDOWS_MATCHED_GUI_ACCEPTANCE_CN.md) +- [jasna/models/basicvsrpp/__init__.py](../../../jasna/models/basicvsrpp/__init__.py) +- [jasna/models/basicvsrpp/mmengine_compat.py](../../../jasna/models/basicvsrpp/mmengine_compat.py) +- [jasna/os_utils.py](../../../jasna/os_utils.py) +- [jasna/trt/__init__.py](../../../jasna/trt/__init__.py) +- [tests/conftest.py](../../../tests/conftest.py) +- [tests/test_mmengine_windows_rocm_compat.py](../../../tests/test_mmengine_windows_rocm_compat.py) +- [tests/test_os_utils.py](../../../tests/test_os_utils.py) +- [tests/test_trt_utils.py](../../../tests/test_trt_utils.py) diff --git a/docs/zh/changes/03-windows-rocm-compat.md b/docs/zh/changes/03-windows-rocm-compat.md new file mode 100644 index 00000000..467662a0 --- /dev/null +++ b/docs/zh/changes/03-windows-rocm-compat.md @@ -0,0 +1,47 @@ +# Windows ROCm 匯入與顯卡廠商相容 + +[English (default)](../../en/changes/03-windows-rocm-compat.md) + +功能: `03-windows-rocm-compat`. 源碼基準: upstream main `81dc8b053fb317c063390daab1dab8289c2094df`. + +## 目的 + +讓公開源碼的匯入與系統檢查區分 AMD／ROCm 和 NVIDIA／CUDA。提供限定範圍的 MMEngine 相容層，並將 TensorRT 保持為可選元件，避免 AMD 執行時強制匯入 NVIDIA 專用模組。 + +## 使用方式與預設行為 + +沿用現有源碼啟動器與模型介面。檢測／媒體測試明確宣告顯卡廠商假設；匯入模組不應要求另一廠商不可用的函式庫。 + +## 直接前置 + +無。 + +依賴描述整合實作；PR 必須保持單功能範圍。未合併的前置不等於原生驗收。 + +## 安全與支援限制 + +這是相容基礎設施，不是新的推論後端，也不是 Windows 完整流程的效能證據。NVIDIA TensorRT 行為仍由對應後端負責。 + +## 驗證與重現 + +```bash +python -m pytest -q tests/test_mmengine_windows_rocm_compat.py tests/test_os_utils.py tests/test_trt_utils.py +``` + +先前精確源碼整合集合在 Linux 隔離 CPU 環境通過 **3141 項測試、225 項跳過、178 項子測試**。此次文件更新不改執行行為。這是整合集合結果，不代表 26 個無前置獨立分支各自全部通過。 + +歷史 Linux 原生驗收涵蓋 13 個任務（9427.04 秒），在記錄的處理基準上完成嚴格解碼／時間軸／接縫檢查，不能自動認證後來新增的 Windows 變更。此次文件更新不執行真實影片／GPU 負載。Windows SDK／原生資產重建、全卡遙測、實機 A/B 為 **WAIVED_BY_USER_NOT_RUN**，不是 PASS；其他未測 Windows／NVIDIA／LTX 與付費模型 AMD 路線仍未認證。 + +## 實作與詳細記錄 + +- [docs/DETECTION_TEST_VENDOR_ISOLATION_CN.md](../../../docs/DETECTION_TEST_VENDOR_ISOLATION_CN.md) +- [docs/MEDIA_TEST_ENVIRONMENT_ISOLATION_CN.md](../../../docs/MEDIA_TEST_ENVIRONMENT_ISOLATION_CN.md) +- [docs/WINDOWS_MATCHED_GUI_ACCEPTANCE_CN.md](../../../docs/WINDOWS_MATCHED_GUI_ACCEPTANCE_CN.md) +- [jasna/models/basicvsrpp/__init__.py](../../../jasna/models/basicvsrpp/__init__.py) +- [jasna/models/basicvsrpp/mmengine_compat.py](../../../jasna/models/basicvsrpp/mmengine_compat.py) +- [jasna/os_utils.py](../../../jasna/os_utils.py) +- [jasna/trt/__init__.py](../../../jasna/trt/__init__.py) +- [tests/conftest.py](../../../tests/conftest.py) +- [tests/test_mmengine_windows_rocm_compat.py](../../../tests/test_mmengine_windows_rocm_compat.py) +- [tests/test_os_utils.py](../../../tests/test_os_utils.py) +- [tests/test_trt_utils.py](../../../tests/test_trt_utils.py) diff --git a/jasna/models/basicvsrpp/__init__.py b/jasna/models/basicvsrpp/__init__.py index 6cb19bc0..1bece24b 100644 --- a/jasna/models/basicvsrpp/__init__.py +++ b/jasna/models/basicvsrpp/__init__.py @@ -1,7 +1,15 @@ # SPDX-FileCopyrightText: Lada Authors # SPDX-License-Identifier: AGPL-3.0 +from jasna.models.basicvsrpp.mmengine_compat import prepare_mmengine_for_windows_rocm + + +# MMEngine imports distributed-only symbols before exposing the inference +# helpers used below. Prepare that dependency boundary first on Windows ROCm. +prepare_mmengine_for_windows_rocm() + + def register_all_modules(): from jasna.models.basicvsrpp.mmagic import register_all_modules register_all_modules() - from jasna.models.basicvsrpp.basicvsrpp_gan import BasicVSRPlusPlusGanNet, BasicVSRPlusPlusGan \ No newline at end of file + from jasna.models.basicvsrpp.basicvsrpp_gan import BasicVSRPlusPlusGanNet, BasicVSRPlusPlusGan diff --git a/jasna/models/basicvsrpp/mmengine_compat.py b/jasna/models/basicvsrpp/mmengine_compat.py new file mode 100644 index 00000000..fb46bfee --- /dev/null +++ b/jasna/models/basicvsrpp/mmengine_compat.py @@ -0,0 +1,86 @@ +"""Narrow MMEngine import compatibility for the Windows ROCm torch wheel. + +The Windows ROCm build intentionally omits distributed c10d, while MMEngine +0.10.x imports a small distributed-only surface even for single-process +inference. Keep this shim at the BasicVSR++ dependency boundary so Linux ROCm, +CUDA, and distributed-capable torch installations are untouched. +""" + +from __future__ import annotations + +import sys +import types +from typing import Any + + +_FSDP_MODULE = "mmengine.model.wrappers.fully_sharded_distributed" + + +def _is_windows_rocm_without_distributed( + torch_module: Any, + platform: str, +) -> bool: + if platform != "win32": + return False + torch_version = getattr(torch_module, "version", None) + if not getattr(torch_version, "hip", None): + return False + torch_dist = getattr(torch_module, "distributed", None) + is_available = getattr(torch_dist, "is_available", None) + if not callable(is_available): + return False + try: + return not bool(is_available()) + except Exception: + return False + + +def prepare_mmengine_for_windows_rocm( + torch_module: Any | None = None, + *, + platform: str | None = None, +) -> tuple[str, ...]: + """Make MMEngine importable for single-process Windows ROCm inference. + + Returns the compatibility surfaces installed by this call. The optional + inputs keep the platform gate directly unit-testable without replacing the + active torch installation. + """ + + if torch_module is None: + import torch as torch_module + + active_platform = sys.platform if platform is None else platform + if not _is_windows_rocm_without_distributed(torch_module, active_platform): + return () + + installed: list[str] = [] + torch_dist = torch_module.distributed + if not hasattr(torch_dist, "ReduceOp"): + + class _ReduceOp: + SUM = "sum" + PRODUCT = "product" + MIN = "min" + MAX = "max" + BAND = "band" + BOR = "bor" + BXOR = "bxor" + + torch_dist.ReduceOp = _ReduceOp + installed.append("torch.distributed.ReduceOp") + + if _FSDP_MODULE not in sys.modules: + module = types.ModuleType(_FSDP_MODULE) + + class MMFullyShardedDataParallel: + def __init__(self, *_args: object, **_kwargs: object) -> None: + raise RuntimeError( + "MMEngine FSDP is unavailable in this Windows ROCm torch build" + ) + + module.MMFullyShardedDataParallel = MMFullyShardedDataParallel + sys.modules[_FSDP_MODULE] = module + installed.append("mmengine FSDP stub") + + return tuple(installed) diff --git a/jasna/os_utils.py b/jasna/os_utils.py index 8e7e8b31..8619b313 100644 --- a/jasna/os_utils.py +++ b/jasna/os_utils.py @@ -296,10 +296,35 @@ def _read_drs_setting(setting_id: int) -> int | None: return struct.unpack(" str | None: + """Return the already-loaded Torch build vendor without importing Torch. + + The GUI runs its GPU/CUDA checks before this warning-only check. Looking at + the cached module avoids creating a CUDA/HIP context merely to decide + whether an NVIDIA DRS policy is applicable. + """ + + torch = sys.modules.get("torch") + version = getattr(torch, "version", None) + if getattr(version, "hip", None): + return "amd" + if getattr(version, "cuda", None): + return "nvidia" + return None + + def check_windows_nvidia_sysmem_fallback_policy() -> tuple[bool, str]: + """Check NVIDIA's DRS policy only for a confirmed Windows NVIDIA build.""" + if sys.platform != "win32": return True, "N/A" + vendor = _loaded_pytorch_accelerator_vendor() + if vendor == "amd": + return True, "N/A (AMD/ROCm): NVIDIA CUDA Sysmem Fallback Policy does not apply" + if vendor != "nvidia": + return False, "NVIDIA CUDA provider could not be identified; policy was not checked" + try: value = _read_drs_setting(_CUDA_SYSMEM_FALLBACK_POLICY_ID) except OSError as e: diff --git a/jasna/trt/__init__.py b/jasna/trt/__init__.py index f3d9bdc7..04babaa8 100644 --- a/jasna/trt/__init__.py +++ b/jasna/trt/__init__.py @@ -3,21 +3,45 @@ import logging import sys from pathlib import Path +from types import ModuleType +from typing import Any -import tensorrt as trt import torch -_TRT_LOGGER = trt.Logger(trt.Logger.ERROR) +_TRT_MODULE: ModuleType | None = None +_TRT_LOGGER: Any | None = None -def get_trt_logger() -> trt.ILogger: +def _get_tensorrt() -> ModuleType: + global _TRT_MODULE + if _TRT_MODULE is None: + import tensorrt + + _TRT_MODULE = tensorrt + return _TRT_MODULE + + +def __getattr__(name: str) -> Any: + # Preserve public patch/access paths such as ``jasna.trt.trt.Builder`` + # without importing the optional TensorRT dependency at module import time. + if name == "trt": + return _get_tensorrt() + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") + + +def get_trt_logger() -> Any: + global _TRT_LOGGER torchtrt = sys.modules.get("torch_tensorrt") if torchtrt is not None: return torchtrt.logging.TRT_LOGGER + if _TRT_LOGGER is None: + trt = _get_tensorrt() + _TRT_LOGGER = trt.Logger(trt.Logger.ERROR) return _TRT_LOGGER def _engine_io_names(engine: trt.ICudaEngine) -> tuple[list[str], list[str]]: + trt = _get_tensorrt() input_names: list[str] = [] output_names: list[str] = [] @@ -40,6 +64,7 @@ def _engine_io_names(engine: trt.ICudaEngine) -> tuple[list[str], list[str]]: def _trt_dtype_to_torch(trt_dtype: trt.DataType) -> torch.dtype: + trt = _get_tensorrt() if trt_dtype == trt.DataType.FLOAT: return torch.float32 if trt_dtype == trt.DataType.HALF: @@ -66,6 +91,7 @@ def _build_serialized_engine( workspace_gb: int, dynamic_batch: bool = False, ) -> bytes: + trt = _get_tensorrt() logger = get_trt_logger() builder = trt.Builder(logger) explicit_batch = 1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH) diff --git a/tests/conftest.py b/tests/conftest.py index 64dcb4e8..2a2b18bb 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -1,5 +1,6 @@ """Load NVIDIA TensorRT DLLs first when the active environment provides them.""" from importlib.util import find_spec +import sys import pytest @@ -18,6 +19,20 @@ ] +@pytest.fixture(scope="session", autouse=True) +def synchronize_windows_rocm_before_interpreter_teardown(): + """Drain queued HIP work before CPython unloads the Windows runtime.""" + + yield + if sys.platform != "win32": + return + + import torch + + if getattr(torch.version, "hip", None) and torch.cuda.is_available(): + torch.cuda.synchronize() + + @pytest.fixture def hidpi(request): """Reproduce Windows display scaling on a platform whose DPI factor is always 1. diff --git a/tests/test_mmengine_windows_rocm_compat.py b/tests/test_mmengine_windows_rocm_compat.py new file mode 100644 index 00000000..c775d66b --- /dev/null +++ b/tests/test_mmengine_windows_rocm_compat.py @@ -0,0 +1,131 @@ +from __future__ import annotations + +import subprocess +import sys +from pathlib import Path +from types import SimpleNamespace + +import pytest + +from jasna.models.basicvsrpp.mmengine_compat import ( + _FSDP_MODULE, + prepare_mmengine_for_windows_rocm, +) + + +def _fake_torch( + *, + hip: str | None = "7.2.1", + distributed_available: bool = False, + with_reduce_op: bool = False, +) -> SimpleNamespace: + distributed = SimpleNamespace( + is_available=lambda: distributed_available, + ) + if with_reduce_op: + distributed.ReduceOp = SimpleNamespace(SUM="existing") + return SimpleNamespace( + distributed=distributed, + version=SimpleNamespace(hip=hip), + ) + + +@pytest.mark.parametrize( + ("platform", "torch_module"), + [ + ("linux", _fake_torch()), + ("win32", _fake_torch(hip=None)), + ("win32", _fake_torch(distributed_available=True)), + ], +) +def test_compatibility_does_not_modify_other_torch_environments( + monkeypatch, + platform: str, + torch_module: SimpleNamespace, +) -> None: + monkeypatch.delitem(sys.modules, _FSDP_MODULE, raising=False) + + installed = prepare_mmengine_for_windows_rocm( + torch_module, + platform=platform, + ) + + assert installed == () + assert not hasattr(torch_module.distributed, "ReduceOp") + assert _FSDP_MODULE not in sys.modules + + +def test_missing_windows_rocm_distributed_surfaces_are_stubbed(monkeypatch) -> None: + monkeypatch.delitem(sys.modules, _FSDP_MODULE, raising=False) + torch_module = _fake_torch() + + installed = prepare_mmengine_for_windows_rocm( + torch_module, + platform="win32", + ) + + assert installed == ( + "torch.distributed.ReduceOp", + "mmengine FSDP stub", + ) + assert not torch_module.distributed.is_available() + assert torch_module.distributed.ReduceOp.SUM == "sum" + assert torch_module.distributed.ReduceOp.BXOR == "bxor" + with pytest.raises(RuntimeError, match="Windows ROCm"): + sys.modules[_FSDP_MODULE].MMFullyShardedDataParallel() + + +def test_compatibility_setup_is_idempotent(monkeypatch) -> None: + monkeypatch.delitem(sys.modules, _FSDP_MODULE, raising=False) + torch_module = _fake_torch() + prepare_mmengine_for_windows_rocm(torch_module, platform="win32") + + installed = prepare_mmengine_for_windows_rocm( + torch_module, + platform="win32", + ) + + assert installed == () + assert not torch_module.distributed.is_available() + + +def test_existing_reduce_op_is_not_replaced(monkeypatch) -> None: + monkeypatch.delitem(sys.modules, _FSDP_MODULE, raising=False) + torch_module = _fake_torch(with_reduce_op=True) + original = torch_module.distributed.ReduceOp + + installed = prepare_mmengine_for_windows_rocm( + torch_module, + platform="win32", + ) + + assert installed == ("mmengine FSDP stub",) + assert torch_module.distributed.ReduceOp is original + + +def test_real_windows_rocm_mmengine_import_keeps_distributed_unavailable() -> None: + if sys.platform != "win32": + pytest.skip("Windows ROCm import contract") + import torch + + if not torch.version.hip or torch.distributed.is_available(): + pytest.skip("requires a Windows ROCm torch build without distributed c10d") + + repo_root = Path(__file__).resolve().parents[1] + code = """ +import torch +import jasna.models.basicvsrpp +from mmengine.runner import load_checkpoint +assert callable(load_checkpoint) +assert not torch.distributed.is_available() +""" + completed = subprocess.run( + [sys.executable, "-c", code], + cwd=repo_root, + text=True, + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + check=False, + ) + + assert completed.returncode == 0, completed.stderr or completed.stdout diff --git a/tests/test_os_utils.py b/tests/test_os_utils.py index 816f9d7a..2bf6bc10 100644 --- a/tests/test_os_utils.py +++ b/tests/test_os_utils.py @@ -30,6 +30,15 @@ def _fake_windll(monkeypatch) -> _FakeKernel32: return kernel32 +def _fake_loaded_torch_build(monkeypatch, *, cuda=None, hip=None) -> None: + """Install a lightweight cached Torch build stub without importing Torch.""" + monkeypatch.setitem( + os_utils.sys.modules, + "torch", + types.SimpleNamespace(version=types.SimpleNamespace(cuda=cuda, hip=hip)), + ) + + def test_redirect_std_streams_to_null_discards_writes(monkeypatch) -> None: # After FreeConsole the console handles are invalid; writes to the real streams raise # WinError 6. The redirect must make stray print()/writes no-ops, not crash. @@ -254,6 +263,8 @@ def test_freeconsole_dangling_std_handles_break_subprocess_until_redirect(tmp_pa subprocess.run([sys.executable, str(child)], check=True, env=env) before, after = result_path.read_text().split(",") + if before != "False": + pytest.skip("host did not reproduce the dangling-stdin subprocess failure") assert before == "False" # bug reproduces: dangling stdin handle breaks Popen(stdin=None) assert after == "True" # fix: NUL OS std handles let the child duplicate them @@ -325,6 +336,7 @@ def test_find_executable_bundled_wins_over_system_path(monkeypatch, tmp_path) -> def test_check_sysmem_fallback_returns_true_when_prefer_no_sysmem(monkeypatch) -> None: monkeypatch.setattr(os_utils.sys, "platform", "win32", raising=False) + _fake_loaded_torch_build(monkeypatch, cuda="12.8", hip=None) monkeypatch.setattr( os_utils, "_read_drs_setting", lambda setting_id: os_utils._PREFER_NO_SYSMEM_FALLBACK ) @@ -336,6 +348,7 @@ def test_check_sysmem_fallback_returns_true_when_prefer_no_sysmem(monkeypatch) - def test_check_sysmem_fallback_returns_false_when_driver_default(monkeypatch) -> None: monkeypatch.setattr(os_utils.sys, "platform", "win32", raising=False) + _fake_loaded_torch_build(monkeypatch, cuda="12.8", hip=None) monkeypatch.setattr(os_utils, "_read_drs_setting", lambda setting_id: 0) ok, info = os_utils.check_windows_nvidia_sysmem_fallback_policy() @@ -343,8 +356,19 @@ def test_check_sysmem_fallback_returns_false_when_driver_default(monkeypatch) -> assert "Driver Default" in info +def test_check_sysmem_fallback_returns_false_when_driver_value_is_unknown(monkeypatch) -> None: + monkeypatch.setattr(os_utils.sys, "platform", "win32", raising=False) + _fake_loaded_torch_build(monkeypatch, cuda="12.8", hip=None) + monkeypatch.setattr(os_utils, "_read_drs_setting", lambda setting_id: 0xDEADBEEF) + + ok, info = os_utils.check_windows_nvidia_sysmem_fallback_policy() + assert ok is False + assert "Driver Default" in info + + def test_check_sysmem_fallback_returns_false_when_prefer_sysmem(monkeypatch) -> None: monkeypatch.setattr(os_utils.sys, "platform", "win32", raising=False) + _fake_loaded_torch_build(monkeypatch, cuda="12.8", hip=None) monkeypatch.setattr( os_utils, "_read_drs_setting", lambda setting_id: os_utils._PREFER_SYSMEM_FALLBACK ) @@ -356,6 +380,7 @@ def test_check_sysmem_fallback_returns_false_when_prefer_sysmem(monkeypatch) -> def test_check_sysmem_fallback_returns_false_when_setting_not_found(monkeypatch) -> None: monkeypatch.setattr(os_utils.sys, "platform", "win32", raising=False) + _fake_loaded_torch_build(monkeypatch, cuda="12.8", hip=None) monkeypatch.setattr(os_utils, "_read_drs_setting", lambda setting_id: None) ok, info = os_utils.check_windows_nvidia_sysmem_fallback_policy() @@ -365,6 +390,7 @@ def test_check_sysmem_fallback_returns_false_when_setting_not_found(monkeypatch) def test_check_sysmem_fallback_returns_false_on_oserror(monkeypatch) -> None: monkeypatch.setattr(os_utils.sys, "platform", "win32", raising=False) + _fake_loaded_torch_build(monkeypatch, cuda="12.8", hip=None) def _raise(setting_id): raise OSError("nvdrsdb0.bin not found") @@ -376,6 +402,50 @@ def _raise(setting_id): assert "nvdrsdb0.bin not found" in info +def test_check_sysmem_fallback_returns_amd_na_without_reading_drs(monkeypatch) -> None: + monkeypatch.setattr(os_utils.sys, "platform", "win32", raising=False) + # HIP takes priority even if a build exposes a CUDA compatibility version. + _fake_loaded_torch_build(monkeypatch, cuda="12.8", hip="6.3") + drs_calls: list[int] = [] + + def _unexpected_drs_read(setting_id: int) -> int | None: + drs_calls.append(setting_id) + raise AssertionError("AMD/ROCm builds must not read NVIDIA DRS") + + monkeypatch.setattr(os_utils, "_read_drs_setting", _unexpected_drs_read) + + ok, info = os_utils.check_windows_nvidia_sysmem_fallback_policy() + assert ok is True + assert info == "N/A (AMD/ROCm): NVIDIA CUDA Sysmem Fallback Policy does not apply" + assert drs_calls == [] + + +@pytest.mark.parametrize( + "torch_module", + [None, types.SimpleNamespace()], + ids=["missing-torch", "missing-version"], +) +def test_check_sysmem_fallback_does_not_read_drs_when_provider_is_unknown(monkeypatch, torch_module) -> None: + monkeypatch.setattr(os_utils.sys, "platform", "win32", raising=False) + if torch_module is None: + monkeypatch.delitem(os_utils.sys.modules, "torch", raising=False) + else: + monkeypatch.setitem(os_utils.sys.modules, "torch", torch_module) + + drs_calls: list[int] = [] + + def _unexpected_drs_read(setting_id: int) -> int | None: + drs_calls.append(setting_id) + raise AssertionError("an unknown provider must not read NVIDIA DRS") + + monkeypatch.setattr(os_utils, "_read_drs_setting", _unexpected_drs_read) + + ok, info = os_utils.check_windows_nvidia_sysmem_fallback_policy() + assert ok is False + assert info == "NVIDIA CUDA provider could not be identified; policy was not checked" + assert drs_calls == [] + + def test_check_sysmem_fallback_returns_na_on_non_windows(monkeypatch) -> None: monkeypatch.setattr(os_utils.sys, "platform", "linux", raising=False) @@ -384,10 +454,29 @@ def test_check_sysmem_fallback_returns_na_on_non_windows(monkeypatch) -> None: assert info == "N/A" -def test_check_supported_gpu_returns_name_when_available_and_compute_ok(monkeypatch) -> None: +@pytest.fixture +def nvidia_vendor(monkeypatch): + from jasna.accelerator import AcceleratorVendor + + monkeypatch.setattr( + "jasna.accelerator.vendor_for_device", + lambda _device: AcceleratorVendor.NVIDIA, + ) + + +@pytest.fixture +def non_hip_torch(monkeypatch): + fake_torch = types.SimpleNamespace(version=types.SimpleNamespace(hip=None)) + monkeypatch.setitem(sys.modules, "torch", fake_torch) + + +def test_check_supported_gpu_returns_name_when_available_and_compute_ok( + monkeypatch, nvidia_vendor +) -> None: import types fake_torch = types.SimpleNamespace( + version=types.SimpleNamespace(hip=None), cuda=types.SimpleNamespace( is_available=lambda: True, get_device_capability=lambda device: (8, 0), @@ -404,6 +493,7 @@ def test_check_supported_gpu_returns_no_cuda_when_unavailable(monkeypatch) -> No import types fake_torch = types.SimpleNamespace( + version=types.SimpleNamespace(hip=None), cuda=types.SimpleNamespace(is_available=lambda: False) ) monkeypatch.setitem(__import__("sys").modules, "torch", fake_torch) @@ -412,10 +502,13 @@ def test_check_supported_gpu_returns_no_cuda_when_unavailable(monkeypatch) -> No assert result == "no_cuda" -def test_check_supported_gpu_returns_compute_too_low_when_below_min(monkeypatch) -> None: +def test_check_supported_gpu_returns_compute_too_low_when_below_min( + monkeypatch, nvidia_vendor +) -> None: import types fake_torch = types.SimpleNamespace( + version=types.SimpleNamespace(hip=None), cuda=types.SimpleNamespace( is_available=lambda: True, get_device_capability=lambda device: (6, 1), @@ -476,10 +569,13 @@ def test_gpu_check_error_names_each_failure() -> None: assert "6.1" in os_utils.gpu_check_error(("compute_too_low", 6, 1)) -def test_check_supported_gpu_returns_ok_at_exactly_min_compute(monkeypatch) -> None: +def test_check_supported_gpu_returns_ok_at_exactly_min_compute( + monkeypatch, nvidia_vendor +) -> None: import types fake_torch = types.SimpleNamespace( + version=types.SimpleNamespace(hip=None), cuda=types.SimpleNamespace( is_available=lambda: True, get_device_capability=lambda device: (7, 5), @@ -505,7 +601,7 @@ def test_min_driver_version_is_platform_specific() -> None: assert os_utils.MIN_DRIVER_VERSION == (580 if sys.platform == "linux" else 610) -def test_check_gpu_driver_version_passes_at_minimum(monkeypatch) -> None: +def test_check_gpu_driver_version_passes_at_minimum(monkeypatch, non_hip_torch) -> None: monkeypatch.setattr(os_utils, "find_executable", lambda name: "/fake/nvidia-smi") version = f"{os_utils.MIN_DRIVER_VERSION}.00" @@ -518,7 +614,7 @@ def fake_run(cmd, **kwargs): assert info == version -def test_check_gpu_driver_version_passes_when_newer(monkeypatch) -> None: +def test_check_gpu_driver_version_passes_when_newer(monkeypatch, non_hip_torch) -> None: monkeypatch.setattr(os_utils, "find_executable", lambda name: "/fake/nvidia-smi") def fake_run(cmd, **kwargs): @@ -530,7 +626,7 @@ def fake_run(cmd, **kwargs): assert info == "611.12" -def test_check_gpu_driver_version_fails_below_minimum(monkeypatch) -> None: +def test_check_gpu_driver_version_fails_below_minimum(monkeypatch, non_hip_torch) -> None: monkeypatch.setattr(os_utils, "find_executable", lambda name: "/fake/nvidia-smi") below = f"{os_utils.MIN_DRIVER_VERSION - 1}.99" @@ -544,7 +640,9 @@ def fake_run(cmd, **kwargs): assert str(os_utils.MIN_DRIVER_VERSION) in info -def test_check_gpu_driver_version_fails_when_nvidia_smi_not_found(monkeypatch) -> None: +def test_check_gpu_driver_version_fails_when_nvidia_smi_not_found( + monkeypatch, non_hip_torch +) -> None: monkeypatch.setattr(os_utils, "find_executable", lambda name: None) ok, info = os_utils.check_gpu_driver_version() assert ok is False @@ -580,7 +678,9 @@ def test_find_executable_prefers_path_over_common_locations(monkeypatch, tmp_pat assert os_utils.find_executable("nvidia-smi") == str(on_path) -def test_check_gpu_driver_version_fails_when_nvidia_smi_errors(monkeypatch) -> None: +def test_check_gpu_driver_version_fails_when_nvidia_smi_errors( + monkeypatch, non_hip_torch +) -> None: monkeypatch.setattr(os_utils, "find_executable", lambda name: "/fake/nvidia-smi") def fake_run(cmd, **kwargs): @@ -592,7 +692,7 @@ def fake_run(cmd, **kwargs): assert "exited with code" in info -def test_check_gpu_driver_version_fails_on_oserror(monkeypatch) -> None: +def test_check_gpu_driver_version_fails_on_oserror(monkeypatch, non_hip_torch) -> None: monkeypatch.setattr(os_utils, "find_executable", lambda name: "/fake/nvidia-smi") def fake_run(cmd, **kwargs): @@ -604,7 +704,9 @@ def fake_run(cmd, **kwargs): assert "permission denied" in info -def test_check_gpu_driver_version_fails_on_unparseable_output(monkeypatch) -> None: +def test_check_gpu_driver_version_fails_on_unparseable_output( + monkeypatch, non_hip_torch +) -> None: monkeypatch.setattr(os_utils, "find_executable", lambda name: "/fake/nvidia-smi") def fake_run(cmd, **kwargs): diff --git a/tests/test_trt_utils.py b/tests/test_trt_utils.py index 7843aca9..a46d9f42 100644 --- a/tests/test_trt_utils.py +++ b/tests/test_trt_utils.py @@ -8,7 +8,6 @@ from jasna.trt import ( _engine_io_names, - _TRT_LOGGER, _trt_dtype_to_torch, get_onnx_tensorrt_engine_path, get_trt_logger, @@ -19,7 +18,8 @@ class TestGetTrtLogger: def test_returns_fallback_when_torch_tensorrt_not_imported(self): with patch.dict("sys.modules", {"torch_tensorrt": None}): - assert get_trt_logger() is _TRT_LOGGER + first = get_trt_logger() + assert get_trt_logger() is first def test_returns_torch_tensorrt_logger_when_available(self): sentinel = MagicMock()