diff --git a/README.md b/README.md index c98666b..d7ae0d8 100644 --- a/README.md +++ b/README.md @@ -19,7 +19,7 @@ SnapForm 是一个 Blender 扩展，可以把一张照片、一段视频或一 最快的 CUDA/MPS，否则 CPU。下载超时会自动重试国内镜像，无需手动配置 5. 打开 3D 视图右侧的 `SnapForm` 面板，点击 `Install Environment` -DA3、BEN2 源码包和 uv 已内置在扩展里，安装时直接从本地解包，不需要访问 GitHub。首次使用仍会自动 +DA3、MoGe-2、utils3d、BEN2 源码包和 uv 已内置在扩展里，安装时直接从本地解包，不需要访问 GitHub。首次使用仍会自动 下载运行环境依赖（PyTorch 等），随后自动下载 `DA3 Small` 和 `BEN2 Base`。其他 AI 模型可在 `Models` panel 中按需下载，请保持网络畅通，首次安装需要一些时间。 @@ -30,7 +30,7 @@ DA3、BEN2 源码包和 uv 已内置在扩展里，安装时直接从本地解 1. 在 3D 视图中选中一个 Image Empty（空物体，显示方式为 Image） 2. 右键打开 `SnapForm` 子菜单，选择 `Convert to Depth Plane` 3. 在对话框中选择： - - `Depth from`：`Original` 直接用原图，`AI Generated` 让 AI 估算深度 + - `Depth from`：`Original` 直接用原图，`DA3` 使用现有 DA3 流程，`MoGe-2` 生成 metric depth、normal 与有效区域 mask - `Resolution Scale` 使用 `1–100%`，按原图长边计算 Process Res；结果会限制在 `252–1008` 并自动对齐到 14 的倍数 - `Remove Background` 默认关闭；开启后在转换期间使用 BEN2 删除背景 - `Unshaded` 保留为直接显示贴图颜色的选项 diff --git a/docs/technical.md b/docs/technical.md index a33b08e..5f11c76 100644 --- a/docs/technical.md +++ b/docs/technical.md @@ -6,24 +6,24 @@ `SnapForm` 是一套从图片快速构建空间结构的 Blender Extension。目前以 [Depth Anything 3](https://github.com/ByteDance-Seed/Depth-Anything-3) 作为首个 AI 能力，在独立 Environment 中运行 DA3，可将官方导出结果保存在独立 job 目录， -也可把 Image Empty 转换成带几何与材质置换的立体面片。Debug 面板还可运行 +也可把 Image Empty 转换成带几何与材质置换的立体面片。MoGe-2 作为独立深度来源，保留其 metric depth、normal、mask 与 intrinsics 输出，不复用 DA3 Prediction 后处理。Debug 面板还可运行 [BEN2](https://github.com/PramaLLC/BEN2) 背景剔除功能测试。 ## 当前功能 - 使用 Blender Extension manifest，可通过 `Install from Disk` 安装 -- 在独立 Environment 中通过常驻 Service 运行 Torch、DA3 与 BEN2，避免污染 Blender Python 进程 +- 在独立 Environment 中通过常驻 Service 运行 Torch、DA3、MoGe-2 与 BEN2，避免污染 Blender Python 进程 - Blender 使用内置 `urllib.request` 访问仅监听 `127.0.0.1` 的 Service；请求带随机 Token - 常驻 Service 复用 Torch、CUDA、BEN2 与最近使用的 DA3 模型，连续执行时不再重复加载 - 多个 Blender 实例自动连接同一个共享 Service；`AI Service` 面板显示状态、端口、客户端与模型缓存，并提供安装、启动、清理、重启和停止控制 - Environment、模型和任务结果默认保存在 `~/.snapform`，可由多个 Blender 版本共享 -- 支持 DA3 Small、Base、Large 1.1、Giant 1.1 和 Nested Giant Large 1.1 +- 当前启用 DA3 Small 和 Base；Large 1.1、Giant 1.1 与 Nested Giant Large 1.1 暂停使用 - 支持 BEN2 Base 单图背景剔除测试，输出透明 PNG 与推理耗时 - 支持单张 JPG、JPEG、PNG、WEBP 图片或常见视频文件输入 - 视频按顺序截取前若干帧，可设置参与重建的最大帧数 - `Models` panel 默认折叠，展开后可逐项查看和下载所有可用模型 - Environment 已安装时不额外占用面板空间；模型列表每个模型只占一行 -- `Debug` 默认折叠，内部的 `DA3 Output Test` 与 `BEN2 Removal Test` 也分别默认折叠 +- `Debug` 默认折叠，内部的 `DA3 Output Test`、`MoGe-2 API Test` 与 `BEN2 Removal Test` 也分别默认折叠 - 使用一个输出类型下拉菜单和一个 `Generate Output` 按钮 - 只显示当前格式的专属参数；GLB 与 PLY 设置不会同时出现 - Debug 的模型选项按格式过滤；PLY 只提供两个带 Gaussian head 的 Giant 模型 @@ -103,38 +103,41 @@ Cyclic 沿用原 Empty 的设置。序列帧无法被 Blender 整体 Pack，因 1. 在 3D View 选中 `Object Type = Empty`、`Display As = Image` 的对象。 2. 右键打开 `SnapForm` 子菜单并选择 `Convert to Depth Plane`，通过 `Depth from` 选择深度来源： - - `AI Generated`：显示 Model、Use Ray Pose 和 Resolution Scale。Process Res 先以原图长边乘以 + - `DA3`：显示 Model、Use Ray Pose 和 Resolution Scale。Process Res 先以原图长边乘以 Resolution Scale，再钳制到 `252–1008`，最后对齐到最接近的 `14` 倍数。Resolution Scale 以百分比 显示，范围为 `1–100%`，默认 `100%`；界面只在标签中实时显示最终 Process Res。 转换导出 `mini_npz`，使用前需要在 `Models` panel 下载所选模型，并固定使用 `process_res_method="upper_bound_resize"`；设备和参考视图策略 继续使用场景设置。Model 下拉菜单只列出已经下载的 DA3 模型；若没有可用模型，对话框会引导到 `Models` panel 下载。 - - `Original`：直接把原图作为深度贴图，不显示估算参数。仅在开启 `Remove Background` 时运行 BEN2 + - `MoGe-2`：使用 Base Normal 模型输出 metric depth、normal、validity mask 与 normalized intrinsics。正式转换的 `results.npz` 不保存体积较大的 points；Debug 测试会保留 points，并另外写出 depth、normal、mask 预览。 + - `Original`：直接把原图作为深度贴图，不显示估算参数。仅在开启 `Remove Background` 时运行 BEN2 取得前景 Alpha；视频或序列帧按需同时抽取原始帧与前景序列，把原始帧作为 `.depth` 序列。 - `Original` 是默认深度来源。选择 `AI Generated` 时，其 Model、Use Ray Pose 与 Resolution Scale + `Original` 是默认深度来源。选择 `DA3` 时，其 Model、Use Ray Pose 与 Resolution Scale 参数会集中显示在独立边框中。`Remove Background` 与 `Unshaded` 直接显示，不再使用 Options 折叠； `Keep Original` 默认关闭，并在分割线下方保持为最后一项。仅当开启 `Remove Background` 时才需要 `BEN2 Base`，并把 `.foreground` 的 Alpha 作为独立 Mask，不在 Service 中改写深度值。保存 `.foreground` 时保留 BEN2 Alpha，但 RGB 使用未经修改的原始输入； 因此材质直接用 `.foreground` 同时提供颜色和 Mask，不再额外加载原图 Image Texture。 - Image Empty 引用视频或序列帧时，AI Generated 会按 Empty 的 Image Frames 数值逐帧估计 + Image Empty 引用视频或序列帧时，DA3 与 MoGe-2 会按 Empty 的 Image Frames 数值逐帧估计 深度（未设置时为完整帧数，显式设为 1 时只处理 1 帧），起始帧跟随 Empty 的 Offset，与背景蒙版范围一致； Offset 超出视频长度时同样把窗口收进视频内。材质节点与置换贴图的 Start 保持与原 Empty 一致；Original 模式的视频或序列帧使用 `frame.0001.jpg` 等原图序列作为颜色与深度， `foreground.0001.png` 只提供前景 Alpha。 - AI Generated 的深度仍以 EXR 序列帧（`depth.0001.exr` 起）保存在 job 目录根下。 + DA3 的深度仍以 EXR 序列帧（`depth.0001.exr` 起）保存在 job 目录根下；MoGe-2 分别使用 `moge_depth`、`moge_mask` 与 `moge_normal` EXR 序列。 `Unshaded`、`Remove Background` 与 `Keep Original` 均默认关闭，因此默认删除原对象。 -3. AI Generated 完成后会自动创建面片。面片几何先由 `mini_npz` 的相机内参与深度确定，再通过对象 +3. DA3 或 MoGe-2 完成后会自动创建面片。DA3 面片几何先由 `mini_npz` 的相机内参与深度确定，再通过对象 Transform 匹配原 Image Empty 的可见尺寸、位置、旋转和 Offset。Original 模式直接使用源 Empty - 的可见尺寸、Offset 与 Transform。AI Generated 会把用于匹配画面尺寸的缩放预乘到 `.depth` 数值， + 的可见尺寸、Offset 与 Transform。DA3 会把用于匹配画面尺寸的缩放预乘到 `.depth` 数值， XY 缩放烘焙到 Mesh，Z 不参与平面匹配；最终对象 Scale 为 `(1, 1, 1)`，视觉深度与旧逻辑一致。 关闭 `Keep Original` 时，成功生成后删除源 Empty。生成面片 Object 使用源 Empty 的原名；开启 `Keep Original` 时不处理 同名冲突，由 Blender 自动为生成面片追加编号。 4. Blender 只创建保留图像纵横比的最低密度基础网格，再把 `O DA3 Depth Plane` 作为 Geometry Nodes 修改器添加到对象。节点组与 `O DA3 Surface` 一同保存在 `O_SnapForm.blend`，关闭 - `show_group_selector`，不为每个对象复制节点组。修改器公开以下参数： + `show_group_selector`，不为每个对象复制节点组。MoGe-2 使用单独的 `O MoGe2 Depth Plane`，只保留 + 网格细分、metric depth 位移、mask 删点和材质设置；材质使用 `O MoGe2 Depth Plane Layer`，直接接入 + 模型 normal，不携带 DA3 的材质置换与背景深度分支。DA3 修改器公开以下参数： - `Displacement`：位于 Geometry 之后的第一项，菜单包含 `Geometry`（默认）与 `Material`。`Geometry` 在节点组内沿法线修改网格，`Material` 保持网格平面并启用材质置换； - `Subdivision`：位于第二项，默认 `5`。节点读取基础平面的 Bounding Box，短边分段数为 `2^Subdivision`， @@ -259,7 +262,7 @@ Height = sqrt(Area / Aspect) 参考深度和尺寸。 材质始终导入深度 Image Texture，并连接到 Displacement 节点和 Material Output；材质 Surface -Displacement 设置为 `Displacement Only`，AI Generated 与 Original 都使用 `Midlevel = 0`、`Scale = 1`。 +Displacement 设置为 `Displacement Only`，DA3 与 Original 都使用 `Midlevel = 0`、`Scale = 1`。MoGe-2 的位移只在专用 Geometry Nodes 中处理。 `o_material_displacement` 为 `False` 时，材质把 Height 乘以零，置换只在几何节点中生效；为 `True` 时， 几何节点关闭结构置换，材质 Height 恢复。色彩贴图的 Alpha 二值结果同时用于选择原始深度或有符号的 `Background Depth`，使两种置换路径保持一致；未移除背景的普通不透明图片因此全部视为前景。 @@ -408,10 +411,10 @@ Python 3.12 与全部依赖。`uv` 会缓存在扩展存储目录下，后续安 `SNAPFORM_TORCH_INDEX`（PyTorch 索引）、`HF_ENDPOINT`（Hugging Face 镜像）、 `HTTP(S)_PROXY`/`ALL_PROXY`（代理）。未指定 PyPI 镜像时， pip 安装失败会自动依次重试镜像（清华、阿里云）；模型下载不预探测，官方 Hugging Face -失败后自动回退 ModelScope（`modelscope.cn`，仓库名与 Hugging Face 一致）。 +失败后自动回退 OmooLab R2（`models.omoolab.xyz`），并校验文件大小和 SHA-256。 GitHub 不再参与安装：DA3、BEN2、uv 已内置，缺少内置文件会直接报错提示重新安装，不再下载兜底。 -单独下载模型时遵循同样的顺序：Hugging Face（可用 `HF_ENDPOINT` 指定其他镜像）失败后自动切换到 -ModelScope。 +单独下载模型时遵循同样的顺序：Hugging Face（可用 `HF_ENDPOINT` 指定其他入口）失败后自动切换到 +OmooLab R2，不再使用 ModelScope。 Windows 上代理若只在系统设置里而未写入环境变量，安装程序会尝试自动从系统代理注入子进程。 - PyTorch 2.9.1 diff --git a/patches/moge/pyproject.toml b/patches/moge/pyproject.toml new file mode 100644 index 0000000..c6b7a3a --- /dev/null +++ b/patches/moge/pyproject.toml @@ -0,0 +1,22 @@ +[build-system] +requires = ["setuptools>=61.0", "wheel"] +build-backend = "setuptools.build_meta" + +[project] +name = "moge" +version = "2.0.0" +description = "MoGe monocular geometry estimation" +readme = "README.md" +license = {text = "MIT"} +dependencies = [ + "huggingface_hub", + "numpy", + "pillow", + "torch>=2.0.0", + "utils3d", +] +requires-python = ">=3.9" + +[tool.setuptools.packages.find] +where = ["."] +include = ["moge*"] diff --git a/patches/utils3d/pyproject.toml b/patches/utils3d/pyproject.toml new file mode 100644 index 0000000..d0898ed --- /dev/null +++ b/patches/utils3d/pyproject.toml @@ -0,0 +1,16 @@ +[build-system] +requires = ["setuptools>=61.0", "wheel"] +build-backend = "setuptools.build_meta" + +[project] +name = "utils3d" +version = "1.7" +description = "Utilities for 3D computer vision" +readme = "README.md" +license = {text = "MIT"} +dependencies = ["numpy", "scipy"] +requires-python = ">=3.8" + +[tool.setuptools.packages.find] +where = ["."] +include = ["utils3d*"] diff --git a/src/snapform/__init__.py b/src/snapform/__init__.py index 9602c25..cb4c813 100644 --- a/src/snapform/__init__.py +++ b/src/snapform/__init__.py @@ -5,6 +5,7 @@ from .image_depth import ( ApplyRemovedBackground, BuildDepthPlane, BuildOriginalDepthPlane, + BuildMoge2DepthPlane, ConvertImageToDepthPlane, RemoveImageBackground, SnapFormImageMenu, @@ -24,6 +25,7 @@ from .panel import ( DebugPanel, AiServicePanel, ModelsPanel, + Moge2ApiDebugPanel, ) from .preferences import SnapFormPreferences from .properties import SnapFormSettings @@ -41,6 +43,7 @@ CLASSES = ( RemoveImageBackground, BuildDepthPlane, BuildOriginalDepthPlane, + BuildMoge2DepthPlane, ConvertImageToDepthPlane, BuildImageMesh, ConvertImageToMesh, @@ -48,6 +51,7 @@ CLASSES = ( ModelsPanel, DebugPanel, Da3OutputDebugPanel, + Moge2ApiDebugPanel, Ben2RemovalDebugPanel, ) diff --git a/src/snapform/assets/O_SnapForm.blend b/src/snapform/assets/O_SnapForm.blend index 14d377f..4e6b0f4 100644 --- a/src/snapform/assets/O_SnapForm.blend +++ b/src/snapform/assets/O_SnapForm.blend @@ -1,3 +1,3 @@ version https://git-lfs.github.com/spec/v1 -oid sha256:79f05ffe742cc2b4a8c184b4e829d4fc4df5824f0d11ca12d4d1d566b93ec6be -size 917173 +oid sha256:a762ec6c2dfb8e3a1d13bc3a937a2ce30d69195f91b05cde45697f38170b220d +size 1079655 diff --git a/src/snapform/client.py b/src/snapform/client.py index 7995847..0dc3d98 100644 --- a/src/snapform/client.py +++ b/src/snapform/client.py @@ -29,6 +29,7 @@ from .models import ( DA3_COMMIT = "3d835ec1a5802d64a8b8b15f817a1ab54809bfe4" BEN2_COMMIT = "2c99a5da477b5523585bfa5c893888a6e818a8f6" MANIFEST_FILENAME = "manifest.json" +MANIFEST_REVISION = 5 SERVICE_ADDRESS = "127.0.0.1" DEFAULT_STORAGE_ROOT = Path.home() / ".snapform" JOB_FOLDER_NAMES = { @@ -36,6 +37,8 @@ JOB_FOLDER_NAMES = { "DOWNLOAD_MODEL": "download-model", "RECONSTRUCT": "export-reconstruction", "IMAGE_DEPTH": "generate-depth-plane", + "MOGE2_DEPTH": "generate-moge2-depth-plane", + "MOGE2_DEBUG": "test-moge2-api", "IMAGE_MESH": "generate-image-mesh", "REMOVE_BACKGROUND": "remove-background", "IMAGE_BACKGROUND": "remove-image-background", @@ -109,7 +112,10 @@ def env_installed(): manifest = json.loads(ready_path.read_text(encoding="utf-8")) except (json.JSONDecodeError, OSError): return False - return isinstance(manifest, dict) + return ( + isinstance(manifest, dict) + and manifest.get("manifest_revision") == MANIFEST_REVISION + ) def models_directory(): @@ -121,6 +127,13 @@ def models_directory(): def model_directory(model_key): root = models_directory() directory = root / model_directory_name(model_key) + old_moge_directory = root / "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 = root / model_key.lower() if ( not directory.is_dir() @@ -132,9 +145,11 @@ def model_directory(model_key): def model_ready(model_key): + model = get_downloadable_model(model_key) path = model_directory(model_key) - return (path / "config.json").is_file() and any( - path.glob("*.safetensors") + return all( + any(path.glob(pattern)) + for pattern in model.ready_patterns ) @@ -743,6 +758,14 @@ def download_model_command(model_key, status_path): ] for pattern in model.download_patterns: command.extend(["--allow-pattern", pattern]) + for pattern in model.ready_patterns: + command.extend(["--required-pattern", pattern]) + if model.r2_directory: + command.extend(["--r2-directory", model.r2_directory]) + for filename, size, checksum in model.r2_files: + command.extend( + ["--r2-file", f"{filename}|{size}|{checksum}"] + ) return command @@ -875,6 +898,55 @@ def reconstruct_request( ) +def moge_request( + settings, + status_path, + result_directory, + input_path=None, + resolution_level=None, + include_points=False, + remove_background=False, + start_frame=1, + device=None, + model_key=None, +): + from .models import BEN2_MODEL, MOGE2_MODEL + + selected_model_key = model_key or MOGE2_MODEL.key + + parameters = { + "input": str( + Path(input_path) + if input_path is not None + else Path(bpy.path.abspath(settings.input_path)) + ), + "model_dir": str(model_directory(selected_model_key)), + "output_dir": str(result_directory), + "status": str(status_path), + "resolution_level": int( + settings.moge_resolution_level + if resolution_level is None + else resolution_level + ), + "include_points": bool(include_points), + "video_frames": settings.video_frames, + "start_frame": start_frame, + "device": (device or settings.device).lower(), + "remove_background_model_dir": ( + str(model_directory(BEN2_MODEL.key)) + if remove_background + else "" + ), + } + return JobRequest( + { + "job_id": job_id_from_status(status_path), + "command": "moge", + "parameters": parameters, + } + ) + + def creation_flags(): if os.name != "nt": return 0 diff --git a/src/snapform/image_depth.py b/src/snapform/image_depth.py index 1c5920d..1d9215a 100644 --- a/src/snapform/image_depth.py +++ b/src/snapform/image_depth.py @@ -13,12 +13,14 @@ ORIGINAL_DISPLACEMENT_STRENGTH = 1.0 ORIGINAL_DISPLACEMENT_MIDLEVEL = 0.0 BUMP_BASE_DISTANCE = 0.001 DEPTH_PLANE_NODE_GROUP_NAME = "O DA3 Depth Plane" +MOGE_DEPTH_PLANE_NODE_GROUP_NAME = "O MoGe2 Depth Plane" NODE_ASSET_NAME = "O_SnapForm.blend" MATERIAL_DISPLACEMENT_ATTRIBUTE = "o_material_displacement" BACKGROUND_DEPTH_ATTRIBUTE = "o_background_depth" DEPTH_SCALE_ATTRIBUTE = "o_depth_scale" SURFACE_MATERIAL_NODE_GROUP_NAME = "O DA3 Surface Layer" DEPTH_PLANE_MATERIAL_NODE_GROUP_NAME = "O DA3 Depth Plane Layer" +MOGE_MATERIAL_NODE_GROUP_NAME = "O MoGe2 Depth Plane Layer" DA3_PATCH_SIZE = 14 DA3_MIN_PROCESS_RES = 252 DA3_MAX_PROCESS_RES = 1008 @@ -34,6 +36,8 @@ IMAGE_NAME_SUFFIXES = ( ".bmp", ".tga", ) + + VIDEO_NAME_SUFFIXES = ( ".mp4", ".mov", @@ -48,17 +52,12 @@ VIDEO_NAME_SUFFIXES = ( def image_depth_model_items(owner, context): - from .models import model_enum_items - from .client import model_ready + from .models import depth_plane_model_enum_items - return [ - item - for item in model_enum_items(owner, context) - if model_ready(item[0]) - ] + return depth_plane_model_enum_items(owner, context) -_last_depth_plane_model = "BASE" +_last_depth_plane_model = "MOGE2_VITB_NORMAL" def is_image_empty(candidate): @@ -800,6 +799,264 @@ def create_depth_image(name, depth_values): return image +def create_moge_mask_image(name, mask_values): + import numpy as np + + height, width = mask_values.shape + image = bpy.data.images.new( + name, + width=width, + height=height, + alpha=True, + float_buffer=True, + ) + image.colorspace_settings.name = "Non-Color" + values = np.flipud(mask_values.astype(np.float32)) + rgba = np.ones((height, width, 4), dtype=np.float32) + rgba[..., :3] = values[..., None] + rgba[..., 3] = values + image.pixels.foreach_set(rgba.ravel()) + image.pack() + return image + + +def create_moge_normal_image(name, normal_values): + import numpy as np + + converted = normal_values.astype(np.float32, copy=True) + converted *= np.array([1.0, -1.0, -1.0], dtype=np.float32) + converted = np.clip(converted * 0.5 + 0.5, 0.0, 1.0) + height, width = converted.shape[:2] + image = bpy.data.images.new( + name, + width=width, + height=height, + alpha=False, + float_buffer=True, + ) + image.colorspace_settings.name = "Non-Color" + rgba = np.ones((height, width, 4), dtype=np.float32) + rgba[..., :3] = np.flipud(converted) + image.pixels.foreach_set(rgba.ravel()) + image.pack() + return image + + +def create_moge_sequence_image(name, frames, scene, directory, prefix, kind): + import numpy as np + + directory = Path(directory) + directory.mkdir(parents=True, exist_ok=True) + height, width = frames[0].shape[:2] + image = bpy.data.images.new( + name, + width=width, + height=height, + alpha=True, + float_buffer=True, + ) + image.colorspace_settings.name = "Non-Color" + image_settings = scene.render.image_settings + original_format = image_settings.file_format + paths = [] + try: + image_settings.file_format = "OPEN_EXR" + for index, values in enumerate(frames): + rgba = np.ones((height, width, 4), dtype=np.float32) + if kind == "depth": + flipped = np.flipud(values.astype(np.float32)) + rgba[..., :3] = flipped[..., None] + elif kind == "mask": + flipped = np.flipud(values.astype(np.float32)) + rgba[..., :3] = flipped[..., None] + rgba[..., 3] = flipped + else: + converted = values.astype(np.float32, copy=True) + converted *= np.array([1.0, -1.0, -1.0], dtype=np.float32) + rgba[..., :3] = np.flipud( + np.clip(converted * 0.5 + 0.5, 0.0, 1.0) + ) + image.pixels.foreach_set(rgba.ravel()) + path = directory / f"{prefix}.{index + 1:04d}.exr" + image.save_render(str(path), scene=scene) + paths.append(path) + finally: + image_settings.file_format = original_format + bpy.data.images.remove(image, do_unlink=True) + sequence = bpy.data.images.load(str(paths[0]), check_existing=False) + force_sequence_source(sequence) + sequence.name = name + sequence.colorspace_settings.name = "Non-Color" + pack_sequence_if_possible(sequence, len(paths)) + return sequence + + +def moge_material_node_group(): + existing = bpy.data.node_groups.get(MOGE_MATERIAL_NODE_GROUP_NAME) + if existing is not None: + return existing + group = bpy.data.node_groups.new( + MOGE_MATERIAL_NODE_GROUP_NAME, + "ShaderNodeTree", + ) + for name, socket_type in ( + ("Color", "NodeSocketColor"), + ("Alpha", "NodeSocketFloat"), + ("Depth", "NodeSocketFloat"), + ("Bump Scale", "NodeSocketFloat"), + ): + socket = group.interface.new_socket( + name=name, + in_out="INPUT", + socket_type=socket_type, + ) + if name == "Alpha": + socket.default_value = 1.0 + if name == "Bump Scale": + socket.default_value = 1.0 + for name, socket_type in ( + ("Color", "NodeSocketColor"), + ("Roughness", "NodeSocketFloat"), + ("Alpha", "NodeSocketFloat"), + ("Normal", "NodeSocketVector"), + ("Displacement", "NodeSocketVector"), + ): + group.interface.new_socket( + name=name, + in_out="OUTPUT", + socket_type=socket_type, + ) + nodes = group.nodes + links = group.links + group_input = nodes.new("NodeGroupInput") + group_input.location = (-900.0, 100.0) + bump = nodes.new("ShaderNodeBump") + bump.location = (250.0, -100.0) + bump_distance = nodes.new("ShaderNodeMath") + bump_distance.operation = "MULTIPLY" + bump_distance.inputs[1].default_value = BUMP_BASE_DISTANCE + bump_distance.location = (0.0, -100.0) + background = nodes.new("ShaderNodeAttribute") + background.attribute_type = "GEOMETRY" + background.attribute_name = BACKGROUND_DEPTH_ATTRIBUTE + background.location = (-650.0, -400.0) + material_mode = nodes.new("ShaderNodeAttribute") + material_mode.attribute_type = "GEOMETRY" + material_mode.attribute_name = MATERIAL_DISPLACEMENT_ATTRIBUTE + material_mode.location = (-650.0, -600.0) + strength = nodes.new("ShaderNodeAttribute") + strength.attribute_type = "GEOMETRY" + strength.attribute_name = DEPTH_SCALE_ATTRIBUTE + strength.location = (-400.0, -600.0) + clamp_depth = nodes.new("ShaderNodeMath") + clamp_depth.operation = "MAXIMUM" + clamp_depth.location = (-350.0, -350.0) + enable_depth = nodes.new("ShaderNodeMath") + enable_depth.operation = "MULTIPLY" + enable_depth.location = (-100.0, -350.0) + scale_depth = nodes.new("ShaderNodeMath") + scale_depth.operation = "MULTIPLY" + scale_depth.location = (150.0, -350.0) + displacement = nodes.new("ShaderNodeDisplacement") + displacement.inputs["Midlevel"].default_value = 0.0 + displacement.inputs["Scale"].default_value = 1.0 + displacement.location = (400.0, -350.0) + group_output = nodes.new("NodeGroupOutput") + group_output.location = (300.0, 100.0) + links.new(group_input.outputs["Color"], group_output.inputs["Color"]) + links.new(group_input.outputs["Alpha"], group_output.inputs["Alpha"]) + links.new(group_input.outputs["Color"], bump.inputs["Height"]) + links.new(group_input.outputs["Bump Scale"], bump_distance.inputs[0]) + links.new(bump_distance.outputs[0], bump.inputs["Distance"]) + links.new(bump.outputs["Normal"], group_output.inputs["Normal"]) + links.new(group_input.outputs["Depth"], clamp_depth.inputs[0]) + links.new(background.outputs["Fac"], clamp_depth.inputs[1]) + links.new(clamp_depth.outputs[0], enable_depth.inputs[0]) + links.new(material_mode.outputs["Fac"], enable_depth.inputs[1]) + links.new(enable_depth.outputs[0], scale_depth.inputs[0]) + links.new(strength.outputs["Fac"], scale_depth.inputs[1]) + links.new(scale_depth.outputs[0], displacement.inputs["Height"]) + links.new(displacement.outputs["Displacement"], group_output.inputs["Displacement"]) + group_output.inputs["Roughness"].default_value = 0.5 + group.color_tag = "SHADER" + group.description = "Apply MoGe-2 color and derive bump detail from it." + group.use_fake_user = True + return group + + +def create_moge_material( + name, + color_image, + depth_image, + depth_scale, + unshaded=False, + frame_count=1, + frame_start=1, + use_cyclic=False, +): + material = bpy.data.materials.new(name) + material.use_nodes = True + if hasattr(material, "surface_render_method"): + material.surface_render_method = "DITHERED" + nodes = material.node_tree.nodes + nodes.clear() + links = material.node_tree.links + uv = nodes.new("ShaderNodeUVMap") + uv.uv_map = "UVMap" + uv.location = (-900.0, 200.0) + color = nodes.new("ShaderNodeTexImage") + color.image = color_image + color.location = (-650.0, 300.0) + depth = nodes.new("ShaderNodeTexImage") + depth.image = depth_image + depth.interpolation = "Linear" + depth.location = (-650.0, -100.0) + if frame_count > 1: + configure_image_user( + color.image_user, + frame_count, + frame_start=frame_start, + use_cyclic=use_cyclic, + ) + configure_image_user( + depth.image_user, + frame_count, + frame_start=frame_start, + use_cyclic=use_cyclic, + ) + layer = nodes.new("ShaderNodeGroup") + layer.node_tree = moge_material_node_group() + layer.name = MOGE_MATERIAL_NODE_GROUP_NAME + layer.location = (-300.0, 200.0) + output = nodes.new("ShaderNodeOutputMaterial") + output.location = (500.0, 200.0) + links.new(uv.outputs["UV"], color.inputs["Vector"]) + links.new(uv.outputs["UV"], depth.inputs["Vector"]) + links.new(color.outputs["Color"], layer.inputs["Color"]) + links.new(color.outputs["Alpha"], layer.inputs["Alpha"]) + links.new(depth.outputs["Color"], layer.inputs["Depth"]) + layer.inputs["Bump Scale"].default_value = depth_scale + if unshaded: + shader = nodes.new("ShaderNodeEmission") + transparent = nodes.new("ShaderNodeBsdfTransparent") + mix = nodes.new("ShaderNodeMixShader") + links.new(layer.outputs["Color"], shader.inputs["Color"]) + links.new(layer.outputs["Alpha"], mix.inputs["Fac"]) + links.new(transparent.outputs["BSDF"], mix.inputs[1]) + links.new(shader.outputs["Emission"], mix.inputs[2]) + links.new(mix.outputs["Shader"], output.inputs["Surface"]) + return material + shader = nodes.new("ShaderNodeBsdfPrincipled") + links.new(layer.outputs["Color"], shader.inputs["Base Color"]) + links.new(layer.outputs["Alpha"], shader.inputs["Alpha"]) + links.new(layer.outputs["Roughness"], shader.inputs["Roughness"]) + links.new(layer.outputs["Normal"], shader.inputs["Normal"]) + links.new(shader.outputs["BSDF"], output.inputs["Surface"]) + links.new(layer.outputs["Displacement"], output.inputs["Displacement"]) + set_displacement_only(material) + return material + + def grid_segments(bounds, subdivision_level): if subdivision_level < 0: raise ValueError("Subdivision must not be negative") @@ -1193,8 +1450,8 @@ def node_asset_path(): return Path(__file__).with_name("assets") / NODE_ASSET_NAME -def load_depth_plane_node_group(): - existing = bpy.data.node_groups.get(DEPTH_PLANE_NODE_GROUP_NAME) +def load_depth_plane_node_group(group_name=DEPTH_PLANE_NODE_GROUP_NAME): + existing = bpy.data.node_groups.get(group_name) if existing is not None: return existing asset_path = node_asset_path() @@ -1204,26 +1461,26 @@ def load_depth_plane_node_group(): if not use_pack: directory = str(asset_path) + "/NodeTree/" result = bpy.ops.wm.append( - filepath=directory + DEPTH_PLANE_NODE_GROUP_NAME, + filepath=directory + group_name, directory=directory, - filename=DEPTH_PLANE_NODE_GROUP_NAME, + filename=group_name, link=False, do_reuse_local_id=True, ) if "FINISHED" not in result: raise RuntimeError("Unable to append the Geometry Nodes asset") - node_group = bpy.data.node_groups.get(DEPTH_PLANE_NODE_GROUP_NAME) + node_group = bpy.data.node_groups.get(group_name) if node_group is None: raise RuntimeError("Appended Geometry Nodes asset is unavailable") node_group.use_fake_user = True return node_group with bpy.data.libraries.load(str(asset_path), link=True) as (source, target): - if DEPTH_PLANE_NODE_GROUP_NAME not in source.node_groups: + if group_name not in source.node_groups: raise RuntimeError( "Geometry Nodes asset is missing " - f"{DEPTH_PLANE_NODE_GROUP_NAME}" + f"{group_name}" ) - target.node_groups = [DEPTH_PLANE_NODE_GROUP_NAME] + target.node_groups = [group_name] linked_group = target.node_groups[0] node_group = bpy.data.pack_linked_ids_hierarchy(root_id=linked_group) if linked_group != node_group and linked_group.users == 0: @@ -1278,6 +1535,42 @@ def add_depth_plane_processing( return modifier +def add_moge_depth_plane_processing( + plane_object, + depth_image, + material, + frame_start=1, + frame_duration=1, + use_cyclic=False, + animate_frame=False, +): + node_group = load_depth_plane_node_group( + MOGE_DEPTH_PLANE_NODE_GROUP_NAME + ) + modifier = plane_object.modifiers.new( + name=MOGE_DEPTH_PLANE_NODE_GROUP_NAME, + type="NODES", + ) + modifier.node_group = node_group + modifier.show_group_selector = False + for name, value in ( + ("Depth Image", depth_image), + ("Material", material), + ("Frame", 1), + ("Frame Start", frame_start), + ("Frames", max(int(frame_duration or 1), 1)), + ("Cyclic", use_cyclic), + ): + modifier[modifier_input_identifier(node_group, name)] = value + if animate_frame: + frame_identifier = modifier_input_identifier(node_group, "Frame") + driver = modifier.driver_add(f'["{frame_identifier}"]') + driver.driver.type = "SCRIPTED" + driver.driver.expression = "frame" + plane_object.update_tag(refresh={"DATA"}) + return modifier + + def finish_depth_plane( context, source_object, @@ -1573,6 +1866,152 @@ class BuildDepthPlane(bpy.types.Operator): return {"FINISHED"} +def build_moge2_depth_plane( + context, + source_object, + npz_path, + delete_source, + unshaded=False, + color_image=None, + foreground_image=None, + data_frames_directory=None, +): + import numpy as np + + try: + with np.load(npz_path, allow_pickle=False) as archive: + depth = np.asarray(archive["depth"], dtype=np.float32) + mask = np.asarray(archive["mask"], dtype=bool) + intrinsics = np.asarray(archive["intrinsics"], dtype=np.float32) + except (KeyError, OSError) as error: + raise ValueError(f"Invalid MoGe-2 NPZ: {npz_path}") from error + if depth.ndim == 2: + depth = depth[None, ...] + mask = mask[None, ...] + if intrinsics.ndim == 2: + intrinsics = intrinsics[None, ...] + if not (len(depth) == len(mask) == len(intrinsics)): + raise ValueError("MoGe-2 arrays have different frame counts") + + first_valid = mask[0] & np.isfinite(depth[0]) & (depth[0] > 0.0) + if not first_valid.any(): + raise ValueError("MoGe-2 mask contains no valid depth") + reference_depth = float(np.median(depth[0][first_valid])) + height, width = depth[0].shape + pixel_intrinsics = intrinsics[0].copy() + pixel_intrinsics[0] *= width + pixel_intrinsics[1] *= height + bounds, _ = camera_plane_bounds( + source_object, + depth[0], + pixel_intrinsics, + 1.0, + reference_depth=reference_depth, + ) + displacement = [reference_depth - frame for frame in depth] + source_frame_start, _, _, source_use_cyclic = image_user_settings( + source_object + ) + mesh = create_plane_mesh(source_object, bounds, None) + plane_object = bpy.data.objects.new(source_object.name, mesh) + context.collection.objects.link(plane_object) + depth_scale = apply_plane_transform(plane_object, source_object, bounds) + apply_object_scale_to_mesh(plane_object) + depth_image = create_depth_image( + f"{depth_data_name(source_object)}.moge2", + displacement[0] * depth_scale, + ) + if len(depth) > 1 and data_frames_directory is not None: + bpy.data.images.remove(depth_image, do_unlink=True) + depth_image = create_moge_sequence_image( + f"{depth_data_name(source_object)}.moge2", + [frame * depth_scale for frame in displacement], + context.scene, + data_frames_directory, + "moge_depth", + "depth", + ) + material = create_moge_material( + source_object.name, + foreground_image or color_image or source_object.data, + depth_image, + depth_scale, + unshaded=unshaded, + frame_count=len(depth), + frame_start=source_frame_start, + use_cyclic=source_use_cyclic, + ) + mesh.materials.append(material) + modifier = add_moge_depth_plane_processing( + plane_object, + depth_image, + material, + frame_start=source_frame_start, + frame_duration=len(depth), + use_cyclic=source_use_cyclic, + animate_frame=is_animated_image(source_object.data), + ) + plane_object["snapform_reference_depth"] = reference_depth + plane_object["snapform_depth_scale"] = depth_scale + plane_object["snapform_frame_count"] = len(depth) + return finish_depth_plane( + context, + source_object, + plane_object, + delete_source, + ) + + +class BuildMoge2DepthPlane(bpy.types.Operator): + bl_idname = "snapform.build_moge2_depth_plane" + bl_label = "Build MoGe-2 Depth Plane" + bl_options = {"INTERNAL"} + + snapform_source_object_name: bpy.props.StringProperty(options={"HIDDEN"}) + snapform_npz_path: bpy.props.StringProperty(options={"HIDDEN"}) + snapform_background_path: bpy.props.StringProperty(options={"HIDDEN"}) + snapform_color_frames_dir: bpy.props.StringProperty(options={"HIDDEN"}) + snapform_data_frames_dir: bpy.props.StringProperty(options={"HIDDEN"}) + snapform_delete_source: bpy.props.BoolProperty(default=True, options={"HIDDEN"}) + snapform_unshaded: bpy.props.BoolProperty(default=False, options={"HIDDEN"}) + + def execute(self, context): + source_object = bpy.data.objects.get(self.snapform_source_object_name) + if not is_image_empty(source_object): + self.report({"ERROR"}, "The source image empty is no longer available") + return {"CANCELLED"} + try: + foreground_image = None + if self.snapform_background_path: + foreground_image = create_removed_background_image( + source_object, + self.snapform_background_path, + assign_to_source=False, + ) + color_image = None + if self.snapform_color_frames_dir: + color_image = create_sequence_image( + f"{image_base_name(source_object.data)}.frame", + self.snapform_color_frames_dir, + ) + build_moge2_depth_plane( + context, + source_object, + Path(self.snapform_npz_path), + self.snapform_delete_source, + unshaded=self.snapform_unshaded, + color_image=color_image, + foreground_image=foreground_image, + data_frames_directory=( + self.snapform_data_frames_dir or None + ), + ) + except (OSError, RuntimeError, ValueError) as error: + self.report({"ERROR"}, str(error)) + return {"CANCELLED"} + return {"FINISHED"} + + class BuildOriginalDepthPlane(bpy.types.Operator): bl_idname = "snapform.build_original_depth_plane" bl_label = "Build Original Depth Plane" @@ -1746,19 +2185,6 @@ class ConvertImageToDepthPlane(bpy.types.Operator): bl_description = "Convert an image to a displaced depth plane" bl_options = {"REGISTER"} - snapform_depth_source: bpy.props.EnumProperty( - name="Depth from", - description="Source used to displace the plane", - items=( - ( - "ESTIMATE", - "AI Generated", - "Generate depth with an AI model", - ), - ("ORIGINAL", "Original", "Use the original image as depth"), - ), - default="ORIGINAL", - ) snapform_unshaded: bpy.props.BoolProperty( name="Unshaded", description="Display the source color through an Emission shader", @@ -1771,8 +2197,9 @@ class ConvertImageToDepthPlane(bpy.types.Operator): ) snapform_model: bpy.props.EnumProperty( name="Model", - description="Model used to estimate image depth", + description="Model or original image used to build depth", items=image_depth_model_items, + default=1, ) snapform_use_ray_pose: bpy.props.BoolProperty( name="Use Ray Pose", @@ -1780,6 +2207,13 @@ class ConvertImageToDepthPlane(bpy.types.Operator): default=True, ) snapform_process_res_factor: process_resolution_scale_property() + snapform_moge_resolution_level: bpy.props.IntProperty( + name="Resolution Level", + description="MoGe-2 inference resolution level; higher is finer and slower", + default=5, + min=0, + max=9, + ) snapform_keep_original: bpy.props.BoolProperty( name="Keep Original", description="Keep the source Image Empty after conversion", @@ -1789,14 +2223,12 @@ class ConvertImageToDepthPlane(bpy.types.Operator): @classmethod def poll(cls, context): from .operators import active_job - from .client import service_busy, model_ready - from .models import MODELS + from .client import service_busy return ( is_image_empty(context.object) and active_job() is None and not service_busy() - and any(model_ready(model_key) for model_key in MODELS) ) def invoke(self, context, _event): @@ -1816,17 +2248,19 @@ class ConvertImageToDepthPlane(bpy.types.Operator): def draw(self, context): layout = self.layout - layout.prop(self, "snapform_depth_source") - if self.snapform_depth_source == "ESTIMATE": + layout.prop(self, "snapform_model") + if self.snapform_model in {"SMALL", "BASE"}: estimate_box = layout.box() - available_models = image_depth_model_items(self, context) - if not available_models: - estimate_box.label(text="No DA3 models installed", icon="ERROR") - estimate_box.label(text="Download one in Models") - else: - estimate_box.prop(self, "snapform_model") - estimate_box.prop(self, "snapform_use_ray_pose") - draw_process_resolution(estimate_box, self, context.object.data) + estimate_box.prop(self, "snapform_use_ray_pose") + draw_process_resolution(estimate_box, self, context.object.data) + elif self.snapform_model.startswith("MOGE2_"): + from .client import model_ready + + moge_box = layout.box() + if not model_ready(self.snapform_model): + moge_box.label(text="MoGe-2 model is not installed", icon="ERROR") + moge_box.label(text="Download it in Models") + moge_box.prop(self, "snapform_moge_resolution_level") layout.prop(self, "snapform_remove_background") layout.prop(self, "snapform_unshaded") layout.separator(type="LINE") @@ -1835,6 +2269,7 @@ class ConvertImageToDepthPlane(bpy.types.Operator): def execute(self, context): from .operators import ( IMAGE_DEPTH_JOB, + MOGE2_DEPTH_JOB, ORIGINAL_BACKGROUND_JOB, ORIGINAL_FRAMES_JOB, ) @@ -1844,7 +2279,7 @@ class ConvertImageToDepthPlane(bpy.types.Operator): context.scene.snapform_settings.video_frames = ( default_extract_frames(source_object) ) - if self.snapform_depth_source == "ORIGINAL": + if self.snapform_model == "ORIGINAL": try: source_path, delete_input = prepare_image_input( source_object.data, @@ -1871,6 +2306,34 @@ class ConvertImageToDepthPlane(bpy.types.Operator): remove_temporary_input(source_path) return {"CANCELLED"} + if self.snapform_model.startswith("MOGE2_"): + try: + source_path, delete_input = prepare_image_input( + source_object.data, + context.scene, + ) + except RuntimeError as error: + self.report({"ERROR"}, str(error)) + return {"CANCELLED"} + result = bpy.ops.snapform.run_job( + job_type=MOGE2_DEPTH_JOB, + snapform_input_path=str(source_path), + snapform_source_object_name=source_object.name, + snapform_moge_resolution_level=( + self.snapform_moge_resolution_level + ), + snapform_model_key=self.snapform_model, + snapform_delete_source=not self.snapform_keep_original, + snapform_delete_input=delete_input, + snapform_unshaded=self.snapform_unshaded, + snapform_remove_background=self.snapform_remove_background, + ) + if "CANCELLED" not in result: + return {"FINISHED"} + if delete_input: + remove_temporary_input(source_path) + return {"CANCELLED"} + image = source_object.data try: process_res = scaled_process_resolution( diff --git a/src/snapform/install_env.py b/src/snapform/install_env.py index cca7cdd..366c236 100644 --- a/src/snapform/install_env.py +++ b/src/snapform/install_env.py @@ -20,7 +20,7 @@ SCRIPT_DIRECTORY = Path(__file__).resolve().parent if str(SCRIPT_DIRECTORY) not in sys.path: sys.path.insert(0, str(SCRIPT_DIRECTORY)) -from models import BEN2_MODEL, MODELS +from models import BEN2_MODEL, DOWNLOADABLE_MODELS, MOGE2_MODEL PYPI_MIRRORS = ( @@ -34,12 +34,13 @@ XFORMERS_VERSION = "0.0.33.post2" ADDICT_VERSION = "2.4.0" FASTAPI_VERSION = "0.139.2" UVICORN_VERSION = "0.51.0" +MANIFEST_REVISION = 5 TORCH_INDEXES = { "cu128": "https://download.pytorch.org/whl/cu128", "cpu": "https://download.pytorch.org/whl/cpu", } _resolved_pypi_index = None -DEFAULT_MODEL_KEYS = ("SMALL", BEN2_MODEL.key) +DEFAULT_MODEL_KEYS = (MOGE2_MODEL.key, BEN2_MODEL.key) RUNTIME_STAGE_COUNT = 3 ENVIRONMENT_SLICE = (0.02, 0.06) PYTORCH_SLICE = (0.06, 0.48) @@ -61,6 +62,7 @@ CHECK_CODE = ( "import json; import torch; " "from depth_anything_3.api import DepthAnything3; " "from ben2 import BEN_Base; " + "from moge.model.v2 import MoGeModel; " "import fastapi, uvicorn; " "mps = getattr(torch.backends, 'mps', None); " "print('SNAPFORM_DEVICES=' + json.dumps({" @@ -70,6 +72,14 @@ CHECK_CODE = ( ) +def configure_standard_streams(): + """Windows 图形进程可能继承 GBK stdout，安装日志统一写为 UTF-8。""" + for stream in (sys.stdout, sys.stderr): + reconfigure = getattr(stream, "reconfigure", None) + if callable(reconfigure): + reconfigure(encoding="utf-8", errors="backslashreplace") + + def write_status( path, progress, @@ -98,6 +108,8 @@ def write_status( def child_environment(environment=None): """Inherit the current environment and inject system/registry proxy settings into child processes (uv only reads proxy env variables).""" process_environment = os.environ.copy() + process_environment["PYTHONIOENCODING"] = "utf-8:backslashreplace" + process_environment["PYTHONUTF8"] = "1" for key, value in urllib.request.getproxies().items(): env_key = f"{key.upper()}_PROXY" if env_key not in process_environment: @@ -334,6 +346,14 @@ def install_stages(args, venv_dir, python, uv, vendor_dir=None): "ben2", "ben2.tar.gz", ) + utils3d_requirement = archive_requirement( + "utils3d", + "utils3d.tar.gz", + ) + moge_requirement = archive_requirement( + "moge", + "moge.tar.gz", + ) torch_command = [ uv, "pip", @@ -371,6 +391,8 @@ def install_stages(args, venv_dir, python, uv, vendor_dir=None): da3_requirement, f"addict=={ADDICT_VERSION}", ben2_requirement, + utils3d_requirement, + moge_requirement, f"fastapi=={FASTAPI_VERSION}", f"uvicorn=={UVICORN_VERSION}", "setuptools", @@ -398,7 +420,7 @@ def install_stages(args, venv_dir, python, uv, vendor_dir=None): 2, ("AI components", xformers_command, None, "install"), ) - models = {**MODELS, BEN2_MODEL.key: BEN2_MODEL} + models = DOWNLOADABLE_MODELS worker_script = Path(__file__).with_name("worker.py") models_dir = Path(args.models_dir) for model_key in DEFAULT_MODEL_KEYS: @@ -418,6 +440,16 @@ def install_stages(args, venv_dir, python, uv, vendor_dir=None): ] for pattern in model.download_patterns: download_command.extend(["--allow-pattern", pattern]) + for pattern in model.ready_patterns: + download_command.extend(["--required-pattern", pattern]) + if model.r2_directory: + download_command.extend( + ["--r2-directory", model.r2_directory] + ) + for filename, size, checksum in model.r2_files: + download_command.extend( + ["--r2-file", f"{filename}|{size}|{checksum}"] + ) stages.append( ( model.label, @@ -437,7 +469,10 @@ def env_installed(venv_dir, manifest_path=None): payload = json.loads(ready_path.read_text(encoding="utf-8")) except (json.JSONDecodeError, OSError): return False - return isinstance(payload, dict) + return ( + isinstance(payload, dict) + and payload.get("manifest_revision") == MANIFEST_REVISION + ) def installed_devices(args): @@ -557,24 +592,79 @@ def bundled_uv_archive(): return path if path.is_file() else None +def vendor_overlay_directory(filename): + names = { + "depth-anything-3.tar.gz": "depth-anything-3", + "moge.tar.gz": "moge", + "utils3d.tar.gz": "utils3d", + } + overlay_name = names.get(filename) + if overlay_name is None: + return None + project_root = Path(__file__).resolve().parents[2] + overlay = project_root / "patches" / overlay_name + return overlay if overlay.is_dir() else None + + +def copy_archive_with_overlay(source, target, overlay): + """开发目录使用原始 vendor 时，应用与正式打包相同的源码覆盖。""" + with tarfile.open(source, "r:gz") as input_archive: + members = input_archive.getmembers() + top_level = members[0].name.split("/", 1)[0] + with tarfile.open(target, "w:gz") as output_archive: + for member in members: + if member.isdir() or member.islnk() or member.issym(): + output_archive.addfile(member) + continue + relative = member.name[len(top_level) + 1 :] + replacement = overlay / relative + if replacement.is_file(): + info = tarfile.TarInfo(member.name) + info.size = replacement.stat().st_size + info.mode = member.mode + with replacement.open("rb") as stream: + output_archive.addfile(info, stream) + continue + with input_archive.extractfile(member) as stream: + output_archive.addfile(member, stream) + + def copy_vendor_inputs(bundled_dir, vendor_dir): """Copy bundled installation inputs into a temporary directory.""" - for name in ("depth-anything-3.tar.gz", "ben2.tar.gz"): + for name in ( + "depth-anything-3.tar.gz", + "ben2.tar.gz", + "utils3d.tar.gz", + "moge.tar.gz", + ): bundled = Path(bundled_dir) / name if not bundled.is_file(): continue target = Path(vendor_dir) / name - if target.is_file() and filecmp.cmp(bundled, target, shallow=False): + overlay = vendor_overlay_directory(name) + if overlay is None and target.is_file() and filecmp.cmp( + bundled, + target, + shallow=False, + ): continue Path(vendor_dir).mkdir(parents=True, exist_ok=True) - shutil.copy2(bundled, target) + if overlay is not None and tarfile.is_tarfile(bundled): + copy_archive_with_overlay(bundled, target, overlay) + else: + shutil.copy2(bundled, target) def remove_legacy_vendor(storage_root): """Remove installation inputs persisted by older development builds.""" root = Path(storage_root) shutil.rmtree(root / "vendor", ignore_errors=True) - for name in ("depth-anything-3.tar.gz", "ben2.tar.gz"): + for name in ( + "depth-anything-3.tar.gz", + "ben2.tar.gz", + "utils3d.tar.gz", + "moge.tar.gz", + ): (root / name).unlink(missing_ok=True) @@ -799,6 +889,7 @@ def install(args): ready_path.write_text( json.dumps( { + "manifest_revision": MANIFEST_REVISION, "python": "3.12", "torch": TORCH_VERSION, "xformers": XFORMERS_VERSION, @@ -870,6 +961,7 @@ def parse_args(): def main(): + configure_standard_streams() args = parse_args() try: install(args) diff --git a/src/snapform/models.py b/src/snapform/models.py index f5c19a8..35ed118 100644 --- a/src/snapform/models.py +++ b/src/snapform/models.py @@ -11,6 +11,10 @@ class ModelSpec: 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 = { @@ -22,6 +26,11 @@ MODELS = { license_name="Apache 2.0", description="Fastest any-view model with Apache 2.0 weights", supports_gs=False, + r2_directory="DA3Small", + r2_files=( + ("config.json", 1202, "a486e29e82b7ab4a7d4cefc1ea4526cfe2ae438a572c8ca98917cfbcde7447d2"), + ("model.safetensors", 137248940, "364492e38a3a06d221ac75da7f6621ada3f2361cd24fde11ba79091e9f40efcf"), + ), ), "BASE": ModelSpec( key="BASE", @@ -31,6 +40,12 @@ MODELS = { license_name="Apache 2.0", description="Balanced any-view model with Apache 2.0 weights", supports_gs=False, + r2_directory="DA3Base", + r2_files=( + ("config.json", 1205, "5e34115ebc17bd2d8d43033c5f72e9446ac8833fd61d3fa160b7e67e0bb5b7b5"), + ("configuration.json", 74, "b684d553257b2ca379fd5dfdf96a6e4a633eec3b969c97119c9875465a7c07a8"), + ("model.safetensors", 541518028, "e01067dc1659613083d9145a9a2547ccdbe6ccbbf83c4fe7b3e8a4e2bdae78b5"), + ), ), "LARGE_1_1": ModelSpec( key="LARGE_1_1", @@ -40,6 +55,7 @@ MODELS = { license_name="CC BY-NC 4.0", description="Higher quality non-commercial model with larger VRAM use", supports_gs=False, + enabled=False, ), "GIANT_1_1": ModelSpec( key="GIANT_1_1", @@ -49,6 +65,7 @@ MODELS = { 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", @@ -58,6 +75,7 @@ MODELS = { license_name="CC BY-NC 4.0", description="Non-commercial nested model with metric depth and Gaussian head", supports_gs=True, + enabled=False, ), } @@ -70,8 +88,51 @@ BEN2_MODEL = ModelSpec( description="Fast high-quality background removal model", supports_gs=False, download_patterns=("config.json", "model.safetensors"), + r2_directory="BEN2Base", + r2_files=( + ("config.json", 124, "01ac876925577b52fa7c6af3f5b1d1420bc9b40016b6c25e2de89b38d8fff865"), + ("model.safetensors", 380577976, "ea8b7907176a09667c86343dc7d00de6a6d871076cb90bb5f753618fd6fb3ebb"), + ), ) -DOWNLOADABLE_MODELS = {**MODELS, BEN2_MODEL.key: BEN2_MODEL} +def moge2_model(key, label, directory_name, repository, description, r2_directory, size, checksum): + return ModelSpec( + key=key, + label=label, + directory_name=directory_name, + repository=repository, + license_name="MIT", + description=description, + supports_gs=False, + download_patterns=("model.pt",), + ready_patterns=("model.pt",), + r2_directory=r2_directory, + r2_files=(("model.pt", size, checksum),), + ) + + +MOGE2_MODELS = { + "MOGE2_VITS_NORMAL": moge2_model( + "MOGE2_VITS_NORMAL", "MoGe-2 ViT-S Normal", "moge2_vits_normal", + "Ruicheng/moge-2-vits-normal", "Fastest MoGe-2 model", "MoGe2ViTSNormal", + 140550416, "79a16621928c2bf0ed04659218c55c01075e950507f40bb3332fb4c873d3e1dc", + ), + "MOGE2_VITB_NORMAL": moge2_model( + "MOGE2_VITB_NORMAL", "MoGe-2 ViT-B Normal", "moge2_vitb_normal", + "Ruicheng/moge-2-vitb-normal", "Balanced MoGe-2 model", "MoGe2ViTBNormal", + 419110160, "16b8110e86d5dc5a849db120ca96ef3a223fd30b0c9146d1d81db504073da5f6", + ), + "MOGE2_VITL_NORMAL": moge2_model( + "MOGE2_VITL_NORMAL", "MoGe-2 ViT-L Normal", "moge2_vitl_normal", + "Ruicheng/moge-2-vitl-normal", "Highest-quality MoGe-2 model", "MoGe2ViTLNormal", + 1323815904, "280741fd09bc3f403ccff9967784c2a391b52d2c0742ae3efdb21d9f90cc1a01", + ), +} +MOGE2_MODEL = MOGE2_MODELS["MOGE2_VITB_NORMAL"] +DOWNLOADABLE_MODELS = { + **MODELS, + BEN2_MODEL.key: BEN2_MODEL, + **MOGE2_MODELS, +} MODEL_NUMBERS = { model_key: number for number, model_key in enumerate(MODELS) } @@ -88,7 +149,8 @@ def enum_items(supports_gs=None, show_download_status=False): models = [ model for model in MODELS.values() - if supports_gs is None or model.supports_gs == supports_gs + if model.enabled + and (supports_gs is None or model.supports_gs == supports_gs) ] download_status = ( tuple(model_ready(model.key) for model in models) @@ -130,6 +192,42 @@ def gs_model_enum_items(_owner, _context): return enum_items(supports_gs=True, show_download_status=True) +def moge2_model_enum_items(_owner, _context): + from .client import model_ready + + return [ + ( + model.key, + model.label if model_ready(model.key) else f"{model.label} (Missing)", + f"{model.description}; license: {model.license_name}", + index, + ) + for index, model in enumerate(MOGE2_MODELS.values()) + ] + + +def depth_plane_model_enum_items(_owner, _context): + from .client import model_ready + + models = [ + *MOGE2_MODELS.values(), + *(model for model in MODELS.values() if model.enabled), + ] + items = [] + for index, model in enumerate(models): + label = model.label if model_ready(model.key) else f"{model.label} (Missing)" + items.append( + ( + model.key, + label, + f"{model.description}; license: {model.license_name}", + index, + ) + ) + items.append(("ORIGINAL", "Original", "Use the original image as depth", len(items))) + return items + + def selected_model_key(settings, export_format=None): selected_format = export_format or settings.export_format if selected_format == "gs_ply": diff --git a/src/snapform/operators.py b/src/snapform/operators.py index 0b7b33b..4609a49 100644 --- a/src/snapform/operators.py +++ b/src/snapform/operators.py @@ -1,3 +1,4 @@ +import os import shutil import subprocess import threading @@ -9,6 +10,7 @@ import bpy from . import client from .models import ( BEN2_MODEL, + MOGE2_MODEL, get_downloadable_model, get_model, selected_model_key, @@ -17,6 +19,8 @@ from .models import ( RECONSTRUCTION_JOB = "RECONSTRUCT" IMAGE_DEPTH_JOB = "IMAGE_DEPTH" +MOGE2_DEPTH_JOB = "MOGE2_DEPTH" +MOGE2_DEBUG_JOB = "MOGE2_DEBUG" IMAGE_MESH_JOB = "IMAGE_MESH" BACKGROUND_REMOVAL_JOB = "REMOVE_BACKGROUND" IMAGE_BACKGROUND_JOB = "IMAGE_BACKGROUND" @@ -26,6 +30,7 @@ DA3_IMAGE_JOBS = frozenset({IMAGE_DEPTH_JOB, IMAGE_MESH_JOB}) IMAGE_EMPTY_JOBS = frozenset( { *DA3_IMAGE_JOBS, + MOGE2_DEPTH_JOB, IMAGE_BACKGROUND_JOB, ORIGINAL_BACKGROUND_JOB, ORIGINAL_FRAMES_JOB, @@ -34,6 +39,7 @@ IMAGE_EMPTY_JOBS = frozenset( PRODUCTION_AI_JOBS = frozenset( { *DA3_IMAGE_JOBS, + MOGE2_DEPTH_JOB, IMAGE_BACKGROUND_JOB, ORIGINAL_BACKGROUND_JOB, } @@ -156,6 +162,8 @@ class RunJob(bpy.types.Operator): ("DOWNLOAD_MODEL", "Download Model", ""), (RECONSTRUCTION_JOB, "Export DA3 Result", ""), (IMAGE_DEPTH_JOB, "Convert Image to Depth Plane", ""), + (MOGE2_DEPTH_JOB, "Convert Image with MoGe-2", ""), + (MOGE2_DEBUG_JOB, "Test MoGe-2 API", ""), (IMAGE_MESH_JOB, "Convert Image to Surface", ""), (BACKGROUND_REMOVAL_JOB, "Test BEN2 Background Removal", ""), (IMAGE_BACKGROUND_JOB, "Remove Image Empty Background", ""), @@ -210,6 +218,12 @@ class RunJob(bpy.types.Operator): default=False, options={"HIDDEN", "SKIP_SAVE"}, ) + snapform_moge_resolution_level: bpy.props.IntProperty( + default=5, + min=0, + max=9, + options={"HIDDEN", "SKIP_SAVE"}, + ) def execute(self, context): global _active_job @@ -236,6 +250,8 @@ class RunJob(bpy.types.Operator): export_format = "mini_npz" elif self.job_type == IMAGE_MESH_JOB: export_format = "mesh_data" + elif self.job_type in {MOGE2_DEPTH_JOB, MOGE2_DEBUG_JOB}: + export_format = "moge_npz" job = ActiveJob( kind=self.job_type, @@ -300,12 +316,20 @@ class RunJob(bpy.types.Operator): if job.kind == "INSTALL_ENV": client.stop_service() job.log_file = job.log_path.open("w", encoding="utf-8") + process_environment = None + if job.kind == "INSTALL_ENV": + process_environment = os.environ.copy() + process_environment["PYTHONIOENCODING"] = ( + "utf-8:backslashreplace" + ) + process_environment["PYTHONUTF8"] = "1" job.process = subprocess.Popen( command, stdout=job.log_file, stderr=subprocess.STDOUT, cwd=working_directory, creationflags=creation_flags, + env=process_environment, ) job.started = True if job.cancelled: @@ -401,6 +425,35 @@ class RunJob(bpy.types.Operator): ), paths, ) + if self.job_type in {MOGE2_DEPTH_JOB, MOGE2_DEBUG_JOB}: + model_key = self.snapform_model_key or MOGE2_MODEL.key + model = get_downloadable_model(model_key) + if not client.model_ready(model_key): + raise RuntimeError(f"Download {model.label} first") + if remove_background and not client.model_ready(BEN2_MODEL.key): + raise RuntimeError(f"Download {BEN2_MODEL.label} first") + return ( + client.moge_request( + settings, + paths["status"], + paths["result"], + input_path=input_path, + resolution_level=( + self.snapform_moge_resolution_level + if self.job_type == MOGE2_DEPTH_JOB + else settings.moge_resolution_level + ), + include_points=self.job_type == MOGE2_DEBUG_JOB, + remove_background=( + self.job_type == MOGE2_DEPTH_JOB + and remove_background + ), + start_frame=start_frame, + device=self._task_device(settings), + model_key=model_key, + ), + paths, + ) if self.job_type not in { RECONSTRUCTION_JOB, IMAGE_DEPTH_JOB, @@ -478,6 +531,8 @@ class RunJob(bpy.types.Operator): return f"Starting {settings.export_format} export" if self.job_type == IMAGE_DEPTH_JOB: return "Starting depth estimation" + if self.job_type in {MOGE2_DEPTH_JOB, MOGE2_DEBUG_JOB}: + return "Starting MoGe-2 inference" if self.job_type == IMAGE_MESH_JOB: return "Starting perspective surface generation" if self.job_type == BACKGROUND_REMOVAL_JOB: @@ -626,6 +681,7 @@ class RunJob(bpy.types.Operator): if job.kind in { RECONSTRUCTION_JOB, BACKGROUND_REMOVAL_JOB, + MOGE2_DEBUG_JOB, }: try: open_job_directory(job.result_directory) @@ -635,8 +691,10 @@ class RunJob(bpy.types.Operator): message = ( f"Exported {job.export_format} to {job.result_directory}" ) - else: + elif job.kind == BACKGROUND_REMOVAL_JOB: message = f"Removed background to {job.result_directory}" + else: + message = f"Exported MoGe-2 data to {job.result_directory}" self.report({"INFO"}, message) if job.kind == IMAGE_DEPTH_JOB: try: @@ -653,6 +711,23 @@ class RunJob(bpy.types.Operator): self._redraw_ui(context, force=True) return {"CANCELLED"} status["message"] = "Depth plane is ready" + if job.kind == MOGE2_DEPTH_JOB: + try: + settings.status = "Building MoGe-2 depth plane" + settings.progress = 0.96 + self._redraw_ui(context, force=True) + self._build_moge2_depth_plane(job) + except (OSError, RuntimeError, ValueError) as import_error: + settings.status = "MoGe-2 depth import failed" + settings.progress = 0.0 + self.report({"ERROR"}, str(import_error)) + client.mark_service_job_complete( + client.job_id_from_status(job.status_path) + ) + close_active_job() + self._redraw_ui(context, force=True) + return {"CANCELLED"} + status["message"] = "MoGe-2 depth plane is ready" if job.kind == IMAGE_MESH_JOB: try: settings.status = "Building perspective surface" @@ -790,6 +865,35 @@ class RunJob(bpy.types.Operator): if "FINISHED" not in undo_result: raise RuntimeError("Unable to create the depth plane Undo step") + def _build_moge2_depth_plane(self, job): + result_directory = job.result_directory + is_sequence = bool(list(result_directory.glob("frame.*"))) + foreground_sequence = bool( + list(result_directory.glob("foreground.*.png")) + ) + result = bpy.ops.snapform.build_moge2_depth_plane( + snapform_source_object_name=job.source_object_name, + snapform_npz_path=str(result_directory / "results.npz"), + snapform_background_path=( + str(result_directory) + if job.remove_background and foreground_sequence + else str(result_directory / "foreground.png") + if job.remove_background + else "" + ), + snapform_color_frames_dir=( + str(result_directory) if is_sequence else "" + ), + snapform_data_frames_dir=str(result_directory), + snapform_delete_source=job.delete_source, + snapform_unshaded=job.unshaded, + ) + if "FINISHED" not in result: + raise RuntimeError("Unable to build the MoGe-2 depth plane") + undo_result = bpy.ops.ed.undo_push(message="Convert to Depth Plane") + if "FINISHED" not in undo_result: + raise RuntimeError("Unable to create the depth plane Undo step") + def _apply_removed_background(self, job, push_undo=True): result_directory = job.result_directory foreground_paths = sorted(result_directory.glob("foreground.*.png")) diff --git a/src/snapform/panel.py b/src/snapform/panel.py index 1340f6d..3a2de93 100644 --- a/src/snapform/panel.py +++ b/src/snapform/panel.py @@ -1,6 +1,6 @@ import bpy -from .models import BEN2_MODEL, MODELS, selected_model_key +from .models import BEN2_MODEL, MODELS, MOGE2_MODELS, selected_model_key from .operators import active_job from . import client from .client import model_ready, env_installed @@ -126,6 +126,9 @@ class ModelsPanel(bpy.types.Panel): running = active_job() is not None or client.service_busy() for model in MODELS.values(): + if model.enabled: + self._draw_model(layout, model, running) + for model in MOGE2_MODELS.values(): self._draw_model(layout, model, running) self._draw_model(layout, BEN2_MODEL, running) @@ -236,6 +239,36 @@ class Da3OutputDebugPanel(bpy.types.Panel): operator.job_type = "RECONSTRUCT" +class Moge2ApiDebugPanel(bpy.types.Panel): + bl_label = "MoGe-2 API Test" + bl_idname = "SNAPFORM_PT_debug_moge2_api" + bl_parent_id = DebugPanel.bl_idname + bl_space_type = "VIEW_3D" + bl_region_type = "UI" + bl_category = "SnapForm" + bl_options = {"DEFAULT_CLOSED"} + bl_order = 1 + + def draw(self, context): + layout = self.layout + settings = context.scene.snapform_settings + layout.prop(settings, "input_path") + layout.prop(settings, "moge_resolution_level") + layout.prop(settings, "moge_model") + run_row = layout.row() + run_row.enabled = ( + active_job() is None + and not client.service_busy() + and model_ready(settings.moge_model) + ) + operator = run_row.operator( + "snapform.run_job", + text="Generate MoGe-2 Output", + icon="EXPORT", + ) + operator.job_type = "MOGE2_DEBUG" + + class Ben2RemovalDebugPanel(bpy.types.Panel): bl_label = "BEN2 Removal Test" bl_idname = "SNAPFORM_PT_debug_ben2_removal" diff --git a/src/snapform/properties.py b/src/snapform/properties.py index 19e5028..9a03aa4 100644 --- a/src/snapform/properties.py +++ b/src/snapform/properties.py @@ -1,6 +1,6 @@ import bpy -from .models import gs_model_enum_items, model_enum_items +from .models import gs_model_enum_items, model_enum_items, moge2_model_enum_items from .client import device_enum_items @@ -65,6 +65,19 @@ class SnapFormSettings(bpy.types.PropertyGroup): ), default="upper_bound_resize", ) + moge_resolution_level: bpy.props.IntProperty( + name="Resolution Level", + description="MoGe-2 inference resolution level; higher is finer and slower", + default=5, + min=0, + max=9, + ) + moge_model: bpy.props.EnumProperty( + name="Model", + description="MoGe-2 model used for depth estimation", + items=moge2_model_enum_items, + default=1, + ) conf_thresh_percentile: bpy.props.FloatProperty( name="Confidence Threshold", description="Remove the lowest-confidence percentile of points or surface faces", diff --git a/src/snapform/service.py b/src/snapform/service.py index c3e625d..8388a1b 100644 --- a/src/snapform/service.py +++ b/src/snapform/service.py @@ -20,6 +20,9 @@ import worker CLIENT_TIMEOUT_SECONDS = 10.0 +JOB_COMMANDS = frozenset( + {"remove-background", "reconstruct", "extract-frames", "moge"} +) IDLE_SHUTDOWN_SECONDS = 20.0 @@ -78,6 +81,8 @@ class ModelCache: self.ben2_key = None self.da3_model = None self.da3_key = None + self.moge_model = None + self.moge_key = None def get_ben2(self, model_dir, device): key = (str(Path(model_dir).resolve()), device.type) @@ -88,6 +93,7 @@ class ModelCache: return self._load_ben2(model_dir, device, key) except self.torch.cuda.OutOfMemoryError: self.release_da3() + self.release_moge() return self._load_ben2(model_dir, device, key) def _load_ben2(self, model_dir, device, key): @@ -111,6 +117,7 @@ class ModelCache: return self._load_da3(model_dir, device, key) except self.torch.cuda.OutOfMemoryError: self.release_ben2() + self.release_moge() return self._load_da3(model_dir, device, key) def _load_da3(self, model_dir, device, key): @@ -123,6 +130,30 @@ class ModelCache: self.da3_key = key return model + def get_moge(self, model_dir, device): + key = (str(Path(model_dir).resolve()), device.type) + if self.moge_model is not None and self.moge_key == key: + return self.moge_model + self.release_moge() + try: + return self._load_moge(model_dir, device, key) + except self.torch.cuda.OutOfMemoryError: + self.release_da3() + self.release_ben2() + return self._load_moge(model_dir, device, key) + + def _load_moge(self, model_dir, device, key): + worker.prepare_moge_runtime() + from moge.model.v2 import MoGeModel + + checkpoint = Path(model_dir) / "model.pt" + model = MoGeModel.from_pretrained(checkpoint) + model = model.to(device=device) + model.eval() + self.moge_model = model + self.moge_key = key + return model + def release_ben2(self): self.ben2_model = None self.ben2_key = None @@ -133,11 +164,18 @@ class ModelCache: self.da3_key = None self._release_cuda_cache() + def release_moge(self): + self.moge_model = None + self.moge_key = None + self._release_cuda_cache() + def close(self): self.ben2_model = None self.da3_model = None + self.moge_model = None self.ben2_key = None self.da3_key = None + self.moge_key = None self._release_cuda_cache() def status(self): @@ -150,6 +188,11 @@ class ModelCache: if self.da3_key is not None else "" ), + "moge2": ( + Path(getattr(self, "moge_key", ("",))[0]).name + if getattr(self, "moge_key", None) is not None + else "" + ), } def _release_cuda_cache(self): @@ -281,6 +324,12 @@ class ServiceState: args, cancel_check=cancel_check, ) + elif payload["command"] == "moge": + worker.run_moge( + args, + model_cache=self.models, + cancel_check=cancel_check, + ) else: worker.reconstruct( args, @@ -429,11 +478,7 @@ def create_app(token, instance_id="test"): client_id = str(payload.get("client_id", "")) command = payload.get("command") parameters = payload.get("parameters") - if not job_id or command not in { - "remove-background", - "reconstruct", - "extract-frames", - }: + if not job_id or command not in JOB_COMMANDS: raise HTTPException(status_code=422, detail="Invalid job request") if not isinstance(parameters, dict) or not parameters.get("status"): raise HTTPException(status_code=422, detail="Missing job parameters") diff --git a/src/snapform/worker.py b/src/snapform/worker.py index c5c04d7..009e024 100644 --- a/src/snapform/worker.py +++ b/src/snapform/worker.py @@ -1,5 +1,5 @@ import argparse -import fnmatch +import hashlib import json import os import shutil @@ -7,7 +7,7 @@ import tempfile import threading import time from pathlib import Path -from urllib.parse import quote, urlencode +from urllib.parse import quote from urllib.request import Request, urlopen @@ -33,9 +33,8 @@ EXPORT_FORMATS = ( "mesh_data", ) HF_ENDPOINT = "https://huggingface.co" -MODELSCOPE_ENDPOINT = "https://modelscope.cn" -MODELSCOPE_REVISION = "master" -MODELSCOPE_TIMEOUT = 60.0 +R2_ENDPOINT = "https://models.omoolab.xyz" +DOWNLOAD_TIMEOUT = 60.0 class TaskCancelled(RuntimeError): @@ -925,21 +924,23 @@ def download_model(args): "stages": getattr(args, "stages", None), "stage_label": getattr(args, "stage_label", None), } - from huggingface_hub import snapshot_download - - progress_class = create_download_progress_class( - args.status, - args.repository, - start=start, - span=span, - label=label, - **stage_kwargs, - ) failures = [] - for endpoint in download_endpoints(): - if not hf_endpoint_reachable(endpoint): - failures.append((endpoint, "endpoint unreachable")) - continue + try: + from huggingface_hub import snapshot_download + except Exception as error: + snapshot_download = None + failures.append((HF_ENDPOINT, error)) + progress_class = None + if snapshot_download: + progress_class = create_download_progress_class( + args.status, + args.repository, + start=start, + span=span, + label=label, + **stage_kwargs, + ) + for endpoint in download_endpoints() if snapshot_download else (): os.environ["HF_ENDPOINT"] = endpoint message = f"Downloading {label}" if endpoint != HF_ENDPOINT: @@ -953,7 +954,10 @@ def download_model(args): allow_patterns=args.allow_pattern or None, endpoint=endpoint, ) - _verify_model_complete(model_dir) + _verify_model_complete( + model_dir, + getattr(args, "required_pattern", None), + ) except Exception as error: failures.append((endpoint, error)) continue @@ -964,15 +968,22 @@ def download_model(args): **stage_kwargs, ) return - # Hugging Face 官方源与已配置镜像均失败时，回退 ModelScope（仓库名与 Hugging Face 一致） + # Hugging Face 失败后使用由项目维护的静态模型镜像。 + r2_directory = getattr(args, "r2_directory", "") + r2_files = getattr(args, "r2_file", []) + if not r2_directory or not r2_files: + detail = "; ".join( + f"{source}: {error}" for source, error in failures + ) + raise RuntimeError(f"All download sources failed: {detail}") try: write_status( args.status, start, - f"Downloading {label} | ModelScope mirror", + f"Downloading {label} | OmooLab mirror", **stage_kwargs, ) - download_from_modelscope( + download_from_r2( args, model_dir, start=start, @@ -980,9 +991,12 @@ def download_model(args): label=label, stage_kwargs=stage_kwargs, ) - _verify_model_complete(model_dir) + _verify_model_complete( + model_dir, + getattr(args, "required_pattern", None), + ) except Exception as error: - failures.append((MODELSCOPE_ENDPOINT, error)) + failures.append((R2_ENDPOINT, error)) else: write_status( args.status, @@ -998,33 +1012,23 @@ def download_model(args): def download_endpoints(): - """Hugging Face download candidates: configured mirror first, official endpoint otherwise.""" + """Return the configured Hugging Face endpoint or official source.""" endpoint = os.environ.get("HF_ENDPOINT") or HF_ENDPOINT return [endpoint.rstrip("/")] -def hf_endpoint_reachable(endpoint, timeout=10.0): - """Quick connectivity probe so slow or blocked Hugging Face endpoints fail fast.""" - try: - request = Request( - f"{endpoint.rstrip('/')}/api/models", - method="HEAD", - headers={"User-Agent": "SnapForm"}, - ) - with urlopen(request, timeout=timeout): - return True - except Exception: - return False - - -def _verify_model_complete(model_dir): +def _verify_model_complete(model_dir, required_patterns=None): """Refuse to report success when the required model files are missing.""" model_dir = Path(model_dir) - missing = [] - if not (model_dir / "config.json").is_file(): - missing.append("config.json") - if not any(model_dir.glob("*.safetensors")): - missing.append("*.safetensors") + required_patterns = required_patterns or ( + "config.json", + "*.safetensors", + ) + missing = [ + pattern + for pattern in required_patterns + if not any(model_dir.glob(pattern)) + ] if missing: raise RuntimeError( "Model download reported success but files are missing: " @@ -1032,59 +1036,42 @@ def _verify_model_complete(model_dir): ) -def _hidden_repo_path(path): - return any(part.startswith(".") for part in Path(path).parts) - - -def _model_file_allowed(path, patterns): - if not patterns: - return True - return any(fnmatch.fnmatch(path, pattern) for pattern in patterns) - - -def modelscope_file_list(repository, revision=MODELSCOPE_REVISION): - """List (path, size) file entries of a ModelScope repository.""" - url = ( - f"{MODELSCOPE_ENDPOINT}/api/v1/models/{repository}/repo/files?" - + urlencode({"Revision": revision, "Recursive": "true"}) - ) - request = Request(url, headers={"User-Agent": "SnapForm"}) - with urlopen(request, timeout=MODELSCOPE_TIMEOUT) as response: - payload = json.loads(response.read().decode("utf-8")) - entries = [] - for entry in payload.get("Data", {}).get("Files", []): - path = entry.get("Path") - if path and entry.get("Type") == "blob" and not _hidden_repo_path(path): - entries.append((path, int(entry.get("Size") or 0))) - return entries - - -def _download_modelscope_file( - repository, - path, +def _download_r2_file( + directory, + filename, destination, + expected_size, + expected_checksum, progress=None, retries=2, ): - """Download one ModelScope file atomically with a per-chunk progress callback.""" - url = ( - f"{MODELSCOPE_ENDPOINT}/models/{repository}/resolve/" - f"{MODELSCOPE_REVISION}/{quote(path)}" - ) + """Download and verify one OmooLab R2 model file atomically.""" + url = f"{R2_ENDPOINT}/{quote(directory)}/{quote(filename)}" request = Request(url, headers={"User-Agent": "SnapForm"}) last_error = None for _attempt in range(retries): temporary = destination.with_name(f"{destination.name}.part") try: - with urlopen(request, timeout=MODELSCOPE_TIMEOUT) as response: + checksum = hashlib.sha256() + downloaded = 0 + with urlopen(request, timeout=DOWNLOAD_TIMEOUT) as response: with open(temporary, "wb") as output: while True: chunk = response.read(1024 * 1024) if not chunk: break output.write(chunk) + checksum.update(chunk) + downloaded += len(chunk) if progress is not None: progress(len(chunk)) + if downloaded != expected_size: + raise RuntimeError( + f"Size mismatch for {filename}: " + f"{downloaded} != {expected_size}" + ) + if checksum.hexdigest().lower() != expected_checksum.lower(): + raise RuntimeError(f"SHA-256 mismatch for {filename}") temporary.replace(destination) return except Exception as error: @@ -1096,7 +1083,7 @@ def _download_modelscope_file( raise last_error -def download_from_modelscope( +def download_from_r2( args, model_dir, start=0.05, @@ -1104,19 +1091,12 @@ def download_from_modelscope( label=None, stage_kwargs=None, ): - """Download a model from ModelScope using only standard library requests.""" + """Download the declared model files from OmooLab R2.""" repository = args.repository - files = modelscope_file_list(repository) - patterns = args.allow_pattern or None - selected = [ - (path, size) - for path, size in files - if _model_file_allowed(path, patterns) - ] - if not selected: - raise RuntimeError( - f"No matching files in ModelScope repository: {repository}" - ) + selected = [] + for value in args.r2_file: + filename, size, checksum = value.split("|", 2) + selected.append((filename, int(size), checksum)) reporter = DownloadProgressReporter( args.status, repository, @@ -1125,7 +1105,7 @@ def download_from_modelscope( label=label or repository, **(stage_kwargs or {}), ) - total_size = sum(size for _path, size in selected) + total_size = sum(size for _path, size, _checksum in selected) downloaded = 0 def on_chunk(count): @@ -1133,13 +1113,15 @@ def download_from_modelscope( downloaded += count reporter.report(downloaded, total_size) - for path, _size in selected: + for path, size, checksum in selected: destination = Path(model_dir) / path destination.parent.mkdir(parents=True, exist_ok=True) - _download_modelscope_file( - repository, + _download_r2_file( + args.r2_directory, path, destination, + size, + checksum, progress=on_chunk, ) @@ -1500,6 +1482,168 @@ def flatten_exported_paths( return flattened_paths +def _save_moge_preview(output_directory, index, depth, normal, mask): + import numpy as np + from PIL import Image + + suffix = "" if index is None else f".{index + 1:04d}" + valid = mask & np.isfinite(depth) & (depth > 0) + preview = np.zeros(depth.shape, dtype=np.uint8) + if np.any(valid): + low, high = np.percentile(depth[valid], (2.0, 98.0)) + span = max(float(high - low), 1e-6) + preview[valid] = np.clip( + (depth[valid] - low) / span * 255.0, + 0.0, + 255.0, + ).astype(np.uint8) + Image.fromarray(preview, mode="L").save( + output_directory / f"depth{suffix}.png" + ) + normal_rgb = np.clip( + (normal * np.array([1.0, -1.0, -1.0], dtype=np.float32) + 1.0) + * 127.5, + 0.0, + 255.0, + ).astype(np.uint8) + Image.fromarray(normal_rgb, mode="RGB").save( + output_directory / f"normal{suffix}.png" + ) + Image.fromarray(mask.astype(np.uint8) * 255, mode="L").save( + output_directory / f"mask{suffix}.png" + ) + + +def run_moge(args, model_cache=None, cancel_check=None): + import numpy as np + from PIL import Image + + prepare_moge_runtime() + write_status(args.status, 0.02, "Starting MoGe-2") + if model_cache is None: + import torch + from moge.model.v2 import MoGeModel + else: + torch = model_cache.torch + output_directory = Path(args.output_dir) + output_directory.mkdir(parents=True, exist_ok=True) + images, temporary_directory = prepare_input( + args.input, + args.video_frames, + output_directory, + start_frame=getattr(args, "start_frame", 1), + cancel_check=cancel_check, + ) + try: + device = select_device(args.device, torch) + write_status(args.status, 0.08, f"Loading MoGe-2 on {device.type}") + if model_cache is None: + checkpoint = Path(args.model_dir) / "model.pt" + model = MoGeModel.from_pretrained(checkpoint) + model = model.to(device=device) + model.eval() + else: + model = model_cache.get_moge(args.model_dir, device) + + depths = [] + normals = [] + masks = [] + points = [] + intrinsics = [] + for index, image_path in enumerate(images): + check_cancelled(cancel_check) + write_status( + args.status, + 0.15 + 0.68 * index / len(images), + f"Estimating MoGe-2 geometry ({index + 1}/{len(images)})", + ) + with Image.open(image_path) as opened: + image = np.asarray(opened.convert("RGB"), dtype=np.float32) / 255.0 + image_tensor = torch.from_numpy(image).permute(2, 0, 1).to(device) + with torch.inference_mode(): + prediction = model.infer( + image_tensor, + resolution_level=int(args.resolution_level), + force_projection=True, + apply_mask=False, + use_fp16=device.type == "cuda", + ) + depth = prediction["depth"].detach().float().cpu().numpy() + normal = prediction["normal"].detach().float().cpu().numpy() + mask = prediction["mask"].detach().cpu().numpy().astype(bool) + point = prediction["points"].detach().float().cpu().numpy() + intrinsic = prediction["intrinsics"].detach().float().cpu().numpy() + depths.append(depth) + normals.append(normal) + masks.append(mask) + intrinsics.append(intrinsic) + if bool(getattr(args, "include_points", False)): + points.append(point) + _save_moge_preview( + output_directory, + None if len(images) == 1 else index, + depth, + normal, + mask, + ) + + shapes = {array.shape for array in depths} + if len(shapes) != 1: + raise RuntimeError("MoGe-2 frame sizes must match") + result = { + "depth": np.stack(depths).astype(np.float32), + "normal": np.stack(normals).astype(np.float32), + "mask": np.stack(masks), + "intrinsics": np.stack(intrinsics).astype(np.float32), + } + if points: + result["points"] = np.stack(points).astype(np.float32) + np.savez(output_directory / "results.npz", **result) + + remove_background_model_dir = getattr( + args, "remove_background_model_dir", "" + ) + if remove_background_model_dir: + write_status(args.status, 0.87, "Removing image background") + if len(images) == 1: + foreground, _, _ = infer_foreground( + images[0], + remove_background_model_dir, + device, + torch, + model_cache=model_cache, + cancel_check=cancel_check, + ) + foreground.save(output_directory / "foreground.png") + else: + infer_foreground_frames( + images, + remove_background_model_dir, + device, + torch, + output_directory, + args.status, + 0.87, + 0.08, + model_cache=model_cache, + cancel_check=cancel_check, + ) + copy_frames(images, output_directory) + write_status( + args.status, + 1.0, + f"Exported MoGe-2 data for {len(images)} frame(s)", + ) + finally: + if temporary_directory is not None: + shutil.rmtree(temporary_directory, ignore_errors=True) + + +def prepare_moge_runtime(): + """Use PyTorch SDPA because xFormers does not support every CUDA GPU.""" + os.environ["XFORMERS_DISABLED"] = "1" + + def reconstruct(args, model_cache=None, cancel_check=None): write_status(args.status, 0.02, "Starting Environment") check_cancelled(cancel_check) @@ -1686,6 +1830,9 @@ def parser(): download.add_argument("--model-dir", required=True) download.add_argument("--status", required=True) download.add_argument("--allow-pattern", action="append", default=[]) + download.add_argument("--required-pattern", action="append", default=[]) + download.add_argument("--r2-directory", default="") + download.add_argument("--r2-file", action="append", default=[]) download.add_argument("--stage", type=int, default=None) download.add_argument("--stages", type=int, default=None) download.add_argument("--stage-label", default=None) diff --git a/tests/test_addon.py b/tests/test_addon.py index 3003b5f..0d0575b 100644 --- a/tests/test_addon.py +++ b/tests/test_addon.py @@ -138,6 +138,15 @@ class AddonTest(unittest.TestCase): self.addon.BuildDepthPlane.bl_options, {"INTERNAL"}, ) + self.assertIn( + "snapform_color_frames_dir", + self.addon.BuildDepthPlane.__annotations__, + ) + self.assertIn( + "snapform_depth_frames_dir", + self.addon.BuildDepthPlane.__annotations__, + ) + self.assertIn("execute", self.addon.BuildDepthPlane.__dict__) self.assertFalse(hasattr(self.addon, "ImportDepthPlane")) def test_registers_perspective_mesh_operators(self): @@ -553,7 +562,7 @@ class AddonTest(unittest.TestCase): def test_original_depth_background_starts_independent_job(self): operator = self.addon.ConvertImageToDepthPlane() - operator.snapform_depth_source = "ORIGINAL" + operator.snapform_model = "ORIGINAL" operator.snapform_unshaded = False operator.snapform_remove_background = True operator.snapform_keep_original = True @@ -586,7 +595,7 @@ class AddonTest(unittest.TestCase): def test_original_depth_animation_extracts_frames_and_background(self): operator = self.addon.ConvertImageToDepthPlane() - operator.snapform_depth_source = "ORIGINAL" + operator.snapform_model = "ORIGINAL" operator.snapform_unshaded = False operator.snapform_remove_background = True operator.snapform_keep_original = True @@ -630,7 +639,7 @@ class AddonTest(unittest.TestCase): def test_original_depth_without_background_uses_frame_extraction(self): operator = self.addon.ConvertImageToDepthPlane() - operator.snapform_depth_source = "ORIGINAL" + operator.snapform_model = "ORIGINAL" operator.snapform_unshaded = False operator.snapform_remove_background = False operator.snapform_keep_original = False @@ -664,10 +673,7 @@ class AddonTest(unittest.TestCase): def test_conversion_dialog_exposes_requested_settings(self): operator = self.addon.ConvertImageToDepthPlane - self.assertNotIn( - "default", - operator.__annotations__["snapform_model"], - ) + self.assertEqual(operator.__annotations__["snapform_model"]["default"], 1) self.assertTrue(operator.__annotations__["snapform_use_ray_pose"]["default"]) self.assertEqual( operator.__annotations__["snapform_process_res_factor"]["default"], @@ -693,23 +699,16 @@ class AddonTest(unittest.TestCase): self.assertEqual( set(operator.__annotations__), { - "snapform_depth_source", "snapform_unshaded", "snapform_model", "snapform_use_ray_pose", "snapform_process_res_factor", + "snapform_moge_resolution_level", "snapform_remove_background", "snapform_keep_original", }, ) - depth_sources = operator.__annotations__["snapform_depth_source"] - self.assertEqual(depth_sources["default"], "ORIGINAL") - self.assertEqual( - [item[0] for item in depth_sources["items"]], - ["ESTIMATE", "ORIGINAL"], - ) - self.assertEqual(depth_sources["items"][0][1], "AI Generated") - self.assertEqual(depth_sources["name"], "Depth from") + self.assertEqual(operator.__annotations__["snapform_model"]["name"], "Model") self.assertEqual( operator.__annotations__["snapform_unshaded"]["name"], "Unshaded", @@ -724,11 +723,12 @@ class AddonTest(unittest.TestCase): self.assertEqual( [item[0] for item in model_items], [ + "MOGE2_VITS_NORMAL", + "MOGE2_VITB_NORMAL", + "MOGE2_VITL_NORMAL", "SMALL", "BASE", - "LARGE_1_1", - "GIANT_1_1", - "NESTED_GIANT_LARGE_1_1", + "ORIGINAL", ], ) @@ -758,7 +758,7 @@ class AddonTest(unittest.TestCase): result = operator.invoke(context, None) self.assertIs(result, operator) - self.assertEqual(operator.snapform_model, "GIANT_1_1") + self.assertNotEqual(getattr(operator, "snapform_model", ""), "GIANT_1_1") def test_remove_background_defaults_to_empty_image_frames(self): source_object = SimpleNamespace( @@ -808,7 +808,6 @@ class AddonTest(unittest.TestCase): ) context = SimpleNamespace(object=source_object, scene=scene) operator = self.addon.ConvertImageToDepthPlane() - operator.snapform_depth_source = "ESTIMATE" operator.snapform_model = "BASE" operator.snapform_use_ray_pose = True operator.snapform_process_res_factor = 50.0 @@ -839,7 +838,7 @@ class AddonTest(unittest.TestCase): self.assertTrue(calls[0]["snapform_remove_background"]) self.assertNotIn("snapform_background_depth_distance", calls[0]) - def test_conversion_dialog_only_lists_installed_models(self): + def test_conversion_dialog_marks_missing_models(self): items_callback = self.addon.ConvertImageToDepthPlane.__annotations__[ "snapform_model" ]["items"] @@ -850,7 +849,20 @@ class AddonTest(unittest.TestCase): ): items = items_callback(None, None) - self.assertEqual([item[0] for item in items], ["SMALL", "BASE"]) + self.assertEqual( + [item[0] for item in items], + [ + "MOGE2_VITS_NORMAL", + "MOGE2_VITB_NORMAL", + "MOGE2_VITL_NORMAL", + "SMALL", + "BASE", + "ORIGINAL", + ], + ) + labels = {item[0]: item[1] for item in items} + self.assertEqual(labels["BASE"], "DA3 Base") + self.assertEqual(labels["MOGE2_VITB_NORMAL"], "MoGe-2 ViT-B Normal (Missing)") def test_update_status_formats_stage_progress(self): operator = self.addon.RunJob() @@ -936,7 +948,6 @@ class AddonTest(unittest.TestCase): def test_packed_image_is_exported_without_unpacking(self): operator = self.addon.ConvertImageToDepthPlane() - operator.snapform_depth_source = "ESTIMATE" operator.snapform_model = "GIANT_1_1" operator.snapform_use_ray_pose = True operator.snapform_process_res_factor = 50.0 @@ -1869,7 +1880,7 @@ class AddonTest(unittest.TestCase): self.assertEqual( identifiers, - ["GIANT_1_1", "NESTED_GIANT_LARGE_1_1"], + [], ) def test_device_enums_list_only_supported_devices(self): @@ -1966,7 +1977,7 @@ class AddonTest(unittest.TestCase): "model_ready", return_value=False, ): - self.assertFalse( + self.assertTrue( self.addon.ConvertImageToDepthPlane.poll(context) ) @@ -1977,6 +1988,7 @@ class AddonTest(unittest.TestCase): ) panels = ( self.addon.Da3OutputDebugPanel, + self.addon.Moge2ApiDebugPanel, self.addon.Ben2RemovalDebugPanel, ) for panel in panels: @@ -1984,7 +1996,7 @@ class AddonTest(unittest.TestCase): self.assertIn("DEFAULT_CLOSED", panel.bl_options) self.assertEqual( [panel.bl_label for panel in panels], - ["DA3 Output Test", "BEN2 Removal Test"], + ["DA3 Output Test", "MoGe-2 API Test", "BEN2 Removal Test"], ) def test_service_panel_exposes_shared_service_state_and_controls(self): diff --git a/tests/test_download_progress.py b/tests/test_download_progress.py index e96af51..9d69a90 100644 --- a/tests/test_download_progress.py +++ b/tests/test_download_progress.py @@ -168,7 +168,7 @@ class DownloadProgressTest(unittest.TestCase): self.assertFalse(status_path.exists()) - def test_download_endpoints_keeps_direct_endpoint(self): + def test_download_endpoints_uses_direct_endpoint(self): with patch.dict(os.environ, {"HF_ENDPOINT": ""}): endpoints = self.worker.download_endpoints() diff --git a/tests/test_install_env.py b/tests/test_install_env.py index 9dd7463..a01fb48 100644 --- a/tests/test_install_env.py +++ b/tests/test_install_env.py @@ -10,7 +10,7 @@ import unittest import zipfile from pathlib import Path from types import SimpleNamespace -from unittest.mock import patch +from unittest.mock import Mock, patch PROJECT_ROOT = Path(__file__).resolve().parents[1] @@ -32,6 +32,8 @@ class InstallEnvTest(unittest.TestCase): vendor.mkdir() (vendor / "depth-anything-3.tar.gz").write_bytes(b"da3") (vendor / "ben2.tar.gz").write_bytes(b"ben2") + (vendor / "utils3d.tar.gz").write_bytes(b"utils3d") + (vendor / "moge.tar.gz").write_bytes(b"moge") return vendor def test_installs_da3_from_archives_without_gsplat(self): @@ -155,19 +157,13 @@ class InstallEnvTest(unittest.TestCase): if any( repository in command for repository in ( - "depth-anything/DA3-SMALL", - "depth-anything/DA3-BASE", + "Ruicheng/moge-2-vitb-normal", "PramaLLC/BEN2", ) ) ] self.assertEqual(len(default_downloads), 2) - self.assertFalse( - any( - "depth-anything/DA3-BASE" in map(str, command) - for command in default_downloads - ) - ) + self.assertTrue(any("Ruicheng/moge-2-vitb-normal" in map(str, command) for command in default_downloads)) ben2_download = next( command for command in default_downloads @@ -663,6 +659,32 @@ class InstallEnvTest(unittest.TestCase): env["PIP_INDEX_URL"], "https://pypi.tuna.tsinghua.edu.cn/simple", ) + self.assertEqual( + env["PYTHONIOENCODING"], + "utf-8:backslashreplace", + ) + self.assertEqual(env["PYTHONUTF8"], "1") + + def test_standard_streams_are_reconfigured_for_unicode_logs(self): + bootstrap = load_install_env() + stdout = Mock() + stderr = Mock() + + with patch.object(bootstrap.sys, "stdout", stdout), patch.object( + bootstrap.sys, + "stderr", + stderr, + ): + bootstrap.configure_standard_streams() + + stdout.reconfigure.assert_called_once_with( + encoding="utf-8", + errors="backslashreplace", + ) + stderr.reconfigure.assert_called_once_with( + encoding="utf-8", + errors="backslashreplace", + ) def test_ready_json_records_installed_packages(self): bootstrap = load_install_env() @@ -706,6 +728,10 @@ class InstallEnvTest(unittest.TestCase): ["numpy==1.26.4", "torch==2.9.1+cu128"], ) self.assertEqual(ready["devices"], ["cuda", "cpu"]) + self.assertEqual( + ready["manifest_revision"], + bootstrap.MANIFEST_REVISION, + ) def test_broken_env_is_cleaned_before_reinstall(self): bootstrap = load_install_env() @@ -752,7 +778,9 @@ class InstallEnvTest(unittest.TestCase): (venv_dir / "Scripts").mkdir() (venv_dir / "Scripts" / "python.exe").write_text("python") (root / "manifest.json").write_text( - json.dumps({}), + json.dumps( + {"manifest_revision": bootstrap.MANIFEST_REVISION} + ), encoding="utf-8", ) (venv_dir / "marker.txt").write_text("keep") @@ -786,7 +814,7 @@ class InstallEnvTest(unittest.TestCase): def fake_run(command, environment=None): commands.append(command) - if any("DA3-SMALL" in str(value) for value in command): + if any("moge-2-vitb-normal" in str(value) for value in command): raise subprocess.CalledProcessError(1, command) with tempfile.TemporaryDirectory() as directory: @@ -880,6 +908,8 @@ class InstallEnvTest(unittest.TestCase): bundled.mkdir() (bundled / "depth-anything-3.tar.gz").write_bytes(b"da3") (bundled / "ben2.tar.gz").write_bytes(b"ben2") + (bundled / "utils3d.tar.gz").write_bytes(b"utils3d") + (bundled / "moge.tar.gz").write_bytes(b"moge") target, archive_name = bootstrap.uv_download_target() archive = bundled / f"uv-{target}.{archive_name}" if archive_name.endswith("zip"): @@ -1025,6 +1055,8 @@ class InstallEnvTest(unittest.TestCase): bundled.mkdir() (bundled / "depth-anything-3.tar.gz").write_bytes(b"da3") (bundled / "ben2.tar.gz").write_bytes(b"ben2") + (bundled / "utils3d.tar.gz").write_bytes(b"utils3d") + (bundled / "moge.tar.gz").write_bytes(b"moge") vendor = root / "vendor" bootstrap.copy_vendor_inputs(bundled, vendor) @@ -1046,6 +1078,8 @@ class InstallEnvTest(unittest.TestCase): bundled.mkdir() (bundled / "depth-anything-3.tar.gz").write_bytes(b"new") (bundled / "ben2.tar.gz").write_bytes(b"ben2") + (bundled / "utils3d.tar.gz").write_bytes(b"utils3d") + (bundled / "moge.tar.gz").write_bytes(b"moge") vendor = root / "vendor" vendor.mkdir() (vendor / "depth-anything-3.tar.gz").write_bytes(b"old") @@ -1086,6 +1120,8 @@ class InstallEnvTest(unittest.TestCase): vendor.mkdir() (vendor / "depth-anything-3.tar.gz").write_bytes(b"da3") (vendor / "ben2.tar.gz").write_bytes(b"ben2") + (vendor / "utils3d.tar.gz").write_bytes(b"utils3d") + (vendor / "moge.tar.gz").write_bytes(b"moge") uv_path = root / "uv" uv_path.write_bytes(b"uv") args = SimpleNamespace( @@ -1126,6 +1162,8 @@ class InstallEnvTest(unittest.TestCase): vendor.mkdir() (vendor / "depth-anything-3.tar.gz").write_bytes(b"da3") (vendor / "ben2.tar.gz").write_bytes(b"ben2") + (vendor / "utils3d.tar.gz").write_bytes(b"utils3d") + (vendor / "moge.tar.gz").write_bytes(b"moge") uv_path = root / "uv" uv_path.write_bytes(b"uv") args = SimpleNamespace( @@ -1150,9 +1188,15 @@ class InstallEnvTest(unittest.TestCase): for label, command, _environment, kind in stages if kind == "download" } - self.assertIn("da3_small", " ".join(map(str, downloads["DA3 Small"]))) + self.assertIn( + "moge2_vitb_normal", + " ".join(map(str, downloads["MoGe-2 ViT-B Normal"])), + ) self.assertIn("ben2_base", " ".join(map(str, downloads["BEN2 Base"]))) - self.assertNotIn("DA3 Base", downloads) + self.assertEqual( + set(downloads), + {"MoGe-2 ViT-B Normal", "BEN2 Base"}, + ) def test_install_does_not_verify_env(self): bootstrap = load_install_env() diff --git a/tests/test_models.py b/tests/test_models.py index 4b258aa..9770903 100644 --- a/tests/test_models.py +++ b/tests/test_models.py @@ -34,11 +34,18 @@ class ModelTest(unittest.TestCase): model = self.models.get_model("LARGE_1_1") self.assertEqual(model.license_name, "CC BY-NC 4.0") + self.assertFalse(model.enabled) def test_giant_models_expose_gaussian_head(self): for model_key in ("GIANT_1_1", "NESTED_GIANT_LARGE_1_1"): with self.subTest(model_key=model_key): self.assertTrue(self.models.get_model(model_key).supports_gs) + self.assertFalse(self.models.get_model(model_key).enabled) + + def test_only_uploaded_da3_models_are_enabled(self): + enabled = [key for key, model in self.models.MODELS.items() if model.enabled] + + self.assertEqual(enabled, ["SMALL", "BASE"]) def test_unknown_model_is_rejected(self): with self.assertRaisesRegex(ValueError, "Unknown DA3 model"): @@ -72,3 +79,22 @@ class ModelTest(unittest.TestCase): self.models.model_directory_name("BEN2_BASE"), "ben2_base", ) + + def test_moge2_uses_single_checkpoint_and_mit_license(self): + model = self.models.get_downloadable_model("MOGE2_VITB_NORMAL") + + self.assertEqual(model.repository, "Ruicheng/moge-2-vitb-normal") + self.assertEqual(model.license_name, "MIT") + self.assertEqual(model.download_patterns, ("model.pt",)) + self.assertEqual(model.ready_patterns, ("model.pt",)) + self.assertFalse(model.supports_gs) + + def test_moge2_models_match_uploaded_r2_variants(self): + self.assertEqual( + list(self.models.MOGE2_MODELS), + ["MOGE2_VITS_NORMAL", "MOGE2_VITB_NORMAL", "MOGE2_VITL_NORMAL"], + ) + self.assertEqual( + [model.r2_directory for model in self.models.MOGE2_MODELS.values()], + ["MoGe2ViTSNormal", "MoGe2ViTBNormal", "MoGe2ViTLNormal"], + ) diff --git a/tests/test_service.py b/tests/test_service.py index a477608..b630617 100644 --- a/tests/test_service.py +++ b/tests/test_service.py @@ -131,6 +131,9 @@ class ServiceTest(unittest.TestCase): state.reserve("other", "second", "reconstruct", "other.json") ) + def test_service_accepts_moge_command(self): + self.assertIn("moge", self.service.JOB_COMMANDS) + def test_cancelling_job_keeps_service_and_model_cache_alive(self): fake_models = SimpleNamespace( status=lambda: {"ben2": "BEN2_BASE", "da3": ""}, diff --git a/tests/test_worker.py b/tests/test_worker.py index c08f723..e07f196 100644 --- a/tests/test_worker.py +++ b/tests/test_worker.py @@ -436,7 +436,13 @@ class WorkerTest(unittest.TestCase): with self.assertRaisesRegex(RuntimeError, "Unsupported input type"): self.worker.prepare_input(input_path, 16, directory) - def test_download_model_skips_unreachable_endpoint(self): + def test_moge_runtime_disables_unsupported_xformers_attention(self): + with patch.dict(os.environ, {}, clear=True): + self.worker.prepare_moge_runtime() + + self.assertEqual(os.environ["XFORMERS_DISABLED"], "1") + + def test_download_model_uses_actual_request_without_preflight(self): calls = [] fake_hugging_face = ModuleType("huggingface_hub") fake_hugging_face.snapshot_download = lambda **options: calls.append( @@ -457,13 +463,6 @@ class WorkerTest(unittest.TestCase): "create_download_progress_class", return_value=object, ), patch.dict(os.environ, {"HF_ENDPOINT": ""}), patch.object( - self.worker, - "hf_endpoint_reachable", - return_value=False, - ), patch.object( - self.worker, - "download_from_modelscope", - ) as mocked_modelscope, patch.object( self.worker, "_verify_model_complete", return_value=None, @@ -472,8 +471,7 @@ class WorkerTest(unittest.TestCase): args.status = str(Path(directory) / "status.json") self.worker.download_model(args) - self.assertEqual(calls, []) - self.assertEqual(mocked_modelscope.call_count, 1) + self.assertEqual(len(calls), 1) def test_download_model_refuses_success_when_files_missing(self): fake_hugging_face = ModuleType("huggingface_hub") @@ -492,10 +490,7 @@ class WorkerTest(unittest.TestCase): self.worker, "create_download_progress_class", return_value=object, - ), patch.dict(os.environ, {"HF_ENDPOINT": ""}), patch.object( - self.worker, - "download_from_modelscope", - ): + ), patch.dict(os.environ, {"HF_ENDPOINT": ""}): args.model_dir = str(Path(directory) / "model") args.status = str(Path(directory) / "status.json") with self.assertRaisesRegex( @@ -527,39 +522,35 @@ class WorkerTest(unittest.TestCase): return_value=object, ), patch.dict(os.environ, {"HF_ENDPOINT": ""}), patch.object( self.worker, - "download_from_modelscope", + "download_from_r2", side_effect=OSError("mirror down"), ): + args.r2_directory = "DA3Small" + args.r2_file = ["config.json|100|abc"] args.model_dir = str(Path(directory) / "model") args.status = str(Path(directory) / "status.json") - with self.assertRaisesRegex(RuntimeError, "modelscope.cn"): + with self.assertRaisesRegex(RuntimeError, "models.omoolab.xyz"): self.worker.download_model(args) - def test_download_from_modelscope_downloads_selected_files(self): + def test_download_from_r2_downloads_declared_files(self): args = SimpleNamespace( repository="depth-anything/DA3-SMALL", model_dir="model", status="status.json", allow_pattern=["config.json"], + r2_directory="DA3Small", + r2_file=["config.json|100|abc"], ) downloaded = [] with tempfile.TemporaryDirectory() as directory, patch.object( self.worker, - "modelscope_file_list", - return_value=[ - ("config.json", 100), - ("model.safetensors", 200), - (".gitattributes", 10), - ], - ), patch.object( - self.worker, - "_download_modelscope_file", + "_download_r2_file", side_effect=lambda *args, **kwargs: downloaded.append(args[1]), ), patch.object(self.worker, "write_status"): args.model_dir = str(Path(directory) / "model") args.status = str(Path(directory) / "status.json") - self.worker.download_from_modelscope(args, args.model_dir) + self.worker.download_from_r2(args, args.model_dir) self.assertEqual(downloaded, ["config.json"]) diff --git a/tools/build_geometry_nodes_asset.py b/tools/build_geometry_nodes_asset.py index db99278..466a3ec 100644 --- a/tools/build_geometry_nodes_asset.py +++ b/tools/build_geometry_nodes_asset.py @@ -5,6 +5,7 @@ import bpy GROUP_NAME = "O DA3 Surface" DEPTH_PLANE_GROUP_NAME = "O DA3 Depth Plane" +MOGE_DEPTH_PLANE_GROUP_NAME = "O MoGe2 Depth Plane" MATERIAL_DISPLACEMENT_ATTRIBUTE = "o_material_displacement" BACKGROUND_DEPTH_ATTRIBUTE = "o_background_depth" DEPTH_SCALE_ATTRIBUTE = "o_depth_scale" @@ -689,7 +690,7 @@ def build_depth_plane_group(): "INPUT", "NodeSocketInt", ) - subdivision.default_value = 5 + subdivision.default_value = 6 subdivision.min_value = 0 subdivision.max_value = 10 displacement_strength = new_socket( @@ -708,7 +709,7 @@ def build_depth_plane_group(): "INPUT", "NodeSocketFloat", ) - background_depth.default_value = -1.0 + background_depth.default_value = 1.0 background_depth.min_value = -10000.0 background_depth.max_value = 10000.0 background_depth.subtype = "DISTANCE" @@ -786,6 +787,14 @@ def build_depth_plane_group(): (350.0, 250.0), {"Background Depth"}, ) + background_depth_value = nodes.new("ShaderNodeMath") + background_depth_value.operation = "MULTIPLY" + background_depth_value.inputs[1].default_value = -1.0 + background_depth_value.location = (550.0, 250.0) + links.new( + background_input.outputs["Background Depth"], + background_depth_value.inputs[0], + ) displacement_input = new_group_input( nodes, (950.0, -150.0), @@ -906,7 +915,7 @@ def build_depth_plane_group(): clamp_foreground.location = (600.0, 650.0) links.new(depth_texture.outputs["Color"], clamp_foreground.inputs[0]) links.new( - background_input.outputs["Background Depth"], + background_depth_value.outputs[0], clamp_foreground.inputs[1], ) foreground_mask = new_compare( @@ -925,7 +934,7 @@ def build_depth_plane_group(): masked_depth.location = (800.0, 650.0) links.new(foreground_mask.outputs["Result"], masked_depth.inputs[0]) links.new( - background_input.outputs["Background Depth"], + background_depth_value.outputs[0], masked_depth.inputs[2], ) links.new(clamp_foreground.outputs[0], masked_depth.inputs[3]) @@ -989,7 +998,7 @@ def build_depth_plane_group(): ( BACKGROUND_DEPTH_ATTRIBUTE, "FLOAT", - attributes_input.outputs["Background Depth"], + background_depth_value.outputs[0], ), ( DEPTH_SCALE_ATTRIBUTE, @@ -1020,11 +1029,225 @@ def build_depth_plane_group(): return group +def build_moge_depth_plane_group(): + group = bpy.data.node_groups.new( + MOGE_DEPTH_PLANE_GROUP_NAME, + "GeometryNodeTree", + ) + group.is_modifier = True + interface = group.interface + new_socket(interface, "Geometry", "INPUT", "NodeSocketGeometry") + displacement = new_socket( + interface, "Displacement", "INPUT", "NodeSocketMenu" + ) + subdivision = new_socket( + interface, "Subdivision", "INPUT", "NodeSocketInt" + ) + subdivision.default_value = 6 + subdivision.min_value = 0 + subdivision.max_value = 10 + depth_scale = new_socket( + interface, "Depth Scale", "INPUT", "NodeSocketFloat" + ) + depth_scale.default_value = 1.0 + depth_scale.structure_type = "SINGLE" + background_depth = new_socket( + interface, "Background Depth", "INPUT", "NodeSocketFloat" + ) + background_depth.default_value = 1.0 + background_depth.min_value = -10000.0 + background_depth.max_value = 10000.0 + background_depth.subtype = "DISTANCE" + background_depth.structure_type = "SINGLE" + data_panel = interface.new_panel(name="Data", default_closed=True) + depth_image = new_socket( + interface, "Depth Image", "INPUT", "NodeSocketImage", data_panel + ) + material = new_socket( + interface, "Material", "INPUT", "NodeSocketMaterial", data_panel + ) + for socket in (depth_image, material): + socket.hide_in_modifier = True + frame_sockets = [] + for name, socket_type, default in ( + ("Frame", "NodeSocketInt", 1), + ("Frame Start", "NodeSocketInt", 1), + ("Frames", "NodeSocketInt", 1), + ("Cyclic", "NodeSocketBool", False), + ): + socket = new_socket(interface, name, "INPUT", socket_type, data_panel) + socket.default_value = default + socket.hide_in_modifier = True + frame_sockets.append(socket) + new_socket(interface, "Geometry", "OUTPUT", "NodeSocketGeometry") + + nodes = group.nodes + links = group.links + geometry_input = new_group_input(nodes, (-1600.0, 650.0), {"Geometry"}) + controls = new_group_input( + nodes, + (-1600.0, 150.0), + {"Subdivision", "Displacement", "Depth Scale", "Background Depth"}, + ) + background_depth_value = nodes.new("ShaderNodeMath") + background_depth_value.operation = "MULTIPLY" + background_depth_value.inputs[1].default_value = -1.0 + background_depth_value.location = (-1350.0, -150.0) + links.new( + controls.outputs["Background Depth"], + background_depth_value.inputs[0], + ) + images = new_group_input( + nodes, (-500.0, 1100.0), {"Depth Image"} + ) + frames = new_group_input( + nodes, + (-500.0, 1500.0), + {"Frame", "Frame Start", "Frames", "Cyclic"}, + ) + material_input = new_group_input(nodes, (900.0, 200.0), {"Material"}) + output = nodes.new("NodeGroupOutput") + output.location = (1700.0, 650.0) + + bounds = nodes.new("GeometryNodeBoundBox") + bounds.location = (-1400.0, 650.0) + links.new(geometry_input.outputs["Geometry"], bounds.inputs["Geometry"]) + dimensions = nodes.new("ShaderNodeVectorMath") + dimensions.operation = "SUBTRACT" + dimensions.location = (-1200.0, 650.0) + links.new(bounds.outputs["Max"], dimensions.inputs[0]) + links.new(bounds.outputs["Min"], dimensions.inputs[1]) + separate = nodes.new("ShaderNodeSeparateXYZ") + separate.location = (-1000.0, 650.0) + links.new(dimensions.outputs["Vector"], separate.inputs[0]) + segments = nodes.new("ShaderNodeMath") + segments.operation = "POWER" + segments.inputs[0].default_value = 2.0 + segments.location = (-1200.0, 100.0) + links.new(controls.outputs["Subdivision"], segments.inputs[1]) + vertices = nodes.new("ShaderNodeMath") + vertices.operation = "ADD" + vertices.inputs[1].default_value = 1.0 + vertices.location = (-1000.0, 100.0) + links.new(segments.outputs[0], vertices.inputs[0]) + grid = nodes.new("GeometryNodeMeshGrid") + grid.location = (-750.0, 650.0) + links.new(separate.outputs["X"], grid.inputs["Size X"]) + links.new(separate.outputs["Y"], grid.inputs["Size Y"]) + links.new(vertices.outputs[0], grid.inputs["Vertices X"]) + links.new(vertices.outputs[0], grid.inputs["Vertices Y"]) + store_uv = nodes.new("GeometryNodeStoreNamedAttribute") + store_uv.data_type = "FLOAT2" + store_uv.domain = "CORNER" + store_uv.inputs["Name"].default_value = "UVMap" + store_uv.location = (-500.0, 650.0) + links.new(grid.outputs["Mesh"], store_uv.inputs["Geometry"]) + links.new(grid.outputs["UV Map"], store_uv.inputs["Value"]) + + center = nodes.new("ShaderNodeVectorMath") + center.operation = "ADD" + center.location = (-1000.0, 900.0) + links.new(bounds.outputs["Min"], center.inputs[0]) + links.new(bounds.outputs["Max"], center.inputs[1]) + half = nodes.new("ShaderNodeVectorMath") + half.operation = "SCALE" + half.inputs[3].default_value = 0.5 + half.location = (-750.0, 900.0) + links.new(center.outputs["Vector"], half.inputs[0]) + transform = nodes.new("GeometryNodeTransform") + transform.location = (-250.0, 650.0) + links.new(store_uv.outputs["Geometry"], transform.inputs["Geometry"]) + links.new(half.outputs["Vector"], transform.inputs["Translation"]) + + depth_texture = nodes.new("GeometryNodeImageTexture") + depth_texture.interpolation = "Linear" + depth_texture.extension = "EXTEND" + depth_texture.location = (-250.0, 1100.0) + links.new(images.outputs["Depth Image"], depth_texture.inputs["Image"]) + links.new(grid.outputs["UV Map"], depth_texture.inputs["Vector"]) + frame = build_sequence_frame( + nodes, links, frames.outputs["Frame"], frames, 1800.0 + ) + links.new(frame, depth_texture.inputs["Frame"]) + + clamp_depth = nodes.new("ShaderNodeMath") + clamp_depth.operation = "MAXIMUM" + clamp_depth.location = (0.0, 750.0) + links.new(depth_texture.outputs["Color"], clamp_depth.inputs[0]) + links.new(background_depth_value.outputs[0], clamp_depth.inputs[1]) + displacement_menu = nodes.new("GeometryNodeMenuSwitch") + displacement_menu.data_type = "FLOAT" + displacement_menu.enum_items.clear() + displacement_menu.enum_items.new("Geometry") + displacement_menu.enum_items.new("Material") + displacement_menu.inputs["Geometry"].default_value = 1.0 + displacement_menu.inputs["Material"].default_value = 0.0 + displacement_menu.location = (0.0, 300.0) + links.new(controls.outputs["Displacement"], displacement_menu.inputs["Menu"]) + displacement.default_value = "Geometry" + material_mode = nodes.new("ShaderNodeMath") + material_mode.operation = "SUBTRACT" + material_mode.inputs[0].default_value = 1.0 + material_mode.location = (250.0, 250.0) + links.new(displacement_menu.outputs[0], material_mode.inputs[1]) + geometry_strength = nodes.new("ShaderNodeMath") + geometry_strength.operation = "MULTIPLY" + geometry_strength.location = (250.0, 500.0) + links.new(displacement_menu.outputs[0], geometry_strength.inputs[0]) + links.new(controls.outputs["Depth Scale"], geometry_strength.inputs[1]) + scaled_depth = nodes.new("ShaderNodeMath") + scaled_depth.operation = "MULTIPLY" + scaled_depth.location = (250.0, 750.0) + links.new(clamp_depth.outputs[0], scaled_depth.inputs[0]) + links.new(geometry_strength.outputs[0], scaled_depth.inputs[1]) + normal = nodes.new("GeometryNodeInputNormal") + normal.location = (0.0, 500.0) + offset = nodes.new("ShaderNodeVectorMath") + offset.operation = "SCALE" + offset.location = (250.0, 700.0) + links.new(normal.outputs["Normal"], offset.inputs[0]) + links.new(scaled_depth.outputs[0], offset.inputs[3]) + set_position = nodes.new("GeometryNodeSetPosition") + set_position.location = (500.0, 650.0) + links.new(transform.outputs["Geometry"], set_position.inputs["Geometry"]) + links.new(offset.outputs["Vector"], set_position.inputs["Offset"]) + + set_material = nodes.new("GeometryNodeSetMaterial") + set_material.location = (750.0, 650.0) + links.new(set_position.outputs["Geometry"], set_material.inputs["Geometry"]) + links.new(material_input.outputs["Material"], set_material.inputs["Material"]) + smooth = nodes.new("GeometryNodeSetShadeSmooth") + smooth.domain = "FACE" + smooth.inputs["Shade Smooth"].default_value = True + smooth.location = (1000.0, 650.0) + links.new(set_material.outputs["Geometry"], smooth.inputs["Geometry"]) + geometry = smooth.outputs["Geometry"] + for index, (name, data_type, value) in enumerate(( + (MATERIAL_DISPLACEMENT_ATTRIBUTE, "BOOLEAN", material_mode.outputs[0]), + (BACKGROUND_DEPTH_ATTRIBUTE, "FLOAT", background_depth_value.outputs[0]), + (DEPTH_SCALE_ATTRIBUTE, "FLOAT", controls.outputs["Depth Scale"]), + )): + store = nodes.new("GeometryNodeStoreNamedAttribute") + store.data_type = data_type + store.domain = "POINT" + store.inputs["Name"].default_value = name + store.location = (1200.0 + index * 200.0, 650.0) + links.new(geometry, store.inputs["Geometry"]) + links.new(value, store.inputs["Value"]) + geometry = store.outputs["Geometry"] + links.new(geometry, output.inputs["Geometry"]) + group.color_tag = "GEOMETRY" + group.description = "Build a MoGe-2 depth plane with geometry or material displacement." + group.use_fake_user = True + return group + + def main(): bpy.ops.wm.read_factory_settings(use_empty=True) bpy.context.preferences.filepaths.save_version = 0 build_group() build_depth_plane_group() + build_moge_depth_plane_group() OUTPUT_PATH.parent.mkdir(parents=True, exist_ok=True) bpy.ops.wm.save_as_mainfile(filepath=str(OUTPUT_PATH), check_existing=False) diff --git a/tools/validate_depth_plane_asset.py b/tools/validate_depth_plane_asset.py index dc64a1a..2fb44c8 100644 --- a/tools/validate_depth_plane_asset.py +++ b/tools/validate_depth_plane_asset.py @@ -29,7 +29,15 @@ def main(): ) assert background_depth.subtype == "DISTANCE" assert background_depth.structure_type == "SINGLE" - assert background_depth.default_value == -1.0 + assert background_depth.default_value == 1.0 + subdivision = next( + item + for item in group.interface.items_tree + if item.item_type == "SOCKET" + and item.in_out == "INPUT" + and item.name == "Subdivision" + ) + assert subdivision.default_value == 6 frame_input = next( item for item in group.interface.items_tree diff --git a/tools/validate_moge_depth_plane.py b/tools/validate_moge_depth_plane.py new file mode 100644 index 0000000..240b9b8 --- /dev/null +++ b/tools/validate_moge_depth_plane.py @@ -0,0 +1,79 @@ +import sys +import tempfile +from pathlib import Path + +import bpy +import numpy as np + + +PROJECT_ROOT = Path(__file__).parents[1] +sys.path.insert(0, str(PROJECT_ROOT / "src")) + +from snapform.image_depth import ( + MOGE_DEPTH_PLANE_NODE_GROUP_NAME, + MOGE_MATERIAL_NODE_GROUP_NAME, + build_moge2_depth_plane, +) + + +def main(): + bpy.ops.wm.read_factory_settings(use_empty=True) + image = bpy.data.images.new("MoGe Test", width=64, height=48, alpha=True) + source = bpy.data.objects.new("MoGe Test", None) + source.empty_display_type = "IMAGE" + source.data = image + bpy.context.collection.objects.link(source) + depth = np.linspace(1.0, 2.0, 48 * 64, dtype=np.float32).reshape(1, 48, 64) + normal = np.zeros((1, 48, 64, 3), dtype=np.float32) + normal[..., 2] = -1.0 + mask = np.ones((1, 48, 64), dtype=bool) + mask[:, :5] = False + intrinsics = np.array( + [[[1.0, 0.0, 0.5], [0.0, 1.0, 0.5], [0.0, 0.0, 1.0]]], + dtype=np.float32, + ) + with tempfile.TemporaryDirectory() as temporary: + path = Path(temporary) / "results.npz" + np.savez( + path, + depth=depth, + normal=normal, + mask=mask, + intrinsics=intrinsics, + ) + plane = build_moge2_depth_plane( + bpy.context, + source, + path, + delete_source=True, + ) + assert plane.modifiers[0].node_group.name == MOGE_DEPTH_PLANE_NODE_GROUP_NAME + assert MOGE_MATERIAL_NODE_GROUP_NAME in bpy.data.node_groups + geometry_group = plane.modifiers[0].node_group + inputs = { + item.name + for item in geometry_group.interface.items_tree + if getattr(item, "in_out", None) == "INPUT" + } + assert "Displacement" in inputs + assert "Background Depth" in inputs + assert "Mask Image" not in inputs + assert "Use Mask" not in inputs + sockets = { + item.name: item + for item in geometry_group.interface.items_tree + if getattr(item, "in_out", None) == "INPUT" + } + assert sockets["Subdivision"].default_value == 6 + assert sockets["Depth Scale"].structure_type == "SINGLE" + assert sockets["Background Depth"].structure_type == "SINGLE" + assert sockets["Background Depth"].subtype == "DISTANCE" + assert sockets["Background Depth"].default_value == 1.0 + material_group = bpy.data.node_groups[MOGE_MATERIAL_NODE_GROUP_NAME] + assert not any(node.bl_idname == "ShaderNodeNormalMap" for node in material_group.nodes) + assert any(node.bl_idname == "ShaderNodeBump" for node in material_group.nodes) + print("MoGe-2 depth plane validation passed") + + +if __name__ == "__main__": + main() diff --git a/vendor.lock.toml b/vendor.lock.toml index a74fec4..0676be0 100644 --- a/vendor.lock.toml +++ b/vendor.lock.toml @@ -13,6 +13,20 @@ filename = "ben2.tar.gz" url = "https://codeload.github.com/PramaLLC/BEN2/tar.gz/2c99a5da477b5523585bfa5c893888a6e818a8f6" sha256 = "01fd26716f3ff3a3f18ad2336c768af5b759edfc7c38bb3bf733b7c86eee6282" +[artifacts.moge] +revision = "925b8ed835a7a9cdb7578ba15c658a0afc969030" +filename = "moge.tar.gz" +url = "https://codeload.github.com/microsoft/MoGe/tar.gz/925b8ed835a7a9cdb7578ba15c658a0afc969030" +sha256 = "c6d90b2564d643f0f3ae65b5ffe163abaea3b55d194c6e2ba9976d67f552bd26" +patches = "moge" + +[artifacts.utils3d] +revision = "3fab839f0be9931dac7c8488eb0e1600c236e183" +filename = "utils3d.tar.gz" +url = "https://codeload.github.com/EasternJournalist/utils3d/tar.gz/3fab839f0be9931dac7c8488eb0e1600c236e183" +sha256 = "f1f167b1f68809bea391a66dd0d2ad0e14f7382cbddf087b655181c31a285abb" +patches = "utils3d" + [uv] version = "0.11.3"