From e8129fe495f643c50039963a0e6cedbb03cdead8 Mon Sep 17 00:00:00 2001 From: nan Date: Thu, 10 Sep 2026 20:50:44 +0800 Subject: [PATCH] Implement BiRefNet and HAT Sharper models with background removal capabilities - Added validation documentation for AI model tiers. - Introduced `background.py` for background removal logic, including inference functions for alpha and foreground extraction. - Created `onnx_birefnet.py` to handle BiRefNet model sessions and alpha inference. - Developed tests for background models, ensuring proper dispatch and session reuse. - Added tests for BiRefNet ONNX model to validate alpha output and error handling for invalid values. - Implemented model export scripts for BiRefNet and HAT Sharper, ensuring proper weight compression and licensing. - Created `deform_sampling.py` for deformable convolution operations used in BiRefNet. - Added utility functions for compressing ONNX weights without altering graph arithmetic. - Introduced a reference plane node for depth field adjustments in geometry processing. --- .../simplify-ai-model-tiers/.openspec.yaml | 2 + .../changes/simplify-ai-model-tiers/design.md | 117 ++++++++ .../mirror-validation.json | 53 ++++ .../simplify-ai-model-tiers/proposal.md | 38 +++ .../retirement-results.json | 267 ++++++++++++++++++ .../simplify-ai-model-tiers/retirement.json | 131 +++++++++ .../specs/ai-model-tiers/spec.md | 42 +++ .../specs/background-removal-models/spec.md | 59 ++++ .../specs/hat-upscale/spec.md | 49 ++++ .../specs/model-asset-retirement/spec.md | 38 +++ .../changes/simplify-ai-model-tiers/tasks.md | 41 +++ .../simplify-ai-model-tiers/validation.md | 94 ++++++ pyproject.toml | 5 +- src/anyimage/assets/O_AnyImage.blend | 4 +- src/anyimage/blender_manifest.toml | 2 +- src/anyimage/operators/ai_setup.py | 8 +- .../convert_to_panorama/operators.py | 7 +- .../operators/convert_to_plane/operators.py | 8 +- .../operators/cutout_tool/operators.py | 7 +- src/anyimage/operators/remove_background.py | 9 +- src/anyimage/operators/upscale.py | 3 +- src/anyimage/preferences.py | 46 ++- src/anyimage/properties.py | 28 +- src/anyimage/runtime.py | 13 +- src/anyimage/server/jobs/remove_background.py | 12 +- src/anyimage/server/model_catalog.py | 129 ++++----- src/anyimage/server/model_manager.py | 57 ++-- src/anyimage/server/models/__init__.py | 4 +- .../server/models/{ben2.py => background.py} | 30 +- src/anyimage/server/models/onnx_birefnet.py | 41 +++ src/anyimage/server/models/onnx_pisa_sr.py | 153 ---------- src/anyimage/server/models/onnx_upscale.py | 18 +- src/anyimage/server/models/upscale.py | 6 +- .../operators/test_ai_image_operators.py | 12 +- .../operators/test_ai_input_failures.py | 2 +- tests/anyimage/operators/test_ai_setup.py | 29 +- .../operators/test_cutout_operators.py | 8 +- .../operators/test_panorama_conversion.py | 3 +- .../operators/test_texture_node_actions.py | 6 +- tests/anyimage/operators/test_upscale_undo.py | 2 +- tests/anyimage/server/support.py | 4 +- .../anyimage/server/test_background_models.py | 75 +++++ tests/anyimage/server/test_job_server.py | 4 +- .../server/test_model_birefnet_onnx.py | 36 +++ tests/anyimage/server/test_model_catalog.py | 26 +- tests/anyimage/server/test_model_inference.py | 5 +- tests/anyimage/server/test_model_manager.py | 51 +++- tests/anyimage/server/test_model_pisa_sr.py | 125 -------- .../server/test_model_upscale_onnx.py | 12 + tests/anyimage/server/test_server_jobs.py | 16 +- tests/anyimage/test_preferences.py | 29 +- tests/support/blender.py | 6 + tests/tools/models/test_cli.py | 8 +- tests/tools/nodes/test_depth_relief_planes.py | 63 ++++- tests/tools/nodes/test_shadeless.py | 2 + tests/tools/test_pack.py | 22 +- tools/models/birefnet_export.py | 108 +++++++ tools/models/cli.py | 32 +-- tools/models/deform_sampling.py | 47 +++ tools/models/exports.py | 95 ------- tools/models/hat_export.py | 95 +++++++ .../licenses/pisa-sr/LICENSE-APACHE.txt | 202 ------------- .../models/licenses/pisa-sr/LICENSE-MODEL.txt | 84 ------ tools/models/licenses/pisa-sr/NOTICE.txt | 21 -- tools/models/onnx_weights.py | 49 ++++ tools/models/pisa_sr_export.py | 198 ------------- tools/nodes/common/reference_plane.py | 31 ++ tools/nodes/groups/image_depth_balloon.py | 32 +-- tools/nodes/groups/image_relief_plane.py | 21 +- uv.lock | 16 +- 70 files changed, 1890 insertions(+), 1208 deletions(-) create mode 100644 openspec/changes/simplify-ai-model-tiers/.openspec.yaml create mode 100644 openspec/changes/simplify-ai-model-tiers/design.md create mode 100644 openspec/changes/simplify-ai-model-tiers/mirror-validation.json create mode 100644 openspec/changes/simplify-ai-model-tiers/proposal.md create mode 100644 openspec/changes/simplify-ai-model-tiers/retirement-results.json create mode 100644 openspec/changes/simplify-ai-model-tiers/retirement.json create mode 100644 openspec/changes/simplify-ai-model-tiers/specs/ai-model-tiers/spec.md create mode 100644 openspec/changes/simplify-ai-model-tiers/specs/background-removal-models/spec.md create mode 100644 openspec/changes/simplify-ai-model-tiers/specs/hat-upscale/spec.md create mode 100644 openspec/changes/simplify-ai-model-tiers/specs/model-asset-retirement/spec.md create mode 100644 openspec/changes/simplify-ai-model-tiers/tasks.md create mode 100644 openspec/changes/simplify-ai-model-tiers/validation.md rename src/anyimage/server/models/{ben2.py => background.py} (76%) create mode 100644 src/anyimage/server/models/onnx_birefnet.py delete mode 100644 src/anyimage/server/models/onnx_pisa_sr.py create mode 100644 tests/anyimage/server/test_background_models.py create mode 100644 tests/anyimage/server/test_model_birefnet_onnx.py delete mode 100644 tests/anyimage/server/test_model_pisa_sr.py create mode 100644 tools/models/birefnet_export.py create mode 100644 tools/models/deform_sampling.py create mode 100644 tools/models/hat_export.py delete mode 100644 tools/models/licenses/pisa-sr/LICENSE-APACHE.txt delete mode 100644 tools/models/licenses/pisa-sr/LICENSE-MODEL.txt delete mode 100644 tools/models/licenses/pisa-sr/NOTICE.txt create mode 100644 tools/models/onnx_weights.py delete mode 100644 tools/models/pisa_sr_export.py create mode 100644 tools/nodes/common/reference_plane.py diff --git a/openspec/changes/simplify-ai-model-tiers/.openspec.yaml b/openspec/changes/simplify-ai-model-tiers/.openspec.yaml new file mode 100644 index 0000000..e8cda9e --- /dev/null +++ b/openspec/changes/simplify-ai-model-tiers/.openspec.yaml @@ -0,0 +1,2 @@ +schema: spec-driven +created: 2026-09-10 diff --git a/openspec/changes/simplify-ai-model-tiers/design.md b/openspec/changes/simplify-ai-model-tiers/design.md new file mode 100644 index 0000000..fd5105a --- /dev/null +++ b/openspec/changes/simplify-ai-model-tiers/design.md @@ -0,0 +1,117 @@ +## Context + +模型矩阵见 [proposal.md](proposal.md)。当前模型清单已有 MoGe-3，Upscale 已统一最终 2×，背景移除虽有偏好项，Operator 和缓存仍直接绑定 BEN2。`model prepare` / `model sync` 从清单生成文件集合，远端上传采用逐文件复制。 + +相关在途 change 包括 `add-moge3-model`、`add-pisa-sr-two-times-upscale`、`add-ben2-hdr-support` 和 `organize-tool-cli`。实施基于工作区实际调用链，承接其有效行为；模型成员和默认值以本提案为准。 + +## Goals + +- 用同一模型清单驱动偏好、初始化、Job 参数、模型下载和镜像资产。 +- 将各模型差异限制在适配器与少量推理配置中，复用现有图像准备、Alpha 合成、缓存和事务处理。 +- 以可复现资产和 CPU / DirectML 真实推理验证完成接入，保持三档与资源需求的说明准确。 + +## Decisions + +### HR-matting 设备决议（2026-09-10） +用户确认保留 HR-matting，暂时仅支持 CPU。2048 输入的 ONNX 在 CPU 上通过原始实现对照，DirectML 在最后解码阶段产生明显误差，多种分块导出未通过验证。正式版本采用已验证的原始 GridSample 导出结构，HR 选项标明 CPU，进度与会话缓存使用实际 CPU 设备。Lite 和 BEN2 继续使用所选设备。 + +### 1. 固定三类各三款，默认入口统一 + +模型键建议如下；已有保留模型沿用当前键： + +| 分类 | 模型键，按界面顺序 | 默认 | +|---|---|---| +| Geometry | `MOGE2_VITS_NORMAL`、`MOGE2_VITB_NORMAL`、`MOGE3_VITL` | 第一项 | +| Upscale | `REALESRGAN_GENERAL_WDN_X4V3`、`REALESRGAN_X4PLUS`、`HAT_GAN_X4_SHARPER` | 第一项 | +| Background | `BIREFNET_LITE`、`BEN2_BASE`、`BIREFNET_HR_MATTING` | 第一项 | + +界面显示用途与模型名：Depth 为快速/均衡/精细，Upscale 为快速/通用/精细，背景移除为快速/通用/精细抠图。说明展示真实模型特点，下载体积来自文件清单；档位名称不作为逐图质量保证。 + +扩展偏好读取入口集中在 `preferences.py`。初始化所需模型从三个类别的默认值派生；新安装只准备 MoGe-2 S、WDN 和 BiRefNet Lite。读取持久化选择时，保留有效键，类别外的键归一为该类别默认值。外部提交未知或退役模型键的 Job 明确失败，避免静默使用不同模型。 + +背景模型资源分类改为职责名称 `background`，同步修改生产者和消费者。依赖 BEN2 名称的公共常量、缓存接口和进度文案按实际职责重命名；BEN2 专用 ONNX 适配器继续明确命名。 + +### 2. 背景移除共享编排，按模型适配 Alpha + +```mermaid +flowchart TD + A[扩展偏好选择模型] --> B[公共输入准备与模型就绪检查] + B --> C[背景移除 Job] + C --> D{模型适配器} + D --> E[BiRefNet Lite 1024] + D --> F[BEN2 当前配置] + D --> G[BiRefNet HR-matting 2048] + E --> H[原图尺寸的有限连续 Alpha] + F --> H + G --> H + H --> I[现有 Alpha 合成与前景输出] + I --> J[主线程更新图片或 HDR Alpha] +``` + +公共编排从 `models/ben2.py` 调整为背景移除职责模块。Lite 和 HR-matting 使用一个 BiRefNet 适配模块及各自配置，保留官方 RGB 归一化、输入尺寸、输出激活和缩放约定。BiRefNet 的概率 Alpha 按官方语义处理，避免套用 BEN2 的输出归一化造成半透明区域失真。 + +共享输出为原始分析图片尺寸的 `float32` Alpha，范围 `[0,1]`，保留连续透明度。静态前景、多帧前景和 HDR 的 Alpha-only 路径复用当前协议；原图颜色与已有 Alpha 的合成遵循现有行为。Blender 输入准备复用 `common`，Server 复用自身 `media` 和模型工具。 + +缓存按模型、目录及设备识别，会话在切换和关闭时释放。同模型多帧复用会话，取消和推理失败继续遵守现有 Job 清理流程。 + +BiRefNet 先在隔离导出环境固定上游版本、checkpoint、许可证和预处理，再确定可通过数值验证的存储精度。Lite 约 90 MB、HR-matting 约 444 MB 仅作为评估参考；正式大小和哈希以导出文件为准。默认 Lite 需在同一设备与样图上证明比 BEN2 更低的下载体积，且热推理耗时或峰值占用至少一项更低，并记录另一项和边缘效果。 + +### 3. HAT Sharper 使用压缩权重、FP32 运算 + +HAT 上游源码固定到 `1638a9a822581657811867bf670717f8371fc3e5`，使用官方 Sharper 检查点中的 `params_ema`。离线导出关闭常量折叠，以免预展开注意力偏置扩大文件；浮点权重存为 FP16，在原图运算前 Cast 回 FP32。输入输出保持 FP32，正式 Runtime 沿用 ONNX。 + +实验记录： + +| 项目 | Sharper 实测 | +|---|---| +| 参数量 | 20,772,507 | +| FP32 ONNX | 87,221,721 bytes | +| 压缩权重 ONNX | 45,874,438 bytes | +| DirectML，RTX 5070 Ti，256×256 输入 | 热推理中位数约 0.563 秒 | +| 与 FP32 PyTorch 对照，4 张裁片 | 最大样本 MAE 0.000190，最大像素误差 0.007370 | +| 执行检查 | CPU 与 DirectML 通过；DirectML profile 的计算事件均在 GPU | + +检查点 SHA-256：`5800b67136006eb8cab3b4ed7c8d73b6a195bb18e6cc709b674f9aa069c00271`。实验压缩模型 SHA-256：`032d8b03c704220a08d91465a822aca252c1d6e63db20c3cf23c5fb67f3d2366`。正式导出器规范化后重新记录产物哈希，并执行同样的精度校验。 + +参考：[本地验证报告](../../../.model-cache/hat-validation/report.md)、[官方 HAT](https://github.com/XPixelGroup/HAT)。本地报告位于忽略目录，上表和哈希作为 change 的持久验证摘要。 + +全 FP16 和已试的混合精度运算在 DirectML 中产生 NaN，因此选择已经验证的 FP32 运算方案。该选择优化磁盘和下载成本，运行内存按 FP32 计算评估。Sharpness 是用户已比较确认的视觉偏好，产品说明应表达细节增强，而非原始噪点保真。 + +### 4. Upscale 沿用 2×，HAT 使用独立分块参数 + +共享超分路径仍为：RGB 推理 4× → Lanczos 缩至 2× → 合成按现有规则缩放的原始 Alpha。多帧、输入尺寸上限、取消与错误处理保持同一流程。 + +HAT 采用固定 256×256 ONNX 输入，通过 padding 与分块支持任意合法图片尺寸。先以已测的 48 像素边距、160 像素有效步长作为正式候选；两个 Real-ESRGAN 保持各自当前配置。参数通过模型明确选择，共享拼接实现。 + +在 448×224 样图上，Sharper 的接缝梯度误差从边距 16 的 4.17 降到边距 48 的 1.44（8-bit 单位），代价是该图从 2 块增加到 6 块。实施时补充自然纹理跨缝、细线、奇数尺寸、小图和透明图片验证；可见接缝存在时先修正分块，验证通过后确定发布资产。 + +### 5. 从新模型发布推进退役资产清理 + +```mermaid +flowchart TD + A[锁定上游版本与导出配置] --> B[生成并验证三款新模型] + B --> C[九款正式清单与初始化更新] + C --> D[model prepare 校验本地资产] + D --> E[model sync 发布并回读校验] + E --> F[列出退役本地路径与 R2 对象] + F --> G[按核对过的清单删除退役对象] + G --> H[检查新资产可下载及退役对象不存在] +``` + +退役范围为 DA3 全部现存变体、`moge-2-vitl-normal-onnx`、`safmn-real-x4` 和 `pisa-sr`。DA3 从实际本地目录及远端对象清单识别，覆盖已知 BASE、SMALL 和发现的其他 DA3 变体。清理对应专用源码、导出器、构建登记、依赖和开发源权重缓存。 + +R2 逐文件上传不会自动清理旧对象。本次退役作为有明确对象清单的一次性迁移：保存删除前清单，确认完整前缀属于该模型，先验证新资源再删除并复查。`model sync` 继续按正式清单发布；其他远端对象维持原样。用户自行配置的模型缓存按显式列出的模型目录操作，用户图片与 Job 产物保持独立。 + +## Risks / Trade-offs + +- Lite 已通过 CPU / DirectML 对照；HR 固定使用 CPU → 逐图误差、边缘效果和资源成本记录于本 change，设备限制明确显示在选项中。 +- HAT 比当前 x4plus 更慢，较宽边距增加推理次数 → 默认继续使用 WDN，记录热推理与整图耗时，验证取消响应和资源释放。 +- HR-matting 高分辨率 CPU 峰值工作集约 19.3 GB → 标明 CPU 与高资源成本，保留 Lite 和 BEN2 供日常使用。 +- 多个相关 change 在工作区并行演进 → 开始实施时核对真实调用链和改动，保持 MoGe-3、HDR 与 2× 输出的行为覆盖；本提案只覆盖模型调整。 +- 远端删除不可依赖重新同步自动恢复 → 发布前固定可复现来源、大小和哈希；清理失败时保留明确的未完成对象清单，修复后逐项继续。 + +## Migration Plan + +先完成新模型资产验证，再统一模型选择和生产调用，最后发布并回读新镜像、清理退役模型。有效的已有模型偏好保留，退役选择归一到类别默认值；首次安装改用三个轻量默认模型。 + +实施运行相关行为测试和全量测试，确认正常退出，并检查退役引用、依赖与打包清单。验证记录写入本 change；用户文档构建和扩展打包保持独立于本次实施任务。 diff --git a/openspec/changes/simplify-ai-model-tiers/mirror-validation.json b/openspec/changes/simplify-ai-model-tiers/mirror-validation.json new file mode 100644 index 0000000..e056234 --- /dev/null +++ b/openspec/changes/simplify-ai-model-tiers/mirror-validation.json @@ -0,0 +1,53 @@ +[ + { + "model": "BIREFNET_LITE", + "ready": true, + "files": [ + { + "url": "https://models.omoolab.xyz/birefnet-lite/model.onnx", + "bytes": 112968598, + "sha256": "6e2a9844c57b080e95a402722ddcae5660e6f254a04d3e003347f06f587b71fe" + }, + { + "url": "https://models.omoolab.xyz/birefnet-lite/LICENSE.txt", + "bytes": 1066, + "sha256": "92a7089e0915fc32bc40067560b398f1e6a7a5958abd7d04eda393629a5acefb" + } + ], + "execution": "passed" + }, + { + "model": "BIREFNET_HR_MATTING", + "ready": true, + "files": [ + { + "url": "https://models.omoolab.xyz/birefnet-hr-matting/model.onnx", + "bytes": 529467341, + "sha256": "b3abc34843a4d96f669ca9fb1a2a16b3417ab8b91c987573fdbda36d20514133" + }, + { + "url": "https://models.omoolab.xyz/birefnet-hr-matting/LICENSE.txt", + "bytes": 1066, + "sha256": "92a7089e0915fc32bc40067560b398f1e6a7a5958abd7d04eda393629a5acefb" + } + ], + "execution": "passed" + }, + { + "model": "HAT_GAN_X4_SHARPER", + "ready": true, + "files": [ + { + "url": "https://models.omoolab.xyz/hat-gan-x4-sharper/model.onnx", + "bytes": 45874438, + "sha256": "032d8b03c704220a08d91465a822aca252c1d6e63db20c3cf23c5fb67f3d2366" + }, + { + "url": "https://models.omoolab.xyz/hat-gan-x4-sharper/LICENSE.txt", + "bytes": 11342, + "sha256": "a91d57ebad8955a1757be7891ca2da8e07493e662f0e34e1e1926571e05f37fc" + } + ], + "execution": "passed" + } +] \ No newline at end of file diff --git a/openspec/changes/simplify-ai-model-tiers/proposal.md b/openspec/changes/simplify-ai-model-tiers/proposal.md new file mode 100644 index 0000000..714518a --- /dev/null +++ b/openspec/changes/simplify-ai-model-tiers/proposal.md @@ -0,0 +1,38 @@ +## Why + +当前模型选项的资源梯度和用途不够清晰，PiSA-SR 的下载体积与实际效果不符合使用需求，背景移除也缺少轻量默认模型。HAT Sharper 已通过本地效果比较和 ONNX 验证，适合将三个 AI 功能统一整理为各三档，并完成退役模型的资产清理。 + +## What Changes + +| 功能 | 第一档，默认 | 第二档 | 第三档 | +|---|---|---|---| +| Depth | MoGe-2 ViT-S Normal | MoGe-2 ViT-B Normal | MoGe-3 ViT-L | +| Upscale | Real-ESRGAN General WDN x4v3 | Real-ESRGAN x4plus | Real HAT GAN x4 Sharper | +| Remove Background | BiRefNet Lite | BEN2 | BiRefNet HR-matting | + +- **BREAKING**：移除 MoGe-2 ViT-L Normal、SAFMN、PiSA-SR 的正式选项和实现；清理 DA3 的现行残留及对应本地、R2 模型资产。 +- **BREAKING**：背景移除默认模型改为 BiRefNet Lite，模型选择、AI 环境初始化、下载和缓存统一按背景移除职责组织。 +- 接入 HAT Sharper，采用已验证的 FP16 权重存储、FP32 运算 ONNX 路径，沿用最终 2× 输出，完成适配 HAT 的分块验证。 +- 接入 BiRefNet Lite 与 HR-matting，按官方预处理和 Alpha 输出语义验证 Lite 的 CPU / DirectML 和 HR-matting 的 CPU、效果、体积和资源占用；HR 选项明确标注 CPU。 +- 九款正式模型统一驱动界面、下载、就绪判定和 `model prepare` / `model sync`；清理退役源码、导出器、专用依赖、缓存及远端对象。 + +## Capabilities + +### New Capabilities + +- `ai-model-tiers`：三类各三款模型、默认选择及初始化下载行为。 +- `background-removal-models`：三款背景模型的选择、推理、Alpha 契约和资源生命周期。 +- `hat-upscale`：HAT Sharper 的小体积 ONNX、2× 输出和分块处理。 +- `model-asset-retirement`：正式模型资产准备、镜像校验和退役对象清理。 + +### Modified Capabilities + +当前 `openspec/specs` 尚无已发布 capability。本提案承接已有 MoGe-3、2× Upscale 和 HDR 输入实现，以本次模型矩阵确定后续正式选项。 + +## Impact + +- `server/model_catalog.py`、`preferences.py`、`properties.py`、AI Setup 与相关 Operator：模型分类、默认值和下载选项。 +- `server/models`、`server/model_manager.py`、背景移除 Job：模型适配、推理分派、会话复用与释放。 +- `tools/models`、`pyproject.toml`、`uv.lock` 及相关测试：可复现导出、专用依赖清理和资产校验。 +- 项目模型目录、开发缓存与 `omoolab-r2:ai-models`：新资产发布和退役对象清理。 +- 本次生成规划文件；实施验证结果记录在 change 内，运行时依赖沿用现有 ONNX 体系。 diff --git a/openspec/changes/simplify-ai-model-tiers/retirement-results.json b/openspec/changes/simplify-ai-model-tiers/retirement-results.json new file mode 100644 index 0000000..a60d1c5 --- /dev/null +++ b/openspec/changes/simplify-ai-model-tiers/retirement-results.json @@ -0,0 +1,267 @@ +{ + "results": [ + { + "kind": "r2", + "path": "DA3-BASE/model.onnx", + "removed": true + }, + { + "kind": "r2", + "path": "DA3-SMALL/model.onnx", + "removed": true + }, + { + "kind": "r2", + "path": "moge-2-vitl-normal-onnx/model.onnx", + "removed": true + }, + { + "kind": "r2", + "path": "pisa-sr/LICENSE-APACHE.txt", + "removed": true + }, + { + "kind": "r2", + "path": "pisa-sr/LICENSE-MODEL.txt", + "removed": true + }, + { + "kind": "r2", + "path": "pisa-sr/NOTICE.txt", + "removed": true + }, + { + "kind": "r2", + "path": "pisa-sr/decoder.onnx", + "removed": true + }, + { + "kind": "r2", + "path": "pisa-sr/encoder.onnx", + "removed": true + }, + { + "kind": "r2", + "path": "pisa-sr/noise.npy", + "removed": true + }, + { + "kind": "r2", + "path": "pisa-sr/unet-full.onnx", + "removed": true + }, + { + "kind": "r2", + "path": "pisa-sr/unet-pixel.onnx", + "removed": true + }, + { + "kind": "r2", + "path": "safmn-real-x4/safmn-real-x4.onnx", + "removed": true + }, + { + "kind": "local", + "path": "D:\\Projects\\AnyImage\\models\\DA3-BASE", + "removed": true + }, + { + "kind": "local", + "path": "D:\\Projects\\AnyImage\\models\\DA3-SMALL", + "removed": true + }, + { + "kind": "local", + "path": "D:\\Projects\\AnyImage\\models\\moge-2-vitl-normal-onnx", + "removed": true + }, + { + "kind": "local", + "path": "D:\\Projects\\AnyImage\\models\\safmn-real-x4", + "removed": true + }, + { + "kind": "local", + "path": "D:\\Projects\\AnyImage\\models\\pisa-sr", + "removed": true + }, + { + "kind": "local", + "path": "C:\\Users\\icrdr\\.anyimage\\models\\moge-2-vitl-normal-onnx", + "removed": true + }, + { + "kind": "local", + "path": "C:\\Users\\icrdr\\.anyimage\\models\\safmn-real-x4", + "removed": true + }, + { + "kind": "local", + "path": "C:\\Users\\icrdr\\.anyimage\\models\\pisa-sr", + "removed": true + }, + { + "kind": "local", + "path": "C:\\Users\\icrdr\\.cache\\huggingface\\hub\\models--depth-anything--DA3-BASE", + "removed": true + }, + { + "kind": "local", + "path": "C:\\Users\\icrdr\\.cache\\huggingface\\hub\\models--depth-anything--DA3-GIANT-1.1", + "removed": true + }, + { + "kind": "local", + "path": "C:\\Users\\icrdr\\.cache\\huggingface\\hub\\.locks\\models--depth-anything--DA3-GIANT-1.1", + "removed": true + }, + { + "kind": "local", + "path": "C:\\Users\\icrdr\\.cache\\huggingface\\hub\\models--depth-anything--DA3-LARGE-1.1", + "removed": true + }, + { + "kind": "local", + "path": "C:\\Users\\icrdr\\.cache\\huggingface\\hub\\models--depth-anything--DA3-SMALL", + "removed": true + }, + { + "kind": "local", + "path": "C:\\Users\\icrdr\\.cache\\huggingface\\hub\\models--depth-anything--DA3NESTED-GIANT-LARGE-1.1", + "removed": true + }, + { + "kind": "local", + "path": "C:\\Users\\icrdr\\.cache\\huggingface\\hub\\models--mindqtrl--DA3-LARGE-1.1-ONNX", + "removed": true + }, + { + "kind": "local", + "path": "C:\\Users\\icrdr\\.cache\\huggingface\\hub\\.locks\\models--mindqtrl--DA3-LARGE-1.1-ONNX", + "removed": true + }, + { + "kind": "local", + "path": "C:\\Users\\icrdr\\.cache\\huggingface\\hub\\models--Ruicheng--moge-2-vitl", + "removed": true + }, + { + "kind": "local", + "path": "C:\\Users\\icrdr\\.cache\\huggingface\\hub\\models--Ruicheng--moge-2-vitl-normal", + "removed": true + }, + { + "kind": "local", + "path": "C:\\Users\\icrdr\\.cache\\huggingface\\hub\\models--Ruicheng--moge-2-vitl-normal-onnx", + "removed": true + } + ], + "remote_remaining": [], + "retained_changed": [], + "retained_objects": [ + { + "Path": "BEN2-ONNX/onnx/model_fp16.onnx", + "Name": "model_fp16.onnx", + "Size": 219121675, + "MimeType": "application/octet-stream", + "ModTime": "2026-08-14T09:16:42.426271600+08:00", + "IsDir": false + }, + { + "Path": "birefnet-hr-matting/LICENSE.txt", + "Name": "LICENSE.txt", + "Size": 1066, + "MimeType": "text/plain; charset=utf-8", + "ModTime": "2026-09-10T18:25:27.626890000+08:00", + "IsDir": false + }, + { + "Path": "birefnet-hr-matting/model.onnx", + "Name": "model.onnx", + "Size": 529467341, + "MimeType": "application/octet-stream", + "ModTime": "2026-09-10T18:25:26.979196800+08:00", + "IsDir": false + }, + { + "Path": "birefnet-lite/LICENSE.txt", + "Name": "LICENSE.txt", + "Size": 1066, + "MimeType": "text/plain; charset=utf-8", + "ModTime": "2026-09-10T18:14:05.058461000+08:00", + "IsDir": false + }, + { + "Path": "birefnet-lite/model.onnx", + "Name": "model.onnx", + "Size": 112968598, + "MimeType": "application/octet-stream", + "ModTime": "2026-09-10T18:14:04.582138300+08:00", + "IsDir": false + }, + { + "Path": "hat-gan-x4-sharper/LICENSE.txt", + "Name": "LICENSE.txt", + "Size": 11342, + "MimeType": "text/plain; charset=utf-8", + "ModTime": "2026-09-10T16:58:24.978566500+08:00", + "IsDir": false + }, + { + "Path": "hat-gan-x4-sharper/model.onnx", + "Name": "model.onnx", + "Size": 45874438, + "MimeType": "application/octet-stream", + "ModTime": "2026-09-10T16:58:24.123871100+08:00", + "IsDir": false + }, + { + "Path": "moge-2-vitb-normal-onnx/model.onnx", + "Name": "model.onnx", + "Size": 419411850, + "MimeType": "application/octet-stream", + "ModTime": "2026-08-14T00:42:46.699086800+08:00", + "IsDir": false + }, + { + "Path": "moge-2-vits-normal-onnx/model.onnx", + "Name": "model.onnx", + "Size": 140852051, + "MimeType": "application/octet-stream", + "ModTime": "2026-08-14T00:27:11.713132700+08:00", + "IsDir": false + }, + { + "Path": "moge-3-vitl-onnx/backbone.onnx", + "Name": "backbone.onnx", + "Size": 1324190318, + "MimeType": "application/octet-stream", + "ModTime": "2026-09-10T10:38:13.252671700+08:00", + "IsDir": false + }, + { + "Path": "moge-3-vitl-onnx/refiner.onnx", + "Name": "refiner.onnx", + "Size": 157729628, + "MimeType": "application/octet-stream", + "ModTime": "2026-09-10T10:38:07.481681400+08:00", + "IsDir": false + }, + { + "Path": "realesr-general-wdn-x4v3/realesr-general-wdn-x4v3.onnx", + "Name": "realesr-general-wdn-x4v3.onnx", + "Size": 4867419, + "MimeType": "application/octet-stream", + "ModTime": "2026-08-23T00:00:31.076651000+08:00", + "IsDir": false + }, + { + "Path": "realesrgan-x4plus/realesrgan-x4plus.onnx", + "Name": "realesrgan-x4plus.onnx", + "Size": 67051618, + "MimeType": "application/octet-stream", + "ModTime": "2026-08-23T00:02:33.710120800+08:00", + "IsDir": false + } + ] +} diff --git a/openspec/changes/simplify-ai-model-tiers/retirement.json b/openspec/changes/simplify-ai-model-tiers/retirement.json new file mode 100644 index 0000000..555dee8 --- /dev/null +++ b/openspec/changes/simplify-ai-model-tiers/retirement.json @@ -0,0 +1,131 @@ +{ + "remote": "omoolab-r2:ai-models", + "local": [ + { + "path": "D:\\Projects\\AnyImage\\models\\DA3-BASE", + "owner": "D:\\Projects\\AnyImage\\models" + }, + { + "path": "D:\\Projects\\AnyImage\\models\\DA3-SMALL", + "owner": "D:\\Projects\\AnyImage\\models" + }, + { + "path": "D:\\Projects\\AnyImage\\models\\moge-2-vitl-normal-onnx", + "owner": "D:\\Projects\\AnyImage\\models" + }, + { + "path": "D:\\Projects\\AnyImage\\models\\safmn-real-x4", + "owner": "D:\\Projects\\AnyImage\\models" + }, + { + "path": "D:\\Projects\\AnyImage\\models\\pisa-sr", + "owner": "D:\\Projects\\AnyImage\\models" + }, + { + "path": "C:\\Users\\icrdr\\.anyimage\\models\\moge-2-vitl-normal-onnx", + "owner": "C:\\Users\\icrdr\\.anyimage\\models" + }, + { + "path": "C:\\Users\\icrdr\\.anyimage\\models\\safmn-real-x4", + "owner": "C:\\Users\\icrdr\\.anyimage\\models" + }, + { + "path": "C:\\Users\\icrdr\\.anyimage\\models\\pisa-sr", + "owner": "C:\\Users\\icrdr\\.anyimage\\models" + }, + { + "path": "C:\\Users\\icrdr\\.cache\\huggingface\\hub\\models--depth-anything--DA3-BASE", + "owner": "C:\\Users\\icrdr\\.cache\\huggingface\\hub" + }, + { + "path": "C:\\Users\\icrdr\\.cache\\huggingface\\hub\\models--depth-anything--DA3-GIANT-1.1", + "owner": "C:\\Users\\icrdr\\.cache\\huggingface\\hub" + }, + { + "path": "C:\\Users\\icrdr\\.cache\\huggingface\\hub\\.locks\\models--depth-anything--DA3-GIANT-1.1", + "owner": "C:\\Users\\icrdr\\.cache\\huggingface\\hub\\.locks" + }, + { + "path": "C:\\Users\\icrdr\\.cache\\huggingface\\hub\\models--depth-anything--DA3-LARGE-1.1", + "owner": "C:\\Users\\icrdr\\.cache\\huggingface\\hub" + }, + { + "path": "C:\\Users\\icrdr\\.cache\\huggingface\\hub\\models--depth-anything--DA3-SMALL", + "owner": "C:\\Users\\icrdr\\.cache\\huggingface\\hub" + }, + { + "path": "C:\\Users\\icrdr\\.cache\\huggingface\\hub\\models--depth-anything--DA3NESTED-GIANT-LARGE-1.1", + "owner": "C:\\Users\\icrdr\\.cache\\huggingface\\hub" + }, + { + "path": "C:\\Users\\icrdr\\.cache\\huggingface\\hub\\models--mindqtrl--DA3-LARGE-1.1-ONNX", + "owner": "C:\\Users\\icrdr\\.cache\\huggingface\\hub" + }, + { + "path": "C:\\Users\\icrdr\\.cache\\huggingface\\hub\\.locks\\models--mindqtrl--DA3-LARGE-1.1-ONNX", + "owner": "C:\\Users\\icrdr\\.cache\\huggingface\\hub\\.locks" + }, + { + "path": "C:\\Users\\icrdr\\.cache\\huggingface\\hub\\models--Ruicheng--moge-2-vitl", + "owner": "C:\\Users\\icrdr\\.cache\\huggingface\\hub" + }, + { + "path": "C:\\Users\\icrdr\\.cache\\huggingface\\hub\\models--Ruicheng--moge-2-vitl-normal", + "owner": "C:\\Users\\icrdr\\.cache\\huggingface\\hub" + }, + { + "path": "C:\\Users\\icrdr\\.cache\\huggingface\\hub\\models--Ruicheng--moge-2-vitl-normal-onnx", + "owner": "C:\\Users\\icrdr\\.cache\\huggingface\\hub" + } + ], + "objects": [ + { + "path": "DA3-BASE/model.onnx", + "size": 412856697 + }, + { + "path": "DA3-SMALL/model.onnx", + "size": 105484029 + }, + { + "path": "moge-2-vitl-normal-onnx/model.onnx", + "size": 1324265014 + }, + { + "path": "pisa-sr/LICENSE-APACHE.txt", + "size": 11358 + }, + { + "path": "pisa-sr/LICENSE-MODEL.txt", + "size": 14206 + }, + { + "path": "pisa-sr/NOTICE.txt", + "size": 1081 + }, + { + "path": "pisa-sr/decoder.onnx", + "size": 99084589 + }, + { + "path": "pisa-sr/encoder.onnx", + "size": 68426712 + }, + { + "path": "pisa-sr/noise.npy", + "size": 32896 + }, + { + "path": "pisa-sr/unet-full.onnx", + "size": 1732713007 + }, + { + "path": "pisa-sr/unet-pixel.onnx", + "size": 1732713007 + }, + { + "path": "safmn-real-x4/safmn-real-x4.onnx", + "size": 22657477 + } + ] +} diff --git a/openspec/changes/simplify-ai-model-tiers/specs/ai-model-tiers/spec.md b/openspec/changes/simplify-ai-model-tiers/specs/ai-model-tiers/spec.md new file mode 100644 index 0000000..7128e9b --- /dev/null +++ b/openspec/changes/simplify-ai-model-tiers/specs/ai-model-tiers/spec.md @@ -0,0 +1,42 @@ +## Purpose + +为深度生成、图片超分和背景移除提供各三款用途明确的正式模型，使偏好选择、首次安装、模型下载和服务就绪状态表达相同的模型集合，并让新用户以轻量模型开始使用。 + +## ADDED Requirements + +### Requirement: 三类模型采用固定三档集合 +系统 SHALL 按以下顺序提供九款正式模型：Depth 为 MoGe-2 ViT-S Normal、MoGe-2 ViT-B Normal、MoGe-3 ViT-L；Upscale 为 Real-ESRGAN General WDN x4v3、Real-ESRGAN x4plus、Real HAT GAN x4 Sharper；Background 为 BiRefNet Lite、BEN2、BiRefNet HR-matting。各类别第一项 SHALL 为新偏好的默认值。 + +#### Scenario: 查看模型选项 +- **WHEN** 用户查看偏好或模型下载列表 +- **THEN** 每个类别恰好显示上述三款模型，顺序一致，模型名、用途、许可证和下载体积对应实际资源 + +#### Scenario: 服务离线时显示目录 +- **WHEN** 服务尚未启动而界面需要展示模型 +- **THEN** 本地目录与服务启动后返回的正式模型集合一致 + +### Requirement: 首次安装仅要求三个默认模型 +系统 SHALL 将 MoGe-2 ViT-S Normal、General WDN x4v3 和 BiRefNet Lite 作为默认环境初始化所需的模型集合，其他模型按用户选择下载。 + +#### Scenario: 新安装完成初始化 +- **WHEN** 推理环境和三个默认模型都已就绪，其他六款尚未下载 +- **THEN** AI 初始化状态为就绪，默认功能可运行 + +#### Scenario: 选择未下载的高档模型 +- **WHEN** 用户执行功能时选中的非默认模型尚未就绪 +- **THEN** 系统指出需要下载该模型，不使用其他模型替代本次请求 + +### Requirement: 有效偏好保留且退役选择明确处理 +系统 SHALL 保留仍属于所选类别的持久化模型键，并将退役或无效的偏好值归一为类别默认值；服务收到未知或退役模型键时 SHALL 返回明确错误。 + +#### Scenario: 用户已有 BEN2 偏好 +- **WHEN** 更新后读取有效的 BEN2 背景移除选择 +- **THEN** 继续选择 BEN2 + +#### Scenario: 用户已有 PiSA-SR 偏好 +- **WHEN** 更新后读取已退役的 PiSA-SR 超分选择 +- **THEN** 当前有效选择为 General WDN x4v3，界面与后续请求一致 + +#### Scenario: 旧请求直接指定退役模型 +- **WHEN** 服务收到指定 DA3、MoGe-2 L、SAFMN 或 PiSA-SR 的请求 +- **THEN** 返回未知或不可用模型错误，不加载或下载退役资源 diff --git a/openspec/changes/simplify-ai-model-tiers/specs/background-removal-models/spec.md b/openspec/changes/simplify-ai-model-tiers/specs/background-removal-models/spec.md new file mode 100644 index 0000000..2ce8d9b --- /dev/null +++ b/openspec/changes/simplify-ai-model-tiers/specs/background-removal-models/spec.md @@ -0,0 +1,59 @@ +## Purpose + +让用户通过同一个背景移除功能选择轻量、通用或精细抠图模型，并在静态图、多帧图与 HDR 图片中得到尺寸、透明度和资源生命周期一致的结果，保留已有图片编辑语义。 + +## ADDED Requirements + +### Requirement: 背景模型选择贯穿真实执行 +系统 SHALL 使用当前选择的 BiRefNet Lite、BEN2 或 BiRefNet HR-matting 执行所有支持的背景移除输入路径，模型名称 SHALL 与下载、进度和实际加载的资源一致。 + +#### Scenario: 切换到精细抠图 +- **WHEN** 用户选择已经下载的 BiRefNet HR-matting 并执行背景移除 +- **THEN** 使用该模型推理，输出和进度不再固定指向 BEN2 + +### Requirement: 输出有限的连续透明度 +每款背景模型 SHALL 产出与分析图片尺寸一致、有限且位于 `[0,1]` 的连续 Alpha。模型适配 SHALL 保持其官方输出语义，支持发丝、毛绒和半透明边缘的连续透明度。 + +#### Scenario: 半透明边缘 +- **WHEN** HR-matting 对包含薄纱或细发丝的图片输出中间透明度 +- **THEN** 后处理保留有效的中间 Alpha 值，不将其强制二值化或按单图极值重新拉伸 + +#### Scenario: 无效预测 +- **WHEN** 模型返回错误形状或包含 NaN、Inf 的透明度 +- **THEN** 任务明确失败，并保留可诊断的模型错误信息 + +### Requirement: 保持已有前景和 HDR 编辑行为 +系统 SHALL 沿用既有前景 RGB 与源 Alpha 合成规则，支持静态前景、多帧前景及 HDR 的 Alpha-only 输出；HDR 结果 SHALL 保持原始浮点 RGB。 + +#### Scenario: HDR 背景移除 +- **WHEN** 用户对浮点 HDR 图片使用任意背景模型 +- **THEN** 用分析输入预测 Alpha，将结果应用到原始图片，并保留原始 HDR RGB 数值 + +#### Scenario: 多帧透明前景 +- **WHEN** 用户对支持的多帧输入执行背景移除 +- **THEN** 每帧使用选定模型，保持帧数、编号、尺寸与已有透明度合成约定 + +### Requirement: 模型资源可复用且可释放 +系统 SHALL 为相同模型和设备复用会话，在模型切换、资源关闭时释放旧会话，并使取消与失败遵循既有任务清理行为。 + +#### Scenario: 连续帧复用 +- **WHEN** 多帧任务连续处理同一模型和设备 +- **THEN** 各帧复用已加载会话，取消后停止后续推理并清理未完成任务产物 + +#### Scenario: 切换背景模型 +- **WHEN** 用户从 Lite 切换到 BEN2 或 HR-matting +- **THEN** 实际推理使用新选择，并释放原先背景模型占用的会话资源 + +### Requirement: 新背景模型经过真实运行验收 +BiRefNet Lite SHALL 在发布前通过 CPU / DirectML 数值和图片验证，HR-matting SHALL 通过 CPU 验证；记录真实下载体积、加载与热推理耗时、峰值资源占用及边缘效果。Lite SHALL 比 BEN2 下载体积更小，并在相同测试条件下至少降低热推理耗时或峰值占用中的一项。 + +#### Scenario: 确定轻量默认模型 +- **WHEN** Lite 的正式资产准备发布 +- **THEN** 存在与 BEN2 的同条件对照和质量检查记录，且其轻量收益满足要求 + +### Requirement: HR-matting 使用 CPU +HR-matting SHALL 暂时固定使用 CPU。模型选项 SHALL 标明 CPU，进度 SHALL 显示实际设备；会话缓存 SHALL 按实际 CPU 设备复用。 + +#### Scenario: GPU 偏好下选择 HR-matting +- **WHEN** 全局设备为 DirectML，用户选择 HR-matting +- **THEN** 创建 CPU 会话并显示 CPU 进度，保持 HR 的 2048 像素输入及连续 Alpha 语义 diff --git a/openspec/changes/simplify-ai-model-tiers/specs/hat-upscale/spec.md b/openspec/changes/simplify-ai-model-tiers/specs/hat-upscale/spec.md new file mode 100644 index 0000000..c32fdfb --- /dev/null +++ b/openspec/changes/simplify-ai-model-tiers/specs/hat-upscale/spec.md @@ -0,0 +1,49 @@ +## Purpose + +为图片超分增加用户已比较确认的 HAT Sharper 细节增强选项，提供可验证的小体积模型资源，并使最终两倍尺寸、透明度、多帧处理和大图分块行为与现有超分功能一致。 + +## ADDED Requirements + +### Requirement: HAT Sharper 提供可验证的小体积资源 +系统 SHALL 提供官方 Real HAT GAN x4 Sharper 对应的模型资产，完整记录上游来源、文件大小及校验信息。发布文件 SHALL 小于 60 MB，保持 FP32 输入输出，并通过 CPU / DirectML 与原始 FP32 参考的数值对照。 + +#### Scenario: 校验发布模型 +- **WHEN** 对预先固定的样图执行正式导出资产 +- **THEN** 输出全部有限，未裁剪的 `[0,1]` 标度输出误差满足 MAE ≤ 0.0005、P99 ≤ 0.005、最大误差 ≤ 0.02 + +#### Scenario: 实验精度路径失败 +- **WHEN** 导出或运行路径在验收中产生 NaN、Inf 或超过预设误差 +- **THEN** 该资产不能登记为可发布模型，也不能标记验收完成 + +### Requirement: 三款超分均保持最终两倍输出 +系统 SHALL 对限制后的分析输入保持宽高各 2 倍的最终输出；HAT 与现有两款 Real-ESRGAN SHALL 保持相同的透明度合成和帧输出约定。 + +#### Scenario: 普通与透明静态图片 +- **WHEN** 用户用 HAT Sharper 放大宽 W、高 H 的分析输入 +- **THEN** 最终图片尺寸为 `2W×2H`，原始 Alpha 按现有超分规则缩放并合成 + +#### Scenario: 多帧输入 +- **WHEN** 使用 HAT Sharper 处理支持的多帧输入 +- **THEN** 每帧输出保持两倍尺寸及既有编号，帧间复用模型资源 + +### Requirement: 分块支持合法图片尺寸且控制边界误差 +HAT SHALL 支持小图、奇数尺寸和跨多个分块的输入，输出覆盖全部像素；分块配置 SHALL 通过纹理、细线和透明边缘样图检查，并相对完整参考记录接缝误差。 + +#### Scenario: 奇数尺寸与边缘块 +- **WHEN** 输入宽高不是分块步长的整数倍 +- **THEN** 结果保持准确的两倍尺寸，边缘完整，没有空白条带或重复像素区 + +#### Scenario: 自然纹理跨越接缝 +- **WHEN** 对固定跨缝样图检查最终 2× 输出 +- **THEN** 没有由拼接造成的可见直线接缝，并记录相对整图参考的接缝和内部误差 + +### Requirement: 推理取消和资源错误可处理 +系统 SHALL 沿用超分任务的取消、会话复用和资源释放机制；资源不足时 SHALL 报告当前模型及可操作的调整建议。 + +#### Scenario: 分块间取消 +- **WHEN** 用户取消正在处理多块图片的 HAT 任务 +- **THEN** 后续块停止推理，未完成产物按现有任务规则清理 + +#### Scenario: 显存或内存不足 +- **WHEN** 当前设备无法完成 HAT 推理 +- **THEN** 返回包含模型和设备上下文的资源错误，提示减小分析输入或选择更轻量模型 diff --git a/openspec/changes/simplify-ai-model-tiers/specs/model-asset-retirement/spec.md b/openspec/changes/simplify-ai-model-tiers/specs/model-asset-retirement/spec.md new file mode 100644 index 0000000..73a4e69 --- /dev/null +++ b/openspec/changes/simplify-ai-model-tiers/specs/model-asset-retirement/spec.md @@ -0,0 +1,38 @@ +## Purpose + +确保正式模型目录、本地资产准备和远端镜像对应同一套可验证资源，并在模型退役时明确清理其源码登记、开发权重缓存及远端对象，使用户下载和开发工具持续使用有效模型。 + +## ADDED Requirements + +### Requirement: 资产准备与发布跟随正式清单 +`model prepare` 和 `model sync` SHALL 覆盖九款正式模型的完整文件集合，对每个文件验证大小和 SHA-256；缺失、损坏或无法验证的资产 SHALL 明确失败。 + +#### Scenario: 新模型资产损坏 +- **WHEN** HAT 或 BiRefNet 的本地文件与登记大小或哈希不一致 +- **THEN** 准备或发布失败，该模型不被标记为完整就绪 + +#### Scenario: 新镜像验收 +- **WHEN** 新模型上传到 R2 +- **THEN** 使用实际下载来源回读并校验完整文件，记录每款模型的来源、大小、哈希和就绪结果 + +### Requirement: 退役资源在各发布入口一致移除 +DA3 全部现存变体、MoGe-2 ViT-L Normal、SAFMN 和 PiSA-SR SHALL 从正式模型集合、自动准备、下载和生产推理入口退役；对应专用源码、导出登记与专用构建依赖 SHALL 清理。 + +#### Scenario: 执行正式准备命令 +- **WHEN** 完成本次模型调整后运行模型准备或同步 +- **THEN** 只准备和发布九款正式模型，不下载、导出或上传退役模型 + +### Requirement: 远端和本地退役清理采用明确对象清单 +本次退役清理 SHALL 在新模型通过发布验收后执行。清理前 SHALL 列出并保存归属于退役模型的本地目录、开发缓存文件与 R2 对象键，仅操作核对过的目标，并在结束后复查对象存在性。 + +#### Scenario: 清理 DA3 远端残留 +- **WHEN** R2 中仍存在未列入正式模型清单的 DA3 变体 +- **THEN** 将归属已确认的 DA3 对象纳入本次退役清单、删除并复查，不能将停止上传视为删除完成 + +#### Scenario: 相邻目录包含其他资源 +- **WHEN** 模型目录或桶内同时存在保留模型、其他对象或用户产物 +- **THEN** 清理范围严格限定于已列出的退役目标，其他资源保持原样 + +#### Scenario: 部分对象清理失败 +- **WHEN** 本地或 R2 删除未能完整完成 +- **THEN** 记录成功与失败目标，报告未完成项，并在复查全部通过前保持退役任务未完成 diff --git a/openspec/changes/simplify-ai-model-tiers/tasks.md b/openspec/changes/simplify-ai-model-tiers/tasks.md new file mode 100644 index 0000000..c69145b --- /dev/null +++ b/openspec/changes/simplify-ai-model-tiers/tasks.md @@ -0,0 +1,41 @@ +## 1. 核对基础与验证 BiRefNet + +- [x] 1.1 核对当前 MoGe-3、2× Upscale、HDR 和工具 CLI 的真实调用链，记录受影响入口及工作区已有修改，确认模型矩阵与相关 change 的衔接。 +- [x] 1.2 固定 BiRefNet Lite 的上游版本、权重和许可证，完成隔离环境的可复现导出；交付通过结构检查的 ONNX、来源、大小和 SHA-256。 +- [x] 1.3 固定 BiRefNet HR-matting 的上游版本、权重和许可证，完成隔离环境的可复现导出；交付通过结构检查的 ONNX、来源、大小和 SHA-256。 +- [x] 1.4 在运行前固定 BiRefNet 数值容差和样图，对 Lite 执行原始实现与 CPU / DirectML 的逐图比较，对 HR-matting 执行 CPU 比较并记录 DirectML 失败及用户确认的 CPU 限制；验证连续 Alpha、尺寸、有限值及官方预处理一致性，记录结果。 +- [x] 1.5 用同一组发丝、毛绒、细杆、薄纱与普通主体样图比较 Lite、BEN2 和 HR-matting，记录下载体积、加载/热推理耗时、峰值内存/显存及边缘效果；验证 Lite 满足轻量默认门槛。 + +## 2. 固化 HAT Sharper 资产 + +- [x] 2.1 将已验证的 Sharper 导出流程整理到 `tools/models`，固定源码和 checkpoint，采用 FP16 权重存储、FP32 运算；验证完整资产小于 60 MB 并记录可复现来源与哈希。 +- [x] 2.2 对正式导出资产复跑 CPU / DirectML 数值对照，检查全部有限且符合 `hat-upscale` 预设门槛；交付样图、误差、加载和热推理记录。 +- [x] 2.3 以 256 像素 tile、48 像素边距验证 HAT 拼接，覆盖跨缝纹理、细线、小图、奇数尺寸和透明图；交付与整图参考的对照，确认最终 2× 无可见拼接直线。 + +## 3. 统一三档目录与默认设置 + +- [x] 3.1 将正式目录调整为九款模型及对应分类、用途和真实资源信息；更新目录与下载测试，验证每类恰好三款且保留顺序一致。 +- [x] 3.2 将默认模型及扩展偏好读取统一到对应入口，背景默认改为 Lite；测试有效旧选择保留、退役偏好归一和直接提交退役 Job 键报错。 +- [x] 3.3 更新 AI Setup、默认所需模型、离线目录与下载界面；测试仅三个默认模型就绪时初始化完成，非默认模型按需提示下载。 + +## 4. 接入背景移除与超分 + +- [x] 4.1 将背景移除公共编排、模型资源分类和缓存接口按业务职责整理，保留各模型专用适配器；运行会话复用、模型/设备切换和关闭释放测试。 +- [x] 4.2 接入 BiRefNet 官方输入和 Alpha 后处理，更新 Job、Operator、进度文案与实际模型选择；测试三款真实分派、有限连续 Alpha 和无效预测失败。 +- [x] 4.3 让三款背景模型贯穿静态、多帧、源 Alpha 和 HDR Alpha-only 流程；运行对应行为测试，验证 HDR RGB 保持以及取消/失败清理。 +- [x] 4.4 将 HAT Sharper 接入共享超分及模型专属分块配置，保持两款 Real-ESRGAN 的已有配置；测试准确 2× 尺寸、透明度、多帧、取消和资源错误。 +- [x] 4.5 运行真实背景移除与 HAT Job，验证模型切换、重复执行和服务关闭后的资源状态；记录正式输入准备和输出路径的端到端结果。 + +## 5. 清理退役实现与准备发布 + +- [x] 5.1 清理 DA3 现行残留及 MoGe-2 L、SAFMN、PiSA-SR 的专用实现、旧导出器、构建登记和测试；沿真实调用链检查旧引用，确认保留模型仍有有效行为覆盖。 +- [x] 5.2 清理退役模型的专用依赖并同步必要锁文件，核对 Manifest 和打包清单；运行依赖及打包内容测试，确认新增导出依赖隔离于正式 Runtime。 +- [x] 5.3 将三款新资产纳入 `model prepare` / `model sync`，执行 `uv run model prepare`；验证九款模型完整文件集合、损坏文件报错和校验信息。 +- [x] 5.4 运行全部相关测试及 `uv run pytest`，确认正常退出；检查最终差异和注册清单，确保原有有效功能与其他工作区修改完整。 + +## 6. 镜像发布与退役对象清理 + +- [x] 6.1 执行 `uv run model sync` 发布正式资产，从实际下载地址回读新模型并校验大小和哈希；确认下载后的模型可就绪和执行。 +- [x] 6.2 枚举 DA3 全部现存变体、MoGe-2 L、SAFMN、PiSA-SR 的本地模型目录、开发缓存及 R2 对象；交付具体删除清单，验证本地绝对路径和远端前缀归属。 +- [x] 6.3 在新资产验收后按清单删除退役对象，逐项复查本地与远端状态；确认保留模型和其他资源完整，失败项保留为未完成任务。 +- [x] 6.4 将实际模型矩阵、导出来源与哈希、真实推理和质量记录、测试结果及镜像清理结果写入本 change 的验证记录，确认各项任务的验收证据齐全。 diff --git a/openspec/changes/simplify-ai-model-tiers/validation.md b/openspec/changes/simplify-ai-model-tiers/validation.md new file mode 100644 index 0000000..0f00791 --- /dev/null +++ b/openspec/changes/simplify-ai-model-tiers/validation.md @@ -0,0 +1,94 @@ +# 实施验证 + +## 基线与入口 + +开始实施时 `git status --short` 仅有本 change 未跟踪，生产代码无已有未提交修改。 + +- Geometry：Operator → Geometry Job → ModelManager → MoGe-2 / MoGe-3 ONNX；保留 MoGe-3 三步细化。 +- Upscale：Operator → Upscale Job → `models/upscale.py` → ONNX 分块 → Lanczos 2× 与源 Alpha 合成。 +- 背景移除：Operator 当前硬编码 BEN2；Job 调用 `models/ben2.py`，HDR 已使用 Alpha-only 返回，输入准备复用 `common`。 +- 模型入口：`preferences.py`、`properties.py`、`server/model_catalog.py`、`server/model_manager.py`；背景资源分类当前为 `ben2`。 +- 资产入口：`uv run model prepare` / `uv run model sync` → `tools/models/cli.py`；上传为逐文件 `rclone copyto`。 +- 当前生产源码未发现 DA3 实现；本地及 R2 残余仍需枚举。历史 OpenSpec 记录不作为现行调用链。 + +衔接 `add-moge3-model`、`add-pisa-sr-two-times-upscale`、`add-ben2-hdr-support`、`organize-tool-cli`，模型成员以本 change 为准。 + +## BiRefNet 来源 + +| 模型 | 官方权重与源码仓库 | 固定 revision | +|---|---|---| +| Lite | ZhengPeng7/BiRefNet_lite | aa62cd87eafb9cc43056d08ef3615a14628b831d | +| HR-matting | ZhengPeng7/BiRefNet_HR-matting | 5d6b6f8adcb5b417c871b1d84ceaae9871355b7f | + +两款模型卡标注 MIT，正式导出结构检查通过，许可证随模型发布。 + +Lite 固定权重 SHA-256：`4417d89795250e698c3cb0ae8df15743810065f646f48a694fdfa7ca052d0815`；HR：`a5a4de698739ea5e0e8bbab28e1b293dde95092b87a442d566cbc585c53cef55`。许可证固定到 BiRefNet revision `ebcc0bc8ec7fe919cec829f2dea656b3078acddc`，1,066 bytes，SHA-256 `92a7089e0915fc32bc40067560b398f1e6a7a5958abd7d04eda393629a5acefb`。 + +数值测试开始前固定 Alpha 容差（激活后、浮点范围 0–1）：FP32 ONNX 对原始 FP32 实现的 MAE ≤ 0.0002、P99 ≤ 0.002、Max ≤ 0.01；压缩权重资产 MAE ≤ 0.002、P99 ≤ 0.02、Max ≤ 0.1。所有输出必须有限、尺寸准确、处于 0–1。容差不得根据测得误差放宽。 + +固定样图使用作者官方演示中的 Helicopter、Jewelry、My_Love、My_MiSheng、Windmill，另补发丝、毛绒及薄纱覆盖；记录来源与逐图结果。性能测量使用同一批源图、同一 Runtime/设备并分开运行模型。 + +补充 Alpha Matting Evaluation 官方测试集中的 `net.png`（薄网纱）和 `plasticbag.png`（透明塑料）；来源 `https://www.alphamatting.com/datasets.php`，仅用于本地评估。 + +Lite 通过七图 CPU / DirectML 对照的正式资产：112,968,598 bytes，SHA-256 `6e2a9844c57b080e95a402722ddcae5660e6f254a04d3e003347f06f587b71fe`。FP16 权重存储与完全无损的静态常量压缩，FP32 运算。最大误差 0.041951，符合预设门槛。CPU 热推理约 2.85 秒；同输入 BEN2 CPU 约 3.84–3.90 秒，Lite 下载成本更低且 CPU 更快。正式 DirectML 与端到端结果见下表。 + +HR 的 DirectML 最后一层解码器有稳定偏差；用户于 2026-09-10 确认保留 HR，暂时仅支持 CPU。正式资产使用原始 GridSample 结构和 FP16 权重存储，不采用失败的分块或混合精度实验。HR 选项、实际 Job 设备和缓存键均明确使用 CPU。 + +## 公共背景编排 + +`models/background.py`、`get_background` / `release_background` 与资源分类 `background` 已统一。保留 `onnx_ben2.py` 专用适配器。三款模型已通过静态、多帧、连续 Alpha、源 Alpha 合成和缓存切换测试；既有 HDR RGB 保持、取消和失败清理测试通过。 + +## 正式背景模型验收 + +正式导出命令：`uv run --script tools/models/birefnet_export.py --variant lite --destination models/birefnet-lite`；HR 使用 `--variant hr-matting --destination models/birefnet-hr-matting`。导出器固定上游代码、权重与许可证。Lite 重复导出哈希完全相同。 + +| 模型 | ONNX bytes | SHA-256 | +|---|---:|---| +| Lite | 112968598 | `6e2a9844c57b080e95a402722ddcae5660e6f254a04d3e003347f06f587b71fe` | +| HR-matting | 529467341 | `b3abc34843a4d96f669ca9fb1a2a16b3417ab8b91c987573fdbda36d20514133` | + +ORT 1.24.4；DirectML 使用 RTX 5070 Ti。Lite 正式资产的七图 CPU / DirectML 误差通过既定门槛；最终 DirectML 记录确认实际加载 DmlExecutionProvider。HR 七图 CPU 与原始 FP32 对照全部通过，最差 MAE 0.000001987、P99 0.000020504、Max 0.000103504；输出均为有限的 2048×2048 连续 Alpha，生产适配器恢复源图尺寸。 + +| 模型 / 设备 | 加载秒 | 七图热推理秒 | 峰值工作集 GB | 峰值进程提交 GB | +|---|---:|---:|---:|---:| +| Lite / DirectML | 2.728 | 0.05895–0.05961 | 0.802 | 6.729 | +| BEN2 / DirectML | 1.154 | 0.11591–0.11801 | 0.604 | 6.519 | +| HR / CPU | 4.318 | 21.718–23.437 | 19.278 | 22.596 | + +Lite 与 BEN2 用同一七图、同一 Runtime 配置串行运行，每图三次，表中取每图最后一次。HR 每图两次，CPU 配置为八线程。工作集与提交量来自 Windows 进程计数器，不将其表述为独立显存测量。Lite 满足体积更小且热推理更快的默认门槛；其加载和峰值内存并不全面优于 BEN2。 + +视觉检查 `.model-cache/birefnet/quality-1.png`、`quality-2.png`：三款均能提取人物、猫、直升机和风车细杆；HR 更明显保留薄网纱和透明塑料的中间透明度，猫须细节更完整。普通珠宝场景 HR 会残留叶片背景，Lite / BEN2 的主体分离更干净。第三档表示高分辨率 matting 用途，不保证对所有主体都更好。源图与逐图 Alpha/PNG、耗时/误差 JSON 保存在 `.model-cache/birefnet`。 + +## 集成与回归 + +- 三类各三款，共九款；默认集合为 MoGe S、Lite、WDN。默认值由目录统一提供，偏好读取集中于 `preferences.py`。保留模型沿用原 Blender 枚举编号，有效 BEN2 / x4plus / MoGe 选择保留，退役偏好归一到默认，直接提交退役 Job 键报错。 +- `uv run model prepare` 成功，九款模型共 13 文件全部通过大小和 SHA-256 校验。 +- 真实 JobServer：Lite 连续两次、BEN2、HR CPU、HAT 连续两次全部成功，Alpha 源尺寸/有限范围及 HAT RGBA 准确 2× 输出检查通过。Lite 第二次加载 0 ms；HR 完整 Job 24.08 秒；HAT 两次完整 Job 3.41 / 0.60 秒。服务关闭后全部资源键为空。记录：`.model-cache/tier-job-validation/results.json`。 +- `uv run pytest -q --tb=short`：1227 passed，51 subtests passed，178.55 秒，正常退出。随后补齐退役偏好持久化归一，对应偏好与初始化测试再次通过：23 passed、2 subtests passed。曾因本地 bpy 安装缺少资源而失败，重新安装锁定的 bpy 4.5.3 后通过；渲染测试明确启用 Cycles。`pyproject.toml`、`uv.lock`、Manifest 未改变，新增模型导出依赖位于隔离脚本，正式 Runtime 无新增依赖。 +- 实施过程中另有节点源文件、节点测试和 `.blend` 资产并行修改；保留这些变化，本任务未修改对应节点逻辑。 + +## 镜像与退役清单 + +`uv run model sync` 正常完成。三款新模型共六个文件从 `https://models.omoolab.xyz` 实际地址回读，大小、SHA-256 和就绪检查全部通过；回读资产依次通过 Lite DirectML、HR CPU、HAT DirectML 生产适配器推理，关闭后缓存为空。完整 URL、大小、哈希与执行结果见 `mirror-validation.json`。 + +精确清理清单见 `retirement.json`：19 个本地模型/专用缓存目录、12 个 R2 对象。每个本地绝对路径已检查属于其声明的模型或缓存根目录。全部逐项删除并复查不存在，失败项为零；镜像保留的 13 文件大小与删除前一致，项目模型目录只剩九款。用户目录保留的 BEN2、WDN、x4plus 哈希复核通过。结果见 `retirement-results.json`。 + +清理后的 `uv run model prepare` 再次成功，九款模型共 13 文件大小和 SHA-256 全部通过。OpenSpec strict 校验通过，本次修改的 `git diff --check` 通过。 + +## HAT Sharper 正式导出 + +`uv run --script tools/models/hat_export.py --destination .model-cache/hat-sharper-export --checkpoint <已验证官方权重>` 成功退出。 + +- 上游：XPixelGroup/HAT，revision `1638a9a822581657811867bf670717f8371fc3e5`。 +- Checkpoint SHA-256：`5800b67136006eb8cab3b4ed7c8d73b6a195bb18e6cc709b674f9aa069c00271`，严格加载 `params_ema`。 +- ONNX：45,874,438 bytes，SHA-256 `032d8b03c704220a08d91465a822aca252c1d6e63db20c3cf23c5fb67f3d2366`，与原实验资产完全相同。 +- Apache 2.0 许可证随资产保存，完整资产小于 60 MB。结构检查通过。 +- 导出时去除 BasicSR 训练注册器，初始化使用等价 PyTorch 工具；官方 forward 保持不变。FP16 initializer 在 FP32 图计算前 Cast 恢复。 + +正式资产 ORT 1.24.4、CPU / DirectML 各 4 张 256×256 样图（bread、village、indoor、indoor_grain）逐图对照全部通过；每图执行两次。最差 MAE 0.000190、P99 0.000613、Max 0.007376，全部有限。CPU 加载 1.19 秒、热推理 14.43–14.90 秒；本次 DirectML 热推理 0.787–0.811 秒。性能为当前机器本次运行记录，不将历史不同负载下的耗时作为同条件对照。 + +详细记录：`.model-cache/hat-sharper-export/numerics.json`；复测脚本 `.model-cache/validate_hat_formal.py`；原始 PyTorch 参考与样图见 `.model-cache/hat-validation` 和原实验临时目录。 + +HAT 拼接验收：321×163 自然屋顶纹理、跨缝细线、渐变 Alpha，另有 7×5 和 1×1 小图。原始 PyTorch 整图使用与分块一致的反射边界，DirectML 用 256 tile / 48 border，均缩至最终 2×。检查整图参考、分块结果和 8 倍差异图，未见拼接直线。细线图接缝带 MAE 为 0.223（8-bit），与整图 MAE 0.213 接近。 + +对照图和逐图记录：`.model-cache/hat-sharper-export/seams/*-comparison.png`、`results.json`；脚本 `.model-cache/hat_seams.py`。参考图不是不同构图或不同缩放的替代品。 diff --git a/pyproject.toml b/pyproject.toml index 2476b04..caeda83 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -5,7 +5,7 @@ description = "Fast image-guided structure building tools for Blender" readme = "README.md" requires-python = ">=3.11,<3.12" license = "GPL-3.0-or-later" -dependencies = ["blendjob==0.1.16", "numpy>=1.26,<2.0", "scipy==1.15.3"] +dependencies = ["blendjob==0.1.17", "numpy>=1.26,<2.0", "scipy==1.15.3"] [dependency-groups] dev = [ @@ -38,6 +38,9 @@ build-backend = "hatchling.build" [tool.hatch.build.targets.wheel] packages = ["tools"] +[tool.uv.sources] +blendjob = { path = "src/anyimage/wheels/blendjob-0.1.17-py3-none-any.whl" } + [tool.pytest.ini_options] # Keep native asserts: bpy 4.5 can crash at shutdown after assertion rewriting. addopts = "--assert=plain" diff --git a/src/anyimage/assets/O_AnyImage.blend b/src/anyimage/assets/O_AnyImage.blend index 8ea94a7..958068e 100644 --- a/src/anyimage/assets/O_AnyImage.blend +++ b/src/anyimage/assets/O_AnyImage.blend @@ -1,3 +1,3 @@ version https://git-lfs.github.com/spec/v1 -oid sha256:3f384eed2b5de9bcc74b081e66a202891e00f262013adb6388dcf5896bf1e8e5 -size 6808659 +oid sha256:36c4cc107a924038dc1553045834194061a072e30afea8596f3372d8286b70f6 +size 7032020 diff --git a/src/anyimage/blender_manifest.toml b/src/anyimage/blender_manifest.toml index acffd5a..c61bdce 100644 --- a/src/anyimage/blender_manifest.toml +++ b/src/anyimage/blender_manifest.toml @@ -9,7 +9,7 @@ type = "add-on" blender_version_min = "4.5.0" platforms = ["windows-x64"] -wheels = ["./wheels/blendjob-0.1.16-py3-none-any.whl", "./wheels/scipy-1.15.3-cp311-cp311-win_amd64.whl", "./wheels/scipy-1.15.3-cp313-cp313-win_amd64.whl"] +wheels = ["./wheels/blendjob-0.1.17-py3-none-any.whl", "./wheels/scipy-1.15.3-cp311-cp311-win_amd64.whl", "./wheels/scipy-1.15.3-cp313-cp313-win_amd64.whl"] license = ["SPDX:GPL-3.0-or-later"] [permissions] diff --git a/src/anyimage/operators/ai_setup.py b/src/anyimage/operators/ai_setup.py index a311a8d..f9c91cb 100644 --- a/src/anyimage/operators/ai_setup.py +++ b/src/anyimage/operators/ai_setup.py @@ -2,8 +2,8 @@ import bpy from ..preferences import ADDON_PACKAGE from ..runtime import JobOperatorBase, runtime +from ..server.model_catalog import DEFAULT_MODEL_KEYS from ..properties import ( - REQUIRED_MODEL_KEYS, ai_status, model_download_size, model_record, @@ -42,7 +42,7 @@ class SetupAIEnvironment(bpy.types.Operator): if not status["environment_ready"]: required_size = sum( model_download_size(model_key) - for model_key in REQUIRED_MODEL_KEYS + for model_key in DEFAULT_MODEL_KEYS ) layout.label( text="AnyImage AI features run in an isolated Python environment." @@ -52,7 +52,7 @@ class SetupAIEnvironment(bpy.types.Operator): ) layout.label( text=( - "It also downloads MoGe-2 ViT-S Normal, BEN2 Base and " + "It also downloads MoGe-2 ViT-S Normal, BiRefNet Lite and " "Real-ESRGAN General WDN x4v3 " f"({format_download_size(required_size)})." ) @@ -154,7 +154,7 @@ class DownloadRequiredModels(JobOperatorBase, bpy.types.Operator): require_environment() if not bpy.app.online_access: raise RuntimeError("Online Access is disabled in Blender Preferences") - return {"models": list(REQUIRED_MODEL_KEYS)} + return {"models": list(DEFAULT_MODEL_KEYS)} class ClearModels(bpy.types.Operator): diff --git a/src/anyimage/operators/convert_to_panorama/operators.py b/src/anyimage/operators/convert_to_panorama/operators.py index 25576af..997bedf 100644 --- a/src/anyimage/operators/convert_to_panorama/operators.py +++ b/src/anyimage/operators/convert_to_panorama/operators.py @@ -6,7 +6,8 @@ from ...common.ai import invoke_ai_setup_if_needed, production_device, report_ai from ...common.image import is_image_empty, video_total_frames from ...common.color_image import prepare_material_color_input, material_analysis_input, cleanup_material_color_input from ...preferences import addon_preferences, configured_max_ai_input_size -from ...properties import DEFAULT_MOGE2_MODEL_KEY +from ...preferences import configured_geometry_model +from ...server.model_catalog import DEFAULT_GEOMETRY_MODEL_KEY from ...runtime import JobOperatorBase from ...server.media.resolution import DEFAULT_MAX_AI_INPUT_SIZE, MAX_MAX_AI_INPUT_SIZE, MIN_MAX_AI_INPUT_SIZE from .object import create_panorama_from_result @@ -26,7 +27,7 @@ class GeneratePanorama(JobOperatorBase, bpy.types.Operator): source_identity: bpy.props.StringProperty(options=_HIDDEN) image_identity: bpy.props.StringProperty(options=_HIDDEN) mesh_detail: bpy.props.IntProperty(default=7, min=0, max=8, options=_HIDDEN) - model: bpy.props.StringProperty(default=DEFAULT_MOGE2_MODEL_KEY, options=_HIDDEN) + model: bpy.props.StringProperty(default=DEFAULT_GEOMETRY_MODEL_KEY, options=_HIDDEN) resolution_level: bpy.props.IntProperty(default=5, min=0, max=9, options=_HIDDEN) max_input_size: bpy.props.IntProperty(default=DEFAULT_MAX_AI_INPUT_SIZE, min=MIN_MAX_AI_INPUT_SIZE, max=MAX_MAX_AI_INPUT_SIZE, options=_HIDDEN) @@ -82,7 +83,7 @@ class ConvertToPanorama(bpy.types.Operator): input_path=str(path), color_path=str(color_path), source_object_name=source.name, source_identity=str(source.as_pointer()), image_identity=str(image.as_pointer()), mesh_detail=self.mesh_detail, - model=getattr(preferences, "geometry_model", DEFAULT_MOGE2_MODEL_KEY), + model=configured_geometry_model(context), resolution_level=int(getattr(preferences, "geometry_resolution_level", 5)), max_input_size=configured_max_ai_input_size(context), ) diff --git a/src/anyimage/operators/convert_to_plane/operators.py b/src/anyimage/operators/convert_to_plane/operators.py index c979610..35ddf47 100644 --- a/src/anyimage/operators/convert_to_plane/operators.py +++ b/src/anyimage/operators/convert_to_plane/operators.py @@ -1,7 +1,7 @@ import bpy -from ...preferences import addon_preferences, configured_max_ai_input_size -from ...properties import DEFAULT_MOGE2_MODEL_KEY +from ...preferences import addon_preferences, configured_max_ai_input_size, configured_geometry_model +from ...server.model_catalog import DEFAULT_GEOMETRY_MODEL_KEY from ...runtime import JobOperatorBase from ...server.media.resolution import ( DEFAULT_MAX_AI_INPUT_SIZE, @@ -37,7 +37,7 @@ _HIDDEN = {"HIDDEN", "SKIP_SAVE"} def depth_plane_settings(context): preferences = addon_preferences(context) return ( - getattr(preferences, "geometry_model", DEFAULT_MOGE2_MODEL_KEY), + configured_geometry_model(context), int(getattr(preferences, "geometry_resolution_level", 5)), ) @@ -59,7 +59,7 @@ class GenerateDepthPlane(JobOperatorBase, bpy.types.Operator): options=_HIDDEN, ) mesh_detail: bpy.props.IntProperty(default=6, min=0, max=10, options=_HIDDEN) - model: bpy.props.StringProperty(default=DEFAULT_MOGE2_MODEL_KEY, options=_HIDDEN) + model: bpy.props.StringProperty(default=DEFAULT_GEOMETRY_MODEL_KEY, options=_HIDDEN) resolution_level: bpy.props.IntProperty(default=5, min=0, max=9, options=_HIDDEN) max_input_size: bpy.props.IntProperty( default=DEFAULT_MAX_AI_INPUT_SIZE, diff --git a/src/anyimage/operators/cutout_tool/operators.py b/src/anyimage/operators/cutout_tool/operators.py index b993a1d..1a274b8 100644 --- a/src/anyimage/operators/cutout_tool/operators.py +++ b/src/anyimage/operators/cutout_tool/operators.py @@ -6,15 +6,16 @@ import bpy from ...preferences import ( addon_preferences, + configured_geometry_model, configured_max_ai_input_size, configured_max_cutout_mesh_resolution, ) from ...properties import ( - DEFAULT_MOGE2_MODEL_KEY, cutout_alpha_threshold_property, cutout_fine_outline_property, ) from ...runtime import JobOperatorBase +from ...server.model_catalog import DEFAULT_GEOMETRY_MODEL_KEY from ...server.media.resolution import ( DEFAULT_MAX_AI_INPUT_SIZE, MAX_MAX_AI_INPUT_SIZE, @@ -169,7 +170,7 @@ def _generated_color_result(input_path, source_bounds): def cutout_geometry_settings(context): preferences = addon_preferences(context) - model = getattr(preferences, "geometry_model", DEFAULT_MOGE2_MODEL_KEY) + model = configured_geometry_model(context) return ( model, int(getattr(preferences, "geometry_resolution_level", 5)), @@ -196,7 +197,7 @@ class CutoutSelectionToShape(JobOperatorBase, bpy.types.Operator): options=_HIDDEN, ) model: bpy.props.StringProperty( - default=DEFAULT_MOGE2_MODEL_KEY, + default=DEFAULT_GEOMETRY_MODEL_KEY, options=_HIDDEN, ) resolution_level: bpy.props.IntProperty( diff --git a/src/anyimage/operators/remove_background.py b/src/anyimage/operators/remove_background.py index 3b593b0..34bdf75 100644 --- a/src/anyimage/operators/remove_background.py +++ b/src/anyimage/operators/remove_background.py @@ -1,6 +1,6 @@ import bpy -from ..properties import BEN2_MODEL_KEY +from ..preferences import configured_background_model from ..runtime import JobOperatorBase from ..common.ai import ( require_input_path, @@ -22,7 +22,7 @@ _HIDDEN = {"HIDDEN", "SKIP_SAVE"} class RemoveImageBackground(bpy.types.Operator): bl_idname = "anyimage.remove_image_background" bl_label = "Remove Background" - bl_description = "Remove the image background with BEN2" + bl_description = "Remove the image background with the selected AI model" bl_options = {"REGISTER"} @classmethod @@ -84,11 +84,12 @@ class RunBackgroundRemoval(JobOperatorBase, bpy.types.Operator): def request(self, context): require_environment() - require_model(BEN2_MODEL_KEY) + model = configured_background_model(context) + require_model(model) settings = context.scene.anyimage_settings parameters = { "input": str(require_input_path(self.input_path)), - "model": BEN2_MODEL_KEY, + "model": model, "device": production_device().lower(), "video_frames": settings.video_frames, "start_frame": self._image_target.start_frame, diff --git a/src/anyimage/operators/upscale.py b/src/anyimage/operators/upscale.py index fe53a16..2456293 100644 --- a/src/anyimage/operators/upscale.py +++ b/src/anyimage/operators/upscale.py @@ -1,7 +1,8 @@ import bpy from ..preferences import configured_max_ai_input_size -from ..properties import configured_upscale_model, upscale_model_label +from ..preferences import configured_upscale_model +from ..properties import upscale_model_label from ..server.media.resolution import ( DEFAULT_MAX_AI_INPUT_SIZE, MAX_MAX_AI_INPUT_SIZE, diff --git a/src/anyimage/preferences.py b/src/anyimage/preferences.py index 8ebd1a2..1432083 100644 --- a/src/anyimage/preferences.py +++ b/src/anyimage/preferences.py @@ -8,12 +8,14 @@ from .server.media.resolution import ( MIN_MAX_AI_INPUT_SIZE, ) from .server.model_catalog import ( + BACKGROUND_MODELS, + MOGE_MODELS, + DEFAULT_BACKGROUND_MODEL_KEY, + DEFAULT_GEOMETRY_MODEL_KEY, DEFAULT_UPSCALE_MODEL_KEY, UPSCALE_MODELS, ) -DEFAULT_MOGE2_MODEL_KEY = "MOGE2_VITS_NORMAL" -DEFAULT_BACKGROUND_MODEL_KEY = "BEN2_BASE" DEFAULT_MAX_CUTOUT_MESH_RESOLUTION = 1024 MIN_CUTOUT_MESH_RESOLUTION = 4 MAX_CUTOUT_MESH_RESOLUTION = 1024 @@ -24,16 +26,13 @@ MAX_FRAME_RESOLUTION = 8192 ADDON_PACKAGE = __package__ DEFAULT_STORAGE_ROOT = Path.home() / ".anyimage" UPSCALE_MODEL_ITEMS = tuple( - (model.key, model.label, model.description, index) - for index, model in enumerate(UPSCALE_MODELS.values()) + (model.key, model.label, model.description, number) + for model, number in zip(UPSCALE_MODELS.values(), (0, 2, 4)) ) -BACKGROUND_MODEL_ITEMS = ( - ( - DEFAULT_BACKGROUND_MODEL_KEY, - "BEN2 Base", - "Model used for AI Remove Background", - ), +BACKGROUND_MODEL_ITEMS = tuple( + (model.key, model.label, model.description, number) + for model, number in zip(BACKGROUND_MODELS.values(), (1, 0, 2)) ) BALLOON_PROFILE_ALGORITHM_ITEMS = ( @@ -61,6 +60,27 @@ def configured_storage_root(): return DEFAULT_STORAGE_ROOT +def _configured_model(context, property_name, models, default): + preferences = addon_preferences(context) + model = getattr(preferences, property_name, default) + if model not in models: + setattr(preferences, property_name, default) + return default + return model + + +def configured_geometry_model(context=None): + return _configured_model(context, "geometry_model", MOGE_MODELS, DEFAULT_GEOMETRY_MODEL_KEY) + + +def configured_background_model(context=None): + return _configured_model(context, "background_model", BACKGROUND_MODELS, DEFAULT_BACKGROUND_MODEL_KEY) + + +def configured_upscale_model(context=None): + return _configured_model(context, "upscale_model", UPSCALE_MODELS, DEFAULT_UPSCALE_MODEL_KEY) + + def configured_material_view_adaptation(context=None): return bool( getattr(addon_preferences(context), "adapt_material_to_view_transform", True) @@ -239,7 +259,7 @@ class AnyImagePreferences(bpy.types.AddonPreferences): ai, "geometry_model", models, - DEFAULT_MOGE2_MODEL_KEY, + DEFAULT_GEOMETRY_MODEL_KEY, "AI Generate Depth Map", ) self._draw_model_property( @@ -286,8 +306,10 @@ class AnyImagePreferences(bpy.types.AddonPreferences): ready_label, ): model_key = getattr(self, property_name, default_model_key) - if model_key not in models: + choices = {"geometry_model": MOGE_MODELS, "background_model": BACKGROUND_MODELS, "upscale_model": UPSCALE_MODELS} + if model_key not in choices[property_name]: model_key = default_model_key + setattr(self, property_name, model_key) split = layout.split(factor=0.58, align=True) split.column(align=True).prop(self, property_name, text="") self._draw_model_status( diff --git a/src/anyimage/properties.py b/src/anyimage/properties.py index 7113243..a460c85 100644 --- a/src/anyimage/properties.py +++ b/src/anyimage/properties.py @@ -3,18 +3,9 @@ import sys import bpy -from .preferences import addon_preferences from .runtime import runtime from .server import model_catalog as shared_model_catalog -BEN2_MODEL_KEY = "BEN2_BASE" -REALESRGAN_GENERAL_WDN_MODEL_KEY = shared_model_catalog.DEFAULT_UPSCALE_MODEL_KEY -DEFAULT_MOGE2_MODEL_KEY = "MOGE2_VITS_NORMAL" -REQUIRED_MODEL_KEYS = ( - DEFAULT_MOGE2_MODEL_KEY, - BEN2_MODEL_KEY, - REALESRGAN_GENERAL_WDN_MODEL_KEY, -) DEVICE_LABELS = { @@ -87,7 +78,7 @@ def ai_status(): environment_ready = runtime.environment_ready() ready_models = {model["key"] for model in model_catalog() if model.get("ready")} missing_models = tuple( - model_key for model_key in REQUIRED_MODEL_KEYS if model_key not in ready_models + model_key for model_key in shared_model_catalog.DEFAULT_MODEL_KEYS if model_key not in ready_models ) return { "environment_ready": environment_ready, @@ -127,14 +118,15 @@ def model_record(model_key): def _model_items(models): + numbers = {"MOGE2_VITS_NORMAL": 0, "MOGE2_VITB_NORMAL": 1, "MOGE3_VITL": 3} return [ ( model["key"], model["label"], f"{model['description']}; license: {model['license']}", - index, + numbers[model["key"]], ) - for index, model in enumerate(models) + for model in models ] @@ -143,18 +135,6 @@ def geometry_model_enum_items(_owner, _context): return _model_items(models) -def configured_upscale_model(context=None): - preferences = addon_preferences(context) - model = getattr( - preferences, - "upscale_model", - shared_model_catalog.DEFAULT_UPSCALE_MODEL_KEY, - ) - if model in shared_model_catalog.UPSCALE_MODELS: - return model - return shared_model_catalog.DEFAULT_UPSCALE_MODEL_KEY - - def upscale_model_label(model_key): model = shared_model_catalog.UPSCALE_MODELS.get(model_key) return model.label if model is not None else model_key diff --git a/src/anyimage/runtime.py b/src/anyimage/runtime.py index facaf3a..d4488ee 100644 --- a/src/anyimage/runtime.py +++ b/src/anyimage/runtime.py @@ -3,6 +3,7 @@ from pathlib import Path from blendjob import JobRuntime from .preferences import configured_storage_root +from .server.model_catalog import DEFAULT_MODEL_KEYS, get_downloadable_model ENVIRONMENT = { @@ -24,19 +25,11 @@ ENVIRONMENT = { def post_install(job_runtime): errors = [] - models = ( - ("MOGE2_VITS_NORMAL", "MoGe-2 ViT-S Normal"), - ("BEN2_BASE", "BEN2 Base"), - ( - "REALESRGAN_GENERAL_WDN_X4V3", - "Real-ESRGAN General WDN x4v3", - ), - ) - for model, label in models: + for model in DEFAULT_MODEL_KEYS: try: job_runtime.request("download-model", {"model": model}) except RuntimeError as error: - errors.append(f"{label}: {error}") + errors.append(f"{get_downloadable_model(model).label}: {error}") if errors: raise RuntimeError("; ".join(errors)) diff --git a/src/anyimage/server/jobs/remove_background.py b/src/anyimage/server/jobs/remove_background.py index 908582c..9dbf48d 100644 --- a/src/anyimage/server/jobs/remove_background.py +++ b/src/anyimage/server/jobs/remove_background.py @@ -2,12 +2,15 @@ import time from ..media import input_frames from ..media.input import open_image -from ..models.ben2 import infer_alpha, infer_foreground, infer_foreground_frames +from ..model_catalog import BACKGROUND_MODELS +from ..models.background import infer_alpha, infer_foreground, infer_foreground_frames, inference_device def run(context, parameters): - """Run BEN2 and write transparent foreground images into the Job directory.""" + """Write foreground images or continuous alpha into the Job directory.""" total_started = time.perf_counter() + if parameters["model"] not in BACKGROUND_MODELS: + raise ValueError(f"Unknown background removal model: {parameters['model']}") context.progress(0.02, "Starting Environment") context.check_cancelled() @@ -16,12 +19,13 @@ def run(context, parameters): device = parameters["device"] model_manager = context.resource("model_manager") model_directory = model_manager.directory(parameters["model"]) + device = inference_device(model_directory, device) output_kind = parameters.get("output_kind", "foreground") if output_kind not in {"foreground", "alpha"}: raise ValueError(f"Unknown background removal output: {output_kind}") with input_frames(context, parameters) as images: - context.progress(0.12, f"Loading BEN2 for {device.upper()}") + context.progress(0.12, f"Loading {parameters['model']} for {device.upper()}") if output_kind == "alpha": import numpy as np @@ -53,7 +57,7 @@ def run(context, parameters): compress_level=1, ) if not output_path.is_file(): - raise RuntimeError("BEN2 did not write a foreground image") + raise RuntimeError("Background removal did not write a foreground image") else: load_ms, inference_ms = infer_foreground_frames( images, diff --git a/src/anyimage/server/model_catalog.py b/src/anyimage/server/model_catalog.py index 67e6d90..d43b835 100644 --- a/src/anyimage/server/model_catalog.py +++ b/src/anyimage/server/model_catalog.py @@ -16,14 +16,31 @@ class ModelSpec: r2_files: tuple[tuple[str, int, str], ...] = () +BIREFNET_LITE = ModelSpec( + key="BIREFNET_LITE", + label="Fast | BiRefNet Lite", + directory_name="birefnet-lite", + repository="ZhengPeng7/BiRefNet_lite", + license_name="MIT", + description="Compact 1024px background removal with continuous alpha", + download_patterns=("model.onnx", "LICENSE.txt"), + ready_patterns=("model.onnx", "LICENSE.txt"), + r2_directory="birefnet-lite", + r2_files=( + ("model.onnx", 112968598, "6e2a9844c57b080e95a402722ddcae5660e6f254a04d3e003347f06f587b71fe"), + ("LICENSE.txt", 1066, "92a7089e0915fc32bc40067560b398f1e6a7a5958abd7d04eda393629a5acefb"), + ), +) + + BEN2_MODEL = ModelSpec( key="BEN2_BASE", - label="BEN2 Base", + label="General | BEN2 Base", directory_name="BEN2-ONNX", repository="onnx-community/BEN2-ONNX", license_name="Apache 2.0", huggingface_repository="onnx-community/BEN2-ONNX", - description="Fast high-quality FP16 ONNX background removal model", + description="General-purpose FP16 ONNX background removal model", download_patterns=("onnx/model_fp16.onnx",), ready_patterns=("onnx/model_fp16.onnx",), r2_directory="BEN2-ONNX", @@ -37,7 +54,7 @@ BEN2_MODEL = ModelSpec( REALESRGAN_GENERAL_WDN_X4V3 = ModelSpec( key="REALESRGAN_GENERAL_WDN_X4V3", - label="Real-ESRGAN General WDN x4v3", + label="Fast | Real-ESRGAN General WDN x4v3", directory_name="realesr-general-wdn-x4v3", repository="xinntao/Real-ESRGAN", license_name="BSD 3-Clause", @@ -52,26 +69,10 @@ REALESRGAN_GENERAL_WDN_X4V3 = ModelSpec( ),), ) -SAFMN_MODEL = ModelSpec( - key="SAFMN_REAL_X4", - label="SAFMN Real x4", - directory_name="safmn-real-x4", - repository="sunny2109/SAFMN", - license_name="Apache 2.0", - description="Balanced 2x output from a real-world 4x ONNX model", - download_patterns=("safmn-real-x4.onnx",), - ready_patterns=("safmn-real-x4.onnx",), - r2_directory="safmn-real-x4", - r2_files=(( - "safmn-real-x4.onnx", - 22657477, - "8f13867b66c0300941b902bf6a5c6e0b545794faabe8b0478bf55668bc517621", - ),), -) REALESRGAN_X4PLUS = ModelSpec( key="REALESRGAN_X4PLUS", - label="Real-ESRGAN x4plus", + label="General | Real-ESRGAN x4plus", directory_name="realesrgan-x4plus", repository="jonathanst29/tinier-upscale-models", license_name="BSD 3-Clause", @@ -88,25 +89,21 @@ REALESRGAN_X4PLUS = ModelSpec( ) -PISA_SR = ModelSpec( - key="PISA_SR", - label="PiSA-SR", - directory_name="pisa-sr", - repository="csslc/PiSA-SR", - license_name="Apache 2.0 / CreativeML Open RAIL++-M", - description="Generative 2x upscale with fixed detail strength 0.7", - download_patterns=("encoder.onnx", "decoder.onnx", "unet-pixel.onnx", "unet-full.onnx", "noise.npy", "LICENSE-APACHE.txt", "LICENSE-MODEL.txt", "NOTICE.txt"), - ready_patterns=("encoder.onnx", "decoder.onnx", "unet-pixel.onnx", "unet-full.onnx", "noise.npy", "LICENSE-APACHE.txt", "LICENSE-MODEL.txt", "NOTICE.txt"), - r2_directory="pisa-sr", + + +HAT_GAN_X4_SHARPER = ModelSpec( + key="HAT_GAN_X4_SHARPER", + label="Detailed | HAT Sharper", + directory_name="hat-gan-x4-sharper", + repository="XPixelGroup/HAT", + license_name="Apache 2.0", + description="Detailed texture enhancement with 2x output; higher processing cost", + download_patterns=("model.onnx", "LICENSE.txt"), + ready_patterns=("model.onnx", "LICENSE.txt"), + r2_directory="hat-gan-x4-sharper", r2_files=( - ("encoder.onnx", 68426712, "4b93e4a3d81c3349a78ab96619501e39ee8f815d14a20b73fc4f969e30a1ceeb"), - ("decoder.onnx", 99084589, "c9021384c402e267654907527fa5f0ee86cab1814bfd96233216a15ae8e83c4c"), - ("unet-pixel.onnx", 1732713007, "4f381574ad584050d0534c8cca794fb5828cb8e99e148f8efc1aaac20833c0b1"), - ("unet-full.onnx", 1732713007, "3b0fd44ff1093608a65865c3f26d0060775337e1fc117d8ac29aacbe748d3f9c"), - ("noise.npy", 32896, "a94ed01e2e963022cfa8ec01e04d2e96e599408ef0c385878347eaacd25a663b"), - ("LICENSE-APACHE.txt", 11358, "cfc7749b96f63bd31c3c42b5c471bf756814053e847c10f3eb003417bc523d30"), - ("LICENSE-MODEL.txt", 14206, "e514164be7f2ca2c0fb39e7283b01ab4a06d802837eaf171aa4645cfd7ee1157"), - ("NOTICE.txt", 1081, "ccb55fcb98cce42401368b88e3ccd555f4ff49ac1347c6ce063061dd48726050"), + ("model.onnx", 45874438, "032d8b03c704220a08d91465a822aca252c1d6e63db20c3cf23c5fb67f3d2366"), + ("LICENSE.txt", 11342, "a91d57ebad8955a1757be7891ca2da8e07493e662f0e34e1e1926571e05f37fc"), ), ) @@ -115,9 +112,8 @@ UPSCALE_MODELS = { model.key: model for model in ( REALESRGAN_GENERAL_WDN_X4V3, - SAFMN_MODEL, REALESRGAN_X4PLUS, - PISA_SR, + HAT_GAN_X4_SHARPER, ) } @@ -128,7 +124,7 @@ DEFAULT_UPSCALE_MODEL_KEY = REALESRGAN_GENERAL_WDN_X4V3.key MOGE2_MODELS = { "MOGE2_VITS_NORMAL": ModelSpec( key="MOGE2_VITS_NORMAL", - label="MoGe-2 ViT-S Normal", + label="Fast | MoGe-2 ViT-S Normal", directory_name="moge-2-vits-normal-onnx", repository="Ruicheng/moge-2-vits-normal-onnx", license_name="MIT", @@ -145,7 +141,7 @@ MOGE2_MODELS = { ), "MOGE2_VITB_NORMAL": ModelSpec( key="MOGE2_VITB_NORMAL", - label="MoGe-2 ViT-B Normal", + label="Balanced | MoGe-2 ViT-B Normal", directory_name="moge-2-vitb-normal-onnx", repository="Ruicheng/moge-2-vitb-normal-onnx", license_name="MIT", @@ -160,27 +156,10 @@ MOGE2_MODELS = { "bbf14e07a30f11e69d36ab861590123f5598ababcbc8946a063eb4a966f35a21", ),), ), - "MOGE2_VITL_NORMAL": ModelSpec( - key="MOGE2_VITL_NORMAL", - label="MoGe-2 ViT-L Normal", - directory_name="moge-2-vitl-normal-onnx", - repository="Ruicheng/moge-2-vitl-normal-onnx", - license_name="MIT", - huggingface_repository="Ruicheng/moge-2-vitl-normal-onnx", - description="Highest-quality MoGe-2 model using ONNX Runtime", - download_patterns=("model.onnx",), - ready_patterns=("model.onnx",), - r2_directory="moge-2-vitl-normal-onnx", - r2_files=(( - "model.onnx", - 1324265014, - "afbc4ccc3450298f3afb35b90f015f4c4f552dea21dc6470d5f7b78b77e2d751", - ),), - ), } MOGE3_VITL = ModelSpec( key="MOGE3_VITL", - label="MoGe-3 ViT-L", + label="Detailed | MoGe-3 ViT-L", directory_name="moge-3-vitl-onnx", repository="Ruicheng/moge-3-vitl", license_name="MIT", @@ -195,9 +174,33 @@ MOGE3_VITL = ModelSpec( ) MOGE_MODELS = {**MOGE2_MODELS, MOGE3_VITL.key: MOGE3_VITL} +BIREFNET_HR_MATTING = ModelSpec( + key="BIREFNET_HR_MATTING", + label="Detailed | BiRefNet HR-matting (CPU)", + directory_name="birefnet-hr-matting", + repository="ZhengPeng7/BiRefNet_HR-matting", + license_name="MIT", + description="2048px matting with continuous alpha; CPU only, high memory and processing cost", + download_patterns=("model.onnx", "LICENSE.txt"), + ready_patterns=("model.onnx", "LICENSE.txt"), + r2_directory="birefnet-hr-matting", + r2_files=( + ("model.onnx", 529467341, "b3abc34843a4d96f669ca9fb1a2a16b3417ab8b91c987573fdbda36d20514133"), + ("LICENSE.txt", 1066, "92a7089e0915fc32bc40067560b398f1e6a7a5958abd7d04eda393629a5acefb"), + ), +) + +BACKGROUND_MODELS = {model.key: model for model in (BIREFNET_LITE, BEN2_MODEL, BIREFNET_HR_MATTING)} +DEFAULT_GEOMETRY_MODEL_KEY = next(iter(MOGE_MODELS)) +DEFAULT_BACKGROUND_MODEL_KEY = next(iter(BACKGROUND_MODELS)) +DEFAULT_MODEL_KEYS = ( + DEFAULT_GEOMETRY_MODEL_KEY, + DEFAULT_BACKGROUND_MODEL_KEY, + DEFAULT_UPSCALE_MODEL_KEY, +) DOWNLOADABLE_MODELS = { - BEN2_MODEL.key: BEN2_MODEL, + **BACKGROUND_MODELS, **UPSCALE_MODELS, **MOGE_MODELS, } @@ -220,8 +223,8 @@ def model_record(model, ready=False): "family": ( "moge" if model.key in MOGE_MODELS - else "ben2" - if model.key == BEN2_MODEL.key + else "background" + if model.key in BACKGROUND_MODELS else "upscale" ), "ready": bool(ready), diff --git a/src/anyimage/server/model_manager.py b/src/anyimage/server/model_manager.py index f3d5d88..b27ba1f 100644 --- a/src/anyimage/server/model_manager.py +++ b/src/anyimage/server/model_manager.py @@ -12,28 +12,30 @@ class ModelSessionCache: """Cache ONNX sessions shared by Jobs in one Server process.""" def __init__(self): - self.ben2_model = None - self.ben2_key = None + self.background_model = None + self.background_key = None self.moge_model = None self.moge_key = None self.upscale_model = None self.upscale_key = None - def get_ben2(self, model_dir, device): - requested_device = getattr(device, "type", device) - key = (str(Path(model_dir).resolve()), requested_device, "onnx") - if self.ben2_model is not None and self.ben2_key == key: - return self.ben2_model, 0.0 - self.release_ben2() - return self._load_ben2(model_dir, requested_device, key) + def get_background(self, model_dir, device): + from .models.background import inference_device - def _load_ben2(self, model_dir, requested_device, key): - from .models.onnx_ben2 import create_session + requested_device = inference_device(model_dir, device) + key = (str(Path(model_dir).resolve()), requested_device, "onnx") + if self.background_model is not None and self.background_key == key: + return self.background_model, 0.0 + self.release_background() + return self._load_background(model_dir, requested_device, key) + + def _load_background(self, model_dir, requested_device, key): + from .models.background import model_adapter started = time.perf_counter() - model = create_session(model_dir, requested_device) - self.ben2_model = model - self.ben2_key = key + model = model_adapter(model_dir).create_session(model_dir, requested_device) + self.background_model = model + self.background_key = key elapsed_ms = (time.perf_counter() - started) * 1000.0 return model, elapsed_ms @@ -56,9 +58,6 @@ class ModelSessionCache: def get_upscale(self, model_dir, requested_device): from .models.onnx_upscale import create_session - if Path(model_dir).name == "pisa-sr": - from .models.onnx_pisa_sr import PisaSession as create_session - key = (str(Path(model_dir).resolve()), requested_device, "onnx") if self.upscale_model is not None and self.upscale_key == key: return self.upscale_model @@ -68,9 +67,9 @@ class ModelSessionCache: self.upscale_key = key return model - def release_ben2(self): - self.ben2_model = None - self.ben2_key = None + def release_background(self): + self.background_model = None + self.background_key = None self._release_runtime_cache() def release_geometry(self): @@ -79,21 +78,15 @@ class ModelSessionCache: self._release_runtime_cache() def release_upscale(self): - close = getattr(self.upscale_model, "close", None) - if close is not None: - close() self.upscale_model = None self.upscale_key = None self._release_runtime_cache() def close(self): - close = getattr(self.upscale_model, "close", None) - if close is not None: - close() - self.ben2_model = None + self.background_model = None self.moge_model = None self.upscale_model = None - self.ben2_key = None + self.background_key = None self.moge_key = None self.upscale_key = None self._release_runtime_cache() @@ -101,7 +94,7 @@ class ModelSessionCache: def snapshot(self): geometry_key = self.moge_key return { - "ben2": Path(self.ben2_key[0]).name if self.ben2_key else "", + "background": Path(self.background_key[0]).name if self.background_key else "", "geometry": Path(geometry_key[0]).name if geometry_key else "", "upscale": (Path(self.upscale_key[0]).name if self.upscale_key else ""), } @@ -125,7 +118,7 @@ class ModelManager: def ready(self, model_key): model = model_catalog.get_downloadable_model(model_key) directory = self.directory(model_key) - if model_key in {"PISA_SR", "MOGE3_VITL"}: + if model_key in {"MOGE3_VITL", "HAT_GAN_X4_SHARPER", *model_catalog.BACKGROUND_MODELS}: for filename, size, checksum in model.r2_files: path = directory / filename try: @@ -148,8 +141,8 @@ class ModelManager: return True return all(any(directory.glob(pattern)) for pattern in model.ready_patterns) - def get_ben2(self, model_dir, device): - return self.sessions.get_ben2(model_dir, device) + def get_background(self, model_dir, device): + return self.sessions.get_background(model_dir, device) def get_moge(self, model_dir, device): diff --git a/src/anyimage/server/models/__init__.py b/src/anyimage/server/models/__init__.py index ebc8c63..906b931 100644 --- a/src/anyimage/server/models/__init__.py +++ b/src/anyimage/server/models/__init__.py @@ -1,7 +1,7 @@ -from . import ben2, moge, upscale +from . import background, moge, upscale __all__ = ( - "ben2", + "background", "moge", "upscale", ) diff --git a/src/anyimage/server/models/ben2.py b/src/anyimage/server/models/background.py similarity index 76% rename from src/anyimage/server/models/ben2.py rename to src/anyimage/server/models/background.py index 6c74acb..ee74aca 100644 --- a/src/anyimage/server/models/ben2.py +++ b/src/anyimage/server/models/background.py @@ -1,10 +1,26 @@ import time +from importlib import import_module from pathlib import Path from ..media import check_cancelled from ..media.input import open_image, validate_input_path +def inference_device(model_dir, requested_device): + if Path(model_dir).name == "birefnet-hr-matting": + return "cpu" + return getattr(requested_device, "type", requested_device) + + +def model_adapter(model_dir): + name = Path(model_dir).name + if name == "BEN2-ONNX": + return import_module(".onnx_ben2", __package__) + if name in {"birefnet-lite", "birefnet-hr-matting"}: + return import_module(".onnx_birefnet", __package__) + raise ValueError(f"Unknown background model directory: {name}") + + def infer_alpha( image, model_dir, @@ -16,16 +32,16 @@ def infer_alpha( from PIL import Image check_cancelled(cancel_check) - from .onnx_ben2 import create_session, infer_alpha as predict_alpha + adapter = model_adapter(model_dir) owns_model = model_cache is None - requested_device = getattr(device, "type", device) + requested_device = inference_device(model_dir, device) if owns_model: load_started = time.perf_counter() - model = create_session(model_dir, requested_device) + model = adapter.create_session(model_dir, requested_device) load_ms = (time.perf_counter() - load_started) * 1000.0 else: - model, load_ms = model_cache.get_ben2(model_dir, device) + model, load_ms = model_cache.get_background(model_dir, requested_device) check_cancelled(cancel_check) source_rgba = image.convert("RGBA") @@ -36,7 +52,7 @@ def infer_alpha( source_rgba.getchannel("A"), ) inference_started = time.perf_counter() - alpha = predict_alpha( + alpha = adapter.infer_alpha( model, image, release_memory=release_memory, @@ -90,7 +106,7 @@ def infer_foreground_frames( output_prefix="foreground", preserve_input_alpha=False, ): - """Run BEN2 on every frame and write numbered transparent PNG files.""" + """Remove each frame's background and write numbered transparent PNG files.""" total_load_ms = 0.0 total_inference_ms = 0.0 for index, image_path in enumerate(images): @@ -115,7 +131,7 @@ def infer_foreground_frames( compress_level=1, ) if not output_path.is_file(): - raise RuntimeError(f"BEN2 did not write: {output_path}") + raise RuntimeError(f"Background removal did not write: {output_path}") total_load_ms += load_ms total_inference_ms += inference_ms return total_load_ms, total_inference_ms diff --git a/src/anyimage/server/models/onnx_birefnet.py b/src/anyimage/server/models/onnx_birefnet.py new file mode 100644 index 0000000..95c8688 --- /dev/null +++ b/src/anyimage/server/models/onnx_birefnet.py @@ -0,0 +1,41 @@ +"""BiRefNet probability alpha with official RGB normalization.""" + +from pathlib import Path + +from .onnx_runtime import create_session as create_runtime_session +from .onnx_runtime import run_session + + +def create_session(model_directory, requested_device="auto"): + return create_runtime_session(Path(model_directory) / "model.onnx", requested_device) + + +def preprocess(image, size): + import numpy as np + from PIL import Image + + resized = image.convert("RGB").resize((size, size), Image.Resampling.BILINEAR) + array = np.asarray(resized, dtype=np.float32) / 255.0 + array = (array - np.array((0.485, 0.456, 0.406), dtype=np.float32)) / np.array( + (0.229, 0.224, 0.225), dtype=np.float32, + ) + return np.ascontiguousarray(array.transpose(2, 0, 1)[None]) + + +def infer_alpha(session, image, *, release_memory=True): + import numpy as np + from PIL import Image + + input_spec = session.get_inputs()[0] + size = input_spec.shape[-1] + if size not in (1024, 2048) or input_spec.shape != [1, 3, size, size]: + raise RuntimeError("BiRefNet returned an invalid input contract") + alpha = np.asarray(run_session( + session, [session.get_outputs()[0].name], + {input_spec.name: preprocess(image, size)}, release_memory=release_memory, + )[0], dtype=np.float32) + if (alpha.shape != (1, 1, size, size) or not np.isfinite(alpha).all() + or alpha.min() < 0 or alpha.max() > 1): + raise RuntimeError("BiRefNet returned an invalid alpha mask") + mask = Image.fromarray(alpha[0, 0]).resize(image.size, Image.Resampling.BILINEAR) + return np.clip(np.asarray(mask, dtype=np.float32), 0.0, 1.0) diff --git a/src/anyimage/server/models/onnx_pisa_sr.py b/src/anyimage/server/models/onnx_pisa_sr.py deleted file mode 100644 index 0823e3b..0000000 --- a/src/anyimage/server/models/onnx_pisa_sr.py +++ /dev/null @@ -1,153 +0,0 @@ -"""Fixed-strength PiSA-SR inference with one resident ONNX component.""" - -import gc -from pathlib import Path -from tempfile import TemporaryDirectory - -import numpy as np -from PIL import Image - -from .onnx_runtime import create_session, run_session -from .onnx_upscale import RESOURCE_ERROR_TOKENS, UpscaleResourceError - -TILE_SIZE = 512 -OVERLAP = 128 -LATENT_SCALE = np.float16(0.18215) -SEMANTIC_STRENGTH = np.float16(0.7) - - -def tile_starts(length): - padded = max(TILE_SIZE, (length + 7) // 8 * 8) - starts = list(range(0, padded - TILE_SIZE + 1, TILE_SIZE - OVERLAP)) - if starts[-1] != padded - TILE_SIZE: - starts.append(padded - TILE_SIZE) - return starts, padded - - -def correct_color(pixels, source): - target = pixels.astype(np.float32) / 255 - reference = np.asarray(source, dtype=np.float32) / 255 - axes = (0, 1) - deviation = np.sqrt(target.var(axis=axes, ddof=1) + 1e-5) - reference_deviation = np.sqrt(reference.var(axis=axes, ddof=1) + 1e-5) - corrected = (target - target.mean(axis=axes)) / deviation - corrected = corrected * reference_deviation + reference.mean(axis=axes) - return Image.fromarray((corrected.clip(0, 1) * 255).astype(np.uint8)) - - -class PisaSession: - def __init__(self, directory, device): - if device == "coreml": - raise RuntimeError("PiSA-SR CoreML is not validated; select CPU for this model") - self.directory = Path(directory) - self.device = device - self.session = None - self.component = None - self.noise = np.load(self.directory / "noise.npy", allow_pickle=False) - if self.noise.shape != (1, 4, 64, 64) or not np.isfinite(self.noise).all(): - raise ValueError("Invalid PiSA-SR noise asset") - - def close(self): - self.session = None - self.component = None - gc.collect() - - def activate(self, component): - if self.component == component: - return - self.close() - self.session = create_session(self.directory / f"{component}.onnx", self.device) - self.session.disable_fallback() - if "CoreMLExecutionProvider" in self.session.get_providers(): - self.close() - raise RuntimeError("PiSA-SR CoreML is not validated; select CPU for this model") - self.component = component - - def run(self, value, final=False): - try: - result = run_session( - self.session, ["output"], - {"input": np.ascontiguousarray(value, dtype=np.float16)}, - release_memory=final, - )[0] - except Exception as error: - if "CUDAExecutionProvider" in self.session.get_providers(): - raise RuntimeError( - "PiSA-SR failed on CUDA; select DirectML or CPU. " + str(error) - ) from error - raise - if not np.isfinite(result).all(): - raise RuntimeError("PiSA-SR returned non-finite pixels on the selected device") - return result - - -def infer(model, image, cancel_check=None, *, release_memory=True): - def check(): - if cancel_check is not None and cancel_check(): - raise RuntimeError("Task cancelled") - - source = image.convert("RGB").resize( - (image.width * 2, image.height * 2), Image.Resampling.BICUBIC, - ) - rows, height = tile_starts(source.height) - columns, width = tile_starts(source.width) - coordinates = [(y, x) for y in rows for x in columns] - array = np.asarray(source) - array = np.pad(array, ((0, height-source.height), (0, width-source.width), (0, 0)), - mode="reflect" if min(source.size) > 1 else "edge") - try: - with TemporaryDirectory(prefix="anyimage-pisa-") as temporary: - directory = Path(temporary) - check() - model.activate("encoder") - for index, (y, x) in enumerate(coordinates): - check() - tile = array[y:y+TILE_SIZE, x:x+TILE_SIZE].astype(np.float32) / 255 - tensor = (tile.transpose(2, 0, 1)[None] * 2 - 1).astype(np.float16) - mean, logvar = np.split(model.run(tensor, index == len(coordinates)-1), 2, axis=1) - noise = np.take(np.take(model.noise, (np.arange(64)+y//8) % 64, axis=2), - (np.arange(64)+x//8) % 64, axis=3) - std = np.exp(logvar.clip(-30, 20) * np.float16(0.5)).astype(np.float16) - latent = ((mean + std * noise) * LATENT_SCALE).astype(np.float16) - np.save(directory / f"latent-{index}.npy", latent) - for component in ("unet-pixel", "unet-full"): - check() - model.activate(component) - for index in range(len(coordinates)): - check() - latent = np.load(directory / f"latent-{index}.npy") - prediction = model.run(latent, index == len(coordinates)-1) - pixel_path = directory / f"pixel-{index}.npy" - if component == "unet-pixel": - np.save(pixel_path, prediction) - else: - pixel = np.load(pixel_path) - corrected = latent - (pixel + SEMANTIC_STRENGTH * (prediction-pixel)) - np.save(directory / f"latent-{index}.npy", corrected) - pixel_path.unlink() - check() - model.activate("decoder") - result = np.zeros((height, width, 3), dtype=np.float32) - weights = np.zeros((height, width, 1), dtype=np.float32) - ramp = np.minimum(np.arange(TILE_SIZE)+1, np.arange(TILE_SIZE, 0, -1)) - ramp = np.minimum(ramp / OVERLAP, 1).astype(np.float32) - blend = (ramp[:, None] * ramp[None, :])[..., None] - for index, (y, x) in enumerate(coordinates): - check() - latent = np.load(directory / f"latent-{index}.npy") - decoded = model.run(latent, release_memory and index == len(coordinates)-1) - pixels = decoded[0].transpose(1, 2, 0).astype(np.float32) - result[y:y+TILE_SIZE, x:x+TILE_SIZE] += pixels * blend - weights[y:y+TILE_SIZE, x:x+TILE_SIZE] += blend - check() - result = (result / weights)[:source.height, :source.width] - pixels = ((result.clip(-1, 1)*0.5+0.5)*255).clip(0, 255).astype(np.uint8) - return correct_color(pixels, source) - except (MemoryError, UnicodeDecodeError) as error: - model.close() - raise UpscaleResourceError("PiSA-SR exhausted available inference memory") from error - except Exception as error: - model.close() - if any(token in str(error).lower() for token in RESOURCE_ERROR_TOKENS): - raise UpscaleResourceError("PiSA-SR exhausted available inference memory") from error - raise diff --git a/src/anyimage/server/models/onnx_upscale.py b/src/anyimage/server/models/onnx_upscale.py index 3d7049b..507360e 100644 --- a/src/anyimage/server/models/onnx_upscale.py +++ b/src/anyimage/server/models/onnx_upscale.py @@ -79,6 +79,8 @@ def _run(session, input_tensor, *, release_memory=True): raise RuntimeError( f"Upscale model returned an unexpected output shape: {prediction.shape}" ) + if not np.isfinite(prediction).all(): + raise RuntimeError("Upscale model returned non-finite pixels") return prediction @@ -89,32 +91,34 @@ def _prediction_pixels(prediction): return np.rint(pixels * 255.0).astype(np.uint8) -def _pad_source(source, extra_height, extra_width): +def _pad_source(source, extra_height, extra_width, tile_border): import numpy as np mode = "reflect" if source.shape[0] > 1 and source.shape[1] > 1 else "edge" return np.pad( source, ( - (TILE_BORDER, TILE_BORDER + extra_height), - (TILE_BORDER, TILE_BORDER + extra_width), + (tile_border, tile_border + extra_height), + (tile_border, tile_border + extra_width), (0, 0), ), mode=mode, ) -def infer(session, image, cancel_check=None, *, release_memory=True): +def infer(session, image, cancel_check=None, *, release_memory=True, tile_border=TILE_BORDER): import numpy as np from PIL import Image try: source = np.asarray(image.convert("RGB"), dtype=np.uint8) height, width = source.shape[:2] - core_size = TILE_SIZE - 2 * TILE_BORDER + if not 0 <= tile_border < TILE_SIZE // 2: + raise ValueError("Invalid upscale tile border") + core_size = TILE_SIZE - 2 * tile_border extra_height = (-height) % core_size extra_width = (-width) % core_size - padded = _pad_source(source, extra_height, extra_width) + padded = _pad_source(source, extra_height, extra_width, tile_border) output = np.empty((height * SCALE, width * SCALE, 3), dtype=np.uint8) for top in range(0, height, core_size): core_height = min(core_size, height - top) @@ -138,7 +142,7 @@ def infer(session, image, cancel_check=None, *, release_memory=True): release_memory=final_tile, ) ) - border = TILE_BORDER * SCALE + border = tile_border * SCALE output[ top * SCALE : (top + core_height) * SCALE, left * SCALE : (left + core_width) * SCALE, diff --git a/src/anyimage/server/models/upscale.py b/src/anyimage/server/models/upscale.py index fdebe07..664fb1a 100644 --- a/src/anyimage/server/models/upscale.py +++ b/src/anyimage/server/models/upscale.py @@ -1,5 +1,6 @@ from ..media.input import open_image from ..media.resolution import DEFAULT_MAX_AI_INPUT_SIZE, limit_image +from ..model_catalog import UPSCALE_MODELS def _source_alpha(image): @@ -19,8 +20,8 @@ def infer_one( from .onnx_upscale import UpscaleResourceError, infer - if parameters["model"] == "PISA_SR": - from .onnx_pisa_sr import infer + if parameters["model"] not in UPSCALE_MODELS: + raise ValueError(f"Unknown upscale model: {parameters['model']}") model_directory = model_manager.directory(parameters["model"]) model = model_manager.get_upscale( @@ -39,6 +40,7 @@ def infer_one( source, cancel_check=cancel_check, release_memory=release_memory, + tile_border=48 if parameters["model"] == "HAT_GAN_X4_SHARPER" else 16, ) prediction = prediction.resize( (source.width * 2, source.height * 2), diff --git a/tests/anyimage/operators/test_ai_image_operators.py b/tests/anyimage/operators/test_ai_image_operators.py index c0a938c..c03062a 100644 --- a/tests/anyimage/operators/test_ai_image_operators.py +++ b/tests/anyimage/operators/test_ai_image_operators.py @@ -22,7 +22,7 @@ class AiImageOperatorsTest(BlenderTestCase): operator = operator_type() operator.input_path = "input.png" operator._image_target = SimpleNamespace(start_frame=4) - operator.model = "SAFMN_REAL_X4" + operator.model = "HAT_GAN_X4_SHARPER" operator.max_input_size = 1024 context = SimpleNamespace( scene=SimpleNamespace(anyimage_settings=SimpleNamespace(video_frames=7)) @@ -37,12 +37,12 @@ class AiImageOperatorsTest(BlenderTestCase): ): request = operator.request(context) - require_model.assert_called_once_with("SAFMN_REAL_X4") + require_model.assert_called_once_with("HAT_GAN_X4_SHARPER") self.assertEqual( request, { "input": "input.png", - "model": "SAFMN_REAL_X4", + "model": "HAT_GAN_X4_SHARPER", "device": "cuda", "video_frames": 7, "start_frame": 4, @@ -108,12 +108,12 @@ class AiImageOperatorsTest(BlenderTestCase): result = SimpleNamespace(directory=Path("result")) operator = self.upscale.UpscaleImage() operator._image_target = Mock() - operator.model = "SAFMN_REAL_X4" + operator.model = "HAT_GAN_X4_SHARPER" with ( patch.object( self.upscale, "upscale_model_label", - return_value="SAFMN Real x4", + return_value="HAT Sharper", ), ): message = operator.response(None, result) @@ -122,4 +122,4 @@ class AiImageOperatorsTest(BlenderTestCase): Path("result"), frame_prefix="upscale.", ) - self.assertEqual(message, "Image is upscaled with SAFMN Real x4") + self.assertEqual(message, "Image is upscaled with HAT Sharper") diff --git a/tests/anyimage/operators/test_ai_input_failures.py b/tests/anyimage/operators/test_ai_input_failures.py index 6d006d4..86ba9c0 100644 --- a/tests/anyimage/operators/test_ai_input_failures.py +++ b/tests/anyimage/operators/test_ai_input_failures.py @@ -30,7 +30,7 @@ class AiInputFailuresTest(BlenderTestCase): stack.enter_context(patch.object(self.runtime.runtime, "active_job", None)) submit = stack.enter_context(patch.object(self.runtime.runtime, "begin_job")) if module == self.upscale: - stack.enter_context(patch.object(module, "configured_upscale_model", return_value="SAFMN_REAL_X4")) + stack.enter_context(patch.object(module, "configured_upscale_model", return_value="HAT_GAN_X4_SHARPER")) stack.enter_context(patch.object(module, "configured_max_ai_input_size", return_value=2048)) self.assertEqual(operator.execute(context), {"CANCELLED"}) submit.assert_not_called() diff --git a/tests/anyimage/operators/test_ai_setup.py b/tests/anyimage/operators/test_ai_setup.py index 350e186..4f135de 100644 --- a/tests/anyimage/operators/test_ai_setup.py +++ b/tests/anyimage/operators/test_ai_setup.py @@ -31,7 +31,7 @@ class AiSetupTest(BlenderTestCase): calls, [ ("download-model", {"model": "MOGE2_VITS_NORMAL"}), - ("download-model", {"model": "BEN2_BASE"}), + ("download-model", {"model": "BIREFNET_LITE"}), ( "download-model", {"model": "REALESRGAN_GENERAL_WDN_X4V3"}, @@ -50,9 +50,9 @@ class AiSetupTest(BlenderTestCase): with self.assertRaisesRegex( RuntimeError, ( - "MoGe-2 ViT-S Normal: network unavailable; " - "BEN2 Base: network unavailable; " - "Real-ESRGAN General WDN x4v3: network unavailable" + "Fast \\| MoGe-2 ViT-S Normal: network unavailable; " + "Fast \\| BiRefNet Lite: network unavailable; " + "Fast \\| Real-ESRGAN General WDN x4v3: network unavailable" ), ): self.runtime.post_install(SimpleNamespace(request=request)) @@ -61,7 +61,7 @@ class AiSetupTest(BlenderTestCase): calls, [ "MOGE2_VITS_NORMAL", - "BEN2_BASE", + "BIREFNET_LITE", "REALESRGAN_GENERAL_WDN_X4V3", ], ) @@ -77,7 +77,7 @@ class AiSetupTest(BlenderTestCase): { "models": [ "MOGE2_VITS_NORMAL", - "BEN2_BASE", + "BIREFNET_LITE", "REALESRGAN_GENERAL_WDN_X4V3", ] }, @@ -93,7 +93,7 @@ class AiSetupTest(BlenderTestCase): "model_catalog", return_value=( {"key": "MOGE2_VITS_NORMAL", "ready": True}, - {"key": "BEN2_BASE", "ready": False}, + {"key": "BIREFNET_LITE", "ready": False}, { "key": "REALESRGAN_GENERAL_WDN_X4V3", "ready": True, @@ -103,7 +103,7 @@ class AiSetupTest(BlenderTestCase): ): status = properties.ai_status() self.assertTrue(status["environment_ready"]) - self.assertEqual(status["missing_models"], ("BEN2_BASE",)) + self.assertEqual(status["missing_models"], ("BIREFNET_LITE",)) self.assertFalse(status["ready"]) self.assertEqual(properties.ai_setup_label(status), "Download Required Models…") @@ -114,7 +114,7 @@ class AiSetupTest(BlenderTestCase): "model_catalog", return_value=tuple( {"key": key, "ready": True} - for key in properties.REQUIRED_MODEL_KEYS + for key in properties.shared_model_catalog.DEFAULT_MODEL_KEYS ), ), ): @@ -148,7 +148,7 @@ class AiSetupTest(BlenderTestCase): "ai_status", return_value={ "environment_ready": False, - "missing_models": ("MOGE2_VITS_NORMAL", "BEN2_BASE"), + "missing_models": ("MOGE2_VITS_NORMAL", "BIREFNET_LITE"), "ready": False, }, ): @@ -159,7 +159,7 @@ class AiSetupTest(BlenderTestCase): "ai_status", return_value={ "environment_ready": True, - "missing_models": ("BEN2_BASE",), + "missing_models": ("BIREFNET_LITE",), "ready": False, }, ): @@ -185,13 +185,14 @@ class AiSetupTest(BlenderTestCase): with tempfile.TemporaryDirectory() as directory: root = Path(directory) moge = root / "models" / "moge-2-vits-normal-onnx" - ben2 = root / "models" / "BEN2-ONNX" / "onnx" + lite = root / "models" / "birefnet-lite" wdn = root / "models" / "realesr-general-wdn-x4v3" moge.mkdir(parents=True) - ben2.mkdir(parents=True) + lite.mkdir(parents=True) wdn.mkdir(parents=True) (moge / "model.onnx").touch() - (ben2 / "model_fp16.onnx").touch() + (lite / "model.onnx").touch() + (lite / "LICENSE.txt").touch() (wdn / "realesr-general-wdn-x4v3.onnx").touch() with ( diff --git a/tests/anyimage/operators/test_cutout_operators.py b/tests/anyimage/operators/test_cutout_operators.py index e410acc..bd5a870 100644 --- a/tests/anyimage/operators/test_cutout_operators.py +++ b/tests/anyimage/operators/test_cutout_operators.py @@ -249,7 +249,7 @@ class CutoutOperatorsTest(BlenderTestCase): operator.color_path = __file__ operator.source_object_name = "Source" operator.shape = "DEPTH_SOLID" - operator.model = self.anyimage.properties.DEFAULT_MOGE2_MODEL_KEY + operator.model = self.anyimage.preferences.DEFAULT_GEOMETRY_MODEL_KEY operator.resolution_level = 5 operator.max_ai_input_size = 2048 operator.generate_normal = True @@ -288,10 +288,8 @@ class CutoutOperatorsTest(BlenderTestCase): geometry_model="MOGE2_VITB_NORMAL", geometry_resolution_level=7, ) - with patch.object( - self.cutout_main, - "addon_preferences", - return_value=preferences, + with patch.object(self.cutout_main, "addon_preferences", return_value=preferences), patch.object( + self.anyimage.preferences, "addon_preferences", return_value=preferences ): settings = self.cutout_main.cutout_geometry_settings(None) diff --git a/tests/anyimage/operators/test_panorama_conversion.py b/tests/anyimage/operators/test_panorama_conversion.py index bac07a5..9dfd303 100644 --- a/tests/anyimage/operators/test_panorama_conversion.py +++ b/tests/anyimage/operators/test_panorama_conversion.py @@ -43,11 +43,12 @@ def test_entry_submits_current_alpha_preferences_and_cleans_cancelled_input(sour return {"CANCELLED"} monkeypatch.setattr(operators, "invoke_ai_setup_if_needed", lambda _operator: None) monkeypatch.setattr(operators, "addon_preferences", lambda _: SimpleNamespace(geometry_model="MOGE2_VITB", geometry_resolution_level=2)) + monkeypatch.setattr(operators, "configured_geometry_model", lambda _: "MOGE2_VITB_NORMAL") monkeypatch.setattr(operators, "configured_max_ai_input_size", lambda _: 1024) monkeypatch.setattr(operators, "bpy", SimpleNamespace(ops=SimpleNamespace(anyimage=SimpleNamespace(generate_panorama=submit)))) operator = SimpleNamespace(mesh_detail=4, report=Mock()) assert operators.ConvertToPanorama.execute(operator, bpy.context) == {"CANCELLED"} - assert submitted["model"] == "MOGE2_VITB" and submitted["resolution_level"] == 2 + assert submitted["model"] == "MOGE2_VITB_NORMAL" and submitted["resolution_level"] == 2 assert submitted["max_input_size"] == 1024 and submitted["mesh_detail"] == 4 assert submitted["color_path"] == submitted["input_path"] from pathlib import Path diff --git a/tests/anyimage/operators/test_texture_node_actions.py b/tests/anyimage/operators/test_texture_node_actions.py index 538840b..08f2a3d 100644 --- a/tests/anyimage/operators/test_texture_node_actions.py +++ b/tests/anyimage/operators/test_texture_node_actions.py @@ -113,7 +113,7 @@ def test_texture_action_undo_redo_restores_image_and_binding(texture, tmp_path, node = bpy.data.materials["Texture material"].node_tree.nodes["Source texture"] operator._image_target = ImageEditTarget.capture(texture_context(node)) if operation == upscale.UpscaleImage: - operator.model = "PISA_SR" + operator.model = "HAT_GAN_X4_SHARPER" operator.response(context, SimpleNamespace(directory=tmp_path)) return {"FINISHED"} @@ -175,7 +175,7 @@ def test_texture_job_request_and_completion(texture, tmp_path, operation, module stack.enter_context(patch.object(runtime, "update_ui")) stack.enter_context(patch.object(runtime, "redraw_ui")) if operation == upscale.UpscaleImage: - stack.enter_context(patch.object(upscale, "configured_upscale_model", return_value="SAFMN_REAL_X4")) + stack.enter_context(patch.object(upscale, "configured_upscale_model", return_value="HAT_GAN_X4_SHARPER")) stack.enter_context(patch.object(upscale, "configured_max_ai_input_size", return_value=2048)) assert original_execute(operator, context) == {"RUNNING_MODAL"} job = runtime.active_job @@ -185,7 +185,7 @@ def test_texture_job_request_and_completion(texture, tmp_path, operation, module assert kind == operation.job_type assert parameters["start_frame"] == 1 assert parameters["device"] == "cpu" - assert parameters["model"] == ("SAFMN_REAL_X4" if operation == upscale.UpscaleImage else "BEN2_BASE") + assert parameters["model"] == ("HAT_GAN_X4_SHARPER" if operation == upscale.UpscaleImage else "BIREFNET_LITE") input_path = operator.input_path assert operator.delete_input # Selection changes must not redirect completion. diff --git a/tests/anyimage/operators/test_upscale_undo.py b/tests/anyimage/operators/test_upscale_undo.py index 7691fc4..5f2f9b1 100644 --- a/tests/anyimage/operators/test_upscale_undo.py +++ b/tests/anyimage/operators/test_upscale_undo.py @@ -23,7 +23,7 @@ def test_upscale_result_replacement_undo_restores_source(tmp_path): def execute(operator, context): operator._image_target = ImageEditTarget.capture(SimpleNamespace(object=owner)) - operator.model = "PISA_SR" + operator.model = "HAT_GAN_X4_SHARPER" operator.response(context, SimpleNamespace(directory=tmp_path)) return {"FINISHED"} diff --git a/tests/anyimage/server/support.py b/tests/anyimage/server/support.py index 164fd4a..5468356 100644 --- a/tests/anyimage/server/support.py +++ b/tests/anyimage/server/support.py @@ -57,13 +57,13 @@ def _make_fake_av(frame_count): class FakeModels: - def __init__(self, directory="model"): + def __init__(self, directory="BEN2-ONNX"): self.model_directory = Path(directory) def directory(self, _model): return self.model_directory - def get_ben2(self, _model_directory, _device): + def get_background(self, _model_directory, _device): from server.models.onnx_ben2 import create_session return create_session(self.model_directory, "cpu"), 0.0 diff --git a/tests/anyimage/server/test_background_models.py b/tests/anyimage/server/test_background_models.py new file mode 100644 index 0000000..06b963d --- /dev/null +++ b/tests/anyimage/server/test_background_models.py @@ -0,0 +1,75 @@ +import hashlib +from unittest.mock import patch + +import numpy as np +import pytest +from PIL import Image + +from server import model_catalog +from server.jobs import remove_background +from server.model_manager import ModelManager +from server.models import background +from tests.anyimage.server.support import FakeJobContext + + +@pytest.mark.parametrize("key", model_catalog.BACKGROUND_MODELS) +@pytest.mark.parametrize("kind", ["foreground", "frames", "alpha"]) +def test_models_dispatch_keep_alpha_and_reuse_sessions(tmp_path, key, kind): + source = tmp_path / "source.png" + Image.new("RGBA", (7, 5), (60, 120, 180, 128)).save(source) + if kind == "frames": + source = tmp_path / "frames" + source.mkdir() + for index in (1, 2): + Image.new("RGBA", (7, 5), (60, 120, 180, 128)).save(source / f"{index}.png") + manager = ModelManager(tmp_path / "models") + context = FakeJobContext(tmp_path / "job", "remove-background", manager) + adapter = background.model_adapter(manager.directory(key)) + parameters = {"input": str(source), "model": key, "device": "directml", "video_frames": 2} + if kind == "alpha": + parameters["output_kind"] = kind + with patch.object(adapter, "create_session", return_value=object()) as create, patch.object( + adapter, "infer_alpha", return_value=np.full((5, 7), .4, np.float32) + ) as infer: + remove_background.run(context, parameters) + create.assert_called_once_with(manager.directory(key), "cpu" if key == "BIREFNET_HR_MATTING" else "directml") + assert infer.call_count == (2 if kind == "frames" else 1) + if kind == "alpha": + np.testing.assert_allclose(np.load(context.directory / "alpha.npy"), .4) + else: + outputs = sorted(context.directory.glob("*.png")) + assert len(outputs) == infer.call_count + with Image.open(outputs[0]) as result: + assert result.size == (7, 5) + assert result.getpixel((0, 0)) == (60, 120, 180, 51) + manager.close() + assert manager.snapshot()["loaded"]["background"] == "" + + +@pytest.mark.parametrize("key", ["BIREFNET_LITE", "BIREFNET_HR_MATTING", "HAT_GAN_X4_SHARPER"]) +def test_readiness_checks_model_and_license_corruption(tmp_path, key): + from dataclasses import replace + + content = b"verified" + files = tuple((name, len(content), hashlib.sha256(content).hexdigest()) for name in ("model.onnx", "LICENSE.txt")) + spec = replace(model_catalog.get_downloadable_model(key), r2_files=files) + directory = tmp_path / spec.directory_name + directory.mkdir() + manager = ModelManager(tmp_path) + with patch.object(model_catalog, "get_downloadable_model", return_value=spec): + assert not manager.ready(key) + for name, _, _ in files: + (directory / name).write_bytes(content) + assert manager.ready(key) + (directory / "LICENSE.txt").write_bytes(b"damaged!") + assert not manager.ready(key) + + +@pytest.mark.parametrize("key", ["PISA_SR", "SAFMN_REAL_X4", "MOGE2_VITL_NORMAL"]) +def test_retired_job_keys_fail_before_loading(key): + from server.models import upscale + + with pytest.raises(ValueError, match="Unknown upscale model"): + upscale.infer_one(None, {"model": key}, "unused.png") + with pytest.raises(ValueError, match="Unknown background removal model"): + remove_background.run(None, {"model": key}) diff --git a/tests/anyimage/server/test_job_server.py b/tests/anyimage/server/test_job_server.py index 8617a65..25415b6 100644 --- a/tests/anyimage/server/test_job_server.py +++ b/tests/anyimage/server/test_job_server.py @@ -63,7 +63,7 @@ class JobServerTest(ServerTestCase): def test_cancelling_job_keeps_server_and_model_cache_alive(self): fake_models = SimpleNamespace( - snapshot=lambda: {"loaded": {"ben2": "BEN2_BASE", "geometry": ""}}, + snapshot=lambda: {"loaded": {"background": "BEN2_BASE", "geometry": ""}}, clear=lambda: None, close=lambda: None, ) @@ -87,7 +87,7 @@ class JobServerTest(ServerTestCase): self.assertEqual(status["state"], "cancelled") self.assertFalse(server.shutdown_event.is_set()) self.assertEqual( - server.snapshot("instance")["resources"]["model_manager"]["loaded"]["ben2"], + server.snapshot("instance")["resources"]["model_manager"]["loaded"]["background"], "BEN2_BASE", ) diff --git a/tests/anyimage/server/test_model_birefnet_onnx.py b/tests/anyimage/server/test_model_birefnet_onnx.py new file mode 100644 index 0000000..caf7701 --- /dev/null +++ b/tests/anyimage/server/test_model_birefnet_onnx.py @@ -0,0 +1,36 @@ +from types import SimpleNamespace +from unittest.mock import patch + +import numpy as np +import pytest +from PIL import Image + +from server.models import onnx_birefnet + + +@pytest.mark.parametrize("size", [1024, 2048]) +def test_probability_alpha_keeps_transparency_and_source_dimensions(size): + session = SimpleNamespace( + get_inputs=lambda: [SimpleNamespace(name="image", shape=[1, 3, size, size])], + get_outputs=lambda: [SimpleNamespace(name="alpha")], + ) + mask = np.full((1, 1, size, size), 0.37, dtype=np.float32) + with patch.object(onnx_birefnet, "run_session", return_value=[mask]) as run: + result = onnx_birefnet.infer_alpha(session, Image.new("RGB", (31, 17), (64, 128, 192)), release_memory=False) + assert result.shape == (17, 31) + assert result.dtype == np.float32 + np.testing.assert_allclose(result, 0.37, atol=1e-6) + assert run.call_args.kwargs["release_memory"] is False + feed = run.call_args.args[2]["image"] + np.testing.assert_allclose(feed[0, :, 0, 0], (np.array([64, 128, 192]) / 255 - [.485, .456, .406]) / [.229, .224, .225], atol=1e-6) + + +@pytest.mark.parametrize("value", [np.nan, np.inf, -0.1, 1.1]) +def test_invalid_alpha_is_rejected(value): + session = SimpleNamespace( + get_inputs=lambda: [SimpleNamespace(name="image", shape=[1, 3, 1024, 1024])], + get_outputs=lambda: [SimpleNamespace(name="alpha")], + ) + with patch.object(onnx_birefnet, "run_session", return_value=[np.full((1, 1, 1024, 1024), value)]): + with pytest.raises(RuntimeError, match="invalid alpha"): + onnx_birefnet.infer_alpha(session, Image.new("RGB", (2, 3))) diff --git a/tests/anyimage/server/test_model_catalog.py b/tests/anyimage/server/test_model_catalog.py index daf4752..3075f8d 100644 --- a/tests/anyimage/server/test_model_catalog.py +++ b/tests/anyimage/server/test_model_catalog.py @@ -26,6 +26,16 @@ class ModelCatalogTest(unittest.TestCase): with self.assertRaisesRegex(ValueError, "Unknown model"): self.models.get_downloadable_model("UNKNOWN") + def test_three_tiers_and_retired_keys(self): + self.assertEqual(list(self.models.MOGE_MODELS), ["MOGE2_VITS_NORMAL", "MOGE2_VITB_NORMAL", "MOGE3_VITL"]) + self.assertEqual(list(self.models.BACKGROUND_MODELS), ["BIREFNET_LITE", "BEN2_BASE", "BIREFNET_HR_MATTING"]) + self.assertEqual(list(self.models.UPSCALE_MODELS), ["REALESRGAN_GENERAL_WDN_X4V3", "REALESRGAN_X4PLUS", "HAT_GAN_X4_SHARPER"]) + self.assertEqual(len(self.models.DOWNLOADABLE_MODELS), 9) + self.assertEqual(self.models.DEFAULT_MODEL_KEYS, ("MOGE2_VITS_NORMAL", "BIREFNET_LITE", "REALESRGAN_GENERAL_WDN_X4V3")) + for key in ("DA3_BASE", "DA3_SMALL", "MOGE2_VITL_NORMAL", "SAFMN_REAL_X4", "PISA_SR"): + with self.assertRaisesRegex(ValueError, "Unknown model"): + self.models.get_downloadable_model(key) + def test_ben2_download_only_includes_required_hugging_face_files(self): model = self.models.get_downloadable_model("BEN2_BASE") @@ -45,15 +55,15 @@ class ModelCatalogTest(unittest.TestCase): def test_upscale_models_use_verified_onnx_artifacts(self): expected = { "REALESRGAN_GENERAL_WDN_X4V3": ("realesr-general-wdn-x4v3.onnx", 4867419, "b021a837"), - "SAFMN_REAL_X4": ("safmn-real-x4.onnx", 22657477, "8f13867b"), + "HAT_GAN_X4_SHARPER": ("model.onnx", 45874438, "032d8b03"), "REALESRGAN_X4PLUS": ("realesrgan-x4plus.onnx", 67051618, "4ed6a45a"), } - self.assertEqual(set(self.models.UPSCALE_MODELS), set(expected) | {"PISA_SR"}) + self.assertEqual(set(self.models.UPSCALE_MODELS), set(expected)) for key, (filename, size, checksum_prefix) in expected.items(): with self.subTest(model=key): model = self.models.get_downloadable_model(key) - self.assertEqual(model.download_patterns, (filename,)) - self.assertEqual(model.ready_patterns, (filename,)) + self.assertIn(filename, model.download_patterns) + self.assertIn(filename, model.ready_patterns) self.assertEqual(model.r2_files[0][:2], (filename, size)) self.assertTrue(model.r2_files[0][2].startswith(checksum_prefix)) self.assertEqual(self.models.model_record(model)["family"], "upscale") @@ -61,14 +71,13 @@ class ModelCatalogTest(unittest.TestCase): self.assertEqual(wdn.repository, "xinntao/Real-ESRGAN") self.assertEqual(wdn.huggingface_repository, "") self.assertEqual( - self.models.UPSCALE_MODELS["SAFMN_REAL_X4"].huggingface_repository, "" + self.models.UPSCALE_MODELS["HAT_GAN_X4_SHARPER"].huggingface_repository, "" ) def test_moge2_models_use_official_onnx_models(self): for key, variant in (("MOGE2_VITS_NORMAL", "vits"), - ("MOGE2_VITB_NORMAL", "vitb"), - ("MOGE2_VITL_NORMAL", "vitl")): + ("MOGE2_VITB_NORMAL", "vitb")): with self.subTest(key=key): model = self.models.get_downloadable_model(key) self.assertEqual( @@ -84,13 +93,12 @@ class ModelCatalogTest(unittest.TestCase): def test_moge2_models_match_uploaded_r2_variants(self): self.assertEqual( list(self.models.MOGE2_MODELS), - ["MOGE2_VITS_NORMAL", "MOGE2_VITB_NORMAL", "MOGE2_VITL_NORMAL"], + ["MOGE2_VITS_NORMAL", "MOGE2_VITB_NORMAL"], ) self.assertEqual( [model.r2_directory for model in self.models.MOGE2_MODELS.values()], [ "moge-2-vits-normal-onnx", "moge-2-vitb-normal-onnx", - "moge-2-vitl-normal-onnx", ], ) diff --git a/tests/anyimage/server/test_model_inference.py b/tests/anyimage/server/test_model_inference.py index 1173dc1..9890909 100644 --- a/tests/anyimage/server/test_model_inference.py +++ b/tests/anyimage/server/test_model_inference.py @@ -93,7 +93,7 @@ class ModelInferenceTest(unittest.TestCase): upscale = load_server_module("models.upscale") -@pytest.mark.parametrize("key", ["REALESRGAN_GENERAL_WDN_X4V3", "SAFMN_REAL_X4", "REALESRGAN_X4PLUS"]) +@pytest.mark.parametrize("key", ["REALESRGAN_GENERAL_WDN_X4V3", "HAT_GAN_X4_SHARPER", "REALESRGAN_X4PLUS"]) def test_upscale_models_resize_prediction_then_restore_alpha(tmp_path, key): rng = np.random.default_rng(1) source = Image.fromarray(rng.integers(0,256,(5,7,4),dtype=np.uint8)) @@ -101,7 +101,8 @@ def test_upscale_models_resize_prediction_then_restore_alpha(tmp_path, key): source.save(path) prediction = Image.fromarray(rng.integers(0,256,(20,28,3),dtype=np.uint8)) manager = SimpleNamespace(directory=lambda _: tmp_path, get_upscale=lambda *args: object()) - with patch("server.models.onnx_upscale.infer", return_value=prediction): + with patch("server.models.onnx_upscale.infer", return_value=prediction) as infer: output = upscale.infer_one(manager, {"model":key,"device":"cpu"}, path) + assert infer.call_args.kwargs["tile_border"] == (48 if key == "HAT_GAN_X4_SHARPER" else 16) np.testing.assert_array_equal(output.convert("RGB"), prediction.resize((14,10), Image.Resampling.LANCZOS)) np.testing.assert_array_equal(output.getchannel("A"), source.getchannel("A").resize((14,10), Image.Resampling.LANCZOS)) diff --git a/tests/anyimage/server/test_model_manager.py b/tests/anyimage/server/test_model_manager.py index fedb4f0..7fb7809 100644 --- a/tests/anyimage/server/test_model_manager.py +++ b/tests/anyimage/server/test_model_manager.py @@ -6,6 +6,33 @@ from unittest.mock import patch from tests.anyimage.server.runtime_support import ServerTestCase class ModelManagerTest(ServerTestCase): + def test_background_cache_switches_model_directory_and_device_and_closes(self): + from server.models import background + + cache = self.create_cache() + loads = [] + + def create(directory, device): + result = object() + loads.append((str(directory), device, result)) + return result + + with patch.object(background, "model_adapter", return_value=SimpleNamespace(create_session=create)): + first, _ = cache.get_background("a/birefnet-lite", "cpu") + assert cache.get_background("a/birefnet-lite", "cpu") == (first, 0.0) + second, _ = cache.get_background("a/BEN2-ONNX", "cpu") + third, _ = cache.get_background("a/birefnet-hr-matting", "cpu") + fourth, _ = cache.get_background("b/birefnet-hr-matting", "cpu") + assert cache.get_background("b/birefnet-hr-matting", "directml") == (fourth, 0.0) + fifth, _ = cache.get_background("b/birefnet-lite", "directml") + assert len(loads) == 5 + assert len({id(item) for item in (first, second, third, fourth, fifth)}) == 5 + assert cache.snapshot()["background"] == "birefnet-lite" + assert loads[2][1] == loads[3][1] == "cpu" + cache.close() + assert cache.background_model is None + assert cache.snapshot()["background"] == "" + def test_required_model_download_skips_ready_model(self): manager = self.ModelManager(Path(".")) downloads = [] @@ -39,7 +66,7 @@ class ModelManagerTest(ServerTestCase): snapshot, { "loaded": { - "ben2": "", + "background": "", "geometry": "", "upscale": "", } @@ -49,7 +76,7 @@ class ModelManagerTest(ServerTestCase): self.assertNotIn("model_catalog", sys.modules) - def test_ben2_model_is_reused_for_matching_runtime_key(self): + def test_background_model_is_reused_for_matching_runtime_key(self): loads = [] ben2 = ModuleType("server.models.onnx_ben2") @@ -59,12 +86,12 @@ class ModelManagerTest(ServerTestCase): cache = self.create_cache() device = SimpleNamespace(type="cuda") with patch.dict(sys.modules, {"server.models.onnx_ben2": ben2}): - first, _load_ms = cache.get_ben2("model", device) - second, cached_load_ms = cache.get_ben2("model", device) + first, _load_ms = cache.get_background("BEN2-ONNX", device) + second, cached_load_ms = cache.get_background("BEN2-ONNX", device) self.assertIs(first, second) - self.assertEqual(loads, [("model", "cuda")]) - self.assertEqual(cache.snapshot()["ben2"], "model") + self.assertEqual(loads, [("BEN2-ONNX", "cuda")]) + self.assertEqual(cache.snapshot()["background"], "BEN2-ONNX") self.assertEqual(cached_load_ms, 0.0) @@ -73,8 +100,8 @@ class ModelManagerTest(ServerTestCase): ben2 = object() old_moge = object() upscale = object() - cache.ben2_model = ben2 - cache.ben2_key = ("ben2", "cuda") + cache.background_model = ben2 + cache.background_key = ("ben2", "cuda") cache.moge_model = old_moge cache.moge_key = ("old-moge", "cuda") cache.upscale_model = upscale @@ -85,7 +112,7 @@ class ModelManagerTest(ServerTestCase): with patch.dict(sys.modules, {"server.models.onnx_moge2": fake_moge}): self.assertIs(cache.get_moge("new-moge", "cuda"), new_moge) - self.assertIs(cache.ben2_model, ben2) + self.assertIs(cache.background_model, ben2) self.assertIs(cache.upscale_model, upscale) self.assertEqual(cache.snapshot()["geometry"], "new-moge") @@ -160,11 +187,11 @@ class ModelManagerTest(ServerTestCase): "server.models.onnx_upscale": upscale, }, ): - first_ben2, _load_ms = cache.get_ben2("ben2", "cuda") + first_ben2, _load_ms = cache.get_background("BEN2-ONNX", "cuda") assert_reused( first_ben2, - lambda: cache.get_ben2("ben2", "cuda")[0], - "release_ben2", + lambda: cache.get_background("BEN2-ONNX", "cuda")[0], + "release_background", ) first_moge2 = cache.get_moge("moge2", "cuda") diff --git a/tests/anyimage/server/test_model_pisa_sr.py b/tests/anyimage/server/test_model_pisa_sr.py deleted file mode 100644 index c32d76b..0000000 --- a/tests/anyimage/server/test_model_pisa_sr.py +++ /dev/null @@ -1,125 +0,0 @@ -import hashlib -from types import SimpleNamespace -from unittest.mock import patch -import numpy as np -import pytest -from PIL import Image - -from server.models import onnx_pisa_sr as pisa -from server.model_manager import ModelManager, ModelSessionCache -from server import model_catalog - - -class FakePisa: - def __init__(self): - self.noise = np.zeros((1, 4, 64, 64), dtype=np.float16) - self.stages = [] - self.closed = False - self.decoder_inputs = [] - - def activate(self, component): - self.component = component - self.stages.append(component) - - def run(self, tensor, final=False): - if self.component == "encoder": - return np.zeros((1,8,64,64), dtype=np.float16) - if self.component == "unet-pixel": - return np.full_like(tensor, .1) - if self.component == "unet-full": - return np.full_like(tensor, .2) - self.decoder_inputs.append(tensor.copy()) - return np.zeros((1,3,512,512), dtype=np.float16) - - def close(self): - self.closed = True - - -@pytest.mark.parametrize("size", [(1,1), (7,5), (256,256), (513,257)]) -def test_preset_output_size_and_prediction_combination(size): - model = FakePisa() - output = pisa.infer(model, Image.new("RGB", size, (32,64,96))) - assert output.size == (size[0]*2, size[1]*2) - assert model.stages == ["encoder", "unet-pixel", "unet-full", "decoder"] - for latent in model.decoder_inputs: - np.testing.assert_allclose(latent, -.17, atol=.0002) - assert np.isfinite(np.asarray(output)).all() - - -def test_cancel_cleans_intermediate_files_and_session(tmp_path): - model = FakePisa() - original = pisa.TemporaryDirectory - def temporary(**kwargs): - return original(dir=tmp_path, **kwargs) - checks = 0 - def cancel(): - nonlocal checks - checks += 1 - return checks == 4 - with patch.object(pisa, "TemporaryDirectory", side_effect=temporary): - with pytest.raises(RuntimeError, match="cancelled"): - pisa.infer(model, Image.new("RGB", (256,256)), cancel) - assert model.closed - assert not list(tmp_path.iterdir()) - - -def test_stage_releases_previous_session_before_loading(tmp_path): - np.save(tmp_path / "noise.npy", np.zeros((1,4,64,64), dtype=np.float16)) - model = pisa.PisaSession(tmp_path, "cpu") - loaded = [] - def create(path, device): - assert model.session is None - loaded.append(path.stem) - return SimpleNamespace(get_providers=lambda: ["CPUExecutionProvider"], disable_fallback=lambda: None) - with patch.object(pisa, "create_session", side_effect=create): - model.activate("unet-pixel") - model.activate("unet-full") - assert loaded == ["unet-pixel", "unet-full"] - - -def test_resource_error_closes_pipeline(): - model = FakePisa() - with patch.object(model, "run", side_effect=RuntimeError("out of memory")): - with pytest.raises(pisa.UpscaleResourceError): - pisa.infer(model, Image.new("RGB", (256,256))) - assert model.closed - - -def test_pisa_readiness_checks_every_asset_and_corruption(tmp_path): - content = b"verified" - digest = hashlib.sha256(content).hexdigest() - files = tuple((name, len(content), digest) for name in ("encoder.onnx", "noise.npy")) - spec = SimpleNamespace(directory_name="pisa-sr", r2_files=files) - directory = tmp_path / "pisa-sr" - directory.mkdir() - manager = ModelManager(tmp_path) - with patch.object(model_catalog, "get_downloadable_model", return_value=spec): - assert not manager.ready("PISA_SR") - for name, _, _ in files: - (directory / name).write_bytes(content) - assert manager.ready("PISA_SR") - (directory / "noise.npy").write_bytes(b"damaged!") - assert not manager.ready("PISA_SR") - - -def test_pisa_cache_reuse_and_clear(tmp_path): - cache = ModelSessionCache() - with patch.object(pisa, "PisaSession", return_value=object()) as create: - first = cache.get_upscale(tmp_path / "pisa-sr", "cpu") - assert cache.get_upscale(tmp_path / "pisa-sr", "cpu") is first - assert create.call_count == 1 - cache.close() - assert cache.upscale_model is None - - -def test_coreml_returns_explicit_support_reason(tmp_path): - with pytest.raises(RuntimeError, match="select CPU"): - pisa.PisaSession(tmp_path, "coreml") - - -def test_clear_closes_pisa_components_even_with_external_reference(): - cache = ModelSessionCache() - model = FakePisa() - cache.upscale_model = model - cache.close() - assert model.closed diff --git a/tests/anyimage/server/test_model_upscale_onnx.py b/tests/anyimage/server/test_model_upscale_onnx.py index 8c6d6a0..117b683 100644 --- a/tests/anyimage/server/test_model_upscale_onnx.py +++ b/tests/anyimage/server/test_model_upscale_onnx.py @@ -26,6 +26,18 @@ class FakeSession: class UpscaleOnnxTest(unittest.TestCase): + def test_wider_context_preserves_pixels_across_odd_and_small_images(self): + for width, height in [(1, 1), (7, 5), (321, 163)]: + pixels = np.random.default_rng(7).integers(0, 256, (height, width, 3), dtype=np.uint8) + result = onnx_upscale.infer(FakeSession(), Image.fromarray(pixels), tile_border=48) + expected = np.repeat(np.repeat(pixels, 4, axis=0), 4, axis=1) + np.testing.assert_array_equal(np.asarray(result), expected) + + def test_nonfinite_pixels_fail_before_conversion(self): + with patch.object(onnx_upscale, "run_session", return_value=[np.full((1, 3, 1024, 1024), np.nan)]): + with self.assertRaisesRegex(RuntimeError, "non-finite"): + onnx_upscale.infer(FakeSession(), Image.new("RGB", (7, 5))) + def test_preprocess_matches_dynamic_rgb_layout(self): image = Image.new("RGB", (7, 5), (64, 128, 192)) diff --git a/tests/anyimage/server/test_server_jobs.py b/tests/anyimage/server/test_server_jobs.py index a759ee9..af0aceb 100644 --- a/tests/anyimage/server/test_server_jobs.py +++ b/tests/anyimage/server/test_server_jobs.py @@ -16,7 +16,7 @@ class ServerJobsTest(unittest.TestCase): cls.background = load_server_module("jobs.remove_background") - cls.ben2 = load_server_module("models.ben2") + cls.background_model = load_server_module("models.background") cls.onnx_upscale = load_server_module("models.onnx_upscale") @@ -96,7 +96,7 @@ class ServerJobsTest(unittest.TestCase): patch.object(onnx_ben2, "infer_alpha", return_value=np.full((1, 2), 128 / 255, np.float32)) as infer, patch.object(Image.Image, "save", autospec=True, side_effect=Image.Image.save) as save, ): - self.background.run(context, {"input": str(input_path), "model": "BEN2", "device": "cpu"}) + self.background.run(context, {"input": str(input_path), "model": "BEN2_BASE", "device": "cpu"}) with Image.open(root / "result" / "image.0001.png") as image: self.assertEqual(image.getpixel((0, 0)), (255, 64, 32, 64)) self.assertEqual(infer.call_args.kwargs, {"release_memory": True}) @@ -126,15 +126,15 @@ class ServerJobsTest(unittest.TestCase): patch.object(onnx_ben2, "create_session", return_value=object()), patch.object(onnx_ben2, "infer_alpha", side_effect=infer), ): - result, _load_ms, _inference_ms = self.ben2.infer_foreground( + result, _load_ms, _inference_ms = self.background_model.infer_foreground( input_path, - "model", + "BEN2-ONNX", "cpu", preserve_input_alpha=True, ) - raw_result, _load_ms, _inference_ms = self.ben2.infer_foreground( + raw_result, _load_ms, _inference_ms = self.background_model.infer_foreground( input_path, - "model", + "BEN2-ONNX", "cpu", ) @@ -166,7 +166,7 @@ class ServerJobsTest(unittest.TestCase): patch.object(onnx_ben2, "create_session", return_value=object()), patch.object(onnx_ben2, "infer_alpha", return_value=np.ones((1, 2), np.float32)) as infer, ): - self.background.run(context, {"input": str(frames), "model": "BEN2", "device": "cpu", "video_frames": 8}) + self.background.run(context, {"input": str(frames), "model": "BEN2_BASE", "device": "cpu", "video_frames": 8}) self.assertEqual([call.kwargs["release_memory"] for call in infer.call_args_list], [False, True]) for index in (1, 2): with Image.open(context.directory / f"image.{index:04d}.png") as image: @@ -186,7 +186,7 @@ class ServerJobsTest(unittest.TestCase): patch.object(onnx_ben2, "create_session", return_value=object()), patch.object(onnx_ben2, "infer_alpha", return_value=alpha), ): - result = self.background.run(context, {"input": str(path), "model": "BEN2", "device": "cpu", "output_kind": "alpha"}) + result = self.background.run(context, {"input": str(path), "model": "BEN2_BASE", "device": "cpu", "output_kind": "alpha"}) self.assertEqual(result, {"alpha": "alpha.npy"}) np.testing.assert_array_equal(np.load(context.directory / result["alpha"], allow_pickle=False), alpha) self.assertFalse(list(context.directory.glob("*.png"))) diff --git a/tests/anyimage/test_preferences.py b/tests/anyimage/test_preferences.py index 238e92d..18bb673 100644 --- a/tests/anyimage/test_preferences.py +++ b/tests/anyimage/test_preferences.py @@ -118,6 +118,23 @@ class PreferencesTest(BlenderTestCase): geometry = properties.geometry_model_enum_items(None, None) self.assertEqual(tuple(item[0] for item in geometry), expected) + def test_model_preferences_keep_valid_choices_and_normalize_retired_keys(self): + module = self.anyimage.preferences + stored = self.fake_bpy.context.preferences.addons["anyimage"].preferences + for attribute, reader, valid, retired, default in ( + ("geometry_model", module.configured_geometry_model, "MOGE3_VITL", "MOGE2_VITL_NORMAL", "MOGE2_VITS_NORMAL"), + ("background_model", module.configured_background_model, "BEN2_BASE", "UNKNOWN", "BIREFNET_LITE"), + ("upscale_model", module.configured_upscale_model, "REALESRGAN_X4PLUS", "PISA_SR", "REALESRGAN_GENERAL_WDN_X4V3"), + ): + setattr(stored, attribute, valid) + self.assertEqual(reader(), valid) + setattr(stored, attribute, retired) + self.assertEqual(reader(), default) + self.assertEqual(getattr(stored, attribute), default) + self.assertEqual({item[0]: item[-1] for item in module.BACKGROUND_MODEL_ITEMS}["BEN2_BASE"], 0) + self.assertEqual({item[0]: item[-1] for item in module.UPSCALE_MODEL_ITEMS}["REALESRGAN_X4PLUS"], 2) + self.assertEqual({item[0]: item[-1] for item in self.anyimage.properties.geometry_model_enum_items(None, None)}["MOGE3_VITL"], 3) + def test_global_ai_input_size_defaults_and_clamps(self): preference_module = self.anyimage.preferences @@ -212,7 +229,7 @@ class PreferencesTest(BlenderTestCase): preferences.layout = Layout() preferences.geometry_model = "MOGE2_VITS_NORMAL" preferences.background_model = "BEN2_BASE" - preferences.upscale_model = "SAFMN_REAL_X4" + preferences.upscale_model = "HAT_GAN_X4_SHARPER" catalog = tuple( { "key": key, @@ -225,7 +242,7 @@ class PreferencesTest(BlenderTestCase): ("MOGE2_VITS_NORMAL", "moge"), ("BEN2_BASE", "ben2"), ("REALESRGAN_GENERAL_WDN_X4V3", "upscale"), - ("SAFMN_REAL_X4", "upscale"), + ("HAT_GAN_X4_SHARPER", "upscale"), ) ) with ( @@ -249,7 +266,7 @@ class PreferencesTest(BlenderTestCase): self.assertEqual( [download.model for download in downloads], - ["MOGE2_VITS_NORMAL", "BEN2_BASE", "SAFMN_REAL_X4"], + ["MOGE2_VITS_NORMAL", "BEN2_BASE", "HAT_GAN_X4_SHARPER"], ) self.assertEqual( [ @@ -409,7 +426,7 @@ class PreferencesTest(BlenderTestCase): operator.analysis_input_path = __file__ operator.source_object_name = "Source" operator.mesh_detail = 0 - operator.model = self.anyimage.properties.DEFAULT_MOGE2_MODEL_KEY + operator.model = self.anyimage.preferences.DEFAULT_GEOMETRY_MODEL_KEY operator.resolution_level = 5 operator.max_input_size = 1536 context = SimpleNamespace( @@ -450,7 +467,7 @@ class PreferencesTest(BlenderTestCase): patch.object( self.upscale, "configured_upscale_model", - return_value="SAFMN_REAL_X4", + return_value="HAT_GAN_X4_SHARPER", ), patch.object( self.upscale, @@ -471,5 +488,5 @@ class PreferencesTest(BlenderTestCase): result = operator.execute(context) self.assertEqual(result, {"RUNNING_MODAL"}) - self.assertEqual(operator.model, "SAFMN_REAL_X4") + self.assertEqual(operator.model, "HAT_GAN_X4_SHARPER") self.assertEqual(operator.max_input_size, 1024) diff --git a/tests/support/blender.py b/tests/support/blender.py index 7600b5e..8d4a36a 100644 --- a/tests/support/blender.py +++ b/tests/support/blender.py @@ -164,6 +164,12 @@ class BlenderTestCase(unittest.TestCase): def setUp(self): + timer_callbacks = set() + self.fake_bpy.app.timers = SimpleNamespace( + is_registered=timer_callbacks.__contains__, + register=lambda callback, **options: timer_callbacks.add(callback), + unregister=timer_callbacks.remove, + ) preferences = SimpleNamespace( storage_root="", device="CUDA", diff --git a/tests/tools/models/test_cli.py b/tests/tools/models/test_cli.py index cd768e5..d6eb53e 100644 --- a/tests/tools/models/test_cli.py +++ b/tests/tools/models/test_cli.py @@ -50,10 +50,10 @@ def test_valid_model_cache_skips_all_preparation(): @pytest.mark.parametrize("valid", [True, False]) -def test_pisa_export_validates_all_outputs_before_replacing(tmp_path, valid): +def test_isolated_export_validates_all_outputs_before_replacing(tmp_path, valid): module = load_sync_models() digest = module.hashlib.sha256(b"new").hexdigest() - models = tuple(module.ModelFile("PISA_SR", "pisa-sr", name, 3, digest) for name in ("a.onnx", "b.onnx")) + models = tuple(module.ModelFile("HAT_GAN_X4_SHARPER", "hat-gan-x4-sharper", name, 3, digest) for name in ("a.onnx", "b.onnx")) def export(command, **kwargs): destination = Path(command[-1]) @@ -71,7 +71,7 @@ def test_pisa_export_validates_all_outputs_before_replacing(tmp_path, valid): module.prepare_model(models[0]) assert not models[0].destination.exists() assert models[1].destination.read_bytes() == b"old" - assert not list(tmp_path.glob(".pisa-sr-*")) + assert not list(tmp_path.glob(".hat-gan-x4-sharper-*")) def test_local_export_failure_cleans_temporary_output(tmp_path): @@ -120,7 +120,7 @@ class SyncModelsTests(unittest.TestCase): builds = self.sync.MODEL_BUILDS self.assertEqual( set(builds), - {"REALESRGAN_GENERAL_WDN_X4V3", "SAFMN_REAL_X4"}, + {"REALESRGAN_GENERAL_WDN_X4V3"}, ) for build in builds.values(): self.assertTrue(build.source_url.startswith("https://github.com/")) diff --git a/tests/tools/nodes/test_depth_relief_planes.py b/tests/tools/nodes/test_depth_relief_planes.py index 93fe566..2a1feee 100644 --- a/tests/tools/nodes/test_depth_relief_planes.py +++ b/tests/tools/nodes/test_depth_relief_planes.py @@ -235,10 +235,67 @@ def test_surface_interface_and_node_structure(surface): validate_node_structure(obj.modifiers[0].node_group) +@pytest.mark.parametrize("thickness", [0.0, 0.4]) +@pytest.mark.parametrize("direction", [(0.2, -0.1, 1), (-0.2, 0.1, 1), (0.4, -0.2, 2), (0, 0, 0), (1, 0, 0)]) +def test_relief_direction_adjusts_depth_and_keeps_base(relief, thickness, direction): + obj, set_value, _, _, image = relief + pixels = np.ones((8, 16, 4), dtype=np.float32) + pixels[..., 2] = 2 + image.pixels.foreach_set(pixels.ravel()) + image.update() + set_value("Thickness", thickness) + set_value("Reference Depth", 2.0) + layers = evaluated(obj)[0] + weights = np.clip(layers[:, 2] / thickness, 0, 1) if thickness else np.ones(len(layers)) + set_value("Depth Scale", 0.5) + base = evaluated(obj) + set_value("Depth Direction", direction) + result = evaluated(obj) + front = base[0][:, 2] > 0 + expected = base[0].copy() + if direction[2]: + expected[front, 2] += 0.5 * weights[front] * (expected[front, :2] @ np.array(direction[:2])) / direction[2] + np.testing.assert_allclose(result[0], expected, atol=1e-6) + assert result[1] == base[1] + np.testing.assert_array_equal(result[2], base[2]) + + +@pytest.mark.parametrize("thickness", [0.0, 0.4]) +def test_relief_smooth_preserves_topology_uv_and_materials(relief, thickness): + obj, set_value, _, _, image = relief + pixels = np.asarray(image.pixels[:], dtype=np.float32).reshape(8, 16, 4) + pixels[..., 2] += (np.arange(16) % 2) * 0.5 + image.pixels.foreach_set(pixels.ravel()) + image.update() + set_value("Thickness", thickness) + set_value("Reference Depth", 2.0) + set_value("Depth Scale", 1.0) + set_value("Depth Direction", (0.1, -0.1, 1)) + base = evaluated(obj) + set_value("Smooth", 3) + result = evaluated(obj) + assert np.isfinite(result[0]).all() + assert np.max(abs(result[0] - base[0])) > 1e-4 + assert result[1] == base[1] + np.testing.assert_array_equal(result[2], base[2]) + assert result[3:] == base[3:] + set_value("Smooth Weight", 0.0) + np.testing.assert_allclose(evaluated(obj)[0], base[0], atol=1e-6) + set_value("Smooth Weight", 1.0) + set_value("Smooth", 0) + np.testing.assert_allclose(evaluated(obj)[0], base[0], atol=1e-6) + + def test_relief_interface_and_node_structure(relief): obj, _, _, inputs, _ = relief - assert list(inputs) == ["Geometry", "Subdivide", "Thickness", "Depth Scale", - "Reference Depth", "Uniform Scale", "Depth Image"] + assert list(inputs) == ["Geometry", "Subdivide", "Thickness", "Depth Scale", "Smooth", + "Reference Depth", "Smooth Weight", "Depth Direction", "Uniform Scale", "Depth Image"] + assert inputs["Smooth"].default_value == 0 + assert (inputs["Smooth"].min_value, inputs["Smooth"].max_value) == (0, 20) + assert inputs["Smooth Weight"].default_value == 1 + assert inputs["Smooth Weight"].subtype == "FACTOR" + assert tuple(inputs["Depth Direction"].default_value) == (0, 0, 1) + assert inputs["Depth Direction"].subtype == "DIRECTION" assert inputs["Thickness"].default_value == 0.5 assert inputs["Subdivide"].default_value == 6 assert inputs["Depth Scale"].default_value == 0.5 @@ -356,6 +413,8 @@ def test_plane_interface_panels(source, kind): controls = ["Reference Depth"] if kind == "DEPTH": controls += ["Normal Smooth", "Smooth Weight", "Validity Threshold"] + else: + controls += ["Smooth Weight", "Depth Direction"] assert [item.name for item in panels[0].interface_items] == controls assert [item.name for item in panels[1].interface_items] == ["Uniform Scale", "Depth Image"] assert all(panel.default_closed for panel in panels) diff --git a/tests/tools/nodes/test_shadeless.py b/tests/tools/nodes/test_shadeless.py index 26917e1..794e849 100644 --- a/tests/tools/nodes/test_shadeless.py +++ b/tests/tools/nodes/test_shadeless.py @@ -1,6 +1,7 @@ """Render the bundled Shadeless group's opacity and HDR color endpoints.""" import bpy +import addon_utils import numpy as np from anyimage.common.node import load_node_group @@ -8,6 +9,7 @@ from anyimage.common.node import load_node_group def test_shadeless_alpha_is_mixed_inside_the_group(tmp_path): bpy.ops.wm.read_factory_settings(use_empty=True) + addon_utils.enable("cycles", default_set=False) group = load_node_group("O Shadeless") inputs = {s.name: s for s in group.interface.items_tree if s.item_type == "SOCKET" and s.in_out == "INPUT"} assert list(inputs) == ["Color", "Alpha"] and inputs["Alpha"].default_value == 1 diff --git a/tests/tools/test_pack.py b/tests/tools/test_pack.py index ee09293..54c2ba5 100644 --- a/tests/tools/test_pack.py +++ b/tests/tools/test_pack.py @@ -47,7 +47,7 @@ class ExtensionPackagingTest(unittest.TestCase): records = {} with tempfile.TemporaryDirectory() as directory: directory_path = Path(directory) - wheel_path = directory_path / "blendjob-0.1.16-py3-none-any.whl" + wheel_path = directory_path / "blendjob-0.1.17-py3-none-any.whl" wheel_path.write_bytes(b"wheel") original_dist = pack.DIST_DIR pack.DIST_DIR = directory_path / "dist" @@ -141,7 +141,7 @@ class ExtensionPackagingTest(unittest.TestCase): self.assertNotIn("server/selection/polygon_mask.py", names) self.assertNotIn("server/jobs/prepare_geometry_data.py", names) self.assertNotIn("server/geometry/cutout.py", names) - self.assertIn("server/models/ben2.py", names) + self.assertIn("server/models/background.py", names) self.assertNotIn("server/models/da3.py", names) self.assertIn("server/models/moge.py", names) self.assertNotIn("server/jobs/debug_predictions.py", names) @@ -150,6 +150,8 @@ class ExtensionPackagingTest(unittest.TestCase): self.assertIn("server/model_catalog.py", names) self.assertIn("server/model_download.py", names) self.assertIn("server/models/onnx_ben2.py", names) + self.assertIn("server/models/onnx_birefnet.py", names) + self.assertNotIn("server/models/onnx_pisa_sr.py", names) self.assertNotIn("server/models/onnx_da3.py", names) self.assertIn("server/models/onnx_moge2.py", names) self.assertIn("server/models/onnx_moge3.py", names) @@ -159,7 +161,7 @@ class ExtensionPackagingTest(unittest.TestCase): self.assertIn("operators/cutout_tool/polygon.py", names) self.assertIn("licenses/scikit-image.txt", names) self.assertIn( - "wheels/blendjob-0.1.16-py3-none-any.whl", + "wheels/blendjob-0.1.17-py3-none-any.whl", names, ) self.assertFalse(any(name.startswith("blendjob/") for name in names)) @@ -177,7 +179,7 @@ class ExtensionPackagingTest(unittest.TestCase): self.assertEqual( pack.project_dependencies(), [ - "blendjob==0.1.16", + "blendjob==0.1.17", "numpy>=1.26,<2.0", "scipy==1.15.3", ], @@ -188,7 +190,7 @@ class ExtensionPackagingTest(unittest.TestCase): wheels = [ Path(name) for name in ( - "blendjob-0.1.16-py3-none-any.whl", + "blendjob-0.1.17-py3-none-any.whl", "numpy-1.26.4-cp311-cp311-win_amd64.whl", "scipy-1.15.3-cp311-cp311-win_amd64.whl", ) @@ -217,12 +219,12 @@ class ExtensionPackagingTest(unittest.TestCase): if target.name == "3.13": constraints = Path(command[command.index("--constraint") + 1]) self.assertEqual( - constraints.read_text(), "blendjob==0.1.16" + constraints.read_text(), "blendjob==0.1.17" ) else: self.assertNotIn("--constraint", command) for name in ( - "blendjob-0.1.16-py3-none-any.whl", + "blendjob-0.1.17-py3-none-any.whl", f"scipy-1.15.3-{tag}-{tag}-{config.pip_platform}.whl", f"numpy-2.1.0-{tag}-{tag}-{config.pip_platform}.whl", ): @@ -238,7 +240,7 @@ class ExtensionPackagingTest(unittest.TestCase): ("numpy>=1.26,<2.0", "numpy>=2.1,<3.0"), ): command = call.args[0] - self.assertIn("blendjob==0.1.16", command) + self.assertIn("blendjob==0.1.17", command) self.assertIn(str(pack.SOURCE_DIR / "wheels"), command) self.assertIn(f"--python-version={version}", command) self.assertIn("--implementation=cp", command) @@ -272,7 +274,7 @@ class ExtensionPackagingTest(unittest.TestCase): def fake_run(command, check): target = Path(command[command.index("--dest") + 1]) if target.name == "3.11": - (target / "blendjob-0.1.16-py3-none-any.whl").write_bytes(b"wheel") + (target / "blendjob-0.1.17-py3-none-any.whl").write_bytes(b"wheel") with patch.object(pack.subprocess, "run", side_effect=fake_run): with self.assertRaisesRegex(RuntimeError, "Python 3.13"): @@ -341,7 +343,7 @@ class ExtensionPackagingTest(unittest.TestCase): target = Path(command[command.index("--dest") + 1]) if target.name == "3.13": raise RuntimeError("download failed") - (target / "blendjob-0.1.16-py3-none-any.whl").write_bytes(b"wheel") + (target / "blendjob-0.1.17-py3-none-any.whl").write_bytes(b"wheel") with ( patch.object(pack, "SOURCE_DIR", source), diff --git a/tools/models/birefnet_export.py b/tools/models/birefnet_export.py new file mode 100644 index 0000000..ab7a1d0 --- /dev/null +++ b/tools/models/birefnet_export.py @@ -0,0 +1,108 @@ +# /// script +# requires-python = ">=3.11,<3.12" +# dependencies = [ +# "torch==2.13.0", "torchvision==0.28.0", "onnx==1.22.0", +# "numpy==2.4.6", "transformers==4.57.6", "timm==1.0.25", +# "kornia==0.8.2", "einops==0.8.2", "safetensors==0.7.0", +# ] +# [tool.uv.sources] +# torch = { index = "pytorch" } +# torchvision = { index = "pytorch" } +# [[tool.uv.index]] +# name = "pytorch" +# url = "https://download.pytorch.org/whl/cu130" +# explicit = true +# /// +"""Export pinned official BiRefNet weights to fixed-resolution ONNX alpha.""" + +import argparse +import hashlib +import json +import sys +from pathlib import Path +from urllib.request import urlopen + +import onnx +import torch +from huggingface_hub import hf_hub_download +from transformers import AutoModelForImageSegmentation + +from onnx_weights import compress_weights, compress_exact_constants +from deform_sampling import deform_conv, validate_deform + + +SOURCES = { + "lite": ("ZhengPeng7/BiRefNet_lite", "aa62cd87eafb9cc43056d08ef3615a14628b831d", 1024), + "hr-matting": ("ZhengPeng7/BiRefNet_HR-matting", "5d6b6f8adcb5b417c871b1d84ceaae9871355b7f", 2048), +} +WEIGHT_HASHES = { + "lite": "4417d89795250e698c3cb0ae8df15743810065f646f48a694fdfa7ca052d0815", + "hr-matting": "a5a4de698739ea5e0e8bbab28e1b293dde95092b87a442d566cbc585c53cef55", +} +LICENSE_REVISION = "ebcc0bc8ec7fe919cec829f2dea656b3078acddc" + + +class AlphaModel(torch.nn.Module): + def __init__(self, model): + super().__init__() + self.model = model + + def forward(self, image): + return self.model(image)[-1].sigmoid() + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--variant", choices=SOURCES, required=True) + parser.add_argument("--destination", type=Path, required=True) + parser.add_argument("--checkpoint", type=Path) + parser.add_argument("--fp32-output", type=Path) + args = parser.parse_args() + repository, revision, size = SOURCES[args.variant] + checkpoint = args.checkpoint / "model.safetensors" if args.checkpoint else Path( + hf_hub_download(repository, "model.safetensors", revision=revision) + ) + with checkpoint.open("rb") as stream: + if hashlib.file_digest(stream, "sha256").hexdigest() != WEIGHT_HASHES[args.variant]: + raise RuntimeError("BiRefNet checkpoint SHA-256 mismatch") + torch.set_num_threads(8) + model = AutoModelForImageSegmentation.from_pretrained( + str(args.checkpoint) if args.checkpoint else repository, + revision=revision, + trust_remote_code=True, + ).eval() + validate_deform() + sys.modules[type(model).__module__].deform_conv2d = deform_conv + device = "cuda" if torch.cuda.is_available() else "cpu" + model = AlphaModel(model).to(device) + args.destination.mkdir(parents=True, exist_ok=True) + output = args.destination / "model.onnx" + with torch.inference_mode(): + torch.onnx.export( + model, torch.zeros(1, 3, size, size, device=device), + str(output), + input_names=["image"], output_names=["alpha"], + opset_version=17, dynamo=False, do_constant_folding=False, + ) + graph = onnx.load(output) + if args.fp32_output: + onnx.save(graph, args.fp32_output) + compress_weights(graph) + compress_exact_constants(graph) + onnx.checker.check_model(graph) + onnx.save(graph, output) + license_url = f"https://raw.githubusercontent.com/ZhengPeng7/BiRefNet/{LICENSE_REVISION}/LICENSE" + license_text = urlopen(license_url, timeout=60).read() + if hashlib.sha256(license_text).hexdigest() != "92a7089e0915fc32bc40067560b398f1e6a7a5958abd7d04eda393629a5acefb": + raise RuntimeError("BiRefNet license SHA-256 mismatch") + (args.destination / "LICENSE.txt").write_bytes(license_text) + with output.open("rb") as stream: + checksum = hashlib.file_digest(stream, "sha256").hexdigest() + record = {"repository": repository, "revision": revision, "input_size": size, + "bytes": output.stat().st_size, "sha256": checksum, + "checkpoint_sha256": WEIGHT_HASHES[args.variant]} + print(json.dumps(record), flush=True) + + +if __name__ == "__main__": + main() diff --git a/tools/models/cli.py b/tools/models/cli.py index 8ae06df..14c9cd8 100644 --- a/tools/models/cli.py +++ b/tools/models/cli.py @@ -63,18 +63,13 @@ MODEL_BUILDS = { ), exporter="realesrgan_wdn_x4", ), - "SAFMN_REAL_X4": BuildSpec( - source_url=( - "https://github.com/sunny2109/SAFMN/releases/download/" - "v0.1.0/SAFMN_L_Real_LSDIR_x4-v2.pth" - ), - source_filename="SAFMN_L_Real_LSDIR_x4-v2.pth", - source_size=45002097, - source_sha256=( - "f1ac0ee3ee143fbbc49aff6584cb7b48b71ba8d1961dfb8ad076b98ca4799ae2" - ), - exporter="safmn_real_x4", - ), +} + +ISOLATED_EXPORTERS = { + "MOGE3_VITL": ("moge3_export.py", ()), + "HAT_GAN_X4_SHARPER": ("hat_export.py", ()), + "BIREFNET_LITE": ("birefnet_export.py", ("--variant", "lite")), + "BIREFNET_HR_MATTING": ("birefnet_export.py", ("--variant", "hr-matting")), } @@ -194,21 +189,24 @@ def prepare_model(model): if model.build: build_model(model) return - if model.key in {"PISA_SR", "MOGE3_VITL"}: + if model.key in ISOLATED_EXPORTERS: model.destination.parent.mkdir(parents=True, exist_ok=True) - exporter = "pisa_sr_export.py" if model.key == "PISA_SR" else "moge3_export.py" + exporter, arguments = ISOLATED_EXPORTERS[model.key] with TemporaryDirectory(prefix=f".{model.folder}-", dir=MODELS_DIR) as directory: subprocess.run( ["uv", "run", "--script", str(Path(__file__).with_name(exporter)), - "--destination", directory], + *arguments, "--destination", directory], cwd=PROJECT_ROOT, check=True, ) models = [item for item in MODEL_FILES if item.key == model.key] for item in models: if not verify_file(Path(directory) / item.filename, item.size, item.sha256): raise RuntimeError(f"{model.key} export differs from the published catalog") - for path in Path(directory).iterdir(): - path.replace(model.destination.parent / path.name) + for item in models: + path = Path(directory) / item.filename + destination = MODELS_DIR / item.folder / item.filename + destination.parent.mkdir(parents=True, exist_ok=True) + path.replace(destination) return if not model.source_url: raise RuntimeError(f"No upstream source or builder for {model.key}") diff --git a/tools/models/deform_sampling.py b/tools/models/deform_sampling.py new file mode 100644 index 0000000..c17ff04 --- /dev/null +++ b/tools/models/deform_sampling.py @@ -0,0 +1,47 @@ +"""Express BiRefNet's single-group deformable convolution with GridSample.""" + +import torch +from torch.nn import functional as F + + +def deform_conv(input, offset, weight, bias=None, stride=(1, 1), padding=0, mask=None): + batch, channels, height, width = map(int, input.shape) + output_height, output_width = map(int, offset.shape[-2:]) + kernel_height, kernel_width = map(int, weight.shape[-2:]) + count = kernel_height * kernel_width + y, x = torch.meshgrid( + torch.arange(output_height, dtype=input.dtype, device=input.device), + torch.arange(output_width, dtype=input.dtype, device=input.device), indexing="ij", + ) + ky, kx = torch.meshgrid( + torch.arange(kernel_height, dtype=input.dtype, device=input.device), + torch.arange(kernel_width, dtype=input.dtype, device=input.device), indexing="ij", + ) + offsets = offset.reshape(batch, count, 2, output_height, output_width) + sy = y[None, None] * stride[0] - padding + ky.reshape(1, count, 1, 1) + offsets[:, :, 0] + sx = x[None, None] * stride[1] - padding + kx.reshape(1, count, 1, 1) + offsets[:, :, 1] + coordinates = torch.stack((sx * (2.0 / (width - 1)) - 1, sy * (2.0 / (height - 1)) - 1), -1) + grid = coordinates.permute(0, 2, 3, 1, 4).reshape(batch, output_height, output_width * count, 2) + sampled = F.grid_sample(input, grid, mode="bilinear", padding_mode="zeros", align_corners=True) + sampled = sampled.reshape(batch, channels, output_height, output_width, count).permute(0, 1, 4, 2, 3) + sampled = (sampled * mask[:, None]).reshape(batch, channels * count, output_height, output_width) + return F.conv2d(sampled, weight.reshape(int(weight.shape[0]), channels * count, 1, 1), bias) + + +def validate_deform(): + from torchvision.ops import deform_conv2d + + torch.manual_seed(12) + for kernel in (1, 3, 7): + image = torch.randn(1, 4, 17, 19) + offset = torch.randn(1, kernel * kernel * 2, 17, 19) * 3 + mask = torch.rand(1, kernel * kernel, 17, 19) * 2 + weights = torch.randn(5, 4, kernel, kernel) + bias = torch.randn(5) + reference = deform_conv2d(image, offset, weights, bias, padding=kernel // 2, mask=mask) + actual = deform_conv(image, offset, weights, bias, padding=kernel // 2, mask=mask) + torch.testing.assert_close(actual, reference, atol=1e-4, rtol=1e-4) + + +if __name__ == "__main__": + validate_deform() diff --git a/tools/models/exports.py b/tools/models/exports.py index b3fc3aa..a8ef920 100644 --- a/tools/models/exports.py +++ b/tools/models/exports.py @@ -98,104 +98,9 @@ def export_realesrgan_wdn_x4(checkpoint_path, destination): _validate_export(onnx, onnxruntime, torch, model, sample, destination) -def export_safmn_real_x4(checkpoint_path, destination): - onnx, onnxruntime, torch = _require_build_dependencies() - nn = torch.nn - functional = torch.nn.functional - - class LayerNorm(nn.Module): - def __init__(self, channels): - super().__init__() - self.weight = nn.Parameter(torch.ones(channels)) - self.bias = nn.Parameter(torch.zeros(channels)) - self.eps = 1e-6 - - def forward(self, value): - mean = value.mean(1, keepdim=True) - variance = (value - mean).pow(2).mean(1, keepdim=True) - value = (value - mean) / torch.sqrt(variance + self.eps) - return self.weight[:, None, None] * value + self.bias[:, None, None] - - class SAFM(nn.Module): - def __init__(self, channels): - super().__init__() - chunk_channels = channels // 4 - self.mfr = nn.ModuleList( - nn.Conv2d(chunk_channels, chunk_channels, 3, 1, 1, groups=chunk_channels) - for _ in range(4) - ) - self.aggr = nn.Conv2d(channels, channels, 1) - self.act = nn.GELU() - - def forward(self, value): - height, width = value.shape[-2:] - outputs = [] - for level, chunk in enumerate(value.chunk(4, dim=1)): - if level: - size = (height // 2**level, width // 2**level) - chunk = functional.adaptive_max_pool2d(chunk, size) - chunk = self.mfr[level](chunk) - if level: - chunk = functional.interpolate( - chunk, - size=(height, width), - mode="nearest", - ) - outputs.append(chunk) - return self.act(self.aggr(torch.cat(outputs, dim=1))) * value - - class CCM(nn.Module): - def __init__(self, channels): - super().__init__() - self.ccm = nn.Sequential( - nn.Conv2d(channels, channels * 2, 3, 1, 1), - nn.GELU(), - nn.Conv2d(channels * 2, channels, 1), - ) - - def forward(self, value): - return self.ccm(value) - - class AttBlock(nn.Module): - def __init__(self, channels): - super().__init__() - self.norm1 = LayerNorm(channels) - self.norm2 = LayerNorm(channels) - self.safm = SAFM(channels) - self.ccm = CCM(channels) - - def forward(self, value): - value = self.safm(self.norm1(value)) + value - return self.ccm(self.norm2(value)) + value - - class SAFMN(nn.Module): - def __init__(self): - super().__init__() - self.to_feat = nn.Conv2d(3, 128, 3, 1, 1) - self.feats = nn.Sequential(*(AttBlock(128) for _ in range(16))) - self.to_img = nn.Sequential( - nn.Conv2d(128, 3 * 4**2, 3, 1, 1), - nn.PixelShuffle(4), - ) - - def forward(self, value): - value = self.to_feat(value) - return self.to_img(self.feats(value) + value) - - model = SAFMN() - model.load_state_dict( - _load_weights(torch, checkpoint_path, "params_ema"), - strict=True, - ) - sample = torch.rand(1, 3, 256, 256) - _export(torch, model, sample, destination) - _validate_export(onnx, onnxruntime, torch, model, sample, destination) - - def export_model(exporter, checkpoint_path, destination): exporters = { "realesrgan_wdn_x4": export_realesrgan_wdn_x4, - "safmn_real_x4": export_safmn_real_x4, } try: export = exporters[exporter] diff --git a/tools/models/hat_export.py b/tools/models/hat_export.py new file mode 100644 index 0000000..af2cc49 --- /dev/null +++ b/tools/models/hat_export.py @@ -0,0 +1,95 @@ +# /// script +# requires-python = ">=3.11,<3.12" +# dependencies = ["torch==2.13.0", "onnx==1.19.0", "numpy==1.26.4", "einops==0.8.2", "gdown==6.2.0"] +# [tool.uv.sources] +# torch = { index = "pytorch" } +# [[tool.uv.index]] +# name = "pytorch" +# url = "https://download.pytorch.org/whl/cpu" +# explicit = true +# /// +"""Export the official HAT Sharper checkpoint with FP16 weight storage.""" + +import argparse +import hashlib +import json +from pathlib import Path +from urllib.request import urlopen + +import gdown +import onnx +import torch + +from onnx_weights import compress_weights + + +REVISION = "1638a9a822581657811867bf670717f8371fc3e5" +CHECKPOINT_SHA256 = "5800b67136006eb8cab3b4ed7c8d73b6a195bb18e6cc709b674f9aa069c00271" +SOURCE_URL = f"https://raw.githubusercontent.com/XPixelGroup/HAT/{REVISION}" + + +def load_model(checkpoint, cache): + architecture = cache / "hat_arch.py" + if not architecture.exists(): + architecture.write_bytes(urlopen(SOURCE_URL + "/hat/archs/hat_arch.py", timeout=60).read()) + if hashlib.sha256(architecture.read_bytes()).hexdigest() != "81d8cecf491975246c9ebb20480898c656f59de020f38b05daa68741e426117f": + raise RuntimeError("HAT architecture SHA-256 mismatch") + # Export only: remove BasicSR's training registry and use the equivalent + # PyTorch initialization utilities. The official forward remains unchanged. + source = architecture.read_text(encoding="utf-8") + source = source.replace("from basicsr.utils.registry import ARCH_REGISTRY", "") + source = source.replace("@ARCH_REGISTRY.register()", "") + source = source.replace( + "from basicsr.archs.arch_util import to_2tuple, trunc_normal_", + "from torch.nn.modules.utils import _pair as to_2tuple\nfrom torch.nn.init import trunc_normal_", + ) + namespace = {"__name__": "hat_arch"} + exec(compile(source, str(architecture), "exec"), namespace) + model = namespace["HAT"]( + upscale=4, in_chans=3, img_size=64, window_size=16, + compress_ratio=3, squeeze_factor=30, conv_scale=0.01, + overlap_ratio=0.5, img_range=1., depths=[6] * 6, + embed_dim=180, num_heads=[6] * 6, mlp_ratio=2, + upsampler="pixelshuffle", resi_connection="1conv", + ).eval() + with checkpoint.open("rb") as stream: + if hashlib.file_digest(stream, "sha256").hexdigest() != CHECKPOINT_SHA256: + raise RuntimeError("HAT Sharper checkpoint SHA-256 mismatch") + model.load_state_dict(torch.load(checkpoint, map_location="cpu", weights_only=True)["params_ema"], strict=True) + return model + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--destination", type=Path, required=True) + parser.add_argument("--checkpoint", type=Path) + args = parser.parse_args() + cache = Path(__file__).resolve().parents[2] / ".model-cache" / "hat-sharper" + cache.mkdir(parents=True, exist_ok=True) + checkpoint = args.checkpoint or cache / "Real_HAT_GAN_sharper.pth" + if not checkpoint.exists(): + gdown.download(id="1EioFq5-mKmv1uqta_Byd9cgXp9SU3zjj", output=str(checkpoint)) + torch.set_num_threads(8) + torch.manual_seed(0) + model = load_model(checkpoint, cache) + args.destination.mkdir(parents=True, exist_ok=True) + output = args.destination / "model.onnx" + with torch.inference_mode(): + torch.onnx.export(model, torch.zeros(1, 3, 256, 256), str(output), + opset_version=17, input_names=["input"], output_names=["output"], + do_constant_folding=False, dynamo=False) + graph = onnx.load(output) + compress_weights(graph) + onnx.checker.check_model(graph) + onnx.save(graph, output) + (args.destination / "LICENSE.txt").write_bytes(urlopen(SOURCE_URL + "/LICENSE", timeout=60).read()) + if sum(path.stat().st_size for path in args.destination.iterdir()) >= 60_000_000: + raise RuntimeError("HAT asset exceeds 60 MB") + with output.open("rb") as stream: + digest = hashlib.file_digest(stream, "sha256").hexdigest() + print(json.dumps({"bytes": output.stat().st_size, "sha256": digest, + "revision": REVISION, "checkpoint_sha256": CHECKPOINT_SHA256}), flush=True) + + +if __name__ == "__main__": + main() diff --git a/tools/models/licenses/pisa-sr/LICENSE-APACHE.txt b/tools/models/licenses/pisa-sr/LICENSE-APACHE.txt deleted file mode 100644 index d645695..0000000 --- a/tools/models/licenses/pisa-sr/LICENSE-APACHE.txt +++ /dev/null @@ -1,202 +0,0 @@ - - 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 - Contributor provides its Contributions) 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 the Work and assume any - risks associated with Your exercise of permissions under this License. - - 8. Limitation of Liability. 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 any Contributor be - liable to You for damages, including any direct, indirect, special, - incidental, or consequential damages of any character arising as a - result of this License or out of the use or inability to use the - Work (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 such Contributor - has been advised of the possibility of such damages. - - 9. Accepting Warranty or Additional Liability. While redistributing - the Work or Derivative Works thereof, You may choose to offer, - and charge a fee for, acceptance of support, warranty, indemnity, - or other liability obligations and/or rights consistent with this - License. However, in accepting such obligations, You may act only - on Your own behalf and on Your sole responsibility, not on behalf - of any other Contributor, and only if You agree to indemnify, - defend, and hold each Contributor harmless for any liability - incurred by, or claims asserted against, such Contributor by reason - of your accepting any such warranty or additional liability. - - END OF TERMS AND CONDITIONS - - APPENDIX: How to apply the Apache License to your work. - - To apply the Apache License to your work, attach the following - boilerplate notice, with the fields enclosed by brackets "[]" - replaced with your own identifying information. (Don't include - the brackets!) The text should be enclosed in the appropriate - comment syntax for the file format. We also recommend that a - file or class name and description of purpose be included on the - same "printed page" as the copyright notice for easier - identification within third-party archives. - - Copyright [yyyy] [name of copyright owner] - - Licensed under the Apache License, Version 2.0 (the "License"); - you may not use this file except in compliance with the License. - You may obtain a copy of the License at - - http://www.apache.org/licenses/LICENSE-2.0 - - Unless required by applicable law or agreed to in writing, software - distributed under the License is distributed on an "AS IS" BASIS, - WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - See the License for the specific language governing permissions and - limitations under the License. diff --git a/tools/models/licenses/pisa-sr/LICENSE-MODEL.txt b/tools/models/licenses/pisa-sr/LICENSE-MODEL.txt deleted file mode 100644 index 9684533..0000000 --- a/tools/models/licenses/pisa-sr/LICENSE-MODEL.txt +++ /dev/null @@ -1,84 +0,0 @@ -Copyright (c) 2022 Stability AI and contributors - -CreativeML Open RAIL++-M License -dated November 24, 2022 - -Section I: PREAMBLE - -Multimodal generative models are being widely adopted and used, and have the potential to transform the way artists, among other individuals, conceive and benefit from AI or ML technologies as a tool for content creation. - -Notwithstanding the current and potential benefits that these artifacts can bring to society at large, there are also concerns about potential misuses of them, either due to their technical limitations or ethical considerations. - -In short, this license strives for both the open and responsible downstream use of the accompanying model. When it comes to the open character, we took inspiration from open source permissive licenses regarding the grant of IP rights. Referring to the downstream responsible use, we added use-based restrictions not permitting the use of the Model in very specific scenarios, in order for the licensor to be able to enforce the license in case potential misuses of the Model may occur. At the same time, we strive to promote open and responsible research on generative models for art and content generation. - -Even though downstream derivative versions of the model could be released under different licensing terms, the latter will always have to include - at minimum - the same use-based restrictions as the ones in the original license (this license). We believe in the intersection between open and responsible AI development; thus, this License aims to strike a balance between both in order to enable responsible open-science in the field of AI. - -This License governs the use of the model (and its derivatives) and is informed by the model card associated with the model. - -NOW THEREFORE, You and Licensor agree as follows: - -1. Definitions - -- "License" means the terms and conditions for use, reproduction, and Distribution as defined in this document. -- "Data" means a collection of information and/or content extracted from the dataset used with the Model, including to train, pretrain, or otherwise evaluate the Model. The Data is not licensed under this License. -- "Output" means the results of operating a Model as embodied in informational content resulting therefrom. -- "Model" means any accompanying machine-learning based assemblies (including checkpoints), consisting of learnt weights, parameters (including optimizer states), corresponding to the model architecture as embodied in the Complementary Material, that have been trained or tuned, in whole or in part on the Data, using the Complementary Material. -- "Derivatives of the Model" means all modifications to the Model, works based on the Model, or any other model which is created or initialized by transfer of patterns of the weights, parameters, activations or output of the Model, to the other model, in order to cause the other model to perform similarly to the Model, including - but not limited to - distillation methods entailing the use of intermediate data representations or methods based on the generation of synthetic data by the Model for training the other model. -- "Complementary Material" means the accompanying source code and scripts used to define, run, load, benchmark or evaluate the Model, and used to prepare data for training or evaluation, if any. This includes any accompanying documentation, tutorials, examples, etc, if any. -- "Distribution" means any transmission, reproduction, publication or other sharing of the Model or Derivatives of the Model to a third party, including providing the Model as a hosted service made available by electronic or other remote means - e.g. API-based or web access. -- "Licensor" means the copyright owner or entity authorized by the copyright owner that is granting the License, including the persons or entities that may have rights in the Model and/or distributing the Model. -- "You" (or "Your") means an individual or Legal Entity exercising permissions granted by this License and/or making use of the Model for whichever purpose and in any field of use, including usage of the Model in an end-use application - e.g. chatbot, translator, image generator. -- "Third Parties" means individuals or legal entities that are not under common control with Licensor or You. -- "Contribution" means any work of authorship, including the original version of the Model and any modifications or additions to that Model or Derivatives of the Model thereof, that is intentionally submitted to Licensor for inclusion in the Model 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 Model, but excluding communication that is conspicuously marked or otherwise designated in writing by the copyright owner as "Not a Contribution." -- "Contributor" means Licensor and any individual or Legal Entity on behalf of whom a Contribution has been received by Licensor and subsequently incorporated within the Model. - -Section II: INTELLECTUAL PROPERTY RIGHTS - -Both copyright and patent grants apply to the Model, Derivatives of the Model and Complementary Material. The Model and Derivatives of the Model are subject to additional terms as described in Section III. - -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, publicly display, publicly perform, sublicense, and distribute the Complementary Material, the Model, and Derivatives of the Model. -3. Grant of Patent License. Subject to the terms and conditions of this License and where and as applicable, each Contributor hereby grants to You a perpetual, worldwide, non-exclusive, no-charge, royalty-free, irrevocable (except as stated in this paragraph) patent license to make, have made, use, offer to sell, sell, import, and otherwise transfer the Model and the Complementary Material, 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 Model 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 Model and/or Complementary Material or a Contribution incorporated within the Model and/or Complementary Material constitutes direct or contributory patent infringement, then any patent licenses granted to You under this License for the Model and/or Work shall terminate as of the date such litigation is asserted or filed. - -Section III: CONDITIONS OF USAGE, DISTRIBUTION AND REDISTRIBUTION - -4. Distribution and Redistribution. You may host for Third Party remote access purposes (e.g. software-as-a-service), reproduce and distribute copies of the Model or Derivatives of the Model thereof in any medium, with or without modifications, provided that You meet the following conditions: -Use-based restrictions as referenced in paragraph 5 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 the Model or Derivatives of the Model, and You shall give notice to subsequent users You Distribute to, that the Model or Derivatives of the Model are subject to paragraph 5. This provision does not apply to the use of Complementary Material. -You must give any Third Party recipients of the Model or Derivatives of the Model a copy of this License; -You must cause any modified files to carry prominent notices stating that You changed the files; -You must retain all copyright, patent, trademark, and attribution notices excluding those notices that do not pertain to any part of the Model, Derivatives of the Model. -You may add Your own copyright statement to Your modifications and may provide additional or different license terms and conditions - respecting paragraph 4.a. - for use, reproduction, or Distribution of Your modifications, or for any such Derivatives of the Model as a whole, provided Your use, reproduction, and Distribution of the Model otherwise complies with the conditions stated in this License. -5. Use-based restrictions. The restrictions set forth in Attachment A are considered Use-based restrictions. Therefore You cannot use the Model and the Derivatives of the Model for the specified restricted uses. You may use the Model subject to this License, including only for lawful purposes and in accordance with the License. Use may include creating any content with, finetuning, updating, running, training, evaluating and/or reparametrizing the Model. You shall require all of Your users who use the Model or a Derivative of the Model to comply with the terms of this paragraph (paragraph 5). -6. The Output You Generate. Except as set forth herein, Licensor claims no rights in the Output You generate using the Model. You are accountable for the Output you generate and its subsequent uses. No use of the output can contravene any provision as stated in the License. - -Section IV: OTHER PROVISIONS - -7. Updates and Runtime Restrictions. To the maximum extent permitted by law, Licensor reserves the right to restrict (remotely or otherwise) usage of the Model in violation of this License. -8. Trademarks and related. Nothing in this License permits You to make use of Licensors’ 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 Licensors. -9. Disclaimer of Warranty. Unless required by applicable law or agreed to in writing, Licensor provides the Model and the Complementary Material (and each Contributor provides its Contributions) 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 the Model, Derivatives of the Model, and the Complementary Material and assume any risks associated with Your exercise of permissions under this License. -10. Limitation of Liability. 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 any Contributor be liable to You for damages, including any direct, indirect, special, incidental, or consequential damages of any character arising as a result of this License or out of the use or inability to use the Model and the Complementary Material (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 such Contributor has been advised of the possibility of such damages. -11. Accepting Warranty or Additional Liability. While redistributing the Model, Derivatives of the Model and the Complementary Material thereof, You may choose to offer, and charge a fee for, acceptance of support, warranty, indemnity, or other liability obligations and/or rights consistent with this License. However, in accepting such obligations, You may act only on Your own behalf and on Your sole responsibility, not on behalf of any other Contributor, and only if You agree to indemnify, defend, and hold each Contributor harmless for any liability incurred by, or claims asserted against, such Contributor by reason of your accepting any such warranty or additional liability. -12. If any provision of this License 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 - -You agree not to use the Model or Derivatives of the Model: - -- In any way that violates any applicable national, federal, state, local or international law or regulation; -- For the purpose of exploiting, harming or attempting to exploit or harm minors in any way; -- To generate or disseminate verifiably false information and/or content with the purpose of harming others; -- To generate or disseminate personal identifiable information that can be used to harm an individual; -- To defame, disparage or otherwise harass others; -- For fully automated decision making that adversely impacts an individual’s legal rights or otherwise creates or modifies a binding, enforceable obligation; -- 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; -- 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; -- For any use intended to or which has the effect of discriminating against individuals or groups based on legally protected characteristics or categories; -- To provide medical advice and medical results interpretation; -- 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). - diff --git a/tools/models/licenses/pisa-sr/NOTICE.txt b/tools/models/licenses/pisa-sr/NOTICE.txt deleted file mode 100644 index 397e76b..0000000 --- a/tools/models/licenses/pisa-sr/NOTICE.txt +++ /dev/null @@ -1,21 +0,0 @@ -PiSA-SR ONNX models for AnyImage - -PiSA-SR: Lingchen Sun and contributors, https://github.com/csslc/PiSA-SR -Source commit: 9c76e12e261d1980148738a0cff41e5e27d18c9f -PiSA-SR is released under Apache License 2.0; see LICENSE-APACHE.txt. - -Stable Diffusion 2.1: Copyright (c) 2022 Stability AI and contributors. -Base weights: sd2-community/stable-diffusion-2-1-base -Revision: 4e63672c03103b6c636b8fb4119ba982469b2955 -Base weights and derivatives are subject to CreativeML Open RAIL++-M, -including paragraph 5 and Attachment A use restrictions. All users must -comply with these terms; see the accompanying LICENSE-MODEL.txt. -License copy: sd2-community/stable-diffusion-2, LICENSE-MODEL. - -AnyImage modifications: merged pixel and semantic LoRA adapters, FP16 -ONNX conversion, fixed empty prompt and timestep, VAE encoder/decoder -separation. Runtime combines pixel/full predictions at semantic strength 0.7. -noise.npy contains a fixed seed 42 latent noise field. - -ONNX architecture derives from Hugging Face Diffusers and PiSA-SR. -Copyright 2023 The HuggingFace Team. All rights reserved. diff --git a/tools/models/onnx_weights.py b/tools/models/onnx_weights.py new file mode 100644 index 0000000..2be3e43 --- /dev/null +++ b/tools/models/onnx_weights.py @@ -0,0 +1,49 @@ +"""Compress stored ONNX weights without changing graph arithmetic.""" + +import numpy as np +from onnx import TensorProto, helper, numpy_helper + + +def compress_weights(model): + casts = [] + for initializer in model.graph.initializer: + if initializer.data_type != TensorProto.FLOAT: + continue + name = initializer.name + values = numpy_helper.to_array(initializer) + if not np.isfinite(values).all() or np.max(np.abs(values), initial=0) > np.finfo(np.float16).max: + continue + array = values.astype(np.float16) + stored_name = name + "__fp16_storage" + initializer.CopyFrom(numpy_helper.from_array(array, stored_name)) + casts.append(helper.make_node( + "Cast", [stored_name], [name], to=TensorProto.FLOAT, + name=stored_name + "_restore", + )) + nodes = list(model.graph.node) + del model.graph.node[:] + model.graph.node.extend(casts + nodes) + + +def compress_exact_constants(model): + """Compress only constant tensors that round-trip without any value change.""" + nodes = [] + for node in model.graph.node: + nodes.append(node) + if node.op_type != "Constant": + continue + attribute = next((item for item in node.attribute if item.name == "value"), None) + if attribute is None or attribute.t.data_type != TensorProto.FLOAT: + continue + values = numpy_helper.to_array(attribute.t) + if not np.isfinite(values).all() or np.max(np.abs(values), initial=0) > np.finfo(np.float16).max: + continue + stored = values.astype(np.float16) + if values.nbytes < 64 or not np.array_equal(values, stored.astype(np.float32)): + continue + name = node.output[0] + node.output[0] = name + "__fp16_constant" + attribute.t.CopyFrom(numpy_helper.from_array(stored)) + nodes.append(helper.make_node("Cast", [node.output[0]], [name], to=TensorProto.FLOAT)) + del model.graph.node[:] + model.graph.node.extend(nodes) diff --git a/tools/models/pisa_sr_export.py b/tools/models/pisa_sr_export.py deleted file mode 100644 index 15531d4..0000000 --- a/tools/models/pisa_sr_export.py +++ /dev/null @@ -1,198 +0,0 @@ -# /// script -# requires-python = "==3.11.*" -# dependencies = [ -# "torch==2.7.1", "torchvision==0.22.1", "diffusers==0.25.0", -# "transformers==4.28.1", "peft==0.9.0", "accelerate==0.27.2", -# "huggingface-hub==0.25.2", "numpy==1.26.4", "onnx==1.19.0", -# ] -# [tool.uv.sources] -# torch = { index = "pytorch" } -# torchvision = { index = "pytorch" } -# [[tool.uv.index]] -# name = "pytorch" -# url = "https://download.pytorch.org/whl/cu128" -# explicit = true -# /// -"""Export PiSA-SR components from pinned upstream weights on CUDA.""" - -import argparse -import hashlib -import json -import sys -from pathlib import Path - -SOURCE_COMMIT = "9c76e12e261d1980148738a0cff41e5e27d18c9f" -BASE_REPOSITORY = "sd2-community/stable-diffusion-2-1-base" -BASE_REVISION = "4e63672c03103b6c636b8fb4119ba982469b2955" -CHECKPOINT_SHA256 = "13502f7bde5e4e585f84f24c3d904d9c848e068bf89f984046fe87da81a391f6" -CHECKPOINT_URL = "https://drive.google.com/uc?export=download&id=1PWFM5zUd6lNMOyKX64anvUXhKEOCrf8L" - - -def finalize_assets(destination): - import shutil - import onnx - - for path in destination.glob("*.onnx"): - model = onnx.load(path) - onnx.helper.set_model_props(model, { - "license": "Apache-2.0; CreativeML Open RAIL++-M (LICENSE-MODEL.txt)", - "modifications": "AnyImage: merged PiSA-SR LoRA; FP16; fixed empty prompt and timestep; split VAE. See NOTICE.txt.", - "source_commit": SOURCE_COMMIT, - }) - onnx.save(model, path) - for path in (Path(__file__).parent / "licenses/pisa-sr").glob("*.txt"): - shutil.copyfile(path, destination / path.name) - - -def prepare_sources(cache): - import io - import zipfile - from urllib.request import urlopen - from huggingface_hub import snapshot_download - - source = cache / f"PiSA-SR-{SOURCE_COMMIT}" - if not source.exists(): - url = f"https://codeload.github.com/csslc/PiSA-SR/zip/{SOURCE_COMMIT}" - with urlopen(url, timeout=60) as response: - archive = zipfile.ZipFile(io.BytesIO(response.read())) - for member in archive.infolist(): - destination = (cache / member.filename).resolve() - if not destination.is_relative_to(cache.resolve()): - raise ValueError("Invalid upstream archive path") - archive.extractall(cache) - checkpoint = cache / "pisa_sr.pkl" - if not checkpoint.exists(): - with urlopen(CHECKPOINT_URL, timeout=60) as response: - checkpoint.write_bytes(response.read()) - if hashlib.sha256(checkpoint.read_bytes()).hexdigest() != CHECKPOINT_SHA256: - raise ValueError("PiSA-SR checkpoint checksum mismatch") - base = snapshot_download( - BASE_REPOSITORY, revision=BASE_REVISION, local_dir=cache / "base-model", - allow_patterns=["*/*.json", "tokenizer/*", "*/*.fp16.safetensors"], - ) - return source, Path(base), checkpoint - - -def export_components(source, base, checkpoint, destination): - import gc - import numpy as np - import torch - from peft import LoraConfig - from diffusers.utils import set_weights_and_activate_adapters - from transformers import AutoTokenizer, CLIPTextModel - - if hashlib.sha256(checkpoint.read_bytes()).hexdigest() != CHECKPOINT_SHA256: - raise ValueError("PiSA-SR checkpoint checksum mismatch") - sys.path.insert(0, str(source)) - from src.models.autoencoder_kl import AutoencoderKL - from src.models.unet_2d_condition import UNet2DConditionModel - - class Residual(torch.nn.Module): - def __init__(self, model, prompt): - super().__init__() - self.model = model - self.register_buffer("prompt", prompt) - self.register_buffer("timestep", torch.tensor([1], dtype=torch.long)) - - def forward(self, latent): - return self.model(latent, self.timestep, encoder_hidden_states=self.prompt).sample - - class Encoder(torch.nn.Module): - def __init__(self, model): - super().__init__() - self.encoder = model.encoder - self.quant_conv = model.quant_conv - - def forward(self, image): - return self.quant_conv(self.encoder(image)) - - class Decoder(torch.nn.Module): - def __init__(self, model): - super().__init__() - self.decoder = model.decoder - self.post_quant_conv = model.post_quant_conv - - def forward(self, latent): - return self.decoder(self.post_quant_conv(latent / 0.18215)).clamp(-1, 1) - - def export(wrapper, component, tensor): - wrapper.eval().cuda() - with torch.inference_mode(): - torch.onnx.export( - wrapper, (tensor.cuda(),), str(destination / f"{component}.onnx"), - opset_version=17, input_names=["input"], output_names=["output"], - do_constant_folding=True, dynamo=False, - ) - wrapper.cpu() - torch.cuda.empty_cache() - - destination.mkdir(parents=True, exist_ok=True) - torch.set_num_threads(8) - torch.backends.cuda.matmul.allow_tf32 = False - tokenizer = AutoTokenizer.from_pretrained(base, subfolder="tokenizer") - text = CLIPTextModel.from_pretrained( - base, subfolder="text_encoder", variant="fp16", torch_dtype=torch.float16, - ).eval().cuda() - with torch.inference_mode(): - tokens = tokenizer("", max_length=tokenizer.model_max_length, - padding="max_length", truncation=True, return_tensors="pt") - prompt = text(tokens.input_ids.cuda())[0].cpu() - del text - torch.cuda.empty_cache() - state = torch.load(checkpoint, map_location="cpu", weights_only=True) - latent = torch.zeros((1, 4, 64, 64), dtype=torch.float16) - unet = UNet2DConditionModel.from_pretrained(base, subfolder="unet", variant="fp16") - for strength, component in (("pix", "unet-pixel"), ("sem", "unet-full")): - adapters = [] - for part in ("encoder", "decoder", "others"): - key = f"{part}_modules" - config = LoraConfig( - r=state[f"lora_rank_unet_{strength}"], init_lora_weights="gaussian", - target_modules=state[f"unet_lora_{key}_{strength}"], - ) - name = f"default_{part}_{strength}" - unet.add_adapter(config, adapter_name=name) - adapters.append(name) - for name, parameter in unet.named_parameters(): - if "lora" in name: - parameter.data.copy_(state["state_dict_unet"][name]) - set_weights_and_activate_adapters(unet, adapters, [1.0]*3) - unet.merge_and_unload() - # Export a half-precision copy; retain FP32 weights for the semantic merge. - import copy - wrapper = Residual(copy.deepcopy(unet).half(), prompt) - export(wrapper, component, latent) - del wrapper - gc.collect() - del unet, state - gc.collect() - vae = AutoencoderKL.from_pretrained(base, subfolder="vae", variant="fp16", torch_dtype=torch.float16) - export(Encoder(vae), "encoder", torch.zeros((1,3,512,512), dtype=torch.float16)) - export(Decoder(vae), "decoder", latent) - torch.manual_seed(42) - np.save(destination / "noise.npy", torch.randn((1,4,64,64), device="cuda", dtype=torch.float16).cpu().numpy()) - finalize_assets(destination) - files = [] - for path in sorted(destination.iterdir()): - if path.suffix in (".onnx", ".npy", ".txt"): - with path.open("rb") as stream: - checksum = hashlib.file_digest(stream, "sha256").hexdigest() - files.append((path.name, path.stat().st_size, checksum)) - (destination / "provenance.json").write_text(json.dumps({ - "source_commit": SOURCE_COMMIT, "base_repository": BASE_REPOSITORY, - "base_revision": BASE_REVISION, "checkpoint_sha256": CHECKPOINT_SHA256, - "licenses": ["Apache-2.0", "CreativeML Open RAIL++-M"], "files": files, - }, indent=2)) - - -def main(): - parser = argparse.ArgumentParser() - parser.add_argument("--cache", type=Path, default=Path(".model-cache/pisa-sr")) - parser.add_argument("--destination", type=Path, default=Path("models/pisa-sr")) - args = parser.parse_args() - args.cache.mkdir(parents=True, exist_ok=True) - export_components(*prepare_sources(args.cache), args.destination) - - -if __name__ == "__main__": - main() diff --git a/tools/nodes/common/reference_plane.py b/tools/nodes/common/reference_plane.py new file mode 100644 index 0000000..7c7a3fc --- /dev/null +++ b/tools/nodes/common/reference_plane.py @@ -0,0 +1,31 @@ +"""Build the direction-adjusted reference depth field.""" +from .nodes import math_node + + +def reference_plane_depth(group, direction, reference_depth): + nodes = group.nodes + links = group.links + normal = nodes.new("ShaderNodeVectorMath") + normal.operation = "NORMALIZE" + links.new(direction, normal.inputs[0]) + normal_components = nodes.new("ShaderNodeSeparateXYZ") + links.new(normal.outputs["Vector"], normal_components.inputs["Vector"]) + position = nodes.new("GeometryNodeInputPosition") + position_components = nodes.new("ShaderNodeSeparateXYZ") + links.new(position.outputs["Position"], position_components.inputs["Vector"]) + x_term = math_node(nodes, "MULTIPLY") + links.new(position_components.outputs["X"], x_term.inputs[0]) + links.new(normal_components.outputs["X"], x_term.inputs[1]) + y_term = math_node(nodes, "MULTIPLY") + links.new(position_components.outputs["Y"], y_term.inputs[0]) + links.new(normal_components.outputs["Y"], y_term.inputs[1]) + point_dot = math_node(nodes, "ADD") + links.new(x_term.outputs[0], point_dot.inputs[0]) + links.new(y_term.outputs[0], point_dot.inputs[1]) + divide = math_node(nodes, "DIVIDE") + links.new(point_dot.outputs[0], divide.inputs[0]) + links.new(normal_components.outputs["Z"], divide.inputs[1]) + plane = math_node(nodes, "ADD") + links.new(reference_depth, plane.inputs[0]) + links.new(divide.outputs[0], plane.inputs[1]) + return plane.outputs[0] diff --git a/tools/nodes/groups/image_depth_balloon.py b/tools/nodes/groups/image_depth_balloon.py index 758c934..3868b59 100644 --- a/tools/nodes/groups/image_depth_balloon.py +++ b/tools/nodes/groups/image_depth_balloon.py @@ -7,6 +7,7 @@ from ..common.nodes import ( menu_switch, read_float_attribute, store_float_attribute, prepare_node_group, ) from ..common.smoothing import expand_smooth +from ..common.reference_plane import reference_plane_depth from ..common.cutout import _extrude_layers, _island_cleanup, _position_field, _profile_field, _remove_zero_thickness, _set_layer_z, _shade_output @@ -29,35 +30,6 @@ def _set_geometry_z(group, geometry, z_value): return set_position.outputs["Geometry"] -def _plane_depth(group, direction, reference_depth): - nodes = group.nodes - links = group.links - normal = nodes.new("ShaderNodeVectorMath") - normal.operation = "NORMALIZE" - links.new(direction, normal.inputs[0]) - normal_components = nodes.new("ShaderNodeSeparateXYZ") - links.new(normal.outputs["Vector"], normal_components.inputs["Vector"]) - position = nodes.new("GeometryNodeInputPosition") - position_components = nodes.new("ShaderNodeSeparateXYZ") - links.new(position.outputs["Position"], position_components.inputs["Vector"]) - x_term = math_node(nodes, "MULTIPLY") - links.new(position_components.outputs["X"], x_term.inputs[0]) - links.new(normal_components.outputs["X"], x_term.inputs[1]) - y_term = math_node(nodes, "MULTIPLY") - links.new(position_components.outputs["Y"], y_term.inputs[0]) - links.new(normal_components.outputs["Y"], y_term.inputs[1]) - point_dot = math_node(nodes, "ADD") - links.new(x_term.outputs[0], point_dot.inputs[0]) - links.new(y_term.outputs[0], point_dot.inputs[1]) - divide = math_node(nodes, "DIVIDE") - links.new(point_dot.outputs[0], divide.inputs[0]) - links.new(normal_components.outputs["Z"], divide.inputs[1]) - plane = math_node(nodes, "ADD") - links.new(reference_depth, plane.inputs[0]) - links.new(divide.outputs[0], plane.inputs[1]) - return plane.outputs[0] - - def _build_depth_balloon( group, geometry, @@ -78,7 +50,7 @@ def _build_depth_balloon( metric_depth = math_node(nodes, "MULTIPLY") links.new(separate_camera.outputs["Z"], metric_depth.inputs[0]) links.new(uniform_scale, metric_depth.inputs[1]) - plane_depth = _plane_depth( + plane_depth = reference_plane_depth( group, depth_direction, depth_reference, diff --git a/tools/nodes/groups/image_relief_plane.py b/tools/nodes/groups/image_relief_plane.py index ed70960..01350fa 100644 --- a/tools/nodes/groups/image_relief_plane.py +++ b/tools/nodes/groups/image_relief_plane.py @@ -3,7 +3,9 @@ import bpy from ..common.nodes import group_input as _group_input, interface_socket as _interface_socket -from ..common.nodes import prepare_node_group +from ..common.nodes import prepare_node_group, float_input, vector_input +from ..common.reference_plane import reference_plane_depth +from ..common.smoothing import expand_smooth RELIEF_PLANE_NODE_GROUP_NAME = "O Image Relief Plane" @@ -38,6 +40,8 @@ def build_image_relief_plane_group(image_plane_group): depth_scale = _interface_socket(group, "Depth Scale", "INPUT", "NodeSocketFloat") depth_scale.default_value = 0.5 depth_scale.min_value = 0.0 + smooth = _interface_socket(group, "Smooth", "INPUT", "NodeSocketInt") + smooth.default_value, smooth.min_value, smooth.max_value = 0, 0, 20 options = group.interface.new_panel(name="Options", default_closed=True) base_plane_depth = _interface_socket( group, @@ -49,6 +53,8 @@ def build_image_relief_plane_group(image_plane_group): base_plane_depth.default_value = 1.0 base_plane_depth.min_value = 0.0 base_plane_depth.subtype = "DISTANCE" + float_input(group, "Smooth Weight", 1.0, maximum=1, subtype="FACTOR", parent=options) + vector_input(group, "Depth Direction", (0, 0, 1), subtype="DIRECTION", parent=options) data = group.interface.new_panel(name="Data", default_closed=True) uniform_scale = _interface_socket( group, @@ -78,7 +84,7 @@ def build_image_relief_plane_group(image_plane_group): ) shape_input = _group_input( nodes, - {"Thickness", "Depth Scale", "Reference Depth"}, + {"Thickness", "Depth Scale"}, ) image_input = _group_input( nodes, @@ -114,7 +120,11 @@ def build_image_relief_plane_group(image_plane_group): links.new(depth_components.outputs["Z"], metric_depth.inputs[0]) links.new(image_input.outputs["Uniform Scale"], metric_depth.inputs[1]) relative_depth = _math(nodes, "SUBTRACT") - links.new(shape_input.outputs["Reference Depth"], relative_depth.inputs[0]) + reference_input = _group_input(nodes, {"Depth Direction", "Reference Depth"}) + plane_depth = reference_plane_depth( + group, reference_input.outputs["Depth Direction"], reference_input.outputs["Reference Depth"], + ) + links.new(plane_depth, relative_depth.inputs[0]) links.new(metric_depth.outputs[0], relative_depth.inputs[1]) full_offset = _math(nodes, "MULTIPLY") links.new(relative_depth.outputs[0], full_offset.inputs[0]) @@ -167,7 +177,10 @@ def build_image_relief_plane_group(image_plane_group): shade_smooth = nodes.new("GeometryNodeSetShadeSmooth") shade_smooth.domain = "FACE" shade_smooth.inputs["Shade Smooth"].default_value = True - links.new(set_position.outputs["Geometry"], shade_smooth.inputs["Geometry"]) + geometry = expand_smooth( + group, set_position.outputs["Geometry"], iterations_name="Smooth", weight_name="Smooth Weight", + ) + links.new(geometry, shade_smooth.inputs["Geometry"]) links.new(shade_smooth.outputs["Geometry"], group_output.inputs["Geometry"]) group.color_tag = "GEOMETRY" group.description = ( diff --git a/uv.lock b/uv.lock index a6a06ac..d1080f1 100644 --- a/uv.lock +++ b/uv.lock @@ -41,7 +41,7 @@ dev = [ [package.metadata] requires-dist = [ - { name = "blendjob", specifier = "==0.1.16" }, + { name = "blendjob", path = "src/anyimage/wheels/blendjob-0.1.17-py3-none-any.whl" }, { name = "numpy", specifier = ">=1.26,<2.0" }, { name = "scipy", specifier = "==1.15.3" }, ] @@ -114,13 +114,19 @@ wheels = [ [[package]] name = "blendjob" -version = "0.1.16" -source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/33/97/bcbf9be2effd238a3caaa647fc202d7e29ff7370c1ca22d6a78e48fa4134/blendjob-0.1.16.tar.gz", hash = "sha256:16ead42a046f9c8eb91bc4c3d4803af366f2b56b63bd16f302190ec60ef23a99", size = 21571, upload-time = "2026-09-02T02:50:03.715Z" } +version = "0.1.17" +source = { path = "src/anyimage/wheels/blendjob-0.1.17-py3-none-any.whl" } wheels = [ - { url = "https://files.pythonhosted.org/packages/db/ac/142988a33d0a2df6335a0e08f5e9a0a9215996810a77f9724325f9174301/blendjob-0.1.16-py3-none-any.whl", hash = "sha256:39fb736f0c2b77321a672b2b023e4b1d0d49c2d567bdd8f6095b6efe5080f982", size = 26931, upload-time = "2026-09-02T02:50:02.777Z" }, + { filename = "blendjob-0.1.17-py3-none-any.whl", hash = "sha256:9173320611ae77ce8852c6f9e96ddcec1b455ac1e2c0d1f886041e5a33e0cc86" }, ] +[package.metadata] +requires-dist = [ + { name = "fastapi", marker = "extra == 'server'", specifier = "==0.139.2" }, + { name = "uvicorn", marker = "extra == 'server'", specifier = "==0.51.0" }, +] +provides-extras = ["server"] + [[package]] name = "bpy" version = "4.5.3"