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.
This commit is contained in:
@@ -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` 和打包测试
|
||||
|
||||
36
README.md
36
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 镜像之间自动回退；某一来源暂时断线不会阻止后续模型尝试另一来源
|
||||
|
||||
@@ -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
|
||||
<job-id>/
|
||||
├── color.0001.jpg/png
|
||||
├── image.0001.jpg/png
|
||||
├── depth.0001.exr
|
||||
└── results.npz
|
||||
```
|
||||
|
||||
@@ -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.<repository>.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 模型仓库决定
|
||||
|
||||
12
docs/faq.md
12
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。
|
||||
|
||||
@@ -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)。
|
||||
|
||||
@@ -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`。
|
||||
|
||||
@@ -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 供纹理过滤使用。未形成有效区域的单击会静默结束；有效选区没有可见像素时只显示警告，不修改图片。
|
||||
|
||||
<div class="screenshot-placeholder" markdown>
|
||||
**截图占位：Polyline 与反选**
|
||||
|
||||
@@ -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 的基本工作方式。
|
||||
|
||||
@@ -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，并显示安装按钮。
|
||||
|
||||
@@ -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-id>/job.log`。
|
||||
- Environment 安装问题查看 Storage Root 下的 `install.log`。
|
||||
- 修改 Storage Root 后，确认你检查的是当前目录中的日志。
|
||||
|
||||
@@ -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 文件格式。
|
||||
|
||||
@@ -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。正式功能会等默认必需模型全部就绪后启用。
|
||||
|
||||
### 透明边缘需要调整
|
||||
|
||||
|
||||
@@ -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-id>/job.log`，因此可先从主线定位 Job ID，再查看单个任务。
|
||||
|
||||
Blender 正常退出或禁用扩展时会停止所属 Server；Server 也监视父进程，在异常退出后自行结束。
|
||||
|
||||
Environment 安装发生在 Server 启动前，使用 `install.log` 记录过程。安装完成后的模型下载是普通 Job，与推理共享队列、进度和取消机制。
|
||||
|
||||
@@ -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 目录中，便于调试和重新检查。
|
||||
|
||||
|
||||
23
docs/upscale.md
Normal file
23
docs/upscale.md
Normal file
@@ -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。
|
||||
@@ -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，并按源图中的选区范围近似更新显示位置与尺寸。
|
||||
|
||||
@@ -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
|
||||
|
||||
2
pack.py
2
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}",
|
||||
|
||||
@@ -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 = [
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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"
|
||||
|
||||
@@ -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()
|
||||
|
||||
@@ -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),
|
||||
}
|
||||
@@ -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,
|
||||
|
||||
@@ -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")
|
||||
|
||||
|
||||
|
||||
@@ -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,
|
||||
)
|
||||
|
||||
@@ -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():
|
||||
|
||||
@@ -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}"
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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")
|
||||
|
||||
@@ -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
|
||||
|
||||
|
||||
|
||||
@@ -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
|
||||
|
||||
|
||||
|
||||
@@ -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"
|
||||
|
||||
|
||||
@@ -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"}
|
||||
|
||||
|
||||
@@ -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)})."
|
||||
)
|
||||
)
|
||||
|
||||
158
src/anyimage/operators/upscale.py
Normal file
158
src/anyimage/operators/upscale.py
Normal file
@@ -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)
|
||||
@@ -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
|
||||
|
||||
@@ -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"):
|
||||
|
||||
@@ -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",
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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",)
|
||||
|
||||
92
src/anyimage/server/app.py
Normal file
92
src/anyimage/server/app.py
Normal file
@@ -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)
|
||||
|
||||
@@ -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": "."}
|
||||
|
||||
@@ -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(
|
||||
@@ -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"
|
||||
)
|
||||
|
||||
34
src/anyimage/server/jobs/upscale.py
Normal file
34
src/anyimage/server/jobs/upscale.py
Normal file
@@ -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": "."}
|
||||
@@ -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
|
||||
|
||||
@@ -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",),
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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",
|
||||
)
|
||||
|
||||
@@ -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")
|
||||
142
src/anyimage/server/models/onnx_upscale.py
Normal file
142
src/anyimage/server/models/onnx_upscale.py
Normal file
@@ -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
|
||||
|
||||
|
||||
@@ -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)
|
||||
76
src/anyimage/server/models/upscale.py
Normal file
76
src/anyimage/server/models/upscale.py
Normal file
@@ -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
|
||||
@@ -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):
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
@@ -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()
|
||||
|
||||
@@ -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):
|
||||
|
||||
@@ -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()
|
||||
|
||||
@@ -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()
|
||||
@@ -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")
|
||||
|
||||
|
||||
140
tests/test_upscale_onnx.py
Normal file
140
tests/test_upscale_onnx.py
Normal file
@@ -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()
|
||||
@@ -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))
|
||||
|
||||
101
uv.lock
generated
101
uv.lock
generated
@@ -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"
|
||||
|
||||
Reference in New Issue
Block a user