diff --git a/invokeai/app/api/dependencies.py b/invokeai/app/api/dependencies.py
index 3092f5ab71a..9db1ffe386b 100644
--- a/invokeai/app/api/dependencies.py
+++ b/invokeai/app/api/dependencies.py
@@ -6,6 +6,7 @@
import torch
from invokeai.app.services.app_settings import AppSettingsService
+from invokeai.app.services.asset_files.asset_files_disk import DiskAssetFileStorage
from invokeai.app.services.auth.token_service import set_jwt_secret
from invokeai.app.services.board_image_records.board_image_records_sqlite import SqliteBoardImageRecordStorage
from invokeai.app.services.board_images.board_images_default import BoardImagesService
@@ -108,6 +109,7 @@ def initialize(
raise ValueError("Output folder is not set")
image_files = DiskImageFileStorage(f"{output_folder}/images")
+ asset_files = DiskAssetFileStorage(output_folder / "assets")
model_images_folder = config.models_path
style_presets_folder = config.style_presets_path
@@ -192,6 +194,7 @@ def initialize(
users = UserService(db=db)
services = InvocationServices(
+ asset_files=asset_files,
board_image_records=board_image_records,
board_images=board_images,
board_records=board_records,
diff --git a/invokeai/app/api/routers/assets.py b/invokeai/app/api/routers/assets.py
new file mode 100644
index 00000000000..55dbbd306e4
--- /dev/null
+++ b/invokeai/app/api/routers/assets.py
@@ -0,0 +1,35 @@
+from fastapi import APIRouter, HTTPException, Path
+from fastapi.responses import FileResponse
+
+from invokeai.app.api.dependencies import ApiDependencies
+
+ASSET_MAX_AGE = 31536000
+
+assets_router = APIRouter(prefix="/v1/assets", tags=["assets"])
+
+
+@assets_router.get(
+ "/i/{asset_name}",
+ operation_id="get_asset",
+ responses={
+ 200: {"description": "The 3D asset was fetched successfully"},
+ 404: {"description": "The 3D asset could not be found"},
+ },
+ status_code=200,
+)
+async def get_asset(
+ asset_name: str = Path(description="The name of the 3D asset file to get"),
+) -> FileResponse:
+ """Gets a 3D asset file (e.g. a Gaussian-splat .ply)."""
+ try:
+ path = ApiDependencies.invoker.services.asset_files.get_path(asset_name)
+ response = FileResponse(
+ path,
+ media_type="application/octet-stream",
+ filename=asset_name,
+ content_disposition_type="inline",
+ )
+ response.headers["Cache-Control"] = f"max-age={ASSET_MAX_AGE}"
+ return response
+ except Exception:
+ raise HTTPException(status_code=404)
diff --git a/invokeai/app/api_app.py b/invokeai/app/api_app.py
index ed02d75eae3..87c49b3a348 100644
--- a/invokeai/app/api_app.py
+++ b/invokeai/app/api_app.py
@@ -17,6 +17,7 @@
from invokeai.app.api.no_cache_staticfiles import NoCacheStaticFiles
from invokeai.app.api.routers import (
app_info,
+ assets,
auth,
board_images,
boards,
@@ -179,6 +180,7 @@ async def dispatch(self, request: Request, call_next: RequestResponseEndpoint):
app.include_router(download_queue.download_queue_router, prefix="/api")
app.include_router(image_moves.image_moves_router, prefix="/api")
app.include_router(images.images_router, prefix="/api")
+app.include_router(assets.assets_router, prefix="/api")
app.include_router(boards.boards_router, prefix="/api")
app.include_router(board_images.board_images_router, prefix="/api")
app.include_router(virtual_boards.virtual_boards_router, prefix="/api")
diff --git a/invokeai/app/invocations/fields.py b/invokeai/app/invocations/fields.py
index 4418c86371a..1f88123de4d 100644
--- a/invokeai/app/invocations/fields.py
+++ b/invokeai/app/invocations/fields.py
@@ -241,6 +241,12 @@ class ImageField(BaseModel):
image_name: str = Field(description="The name of the image")
+class Asset3DField(BaseModel):
+ """A 3D asset primitive field"""
+
+ asset_name: str = Field(description="The name of the 3D asset file")
+
+
class BoardField(BaseModel):
"""A board primitive field"""
diff --git a/invokeai/app/invocations/image_to_3d.py b/invokeai/app/invocations/image_to_3d.py
new file mode 100644
index 00000000000..2be75cc0b0e
--- /dev/null
+++ b/invokeai/app/invocations/image_to_3d.py
@@ -0,0 +1,96 @@
+from pathlib import Path
+
+import torch
+
+from invokeai.app.invocations.baseinvocation import BaseInvocation, Classification, invocation
+from invokeai.app.invocations.fields import ImageField, InputField, WithBoard, WithMetadata
+from invokeai.app.invocations.primitives import Asset3DOutput
+from invokeai.app.services.session_processor.session_processor_common import CanceledException
+from invokeai.app.services.shared.invocation_context import InvocationContext
+from invokeai.app.util.misc import uuid_string
+from invokeai.backend.image_util.triposplat.triposplat_model import TripoSplatModel
+
+# HuggingFace repo id. The repo holds five safetensors checkpoints in subfolders (diffusion_models/, vae/,
+# clip_vision/, background_removal/). NOTE: we fetch it with huggingface_hub.snapshot_download rather than
+# InvokeAI's load_remote_model, because the latter's filter_files() only keeps weight files whose name
+# matches `model.*` — which excludes ALL of TripoSplat's checkpoints (it would resolve "0 files").
+# License: MIT (code + weights).
+TRIPOSPLAT_SOURCE = "VAST-AI/TripoSplat"
+
+# Bounds enforced by the upstream pipeline (TripoSplatPipeline._validate_num_gaussians).
+TRIPOSPLAT_MIN_GAUSSIANS = 32768
+TRIPOSPLAT_MAX_GAUSSIANS = 262144
+
+
+@invocation(
+ "image_to_3d",
+ title="Image to 3D (TripoSplat)",
+ tags=["3d", "image", "gaussian", "splat"],
+ category="3d",
+ version="1.0.0",
+ classification=Classification.Prototype,
+)
+class ImageTo3DInvocation(BaseInvocation, WithMetadata, WithBoard):
+ """Generates a 3D Gaussian splat (.ply) from a single image using TripoSplat."""
+
+ image: ImageField = InputField(description="The image to convert to a 3D Gaussian splat.")
+ remove_background: bool = InputField(
+ default=True,
+ description="Isolate the subject by removing the background (BiRefNet). Disable to feed the full image as-is.",
+ )
+ num_gaussians: int = InputField(
+ default=TRIPOSPLAT_MAX_GAUSSIANS,
+ ge=TRIPOSPLAT_MIN_GAUSSIANS,
+ le=TRIPOSPLAT_MAX_GAUSSIANS,
+ description="Number of 3D Gaussians to generate. Higher is more detailed but larger and slower to render.",
+ )
+ steps: int = InputField(default=20, ge=1, le=100, description="Number of flow-matching sampler steps.")
+ guidance_scale: float = InputField(
+ default=3.0, ge=0.0, description="Classifier-free guidance scale. Values <= 1.0 disable guidance."
+ )
+ seed: int = InputField(default=42, ge=0, description="Seed for reproducible generation.")
+
+ @torch.no_grad()
+ def invoke(self, context: InvocationContext) -> Asset3DOutput:
+ if self.remove_background:
+ # RGBA honors an existing alpha cutout; a fully-opaque input falls back to TripoSplat's own
+ # BiRefNet background removal (which isolates the single most salient subject).
+ image = context.images.get_pil(self.image.image_name, mode="RGBA")
+ else:
+ # Force-skip BiRefNet by handing the model a uniformly near-opaque alpha — TripoSplat only runs
+ # background removal when there's no real alpha channel, so this feeds the whole frame as-is.
+ image = context.images.get_pil(self.image.image_name, mode="RGB").convert("RGBA")
+ image.putalpha(254)
+
+ # Fetch the full repo snapshot directly (see TRIPOSPLAT_SOURCE note re: load_remote_model). Cached
+ # by huggingface_hub after the first download; subsequent runs resolve instantly.
+ from huggingface_hub import snapshot_download
+
+ context.util.signal_progress("Downloading TripoSplat weights (first run only)")
+ weights_dir = Path(snapshot_download(repo_id=TRIPOSPLAT_SOURCE))
+
+ loaded_model = context.models.load_local_model(weights_dir, TripoSplatModel.load_model)
+
+ with loaded_model as model:
+ assert isinstance(model, TripoSplatModel)
+
+ def _on_step(step: int, total: int) -> None:
+ if context.util.is_canceled():
+ raise CanceledException
+ context.util.signal_progress("Generating 3D Gaussian splat", step / total)
+
+ gaussian, _prepared = model._pipe.run(
+ image,
+ seed=self.seed,
+ steps=self.steps,
+ guidance_scale=self.guidance_scale,
+ num_gaussians=self.num_gaussians,
+ show_progress=False,
+ callback=_on_step,
+ )
+ ply_bytes = gaussian.to_ply_bytes()
+
+ asset_name = f"{uuid_string()}.ply"
+ # asset_files is a custom service with no curated InvocationContext accessor, so reach it directly.
+ context._services.asset_files.save(asset_name, ply_bytes)
+ return Asset3DOutput.build(asset_name=asset_name)
diff --git a/invokeai/app/invocations/primitives.py b/invokeai/app/invocations/primitives.py
index 6249de0cd8e..6ba30285edf 100644
--- a/invokeai/app/invocations/primitives.py
+++ b/invokeai/app/invocations/primitives.py
@@ -13,6 +13,7 @@
from invokeai.app.invocations.constants import LATENT_SCALE_FACTOR
from invokeai.app.invocations.fields import (
AnimaConditioningField,
+ Asset3DField,
BoundingBoxField,
CogView4ConditioningField,
ColorField,
@@ -253,6 +254,17 @@ def build(cls, image_dto: ImageDTO) -> "ImageOutput":
)
+@invocation_output("asset_3d_output")
+class Asset3DOutput(BaseInvocationOutput):
+ """Base class for nodes that output a single 3D asset"""
+
+ asset: Asset3DField = OutputField(description="The output 3D asset")
+
+ @classmethod
+ def build(cls, asset_name: str) -> "Asset3DOutput":
+ return cls(asset=Asset3DField(asset_name=asset_name))
+
+
@invocation_output("image_collection_output")
class ImageCollectionOutput(BaseInvocationOutput):
"""Base class for nodes that output a collection of images"""
diff --git a/invokeai/app/services/asset_files/__init__.py b/invokeai/app/services/asset_files/__init__.py
new file mode 100644
index 00000000000..e69de29bb2d
diff --git a/invokeai/app/services/asset_files/asset_files_base.py b/invokeai/app/services/asset_files/asset_files_base.py
new file mode 100644
index 00000000000..ddf3cc38c5e
--- /dev/null
+++ b/invokeai/app/services/asset_files/asset_files_base.py
@@ -0,0 +1,26 @@
+from abc import ABC, abstractmethod
+from pathlib import Path
+
+
+class AssetFilesServiceBase(ABC):
+ """Stores and serves binary 3D asset files (e.g. Gaussian-splat .ply / .splat)."""
+
+ @abstractmethod
+ def save(self, asset_name: str, data: bytes) -> None:
+ """Saves a 3D asset file."""
+ pass
+
+ @abstractmethod
+ def get_path(self, asset_name: str) -> Path:
+ """Gets the path to a 3D asset file."""
+ pass
+
+ @abstractmethod
+ def get_url(self, asset_name: str) -> str:
+ """Gets the URL of a 3D asset file."""
+ pass
+
+ @abstractmethod
+ def delete(self, asset_name: str) -> None:
+ """Deletes a 3D asset file."""
+ pass
diff --git a/invokeai/app/services/asset_files/asset_files_common.py b/invokeai/app/services/asset_files/asset_files_common.py
new file mode 100644
index 00000000000..ce10ba70621
--- /dev/null
+++ b/invokeai/app/services/asset_files/asset_files_common.py
@@ -0,0 +1,22 @@
+class AssetFileNotFoundException(Exception):
+ """Raised when a 3D asset file is not found"""
+
+ def __init__(self, message: str = "Asset file not found"):
+ self.message = message
+ super().__init__(self.message)
+
+
+class AssetFileSaveException(Exception):
+ """Raised when a 3D asset file cannot be saved"""
+
+ def __init__(self, message: str = "Asset file cannot be saved"):
+ self.message = message
+ super().__init__(self.message)
+
+
+class AssetFileDeleteException(Exception):
+ """Raised when a 3D asset file cannot be deleted"""
+
+ def __init__(self, message: str = "Asset file cannot be deleted"):
+ self.message = message
+ super().__init__(self.message)
diff --git a/invokeai/app/services/asset_files/asset_files_disk.py b/invokeai/app/services/asset_files/asset_files_disk.py
new file mode 100644
index 00000000000..7974846fb7f
--- /dev/null
+++ b/invokeai/app/services/asset_files/asset_files_disk.py
@@ -0,0 +1,62 @@
+from pathlib import Path
+
+from invokeai.app.services.asset_files.asset_files_base import AssetFilesServiceBase
+from invokeai.app.services.asset_files.asset_files_common import (
+ AssetFileDeleteException,
+ AssetFileNotFoundException,
+ AssetFileSaveException,
+)
+from invokeai.app.services.invoker import Invoker
+
+
+class DiskAssetFileStorage(AssetFilesServiceBase):
+ """Stores 3D asset files (Gaussian-splat .ply / .splat) on disk."""
+
+ def __init__(self, asset_files_folder: Path):
+ self._asset_files_folder = asset_files_folder
+ self._validate_storage_folder()
+
+ def start(self, invoker: Invoker) -> None:
+ self._invoker = invoker
+
+ def save(self, asset_name: str, data: bytes) -> None:
+ try:
+ self._validate_storage_folder()
+ path = self._resolve_path(asset_name)
+ path.write_bytes(data)
+ except Exception as e:
+ raise AssetFileSaveException from e
+
+ def get_path(self, asset_name: str) -> Path:
+ path = self._resolve_path(asset_name)
+ if not path.exists():
+ raise AssetFileNotFoundException
+ return path
+
+ def get_url(self, asset_name: str) -> str:
+ return self._invoker.services.urls.get_asset_url(asset_name)
+
+ def delete(self, asset_name: str) -> None:
+ try:
+ path = self._resolve_path(asset_name)
+ if not path.exists():
+ raise AssetFileNotFoundException
+ path.unlink()
+ except AssetFileNotFoundException as e:
+ raise AssetFileNotFoundException from e
+ except Exception as e:
+ raise AssetFileDeleteException from e
+
+ def _resolve_path(self, asset_name: str) -> Path:
+ """Resolves an asset name to a path, guarding against path traversal."""
+ # Only bare filenames are allowed - reject any path separators / traversal.
+ if not asset_name or asset_name != Path(asset_name).name:
+ raise ValueError(f"Invalid asset name: {asset_name}")
+ path = (self._asset_files_folder / asset_name).resolve()
+ if not path.is_relative_to(self._asset_files_folder.resolve()):
+ raise ValueError(f"Invalid asset name: {asset_name}")
+ return path
+
+ def _validate_storage_folder(self) -> None:
+ """Creates the storage folder if it does not exist."""
+ self._asset_files_folder.mkdir(parents=True, exist_ok=True)
diff --git a/invokeai/app/services/invocation_services.py b/invokeai/app/services/invocation_services.py
index 54f9d82b786..13cb4bd42b9 100644
--- a/invokeai/app/services/invocation_services.py
+++ b/invokeai/app/services/invocation_services.py
@@ -12,6 +12,7 @@
import torch
+ from invokeai.app.services.asset_files.asset_files_base import AssetFilesServiceBase
from invokeai.app.services.board_image_records.board_image_records_base import BoardImageRecordStorageBase
from invokeai.app.services.board_images.board_images_base import BoardImagesServiceABC
from invokeai.app.services.board_records.board_records_base import BoardRecordStorageBase
@@ -49,6 +50,7 @@ class InvocationServices:
def __init__(
self,
+ asset_files: "AssetFilesServiceBase",
board_images: "BoardImagesServiceABC",
board_image_records: "BoardImageRecordStorageBase",
boards: "BoardServiceABC",
@@ -82,6 +84,7 @@ def __init__(
users: "UserServiceBase",
image_moves: "ImageMoveService | None" = None,
):
+ self.asset_files = asset_files
self.board_images = board_images
self.board_image_records = board_image_records
self.boards = boards
diff --git a/invokeai/app/services/urls/urls_base.py b/invokeai/app/services/urls/urls_base.py
index a5602abb3b4..1621609cca3 100644
--- a/invokeai/app/services/urls/urls_base.py
+++ b/invokeai/app/services/urls/urls_base.py
@@ -23,3 +23,8 @@ def get_style_preset_image_url(self, style_preset_id: str) -> str:
def get_workflow_thumbnail_url(self, workflow_id: str) -> str:
"""Gets the URL for a workflow thumbnail"""
pass
+
+ @abstractmethod
+ def get_asset_url(self, asset_name: str) -> str:
+ """Gets the URL for a 3D asset file"""
+ pass
diff --git a/invokeai/app/services/urls/urls_default.py b/invokeai/app/services/urls/urls_default.py
index 2e4f36d9d51..3bd9d635a07 100644
--- a/invokeai/app/services/urls/urls_default.py
+++ b/invokeai/app/services/urls/urls_default.py
@@ -25,3 +25,10 @@ def get_style_preset_image_url(self, style_preset_id: str) -> str:
def get_workflow_thumbnail_url(self, workflow_id: str) -> str:
return f"{self._base_url}/workflows/i/{workflow_id}/thumbnail"
+
+ def get_asset_url(self, asset_name: str) -> str:
+ import os
+
+ asset_basename = os.path.basename(asset_name)
+ # This path is determined by the route in invokeai/app/api/routers/assets.py
+ return f"{self._base_url}/assets/i/{asset_basename}"
diff --git a/invokeai/backend/image_util/triposplat/LICENSE b/invokeai/backend/image_util/triposplat/LICENSE
new file mode 100644
index 00000000000..afd688304b5
--- /dev/null
+++ b/invokeai/backend/image_util/triposplat/LICENSE
@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2026 VAST
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.
diff --git a/invokeai/backend/image_util/triposplat/__init__.py b/invokeai/backend/image_util/triposplat/__init__.py
new file mode 100644
index 00000000000..9d6cf80f6ab
--- /dev/null
+++ b/invokeai/backend/image_util/triposplat/__init__.py
@@ -0,0 +1,13 @@
+"""Vendored TripoSplat inference code — single image to 3D Gaussian splat.
+
+Source: https://github.com/VAST-AI-Research/TripoSplat (MIT License, see LICENSE in this directory).
+
+The only change from upstream is that the two cross-module imports were made package-relative so the
+two files import correctly as a package:
+ - triposplat.py: `from model import (...)` -> `from .model import (...)`
+ - model.py: `from triposplat import _build_gaussians` -> `from .triposplat import _build_gaussians`
+
+`triposplat_model.py` is NOT vendored — it is InvokeAI's RawModel wrapper around the pipeline.
+
+Import `TripoSplatPipeline` from `.triposplat` lazily (it pulls in torch/torchvision).
+"""
diff --git a/invokeai/backend/image_util/triposplat/model.py b/invokeai/backend/image_util/triposplat/model.py
new file mode 100644
index 00000000000..071c56f7a34
--- /dev/null
+++ b/invokeai/backend/image_util/triposplat/model.py
@@ -0,0 +1,1726 @@
+from typing import Optional
+import math
+import re
+
+import numpy as np
+import safetensors.torch
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torchvision.ops import deform_conv2d
+
+
+# ---------------------------------------------------------------------------
+# DINOv3 ViT-H/16+
+# ---------------------------------------------------------------------------
+
+def _rotate_half(x: torch.Tensor) -> torch.Tensor:
+ x1, x2 = x.chunk(2, dim=-1)
+ return torch.cat((-x2, x1), dim=-1)
+
+
+class DinoV3PatchEmbed(nn.Module):
+ def __init__(self, patch_size=16, in_chans=3, embed_dim=1280):
+ super().__init__()
+ self.proj = nn.Conv2d(in_chans, embed_dim, kernel_size=patch_size, stride=patch_size, bias=True)
+
+ def forward(self, x):
+ return self.proj(x).flatten(2).transpose(1, 2)
+
+
+class DinoV3RotaryEmbedding2D(nn.Module):
+ def __init__(self, dim: int, base: float = 100.0):
+ super().__init__()
+ inv_freq = 1.0 / (base ** torch.arange(0, 1, 4.0 / dim, dtype=torch.float32))
+ self.register_buffer("inv_freq", inv_freq, persistent=False)
+
+ def forward(self, height: int, width: int, device: torch.device, dtype: torch.dtype):
+ coords_h = torch.arange(0.5, height, dtype=torch.float32, device=device) / height
+ coords_w = torch.arange(0.5, width, dtype=torch.float32, device=device) / width
+ coords = torch.stack(torch.meshgrid(coords_h, coords_w, indexing="ij"), dim=-1)
+ coords = (2.0 * coords - 1.0).flatten(0, 1)
+ angles = (2 * math.pi * coords[:, :, None] * self.inv_freq[None, None, :]).flatten(1, 2).tile(2)
+ cos = angles.cos().unsqueeze(0).unsqueeze(0)
+ sin = angles.sin().unsqueeze(0).unsqueeze(0)
+ return cos.to(dtype=dtype), sin.to(dtype=dtype)
+
+
+class DinoV3Attention(nn.Module):
+ def __init__(self, dim: int, num_heads: int, qkv_bias: tuple = (True, False, True)):
+ super().__init__()
+ self.num_heads = num_heads
+ self.head_dim = dim // num_heads
+ q_bias, k_bias, v_bias = qkv_bias
+ self.q_proj = nn.Linear(dim, dim, bias=q_bias)
+ self.k_proj = nn.Linear(dim, dim, bias=k_bias)
+ self.v_proj = nn.Linear(dim, dim, bias=v_bias)
+ self.o_proj = nn.Linear(dim, dim, bias=True)
+
+ def forward(self, x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor,
+ num_prefix_tokens: int = 0) -> torch.Tensor:
+ B, N, C = x.shape
+ q = self.q_proj(x).reshape(B, N, self.num_heads, self.head_dim).transpose(1, 2)
+ k = self.k_proj(x).reshape(B, N, self.num_heads, self.head_dim).transpose(1, 2)
+ v = self.v_proj(x).reshape(B, N, self.num_heads, self.head_dim).transpose(1, 2)
+ if num_prefix_tokens > 0:
+ q_pre, q_pat = q.split((num_prefix_tokens, N - num_prefix_tokens), dim=-2)
+ k_pre, k_pat = k.split((num_prefix_tokens, N - num_prefix_tokens), dim=-2)
+ q = torch.cat((q_pre, q_pat * cos + _rotate_half(q_pat) * sin), dim=-2)
+ k = torch.cat((k_pre, k_pat * cos + _rotate_half(k_pat) * sin), dim=-2)
+ else:
+ q = q * cos + _rotate_half(q) * sin
+ k = k * cos + _rotate_half(k) * sin
+ out = F.scaled_dot_product_attention(q, k, v)
+ return self.o_proj(out.transpose(1, 2).reshape(B, N, C))
+
+
+class DinoV3MLP(nn.Module):
+ def __init__(self, dim: int, hidden_dim: int, bias: bool = True):
+ super().__init__()
+ self.gate_proj = nn.Linear(dim, hidden_dim, bias=bias)
+ self.up_proj = nn.Linear(dim, hidden_dim, bias=bias)
+ self.down_proj = nn.Linear(hidden_dim, dim, bias=bias)
+
+ def forward(self, x: torch.Tensor) -> torch.Tensor:
+ return self.down_proj(F.silu(self.gate_proj(x)) * self.up_proj(x))
+
+
+class DinoV3Block(nn.Module):
+ def __init__(self, dim: int, num_heads: int, mlp_ratio: float = 4.0,
+ qkv_bias: tuple = (True, False, True), layerscale_init: float = 1.0,
+ mlp_bias: bool = True, eps: float = 1e-5):
+ super().__init__()
+ self.norm1 = nn.LayerNorm(dim, eps=eps)
+ self.attn = DinoV3Attention(dim, num_heads, qkv_bias=qkv_bias)
+ self.ls1 = nn.Parameter(torch.ones(dim) * layerscale_init)
+ self.norm2 = nn.LayerNorm(dim, eps=eps)
+ self.mlp = DinoV3MLP(dim, int(dim * mlp_ratio), bias=mlp_bias)
+ self.ls2 = nn.Parameter(torch.ones(dim) * layerscale_init)
+
+ def forward(self, x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor,
+ num_prefix_tokens: int = 0) -> torch.Tensor:
+ x = x + self.ls1 * self.attn(self.norm1(x), cos, sin, num_prefix_tokens=num_prefix_tokens)
+ x = x + self.ls2 * self.mlp(self.norm2(x))
+ return x
+
+
+class DinoV3ViT(nn.Module):
+ def __init__(self, hidden_size: int = 1280, num_heads: int = 20, num_layers: int = 32,
+ patch_size: int = 16, num_register_tokens: int = 4,
+ intermediate_size: int = 5120, layerscale_init: float = 1.0,
+ query_bias: bool = True, key_bias: bool = False, value_bias: bool = True,
+ mlp_bias: bool = True, rope_theta: float = 100.0, layer_norm_eps: float = 1e-5):
+ super().__init__()
+ self.patch_size = patch_size
+ self.num_register_tokens = num_register_tokens
+ self.patch_embed = DinoV3PatchEmbed(patch_size=patch_size, embed_dim=hidden_size)
+ self.cls_token = nn.Parameter(torch.zeros(1, 1, hidden_size))
+ self.register_tokens = nn.Parameter(torch.zeros(1, num_register_tokens, hidden_size))
+ self.rope = DinoV3RotaryEmbedding2D(dim=hidden_size // num_heads, base=rope_theta)
+ qkv_bias = (query_bias, key_bias, value_bias)
+ self.blocks = nn.ModuleList([
+ DinoV3Block(hidden_size, num_heads, mlp_ratio=intermediate_size / hidden_size,
+ qkv_bias=qkv_bias, layerscale_init=layerscale_init,
+ mlp_bias=mlp_bias, eps=layer_norm_eps)
+ for _ in range(num_layers)
+ ])
+ self.norm = nn.LayerNorm(hidden_size, eps=layer_norm_eps)
+
+ @property
+ def device(self) -> torch.device:
+ return self.cls_token.device
+
+ def forward(self, pixel_values: torch.Tensor) -> torch.Tensor:
+ B, _, H, W = pixel_values.shape
+ x = self.patch_embed(pixel_values)
+ hp, wp = H // self.patch_size, W // self.patch_size
+ cos, sin = self.rope(hp, wp, x.device, x.dtype)
+ x = torch.cat([self.cls_token.expand(B, -1, -1),
+ self.register_tokens.expand(B, -1, -1), x], dim=1)
+ num_prefix = 1 + self.num_register_tokens
+ for block in self.blocks:
+ x = block(x, cos, sin, num_prefix_tokens=num_prefix)
+ return self.norm(x)
+
+ def load_safetensors(self, path: str) -> None:
+ state_dict = safetensors.torch.load_file(path)
+ our_sd = self.state_dict()
+ loaded = {}
+ for hf_key in state_dict:
+ k = (hf_key
+ .replace("embeddings.patch_embeddings.", "patch_embed.proj.")
+ .replace("embeddings.cls_token", "cls_token")
+ .replace("embeddings.mask_token", "mask_token")
+ .replace("embeddings.register_tokens", "register_tokens"))
+ m = re.match(r"layer\.(\d+)\.(.+)", k)
+ if m:
+ rest = m.group(2)
+ for proj in ["q_proj", "k_proj", "v_proj", "o_proj"]:
+ rest = rest.replace(f"attention.{proj}", f"attn.{proj}")
+ rest = (rest.replace("layer_scale1.lambda1", "ls1")
+ .replace("layer_scale2.lambda1", "ls2"))
+ k = f"blocks.{m.group(1)}.{rest}"
+ if k in our_sd:
+ assert state_dict[hf_key].shape == our_sd[k].shape, \
+ f"Shape mismatch {k}: {state_dict[hf_key].shape} vs {our_sd[k].shape}"
+ loaded[k] = state_dict[hf_key]
+ check_sd = {k: v for k, v in our_sd.items() if k != "mask_token"}
+ missing = set(check_sd) - set(loaded)
+ unexpected = set(loaded) - set(check_sd)
+ if missing:
+ raise KeyError(f"[DINOv3] Missing keys: {missing}")
+ if unexpected:
+ raise KeyError(f"[DINOv3] Unexpected keys: {unexpected}")
+ self.load_state_dict(loaded, strict=True)
+
+
+# ---------------------------------------------------------------------------
+# Flux2 VAE Encoder
+# ---------------------------------------------------------------------------
+
+class Flux2ResnetBlock(nn.Module):
+ def __init__(self, in_channels, out_channels, use_shortcut=False):
+ super().__init__()
+ self.norm1 = nn.GroupNorm(32, in_channels, eps=1e-6)
+ self.conv1 = nn.Conv2d(in_channels, out_channels, 3, 1, 1)
+ self.norm2 = nn.GroupNorm(32, out_channels, eps=1e-6)
+ self.conv2 = nn.Conv2d(out_channels, out_channels, 3, 1, 1)
+ self.conv_shortcut = nn.Conv2d(in_channels, out_channels, 1, 1, 0) if use_shortcut else None
+
+ def forward(self, x):
+ h = F.silu(self.norm1(x))
+ h = F.silu(self.norm2(self.conv1(h)))
+ h = self.conv2(h)
+ return h + (self.conv_shortcut(x) if self.conv_shortcut is not None else x)
+
+
+class Flux2Downsampler(nn.Module):
+ def __init__(self, channels):
+ super().__init__()
+ self.conv = nn.Conv2d(channels, channels, 3, 2, 0)
+
+ def forward(self, x):
+ return self.conv(F.pad(x, (0, 1, 0, 1)))
+
+
+class Flux2Attention(nn.Module):
+ def __init__(self, channels):
+ super().__init__()
+ self.group_norm = nn.GroupNorm(32, channels, eps=1e-6)
+ self.to_q = nn.Linear(channels, channels)
+ self.to_k = nn.Linear(channels, channels)
+ self.to_v = nn.Linear(channels, channels)
+ self.to_out = nn.ModuleList([nn.Linear(channels, channels), nn.Identity()])
+
+ def forward(self, x):
+ B, C, H, W = x.shape
+ h = self.group_norm(x).reshape(B, C, H * W).transpose(1, 2)
+ q = self.to_q(h).reshape(B, -1, 1, C).permute(0, 2, 1, 3)
+ k = self.to_k(h).reshape(B, -1, 1, C).permute(0, 2, 1, 3)
+ v = self.to_v(h).reshape(B, -1, 1, C).permute(0, 2, 1, 3)
+ out = F.scaled_dot_product_attention(q, k, v)
+ out = self.to_out[0](out.permute(0, 2, 1, 3).reshape(B, -1, C))
+ return x + out.transpose(1, 2).reshape(B, C, H, W)
+
+
+class Flux2Encoder(nn.Module):
+ def __init__(self):
+ super().__init__()
+ self.conv_in = nn.Conv2d(3, 128, 3, 1, 1)
+ self.down_0_resnets = nn.ModuleList([Flux2ResnetBlock(128, 128), Flux2ResnetBlock(128, 128)])
+ self.down_0_sampler = Flux2Downsampler(128)
+ self.down_1_resnets = nn.ModuleList([Flux2ResnetBlock(128, 256, use_shortcut=True), Flux2ResnetBlock(256, 256)])
+ self.down_1_sampler = Flux2Downsampler(256)
+ self.down_2_resnets = nn.ModuleList([Flux2ResnetBlock(256, 512, use_shortcut=True), Flux2ResnetBlock(512, 512)])
+ self.down_2_sampler = Flux2Downsampler(512)
+ self.down_3_resnets = nn.ModuleList([Flux2ResnetBlock(512, 512), Flux2ResnetBlock(512, 512)])
+ self.mid_attn = Flux2Attention(512)
+ self.mid_resnets = nn.ModuleList([Flux2ResnetBlock(512, 512), Flux2ResnetBlock(512, 512)])
+ self.conv_norm_out = nn.GroupNorm(32, 512, eps=1e-6)
+ self.conv_out = nn.Conv2d(512, 64, 3, 1, 1)
+
+ def forward(self, x):
+ x = self.conv_in(x)
+ for r in self.down_0_resnets: x = r(x)
+ x = self.down_0_sampler(x)
+ for r in self.down_1_resnets: x = r(x)
+ x = self.down_1_sampler(x)
+ for r in self.down_2_resnets: x = r(x)
+ x = self.down_2_sampler(x)
+ for r in self.down_3_resnets: x = r(x)
+ x = self.mid_resnets[0](x)
+ x = self.mid_attn(x)
+ x = self.mid_resnets[1](x)
+ return self.conv_out(F.silu(self.conv_norm_out(x)))
+
+
+class Flux2VAEEncoder(nn.Module):
+ def __init__(self):
+ super().__init__()
+ self.encoder = Flux2Encoder()
+ self.quant_conv = nn.Conv2d(64, 64, 1, 1, 0)
+ self.bn = nn.BatchNorm1d(128, eps=1e-5, momentum=0.1, affine=False, track_running_stats=True)
+
+ def load_safetensors(self, path: str):
+ sd = safetensors.torch.load_file(path)
+ remapped = {}
+ for k, v in sd.items():
+ # Skip the decoder half of a full Flux2-VAE ckpt — we only need the encoder.
+ if k.startswith(("decoder.", "post_quant_conv.")):
+ continue
+ # Comfy / diffusers-style naming → our flattened naming.
+ m = re.match(r"encoder\.down_blocks\.(\d+)\.resnets\.(\d+)\.(.+)", k)
+ if m:
+ remapped[f"encoder.down_{m.group(1)}_resnets.{m.group(2)}.{m.group(3)}"] = v
+ continue
+ m = re.match(r"encoder\.down_blocks\.(\d+)\.downsamplers\.0\.(.+)", k)
+ if m:
+ remapped[f"encoder.down_{m.group(1)}_sampler.{m.group(2)}"] = v
+ continue
+ m = re.match(r"encoder\.mid_block\.resnets\.(\d+)\.(.+)", k)
+ if m:
+ remapped[f"encoder.mid_resnets.{m.group(1)}.{m.group(2)}"] = v
+ continue
+ m = re.match(r"encoder\.mid_block\.attentions\.0\.(.+)", k)
+ if m:
+ remapped[f"encoder.mid_attn.{m.group(1)}"] = v
+ continue
+ remapped[k] = v
+ missing, unexpected = self.load_state_dict(remapped, strict=False)
+ if missing:
+ raise KeyError(f"[VAE] Missing keys: {missing}")
+ if unexpected:
+ raise KeyError(f"[VAE] Unexpected keys: {unexpected}")
+
+ def encode(self, images, deterministic: bool = True, generator: torch.Generator = None):
+ moments = self.quant_conv(self.encoder(images))
+ mean, logvar = moments.chunk(2, dim=1)
+ if deterministic:
+ latents = mean
+ else:
+ noise = torch.randn(mean.shape, dtype=mean.dtype, device=mean.device, generator=generator)
+ latents = mean + torch.exp(0.5 * logvar) * noise
+ B, C, H, W = latents.shape
+ latents = latents.view(B, C, H // 2, 2, W // 2, 2).permute(0, 1, 3, 5, 2, 4)
+ latents = latents.reshape(B, C * 4, H // 2, W // 2)
+ bn_mean = self.bn.running_mean.view(1, -1, 1, 1).to(latents.device, latents.dtype)
+ bn_std = torch.sqrt(self.bn.running_var.view(1, -1, 1, 1) + self.bn.eps).to(latents.device, latents.dtype)
+ return ((latents - bn_mean) / bn_std).to(torch.float32).flatten(2).transpose(1, 2).contiguous()
+
+ return rgba
+
+
+# ---------------------------------------------------------------------------
+# BiRefNet background removal (Swin-L + ASPP-deformable decoder)
+# ---------------------------------------------------------------------------
+
+# -- timm-style helpers, inlined to avoid a timm dependency --------------------
+
+def _trunc_normal_(tensor: torch.Tensor, mean: float = 0.0, std: float = 1.0,
+ a: float = -2.0, b: float = 2.0) -> torch.Tensor:
+ # Initialization helper — only used at __init__ time, the released ckpt
+ # overwrites everything in load_safetensors so the exact distribution here
+ # is unimportant.
+ with torch.no_grad():
+ tensor.normal_(mean, std).clamp_(mean + a * std, mean + b * std)
+ return tensor
+
+
+# -- Swin Transformer (Swin-Large preset) -------------------------------------
+
+class _SwinMlp(nn.Module):
+ def __init__(self, in_features, hidden_features=None, out_features=None):
+ super().__init__()
+ hidden_features = hidden_features or in_features
+ out_features = out_features or in_features
+ self.fc1 = nn.Linear(in_features, hidden_features)
+ self.act = nn.GELU()
+ self.fc2 = nn.Linear(hidden_features, out_features)
+
+ def forward(self, x):
+ return self.fc2(self.act(self.fc1(x)))
+
+
+def _window_partition(x, window_size):
+ B, H, W, C = x.shape
+ x = x.view(B, H // window_size, window_size, W // window_size, window_size, C)
+ return x.permute(0, 1, 3, 2, 4, 5).contiguous().view(-1, window_size, window_size, C)
+
+
+def _window_reverse(windows, window_size, H, W):
+ B = int(windows.shape[0] / (H * W / window_size / window_size))
+ x = windows.view(B, H // window_size, W // window_size, window_size, window_size, -1)
+ return x.permute(0, 1, 3, 2, 4, 5).contiguous().view(B, H, W, -1)
+
+
+class _WindowAttention(nn.Module):
+ def __init__(self, dim, window_size, num_heads):
+ super().__init__()
+ self.dim = dim
+ self.window_size = window_size # (Wh, Ww)
+ self.num_heads = num_heads
+ head_dim = dim // num_heads
+ self.scale = head_dim ** -0.5
+
+ self.relative_position_bias_table = nn.Parameter(
+ torch.zeros((2 * window_size[0] - 1) * (2 * window_size[1] - 1), num_heads))
+ coords_h = torch.arange(window_size[0])
+ coords_w = torch.arange(window_size[1])
+ coords = torch.stack(torch.meshgrid([coords_h, coords_w], indexing="ij"))
+ coords_flatten = torch.flatten(coords, 1)
+ relative_coords = coords_flatten[:, :, None] - coords_flatten[:, None, :]
+ relative_coords = relative_coords.permute(1, 2, 0).contiguous()
+ relative_coords[:, :, 0] += window_size[0] - 1
+ relative_coords[:, :, 1] += window_size[1] - 1
+ relative_coords[:, :, 0] *= 2 * window_size[1] - 1
+ self.register_buffer("relative_position_index", relative_coords.sum(-1))
+
+ self.qkv = nn.Linear(dim, dim * 3, bias=True)
+ self.proj = nn.Linear(dim, dim)
+ _trunc_normal_(self.relative_position_bias_table, std=0.02)
+
+ def forward(self, x, mask=None):
+ B_, N, C = x.shape
+ qkv = self.qkv(x).reshape(B_, N, 3, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4)
+ q, k, v = qkv[0], qkv[1], qkv[2]
+ q = q * self.scale
+ attn = q @ k.transpose(-2, -1)
+ bias = self.relative_position_bias_table[self.relative_position_index.view(-1)].view(
+ self.window_size[0] * self.window_size[1],
+ self.window_size[0] * self.window_size[1], -1)
+ attn = attn + bias.permute(2, 0, 1).contiguous().unsqueeze(0)
+ if mask is not None:
+ nW = mask.shape[0]
+ attn = attn.view(B_ // nW, nW, self.num_heads, N, N) + mask.unsqueeze(1).unsqueeze(0)
+ attn = attn.view(-1, self.num_heads, N, N)
+ attn = attn.softmax(dim=-1)
+ x = (attn @ v).transpose(1, 2).reshape(B_, N, C)
+ return self.proj(x)
+
+
+class _SwinBlock(nn.Module):
+ def __init__(self, dim, num_heads, window_size, shift_size, mlp_ratio=4.0):
+ super().__init__()
+ self.dim = dim
+ self.window_size = window_size
+ self.shift_size = shift_size
+ self.norm1 = nn.LayerNorm(dim)
+ self.attn = _WindowAttention(dim, (window_size, window_size), num_heads)
+ self.norm2 = nn.LayerNorm(dim)
+ self.mlp = _SwinMlp(dim, int(dim * mlp_ratio))
+ self.H = None
+ self.W = None
+
+ def forward(self, x, mask_matrix):
+ B, L, C = x.shape
+ H, W = self.H, self.W
+ shortcut = x
+ x = self.norm1(x).view(B, H, W, C)
+ pad_r = (self.window_size - W % self.window_size) % self.window_size
+ pad_b = (self.window_size - H % self.window_size) % self.window_size
+ x = F.pad(x, (0, 0, 0, pad_r, 0, pad_b))
+ _, Hp, Wp, _ = x.shape
+ if self.shift_size > 0:
+ shifted_x = torch.roll(x, shifts=(-self.shift_size, -self.shift_size), dims=(1, 2))
+ attn_mask = mask_matrix
+ else:
+ shifted_x = x
+ attn_mask = None
+ x_windows = _window_partition(shifted_x, self.window_size).view(
+ -1, self.window_size * self.window_size, C)
+ attn_windows = self.attn(x_windows, mask=attn_mask).view(
+ -1, self.window_size, self.window_size, C)
+ shifted_x = _window_reverse(attn_windows, self.window_size, Hp, Wp)
+ if self.shift_size > 0:
+ x = torch.roll(shifted_x, shifts=(self.shift_size, self.shift_size), dims=(1, 2))
+ else:
+ x = shifted_x
+ if pad_r > 0 or pad_b > 0:
+ x = x[:, :H, :W, :].contiguous()
+ x = x.view(B, H * W, C)
+ x = shortcut + x
+ x = x + self.mlp(self.norm2(x))
+ return x
+
+
+class _PatchMerging(nn.Module):
+ def __init__(self, dim):
+ super().__init__()
+ self.dim = dim
+ self.reduction = nn.Linear(4 * dim, 2 * dim, bias=False)
+ self.norm = nn.LayerNorm(4 * dim)
+
+ def forward(self, x, H, W):
+ B, L, C = x.shape
+ x = x.view(B, H, W, C)
+ if H % 2 == 1 or W % 2 == 1:
+ x = F.pad(x, (0, 0, 0, W % 2, 0, H % 2))
+ x0 = x[:, 0::2, 0::2, :]
+ x1 = x[:, 1::2, 0::2, :]
+ x2 = x[:, 0::2, 1::2, :]
+ x3 = x[:, 1::2, 1::2, :]
+ x = torch.cat([x0, x1, x2, x3], -1).view(B, -1, 4 * C)
+ return self.reduction(self.norm(x))
+
+
+class _SwinBasicLayer(nn.Module):
+ def __init__(self, dim, depth, num_heads, window_size, mlp_ratio=4.0, downsample=True):
+ super().__init__()
+ self.window_size = window_size
+ self.shift_size = window_size // 2
+ self.depth = depth
+ self.blocks = nn.ModuleList([
+ _SwinBlock(dim=dim, num_heads=num_heads, window_size=window_size,
+ shift_size=0 if (i % 2 == 0) else window_size // 2,
+ mlp_ratio=mlp_ratio)
+ for i in range(depth)
+ ])
+ self.downsample = _PatchMerging(dim) if downsample else None
+
+ def forward(self, x, H, W):
+ Hp = int(math.ceil(H / self.window_size)) * self.window_size
+ Wp = int(math.ceil(W / self.window_size)) * self.window_size
+ img_mask = torch.zeros((1, Hp, Wp, 1), device=x.device)
+ h_slices = (slice(0, -self.window_size),
+ slice(-self.window_size, -self.shift_size),
+ slice(-self.shift_size, None))
+ w_slices = (slice(0, -self.window_size),
+ slice(-self.window_size, -self.shift_size),
+ slice(-self.shift_size, None))
+ cnt = 0
+ for h in h_slices:
+ for w in w_slices:
+ img_mask[:, h, w, :] = cnt
+ cnt += 1
+ mask_windows = _window_partition(img_mask, self.window_size).view(-1, self.window_size ** 2)
+ attn_mask = mask_windows.unsqueeze(1) - mask_windows.unsqueeze(2)
+ attn_mask = attn_mask.masked_fill(attn_mask != 0, float(-100.0)) \
+ .masked_fill(attn_mask == 0, float(0.0)).to(x.dtype)
+ for blk in self.blocks:
+ blk.H, blk.W = H, W
+ x = blk(x, attn_mask)
+ if self.downsample is not None:
+ x_down = self.downsample(x, H, W)
+ Wh, Ww = (H + 1) // 2, (W + 1) // 2
+ return x, H, W, x_down, Wh, Ww
+ return x, H, W, x, H, W
+
+
+class _SwinPatchEmbed(nn.Module):
+ def __init__(self, patch_size=4, in_channels=3, embed_dim=192):
+ super().__init__()
+ self.patch_size = (patch_size, patch_size)
+ self.proj = nn.Conv2d(in_channels, embed_dim, kernel_size=patch_size, stride=patch_size)
+ self.norm = nn.LayerNorm(embed_dim)
+ self.embed_dim = embed_dim
+
+ def forward(self, x):
+ _, _, H, W = x.shape
+ if W % self.patch_size[1] != 0:
+ x = F.pad(x, (0, self.patch_size[1] - W % self.patch_size[1]))
+ if H % self.patch_size[0] != 0:
+ x = F.pad(x, (0, 0, 0, self.patch_size[0] - H % self.patch_size[0]))
+ x = self.proj(x)
+ Wh, Ww = x.size(2), x.size(3)
+ x = x.flatten(2).transpose(1, 2)
+ x = self.norm(x)
+ return x.transpose(1, 2).view(-1, self.embed_dim, Wh, Ww)
+
+
+class _SwinLarge(nn.Module):
+ """Swin-Large backbone matching the BiRefNet HF release.
+
+ embed_dim=192, depths=[2,2,18,2], num_heads=[6,12,24,48], window_size=12.
+ """
+ def __init__(self):
+ super().__init__()
+ embed_dim = 192
+ depths = [2, 2, 18, 2]
+ num_heads = [6, 12, 24, 48]
+ window_size = 12
+ self.num_layers = len(depths)
+ self.embed_dim = embed_dim
+ self.patch_embed = _SwinPatchEmbed(patch_size=4, in_channels=3, embed_dim=embed_dim)
+ self.layers = nn.ModuleList([
+ _SwinBasicLayer(
+ dim=int(embed_dim * 2 ** i),
+ depth=depths[i],
+ num_heads=num_heads[i],
+ window_size=window_size,
+ downsample=(i < self.num_layers - 1),
+ ) for i in range(self.num_layers)
+ ])
+ num_features = [int(embed_dim * 2 ** i) for i in range(self.num_layers)]
+ self.num_features = num_features
+ for i in range(self.num_layers):
+ self.add_module(f"norm{i}", nn.LayerNorm(num_features[i]))
+
+ def forward(self, x):
+ x = self.patch_embed(x)
+ Wh, Ww = x.size(2), x.size(3)
+ x = x.flatten(2).transpose(1, 2)
+ outs = []
+ for i in range(self.num_layers):
+ x_out, H, W, x, Wh, Ww = self.layers[i](x, Wh, Ww)
+ norm_layer = getattr(self, f"norm{i}")
+ x_out = norm_layer(x_out)
+ out = x_out.view(-1, H, W, self.num_features[i]).permute(0, 3, 1, 2).contiguous()
+ outs.append(out)
+ return tuple(outs)
+
+
+# -- ASPP-Deformable -----------------------------------------------------------
+
+class _DeformableConv2d(nn.Module):
+ def __init__(self, in_channels, out_channels, kernel_size=3, stride=1, padding=1, bias=False):
+ super().__init__()
+ if isinstance(kernel_size, int):
+ kernel_size = (kernel_size, kernel_size)
+ self.stride = (stride, stride) if isinstance(stride, int) else stride
+ self.padding = padding
+ self.offset_conv = nn.Conv2d(in_channels, 2 * kernel_size[0] * kernel_size[1],
+ kernel_size=kernel_size, stride=stride, padding=padding, bias=True)
+ self.modulator_conv = nn.Conv2d(in_channels, 1 * kernel_size[0] * kernel_size[1],
+ kernel_size=kernel_size, stride=stride, padding=padding, bias=True)
+ self.regular_conv = nn.Conv2d(in_channels, out_channels,
+ kernel_size=kernel_size, stride=stride, padding=padding, bias=bias)
+
+ def forward(self, x):
+ offset = self.offset_conv(x)
+ modulator = 2.0 * torch.sigmoid(self.modulator_conv(x))
+ return deform_conv2d(
+ input=x, offset=offset,
+ weight=self.regular_conv.weight, bias=self.regular_conv.bias,
+ padding=self.padding, mask=modulator, stride=self.stride,
+ )
+
+
+class _ASPPModuleDeformable(nn.Module):
+ def __init__(self, in_channels, planes, kernel_size, padding):
+ super().__init__()
+ self.atrous_conv = _DeformableConv2d(in_channels, planes, kernel_size=kernel_size,
+ stride=1, padding=padding, bias=False)
+ self.bn = nn.BatchNorm2d(planes)
+ self.relu = nn.ReLU(inplace=True)
+
+ def forward(self, x):
+ return self.relu(self.bn(self.atrous_conv(x)))
+
+
+class _ASPPDeformable(nn.Module):
+ def __init__(self, in_channels, out_channels=None, parallel_block_sizes=(1, 3, 7)):
+ super().__init__()
+ if out_channels is None:
+ out_channels = in_channels
+ inter = 256
+ self.aspp1 = _ASPPModuleDeformable(in_channels, inter, 1, padding=0)
+ self.aspp_deforms = nn.ModuleList([
+ _ASPPModuleDeformable(in_channels, inter, k, padding=k // 2)
+ for k in parallel_block_sizes
+ ])
+ self.global_avg_pool = nn.Sequential(
+ nn.AdaptiveAvgPool2d((1, 1)),
+ nn.Conv2d(in_channels, inter, 1, stride=1, bias=False),
+ nn.BatchNorm2d(inter),
+ nn.ReLU(inplace=True),
+ )
+ self.conv1 = nn.Conv2d(inter * (2 + len(self.aspp_deforms)), out_channels, 1, bias=False)
+ self.bn1 = nn.BatchNorm2d(out_channels)
+ self.relu = nn.ReLU(inplace=True)
+
+ def forward(self, x):
+ x1 = self.aspp1(x)
+ x_aspp_deforms = [m(x) for m in self.aspp_deforms]
+ x5 = self.global_avg_pool(x)
+ x5 = F.interpolate(x5, size=x1.size()[2:], mode='bilinear', align_corners=True)
+ y = torch.cat((x1, *x_aspp_deforms, x5), dim=1)
+ return self.relu(self.bn1(self.conv1(y)))
+
+
+# -- Decoder blocks ------------------------------------------------------------
+
+class _BasicDecBlk(nn.Module):
+ def __init__(self, in_channels, out_channels, inter_channels=64):
+ super().__init__()
+ self.conv_in = nn.Conv2d(in_channels, inter_channels, 3, 1, padding=1)
+ self.bn_in = nn.BatchNorm2d(inter_channels)
+ self.relu_in = nn.ReLU(inplace=True)
+ self.dec_att = _ASPPDeformable(in_channels=inter_channels)
+ self.conv_out = nn.Conv2d(inter_channels, out_channels, 3, 1, padding=1)
+ self.bn_out = nn.BatchNorm2d(out_channels)
+
+ def forward(self, x):
+ x = self.relu_in(self.bn_in(self.conv_in(x)))
+ x = self.dec_att(x)
+ x = self.bn_out(self.conv_out(x))
+ return x
+
+
+class _BasicLatBlk(nn.Module):
+ def __init__(self, in_channels, out_channels):
+ super().__init__()
+ self.conv = nn.Conv2d(in_channels, out_channels, 1, 1, 0)
+
+ def forward(self, x):
+ return self.conv(x)
+
+
+class _SimpleConvs(nn.Module):
+ def __init__(self, in_channels, out_channels, inter_channels=64):
+ super().__init__()
+ self.conv1 = nn.Conv2d(in_channels, inter_channels, 3, 1, 1)
+ self.conv_out = nn.Conv2d(inter_channels, out_channels, 3, 1, 1)
+
+ def forward(self, x):
+ return self.conv_out(self.conv1(x))
+
+
+# -- Image → patch-stack helper -----------------------------------------------
+
+def _image2patches(image, patch_ref):
+ """`einops` rearrange 'b c (hg h) (wg w) -> b (c hg wg) h w' replacement.
+
+ Splits `image` into hg×wg non-overlapping patches and stacks them along
+ the channel axis. `hg`/`wg` are inferred from image and patch_ref sizes.
+ """
+ b, c, h_full, w_full = image.shape
+ hg, wg = h_full // patch_ref.shape[-2], w_full // patch_ref.shape[-1]
+ h, w = h_full // hg, w_full // wg
+ # (b, c, hg*h, wg*w) -> (b, c, hg, h, wg, w) -> (b, c, hg, wg, h, w) -> (b, c*hg*wg, h, w)
+ return image.view(b, c, hg, h, wg, w).permute(0, 1, 2, 4, 3, 5).reshape(b, c * hg * wg, h, w)
+
+
+# -- Decoder + top-level BiRefNet ---------------------------------------------
+
+class _BiRefNetDecoder(nn.Module):
+ def __init__(self, channels=(3072, 1536, 768, 384)):
+ super().__init__()
+ c = channels # high-to-low resolution channel counts
+ # input-modulator blocks (one per resolution; channels are
+ # `3 * patch_grid**2`, see _image2patches docstring).
+ self.ipt_blk5 = _SimpleConvs(2 ** 10 * 3, c[0] // 8, inter_channels=64)
+ self.ipt_blk4 = _SimpleConvs(2 ** 8 * 3, c[0] // 8, inter_channels=64)
+ self.ipt_blk3 = _SimpleConvs(2 ** 6 * 3, c[1] // 8, inter_channels=64)
+ self.ipt_blk2 = _SimpleConvs(2 ** 4 * 3, c[2] // 8, inter_channels=64)
+ self.ipt_blk1 = _SimpleConvs(2 ** 0 * 3, c[3] // 8, inter_channels=64)
+
+ self.decoder_block4 = _BasicDecBlk(c[0] + c[0] // 8, c[1])
+ self.decoder_block3 = _BasicDecBlk(c[1] + c[0] // 8, c[2])
+ self.decoder_block2 = _BasicDecBlk(c[2] + c[1] // 8, c[3])
+ self.decoder_block1 = _BasicDecBlk(c[3] + c[2] // 8, c[3] // 2)
+ self.conv_out1 = nn.Sequential(nn.Conv2d(c[3] // 2 + c[3] // 8, 1, 1, 1, 0))
+
+ self.lateral_block4 = _BasicLatBlk(c[1], c[1])
+ self.lateral_block3 = _BasicLatBlk(c[2], c[2])
+ self.lateral_block2 = _BasicLatBlk(c[3], c[3])
+
+ # multi-scale supervision heads (training only — kept for state_dict
+ # parity with the released checkpoint; not consumed at inference).
+ self.conv_ms_spvn_4 = nn.Conv2d(c[1], 1, 1, 1, 0)
+ self.conv_ms_spvn_3 = nn.Conv2d(c[2], 1, 1, 1, 0)
+ self.conv_ms_spvn_2 = nn.Conv2d(c[3], 1, 1, 1, 0)
+
+ # gradient-decoder-triggering (gdt) attention: used at inference to
+ # gate p4/p3/p2.
+ _N = 16
+ def _gdt_branch(in_c):
+ return nn.Sequential(nn.Conv2d(in_c, _N, 3, 1, 1), nn.BatchNorm2d(_N), nn.ReLU(inplace=True))
+ self.gdt_convs_4 = _gdt_branch(c[1])
+ self.gdt_convs_3 = _gdt_branch(c[2])
+ self.gdt_convs_2 = _gdt_branch(c[3])
+
+ def _head_1x1():
+ return nn.Sequential(nn.Conv2d(_N, 1, 1, 1, 0))
+ # multi-scale supervision heads on the gdt branch (training only)
+ self.gdt_convs_pred_4 = _head_1x1()
+ self.gdt_convs_pred_3 = _head_1x1()
+ self.gdt_convs_pred_2 = _head_1x1()
+ # attention heads
+ self.gdt_convs_attn_4 = _head_1x1()
+ self.gdt_convs_attn_3 = _head_1x1()
+ self.gdt_convs_attn_2 = _head_1x1()
+
+ def forward(self, x, x1, x2, x3, x4):
+ x4 = torch.cat((x4, self.ipt_blk5(_image2patches(x, x4))), 1)
+ p4 = self.decoder_block4(x4)
+ p4 = p4 * self.gdt_convs_attn_4(self.gdt_convs_4(p4)).sigmoid()
+ _p4 = F.interpolate(p4, size=x3.shape[2:], mode='bilinear', align_corners=True)
+ _p3 = _p4 + self.lateral_block4(x3)
+
+ _p3 = torch.cat((_p3, self.ipt_blk4(_image2patches(x, _p3))), 1)
+ p3 = self.decoder_block3(_p3)
+ p3 = p3 * self.gdt_convs_attn_3(self.gdt_convs_3(p3)).sigmoid()
+ _p3 = F.interpolate(p3, size=x2.shape[2:], mode='bilinear', align_corners=True)
+ _p2 = _p3 + self.lateral_block3(x2)
+
+ _p2 = torch.cat((_p2, self.ipt_blk3(_image2patches(x, _p2))), 1)
+ p2 = self.decoder_block2(_p2)
+ p2 = p2 * self.gdt_convs_attn_2(self.gdt_convs_2(p2)).sigmoid()
+ _p2 = F.interpolate(p2, size=x1.shape[2:], mode='bilinear', align_corners=True)
+ _p1 = _p2 + self.lateral_block2(x1)
+
+ _p1 = torch.cat((_p1, self.ipt_blk2(_image2patches(x, _p1))), 1)
+ _p1 = self.decoder_block1(_p1)
+ _p1 = F.interpolate(_p1, size=x.shape[2:], mode='bilinear', align_corners=True)
+
+ _p1 = torch.cat((_p1, self.ipt_blk1(_image2patches(x, _p1))), 1)
+ return self.conv_out1(_p1)
+
+
+class BiRefNet(nn.Module):
+ """BiRefNet (ZhengPeng7/BiRefNet) with Swin-L backbone, multi-scale input
+ concatenation, ASPP-deformable squeeze block, and the 4-level
+ input-modulating decoder used in the v1 release.
+
+ `forward(x)` returns a single 1-channel alpha map in `[0, 1]` (post-sigmoid).
+ `remove_background(pil_img)` is the PIL helper used by the pipeline —
+ accepts a PIL RGB image and returns an RGBA copy with the predicted matte
+ in the alpha channel.
+ """
+
+ INPUT_SIZE = (1024, 1024)
+ # backbone channel counts post mul_scl_ipt='cat' (doubled from raw Swin-L)
+ _CHANNELS = (3072, 1536, 768, 384)
+ # ImageNet normalization used by the BiRefNet recipe
+ _NORM_MEAN = (0.485, 0.456, 0.406)
+ _NORM_STD = (0.229, 0.224, 0.225)
+
+ def __init__(self):
+ super().__init__()
+ self.bb = _SwinLarge()
+ cxt = list(self._CHANNELS[1:][::-1][-3:]) # = [384, 768, 1536]
+ self.squeeze_module = nn.Sequential(
+ _BasicDecBlk(self._CHANNELS[0] + sum(cxt), self._CHANNELS[0])
+ )
+ self.decoder = _BiRefNetDecoder(channels=self._CHANNELS)
+
+ @property
+ def device(self):
+ return next(self.parameters()).device
+
+ @property
+ def dtype(self):
+ return next(self.parameters()).dtype
+
+ def _forward_enc(self, x):
+ x1, x2, x3, x4 = self.bb(x)
+ # mul_scl_ipt='cat': re-run backbone at half resolution, concat features
+ B, C, H, W = x.shape
+ x1_, x2_, x3_, x4_ = self.bb(F.interpolate(x, size=(H // 2, W // 2),
+ mode='bilinear', align_corners=True))
+ x1 = torch.cat([x1, F.interpolate(x1_, size=x1.shape[2:], mode='bilinear', align_corners=True)], 1)
+ x2 = torch.cat([x2, F.interpolate(x2_, size=x2.shape[2:], mode='bilinear', align_corners=True)], 1)
+ x3 = torch.cat([x3, F.interpolate(x3_, size=x3.shape[2:], mode='bilinear', align_corners=True)], 1)
+ x4 = torch.cat([x4, F.interpolate(x4_, size=x4.shape[2:], mode='bilinear', align_corners=True)], 1)
+ # cxt: upsample x1/x2/x3 to x4 spatial and concat for the squeeze input
+ x4 = torch.cat([
+ F.interpolate(x1, size=x4.shape[2:], mode='bilinear', align_corners=True),
+ F.interpolate(x2, size=x4.shape[2:], mode='bilinear', align_corners=True),
+ F.interpolate(x3, size=x4.shape[2:], mode='bilinear', align_corners=True),
+ x4,
+ ], 1)
+ return x1, x2, x3, x4
+
+ def forward(self, x):
+ x1, x2, x3, x4 = self._forward_enc(x)
+ x4 = self.squeeze_module(x4)
+ logits = self.decoder(x, x1, x2, x3, x4)
+ return torch.sigmoid(logits)
+
+ def load_safetensors(self, path: str) -> None:
+ sd = safetensors.torch.load_file(path)
+ # The decoder's gdt_convs_pred_* / conv_ms_spvn_* heads are training-only
+ # but are kept as submodules for state_dict parity. strict=True works.
+ missing, unexpected = self.load_state_dict(sd, strict=False)
+ if unexpected:
+ raise KeyError(f"[birefnet] unexpected keys (e.g. {unexpected[:3]})")
+ if missing:
+ raise KeyError(f"[birefnet] missing keys (e.g. {missing[:3]})")
+
+ @torch.no_grad()
+ def remove_background(self, image) -> "Image.Image":
+ from PIL import Image
+ if image.mode != "RGB":
+ image = image.convert("RGB")
+ W, H = image.size
+ arr = np.array(image, dtype=np.float32) / 255.0
+ t = torch.from_numpy(arr).permute(2, 0, 1).unsqueeze(0)
+ t = F.interpolate(t, size=self.INPUT_SIZE, mode='bilinear', align_corners=True)
+ mean = torch.tensor(self._NORM_MEAN).view(1, 3, 1, 1)
+ std = torch.tensor(self._NORM_STD).view(1, 3, 1, 1)
+ t = ((t - mean) / std).to(device=self.device, dtype=self.dtype)
+ alpha = self.forward(t)
+ alpha = F.interpolate(alpha.float(), size=(H, W), mode='bilinear', align_corners=True)[0, 0]
+ a = (alpha.clamp(0, 1) * 255).to(torch.uint8).cpu().numpy()
+ rgba = image.copy()
+ # No explicit mode: Pillow infers "L" from the 2D uint8 array (mode= is deprecated, removed in Pillow 13)
+ rgba.putalpha(Image.fromarray(a))
+ return rgba
+
+
+# ---------------------------------------------------------------------------
+# Shared transformer helpers
+# ---------------------------------------------------------------------------
+
+class LayerNorm32(nn.LayerNorm):
+ def forward(self, x):
+ origin_dtype = x.dtype
+ return F.layer_norm(
+ x.float(),
+ self.normalized_shape,
+ self.weight.float() if self.weight is not None else None,
+ self.bias.float() if self.bias is not None else None,
+ self.eps,
+ ).to(origin_dtype)
+
+
+class MultiHeadRMSNorm(nn.Module):
+ def __init__(self, dim, heads):
+ super().__init__()
+ self.scale = dim ** 0.5
+ self.gamma = nn.Parameter(torch.ones(heads, dim))
+
+ def forward(self, x):
+ origin_dtype = x.dtype
+ return (F.normalize(x.float(), dim=-1) * self.gamma.float() * self.scale).to(origin_dtype)
+
+
+def apply_rotary_emb(hidden_states, freqs):
+ x_rotated = torch.view_as_complex(hidden_states.float().reshape(*hidden_states.shape[:-1], -1, 2))
+ x_rotated = x_rotated * freqs
+ x_out = torch.view_as_real(x_rotated).reshape(*x_rotated.shape[:-1], -1)
+ return x_out.type_as(hidden_states)
+
+
+def clamp_mul(x, f):
+ f_t = f.tanh()
+ return x * f_t + x.detach() * (f - f_t)
+
+
+def scaled_dot_product_attention(qkv=None, q=None, k=None, v=None, kv=None):
+ if qkv is not None:
+ q, k, v = qkv.unbind(dim=2)
+ elif kv is not None:
+ k, v = kv.unbind(dim=2)
+ q, k, v = q.permute(0, 2, 1, 3), k.permute(0, 2, 1, 3), v.permute(0, 2, 1, 3)
+ return F.scaled_dot_product_attention(q, k, v).permute(0, 2, 1, 3)
+
+
+# ---------------------------------------------------------------------------
+# Positional embeddings
+# ---------------------------------------------------------------------------
+
+class RePo3DRotaryEmbedding(nn.Module):
+ def __init__(self, model_channels, num_heads, head_dim, repo_hidden_ratio=0.125, max_freq=16.0):
+ super().__init__()
+ self.num_heads = num_heads
+ self.head_dim = head_dim
+ repo_hidden_size = int(model_channels * repo_hidden_ratio)
+ self.norm = LayerNorm32(model_channels)
+ self.gate_map = nn.Linear(model_channels, repo_hidden_size, bias=False)
+ self.content_map = nn.Linear(model_channels, repo_hidden_size, bias=False)
+ self.act = nn.SiLU()
+ self.final_map = nn.Linear(repo_hidden_size, 3 * num_heads, bias=False)
+ self.dim_0 = 2 * (head_dim // 6)
+ self.dim_1 = 2 * (head_dim // 6)
+ self.dim_2 = head_dim - self.dim_0 - self.dim_1
+ dims = [self.dim_0, self.dim_1, self.dim_2]
+ freqs_list = []
+ for d in dims:
+ freq_dim = d // 2
+ freqs_list.append(torch.linspace(1.0, float(max_freq), steps=freq_dim, dtype=torch.float32))
+ self.freqs_0 = nn.Parameter(freqs_list[0])
+ self.freqs_1 = nn.Parameter(freqs_list[1])
+ self.freqs_2 = nn.Parameter(freqs_list[2])
+
+ def forward(self, hidden_states):
+ h = self.norm(hidden_states)
+ feat = self.act(self.gate_map(h)) * self.content_map(h)
+ out = self.final_map(feat)
+ B, L, _ = out.shape
+ delta_pos = out.reshape(B, L, self.num_heads, 3)
+ ang_0 = clamp_mul(delta_pos[..., 0].unsqueeze(-1), self.freqs_0) * torch.pi
+ ang_1 = clamp_mul(delta_pos[..., 1].unsqueeze(-1), self.freqs_1) * torch.pi
+ ang_2 = clamp_mul(delta_pos[..., 2].unsqueeze(-1), self.freqs_2) * torch.pi
+ ang = torch.cat([ang_0, ang_1, ang_2], dim=-1).float() # fp32 needed for torch.polar → complex64
+ return torch.polar(torch.ones_like(ang), ang).type(torch.complex64)
+
+
+class PcdAbsolutePositionEmbedder(nn.Module):
+ def __init__(self, channels: int, in_channels: int = 3, max_res: int = 16):
+ super().__init__()
+ self.channels = channels
+ self.in_channels = in_channels
+ self.max_res = max_res
+ self.freq_dim = channels // in_channels // 2
+
+ def _freqs(self, device):
+ freqs_2exp = torch.arange(self.max_res, dtype=torch.float32, device=device)
+ res_dim = max(0, self.freq_dim - self.max_res)
+ freqs_res = (torch.arange(res_dim, dtype=torch.float32, device=device) / max(res_dim, 1) * self.max_res
+ if res_dim > 0 else torch.empty(0, device=device))
+ freqs = torch.cat([freqs_2exp, freqs_res], dim=0)[:self.freq_dim]
+ return torch.pow(2.0, freqs)
+
+ def forward(self, x: torch.Tensor) -> torch.Tensor:
+ orig_dtype = x.dtype
+ x = x.float()
+ *dims, D = x.shape
+ out = torch.outer(x.reshape(-1), self._freqs(x.device)) * 2 * torch.pi
+ out = torch.cat([out.sin(), out.cos()], dim=-1).reshape(*dims, -1)
+ if out.shape[-1] < self.channels:
+ out = torch.cat([out, torch.zeros(*dims, self.channels - out.shape[-1],
+ device=out.device, dtype=out.dtype)], dim=-1)
+ return out.to(orig_dtype)
+
+
+class PcdAbsolutePositionEmbedderV2(nn.Module):
+ def __init__(self, channels: int, in_channels: int = 3, max_res: int = 10):
+ super().__init__()
+ self.channels = channels
+ self.in_channels = in_channels
+ self.max_res = max_res
+ self.freq_dim = channels // in_channels // 2
+
+ def _freqs(self, device):
+ logs = torch.linspace(0.0, float(self.max_res), steps=self.freq_dim, dtype=torch.float32, device=device)
+ return torch.pow(2.0, logs)
+
+ def forward(self, x: torch.Tensor) -> torch.Tensor:
+ orig_dtype = x.dtype
+ x = x.float()
+ N, D = x.shape
+ ang = x.unsqueeze(-1) * self._freqs(x.device) * torch.pi
+ embed = torch.cat([torch.sin(ang), torch.cos(ang)], dim=-1).reshape(N, -1)
+ if embed.shape[1] < self.channels:
+ embed = torch.cat([embed, torch.zeros(N, self.channels - embed.shape[1],
+ device=embed.device, dtype=embed.dtype)], dim=-1)
+ return embed.to(orig_dtype)
+
+
+# ---------------------------------------------------------------------------
+# Transformer building blocks
+# ---------------------------------------------------------------------------
+
+class FeedForwardNet(nn.Module):
+ def __init__(self, channels, mlp_ratio=4.0, channels_out=None):
+ super().__init__()
+ self.mlp = nn.Sequential(
+ nn.Linear(channels, int(channels * mlp_ratio)),
+ nn.GELU(approximate="tanh"),
+ nn.Linear(int(channels * mlp_ratio), channels if channels_out is None else channels_out),
+ )
+
+ def forward(self, x):
+ return self.mlp(x)
+
+
+class MLP(nn.Module):
+ def __init__(self, channels: int, inner_channels: int, channels_out: Optional[int] = None,
+ mlp_layer_num: int = 2):
+ super().__init__()
+ layers = []
+ for i in range(mlp_layer_num - 1):
+ layers.append(nn.Linear(channels if i == 0 else inner_channels, inner_channels))
+ layers.append(nn.GELU(approximate="tanh"))
+ layers.append(nn.Linear(inner_channels, channels if channels_out is None else channels_out))
+ self.mlp = nn.Sequential(*layers)
+
+ def forward(self, x: torch.Tensor) -> torch.Tensor:
+ return self.mlp(x)
+
+
+class RopeMultiHeadAttention(nn.Module):
+ def __init__(self, channels, num_heads, ctx_channels=None, type="self",
+ attn_mode="full", qkv_bias=True, qk_rms_norm=False, use_rope=False):
+ super().__init__()
+ self.channels = channels
+ self.num_heads = num_heads
+ self.head_dim = channels // num_heads
+ self.ctx_channels = ctx_channels if ctx_channels is not None else channels
+ self._type = type
+ self.qk_rms_norm = qk_rms_norm
+ self.use_rope = use_rope
+ if self._type == "self":
+ self.qkv = nn.Linear(channels, channels * 3, bias=qkv_bias)
+ else:
+ self.q = nn.Linear(channels, channels, bias=qkv_bias)
+ self.kv = nn.Linear(self.ctx_channels, channels * 2, bias=qkv_bias)
+ if self.qk_rms_norm:
+ self.q_norm = MultiHeadRMSNorm(self.head_dim, num_heads)
+ self.k_norm = MultiHeadRMSNorm(self.head_dim, num_heads)
+ self.out = nn.Linear(channels, channels)
+
+ def forward(self, x, context=None, rope_emb=None):
+ B, L, C = x.shape
+ if self._type == "self":
+ qkv = self.qkv(x).reshape(B, L, 3, self.num_heads, self.head_dim)
+ q, k, v = qkv.unbind(2)
+ if self.use_rope:
+ q = apply_rotary_emb(q, rope_emb)
+ k = apply_rotary_emb(k, rope_emb)
+ else:
+ q = self.q(x).reshape(B, L, self.num_heads, self.head_dim)
+ if context is None:
+ raise ValueError("Context must be provided for cross attention")
+ kv = self.kv(context).reshape(B, context.shape[1], 2, self.num_heads, self.head_dim)
+ k, v = kv.unbind(2)
+ if self.qk_rms_norm:
+ q = self.q_norm(q)
+ k = self.k_norm(k)
+ h = scaled_dot_product_attention(q=q, k=k, v=v)
+ return self.out(h.reshape(B, L, C))
+
+
+class MultiHeadAttention(nn.Module):
+ def __init__(self, channels, num_heads, ctx_channels=None, type="self",
+ attn_mode="full", qkv_bias=True, qk_rms_norm=False):
+ super().__init__()
+ assert channels % num_heads == 0
+ assert type in ["self", "cross"]
+ assert attn_mode == "full"
+ self.channels = channels
+ self.head_dim = channels // num_heads
+ self.ctx_channels = ctx_channels if ctx_channels is not None else channels
+ self.num_heads = num_heads
+ self._type = type
+ self.qk_rms_norm = qk_rms_norm
+ if self._type == "self":
+ self.to_qkv = nn.Linear(channels, channels * 3, bias=qkv_bias)
+ else:
+ self.to_q = nn.Linear(channels, channels, bias=qkv_bias)
+ self.to_kv = nn.Linear(self.ctx_channels, channels * 2, bias=qkv_bias)
+ if self.qk_rms_norm:
+ self.q_rms_norm = MultiHeadRMSNorm(self.head_dim, num_heads)
+ self.k_rms_norm = MultiHeadRMSNorm(self.head_dim, num_heads)
+ self.to_out = nn.Linear(channels, channels)
+
+ def forward(self, x, context=None):
+ B, L, C = x.shape
+ if self._type == "self":
+ qkv = self.to_qkv(x).reshape(B, L, 3, self.num_heads, -1)
+ if self.qk_rms_norm:
+ q, k, v = qkv.unbind(dim=2)
+ q = self.q_rms_norm(q)
+ k = self.k_rms_norm(k)
+ qkv = torch.stack([q, k, v], dim=2)
+ h = scaled_dot_product_attention(qkv=qkv)
+ else:
+ Lkv = context.shape[1]
+ q = self.to_q(x).reshape(B, L, self.num_heads, -1)
+ kv = self.to_kv(context).reshape(B, Lkv, 2, self.num_heads, -1)
+ if self.qk_rms_norm:
+ q = self.q_rms_norm(q)
+ k, v = kv.unbind(dim=2)
+ k = self.k_rms_norm(k)
+ h = scaled_dot_product_attention(q=q, k=k, v=v)
+ else:
+ h = scaled_dot_product_attention(q=q, kv=kv)
+ return self.to_out(h.reshape(B, L, -1))
+
+
+class UnifiedTransformerBlock(nn.Module):
+ def __init__(self, channels, num_heads, mlp_ratio=4.0, attn_mode="full",
+ use_checkpoint=False, use_rope=False, qk_rms_norm=False, qkv_bias=True,
+ modulation=True, share_mod=False, use_shift_table=False):
+ super().__init__()
+ self.modulation = modulation
+ self.share_mod = share_mod
+ self.norm1 = LayerNorm32(channels, elementwise_affine=not modulation, eps=1e-6)
+ self.norm2 = LayerNorm32(channels, elementwise_affine=not modulation, eps=1e-6)
+ self.attn = RopeMultiHeadAttention(channels, num_heads=num_heads, type="self",
+ attn_mode=attn_mode, qkv_bias=qkv_bias,
+ use_rope=use_rope, qk_rms_norm=qk_rms_norm)
+ self.mlp = FeedForwardNet(channels, mlp_ratio=mlp_ratio)
+ if modulation:
+ if not share_mod:
+ self.adaLN_modulation = nn.Sequential(nn.SiLU(), nn.Linear(channels, 6 * channels, bias=True))
+ self.shift_table = nn.Parameter(torch.randn(1, 6 * channels) / channels ** 0.5) if use_shift_table else None
+
+ def forward(self, x, mod=None, rotary_emb=None):
+ if self.modulation:
+ if not self.share_mod:
+ mod = self.adaLN_modulation(mod)
+ if hasattr(self, 'shift_table') and self.shift_table is not None:
+ mod = mod + self.shift_table.type(mod.dtype)
+ shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = mod.chunk(6, dim=1)
+ h = self.norm1(x)
+ h = h * (1 + scale_msa.unsqueeze(1)) + shift_msa.unsqueeze(1)
+ h = self.attn(h, rope_emb=rotary_emb)
+ x = x + h * gate_msa.unsqueeze(1)
+ h = self.norm2(x)
+ h = h * (1 + scale_mlp.unsqueeze(1)) + shift_mlp.unsqueeze(1)
+ x = x + self.mlp(h) * gate_mlp.unsqueeze(1)
+ else:
+ x = x + self.attn(self.norm1(x), rope_emb=rotary_emb)
+ x = x + self.mlp(self.norm2(x))
+ return x
+
+
+# ---------------------------------------------------------------------------
+# Quasi-random sampling utilities
+# ---------------------------------------------------------------------------
+
+PRIMES = [2, 3, 5, 7, 11, 13, 17, 19, 23, 29, 31, 37, 41, 43, 47, 53]
+
+
+def radical_inverse(base, n):
+ val = 0
+ inv_base = 1.0 / base
+ inv_base_n = inv_base
+ while n > 0:
+ digit = n % base
+ val += digit * inv_base_n
+ n //= base
+ inv_base_n *= inv_base
+ return val
+
+
+def halton_sequence(dim, n):
+ return [radical_inverse(PRIMES[dim], n) for dim in range(dim)]
+
+
+def hammersley_sequence(dim, n, num_samples):
+ return [n / num_samples] + halton_sequence(dim - 1, n)
+
+
+@torch.no_grad()
+def sample_probs(probs, counts, algo="systematic"):
+ batch_shape = counts.shape
+ B = counts.numel()
+ P = probs.size(-1)
+ device = probs.device
+ probs = probs.view(B, P)
+ counts = counts.view(B)
+
+ probs = probs.to(torch.float32).clamp_min_(0)
+ row_sums = probs.sum(1, keepdim=True)
+ zero_mask = row_sums.eq(0)
+ probs = probs / row_sums.clamp_min_(1)
+ if zero_mask.any():
+ probs = probs.clone()
+ probs[zero_mask.expand_as(probs)] = 1.0 / P
+
+ counts = counts.to(device=device, dtype=torch.long)
+ out = torch.zeros(B, P, dtype=torch.long, device=device)
+ cdf = probs.cumsum(dim=1).clamp(max=1.0 - 1e-12)
+ unique_n, inv = counts.unique(sorted=False, return_inverse=True)
+ for i, n in enumerate(unique_n.tolist()):
+ if n == 0:
+ continue
+ rows = (inv == i).nonzero(as_tuple=False).squeeze(1)
+ r = rows.numel()
+ U0 = torch.rand(r, 1, device=device) / float(n)
+ grid = torch.arange(n, device=device, dtype=torch.float32)[None, :] / float(n)
+ us = (U0 + grid).clamp(max=1.0 - 1e-12)
+ cdf_rows = cdf.index_select(0, rows)
+ idx = torch.searchsorted(cdf_rows, us).clamp_max(probs.size(1) - 1)
+ buf = torch.zeros(r, P, dtype=torch.float32, device=device)
+ buf.scatter_add_(1, idx, torch.ones_like(idx, dtype=buf.dtype))
+ out.index_copy_(0, rows, buf.to(torch.long))
+
+ return out.view(*batch_shape, P)
+
+
+# ---------------------------------------------------------------------------
+# VAE decoders
+# ---------------------------------------------------------------------------
+
+class LevelEmbedder(nn.Module):
+ def __init__(self, hidden_size, frequency_embedding_size=256, max_period=1024):
+ super().__init__()
+ self.mlp = nn.Sequential(
+ nn.Linear(frequency_embedding_size, hidden_size, bias=True),
+ nn.SiLU(),
+ nn.Linear(hidden_size, hidden_size, bias=True),
+ )
+ self.frequency_embedding_size = frequency_embedding_size
+ self.max_period = max_period
+
+ @staticmethod
+ def level_embedding(t, dim, max_period=1024):
+ half = dim // 2
+ freqs = torch.exp(-np.log(max_period) * torch.arange(start=0, end=half, dtype=torch.float32) / half).to(device=t.device)
+ args = t[:, None].float() * freqs[None] * 2 * torch.pi
+ embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1)
+ if dim % 2:
+ embedding = torch.cat([embedding, torch.zeros_like(embedding[:, :1])], dim=-1)
+ return embedding
+
+ def forward(self, t):
+ emb = self.level_embedding(t, self.frequency_embedding_size, self.max_period)
+ return self.mlp(emb.to(self.mlp[0].weight.dtype))
+
+
+class ModulatedTransformerCrossOnlyBlock(nn.Module):
+ def __init__(self, channels, ctx_channels, num_heads, mlp_ratio=4.0, share_mod=False,
+ qk_rms_norm_cross=True, qkv_bias=True):
+ super().__init__()
+ self.share_mod = share_mod
+ self.norm1 = LayerNorm32(channels, elementwise_affine=False, eps=1e-6)
+ self.norm2 = LayerNorm32(channels, elementwise_affine=False, eps=1e-6)
+ self.cross_attn = MultiHeadAttention(channels, ctx_channels=ctx_channels, num_heads=num_heads,
+ type="cross", attn_mode="full", qkv_bias=qkv_bias,
+ qk_rms_norm=qk_rms_norm_cross)
+ self.mlp = FeedForwardNet(channels, mlp_ratio=mlp_ratio)
+ if not share_mod:
+ self.adaLN_modulation = nn.Sequential(nn.SiLU(), nn.Linear(channels, 6 * channels, bias=True))
+
+ def forward(self, x, mod, context):
+ if self.share_mod:
+ shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = mod.chunk(6, dim=1)
+ else:
+ shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.adaLN_modulation(mod).chunk(6, dim=1)
+ h = self.norm1(x) * (1 + scale_msa.unsqueeze(1)) + shift_msa.unsqueeze(1)
+ x = x + self.cross_attn(h, context) * gate_msa.unsqueeze(1)
+ h = self.norm2(x) * (1 + scale_mlp.unsqueeze(1)) + shift_mlp.unsqueeze(1)
+ x = x + self.mlp(h) * gate_mlp.unsqueeze(1)
+ return x
+
+
+class ModulatedCrossOnlyTransformerBase(nn.Module):
+ def __init__(self, in_channels, model_channels, cond_channels, num_blocks, num_heads=None,
+ num_head_channels=64, mlp_ratio=4.0, share_mod=False, additional_level_embed=False,
+ qk_rms_norm_cross=True):
+ super().__init__()
+ self.model_channels = model_channels
+ self.cond_channels = cond_channels
+ self.num_blocks = num_blocks
+ self.num_heads = num_heads or model_channels // num_head_channels
+ self.mlp_ratio = mlp_ratio
+ self.share_mod = share_mod
+ self.qk_rms_norm_cross = qk_rms_norm_cross
+
+ self.input_layer = nn.Linear(in_channels, model_channels)
+ self.l_embedder = LevelEmbedder(model_channels)
+ self.l_embedder2 = LevelEmbedder(model_channels, max_period=100) if additional_level_embed else None
+ if share_mod:
+ self.adaLN_modulation = nn.Sequential(
+ nn.SiLU(), nn.Linear(model_channels, 6 * model_channels, bias=True))
+ if cond_channels is not None:
+ self.blocks = nn.ModuleList([
+ ModulatedTransformerCrossOnlyBlock(
+ model_channels, ctx_channels=cond_channels, num_heads=self.num_heads,
+ mlp_ratio=self.mlp_ratio, qk_rms_norm_cross=self.qk_rms_norm_cross,
+ share_mod=self.share_mod)
+ for _ in range(num_blocks)
+ ])
+
+ @property
+ def dtype(self) -> torch.dtype:
+ return next(self.parameters()).dtype
+
+ @property
+ def device(self) -> torch.device:
+ return next(self.parameters()).device
+
+ def forward(self, x, l, cond, l2=None):
+ h = self.input_layer(x)
+ l_emb = self.l_embedder(l)
+ if self.l_embedder2 is not None and l2 is not None:
+ l_emb = l_emb + self.l_embedder2(l2)
+ if self.share_mod:
+ l_emb = self.adaLN_modulation(l_emb)
+ for block in self.blocks:
+ h = block(h, l_emb, cond)
+ return h
+
+
+class OctreeProbabilityFixedlenDecoder(ModulatedCrossOnlyTransformerBase):
+ def __init__(self, model_channels, cond_channels, num_blocks, num_heads=None,
+ num_head_channels=64, mlp_ratio=4.0, share_mod=False,
+ additional_level_embed=False, qk_rms_norm_cross=True, *,
+ no_norm=False):
+ super().__init__(
+ in_channels=model_channels, model_channels=model_channels,
+ cond_channels=cond_channels, num_blocks=num_blocks,
+ num_heads=num_heads, num_head_channels=num_head_channels,
+ mlp_ratio=mlp_ratio, share_mod=share_mod,
+ additional_level_embed=additional_level_embed,
+ qk_rms_norm_cross=qk_rms_norm_cross,
+ )
+ self.out_proj = nn.Linear(self.model_channels, 8)
+ self.no_norm = no_norm
+ self.in_proj = nn.Linear(3, self.model_channels)
+ self.pos_embedder = PcdAbsolutePositionEmbedderV2(channels=model_channels, in_channels=3)
+
+ def forward(self, x, l, cond, l2=None):
+ d = self.dtype
+ B, L, C = x.shape
+ h = self.in_proj(x.to(d)) + self.pos_embedder(x.reshape(-1, 3)).reshape(B, L, -1).to(d)
+ if l2 is not None:
+ l2 = torch.log2(l2)
+ h = super().forward(h, l, cond.to(d), l2)
+ h = F.layer_norm(h.float(), h.shape[-1:]).to(d) if not self.no_norm else h / (1 + 2 * self.num_blocks) ** 0.5
+ logits = self.out_proj(h)
+ return {"logits": logits, "probs": torch.softmax(logits, dim=-1)}
+
+ @staticmethod
+ def sample(model, cond, num_points, level, temperature=1.0, algo="systematic"):
+ B = cond.shape[0]
+ device = cond.device
+ child_offset = torch.tensor([[i, j, k] for k in [0, 1] for j in [0, 1] for i in [0, 1]],
+ dtype=torch.long, device=device)
+ prev_coords_int = torch.zeros(B, 1, 3, dtype=torch.long, device=device)
+ prev_counts = torch.full((B, 1), num_points, dtype=torch.long, device=device)
+ prev_log_probs = torch.zeros(B, 1, dtype=torch.float32, device=device)
+ batch_indices_range = torch.arange(B, device=device).unsqueeze(1)
+ num_tensor = torch.full((B,), num_points, dtype=torch.long, device=device)
+
+ for lv in range(1, level + 1):
+ res_p = 1 << (lv - 1)
+ res = 1 << lv
+ parent_coords_norm = (prev_coords_int.to(torch.float32) + 0.5) / res_p
+ res_tensor = torch.full((B,), res, dtype=torch.long, device=device)
+ pred_logits = model(parent_coords_norm, res_tensor, cond, num_tensor)["logits"] / temperature
+ pred_probs = torch.softmax(pred_logits, dim=-1)
+ pred_log_probs = torch.log_softmax(pred_logits, dim=-1)
+ sampled = sample_probs(pred_probs, prev_counts, algo=algo).flatten(1, 2)
+ pred_log_probs = pred_log_probs.flatten(1, 2)
+ prev_log_probs_expanded = prev_log_probs.repeat_interleave(8, dim=1)
+ child_coords_int = (prev_coords_int[:, :, None, :] * 2 + child_offset[None, None, :, :]).flatten(1, 2)
+ mask = sampled > 0
+ max_valid = mask.sum(dim=1).max().item()
+ scatter_indices = mask.cumsum(dim=1) - 1
+ valid_scatter_indices = scatter_indices[mask]
+ valid_batch_indices = batch_indices_range.expand_as(mask)[mask]
+ next_prev_coords_int = torch.zeros(B, max_valid, 3, dtype=child_coords_int.dtype, device=device)
+ next_prev_coords_int[valid_batch_indices, valid_scatter_indices] = child_coords_int[mask]
+ next_prev_counts = torch.zeros(B, max_valid, dtype=sampled.dtype, device=device)
+ next_prev_counts[valid_batch_indices, valid_scatter_indices] = sampled[mask]
+ next_prev_log_probs = torch.zeros(B, max_valid, dtype=prev_log_probs.dtype, device=device)
+ next_prev_log_probs[valid_batch_indices, valid_scatter_indices] = (prev_log_probs_expanded + pred_log_probs)[mask]
+ prev_coords_int = next_prev_coords_int
+ prev_counts = next_prev_counts
+ prev_log_probs = next_prev_log_probs
+
+ res = 1 << level
+ prev_log_probs = torch.repeat_interleave(prev_log_probs.flatten(0, 1), prev_counts.flatten(0, 1), dim=0).reshape(B, num_points)
+ coords_int = torch.repeat_interleave(prev_coords_int.flatten(0, 1), prev_counts.flatten(0, 1), dim=0).reshape(B, num_points, -1)
+ coords_norm = (coords_int.to(torch.float32) + torch.rand_like(coords_int, dtype=torch.float32)) / res
+ return {"points": coords_norm, "log_probs": prev_log_probs}
+
+
+class TransformerCrossBlock(nn.Module):
+ def __init__(self, channels, ctx_channels, num_heads, mlp_ratio=4.0, attn_mode="full",
+ qk_rms_norm=True, qk_rms_norm_cross=True, qkv_bias=True):
+ super().__init__()
+ self.norm1 = LayerNorm32(channels, elementwise_affine=False, eps=1e-6)
+ self.norm2 = LayerNorm32(channels, elementwise_affine=True, eps=1e-6)
+ self.norm3 = LayerNorm32(channels, elementwise_affine=False, eps=1e-6)
+ self.self_attn = MultiHeadAttention(channels, num_heads=num_heads, type="self",
+ attn_mode=attn_mode, qkv_bias=qkv_bias,
+ qk_rms_norm=qk_rms_norm)
+ self.cross_attn = MultiHeadAttention(channels, ctx_channels=ctx_channels, num_heads=num_heads,
+ type="cross", attn_mode="full", qkv_bias=qkv_bias,
+ qk_rms_norm=qk_rms_norm_cross)
+ self.mlp = FeedForwardNet(channels, mlp_ratio=mlp_ratio)
+
+ def forward(self, x, context):
+ x = x + self.self_attn(self.norm1(x))
+ x = x + self.cross_attn(self.norm2(x), context)
+ x = x + self.mlp(self.norm3(x))
+ return x
+
+
+class TransformerBase(nn.Module):
+ def __init__(self, in_channels, model_channels, cond_channels, num_blocks, num_heads=None,
+ num_head_channels=64, mlp_ratio=4.0, attn_mode="full", window_num=None,
+ qk_rms_norm=True, qk_rms_norm_cross=True):
+ super().__init__()
+ self.model_channels = model_channels
+ self.cond_channels = cond_channels
+ self.num_blocks = num_blocks
+ self.num_heads = num_heads or model_channels // num_head_channels
+ self.mlp_ratio = mlp_ratio
+ self.input_layer = nn.Linear(in_channels, model_channels)
+ if cond_channels is not None:
+ self.blocks = nn.ModuleList([
+ TransformerCrossBlock(model_channels, ctx_channels=cond_channels,
+ num_heads=self.num_heads, mlp_ratio=self.mlp_ratio,
+ attn_mode="full", qk_rms_norm=qk_rms_norm,
+ qk_rms_norm_cross=qk_rms_norm_cross)
+ for _ in range(num_blocks)
+ ])
+
+ @property
+ def dtype(self) -> torch.dtype:
+ return next(self.parameters()).dtype
+
+ def forward(self, x, cond=None, l=None, cond2=None):
+ h = self.input_layer(x)
+ for block in self.blocks:
+ h = block(h, cond)
+ return h
+
+
+class FixedlenDecoder(TransformerBase):
+ def __init__(self, in_channels, model_channels, cond_channels, num_blocks, num_heads=None,
+ num_head_channels=64, mlp_ratio=4.0, attn_mode="full", window_num=None,
+ qk_rms_norm=True, qk_rms_norm_cross=True):
+ super().__init__(in_channels=model_channels, model_channels=model_channels,
+ cond_channels=cond_channels, num_blocks=num_blocks,
+ num_heads=num_heads, num_head_channels=num_head_channels,
+ mlp_ratio=mlp_ratio, attn_mode=attn_mode, window_num=window_num,
+ qk_rms_norm=qk_rms_norm, qk_rms_norm_cross=qk_rms_norm_cross)
+ self.in_proj = nn.Linear(in_channels, model_channels)
+ self.pos_embedder = PcdAbsolutePositionEmbedderV2(channels=model_channels, in_channels=3)
+
+ def forward(self, x=None, cond=None):
+ pcd = x["points"]
+ d = self.dtype
+ B, L, C = pcd.shape
+ h = self.in_proj(pcd.to(d)) + self.pos_embedder(pcd.reshape(-1, 3)).reshape(B, L, -1).to(d)
+ return super().forward(h, cond.to(d))
+
+
+class ElasticGaussianFixedlenDecoder(FixedlenDecoder):
+ def __init__(self, in_channels, model_channels, cond_channels, num_blocks, num_heads=None,
+ num_head_channels=64, mlp_ratio=4.0, attn_mode="full", window_num=None,
+ *, no_norm=False, representation_config=None,
+ use_learned_offset_scale=True, use_per_offset=True,
+ qk_rms_norm=True, qk_rms_norm_cross=True):
+ self.rep_config = representation_config
+ self.use_learned_offset_scale = use_learned_offset_scale
+ self.use_per_offset = use_per_offset
+ self.out_channels = self._calc_layout()
+ super().__init__(in_channels=in_channels, model_channels=model_channels,
+ cond_channels=cond_channels, num_blocks=num_blocks,
+ num_heads=num_heads, num_head_channels=num_head_channels,
+ mlp_ratio=mlp_ratio, attn_mode=attn_mode, window_num=window_num,
+ qk_rms_norm=qk_rms_norm, qk_rms_norm_cross=qk_rms_norm_cross)
+ self.out_proj = nn.Linear(model_channels, self.out_channels)
+ self.no_norm = no_norm
+ self._build_perturbation()
+
+ def _calc_layout(self):
+ ng = self.rep_config['num_gaussians']
+ self.layout = {
+ '_xyz': {'shape': (ng, 3), 'size': ng * 3},
+ '_features_dc': {'shape': (ng, 1, 3), 'size': ng * 3},
+ '_scaling': {'shape': (ng, 3), 'size': ng * 3},
+ '_rotation': {'shape': (ng, 4), 'size': ng * 4},
+ '_opacity': {'shape': (ng, 1), 'size': ng},
+ }
+ if self.use_learned_offset_scale and self.use_per_offset:
+ self.layout['_offset_scale'] = {'shape': (ng, 1), 'size': ng}
+ start = 0
+ for k, v in self.layout.items():
+ v['range'] = (start, start + v['size'])
+ start += v['size']
+ return start
+
+ def _build_perturbation(self):
+ ng = self.rep_config['num_gaussians']
+ perturbation = torch.tensor([hammersley_sequence(3, i, ng) for i in range(ng)]).float()
+ perturbation = torch.atanh((perturbation * 2 - 1) / self.rep_config['perturbe_size'])
+ self.register_buffer('points_offset_perturbation', perturbation)
+ if self.use_learned_offset_scale:
+ base = torch.tensor(self.rep_config['offset_scale'])
+ self.register_buffer('base_offset_scale', torch.log(torch.exp(base) - 1.0))
+
+ def _get_offset(self, h):
+ B = h.shape[0]
+ if self.use_learned_offset_scale:
+ r = self.layout['_offset_scale']['range']
+ _offset_scale = F.softplus(
+ h[:, :, r[0]:r[1]].reshape(B, -1, *self.layout['_offset_scale']['shape'])
+ + self.base_offset_scale)
+
+ r = self.layout['_xyz']['range']
+ offset = h[:, :, r[0]:r[1]].reshape(B, -1, *self.layout['_xyz']['shape'])
+ offset = offset * self.rep_config['lr']['_xyz']
+ if self.rep_config['perturb_offset']:
+ offset = offset + self.points_offset_perturbation
+ offset = torch.tanh(offset) * 0.5 * self.rep_config['perturbe_size']
+ offset = offset * (_offset_scale if self.use_learned_offset_scale else self.rep_config['offset_scale'])
+ return offset
+
+ def forward(self, x=None, cond=None):
+ h = super().forward(x, cond)
+ h = F.layer_norm(h.float(), h.shape[-1:]).to(h.dtype) if not self.no_norm else h / (1 + 3 * self.num_blocks) ** 0.5
+ return {"features": self.out_proj(h)}
+
+
+# ---------------------------------------------------------------------------
+# Flow matching denoiser
+# ---------------------------------------------------------------------------
+
+class TimestepEmbedder(nn.Module):
+ def __init__(self, hidden_size, frequency_embedding_size=256):
+ super().__init__()
+ self.mlp = nn.Sequential(
+ nn.Linear(frequency_embedding_size, hidden_size, bias=True),
+ nn.SiLU(),
+ nn.Linear(hidden_size, hidden_size, bias=True),
+ )
+ self.frequency_embedding_size = frequency_embedding_size
+
+ @staticmethod
+ def timestep_embedding(t, dim, max_period=10000):
+ half = dim // 2
+ freqs = torch.exp(-np.log(max_period) * torch.arange(start=0, end=half, dtype=torch.float32) / half).to(device=t.device)
+ args = t[:, None].float() * freqs[None]
+ embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1)
+ if dim % 2:
+ embedding = torch.cat([embedding, torch.zeros_like(embedding[:, :1])], dim=-1)
+ return embedding
+
+ def forward(self, t):
+ emb = self.timestep_embedding(t, self.frequency_embedding_size)
+ return self.mlp(emb.to(self.mlp[0].weight.dtype))
+
+
+class LatentSeqMMFlowModel(nn.Module):
+ def __init__(self, q_token_length, in_channels, model_channels, cond_channels,
+ out_channels, num_blocks, num_refiner_blocks=2, num_heads=None,
+ num_head_channels=64, cam_channels=None, cond2_channels=None,
+ mlp_ratio=4, share_mod=True, qk_rms_norm=False, use_shift_table=False):
+ super().__init__()
+ self.q_token_length = q_token_length
+ self.in_channels = in_channels
+ self.cam_channels = cam_channels
+ self.model_channels = model_channels
+ self.cond_channels = cond_channels
+ self.cond2_channels = cond2_channels
+ self.out_channels = out_channels
+ self.num_blocks = num_blocks
+ self.num_refiner_blocks = num_refiner_blocks
+ self.num_heads = num_heads or model_channels // num_head_channels
+ self.mlp_ratio = mlp_ratio
+ self.share_mod = share_mod
+ self.qk_rms_norm = qk_rms_norm
+ self.use_shift_table = use_shift_table
+
+ self.t_embedder = TimestepEmbedder(model_channels)
+ if share_mod:
+ self.adaLN_modulation = nn.Sequential(
+ nn.SiLU(), nn.Linear(model_channels, 6 * model_channels, bias=True))
+
+ self.input_layer = nn.Linear(in_channels, model_channels)
+ self.cond_embedder = nn.Linear(cond_channels, model_channels)
+ self.cond_embedder2 = nn.Linear(cond2_channels, model_channels) if cond2_channels is not None else None
+
+ sobol_seq = torch.quasirandom.SobolEngine(dimension=3, scramble=True, seed=123).draw(q_token_length)
+ self.pos_pe = sobol_seq.unsqueeze(0)
+ self.pos_embedder = PcdAbsolutePositionEmbedder(model_channels)
+ self.noise_repo_layers = nn.ModuleList([
+ RePo3DRotaryEmbedding(model_channels, num_heads=self.num_heads, head_dim=num_head_channels)
+ for _ in range(num_refiner_blocks)])
+ self.context_repo_layers = nn.ModuleList([
+ RePo3DRotaryEmbedding(model_channels, num_heads=self.num_heads, head_dim=num_head_channels)
+ for _ in range(num_refiner_blocks)])
+ self.repo_layers = nn.ModuleList([
+ RePo3DRotaryEmbedding(model_channels, num_heads=self.num_heads, head_dim=num_head_channels)
+ for _ in range(num_blocks)])
+
+ block_kwargs = dict(num_heads=self.num_heads, mlp_ratio=self.mlp_ratio, attn_mode='full',
+ use_rope=True, qk_rms_norm=self.qk_rms_norm,
+ use_shift_table=self.use_shift_table)
+ self.noise_refiner = nn.ModuleList([
+ UnifiedTransformerBlock(model_channels, modulation=True, share_mod=self.share_mod, **block_kwargs)
+ for _ in range(num_refiner_blocks)])
+ self.context_refiner = nn.ModuleList([
+ UnifiedTransformerBlock(model_channels, modulation=False, **block_kwargs)
+ for _ in range(num_refiner_blocks)])
+ if self.cam_channels is not None:
+ self.cam_refiner = MLP(self.cam_channels, model_channels, model_channels,
+ mlp_layer_num=num_refiner_blocks)
+ self.blocks = nn.ModuleList([
+ UnifiedTransformerBlock(model_channels, modulation=True, share_mod=self.share_mod, **block_kwargs)
+ for _ in range(num_blocks)])
+ self.shift_table = nn.Parameter(torch.randn(1, 2, model_channels) / model_channels**0.5) if use_shift_table else None
+ self.out_layer = nn.Linear(model_channels, out_channels)
+ if cam_channels is not None:
+ self.cam_out_layer = nn.Linear(model_channels, cam_channels)
+
+ @property
+ def dtype(self) -> torch.dtype:
+ return next(self.parameters()).dtype
+
+ @property
+ def device(self) -> torch.device:
+ return next(self.parameters()).device
+
+ def load_safetensors(self, path: str) -> None:
+ self.load_state_dict(safetensors.torch.load_file(path), strict=True)
+
+ def forward(self, x_t, t, cond):
+ d = self.dtype
+ z = x_t['latent'].to(d)
+ feat1 = cond['feature1'].to(d)
+ feat2 = cond['feature2'].to(d) if self.cond_embedder2 is not None else None
+ self.pos_pe = self.pos_pe.to(z.device)
+
+ h_x = self.input_layer(z)
+ h_cond = self.cond_embedder(feat1)
+ if feat2 is not None:
+ h_cond = h_cond + self.cond_embedder2(feat2)
+ t_emb = self.t_embedder(t)
+ t_mod = self.adaLN_modulation(t_emb) if self.share_mod else t_emb
+
+ h_x = h_x + self.pos_embedder(self.pos_pe).to(d)
+
+ for i, block in enumerate(self.noise_refiner):
+ h_x = block(h_x, mod=t_mod, rotary_emb=self.noise_repo_layers[i](h_x))
+
+ for i, block in enumerate(self.context_refiner):
+ h_cond = block(h_cond, mod=None, rotary_emb=self.context_repo_layers[i](h_cond))
+
+ if self.cam_channels is not None:
+ cam = x_t.get('camera').to(d)
+ h_cam = self.cam_refiner(cam)
+
+ h = torch.cat([h_x, h_cond], dim=1)
+ if self.cam_channels is not None:
+ h = torch.cat([h, h_cam], dim=1)
+
+ for i, block in enumerate(self.blocks):
+ h = block(h, mod=t_mod, rotary_emb=self.repo_layers[i](h))
+
+ h_x = F.layer_norm(h[:, :z.shape[1]].float(), h.shape[-1:]).type(d)
+ if self.cam_channels is not None:
+ h_cam = F.layer_norm(h[:, -cam.shape[1]:].float(), h.shape[-1:]).type(d)
+
+ if self.use_shift_table:
+ shift, scale = (self.shift_table + t_emb.unsqueeze(1)).chunk(2, dim=1)
+ h_x = h_x * (1 + scale) + shift
+ if self.cam_channels is not None:
+ h_cam = h_cam * (1 + scale) + shift
+
+ out = {'latent': self.out_layer(h_x)}
+ if self.cam_channels is not None:
+ out['camera'] = self.cam_out_layer(h_cam)
+ return out
+
+
+# ---------------------------------------------------------------------------
+# OctreeGaussianDecoder
+# ---------------------------------------------------------------------------
+
+class OctreeGaussianDecoder(nn.Module):
+ _MAX_VOXEL_LEVEL = 8
+
+ def __init__(self, octree_args: dict, gs_args: dict):
+ super().__init__()
+ self.octree = OctreeProbabilityFixedlenDecoder(**octree_args)
+ self.gs = ElasticGaussianFixedlenDecoder(**gs_args)
+
+ def load_safetensors(self, path: str) -> None:
+ self.load_state_dict(safetensors.torch.load_file(path), strict=True)
+
+ @property
+ def gaussians_per_point(self) -> int:
+ return self.gs.rep_config['num_gaussians']
+
+ @torch.no_grad()
+ def decode(self, latent: torch.Tensor, num_gaussians: int):
+ from .triposplat import _build_gaussians # local import: avoid model.py ↔ triposplat.py cycle
+ num_decoder_tokens = max(1, num_gaussians // self.gaussians_per_point)
+ points_pred = OctreeProbabilityFixedlenDecoder.sample(
+ self.octree, latent,
+ num_points=num_decoder_tokens, level=self._MAX_VOXEL_LEVEL,
+ temperature=1.0, algo='systematic',
+ )
+ pred = self.gs(x=points_pred, cond=latent)
+ return _build_gaussians(self.gs, points_pred, pred)[0]
diff --git a/invokeai/backend/image_util/triposplat/triposplat.py b/invokeai/backend/image_util/triposplat/triposplat.py
new file mode 100644
index 00000000000..b014bee65f0
--- /dev/null
+++ b/invokeai/backend/image_util/triposplat/triposplat.py
@@ -0,0 +1,599 @@
+import numpy as np
+import torch
+import torch.nn.functional as F
+import safetensors.torch
+from PIL import Image, ImageFilter
+from torchvision import transforms
+from tqdm.auto import tqdm
+
+from .model import (
+ DinoV3ViT, Flux2VAEEncoder, BiRefNet,
+ OctreeProbabilityFixedlenDecoder, ElasticGaussianFixedlenDecoder,
+ LatentSeqMMFlowModel, OctreeGaussianDecoder,
+)
+
+
+# ---------------------------------------------------------------------------
+# Gaussian
+# ---------------------------------------------------------------------------
+
+class Gaussian:
+ def __init__(self, aabb: list, sh_degree: int = 0, mininum_kernel_size: float = 0.0,
+ scaling_bias: float = 0.01, opacity_bias: float = 0.1,
+ scaling_activation: str = "exp", device='cuda'):
+ self.sh_degree = sh_degree
+ self.mininum_kernel_size = mininum_kernel_size
+ self.scaling_bias = scaling_bias
+ self.opacity_bias = opacity_bias
+ self.device = device
+ self.aabb = torch.tensor(aabb, dtype=torch.float32, device=device)
+
+ if scaling_activation == "exp":
+ self._scaling_activation = torch.exp
+ self._inverse_scaling_activation = torch.log
+ elif scaling_activation == "softplus":
+ self._scaling_activation = F.softplus
+ self._inverse_scaling_activation = lambda x: x + torch.log(-torch.expm1(-x))
+
+ self._opacity_activation = torch.sigmoid
+ self._inverse_opacity_activation = lambda x: torch.log(x / (1 - x))
+
+ self.scale_bias = self._inverse_scaling_activation(torch.tensor(self.scaling_bias)).to(self.device)
+ self.rots_bias = torch.zeros(4, device=self.device)
+ self.rots_bias[0] = 1
+ self.opacity_bias_val = self._inverse_opacity_activation(torch.tensor(self.opacity_bias)).to(self.device)
+
+ self._storage = {}
+
+ def _get_store(self, name):
+ return self._storage.get(name)
+
+ def _set_store(self, name, value):
+ self._storage[name] = value
+
+ @property
+ def _xyz(self):
+ return self._get_store("_xyz")
+ @_xyz.setter
+ def _xyz(self, value):
+ if value is None:
+ self._set_store("_xyz", None); self._set_store("xyz", None); return
+ self._set_store("_xyz", value)
+ self._set_store("xyz", value * self.aabb[None, 3:] + self.aabb[None, :3])
+
+ @property
+ def get_xyz(self):
+ return self._get_store("xyz")
+
+ @property
+ def _features_dc(self):
+ return self._get_store("_features_dc")
+ @_features_dc.setter
+ def _features_dc(self, value):
+ self._set_store("_features_dc", value)
+
+ @property
+ def _opacity(self):
+ return self._get_store("_opacity")
+ @_opacity.setter
+ def _opacity(self, value):
+ if value is None:
+ self._set_store("_opacity", None); self._set_store("opacity", None); return
+ self._set_store("_opacity", value)
+ self._set_store("opacity", self._opacity_activation(value + self.opacity_bias_val))
+
+ @property
+ def get_opacity(self):
+ return self._get_store("opacity")
+
+ @property
+ def _scaling(self):
+ return self._get_store("_scaling")
+ @_scaling.setter
+ def _scaling(self, value):
+ if value is None:
+ self._set_store("_scaling", None); self._set_store("scaling", None); return
+ self._set_store("_scaling", value)
+ s = self._scaling_activation(value + self.scale_bias)
+ s = torch.square(s) + self.mininum_kernel_size ** 2
+ self._set_store("scaling", torch.sqrt(s))
+
+ @property
+ def get_scaling(self):
+ return self._get_store("scaling")
+
+ @property
+ def _rotation(self):
+ return self._get_store("_rotation")
+ @_rotation.setter
+ def _rotation(self, value):
+ self._set_store("_rotation", value)
+
+ def construct_list_of_attributes(self):
+ l = ['x', 'y', 'z', 'nx', 'ny', 'nz']
+ dc = self._features_dc
+ for i in range(dc.shape[1] * dc.shape[2]):
+ l.append(f'f_dc_{i}')
+ l.append('opacity')
+ for i in range(self._scaling.shape[1]):
+ l.append(f'scale_{i}')
+ for i in range(self._rotation.shape[1]):
+ l.append(f'rot_{i}')
+ return l
+
+ _DEFAULT_TRANSFORM = [[1, 0, 0], [0, 0, -1], [0, 1, 0]]
+
+ def _get_ply_data(self, transform=None):
+ xyz = self.get_xyz.detach().cpu().numpy()
+ normals = np.zeros_like(xyz)
+ f_dc = self._features_dc.detach().transpose(1, 2).flatten(start_dim=1).contiguous().cpu().numpy()
+ opacities = self._inverse_opacity_activation(self.get_opacity).detach().cpu().numpy()
+ scale = torch.log(self.get_scaling).detach().cpu().numpy()
+ rotation = (self._rotation + self.rots_bias[None, :]).detach().cpu().numpy()
+ if transform is not None:
+ transform = np.array(transform)
+ xyz = np.matmul(xyz, transform.T)
+ R_mat = _quat_to_matrix(rotation)
+ R_mat = np.matmul(transform, R_mat)
+ rotation = _matrix_to_quat(R_mat)
+ return xyz, normals, f_dc, opacities, scale, rotation
+
+ def _transformed_xyz_rot(self, transform=None):
+ if transform is None:
+ transform = self._DEFAULT_TRANSFORM
+ transform = np.array(transform, dtype=np.float32)
+ xyz = self.get_xyz.detach().cpu().numpy().astype(np.float32)
+ rotation = (self._rotation + self.rots_bias[None, :]).detach().cpu().numpy()
+ xyz = np.matmul(xyz, transform.T)
+ R_mat = _quat_to_matrix(rotation)
+ R_mat = np.matmul(transform, R_mat)
+ rotation = _matrix_to_quat(R_mat)
+ return xyz, rotation
+
+ def to_ply_bytes(self, transform=None) -> bytes:
+ if transform is None:
+ transform = self._DEFAULT_TRANSFORM
+ xyz, normals, f_dc, opacities, scale, rotation = self._get_ply_data(transform=transform)
+ dtype_full = [(attr, 'f4') for attr in self.construct_list_of_attributes()]
+ elements = np.empty(xyz.shape[0], dtype=dtype_full)
+ elements[:] = list(map(tuple, np.concatenate((xyz, normals, f_dc, opacities, scale, rotation), axis=1)))
+ return _binary_ply_bytes(elements, dtype_full)
+
+ def to_splat_bytes(self, transform=None) -> bytes:
+ if transform is None:
+ transform = self._DEFAULT_TRANSFORM
+ xyz, rotation = self._transformed_xyz_rot(transform=transform)
+ scale = self.get_scaling.detach().cpu().numpy().astype(np.float32)
+ opacity = self.get_opacity.detach().cpu().numpy()
+ f_dc = self._features_dc.detach().cpu().numpy()
+ C0 = 0.28209479177387814
+ # .splat packs color as 4 bytes RGBA: RGB from the SH DC term, A from opacity.
+ rgb = np.clip((f_dc[:, 0, :] * C0 + 0.5) * 255, 0, 255).astype(np.uint8)
+ alpha = np.clip(opacity[:, 0:1] * 255, 0, 255).astype(np.uint8)
+ rgba = np.concatenate([rgb, alpha], axis=1)
+ rot = rotation / np.linalg.norm(rotation, axis=-1, keepdims=True)
+ rot_u8 = np.clip(rot * 128 + 128, 0, 255).astype(np.uint8)
+ order = np.argsort(-opacity[:, 0] * np.prod(scale, axis=-1))
+ xyz, scale, rgba, rot_u8 = xyz[order], scale[order], rgba[order], rot_u8[order]
+ # Per-splat record is exactly 32 bytes: xyz(12) + scale(12) + rgba(4) + rot(4).
+ data = np.concatenate([
+ xyz.astype(np.float32).view(np.uint8).reshape(-1, 12),
+ scale.astype(np.float32).view(np.uint8).reshape(-1, 12),
+ rgba.reshape(-1, 4),
+ rot_u8.reshape(-1, 4),
+ ], axis=1).reshape(-1)
+ return data.tobytes()
+
+ def save_ply(self, path, transform=None):
+ with open(path, 'wb') as f:
+ f.write(self.to_ply_bytes(transform=transform))
+
+ def save_splat(self, path, transform=None):
+ with open(path, 'wb') as f:
+ f.write(self.to_splat_bytes(transform=transform))
+
+
+def _binary_ply_bytes(elements, dtype_full) -> bytes:
+ num_vertices = len(elements)
+ header = "ply\nformat binary_little_endian 1.0\n"
+ header += f"element vertex {num_vertices}\n"
+ type_map = {'f4': 'float', 'u1': 'uchar', 'i4': 'int'}
+ for name, t in dtype_full:
+ header += f"property {type_map.get(t, t)} {name}\n"
+ header += "end_header\n"
+ return header.encode('ascii') + elements.tobytes()
+
+
+def _quat_to_matrix(q):
+ q = q / np.linalg.norm(q, axis=-1, keepdims=True)
+ w, x, y, z = q[:, 0], q[:, 1], q[:, 2], q[:, 3]
+ R = np.stack([
+ 1 - 2*(y*y + z*z), 2*(x*y - w*z), 2*(x*z + w*y),
+ 2*(x*y + w*z), 1 - 2*(x*x + z*z), 2*(y*z - w*x),
+ 2*(x*z - w*y), 2*(y*z + w*x), 1 - 2*(x*x + y*y),
+ ], axis=-1).reshape(-1, 3, 3)
+ return R
+
+
+def _matrix_to_quat(R):
+ trace = R[:, 0, 0] + R[:, 1, 1] + R[:, 2, 2]
+ q = np.zeros((R.shape[0], 4), dtype=R.dtype)
+ s = np.sqrt(np.maximum(trace + 1, 0)) * 2
+ q[:, 0] = 0.25 * s
+ q[:, 1] = (R[:, 2, 1] - R[:, 1, 2]) / np.where(s != 0, s, 1)
+ q[:, 2] = (R[:, 0, 2] - R[:, 2, 0]) / np.where(s != 0, s, 1)
+ q[:, 3] = (R[:, 1, 0] - R[:, 0, 1]) / np.where(s != 0, s, 1)
+ m01 = (R[:, 0, 0] >= R[:, 1, 1]) & (R[:, 0, 0] >= R[:, 2, 2]) & (s == 0)
+ s1 = np.sqrt(np.maximum(1 + R[:, 0, 0] - R[:, 1, 1] - R[:, 2, 2], 0)) * 2
+ q[m01, 0] = (R[m01, 2, 1] - R[m01, 1, 2]) / s1[m01]
+ q[m01, 1] = 0.25 * s1[m01]
+ q[m01, 2] = (R[m01, 0, 1] + R[m01, 1, 0]) / s1[m01]
+ q[m01, 3] = (R[m01, 0, 2] + R[m01, 2, 0]) / s1[m01]
+ m11 = (R[:, 1, 1] > R[:, 0, 0]) & (R[:, 1, 1] >= R[:, 2, 2]) & (s == 0)
+ s2 = np.sqrt(np.maximum(1 + R[:, 1, 1] - R[:, 0, 0] - R[:, 2, 2], 0)) * 2
+ q[m11, 0] = (R[m11, 0, 2] - R[m11, 2, 0]) / s2[m11]
+ q[m11, 1] = (R[m11, 0, 1] + R[m11, 1, 0]) / s2[m11]
+ q[m11, 2] = 0.25 * s2[m11]
+ q[m11, 3] = (R[m11, 1, 2] + R[m11, 2, 1]) / s2[m11]
+ m21 = (R[:, 2, 2] > R[:, 0, 0]) & (R[:, 2, 2] > R[:, 1, 1]) & (s == 0)
+ s3 = np.sqrt(np.maximum(1 + R[:, 2, 2] - R[:, 0, 0] - R[:, 1, 1], 0)) * 2
+ q[m21, 0] = (R[m21, 1, 0] - R[m21, 0, 1]) / s3[m21]
+ q[m21, 1] = (R[m21, 0, 2] + R[m21, 2, 0]) / s3[m21]
+ q[m21, 2] = (R[m21, 1, 2] + R[m21, 2, 1]) / s3[m21]
+ q[m21, 3] = 0.25 * s3[m21]
+ return q / np.linalg.norm(q, axis=-1, keepdims=True)
+
+
+def _build_gaussians(decoder: ElasticGaussianFixedlenDecoder, points_pred: dict, pred: dict):
+ x = points_pred
+ offset = decoder._get_offset(pred['features'])
+ h = pred["features"]
+ ret = []
+ for i in range(h.shape[0]):
+ g = Gaussian(
+ sh_degree=0,
+ aabb=[-0.5, -0.5, -0.5, 1.0, 1.0, 1.0],
+ mininum_kernel_size=decoder.rep_config['filter_kernel_size_3d'],
+ scaling_bias=decoder.rep_config['scaling_bias'],
+ opacity_bias=decoder.rep_config['opacity_bias'],
+ scaling_activation=decoder.rep_config['scaling_activation'],
+ device=x['points'].device,
+ )
+ _x = x["points"][i, :, None, :]
+ for k, v in decoder.layout.items():
+ if k == '_xyz':
+ setattr(g, k, (offset[i] + _x).flatten(0, 1))
+ elif k in ('_xyz_center', '_offset_scale'):
+ continue
+ else:
+ feats = h[i][:, v['range'][0]:v['range'][1]].reshape(-1, *v['shape']).flatten(0, 1)
+ setattr(g, k, feats * decoder.rep_config['lr'][k])
+ ret.append(g)
+ return ret
+
+
+# ---------------------------------------------------------------------------
+# Euler flow sampler
+# ---------------------------------------------------------------------------
+
+class FlowEulerCfgSampler:
+ def __init__(self, sigma_min: float = 1e-5):
+ self.sigma_min = sigma_min
+
+ def _get_batch_size(self, x_t):
+ return next(iter(x_t.values())).shape[0] if isinstance(x_t, dict) else x_t.shape[0]
+
+ def _get_device(self, x_t):
+ return next(iter(x_t.values())).device if isinstance(x_t, dict) else x_t.device
+
+ def _inference_model(self, model, x_t, t, cond=None):
+ batch = self._get_batch_size(x_t)
+ device = self._get_device(x_t)
+ t_scaled = torch.tensor([1000 * t] * batch, device=device, dtype=torch.float32)
+ if isinstance(cond, dict):
+ for k, v in cond.items():
+ if isinstance(v, torch.Tensor) and v.shape[0] == 1 and batch > 1:
+ cond[k] = v.repeat(batch, *([1] * (len(v.shape) - 1)))
+ elif cond is not None and cond.shape[0] == 1 and batch > 1:
+ cond = cond.repeat(batch, *([1] * (len(cond.shape) - 1)))
+ return model(x_t, t_scaled, cond)
+
+ def _cfg_prediction(self, model, x_t, t, cond, neg_cond, guidance_scale):
+ # Diffusers-style convention: guidance_scale == 1 (or <= 1, or None) means no CFG —
+ # only the conditional pass runs, halving the per-step cost. > 1 enables CFG and
+ # blends as `pred = s * cond + (1 - s) * uncond = s * cond - (s - 1) * uncond`.
+ pred_v = self._inference_model(model, x_t, t, cond)
+ if isinstance(guidance_scale, dict):
+ if not any(s > 1 for s in guidance_scale.values()):
+ return pred_v
+ neg_pred_v = self._inference_model(model, x_t, t, neg_cond)
+ for key in pred_v:
+ s = guidance_scale.get(key, 1.0)
+ if s > 1:
+ pred_v[key] = s * pred_v[key] - (s - 1) * neg_pred_v[key]
+ return pred_v
+ if guidance_scale is None or guidance_scale <= 1:
+ return pred_v
+ neg_pred_v = self._inference_model(model, x_t, t, neg_cond)
+ for key in pred_v:
+ pred_v[key] = guidance_scale * pred_v[key] - (guidance_scale - 1) * neg_pred_v[key]
+ return pred_v
+
+ @torch.no_grad()
+ def sample(self, model, noise, cond, neg_cond, steps=50, shift=1.0,
+ guidance_scale=None, show_progress=False, callback=None):
+ sample = noise
+ t_seq = shift * np.linspace(1, 0, steps + 1) / (1 + (shift - 1) * np.linspace(1, 0, steps + 1))
+ t_pairs = list(zip(t_seq[:-1], t_seq[1:]))
+ iterator = tqdm(t_pairs, desc="Sampling", total=steps) if show_progress else t_pairs
+ for i, (t, t_prev) in enumerate(iterator):
+ x_t = {k: v.clone() for k, v in sample.items()} if isinstance(sample, dict) else sample.clone()
+ pred_v = self._cfg_prediction(model, x_t, t, cond, neg_cond, guidance_scale)
+ dt = t - t_prev
+ if isinstance(sample, dict):
+ for key in sample:
+ sample[key] = sample[key] - pred_v[key] * dt
+ else:
+ sample = sample - pred_v * dt
+ if callback is not None:
+ callback(i + 1, steps)
+ return sample
+
+
+# ---------------------------------------------------------------------------
+# Component loaders
+# ---------------------------------------------------------------------------
+
+def _place(m, device, dtype):
+ if device is not None or dtype is not None:
+ m = m.to(device=device, dtype=dtype)
+ return m.eval()
+
+
+def load_dinov3(path: str, device=None, dtype=None) -> DinoV3ViT:
+ m = DinoV3ViT()
+ m.load_safetensors(path)
+ return _place(m, device, dtype)
+
+
+def load_vae_encoder(path: str, device=None, dtype=None) -> Flux2VAEEncoder:
+ m = Flux2VAEEncoder()
+ m.load_safetensors(path)
+ return _place(m, device, dtype)
+
+
+def load_rmbg(path: str, device=None, dtype=None) -> BiRefNet:
+ m = BiRefNet()
+ m.load_safetensors(path)
+ return _place(m, device, dtype)
+
+
+FLOW_MODEL_ARGS = dict(
+ q_token_length=8192, in_channels=16, cam_channels=5, out_channels=16,
+ model_channels=1024, cond_channels=1280, cond2_channels=128,
+ num_refiner_blocks=2, num_blocks=24, num_heads=16, mlp_ratio=4,
+ qk_rms_norm=True, share_mod=True, use_shift_table=True,
+)
+
+
+def load_flow_model(path: str, device=None, dtype=None) -> LatentSeqMMFlowModel:
+ m = LatentSeqMMFlowModel(**FLOW_MODEL_ARGS)
+ m.load_safetensors(path)
+ return _place(m, device, dtype)
+
+
+OCTREE_DECODER_ARGS = dict(
+ model_channels=1024, cond_channels=16,
+ num_blocks=4, num_heads=16, mlp_ratio=4, share_mod=True,
+)
+
+GS_DECODER_ARGS = dict(
+ in_channels=3, model_channels=1024, cond_channels=16,
+ attn_mode="full", num_blocks=16, num_heads=16, mlp_ratio=4,
+ use_learned_offset_scale=True, use_per_offset=True,
+ representation_config=dict(
+ lr=dict(_xyz=1.0, _features_dc=1.0, _opacity=1.0, _scaling=1.0, _rotation=0.1),
+ perturb_offset=True, perturbe_size=1.5, offset_scale=0.05, num_gaussians=32,
+ filter_kernel_size_3d=0.0009, scaling_bias=0.004, opacity_bias=0.1,
+ scaling_activation="softplus",
+ ),
+)
+
+
+def load_decoder(path: str, device=None, dtype=None) -> OctreeGaussianDecoder:
+ m = OctreeGaussianDecoder(OCTREE_DECODER_ARGS, GS_DECODER_ARGS)
+ m.load_safetensors(path)
+ return _place(m, device, dtype)
+
+
+# ---------------------------------------------------------------------------
+# Pipeline stages
+# ---------------------------------------------------------------------------
+
+_CANVAS_SIZE = 1024
+
+
+def _image_to_pil(image) -> Image.Image:
+ if isinstance(image, Image.Image):
+ return image
+ if isinstance(image, (str, bytes)) or hasattr(image, "__fspath__"):
+ return Image.open(image)
+ if isinstance(image, torch.Tensor):
+ t = image.detach().cpu()
+ if t.ndim == 4:
+ assert t.shape[0] == 1, (
+ f"batched image input is not supported (got B={t.shape[0]}); "
+ "pass one image at a time"
+ )
+ t = t[0]
+ arr = (t.clamp(0, 1) * 255).to(torch.uint8).numpy()
+ # No explicit mode: Pillow infers RGB/RGBA from the uint8 HWC shape (mode= is deprecated, removed in Pillow 13)
+ return Image.fromarray(arr)
+ raise TypeError(f"unsupported image type: {type(image)}")
+
+
+def preprocess_image(image, rmbg: BiRefNet, erode_radius: int = 1) -> Image.Image:
+ image = _image_to_pil(image)
+ size = _CANVAS_SIZE
+ w, h = image.size
+ s = size / min(w, h)
+ image = image.resize((max(1, int(round(w * s))), max(1, int(round(h * s)))), Image.LANCZOS)
+ has_real_alpha = (image.mode == "RGBA"
+ and np.array(image.getchannel(3), dtype=np.int32).min() < 255)
+ if not has_real_alpha:
+ image = rmbg.remove_background(image.convert("RGB"))
+ if erode_radius > 0:
+ image.putalpha(image.getchannel(3).filter(ImageFilter.MinFilter(2 * erode_radius + 1)))
+ alpha = np.array(image.getchannel(3))
+ ys, xs = np.nonzero(alpha)
+ bbox = [xs.min(), ys.min(), xs.max(), ys.max()]
+ cx, cy = (bbox[0] + bbox[2]) / 2, (bbox[1] + bbox[3]) / 2
+ half = max(bbox[2] - bbox[0], bbox[3] - bbox[1]) / 2 * 1.2
+ image = image.crop([int(cx - half), int(cy - half), int(cx + half), int(cy + half)])
+ image = image.resize((size, size), Image.LANCZOS)
+ bg = Image.new("RGB", (size, size), (0, 0, 0))
+ bg.paste(image, mask=image.split()[3])
+ return bg
+
+
+_DINOV3_NORMALIZE = transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
+
+
+@torch.no_grad()
+def encode_image(image: Image.Image, dinov3: DinoV3ViT, vae_encoder: Flux2VAEEncoder,
+ generator: torch.Generator = None) -> dict:
+ device = next(dinov3.parameters()).device
+ img_tensor = transforms.ToTensor()(image).unsqueeze(0).to(device=device, dtype=torch.float32)
+ img_normed = _DINOV3_NORMALIZE(img_tensor)
+ dinov3_dtype = next(dinov3.parameters()).dtype
+ vae_dtype = next(vae_encoder.parameters()).dtype
+ dinov3_feat = dinov3(pixel_values=img_normed.to(dinov3_dtype))
+ dinov3_feat = F.layer_norm(dinov3_feat.float(), dinov3_feat.shape[-1:])
+ vae_feat = vae_encoder.encode(img_tensor.to(vae_dtype) * 2 - 1,
+ deterministic=False, generator=generator)
+ # pad 5 zero tokens so feature2's token length matches feature1's (cls + 4 registers + patches)
+ zero_reg = torch.zeros(vae_feat.shape[0], 5, vae_feat.shape[2],
+ dtype=vae_feat.dtype, device=vae_feat.device)
+ vae_feat = torch.cat([zero_reg, vae_feat], dim=1)
+ return {'feature1': dinov3_feat, 'feature2': vae_feat}
+
+
+@torch.no_grad()
+def sample_latent(flow_model: LatentSeqMMFlowModel, cond: dict,
+ steps: int = 50, guidance_scale: float = 7.0, shift: float = 3.0,
+ generator: torch.Generator = None,
+ show_progress: bool = False, callback=None) -> dict:
+ device = flow_model.device
+ neg_cond = {k: torch.zeros_like(v) for k, v in cond.items()}
+ noise = {'latent': torch.randn(1, flow_model.q_token_length, flow_model.in_channels,
+ device=device, generator=generator)}
+ if flow_model.cam_channels is not None:
+ noise['camera'] = torch.randn(1, 1, flow_model.cam_channels,
+ device=device, generator=generator)
+ sampler = FlowEulerCfgSampler()
+ return sampler.sample(flow_model, noise, cond=cond, neg_cond=neg_cond,
+ steps=steps, guidance_scale=guidance_scale, shift=shift,
+ show_progress=show_progress, callback=callback)
+
+
+# ---------------------------------------------------------------------------
+# Pipeline
+# ---------------------------------------------------------------------------
+
+class TripoSplatPipeline:
+ def __init__(self, ckpt_path: str, decoder_path: str, dinov3_path: str,
+ flux2_vae_encoder_path: str, rmbg_path: str, device: str = "cuda"):
+ self._device = torch.device(device)
+ self.dinov3 = load_dinov3 (dinov3_path, device=self._device, dtype=torch.bfloat16)
+ self.vae_encoder = load_vae_encoder (flux2_vae_encoder_path, device=self._device, dtype=torch.bfloat16)
+ self.rmbg = load_rmbg (rmbg_path, device=self._device, dtype=torch.float16)
+ self.flow_model = load_flow_model (ckpt_path, device=self._device, dtype=torch.float16)
+ self.decoder = load_decoder (decoder_path, device=self._device, dtype=torch.float16)
+
+ def preprocess_image(self, image, erode_radius: int = 1) -> Image.Image:
+ return preprocess_image(image, self.rmbg, erode_radius=erode_radius)
+
+ def encode_image(self, image: Image.Image, generator: torch.Generator = None) -> dict:
+ return encode_image(image, self.dinov3, self.vae_encoder, generator=generator)
+
+ def sample_latent(self, cond: dict, steps: int = 50, guidance_scale: float = 7.0,
+ shift: float = 3.0, generator: torch.Generator = None,
+ show_progress: bool = False, callback=None) -> dict:
+ return sample_latent(self.flow_model, cond, steps=steps, guidance_scale=guidance_scale,
+ shift=shift, generator=generator,
+ show_progress=show_progress, callback=callback)
+
+ def decode_latent(self, latent: torch.Tensor, num_gaussians: int = 262144):
+ return self.decoder.decode(latent, num_gaussians=num_gaussians)
+
+ _NUM_GAUSSIANS_MIN = 32768
+ _NUM_GAUSSIANS_MAX = 262144
+
+ def _validate_num_gaussians(self, n: int) -> int:
+ assert self._NUM_GAUSSIANS_MIN <= n <= self._NUM_GAUSSIANS_MAX, (
+ f"num_gaussians must be in [{self._NUM_GAUSSIANS_MIN}, {self._NUM_GAUSSIANS_MAX}], got {n}"
+ )
+ gpp = self.decoder.gaussians_per_point
+ if n % gpp == 0:
+ return n
+ rounded = round(n / gpp) * gpp
+ print(f"[TripoSplatPipeline] num_gaussians={n} is not a multiple of {gpp}; rounding to {rounded}")
+ return rounded
+
+ @torch.no_grad()
+ def run(self, image, seed: int = 42, steps: int = 20, guidance_scale: float = 3.0,
+ shift: float = 3.0, num_gaussians=262144, erode_radius: int = 1,
+ show_progress: bool = False, callback=None):
+ """
+ Args:
+ image: Input image. Accepts a file path / PIL.Image / torch.Tensor
+ (`[1,H,W,C]` or `[H,W,C]`, float in `[0, 1]`, optional alpha
+ channel as the 4th channel).
+ seed: RNG seed for the VAE encoder's stochastic latent sampling and
+ the initial flow-matching noise. Same seed → same output.
+ steps: Number of Euler integrator steps in the flow-matching sampler.
+ More steps → better fidelity, linear runtime cost.
+ Recommend: 10~20.
+ guidance_scale: Classifier-free-guidance strength (diffusers
+ convention). `≤ 1.0` disables CFG. Higher → more detail,
+ stronger adherence to the input image; too high can cause color
+ oversaturation.
+ Recommend: 3.0.
+ shift: Flow-matching timestep schedule shift. `1.0` gives a uniform
+ schedule; `>1.0` allocates more steps to the early/high-noise end.
+ Recommend: 3.0.
+ num_gaussians: Target Gaussian-splat count. An `int` returns a
+ single `Gaussian`. A `list` / `tuple` of ints returns a
+ `list[Gaussian]`. Each count is rounded to the nearest multiple
+ of 32. More gaussians → more detail but higher rendering and
+ storage cost.
+ Recommend: 32768~262144.
+ erode_radius: Pixel radius used to erode the alpha matte after
+ background removal, to avoid segmentation-border bleed before
+ compositing on black. `0` disables; `1` is a 3×3 minimum filter.
+ Recommend: 1.
+ show_progress: Print a `tqdm` progress bar over sampler steps.
+ callback: Optional `fn(step, total)` invoked after each sampler step.
+ Useful for external progress UIs (e.g. ComfyUI's
+ `ProgressBar.update`).
+
+ Returns:
+ `(gaussian, prepared_image)` for an `int` `num_gaussians`, or
+ `(list_of_gaussians, prepared_image)` for a `list` / `tuple`. The
+ second element is the RGB composite the encoders actually saw —
+ useful for display / debugging.
+ """
+ if isinstance(num_gaussians, (list, tuple)):
+ counts = [self._validate_num_gaussians(n) for n in num_gaussians]
+ else:
+ counts = [self._validate_num_gaussians(num_gaussians)]
+
+ gen = torch.Generator(device=self._device).manual_seed(seed)
+ prepared = self.preprocess_image(image, erode_radius=erode_radius)
+ cond = self.encode_image(prepared, generator=gen)
+ out = self.sample_latent(cond, steps=steps, guidance_scale=guidance_scale, shift=shift,
+ generator=gen, show_progress=show_progress, callback=callback)
+ gaussians = [self.decode_latent(out['latent'], num_gaussians=n) for n in counts]
+ if isinstance(num_gaussians, (list, tuple)):
+ return gaussians, prepared
+ return gaussians[0], prepared
diff --git a/invokeai/backend/image_util/triposplat/triposplat_model.py b/invokeai/backend/image_util/triposplat/triposplat_model.py
new file mode 100644
index 00000000000..96359ab7e57
--- /dev/null
+++ b/invokeai/backend/image_util/triposplat/triposplat_model.py
@@ -0,0 +1,62 @@
+"""InvokeAI's model-manager wrapper for the vendored TripoSplat pipeline (this file is NOT vendored)."""
+
+from pathlib import Path
+from typing import Optional
+
+import torch
+
+from invokeai.backend.raw_model import RawModel
+
+
+class TripoSplatModel(RawModel):
+ """Wraps the vendored TripoSplatPipeline so the model manager's memory cache can move it between CPU
+ and GPU.
+
+ The pipeline holds five sub-modules with mixed dtypes (DINOv3 + Flux2 VAE are bfloat16; rmbg / flow /
+ decoder are float16), so we move device only and never re-cast dtype.
+ """
+
+ def __init__(self, pipe: object):
+ self._pipe = pipe
+
+ @property
+ def _submodules(self) -> tuple:
+ p = self._pipe
+ return (p.dinov3, p.vae_encoder, p.rmbg, p.flow_model, p.decoder)
+
+ def to(self, device: Optional[torch.device] = None, dtype: Optional[torch.dtype] = None) -> None:
+ # Only CPU/CUDA are supported (mirrors GroundingDinoPipeline's MPS guard). dtype is intentionally
+ # ignored: the sub-modules require their original mixed precision.
+ if device is None or device.type not in {"cpu", "cuda"}:
+ return
+ for module in self._submodules:
+ module.to(device=device)
+ self._pipe._device = torch.device(device)
+
+ def calc_size(self) -> int:
+ from invokeai.backend.model_manager.load.model_util import calc_module_size
+
+ return sum(calc_module_size(module) for module in self._submodules)
+
+ @staticmethod
+ def load_model(model_path: Path) -> "TripoSplatModel":
+ from invokeai.backend.image_util.triposplat.triposplat import TripoSplatPipeline
+
+ def _find(filename: str) -> str:
+ # The download cache may hand us the repo dir or a parent; locate each checkpoint by name so
+ # we are robust to the exact directory layout returned by download_and_cache_model.
+ matches = list(model_path.rglob(filename))
+ if not matches:
+ raise FileNotFoundError(f"TripoSplat checkpoint '{filename}' not found under {model_path}")
+ return str(matches[0])
+
+ # Construct on CPU; the model cache moves it to GPU on lock via .to().
+ pipe = TripoSplatPipeline(
+ ckpt_path=_find("triposplat_fp16.safetensors"),
+ decoder_path=_find("triposplat_vae_decoder_fp16.safetensors"),
+ dinov3_path=_find("dino_v3_vit_h.safetensors"),
+ flux2_vae_encoder_path=_find("flux2-vae.safetensors"),
+ rmbg_path=_find("birefnet.safetensors"),
+ device="cpu",
+ )
+ return TripoSplatModel(pipe)
diff --git a/invokeai/backend/model_manager/load/model_util.py b/invokeai/backend/model_manager/load/model_util.py
index 0c6c8f7f4ab..9d7569c3483 100644
--- a/invokeai/backend/model_manager/load/model_util.py
+++ b/invokeai/backend/model_manager/load/model_util.py
@@ -15,6 +15,7 @@
from invokeai.backend.image_util.depth_anything.depth_anything_pipeline import DepthAnythingPipeline
from invokeai.backend.image_util.grounding_dino.grounding_dino_pipeline import GroundingDinoPipeline
from invokeai.backend.image_util.segment_anything.segment_anything_pipeline import SegmentAnythingPipeline
+from invokeai.backend.image_util.triposplat.triposplat_model import TripoSplatModel
from invokeai.backend.ip_adapter.ip_adapter import IPAdapter
from invokeai.backend.model_manager.taxonomy import AnyModel
from invokeai.backend.onnx.onnx_runtime import IAIOnnxRuntimeModel
@@ -49,6 +50,7 @@ def calc_model_size_by_data(logger: logging.Logger, model: AnyModel) -> int:
GroundingDinoPipeline,
SegmentAnythingPipeline,
DepthAnythingPipeline,
+ TripoSplatModel,
),
):
return model.calc_size()
diff --git a/invokeai/frontend/web/openapi.json b/invokeai/frontend/web/openapi.json
index be3ff4059e1..cc41ac98448 100644
--- a/invokeai/frontend/web/openapi.json
+++ b/invokeai/frontend/web/openapi.json
@@ -5510,6 +5510,50 @@
]
}
},
+ "/api/v1/assets/i/{asset_name}": {
+ "get": {
+ "tags": ["assets"],
+ "summary": "Get Asset",
+ "description": "Gets a 3D asset file (e.g. a Gaussian-splat .ply).",
+ "operationId": "get_asset",
+ "parameters": [
+ {
+ "name": "asset_name",
+ "in": "path",
+ "required": true,
+ "schema": {
+ "type": "string",
+ "description": "The name of the 3D asset file to get",
+ "title": "Asset Name"
+ },
+ "description": "The name of the 3D asset file to get"
+ }
+ ],
+ "responses": {
+ "200": {
+ "description": "The 3D asset was fetched successfully",
+ "content": {
+ "application/json": {
+ "schema": {}
+ }
+ }
+ },
+ "404": {
+ "description": "The 3D asset could not be found"
+ },
+ "422": {
+ "description": "Validation Error",
+ "content": {
+ "application/json": {
+ "schema": {
+ "$ref": "#/components/schemas/HTTPValidationError"
+ }
+ }
+ }
+ }
+ }
+ }
+ },
"/api/v1/boards/": {
"post": {
"tags": ["boards"],
@@ -11973,6 +12017,41 @@
"$ref": "#/components/schemas/ImageOutput"
}
},
+ "Asset3DField": {
+ "description": "A 3D asset primitive field",
+ "properties": {
+ "asset_name": {
+ "description": "The name of the 3D asset file",
+ "title": "Asset Name",
+ "type": "string"
+ }
+ },
+ "required": ["asset_name"],
+ "title": "Asset3DField",
+ "type": "object"
+ },
+ "Asset3DOutput": {
+ "class": "output",
+ "description": "Base class for nodes that output a single 3D asset",
+ "properties": {
+ "asset": {
+ "$ref": "#/components/schemas/Asset3DField",
+ "description": "The output 3D asset",
+ "field_kind": "output",
+ "ui_hidden": false
+ },
+ "type": {
+ "const": "asset_3d_output",
+ "default": "asset_3d_output",
+ "field_kind": "node_attribute",
+ "title": "type",
+ "type": "string"
+ }
+ },
+ "required": ["output_meta", "asset", "type", "type"],
+ "title": "Asset3DOutput",
+ "type": "object"
+ },
"BaseMetadata": {
"properties": {
"name": {
@@ -29139,6 +29218,9 @@
{
"$ref": "#/components/schemas/ImageScaleInvocation"
},
+ {
+ "$ref": "#/components/schemas/ImageTo3DInvocation"
+ },
{
"$ref": "#/components/schemas/ImageToLatentsInvocation"
},
@@ -29622,6 +29704,9 @@
{
"$ref": "#/components/schemas/AnimaModelLoaderOutput"
},
+ {
+ "$ref": "#/components/schemas/Asset3DOutput"
+ },
{
"$ref": "#/components/schemas/BooleanCollectionOutput"
},
@@ -34937,6 +35022,157 @@
"$ref": "#/components/schemas/ImageOutput"
}
},
+ "ImageTo3DInvocation": {
+ "category": "3d",
+ "class": "invocation",
+ "classification": "prototype",
+ "description": "Generates a 3D Gaussian splat (.ply) from a single image using TripoSplat.",
+ "node_pack": "invokeai",
+ "properties": {
+ "board": {
+ "anyOf": [
+ {
+ "$ref": "#/components/schemas/BoardField"
+ },
+ {
+ "type": "null"
+ }
+ ],
+ "default": null,
+ "description": "The board to save the image to",
+ "field_kind": "internal",
+ "input": "direct",
+ "orig_required": false,
+ "ui_hidden": false
+ },
+ "metadata": {
+ "anyOf": [
+ {
+ "$ref": "#/components/schemas/MetadataField"
+ },
+ {
+ "type": "null"
+ }
+ ],
+ "default": null,
+ "description": "Optional metadata to be saved with the image",
+ "field_kind": "internal",
+ "input": "connection",
+ "orig_required": false,
+ "ui_hidden": false
+ },
+ "id": {
+ "description": "The id of this instance of an invocation. Must be unique among all instances of invocations.",
+ "field_kind": "node_attribute",
+ "title": "Id",
+ "type": "string"
+ },
+ "is_intermediate": {
+ "default": false,
+ "description": "Whether or not this is an intermediate invocation.",
+ "field_kind": "node_attribute",
+ "input": "direct",
+ "orig_required": true,
+ "title": "Is Intermediate",
+ "type": "boolean",
+ "ui_hidden": false,
+ "ui_type": "IsIntermediate"
+ },
+ "use_cache": {
+ "default": true,
+ "description": "Whether or not to use the cache",
+ "field_kind": "node_attribute",
+ "title": "Use Cache",
+ "type": "boolean"
+ },
+ "image": {
+ "anyOf": [
+ {
+ "$ref": "#/components/schemas/ImageField"
+ },
+ {
+ "type": "null"
+ }
+ ],
+ "default": null,
+ "description": "The image to convert to a 3D Gaussian splat.",
+ "field_kind": "input",
+ "input": "any",
+ "orig_required": true
+ },
+ "remove_background": {
+ "default": true,
+ "description": "Isolate the subject by removing the background (BiRefNet). Disable to feed the full image as-is.",
+ "field_kind": "input",
+ "input": "any",
+ "orig_default": true,
+ "orig_required": false,
+ "title": "Remove Background",
+ "type": "boolean"
+ },
+ "num_gaussians": {
+ "default": 262144,
+ "description": "Number of 3D Gaussians to generate. Higher is more detailed but larger and slower to render.",
+ "field_kind": "input",
+ "input": "any",
+ "maximum": 262144,
+ "minimum": 32768,
+ "orig_default": 262144,
+ "orig_required": false,
+ "title": "Num Gaussians",
+ "type": "integer"
+ },
+ "steps": {
+ "default": 20,
+ "description": "Number of flow-matching sampler steps.",
+ "field_kind": "input",
+ "input": "any",
+ "maximum": 100,
+ "minimum": 1,
+ "orig_default": 20,
+ "orig_required": false,
+ "title": "Steps",
+ "type": "integer"
+ },
+ "guidance_scale": {
+ "default": 3.0,
+ "description": "Classifier-free guidance scale. Values <= 1.0 disable guidance.",
+ "field_kind": "input",
+ "input": "any",
+ "minimum": 0.0,
+ "orig_default": 3.0,
+ "orig_required": false,
+ "title": "Guidance Scale",
+ "type": "number"
+ },
+ "seed": {
+ "default": 42,
+ "description": "Seed for reproducible generation.",
+ "field_kind": "input",
+ "input": "any",
+ "minimum": 0,
+ "orig_default": 42,
+ "orig_required": false,
+ "title": "Seed",
+ "type": "integer"
+ },
+ "type": {
+ "const": "image_to_3d",
+ "default": "image_to_3d",
+ "field_kind": "node_attribute",
+ "title": "type",
+ "type": "string"
+ }
+ },
+ "required": ["type", "id"],
+ "tags": ["3d", "image", "gaussian", "splat"],
+ "title": "Image to 3D (TripoSplat)",
+ "type": "object",
+ "version": "1.0.0",
+ "output": {
+ "$ref": "#/components/schemas/Asset3DOutput"
+ }
+ },
"ImageToLatentsInvocation": {
"category": "latents",
"class": "invocation",
@@ -36835,6 +37071,9 @@
{
"$ref": "#/components/schemas/ImageScaleInvocation"
},
+ {
+ "$ref": "#/components/schemas/ImageTo3DInvocation"
+ },
{
"$ref": "#/components/schemas/ImageToLatentsInvocation"
},
@@ -37275,6 +37514,9 @@
{
"$ref": "#/components/schemas/AnimaModelLoaderOutput"
},
+ {
+ "$ref": "#/components/schemas/Asset3DOutput"
+ },
{
"$ref": "#/components/schemas/BooleanCollectionOutput"
},
@@ -37985,6 +38227,9 @@
{
"$ref": "#/components/schemas/ImageScaleInvocation"
},
+ {
+ "$ref": "#/components/schemas/ImageTo3DInvocation"
+ },
{
"$ref": "#/components/schemas/ImageToLatentsInvocation"
},
@@ -38753,6 +38998,9 @@
"image_panel_layout": {
"$ref": "#/components/schemas/ImagePanelCoordinateOutput"
},
+ "image_to_3d": {
+ "$ref": "#/components/schemas/Asset3DOutput"
+ },
"img_blur": {
"$ref": "#/components/schemas/ImageOutput"
},
@@ -39330,6 +39578,7 @@
"image_generator",
"image_mask_to_tensor",
"image_panel_layout",
+ "image_to_3d",
"img_blur",
"img_chan",
"img_channel_multiply",
@@ -39910,6 +40159,9 @@
{
"$ref": "#/components/schemas/ImageScaleInvocation"
},
+ {
+ "$ref": "#/components/schemas/ImageTo3DInvocation"
+ },
{
"$ref": "#/components/schemas/ImageToLatentsInvocation"
},
@@ -40809,6 +41061,9 @@
{
"$ref": "#/components/schemas/ImageScaleInvocation"
},
+ {
+ "$ref": "#/components/schemas/ImageTo3DInvocation"
+ },
{
"$ref": "#/components/schemas/ImageToLatentsInvocation"
},
diff --git a/invokeai/frontend/web/package.json b/invokeai/frontend/web/package.json
index 6c7ea3f65ca..b7960dee3a9 100644
--- a/invokeai/frontend/web/package.json
+++ b/invokeai/frontend/web/package.json
@@ -51,6 +51,7 @@
"@observ33r/object-equals": "^1.1.6",
"@reduxjs/toolkit": "2.8.2",
"@roarr/browser-log-writer": "^1.3.0",
+ "@sparkjsdev/spark": "^2.1.0",
"@xyflow/react": "^12.10.0",
"ag-psd": "^28.5.1",
"async-mutex": "^0.5.0",
@@ -102,6 +103,7 @@
"serialize-error": "^12.0.0",
"socket.io-client": "^4.8.3",
"stable-hash": "^0.0.6",
+ "three": "^0.180.0",
"use-debounce": "^10.0.6",
"use-device-pixel-ratio": "^1.1.2",
"uuid": "^11.1.0",
@@ -121,6 +123,7 @@
"@types/node": "^22.19.3",
"@types/react": "^19.2.14",
"@types/react-dom": "^19.2.3",
+ "@types/three": "^0.180.0",
"@types/uuid": "^10.0.0",
"@typescript-eslint/eslint-plugin": "^8.59.2",
"@typescript-eslint/parser": "^8.59.2",
diff --git a/invokeai/frontend/web/pnpm-lock.yaml b/invokeai/frontend/web/pnpm-lock.yaml
index 4901ad00405..8897e76c3e5 100644
--- a/invokeai/frontend/web/pnpm-lock.yaml
+++ b/invokeai/frontend/web/pnpm-lock.yaml
@@ -41,6 +41,9 @@ importers:
'@roarr/browser-log-writer':
specifier: ^1.3.0
version: 1.3.0
+ '@sparkjsdev/spark':
+ specifier: ^2.1.0
+ version: 2.1.0(three@0.180.0)
'@xyflow/react':
specifier: ^12.10.0
version: 12.10.0(@types/react@19.2.14)(immer@10.2.0)(react-dom@19.2.6(react@19.2.6))(react@19.2.6)
@@ -194,6 +197,9 @@ importers:
stable-hash:
specifier: ^0.0.6
version: 0.0.6
+ three:
+ specifier: ^0.180.0
+ version: 0.180.0
use-debounce:
specifier: ^10.0.6
version: 10.0.6(react@19.2.6)
@@ -234,6 +240,9 @@ importers:
'@types/react-dom':
specifier: ^19.2.3
version: 19.2.3(@types/react@19.2.14)
+ '@types/three':
+ specifier: ^0.180.0
+ version: 0.180.0
'@types/uuid':
specifier: ^10.0.0
version: 10.0.0
@@ -742,6 +751,9 @@ packages:
resolution: {integrity: sha512-mepCf/e9+SKYy1d02/UkvSy6+6MoyXhVxP8lLDfA7BPE1X1d4dR0sZznmbM8/XVJ1GPM+Svnx7Xj6ZweByWUkw==}
engines: {node: '>17.0.0'}
+ '@dimforge/rapier3d-compat@0.12.0':
+ resolution: {integrity: sha512-uekIGetywIgopfD97oDL5PfeezkFpNhwlzlaEYNOA0N6ghdsOvh/HYjSMek5Q2O1PYvRSDFcqFVJl4r4ZBwOow==}
+
'@dmsnell/diff-match-patch@1.1.0':
resolution: {integrity: sha512-yejLPmM5pjsGvxS9gXablUSbInW7H976c/FJ4iQxWIm7/38xBySRemTPDe34lhg1gVLbJntX0+sH0jYfU+PN9A==}
@@ -1680,6 +1692,11 @@ packages:
'@socket.io/component-emitter@3.1.2':
resolution: {integrity: sha512-9BCxFwvbGg/RsZK9tjXd8s4UcwR0MWeFQ1XEKIQVVvAGJyINdrqKMcTRyLoK8Rse1GjzLV9cwjWV1olXRWEXVA==}
+ '@sparkjsdev/spark@2.1.0':
+ resolution: {integrity: sha512-BRw+MuMzx0B3K8fDLQygt2OHEhYUV+41RX7btq9pZ3rCVrq42o57jW34VAIvC7JO/84DJh/1AutACV9ym6BfVg==}
+ peerDependencies:
+ three: '>=0.180.0'
+
'@standard-schema/spec@1.1.0':
resolution: {integrity: sha512-l2aFy5jALhniG5HgqrD6jXLi/rUWrKvqN/qJx6yoJsgKhblVd+iqqU4RCXavm/jPityDo5TCvKMnpjKnOriy0w==}
@@ -1860,6 +1877,9 @@ packages:
peerDependencies:
'@testing-library/dom': '>=7.21.4'
+ '@tweenjs/tween.js@23.1.3':
+ resolution: {integrity: sha512-vJmvvwFxYuGnF2axRtPYocag6Clbb5YS7kLL+SO/TeVFzHqDIWrNKYtcsPMibjDx9O+bu+psAy9NKfWklassUA==}
+
'@tybys/wasm-util@0.10.1':
resolution: {integrity: sha512-9tTaPJLSiejZKx+Bmog4uSubteqTvFrVrURwkmHixBo0G4seD0zUxp98E1DzUBJxLQ3NPwXrGKDiVjwx/DpPsg==}
@@ -1963,12 +1983,21 @@ packages:
'@types/resolve@1.20.6':
resolution: {integrity: sha512-A4STmOXPhMUtHH+S6ymgE2GiBSMqf4oTvcQZMcHzokuTLVYzXTB8ttjcgxOVaAp2lGwEdzZ0J+cRbbeevQj1UQ==}
+ '@types/stats.js@0.17.4':
+ resolution: {integrity: sha512-jIBvWWShCvlBqBNIZt0KAshWpvSjhkwkEu4ZUcASoAvhmrgAUI2t1dXrjSL4xXVLB4FznPrIsX3nKXFl/Dt4vA==}
+
+ '@types/three@0.180.0':
+ resolution: {integrity: sha512-ykFtgCqNnY0IPvDro7h+9ZeLY+qjgUWv+qEvUt84grhenO60Hqd4hScHE7VTB9nOQ/3QM8lkbNE+4vKjEpUxKg==}
+
'@types/use-sync-external-store@0.0.6':
resolution: {integrity: sha512-zFDAD+tlpf2r4asuHEj0XH6pY6i0g5NeAHPn+15wk3BV6JA69eERFXC1gyGThDkVa1zCyKr5jox1+2LbV/AMLg==}
'@types/uuid@10.0.0':
resolution: {integrity: sha512-7gqG38EyHgyP1S+7+xomFtL+ZNHcKv6DwNaCZmJmo1vgMugyF3TCnXVg4t1uk89mLNwnLtnY3TpOpCOyp1/xHQ==}
+ '@types/webxr@0.5.24':
+ resolution: {integrity: sha512-h8fgEd/DpoS9CBrjEQXR+dIDraopAEfu4wYVNY2tEPwk60stPWhvZMf4Foo5FakuQ7HFZoa8WceaWFervK2Ovg==}
+
'@typescript-eslint/eslint-plugin@8.59.2':
resolution: {integrity: sha512-j/bwmkBvHUtPNxzuWe5z6BEk3q54YRyGlBXkSsmfoih7zNrBvl5A9A98anlp/7JbyZcWIJ8KXo/3Tq/DjFLtuQ==}
engines: {node: ^18.18.0 || ^20.9.0 || >=21.1.0}
@@ -2129,6 +2158,9 @@ packages:
'@webcontainer/env@1.1.1':
resolution: {integrity: sha512-6aN99yL695Hi9SuIk1oC88l9o0gmxL1nGWWQ/kNy81HigJ0FoaoTXpytCj6ItzgyCEwA9kF1wixsTuv5cjsgng==}
+ '@webgpu/types@0.1.71':
+ resolution: {integrity: sha512-mMy8/ODcKhab808co15eW+yN+HgXoQxRQHTiBV9Mrvl1r0ufnid7YOcI+gi4eUWSWl9ezD6TW2KXccrL8HCh2A==}
+
'@xobotyi/scrollbar-width@1.9.5':
resolution: {integrity: sha512-N8tkAACJx2ww8vFMneJmaAgmjAG1tnVBZJRLRcx061tmsLRZHSEZSLuGWnwPtunsSLvSqXQ2wfp7Mgqg1I+2dQ==}
@@ -3045,6 +3077,7 @@ packages:
glob@10.5.0:
resolution: {integrity: sha512-DfXN8DfhJ7NH3Oe7cFmu3NCu1wKbkReJ8TorzSAFbSKrlNaQSKfIzqYqVY8zlbs2NLBbWpRiU52GX2PbaBVNkg==}
+ deprecated: Old versions of glob are not supported, and contain widely publicized security vulnerabilities, which have been fixed in the current version. Please update. Support for old versions may be purchased (at exorbitant rates) by contacting i@izs.me
hasBin: true
glob@13.0.6:
@@ -3591,6 +3624,9 @@ packages:
resolution: {integrity: sha512-8q7VEgMJW4J8tcfVPy8g09NcQwZdbwFEqhe/WZkoIzjn/3TGDwtOCYtXGxA3O8tPzpczCCDgv+P2P5y00ZJOOg==}
engines: {node: '>= 8'}
+ meshoptimizer@0.22.0:
+ resolution: {integrity: sha512-IebiK79sqIy+E4EgOr+CAw+Ke8hAspXKzBd0JdgEmPHiAwmvEj2S4h1rfvo+o/BnfEYd/jAOg5IeeIjzlzSnDg==}
+
micromatch@4.0.8:
resolution: {integrity: sha512-PXwfBhYu0hBCPw8Dn0E+WDYb7af3dSLVWKi3HGv84IdF4TyFoC0ysxFd0Goxw7nSv4T/PzEJQxsYsEiFCKo2BA==}
engines: {node: '>=8.6'}
@@ -4528,6 +4564,9 @@ packages:
resolution: {integrity: sha512-ot0WnXS9fgdkgIcePe6RHNk1WA8+muPa6cSjeR3V8K27q9BB1rTE3R1p7Hv0z1ZyAc8s6Vvv8DIyWf681MAt0w==}
engines: {node: '>= 0.4'}
+ three@0.180.0:
+ resolution: {integrity: sha512-o+qycAMZrh+TsE01GqWUxUIKR1AL0S8pq7zDkYOQw8GqfX8b8VoCKYUoHbhiX5j+7hr8XsuHDVU6+gkQJQKg9w==}
+
throttle-debounce@3.0.1:
resolution: {integrity: sha512-dTEWWNu6JmeVXY0ZYoPuH5cRIwc0MeGbJwah9KUNYSJwommQpCzTySTpEe8Gs1J23aeWEuAobe4Ag7EHVt/LOg==}
engines: {node: '>=10'}
@@ -5545,6 +5584,8 @@ snapshots:
'@dagrejs/graphlib@2.2.4': {}
+ '@dimforge/rapier3d-compat@0.12.0': {}
+
'@dmsnell/diff-match-patch@1.1.0': {}
'@emnapi/core@1.10.0':
@@ -6338,6 +6379,11 @@ snapshots:
'@socket.io/component-emitter@3.1.2': {}
+ '@sparkjsdev/spark@2.1.0(three@0.180.0)':
+ dependencies:
+ fflate: 0.8.2
+ three: 0.180.0
+
'@standard-schema/spec@1.1.0': {}
'@standard-schema/utils@0.3.0': {}
@@ -6497,7 +6543,7 @@ snapshots:
'@testing-library/dom@10.4.0':
dependencies:
- '@babel/code-frame': 7.29.0
+ '@babel/code-frame': 7.29.7
'@babel/runtime': 7.29.2
'@types/aria-query': 5.0.4
aria-query: 5.3.0
@@ -6519,6 +6565,8 @@ snapshots:
dependencies:
'@testing-library/dom': 10.4.0
+ '@tweenjs/tween.js@23.1.3': {}
+
'@tybys/wasm-util@0.10.1':
dependencies:
tslib: 2.8.1
@@ -6631,10 +6679,24 @@ snapshots:
'@types/resolve@1.20.6': {}
+ '@types/stats.js@0.17.4': {}
+
+ '@types/three@0.180.0':
+ dependencies:
+ '@dimforge/rapier3d-compat': 0.12.0
+ '@tweenjs/tween.js': 23.1.3
+ '@types/stats.js': 0.17.4
+ '@types/webxr': 0.5.24
+ '@webgpu/types': 0.1.71
+ fflate: 0.8.2
+ meshoptimizer: 0.22.0
+
'@types/use-sync-external-store@0.0.6': {}
'@types/uuid@10.0.0': {}
+ '@types/webxr@0.5.24': {}
+
'@typescript-eslint/eslint-plugin@8.59.2(@typescript-eslint/parser@8.59.2(eslint@9.39.2(jiti@2.6.1))(typescript@5.9.3))(eslint@9.39.2(jiti@2.6.1))(typescript@5.9.3)':
dependencies:
'@eslint-community/regexpp': 4.12.2
@@ -6875,6 +6937,8 @@ snapshots:
'@webcontainer/env@1.1.1': {}
+ '@webgpu/types@0.1.71': {}
+
'@xobotyi/scrollbar-width@1.9.5': {}
'@xyflow/react@12.10.0(@types/react@19.2.14)(immer@10.2.0)(react-dom@19.2.6(react@19.2.6))(react@19.2.6)':
@@ -8468,6 +8532,8 @@ snapshots:
merge2@1.4.1: {}
+ meshoptimizer@0.22.0: {}
+
micromatch@4.0.8:
dependencies:
braces: 3.0.3
@@ -9524,6 +9590,8 @@ snapshots:
supports-preserve-symlinks-flag@1.0.0: {}
+ three@0.180.0: {}
+
throttle-debounce@3.0.1: {}
tiny-invariant@1.3.3: {}
diff --git a/invokeai/frontend/web/src/features/controlLayers/components/RasterLayer/RasterLayerMenuItems.tsx b/invokeai/frontend/web/src/features/controlLayers/components/RasterLayer/RasterLayerMenuItems.tsx
index 3c8081d9214..d1ff47d6319 100644
--- a/invokeai/frontend/web/src/features/controlLayers/components/RasterLayer/RasterLayerMenuItems.tsx
+++ b/invokeai/frontend/web/src/features/controlLayers/components/RasterLayer/RasterLayerMenuItems.tsx
@@ -1,6 +1,7 @@
import { MenuDivider } from '@invoke-ai/ui-library';
import { IconMenuItemGroup } from 'common/components/IconMenuItem';
import { CanvasEntityMenuItemsArrange } from 'features/controlLayers/components/common/CanvasEntityMenuItemsArrange';
+import { CanvasEntityMenuItemsConvertTo3D } from 'features/controlLayers/components/common/CanvasEntityMenuItemsConvertTo3D';
import { CanvasEntityMenuItemsCropToBbox } from 'features/controlLayers/components/common/CanvasEntityMenuItemsCropToBbox';
import { CanvasEntityMenuItemsDelete } from 'features/controlLayers/components/common/CanvasEntityMenuItemsDelete';
import { CanvasEntityMenuItemsDuplicate } from 'features/controlLayers/components/common/CanvasEntityMenuItemsDuplicate';
@@ -28,6 +29,7 @@ export const RasterLayerMenuItems = memo(() => {
+
diff --git a/invokeai/frontend/web/src/features/controlLayers/components/SplatOverlay/CanvasSplatOverlay.tsx b/invokeai/frontend/web/src/features/controlLayers/components/SplatOverlay/CanvasSplatOverlay.tsx
new file mode 100644
index 00000000000..7add040fd77
--- /dev/null
+++ b/invokeai/frontend/web/src/features/controlLayers/components/SplatOverlay/CanvasSplatOverlay.tsx
@@ -0,0 +1,390 @@
+import { Box, Button, ButtonGroup, Flex, Spinner, Text } from '@invoke-ai/ui-library';
+import { useStore } from '@nanostores/react';
+import { logger } from 'app/logging/logger';
+import { useAppDispatch } from 'app/store/storeHooks';
+import type { RectCorner } from 'features/controlLayers/components/SplatOverlay/rectTransforms';
+import { moveRect, resizeRectFromCorner } from 'features/controlLayers/components/SplatOverlay/rectTransforms';
+import type { SplatScene } from 'features/controlLayers/components/SplatOverlay/splatScene';
+import type { SplatRect } from 'features/controlLayers/components/SplatOverlay/state';
+import {
+ $splatOverlay,
+ clearSplatOverlay,
+ updateSplatOverlayRect,
+} from 'features/controlLayers/components/SplatOverlay/state';
+import { useCanvasManager } from 'features/controlLayers/contexts/CanvasManagerProviderGate';
+import { rasterLayerAdded } from 'features/controlLayers/store/canvasSlice';
+import { imageDTOToImageObject } from 'features/controlLayers/store/util';
+import type { PointerEvent as ReactPointerEvent } from 'react';
+import { lazy, memo, Suspense, useCallback, useEffect, useRef, useState } from 'react';
+import { uploadImage } from 'services/api/endpoints/images';
+
+const log = logger('canvas');
+
+// Lazy so the three.js + Spark chunk is only fetched when the overlay opens (and a load failure there
+// can't blank the rest of the app). `type`-only import of SplatScene above is erased, so it stays lazy.
+const SplatViewer = lazy(() => import('features/controlLayers/components/SplatOverlay/SplatViewer'));
+
+// Chrome sizes in screen px (divided by stage scale so they're constant at any zoom).
+const OUTLINE_WIDTH_PX = 2;
+const EDGE_HIT_PX = 10;
+const HANDLE_SIZE_PX = 12;
+const MIN_RECT_SIZE = 16; // world px
+
+const CORNERS: { corner: RectCorner; cursor: string }[] = [
+ { corner: 'nw', cursor: 'nwse-resize' },
+ { corner: 'ne', cursor: 'nesw-resize' },
+ { corner: 'sw', cursor: 'nesw-resize' },
+ { corner: 'se', cursor: 'nwse-resize' },
+];
+
+type RotateMode = 'orbit' | 'object';
+
+const CONTROL_HINTS: Record = {
+ orbit: [
+ 'Drag to orbit view',
+ 'Right-drag to pan (moves the rotation point)',
+ 'Scroll to zoom',
+ 'Edges to move',
+ 'Corners to resize',
+ ],
+ object: [
+ 'Drag to rotate object',
+ 'Right-drag to pan (moves the rotation point)',
+ 'Scroll to zoom',
+ 'Edges to move',
+ 'Corners to resize',
+ ],
+};
+
+type DragState = {
+ pointerId: number;
+ mode: 'move' | RectCorner;
+ startPointer: { x: number; y: number };
+ startRect: SplatRect;
+};
+
+/**
+ * The in-canvas 3D (Gaussian-splat) overlay: a transparent three.js viewport pinned to a world-space
+ * footprint rect, so the splat composites live over the actual canvas content. Dragging inside the frame
+ * either orbits the view or rotates the object itself (any axis, including roll) — a toolbar toggle picks
+ * the mode. Drag the frame edges to move it; drag the corners to resize (shift = keep aspect). The stage
+ * stays interactive outside the frame, so you can pan/zoom the canvas for precise placement. Commit bakes
+ * the current framing to a new raster layer at the current rect.
+ */
+export const CanvasSplatOverlay = memo(() => {
+ const canvasManager = useCanvasManager();
+ const state = useStore($splatOverlay);
+ const stageAttrs = useStore(canvasManager.stage.$stageAttrs);
+ const dispatch = useAppDispatch();
+ const sceneRef = useRef(null);
+ const dragRef = useRef(null);
+ const [isCommitting, setIsCommitting] = useState(false);
+ const [rotateMode, setRotateMode] = useState('orbit');
+ // The scene remounts on asset/session changes, so the current mode is kept in a ref and re-applied
+ // whenever a scene arrives.
+ const rotateModeRef = useRef(rotateMode);
+
+ const onSceneReady = useCallback((scene: SplatScene | null) => {
+ sceneRef.current = scene;
+ scene?.setRotateObjectMode(rotateModeRef.current === 'object');
+ }, []);
+
+ const applyRotateMode = useCallback((mode: RotateMode) => {
+ rotateModeRef.current = mode;
+ setRotateMode(mode);
+ sceneRef.current?.setRotateObjectMode(mode === 'object');
+ }, []);
+ const onOrbitMode = useCallback(() => applyRotateMode('orbit'), [applyRotateMode]);
+ const onObjectMode = useCallback(() => applyRotateMode('object'), [applyRotateMode]);
+
+ const getWorldPoint = useCallback(
+ (e: ReactPointerEvent) => {
+ const containerRect = canvasManager.stage.container.getBoundingClientRect();
+ return {
+ x: (e.clientX - containerRect.left - stageAttrs.x) / stageAttrs.scale,
+ y: (e.clientY - containerRect.top - stageAttrs.y) / stageAttrs.scale,
+ };
+ },
+ [canvasManager.stage.container, stageAttrs.x, stageAttrs.y, stageAttrs.scale]
+ );
+
+ const startDrag = useCallback(
+ (e: ReactPointerEvent, mode: DragState['mode']) => {
+ const current = $splatOverlay.get();
+ if (e.button !== 0 || !current) {
+ return;
+ }
+ e.preventDefault();
+ e.stopPropagation();
+ dragRef.current = {
+ pointerId: e.pointerId,
+ mode,
+ startPointer: getWorldPoint(e),
+ startRect: current.rect,
+ };
+ e.currentTarget.setPointerCapture(e.pointerId);
+ },
+ [getWorldPoint]
+ );
+
+ const onDragMove = useCallback(
+ (e: ReactPointerEvent) => {
+ const drag = dragRef.current;
+ if (!drag || drag.pointerId !== e.pointerId) {
+ return;
+ }
+ e.preventDefault();
+ const point = getWorldPoint(e);
+ const dx = point.x - drag.startPointer.x;
+ const dy = point.y - drag.startPointer.y;
+ const next =
+ drag.mode === 'move'
+ ? moveRect(drag.startRect, dx, dy)
+ : resizeRectFromCorner(drag.startRect, drag.mode, dx, dy, {
+ keepAspect: e.shiftKey,
+ minSize: MIN_RECT_SIZE,
+ });
+ updateSplatOverlayRect(next);
+ },
+ [getWorldPoint]
+ );
+
+ const endDrag = useCallback((e: ReactPointerEvent) => {
+ const drag = dragRef.current;
+ if (!drag || drag.pointerId !== e.pointerId) {
+ return;
+ }
+ e.currentTarget.releasePointerCapture(e.pointerId);
+ dragRef.current = null;
+ }, []);
+
+ const handleCommit = useCallback(() => {
+ const scene = sceneRef.current;
+ const current = $splatOverlay.get();
+ if (!scene || current?.status !== 'ready') {
+ return;
+ }
+ const { rect } = current;
+ setIsCommitting(true);
+ void (async () => {
+ try {
+ // Render at the footprint the user framed, so the baked layer lands at the right size + position.
+ const blob = await scene.capture(Math.max(1, Math.round(rect.width)), Math.max(1, Math.round(rect.height)));
+ if (!blob) {
+ throw new Error('Failed to capture splat');
+ }
+ const file = new File([blob], 'splat.png', { type: 'image/png' });
+ const imageDTO = await uploadImage({ file, image_category: 'general', is_intermediate: true, silent: true });
+ dispatch(
+ rasterLayerAdded({
+ overrides: {
+ objects: [imageDTOToImageObject(imageDTO)],
+ position: { x: Math.floor(rect.x), y: Math.floor(rect.y) },
+ },
+ isSelected: true,
+ })
+ );
+ clearSplatOverlay();
+ } catch (error) {
+ log.error({ error: String(error) }, 'Failed to commit 3D layer to canvas');
+ setIsCommitting(false);
+ }
+ })();
+ }, [dispatch]);
+
+ // This component stays mounted across open/close (it just renders null when closed), so per-session UI
+ // state would persist between sessions. Reset it whenever we're not in an active "ready" session — i.e.
+ // after a successful commit (status -> null) or when a new conversion starts (status -> loading).
+ const status = state?.status;
+ useEffect(() => {
+ if (status !== 'ready') {
+ setIsCommitting(false);
+ rotateModeRef.current = 'orbit';
+ setRotateMode('orbit');
+ }
+ }, [status]);
+
+ // Escape cancels the session (capture phase so canvas hotkeys don't also react).
+ const isOpen = state !== null;
+ useEffect(() => {
+ if (!isOpen) {
+ return;
+ }
+ const onKeyDown = (e: KeyboardEvent) => {
+ if (e.key === 'Escape') {
+ e.stopPropagation();
+ clearSplatOverlay();
+ }
+ };
+ window.addEventListener('keydown', onKeyDown, true);
+ return () => {
+ window.removeEventListener('keydown', onKeyDown, true);
+ };
+ }, [isOpen]);
+
+ if (!state) {
+ return null;
+ }
+
+ const { rect } = state;
+ const scale = stageAttrs.scale || 1;
+ // Chrome sized in world units so it renders at a constant screen size regardless of stage zoom.
+ const outlineWidth = OUTLINE_WIDTH_PX / scale;
+ const edgeHit = EDGE_HIT_PX / scale;
+ const edgeHalf = edgeHit / 2;
+ const handleSize = HANDLE_SIZE_PX / scale;
+ const handleHalf = handleSize / 2;
+ const chromeEnabled = !isCommitting;
+
+ return (
+
+ {/* World-pinned layer: children are positioned in canvas/world coordinates. */}
+
+
+ {state.status === 'ready' && (
+
+
+
+
+
+ )}
+ {state.status === 'loading' && (
+
+
+
+ )}
+ {/* Dashed footprint outline */}
+
+ {/* Edge strips: drag to move the footprint. Rendered after the viewer so they win near the border. */}
+ {chromeEnabled &&
+ (
+ [
+ { key: 'n', style: { top: -edgeHalf, left: edgeHalf, right: edgeHalf, height: edgeHit } },
+ { key: 's', style: { bottom: -edgeHalf, left: edgeHalf, right: edgeHalf, height: edgeHit } },
+ { key: 'w', style: { left: -edgeHalf, top: edgeHalf, bottom: edgeHalf, width: edgeHit } },
+ { key: 'e', style: { right: -edgeHalf, top: edgeHalf, bottom: edgeHalf, width: edgeHit } },
+ ] as const
+ ).map(({ key, style }) => (
+ startDrag(e, 'move')}
+ onPointerMove={onDragMove}
+ onPointerUp={endDrag}
+ onPointerCancel={endDrag}
+ />
+ ))}
+ {/* Corner handles: drag to resize (shift = keep aspect). */}
+ {chromeEnabled &&
+ CORNERS.map(({ corner, cursor }) => (
+ startDrag(e, corner)}
+ onPointerMove={onDragMove}
+ onPointerUp={endDrag}
+ onPointerCancel={endDrag}
+ />
+ ))}
+
+
+ {/* Screen-space toolbar: always visible regardless of pan/zoom. Wraps on narrow panels — the hint
+ breaks between segments and the button pair drops to its own row rather than clipping. */}
+
+ {state.status === 'loading' ? (
+
+
+
+ Generating 3D…
+
+
+ ) : (
+ <>
+
+
+
+
+
+ {CONTROL_HINTS[rotateMode].map((hint, i) => (
+
+ {hint}
+ {i < CONTROL_HINTS[rotateMode].length - 1 ? ' ·' : ''}
+
+ ))}
+
+
+ )}
+
+
+
+
+
+
+ );
+});
+CanvasSplatOverlay.displayName = 'CanvasSplatOverlay';
diff --git a/invokeai/frontend/web/src/features/controlLayers/components/SplatOverlay/SplatViewer.tsx b/invokeai/frontend/web/src/features/controlLayers/components/SplatOverlay/SplatViewer.tsx
new file mode 100644
index 00000000000..009b7620cf3
--- /dev/null
+++ b/invokeai/frontend/web/src/features/controlLayers/components/SplatOverlay/SplatViewer.tsx
@@ -0,0 +1,61 @@
+import { Box } from '@invoke-ai/ui-library';
+import { logger } from 'app/logging/logger';
+import { SplatScene } from 'features/controlLayers/components/SplatOverlay/splatScene';
+import { memo, useCallback, useEffect, useRef, useState } from 'react';
+
+const log = logger('canvas');
+
+type Props = {
+ assetUrl: string;
+ /** The canvas stage scale, folded into the renderer's pixel ratio so the viewport stays crisp at any zoom. */
+ stageScale: number;
+ /** Hand the live SplatScene to the parent so it can capture on commit; called with null on unmount. */
+ onSceneReady: (scene: SplatScene | null) => void;
+};
+
+/**
+ * The heavy three.js + Spark viewer, filling the splat overlay's footprint frame. Loaded lazily (React.lazy)
+ * by CanvasSplatOverlay so the `three` / `@sparkjsdev/spark` chunk is only fetched when the 3D overlay opens.
+ * Default export required.
+ */
+const SplatViewer = ({ assetUrl, stageScale, onSceneReady }: Props) => {
+ const containerRef = useRef(null);
+ const [scene, setScene] = useState(null);
+
+ useEffect(() => {
+ const container = containerRef.current;
+ if (!container) {
+ return;
+ }
+ const scene = new SplatScene(container);
+ setScene(scene);
+ onSceneReady(scene);
+ scene.loadFromUrl(assetUrl).catch((error: unknown) => {
+ log.error({ error: String(error) }, 'Failed to load splat');
+ });
+ return () => {
+ onSceneReady(null);
+ setScene(null);
+ scene.dispose();
+ };
+ }, [assetUrl, onSceneReady]);
+
+ useEffect(() => {
+ scene?.setStageScale(stageScale);
+ }, [scene, stageScale]);
+
+ const onPointerEnter = useCallback(() => scene?.setGizmoVisible(true), [scene]);
+ const onPointerLeave = useCallback(() => scene?.setGizmoVisible(false), [scene]);
+
+ return (
+
+ );
+};
+
+export default memo(SplatViewer);
diff --git a/invokeai/frontend/web/src/features/controlLayers/components/SplatOverlay/rectTransforms.test.ts b/invokeai/frontend/web/src/features/controlLayers/components/SplatOverlay/rectTransforms.test.ts
new file mode 100644
index 00000000000..ca84d486e61
--- /dev/null
+++ b/invokeai/frontend/web/src/features/controlLayers/components/SplatOverlay/rectTransforms.test.ts
@@ -0,0 +1,72 @@
+import { describe, expect, it } from 'vitest';
+
+import { moveRect, resizeRectFromCorner } from './rectTransforms';
+
+const START = { x: 100, y: 200, width: 400, height: 300 };
+const OPTS = { keepAspect: false, minSize: 16 };
+
+describe('moveRect', () => {
+ it('translates by the delta without changing size', () => {
+ expect(moveRect(START, 25, -50)).toEqual({ x: 125, y: 150, width: 400, height: 300 });
+ });
+});
+
+describe('resizeRectFromCorner', () => {
+ it('se drag grows right/down, anchoring the top-left corner', () => {
+ expect(resizeRectFromCorner(START, 'se', 40, 30, OPTS)).toEqual({ x: 100, y: 200, width: 440, height: 330 });
+ });
+
+ it('nw drag anchors the bottom-right corner', () => {
+ const result = resizeRectFromCorner(START, 'nw', -40, -30, OPTS);
+ expect(result).toEqual({ x: 60, y: 170, width: 440, height: 330 });
+ // Bottom-right corner unchanged
+ expect(result.x + result.width).toBe(START.x + START.width);
+ expect(result.y + result.height).toBe(START.y + START.height);
+ });
+
+ it('ne drag anchors the bottom-left corner', () => {
+ const result = resizeRectFromCorner(START, 'ne', 40, -30, OPTS);
+ expect(result).toEqual({ x: 100, y: 170, width: 440, height: 330 });
+ expect(result.x).toBe(START.x);
+ expect(result.y + result.height).toBe(START.y + START.height);
+ });
+
+ it('sw drag anchors the top-right corner', () => {
+ const result = resizeRectFromCorner(START, 'sw', -40, 30, OPTS);
+ expect(result).toEqual({ x: 60, y: 200, width: 440, height: 330 });
+ expect(result.x + result.width).toBe(START.x + START.width);
+ expect(result.y).toBe(START.y);
+ });
+
+ it('clamps each axis to minSize independently', () => {
+ const result = resizeRectFromCorner(START, 'se', -1000, -1000, OPTS);
+ expect(result.width).toBe(16);
+ expect(result.height).toBe(16);
+ // Anchored corner stays put even when clamped
+ expect(result.x).toBe(START.x);
+ expect(result.y).toBe(START.y);
+ });
+
+ it('nw clamp keeps the bottom-right corner anchored', () => {
+ const result = resizeRectFromCorner(START, 'nw', 1000, 1000, OPTS);
+ expect(result.width).toBe(16);
+ expect(result.height).toBe(16);
+ expect(result.x + result.width).toBe(START.x + START.width);
+ expect(result.y + result.height).toBe(START.y + START.height);
+ });
+
+ it('keepAspect scales uniformly from the dominant axis', () => {
+ // dx is the dominant relative change: 400 -> 600 is 1.5x
+ const result = resizeRectFromCorner(START, 'se', 200, 10, { keepAspect: true, minSize: 16 });
+ expect(result.width).toBeCloseTo(600);
+ expect(result.height).toBeCloseTo(450);
+ expect(result.width / result.height).toBeCloseTo(START.width / START.height);
+ });
+
+ it('keepAspect respects minSize on the smaller axis', () => {
+ const result = resizeRectFromCorner(START, 'se', -1000, -1000, { keepAspect: true, minSize: 16 });
+ expect(result.width).toBeGreaterThanOrEqual(16);
+ expect(result.height).toBeGreaterThanOrEqual(16);
+ expect(result.width / result.height).toBeCloseTo(START.width / START.height);
+ });
+});
diff --git a/invokeai/frontend/web/src/features/controlLayers/components/SplatOverlay/rectTransforms.ts b/invokeai/frontend/web/src/features/controlLayers/components/SplatOverlay/rectTransforms.ts
new file mode 100644
index 00000000000..1df0e423b6b
--- /dev/null
+++ b/invokeai/frontend/web/src/features/controlLayers/components/SplatOverlay/rectTransforms.ts
@@ -0,0 +1,56 @@
+import type { SplatRect } from 'features/controlLayers/components/SplatOverlay/state';
+
+export type RectCorner = 'nw' | 'ne' | 'sw' | 'se';
+
+type ResizeOptions = {
+ /** Preserve the starting rect's aspect ratio (dominant-axis scale). */
+ keepAspect: boolean;
+ /** Minimum width/height in world px. */
+ minSize: number;
+};
+
+/** Translate a rect by a world-space delta. */
+export const moveRect = (start: SplatRect, dx: number, dy: number): SplatRect => ({
+ ...start,
+ x: start.x + dx,
+ y: start.y + dy,
+});
+
+/**
+ * Resize a rect by dragging one of its corners by a world-space delta, keeping the opposite corner
+ * anchored. With `keepAspect`, the axis with the larger relative change drives a uniform scale.
+ */
+export const resizeRectFromCorner = (
+ start: SplatRect,
+ corner: RectCorner,
+ dx: number,
+ dy: number,
+ options: ResizeOptions
+): SplatRect => {
+ const { keepAspect, minSize } = options;
+ const signX = corner === 'nw' || corner === 'sw' ? -1 : 1;
+ const signY = corner === 'nw' || corner === 'ne' ? -1 : 1;
+
+ let width = start.width + signX * dx;
+ let height = start.height + signY * dy;
+
+ if (keepAspect && start.width > 0 && start.height > 0) {
+ const scaleX = width / start.width;
+ const scaleY = height / start.height;
+ // The axis the user pulled further (relative to the rect's size) wins.
+ let scale = Math.abs(scaleX - 1) >= Math.abs(scaleY - 1) ? scaleX : scaleY;
+ scale = Math.max(scale, minSize / start.width, minSize / start.height);
+ width = start.width * scale;
+ height = start.height * scale;
+ } else {
+ width = Math.max(minSize, width);
+ height = Math.max(minSize, height);
+ }
+
+ return {
+ x: signX === -1 ? start.x + start.width - width : start.x,
+ y: signY === -1 ? start.y + start.height - height : start.y,
+ width,
+ height,
+ };
+};
diff --git a/invokeai/frontend/web/src/features/controlLayers/components/SplatOverlay/splatScene.ts b/invokeai/frontend/web/src/features/controlLayers/components/SplatOverlay/splatScene.ts
new file mode 100644
index 00000000000..ebac9c66b79
--- /dev/null
+++ b/invokeai/frontend/web/src/features/controlLayers/components/SplatOverlay/splatScene.ts
@@ -0,0 +1,445 @@
+import { SparkRenderer, SplatMesh } from '@sparkjsdev/spark';
+import * as THREE from 'three';
+import { OrbitControls } from 'three/examples/jsm/controls/OrbitControls.js';
+import { ViewHelper } from 'three/examples/jsm/helpers/ViewHelper.js';
+
+import { remapPointerForViewHelper, VIEW_HELPER_DIM } from './viewHelperPointer';
+
+// Cap the drawing buffer's largest dimension so deep stage zooms don't allocate absurd buffers.
+const MAX_DRAWING_BUFFER_DIM = 4096;
+
+// Object-rotation speed in rotate-object mode, radians per screen px (~0.46°/px: 200px ≈ a quarter turn).
+const OBJECT_ROTATE_SPEED = 0.008;
+
+// Pivot marker: shown while a gesture is active, at the point the gesture rotates around.
+const PIVOT_MARKER_COLOR = 0x66b2ff; // ≈ invokeBlue.300, matches the overlay chrome
+const PIVOT_MARKER_SCREEN_SCALE = 0.008; // world size per unit of camera distance ≈ constant screen size
+const PIVOT_FADE_RATE = 8; // opacity lerp rate per second (~0.3s fade)
+
+const buildPivotMarker = (): {
+ group: THREE.Group;
+ materials: (THREE.MeshBasicMaterial | THREE.LineBasicMaterial)[];
+} => {
+ const dotMaterial = new THREE.MeshBasicMaterial({
+ color: PIVOT_MARKER_COLOR,
+ transparent: true,
+ opacity: 0,
+ depthTest: false,
+ depthWrite: false,
+ });
+ const lineMaterial = new THREE.LineBasicMaterial({
+ color: PIVOT_MARKER_COLOR,
+ transparent: true,
+ opacity: 0,
+ depthTest: false,
+ depthWrite: false,
+ });
+ const dot = new THREE.Mesh(new THREE.SphereGeometry(0.6, 12, 8), dotMaterial);
+ const cross = new THREE.LineSegments(
+ new THREE.BufferGeometry().setFromPoints([
+ new THREE.Vector3(-2.4, 0, 0),
+ new THREE.Vector3(2.4, 0, 0),
+ new THREE.Vector3(0, -2.4, 0),
+ new THREE.Vector3(0, 2.4, 0),
+ new THREE.Vector3(0, 0, -2.4),
+ new THREE.Vector3(0, 0, 2.4),
+ ]),
+ lineMaterial
+ );
+ // Drawn on top (no depth test, late render order): the pivot usually sits inside the splat body,
+ // which would otherwise occlude it.
+ dot.renderOrder = 999;
+ cross.renderOrder = 999;
+ const group = new THREE.Group();
+ group.add(dot, cross);
+ group.visible = false;
+ return { group, materials: [dotMaterial, lineMaterial] };
+};
+
+/**
+ * A transparent three.js + Spark viewport for a single 3D Gaussian splat, designed to sit directly on the
+ * canvas: orbit/zoom/pan camera and a clickable corner navigation gizmo (ViewHelper, shown only while the
+ * pointer is over the viewport). A toggleable rotate-object mode (setRotateObjectMode) makes left-drag
+ * rotate the object itself (all three axes) instead of orbiting — OrbitControls keeps the horizon level,
+ * so poses like a diagonal lean are unreachable by camera orbit alone. Every gesture — orbit, pan, zoom,
+ * and rotate-object — pivots around the same point (the orbit target, movable by panning), and while a
+ * gesture is active a crosshair marker fades in at that point; it is never rendered into captures. The
+ * background is always transparent
+ * so the canvas shows through — the live view is the compositing preview. Capture renders the object alone
+ * at a given pixel size.
+ *
+ * The element hosting this scene is CSS-scaled by the canvas stage transform, so the drawing buffer is
+ * sized at (layout px × devicePixelRatio × stage scale) to stay crisp at any zoom — see setStageScale.
+ *
+ * The corner gizmo requires the renderer's `autoClear` to be off (it draws over the main render via a corner
+ * viewport), so the main pass clears manually each frame.
+ *
+ * Orientation: TripoSplat .ply is -Y-up (3DGS). A parent group is flipped 180° about X so the object stands
+ * upright (+Y up); the splat is yawed 90° so its front faces the camera.
+ */
+export class SplatScene {
+ private readonly container: HTMLElement;
+ private readonly renderer: THREE.WebGLRenderer;
+ private readonly scene: THREE.Scene;
+ private readonly camera: THREE.PerspectiveCamera;
+ private readonly controls: OrbitControls;
+ private readonly spark: SparkRenderer;
+ private readonly viewHelper: ViewHelper;
+ private readonly clock: THREE.Clock;
+ private readonly splatRoot: THREE.Group;
+ private readonly resizeObserver: ResizeObserver;
+ private readonly onPointerDown: (e: PointerEvent) => void;
+ private readonly onPointerUp: (e: PointerEvent) => void;
+ private readonly onObjectRotateStart: (e: PointerEvent) => void;
+ private readonly onObjectRotateMove: (e: PointerEvent) => void;
+ private readonly onObjectRotateEnd: (e: PointerEvent) => void;
+ private readonly onContextMenu: (e: MouseEvent) => void;
+ private pointerDownPos: { x: number; y: number } | null = null;
+ private objectRotate: { pointerId: number; lastX: number; lastY: number } | null = null;
+ private rotateObjectMode = false;
+ private readonly pivotMarker: THREE.Group;
+ private readonly pivotMaterials: (THREE.MeshBasicMaterial | THREE.LineBasicMaterial)[];
+ private pivotOpacity = 0;
+ private pivotTargetOpacity = 0;
+ private mesh: SplatMesh | null = null;
+ private stageScale = 1;
+ private gizmoVisible = false;
+ private disposed = false;
+
+ constructor(container: HTMLElement) {
+ this.container = container;
+ this.clock = new THREE.Clock();
+
+ this.scene = new THREE.Scene();
+
+ this.camera = new THREE.PerspectiveCamera(45, 1, 0.01, 1000);
+ this.camera.position.set(0, 0, 2.2);
+
+ this.renderer = new THREE.WebGLRenderer({
+ alpha: true,
+ antialias: true,
+ preserveDrawingBuffer: true, // required so capture() can read the rendered frame
+ });
+ this.renderer.autoClear = false; // ViewHelper draws over the main render; we clear manually
+ this.renderer.setPixelRatio(window.devicePixelRatio || 1);
+ this.renderer.domElement.style.width = '100%';
+ this.renderer.domElement.style.height = '100%';
+ this.renderer.domElement.style.display = 'block';
+ container.appendChild(this.renderer.domElement);
+
+ this.spark = new SparkRenderer({ renderer: this.renderer });
+ this.scene.add(this.spark);
+
+ this.splatRoot = new THREE.Group();
+ this.splatRoot.rotation.x = Math.PI;
+ this.scene.add(this.splatRoot);
+
+ this.controls = new OrbitControls(this.camera, this.renderer.domElement);
+ this.controls.enableDamping = true;
+ this.controls.dampingFactor = 0.08;
+ this.controls.minDistance = 0.3;
+ this.controls.maxDistance = 50;
+ this.controls.target.set(0, 0, 0);
+
+ this.viewHelper = new ViewHelper(this.camera, this.renderer.domElement);
+ this.viewHelper.center = this.controls.target; // snap orbits around the same point as OrbitControls
+
+ // Pivot marker: fades in while the user orbits/pans/zooms (or rotates the object) so the rotation
+ // center is visible; tick() keeps it positioned, screen-sized, and faded.
+ const pivot = buildPivotMarker();
+ this.pivotMarker = pivot.group;
+ this.pivotMaterials = pivot.materials;
+ this.scene.add(this.pivotMarker);
+ this.controls.addEventListener('start', this.showPivotMarker);
+ this.controls.addEventListener('end', this.hidePivotMarker);
+
+ // Only treat a near-stationary pointerup as a gizmo click (so orbiting doesn't accidentally snap).
+ this.onPointerDown = (e) => {
+ this.pointerDownPos = { x: e.clientX, y: e.clientY };
+ };
+ this.onPointerUp = (e) => {
+ const down = this.pointerDownPos;
+ this.pointerDownPos = null;
+ if (down && this.gizmoVisible && Math.hypot(e.clientX - down.x, e.clientY - down.y) < 5) {
+ // ViewHelper's own hit test breaks while the stage transform CSS-scales this viewport (it mixes
+ // visual and layout coordinates) — remap so axis clicks land at any stage zoom.
+ const el = this.renderer.domElement;
+ const remapped = remapPointerForViewHelper(
+ { rect: el.getBoundingClientRect(), offsetWidth: el.offsetWidth, offsetHeight: el.offsetHeight },
+ e.clientX,
+ e.clientY
+ );
+ this.viewHelper.handleClick(remapped as PointerEvent);
+ }
+ };
+ this.renderer.domElement.addEventListener('pointerdown', this.onPointerDown);
+ this.renderer.domElement.addEventListener('pointerup', this.onPointerUp);
+
+ // In rotate-object mode, left-drag rotates the object itself about its center on the camera's screen
+ // axes. Registered on the container in the capture phase so it claims the gesture before OrbitControls'
+ // own pointerdown runs; wheel-zoom and right-drag pan still reach OrbitControls, and clicks on the
+ // corner gizmo pass through so view snapping keeps working.
+ this.onObjectRotateStart = (e) => {
+ if (!this.rotateObjectMode || e.button !== 0 || this.disposed || this.objectRotate || this.isOverGizmo(e)) {
+ return;
+ }
+ e.preventDefault();
+ e.stopPropagation();
+ this.objectRotate = { pointerId: e.pointerId, lastX: e.clientX, lastY: e.clientY };
+ this.showPivotMarker();
+ window.addEventListener('pointermove', this.onObjectRotateMove);
+ window.addEventListener('pointerup', this.onObjectRotateEnd);
+ window.addEventListener('pointercancel', this.onObjectRotateEnd);
+ };
+ this.onObjectRotateMove = (e) => {
+ const rot = this.objectRotate;
+ if (!rot || e.pointerId !== rot.pointerId) {
+ return;
+ }
+ const dx = e.clientX - rot.lastX;
+ const dy = e.clientY - rot.lastY;
+ rot.lastX = e.clientX;
+ rot.lastY = e.clientY;
+ // Rotate about the camera's current screen axes so the object follows the pointer from any view
+ // angle; successive two-axis increments compose to reach any orientation, including roll. The
+ // rotation is applied about the orbit target — the same pivot that pan moves and orbit uses — so
+ // after panning, the object swings around the panned point instead of always spinning in place.
+ const right = new THREE.Vector3().setFromMatrixColumn(this.camera.matrixWorld, 0);
+ const up = new THREE.Vector3().setFromMatrixColumn(this.camera.matrixWorld, 1);
+ const q = new THREE.Quaternion()
+ .setFromAxisAngle(right, dy * OBJECT_ROTATE_SPEED)
+ .multiply(new THREE.Quaternion().setFromAxisAngle(up, dx * OBJECT_ROTATE_SPEED));
+ const pivot = this.controls.target;
+ this.splatRoot.quaternion.premultiply(q);
+ this.splatRoot.position.sub(pivot).applyQuaternion(q).add(pivot);
+ };
+ this.onObjectRotateEnd = (e) => {
+ if (!this.objectRotate || e.pointerId !== this.objectRotate.pointerId) {
+ return;
+ }
+ this.objectRotate = null;
+ this.hidePivotMarker();
+ window.removeEventListener('pointermove', this.onObjectRotateMove);
+ window.removeEventListener('pointerup', this.onObjectRotateEnd);
+ window.removeEventListener('pointercancel', this.onObjectRotateEnd);
+ };
+ container.addEventListener('pointerdown', this.onObjectRotateStart, true);
+
+ // Right-drag pans the camera, and contextmenu fires on right-button RELEASE — without this, every pan
+ // ends by popping the canvas context menu. preventDefault kills the browser menu; stopPropagation keeps
+ // the event from bubbling to the app's canvas context-menu trigger.
+ this.onContextMenu = (e) => {
+ e.preventDefault();
+ e.stopPropagation();
+ };
+ container.addEventListener('contextmenu', this.onContextMenu);
+
+ this.resizeObserver = new ResizeObserver(this.resize);
+ this.resizeObserver.observe(container);
+ this.resize();
+
+ this.renderer.setAnimationLoop(this.tick);
+ }
+
+ private tick = (): void => {
+ if (this.disposed) {
+ return;
+ }
+ const delta = this.clock.getDelta();
+ if (this.viewHelper.animating) {
+ this.viewHelper.update(delta); // drive the camera during a snap; don't let OrbitControls fight it
+ } else {
+ this.controls.update();
+ }
+ // Pivot marker: track the active gesture's rotation center at a constant screen size, fading around
+ // interactions.
+ this.pivotOpacity += (this.pivotTargetOpacity - this.pivotOpacity) * Math.min(1, delta * PIVOT_FADE_RATE);
+ if (this.pivotTargetOpacity === 0 && this.pivotOpacity < 0.02) {
+ this.pivotOpacity = 0;
+ }
+ this.pivotMarker.visible = this.pivotOpacity > 0;
+ if (this.pivotMarker.visible) {
+ // Every gesture — orbit, pan, zoom, and rotate-object — pivots around the same controls.target.
+ const pivot = this.controls.target;
+ this.pivotMarker.position.copy(pivot);
+ this.pivotMarker.scale.setScalar(
+ Math.max(1e-6, this.camera.position.distanceTo(pivot) * PIVOT_MARKER_SCREEN_SCALE)
+ );
+ for (const material of this.pivotMaterials) {
+ material.opacity = this.pivotOpacity * 0.9;
+ }
+ }
+ this.renderer.setClearColor(0x000000, 0); // transparent — the canvas beneath is the background
+ this.renderer.clear();
+ this.renderer.render(this.scene, this.camera);
+ if (this.gizmoVisible) {
+ this.viewHelper.render(this.renderer);
+ }
+ };
+
+ /**
+ * The canvas stage scale currently applied (via CSS transform) to this viewport's container. Folded into
+ * the renderer's pixel ratio so the drawing buffer matches on-screen pixels — crisp at any zoom.
+ */
+ setStageScale = (scale: number): void => {
+ const next = Math.max(0.001, scale || 1);
+ if (next === this.stageScale) {
+ return;
+ }
+ this.stageScale = next;
+ this.resize();
+ };
+
+ /** Show/hide the corner navigation gizmo (shown only while the pointer is over the viewport). */
+ setGizmoVisible = (visible: boolean): void => {
+ this.gizmoVisible = visible;
+ };
+
+ /** When enabled, left-drag rotates the object itself (any axis) instead of orbiting the camera. */
+ setRotateObjectMode = (enabled: boolean): void => {
+ this.rotateObjectMode = enabled;
+ };
+
+ private showPivotMarker = (): void => {
+ this.pivotTargetOpacity = 1;
+ };
+
+ private hidePivotMarker = (): void => {
+ this.pivotTargetOpacity = 0;
+ };
+
+ /** Whether the pointer is over the ViewHelper's corner viewport (compared in the element's layout space). */
+ private isOverGizmo(e: PointerEvent): boolean {
+ if (!this.gizmoVisible) {
+ return false;
+ }
+ const el = this.renderer.domElement;
+ const rect = el.getBoundingClientRect();
+ const scaleX = rect.width / (el.offsetWidth || 1) || 1;
+ const scaleY = rect.height / (el.offsetHeight || 1) || 1;
+ const layoutX = (e.clientX - rect.left) / scaleX;
+ const layoutY = (e.clientY - rect.top) / scaleY;
+ return layoutX >= el.offsetWidth - VIEW_HELPER_DIM && layoutY >= el.offsetHeight - VIEW_HELPER_DIM;
+ }
+
+ private effectivePixelRatio(w: number, h: number): number {
+ const dpr = window.devicePixelRatio || 1;
+ // Quantize so wheel-zooming the stage doesn't reallocate the drawing buffer on every tick.
+ const quantized = Math.max(0.25, Math.round(dpr * this.stageScale * 4) / 4);
+ return Math.min(quantized, MAX_DRAWING_BUFFER_DIM / Math.max(w, h, 1));
+ }
+
+ resize = (): void => {
+ const w = this.container.clientWidth || 1;
+ const h = this.container.clientHeight || 1;
+ this.renderer.setPixelRatio(this.effectivePixelRatio(w, h));
+ this.renderer.setSize(w, h, false);
+ this.camera.aspect = w / h;
+ this.camera.updateProjectionMatrix();
+ };
+
+ /** Position the camera for a front-on default view that roughly fills the frame; user can orbit from there. */
+ private frameObject(): void {
+ const halfExtent = 0.5;
+ const margin = 1.1;
+ const aspect = this.camera.aspect || 1;
+ const vfov = (this.camera.fov * Math.PI) / 180;
+ const dv = (halfExtent * margin) / Math.tan(vfov / 2);
+ const dist = Math.max(dv, dv / aspect);
+ this.camera.position.set(0, 0, dist);
+ this.camera.updateProjectionMatrix();
+ this.controls.target.set(0, 0, 0);
+ this.controls.update();
+ }
+
+ async loadFromUrl(url: string): Promise {
+ const res = await fetch(url);
+ if (!res.ok) {
+ throw new Error(`Failed to fetch splat: HTTP ${res.status}`);
+ }
+ const bytes = await res.arrayBuffer();
+ if (this.disposed) {
+ return;
+ }
+ const mesh = new SplatMesh({ fileBytes: bytes, fileName: url.split('/').pop() || 'model.ply' });
+ await mesh.initialized;
+ if (this.disposed) {
+ this.disposeMesh(mesh);
+ return;
+ }
+ if (this.mesh) {
+ this.splatRoot.remove(this.mesh);
+ this.disposeMesh(this.mesh);
+ }
+ mesh.rotation.y = Math.PI / 2;
+ this.mesh = mesh;
+ this.splatRoot.add(mesh);
+ this.frameObject();
+ }
+
+ /**
+ * Capture the object alone (no gizmo, transparent background) at the given pixel size, from the user's
+ * current camera (WYSIWYG). Renders at pixelRatio 1 so the blob's intrinsic size equals width×height. The
+ * live render loop is paused during capture to avoid it overwriting the frame before toBlob reads it.
+ */
+ async capture(width: number, height: number): Promise {
+ this.renderer.setAnimationLoop(null);
+ this.pivotMarker.visible = false; // never bake the pivot marker into the layer; tick() restores it
+ this.renderer.setPixelRatio(1);
+ this.renderer.setSize(width, height, false);
+ this.camera.aspect = width / height;
+ this.camera.updateProjectionMatrix();
+ this.renderer.setClearColor(0x000000, 0); // transparent
+ this.renderer.clear();
+ this.renderer.render(this.scene, this.camera);
+ const blob = await new Promise((resolve) => {
+ this.renderer.domElement.toBlob(resolve, 'image/png');
+ });
+ // Restore the live preview (resize() also restores the stage-scale-aware pixel ratio).
+ this.resize();
+ if (!this.disposed) {
+ this.renderer.setAnimationLoop(this.tick);
+ }
+ return blob;
+ }
+
+ private disposeMesh(mesh: SplatMesh): void {
+ (mesh as THREE.Object3D & { dispose?: () => void }).dispose?.();
+ }
+
+ dispose(): void {
+ this.disposed = true;
+ this.renderer.setAnimationLoop(null);
+ this.resizeObserver.disconnect();
+ this.renderer.domElement.removeEventListener('pointerdown', this.onPointerDown);
+ this.renderer.domElement.removeEventListener('pointerup', this.onPointerUp);
+ this.container.removeEventListener('pointerdown', this.onObjectRotateStart, true);
+ this.container.removeEventListener('contextmenu', this.onContextMenu);
+ window.removeEventListener('pointermove', this.onObjectRotateMove);
+ window.removeEventListener('pointerup', this.onObjectRotateEnd);
+ window.removeEventListener('pointercancel', this.onObjectRotateEnd);
+ this.controls.removeEventListener('start', this.showPivotMarker);
+ this.controls.removeEventListener('end', this.hidePivotMarker);
+ this.scene.remove(this.pivotMarker);
+ this.pivotMarker.traverse((obj) => {
+ if (obj instanceof THREE.Mesh || obj instanceof THREE.LineSegments) {
+ obj.geometry.dispose();
+ }
+ });
+ for (const material of this.pivotMaterials) {
+ material.dispose();
+ }
+ this.controls.dispose();
+ this.viewHelper.dispose();
+ if (this.mesh) {
+ this.splatRoot.remove(this.mesh);
+ this.disposeMesh(this.mesh);
+ this.mesh = null;
+ }
+ // SparkRenderer owns sort/LOD web workers and GPU-side state that only its dispose() releases.
+ this.scene.remove(this.spark);
+ this.spark.dispose();
+ this.renderer.dispose();
+ this.renderer.domElement.remove();
+ }
+}
diff --git a/invokeai/frontend/web/src/features/controlLayers/components/SplatOverlay/state.ts b/invokeai/frontend/web/src/features/controlLayers/components/SplatOverlay/state.ts
new file mode 100644
index 00000000000..604da6ae399
--- /dev/null
+++ b/invokeai/frontend/web/src/features/controlLayers/components/SplatOverlay/state.ts
@@ -0,0 +1,43 @@
+import { atom } from 'nanostores';
+
+export type SplatRect = { x: number; y: number; width: number; height: number };
+
+type SplatOverlayState =
+ | { status: 'loading'; sessionId: string; rect: SplatRect }
+ | { status: 'ready'; sessionId: string; assetUrl: string; rect: SplatRect };
+
+/**
+ * Transient state for the in-canvas 3D (Gaussian-splat) overlay. `null` means the overlay is closed.
+ * The `rect` is the overlay's footprint in canvas/world coords — seeded from the source raster layer's
+ * bbox, then moved/resized by the user before committing. The bake step renders at this rect.
+ */
+export const $splatOverlay = atom(null);
+
+/** Move/resize the overlay footprint. No-op if the overlay is closed. */
+export const updateSplatOverlayRect = (rect: SplatRect): void => {
+ const state = $splatOverlay.get();
+ if (!state) {
+ return;
+ }
+ $splatOverlay.set({ ...state, rect });
+};
+
+// Tracks the in-flight generation so closing/cancelling aborts the backend run, not just the overlay UI.
+let activeGenerationAbort: AbortController | null = null;
+
+export const setSplatGenerationAbort = (controller: AbortController | null): void => {
+ activeGenerationAbort = controller;
+};
+
+/** Clears the abort slot only if it still holds `controller` — a newer session may have taken it over. */
+export const clearSplatGenerationAbort = (controller: AbortController): void => {
+ if (activeGenerationAbort === controller) {
+ activeGenerationAbort = null;
+ }
+};
+
+export const clearSplatOverlay = (): void => {
+ activeGenerationAbort?.abort();
+ activeGenerationAbort = null;
+ $splatOverlay.set(null);
+};
diff --git a/invokeai/frontend/web/src/features/controlLayers/components/SplatOverlay/viewHelperPointer.test.ts b/invokeai/frontend/web/src/features/controlLayers/components/SplatOverlay/viewHelperPointer.test.ts
new file mode 100644
index 00000000000..3252791ec0b
--- /dev/null
+++ b/invokeai/frontend/web/src/features/controlLayers/components/SplatOverlay/viewHelperPointer.test.ts
@@ -0,0 +1,75 @@
+import { describe, expect, it } from 'vitest';
+
+import type { ElementBox } from './viewHelperPointer';
+import { remapPointerForViewHelper, VIEW_HELPER_DIM } from './viewHelperPointer';
+
+/** Oracle: the exact NDC mapping three's ViewHelper.handleClick applies to a pointer event. */
+const viewHelperNdc = (box: ElementBox, clientX: number, clientY: number, dim = VIEW_HELPER_DIM) => {
+ const { rect, offsetWidth, offsetHeight } = box;
+ const offsetX = rect.left + (offsetWidth - dim);
+ const offsetY = rect.top + (offsetHeight - dim);
+ return {
+ x: ((clientX - offsetX) / (rect.right - offsetX)) * 2 - 1,
+ y: -((clientY - offsetY) / (rect.bottom - offsetY)) * 2 + 1,
+ };
+};
+
+const makeBox = (left: number, top: number, offsetWidth: number, offsetHeight: number, scale: number): ElementBox => ({
+ rect: {
+ left,
+ top,
+ right: left + offsetWidth * scale,
+ bottom: top + offsetHeight * scale,
+ width: offsetWidth * scale,
+ height: offsetHeight * scale,
+ },
+ offsetWidth,
+ offsetHeight,
+});
+
+/** Visual (client) position of a point given in the element's layout space. */
+const visualPos = (box: ElementBox, layoutX: number, layoutY: number, scale: number) => ({
+ clientX: box.rect.left + layoutX * scale,
+ clientY: box.rect.top + layoutY * scale,
+});
+
+describe('remapPointerForViewHelper', () => {
+ it('is the identity when the element is unscaled', () => {
+ const box = makeBox(10, 20, 400, 300, 1);
+ const remapped = remapPointerForViewHelper(box, 350, 250);
+ expect(remapped.clientX).toBeCloseTo(350);
+ expect(remapped.clientY).toBeCloseTo(250);
+ });
+
+ it.each([2, 0.5, 1.37])('maps a click on the gizmo center to NDC (0, 0) at scale %f', (scale) => {
+ const box = makeBox(10, 20, 400, 300, scale);
+ // Gizmo center in layout space: middle of the bottom-right dim×dim box.
+ const { clientX, clientY } = visualPos(box, 400 - VIEW_HELPER_DIM / 2, 300 - VIEW_HELPER_DIM / 2, scale);
+ const remapped = remapPointerForViewHelper(box, clientX, clientY);
+ const ndc = viewHelperNdc(box, remapped.clientX, remapped.clientY);
+ expect(ndc.x).toBeCloseTo(0);
+ expect(ndc.y).toBeCloseTo(0);
+ });
+
+ it('maps the gizmo corners to NDC extremes at scale 2', () => {
+ const box = makeBox(0, 0, 640, 480, 2);
+ const topLeft = visualPos(box, 640 - VIEW_HELPER_DIM, 480 - VIEW_HELPER_DIM, 2);
+ const bottomRight = visualPos(box, 640, 480, 2);
+ const remappedTl = remapPointerForViewHelper(box, topLeft.clientX, topLeft.clientY);
+ const remappedBr = remapPointerForViewHelper(box, bottomRight.clientX, bottomRight.clientY);
+ const ndcTl = viewHelperNdc(box, remappedTl.clientX, remappedTl.clientY);
+ const ndcBr = viewHelperNdc(box, remappedBr.clientX, remappedBr.clientY);
+ expect(ndcTl.x).toBeCloseTo(-1);
+ expect(ndcTl.y).toBeCloseTo(1);
+ expect(ndcBr.x).toBeCloseTo(1);
+ expect(ndcBr.y).toBeCloseTo(-1);
+ });
+
+ it('documents the bug: without remapping, a scaled gizmo-center click drifts far from its true NDC', () => {
+ const box = makeBox(10, 20, 400, 300, 2);
+ const { clientX, clientY } = visualPos(box, 400 - VIEW_HELPER_DIM / 2, 300 - VIEW_HELPER_DIM / 2, 2);
+ // The true position is the gizmo center, NDC (0, 0); unremapped it lands well off it, missing the axes.
+ const ndc = viewHelperNdc(box, clientX, clientY);
+ expect(Math.hypot(ndc.x, ndc.y)).toBeGreaterThan(0.3);
+ });
+});
diff --git a/invokeai/frontend/web/src/features/controlLayers/components/SplatOverlay/viewHelperPointer.ts b/invokeai/frontend/web/src/features/controlLayers/components/SplatOverlay/viewHelperPointer.ts
new file mode 100644
index 00000000000..a32825d89b3
--- /dev/null
+++ b/invokeai/frontend/web/src/features/controlLayers/components/SplatOverlay/viewHelperPointer.ts
@@ -0,0 +1,41 @@
+/** The fixed corner-viewport size (layout px) hardcoded as `dim` inside three's ViewHelper. */
+export const VIEW_HELPER_DIM = 128;
+
+export type ElementBox = {
+ /** getBoundingClientRect() of the renderer's canvas — visual coords, includes any ancestor CSS scale. */
+ rect: { left: number; top: number; right: number; bottom: number; width: number; height: number };
+ /** Layout (untransformed) size of the canvas element. */
+ offsetWidth: number;
+ offsetHeight: number;
+};
+
+/**
+ * three's ViewHelper.handleClick derives its raycast NDC by mixing visual coordinates
+ * (getBoundingClientRect, event.clientX/Y) with layout sizes (offsetWidth/offsetHeight and its fixed
+ * 128px `dim`). Those spaces only agree when the element is unscaled — but the splat viewport is
+ * CSS-scaled by the canvas stage transform, so at any stage zoom other than 100% the gizmo's click
+ * hotspot drifts off the drawn gizmo and axis clicks miss.
+ *
+ * This inverts ViewHelper's formula: it returns clientX/Y such that handleClick's mixed-space math
+ * yields the NDC of the pointer's true layout-space position. At scale 1 it is the identity.
+ */
+export const remapPointerForViewHelper = (
+ box: ElementBox,
+ clientX: number,
+ clientY: number,
+ dim: number = VIEW_HELPER_DIM
+): { clientX: number; clientY: number } => {
+ const { rect, offsetWidth, offsetHeight } = box;
+ const scaleX = rect.width / (offsetWidth || 1) || 1;
+ const scaleY = rect.height / (offsetHeight || 1) || 1;
+ // Pointer position in the element's layout space — where the click "really" is on the gizmo.
+ const layoutX = (clientX - rect.left) / scaleX;
+ const layoutY = (clientY - rect.top) / scaleY;
+ // ViewHelper's mixed-space gizmo origin and viewport extent (its exact internal expressions).
+ const offsetX = rect.left + (offsetWidth - dim);
+ const offsetY = rect.top + (offsetHeight - dim);
+ return {
+ clientX: offsetX + ((rect.right - offsetX) * (layoutX - (offsetWidth - dim))) / dim,
+ clientY: offsetY + ((rect.bottom - offsetY) * (layoutY - (offsetHeight - dim))) / dim,
+ };
+};
diff --git a/invokeai/frontend/web/src/features/controlLayers/components/common/CanvasEntityMenuItemsConvertTo3D.tsx b/invokeai/frontend/web/src/features/controlLayers/components/common/CanvasEntityMenuItemsConvertTo3D.tsx
new file mode 100644
index 00000000000..7d89cf9be89
--- /dev/null
+++ b/invokeai/frontend/web/src/features/controlLayers/components/common/CanvasEntityMenuItemsConvertTo3D.tsx
@@ -0,0 +1,34 @@
+import { Menu, MenuButton, MenuItem, MenuList } from '@invoke-ai/ui-library';
+import { SubMenuButtonContent, useSubMenu } from 'common/hooks/useSubMenu';
+import { useEntityIdentifierContext } from 'features/controlLayers/contexts/EntityIdentifierContext';
+import { useEntityConvertTo3D } from 'features/controlLayers/hooks/useEntityConvertTo3D';
+import { memo, useCallback } from 'react';
+import { PiCubeFill } from 'react-icons/pi';
+
+export const CanvasEntityMenuItemsConvertTo3D = memo(() => {
+ const subMenu = useSubMenu();
+ const entityIdentifier = useEntityIdentifierContext();
+ const { isDisabled, start } = useEntityConvertTo3D(entityIdentifier);
+ const convertIsolated = useCallback(() => start(true), [start]);
+ const convertFullImage = useCallback(() => start(false), [start]);
+
+ return (
+ } isDisabled={isDisabled}>
+
+
+ );
+});
+
+CanvasEntityMenuItemsConvertTo3D.displayName = 'CanvasEntityMenuItemsConvertTo3D';
diff --git a/invokeai/frontend/web/src/features/controlLayers/hooks/useEntityConvertTo3D.ts b/invokeai/frontend/web/src/features/controlLayers/hooks/useEntityConvertTo3D.ts
new file mode 100644
index 00000000000..f0c761f59d7
--- /dev/null
+++ b/invokeai/frontend/web/src/features/controlLayers/hooks/useEntityConvertTo3D.ts
@@ -0,0 +1,117 @@
+import { NUMPY_RAND_MAX, NUMPY_RAND_MIN } from 'app/constants';
+import { logger } from 'app/logging/logger';
+import randomInt from 'common/util/randomInt';
+import { withResultAsync } from 'common/util/result';
+import {
+ $splatOverlay,
+ clearSplatGenerationAbort,
+ clearSplatOverlay,
+ setSplatGenerationAbort,
+} from 'features/controlLayers/components/SplatOverlay/state';
+import { useCanvasManager } from 'features/controlLayers/contexts/CanvasManagerProviderGate';
+import { useEntityAdapterSafe } from 'features/controlLayers/contexts/EntityAdapterContext';
+import { useCanvasIsBusy } from 'features/controlLayers/hooks/useCanvasIsBusy';
+import { useEntityIsEmpty } from 'features/controlLayers/hooks/useEntityIsEmpty';
+import { useEntityIsLocked } from 'features/controlLayers/hooks/useEntityIsLocked';
+import { getPrefixedId } from 'features/controlLayers/konva/util';
+import type { CanvasEntityIdentifier } from 'features/controlLayers/store/types';
+import { Graph } from 'features/nodes/util/graph/generation/Graph';
+import { useCallback, useMemo } from 'react';
+import { buildV1Url } from 'services/api';
+
+const log = logger('canvas');
+
+/**
+ * Trigger for the "Convert to 3D" raster-layer action: rasterizes the layer, runs the `image_to_3d`
+ * (TripoSplat) node, and opens the splat overlay with the resulting .ply. Mirrors useEntitySegmentAnything.
+ */
+export const useEntityConvertTo3D = (entityIdentifier: CanvasEntityIdentifier | null) => {
+ const canvasManager = useCanvasManager();
+ const adapter = useEntityAdapterSafe(entityIdentifier);
+ const isBusy = useCanvasIsBusy();
+ const isLocked = useEntityIsLocked(entityIdentifier);
+ const isEmpty = useEntityIsEmpty(entityIdentifier);
+
+ const isDisabled = useMemo(() => {
+ if (!entityIdentifier || entityIdentifier.type !== 'raster_layer') {
+ return true;
+ }
+ if (!adapter || isBusy || isLocked || isEmpty) {
+ return true;
+ }
+ return false;
+ }, [entityIdentifier, adapter, isBusy, isLocked, isEmpty]);
+
+ const start = useCallback(
+ (removeBackground: boolean) => {
+ if (isDisabled || !entityIdentifier || entityIdentifier.type !== 'raster_layer') {
+ return;
+ }
+ const adapter = canvasManager.getAdapter(entityIdentifier);
+ if (!adapter) {
+ return;
+ }
+ const rect = adapter.transformer.getRelativeRect();
+ if (rect.width === 0 || rect.height === 0) {
+ return;
+ }
+
+ // Only one overlay session can exist; replacing a session must abort its backend run first, or the
+ // old run becomes uncancelable and its late completion would write into the new session's state.
+ clearSplatOverlay();
+ const sessionId = getPrefixedId('splat_session');
+ const controller = new AbortController();
+ setSplatGenerationAbort(controller);
+ $splatOverlay.set({ status: 'loading', sessionId, rect });
+
+ void (async () => {
+ const result = await withResultAsync(async () => {
+ const imageDTO = await adapter.renderer.rasterize({ rect, attrs: { filters: [], opacity: 1 } });
+ const graph = new Graph(getPrefixedId('canvas_image_to_3d'));
+ const node = graph.addNode({
+ id: getPrefixedId('image_to_3d'),
+ type: 'image_to_3d',
+ image: { image_name: imageDTO.image_name },
+ remove_background: removeBackground,
+ // Random per conversion so re-converting the same layer is a fresh roll: a fixed seed would hit
+ // the invocation cache and return the identical splat (and full determinism isn't available
+ // anyway — TripoSplat's decode stage draws from the unseeded global RNG).
+ seed: randomInt(NUMPY_RAND_MIN, NUMPY_RAND_MAX),
+ });
+ const output = await canvasManager.stateApi.runGraphAndReturnOutput({
+ graph,
+ outputNodeId: node.id,
+ options: { prepend: true, signal: controller.signal },
+ });
+ if (output.type !== 'asset_3d_output') {
+ throw new Error(`Unexpected output type: ${output.type}`);
+ }
+ return buildV1Url(`assets/i/${output.asset.asset_name}`);
+ });
+
+ clearSplatGenerationAbort(controller);
+
+ // Every write below is gated on the overlay still showing *this* session — the user may have
+ // cancelled (state null) or started another conversion (different sessionId) while we generated.
+ const current = $splatOverlay.get();
+ const isCurrentSession = current?.status === 'loading' && current.sessionId === sessionId;
+
+ if (result.isErr()) {
+ log.error({ error: String(result.error) }, 'Failed to convert image to 3D');
+ if (isCurrentSession) {
+ clearSplatOverlay();
+ }
+ return;
+ }
+
+ if (isCurrentSession) {
+ // Use the overlay's *current* rect, not the one captured at start — the frame is movable while loading.
+ $splatOverlay.set({ status: 'ready', sessionId, assetUrl: result.value, rect: current.rect });
+ }
+ })();
+ },
+ [isDisabled, entityIdentifier, canvasManager]
+ );
+
+ return { isDisabled, start } as const;
+};
diff --git a/invokeai/frontend/web/src/features/controlLayers/konva/CanvasStateApiModule.ts b/invokeai/frontend/web/src/features/controlLayers/konva/CanvasStateApiModule.ts
index 26abd908e51..71c1a138eae 100644
--- a/invokeai/frontend/web/src/features/controlLayers/konva/CanvasStateApiModule.ts
+++ b/invokeai/frontend/web/src/features/controlLayers/konva/CanvasStateApiModule.ts
@@ -350,6 +350,25 @@ export class CanvasStateApiModule extends CanvasModuleBase {
return imageDTO;
};
+ /**
+ * Run a graph and return the raw output of a node. Use this for non-image outputs (e.g. the 3D asset
+ * output from the image_to_3d node).
+ *
+ * @param arg.graph The graph to execute.
+ * @param arg.outputNodeId The id of the node whose output will be retrieved.
+ * @param arg.options Optional RunGraphOptions (prepend, signal, timeout, destination).
+ * @returns A promise that resolves to the raw invocation output of the specified node.
+ */
+ runGraphAndReturnOutput = async (arg: {
+ graph: Graph;
+ outputNodeId: string;
+ options?: RunGraphOptions;
+ }): Promise => {
+ const dependencies = buildRunGraphDependencies(this.store.dispatch, this.manager.socket);
+ const { output } = await runGraph({ dependencies, ...arg });
+ return output;
+ };
+
/**
* Helper function to extract ImageDTO from graph execution result.
* Expects the result to be an ImageOutput.
diff --git a/invokeai/frontend/web/src/features/ui/layouts/CanvasWorkspacePanel.tsx b/invokeai/frontend/web/src/features/ui/layouts/CanvasWorkspacePanel.tsx
index 74f0c3f5f45..8510dd902bf 100644
--- a/invokeai/frontend/web/src/features/ui/layouts/CanvasWorkspacePanel.tsx
+++ b/invokeai/frontend/web/src/features/ui/layouts/CanvasWorkspacePanel.tsx
@@ -14,6 +14,7 @@ import { Filter } from 'features/controlLayers/components/Filters/Filter';
import { CanvasHUD } from 'features/controlLayers/components/HUD/CanvasHUD';
import { InvokeCanvasComponent } from 'features/controlLayers/components/InvokeCanvasComponent';
import { SelectObject } from 'features/controlLayers/components/SelectObject/SelectObject';
+import { CanvasSplatOverlay } from 'features/controlLayers/components/SplatOverlay/CanvasSplatOverlay';
import { StagingAreaContextProvider } from 'features/controlLayers/components/StagingArea/context';
import { CanvasTextOverlay } from 'features/controlLayers/components/Text/CanvasTextOverlay';
import { PinnedFillColorPickerOverlay } from 'features/controlLayers/components/Tool/PinnedFillColorPickerOverlay';
@@ -83,6 +84,7 @@ export const CanvasWorkspacePanel = memo(() => {
+
dispatch(queueApi.util.getRunningQueryThunk('getQueueStatus', undefined));
+ const reinvalidateQueueStatusAfter = (inFlightFetch: ReturnType) => {
+ if (!inFlightFetch) {
+ return;
+ }
+ void Promise.resolve(inFlightFetch)
+ .catch(() => undefined)
+ .then(() => {
+ dispatch(queueApi.util.invalidateTags(['SessionQueueStatus']));
+ });
+ };
+
socket.on('queue_item_status_changed', (data) => {
+ const isTerminalStatus = data.status === 'completed' || data.status === 'failed' || data.status === 'canceled';
+ const inFlightQueueStatusFetch = isTerminalStatus ? getInFlightQueueStatusFetch() : undefined;
+
// Sanitized companion event sent to non-owner queue subscribers in multiuser mode. The
// backend sets user_id="redacted" and clears identifiers/error fields. We must not run
// payload-driven cache mutations or per-session side effects (node state reset, progress
@@ -439,6 +465,7 @@ export const setEventListeners = ({ socket, store, setIsConnected }: SetEventLis
{ type: 'SessionQueueItem', id: LIST_ALL_TAG },
];
dispatch(queueApi.util.invalidateTags(tags));
+ reinvalidateQueueStatusAfter(inFlightQueueStatusFetch);
return;
}
@@ -519,6 +546,7 @@ export const setEventListeners = ({ socket, store, setIsConnected }: SetEventLis
tagsToInvalidate.push({ type: 'QueueCountsByDestination', id: destination });
}
dispatch(queueApi.util.invalidateTags(tagsToInvalidate));
+ reinvalidateQueueStatusAfter(inFlightQueueStatusFetch);
if (status === 'completed' || status === 'failed' || status === 'canceled') {
if (status === 'failed' && error_type) {
diff --git a/pyproject.toml b/pyproject.toml
index de665d621e3..50cb45b76d6 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -249,6 +249,8 @@ exclude = [
"invokeai/backend/image_util/normal_bae/", # External code
"invokeai/backend/image_util/pidi/", # External code
"invokeai/backend/image_util/imwatermark/", # External code
+ "invokeai/backend/image_util/triposplat/model.py", # External code (vendored TripoSplat)
+ "invokeai/backend/image_util/triposplat/triposplat.py", # External code (vendored TripoSplat)
]
[tool.ruff.lint]
diff --git a/tests/app/routers/test_image_moves.py b/tests/app/routers/test_image_moves.py
index 697447269c3..23f5475c447 100644
--- a/tests/app/routers/test_image_moves.py
+++ b/tests/app/routers/test_image_moves.py
@@ -48,6 +48,7 @@ def mock_services() -> InvocationServices:
db = create_mock_sqlite_database(configuration, logger)
image_moves = MagicMock()
return InvocationServices(
+ asset_files=None, # type: ignore
board_image_records=SqliteBoardImageRecordStorage(db=db),
board_images=None, # type: ignore
board_records=SqliteBoardRecordStorage(db=db),
diff --git a/tests/app/routers/test_multiuser_authorization.py b/tests/app/routers/test_multiuser_authorization.py
index f1ae8defa1a..e81af84dc74 100644
--- a/tests/app/routers/test_multiuser_authorization.py
+++ b/tests/app/routers/test_multiuser_authorization.py
@@ -85,6 +85,7 @@ def mock_services() -> InvocationServices:
db = create_mock_sqlite_database(configuration, logger)
return InvocationServices(
+ asset_files=None, # type: ignore
board_image_records=SqliteBoardImageRecordStorage(db=db),
board_images=None, # type: ignore
board_records=SqliteBoardRecordStorage(db=db),
diff --git a/tests/app/routers/test_session_queue_workflow_call.py b/tests/app/routers/test_session_queue_workflow_call.py
index d8fd465d8d3..b0d1ecc5f80 100644
--- a/tests/app/routers/test_session_queue_workflow_call.py
+++ b/tests/app/routers/test_session_queue_workflow_call.py
@@ -62,6 +62,7 @@ def mock_services() -> InvocationServices:
db = create_mock_sqlite_database(configuration, logger)
return InvocationServices(
+ asset_files=None, # type: ignore
board_image_records=SqliteBoardImageRecordStorage(db=db),
board_images=None, # type: ignore
board_records=SqliteBoardRecordStorage(db=db),
diff --git a/tests/app/routers/test_workflows_multiuser.py b/tests/app/routers/test_workflows_multiuser.py
index 372a34b870f..14fcb6e06ee 100644
--- a/tests/app/routers/test_workflows_multiuser.py
+++ b/tests/app/routers/test_workflows_multiuser.py
@@ -76,6 +76,7 @@ def mock_services() -> InvocationServices:
db = create_mock_sqlite_database(configuration, logger)
return InvocationServices(
+ asset_files=None, # type: ignore
board_image_records=SqliteBoardImageRecordStorage(db=db),
board_images=None, # type: ignore
board_records=SqliteBoardRecordStorage(db=db),
diff --git a/tests/app/services/test_image_move_startup_safety.py b/tests/app/services/test_image_move_startup_safety.py
index f7f673322ea..b1aea66fb1a 100644
--- a/tests/app/services/test_image_move_startup_safety.py
+++ b/tests/app/services/test_image_move_startup_safety.py
@@ -7,6 +7,7 @@
def _services(**overrides):
services = {
+ "asset_files": object(),
"board_image_records": object(),
"board_images": object(),
"board_records": object(),
diff --git a/tests/backend/model_manager/load/test_model_util.py b/tests/backend/model_manager/load/test_model_util.py
new file mode 100644
index 00000000000..65a88b5a803
--- /dev/null
+++ b/tests/backend/model_manager/load/test_model_util.py
@@ -0,0 +1,27 @@
+import logging
+
+import torch
+
+from invokeai.backend.image_util.triposplat.triposplat_model import TripoSplatModel
+from invokeai.backend.model_manager.load.model_util import calc_model_size_by_data
+
+
+class _FakePipe:
+ """Stands in for TripoSplatPipeline: just the five submodule attributes the wrapper reads."""
+
+ def __init__(self) -> None:
+ module = torch.nn.Linear(4, 4)
+ self.dinov3 = module
+ self.vae_encoder = module
+ self.rmbg = module
+ self.flow_model = module
+ self.decoder = module
+
+
+def test_calc_model_size_by_data_dispatches_to_triposplat_calc_size():
+ # TripoSplatModel is not an nn.Module — if it is missing from calc_model_size_by_data's isinstance
+ # dispatch it silently sizes as 0 and the model cache cannot make room for the multi-GB pipeline.
+ model = TripoSplatModel(_FakePipe())
+ size = calc_model_size_by_data(logging.getLogger(__name__), model)
+ assert size == model.calc_size()
+ assert size > 0
diff --git a/tests/conftest.py b/tests/conftest.py
index 80df0aa93e9..5849bc08c5b 100644
--- a/tests/conftest.py
+++ b/tests/conftest.py
@@ -39,6 +39,7 @@ def mock_services() -> InvocationServices:
# NOTE: none of these are actually called by the test invocations
return InvocationServices(
+ asset_files=None, # type: ignore
board_image_records=SqliteBoardImageRecordStorage(db=db),
board_images=None, # type: ignore
board_records=SqliteBoardRecordStorage(db=db),