From 0cf04a13c6c47ee17896f7303d51290c75fb35ad Mon Sep 17 00:00:00 2001 From: nan Date: Sun, 23 Aug 2026 20:34:35 +0800 Subject: [PATCH] feat: Implement upscale functionality with ONNX models - Added upscale operator to handle image upscaling by 4x using AI models. - Introduced new jobs for upscale processing in the server, including debug and standard upscale jobs. - Created a new ONNX model handler for upscale operations, supporting tiled inference and resource management. - Updated tests to cover new upscale functionality, including input validation and output verification. - Added documentation for the upscale feature, detailing usage and model options. --- AGENTS.md | 1 + README.md | 36 +- docs/depth-plane.md | 4 +- docs/development.md | 12 +- docs/faq.md | 12 +- docs/features.md | 16 +- docs/image-tools-technical.md | 8 +- docs/image-tools.md | 2 +- docs/index.md | 5 +- docs/installation.md | 6 +- docs/models-and-storage.md | 19 +- docs/outputs.md | 9 +- docs/remove-background.md | 4 +- docs/server.md | 18 +- docs/surface.md | 2 +- docs/upscale.md | 23 + docs/workflows.md | 8 +- mkdocs.yml | 1 + pack.py | 2 + pyproject.toml | 2 +- src/anyimage/__init__.py | 13 +- src/anyimage/blender_manifest.toml | 2 +- src/anyimage/menu.py | 27 +- src/anyimage/model_catalog.py | 219 ----- src/anyimage/operators/__init__.py | 11 +- src/anyimage/operators/background.py | 9 +- src/anyimage/operators/balloon/operator.py | 46 +- src/anyimage/operators/common.py | 5 +- src/anyimage/operators/debug.py | 27 +- src/anyimage/operators/depth_plane/build.py | 4 +- src/anyimage/operators/depth_plane/data.py | 11 +- src/anyimage/operators/image.py | 26 +- src/anyimage/operators/image_edit/common.py | 55 +- src/anyimage/operators/image_edit/mask.py | 18 +- .../operators/image_edit/perspective.py | 17 +- src/anyimage/operators/models.py | 3 +- src/anyimage/operators/upscale.py | 158 ++++ src/anyimage/panel.py | 40 +- src/anyimage/preferences.py | 71 +- src/anyimage/properties.py | 46 +- src/anyimage/runtime.py | 17 +- src/anyimage/server/__init__.py | 94 +-- src/anyimage/server/app.py | 92 ++ src/anyimage/server/jobs/background.py | 8 +- .../{debug_realesrgan.py => debug_upscale.py} | 3 +- src/anyimage/server/jobs/foreground.py | 8 +- src/anyimage/server/jobs/upscale.py | 34 + src/anyimage/server/media/input.py | 4 +- src/anyimage/server/model_catalog.py | 76 +- src/anyimage/server/model_download.py | 14 +- src/anyimage/server/model_manager.py | 70 +- src/anyimage/server/models/__init__.py | 4 +- src/anyimage/server/models/onnx_realesrgan.py | 51 -- src/anyimage/server/models/onnx_upscale.py | 142 ++++ src/anyimage/server/models/realesrgan.py | 12 - src/anyimage/server/models/upscale.py | 76 ++ sync_models.py | 96 +-- tests/test_addon.py | 786 +++++++++++++++--- tests/test_download_progress.py | 21 + tests/test_models.py | 59 +- tests/test_pack.py | 22 +- tests/test_realesrgan_onnx.py | 75 -- tests/test_server.py | 62 +- tests/test_upscale_onnx.py | 140 ++++ tests/test_worker.py | 179 +++- uv.lock | 101 ++- 66 files changed, 2156 insertions(+), 1088 deletions(-) create mode 100644 docs/upscale.md delete mode 100644 src/anyimage/model_catalog.py create mode 100644 src/anyimage/operators/upscale.py create mode 100644 src/anyimage/server/app.py rename src/anyimage/server/jobs/{debug_realesrgan.py => debug_upscale.py} (90%) create mode 100644 src/anyimage/server/jobs/upscale.py delete mode 100644 src/anyimage/server/models/onnx_realesrgan.py create mode 100644 src/anyimage/server/models/onnx_upscale.py delete mode 100644 src/anyimage/server/models/realesrgan.py create mode 100644 src/anyimage/server/models/upscale.py delete mode 100644 tests/test_realesrgan_onnx.py create mode 100644 tests/test_upscale_onnx.py diff --git a/AGENTS.md b/AGENTS.md index 629e234..2e7692a 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -69,6 +69,7 @@ uv lock --refresh - 只为特殊、非常规或存在兼容性原因的实现写注释，不复述代码本身 - 删除功能时同时删除无效参数、分支、测试和文档，不保留没有实际消费者的兼容层 - Blender 侧不载入 AI 依赖：UI、Operator 和结果导入放在 `src/anyimage/` 或 `src/anyimage/blender/`，模型下载、ONNX 推理和输出放在 `src/anyimage/server/` +- 扩展级设置及其读取入口统一放在 `preferences.py`；使用 `addon_preferences()` 读取，不把 Blender Add-on Preferences 的定位、默认值或路径解析放进专门管理 BlendJob 的 `runtime.py` - Job 参数使用可序列化的普通字典；Blender 数据只能在主线程的 Operator 响应阶段读写 - 新增或修改 Blender 类型时同步检查 `src/anyimage/__init__.py` 的注册与逆序注销，并补充对应测试 - 依赖版本、wheel 或 Manifest 改变时同步检查 `pyproject.toml`、`uv.lock`、`blender_manifest.toml` 和打包测试 diff --git a/README.md b/README.md index b06e515..d79df13 100644 --- a/README.md +++ b/README.md @@ -12,10 +12,12 @@ AnyImage 是一个 Blender 扩展，可以把一张照片、一段视频或一 - 不运行 AI，把 Image Empty 转换为可实时调整 Thickness 与 Subdivide 的 Mesh Plane - 从 Image Empty 套索提取不规则区域，通过 Poisson 或 Teddy / Chordal Axis 生成封闭 Balloon，并可用 BEN2 分析轮廓、MoGe-2 添加可调深度；源 Empty 保持不变 - 用 Image Box / Lasso / Line / Polyline 直接清理 Image Empty；按 `F` 反选，也可用 BEN2 细化选区 +- Image 编辑会直接更新未共享的图片；图片被多个 Image Empty 共用时，只为当前 Empty 分离编辑结果 - 用 Image Perspective 勾勒四个角点，把透视四边形矫正为矩形裁切图，并沿用源 Image Empty 的空间朝向 - 一键把 Image Empty 转换为立体面片：让 AI 估算深度，或直接用原图作为深度 - 把单张 Image Empty 转换为透视纹理网格，并自动匹配原对象的 Transform 与 Offset - 一键去除图片背景，生成透明前景图 +- 用三档 ONNX 模型将图片、视频或序列帧逐帧放大 4 倍 - 支持视频与序列帧输入，逐帧处理，动画也能得到逐帧对应的立体面片 - 主要转换操作集中在右键菜单，剪贴板图片可直接用系统粘贴快捷键导入 @@ -28,11 +30,11 @@ AnyImage 是一个 Blender 扩展，可以把一张照片、一段视频或一 最快的 DirectML provider；Apple M4 或更新芯片可选 MPS/CoreML，其他 Mac 只提供 CPU 5. 选中 Image Empty，在右键 `AnyImage` 子菜单中点击 `Set Up AI Environment…`；也可以从扩展偏好设置进入 -扩展不内置 uv 或 AI 运行库。`Set Up AI Environment…` 会先说明下载内容、约 600 MB 的下载量、约 1 GB 的磁盘占用和预计等待时间，并可勾选 `Use Mirror`。确认后会下载并校验固定版本的 uv，自动安装轻量的 ONNX Runtime 环境，再下载 `MoGe-2 ViT-S Normal` 和 `BEN2 Base`；不安装 PyTorch，也不下载各模型的 Python 源码包。`Use Mirror` 只切换 uv、Python 和 Python 包的下载源，模型仍使用自己的下载源与备用镜像规则。 +扩展不内置 uv 或 AI 运行库。`Set Up AI Environment…` 会先说明下载内容、约 600 MB 的下载量、约 1 GB 的磁盘占用和预计等待时间，并可勾选 `Use Mirror`。确认后会下载并校验固定版本的 uv，自动安装轻量的 ONNX Runtime 环境，再下载 `MoGe-2 ViT-S Normal`、`BEN2 Base` 和默认高清模型 `Real-ESRGAN General WDN x4v3`；不安装 PyTorch，也不下载各模型的 Python 源码包。`Use Mirror` 只切换 uv、Python 和 Python 包的下载源，模型仍使用自己的下载源与备用镜像规则。 Environment Bootstrap 使用连续的 0–100%：uv 为 1–20%、Python 为 20–50%、packages 为 50–98%，最终验证为 98–100%。三个阶段显示真实的 `1/3`、`2/3`、`3/3`，下载占各阶段的 84–87%；开始与最终验证不显示编号，单阶段状态不显示 `1/1`。MoGe-2 与 BEN2 属于 post-install，各自使用独立进度，不计入 Environment 的这轮 0–100%。任一必需模型在两个来源都下载失败时，安装任务会明确报错；Environment 本身仍会保留。 -选中 Image Empty 后，右键 `AnyImage` 子菜单依次显示 Balloon、三个 AI 功能和环境入口。Environment 或必需模型未就绪时，点击 `Convert to Depth Plane…`、`Convert to Surface…` 或 `Remove Background` 会打开对应准备对话框：Environment 未安装时入口显示 `Set Up AI Environment…`，安装 Environment 与默认模型；Environment 已安装但默认模型缺失时入口显示 `Download Required Models…`，只补齐缺失模型。全部就绪后，菜单入口改为 `AI Environment Settings…`，并打开 AnyImage 扩展偏好设置。 +选中 Image Empty 后，`AnyImage` 会显示在右键菜单最上方。Mesh / Depth / Surface 转换为一组，Remove Background 与 Upscale 作为图片处理归为下一组，环境入口位于最后。Environment 或必需模型未就绪时，点击 `Convert to Depth Plane…`、`Convert to Surface…`、`Remove Background` 或 `Upscale` 会打开对应准备对话框：Environment 未安装时入口显示 `Set Up AI Environment…`，安装 Environment 与默认模型；Environment 已安装但默认模型缺失时入口显示 `Download Required Models…`，只补齐缺失模型。全部就绪后，菜单入口改为 `AI Environment Settings…`，并打开 AnyImage 扩展偏好设置。 其他 AI 模型统一在 AnyImage 扩展偏好设置中按需下载。`Balloon Lasso Tool` 的本地几何功能始终可用；AI 未就绪时，只禁用 Tool Settings 中的 `Contour`、`Depth` 与 Depth Model，并在同处根据状态显示 `Set Up AI Environment…` 或 `Download Required Models…`。共享的 AI 状态检测直接检查 Environment 和 Storage Root 的 `models/` 文件，不会仅因打开菜单或设置而启动 Job Server。 @@ -56,13 +58,13 @@ Environment Bootstrap 使用连续的 0–100%：uv 为 1–20%、Python 为 20 - `Remove Background` 默认关闭；开启后在转换期间使用 BEN2 删除背景 - `Shadeless` 使用 `O Shadeless` 节点组显示不受场景灯光影响的贴图颜色 - `Keep Original` 默认关闭，位于分割线下方；开启后保留源 Image Empty -4. 点击确认，稍等片刻就得到带位移的立体面片。`O DA3 Depth Plane` 与 `O MoGe2 Depth Plane` 都使用 `.color` 的 `Alpha > 0.5` 划分前景范围，范围外由 `Background Depth` 控制。 +4. 点击确认，稍等片刻就得到带位移的立体面片。`O DA3 Depth Plane` 与 `O MoGe2 Depth Plane` 都使用颜色图片的 `Alpha > 0.5` 划分前景范围，范围外由 `Background Depth` 控制。 ### 编辑图片选区 在 3D Viewport 左侧工具栏选择 Image Mask 工具组。Box 拖出矩形，Lasso 手绘轮廓，Line 清理有向直线的一侧，Polyline 逐点构建多边形。点击另一张 Image Empty 时只切换选择；点击当前 Active 图片、空白处或非图片物体时，才开始编辑当前图片。没有 Active 图片时会提示先选择 Image Empty，切换其他工具会取消当前绘制。 -绘制中按 `F` 切换反选：默认保留选区内，反选后保留选区外。Line 按住 `Ctrl` 以 5° 吸附；Polyline 可单击起点圆环或按 `Enter` 完成。`Refine Selection` 只用于 Box / Lasso：先用 BEN2 将几何选区收紧到其中的前景，再应用反选。结果直接替换当前 Empty 的图片，按最终非零 Alpha 紧裁切，并更新显示尺寸与 Image Offset；对象 Transform 不变。透明区保留源 RGB，避免纹理过滤产生黑边。 +绘制中按 `F` 切换反选：默认保留选区内，反选后保留选区外。Line 按住 `Ctrl` 以 5° 吸附；Polyline 可单击起点圆环或按 `Enter` 完成。`Refine Selection` 只用于 Box / Lasso：先用 BEN2 将几何选区收紧到其中的前景，再应用反选。结果直接替换当前 Empty 的图片，按最终非零 Alpha 紧裁切，并更新显示尺寸与 Image Offset；对象 Transform 不变。透明区保留源 RGB，避免纹理过滤产生黑边。未形成有效区域的单击会静默结束；有效选区没有可见像素时只显示警告，不修改图片。 ### 四点透视矫正 @@ -95,7 +97,14 @@ MoGe normal 会转换成适配 Balloon UV 的 tangent-space `normal.png`，由 ### 去除背景 选中 Image Empty 后，右键打开 `AnyImage` 子菜单并选择 `Remove Background`，完成后 Empty 会自动 -换成透明前景图。该操作在子菜单中使用独立分割区域。 +换成透明前景图。 + +### 高清放大 + +选中 Image Empty 后，在右键 `AnyImage` 菜单中点击 `Upscale`。它使用扩展偏好设置中选择的模型进行 4× 放大，支持单图、Movie 和 Image Sequence，并保留源 Alpha。 + + +[阅读高清放大指南](docs/upscale.md) ### 转换为透视网格 @@ -106,12 +115,12 @@ Surface 输出流程构建 DA3 与 MoGe-2 源网格：MoGe-2 使用深度断层 数据。两者再由独立的 Geometry Nodes 修改器完成剔面、快速重拓扑、Relax 与 UV 重投；顶层 `Iterations` 可设为 `0`。开启 Remove Background 时，BEN2 Alpha 只要大于 `0` 就视为前景；三角形三个顶点必须全部为前景，任一顶点 Alpha 为 `0` -都会删除该面并且不再交给 Geometry Nodes；Surface 材质使用 `.color`。匹配相机默认不创建；开启后会作为 Surface 的子级。 +都会删除该面并且不再交给 Geometry Nodes；Surface 材质使用生成的颜色图片。匹配相机默认不创建；开启后会作为 Surface 的子级。 MoGe-2 的像素级 normal 会作为 Object Space normal 贴图接入 `O MoGe2 Surface Layer`，并与原图 Bump 合并；网格本身不写入 custom normal，因此重拓扑降低顶点数后仍能保留细腻的法线细节。 当前此功能仅支持单张图片。 -MoGe-2 ViT-S、ViT-B、ViT-L、DA3 Small、DA3 Base、BEN2 与 Debug 使用的 Real-ESRGAN General x4v3 全部使用 ONNX。模型执行、相机与深度恢复、 +MoGe-2 ViT-S、ViT-B、ViT-L、DA3 Small、DA3 Base、BEN2 与 Upscale 使用的 Real-ESRGAN General WDN x4v3、SAFMN Real x4、Real-ESRGAN x4plus 全部使用 ONNX。模型执行、相机与深度恢复、 点云重投影和网格生成均不依赖 PyTorch。Windows 使用 DirectML 加速；Apple M4 或更新芯片可使用 CoreML 的 `MLProgram` 模型格式，较早或无法识别的 Mac 只提供 CPU；provider 不可用时会回退到 ONNX Runtime CPU， 并在任务状态中显示实际 provider。 @@ -119,7 +128,7 @@ MoGe-2 ViT-S、ViT-B、ViT-L、DA3 Small、DA3 Base、BEN2 与 Debug 使用的 R `Resolution Scale`，MoGe-2 显示 `Resolution Level`， 以及分割线下方的 `Keep Original`。 界面会显示计算所得的实际 Process Res。剔除和重拓扑参数直接在 -Geometry Nodes 修改器中调整；`Shadeless` 与 Depth Plane 共用 `O Shadeless` 材质节点组。颜色贴图始终使用统一的 `.color` 转换产物；Remove Background 只决定是否由 BEN2 修改 Alpha。颜色贴图会按当前场景的 View Transform 选择对应输入色彩空间：AgX、ACES 1.3、ACES 2.0、Filmic、Khronos PBR Neutral 分别使用 `AgX Base sRGB`、`ACES 1.3 sRGB`、`ACES 2.0 sRGB`、`Filmic sRGB`、`Khronos PBR Neutral sRGB`；无法匹配当前 View Transform 时回退 `sRGB`。普通材质的 +Geometry Nodes 修改器中调整；`Shadeless` 与 Depth Plane 共用 `O Shadeless` 材质节点组。颜色贴图始终使用统一的图片转换产物；Remove Background 只决定是否由 BEN2 修改 Alpha。颜色贴图会按当前场景的 View Transform 选择对应输入色彩空间：AgX、ACES 1.3、ACES 2.0、Filmic、Khronos PBR Neutral 分别使用 `AgX Base sRGB`、`ACES 1.3 sRGB`、`ACES 2.0 sRGB`、`Filmic sRGB`、`Khronos PBR Neutral sRGB`；无法匹配当前 View Transform 时回退 `sRGB`。普通材质的 Bump Scale 会在材质组内乘以 `0.001`，并按透视表面的深度空间缩放。Surface 与 Depth Plane 的 Layer Color 同时驱动 Principled BSDF 的 Base Color、Subsurface Radius 和 Emission Color，Layer Normal 同时驱动 Normal 与 Coat Normal；Subsurface、Emission、Coat 的其他设置保持默认。Surface attributes @@ -128,18 +137,19 @@ Layer Color 同时驱动 Principled BSDF 的 Base Color、Subsurface Radius 和 ### 视频与序列帧 -AnyImage 支持视频（mp4、mkv、mov 等）和序列帧，会逐帧处理并生成逐帧对应的立体面片。`.color` 与 `.depth` 按源 Empty Image 的类型载入：普通单图为 `IMAGE`，Movie/Sequence 即使只处理一帧也仍为 `SEQUENCE`。处理范围默认 +AnyImage 支持视频（mp4、mkv、mov 等）和序列帧，会逐帧处理并生成逐帧对应的立体面片。颜色图片与 `.depth` 按源 Empty Image 的类型载入：普通单图为 `IMAGE`，Movie/Sequence 即使只处理一帧也仍为 `SEQUENCE`。处理范围默认 跟随 Empty 上设置的 `Image Frames` 与 `Offset`，不需要额外配置。DA3 的转换任务复用同一个 ONNX -Session 逐帧推理，显存峰值不随视频帧数平方增长。Debug 只保存模型 API 的直接返回：DA3 多帧输入 -保留联合推理，MoGe-2、BEN2 与 Real-ESRGAN General x4v3 按各自的单图 API 逐帧保存。General x4v3 直接输出 4× RGB `prediction*.png`，不做切片、强度混合或图片替换。 +Session 逐帧推理，显存峰值不随视频帧数平方增长。Upscale 也会逐帧生成受输出上限约束的 4× 颜色序列并写回 Empty。Debug 只保存模型 API 的直接返回：DA3 多帧输入 +保留联合推理，MoGe-2、BEN2 与三档高清模型按各自的单图 API 逐帧保存。高清 Debug 输出 4× RGB `prediction*.png`，不做强度混合或图片替换；三个模型都由适配层自动分块，因此接受任意图片宽高。 ## 任务结果保存在哪里 -任务结果保存在 Runtime Storage Root 的 `jobs/` 下，默认是 `~/.anyimage/jobs/`。Job Server 接受任务时自动生成 opaque Job ID，并创建同名文件夹；Blender 不预先申请 ID 或目录。Packed Image 和序列帧的临时输入由 Blender 在任务结束后清理，Server 产物直接写入该 Job 目录。视频或序列帧生成的面片会引用其中的序列帧，请保留对应文件夹，不要随意删除。 +任务结果保存在 Runtime Storage Root 的 `jobs/` 下，默认是 `~/.anyimage/jobs/`。Job Server 接受任务时自动生成 opaque Job ID，并创建同名文件夹；Blender 不预先申请 ID 或目录。Packed Image 和序列帧的临时输入由 Blender 在任务结束后清理，Server 产物直接写入该 Job 目录。普通图片层使用 `image.0001.*` 起的文件名，导入后沿用源 Image 名称并由 Blender 自动添加 `.001`；Upscale 单独使用 `upscale.0001.png` 和 `.upscale` 数据块后缀。视频或序列帧生成的面片会引用其中的序列帧，请保留对应文件夹，不要随意删除。 + +每个任务目录中的 `job.log` 保存任务细节和完整 traceback；Storage Root 下的 `server.log` 只保留带时间、Job ID 与 Job Type 的生命周期主线，以及 Server 自身异常。 ## 常见问题 -- 从 0.12.0–0.12.5 升级到 0.12.6 或更高版本后，请重新执行一次 `Set Up AI Environment…`，重建不再单独安装 BlendJob 的 Server Environment - 首次使用需要下载模型；Environment 已安装但默认模型下载慢或失败时，点击 `Download Required Models…` 可补齐缺失项，也可在 AnyImage 扩展偏好设置中逐项重试 - 安装 Environment 时下载状态会显示平均速度；网络适合国内镜像时在确认对话框勾选 `Use Mirror`，否则保持关闭。选择的来源固定用于 uv、Python 和 Python 包，失败时不会自动切换 - 默认模型下载会在 Hugging Face 与 OmooLab 镜像之间自动回退；某一来源暂时断线不会阻止后续模型尝试另一来源 diff --git a/docs/depth-plane.md b/docs/depth-plane.md index c0290e3..e087a70 100644 --- a/docs/depth-plane.md +++ b/docs/depth-plane.md @@ -38,7 +38,7 @@ DA3 和 MoGe-2 都以未经抠图的 RGB 估计深度。启用 Remove Background 转换会创建以下数据： - 与源 Image Empty 对齐的面片 -- `.depth` 与 `.color` Image datablock +- `.depth` 与沿用源图名称的颜色 Image datablock - 使用 `O DA3 Depth Plane` 或 `O MoGe2 Depth Plane` 的 Geometry Nodes 修改器 - 连接颜色、Alpha、Bump 与材质位移的材质 @@ -54,7 +54,7 @@ Geometry Nodes 提供深度范围、细分、重拓扑、平滑、背景深度 ```text / -├── color.0001.jpg/png +├── image.0001.jpg/png ├── depth.0001.exr └── results.npz ``` diff --git a/docs/development.md b/docs/development.md index 404b98f..e9c73b2 100644 --- a/docs/development.md +++ b/docs/development.md @@ -37,6 +37,12 @@ Mesh Plane、Depth Plane、Surface、Background 和 Balloon 按功能组织。 Server 的模型模块返回 Prediction，几何模块处理数组与网格，输出模块写文件。Blender 数据只在 Operator Response 的主线程阶段创建。 +扩展级设置统一通过 `preferences.addon_preferences()` 读取，不归属 BlendJob Runtime；Storage Root 路径回调也由 Preferences 模块提供，Runtime 只消费回调结果。扩展包标识直接使用根模块中的 `__package__`，以兼容 Blender 实际的 `bl_ext..anyimage` 命名空间；不得按第一个点截断，否则会错误查找 `bl_ext` 并静默回退默认值。Preferences 的设备枚举和面板数据按需导入，避免循环依赖。 + +正式 Upscale 使用仅具备 `UNDO` 的 Job Operator 完成输入准备、任务提交和结果回写，不注册 Last Operation 重做面板。每次 `execute()` 都从 Preferences 读取默认模型与输出长边上限，再写入隐藏的可序列化 Operator 参数。高清结果只通过 `Image.anyimage_upscale_source` 保存原始 Image 引用；Operator 自身仅以隐藏的 `StringProperty` 保存原图名称，并在主线程响应时从 `bpy.data.images` 重新解析，避免保留已被移除的 RNA 指针。再次执行时从原图引用准备输入并替换当前高清数据块；无其他用户的旧高清结果会在新结果配置成功后移除，因此同一 Empty 只保留一个高清版本。模型只处理 RGB；源图 Alpha 经过同样的输入尺寸限制后独立放大并合并回结果。 + +Blender 侧需要落盘的 Packed Image 与浮点深度 EXR 都通过 `Image.save(save_copy=True)` 导出，不调用 `save_render`，避免把 Scene Render 的 View Transform、Look 或 Exposure 耦合到模型输入和 Non-Color 数据；颜色结果载入后继承源 Image 的 Color Space。Server 中的 `PIL.Image.save()` 只负责普通文件编码，不属于 Blender Render 色彩管理。 + Image Box、Lasso、Line、Polyline 与 Perspective 位于 `operators/image_edit/` 功能包。模块职责、交互状态、像素管线、Refine 和 Empty 定位见 [Image 工具技术实现](image-tools-technical.md)。 ## 项目结构 @@ -74,7 +80,7 @@ uv run pytest tests/test_addon.py tests/test_manifest.py tests/test_pack.py -v uv run python pack.py ``` -`pack.py` 为 Windows x64、macOS Apple Silicon 与 Linux x64 构建扩展 ZIP 到 `dist/`。依赖、wheel 或 Manifest 改变后，应同步检查 `pyproject.toml`、`uv.lock`、`blender_manifest.toml` 与打包测试。 +`pack.py` 为 Windows x64、macOS Apple Silicon 与 Linux x64 构建扩展 ZIP 到 `dist/`。下载项目依赖时会把 `src/anyimage/wheels/` 作为优先可用的 wheel 来源，因此尚未发布到 PyPI 的本地 BlendJob 构建也能进入三个平台包。依赖、wheel 或 Manifest 改变后，应同步检查 `pyproject.toml`、`uv.lock`、`blender_manifest.toml` 与打包测试。 CI 分为三条流程： @@ -122,6 +128,8 @@ BLENDER_BIN="/c/Program Files/Blender Foundation/Blender 4.5/blender.exe" - Depth Anything 3 代码采用 Apache 2.0 - BEN2 ONNX 模型采用其仓库声明的 Apache 2.0 License -- Real-ESRGAN General x4v3 ONNX 权重采用其模型仓库声明的 BSD 3-Clause License +- Real-ESRGAN General WDN x4v3 与 Real-ESRGAN x4plus ONNX 权重采用其模型仓库声明的 BSD 3-Clause License +- SAFMN Real x4 ONNX 权重采用其模型仓库声明的 Apache 2.0 License +- 模型目录显式声明可用的 Hugging Face 仓库；未声明时只同步经过大小与 SHA-256 校验的 R2 制品 - AnyImage Extension 代码采用 GPL-3.0-or-later - 模型权重许可证由所选 Hugging Face 模型仓库决定 diff --git a/docs/faq.md b/docs/faq.md index 1a56e00..6aeba50 100644 --- a/docs/faq.md +++ b/docs/faq.md @@ -12,7 +12,7 @@ AnyImage 菜单只会在选中 Image Empty 时出现。确认对象类型是 Emp - Surface 的输入不是单张图片。 - 当前对象不是 Image Empty。 -Environment 或默认模型缺失不会再让三个 AI 功能变灰；点击功能会根据状态打开 **Set Up AI Environment…** 或 **Download Required Models…**。Balloon 的 **Contour**、**Depth** 与 Depth Model 会在 AI 未就绪时单独禁用，同一位置提供安装按钮。 +Environment 或默认模型缺失不会让四个 AI 功能变灰；点击功能会根据状态打开 **Set Up AI Environment…** 或 **Download Required Models…**。Balloon 的 **Contour**、**Depth** 与 Depth Model 会在 AI 未就绪时单独禁用，同一位置提供安装按钮。 ## Environment 安装失败 @@ -31,6 +31,12 @@ Environment 或默认模型缺失不会再让三个 AI 功能变灰；点击功 第一次任务需要启动 Job Server 并把模型载入内存。Server 运行且模型仍在缓存中时，后续同类任务通常会更快。 +## Upscale 很卡或提示内存不足 + +三个高清模型都使用分块推理，单次 ONNX 调用不会把完整图片送进 provider；但完整输出数组、PNG 编码和 Blender 图片载入仍取决于最终分辨率。正式 Upscale 默认将输出长边限制为 4096 px，可在 Preferences 中调整 **Maximum Upscale Size**，最高 8192 px，并始终限制在约 32 MP。DirectML 或其他 provider 无法分配资源时，任务会提示降低该值或改用 CPU，而不会显示底层的 UTF-8 解码异常。 + +对已经高清的结果再次执行不会继续放大该结果。AnyImage 只保留原始 Image 引用，并从原图重新运行所选模型；开始和完成时会在状态栏说明。 + ## Depth Plane 太平或深度方向反了 在 Geometry Nodes 修改器中调整 **Depth Scale**。增大绝对值会加强深度，使用负值可以反转前后关系。 @@ -50,7 +56,3 @@ Environment 或默认模型缺失不会再让三个 AI 功能变灰；点击功 ## 如何完全撤销一次转换 转换完成后直接按 `Ctrl+Z`。AnyImage 会把网格、材质、图片和修改器的创建合并为一次 Blender Undo 操作。 - -## 从旧版本升级后功能不可用 - -如果 Environment 声明或依赖发生变化，请重新执行一次 **Set Up AI Environment…**。从 0.12.0–0.12.5 升级到 0.12.6 或更高版本时需要重建 Environment。 diff --git a/docs/features.md b/docs/features.md index 8efc02b..8a0ac26 100644 --- a/docs/features.md +++ b/docs/features.md @@ -11,6 +11,7 @@ | Depth Plane | 快速制作可调节的立体图片 | 单图、视频、序列帧 | 带颜色、深度与 Geometry Nodes 的面片 | | Surface | 从单张图片恢复透视空间 | 单图 | 相机空间网格、材质与可选匹配相机 | | Remove Background | 提取透明前景 | 单图、视频、序列帧 | 透明图片或透明序列 | +| Upscale | 在安全输出上限内增加图片像素 | 单图、视频、序列帧 | 三档模型生成的高清图片或序列 | | Image Mask | 用 Box、Lasso、Line 或 Polyline 直接清理图片 | 静态 Image Empty | 原位更新并 Pack 的透明裁切图 | | Image Perspective | 四点透视矫正裁切 | 静态 Image Empty | Pack 的矩形裁切图 | | Balloon | 沿轮廓绘制鼓起造型 | 静态 Image Empty | 带原图 UV 的封闭网格 | @@ -22,11 +23,11 @@ Depth Plane 与 Surface 可使用 DA3 或 MoGe-2。DA3 提供相对深度和相 复制图片后，在支持的编辑器中按 `Ctrl+V`（macOS 为 `Cmd+V`）即可粘贴；图片会 Pack 到当前 `.blend`。3D View 的 Object Mode 默认创建 Reference Image，也可在 Redo 面板切换为 Plane；Shader Editor、Compositor 与 Geometry Nodes 创建对应图片节点，Sculpt、Vertex Paint 与 Texture Paint 则设置笔刷 Texture。详见[粘贴剪贴板图片](clipboard-image.md)。 -主要功能位于 Image Empty 的 `AnyImage` 右键菜单。 +主要功能位于 Image Empty 的 `AnyImage` 右键菜单；AnyImage 使用 prepend 注册，显示在右键菜单最上方。 `Convert to Mesh Plane` 位于 `Convert to Depth Plane` 上方，不需要 AI Environment。它生成基础四边形，并由 `O Mesh Plane` 修改器实时处理 `Thickness` 与 `Subdivide`；`Thickness` 为 `0` 时保持单面。操作会立即执行；创建后可在 Last Operation 中调整 `Shadeless` 与 `Keep Original`。 -Balloon 作为不依赖 AI 的基础工具单独排列；三个 AI 功能在缺少 Environment 或默认模型时仍可点击：缺 Environment 时转入 `Set Up AI Environment…`，只缺默认模型时转入 `Download Required Models…`。全部就绪后，菜单底部显示 `AI Environment Settings…`，用于打开 AnyImage 扩展偏好设置。 +Balloon 作为不依赖 AI 的基础工具单独排列；Mesh Plane、Depth Plane 与 Surface 属于转换组，Remove Background 与 Upscale 由分割线归入图片处理组。Depth Plane、Surface、Remove Background 与 Upscale 四个 AI 功能在缺少 Environment 或默认模型时仍可点击：缺 Environment 时转入 `Set Up AI Environment…`，只缺默认模型时转入 `Download Required Models…`。全部就绪后，菜单底部显示 `AI Environment Settings…`，用于打开 AnyImage 扩展偏好设置。Upscale 菜单项同时显示当前图片分辨率，不弹出属性对话框；模型与输出长边上限统一读取 Preferences 的 `Default Upscale Model` 和 `Maximum Upscale Size`，Operator 不注册 Last Operation 重做面板。高清结果再次执行时从保存的原图重新生成并替换当前版本，不级联处理已有结果；名称始终基于原图，透明图片保留放大后的 Alpha。 转换任务完成后，AnyImage 在一个 Undo 步骤中创建网格、图片、材质和修改器。源 Image Empty 是否保留由转换选项决定。 @@ -42,7 +43,7 @@ Balloon 默认在 Blender 内部生成，无需 Environment 或模型。Tool Set ## 动画 -Depth Plane 和 Remove Background 支持 Movie 与 Image Sequence。它们按 Image Empty 的帧范围处理素材，并生成与场景帧对应的颜色或深度序列。 +Depth Plane、Remove Background 和 Upscale 支持 Movie 与 Image Sequence。它们按 Image Empty 的帧范围处理素材，并生成与场景帧对应的颜色或深度序列。 Blender 无法整体 Pack 图片序列，因此动画结果会继续引用 Job 目录。保存或移动 `.blend` 时，应同时保留这些外部文件。详见[动画与序列帧](animation.md)。 @@ -50,7 +51,7 @@ Blender 无法整体 Pack 图片序列，因此动画结果会继续引用 Job AI 任务在独立 Python Environment 中通过 ONNX Runtime 执行。Blender 负责界面、任务提交和结果导入，Job Server 负责模型与文件处理。 -默认 Storage Root 为 `~/.anyimage`，其中保存 Environment、模型、日志和 Job 结果。连续任务会复用已加载的模型 Session；DA3 与 MoGe-2 共享一个深度模型缓存位置，BEN2 与 Real-ESRGAN 可同时保留。 +默认 Storage Root 为 `~/.anyimage`，其中保存 Environment、模型、日志和 Job 结果。连续使用同一模型会复用已加载的 Session；切换到另一模型族时释放旧 Session，使 Server 最多保留当前使用的一套模型。 安装、下载和推理都在后台运行，并通过状态栏提供进度与取消操作。详细设计见[Job Server](server.md)。 @@ -61,9 +62,10 @@ AI 任务在独立 Python Environment 中通过 ONNX Runtime 执行。Blender - MoGe-2 ViT-S、ViT-B、ViT-L Normal - DA3 Small、DA3 Base - BEN2 Base +- Real-ESRGAN General WDN x4v3、SAFMN Real x4、Real-ESRGAN x4plus -Debug 另外提供 Real-ESRGAN General x4v3，用于直接检查 4× RGB 放大能力；它当前不属于正式图片工具。 +Upscale 提供 Real-ESRGAN General WDN x4v3、SAFMN Real x4 和 Real-ESRGAN x4plus 三档 4× RGB 模型。默认 WDN 偏速度和保守细节，SAFMN 居中，x4plus 的细节更明显；三者也保留对应 Debug 入口用于模型直通比较。正式功能会独立放大并合并源 Alpha，默认把输出长边限制在 4096 px，可在 Preferences 提高到 8192 px，并始终受约 32 MP 硬上限约束。 -首次安装会准备 MoGe-2 ViT-S Normal 与 BEN2 Base，其他模型可从 AnyImage 扩展偏好设置按需下载。 +首次安装会准备 MoGe-2 ViT-S Normal、BEN2 Base 与 Real-ESRGAN General WDN x4v3，其他模型可从 AnyImage 扩展偏好设置按需下载。 -Debug 工作流保存模型 API 的直接返回，适合验证模型、设备与输入。General x4v3 对每张输入整图推理并保存 `prediction*.png`，不加入切片、降噪强度、Alpha 合成或 Blender 图片写回。Depth Plane 和 Surface 则在模型结果之上继续生成 Blender 可用的贴图与几何。文件字段见[输出格式](outputs.md)。 +Debug 工作流保存模型 API 的直接返回，适合验证模型、设备与输入。三档高清模型对每张输入保存 `prediction*.png`，不加入降噪强度、Alpha 合成或 Blender 图片写回。三个模型统一由适配层以 256×256、16 像素上下文执行重叠分块，因而都接受任意图片宽高。Depth Plane 和 Surface 则在模型结果之上继续生成 Blender 可用的贴图与几何。文件字段见[输出格式](outputs.md)。 diff --git a/docs/image-tools-technical.md b/docs/image-tools-technical.md index 45fb199..0ee2879 100644 --- a/docs/image-tools-technical.md +++ b/docs/image-tools-technical.md @@ -50,21 +50,25 @@ flowchart TD Box / Lasso 的顺序固定为 **Refine → Invert**。BEN2 只提供内容 Alpha；最终 RGB 始终来自源图。Blender 侧只修改 Alpha，不清理透明区 RGB，以避免 Viewport 纹理过滤产生黑边。Line / Polyline 强制关闭 Refine，只执行几何 mask 与 Invert。 +未形成有效区域的单击直接返回 `CANCELLED`，不报告消息；最终 Alpha 全空时，Operator 报告 `WARNING` 并保持源图不变。其他数据、状态或运行时异常仍报告 `ERROR`。 + Line 的屏幕坐标 Y 轴与图片坐标 Y 轴方向相反，因此像素阶段会反转线段端点，使半透明预览侧与实际保留侧一致；`Ctrl` 只影响屏幕端点的 5° 吸附。零长度点击在投影前直接取消，投影结果不足两个点时也会安全终止。Polyline 接近起点时绘制闭合圆环，并支持单击起点或 `Enter` 完成；完成时显式闭合，不把当前游标位置作为额外顶点。 ## Perspective 管线 四点先按源图片 X/Y 方向规范为凸四边形，输入顺序不参与输出方向。第四点后只生成一次 512 × 512 临时纹理；移动鼠标改变 UI 比例，确认时才进行全分辨率 Homography。 +预览引用在删除前先从 Operator 清空；若 Undo 或 Blender 已移除临时 Image，清理会忽略失效的 RNA 引用，因此切换工具和重复结束都安全。 + 透视重采样使用预乘 Alpha 插值，随后恢复 RGB、清理透明像素并紧裁切。输出最长边限制为 8192 px，总像素限制约 16 MP。 Perspective 不调用 BEN2，也不支持 Invert。它是 deform 工具，不属于 mask 清理语义。 ## Empty 写回 -`replace_empty_image_pixels()` 创建并 Pack 新 Image datablock，然后替换当前 Empty 的 `data`。它不复制对象，也不修改 `matrix_world`。 +`replace_empty_image_pixels()` 先创建并验证完整结果，再根据引用源 Image 的 Image Empty 数量选择提交方式：只有一个引用时，把原 Image 的用户重映射到结果、删除旧 datablock 并恢复原名，最终不留下额外图片；两个或更多引用时，只替换当前 Empty 的 `data`，其他 Empty 继续使用原图。共享图片的新结果沿用源 Image 基础名，不附加 `.mask`、`.perspective` 或其他工具后缀；同名冲突交给 Blender 自动添加 `.001`。两种方式都不复制对象，也不修改 `matrix_world`。 -`align_empty_display_to_bounds()` 根据源画布、旧 `plane_bounds`、结果尺寸和放置 bounds 更新： +`empty_display_alignment()` 在写入像素前，根据源画布、旧 `plane_bounds`、结果尺寸和放置 bounds 计算显示参数；写回函数随后更新： - `empty_display_size`； - `empty_image_offset`。 diff --git a/docs/image-tools.md b/docs/image-tools.md index e03105a..f6d6871 100644 --- a/docs/image-tools.md +++ b/docs/image-tools.md @@ -50,7 +50,7 @@ Image Box、Lasso、Line、Polyline 与 Perspective 直接编辑静态 Image Emp Box / Lasso 启用 **Refine Selection** 后，BEN2 会先把几何选区收紧到其中的前景，再应用反选。例如框住一个苹果并开启反选，最终会清理苹果本身、保留其余内容。Line / Polyline 不提供 Refine，以保持直线边界和明确的裁切方向。 -结果会直接替换当前 Empty 的图片，并把画布紧裁切到最终可见 Alpha。对象 Transform 不变，显示尺寸与 Image Offset 会更新，使剩余像素保持在原图片平面上的位置。透明区只修改 Alpha，保留源 RGB 供纹理过滤使用；完全透明的结果不会写入。 +结果会直接更新当前 Empty 的图片，并把画布紧裁切到最终可见 Alpha。只有一个 Image Empty 使用该图片时，工具更新原 Image；两个或更多 Image Empty 共用它时，工具为当前 Empty 创建独立结果，其他 Empty 保持不变。对象 Transform 不变，显示尺寸与 Image Offset 会更新，使剩余像素保持在原图片平面上的位置。透明区只修改 Alpha，保留源 RGB 供纹理过滤使用。未形成有效区域的单击会静默结束；有效选区没有可见像素时只显示警告，不修改图片。
**截图占位：Polyline 与反选** diff --git a/docs/index.md b/docs/index.md index 0e1db14..3363a38 100644 --- a/docs/index.md +++ b/docs/index.md @@ -9,6 +9,7 @@ AnyImage 是一款 Blender 扩展，可以把照片、视频或序列帧转换 - 用 **Depth Plane** 快速得到可调节的浮雕式面片 - 用 **Surface** 从单张图片生成透视纹理网格 - 用 **Remove Background** 一键获得透明前景 +- 用 **Upscale** 选择三档 ONNX 模型进行 4× 高清放大 - 用 **Image Box / Lasso / Line / Polyline** 直接清理图片，或用 **Image Perspective** 做四点透视矫正 - 处理视频和序列帧，生成可随时间播放的 Depth Plane - 使用 Geometry Nodes 继续调整深度、细分、重拓扑与平滑效果 @@ -34,7 +35,7 @@ AnyImage 是一款 Blender 扩展，可以把照片、视频或序列帧转换 - **选择合适的结果** - 了解 Depth Plane、Surface 和去背景之间的区别。 + 了解 Depth Plane、Surface、去背景和高清放大之间的区别。 [查看工作流对比](workflows.md) @@ -51,7 +52,7 @@ AnyImage 是一款 Blender 扩展，可以把照片、视频或序列帧转换 - Blender 4.5 或更高版本 - Windows x64、macOS Apple Silicon 或 Linux x64 - 首次安装环境与模型时需要网络连接 -- Depth Plane、Surface、Remove Background 与 Balloon 的 AI 增强需要独立 Environment，以及默认的 MoGe-2 和 BEN2 模型；Balloon 基础功能无需安装 +- Depth Plane、Surface、Remove Background、Upscale 与 Balloon 的 AI 增强需要独立 Environment，以及默认的 MoGe-2、BEN2 和 Real-ESRGAN General WDN x4v3 模型；Balloon 基础功能无需安装 !!! tip 不确定从哪里开始时，先使用默认设置将一张人物或物体照片转换为 Depth Plane。它速度快、参数直观，也最适合熟悉 AnyImage 的基本工作方式。 diff --git a/docs/installation.md b/docs/installation.md index b9095d3..22f2e5a 100644 --- a/docs/installation.md +++ b/docs/installation.md @@ -21,6 +21,8 @@ AnyImage 的安装分为两步：先安装 Blender Extension，再安装独立 - **Storage Root**：保存运行环境、模型和任务文件的位置。默认是用户目录下的 `.anyimage`，一般无需修改。 - **Device**：正式转换任务使用的计算设备。AnyImage 默认选择本机可用的最快设备；如果 GPU Provider 不可用，会使用 CPU。 +- **Default Upscale Model**：Upscale 使用的模型，默认是 General WDN x4v3。 +- **Maximum Upscale Size**：Upscale 输出长边上限，默认 4096 px，最高 8192 px。 如果准备把 Storage Root 放到其他磁盘，建议在首次安装 Environment 之前修改。更换目录后，旧目录中的环境和模型不会自动移动。 @@ -30,7 +32,7 @@ AnyImage 的安装分为两步：先安装 Blender Extension，再安装独立 确认安装前，对话框会说明以下内容： -- 安装独立 Python、ONNX Runtime 与支持包，并下载 **MoGe-2 ViT-S Normal** 和 **BEN2 Base**。 +- 安装独立 Python、ONNX Runtime 与支持包，并下载 **MoGe-2 ViT-S Normal**、**BEN2 Base** 和 **Real-ESRGAN General WDN x4v3**。 - 默认模型合计约 360 MB；连同 Environment 预计需要约 600 MB 下载量和约 1 GB 磁盘空间。 - 根据网络速度，首次安装通常需要 5–20 分钟。 - **Use Mirror** 只让 uv、Python 和 Python 包使用国内镜像；模型下载源保持不变。 @@ -44,7 +46,7 @@ AnyImage 的安装分为两步：先安装 Blender Extension，再安装独立 完成后可以用以下方式确认： - **Job Server** 不再显示 “Environment is not installed”。 -- AnyImage 扩展偏好设置中的 MoGe-2 ViT-S Normal 和 BEN2 Base 显示 **Ready**。 +- AnyImage 扩展偏好设置中的 MoGe-2 ViT-S Normal、BEN2 Base 和 Real-ESRGAN General WDN x4v3 显示 **Ready**。 - 选中 Image Empty 后，右键菜单底部显示 **AI Environment Settings…**。 **Convert to Depth Plane…**、**Convert to Surface…** 和 **Remove Background** 在 AI 未就绪时仍可点击，点击后会根据状态转入 **Set Up AI Environment…** 或 **Download Required Models…**。Balloon 的本地套索与几何生成功能不要求安装 AI Environment；未就绪时，Tool Settings 只会禁用 **Contour**、**Depth** 与 Depth Model，并显示安装按钮。 diff --git a/docs/models-and-storage.md b/docs/models-and-storage.md index cca8320..4ce77a5 100644 --- a/docs/models-and-storage.md +++ b/docs/models-and-storage.md @@ -8,14 +8,16 @@ AnyImage 把运行环境、模型和任务结果放在 Blender 之外，多个 B - **Ready** 表示模型文件已经就绪。 - **Download** 用于按需下载模型。 -- 默认安装会准备 MoGe-2 ViT-S Normal 和 BEN2 Base。 -- DA3、更大的 MoGe-2，以及只供 Debug 使用的 Real-ESRGAN General x4v3 需要时再下载即可。 +- 默认安装会准备 MoGe-2 ViT-S Normal、BEN2 Base 和 Real-ESRGAN General WDN x4v3。 +- DA3、更大的 MoGe-2，以及可选的 SAFMN Real x4、Real-ESRGAN x4plus 需要时再下载即可。 +- 可公开取得的 ONNX 制品使用经过大小与 SHA-256 校验的原址，并以 OmooLab R2 为镜像；不会把源码仓库名称推测成 Hugging Face 下载地址。 +- Real-ESRGAN General WDN x4v3 与 SAFMN Real x4 的上游只发布 PyTorch 权重，当前 ONNX 均从官方权重导出并仅从 R2 下载。 较大的模型不一定适合所有图片。建议先用默认模型完成构图，再针对细节与性能需求比较其他模型。 ## Device -扩展偏好设置中的 **Device** 控制 Depth Plane、Surface 和 Remove Background 等正式任务。 +扩展偏好设置中的 **Device** 控制 Depth Plane、Surface、Remove Background 和 Upscale 等正式任务。 - Windows 通常优先使用 DirectML。 - Apple Silicon macOS 通常优先使用 CoreML。 @@ -23,6 +25,10 @@ AnyImage 把运行环境、模型和任务结果放在 Blender 之外，多个 B Debug 面板中的 **Debug Device** 只影响 Debug 任务，不会覆盖正式任务的 Device。 +## Session 缓存 + +连续使用同一模型和 Device 时，Server 会复用已经加载的 ONNX Session。切换到 BEN2、深度或高清等另一模型族时，旧 Session 会被释放，因此 Server 最多保留当前使用的一套模型；如果高清推理遇到 provider 资源错误，也会立即卸载该 Session。这样保留连续操作的速度，同时避免多个模型长期占用显存或系统内存。 + ## Storage Root 默认 Storage Root 是 `~/.anyimage`，主要包含： @@ -31,9 +37,9 @@ Debug 面板中的 **Debug Device** 只影响 Debug 任务，不会覆盖正式 ~/.anyimage/ ├── .venv/ # 独立 Python Environment ├── models/ # 模型文件 -├── jobs/ # 每次任务的结果 +├── jobs/ # 每次任务的结果与 job.log ├── install.log # Environment 安装日志 -└── server.log # Job Server 日志 +└── server.log # 带 Job ID 的任务主线与 Job Server 日志 ``` 更改 Storage Root 相当于切换到另一套环境、模型和任务数据。AnyImage 不会自动移动旧目录。 @@ -48,6 +54,7 @@ Debug 面板中的 **Debug Device** 只影响 Debug 任务，不会覆盖正式 遇到安装或 Server 问题时： -- 在扩展偏好设置中点击 **Open Server Log** 查看运行日志。 +- 在扩展偏好设置中点击 **Open Server Log** 查看带时间、Job ID 与 Job Type 的运行主线。 +- 具体 Job 的 Handler 细节和完整 traceback 查看 `jobs//job.log`。 - Environment 安装问题查看 Storage Root 下的 `install.log`。 - 修改 Storage Root 后，确认你检查的是当前目录中的日志。 diff --git a/docs/outputs.md b/docs/outputs.md index bfc6ae1..279ac47 100644 --- a/docs/outputs.md +++ b/docs/outputs.md @@ -8,13 +8,14 @@ | --- | --- | --- | | `generate-depth` | `results.npz`、深度与颜色序列 | Depth Plane 所需深度、相机和贴图 | | `generate-surface` | `mesh.npz`、`normal.png`、颜色层 | Surface 网格、相机与材质输入 | -| `remove-background` | `color.0001.png` 起 | 带 Alpha 的颜色层 | +| `remove-background` | `image.0001.png` 起 | 带 Alpha 的图片层 | +| `upscale-image` | `upscale.0001.png` 起 | 所选高清模型在输出上限内生成的 4× 高清层，并保留源图 Alpha | | `refine-image-mask` | `selection-input.png`、`image.png`、`bounds.json` | BEN2 细化后的局部内容 mask 与原图坐标 bounds | | `analyze-balloon` | `lasso-input.png`，按选项输出 `foreground.png`、`contour.json`、`depth.npz`、`normal.png` | 套索约束输入、BEN2 轮廓与 MoGe-2 深度、法线 | | `debug-da3` | `prediction.npz` | DA3 API 返回 | | `debug-moge2` | `prediction*.npz` | MoGe-2 API 返回 | | `debug-ben2` | `foreground*.png` | BEN2 API 返回 | -| `debug-realesrgan-general-x4v3` | `prediction*.png` | Real-ESRGAN General x4v3 的 4× RGB API 返回 | +| `debug-upscale` | `prediction*.png` | 请求模型的 4× RGB API 返回 | ## NPZ 字段 @@ -24,10 +25,10 @@ Balloon Depth 的 `depth.npz` 包含相对 `depth_signal`、归一化拟合平 Surface 的 `mesh.npz` 包含 `vertices`、`faces`、`uv`、`intrinsics`、`image_size` 与 `reference_depth`，以及 Geometry Nodes 使用的网格属性。 -Debug 文件保持模型 API 的数据形态，用于比较推理结果；正式 Job 则提供 Blender 导入所需的派生数据。General x4v3 的浮点 RGB 返回只做 `[0, 1]` 限制和 8-bit PNG 序列化，不做额外画质处理。 +Debug 文件保持模型 API 的数据形态，用于比较推理结果；正式 Job 则提供 Blender 导入所需的派生数据。三档高清模型的浮点 RGB 返回只做 `[0, 1]` 限制和 8-bit PNG 序列化，不做色调保护或额外画质处理；三个模型都使用 256×256、16 像素上下文的重叠分块，外边界以镜像像素提供模型上下文。输入含 Alpha 时，Alpha 不进入 RGB 模型，而是在相同输入限制后独立进行 4× 高质量插值并合并回 RGBA 结果。正式 `upscale-image` 还会在推理前应用输出长边与总像素上限，再把结果命名为 `upscale.0001.png` 起的序列并交给 Blender 写回 Image Empty；结果 Image 数据块始终基于原图名称使用 `.upscale` 后缀，重复执行时替换当前高清版本。 ## 图片序列 -正式颜色层统一使用 `color.0001.*` 起始的文件名。深度动画使用 EXR 序列以保存浮点值。源 Empty 为 Movie 或 Sequence 时，Blender 按序列载入这些文件，即使任务只输出一帧。 +正式图片层使用 `image.0001.*` 起始的文件名，Upscale 层单独使用 `upscale.0001.png`。导入 Blender 后沿用源 Image 名称，由 Blender 在同名冲突时自动添加 `.001`；不再添加 `.color`。深度动画使用 EXR 序列以保存浮点值。源 Empty 为 Movie 或 Sequence 时，Blender 按序列载入这些文件，即使任务只输出一帧。 MoGe-2 Surface 的法线单独保存为 `normal.png`。材质在 Blender 端创建，不属于 Job 文件格式。 diff --git a/docs/remove-background.md b/docs/remove-background.md index 6047378..c6e82ec 100644 --- a/docs/remove-background.md +++ b/docs/remove-background.md @@ -8,7 +8,7 @@ Remove Background 使用 BEN2 为 Image Empty 生成透明前景，不创建网 2. 在 3D View 中右键，选择 **AnyImage > Remove Background**。 3. 等待状态栏提示任务完成。 -AnyImage 会创建新的透明颜色图，并把当前 Empty 切换到新图。对象的 Transform、显示尺寸、Offset 和其他显示设置保持不变。 +AnyImage 会生成透明颜色图。原图只被当前 Empty 引用时，结果直接替换原图并沿用名称；原图被多个对象引用时，只为当前 Empty 创建新图，其他对象继续使用原图。对象的 Transform、显示尺寸、Offset 和其他显示设置保持不变。 ## 在转换时去除背景 @@ -29,7 +29,7 @@ Depth Plane 和 Surface 的转换对话框也提供 **Remove Background**： ### Remove Background 是灰色的 -确认 Environment 已安装，且 **MoGe-2 ViT-S Normal** 与 **BEN2 Base** 都显示 Ready。正式功能会等默认必需模型全部就绪后启用。 +确认 Environment 已安装，且 **MoGe-2 ViT-S Normal**、**BEN2 Base** 与 **Real-ESRGAN General WDN x4v3** 都显示 Ready。正式功能会等默认必需模型全部就绪后启用。 ### 透明边缘需要调整 diff --git a/docs/server.md b/docs/server.md index 3b5c088..0e38830 100644 --- a/docs/server.md +++ b/docs/server.md @@ -29,11 +29,11 @@ BlendJob 提供 Runtime、HTTP transport、队列、状态和取消协议。AnyI ~/.anyimage/ ├── .venv/ Python Environment ├── models/ ONNX 模型 -├── jobs/ 任务产物 +├── jobs/ 任务产物与 job.log ├── environment.json 环境声明 ├── manifest.json 已安装环境指纹 ├── install.log 安装日志 -└── server.log Server 日志 +└── server.log 带 Job ID 的任务主线与 Server 日志 ``` 更改 Storage Root 相当于切换整套 Environment、模型和任务数据。目录说明见[模型、设备与存储](models-and-storage.md)。 @@ -48,8 +48,7 @@ BlendJob 提供 Runtime、HTTP transport、队列、状态和取消协议。AnyI | `server/models/` | ONNX Session 与模型 Prediction | | `server/geometry/` | 投影、网格和 Surface 数据转换 | | `server/outputs/` | NPZ、EXR、PNG 与序列文件 | -| `model_catalog.py` | Blender UI 的模型声明与本地就绪检查 | -| `server/model_catalog.py` | Server 的模型下载与推理声明 | +| `server/model_catalog.py` | Blender UI 与 Server 共用的模型声明 | 调用方向从 Job 指向媒体、模型、几何和输出模块。模型层返回内存数据，输出层管理文件，Blender 数据始终由 Operator 的 Response 创建。 @@ -68,14 +67,15 @@ Image Box 与 Lasso 的 `Refine Selection` 使用 `refine-image-mask` Job。它 - MoGe-2 ViT-S、ViT-B、ViT-L Normal - DA3 Small、DA3 Base - BEN2 Base +- Real-ESRGAN General WDN x4v3、SAFMN Real x4、Real-ESRGAN x4plus -Debug 任务另外支持 Real-ESRGAN General x4v3；它只提供模型直通测试，不属于正式功能。 +Upscale 使用上述三档 4× RGB 模型。正式 `upscale-image` Job 把结果写成 `upscale.0001.png` 起的序列并回写为 `.upscale` Image 数据块；对应 Debug Job 保留模型直通输出，便于横向比较。 -首次安装会下载 MoGe-2 ViT-S Normal 与 BEN2 Base。Real-ESRGAN General x4v3 由扩展偏好设置按需下载。模型保存在 `models/`，可由多个 Blender 版本共享。Server 缓存 BEN2、Real-ESRGAN 和最近使用的一个深度模型，以控制内存占用。 +首次安装会下载 MoGe-2 ViT-S Normal、BEN2 Base 与默认的 Real-ESRGAN General WDN x4v3。SAFMN Real x4 和 Real-ESRGAN x4plus 可在扩展偏好设置按需下载。模型保存在 `models/`，可由多个 Blender 版本共享。Server 会复用连续任务所需的同一 Session；切换模型族时释放旧 Session，推理资源异常时立即卸载当前高清 Session。 -模型下载优先使用 Hugging Face，并以 OmooLab 镜像作为回退。AnyImage 用同一个准备 Operator 封装 BlendJob 安装任务，并按状态显示 `Set Up AI Environment…` 或 `Download Required Models…`；`Use Mirror` 属性只影响 Environment Bootstrap，不改变模型下载策略。Blender UI 通过同一个状态函数检查 Environment 指纹与默认模型文件：缺 Environment 时执行完整安装，只有默认模型缺失时仅提交补齐模型任务。 +模型按目录声明的 Hugging Face 仓库下载，并以 OmooLab 镜像作为回退；未声明 Hugging Face 仓库的模型直接使用镜像。AnyImage 用同一个准备 Operator 封装 BlendJob 安装任务，并按状态显示 `Set Up AI Environment…` 或 `Download Required Models…`。 -`debug-realesrgan-general-x4v3` 将每张输入的完整 RGB 动态尺寸张量直接送入 ONNX，并把 4× 返回保存为 `prediction*.png`。该 Job 不实现 tile、强度混合、Alpha 合成或 Blender 结果导入。 +`upscale-image` 与 `debug-upscale` 共用模型推理。三个模型在模型层分块处理任意宽高输入；正式 Job 还应用输出上限、保留 Alpha 并导入 Blender。 ## 任务生命周期 @@ -83,6 +83,8 @@ Server 在第一次提交 Job 或手动 Start 时启动。任务进入单 Worker Job 被接受时，Runtime 创建 opaque ID 和对应目录。进度保存在 Server 内存中，Blender 定期轮询状态。取消会在下载分块或一次 ONNX 调用结束后的检查点生效。 +`server.log` 记录带时间、Job ID 与 Job Type 的排队、开始、成功、取消和失败主线，以及 Server 自身异常。Handler 细节与完整 traceback 写入对应 `jobs//job.log`，因此可先从主线定位 Job ID，再查看单个任务。 + Blender 正常退出或禁用扩展时会停止所属 Server；Server 也监视父进程，在异常退出后自行结束。 Environment 安装发生在 Server 启动前，使用 `install.log` 记录过程。安装完成后的模型下载是普通 Job，与推理共享队列、进度和取消机制。 diff --git a/docs/surface.md b/docs/surface.md index f264b06..c07f30b 100644 --- a/docs/surface.md +++ b/docs/surface.md @@ -57,7 +57,7 @@ DA3 使用 `O DA3 Surface`，MoGe-2 使用 `O MoGe2 Surface`。两者采用相 ## 材质与文件 -颜色统一写为 `color.0001.*` 并导入为 `.color` Image。MoGe-2 还会输出 `normal.png`，用于材质法线。材质名称以源 Image datablock 为基名。 +颜色统一写为 `image.0001.*`，导入时沿用源 Image 名称，由 Blender 自动解决同名冲突。MoGe-2 还会输出 `normal.png`，用于材质法线。材质名称以源 Image datablock 为基名。 单图结果可以 Pack 到 `.blend`。模型原始结果与 `mesh.npz` 保留在 Job 目录中，便于调试和重新检查。 diff --git a/docs/upscale.md b/docs/upscale.md new file mode 100644 index 0000000..68b9663 --- /dev/null +++ b/docs/upscale.md @@ -0,0 +1,23 @@ +# 高清放大 + +Upscale 用 ONNX 模型把 Image Empty 的图片放大 4 倍，不创建网格。它支持单图、Movie 和 Image Sequence，并保留源 Alpha。 + +## 使用 + +1. 在 AnyImage 扩展偏好设置中选择 **Default Upscale Model**。 +2. 按需要设置 **Maximum Upscale Size**；默认 4096 px，最高 8192 px。 +3. 选中 Image Empty，在右键 **AnyImage** 菜单中点击 **Upscale**。 + +较大的输入会先按比例缩小，使输出不超过尺寸和总像素上限。完成后，当前 Empty 会改用名称以 `.upscale` 结尾的结果图片。 + +## 模型选择 + +- **Real-ESRGAN General WDN x4v3**：默认模型，速度快，细节处理较保守。 +- **SAFMN Real x4**：速度与细节较均衡。 +- **Real-ESRGAN x4plus**：生成的细节更明显。 + +后两项可在扩展偏好设置中按需下载。 + +对高清结果再次执行 Upscale 时，AnyImage 会从最初的原图重新生成并替换当前结果，不会连续放大已有高清图。 + +如果显存或内存不足，请降低 **Maximum Upscale Size**，或改用 CPU。 diff --git a/docs/workflows.md b/docs/workflows.md index 74ba53e..291ef52 100644 --- a/docs/workflows.md +++ b/docs/workflows.md @@ -1,6 +1,6 @@ # 选择转换方式 -AnyImage 提供多种日常工作流。选择结果时，先考虑你需要的是普通图片平面、可调节的“深度面片”、带透视结构的“空间网格”，还是只需要透明前景。 +AnyImage 提供多种日常工作流。选择结果时，先考虑你需要的是普通图片平面、可调节的“深度面片”、带透视结构的“空间网格”、透明前景，还是 4× 高清图片。 | 工作流 | 适合场景 | 输入 | 结果 | 主要优点 | | --- | --- | --- | --- | --- | @@ -8,6 +8,7 @@ AnyImage 提供多种日常工作流。选择结果时，先考虑你需要的 | Depth Plane | 人物、产品、插画、动态背景、快速场景搭建 | 图片、视频、序列帧 | 带深度位移的平面 | 快速、可动画、Geometry Nodes 参数直观 | | Surface | 建筑、室内、街景、需要匹配视角的单图 | 单张图片 | 透视纹理网格 | 空间感更强，可创建匹配相机 | | Remove Background | 抠出人物或物体，继续作为 Image Empty 使用 | 图片、视频、序列帧 | 带透明 Alpha 的颜色图 | 不创建网格，保留 Empty 的位置和显示设置 | +| Upscale | 摄影、插画、二次元等素材的 4× 高清放大 | 图片、视频、序列帧 | 4× RGB 颜色图 | 三档 ONNX 模型可选，直接替换 Empty 图片 | | Image Box / Lasso / Line / Polyline | 用几何选区直接清理图片 | 静态图片 | 原位更新的透明裁切图 | `F` 反选，可用 BEN2 细化选区 | | Image Perspective | 海报、屏幕、门窗等四边形内容的透视矫正裁切 | 静态图片 | 原位更新的矩形图片 | 四点任意顺序，交互调整比例 | @@ -35,6 +36,11 @@ Surface 会根据模型估算的相机和深度构建透视网格。相较 Depth [阅读去除背景指南](remove-background.md) +## 什么时候选择 Upscale + +需要增加像素尺寸而不创建网格时，选择 **Upscale**。它支持单图、Movie 和 Image Sequence；模型选择、输出限制与重复执行行为见[高清放大指南](upscale.md)。 + + ## 什么时候使用 Image Perspective 当图片中的矩形物体因为透视变成四边形，并且希望得到一张正视矩形裁切图时，使用 3D Viewport 工具栏中的 **Image Perspective**。它只把四点当作图片裁切与 Homography 输入，不假设当前 View 的焦距与原图拍摄相机一致。四点顺序不影响输出方向；结果 Empty 保持源图 Transform，并按源图中的选区范围近似更新显示位置与尺寸。 diff --git a/mkdocs.yml b/mkdocs.yml index f513eb5..76e7990 100644 --- a/mkdocs.yml +++ b/mkdocs.yml @@ -20,6 +20,7 @@ nav: - Depth Plane: guide-depth-plane.md - Surface: guide-surface.md - 去除背景: remove-background.md + - 高清放大: upscale.md - 视频与序列帧: animation.md - 模型、设备与存储: models-and-storage.md - 常见问题: faq.md diff --git a/pack.py b/pack.py index ee0921d..6bd9db6 100644 --- a/pack.py +++ b/pack.py @@ -48,6 +48,8 @@ def download_wheels(platform, destination): *dependencies, "--dest", str(destination), + "--find-links", + str(SOURCE_DIR / "wheels"), "--only-binary=:all:", f"--python-version={PYTHON_VERSION}", f"--platform={config.pip_platform}", diff --git a/pyproject.toml b/pyproject.toml index b067e0b..e7e2029 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" license = "GPL-3.0-or-later" -dependencies = ["blendjob==0.1.13"] +dependencies = ["blendjob==0.1.14"] [dependency-groups] dev = [ diff --git a/src/anyimage/__init__.py b/src/anyimage/__init__.py index 212e44c..e21a589 100644 --- a/src/anyimage/__init__.py +++ b/src/anyimage/__init__.py @@ -21,8 +21,8 @@ from .panel import ( Da3OutputDebugPanel, DebugPanel, Moge2ApiDebugPanel, - RealesrganGeneralX4v3DebugPanel, ServerPanel, + UpscaleDebugPanel, ) @@ -37,7 +37,7 @@ CLASSES = ( Moge2ApiDebugPanel, Da3OutputDebugPanel, Ben2RemovalDebugPanel, - RealesrganGeneralX4v3DebugPanel, + UpscaleDebugPanel, ) @@ -72,8 +72,11 @@ def register(): bpy.types.Scene.anyimage_settings = bpy.props.PointerProperty( type=AnyImageSettings, ) - bpy.types.VIEW3D_MT_object_context_menu.append(draw_image_context_menu) - bpy.types.OUTLINER_MT_object.append(draw_image_context_menu) + bpy.types.Image.anyimage_upscale_source = bpy.props.PointerProperty( + type=bpy.types.Image, + ) + bpy.types.VIEW3D_MT_object_context_menu.prepend(draw_image_context_menu) + bpy.types.OUTLINER_MT_object.prepend(draw_image_context_menu) clipboard_image.register_keymaps() runtime.register() @@ -91,5 +94,7 @@ def unregister(): bpy.types.OUTLINER_MT_object.remove(draw_image_context_menu) if hasattr(bpy.types.Scene, "anyimage_settings"): del bpy.types.Scene.anyimage_settings + if hasattr(bpy.types.Image, "anyimage_upscale_source"): + del bpy.types.Image.anyimage_upscale_source for class_type in reversed(CLASSES): bpy.utils.unregister_class(class_type) diff --git a/src/anyimage/blender_manifest.toml b/src/anyimage/blender_manifest.toml index aa176ee..0e8d0b1 100644 --- a/src/anyimage/blender_manifest.toml +++ b/src/anyimage/blender_manifest.toml @@ -10,7 +10,7 @@ type = "add-on" blender_version_min = "4.5.0" platforms = ["windows-x64", "macos-arm64", "linux-x64"] license = ["SPDX:GPL-3.0-or-later"] -wheels = ["./wheels/blendjob-0.1.13-py3-none-any.whl"] +wheels = ["./wheels/blendjob-0.1.14-py3-none-any.whl"] [permissions] network = "Access the local Job Server and download environment dependencies and model weights" diff --git a/src/anyimage/menu.py b/src/anyimage/menu.py index ce7631f..446aaa3 100644 --- a/src/anyimage/menu.py +++ b/src/anyimage/menu.py @@ -7,8 +7,11 @@ class AnyImageImageMenu(bpy.types.Menu): bl_idname = "ANYIMAGE_MT_image_context" bl_label = "AnyImage" - def draw(self, _context): - from .properties import ai_setup_label, ai_status + def draw(self, context): + from .properties import ( + ai_setup_label, + ai_status, + ) layout = self.layout layout.operator_context = "INVOKE_DEFAULT" @@ -27,11 +30,17 @@ class AnyImageImageMenu(bpy.types.Menu): text="Convert to Surface…", icon="MESH_GRID", ) + layout.separator() layout.operator( "anyimage.remove_image_background", text="Remove Background", icon="IMAGE_ALPHA", ) + layout.operator( + "anyimage.upscale_image", + text=upscale_menu_label(context), + icon="FULLSCREEN_ENTER", + ) layout.separator() status = ai_status() if status["ready"]: @@ -48,12 +57,24 @@ class AnyImageImageMenu(bpy.types.Menu): ) +def upscale_menu_label(context): + source_object = getattr(context, "object", None) + if not is_image_empty(source_object): + return "Upscale" + width, height = ( + int(value) for value in getattr(source_object.data, "size", (0, 0)) + ) + if width <= 0 or height <= 0: + return "Upscale" + return f"Upscale ({width} × {height})" + + def draw_image_context_menu(self, context): if not is_image_empty(context.object): return - self.layout.separator() self.layout.menu( AnyImageImageMenu.bl_idname, text="AnyImage", icon="PLUGIN", ) + self.layout.separator() diff --git a/src/anyimage/model_catalog.py b/src/anyimage/model_catalog.py deleted file mode 100644 index f63f039..0000000 --- a/src/anyimage/model_catalog.py +++ /dev/null @@ -1,219 +0,0 @@ -from dataclasses import dataclass - - -@dataclass(frozen=True) -class ModelSpec: - key: str - label: str - directory_name: str - repository: str - license_name: str - description: str - supports_gs: bool - download_patterns: tuple[str, ...] = () - ready_patterns: tuple[str, ...] = ("config.json", "*.safetensors") - enabled: bool = True - r2_directory: str = "" - r2_files: tuple[tuple[str, int, str], ...] = () - - -MODELS = { - "SMALL": ModelSpec( - key="SMALL", - label="DA3 Small", - directory_name="DA3-SMALL", - repository="depth-anything/DA3-SMALL", - license_name="Apache 2.0", - description="Fastest any-view model using ONNX Runtime", - supports_gs=False, - download_patterns=("model.onnx",), - ready_patterns=("model.onnx",), - r2_directory="DA3-SMALL", - r2_files=( - ("model.onnx", 105484029, "7a9d81475b3e2334ace0d6d2c2a816b3c93081437e1866cc434ce1789125f631"), - ), - ), - "BASE": ModelSpec( - key="BASE", - label="DA3 Base", - directory_name="DA3-BASE", - repository="depth-anything/DA3-BASE", - license_name="Apache 2.0", - description="Balanced any-view model using ONNX Runtime", - supports_gs=False, - download_patterns=("model.onnx",), - ready_patterns=("model.onnx",), - r2_directory="DA3-BASE", - r2_files=( - ("model.onnx", 412856697, "31f38fc4990539d7902ba267924b9731db1acb00d840c248a60b1865810ff9e6"), - ), - ), - "GIANT_1_1": ModelSpec( - key="GIANT_1_1", - label="DA3 Giant 1.1", - directory_name="da3_giant_1_1", - repository="depth-anything/DA3-GIANT-1.1", - license_name="CC BY-NC 4.0", - description="Non-commercial Giant model with Gaussian head", - supports_gs=True, - enabled=False, - ), - "NESTED_GIANT_LARGE_1_1": ModelSpec( - key="NESTED_GIANT_LARGE_1_1", - label="DA3 Nested Giant Large 1.1", - directory_name="da3_nested_giant_large_1_1", - repository="depth-anything/DA3NESTED-GIANT-LARGE-1.1", - license_name="CC BY-NC 4.0", - description="Non-commercial nested model with metric depth and Gaussian head", - supports_gs=True, - enabled=False, - ), -} - -BEN2_MODEL = ModelSpec( - key="BEN2_BASE", - label="BEN2 Base", - directory_name="BEN2-ONNX", - repository="onnx-community/BEN2-ONNX", - license_name="Apache 2.0", - description="Fast high-quality FP16 ONNX background removal model", - supports_gs=False, - download_patterns=("onnx/model_fp16.onnx",), - ready_patterns=("onnx/model_fp16.onnx",), - r2_directory="BEN2-ONNX", - r2_files=(( - "onnx/model_fp16.onnx", - 219121675, - "dfdc25f421f32a0d1268e0f2ff2153d340e8f1d52d3dd16f5dc33c1ce85cedf1", - ),), -) - - -REALESRGAN_MODEL = ModelSpec( - key="REALESRGAN_GENERAL_X4V3", - label="Real-ESRGAN General x4v3", - directory_name="realesr-general-x4v3", - repository="jonathanst29/tinier-upscale-models", - license_name="BSD 3-Clause", - description="Fast general-purpose 4x image upscaler using ONNX Runtime", - supports_gs=False, - download_patterns=("realesr-general-x4v3.onnx",), - ready_patterns=("realesr-general-x4v3.onnx",), - r2_directory="realesr-general-x4v3", - r2_files=(( - "realesr-general-x4v3.onnx", - 4866396, - "924ebad6532777303582d4ce7811849b88869231a3ae7093e0f21200249df8d5", - ),), -) - - -UPSCALE_MODELS = { - REALESRGAN_MODEL.key: REALESRGAN_MODEL, -} - - -MOGE2_MODELS = { - "MOGE2_VITS_NORMAL": ModelSpec( - key="MOGE2_VITS_NORMAL", - label="MoGe-2 ViT-S Normal", - directory_name="moge-2-vits-normal-onnx", - repository="Ruicheng/moge-2-vits-normal-onnx", - license_name="MIT", - description="Fastest MoGe-2 model using ONNX Runtime", - supports_gs=False, - download_patterns=("model.onnx",), - ready_patterns=("model.onnx",), - r2_directory="moge-2-vits-normal-onnx", - r2_files=(( - "model.onnx", - 140852051, - "24eacb5dc7a2c54c7bc98f7de085ffbed79ad006ea5b664c2c2cdc02ff3a52f0", - ),), - ), - "MOGE2_VITB_NORMAL": ModelSpec( - key="MOGE2_VITB_NORMAL", - label="MoGe-2 ViT-B Normal", - directory_name="moge-2-vitb-normal-onnx", - repository="Ruicheng/moge-2-vitb-normal-onnx", - license_name="MIT", - description="Balanced MoGe-2 model using ONNX Runtime", - supports_gs=False, - download_patterns=("model.onnx",), - ready_patterns=("model.onnx",), - r2_directory="moge-2-vitb-normal-onnx", - r2_files=(( - "model.onnx", - 419411850, - "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", - description="Highest-quality MoGe-2 model using ONNX Runtime", - supports_gs=False, - download_patterns=("model.onnx",), - ready_patterns=("model.onnx",), - r2_directory="moge-2-vitl-normal-onnx", - r2_files=(( - "model.onnx", - 1324265014, - "afbc4ccc3450298f3afb35b90f015f4c4f552dea21dc6470d5f7b78b77e2d751", - ),), - ), -} -MOGE2_MODEL = MOGE2_MODELS["MOGE2_VITS_NORMAL"] -DOWNLOADABLE_MODELS = { - **MODELS, - BEN2_MODEL.key: BEN2_MODEL, - **UPSCALE_MODELS, - **MOGE2_MODELS, -} -MODEL_NUMBERS = { - model_key: number for number, model_key in enumerate(MODELS) -} - - -def get_model(model_key): - try: - return MODELS[model_key] - except KeyError as error: - raise ValueError(f"Unknown DA3 model: {model_key}") from error - - -def get_downloadable_model(model_key): - try: - return DOWNLOADABLE_MODELS[model_key] - except KeyError as error: - raise ValueError(f"Unknown model: {model_key}") from error - - -def model_directory_name(model_key): - """Return the model's folder name inside the shared models directory.""" - return get_downloadable_model(model_key).directory_name - - -def model_record(model, ready=False): - """Return public model metadata exposed by the Models Resource.""" - return { - "key": model.key, - "label": model.label, - "description": model.description, - "license": model.license_name, - "supports_gs": model.supports_gs, - "enabled": model.enabled, - "family": ( - "moge2" - if model.key.startswith("MOGE2_") - else "ben2" - if model.key == BEN2_MODEL.key - else "upscale" - if model.key in UPSCALE_MODELS - else "da3" - ), - "ready": bool(ready), - } diff --git a/src/anyimage/operators/__init__.py b/src/anyimage/operators/__init__.py index 274e9a5..2e8d2ff 100644 --- a/src/anyimage/operators/__init__.py +++ b/src/anyimage/operators/__init__.py @@ -4,12 +4,7 @@ from .balloon import ( ExtractBalloon, AnalyzeBalloon, ) -from .debug import ( - DebugBen2, - DebugDa3, - DebugMoge2, - DebugRealesrganGeneralX4v3, -) +from .debug import DebugBen2, DebugDa3, DebugMoge2, DebugUpscale from .image_edit import CLASSES as IMAGE_EDIT_CLASSES from .depth_plane import ( ConvertToDepthPlane, @@ -27,6 +22,7 @@ from .models import ( ) from .mesh_plane import ConvertToMeshPlane from .surface import ConvertToSurface, GenerateDa3Surface, GenerateMoge2Surface +from .upscale import UpscaleImage CLASSES = ( @@ -37,7 +33,7 @@ CLASSES = ( DebugDa3, DebugMoge2, DebugBen2, - DebugRealesrganGeneralX4v3, + DebugUpscale, GenerateDa3Depth, GenerateMoge2Depth, GenerateDa3Surface, @@ -51,6 +47,7 @@ CLASSES = ( AnalyzeBalloon, ActivateBalloonTool, RemoveImageBackground, + UpscaleImage, ConvertToMeshPlane, ConvertToDepthPlane, ConvertToSurface, diff --git a/src/anyimage/operators/background.py b/src/anyimage/operators/background.py index 0a5c4fc..23e1bf1 100644 --- a/src/anyimage/operators/background.py +++ b/src/anyimage/operators/background.py @@ -79,7 +79,14 @@ def apply_background_response(result, source_object_name): source_object = bpy.data.objects.get(source_object_name) if not is_image_empty(source_object): raise RuntimeError("The source image empty is no longer available") - create_color_image(source_object, result.directory) + source_image = source_object.data + source_name = source_image.name + original_fake_user = source_image.use_fake_user + result_image = create_color_image(source_object, result.directory) + source_image.use_fake_user = original_fake_user + if not original_fake_user and int(source_image.users) == 0: + bpy.data.images.remove(source_image) + result_image.name = source_name push_undo("Remove Background") diff --git a/src/anyimage/operators/balloon/operator.py b/src/anyimage/operators/balloon/operator.py index f88eea8..1ef43c4 100644 --- a/src/anyimage/operators/balloon/operator.py +++ b/src/anyimage/operators/balloon/operator.py @@ -151,24 +151,6 @@ def data_name_matches(name, base_name): return suffix != name and suffix.isdigit() -def material_uses_balloon_texture(material, color_name): - node_tree = getattr(material, "node_tree", None) - if node_tree is None: - return False - has_texture_layer = any( - getattr(node, "type", None) == "GROUP" - and getattr(getattr(node, "node_tree", None), "name", None) - == TEXTURE_MATERIAL_NODE_GROUP_NAME - for node in node_tree.nodes - ) - has_color_image = any( - getattr(node, "type", None) == "TEX_IMAGE" - and getattr(node, "image", None) is not None - and data_name_matches(node.image.name, color_name) - for node in node_tree.nodes - ) - return has_texture_layer and has_color_image - def marked_balloon_material(source_image, normal_required=False): return next( @@ -194,30 +176,8 @@ def mark_balloon_material(material, source_image, normal_image=None): material[BALLOON_NORMAL_IMAGE_PROPERTY] = normal_image -def reusable_balloon_material( - material_name, - color_name, - source_image=None, - normal_required=False, -): - marked = ( - marked_balloon_material(source_image, normal_required) - if source_image is not None - else None - ) - if marked is not None: - return marked - if normal_required: - return None - return next( - ( - material - for material in bpy.data.materials - if data_name_matches(material.name, material_name) - and material_uses_balloon_texture(material, color_name) - ), - None, - ) +def reusable_balloon_material(source_image, normal_required=False): + return marked_balloon_material(source_image, normal_required) def load_balloon_normal_image(path, source_object): @@ -331,8 +291,6 @@ def extract_balloon_object( uses_normal = normal_path is not None material = ( reusable_balloon_material( - material_name, - color_name, image, normal_required=uses_normal, ) diff --git a/src/anyimage/operators/common.py b/src/anyimage/operators/common.py index d7f41bb..27b882b 100644 --- a/src/anyimage/operators/common.py +++ b/src/anyimage/operators/common.py @@ -4,6 +4,7 @@ from pathlib import Path import bpy from .image import extract_frame_range +from ..preferences import ADDON_PACKAGE, addon_preferences from ..runtime import runtime from ..properties import ai_ready, model_ready, model_record @@ -13,7 +14,6 @@ BEN2_VIDEO_SUFFIXES = { ".mp4", ".mov", ".avi", ".mkv", ".webm", ".m4v", ".mts", ".mpg", ".mpeg" } HIDDEN = {"HIDDEN", "SKIP_SAVE"} -ADDON_PACKAGE = __package__.partition(".")[0] def input_path(value): @@ -40,8 +40,7 @@ def source_start_frame(source_object_name): def production_device(): - addon = bpy.context.preferences.addons.get(ADDON_PACKAGE) - return getattr(getattr(addon, "preferences", None), "device", "AUTO") + return getattr(addon_preferences(), "device", "AUTO") def require_environment(): diff --git a/src/anyimage/operators/debug.py b/src/anyimage/operators/debug.py index 238e290..0f30112 100644 --- a/src/anyimage/operators/debug.py +++ b/src/anyimage/operators/debug.py @@ -1,10 +1,7 @@ import bpy from ..runtime import JobOperatorBase -from ..properties import ( - BEN2_MODEL_KEY, - REALESRGAN_MODEL_KEY, -) +from ..properties import BEN2_MODEL_KEY, REALESRGAN_GENERAL_WDN_MODEL_KEY from .common import ( background_input_path, da3_parameters, @@ -94,19 +91,24 @@ class DebugBen2(JobOperatorBase, bpy.types.Operator): return f"Removed background to {result.directory}" -class DebugRealesrganGeneralX4v3(JobOperatorBase, bpy.types.Operator): - bl_idname = "anyimage.debug_realesrgan_general_x4v3" - bl_label = "Debug Real-ESRGAN General x4v3" - job_type = "debug-realesrgan-general-x4v3" - starting_message = "Starting Real-ESRGAN General x4v3 inference" +class DebugUpscale(JobOperatorBase, bpy.types.Operator): + bl_idname = "anyimage.debug_upscale" + bl_label = "Debug Upscale" + job_type = "debug-upscale" + starting_message = "Starting Upscale inference" + + model: bpy.props.StringProperty( + default=REALESRGAN_GENERAL_WDN_MODEL_KEY, + options={"HIDDEN", "SKIP_SAVE"}, + ) def request(self, context): require_environment() settings = context.scene.anyimage_settings - require_model(REALESRGAN_MODEL_KEY) + require_model(self.model) return { "input": str(input_path(settings.input_path)), - "model": REALESRGAN_MODEL_KEY, + "model": self.model, "device": settings.device.lower(), "video_frames": settings.video_frames, "start_frame": 1, @@ -114,5 +116,4 @@ class DebugRealesrganGeneralX4v3(JobOperatorBase, bpy.types.Operator): def response(self, _context, result): open_job_directory(result.directory) - return f"Saved Real-ESRGAN API prediction to {result.directory}" - + return f"Saved upscale API prediction to {result.directory}" diff --git a/src/anyimage/operators/depth_plane/build.py b/src/anyimage/operators/depth_plane/build.py index e439e5b..024a30e 100644 --- a/src/anyimage/operators/depth_plane/build.py +++ b/src/anyimage/operators/depth_plane/build.py @@ -267,10 +267,10 @@ def build_image_depth_plane( frames_image = create_sequence_image( depth_data_name(source_object), depth_frames_directory, - frame_prefix="color.", + frame_prefix="image.", ) frame_count = len( - sorted_frame_paths(depth_frames_directory, "color.") + sorted_frame_paths(depth_frames_directory, "image.") ) ( source_frame_start, diff --git a/src/anyimage/operators/depth_plane/data.py b/src/anyimage/operators/depth_plane/data.py index a3c18df..9d8d2f2 100644 --- a/src/anyimage/operators/depth_plane/data.py +++ b/src/anyimage/operators/depth_plane/data.py @@ -9,7 +9,7 @@ INTRINSICS_KEY_CANDIDATES = ("intrinsics", "ixts") def create_depth_sequence_image( name, depth_frames, - scene, + _scene, frame_directory, frame_prefix="depth", ): @@ -27,11 +27,10 @@ def create_depth_sequence_image( float_buffer=True, ) image.colorspace_settings.name = "Non-Color" - image_settings = scene.render.image_settings - original_file_format = image_settings.file_format + original_file_format = image.file_format frame_paths = [] try: - image_settings.file_format = "OPEN_EXR" + image.file_format = "OPEN_EXR" for frame_index, depth_values in enumerate(depth_frames): # NPZ arrays use a top-left origin, while Blender image pixels start below. depth_pixels = np.flipud(depth_values) @@ -44,12 +43,12 @@ def create_depth_sequence_image( frame_path = ( frame_directory / f"{frame_prefix}.{frame_index + 1:04d}.exr" ) - image.save_render(str(frame_path), scene=scene) + image.save(filepath=str(frame_path), save_copy=True) if not frame_path.is_file(): raise RuntimeError(f"Blender did not write: {frame_path}") frame_paths.append(frame_path) finally: - image_settings.file_format = original_file_format + image.file_format = original_file_format bpy.data.images.remove(image, do_unlink=True) if not frame_paths: raise RuntimeError("Unable to write the depth frame sequence") diff --git a/src/anyimage/operators/image.py b/src/anyimage/operators/image.py index 891345f..6c51437 100644 --- a/src/anyimage/operators/image.py +++ b/src/anyimage/operators/image.py @@ -155,8 +155,8 @@ def copy_depth_image(source_object): return depth_image -def color_data_name(source_object): - return f"{image_base_name(source_object.data)}.color" +def color_data_name(source_object, suffix=".color"): + return f"{image_base_name(source_object.data)}{suffix}" def material_data_name(source_object): @@ -204,18 +204,17 @@ def draw_process_resolution(layout, owner, image): ) -def save_packed_image(image, scene): +def save_packed_image(image, _scene): output_path = Path(tempfile.mkdtemp(prefix="anyimage-input-")) / "input.png" - image_settings = scene.render.image_settings - original_file_format = image_settings.file_format + original_file_format = image.file_format try: - image_settings.file_format = "PNG" - image.save_render(str(output_path), scene=scene) + image.file_format = "PNG" + image.save(filepath=str(output_path), save_copy=True) except (OSError, RuntimeError) as error: output_path.unlink(missing_ok=True) raise RuntimeError("Unable to export the packed image") from error finally: - image_settings.file_format = original_file_format + image.file_format = original_file_format if not output_path.is_file(): raise RuntimeError("Blender did not export the packed image") return output_path @@ -377,14 +376,17 @@ def create_color_image( source_object, output_directory, assign_to_source=True, + *, + frame_prefix="image.", ): output_directory = Path(output_directory) - color_paths = sorted_frame_paths(output_directory, "color.") + color_paths = sorted_frame_paths(output_directory, frame_prefix) if not color_paths: raise RuntimeError( f"Color result has no images: {output_directory}" ) source_image = source_object.data + source_color_space = source_image.colorspace_settings.name animated = is_animated_image(source_image) try: result_image = bpy.data.images.load( @@ -392,16 +394,17 @@ def create_color_image( check_existing=False, ) result_image.alpha_mode = "CHANNEL_PACKED" + result_image.colorspace_settings.name = source_color_space if animated: force_sequence_source(result_image) - except (OSError, RuntimeError) as error: + except (OSError, RuntimeError, TypeError, ValueError) as error: if "result_image" in locals(): bpy.data.images.remove(result_image, do_unlink=True) raise RuntimeError("Unable to load the color result") from error original_fake_user = source_image.use_fake_user try: - result_image.name = color_data_name(source_object) + result_name = color_data_name(source_object) if animated: pack_sequence_if_possible(result_image, len(color_paths)) else: @@ -427,6 +430,7 @@ def create_color_image( source_image.use_fake_user = original_fake_user bpy.data.images.remove(result_image, do_unlink=True) raise RuntimeError("Unable to prepare the color result") from error + result_image.name = result_name return result_image diff --git a/src/anyimage/operators/image_edit/common.py b/src/anyimage/operators/image_edit/common.py index 40aa73b..3aa364b 100644 --- a/src/anyimage/operators/image_edit/common.py +++ b/src/anyimage/operators/image_edit/common.py @@ -13,6 +13,15 @@ from ..common import production_device, require_environment, require_model from ..image import image_base_name, is_image_empty, plane_bounds +class ImageEditWarning(ValueError): + pass + + +def report_image_edit_exception(operator, error): + level = "WARNING" if isinstance(error, ImageEditWarning) else "ERROR" + operator.report({level}, str(error)) + + def image_edit_poll(context): return context.area is not None and context.area.type == "VIEW_3D" @@ -94,7 +103,7 @@ def crop_pixels_to_alpha(pixels, bounds): ) visible_y, visible_x = np.nonzero(rgba[:, :, 3] > 0.0) if not len(visible_x): - raise ValueError("The edit contains no visible image pixels") + raise ImageEditWarning("The edit contains no visible image pixels") content_left = int(visible_x.min()) content_top = int(visible_y.min()) content_right = int(visible_x.max()) + 1 @@ -108,10 +117,10 @@ def crop_pixels_to_alpha(pixels, bounds): ) -def align_empty_display_to_bounds( +def empty_display_alignment( source_object, source_size, output_size, placement_bounds ): - """Place a new image canvas without changing the Empty transform.""" + """Calculate image canvas placement without changing the Empty transform.""" source_width, source_height = (float(value) for value in source_size) output_width, output_height = (float(value) for value in output_size) left, top, right, bottom = (float(value) for value in placement_bounds) @@ -134,21 +143,36 @@ def align_empty_display_to_bounds( display_height = display_size center_x = (desired_x_min + desired_x_max) * 0.5 center_y = (desired_y_min + desired_y_max) * 0.5 - source_object.empty_display_size = display_size - source_object.empty_image_offset = ( + return display_size, ( center_x / display_width - 0.5, center_y / display_height - 0.5, ) + +def write_image_pixels(image, pixels): + image.alpha_mode = "STRAIGHT" + image.pixels.foreach_set(pixels) + image.pack() + image.update() + + def replace_empty_image_pixels( - source_object, pixels, output_size, placement_bounds, suffix=".edit" + source_object, pixels, output_size, placement_bounds, suffix="" ): source_image = source_object.data source_size = tuple(source_image.size) output_width, output_height = (int(value) for value in output_size) if output_width <= 0 or output_height <= 0: raise ValueError("The edited image dimensions must be positive") + if len(pixels) != output_width * output_height * 4: + raise ValueError("The edited pixels do not match the image dimensions") + display_size, image_offset = empty_display_alignment( + source_object, + source_size, + (output_width, output_height), + placement_bounds, + ) image = bpy.data.images.new( f"{image_base_name(source_image)}{suffix}", width=output_width, @@ -157,20 +181,17 @@ def replace_empty_image_pixels( ) try: image.colorspace_settings.name = source_image.colorspace_settings.name - image.alpha_mode = "STRAIGHT" - image.pixels.foreach_set(pixels) - image.pack() - image.update() - align_empty_display_to_bounds( - source_object, - source_size, - (output_width, output_height), - placement_bounds, - ) - source_object.data = image + write_image_pixels(image, pixels) except Exception: bpy.data.images.remove(image, do_unlink=True) raise + source_name = source_image.name + source_object.data = image + if not source_image.use_fake_user and int(source_image.users) == 0: + bpy.data.images.remove(source_image) + image.name = source_name + source_object.empty_display_size = display_size + source_object.empty_image_offset = image_offset return image diff --git a/src/anyimage/operators/image_edit/mask.py b/src/anyimage/operators/image_edit/mask.py index 1b98a51..6a2f329 100644 --- a/src/anyimage/operators/image_edit/mask.py +++ b/src/anyimage/operators/image_edit/mask.py @@ -30,10 +30,12 @@ from .common import ( draw_dashed_line, draw_image_mask_operator_settings, draw_polygon_fill, + ImageEditWarning, image_edit_poll, image_mask_refine_request, polygon_bounds, replace_empty_image_pixels, + report_image_edit_exception, resolve_image_edit_click, ) @@ -141,7 +143,6 @@ def edit_image_object(source_object, selection_mask, invert=False): pixels, (right - left, bottom - top), crop_bounds, - suffix=".mask", ) @@ -389,7 +390,8 @@ class EditImageMask(bpy.types.Operator): else: path = self._screen_polygon() if len(path) < 3: - raise ValueError("Draw a larger image region") + self._finish(context) + return {"CANCELLED"} polygon = screen_path_to_image_pixels( path, context.region, @@ -444,7 +446,7 @@ class EditImageMask(bpy.types.Operator): try: return self._complete(context) except (RuntimeError, TypeError, ValueError) as error: - self.report({"ERROR"}, str(error)) + report_image_edit_exception(self, error) self._finish(context) return {"CANCELLED"} return {"PASS_THROUGH"} @@ -482,7 +484,7 @@ class EditImageMask(bpy.types.Operator): try: return self._complete(context) except (RuntimeError, TypeError, ValueError) as error: - self.report({"ERROR"}, str(error)) + report_image_edit_exception(self, error) self._finish(context) return {"CANCELLED"} @@ -518,7 +520,7 @@ class EditImageMask(bpy.types.Operator): self.invert, ) except (json.JSONDecodeError, RuntimeError, TypeError, ValueError) as error: - self.report({"ERROR"}, str(error)) + report_image_edit_exception(self, error) return {"CANCELLED"} return {"FINISHED"} @@ -559,7 +561,11 @@ class RefineImageMask(JobOperatorBase, bpy.types.Operator): ) finally: bpy.data.images.remove(refined_image, do_unlink=True) - edit_image_object(source_object, mask, self.invert) + try: + edit_image_object(source_object, mask, self.invert) + except ImageEditWarning as error: + self.report({"WARNING"}, str(error)) + return str(error) push_undo("Edit Image Mask") return "Image selection refined" diff --git a/src/anyimage/operators/image_edit/perspective.py b/src/anyimage/operators/image_edit/perspective.py index 9e40b91..b978588 100644 --- a/src/anyimage/operators/image_edit/perspective.py +++ b/src/anyimage/operators/image_edit/perspective.py @@ -24,6 +24,7 @@ from .common import ( image_edit_poll, polygon_bounds, replace_empty_image_pixels, + report_image_edit_exception, resolve_image_edit_click, ) @@ -484,7 +485,6 @@ def create_perspective_result( pixels, output_size, placement_bounds, - suffix=".perspective", ) @@ -592,9 +592,14 @@ class EditImagePerspective(bpy.types.Operator): context.area.tag_redraw() def _remove_preview(self): - if self._preview_image is not None: - bpy.data.images.remove(self._preview_image, do_unlink=True) - self._preview_image = None + preview_image = self._preview_image + self._preview_image = None + if preview_image is None: + return + try: + bpy.data.images.remove(preview_image, do_unlink=True) + except ReferenceError: + pass def _finish(self, context): if self._handle is not None: @@ -771,7 +776,7 @@ class EditImagePerspective(bpy.types.Operator): self._phase = "ASPECT" self._cursor = None except (RuntimeError, TypeError, ValueError) as error: - self.report({"ERROR"}, str(error)) + report_image_edit_exception(self, error) self._finish(context) return {"CANCELLED"} self._set_aspect_status(context) @@ -825,7 +830,7 @@ class EditImagePerspective(bpy.types.Operator): TypeError, ValueError, ) as error: - self.report({"ERROR"}, str(error)) + report_image_edit_exception(self, error) return {"CANCELLED"} return {"FINISHED"} diff --git a/src/anyimage/operators/models.py b/src/anyimage/operators/models.py index 4da3826..f548c02 100644 --- a/src/anyimage/operators/models.py +++ b/src/anyimage/operators/models.py @@ -51,7 +51,8 @@ class SetupAIEnvironment(bpy.types.Operator): ) layout.label( text=( - "It also downloads MoGe-2 ViT-S Normal and BEN2 Base " + "It also downloads MoGe-2 ViT-S Normal, BEN2 Base and " + "Real-ESRGAN General WDN x4v3 " f"({format_download_size(required_size)})." ) ) diff --git a/src/anyimage/operators/upscale.py b/src/anyimage/operators/upscale.py new file mode 100644 index 0000000..832666e --- /dev/null +++ b/src/anyimage/operators/upscale.py @@ -0,0 +1,158 @@ +import bpy + +from ..preferences import ( + DEFAULT_UPSCALE_MAX_OUTPUT_SIZE, + MAX_UPSCALE_OUTPUT_SIZE, +) +from ..properties import ( + configured_upscale_max_output_size, + configured_upscale_model, + upscale_model_label, +) +from ..runtime import JobOperatorBase +from .common import ( + HIDDEN, + cleanup_temporary_input, + input_path, + invoke_ai_setup_if_needed, + production_device, + require_environment, + require_model, + source_start_frame, +) +from .image import ( + create_color_image, + default_extract_frames, + image_base_name, + is_image_empty, + prepare_image_input, + video_total_frames, +) + + +class UpscaleImage(JobOperatorBase, bpy.types.Operator): + bl_idname = "anyimage.upscale_image" + bl_label = "Upscale" + bl_description = "Upscale this Image Empty by 4x with an AI model" + bl_options = {"UNDO"} + job_type = "upscale-image" + starting_message = "Starting Image Empty upscale" + + model: bpy.props.StringProperty( + default="REALESRGAN_GENERAL_WDN_X4V3", + options=HIDDEN, + ) + max_output_size: bpy.props.IntProperty( + name="Max Output Size", + description=( + "Maximum output long edge; larger inputs are resized before 4x " + "upscaling to keep memory use bounded" + ), + default=DEFAULT_UPSCALE_MAX_OUTPUT_SIZE, + min=1024, + max=MAX_UPSCALE_OUTPUT_SIZE, + step=256, + subtype="PIXEL", + options=HIDDEN, + ) + input_path: bpy.props.StringProperty(options=HIDDEN) + source_object_name: bpy.props.StringProperty(options=HIDDEN) + source_image_name: bpy.props.StringProperty(options=HIDDEN) + delete_input: bpy.props.BoolProperty(default=False, options=HIDDEN) + + @classmethod + def poll(cls, context): + from ..runtime import runtime + + return is_image_empty(context.object) and not runtime.server_busy() + + def invoke(self, context, _event): + setup_result = invoke_ai_setup_if_needed() + if setup_result is not None: + return setup_result + return self.execute(context) + + def execute(self, context): + self.model = configured_upscale_model(context) + self.max_output_size = configured_upscale_max_output_size(context) + source_object = context.object + if not is_image_empty(source_object): + self.report({"ERROR"}, "Select an Image Empty to upscale") + return {"CANCELLED"} + current_image = source_object.data + try: + original_image = getattr( + current_image, + "anyimage_upscale_source", + None, + ) + except ReferenceError: + original_image = None + source_image = original_image or current_image + self.source_image_name = source_image.name + self._rebuild_from_original = source_image is not current_image + if self._rebuild_from_original: + self.report( + {"WARNING"}, + "This image was already upscaled; rebuilding from the original", + ) + if video_total_frames(source_image) > 1: + context.scene.anyimage_settings.video_frames = ( + default_extract_frames(source_object) + ) + try: + source_path, delete_input = prepare_image_input( + source_image, + context.scene, + ) + except RuntimeError as error: + self.report({"ERROR"}, str(error)) + return {"CANCELLED"} + + self.input_path = str(source_path) + self.source_object_name = source_object.name + self.delete_input = delete_input + result = super().execute(context) + if "CANCELLED" in result: + cleanup_temporary_input(self.input_path, self.delete_input) + return result + + def request(self, context): + require_environment() + require_model(self.model) + settings = context.scene.anyimage_settings + return { + "input": str(input_path(self.input_path)), + "model": self.model, + "device": production_device().lower(), + "video_frames": settings.video_frames, + "start_frame": source_start_frame(self.source_object_name), + "max_output_size": int(self.max_output_size), + } + + def response(self, _context, result): + source_object = bpy.data.objects.get(self.source_object_name) + if not is_image_empty(source_object): + raise RuntimeError("The source image empty is no longer available") + source_image = bpy.data.images.get(self.source_image_name) + if source_image is None: + raise RuntimeError("The original source image is no longer available") + replaced_image = source_object.data + result_image = create_color_image( + source_object, + result.directory, + frame_prefix="upscale.", + ) + if replaced_image is not source_image: + replaced_image.use_fake_user = False + if int(replaced_image.users) == 0: + bpy.data.images.remove(replaced_image, do_unlink=True) + result_image.name = f"{image_base_name(source_image)}.upscale" + result_image.anyimage_upscale_source = source_image + model_label = upscale_model_label(self.model) + if getattr(self, "_rebuild_from_original", False): + return f"Image is rebuilt from the original with {model_label}" + return f"Image is upscaled with {model_label}" + + def cleanup(self): + cleanup_temporary_input(self.input_path, self.delete_input) diff --git a/src/anyimage/panel.py b/src/anyimage/panel.py index 1d3ba2c..e51be31 100644 --- a/src/anyimage/panel.py +++ b/src/anyimage/panel.py @@ -3,9 +3,9 @@ import bpy from .runtime import runtime from .properties import ( BEN2_MODEL_KEY, - REALESRGAN_MODEL_KEY, model_ready, ) +from .server.model_catalog import UPSCALE_MODELS class ServerPanel(bpy.types.Panel): bl_label = "Job Server" @@ -227,9 +227,9 @@ class Ben2RemovalDebugPanel(bpy.types.Panel): ) -class RealesrganGeneralX4v3DebugPanel(bpy.types.Panel): - bl_label = "Real-ESRGAN General x4v3" - bl_idname = "ANYIMAGE_PT_debug_realesrgan_general_x4v3" +class UpscaleDebugPanel(bpy.types.Panel): + bl_label = "Upscale" + bl_idname = "ANYIMAGE_PT_debug_upscale" bl_parent_id = DebugPanel.bl_idname bl_space_type = "VIEW_3D" bl_region_type = "UI" @@ -243,27 +243,15 @@ class RealesrganGeneralX4v3DebugPanel(bpy.types.Panel): layout.prop(settings, "input_path") input_path = bpy.path.abspath(settings.input_path) if input_path.lower().endswith( - ( - ".mp4", - ".mov", - ".avi", - ".mkv", - ".webm", - ".m4v", - ".mts", - ".mpg", - ".mpeg", - ) + (".mp4", ".mov", ".avi", ".mkv", ".webm", ".m4v", ".mts", ".mpg", ".mpeg") ): layout.prop(settings, "video_frames") - run_row = layout.row() - run_row.enabled = ( - not runtime.server_busy() - and model_ready(REALESRGAN_MODEL_KEY) - ) - run_row.operator( - "anyimage.debug_realesrgan_general_x4v3", - text="Run General x4v3 API", - icon="EXPORT", - ) - + for model in UPSCALE_MODELS.values(): + row = layout.row() + row.enabled = not runtime.server_busy() and model_ready(model.key) + operator = row.operator( + "anyimage.debug_upscale", + text=f"Run {model.label} API", + icon="EXPORT", + ) + operator.model = model.key diff --git a/src/anyimage/preferences.py b/src/anyimage/preferences.py index 476f8e7..4b06c03 100644 --- a/src/anyimage/preferences.py +++ b/src/anyimage/preferences.py @@ -1,16 +1,46 @@ +from pathlib import Path + import bpy -from .runtime import DEFAULT_STORAGE_ROOT, runtime -from .properties import ( - ai_setup_label, - ai_status, - device_enum_items, - model_catalog, +from .server.model_catalog import DEFAULT_UPSCALE_MODEL_KEY, UPSCALE_MODELS + + + +ADDON_PACKAGE = __package__ +DEFAULT_STORAGE_ROOT = Path.home() / ".anyimage" +DEFAULT_UPSCALE_MAX_OUTPUT_SIZE = 4096 +MAX_UPSCALE_OUTPUT_SIZE = 8192 +UPSCALE_MODEL_ITEMS = tuple( + (model.key, model.label, model.description, index) + for index, model in enumerate(UPSCALE_MODELS.values()) ) +def addon_preferences(context=None): + preference_context = ( + context + if getattr(context, "preferences", None) is not None + else bpy.context + ) + addon = preference_context.preferences.addons.get(ADDON_PACKAGE) + return getattr(addon, "preferences", None) + + +def configured_storage_root(): + configured = getattr(addon_preferences(), "storage_root", "") + if configured: + return Path(bpy.path.abspath(configured)) + return DEFAULT_STORAGE_ROOT + + +def device_items(owner, context): + from .properties import device_enum_items + + return device_enum_items(owner, context) + + class AnyImagePreferences(bpy.types.AddonPreferences): - bl_idname = __package__ + bl_idname = ADDON_PACKAGE storage_root: bpy.props.StringProperty( name="Storage Root", @@ -27,14 +57,37 @@ class AnyImagePreferences(bpy.types.AddonPreferences): "Inference device used by production tasks; only devices supported " "on this machine are listed" ), - items=device_enum_items, + items=device_items, default=0, ) + upscale_model: bpy.props.EnumProperty( + name="Default Upscale Model", + description="Model used when starting a new Upscale operation", + items=UPSCALE_MODEL_ITEMS, + default=DEFAULT_UPSCALE_MODEL_KEY, + ) + upscale_max_output_size: bpy.props.IntProperty( + name="Maximum Upscale Size", + description=( + "Maximum output long edge; larger inputs are resized before 4x " + "upscaling to keep memory use bounded" + ), + default=DEFAULT_UPSCALE_MAX_OUTPUT_SIZE, + min=1024, + max=MAX_UPSCALE_OUTPUT_SIZE, + step=256, + subtype="PIXEL", + ) def draw(self, _context): + from .properties import ai_setup_label, ai_status, model_catalog + from .runtime import runtime + layout = self.layout layout.prop(self, "storage_root") layout.prop(self, "device") + layout.prop(self, "upscale_model") + layout.prop(self, "upscale_max_output_size") layout.separator() status = ai_status() layout.label(text=f"Environment: {runtime.environment_directory()}") @@ -74,6 +127,8 @@ class AnyImagePreferences(bpy.types.AddonPreferences): ) def _draw_model(self, layout, model): + from .runtime import runtime + row = layout.row(align=True) row.label(text=model["label"]) if model.get("ready"): diff --git a/src/anyimage/properties.py b/src/anyimage/properties.py index 82bcb27..612d17b 100644 --- a/src/anyimage/properties.py +++ b/src/anyimage/properties.py @@ -6,15 +6,26 @@ import sys import bpy +from .preferences import ( + DEFAULT_UPSCALE_MAX_OUTPUT_SIZE, + MAX_UPSCALE_OUTPUT_SIZE, + addon_preferences, +) from .runtime import runtime -from . import model_catalog as shared_model_catalog +from .server import model_catalog as shared_model_catalog BEN2_MODEL_KEY = "BEN2_BASE" -REALESRGAN_MODEL_KEY = "REALESRGAN_GENERAL_X4V3" +REALESRGAN_GENERAL_WDN_MODEL_KEY = shared_model_catalog.DEFAULT_UPSCALE_MODEL_KEY +SAFMN_MODEL_KEY = shared_model_catalog.SAFMN_MODEL.key +REALESRGAN_X4PLUS_MODEL_KEY = shared_model_catalog.REALESRGAN_X4PLUS.key DEFAULT_DA3_MODEL_KEY = "BASE" DEFAULT_MOGE2_MODEL_KEY = "MOGE2_VITS_NORMAL" -REQUIRED_MODEL_KEYS = (DEFAULT_MOGE2_MODEL_KEY, BEN2_MODEL_KEY) +REQUIRED_MODEL_KEYS = ( + DEFAULT_MOGE2_MODEL_KEY, + BEN2_MODEL_KEY, + REALESRGAN_GENERAL_WDN_MODEL_KEY, +) DEVICE_LABELS = { @@ -214,6 +225,35 @@ def surface_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 configured_upscale_max_output_size(context=None): + preferences = addon_preferences(context) + value = int( + getattr( + preferences, + "upscale_max_output_size", + DEFAULT_UPSCALE_MAX_OUTPUT_SIZE, + ) + ) + return min(max(value, 1024), MAX_UPSCALE_OUTPUT_SIZE) + + +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 + + BALLOON_MODE_ITEMS = ( ( "POISSON", diff --git a/src/anyimage/runtime.py b/src/anyimage/runtime.py index bac1a32..040b1b4 100644 --- a/src/anyimage/runtime.py +++ b/src/anyimage/runtime.py @@ -1,19 +1,8 @@ from pathlib import Path -import bpy - from blendjob import JobRuntime - -DEFAULT_STORAGE_ROOT = Path.home() / ".anyimage" - - -def configured_storage_root(): - addon = bpy.context.preferences.addons.get(__package__.partition(".")[0]) - configured = getattr(getattr(addon, "preferences", None), "storage_root", "") - if configured: - return Path(bpy.path.abspath(configured)) - return DEFAULT_STORAGE_ROOT +from .preferences import configured_storage_root ENVIRONMENT = { @@ -37,6 +26,10 @@ def post_install(job_runtime): 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: try: diff --git a/src/anyimage/server/__init__.py b/src/anyimage/server/__init__.py index 3b933f5..0195b6d 100644 --- a/src/anyimage/server/__init__.py +++ b/src/anyimage/server/__init__.py @@ -1,90 +1,12 @@ -"""AnyImage Job Server without Blender dependencies.""" - -from blendjob import JobServer - -from .jobs import generate_depth as _generate_depth -from .jobs import generate_surface as _generate_surface -from .jobs.background import run as _remove_background -from .jobs.balloon import run as _analyze_balloon -from .jobs.debug_ben2 import run as _debug_ben2 -from .jobs.debug_da3 import run as _debug_da3 -from .jobs.debug_moge2 import run as _debug_moge2 -from .jobs.debug_realesrgan import run as _debug_realesrgan -from .jobs.frames import run as _extract_frames -from .jobs.image_mask import run as _refine_image_mask -from .model_manager import ModelManager, ModelSessionCache +"""Lightweight package boundary shared by Blender and the Job Server.""" -server = JobServer("AnyImage Job Server") +def __getattr__(name): + if name != "server": + raise AttributeError(name) + from .app import server + + return server -@server.resource("model_manager") -def create_model_manager(job_server): - return ModelManager(job_server.storage_root / "models") - - -@server.job("download-model") -def download_model(context, parameters): - manager = context.resource("model_manager") - return manager.download(context, parameters["model"]) - - -@server.job("download-required-models") -def download_required_models(context, parameters): - manager = context.resource("model_manager") - return manager.download_missing(context, parameters["models"]) - - -@server.job("analyze-balloon") -def analyze_balloon(context, parameters): - return _analyze_balloon(context, parameters) - - -@server.job("refine-image-mask") -def refine_image_mask(context, parameters): - return _refine_image_mask(context, parameters) - - -@server.job("remove-background") -def remove_background(context, parameters): - return _remove_background(context, parameters) - - -@server.job("extract-frames") -def extract_frames(context, parameters): - return _extract_frames(context, parameters) - - -@server.job("generate-depth") -def generate_depth(context, parameters): - return _generate_depth(context, parameters) - - -@server.job("generate-surface") -def generate_surface(context, parameters): - return _generate_surface(context, parameters) - - -@server.job("debug-da3") -def debug_da3(context, parameters): - return _debug_da3(context, parameters) - - -@server.job("debug-moge2") -def debug_moge2(context, parameters): - return _debug_moge2(context, parameters) - - -@server.job("debug-ben2") -def debug_ben2(context, parameters): - return _debug_ben2(context, parameters) - - -@server.job("debug-realesrgan-general-x4v3") -def debug_realesrgan(context, parameters): - return _debug_realesrgan(context, parameters) - - - - -__all__ = ("ModelManager", "ModelSessionCache", "server") +__all__ = ("server",) diff --git a/src/anyimage/server/app.py b/src/anyimage/server/app.py new file mode 100644 index 0000000..ea4c6cf --- /dev/null +++ b/src/anyimage/server/app.py @@ -0,0 +1,92 @@ +"""AnyImage Job Server without Blender dependencies.""" + +from blendjob import JobServer + +from .jobs import generate_depth as _generate_depth +from .jobs import generate_surface as _generate_surface +from .jobs.background import run as _remove_background +from .jobs.balloon import run as _analyze_balloon +from .jobs.debug_ben2 import run as _debug_ben2 +from .jobs.debug_da3 import run as _debug_da3 +from .jobs.debug_moge2 import run as _debug_moge2 +from .jobs.debug_upscale import run as _debug_upscale +from .jobs.frames import run as _extract_frames +from .jobs.image_mask import run as _refine_image_mask +from .jobs.upscale import run as _upscale_image +from .model_manager import ModelManager + + +server = JobServer("AnyImage Job Server") + + +@server.resource("model_manager") +def create_model_manager(job_server): + return ModelManager(job_server.storage_root / "models") + + +@server.job("download-model") +def download_model(context, parameters): + manager = context.resource("model_manager") + return manager.download(context, parameters["model"]) + + +@server.job("download-required-models") +def download_required_models(context, parameters): + manager = context.resource("model_manager") + return manager.download_missing(context, parameters["models"]) + + +@server.job("analyze-balloon") +def analyze_balloon(context, parameters): + return _analyze_balloon(context, parameters) + + +@server.job("refine-image-mask") +def refine_image_mask(context, parameters): + return _refine_image_mask(context, parameters) + + +@server.job("remove-background") +def remove_background(context, parameters): + return _remove_background(context, parameters) + + +@server.job("upscale-image") +def upscale_image(context, parameters): + return _upscale_image(context, parameters) + + +@server.job("extract-frames") +def extract_frames(context, parameters): + return _extract_frames(context, parameters) + + +@server.job("generate-depth") +def generate_depth(context, parameters): + return _generate_depth(context, parameters) + + +@server.job("generate-surface") +def generate_surface(context, parameters): + return _generate_surface(context, parameters) + + +@server.job("debug-da3") +def debug_da3(context, parameters): + return _debug_da3(context, parameters) + + +@server.job("debug-moge2") +def debug_moge2(context, parameters): + return _debug_moge2(context, parameters) + + +@server.job("debug-ben2") +def debug_ben2(context, parameters): + return _debug_ben2(context, parameters) + + +@server.job("debug-upscale") +def debug_upscale(context, parameters): + return _debug_upscale(context, parameters) + diff --git a/src/anyimage/server/jobs/background.py b/src/anyimage/server/jobs/background.py index 893a284..539662c 100644 --- a/src/anyimage/server/jobs/background.py +++ b/src/anyimage/server/jobs/background.py @@ -53,7 +53,7 @@ def run(context, parameters): ) context.check_cancelled() context.progress(0.92, "Writing transparent PNG") - output_path = output_directory / "color.0001.png" + output_path = output_directory / "image.0001.png" foreground.convert("RGBA").save( output_path, format="PNG", @@ -72,7 +72,7 @@ def run(context, parameters): 0.7, model_cache=model_manager, cancel_check=context.check_cancelled, - output_prefix="color", + output_prefix="image", preserve_input_alpha=True, ) finally: @@ -86,6 +86,6 @@ def run(context, parameters): f"Background removed{frame_text} in {inference_ms:.0f} ms " f"(model load {load_ms:.0f} ms, total {total_ms:.0f} ms)" ) - print(message) + context.log(message) context.progress(1.0, message) - return {"color": "."} + return {"image": "."} diff --git a/src/anyimage/server/jobs/debug_realesrgan.py b/src/anyimage/server/jobs/debug_upscale.py similarity index 90% rename from src/anyimage/server/jobs/debug_realesrgan.py rename to src/anyimage/server/jobs/debug_upscale.py index f7111c6..927debc 100644 --- a/src/anyimage/server/jobs/debug_realesrgan.py +++ b/src/anyimage/server/jobs/debug_upscale.py @@ -1,5 +1,5 @@ from ..media import copy_frames -from ..models.realesrgan import infer_one +from ..models.upscale import infer_one from .common import input_frames @@ -13,6 +13,7 @@ def run(context, parameters): context.resource("model_manager"), parameters, image, + cancel_check=context.check_cancelled, ) suffix = "" if len(images) == 1 else f".{index + 1:04d}" prediction.save( diff --git a/src/anyimage/server/jobs/foreground.py b/src/anyimage/server/jobs/foreground.py index 540ee41..78a60ad 100644 --- a/src/anyimage/server/jobs/foreground.py +++ b/src/anyimage/server/jobs/foreground.py @@ -25,7 +25,7 @@ def write_optional( cancel_check=context.check_cancelled, preserve_input_alpha=True, ) - path = context.directory / "color.0001.png" + path = context.directory / "image.0001.png" foreground.convert("RGBA").save( path, format="PNG", @@ -43,11 +43,11 @@ def write_optional( progress_span, model_cache=model_manager, cancel_check=context.check_cancelled, - output_prefix="color", + output_prefix="image", preserve_input_alpha=True, ) - paths = sorted(context.directory.glob("color.*.png")) - print( + paths = sorted(context.directory.glob("image.*.png")) + context.log( f"BEN2 inference {inference_ms:.0f} ms; " f"model load {load_ms:.0f} ms" ) diff --git a/src/anyimage/server/jobs/upscale.py b/src/anyimage/server/jobs/upscale.py new file mode 100644 index 0000000..999fbfb --- /dev/null +++ b/src/anyimage/server/jobs/upscale.py @@ -0,0 +1,34 @@ +from ..models.upscale import infer_one +from .common import frame_progress_message, input_frames + + +def run(context, parameters): + context.log( + f"Upscale model: {parameters['model']}; " + f"maximum output size: {parameters.get('max_output_size', 4096)} px" + ) + context.progress(0.02, "Preparing upscale input") + context.directory.mkdir(parents=True, exist_ok=True) + with input_frames(context, parameters) as images: + total = len(images) + for index, image in enumerate(images): + context.check_cancelled() + context.progress( + 0.08 + index / total * 0.84, + frame_progress_message("Upscaling image", index, total), + ) + prediction = infer_one( + context.resource("model_manager"), + parameters, + image, + cancel_check=context.check_cancelled, + preserve_alpha=True, + ) + prediction.save( + context.directory / f"upscale.{index + 1:04d}.png", + format="PNG", + compress_level=1, + ) + context.check_cancelled() + context.progress(1.0, "Image upscale complete") + return {"upscale": "."} diff --git a/src/anyimage/server/media/input.py b/src/anyimage/server/media/input.py index b5df2ce..ea70545 100644 --- a/src/anyimage/server/media/input.py +++ b/src/anyimage/server/media/input.py @@ -144,11 +144,11 @@ def prepare_background_input( def copy_color_images(images, output_directory): - """Copy input images as the canonical numbered color layer.""" + """Copy input images as the canonical numbered image layer.""" for index, image_path in enumerate(images): suffix = Path(image_path).suffix.lower() or ".jpg" destination = ( - Path(output_directory) / f"color.{index + 1:04d}{suffix}" + Path(output_directory) / f"image.{index + 1:04d}{suffix}" ) if Path(image_path).resolve() == destination.resolve(): continue diff --git a/src/anyimage/server/model_catalog.py b/src/anyimage/server/model_catalog.py index f63f039..65accde 100644 --- a/src/anyimage/server/model_catalog.py +++ b/src/anyimage/server/model_catalog.py @@ -10,6 +10,7 @@ class ModelSpec: license_name: str description: str supports_gs: bool + huggingface_repository: str = "" download_patterns: tuple[str, ...] = () ready_patterns: tuple[str, ...] = ("config.json", "*.safetensors") enabled: bool = True @@ -24,6 +25,7 @@ MODELS = { directory_name="DA3-SMALL", repository="depth-anything/DA3-SMALL", license_name="Apache 2.0", + huggingface_repository="depth-anything/DA3-SMALL", description="Fastest any-view model using ONNX Runtime", supports_gs=False, download_patterns=("model.onnx",), @@ -39,6 +41,7 @@ MODELS = { directory_name="DA3-BASE", repository="depth-anything/DA3-BASE", license_name="Apache 2.0", + huggingface_repository="depth-anything/DA3-BASE", description="Balanced any-view model using ONNX Runtime", supports_gs=False, download_patterns=("model.onnx",), @@ -76,6 +79,7 @@ BEN2_MODEL = ModelSpec( 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", supports_gs=False, download_patterns=("onnx/model_fp16.onnx",), @@ -89,30 +93,75 @@ BEN2_MODEL = ModelSpec( ) -REALESRGAN_MODEL = ModelSpec( - key="REALESRGAN_GENERAL_X4V3", - label="Real-ESRGAN General x4v3", - directory_name="realesr-general-x4v3", +REALESRGAN_GENERAL_WDN_X4V3 = ModelSpec( + key="REALESRGAN_GENERAL_WDN_X4V3", + label="Real-ESRGAN General WDN x4v3", + directory_name="realesr-general-wdn-x4v3", + repository="xinntao/Real-ESRGAN", + license_name="BSD 3-Clause", + description="Fast and conservative 4x image upscaler using ONNX Runtime", + supports_gs=False, + download_patterns=("realesr-general-wdn-x4v3.onnx",), + ready_patterns=("realesr-general-wdn-x4v3.onnx",), + r2_directory="realesr-general-wdn-x4v3", + r2_files=(( + "realesr-general-wdn-x4v3.onnx", + 4867419, + "b021a837f7a0d3f287f9a39e367f581f1e6c3a759a08dc0c1c79c22442729774", + ),), +) + +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 real-world 4x image upscaler using ONNX Runtime", + supports_gs=False, + 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", + directory_name="realesrgan-x4plus", repository="jonathanst29/tinier-upscale-models", license_name="BSD 3-Clause", - description="Fast general-purpose 4x image upscaler using ONNX Runtime", + huggingface_repository="jonathanst29/tinier-upscale-models", + description="Detailed general-purpose 4x image upscaler using ONNX Runtime", supports_gs=False, - download_patterns=("realesr-general-x4v3.onnx",), - ready_patterns=("realesr-general-x4v3.onnx",), - r2_directory="realesr-general-x4v3", + download_patterns=("realesrgan-x4plus.onnx",), + ready_patterns=("realesrgan-x4plus.onnx",), + r2_directory="realesrgan-x4plus", r2_files=(( - "realesr-general-x4v3.onnx", - 4866396, - "924ebad6532777303582d4ce7811849b88869231a3ae7093e0f21200249df8d5", + "realesrgan-x4plus.onnx", + 67051618, + "4ed6a45a8185bde6c8d7c7790a6e80e3c5b9f608946f634021068dd0fd0473c8", ),), ) UPSCALE_MODELS = { - REALESRGAN_MODEL.key: REALESRGAN_MODEL, + model.key: model + for model in ( + REALESRGAN_GENERAL_WDN_X4V3, + SAFMN_MODEL, + REALESRGAN_X4PLUS, + ) } +DEFAULT_UPSCALE_MODEL_KEY = REALESRGAN_GENERAL_WDN_X4V3.key + + MOGE2_MODELS = { "MOGE2_VITS_NORMAL": ModelSpec( key="MOGE2_VITS_NORMAL", @@ -120,6 +169,7 @@ MOGE2_MODELS = { directory_name="moge-2-vits-normal-onnx", repository="Ruicheng/moge-2-vits-normal-onnx", license_name="MIT", + huggingface_repository="Ruicheng/moge-2-vits-normal-onnx", description="Fastest MoGe-2 model using ONNX Runtime", supports_gs=False, download_patterns=("model.onnx",), @@ -137,6 +187,7 @@ MOGE2_MODELS = { directory_name="moge-2-vitb-normal-onnx", repository="Ruicheng/moge-2-vitb-normal-onnx", license_name="MIT", + huggingface_repository="Ruicheng/moge-2-vitb-normal-onnx", description="Balanced MoGe-2 model using ONNX Runtime", supports_gs=False, download_patterns=("model.onnx",), @@ -154,6 +205,7 @@ MOGE2_MODELS = { 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", supports_gs=False, download_patterns=("model.onnx",), diff --git a/src/anyimage/server/model_download.py b/src/anyimage/server/model_download.py index fcb40d8..1179c35 100644 --- a/src/anyimage/server/model_download.py +++ b/src/anyimage/server/model_download.py @@ -171,15 +171,15 @@ def download_model(args): def _download_declared_or_fail(args, model_dir, label, start, span): - """Download declared files from Hugging Face, then the R2 mirror.""" - sources = ( - (HF_ENDPOINT, download_from_hf, f"Downloading {label}"), - ( + """Download declared files from Hugging Face when available, then R2.""" + sources = [] + if args.repository: + sources.append((HF_ENDPOINT, download_from_hf, f"Downloading {label}")) + sources.append(( R2_ENDPOINT, download_from_r2, f"Downloading {label} (mirror)", - ), - ) + )) failures = [] for source, downloader, message in sources: args.progress(start, message) @@ -431,6 +431,8 @@ def download_from_hf( ) + + def format_rate(rate): size = float(rate) if size <= 0: diff --git a/src/anyimage/server/model_manager.py b/src/anyimage/server/model_manager.py index 20a081a..1e994c7 100644 --- a/src/anyimage/server/model_manager.py +++ b/src/anyimage/server/model_manager.py @@ -17,15 +17,15 @@ class ModelSessionCache: self.da3_key = None self.moge_model = None self.moge_key = None - self.realesrgan_model = None - self.realesrgan_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() + self.close() return self._load_ben2(model_dir, requested_device, key) def _load_ben2(self, model_dir, requested_device, key): @@ -41,11 +41,9 @@ class ModelSessionCache: def get_da3(self, model_dir, device): requested_device = getattr(device, "type", device) key = (str(Path(model_dir).resolve()), requested_device, "onnx") - if self.moge_model is not None: - self.release_moge() if self.da3_model is not None and self.da3_key == key: return self.da3_model - self.release_da3() + self.close() return self._load_da3(model_dir, requested_device, key) def _load_da3(self, model_dir, requested_device, key): @@ -60,26 +58,24 @@ class ModelSessionCache: from .models.onnx_moge2 import create_session key = (str(Path(model_dir).resolve()), requested_device, "onnx") - if self.da3_model is not None: - self.release_da3() if self.moge_model is not None and self.moge_key == key: return self.moge_model - self.release_moge() + self.close() model = create_session(model_dir, requested_device) self.moge_model = model self.moge_key = key return model - def get_realesrgan(self, model_dir, requested_device): - from .models.onnx_realesrgan import create_session + def get_upscale(self, model_dir, requested_device): + from .models.onnx_upscale import create_session key = (str(Path(model_dir).resolve()), requested_device, "onnx") - if self.realesrgan_model is not None and self.realesrgan_key == key: - return self.realesrgan_model - self.release_realesrgan() + if self.upscale_model is not None and self.upscale_key == key: + return self.upscale_model + self.close() model = create_session(model_dir, requested_device) - self.realesrgan_model = model - self.realesrgan_key = key + self.upscale_model = model + self.upscale_key = key return model def release_ben2(self): @@ -97,20 +93,20 @@ class ModelSessionCache: self.moge_key = None self._release_runtime_cache() - def release_realesrgan(self): - self.realesrgan_model = None - self.realesrgan_key = None + def release_upscale(self): + self.upscale_model = None + self.upscale_key = None self._release_runtime_cache() def close(self): self.ben2_model = None self.da3_model = None self.moge_model = None - self.realesrgan_model = None + self.upscale_model = None self.ben2_key = None self.da3_key = None self.moge_key = None - self.realesrgan_key = None + self.upscale_key = None self._release_runtime_cache() def snapshot(self): @@ -118,9 +114,9 @@ class ModelSessionCache: "ben2": Path(self.ben2_key[0]).name if self.ben2_key else "", "da3": Path(self.da3_key[0]).name if self.da3_key else "", "moge2": Path(self.moge_key[0]).name if self.moge_key else "", - "realesrgan": ( - Path(self.realesrgan_key[0]).name - if self.realesrgan_key + "upscale": ( + Path(self.upscale_key[0]).name + if self.upscale_key else "" ), } @@ -138,22 +134,7 @@ class ModelManager: def directory(self, model_key): model = model_catalog.get_downloadable_model(model_key) - directory = self.models_directory / model.directory_name - old_moge_directory = self.models_directory / "moge2_base_normal" - if ( - model_key == "MOGE2_VITB_NORMAL" - and not directory.is_dir() - and old_moge_directory.is_dir() - ): - old_moge_directory.rename(directory) - legacy_directory = self.models_directory / model_key.lower() - if ( - not directory.is_dir() - and legacy_directory.is_dir() - and legacy_directory != directory - ): - legacy_directory.rename(directory) - return directory + return self.models_directory / model.directory_name def ready(self, model_key): model = model_catalog.get_downloadable_model(model_key) @@ -169,8 +150,11 @@ class ModelManager: def get_moge2(self, model_dir, device): return self.sessions.get_moge2(model_dir, device) - def get_realesrgan(self, model_dir, device): - return self.sessions.get_realesrgan(model_dir, device) + def get_upscale(self, model_dir, device): + return self.sessions.get_upscale(model_dir, device) + + def release_upscale(self): + self.sessions.release_upscale() def snapshot(self): return {"loaded": self.sessions.snapshot()} @@ -201,7 +185,7 @@ class ModelManager: def _download(self, context, model_key, progress): model = model_catalog.get_downloadable_model(model_key) args = SimpleNamespace( - repository=model.repository, + repository=model.huggingface_repository, model_dir=str(self.directory(model_key)), progress=progress, cancel_check=context.check_cancelled, diff --git a/src/anyimage/server/models/__init__.py b/src/anyimage/server/models/__init__.py index 3309119..0eae15b 100644 --- a/src/anyimage/server/models/__init__.py +++ b/src/anyimage/server/models/__init__.py @@ -1,4 +1,4 @@ -from . import ben2, da3, moge2, realesrgan +from . import ben2, da3, moge2, upscale def geometry_adapter(model_key): @@ -9,6 +9,6 @@ __all__ = ( "ben2", "da3", "moge2", - "realesrgan", + "upscale", "geometry_adapter", ) diff --git a/src/anyimage/server/models/onnx_realesrgan.py b/src/anyimage/server/models/onnx_realesrgan.py deleted file mode 100644 index b10cce3..0000000 --- a/src/anyimage/server/models/onnx_realesrgan.py +++ /dev/null @@ -1,51 +0,0 @@ -"""Real-ESRGAN General x4v3 ONNX inference.""" - -from pathlib import Path - -try: - from .onnx_runtime import create_session as create_runtime_session -except ImportError: - from onnx_runtime import create_session as create_runtime_session - - -ONNX_FILENAME = "realesr-general-x4v3.onnx" -SCALE = 4 - - -def create_session(model_directory, requested_device="auto"): - path = Path(model_directory) / ONNX_FILENAME - return create_runtime_session(path, requested_device) - - -def preprocess(image): - import numpy as np - - image_array = np.asarray(image.convert("RGB"), dtype=np.float32) / 255.0 - return np.ascontiguousarray(image_array.transpose(2, 0, 1)[None]) - - -def infer(session, image): - import numpy as np - from PIL import Image - - input_tensor = preprocess(image) - input_name = session.get_inputs()[0].name - output_name = session.get_outputs()[0].name - prediction = np.asarray( - session.run([output_name], {input_name: input_tensor})[0], - dtype=np.float32, - ) - expected_shape = ( - 1, - 3, - input_tensor.shape[2] * SCALE, - input_tensor.shape[3] * SCALE, - ) - if prediction.shape != expected_shape: - raise RuntimeError( - "Real-ESRGAN General x4v3 returned an unexpected output shape: " - f"{prediction.shape}" - ) - pixels = np.clip(prediction[0].transpose(1, 2, 0), 0.0, 1.0) - pixels = np.rint(pixels * 255.0).astype(np.uint8) - return Image.fromarray(pixels, mode="RGB") diff --git a/src/anyimage/server/models/onnx_upscale.py b/src/anyimage/server/models/onnx_upscale.py new file mode 100644 index 0000000..eaa4d15 --- /dev/null +++ b/src/anyimage/server/models/onnx_upscale.py @@ -0,0 +1,142 @@ +"""Tiled ONNX inference shared by the Upscale models.""" + +from pathlib import Path + +try: + from .onnx_runtime import create_session as create_runtime_session +except ImportError: + from onnx_runtime import create_session as create_runtime_session + + +SCALE = 4 +TILE_SIZE = 256 +TILE_BORDER = 16 +RESOURCE_ERROR_TOKENS = ( + "out of memory", + "bad allocation", + "bad_alloc", + "failed to allocate", + "not enough memory", + "insufficient memory", +) + + +class UpscaleResourceError(RuntimeError): + """The execution provider could not allocate resources for one tile.""" + + +def create_session(model_directory, requested_device="auto"): + paths = tuple(Path(model_directory).glob("*.onnx")) + if len(paths) != 1: + raise RuntimeError( + f"Expected exactly one ONNX model in {model_directory}, found {len(paths)}" + ) + path = paths[0] + return create_runtime_session(path, requested_device) + + +def preprocess(image): + import numpy as np + + image_array = np.asarray(image.convert("RGB"), dtype=np.float32) / 255.0 + return np.ascontiguousarray(image_array.transpose(2, 0, 1)[None]) + + +def _run(session, input_tensor): + import numpy as np + + input_name = session.get_inputs()[0].name + output_name = session.get_outputs()[0].name + try: + raw_prediction = session.run( + [output_name], + {input_name: input_tensor}, + )[0] + except (MemoryError, UnicodeDecodeError) as error: + raise UpscaleResourceError( + "the ONNX execution provider failed, usually because GPU or " + "system memory was exhausted" + ) from error + except RuntimeError as error: + if any(token in str(error).lower() for token in RESOURCE_ERROR_TOKENS): + raise UpscaleResourceError( + "the ONNX execution provider failed because GPU or system " + "memory was exhausted" + ) from error + raise + prediction = np.asarray(raw_prediction, dtype=np.float32) + expected_shape = ( + 1, + 3, + input_tensor.shape[2] * SCALE, + input_tensor.shape[3] * SCALE, + ) + if prediction.shape != expected_shape: + raise RuntimeError( + f"Upscale model returned an unexpected output shape: {prediction.shape}" + ) + return prediction + + +def _prediction_pixels(prediction): + import numpy as np + + pixels = np.clip(prediction[0].transpose(1, 2, 0), 0.0, 1.0) + return np.rint(pixels * 255.0).astype(np.uint8) + + +def _pad_source(source, extra_height, extra_width): + 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), + (0, 0), + ), + mode=mode, + ) + + +def infer(session, image, cancel_check=None): + 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 + extra_height = (-height) % core_size + extra_width = (-width) % core_size + padded = _pad_source(source, extra_height, extra_width) + 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) + for left in range(0, width, core_size): + core_width = min(core_size, width - left) + if cancel_check is not None and cancel_check(): + raise RuntimeError("Task cancelled") + tile = Image.fromarray( + padded[top : top + TILE_SIZE, left : left + TILE_SIZE], + mode="RGB", + ) + pixels = _prediction_pixels(_run(session, preprocess(tile))) + border = TILE_BORDER * SCALE + output[ + top * SCALE : (top + core_height) * SCALE, + left * SCALE : (left + core_width) * SCALE, + ] = pixels[ + border : border + core_height * SCALE, + border : border + core_width * SCALE, + ] + return Image.fromarray(output, mode="RGB") + except UpscaleResourceError: + raise + except MemoryError as error: + raise UpscaleResourceError( + "the complete upscale result exceeded available system memory" + ) from error + + diff --git a/src/anyimage/server/models/realesrgan.py b/src/anyimage/server/models/realesrgan.py deleted file mode 100644 index 34feb53..0000000 --- a/src/anyimage/server/models/realesrgan.py +++ /dev/null @@ -1,12 +0,0 @@ -def infer_one(model_manager, parameters, image_path): - from PIL import Image - - from .onnx_realesrgan import infer - - model_directory = model_manager.directory(parameters["model"]) - model = model_manager.get_realesrgan( - model_directory, - parameters["device"], - ) - with Image.open(image_path) as opened: - return infer(model, opened) diff --git a/src/anyimage/server/models/upscale.py b/src/anyimage/server/models/upscale.py new file mode 100644 index 0000000..41232d8 --- /dev/null +++ b/src/anyimage/server/models/upscale.py @@ -0,0 +1,76 @@ +import math + + +SCALE = 4 +MAX_OUTPUT_SIZE = 8192 +MAX_OUTPUT_PIXELS = 32 * 1024 * 1024 + + +def _limited_input(image, max_output_size): + from PIL import Image + + max_output_size = int(max_output_size) + if max_output_size < 1 or max_output_size > MAX_OUTPUT_SIZE: + raise ValueError( + f"Max Output Size must be between 1 and {MAX_OUTPUT_SIZE} pixels" + ) + width, height = image.size + scale = min( + 1.0, + max_output_size / (max(width, height) * SCALE), + math.sqrt(MAX_OUTPUT_PIXELS / (width * height * SCALE * SCALE)), + ) + if scale >= 1.0: + return image + size = ( + max(1, int(width * scale)), + max(1, int(height * scale)), + ) + return image.resize(size, Image.Resampling.LANCZOS) + + +def _source_alpha(image): + if "A" not in image.getbands() and "transparency" not in image.info: + return None + return image.convert("RGBA").getchannel("A") + + +def infer_one( + model_manager, + parameters, + image_path, + cancel_check=None, + preserve_alpha=False, +): + from PIL import Image + + from .onnx_upscale import UpscaleResourceError, infer + + model_directory = model_manager.directory(parameters["model"]) + model = model_manager.get_upscale( + model_directory, + parameters["device"], + ) + try: + with Image.open(image_path) as opened: + source = opened + if "max_output_size" in parameters: + source = _limited_input(source, parameters["max_output_size"]) + source_alpha = _source_alpha(source) if preserve_alpha else None + prediction = infer(model, source, cancel_check=cancel_check) + if source_alpha is None: + return prediction + alpha = source_alpha.resize( + prediction.size, + Image.Resampling.LANCZOS, + ) + prediction = prediction.convert("RGBA") + prediction.putalpha(alpha) + return prediction + except (MemoryError, UnicodeDecodeError, UpscaleResourceError) as error: + model_manager.release_upscale() + device = str(parameters.get("device", "selected device")).upper() + raise RuntimeError( + f"Upscale could not finish on {device}: {error}. " + "Reduce Max Output Size or select CPU, then try again." + ) from error diff --git a/sync_models.py b/sync_models.py index 1a6b3e9..4d9bba1 100644 --- a/sync_models.py +++ b/sync_models.py @@ -12,6 +12,10 @@ from urllib.request import Request, urlopen PROJECT_ROOT = Path(__file__).resolve().parent MODELS_DIR = PROJECT_ROOT / "models" +SOURCE_ROOT = PROJECT_ROOT / "src" / "anyimage" +sys.path.insert(0, str(SOURCE_ROOT)) +from server.model_catalog import DOWNLOADABLE_MODELS + R2_ENDPOINT = "https://models.omoolab.xyz" CHUNK_SIZE = 1024 * 1024 @@ -37,73 +41,31 @@ def r2_url(folder, filename="model.onnx"): return f"{R2_ENDPOINT}/{folder}/{filename}" -MODEL_FILES = ( - ModelFile( - "DA3-SMALL", - 105484029, - "7a9d81475b3e2334ace0d6d2c2a816b3c93081437e1866cc434ce1789125f631", - (r2_url("DA3-SMALL"),), - ), - ModelFile( - "DA3-BASE", - 412856697, - "31f38fc4990539d7902ba267924b9731db1acb00d840c248a60b1865810ff9e6", - (r2_url("DA3-BASE"),), - ), - ModelFile( - "BEN2-ONNX/onnx", - 219121675, - "dfdc25f421f32a0d1268e0f2ff2153d340e8f1d52d3dd16f5dc33c1ce85cedf1", - ( - hf_url("onnx-community/BEN2-ONNX", "onnx/model_fp16.onnx"), - r2_url("BEN2-ONNX", "onnx/model_fp16.onnx"), - ), - "model_fp16.onnx", - ), - ModelFile( - "realesr-general-x4v3", - 4866396, - "924ebad6532777303582d4ce7811849b88869231a3ae7093e0f21200249df8d5", - ( - hf_url( - "jonathanst29/tinier-upscale-models", - "realesr-general-x4v3.onnx", - ), - r2_url( - "realesr-general-x4v3", - "realesr-general-x4v3.onnx", - ), - ), - "realesr-general-x4v3.onnx", - ), - ModelFile( - "moge-2-vits-normal-onnx", - 140852051, - "24eacb5dc7a2c54c7bc98f7de085ffbed79ad006ea5b664c2c2cdc02ff3a52f0", - ( - hf_url("Ruicheng/moge-2-vits-normal-onnx", "model.onnx"), - r2_url("moge-2-vits-normal-onnx"), - ), - ), - ModelFile( - "moge-2-vitb-normal-onnx", - 419411850, - "bbf14e07a30f11e69d36ab861590123f5598ababcbc8946a063eb4a966f35a21", - ( - hf_url("Ruicheng/moge-2-vitb-normal-onnx", "model.onnx"), - r2_url("moge-2-vitb-normal-onnx"), - ), - ), - ModelFile( - "moge-2-vitl-normal-onnx", - 1324265014, - "afbc4ccc3450298f3afb35b90f015f4c4f552dea21dc6470d5f7b78b77e2d751", - ( - hf_url("Ruicheng/moge-2-vitl-normal-onnx", "model.onnx"), - r2_url("moge-2-vitl-normal-onnx"), - ), - ), -) +def catalog_model_files(): + files = [] + for model in DOWNLOADABLE_MODELS.values(): + for filename, size, sha256 in model.r2_files: + relative_path = Path(filename) + folder = str(Path(model.r2_directory) / relative_path.parent) + sources = [] + if model.huggingface_repository: + sources.append( + hf_url(model.huggingface_repository, filename) + ) + sources.append(r2_url(model.r2_directory, filename)) + files.append( + ModelFile( + folder, + size, + sha256, + tuple(sources), + relative_path.name, + ) + ) + return tuple(files) + + +MODEL_FILES = catalog_model_files() def file_sha256(path): diff --git a/tests/test_addon.py b/tests/test_addon.py index 11150ba..a99b682 100644 --- a/tests/test_addon.py +++ b/tests/test_addon.py @@ -39,6 +39,7 @@ class AddonTest(unittest.TestCase): Panel=object, PropertyGroup=object, WorkSpaceTool=object, + Image=type("Image", (), {}), Scene=type("Scene", (), {}), STATUSBAR_HT_header=SimpleNamespace( append=append_status_bar, @@ -46,10 +47,12 @@ class AddonTest(unittest.TestCase): ), VIEW3D_MT_object_context_menu=SimpleNamespace( append=lambda draw: cls.context_menu_draws.append(draw), + prepend=lambda draw: cls.context_menu_draws.insert(0, draw), remove=lambda draw: cls.context_menu_draws.remove(draw), ), OUTLINER_MT_object=SimpleNamespace( append=lambda draw: cls.outliner_menu_draws.append(draw), + prepend=lambda draw: cls.outliner_menu_draws.insert(0, draw), remove=lambda draw: cls.outliner_menu_draws.remove(draw), ), ) @@ -95,6 +98,7 @@ class AddonTest(unittest.TestCase): "anyimage.operators.mesh_plane" ) cls.image_background = importlib.import_module("anyimage.operators.background") + cls.image_upscale = importlib.import_module("anyimage.operators.upscale") cls.image_mesh = importlib.import_module("anyimage.operators.surface") cls.image_edit_common = importlib.import_module( "anyimage.operators.image_edit.common" @@ -145,6 +149,8 @@ class AddonTest(unittest.TestCase): preferences = SimpleNamespace( storage_root="", device="CUDA", + upscale_model="REALESRGAN_GENERAL_WDN_X4V3", + upscale_max_output_size=4096, ) self.fake_bpy.context.preferences.addons = { "anyimage": SimpleNamespace(preferences=preferences) @@ -162,6 +168,8 @@ class AddonTest(unittest.TestCase): ) if hasattr(self.fake_bpy.types.Scene, "anyimage_settings"): del self.fake_bpy.types.Scene.anyimage_settings + if hasattr(self.fake_bpy.types.Image, "anyimage_upscale_source"): + del self.fake_bpy.types.Image.anyimage_upscale_source def test_runtime_and_operator_support_modules_have_explicit_boundaries(self): for name in ( @@ -172,12 +180,6 @@ class AddonTest(unittest.TestCase): "status_bar.py", ): self.assertFalse((SOURCE_ROOT / "anyimage" / name).exists(), name) - self.assertTrue((SOURCE_ROOT / "anyimage" / "runtime.py").is_file()) - self.assertFalse((SOURCE_ROOT / "anyimage" / "job_runtime.py").exists()) - self.assertFalse((SOURCE_ROOT / "anyimage" / "client").exists()) - self.assertFalse((SOURCE_ROOT / "anyimage" / "ui").exists()) - self.assertFalse((SOURCE_ROOT / "anyimage" / "operators.py").exists()) - self.assertFalse((SOURCE_ROOT / "anyimage" / "blender").exists()) def test_mps_requires_m4_or_newer(self): properties = self.addon.properties @@ -228,12 +230,6 @@ class AddonTest(unittest.TestCase): def test_import_does_not_add_addon_directories_to_sys_path(self): self.assertEqual(self.paths_after_import, self.paths_before_import) - self.assertFalse( - any( - name == "anyimage.server" or name.startswith("anyimage.server.") - for name in sys.modules - ) - ) def test_job_request_is_not_part_of_blendjob(self): blendjob = importlib.import_module("blendjob") @@ -260,18 +256,7 @@ class AddonTest(unittest.TestCase): self.assertIn(self.job_runtime.JobOperatorBase, operator.__mro__) self.assertIn(self.fake_bpy.types.Operator, operator.__mro__) - def test_realesrgan_debug_operator_is_a_direct_model_request(self): - operator_type = self.debug.DebugRealesrganGeneralX4v3 - self.assertIn(operator_type, self.addon.CLASSES) - self.assertEqual( - operator_type.bl_idname, - "anyimage.debug_realesrgan_general_x4v3", - ) - self.assertEqual( - operator_type.job_type, - "debug-realesrgan-general-x4v3", - ) - + def test_upscale_debug_operator_accepts_each_catalog_model(self): with tempfile.TemporaryDirectory() as directory: image_path = Path(directory) / "input.png" image_path.touch() @@ -283,33 +268,19 @@ class AddonTest(unittest.TestCase): context = SimpleNamespace( scene=SimpleNamespace(anyimage_settings=settings) ) - with ( - patch.object(self.debug, "require_environment"), - patch.object(self.debug, "require_model") as require_model, - ): - request = operator_type().request(context) - - require_model.assert_called_once_with( - self.addon.properties.REALESRGAN_MODEL_KEY - ) - self.assertEqual( - request, - { - "input": str(image_path), - "model": "REALESRGAN_GENERAL_X4V3", - "device": "cpu", - "video_frames": 3, - "start_frame": 1, - }, - ) - self.assertNotIn("tile", request) - self.assertNotIn("strength", request) - + for model_key in self.addon.properties.shared_model_catalog.UPSCALE_MODELS: + operator = self.debug.DebugUpscale() + operator.model = model_key + with self.subTest(model=model_key), patch.object( + self.debug, "require_environment" + ), patch.object(self.debug, "require_model") as require_model: + request = operator.request(context) + require_model.assert_called_once_with(model_key) + self.assertEqual(request["model"], model_key) + self.assertEqual(request["input"], str(image_path)) + self.assertIn(self.debug.DebugUpscale, self.addon.CLASSES) panel = importlib.import_module("anyimage.panel") - self.assertIn( - panel.RealesrganGeneralX4v3DebugPanel, - self.addon.CLASSES, - ) + self.assertIn(panel.UpscaleDebugPanel, self.addon.CLASSES) def test_default_request_sends_every_operator_property(self): class Example(self.job_runtime.JobOperatorBase): @@ -451,6 +422,10 @@ class AddonTest(unittest.TestCase): [ ("download-model", {"model": "MOGE2_VITS_NORMAL"}), ("download-model", {"model": "BEN2_BASE"}), + ( + "download-model", + {"model": "REALESRGAN_GENERAL_WDN_X4V3"}, + ), ], ) @@ -463,11 +438,22 @@ class AddonTest(unittest.TestCase): with self.assertRaisesRegex( RuntimeError, - "MoGe-2 ViT-S Normal: network unavailable; BEN2 Base: network unavailable", + ( + "MoGe-2 ViT-S Normal: network unavailable; " + "BEN2 Base: network unavailable; " + "Real-ESRGAN General WDN x4v3: network unavailable" + ), ): self.job_runtime.post_install(SimpleNamespace(request=request)) - self.assertEqual(calls, ["MOGE2_VITS_NORMAL", "BEN2_BASE"]) + self.assertEqual( + calls, + [ + "MOGE2_VITS_NORMAL", + "BEN2_BASE", + "REALESRGAN_GENERAL_WDN_X4V3", + ], + ) def test_environment_controller_runs_arbitrary_post_install_callable(self): runtime_module = importlib.import_module("blendjob.runtime") @@ -587,7 +573,13 @@ class AddonTest(unittest.TestCase): self.assertEqual( parameters, - {"models": ["MOGE2_VITS_NORMAL", "BEN2_BASE"]}, + { + "models": [ + "MOGE2_VITS_NORMAL", + "BEN2_BASE", + "REALESRGAN_GENERAL_WDN_X4V3", + ] + }, ) def test_ai_status_checks_environment_and_required_model_files(self): @@ -600,6 +592,10 @@ class AddonTest(unittest.TestCase): return_value=( {"key": "MOGE2_VITS_NORMAL", "ready": True}, {"key": "BEN2_BASE", "ready": False}, + { + "key": "REALESRGAN_GENERAL_WDN_X4V3", + "ready": True, + }, ), ), ): @@ -675,12 +671,13 @@ class AddonTest(unittest.TestCase): def test_preferences_draws_upscale_model_download(self): preference_module = importlib.import_module("anyimage.preferences") downloads = [] + properties = [] class Layout: enabled = True - def prop(self, *_args, **_options): - return None + def prop(self, _owner, name, **_options): + properties.append(name) def separator(self, **_options): return None @@ -699,16 +696,15 @@ class AddonTest(unittest.TestCase): preferences = preference_module.AnyImagePreferences() preferences.layout = Layout() - catalog = ({ - "key": "REALESRGAN_GENERAL_X4V3", - "label": "Real-ESRGAN General x4v3", - "family": "upscale", - "enabled": True, - "ready": False, - },) + catalog = tuple({ + "key": key, "label": key, "family": "upscale", + "enabled": True, "ready": False, + } for key in ( + "REALESRGAN_GENERAL_WDN_X4V3", "SAFMN_REAL_X4", "REALESRGAN_X4PLUS", + )) with ( patch.object( - preference_module, + self.addon.properties, "ai_status", return_value={ "environment_ready": True, @@ -717,16 +713,26 @@ class AddonTest(unittest.TestCase): }, ), patch.object( - preference_module, + self.addon.properties, "model_catalog", return_value=catalog, ), - patch.object(preference_module.runtime, "server_busy", return_value=False), + patch.object(self.addon.runtime, "server_busy", return_value=False), ): preferences.draw(None) - self.assertEqual(len(downloads), 1) - self.assertEqual(downloads[0].model, "REALESRGAN_GENERAL_X4V3") + self.assertEqual( + [download.model for download in downloads], + [model["key"] for model in catalog], + ) + self.assertIn("upscale_model", properties) + self.assertIn("upscale_max_output_size", properties) + self.assertEqual( + preference_module.AnyImagePreferences.__annotations__[ + "upscale_model" + ]["name"], + "Default Upscale Model", + ) def test_mesh_plane_operator_is_registered_with_material_options(self): operator = self.mesh_plane.ConvertToMeshPlane @@ -846,7 +852,15 @@ class AddonTest(unittest.TestCase): "ready": False, }, ): - menu.draw(None) + menu.draw( + SimpleNamespace( + object=SimpleNamespace( + type="EMPTY", + empty_display_type="IMAGE", + data=SimpleNamespace(size=(640, 480)), + ) + ) + ) self.assertEqual( calls, @@ -854,7 +868,9 @@ class AddonTest(unittest.TestCase): ("anyimage.convert_to_mesh_plane", True), ("anyimage.convert_to_depth_plane", True), ("anyimage.convert_to_surface", True), + ("separator", True), ("anyimage.remove_image_background", True), + ("anyimage.upscale_image", True), ("separator", True), ("anyimage.setup_ai_environment", True), ], @@ -864,6 +880,10 @@ class AddonTest(unittest.TestCase): "Download Required Models…", ) + self.assertEqual( + labels["anyimage.upscale_image"], + "Upscale (640 × 480)", + ) def test_image_menu_opens_environment_settings_when_ai_is_ready(self): calls = [] @@ -893,6 +913,366 @@ class AddonTest(unittest.TestCase): self.assertIn(("anyimage.open_ai_environment_settings", True), calls) self.assertNotIn(("anyimage.setup_ai_environment", True), calls) + def test_upscale_operator_defaults_to_wdn_and_sends_selected_model(self): + operator_type = self.image_upscale.UpscaleImage + self.assertIn(operator_type, self.addon.CLASSES) + self.assertEqual( + operator_type.__annotations__["model"]["default"], + "REALESRGAN_GENERAL_WDN_X4V3", + ) + self.assertEqual( + operator_type.__annotations__["max_output_size"]["default"], + 4096, + ) + + self.assertIn(self.job_runtime.JobOperatorBase, operator_type.__mro__) + self.assertEqual(operator_type.bl_options, {"UNDO"}) + operator = operator_type() + operator.input_path = "input.png" + operator.source_object_name = "Source" + operator.model = "SAFMN_REAL_X4" + operator.max_output_size = 6144 + context = SimpleNamespace( + scene=SimpleNamespace( + anyimage_settings=SimpleNamespace(video_frames=7) + ) + ) + with ( + patch.object(self.image_upscale, "require_environment"), + patch.object(self.image_upscale, "require_model") as require_model, + patch.object( + self.image_upscale, "input_path", return_value=Path("input.png") + ), + patch.object( + self.image_upscale, "production_device", return_value="CUDA" + ), + patch.object( + self.image_upscale, "source_start_frame", return_value=4 + ), + ): + request = operator.request(context) + + require_model.assert_called_once_with("SAFMN_REAL_X4") + self.assertEqual( + request, + { + "input": "input.png", + "model": "SAFMN_REAL_X4", + "device": "cuda", + "video_frames": 7, + "start_frame": 4, + "max_output_size": 6144, + }, + ) + + def test_upscale_execute_reads_model_and_size_from_preferences(self): + source_image = SimpleNamespace( + name="Original", + source="FILE", + anyimage_upscale_source=None, + ) + source_object = SimpleNamespace( + name="Source", + type="EMPTY", + empty_display_type="IMAGE", + data=source_image, + ) + context = SimpleNamespace( + object=source_object, + scene=SimpleNamespace( + anyimage_settings=SimpleNamespace(video_frames=1) + ), + ) + operator = self.image_upscale.UpscaleImage() + with ( + patch.object( + self.image_upscale, + "configured_upscale_model", + return_value="SAFMN_REAL_X4", + ), + patch.object( + self.image_upscale, + "configured_upscale_max_output_size", + return_value=6144, + ), + patch.object( + self.image_upscale, + "prepare_image_input", + return_value=(Path("input.png"), False), + ), + patch.object( + self.job_runtime.JobOperatorBase, + "execute", + return_value={"RUNNING_MODAL"}, + ), + ): + result = operator.execute(context) + + self.assertEqual(result, {"RUNNING_MODAL"}) + self.assertEqual(operator.model, "SAFMN_REAL_X4") + self.assertEqual(operator.max_output_size, 6144) + + def test_upscale_operator_opens_setup_when_required_ai_is_missing(self): + expected = {"CANCELLED"} + with patch.object( + self.image_upscale, + "invoke_ai_setup_if_needed", + return_value=expected, + ) as setup: + result = self.image_upscale.UpscaleImage().invoke(None, None) + + setup.assert_called_once_with() + self.assertIs(result, expected) + + def test_upscale_operator_runs_immediately_when_ai_is_ready(self): + operator = self.image_upscale.UpscaleImage() + context = object() + expected = {"FINISHED"} + with ( + patch.object( + self.image_upscale, + "invoke_ai_setup_if_needed", + return_value=None, + ), + patch.object( + operator, + "execute", + return_value=expected, + ) as execute, + ): + result = operator.invoke(context, None) + + execute.assert_called_once_with(context) + self.assertIs(result, expected) + self.assertNotIn("REGISTER", operator.bl_options) + + def test_upscale_response_replaces_source_image(self): + original_image = SimpleNamespace(name="Original") + source_object = SimpleNamespace( + type="EMPTY", + empty_display_type="IMAGE", + data=original_image, + ) + self.fake_bpy.data = SimpleNamespace( + objects=SimpleNamespace( + get=lambda name: source_object if name == "Source" else None + ), + images=SimpleNamespace( + get=lambda name: original_image if name == "Original" else None + ), + ) + operator = self.image_upscale.UpscaleImage() + operator.source_object_name = "Source" + operator.source_image_name = "Original" + operator.model = "REALESRGAN_GENERAL_WDN_X4V3" + result = SimpleNamespace(directory=Path("result")) + result_image = SimpleNamespace(name="", anyimage_upscale_source=None) + with patch.object( + self.image_upscale, + "create_color_image", + return_value=result_image, + ) as create_color: + message = operator.response(None, result) + + create_color.assert_called_once_with( + source_object, + result.directory, + frame_prefix="upscale.", + ) + self.assertEqual(result_image.name, "Original.upscale") + self.assertIs( + result_image.anyimage_upscale_source, + original_image, + ) + self.assertEqual( + message, + "Image is upscaled with Real-ESRGAN General WDN x4v3", + ) + + def test_upscale_response_reports_removed_original_image(self): + source_object = SimpleNamespace( + type="EMPTY", + empty_display_type="IMAGE", + data=object(), + ) + self.fake_bpy.data = SimpleNamespace( + objects=SimpleNamespace(get=lambda _name: source_object), + images=SimpleNamespace(get=lambda _name: None), + ) + operator = self.image_upscale.UpscaleImage() + operator.source_object_name = "Source" + operator.source_image_name = "Removed" + + with self.assertRaisesRegex( + RuntimeError, + "original source image is no longer available", + ): + operator.response(None, SimpleNamespace(directory=Path("result"))) + + def test_repeated_upscale_uses_original_image_and_reports_warning(self): + original_image = SimpleNamespace(name="Original", source="FILE") + current_image = SimpleNamespace( + source="FILE", + anyimage_upscale_source=original_image, + ) + source_object = SimpleNamespace( + name="Source", + type="EMPTY", + empty_display_type="IMAGE", + data=current_image, + ) + context = SimpleNamespace( + object=source_object, + scene=SimpleNamespace( + anyimage_settings=SimpleNamespace(video_frames=1) + ), + ) + reports = [] + operator = self.image_upscale.UpscaleImage() + operator.report = lambda levels, message: reports.append((levels, message)) + with ( + patch.object( + self.image_upscale, + "prepare_image_input", + return_value=(Path("original.png"), False), + ) as prepare, + patch.object( + self.job_runtime.JobOperatorBase, + "execute", + return_value={"RUNNING_MODAL"}, + ), + ): + result = operator.execute(context) + + self.assertEqual(result, {"RUNNING_MODAL"}) + prepare.assert_called_once_with(original_image, context.scene) + self.assertEqual(operator.source_image_name, "Original") + self.assertIn("rebuilding from the original", reports[0][1]) + + def test_packed_image_export_does_not_use_scene_render_transform(self): + with tempfile.TemporaryDirectory() as directory: + output_directory = Path(directory) + calls = [] + + class Image: + file_format = "JPEG" + + def save(self, **options): + calls.append(options) + Path(options["filepath"]).write_bytes(b"raw image") + + image = Image() + scene = SimpleNamespace( + view_settings=SimpleNamespace(look="AgX - Medium High Contrast") + ) + with patch.object( + self.image_data.tempfile, + "mkdtemp", + return_value=str(output_directory), + ): + output = self.image_data.save_packed_image(image, scene) + + self.assertEqual(output, output_directory / "input.png") + self.assertEqual( + calls, + [{"filepath": str(output), "save_copy": True}], + ) + self.assertEqual(image.file_format, "JPEG") + + def test_color_result_inherits_source_color_space(self): + with tempfile.TemporaryDirectory() as directory: + output_directory = Path(directory) + (output_directory / "image.0001.png").touch() + source_image = SimpleNamespace( + name="source.png", + source="FILE", + use_fake_user=False, + colorspace_settings=SimpleNamespace(name="ACEScg"), + ) + + class ResultImage: + def __init__(self): + self.alpha_mode = "" + self.colorspace_settings = SimpleNamespace(name="sRGB") + self.name = "" + + def pack(self): + pass + + result_image = ResultImage() + source_object = SimpleNamespace( + name="Source", + type="EMPTY", + empty_display_type="IMAGE", + data=source_image, + ) + self.fake_bpy.data = SimpleNamespace( + images=SimpleNamespace( + load=lambda *_args, **_options: result_image, + remove=lambda *_args, **_options: None, + ) + ) + + created = self.image_data.create_color_image( + source_object, + output_directory, + ) + + self.assertIs(created, result_image) + self.assertEqual(result_image.colorspace_settings.name, "ACEScg") + self.assertEqual(result_image.name, "source.color") + self.assertIs(source_object.data, result_image) + + def test_background_result_replaces_only_an_unshared_image(self): + for remaining_users, removed_source, result_name in ( + (0, True, "Source.png"), + (1, False, "Source.color"), + ): + with self.subTest(remaining_users=remaining_users): + source_image = SimpleNamespace( + name="Source.png", + use_fake_user=False, + users=remaining_users, + ) + source_object = SimpleNamespace( + name="Source", + type="EMPTY", + empty_display_type="IMAGE", + data=source_image, + ) + result_image = SimpleNamespace(name="Source.color") + removed = [] + self.fake_bpy.data = SimpleNamespace( + objects={"Source": source_object}, + images=SimpleNamespace( + remove=lambda image: removed.append(image), + ), + ) + + def create_result(target, _directory): + source_image.use_fake_user = True + target.data = result_image + return result_image + + with ( + patch.object( + self.image_data, + "create_color_image", + side_effect=create_result, + ), + patch.object(self.image_background, "push_undo") as push_undo, + ): + self.image_background.apply_background_response( + SimpleNamespace(directory=Path("result")), + "Source", + ) + + self.assertIs(source_object.data, result_image) + self.assertFalse(source_image.use_fake_user) + self.assertEqual(removed, [source_image] if removed_source else []) + self.assertEqual(result_image.name, result_name) + push_undo.assert_called_once_with("Remove Background") + def test_balloon_contour_is_smooth_and_evenly_resampled(self): import numpy as np @@ -970,6 +1350,34 @@ class AddonTest(unittest.TestCase): np.zeros((2, 2), dtype=np.float32), ) + def test_expected_empty_image_edits_are_reported_as_warnings(self): + reports = [] + operator = SimpleNamespace( + report=lambda level, message: reports.append((level, message)) + ) + + self.image_edit_common.report_image_edit_exception( + operator, + self.image_edit_common.ImageEditWarning( + "The edit contains no visible image pixels" + ), + ) + self.image_edit_common.report_image_edit_exception( + operator, + ValueError("Unexpected invalid data"), + ) + + self.assertEqual( + reports, + [ + ( + {"WARNING"}, + "The edit contains no visible image pixels", + ), + ({"ERROR"}, "Unexpected invalid data"), + ], + ) + def test_image_mask_crops_to_visible_alpha_without_padding(self): import numpy as np @@ -1002,28 +1410,132 @@ class AddonTest(unittest.TestCase): expected[1:3, 1:3] = top_down[:, :, 3] np.testing.assert_allclose(mask, expected) - def test_image_edit_aligns_canvas_without_changing_object_transform(self): - matrix = object() - image_object = SimpleNamespace( - matrix_world=matrix, - empty_display_size=4.0, + def test_image_edit_replaces_an_unshared_image_without_leaving_a_copy(self): + class Pixels: + values = None + + def foreach_set(self, values): + self.values = tuple(values) + + source_image = SimpleNamespace( + name="Source.png", + size=(4, 2), + colorspace_settings=SimpleNamespace(name="sRGB"), + use_fake_user=False, + users=0, + ) + source_object = SimpleNamespace( + type="EMPTY", + empty_display_type="IMAGE", + data=source_image, + empty_display_size=1.0, empty_image_offset=(0.0, 0.0), ) + result_image = SimpleNamespace( + name="", + alpha_mode="", + colorspace_settings=SimpleNamespace(name=""), + pixels=Pixels(), + pack=lambda: packed.append(True), + update=lambda: updated.append(True), + ) + packed = [] + updated = [] + removed = [] + self.fake_bpy.data = SimpleNamespace( + objects=[source_object], + images=SimpleNamespace( + new=lambda *_args, **_options: result_image, + remove=lambda image, **_options: removed.append(image), + ), + ) + pixels = tuple(float(index) for index in range(16)) with patch.object( self.image_edit_common, - "plane_bounds", - return_value=(-2.0, 2.0, -1.0, 1.0), + "empty_display_alignment", + return_value=(2.0, (0.25, -0.25)), ): - self.image_edit_common.align_empty_display_to_bounds( - image_object, - (100, 50), - (50, 25), - (50, 0, 100, 25), + result = self.image_edit_common.replace_empty_image_pixels( + source_object, pixels, (2, 2), (0, 0, 2, 2) ) - self.assertIs(image_object.matrix_world, matrix) - self.assertEqual(image_object.empty_display_size, 2.0) - self.assertEqual(image_object.empty_image_offset, (0.0, 0.0)) + self.assertIs(result, result_image) + self.assertIs(source_object.data, result_image) + self.assertEqual(removed, [source_image]) + self.assertEqual(result_image.name, "Source.png") + self.assertEqual(result_image.pixels.values, pixels) + self.assertEqual(packed, [True]) + self.assertEqual(updated, [True]) + self.assertEqual(source_object.empty_display_size, 2.0) + self.assertEqual(source_object.empty_image_offset, (0.25, -0.25)) + + def test_image_edit_forks_an_image_shared_by_multiple_empties(self): + class Pixels: + values = None + + def foreach_set(self, values): + self.values = tuple(values) + + class ResultImage: + colorspace_settings = SimpleNamespace(name="") + pixels = Pixels() + alpha_mode = "" + + def pack(self): + pass + + def update(self): + pass + + source_image = SimpleNamespace( + name="Source.png", + size=(4, 2), + colorspace_settings=SimpleNamespace(name="ACEScg"), + use_fake_user=False, + users=1, + ) + source_object = SimpleNamespace( + type="EMPTY", + empty_display_type="IMAGE", + data=source_image, + empty_display_size=1.0, + empty_image_offset=(0.0, 0.0), + ) + sibling = SimpleNamespace( + type="EMPTY", + empty_display_type="IMAGE", + data=source_image, + ) + result_image = ResultImage() + created = [] + + def new_image(*args, **options): + created.append((args, options)) + return result_image + + self.fake_bpy.data = SimpleNamespace( + objects=[source_object, sibling], + images=SimpleNamespace( + new=new_image, + remove=lambda *_args, **_options: None, + ), + ) + pixels = tuple(float(index) for index in range(16)) + with patch.object( + self.image_edit_common, + "empty_display_alignment", + return_value=(2.0, (0.25, -0.25)), + ): + result = self.image_edit_common.replace_empty_image_pixels( + source_object, pixels, (2, 2), (0, 0, 2, 2) + ) + + self.assertIs(result, result_image) + self.assertIs(source_object.data, result_image) + self.assertIs(sibling.data, source_image) + self.assertEqual(len(created), 1) + self.assertEqual(result_image.colorspace_settings.name, "ACEScg") + self.assertEqual(result_image.pixels.values, pixels) def test_image_edit_tools_and_operators_are_registered(self): import inspect @@ -1099,6 +1611,17 @@ class AddonTest(unittest.TestCase): self.assertIn("math.dist(self._path[0], self._cursor) < 2.0", modal_source) self.assertIn("if len(line) != 2", projection_source) + def test_image_area_click_is_cancelled_without_a_warning(self): + import inspect + + completion_source = inspect.getsource( + self.image_mask.EditImageMask._complete + ) + + self.assertIn("if len(path) < 3", completion_source) + self.assertIn('return {"CANCELLED"}', completion_source) + self.assertNotIn("Draw a larger image region", completion_source) + def test_image_polyline_closes_from_start_or_enter(self): import inspect @@ -1336,6 +1859,20 @@ class AddonTest(unittest.TestCase): self.assertAlmostEqual(ratio, 2.0) + def test_image_perspective_ignores_an_already_removed_preview(self): + operator = self.image_perspective.EditImagePerspective() + operator._preview_image = object() + + def remove(*_args, **_options): + raise ReferenceError("StructRNA of type Image has been removed") + + self.fake_bpy.data = SimpleNamespace( + images=SimpleNamespace(remove=remove) + ) + operator._remove_preview() + + self.assertIsNone(operator._preview_image) + def test_image_perspective_maps_crop_to_source_plane_bounds(self): bounds = self.image_perspective.perspective_placement_bounds( ((10, 20), (30, 20), (30, 40), (10, 40)), @@ -2081,81 +2618,42 @@ class AddonTest(unittest.TestCase): "O Texture Layer", ) - def test_balloon_material_reuse_requires_matching_image_and_node_group(self): + def test_balloon_material_reuse_uses_explicit_source_markers(self): class MarkedMaterial(dict): def __init__(self, name, **properties): super().__init__(properties) self.name = name - self.node_tree = None color_image = SimpleNamespace(name="cow.color") other_image = SimpleNamespace(name="other") - texture_node = SimpleNamespace(type="TEX_IMAGE", image=color_image) - texture_layer = SimpleNamespace( - type="GROUP", - node_tree=SimpleNamespace(name="O Texture Layer"), - ) - compatible = SimpleNamespace( - name="cow.001", - node_tree=SimpleNamespace(nodes=[texture_node, texture_layer]), - ) marked = MarkedMaterial( - "unrelated-name", + "color-material", anyimage_balloon_material=True, anyimage_source_image=color_image, ) marked_other = MarkedMaterial( - "cow", + "other-material", anyimage_balloon_material=True, anyimage_source_image=other_image, ) normal_image = SimpleNamespace(name="cow.normal") marked_normal = MarkedMaterial( - "depth-material", + "normal-material", anyimage_balloon_material=True, anyimage_source_image=color_image, anyimage_balloon_normal_image=normal_image, ) - wrong_group = SimpleNamespace( - name="cow", - node_tree=SimpleNamespace( - nodes=[ - texture_node, - SimpleNamespace( - type="GROUP", - node_tree=SimpleNamespace(name="O DA3 Surface Layer"), - ), - ] - ), - ) self.fake_bpy.data = SimpleNamespace( - objects=SimpleNamespace(get=lambda _name: None), - materials=[ - marked_other, - wrong_group, - compatible, - marked_normal, - marked, - ], + materials=[marked_other, marked_normal, marked], images=[color_image], ) self.assertIs( - self.balloon.reusable_balloon_material( - "cow", - "cow.color", - color_image, - ), + self.balloon.reusable_balloon_material(color_image), marked, ) - self.assertIs( - self.balloon.reusable_balloon_material("cow", "cow.color"), - compatible, - ) self.assertIs( self.balloon.reusable_balloon_material( - "cow", - "cow.color", color_image, normal_required=True, ), @@ -2221,10 +2719,13 @@ class AddonTest(unittest.TestCase): root = Path(directory) moge = root / "models" / "moge-2-vits-normal-onnx" ben2 = root / "models" / "BEN2-ONNX" / "onnx" + wdn = root / "models" / "realesr-general-wdn-x4v3" moge.mkdir(parents=True) ben2.mkdir(parents=True) + wdn.mkdir(parents=True) (moge / "model.onnx").touch() (ben2 / "model_fp16.onnx").touch() + (wdn / "realesr-general-wdn-x4v3.onnx").touch() with ( patch.object( @@ -2310,6 +2811,11 @@ class AddonTest(unittest.TestCase): self.addon.clipboard_image.TrackNativeCopy, self.addon.CLASSES, ) + existing_context_draw = object() + existing_outliner_draw = object() + self.context_menu_draws.append(existing_context_draw) + self.outliner_menu_draws.append(existing_outliner_draw) + self.addon.register() system_classes = self.addon.runtime.operator_classes() @@ -2361,11 +2867,35 @@ class AddonTest(unittest.TestCase): ) self.assertEqual(len(self.status_bar_draws), 1) self.assertEqual(self.status_bar_draws[0]._owner, "anyimage") + self.assertIs( + self.context_menu_draws[0], + self.addon.draw_image_context_menu, + ) + self.assertIs( + self.outliner_menu_draws[0], + self.addon.draw_image_context_menu, + ) + self.assertTrue( + hasattr(self.fake_bpy.types.Image, "anyimage_upscale_source") + ) self.addon.unregister() self.assertEqual(self.registered, []) self.assertEqual(self.registered_tools, []) self.assertEqual(self.status_bar_draws, []) + self.assertEqual(self.context_menu_draws, [existing_context_draw]) + self.assertEqual(self.outliner_menu_draws, [existing_outliner_draw]) + self.assertFalse( + hasattr(self.fake_bpy.types.Image, "anyimage_upscale_source") + ) + + def test_blender_image_exports_do_not_use_scene_render_color_management(self): + source_root = PROJECT_ROOT / "src" / "anyimage" + offenders = [] + for path in source_root.rglob("*.py"): + if ".save_render(" in path.read_text(encoding="utf-8"): + offenders.append(path.relative_to(source_root)) + self.assertEqual(offenders, []) def test_status_bar_reads_runtime_state(self): calls = [] diff --git a/tests/test_download_progress.py b/tests/test_download_progress.py index b947823..2d71931 100644 --- a/tests/test_download_progress.py +++ b/tests/test_download_progress.py @@ -2,6 +2,7 @@ import os import sys import unittest from pathlib import Path +import tempfile from types import SimpleNamespace from unittest.mock import patch @@ -146,6 +147,26 @@ class DownloadProgressTest(unittest.TestCase): self.assertEqual(endpoints, ["https://hf.example"]) + def test_declared_download_skips_hf_when_no_public_onnx_exists(self): + args = SimpleNamespace( + repository="", + progress=lambda *_args: None, + cancel_check=lambda: None, + ) + with tempfile.TemporaryDirectory() as directory, patch.object( + self.download, "download_from_hf" + ) as hugging_face, patch.object( + self.download, "download_from_r2" + ) as mirror, patch.object( + self.download, "_verify_model_complete" + ): + self.download._download_declared_or_fail( + args, Path(directory), "Model", 0.05, 0.90 + ) + + hugging_face.assert_not_called() + mirror.assert_called_once() + if __name__ == "__main__": unittest.main() diff --git a/tests/test_models.py b/tests/test_models.py index b692c51..3d47280 100644 --- a/tests/test_models.py +++ b/tests/test_models.py @@ -20,24 +20,10 @@ class ModelTest(unittest.TestCase): @classmethod def setUpClass(cls): cls.models = load_module( - "anyimage_models", - "src/anyimage/model_catalog.py", - ) - cls.server_models = load_module( "anyimage_server_models", "src/anyimage/server/model_catalog.py", ) - def test_client_and_server_catalogs_share_storage_contract(self): - client = self.models.DOWNLOADABLE_MODELS - server = self.server_models.DOWNLOADABLE_MODELS - - self.assertEqual(set(client), set(server)) - for key in client: - with self.subTest(key=key): - self.assertEqual(client[key].directory_name, server[key].directory_name) - self.assertEqual(client[key].ready_patterns, server[key].ready_patterns) - def test_base_model_is_default_friendly_apache_model(self): model = self.models.get_model("BASE") @@ -89,33 +75,26 @@ class ModelTest(unittest.TestCase): "BEN2-ONNX", ) - def test_realesrgan_general_x4v3_uses_verified_onnx_artifact(self): - model = self.models.get_downloadable_model( - "REALESRGAN_GENERAL_X4V3" - ) - + 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"), + "REALESRGAN_X4PLUS": ("realesrgan-x4plus.onnx", 67051618, "4ed6a45a"), + } + 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.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") + wdn = self.models.UPSCALE_MODELS["REALESRGAN_GENERAL_WDN_X4V3"] + self.assertEqual(wdn.repository, "xinntao/Real-ESRGAN") + self.assertEqual(wdn.huggingface_repository, "") self.assertEqual( - model.repository, - "jonathanst29/tinier-upscale-models", - ) - self.assertEqual(model.license_name, "BSD 3-Clause") - self.assertEqual( - model.download_patterns, - ("realesr-general-x4v3.onnx",), - ) - self.assertEqual(model.ready_patterns, model.download_patterns) - self.assertEqual(model.r2_directory, "realesr-general-x4v3") - self.assertEqual( - model.r2_files, - (( - "realesr-general-x4v3.onnx", - 4866396, - "924ebad6532777303582d4ce7811849b88869231a3ae7093e0f21200249df8d5", - ),), - ) - self.assertEqual( - self.models.model_record(model)["family"], - "upscale", + self.models.UPSCALE_MODELS["SAFMN_REAL_X4"].huggingface_repository, "" ) def test_enabled_da3_models_use_onnx_and_r2_checksums(self): diff --git a/tests/test_pack.py b/tests/test_pack.py index 1d922cc..0378384 100644 --- a/tests/test_pack.py +++ b/tests/test_pack.py @@ -25,7 +25,7 @@ class PackTest(unittest.TestCase): pack = load_pack() with tempfile.TemporaryDirectory() as directory: directory_path = Path(directory) - wheel_path = directory_path / "blendjob-0.1.13-py3-none-any.whl" + wheel_path = directory_path / "blendjob-0.1.14-py3-none-any.whl" wheel_path.write_bytes(b"wheel") original_dist = pack.DIST_DIR pack.DIST_DIR = directory_path / "dist" @@ -95,21 +95,21 @@ class PackTest(unittest.TestCase): self.assertIn("server/jobs/debug_da3.py", names) self.assertIn("server/jobs/debug_moge2.py", names) self.assertIn("server/jobs/debug_ben2.py", names) - self.assertIn("server/jobs/debug_realesrgan.py", names) + self.assertIn("server/jobs/debug_upscale.py", names) self.assertNotIn("server/processing.py", names) self.assertNotIn("server/progress.py", names) - self.assertIn("model_catalog.py", names) + 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_da3.py", names) self.assertIn("server/models/onnx_moge2.py", names) - self.assertIn("server/models/onnx_realesrgan.py", names) - self.assertIn("server/models/realesrgan.py", names) + self.assertIn("server/models/onnx_upscale.py", names) + self.assertIn("server/models/upscale.py", names) self.assertNotIn("server/da3.py", names) self.assertNotIn("server/moge2.py", names) self.assertIn("server/__init__.py", names) self.assertIn( - "wheels/blendjob-0.1.13-py3-none-any.whl", + "wheels/blendjob-0.1.14-py3-none-any.whl", names, ) self.assertFalse( @@ -134,7 +134,7 @@ class PackTest(unittest.TestCase): def test_reads_extension_dependencies_from_pyproject(self): pack = load_pack() - self.assertEqual(pack.project_dependencies(), ["blendjob==0.1.13"]) + self.assertEqual(pack.project_dependencies(), ["blendjob==0.1.14"]) def test_downloads_declared_wheels_for_target_platform(self): pack = load_pack() @@ -143,7 +143,7 @@ class PackTest(unittest.TestCase): def fake_run(command, check): self.assertTrue(check) - (destination / "blendjob-0.1.13-py3-none-any.whl").write_bytes( + (destination / "blendjob-0.1.14-py3-none-any.whl").write_bytes( b"wheel" ) @@ -151,12 +151,14 @@ class PackTest(unittest.TestCase): wheels = pack.download_wheels("linux-x64", destination) command = run.call_args.args[0] - self.assertIn("blendjob==0.1.13", command) + self.assertIn("blendjob==0.1.14", command) + self.assertIn("--find-links", command) + self.assertIn(str(pack.SOURCE_DIR / "wheels"), command) self.assertIn("--python-version=3.11", command) self.assertIn("--platform=manylinux2014_x86_64", command) self.assertEqual( [wheel.name for wheel in wheels], - ["blendjob-0.1.13-py3-none-any.whl"], + ["blendjob-0.1.14-py3-none-any.whl"], ) def test_rejects_unknown_platform(self): pack = load_pack() diff --git a/tests/test_realesrgan_onnx.py b/tests/test_realesrgan_onnx.py deleted file mode 100644 index 262dbbf..0000000 --- a/tests/test_realesrgan_onnx.py +++ /dev/null @@ -1,75 +0,0 @@ -import sys -import unittest -from pathlib import Path -from types import SimpleNamespace - -import numpy as np -from PIL import Image - - -SOURCE_DIR = ( - Path(__file__).resolve().parents[1] - / "src" - / "anyimage" - / "server" - / "models" -) -sys.path.insert(0, str(SOURCE_DIR)) - -import onnx_realesrgan - - -class FakeSession: - def get_inputs(self): - return [SimpleNamespace(name="input")] - - def get_outputs(self): - return [SimpleNamespace(name="output")] - - def run(self, output_names, feeds): - self.output_names = output_names - self.feeds = feeds - source = feeds["input"] - return [np.repeat(np.repeat(source, 4, axis=2), 4, axis=3)] - - -class RealesrganOnnxTest(unittest.TestCase): - def test_preprocess_matches_dynamic_rgb_layout(self): - image = Image.new("RGB", (7, 5), (64, 128, 192)) - - result = onnx_realesrgan.preprocess(image) - - self.assertEqual(result.shape, (1, 3, 5, 7)) - self.assertEqual(result.dtype, np.float32) - self.assertTrue(result.flags.c_contiguous) - np.testing.assert_allclose( - result[0, :, 0, 0], - np.array((64, 128, 192), dtype=np.float32) / 255.0, - ) - - def test_infer_returns_direct_four_x_rgb_prediction(self): - session = FakeSession() - image = Image.new("RGB", (7, 5), (10, 20, 30)) - - result = onnx_realesrgan.infer(session, image) - - self.assertEqual(result.mode, "RGB") - self.assertEqual(result.size, (28, 20)) - self.assertEqual(session.output_names, ["output"]) - self.assertEqual(session.feeds["input"].shape, (1, 3, 5, 7)) - self.assertEqual(result.getpixel((0, 0)), (10, 20, 30)) - - def test_infer_rejects_non_four_x_output(self): - class InvalidSession(FakeSession): - def run(self, _output_names, feeds): - return [feeds["input"]] - - with self.assertRaisesRegex(RuntimeError, "unexpected output shape"): - onnx_realesrgan.infer( - InvalidSession(), - Image.new("RGB", (7, 5)), - ) - - -if __name__ == "__main__": - unittest.main() diff --git a/tests/test_server.py b/tests/test_server.py index 771525f..6313e5f 100644 --- a/tests/test_server.py +++ b/tests/test_server.py @@ -42,14 +42,21 @@ class ServerTest(unittest.TestCase): @classmethod def setUpClass(cls): cls.server_module = load_server_module() + from blendjob import JobServer + from server.model_manager import ModelManager, ModelSessionCache + cls.JobServer = JobServer + + cls.ModelManager, cls.ModelSessionCache = ModelManager, ModelSessionCache + cls.server = cls.server_module.server + cls.app_module = sys.modules["server.app"] def setUp(self): self.jobs_temp = tempfile.TemporaryDirectory() self.addCleanup(self.jobs_temp.cleanup) self.storage_root = Path(self.jobs_temp.name) def test_required_model_download_skips_ready_model(self): - manager = self.server_module.ModelManager(Path(".")) + manager = self.ModelManager(Path(".")) downloads = [] context = SimpleNamespace(progress=lambda *_args: None) with ( @@ -75,7 +82,7 @@ class ServerTest(unittest.TestCase): def test_models_are_owned_by_the_server_resource(self): self.assertFalse((PROJECT_ROOT / "src" / "anyimage" / "models.py").exists()) with tempfile.TemporaryDirectory() as directory: - manager = self.server_module.ModelManager(Path(directory)) + manager = self.ModelManager(Path(directory)) snapshot = manager.snapshot() self.assertEqual( snapshot, @@ -84,7 +91,7 @@ class ServerTest(unittest.TestCase): "ben2": "", "da3": "", "moge2": "", - "realesrgan": "", + "upscale": "", } }, ) @@ -92,8 +99,8 @@ class ServerTest(unittest.TestCase): self.assertNotIn("model_catalog", sys.modules) def create_cache(self): - cache = self.server_module.ModelSessionCache.__new__( - self.server_module.ModelSessionCache + cache = self.ModelSessionCache.__new__( + self.ModelSessionCache ) cache.ben2_model = None cache.ben2_key = None @@ -101,13 +108,13 @@ class ServerTest(unittest.TestCase): cache.da3_key = None cache.moge_model = None cache.moge_key = None - cache.realesrgan_model = None - cache.realesrgan_key = None + cache.upscale_model = None + cache.upscale_key = None return cache def create_server(self): source = self.server_module.server - server = self.server_module.JobServer( + server = self.JobServer( "AnyImage Job Server", storage_root=self.storage_root ) server.handlers.update(source.handlers) @@ -163,7 +170,7 @@ class ServerTest(unittest.TestCase): self.assertEqual(loads, [("model", "cuda")]) self.assertEqual(cache.snapshot()["da3"], "model") - def test_depth_model_cache_keeps_only_da3_or_moge(self): + def test_model_cache_keeps_only_the_most_recent_family(self): cache = self.create_cache() cache.ben2_model = object() cache.ben2_key = ("ben2", "cuda") @@ -181,7 +188,7 @@ class ServerTest(unittest.TestCase): self.assertIs(cache.get_da3("da3", device), da3) self.assertIsNone(cache.moge_model) - self.assertIsNotNone(cache.ben2_model) + self.assertIsNone(cache.ben2_model) moge = object() @@ -191,29 +198,29 @@ class ServerTest(unittest.TestCase): self.assertIs(cache.get_moge2("moge", "cuda"), moge) self.assertIsNone(cache.da3_model) - self.assertIsNotNone(cache.ben2_model) + self.assertIsNone(cache.ben2_model) - def test_realesrgan_model_is_reused_for_matching_runtime_key(self): + def test_upscale_model_is_reused_for_matching_runtime_key(self): loads = [] - realesrgan = ModuleType("server.models.onnx_realesrgan") - realesrgan.create_session = lambda model_dir, device: ( + upscale = ModuleType("server.models.onnx_upscale") + upscale.create_session = lambda model_dir, device: ( loads.append((model_dir, device)) or object() ) cache = self.create_cache() with patch.dict( sys.modules, - {"server.models.onnx_realesrgan": realesrgan}, + {"server.models.onnx_upscale": upscale}, ): - first = cache.get_realesrgan("model", "cpu") - second = cache.get_realesrgan("model", "cpu") + first = cache.get_upscale("model", "cpu") + second = cache.get_upscale("model", "cpu") self.assertIs(first, second) self.assertEqual(loads, [("model", "cpu")]) - self.assertEqual(cache.snapshot()["realesrgan"], "model") + self.assertEqual(cache.snapshot()["upscale"], "model") def test_dedicated_server_runs_only_one_job_at_a_time(self): - server = self.server_module.JobServer("Test", storage_root=self.storage_root) + server = self.JobServer("Test", storage_root=self.storage_root) first_started = threading.Event() release_first = threading.Event() order = [] @@ -250,18 +257,19 @@ class ServerTest(unittest.TestCase): "analyze-balloon", "refine-image-mask", "remove-background", + "upscale-image", "extract-frames", "generate-depth", "generate-surface", "debug-da3", "debug-moge2", "debug-ben2", - "debug-realesrgan-general-x4v3", + "debug-upscale", }, ) def test_queued_job_can_be_cancelled_before_it_runs(self): - server = self.server_module.JobServer("Test", storage_root=self.storage_root) + server = self.JobServer("Test", storage_root=self.storage_root) first_started = threading.Event() release_first = threading.Event() executed = [] @@ -292,7 +300,7 @@ class ServerTest(unittest.TestCase): clear=lambda: None, close=lambda: None, ) - with patch.object(self.server_module, "ModelManager", return_value=fake_models): + with patch.object(self.app_module, "ModelManager", return_value=fake_models): server = self.create_server() started = threading.Event() @@ -317,7 +325,7 @@ class ServerTest(unittest.TestCase): ) def test_job_progress_is_kept_in_server_memory(self): - server = self.server_module.JobServer("Test", storage_root=self.storage_root) + server = self.JobServer("Test", storage_root=self.storage_root) @server.job("example") def example(context, _parameters): @@ -333,7 +341,7 @@ class ServerTest(unittest.TestCase): def test_runtime_storage_root_owns_the_jobs_directory(self): storage_root = self.storage_root - server = self.server_module.JobServer("Test") + server = self.JobServer("Test") server.bind(storage_root) @server.job("example") @@ -354,7 +362,7 @@ class ServerTest(unittest.TestCase): clear=lambda: events.append("clear"), close=lambda: events.append("close"), ) - server = self.server_module.JobServer("Test", storage_root=self.storage_root) + server = self.JobServer("Test", storage_root=self.storage_root) server.add_resource("model_manager", resource) @server.job("example") @@ -382,7 +390,7 @@ class ServerTest(unittest.TestCase): clear=lambda: None, close=lambda: None, ) - with patch.object(self.server_module, "ModelManager", return_value=fake_models): + with patch.object(self.app_module, "ModelManager", return_value=fake_models): server = self.create_server() with tempfile.TemporaryDirectory() as directory: root = Path(directory) @@ -400,7 +408,7 @@ class ServerTest(unittest.TestCase): job_id="job", ) status = self.wait_for_job(context) - self.assertTrue((context.directory / "color.0001.jpg").is_file()) + self.assertTrue((context.directory / "image.0001.jpg").is_file()) self.assertEqual(status["state"], "succeeded") diff --git a/tests/test_upscale_onnx.py b/tests/test_upscale_onnx.py new file mode 100644 index 0000000..17e9da8 --- /dev/null +++ b/tests/test_upscale_onnx.py @@ -0,0 +1,140 @@ +import sys +import unittest +from pathlib import Path +from types import SimpleNamespace + +import numpy as np +from PIL import Image + + +SOURCE_DIR = ( + Path(__file__).resolve().parents[1] + / "src" + / "anyimage" + / "server" + / "models" +) +sys.path.insert(0, str(SOURCE_DIR)) + +import onnx_upscale + + +class FakeSession: + def get_inputs(self): + return [SimpleNamespace(name="input")] + + def get_outputs(self): + return [SimpleNamespace(name="output")] + + def run(self, output_names, feeds): + self.output_names = output_names + self.feeds = feeds + source = feeds["input"] + return [np.repeat(np.repeat(source, 4, axis=2), 4, axis=3)] + + +class UpscaleOnnxTest(unittest.TestCase): + def test_preprocess_matches_dynamic_rgb_layout(self): + image = Image.new("RGB", (7, 5), (64, 128, 192)) + + result = onnx_upscale.preprocess(image) + + self.assertEqual(result.shape, (1, 3, 5, 7)) + self.assertEqual(result.dtype, np.float32) + self.assertTrue(result.flags.c_contiguous) + np.testing.assert_allclose( + result[0, :, 0, 0], + np.array((64, 128, 192), dtype=np.float32) / 255.0, + ) + + def test_infer_returns_direct_four_x_rgb_prediction(self): + session = FakeSession() + image = Image.new("RGB", (7, 5), (10, 20, 30)) + + result = onnx_upscale.infer(session, image) + + self.assertEqual(result.mode, "RGB") + self.assertEqual(result.size, (28, 20)) + self.assertEqual(session.output_names, ["output"]) + self.assertEqual(session.feeds["input"].shape, (1, 3, 256, 256)) + self.assertEqual(result.getpixel((0, 0)), (10, 20, 30)) + + def test_infer_returns_raw_rgb_for_rgba_input(self): + session = FakeSession() + pixels = np.zeros((5, 7, 4), dtype=np.uint8) + pixels[..., :3] = (10, 20, 30) + pixels[..., 3] = np.arange(7, dtype=np.uint8) * 40 + image = Image.fromarray(pixels, mode="RGBA") + + result = onnx_upscale.infer(session, image) + + self.assertEqual(result.mode, "RGB") + self.assertEqual(result.size, (28, 20)) + self.assertEqual(result.getpixel((0, 0)), (10, 20, 30)) + + def test_tiles_use_reflected_context_at_outer_image_edges(self): + class RecordingSession(FakeSession): + def run(self, output_names, feeds): + self.input = feeds["input"].copy() + return super().run(output_names, feeds) + + pixels = np.zeros((5, 7, 3), dtype=np.uint8) + pixels[..., 0] = np.arange(7, dtype=np.uint8) * 20 + session = RecordingSession() + + onnx_upscale.infer( + session, + Image.fromarray(pixels, mode="RGB"), + ) + + red = session.input[0, 0] + self.assertAlmostEqual(red[16, 22], 120 / 255.0) + self.assertAlmostEqual(red[16, 23], 100 / 255.0) + + def test_infer_rejects_non_four_x_output(self): + class InvalidSession(FakeSession): + def run(self, _output_names, feeds): + return [feeds["input"]] + + with self.assertRaisesRegex(RuntimeError, "unexpected output shape"): + onnx_upscale.infer( + InvalidSession(), + Image.new("RGB", (7, 5)), + ) + + def test_all_models_split_large_inputs_into_fixed_tiles(self): + class CountingSession(FakeSession): + def __init__(self): + self.calls = 0 + + def run(self, output_names, feeds): + self.calls += 1 + return super().run(output_names, feeds) + + session = CountingSession() + result = onnx_upscale.infer( + session, + Image.new("RGB", (449, 225), (1, 2, 3)), + ) + + self.assertEqual(result.size, (1796, 900)) + self.assertEqual(session.calls, 6) + self.assertEqual(session.feeds["input"].shape, (1, 3, 256, 256)) + + def test_provider_unicode_error_is_presented_as_resource_error(self): + class InvalidProviderSession(FakeSession): + def run(self, _output_names, _feeds): + raise UnicodeDecodeError("utf-8", b"\xc4", 0, 1, "invalid") + + with self.assertRaisesRegex( + onnx_upscale.UpscaleResourceError, + "provider failed", + ): + onnx_upscale.infer( + InvalidProviderSession(), + Image.new("RGB", (7, 5)), + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_worker.py b/tests/test_worker.py index f8d9dfc..5c34e77 100644 --- a/tests/test_worker.py +++ b/tests/test_worker.py @@ -81,6 +81,9 @@ class FakeModels: def get_da3(self, _model_directory, _device): return object() + def get_upscale(self, _model_directory, _device): + return object() + class FakeJobContext: def __init__(self, directory, job_type, models=None): @@ -88,6 +91,7 @@ class FakeJobContext: self.job_type = job_type self.models = models or FakeModels() self.status = {} + self.logs = [] def resource(self, name): if name != "model_manager": @@ -101,6 +105,9 @@ class FakeJobContext: **details, } + def log(self, message, *, level="INFO"): + self.logs.append((level, message)) + def check_cancelled(self): return None @@ -115,7 +122,8 @@ class WorkerTest(unittest.TestCase): cls.ben2 = load_server_module("models.ben2") cls.da3_model = load_server_module("models.da3") cls.moge2_model = load_server_module("models.moge2") - cls.realesrgan_model = load_server_module("models.realesrgan") + cls.upscale_model = load_server_module("models.upscale") + cls.onnx_upscale = load_server_module("models.onnx_upscale") cls.surface_outputs = load_server_module("outputs.surface") cls.debug_da3_outputs = load_server_module("outputs.debug_da3") cls.debug_moge2_outputs = load_server_module("outputs.debug_moge2") @@ -125,9 +133,10 @@ class WorkerTest(unittest.TestCase): cls.debug_da3_job = load_server_module("jobs.debug_da3") cls.debug_moge2_job = load_server_module("jobs.debug_moge2") cls.debug_ben2_job = load_server_module("jobs.debug_ben2") - cls.debug_realesrgan_job = load_server_module( - "jobs.debug_realesrgan" + cls.debug_upscale_job = load_server_module( + "jobs.debug_upscale" ) + cls.upscale_job = load_server_module("jobs.upscale") cls.frames_job = load_server_module("jobs.frames") cls.foreground_job = load_server_module("jobs.foreground") cls.download = load_server_module("model_download") @@ -276,7 +285,7 @@ class WorkerTest(unittest.TestCase): self.media.copy_color_images([source], output) self.assertEqual( - (output / "color.0001.png").read_bytes(), + (output / "image.0001.png").read_bytes(), b"frame", ) @@ -572,7 +581,7 @@ class WorkerTest(unittest.TestCase): {"input": str(input_path), "model": "BEN2", "device": "cpu"}, ) - self.assertTrue((root / "result" / "color.0001.png").is_file()) + self.assertTrue((root / "result" / "image.0001.png").is_file()) self.assertEqual( inference_options[0][1], @@ -721,10 +730,10 @@ class WorkerTest(unittest.TestCase): result_directory = context.directory self.assertTrue( - (result_directory / "color.0001.png").is_file() + (result_directory / "image.0001.png").is_file() ) self.assertTrue( - (result_directory / "color.0002.png").is_file() + (result_directory / "image.0002.png").is_file() ) self.assertEqual(len(inference_options), 2) @@ -755,7 +764,7 @@ class WorkerTest(unittest.TestCase): for path in result_directory.iterdir() if path.is_file() ), - ["color.0001.jpg", "color.0002.jpg", "color.0003.jpg"], + ["image.0001.jpg", "image.0002.jpg", "image.0003.jpg"], ) self.assertEqual(context.status["progress"], 1.0) self.assertIn("3 frame", context.status["message"]) @@ -890,7 +899,7 @@ class WorkerTest(unittest.TestCase): b"foreground", ) - def test_realesrgan_debug_saves_each_direct_api_prediction(self): + def test_upscale_debug_saves_each_direct_api_prediction(self): from PIL import Image with tempfile.TemporaryDirectory() as directory: @@ -899,17 +908,17 @@ class WorkerTest(unittest.TestCase): frames.mkdir() (frames / "a.png").touch() (frames / "b.png").touch() - context = FakeJobContext(root / "result", "debug-realesrgan") + context = FakeJobContext(root / "result", "debug-upscale") with patch.object( - self.debug_realesrgan_job, + self.debug_upscale_job, "infer_one", return_value=Image.new("RGB", (8, 8), (10, 20, 30)), ) as infer_one: - self.debug_realesrgan_job.run( + self.debug_upscale_job.run( context, { "input": str(frames), - "model": "REALESRGAN_GENERAL_X4V3", + "model": "REALESRGAN_GENERAL_WDN_X4V3", "video_frames": 2, "device": "cpu", }, @@ -922,6 +931,146 @@ class WorkerTest(unittest.TestCase): self.assertEqual(prediction.size, (8, 8)) self.assertEqual(prediction.mode, "RGB") + def test_upscale_job_saves_importable_upscale_sequence(self): + from PIL import Image + + with tempfile.TemporaryDirectory() as directory: + root = Path(directory) + frames = root / "frames" + frames.mkdir() + (frames / "a.png").touch() + (frames / "b.png").touch() + context = FakeJobContext(root / "result", "upscale-image") + with patch.object( + self.upscale_job, + "infer_one", + return_value=Image.new("RGB", (8, 8), (10, 20, 30)), + ) as infer_one: + result = self.upscale_job.run( + context, + { + "input": str(frames), + "model": "REALESRGAN_GENERAL_WDN_X4V3", + "video_frames": 2, + "device": "cpu", + }, + ) + + self.assertEqual(infer_one.call_count, 2) + self.assertEqual(result, {"upscale": "."}) + colors = sorted(context.directory.glob("upscale.*.png")) + self.assertEqual( + [path.name for path in colors], + ["upscale.0001.png", "upscale.0002.png"], + ) + self.assertIn( + ( + "INFO", + "Upscale model: REALESRGAN_GENERAL_WDN_X4V3; " + "maximum output size: 4096 px", + ), + context.logs, + ) + + def test_upscale_job_preserves_source_alpha(self): + from PIL import Image + + with tempfile.TemporaryDirectory() as directory: + root = Path(directory) + source_path = root / "source.png" + pixels = np.zeros((3, 5, 4), dtype=np.uint8) + pixels[..., :3] = (10, 20, 30) + pixels[..., 3] = np.arange(5, dtype=np.uint8) * 60 + Image.fromarray(pixels, mode="RGBA").save(source_path) + context = FakeJobContext(root / "result", "upscale-image") + prediction = Image.new("RGB", (20, 12), (10, 20, 30)) + + with patch.object( + self.onnx_upscale, + "infer", + return_value=prediction, + ): + self.upscale_job.run( + context, + { + "input": str(source_path), + "model": "REALESRGAN_GENERAL_WDN_X4V3", + "video_frames": 1, + "device": "cpu", + }, + ) + + with Image.open(context.directory / "upscale.0001.png") as result: + self.assertEqual(result.mode, "RGBA") + self.assertEqual(result.size, (20, 12)) + alpha = np.asarray(result.getchannel("A")) + self.assertLess( + int(alpha[:, 0].max()), + int(alpha[:, -1].min()), + ) + + def test_upscale_input_is_limited_before_four_x_inference(self): + from PIL import Image + + source = Image.new("RGB", (2000, 1000)) + limited = self.upscale_model._limited_input(source, 4096) + + self.assertEqual(limited.size, (1024, 512)) + square = self.upscale_model._limited_input( + Image.new("RGB", (3000, 3000)), + 8192, + ) + self.assertLessEqual( + square.width * square.height * 16, + self.upscale_model.MAX_OUTPUT_PIXELS, + ) + with self.assertRaisesRegex(ValueError, "between 1 and 8192"): + self.upscale_model._limited_input(source, 9000) + + def test_upscale_resource_error_releases_session_and_is_readable(self): + from PIL import Image + + class Models: + def __init__(self, directory): + self.model_directory = directory + self.released = False + + def directory(self, _model): + return self.model_directory + + def get_upscale(self, _directory, _device): + return object() + + def release_upscale(self): + self.released = True + + with tempfile.TemporaryDirectory() as directory: + path = Path(directory) / "input.png" + Image.new("RGB", (8, 8)).save(path) + models = Models(Path(directory)) + with patch.object( + self.onnx_upscale, + "infer", + side_effect=self.onnx_upscale.UpscaleResourceError( + "provider memory exhausted" + ), + ): + with self.assertRaisesRegex( + RuntimeError, + "Reduce Max Output Size or select CPU", + ): + self.upscale_model.infer_one( + models, + { + "model": "REALESRGAN_GENERAL_WDN_X4V3", + "device": "directml", + "max_output_size": 4096, + }, + path, + ) + + self.assertTrue(models.released) + def test_da3_removes_background_after_depth_using_original_input(self): order = [] inference_inputs = [] @@ -1047,10 +1196,10 @@ class WorkerTest(unittest.TestCase): result_directory = context.directory self.assertTrue( - (result_directory / "color.0001.jpg").is_file() + (result_directory / "image.0001.jpg").is_file() ) self.assertTrue( - (result_directory / "color.0002.jpg").is_file() + (result_directory / "image.0002.jpg").is_file() ) self.assertEqual(len(inference_inputs), 2) self.assertTrue(all(len(paths) == 1 for paths in inference_inputs)) diff --git a/uv.lock b/uv.lock index 5da2a15..a7dc22d 100644 --- a/uv.lock +++ b/uv.lock @@ -15,6 +15,53 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/3e/30/e900b21425a860e195f32e37657aa1f7c7f2b1bfb26f03ca209b90933c06/annotated_doc-0.0.5-py3-none-any.whl", hash = "sha256:117bac03a25ede5df5440e855b32d556049ca169ead221505badf432fed4b101", size = 5302, upload-time = "2026-07-28T13:50:57.239Z" }, ] +[[package]] +name = "anyimage" +version = "0.12.15" +source = { virtual = "." } +dependencies = [ + { name = "blendjob" }, +] + +[package.dev-dependencies] +dev = [ + { name = "av" }, + { name = "bpy" }, + { name = "huggingface-hub" }, + { name = "numpy" }, + { name = "onnxruntime" }, + { name = "pillow" }, + { name = "pip" }, + { name = "pytest" }, + { name = "trimesh" }, +] +docs = [ + { name = "mike" }, + { name = "mkdocs" }, + { name = "mkdocs-material" }, +] + +[package.metadata] +requires-dist = [{ name = "blendjob", specifier = "==0.1.14" }] + +[package.metadata.requires-dev] +dev = [ + { name = "av", specifier = "==17.0.0" }, + { name = "bpy", specifier = "==4.5.3" }, + { name = "huggingface-hub", specifier = "==1.4.1" }, + { name = "numpy", specifier = ">=1.26,<2.0" }, + { name = "onnxruntime", specifier = "==1.24.4" }, + { name = "pillow", specifier = "==12.1.1" }, + { name = "pip", specifier = ">=25,<27" }, + { name = "pytest" }, + { name = "trimesh", specifier = "==4.11.2" }, +] +docs = [ + { name = "mike", specifier = ">=2.1.3,<3" }, + { name = "mkdocs", specifier = ">=1.6.1,<2" }, + { name = "mkdocs-material", specifier = ">=9.6.16,<10" }, +] + [[package]] name = "anyio" version = "4.14.2" @@ -76,11 +123,10 @@ wheels = [ [[package]] name = "blendjob" -version = "0.1.13" -source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/90/ff/2bf54e6e7f70f529bc6a72c7ae1e622ee360a02f7534dd2e1b77059c96a0/blendjob-0.1.13.tar.gz", hash = "sha256:a964905e36fdbf114076384751faa4818f8a2385e59900c06ae403c575f7dcef", size = 21045, upload-time = "2026-08-18T05:38:51.704Z" } +version = "0.1.14" +source = { registry = "src/anyimage/wheels" } wheels = [ - { url = "https://files.pythonhosted.org/packages/86/71/3762d11ec687cd35996317f1ed335b091851de4d65801ff9bddc37a486c1/blendjob-0.1.13-py3-none-any.whl", hash = "sha256:2089042bc769ddd7424cb3829eb83e129be784e5dbdba6d2c0e775833e55d20c", size = 26363, upload-time = "2026-08-18T05:38:50.838Z" }, + { path = "blendjob-0.1.14-py3-none-any.whl" }, ] [[package]] @@ -1060,53 +1106,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/b7/ce/149a00dd41f10bc29e5921b496af8b574d8413afcd5e30dfa0ed46c2cc5e/six-1.17.0-py2.py3-none-any.whl", hash = "sha256:4721f391ed90541fddacab5acf947aa0d3dc7d27b2e1e8eda2be8970586c3274", size = 11050, upload-time = "2024-12-04T17:35:26.475Z" }, ] -[[package]] -name = "anyimage" -version = "0.12.15" -source = { virtual = "." } -dependencies = [ - { name = "blendjob" }, -] - -[package.dev-dependencies] -dev = [ - { name = "av" }, - { name = "bpy" }, - { name = "huggingface-hub" }, - { name = "numpy" }, - { name = "onnxruntime" }, - { name = "pillow" }, - { name = "pip" }, - { name = "pytest" }, - { name = "trimesh" }, -] -docs = [ - { name = "mike" }, - { name = "mkdocs" }, - { name = "mkdocs-material" }, -] - -[package.metadata] -requires-dist = [{ name = "blendjob", specifier = "==0.1.13" }] - -[package.metadata.requires-dev] -dev = [ - { name = "av", specifier = "==17.0.0" }, - { name = "bpy", specifier = "==4.5.3" }, - { name = "huggingface-hub", specifier = "==1.4.1" }, - { name = "numpy", specifier = ">=1.26,<2.0" }, - { name = "onnxruntime", specifier = "==1.24.4" }, - { name = "pillow", specifier = "==12.1.1" }, - { name = "pip", specifier = ">=25,<27" }, - { name = "pytest" }, - { name = "trimesh", specifier = "==4.11.2" }, -] -docs = [ - { name = "mike", specifier = ">=2.1.3,<3" }, - { name = "mkdocs", specifier = ">=1.6.1,<2" }, - { name = "mkdocs-material", specifier = ">=9.6.16,<10" }, -] - [[package]] name = "sympy" version = "1.14.0"