diff --git a/IF_Trellis.py b/IF_Trellis.py
index 2eb6c22..722eb77 100644
--- a/IF_Trellis.py
+++ b/IF_Trellis.py
@@ -193,7 +193,7 @@ class IF_TrellisImageTo3D:
video_path = os.path.join(out_dir, f"{project_name}_preview.mp4")
imageio.mimsave(video_path, video, fps=fps)
full_video_path = os.path.abspath(video_path)
- video_path = get_subpath_after_dir(full_video_path, "output")
+ video_path = os.path.abspath(video_path)
logger.info(f"Full video path: {full_video_path}, Processed video path: {video_path}")
if save_glb:
@@ -308,7 +308,7 @@ class IF_TrellisImageTo3D:
pipeline_params = self.get_pipeline_params(
seed, ss_sampling_steps, ss_guidance_strength,
- slat_sampling_steps, slat_guidance_strength
+ slat_sampling_steps, slat_guidance_strength,
)
# Handle single vs multi mode differently
@@ -316,7 +316,8 @@ class IF_TrellisImageTo3D:
# Take just the first image regardless of how many were input
images = images[0:1]
pil_imgs = self.torch_to_pil_batch(images, masks)
- outputs = model.run(pil_imgs[0], **pipeline_params)
+ outputs = model.run(pil_imgs[0],
+ **pipeline_params)
else:
# In multi mode, treat the whole list as a batch
pil_imgs = self.torch_to_pil_batch(images, masks)
diff --git a/IF_TrellisCheckpointLoader.py b/IF_TrellisCheckpointLoader.py
index 27ec9bd..f69fcaa 100644
--- a/IF_TrellisCheckpointLoader.py
+++ b/IF_TrellisCheckpointLoader.py
@@ -1,261 +1,161 @@
# IF_TrellisCheckpointLoader.py
import os
-import sys
-import importlib
-import torch
import logging
+import torch
import folder_paths
-from huggingface_hub import hf_hub_download, snapshot_download
-from pathlib import Path
-import json
from trellis_model_manager import TrellisModelManager
from trellis.pipelines.trellis_image_to_3d import TrellisImageTo3DPipeline
-from trellis.modules import set_attention_backend
+from trellis.backend_config import (
+ set_attention_backend,
+ set_sparse_backend,
+ get_available_backends,
+ get_available_sparse_backends
+)
from typing import Literal
-from trellis.modules.attention_utils import enable_sage_attention, disable_sage_attention
+from torchvision import transforms
logger = logging.getLogger("IF_Trellis")
-def set_backend(backend: Literal['spconv', 'torchsparse']):
- # Example helper if you wish to call the underlying global set_backend from trellis.modules.sparse:
- from trellis.modules.sparse import set_backend as _set_sparse_backend
- # Also handle spconv algo if desired, e.g. os.environ['SPCONV_ALGO'] = ...
- _set_sparse_backend(backend)
-
-class TrellisConfig:
- """Global configuration for Trellis"""
- def __init__(self):
- self.logger = logger
- self.attention_backend = "sage"
- self.spconv_algo = "implicit_gemm"
- self.smooth_k = True
- self.device = "cuda"
- self.use_fp16 = True
- # Added new configuration dictionary
- self._config = {
- "dinov2_size": "large", # Default model size
- "dinov2_model": "dinov2_vitg14" # Default model name
- }
-
- # Added new methods
- def get(self, key, default=None):
- """Get configuration value with fallback"""
- return self._config.get(key, default)
-
- def set(self, key, value):
- """Set configuration value"""
- self._config[key] = value
-
- def setup_environment(self):
- """Set up all environment variables and backends"""
- import os
- from trellis.modules import set_attention_backend
- from trellis.modules.sparse import set_backend
-
- # Set attention backend
- set_attention_backend(self.attention_backend)
-
- # Set smooth k for sage attention
- os.environ['SAGEATTN_SMOOTH_K'] = '1' if self.smooth_k else '0'
-
- # Set spconv algorithm
- os.environ['SPCONV_ALGO'] = self.spconv_algo
-
- # Always use spconv as backend for now
- set_backend('spconv')
-
- logger.info(f"Environment configured - Backend: spconv, "
- f"Attention: {self.attention_backend}, "
- f"Smooth K: {self.smooth_k}, "
- f"SpConv Algo: {self.spconv_algo}")
-
-# Global config instance
-TRELLIS_CONFIG = TrellisConfig()
-
class IF_TrellisCheckpointLoader:
"""
- Node to manage the loading of the TRELLIS model.
- Follows ComfyUI conventions for model management.
+ Node to manage the loading of the TRELLIS model with lazy backend selection.
"""
def __init__(self):
self.logger = logger
self.model_manager = None
- # Check for available devices
- self.device = self._get_device()
-
- def _get_device(self):
- """Determine the best available device."""
- if torch.cuda.is_available():
- return "cuda"
- elif hasattr(torch.backends, 'mps') and torch.backends.mps.is_available():
- return "mps"
- return "cpu"
-
+ self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+ # We might call these to figure out what's actually installed,
+ # if we want to populate UI dropdowns:
+ self.attn_backends = get_available_backends() # e.g. { 'xformers': True, 'flash_attn': False, ... }
+ self.sparse_backends = get_available_sparse_backends()# e.g. { 'spconv': True, 'torchsparse': True }
+
@classmethod
def INPUT_TYPES(cls):
"""Define input types with device-specific options."""
- device_options = []
- if torch.cuda.is_available():
- device_options.append("cuda")
- if hasattr(torch.backends, 'mps') and torch.backends.mps.is_available():
- device_options.append("mps")
- device_options.append("cpu")
+ # Filter only available backends
+ attn_backends = get_available_backends()
+ sparse_backends = get_available_sparse_backends()
+
+ # e.g. create a list of names that are True:
+ available_attn = [k for k, v in attn_backends.items() if v]
+ if not available_attn:
+ available_attn = ['sdpa'] # fallback
+
+ available_sparse = [k for k, v in sparse_backends.items() if v]
+ if not available_sparse:
+ available_sparse = ['spconv'] # fallback
return {
"required": {
"model_name": (["TRELLIS-image-large"],),
- "dinov2_model": (["dinov2_vitl14_reg", "dinov2_vitg14_reg"], {"default": "dinov2_vitl14_reg", "tooltip": "Select the Dinov2 model to use for the image to 3D conversion. Smaller models work but better results with larger models."}),
+ "dinov2_model": (["dinov2_vitl14_reg"],
+ {"default": "dinov2_vitl14_reg",
+ "tooltip": "Select which Dinov2 model to use."}),
"use_fp16": ("BOOLEAN", {"default": True}),
- "attn_backend": (["sage", "xformers", "flash_attn", "sdpa", "naive"], {"default": "sage", "tooltip": "Select the attention backend to use for the image to 3D conversion. Sage is experimental but faster"}),
- "smooth_k": ("BOOLEAN", {"default": True, "tooltip": "Smooth k for sage attention. This is a hyperparameter that controls the smoothness of the attention distribution. It is a boolean value that determines whether to use smooth k or not. Smooth k is a hyperparameter that controls the smoothness of the attention distribution. It is a boolean value that determines whether to use smooth k or not."}),
- "spconv_algo": (["implicit_gemm", "native"], {"default": "implicit_gemm", "tooltip": "Select the spconv algorithm to use for the image to 3D conversion. Implicit gemm is the best but slower. Native is the fastest but less accurate."}),
- "main_device": (device_options, {"default": device_options[0]}),
+ #
+ # The user picks from the actually installed backends
+ #
+ "attn_backend": (available_attn,
+ {"default": "sdpa" if "sdpa" in available_attn else available_attn[0],
+ "tooltip": "Select attention backend."}),
+ "sparse_backend": (available_sparse,
+ {"default": "spconv" if "spconv" in available_sparse else available_sparse[0],
+ "tooltip": "Select sparse backend."}),
+ "spconv_algo": (["implicit_gemm", "native", "auto"],
+ {"default": "implicit_gemm",
+ "tooltip": "Spconv algorithm. 'implicit_gemm' is slower but more robust."}),
+ "smooth_k": ("BOOLEAN",
+ {"default": True,
+ "tooltip": "Smooth-k for SageAttention. Only relevant if attn_backend=sage."}),
},
}
-
+
RETURN_TYPES = ("TRELLIS_MODEL",)
RETURN_NAMES = ("model",)
FUNCTION = "load_model"
CATEGORY = "ImpactFrames๐ฅ๐๏ธ/Trellis"
- @classmethod
- def _check_backend_availability(cls, backend: str) -> bool:
- """Check if a specific attention backend is available"""
- try:
- if backend == 'sage':
- import sageattention
- elif backend == 'xformers':
- import xformers.ops
- elif backend == 'flash_attn':
- import flash_attn
- elif backend in ['sdpa', 'naive']:
- # These are always available in PyTorch
- pass
- else:
- return False
- return True
- except ImportError:
- return False
-
- @classmethod
- def _initialize_backend(cls, requested_backend: str = None) -> str:
- """Initialize attention backend with fallback logic"""
- # Priority order for backends
- backend_priority = ['sage', 'flash_attn', 'xformers', 'sdpa']
-
- # If a specific backend is requested, try it first
- if requested_backend:
- if cls._check_backend_availability(requested_backend):
- logger.info(f"Using requested attention backend: {requested_backend}")
- return requested_backend
- else:
- logger.warning(f"Requested backend '{requested_backend}' not available, falling back")
-
- # Try backends in priority order
- for backend in backend_priority:
- if cls._check_backend_availability(backend):
- logger.info(f"Using attention backend: {backend}")
- return backend
-
- # Final fallback to SDPA
- logger.info("All optimized attention backends unavailable, using PyTorch SDPA")
- return 'sdpa'
-
- def _setup_environment(self):
+ def _setup_environment(self, attn_backend: str, sparse_backend: str, spconv_algo: str, smooth_k: bool):
"""
- Set up environment variables based on the global TRELLIS_CONFIG.
+ Set up environment variables and backends lazily.
+ This is the main difference: we call our new lazy set_*_backend funcs.
"""
- import os
- from trellis.modules import set_attention_backend
- from trellis.modules.sparse import set_backend
- from trellis.modules.sparse.conv import SPCONV_ALGO
+ # Try attention
+ success = set_attention_backend(attn_backend)
+ if not success:
+ self.logger.warning(f"Failed to set {attn_backend} or not installed, fallback to sdpa.")
- # Set attention backend
- os.environ['ATTN_BACKEND'] = TRELLIS_CONFIG.attention_backend
- set_attention_backend(TRELLIS_CONFIG.attention_backend)
+ # Try sparse
+ success2 = set_sparse_backend(sparse_backend, spconv_algo)
+ if not success2:
+ self.logger.warning(f"Failed to set {sparse_backend} or not installed, fallback to default.")
- # Set smooth k for sage attention
- os.environ['SAGEATTN_SMOOTH_K'] = '1' if TRELLIS_CONFIG.smooth_k else '0'
+ # If user wants SageAttn smooth_k, we set environment var (if they'd want that):
+ os.environ['SAGEATTN_SMOOTH_K'] = '1' if smooth_k else '0'
- # Set spconv algorithm
- os.environ['SPCONV_ALGO'] = TRELLIS_CONFIG.spconv_algo
-
- # Always use spconv as backend for now
- set_backend('spconv')
+ def _initialize_transforms(self):
+ """Initialize image transforms if needed."""
+ return transforms.Compose([
+ transforms.Normalize(
+ mean=[0.485, 0.456, 0.406],
+ std=[0.229, 0.224, 0.225]
+ )
+ ])
- logger.info(f"Environment configured - Backend: spconv, "
- f"Attention: {TRELLIS_CONFIG.attention_backend}, "
- f"Smooth K: {TRELLIS_CONFIG.smooth_k}, "
- f"SpConv Algo: {TRELLIS_CONFIG.spconv_algo}")
-
- def optimize_pipeline(self, pipeline, use_fp16=True, attn_backend='sage'):
- """Apply optimizations to the pipeline if available"""
- if self.device == "cuda":
+ def _optimize_pipeline(self, pipeline, use_fp16: bool = True):
+ """
+ Apply typical optimizations, half-precision, etc.
+ """
+ if self.device.type == "cuda":
try:
if hasattr(pipeline, 'cuda'):
pipeline.cuda()
-
+
if use_fp16:
if hasattr(pipeline, 'enable_attention_slicing'):
- pipeline.enable_attention_slicing()
+ pipeline.enable_attention_slicing(slice_size="auto")
if hasattr(pipeline, 'half'):
pipeline.half()
-
- # Only enable xformers if using xformers backend
- if attn_backend == 'xformers' and hasattr(pipeline, 'enable_xformers_memory_efficient_attention'):
- pipeline.enable_xformers_memory_efficient_attention()
-
except Exception as e:
- logger.warning(f"Some optimizations failed: {str(e)}")
-
+ logger.warning(f"Some pipeline optimizations failed: {str(e)}")
+
return pipeline
- def load_model(self, model_name, dinov2_model="dinov2_vitg14", attn_backend="sage", use_fp16=True,
- smooth_k=True, spconv_algo="implicit_gemm", main_device="cuda"):
- """Load and configure the TRELLIS model."""
+ def load_model(
+ self,
+ model_name: str,
+ dinov2_model: str = "dinov2_vitl14_reg",
+ attn_backend: str = "sdpa",
+ sparse_backend: str = "spconv",
+ spconv_algo: str = "implicit_gemm",
+ use_fp16: bool = True,
+ smooth_k: bool = True,
+ ) -> tuple:
+ """
+ Load and configure the TRELLIS pipeline.
+ This is typically the main function invoked by ComfyUI at node execution time.
+ """
try:
- # Update global config
- TRELLIS_CONFIG.attention_backend = attn_backend
- TRELLIS_CONFIG.spconv_algo = spconv_algo
- TRELLIS_CONFIG.smooth_k = smooth_k
- TRELLIS_CONFIG.device = main_device
- TRELLIS_CONFIG.use_fp16 = use_fp16
- TRELLIS_CONFIG.set("dinov2_model", dinov2_model)
+ # 1) Setup environment + backends
+ self._setup_environment(attn_backend, sparse_backend, spconv_algo, smooth_k)
- # Set up environment
- self._setup_environment()
-
- # Configure attention backend
- set_attention_backend(attn_backend)
- if attn_backend == 'sage':
- enable_sage_attention()
- else:
- disable_sage_attention()
-
- # Get model path
+ # 2) Get model path
model_path = folder_paths.get_full_path("checkpoints", model_name)
if model_path is None:
model_path = os.path.join(folder_paths.models_dir, "checkpoints", model_name)
+ if not os.path.exists(model_path):
+ raise FileNotFoundError(f"Model not found: {model_path}")
- # Create pipeline with specified dinov2 model
- pipeline = TrellisImageTo3DPipeline.from_pretrained(model_path, dinov2_model=dinov2_model)
-
- # Configure pipeline after loading
- pipeline._device = torch.device(main_device)
- pipeline.attention_backend = attn_backend
-
- # Store configuration in pipeline
- pipeline.config = {
- 'device': main_device,
- 'use_fp16': use_fp16,
- 'attention_backend': attn_backend,
- 'dinov2_model': dinov2_model,
- 'spconv_algo': spconv_algo,
- 'smooth_k': smooth_k
- }
+ # 3) Create pipeline with the config
+ pipeline = TrellisImageTo3DPipeline.from_pretrained(
+ model_path,
+ dinov2_model=dinov2_model
+ )
+ pipeline._device = self.device # ensure pipeline uses our same device
- # Apply optimizations
- pipeline = self.optimize_pipeline(pipeline, use_fp16, attn_backend)
+ # 4) Apply optimizations
+ pipeline = self._optimize_pipeline(pipeline, use_fp16)
return (pipeline,)
diff --git a/README.md b/README.md
index 97aa9e5..cd44ea9 100644
--- a/README.md
+++ b/README.md
@@ -1,117 +1,117 @@
-ComfyUI-IF_Trellis
-
-
-
-
-ComfyUI TRELLIS is a large 3D asset generation in various formats, such as Radiance Fields, 3D Gaussians, and meshes. The cornerstone of TRELLIS is a unified Structured LATent (SLAT) representation that allows decoding to different output formats and Rectified Flow Transformers tailored for SLAT as the powerful backbones.
-
-
-
-### Prerequisites
-- **System**: The original code is currently tested only on **Linux**. For windows setup, you may refer to [#3](https://github.com/microsoft/TRELLIS/issues/3)
-- (This comfyui node is following this suggeted installation steps It works but you need to follow the steps as described in part one and two of the guide).
-- Windows users need to use the Win_requirements.txt.
-- Linux Users (Not tested Yet) use the linux_requirements.txt but I am still testing if it works with comfy in Linux I just added the original repo requirements.
-
-- **Hardware**: This repo improve memmory mangement with the PR made by Amorano you will need NVIDIA GPU with at least 8GB.
-- **Software**:
- - The [CUDA Toolkit](https://developer.nvidia.com/cuda-toolkit-archive) is needed to compile certain submodules. The code has been tested with CUDA versions 11.8 and 12.2. This repo use **CUDA 12.4**.
- - [Conda](https://docs.anaconda.com/miniconda/install/#quick-command-line-install) is recommended for managing dependencies.
- - Python version 3.8 or higher is required.
-
- Give unrestricted script access to powershell so venv can work:
-
-- Open an administrator powershell window
-- Type `Set-ExecutionPolicy Unrestricted` and answer A
-- Close admin powershell window
-
-### Installation Steps
-1. Clone the repo:
- ```
- cd ComfyUI/Custom_nodes
- git clone --recurse-submodules https://github.com/if-ai/ComfyUI-IF_Trellis.git
- ```
-## MUST HAVE `--recurse-submodules`
-
-ONLY tested on windows but it should work easier in Linux without any issues or needing such specific stuffs.
-***NOT tested or compatible with PORTABLE comfy embeded python env***
-watch quick overview of setting the env if needed
-[](https://www.youtube.com/watch?v=-vEpuYL9I3g)
-
-You need to set up the environment first
-follow this guide for the first part
-
-Set the VSCode Cpp Envirronment as in the guide
-
-[Installing Triton and Sage Attention Flash Attention](https://ko-fi.com/post/Installing-Triton-and-Sage-Attention-Flash-Attenti-P5P8175434)
-
-
-Setting up ComfyUI with the Xformers, flash attention, Sage-attention(Optional Recommended for Hunyuan and other Video models)
-
-*** You can Also Try This Other Guide ***
-
-[](https://www.youtube.com/watch?v=FjNfDsX-jR0)
-
-[How to Run Micromamba in ComfyUI with Triton, SAGE Attention, Flash Attention, and xFormers for ComfyUI 3D](https://comfyuiblog.com/how-to-run-micromamba-in-comfyui-with-triton-sage-attention-flash-attention-and-x-formers-for-comfyui-3d/)
-
-
-
-
-## ๐ฆ Second part of the installation
-
-Activate youur comfy environment
-```
-(gen) PS D:\ComfyUI\custom_nodes\ComfyUI-IF_Trellis> micromamba activate gen
-```
-If you haven't set your vars or for some reason it can't compile some of this specially `nvdiffrast`
-it doesn't hurt if you do it again now.
-```
-cmd.exe /c "C:\Program Files (x86)\Microsoft Visual Studio\2022\BuildTools\VC\Auxiliary\Build\vcvarsall.bat" x64 "&&" powershell
-```
-You will see some message like this:
-
-**********************************************************************
-** Visual Studio 2019 Developer Command Prompt v16.11.41
-** Copyright (c) 2021 Microsoft Corporation
-**********************************************************************
-[vcvarsall.bat] Environment initialized for: 'x64'
-Windows PowerShell
-Copyright (C) Microsoft Corporation. All rights reserved.
-
-```
-pip install -r win_requirements.txt
-```
-
-```bash
-pip install git+https://github.com/EasternJournalist/utils3d.git@9a4eb15e4021b67b12c460c7057d642626897ec8
-New-Item -ItemType Directory -Force -Path C:\tmp\extensions
-git clone --recurse-submodules https://github.com/JeffreyXiang/diffoctreerast.git C:\tmp\extensions\diffoctreerast
-pip install C:\tmp\extensions\diffoctreerast
-git clone https://github.com/autonomousvision/mip-splatting.git C:\tmp\extensions\mip-splatting
-pip install C:\tmp\extensions\mip-splatting\submodules\diff-gaussian-rasterization\
-pip install kaolin -f https://nvidia-kaolin.s3.us-east-2.amazonaws.com/torch-2.4.0_cu121.html
-git clone https://github.com/NVlabs/nvdiffrast.git C:\tmp\extensions\nvdiffrast
-pip install C:\tmp\extensions\nvdiffrast
-```
-
-
-
-
-
-
-## ๐ Features
-- **High Quality**: It produces diverse 3D assets at high quality with intricate shape and texture details.
-- **Versatility**: It takes text or image prompts and can generate various final 3D representations including but not limited to *Radiance Fields*, *3D Gaussians*, and *meshes*, accommodating diverse downstream requirements.
-- **Flexible Editing**: It allows for easy editings of generated 3D assets, such as generating variants of the same object or local editing of the 3D asset.
-
-
-## ๐ง TODO List
-- [x] Release comfyUI-IF_Trellis
-- [x] 3D Viewport
-- [x] OPT mode
-- [x] Multiviews
-- [x] Sage attn
-- [ ] Installation
-- [ ] colab notebook
-- [ ] MarchingCubes mode
-- [ ] Exposing Texture paranmeters
+ComfyUI-IF_Trellis
+
+
+
+
+ComfyUI TRELLIS is a large 3D asset generation in various formats, such as Radiance Fields, 3D Gaussians, and meshes. The cornerstone of TRELLIS is a unified Structured LATent (SLAT) representation that allows decoding to different output formats and Rectified Flow Transformers tailored for SLAT as the powerful backbones.
+
+
+
+### Prerequisites
+- **System**: The original code is currently tested only on **Linux**. For windows setup, you may refer to [#3](https://github.com/microsoft/TRELLIS/issues/3)
+- (This comfyui node is following this suggeted installation steps It works but you need to follow the steps as described in part one and two of the guide).
+- Windows users need to use the Win_requirements.txt.
+- Linux Users (Not tested Yet) use the linux_requirements.txt but I am still testing if it works with comfy in Linux I just added the original repo requirements.
+
+- **Hardware**: This repo improve memmory mangement with the PR made by Amorano you will need NVIDIA GPU with at least 8GB.
+- **Software**:
+ - The [CUDA Toolkit](https://developer.nvidia.com/cuda-toolkit-archive) is needed to compile certain submodules. The code has been tested with CUDA versions 11.8 and 12.2. This repo use **CUDA 12.4**.
+ - [Conda](https://docs.anaconda.com/miniconda/install/#quick-command-line-install) is recommended for managing dependencies.
+ - Python version 3.8 or higher is required.
+
+ Give unrestricted script access to powershell so venv can work:
+
+- Open an administrator powershell window
+- Type `Set-ExecutionPolicy Unrestricted` and answer A
+- Close admin powershell window
+
+### Installation Steps
+1. Clone the repo:
+ ```
+ cd ComfyUI/Custom_nodes
+ git clone --recurse-submodules https://github.com/if-ai/ComfyUI-IF_Trellis.git
+ ```
+## MUST HAVE `--recurse-submodules`
+
+ONLY tested on windows but it should work easier in Linux without any issues or needing such specific stuffs.
+***NOT tested or compatible with PORTABLE comfy embeded python env***
+watch quick overview of setting the env if needed
+[](https://www.youtube.com/watch?v=-vEpuYL9I3g)
+
+You need to set up the environment first
+follow this guide for the first part
+
+Set the VSCode Cpp Envirronment as in the guide
+
+[Installing Triton and Sage Attention Flash Attention](https://ko-fi.com/post/Installing-Triton-and-Sage-Attention-Flash-Attenti-P5P8175434)
+
+
+Setting up ComfyUI with the Xformers, flash attention, Sage-attention(Optional Recommended for Hunyuan and other Video models)
+
+*** You can Also Try This Other Guide ***
+
+[](https://www.youtube.com/watch?v=FjNfDsX-jR0)
+
+[How to Run Micromamba in ComfyUI with Triton, SAGE Attention, Flash Attention, and xFormers for ComfyUI 3D](https://comfyuiblog.com/how-to-run-micromamba-in-comfyui-with-triton-sage-attention-flash-attention-and-x-formers-for-comfyui-3d/)
+
+
+
+
+## ๐ฆ Second part of the installation
+
+Activate youur comfy environment
+```
+(gen) PS D:\ComfyUI\custom_nodes\ComfyUI-IF_Trellis> micromamba activate gen
+```
+If you haven't set your vars or for some reason it can't compile some of this specially `nvdiffrast`
+it doesn't hurt if you do it again now.
+```
+cmd.exe /c "C:\Program Files (x86)\Microsoft Visual Studio\2022\BuildTools\VC\Auxiliary\Build\vcvarsall.bat" x64 "&&" powershell
+```
+You will see some message like this:
+
+**********************************************************************
+** Visual Studio 2019 Developer Command Prompt v16.11.41
+** Copyright (c) 2021 Microsoft Corporation
+**********************************************************************
+[vcvarsall.bat] Environment initialized for: 'x64'
+Windows PowerShell
+Copyright (C) Microsoft Corporation. All rights reserved.
+
+```
+pip install -r win_requirements.txt
+```
+
+```bash
+pip install git+https://github.com/EasternJournalist/utils3d.git@9a4eb15e4021b67b12c460c7057d642626897ec8
+New-Item -ItemType Directory -Force -Path C:\tmp\extensions
+git clone --recurse-submodules https://github.com/JeffreyXiang/diffoctreerast.git C:\tmp\extensions\diffoctreerast
+pip install C:\tmp\extensions\diffoctreerast
+git clone https://github.com/autonomousvision/mip-splatting.git C:\tmp\extensions\mip-splatting
+pip install C:\tmp\extensions\mip-splatting\submodules\diff-gaussian-rasterization\
+pip install kaolin -f https://nvidia-kaolin.s3.us-east-2.amazonaws.com/torch-2.4.0_cu121.html
+git clone https://github.com/NVlabs/nvdiffrast.git C:\tmp\extensions\nvdiffrast
+pip install C:\tmp\extensions\nvdiffrast
+```
+
+
+
+
+
+
+## ๐ Features
+- **High Quality**: It produces diverse 3D assets at high quality with intricate shape and texture details.
+- **Versatility**: It takes text or image prompts and can generate various final 3D representations including but not limited to *Radiance Fields*, *3D Gaussians*, and *meshes*, accommodating diverse downstream requirements.
+- **Flexible Editing**: It allows for easy editings of generated 3D assets, such as generating variants of the same object or local editing of the 3D asset.
+
+
+## ๐ง TODO List
+- [x] Release comfyUI-IF_Trellis
+- [x] 3D Viewport
+- [x] OPT mode
+- [x] Multiviews
+- [x] Sage attn
+- [ ] Installation
+- [ ] colab notebook
+- [ ] MarchingCubes mode
+- [ ] Exposing Texture paranmeters
diff --git a/trellis/backend_config.py b/trellis/backend_config.py
new file mode 100644
index 0000000..53575ee
--- /dev/null
+++ b/trellis/backend_config.py
@@ -0,0 +1,225 @@
+# trellis/backend_config.py
+from typing import *
+import os
+import logging
+import importlib
+
+# Global variables
+BACKEND = 'spconv' # Default sparse backend
+DEBUG = False # Debug mode flag
+ATTN = 'sdpa' # Default attention backend
+SPCONV_ALGO = 'implicit_gemm' # Default algorithm
+
+def get_spconv_algo() -> str:
+ """Get current spconv algorithm."""
+ global SPCONV_ALGO
+ return SPCONV_ALGO
+
+def set_spconv_algo(algo: Literal['implicit_gemm', 'native', 'auto']) -> bool:
+ """Set spconv algorithm with validation."""
+ global SPCONV_ALGO
+
+ if algo not in ['implicit_gemm', 'native', 'auto']:
+ logger.warning(f"Invalid spconv algorithm: {algo}. Must be 'implicit_gemm', 'native', or 'auto'")
+ return False
+
+ SPCONV_ALGO = algo
+ os.environ['SPCONV_ALGO'] = algo
+ logger.info(f"Set spconv algorithm to: {algo}")
+ return True
+
+logger = logging.getLogger(__name__)
+
+def _try_import_xformers() -> bool:
+ try:
+ import xformers.ops
+ return True
+ except ImportError:
+ return False
+
+def _try_import_flash_attn() -> bool:
+ try:
+ import flash_attn
+ return True
+ except ImportError:
+ return False
+
+def _try_import_sageattention() -> bool:
+ try:
+ import torch.nn.functional as F
+ from sageattention import sageattn
+ F.scaled_dot_product_attention = sageattn
+ #import sageattention
+ return True
+ except ImportError:
+ return False
+
+def _try_import_spconv() -> bool:
+ try:
+ import spconv
+ return True
+ except ImportError:
+ return False
+
+def _try_import_torchsparse() -> bool:
+ try:
+ import torchsparse
+ return True
+ except ImportError:
+ return False
+
+def get_available_backends() -> Dict[str, bool]:
+ """Return dict of available attention backends and their status"""
+ return {
+ 'xformers': _try_import_xformers(),
+ 'flash_attn': _try_import_flash_attn(),
+ 'sage': _try_import_sageattention(),
+ 'naive': True,
+ 'sdpa': True # Always available with PyTorch >= 2.0
+ }
+
+def get_available_sparse_backends() -> Dict[str, bool]:
+ """Return dict of available sparse backends and their status"""
+ return {
+ 'spconv': _try_import_spconv(),
+ 'torchsparse': _try_import_torchsparse()
+ }
+
+def get_attention_backend() -> str:
+ """Get current attention backend"""
+ global ATTN
+ return ATTN
+
+def get_sparse_backend() -> str:
+ """Get current sparse backend"""
+ global BACKEND
+ return BACKEND
+
+def get_debug_mode() -> bool:
+ """Get current debug mode status"""
+ global DEBUG
+ return DEBUG
+
+def __from_env():
+ """Initialize settings from environment variables"""
+ global BACKEND
+ global DEBUG
+ global ATTN
+
+ env_sparse_backend = os.environ.get('SPARSE_BACKEND')
+ env_sparse_debug = os.environ.get('SPARSE_DEBUG')
+ env_sparse_attn = os.environ.get('SPARSE_ATTN_BACKEND')
+
+ if env_sparse_backend is not None and env_sparse_backend in ['spconv', 'torchsparse']:
+ BACKEND = env_sparse_backend
+ if env_sparse_debug is not None:
+ DEBUG = env_sparse_debug == '1'
+ if env_sparse_attn is not None and env_sparse_attn in ['xformers', 'flash_attn', 'sage', 'sdpa', 'naive']:
+ ATTN = env_sparse_attn
+ os.environ['SPARSE_ATTN_BACKEND'] = env_sparse_attn
+ os.environ['ATTN_BACKEND'] = env_sparse_attn
+
+ logger.info(f"[SPARSE] Backend: {BACKEND}, Attention: {ATTN}")
+
+def set_backend(backend: Literal['spconv', 'torchsparse']) -> bool:
+ """Set sparse backend with validation"""
+ global BACKEND
+
+ backend = backend.lower().strip()
+ logger.info(f"Setting sparse backend to: {backend}")
+
+ if backend == 'spconv':
+ try:
+ import spconv
+ BACKEND = 'spconv'
+ os.environ['SPARSE_BACKEND'] = 'spconv'
+ return True
+ except ImportError:
+ logger.warning("spconv not available")
+ return False
+
+ elif backend == 'torchsparse':
+ try:
+ import torchsparse
+ BACKEND = 'torchsparse'
+ os.environ['SPARSE_BACKEND'] = 'torchsparse'
+ return True
+ except ImportError:
+ logger.warning("torchsparse not available")
+ return False
+
+ return False
+
+def set_sparse_backend(backend: Literal['spconv', 'torchsparse'], algo: str = None) -> bool:
+ """Alias for set_backend for backwards compatibility
+
+ Parameters:
+ backend: The sparse backend to use
+ algo: The algorithm to use (only relevant for spconv backend)
+ """
+ # Call set_backend first
+ result = set_backend(backend)
+
+ # If algorithm is provided and backend was set successfully
+ if algo is not None and result:
+ set_spconv_algo(algo)
+
+ return result
+
+def set_debug(debug: bool):
+ """Set debug mode"""
+ global DEBUG
+ DEBUG = debug
+ if debug:
+ os.environ['SPARSE_DEBUG'] = '1'
+ else:
+ os.environ['SPARSE_DEBUG'] = '0'
+
+def set_attn(attn: Literal['xformers', 'flash_attn', 'sage', 'sdpa', 'naive']) -> bool:
+ """Set attention backend with validation"""
+ global ATTN
+
+ attn = attn.lower().strip()
+ logger.info(f"Setting attention backend to: {attn}")
+
+ if attn == 'xformers' and _try_import_xformers():
+ ATTN = 'xformers'
+ os.environ['SPARSE_ATTN_BACKEND'] = 'xformers'
+ os.environ['ATTN_BACKEND'] = 'xformers'
+ return True
+
+ elif attn == 'flash_attn' and _try_import_flash_attn():
+ ATTN = 'flash_attn'
+ os.environ['SPARSE_ATTN_BACKEND'] = 'flash_attn'
+ os.environ['ATTN_BACKEND'] = 'flash_attn'
+ return True
+
+ elif attn == 'sage' and _try_import_sageattention():
+ ATTN = 'sage'
+ os.environ['SPARSE_ATTN_BACKEND'] = 'sage'
+ os.environ['ATTN_BACKEND'] = 'sage'
+ return True
+
+ elif attn == 'sdpa':
+ ATTN = 'sdpa'
+ os.environ['SPARSE_ATTN_BACKEND'] = 'sdpa'
+ os.environ['ATTN_BACKEND'] = 'sdpa'
+ return True
+
+ elif attn == 'naive':
+ ATTN = 'naive'
+ os.environ['SPARSE_ATTN_BACKEND'] = 'naive'
+ os.environ['ATTN_BACKEND'] = 'naive'
+ return True
+
+
+ logger.warning(f"Attention backend {attn} not available")
+ return False
+
+# Add alias for backwards compatibility
+def set_attention_backend(backend: Literal['xformers', 'flash_attn', 'sage', 'sdpa']) -> bool:
+ """Alias for set_attn for backwards compatibility"""
+ return set_attn(backend)
+
+# Initialize from environment variables on module import
+__from_env()
diff --git a/trellis/modules/__init__.py b/trellis/modules/__init__.py
index 5c2cefc..2a16f24 100644
--- a/trellis/modules/__init__.py
+++ b/trellis/modules/__init__.py
@@ -1,18 +1,16 @@
-from .attention_utils import enable_sage_attention, disable_sage_attention
+#from .attention_utils import enable_sage_attention, disable_sage_attention
from .attention import (
- set_attention_backend,
- get_attention_op,
- ATTN_BACKEND,
+ scaled_dot_product_attention,
+ BACKEND,
+ DEBUG,
MultiHeadAttention,
RotaryPositionEmbedder
)
__all__ = [
- 'enable_sage_attention',
- 'disable_sage_attention',
- 'set_attention_backend',
- 'get_attention_op',
- 'ATTN_BACKEND',
+ 'scaled_dot_product_attention',
+ 'BACKEND',
+ 'DEBUG',
'MultiHeadAttention',
'RotaryPositionEmbedder'
]
\ No newline at end of file
diff --git a/trellis/modules/attention/__init__.py b/trellis/modules/attention/__init__.py
index b3cd0b7..27f8c0b 100644
--- a/trellis/modules/attention/__init__.py
+++ b/trellis/modules/attention/__init__.py
@@ -1,65 +1,38 @@
import os
import logging
from typing import Literal
+from trellis.backend_config import (
+ get_attention_backend,
+ get_debug_mode,
+)
+import logging
logger = logging.getLogger(__name__)
-# Global settings
-ATTN_BACKEND = 'flash_attn' # Default backend
-DEBUG = False # Add DEBUG flag here
-BACKEND = 'flash_attn' # Change default to a valid backend
+#ATTN = get_attention_backend()
+BACKEND = get_attention_backend()
+DEBUG = get_debug_mode()
-def set_attention_backend(backend: Literal['xformers', 'flash_attn', 'sdpa', 'sage', 'naive']):
- """Set the global attention backend"""
- global ATTN_BACKEND, BACKEND # Also update BACKEND when setting attention backend
- if backend not in ['xformers', 'flash_attn', 'sdpa', 'sage', 'naive']:
- raise ValueError(f"Unsupported attention backend: {backend}")
- ATTN_BACKEND = backend
- BACKEND = backend # Keep BACKEND in sync with ATTN_BACKEND
- os.environ['ATTN_BACKEND'] = backend
- logger.info(f"[ATTENTION] Set backend to: {backend}")
-
-def get_attention_op():
- """Get the appropriate attention implementation"""
- if ATTN_BACKEND == 'xformers':
- try:
- import xformers.ops
- return xformers.ops.memory_efficient_attention
- except ImportError:
- logger.warning("xformers not available, falling back to naive attention")
- return None
- elif ATTN_BACKEND == 'flash_attn':
- try:
- from flash_attn import flash_attn_func
- return flash_attn_func
- except ImportError:
- logger.warning("flash_attn not available, falling back to naive attention")
- return None
- elif ATTN_BACKEND == 'sage':
- try:
- from sageattention import sageattn
- return sageattn
- except ImportError:
- logger.warning("sageattention not available, falling back to naive attention")
- return None
- elif ATTN_BACKEND == 'sdpa':
- import torch
- if hasattr(torch.nn.functional, 'scaled_dot_product_attention'):
- return torch.nn.functional.scaled_dot_product_attention
+def __from_env():
+ """Read current backend configuration"""
+ #global ATTN
+ global BACKEND
+ global DEBUG
- # Fallback to naive implementation
- logger.warning("Using naive attention implementation")
- return None
+ # Get current settings from central config
+ #ATTN =
+ BACKEND = get_attention_backend()
+ DEBUG = get_debug_mode()
+
+ print(f"[ATTENTION] Using backend: {BACKEND}")
-# Import attention modules after defining globals
from .modules import MultiHeadAttention, RotaryPositionEmbedder
+from .full_attn import scaled_dot_product_attention
__all__ = [
- 'set_attention_backend',
- 'get_attention_op',
- 'ATTN_BACKEND',
- 'DEBUG',
+ 'scaled_dot_product_attention',
'BACKEND',
+ 'DEBUG',
'MultiHeadAttention',
'RotaryPositionEmbedder'
]
diff --git a/trellis/modules/attention/full_attn.py b/trellis/modules/attention/full_attn.py
index 68303dc..c34c91e 100644
--- a/trellis/modules/attention/full_attn.py
+++ b/trellis/modules/attention/full_attn.py
@@ -1,18 +1,36 @@
from typing import *
import torch
import math
-from . import DEBUG, BACKEND
+from trellis.backend_config import (
+ get_attention_backend,
+ get_debug_mode,
+ get_available_backends
+)
+import logging
-if BACKEND == 'xformers':
+logger = logging.getLogger(__name__)
+
+# Get configuration from central config
+BACKEND = get_attention_backend()
+DEBUG = get_debug_mode()
+
+# Get available backends and import if active
+available_backends = get_available_backends()
+
+if BACKEND == "xformers" and available_backends['xformers']:
import xformers.ops as xops
-elif BACKEND == 'flash_attn':
+elif BACKEND == "flash_attn" and available_backends['flash_attn']:
import flash_attn
-elif BACKEND == 'sdpa':
+elif BACKEND == "sage" and available_backends['sage']:
+ import torch.nn.functional as F
+ from sageattention import sageattn
+ F.scaled_dot_product_attention = sageattn
+elif BACKEND == "sdpa":
from torch.nn.functional import scaled_dot_product_attention as sdpa
-elif BACKEND == 'naive':
- pass
+elif BACKEND == "naive":
+ from torch.nn.functional import scaled_dot_product_attention as naive
else:
- raise ValueError(f"Unknown attention backend: {BACKEND}")
+ raise ValueError(f"Unknown attention module: {BACKEND}")
__all__ = [
@@ -137,4 +155,4 @@ def scaled_dot_product_attention(*args, **kwargs):
else:
raise ValueError(f"Unknown attention module: {BACKEND}")
- return out
+ return out
\ No newline at end of file
diff --git a/trellis/modules/repomix-output.txt b/trellis/modules/repomix-output.txt
new file mode 100644
index 0000000..40795da
--- /dev/null
+++ b/trellis/modules/repomix-output.txt
@@ -0,0 +1,2461 @@
+This file is a merged representation of the entire codebase, combining all repository files into a single document.
+Generated by Repomix on: 2025-01-12T00:27:20.764Z
+
+================================================================
+File Summary
+================================================================
+
+Purpose:
+--------
+This file contains a packed representation of the entire repository's contents.
+It is designed to be easily consumable by AI systems for analysis, code review,
+or other automated processes.
+
+File Format:
+------------
+The content is organized as follows:
+1. This summary section
+2. Repository information
+3. Repository structure
+4. Multiple file entries, each consisting of:
+ a. A separator line (================)
+ b. The file path (File: path/to/file)
+ c. Another separator line
+ d. The full contents of the file
+ e. A blank line
+
+Usage Guidelines:
+-----------------
+- This file should be treated as read-only. Any changes should be made to the
+ original repository files, not this packed version.
+- When processing this file, use the file path to distinguish
+ between different files in the repository.
+- Be aware that this file may contain sensitive information. Handle it with
+ the same level of security as you would the original repository.
+
+Notes:
+------
+- Some files may have been excluded based on .gitignore rules and Repomix's
+ configuration.
+- Binary files are not included in this packed representation. Please refer to
+ the Repository Structure section for a complete list of file paths, including
+ binary files.
+
+Additional Info:
+----------------
+
+For more information about Repomix, visit: https://github.com/yamadashy/repomix
+
+================================================================
+Repository Structure
+================================================================
+__init__.py
+attention_utils.py
+attention/__init__.py
+attention/full_attn.py
+attention/modules.py
+norm.py
+sparse/__init__.py
+sparse/attention/__init__.py
+sparse/attention/full_attn.py
+sparse/attention/modules.py
+sparse/attention/serialized_attn.py
+sparse/attention/windowed_attn.py
+sparse/basic.py
+sparse/conv/__init__.py
+sparse/conv/conv_spconv.py
+sparse/conv/conv_torchsparse.py
+sparse/linear.py
+sparse/nonlinearity.py
+sparse/norm.py
+sparse/spatial.py
+spatial.py
+utils.py
+
+================================================================
+Repository Files
+================================================================
+
+================
+File: __init__.py
+================
+import logging
+from trellis.backend_config import (
+ get_attention_backend,
+ get_available_backends
+)
+logger = logging.getLogger(__name__)
+
+# Import modules
+from .attention import MultiHeadAttention, RotaryPositionEmbedder
+
+
+__all__ = [
+ 'get_attention_backend',
+ 'MultiHeadAttention',
+ 'RotaryPositionEmbedder',
+ 'get_available_backends'
+]
+
+================
+File: attention_utils.py
+================
+#sage_attn.py
+import os
+from typing import Optional
+import torch
+import torch.nn.functional as F
+from sageattention import sageattn
+import math
+
+__all__ = ['SageAttention', 'sage_attention']
+
+
+def enable_sage_attention():
+ """
+ Enable SageAttention by replacing PyTorch's scaled_dot_product_attention
+ with sageattn from the SageAttention library.
+ """
+ F.scaled_dot_product_attention = sageattn
+ return True
+
+def disable_sage_attention():
+ """
+ Restore PyTorch's original scaled_dot_product_attention function.
+ """
+ F.scaled_dot_product_attention = torch.nn.functional.scaled_dot_product_attention
+ return True
+
+================
+File: attention/__init__.py
+================
+import os
+import logging
+from typing import Literal
+from trellis.backend_config import (
+ get_attention_backend,
+ get_debug_mode,
+)
+import logging
+
+logger = logging.getLogger(__name__)
+
+BACKEND = get_attention_backend()
+DEBUG = get_debug_mode()
+
+def __from_env():
+ """Read current backend configuration"""
+ global BACKEND
+ global DEBUG
+
+ # Get current settings from central config
+ BACKEND = get_attention_backend()
+ DEBUG = get_debug_mode()
+
+ print(f"[ATTENTION] Using backend: {BACKEND}")
+
+from .modules import MultiHeadAttention, RotaryPositionEmbedder
+from .full_attn import scaled_dot_product_attention
+
+__all__ = [
+ 'MultiHeadAttention',
+ 'RotaryPositionEmbedder',
+ 'scaled_dot_product_attention'
+]
+
+================
+File: attention/full_attn.py
+================
+from typing import *
+import torch
+import math
+import os
+import logging
+from typing import *
+import torch
+import math
+import logging
+from trellis.backend_config import get_attention_backend
+
+logger = logging.getLogger(__name__)
+
+# Get the backend from central config
+ATTN_BACKEND = get_attention_backend()
+
+# Import available backends based on central configuration
+if ATTN_BACKEND == "xformers":
+ import xformers.ops as xops
+elif ATTN_BACKEND == "flash_attn":
+ import flash_attn
+elif ATTN_BACKEND == "sage":
+ import sageattention
+ sageattn = sageattention.sageattn
+ from sageattention.core import sageattn_varlen
+elif ATTN_BACKEND == "sdpa":
+ from torch.nn.functional import scaled_dot_product_attention as sdpa
+
+# Log the active backend
+logger.info(f"Using attention backend: {ATTN_BACKEND}")
+
+__all__ = [
+ 'scaled_dot_product_attention',
+]
+
+
+def _naive_sdpa(q, k, v):
+ """
+ Naive implementation of scaled dot product attention.
+ """
+ q = q.permute(0, 2, 1, 3) # [N, H, L, C]
+ k = k.permute(0, 2, 1, 3) # [N, H, L, C]
+ v = v.permute(0, 2, 1, 3) # [N, H, L, C]
+ scale_factor = 1 / math.sqrt(q.size(-1))
+ attn_weight = q @ k.transpose(-2, -1) * scale_factor
+ attn_weight = torch.softmax(attn_weight, dim=-1)
+ out = attn_weight @ v
+ out = out.permute(0, 2, 1, 3) # [N, L, H, C]
+ return out
+
+
+@overload
+def scaled_dot_product_attention(qkv: torch.Tensor) -> torch.Tensor:
+ """
+ Apply scaled dot product attention.
+
+ Args:
+ qkv (torch.Tensor): A [N, L, 3, H, C] tensor containing Qs, Ks, and Vs.
+ """
+ ...
+
+@overload
+def scaled_dot_product_attention(q: torch.Tensor, kv: torch.Tensor) -> torch.Tensor:
+ """
+ Apply scaled dot product attention.
+
+ Args:
+ q (torch.Tensor): A [N, L, H, C] tensor containing Qs.
+ kv (torch.Tensor): A [N, L, 2, H, C] tensor containing Ks and Vs.
+ """
+ ...
+
+@overload
+def scaled_dot_product_attention(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor) -> torch.Tensor:
+ """
+ Apply scaled dot product attention.
+
+ Args:
+ q (torch.Tensor): A [N, L, H, Ci] tensor containing Qs.
+ k (torch.Tensor): A [N, L, H, Ci] tensor containing Ks.
+ v (torch.Tensor): A [N, L, H, Co] tensor containing Vs.
+
+ Note:
+ k and v are assumed to have the same coordinate map.
+ """
+ ...
+
+def scaled_dot_product_attention(*args, **kwargs):
+ arg_names_dict = {
+ 1: ['qkv'],
+ 2: ['q', 'kv'],
+ 3: ['q', 'k', 'v']
+ }
+ num_all_args = len(args) + len(kwargs)
+ assert num_all_args in arg_names_dict, f"Invalid number of arguments, got {num_all_args}, expected 1, 2, or 3"
+ for key in arg_names_dict[num_all_args][len(args):]:
+ assert key in kwargs, f"Missing argument {key}"
+
+ if num_all_args == 1:
+ qkv = args[0] if len(args) > 0 else kwargs['qkv']
+ assert len(qkv.shape) == 5 and qkv.shape[2] == 3, f"Invalid shape for qkv, got {qkv.shape}, expected [N, L, 3, H, C]"
+ device = qkv.device
+
+ elif num_all_args == 2:
+ q = args[0] if len(args) > 0 else kwargs['q']
+ kv = args[1] if len(args) > 1 else kwargs['kv']
+ assert q.shape[0] == kv.shape[0], f"Batch size mismatch, got {q.shape[0]} and {kv.shape[0]}"
+ assert len(q.shape) == 4, f"Invalid shape for q, got {q.shape}, expected [N, L, H, C]"
+ assert len(kv.shape) == 5, f"Invalid shape for kv, got {kv.shape}, expected [N, L, 2, H, C]"
+ device = q.device
+
+ elif num_all_args == 3:
+ q = args[0] if len(args) > 0 else kwargs['q']
+ k = args[1] if len(args) > 1 else kwargs['k']
+ v = args[2] if len(args) > 2 else kwargs['v']
+ assert q.shape[0] == k.shape[0] == v.shape[0], f"Batch size mismatch, got {q.shape[0]}, {k.shape[0]}, and {v.shape[0]}"
+ assert len(q.shape) == 4, f"Invalid shape for q, got {q.shape}, expected [N, L, H, Ci]"
+ assert len(k.shape) == 4, f"Invalid shape for k, got {k.shape}, expected [N, L, H, Ci]"
+ assert len(v.shape) == 4, f"Invalid shape for v, got {v.shape}, expected [N, L, H, Co]"
+ device = q.device
+
+ # Implementation based on backend
+ if ATTN_BACKEND == "xformers":
+ if num_all_args == 1:
+ q, k, v = qkv.unbind(dim=2)
+ elif num_all_args == 2:
+ k, v = kv.unbind(dim=2)
+ out = xops.memory_efficient_attention(q, k, v)
+ elif ATTN_BACKEND == "flash_attn":
+ if num_all_args == 1:
+ out = flash_attn.flash_attn_qkvpacked_func(qkv)
+ elif num_all_args == 2:
+ out = flash_attn.flash_attn_kvpacked_func(q, kv)
+ elif num_all_args == 3:
+ out = flash_attn.flash_attn_func(q, k, v)
+ elif ATTN_BACKEND == "sdpa":
+ if num_all_args == 1:
+ q, k, v = qkv.unbind(dim=2)
+ elif num_all_args == 2:
+ k, v = kv.unbind(dim=2)
+ q = q.permute(0, 2, 1, 3) # [N, H, L, C]
+ k = k.permute(0, 2, 1, 3) # [N, H, L, C]
+ v = v.permute(0, 2, 1, 3) # [N, H, L, C]
+ if "attn_mask" in kwargs:
+ out = sdpa(q, k, v, attn_mask=kwargs["attn_mask"])
+ else:
+ out = sdpa(q, k, v)
+ out = out.permute(0, 2, 1, 3) # [N, L, H, C]
+ elif ATTN_BACKEND == 'sage':
+ q = q.permute(0, 2, 1, 3).contiguous() # [N, H, L, C]
+ k = k.permute(0, 2, 1, 3).contiguous() # [N, H, L, C]
+ v = v.permute(0, 2, 1, 3).contiguous() # [N, H, L, C]
+ out = sageattn(q, k, v, tensor_layout="HND")
+ out = out.permute(0, 2, 1, 3) # [N, L, H, C]
+ elif ATTN_BACKEND == "naive":
+ if num_all_args == 1:
+ q, k, v = qkv.unbind(dim=2)
+ elif num_all_args == 2:
+ k, v = kv.unbind(dim=2)
+ out = _naive_sdpa(q, k, v)
+ else:
+ raise ValueError(f"Unknown attention module: {ATTN_BACKEND}")
+
+ return out
+
+================
+File: attention/modules.py
+================
+from typing import *
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from .full_attn import scaled_dot_product_attention
+
+
+class MultiHeadRMSNorm(nn.Module):
+ def __init__(self, dim: int, heads: int):
+ super().__init__()
+ self.scale = dim ** 0.5
+ self.gamma = nn.Parameter(torch.ones(heads, dim))
+
+ def forward(self, x: torch.Tensor) -> torch.Tensor:
+ return (F.normalize(x.float(), dim = -1) * self.gamma * self.scale).to(x.dtype)
+
+
+class RotaryPositionEmbedder(nn.Module):
+ def __init__(self, hidden_size: int, in_channels: int = 3):
+ super().__init__()
+ assert hidden_size % 2 == 0, "Hidden size must be divisible by 2"
+ self.hidden_size = hidden_size
+ self.in_channels = in_channels
+ self.freq_dim = hidden_size // in_channels // 2
+ self.freqs = torch.arange(self.freq_dim, dtype=torch.float32) / self.freq_dim
+ self.freqs = 1.0 / (10000 ** self.freqs)
+
+ def _get_phases(self, indices: torch.Tensor) -> torch.Tensor:
+ self.freqs = self.freqs.to(indices.device)
+ phases = torch.outer(indices, self.freqs)
+ phases = torch.polar(torch.ones_like(phases), phases)
+ return phases
+
+ def _rotary_embedding(self, x: torch.Tensor, phases: torch.Tensor) -> torch.Tensor:
+ x_complex = torch.view_as_complex(x.float().reshape(*x.shape[:-1], -1, 2))
+ x_rotated = x_complex * phases
+ x_embed = torch.view_as_real(x_rotated).reshape(*x_rotated.shape[:-1], -1).to(x.dtype)
+ return x_embed
+
+ def forward(self, q: torch.Tensor, k: torch.Tensor, indices: Optional[torch.Tensor] = None) -> Tuple[torch.Tensor, torch.Tensor]:
+ """
+ Args:
+ q (sp.SparseTensor): [..., N, D] tensor of queries
+ k (sp.SparseTensor): [..., N, D] tensor of keys
+ indices (torch.Tensor): [..., N, C] tensor of spatial positions
+ """
+ if indices is None:
+ indices = torch.arange(q.shape[-2], device=q.device)
+ if len(q.shape) > 2:
+ indices = indices.unsqueeze(0).expand(q.shape[:-2] + (-1,))
+
+ phases = self._get_phases(indices.reshape(-1)).reshape(*indices.shape[:-1], -1)
+ if phases.shape[1] < self.hidden_size // 2:
+ phases = torch.cat([phases, torch.polar(
+ torch.ones(*phases.shape[:-1], self.hidden_size // 2 - phases.shape[1], device=phases.device),
+ torch.zeros(*phases.shape[:-1], self.hidden_size // 2 - phases.shape[1], device=phases.device)
+ )], dim=-1)
+ q_embed = self._rotary_embedding(q, phases)
+ k_embed = self._rotary_embedding(k, phases)
+ return q_embed, k_embed
+
+
+class MultiHeadAttention(nn.Module):
+ def __init__(
+ self,
+ channels: int,
+ num_heads: int,
+ ctx_channels: Optional[int]=None,
+ type: Literal["self", "cross"] = "self",
+ attn_mode: Literal["full", "windowed"] = "full",
+ window_size: Optional[int] = None,
+ shift_window: Optional[Tuple[int, int, int]] = None,
+ qkv_bias: bool = True,
+ use_rope: bool = False,
+ qk_rms_norm: bool = False,
+ ):
+ super().__init__()
+ assert channels % num_heads == 0
+ assert type in ["self", "cross"], f"Invalid attention type: {type}"
+ assert attn_mode in ["full", "windowed"], f"Invalid attention mode: {attn_mode}"
+ assert type == "self" or attn_mode == "full", "Cross-attention only supports full attention"
+
+ if attn_mode == "windowed":
+ raise NotImplementedError("Windowed attention is not yet implemented")
+
+ 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.attn_mode = attn_mode
+ self.window_size = window_size
+ self.shift_window = shift_window
+ self.use_rope = use_rope
+ 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)
+
+ if use_rope:
+ self.rope = RotaryPositionEmbedder(channels)
+
+ def forward(self, x: torch.Tensor, context: Optional[torch.Tensor] = None, indices: Optional[torch.Tensor] = None) -> torch.Tensor:
+ B, L, C = x.shape
+ if self._type == "self":
+ qkv = self.to_qkv(x)
+ qkv = qkv.reshape(B, L, 3, self.num_heads, -1)
+ if self.use_rope:
+ q, k, v = qkv.unbind(dim=2)
+ q, k = self.rope(q, k, indices)
+ qkv = torch.stack([q, k, v], dim=2)
+ if self.attn_mode == "full":
+ if self.qk_rms_norm:
+ q, k, v = qkv.unbind(dim=2)
+ q = self.q_rms_norm(q)
+ k = self.k_rms_norm(k)
+ h = scaled_dot_product_attention(q, k, v)
+ else:
+ h = scaled_dot_product_attention(qkv)
+ elif self.attn_mode == "windowed":
+ raise NotImplementedError("Windowed attention is not yet implemented")
+ else:
+ Lkv = context.shape[1]
+ q = self.to_q(x)
+ kv = self.to_kv(context)
+ q = q.reshape(B, L, self.num_heads, -1)
+ kv = kv.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, k, v)
+ else:
+ h = scaled_dot_product_attention(q, kv)
+ h = h.reshape(B, L, -1)
+ h = self.to_out(h)
+ return h
+
+================
+File: norm.py
+================
+import torch
+import torch.nn as nn
+
+
+class LayerNorm32(nn.LayerNorm):
+ def forward(self, x: torch.Tensor) -> torch.Tensor:
+ return super().forward(x.float()).type(x.dtype)
+
+
+class GroupNorm32(nn.GroupNorm):
+ """
+ A GroupNorm layer that converts to float32 before the forward pass.
+ """
+ def forward(self, x: torch.Tensor) -> torch.Tensor:
+ return super().forward(x.float()).type(x.dtype)
+
+
+class ChannelLayerNorm32(LayerNorm32):
+ def forward(self, x: torch.Tensor) -> torch.Tensor:
+ DIM = x.dim()
+ x = x.permute(0, *range(2, DIM), 1).contiguous()
+ x = super().forward(x)
+ x = x.permute(0, DIM-1, *range(1, DIM-1)).contiguous()
+ return x
+
+================
+File: sparse/__init__.py
+================
+from typing import *
+import importlib
+import logging
+from trellis.backend_config import (
+ get_sparse_backend,
+ get_debug_mode,
+)
+
+# Get configuration from central backend config
+BACKEND = get_sparse_backend()
+DEBUG = get_debug_mode()
+
+def __from_env():
+ """Read current backend configuration"""
+ global BACKEND
+ global DEBUG
+
+ # Get current settings from central config
+ BACKEND = get_sparse_backend()
+ DEBUG = get_debug_mode()
+
+ print(f"[SPARSE] Using backend: {BACKEND}")
+
+logger = logging.getLogger(__name__)
+
+# Define module attributes and their corresponding files
+__attributes = {
+ 'SparseTensor': 'basic',
+ 'sparse_batch_broadcast': 'basic',
+ 'sparse_batch_op': 'basic',
+ 'sparse_cat': 'basic',
+ 'sparse_unbind': 'basic',
+ 'SparseGroupNorm': 'norm',
+ 'SparseLayerNorm': 'norm',
+ 'SparseGroupNorm32': 'norm',
+ 'SparseLayerNorm32': 'norm',
+ 'SparseReLU': 'nonlinearity',
+ 'SparseSiLU': 'nonlinearity',
+ 'SparseGELU': 'nonlinearity',
+ 'SparseActivation': 'nonlinearity',
+ 'SparseLinear': 'linear',
+ 'sparse_scaled_dot_product_attention': 'attention.full_attn',
+ 'SerializeMode': 'attention.serialized_attn',
+ 'sparse_serialized_scaled_dot_product_self_attention': 'attention.serialized_attn',
+ 'sparse_windowed_scaled_dot_product_self_attention': 'attention.windowed_attn',
+ 'SparseMultiHeadAttention': 'attention.modules',
+ 'SparseConv3d': 'conv',
+ 'SparseInverseConv3d': 'conv',
+ 'SparseDownsample': 'spatial',
+ 'SparseUpsample': 'spatial',
+ 'SparseSubdivide': 'spatial',
+}
+
+__submodules = [
+ 'transformer', # for blocks.py, modulated.py, etc.
+]
+
+__all__ = list(__attributes.keys()) + __submodules
+
+def __getattr__(name: str):
+ """Lazy import logic for submodules"""
+ if name in __attributes:
+ module_name = __attributes[name]
+ submod = importlib.import_module(f".{module_name}", __name__)
+ return getattr(submod, name)
+ elif name in __submodules:
+ return importlib.import_module(f".{name}", __name__)
+ else:
+ raise AttributeError(f"Module '{__name__}' has no attribute '{name}'.")
+
+logger.debug(f"Initialized trellis.modules.sparse package in lazy mode using backend: {BACKEND}")
+
+# Initialize settings from environment
+__from_env()
+
+# For Pylance/type checking
+if __name__ == '__main__':
+ from .basic import *
+ from .norm import *
+ from .nonlinearity import *
+ from .linear import *
+ from .attention import *
+ from .conv import *
+ from .spatial import *
+ import transformer
+
+================
+File: sparse/attention/__init__.py
+================
+from .full_attn import *
+from .serialized_attn import *
+from .windowed_attn import *
+from .modules import *
+
+# Re-export SerializeMode from serialized_attn
+from .serialized_attn import SerializeMode
+
+__all__ = [
+ 'sparse_scaled_dot_product_attention',
+ 'sparse_serialized_scaled_dot_product_self_attention',
+ 'sparse_windowed_scaled_dot_product_self_attention',
+ 'SparseMultiHeadAttention',
+ 'SerializeMode'
+]
+
+================
+File: sparse/attention/full_attn.py
+================
+from typing import *
+import torch
+import math
+from .. import SparseTensor
+from trellis.backend_config import (
+ get_attention_backend,
+ get_debug_mode,
+ get_available_backends
+)
+import logging
+
+logger = logging.getLogger(__name__)
+
+# Get configuration from central config
+ATTN_BACKEND = get_attention_backend()
+DEBUG = get_debug_mode()
+
+# Get available backends and import if active
+available_backends = get_available_backends()
+
+if ATTN_BACKEND == "xformers" and available_backends['xformers']:
+ import xformers.ops as xops
+elif ATTN_BACKEND == "flash_attn" and available_backends['flash_attn']:
+ import flash_attn
+elif ATTN_BACKEND == "sage" and available_backends['sage']:
+ import sageattention
+ sageattn = sageattention.sageattn
+ from sageattention.core import sageattn_varlen
+elif ATTN_BACKEND == "sdpa":
+ from torch.nn.functional import scaled_dot_product_attention as sdpa
+
+# Log the active backend
+logger.info(f"Using attention backend: {ATTN_BACKEND}")
+
+__all__ = ['sparse_scaled_dot_product_attention']
+
+@overload
+def sparse_scaled_dot_product_attention(qkv: SparseTensor) -> SparseTensor:
+ """
+ Apply scaled dot product attention to a sparse tensor.
+
+ Args:
+ qkv (SparseTensor): A [N, *, 3, H, C] sparse tensor containing Qs, Ks, and Vs.
+ """
+ ...
+
+@overload
+def sparse_scaled_dot_product_attention(q: SparseTensor, kv: Union[SparseTensor, torch.Tensor]) -> SparseTensor:
+ """
+ Apply scaled dot product attention to a sparse tensor.
+
+ Args:
+ q (SparseTensor): A [N, *, H, C] sparse tensor containing Qs.
+ kv (SparseTensor or torch.Tensor): A [N, *, 2, H, C] sparse tensor or a [N, L, 2, H, C] dense tensor containing Ks and Vs.
+ """
+ ...
+
+@overload
+def sparse_scaled_dot_product_attention(q: torch.Tensor, kv: SparseTensor) -> torch.Tensor:
+ """
+ Apply scaled dot product attention to a sparse tensor.
+
+ Args:
+ q (SparseTensor): A [N, L, H, C] dense tensor containing Qs.
+ kv (SparseTensor or torch.Tensor): A [N, *, 2, H, C] sparse tensor containing Ks and Vs.
+ """
+ ...
+
+@overload
+def sparse_scaled_dot_product_attention(q: SparseTensor, k: SparseTensor, v: SparseTensor) -> SparseTensor:
+ """
+ Apply scaled dot product attention to a sparse tensor.
+
+ Args:
+ q (SparseTensor): A [N, *, H, Ci] sparse tensor containing Qs.
+ k (SparseTensor): A [N, *, H, Ci] sparse tensor containing Ks.
+ v (SparseTensor): A [N, *, H, Co] sparse tensor containing Vs.
+
+ Note:
+ k and v are assumed to have the same coordinate map.
+ """
+ ...
+
+@overload
+def sparse_scaled_dot_product_attention(q: SparseTensor, k: torch.Tensor, v: torch.Tensor) -> SparseTensor:
+ """
+ Apply scaled dot product attention to a sparse tensor.
+
+ Args:
+ q (SparseTensor): A [N, *, H, Ci] sparse tensor containing Qs.
+ k (torch.Tensor): A [N, L, H, Ci] dense tensor containing Ks.
+ v (torch.Tensor): A [N, L, H, Co] dense tensor containing Vs.
+ """
+ ...
+
+@overload
+def sparse_scaled_dot_product_attention(q: torch.Tensor, k: SparseTensor, v: SparseTensor) -> torch.Tensor:
+ """
+ Apply scaled dot product attention to a sparse tensor.
+
+ Args:
+ q (torch.Tensor): A [N, L, H, Ci] dense tensor containing Qs.
+ k (SparseTensor): A [N, *, H, Ci] sparse tensor containing Ks.
+ v (SparseTensor): A [N, *, H, Co] sparse tensor containing Vs.
+ """
+ ...
+
+def sparse_scaled_dot_product_attention(*args, **kwargs):
+ arg_names_dict = {
+ 1: ['qkv'],
+ 2: ['q', 'kv'],
+ 3: ['q', 'k', 'v']
+ }
+ num_all_args = len(args) + len(kwargs)
+ assert num_all_args in arg_names_dict, f"Invalid number of arguments, got {num_all_args}, expected 1, 2, or 3"
+ for key in arg_names_dict[num_all_args][len(args):]:
+ assert key in kwargs, f"Missing argument {key}"
+
+ if num_all_args == 1:
+ qkv = args[0] if len(args) > 0 else kwargs['qkv']
+ assert isinstance(qkv, SparseTensor), f"qkv must be a SparseTensor, got {type(qkv)}"
+ assert len(qkv.shape) == 4 and qkv.shape[1] == 3, f"Invalid shape for qkv, got {qkv.shape}, expected [N, *, 3, H, C]"
+ device = qkv.device
+
+ s = qkv
+ q_seqlen = [qkv.layout[i].stop - qkv.layout[i].start for i in range(qkv.shape[0])]
+ kv_seqlen = q_seqlen
+ qkv = qkv.feats # [T, 3, H, C]
+
+ elif num_all_args == 2:
+ q = args[0] if len(args) > 0 else kwargs['q']
+ kv = args[1] if len(args) > 1 else kwargs['kv']
+ assert isinstance(q, SparseTensor) and isinstance(kv, (SparseTensor, torch.Tensor)) or \
+ isinstance(q, torch.Tensor) and isinstance(kv, SparseTensor), \
+ f"Invalid types, got {type(q)} and {type(kv)}"
+ assert q.shape[0] == kv.shape[0], f"Batch size mismatch, got {q.shape[0]} and {kv.shape[0]}"
+ device = q.device
+
+ if isinstance(q, SparseTensor):
+ assert len(q.shape) == 3, f"Invalid shape for q, got {q.shape}, expected [N, *, H, C]"
+ s = q
+ q_seqlen = [q.layout[i].stop - q.layout[i].start for i in range(q.shape[0])]
+ q = q.feats # [T_Q, H, C]
+ else:
+ assert len(q.shape) == 4, f"Invalid shape for q, got {q.shape}, expected [N, L, H, C]"
+ s = None
+ N, L, H, C = q.shape
+ q_seqlen = [L] * N
+ q = q.reshape(N * L, H, C) # [T_Q, H, C]
+
+ if isinstance(kv, SparseTensor):
+ assert len(kv.shape) == 4 and kv.shape[1] == 2, f"Invalid shape for kv, got {kv.shape}, expected [N, *, 2, H, C]"
+ kv_seqlen = [kv.layout[i].stop - kv.layout[i].start for i in range(kv.shape[0])]
+ kv = kv.feats # [T_KV, 2, H, C]
+ else:
+ assert len(kv.shape) == 5, f"Invalid shape for kv, got {kv.shape}, expected [N, L, 2, H, C]"
+ N, L, _, H, C = kv.shape
+ kv_seqlen = [L] * N
+ kv = kv.reshape(N * L, 2, H, C) # [T_KV, 2, H, C]
+
+ elif num_all_args == 3:
+ q = args[0] if len(args) > 0 else kwargs['q']
+ k = args[1] if len(args) > 1 else kwargs['k']
+ v = args[2] if len(args) > 2 else kwargs['v']
+ assert isinstance(q, SparseTensor) and isinstance(k, (SparseTensor, torch.Tensor)) and type(k) == type(v) or \
+ isinstance(q, torch.Tensor) and isinstance(k, SparseTensor) and isinstance(v, SparseTensor), \
+ f"Invalid types, got {type(q)}, {type(k)}, and {type(v)}"
+ assert q.shape[0] == k.shape[0] == v.shape[0], f"Batch size mismatch, got {q.shape[0]}, {k.shape[0]}, and {v.shape[0]}"
+ device = q.device
+
+ if isinstance(q, SparseTensor):
+ assert len(q.shape) == 3, f"Invalid shape for q, got {q.shape}, expected [N, *, H, Ci]"
+ s = q
+ q_seqlen = [q.layout[i].stop - q.layout[i].start for i in range(q.shape[0])]
+ q = q.feats # [T_Q, H, Ci]
+ else:
+ assert len(q.shape) == 4, f"Invalid shape for q, got {q.shape}, expected [N, L, H, Ci]"
+ s = None
+ N, L, H, CI = q.shape
+ q_seqlen = [L] * N
+ q = q.reshape(N * L, H, CI) # [T_Q, H, Ci]
+
+ if isinstance(k, SparseTensor):
+ assert len(k.shape) == 3, f"Invalid shape for k, got {k.shape}, expected [N, *, H, Ci]"
+ assert len(v.shape) == 3, f"Invalid shape for v, got {v.shape}, expected [N, *, H, Co]"
+ kv_seqlen = [k.layout[i].stop - k.layout[i].start for i in range(k.shape[0])]
+ k = k.feats # [T_KV, H, Ci]
+ v = v.feats # [T_KV, H, Co]
+ else:
+ assert len(k.shape) == 4, f"Invalid shape for k, got {k.shape}, expected [N, L, H, Ci]"
+ assert len(v.shape) == 4, f"Invalid shape for v, got {v.shape}, expected [N, L, H, Co]"
+ N, L, H, CI, CO = *k.shape, v.shape[-1]
+ kv_seqlen = [L] * N
+ k = k.reshape(N * L, H, CI) # [T_KV, H, Ci]
+ v = v.reshape(N * L, H, CO) # [T_KV, H, Co]
+
+ if DEBUG:
+ if s is not None:
+ for i in range(s.shape[0]):
+ assert (s.coords[s.layout[i]] == i).all(), f"SparseScaledDotProductSelfAttention: batch index mismatch"
+ if num_all_args in [2, 3]:
+ assert q.shape[:2] == [1, sum(q_seqlen)], f"SparseScaledDotProductSelfAttention: q shape mismatch"
+ if num_all_args == 3:
+ assert k.shape[:2] == [1, sum(kv_seqlen)], f"SparseScaledDotProductSelfAttention: k shape mismatch"
+ assert v.shape[:2] == [1, sum(kv_seqlen)], f"SparseScaledDotProductSelfAttention: v shape mismatch"
+
+ if ATTN_BACKEND == 'xformers':
+ if num_all_args == 1:
+ q, k, v = qkv.unbind(dim=1)
+ elif num_all_args == 2:
+ k, v = kv.unbind(dim=1)
+ q = q.unsqueeze(0)
+ k = k.unsqueeze(0)
+ v = v.unsqueeze(0)
+ mask = xops.fmha.BlockDiagonalMask.from_seqlens(q_seqlen, kv_seqlen)
+ out = xops.memory_efficient_attention(q, k, v, mask)[0]
+ elif ATTN_BACKEND == 'flash_attn':
+ cu_seqlens_q = torch.cat([torch.tensor([0]), torch.cumsum(torch.tensor(q_seqlen), dim=0)]).int().to(device)
+ if num_all_args in [2, 3]:
+ cu_seqlens_kv = torch.cat([torch.tensor([0]), torch.cumsum(torch.tensor(kv_seqlen), dim=0)]).int().to(device)
+ if num_all_args == 1:
+ out = flash_attn.flash_attn_varlen_qkvpacked_func(qkv, cu_seqlens_q, max(q_seqlen))
+ elif num_all_args == 2:
+ out = flash_attn.flash_attn_varlen_kvpacked_func(q, kv, cu_seqlens_q, cu_seqlens_kv, max(q_seqlen), max(kv_seqlen))
+ elif num_all_args == 3:
+ out = flash_attn.flash_attn_varlen_func(q, k, v, cu_seqlens_q, cu_seqlens_kv, max(q_seqlen), max(kv_seqlen))
+ elif ATTN_BACKEND == 'sage':
+ # SAGE branch
+ # Usually you have a [N, L, H, C] or [T, H, C], etc. shape in a sparse representation.
+ # You must reshape or reorder so it's [N, H, L, C].
+ # Then call sageattn(q, k, v, is_causal=..., tensor_layout="HND")
+ # or perhaps sageattn_varlen(...) if your batch sequences differ in length.
+
+ # Example if your 'q' is shaped [N, L, H, C]:
+ q = q.permute(0, 2, 1, 3).contiguous()
+ k = k.permute(0, 2, 1, 3).contiguous()
+ v = v.permute(0, 2, 1, 3).contiguous()
+ out = sageattn(q, k, v, tensor_layout="HND")
+ out = out.permute(0, 2, 1, 3) # back to [N, L, H, C]
+
+ # Replace feats in SparseTensor if needed
+ # e.g. if 's' is your SparseTensor handle:
+ # return s.replace(out)
+ else:
+ raise ValueError(f"Unknown attention module: {ATTN_BACKEND}")
+
+ if s is not None:
+ return s.replace(out)
+ else:
+ return out.reshape(N, L, H, -1)
+
+================
+File: sparse/attention/modules.py
+================
+from typing import *
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from .. import SparseTensor
+from .full_attn import sparse_scaled_dot_product_attention
+from .serialized_attn import SerializeMode, sparse_serialized_scaled_dot_product_self_attention
+from .windowed_attn import sparse_windowed_scaled_dot_product_self_attention
+from ...attention import RotaryPositionEmbedder
+
+
+class SparseMultiHeadRMSNorm(nn.Module):
+ def __init__(self, dim: int, heads: int):
+ super().__init__()
+ self.scale = dim ** 0.5
+ self.gamma = nn.Parameter(torch.ones(heads, dim))
+
+ def forward(self, x: Union[SparseTensor, torch.Tensor]) -> Union[SparseTensor, torch.Tensor]:
+ x_type = x.dtype
+ x = x.float()
+ if isinstance(x, SparseTensor):
+ x = x.replace(F.normalize(x.feats, dim=-1))
+ else:
+ x = F.normalize(x, dim=-1)
+ return (x * self.gamma * self.scale).to(x_type)
+
+
+class SparseMultiHeadAttention(nn.Module):
+ def __init__(
+ self,
+ channels: int,
+ num_heads: int,
+ ctx_channels: Optional[int] = None,
+ type: Literal["self", "cross"] = "self",
+ attn_mode: Literal["full", "serialized", "windowed"] = "full",
+ window_size: Optional[int] = None,
+ shift_sequence: Optional[int] = None,
+ shift_window: Optional[Tuple[int, int, int]] = None,
+ serialize_mode: Optional[SerializeMode] = None,
+ qkv_bias: bool = True,
+ use_rope: bool = False,
+ qk_rms_norm: bool = False,
+ ):
+ super().__init__()
+ assert channels % num_heads == 0
+ assert type in ["self", "cross"], f"Invalid attention type: {type}"
+ assert attn_mode in ["full", "serialized", "windowed"], f"Invalid attention mode: {attn_mode}"
+ assert type == "self" or attn_mode == "full", "Cross-attention only supports full attention"
+ assert type == "self" or use_rope is False, "Rotary position embeddings only supported for self-attention"
+ self.channels = channels
+ self.ctx_channels = ctx_channels if ctx_channels is not None else channels
+ self.num_heads = num_heads
+ self._type = type
+ self.attn_mode = attn_mode
+ self.window_size = window_size
+ self.shift_sequence = shift_sequence
+ self.shift_window = shift_window
+ self.serialize_mode = serialize_mode
+ self.use_rope = use_rope
+ 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 = SparseMultiHeadRMSNorm(channels // num_heads, num_heads)
+ self.k_rms_norm = SparseMultiHeadRMSNorm(channels // num_heads, num_heads)
+
+ self.to_out = nn.Linear(channels, channels)
+
+ if use_rope:
+ self.rope = RotaryPositionEmbedder(channels)
+
+ @staticmethod
+ def _linear(module: nn.Linear, x: Union[SparseTensor, torch.Tensor]) -> Union[SparseTensor, torch.Tensor]:
+ if isinstance(x, SparseTensor):
+ # Reshape features to match expected dimensions
+ feats = x.feats
+ if feats.shape[-1] != module.in_features:
+ feats = feats.view(-1, module.in_features)
+ return x.replace(module(feats))
+ else:
+ # Reshape tensor to match expected dimensions
+ if x.shape[-1] != module.in_features:
+ x = x.view(-1, module.in_features)
+ return module(x)
+
+ @staticmethod
+ def _reshape_chs(x: Union[SparseTensor, torch.Tensor], shape: Tuple[int, ...]) -> Union[SparseTensor, torch.Tensor]:
+ if isinstance(x, SparseTensor):
+ return x.reshape(*shape)
+ else:
+ return x.reshape(*x.shape[:2], *shape)
+
+ def _fused_pre(self, x: Union[SparseTensor, torch.Tensor], num_fused: int) -> Union[SparseTensor, torch.Tensor]:
+ if isinstance(x, SparseTensor):
+ x_feats = x.feats.unsqueeze(0)
+ else:
+ x_feats = x
+ x_feats = x_feats.reshape(*x_feats.shape[:2], num_fused, self.num_heads, -1)
+ return x.replace(x_feats.squeeze(0)) if isinstance(x, SparseTensor) else x_feats
+
+ def _rope(self, qkv: SparseTensor) -> SparseTensor:
+ q, k, v = qkv.feats.unbind(dim=1) # [T, H, C]
+ q, k = self.rope(q, k, qkv.coords[:, 1:])
+ qkv = qkv.replace(torch.stack([q, k, v], dim=1))
+ return qkv
+
+ def forward(self, x: Union[SparseTensor, torch.Tensor], context: Optional[Union[SparseTensor, torch.Tensor]] = None) -> Union[SparseTensor, torch.Tensor]:
+ if self._type == "self":
+ qkv = self._linear(self.to_qkv, x)
+ qkv = self._fused_pre(qkv, num_fused=3)
+ if self.use_rope:
+ qkv = self._rope(qkv)
+ if self.qk_rms_norm:
+ q, k, v = qkv.unbind(dim=1)
+ q = self.q_rms_norm(q)
+ k = self.k_rms_norm(k)
+ qkv = qkv.replace(torch.stack([q.feats, k.feats, v.feats], dim=1))
+ if self.attn_mode == "full":
+ h = sparse_scaled_dot_product_attention(qkv)
+ elif self.attn_mode == "serialized":
+ h = sparse_serialized_scaled_dot_product_self_attention(
+ qkv, self.window_size, serialize_mode=self.serialize_mode, shift_sequence=self.shift_sequence, shift_window=self.shift_window
+ )
+ elif self.attn_mode == "windowed":
+ h = sparse_windowed_scaled_dot_product_self_attention(
+ qkv, self.window_size, shift_window=self.shift_window, qk_rms_norm=self.qk_rms_norm
+ )
+ else:
+ q = self._linear(self.to_q, x)
+ q = self._reshape_chs(q, (self.num_heads, -1))
+ kv = self._linear(self.to_kv, context)
+ kv = self._fused_pre(kv, num_fused=2)
+ if self.qk_rms_norm:
+ q = self.q_rms_norm(q)
+ k, v = kv.unbind(dim=1)
+ k = self.k_rms_norm(k)
+ kv = kv.replace(torch.stack([k.feats, v.feats], dim=1))
+ h = sparse_scaled_dot_product_attention(q, kv)
+ h = self._reshape_chs(h, (-1,))
+ h = self._linear(self.to_out, h)
+ return h
+
+================
+File: sparse/attention/serialized_attn.py
+================
+from typing import *
+from enum import Enum
+import torch
+import math
+from .. import SparseTensor
+from trellis.backend_config import (
+ get_attention_backend,
+ get_debug_mode,
+ get_available_backends
+)
+import logging
+
+logger = logging.getLogger(__name__)
+
+# Get configuration from central config
+ATTN_BACKEND = get_attention_backend()
+DEBUG = get_debug_mode()
+
+# Get available backends and import if active
+available_backends = get_available_backends()
+
+if ATTN_BACKEND == "xformers" and available_backends['xformers']:
+ import xformers.ops as xops
+elif ATTN_BACKEND == "flash_attn" and available_backends['flash_attn']:
+ import flash_attn
+elif ATTN_BACKEND == "sage" and available_backends['sage']:
+ import sageattention
+ sageattn = sageattention.sageattn
+ from sageattention.core import sageattn_varlen
+elif ATTN_BACKEND == "sdpa":
+ from torch.nn.functional import scaled_dot_product_attention as sdpa
+
+# Log the active backend
+logger.info(f"Using attention backend: {ATTN_BACKEND}")
+
+__all__ = [
+ 'SerializeMode',
+ 'sparse_serialized_scaled_dot_product_self_attention',
+]
+
+class SerializeMode(Enum):
+ Z_ORDER = 0
+ Z_ORDER_TRANSPOSED = 1
+ HILBERT = 2
+ HILBERT_TRANSPOSED = 3
+
+SerializeModes = [
+ SerializeMode.Z_ORDER,
+ SerializeMode.Z_ORDER_TRANSPOSED,
+ SerializeMode.HILBERT,
+ SerializeMode.HILBERT_TRANSPOSED
+]
+
+def calc_serialization(
+ tensor: SparseTensor,
+ window_size: int,
+ serialize_mode: SerializeMode = SerializeMode.Z_ORDER,
+ shift_sequence: int = 0,
+ shift_window: Tuple[int, int, int] = (0, 0, 0)
+) -> Tuple[torch.Tensor, torch.Tensor, List[int]]:
+ """
+ Calculate serialization and partitioning for a set of coordinates.
+
+ Args:
+ tensor (SparseTensor): The input tensor.
+ window_size (int): The window size to use.
+ serialize_mode (SerializeMode): The serialization mode to use.
+ shift_sequence (int): The shift of serialized sequence.
+ shift_window (Tuple[int, int, int]): The shift of serialized coordinates.
+
+ Returns:
+ (torch.Tensor, torch.Tensor): Forwards and backwards indices.
+ """
+ fwd_indices = []
+ bwd_indices = []
+ seq_lens = []
+ seq_batch_indices = []
+ offsets = [0]
+
+ if 'vox2seq' not in globals():
+ import vox2seq
+
+ # Serialize the input
+ serialize_coords = tensor.coords[:, 1:].clone()
+ serialize_coords += torch.tensor(shift_window, dtype=torch.int32, device=tensor.device).reshape(1, 3)
+ if serialize_mode == SerializeMode.Z_ORDER:
+ code = vox2seq.encode(serialize_coords, mode='z_order', permute=[0, 1, 2])
+ elif serialize_mode == SerializeMode.Z_ORDER_TRANSPOSED:
+ code = vox2seq.encode(serialize_coords, mode='z_order', permute=[1, 0, 2])
+ elif serialize_mode == SerializeMode.HILBERT:
+ code = vox2seq.encode(serialize_coords, mode='hilbert', permute=[0, 1, 2])
+ elif serialize_mode == SerializeMode.HILBERT_TRANSPOSED:
+ code = vox2seq.encode(serialize_coords, mode='hilbert', permute=[1, 0, 2])
+ else:
+ raise ValueError(f"Unknown serialize mode: {serialize_mode}")
+
+ for bi, s in enumerate(tensor.layout):
+ num_points = s.stop - s.start
+ num_windows = (num_points + window_size - 1) // window_size
+ valid_window_size = num_points / num_windows
+ to_ordered = torch.argsort(code[s.start:s.stop])
+ if num_windows == 1:
+ fwd_indices.append(to_ordered)
+ bwd_indices.append(torch.zeros_like(to_ordered).scatter_(0, to_ordered, torch.arange(num_points, device=tensor.device)))
+ fwd_indices[-1] += s.start
+ bwd_indices[-1] += offsets[-1]
+ seq_lens.append(num_points)
+ seq_batch_indices.append(bi)
+ offsets.append(offsets[-1] + seq_lens[-1])
+ else:
+ # Partition the input
+ offset = 0
+ mids = [(i + 0.5) * valid_window_size + shift_sequence for i in range(num_windows)]
+ split = [math.floor(i * valid_window_size + shift_sequence) for i in range(num_windows + 1)]
+ bwd_index = torch.zeros((num_points,), dtype=torch.int64, device=tensor.device)
+ for i in range(num_windows):
+ mid = mids[i]
+ valid_start = split[i]
+ valid_end = split[i + 1]
+ padded_start = math.floor(mid - 0.5 * window_size)
+ padded_end = padded_start + window_size
+ fwd_indices.append(to_ordered[torch.arange(padded_start, padded_end, device=tensor.device) % num_points])
+ offset += valid_start - padded_start
+ bwd_index.scatter_(0, fwd_indices[-1][valid_start-padded_start:valid_end-padded_start], torch.arange(offset, offset + valid_end - valid_start, device=tensor.device))
+ offset += padded_end - valid_start
+ fwd_indices[-1] += s.start
+ seq_lens.extend([window_size] * num_windows)
+ seq_batch_indices.extend([bi] * num_windows)
+ bwd_indices.append(bwd_index + offsets[-1])
+ offsets.append(offsets[-1] + num_windows * window_size)
+
+ fwd_indices = torch.cat(fwd_indices)
+ bwd_indices = torch.cat(bwd_indices)
+
+ return fwd_indices, bwd_indices, seq_lens, seq_batch_indices
+
+
+ '''fwd_indices = torch.arange(tensor.feats.shape[0], device=tensor.device)
+ bwd_indices = torch.arange(tensor.feats.shape[0], device=tensor.device)
+ seq_lens = [tensor.feats.shape[0]]
+ seq_batch_indices = [0]
+ return fwd_indices, bwd_indices, seq_lens, seq_batch_indices'''
+
+
+def sparse_serialized_scaled_dot_product_self_attention(
+ qkv: SparseTensor,
+ window_size: int,
+ serialize_mode: SerializeMode = SerializeMode.Z_ORDER,
+ shift_sequence: int = 0,
+ shift_window: Tuple[int, int, int] = (0, 0, 0)
+) -> SparseTensor:
+ """
+ Apply serialized scaled dot product self attention to a sparse tensor.
+
+ Args:
+ qkv (SparseTensor): shape [N, *, 3, H, C]
+ window_size (int)
+ serialize_mode (SerializeMode)
+ shift_sequence (int)
+ shift_window (Tuple[int,int,int])
+
+ Returns:
+ SparseTensor with updated .feats
+ """
+ assert len(qkv.shape) == 4 and qkv.shape[1] == 3, (
+ f"Invalid shape for qkv, got {qkv.shape}, expected [N, *, 3, H, C]"
+ )
+
+ # Retrieve or compute serialization indices
+ serialization_cache_name = f'serialization_{serialize_mode}_{window_size}_{shift_sequence}_{shift_window}'
+ serialization_spatial_cache = qkv.get_spatial_cache(serialization_cache_name)
+ if serialization_spatial_cache is None:
+ fwd_indices, bwd_indices, seq_lens, seq_batch_indices = calc_serialization(
+ qkv, window_size, serialize_mode, shift_sequence, shift_window
+ )
+ qkv.register_spatial_cache(serialization_cache_name,
+ (fwd_indices, bwd_indices, seq_lens, seq_batch_indices))
+ else:
+ fwd_indices, bwd_indices, seq_lens, seq_batch_indices = serialization_spatial_cache
+
+ # Gather QKV feats in new order
+ qkv_feats = qkv.feats[fwd_indices] # [M, 3, H, C]
+ M = qkv_feats.shape[0]
+ H = qkv_feats.shape[2]
+ C = qkv_feats.shape[3]
+
+ # Handle uniform vs. variable chunk sizes
+ if all(s == window_size for s in seq_lens):
+ # Uniform chunk size
+ B = len(seq_lens)
+ N = window_size # chunk length
+ assert B * N == M, "Mismatched total tokens (B*N != M)"
+
+ qkv_feats = qkv_feats.reshape(B, N, 3, H, C) # [B, N, 3, H, C]
+
+ if ATTN_BACKEND == 'xformers':
+ q, k, v = qkv_feats.unbind(dim=2) # [B, N, H, C]
+ out = xops.memory_efficient_attention(q, k, v) # [B, N, H, C]
+ out = out.reshape(M, H, C) # [M, H, C]
+ elif ATTN_BACKEND == 'flash_attn':
+ out = flash_attn.flash_attn_qkvpacked_func(qkv_feats) # [B, N, H, C]
+ out = out.reshape(M, H, C)
+ elif ATTN_BACKEND == 'sage':
+ # SageAttention path
+ # Re-permute to [B, 3, H, N, C]
+ qkv_feats = qkv_feats.permute(0, 2, 3, 1, 4).contiguous() # [B, 3, H, N, C]
+ q, k, v = qkv_feats.unbind(dim=1) # each => [B, H, N, C]
+
+ out = sageattn(q, k, v, tensor_layout="HND") # => [B, H, N, C]
+ # Flatten back
+ out = out.permute(0, 2, 1, 3).contiguous() # => [B, N, H, C]
+ out = out.reshape(M, H, C)
+ else:
+ raise ValueError(f"Unknown attention module: {ATTN_BACKEND}")
+
+ else:
+ # Variable-length chunks
+ if ATTN_BACKEND == 'xformers':
+ q, k, v = qkv_feats.unbind(dim=1) # [M, H, C]
+ q = q.unsqueeze(0)
+ k = k.unsqueeze(0)
+ v = v.unsqueeze(0)
+ mask = xops.fmha.BlockDiagonalMask.from_seqlens(seq_lens)
+ out = xops.memory_efficient_attention(q, k, v, mask=mask)[0] # [M, H, C]
+ elif ATTN_BACKEND == 'flash_attn':
+ cu_seqlens = torch.zeros(len(seq_lens)+1,
+ dtype=torch.int32,
+ device=qkv.device)
+ cu_seqlens[1:] = torch.cumsum(torch.tensor(seq_lens,
+ device=qkv.device), dim=0)
+ out = flash_attn.flash_attn_varlen_qkvpacked_func(qkv_feats, cu_seqlens, max(seq_lens)) # => [M, H, C]
+ elif ATTN_BACKEND == 'sage':
+ # SageAttention varlen path
+ # If you want to do "varlen," you can do something like:
+ cu_seqlens = torch.zeros(len(seq_lens)+1, dtype=torch.int32, device=qkv.device)
+ cu_seqlens[1:] = torch.cumsum(torch.tensor(seq_lens, device=qkv.device), dim=0)
+
+ # qkv_feats => [M, 3, H, C], reorder to [3, M, H, C]
+ qkv_feats = qkv_feats.permute(1, 0, 2, 3).contiguous() # [3, M, H, C]
+ q, k, v = qkv_feats[0], qkv_feats[1], qkv_feats[2] # each => [M, H, C]
+
+ # SageAttention varlen typically expects q, k, v => [M, H, D], etc.
+ # For is_causal or not, pass as needed
+ #from sageattention import sageattn_varlen # or define at top
+ max_len = max(seq_lens)
+ out = sageattn_varlen(
+ q, k, v,
+ cu_seqlens, cu_seqlens,
+ max_len, max_len,
+ is_causal=False
+ ) # => [M, H, C]
+ else:
+ raise ValueError(f"Unknown attention module: {ATTN_BACKEND}")
+
+ # Re-map via bwd_indices => original order
+ out = out[bwd_indices] # [T, H, C]
+ return qkv.replace(out)
+
+================
+File: sparse/attention/windowed_attn.py
+================
+from typing import *
+import torch
+import math
+import os
+import logging
+import torch
+from .. import SparseTensor
+from trellis.backend_config import (
+ get_attention_backend,
+ get_debug_mode,
+ get_available_backends
+)
+import logging
+
+logger = logging.getLogger(__name__)
+
+# Get configuration from central config
+ATTN_BACKEND = get_attention_backend()
+DEBUG = get_debug_mode()
+
+# Get available backends and import if active
+available_backends = get_available_backends()
+
+if ATTN_BACKEND == "xformers" and available_backends['xformers']:
+ import xformers.ops as xops
+elif ATTN_BACKEND == "flash_attn" and available_backends['flash_attn']:
+ import flash_attn
+elif ATTN_BACKEND == "sage" and available_backends['sage']:
+ import sageattention
+ sageattn = sageattention.sageattn
+ from sageattention.core import sageattn_varlen
+elif ATTN_BACKEND == "sdpa":
+ from torch.nn.functional import scaled_dot_product_attention as sdpa
+
+# Log the active backend
+logger.info(f"Using attention backend: {ATTN_BACKEND}")
+
+__all__ = [
+ "WindowedMode",
+ 'sparse_windowed_scaled_dot_product_self_attention',
+]
+
+def calc_window_partition(
+ tensor: SparseTensor,
+ window_size: Union[int, Tuple[int, ...]],
+ shift_window: Union[int, Tuple[int, ...]] = 0
+) -> Tuple[torch.Tensor, torch.Tensor, List[int], List[int]]:
+ """
+ Calculate serialization and partitioning for a set of coordinates.
+
+ Args:
+ tensor (SparseTensor): The input tensor.
+ window_size (int): The window size to use.
+ shift_window (Tuple[int, ...]): The shift of serialized coordinates.
+
+ Returns:
+ (torch.Tensor): Forwards indices.
+ (torch.Tensor): Backwards indices.
+ (List[int]): Sequence lengths.
+ (List[int]): Sequence batch indices.
+ """
+ DIM = tensor.coords.shape[1] - 1
+ shift_window = (shift_window,) * DIM if isinstance(shift_window, int) else shift_window
+ window_size = (window_size,) * DIM if isinstance(window_size, int) else window_size
+ shifted_coords = tensor.coords.clone().detach()
+ shifted_coords[:, 1:] += torch.tensor(shift_window, device=tensor.device, dtype=torch.int32).unsqueeze(0)
+
+ MAX_COORDS = shifted_coords[:, 1:].max(dim=0).values.tolist()
+ NUM_WINDOWS = [math.ceil((mc + 1) / ws) for mc, ws in zip(MAX_COORDS, window_size)]
+ OFFSET = torch.cumprod(torch.tensor([1] + NUM_WINDOWS[::-1]), dim=0).tolist()[::-1]
+
+ shifted_coords[:, 1:] //= torch.tensor(window_size, device=tensor.device, dtype=torch.int32).unsqueeze(0)
+ shifted_indices = (shifted_coords * torch.tensor(OFFSET, device=tensor.device, dtype=torch.int32).unsqueeze(0)).sum(dim=1)
+ fwd_indices = torch.argsort(shifted_indices)
+ bwd_indices = torch.empty_like(fwd_indices)
+ bwd_indices[fwd_indices] = torch.arange(fwd_indices.shape[0], device=tensor.device)
+ seq_lens = torch.bincount(shifted_indices)
+ seq_batch_indices = torch.arange(seq_lens.shape[0], device=tensor.device, dtype=torch.int32) // OFFSET[0]
+ mask = seq_lens != 0
+ seq_lens = seq_lens[mask].tolist()
+ seq_batch_indices = seq_batch_indices[mask].tolist()
+
+ return fwd_indices, bwd_indices, seq_lens, seq_batch_indices
+
+
+ '''fwd_indices = torch.arange(tensor.feats.shape[0], device=tensor.device)
+ bwd_indices = torch.arange(tensor.feats.shape[0], device=tensor.device)
+ seq_lens = [tensor.feats.shape[0]]
+ seq_batch_indices = [0]
+ return fwd_indices, bwd_indices, seq_lens, seq_batch_indices'''
+
+
+def sparse_windowed_scaled_dot_product_self_attention(
+ qkv: SparseTensor,
+ window_size: int,
+ shift_window: Tuple[int, int, int] = (0, 0, 0)
+) -> SparseTensor:
+ """
+ Apply windowed scaled dot product self attention to a sparse tensor.
+
+ Args:
+ qkv (SparseTensor): shape [N, *, 3, H, C]
+ window_size (int)
+ shift_window (Tuple[int,int,int])
+
+ Returns:
+ SparseTensor with updated .feats
+ """
+ assert len(qkv.shape) == 4 and qkv.shape[1] == 3, (
+ f"Invalid shape for qkv, got {qkv.shape}, expected [N, *, 3, H, C]"
+ )
+
+ cache_name = f'window_partition_{window_size}_{shift_window}'
+ window_spatial_cache = qkv.get_spatial_cache(cache_name)
+ if window_spatial_cache is None:
+ fwd_indices, bwd_indices, seq_lens, seq_batch_indices = calc_window_partition(
+ qkv, window_size, shift_window
+ )
+ qkv.register_spatial_cache(cache_name, (fwd_indices, bwd_indices, seq_lens, seq_batch_indices))
+ else:
+ fwd_indices, bwd_indices, seq_lens, seq_batch_indices = window_spatial_cache
+
+ M = fwd_indices.shape[0]
+ qkv_feats = qkv.feats[fwd_indices] # [M, 3, H, C]
+ H = qkv_feats.shape[2]
+ C = qkv_feats.shape[3]
+
+ if all(s == window_size for s in seq_lens):
+ # Uniform chunk size
+ B = len(seq_lens)
+ N = window_size
+ assert B * N == M, "Mismatched total tokens"
+
+ qkv_feats = qkv_feats.reshape(B, N, 3, H, C)
+
+ if ATTN_BACKEND == 'xformers':
+ q, k, v = qkv_feats.unbind(dim=2) # [B, N, H, C]
+ out = xops.memory_efficient_attention(q, k, v) # [B, N, H, C]
+ out = out.reshape(M, H, C)
+ elif ATTN_BACKEND == 'flash_attn':
+ out = flash_attn.flash_attn_qkvpacked_func(qkv_feats) # [B, N, H, C]
+ out = out.reshape(M, H, C)
+ elif ATTN_BACKEND == 'sage':
+ # SageAttention uniform path
+ # permute => [B, 3, H, N, C]
+ qkv_feats = qkv_feats.permute(0, 2, 3, 1, 4).contiguous()
+ q, k, v = qkv_feats.unbind(dim=1) # => [B, H, N, C]
+
+ out = sageattn(q, k, v, tensor_layout="HND")
+ # => [B, H, N, C]
+ out = out.permute(0, 2, 1, 3).contiguous() # => [B, N, H, C]
+ out = out.reshape(M, H, C)
+ else:
+ raise ValueError(f"Unknown attention module: {ATTN_BACKEND}")
+
+ else:
+ # Varlen branch
+ if ATTN_BACKEND == 'xformers':
+ q, k, v = qkv_feats.unbind(dim=1) # [M, H, C]
+ q = q.unsqueeze(0)
+ k = k.unsqueeze(0)
+ v = v.unsqueeze(0)
+ mask = xops.fmha.BlockDiagonalMask.from_seqlens(seq_lens)
+ out = xops.memory_efficient_attention(q, k, v, mask=mask)[0] # [M, H, C]
+ elif ATTN_BACKEND == 'flash_attn':
+ cu_seqlens = torch.zeros(len(seq_lens)+1, dtype=torch.int32, device=qkv.device)
+ cu_seqlens[1:] = torch.cumsum(torch.tensor(seq_lens, device=qkv.device), dim=0)
+ out = flash_attn.flash_attn_varlen_qkvpacked_func(qkv_feats, cu_seqlens, max(seq_lens)) # => [M, H, C]
+ elif ATTN_BACKEND == 'sage':
+ # Varlen SageAttention approach
+ #from sageattention import sageattn_varlen # or whichever kernel
+
+ cu_seqlens = torch.zeros(len(seq_lens)+1, dtype=torch.int32, device=qkv.device)
+ cu_seqlens[1:] = torch.cumsum(torch.tensor(seq_lens, device=qkv.device), dim=0)
+
+ # reorder => [3, M, H, C]
+ qkv_feats = qkv_feats.permute(1, 0, 2, 3).contiguous()
+ q, k, v = qkv_feats[0], qkv_feats[1], qkv_feats[2] # each => [M, H, C]
+ max_len = max(seq_lens)
+
+ out = sageattn_varlen(
+ q, k, v,
+ cu_seqlens, cu_seqlens,
+ max_len, max_len,
+ is_causal=False
+ ) # => [M, H, C]
+ else:
+ raise ValueError(f"Unknown attention module: {ATTN_BACKEND}")
+
+ out = out[bwd_indices] # [T, H, C]
+ return qkv.replace(out)
+
+================
+File: sparse/basic.py
+================
+from typing import *
+import torch
+import torch.nn as nn
+from trellis.backend_config import get_sparse_backend, get_debug_mode
+
+# Get configuration from central config
+BACKEND = get_sparse_backend()
+DEBUG = get_debug_mode()
+
+SparseTensorData = None # Lazy import
+
+__all__ = [
+ 'SparseTensor',
+ 'sparse_batch_broadcast',
+ 'sparse_batch_op',
+ 'sparse_cat',
+ 'sparse_unbind',
+]
+
+class SparseTensor:
+ """
+ Sparse tensor with support for both torchsparse and spconv backends.
+
+ Parameters:
+ - feats (torch.Tensor): Features of the sparse tensor.
+ - coords (torch.Tensor): Coordinates of the sparse tensor.
+ - shape (torch.Size): Shape of the sparse tensor.
+ - layout (List[slice]): Layout of the sparse tensor for each batch
+ - data (SparseTensorData): Sparse tensor data used for convolusion
+
+ NOTE:
+ - Data corresponding to a same batch should be contiguous.
+ - Coords should be in [0, 1023]
+ """
+ @overload
+ def __init__(self, feats: torch.Tensor, coords: torch.Tensor, shape: Optional[torch.Size] = None, layout: Optional[List[slice]] = None, **kwargs): ...
+
+ @overload
+ def __init__(self, data, shape: Optional[torch.Size] = None, layout: Optional[List[slice]] = None, **kwargs): ...
+
+ def __init__(self, *args, **kwargs):
+ # Lazy import of sparse tensor backend
+ global SparseTensorData
+ if SparseTensorData is None:
+ import importlib
+ backend = get_sparse_backend() # Get current backend from central config
+ if backend == 'torchsparse':
+ SparseTensorData = importlib.import_module('torchsparse').SparseTensor
+ elif backend == 'spconv':
+ SparseTensorData = importlib.import_module('spconv.pytorch').SparseConvTensor
+
+ method_id = 0
+ if len(args) != 0:
+ method_id = 0 if isinstance(args[0], torch.Tensor) else 1
+ else:
+ method_id = 1 if 'data' in kwargs else 0
+
+ if method_id == 0:
+ feats, coords, shape, layout = args + (None,) * (4 - len(args))
+ if 'feats' in kwargs:
+ feats = kwargs['feats']
+ del kwargs['feats']
+ if 'coords' in kwargs:
+ coords = kwargs['coords']
+ del kwargs['coords']
+ if 'shape' in kwargs:
+ shape = kwargs['shape']
+ del kwargs['shape']
+ if 'layout' in kwargs:
+ layout = kwargs['layout']
+ del kwargs['layout']
+
+ if shape is None:
+ shape = self.__cal_shape(feats, coords)
+ if layout is None:
+ layout = self.__cal_layout(coords, shape[0])
+ if BACKEND == 'torchsparse':
+ self.data = SparseTensorData(feats, coords, **kwargs)
+ elif BACKEND == 'spconv':
+ spatial_shape = list(coords.max(0)[0] + 1)[1:]
+ self.data = SparseTensorData(feats.reshape(feats.shape[0], -1), coords, spatial_shape, shape[0], **kwargs)
+ self.data._features = feats
+ elif method_id == 1:
+ data, shape, layout = args + (None,) * (3 - len(args))
+ if 'data' in kwargs:
+ data = kwargs['data']
+ del kwargs['data']
+ if 'shape' in kwargs:
+ shape = kwargs['shape']
+ del kwargs['shape']
+ if 'layout' in kwargs:
+ layout = kwargs['layout']
+ del kwargs['layout']
+
+ self.data = data
+ if shape is None:
+ shape = self.__cal_shape(self.feats, self.coords)
+ if layout is None:
+ layout = self.__cal_layout(self.coords, shape[0])
+
+ self._shape = shape
+ self._layout = layout
+ self._scale = kwargs.get('scale', (1, 1, 1))
+ self._spatial_cache = kwargs.get('spatial_cache', {})
+
+ if DEBUG:
+ try:
+ assert self.feats.shape[0] == self.coords.shape[0], f"Invalid feats shape: {self.feats.shape}, coords shape: {self.coords.shape}"
+ assert self.shape == self.__cal_shape(self.feats, self.coords), f"Invalid shape: {self.shape}"
+ assert self.layout == self.__cal_layout(self.coords, self.shape[0]), f"Invalid layout: {self.layout}"
+ for i in range(self.shape[0]):
+ assert torch.all(self.coords[self.layout[i], 0] == i), f"The data of batch {i} is not contiguous"
+ except Exception as e:
+ print('Debugging information:')
+ print(f"- Shape: {self.shape}")
+ print(f"- Layout: {self.layout}")
+ print(f"- Scale: {self._scale}")
+ print(f"- Coords: {self.coords}")
+ raise e
+
+ def __cal_shape(self, feats, coords):
+ shape = []
+ shape.append(coords[:, 0].max().item() + 1)
+ shape.extend([*feats.shape[1:]])
+ return torch.Size(shape)
+
+ def __cal_layout(self, coords, batch_size):
+ seq_len = torch.bincount(coords[:, 0], minlength=batch_size)
+ offset = torch.cumsum(seq_len, dim=0)
+ layout = [slice((offset[i] - seq_len[i]).item(), offset[i].item()) for i in range(batch_size)]
+ return layout
+
+ @property
+ def shape(self) -> torch.Size:
+ return self._shape
+
+ def dim(self) -> int:
+ return len(self.shape)
+
+ @property
+ def layout(self) -> List[slice]:
+ return self._layout
+
+ @property
+ def feats(self) -> torch.Tensor:
+ if BACKEND == 'torchsparse':
+ return self.data.F
+ elif BACKEND == 'spconv':
+ return self.data.features
+
+ @feats.setter
+ def feats(self, value: torch.Tensor):
+ if BACKEND == 'torchsparse':
+ self.data.F = value
+ elif BACKEND == 'spconv':
+ self.data.features = value
+
+ @property
+ def coords(self) -> torch.Tensor:
+ if BACKEND == 'torchsparse':
+ return self.data.C
+ elif BACKEND == 'spconv':
+ return self.data.indices
+
+ @coords.setter
+ def coords(self, value: torch.Tensor):
+ if BACKEND == 'torchsparse':
+ self.data.C = value
+ elif BACKEND == 'spconv':
+ self.data.indices = value
+
+ @property
+ def dtype(self):
+ return self.feats.dtype
+
+ @property
+ def device(self):
+ return self.feats.device
+
+ @overload
+ def to(self, dtype: torch.dtype) -> 'SparseTensor': ...
+
+ @overload
+ def to(self, device: Optional[Union[str, torch.device]] = None, dtype: Optional[torch.dtype] = None) -> 'SparseTensor': ...
+
+ def to(self, *args, **kwargs) -> 'SparseTensor':
+ device = None
+ dtype = None
+ if len(args) == 2:
+ device, dtype = args
+ elif len(args) == 1:
+ if isinstance(args[0], torch.dtype):
+ dtype = args[0]
+ else:
+ device = args[0]
+ if 'dtype' in kwargs:
+ assert dtype is None, "to() received multiple values for argument 'dtype'"
+ dtype = kwargs['dtype']
+ if 'device' in kwargs:
+ assert device is None, "to() received multiple values for argument 'device'"
+ device = kwargs['device']
+
+ new_feats = self.feats.to(device=device, dtype=dtype)
+ new_coords = self.coords.to(device=device)
+ return self.replace(new_feats, new_coords)
+
+ def type(self, dtype):
+ new_feats = self.feats.type(dtype)
+ return self.replace(new_feats)
+
+ def cpu(self) -> 'SparseTensor':
+ new_feats = self.feats.cpu()
+ new_coords = self.coords.cpu()
+ return self.replace(new_feats, new_coords)
+
+ def cuda(self) -> 'SparseTensor':
+ new_feats = self.feats.cuda()
+ new_coords = self.coords.cuda()
+ return self.replace(new_feats, new_coords)
+
+ def half(self) -> 'SparseTensor':
+ new_feats = self.feats.half()
+ return self.replace(new_feats)
+
+ def float(self) -> 'SparseTensor':
+ new_feats = self.feats.float()
+ return self.replace(new_feats)
+
+ def detach(self) -> 'SparseTensor':
+ new_coords = self.coords.detach()
+ new_feats = self.feats.detach()
+ return self.replace(new_feats, new_coords)
+
+ def dense(self) -> torch.Tensor:
+ if BACKEND == 'torchsparse':
+ return self.data.dense()
+ elif BACKEND == 'spconv':
+ return self.data.dense()
+
+ def reshape(self, *shape) -> 'SparseTensor':
+ new_feats = self.feats.reshape(self.feats.shape[0], *shape)
+ return self.replace(new_feats)
+
+ def unbind(self, dim: int) -> List['SparseTensor']:
+ return sparse_unbind(self, dim)
+
+ def replace(self, feats: torch.Tensor, coords: Optional[torch.Tensor] = None) -> 'SparseTensor':
+ new_shape = [self.shape[0]]
+ new_shape.extend(feats.shape[1:])
+ if BACKEND == 'torchsparse':
+ new_data = SparseTensorData(
+ feats=feats,
+ coords=self.data.coords if coords is None else coords,
+ stride=self.data.stride,
+ spatial_range=self.data.spatial_range,
+ )
+ new_data._caches = self.data._caches
+ elif BACKEND == 'spconv':
+ new_data = SparseTensorData(
+ self.data.features.reshape(self.data.features.shape[0], -1),
+ self.data.indices,
+ self.data.spatial_shape,
+ self.data.batch_size,
+ self.data.grid,
+ self.data.voxel_num,
+ self.data.indice_dict
+ )
+ new_data._features = feats
+ new_data.benchmark = self.data.benchmark
+ new_data.benchmark_record = self.data.benchmark_record
+ new_data.thrust_allocator = self.data.thrust_allocator
+ new_data._timer = self.data._timer
+ new_data.force_algo = self.data.force_algo
+ new_data.int8_scale = self.data.int8_scale
+ if coords is not None:
+ new_data.indices = coords
+ new_tensor = SparseTensor(new_data, shape=torch.Size(new_shape), layout=self.layout, scale=self._scale, spatial_cache=self._spatial_cache)
+ return new_tensor
+
+ @staticmethod
+ def full(aabb, dim, value, dtype=torch.float32, device=None) -> 'SparseTensor':
+ N, C = dim
+ x = torch.arange(aabb[0], aabb[3] + 1)
+ y = torch.arange(aabb[1], aabb[4] + 1)
+ z = torch.arange(aabb[2], aabb[5] + 1)
+ coords = torch.stack(torch.meshgrid(x, y, z, indexing='ij'), dim=-1).reshape(-1, 3)
+ coords = torch.cat([
+ torch.arange(N).view(-1, 1).repeat(1, coords.shape[0]).view(-1, 1),
+ coords.repeat(N, 1),
+ ], dim=1).to(dtype=torch.int32, device=device)
+ feats = torch.full((coords.shape[0], C), value, dtype=dtype, device=device)
+ return SparseTensor(feats=feats, coords=coords)
+
+ def __merge_sparse_cache(self, other: 'SparseTensor') -> dict:
+ new_cache = {}
+ for k in set(list(self._spatial_cache.keys()) + list(other._spatial_cache.keys())):
+ if k in self._spatial_cache:
+ new_cache[k] = self._spatial_cache[k]
+ if k in other._spatial_cache:
+ if k not in new_cache:
+ new_cache[k] = other._spatial_cache[k]
+ else:
+ new_cache[k].update(other._spatial_cache[k])
+ return new_cache
+
+ def __neg__(self) -> 'SparseTensor':
+ return self.replace(-self.feats)
+
+ def __elemwise__(self, other: Union[torch.Tensor, 'SparseTensor'], op: callable) -> 'SparseTensor':
+ if isinstance(other, torch.Tensor):
+ try:
+ other = torch.broadcast_to(other, self.shape)
+ other = sparse_batch_broadcast(self, other)
+ except:
+ pass
+ if isinstance(other, SparseTensor):
+ other = other.feats
+ new_feats = op(self.feats, other)
+ new_tensor = self.replace(new_feats)
+ if isinstance(other, SparseTensor):
+ new_tensor._spatial_cache = self.__merge_sparse_cache(other)
+ return new_tensor
+
+ def __add__(self, other: Union[torch.Tensor, 'SparseTensor', float]) -> 'SparseTensor':
+ return self.__elemwise__(other, torch.add)
+
+ def __radd__(self, other: Union[torch.Tensor, 'SparseTensor', float]) -> 'SparseTensor':
+ return self.__elemwise__(other, torch.add)
+
+ def __sub__(self, other: Union[torch.Tensor, 'SparseTensor', float]) -> 'SparseTensor':
+ return self.__elemwise__(other, torch.sub)
+
+ def __rsub__(self, other: Union[torch.Tensor, 'SparseTensor', float]) -> 'SparseTensor':
+ return self.__elemwise__(other, lambda x, y: torch.sub(y, x))
+
+ def __mul__(self, other: Union[torch.Tensor, 'SparseTensor', float]) -> 'SparseTensor':
+ return self.__elemwise__(other, torch.mul)
+
+ def __rmul__(self, other: Union[torch.Tensor, 'SparseTensor', float]) -> 'SparseTensor':
+ return self.__elemwise__(other, torch.mul)
+
+ def __truediv__(self, other: Union[torch.Tensor, 'SparseTensor', float]) -> 'SparseTensor':
+ return self.__elemwise__(other, torch.div)
+
+ def __rtruediv__(self, other: Union[torch.Tensor, 'SparseTensor', float]) -> 'SparseTensor':
+ return self.__elemwise__(other, lambda x, y: torch.div(y, x))
+
+ def __getitem__(self, idx):
+ if isinstance(idx, int):
+ idx = [idx]
+ elif isinstance(idx, slice):
+ idx = range(*idx.indices(self.shape[0]))
+ elif isinstance(idx, torch.Tensor):
+ if idx.dtype == torch.bool:
+ assert idx.shape == (self.shape[0],), f"Invalid index shape: {idx.shape}"
+ idx = idx.nonzero().squeeze(1)
+ elif idx.dtype in [torch.int32, torch.int64]:
+ assert len(idx.shape) == 1, f"Invalid index shape: {idx.shape}"
+ else:
+ raise ValueError(f"Unknown index type: {idx.dtype}")
+ else:
+ raise ValueError(f"Unknown index type: {type(idx)}")
+
+ coords = []
+ feats = []
+ for new_idx, old_idx in enumerate(idx):
+ coords.append(self.coords[self.layout[old_idx]].clone())
+ coords[-1][:, 0] = new_idx
+ feats.append(self.feats[self.layout[old_idx]])
+ coords = torch.cat(coords, dim=0).contiguous()
+ feats = torch.cat(feats, dim=0).contiguous()
+ return SparseTensor(feats=feats, coords=coords)
+
+ def register_spatial_cache(self, key, value) -> None:
+ """
+ Register a spatial cache.
+ The spatial cache can be any thing you want to cache.
+ The registery and retrieval of the cache is based on current scale.
+ """
+ scale_key = str(self._scale)
+ if scale_key not in self._spatial_cache:
+ self._spatial_cache[scale_key] = {}
+ self._spatial_cache[scale_key][key] = value
+
+ def get_spatial_cache(self, key=None):
+ """
+ Get a spatial cache.
+ """
+ scale_key = str(self._scale)
+ cur_scale_cache = self._spatial_cache.get(scale_key, {})
+ if key is None:
+ return cur_scale_cache
+ return cur_scale_cache.get(key, None)
+
+
+def sparse_batch_broadcast(input: SparseTensor, other: torch.Tensor) -> torch.Tensor:
+ """
+ Broadcast a 1D tensor to a sparse tensor along the batch dimension then perform an operation.
+
+ Args:
+ input (torch.Tensor): 1D tensor to broadcast.
+ target (SparseTensor): Sparse tensor to broadcast to.
+ op (callable): Operation to perform after broadcasting. Defaults to torch.add.
+ """
+ coords, feats = input.coords, input.feats
+ broadcasted = torch.zeros_like(feats)
+ for k in range(input.shape[0]):
+ broadcasted[input.layout[k]] = other[k]
+ return broadcasted
+
+
+def sparse_batch_op(input: SparseTensor, other: torch.Tensor, op: callable = torch.add) -> SparseTensor:
+ """
+ Broadcast a 1D tensor to a sparse tensor along the batch dimension then perform an operation.
+
+ Args:
+ input (torch.Tensor): 1D tensor to broadcast.
+ target (SparseTensor): Sparse tensor to broadcast to.
+ op (callable): Operation to perform after broadcasting. Defaults to torch.add.
+ """
+ return input.replace(op(input.feats, sparse_batch_broadcast(input, other)))
+
+
+def sparse_cat(inputs: List[SparseTensor], dim: int = 0) -> SparseTensor:
+ """
+ Concatenate a list of sparse tensors.
+
+ Args:
+ inputs (List[SparseTensor]): List of sparse tensors to concatenate.
+ """
+ if dim == 0:
+ start = 0
+ coords = []
+ for input in inputs:
+ coords.append(input.coords.clone())
+ coords[-1][:, 0] += start
+ start += input.shape[0]
+ coords = torch.cat(coords, dim=0)
+ feats = torch.cat([input.feats for input in inputs], dim=0)
+ output = SparseTensor(
+ coords=coords,
+ feats=feats,
+ )
+ else:
+ feats = torch.cat([input.feats for input in inputs], dim=dim)
+ output = inputs[0].replace(feats)
+
+ return output
+
+
+def sparse_unbind(input: SparseTensor, dim: int) -> List[SparseTensor]:
+ """
+ Unbind a sparse tensor along a dimension.
+
+ Args:
+ input (SparseTensor): Sparse tensor to unbind.
+ dim (int): Dimension to unbind.
+ """
+ if dim == 0:
+ return [input[i] for i in range(input.shape[0])]
+ else:
+ feats = input.feats.unbind(dim)
+ return [input.replace(f) for f in feats]
+
+================
+File: sparse/conv/__init__.py
+================
+import os
+import logging
+from trellis.backend_config import get_sparse_backend, get_spconv_algo
+
+logger = logging.getLogger(__name__)
+ # 'auto', 'implicit_gemm', 'native'
+def get_conv_impl():
+ """
+ Called once at import time (or the first time someone references these objects).
+ Decides which conv module to load based on the user's chosen backend.
+ """
+ backend = get_sparse_backend() # e.g. 'spconv' or 'torchsparse'
+ logger.debug(f"[Sparse Conv] Detected user-chosen backend: {backend}")
+
+ if backend == 'spconv':
+ # We also set or fetch the spconv algorithm from backend_config
+ algo = get_spconv_algo()
+ os.environ['SPCONV_ALGO'] = algo # Let the submodule read it if needed
+ # Now import the actual spconv-based conv
+ from .conv_spconv import SparseConv3d, SparseInverseConv3d
+ return SparseConv3d, SparseInverseConv3d
+
+ elif backend == 'torchsparse':
+ from .conv_torchsparse import SparseConv3d, SparseInverseConv3d
+ return SparseConv3d, SparseInverseConv3d
+
+ else:
+ # If user didn't set a valid backend, or there's no available backend
+ raise ImportError(f"Unknown or unsupported sparse backend '{backend}'.")
+
+SparseConv3d, SparseInverseConv3d = get_conv_impl()
+
+__all__ = [
+ "SparseConv3d",
+ "SparseInverseConv3d",
+]
+
+================
+File: sparse/conv/conv_spconv.py
+================
+import torch
+import torch.nn as nn
+from .. import SparseTensor
+import torch
+import torch.nn as nn
+from .. import SparseTensor
+from trellis.backend_config import get_debug_mode, get_spconv_algo
+
+# Get configuration from central config
+DEBUG = get_debug_mode()
+SPCONV_ALGO = get_spconv_algo()
+
+class SparseConv3d(nn.Module):
+ def __init__(self, in_channels, out_channels, kernel_size, stride=1, dilation=1, padding=None, bias=True, indice_key=None):
+ super(SparseConv3d, self).__init__()
+ if 'spconv' not in globals():
+ import spconv.pytorch as spconv
+
+ # Get algorithm from central config
+ algo = None
+ current_algo = get_spconv_algo()
+ if current_algo == 'native':
+ algo = spconv.ConvAlgo.Native
+ elif current_algo == 'implicit_gemm':
+ algo = spconv.ConvAlgo.MaskImplicitGemm
+ if stride == 1 and (padding is None):
+ self.conv = spconv.SubMConv3d(in_channels, out_channels, kernel_size, dilation=dilation, bias=bias, indice_key=indice_key, algo=algo)
+ else:
+ self.conv = spconv.SparseConv3d(in_channels, out_channels, kernel_size, stride=stride, dilation=dilation, padding=padding, bias=bias, indice_key=indice_key, algo=algo)
+ self.stride = tuple(stride) if isinstance(stride, (list, tuple)) else (stride, stride, stride)
+ self.padding = padding
+
+ def forward(self, x: SparseTensor) -> SparseTensor:
+ spatial_changed = any(s != 1 for s in self.stride) or (self.padding is not None)
+ new_data = self.conv(x.data)
+ new_shape = [x.shape[0], self.conv.out_channels]
+ new_layout = None if spatial_changed else x.layout
+
+ if spatial_changed and (x.shape[0] != 1):
+ # spconv was non-1 stride will break the contiguous of the output tensor, sort by the coords
+ fwd = new_data.indices[:, 0].argsort()
+ bwd = torch.zeros_like(fwd).scatter_(0, fwd, torch.arange(fwd.shape[0], device=fwd.device))
+ sorted_feats = new_data.features[fwd]
+ sorted_coords = new_data.indices[fwd]
+ unsorted_data = new_data
+ new_data = spconv.SparseConvTensor(sorted_feats, sorted_coords, unsorted_data.spatial_shape, unsorted_data.batch_size) # type: ignore
+
+ out = SparseTensor(
+ new_data, shape=torch.Size(new_shape), layout=new_layout,
+ scale=tuple([s * stride for s, stride in zip(x._scale, self.stride)]),
+ spatial_cache=x._spatial_cache,
+ )
+
+ if spatial_changed and (x.shape[0] != 1):
+ out.register_spatial_cache(f'conv_{self.stride}_unsorted_data', unsorted_data)
+ out.register_spatial_cache(f'conv_{self.stride}_sort_bwd', bwd)
+
+ return out
+
+
+class SparseInverseConv3d(nn.Module):
+ def __init__(self, in_channels, out_channels, kernel_size, stride=1, dilation=1, bias=True, indice_key=None):
+ super(SparseInverseConv3d, self).__init__()
+ if 'spconv' not in globals():
+ import spconv.pytorch as spconv
+ self.conv = spconv.SparseInverseConv3d(in_channels, out_channels, kernel_size, bias=bias, indice_key=indice_key)
+ self.stride = tuple(stride) if isinstance(stride, (list, tuple)) else (stride, stride, stride)
+
+ def forward(self, x: SparseTensor) -> SparseTensor:
+ spatial_changed = any(s != 1 for s in self.stride)
+ if spatial_changed:
+ # recover the original spconv order
+ data = x.get_spatial_cache(f'conv_{self.stride}_unsorted_data')
+ bwd = x.get_spatial_cache(f'conv_{self.stride}_sort_bwd')
+ data = data.replace_feature(x.feats[bwd])
+ if DEBUG:
+ assert torch.equal(data.indices, x.coords[bwd]), 'Recover the original order failed'
+ else:
+ data = x.data
+
+ new_data = self.conv(data)
+ new_shape = [x.shape[0], self.conv.out_channels]
+ new_layout = None if spatial_changed else x.layout
+ out = SparseTensor(
+ new_data, shape=torch.Size(new_shape), layout=new_layout,
+ scale=tuple([s // stride for s, stride in zip(x._scale, self.stride)]),
+ spatial_cache=x._spatial_cache,
+ )
+ return out
+
+================
+File: sparse/conv/conv_torchsparse.py
+================
+import torch
+import torch.nn as nn
+from .. import SparseTensor
+
+
+class SparseConv3d(nn.Module):
+ def __init__(self, in_channels, out_channels, kernel_size, stride=1, dilation=1, bias=True, indice_key=None):
+ super(SparseConv3d, self).__init__()
+ if 'torchsparse' not in globals():
+ import torchsparse
+ self.conv = torchsparse.nn.Conv3d(in_channels, out_channels, kernel_size, stride, 0, dilation, bias)
+
+ def forward(self, x: SparseTensor) -> SparseTensor:
+ out = self.conv(x.data)
+ new_shape = [x.shape[0], self.conv.out_channels]
+ out = SparseTensor(out, shape=torch.Size(new_shape), layout=x.layout if all(s == 1 for s in self.conv.stride) else None)
+ out._spatial_cache = x._spatial_cache
+ out._scale = tuple([s * stride for s, stride in zip(x._scale, self.conv.stride)])
+ return out
+
+
+class SparseInverseConv3d(nn.Module):
+ def __init__(self, in_channels, out_channels, kernel_size, stride=1, dilation=1, bias=True, indice_key=None):
+ super(SparseInverseConv3d, self).__init__()
+ if 'torchsparse' not in globals():
+ import torchsparse
+ self.conv = torchsparse.nn.Conv3d(in_channels, out_channels, kernel_size, stride, 0, dilation, bias, transposed=True)
+
+ def forward(self, x: SparseTensor) -> SparseTensor:
+ out = self.conv(x.data)
+ new_shape = [x.shape[0], self.conv.out_channels]
+ out = SparseTensor(out, shape=torch.Size(new_shape), layout=x.layout if all(s == 1 for s in self.conv.stride) else None)
+ out._spatial_cache = x._spatial_cache
+ out._scale = tuple([s // stride for s, stride in zip(x._scale, self.conv.stride)])
+ return out
+
+================
+File: sparse/linear.py
+================
+import torch
+import torch.nn as nn
+from . import SparseTensor
+
+__all__ = [
+ 'SparseLinear'
+]
+
+
+class SparseLinear(nn.Linear):
+ def __init__(self, in_features, out_features, bias=True):
+ super(SparseLinear, self).__init__(in_features, out_features, bias)
+
+ def forward(self, input: SparseTensor) -> SparseTensor:
+ return input.replace(super().forward(input.feats))
+
+================
+File: sparse/nonlinearity.py
+================
+import torch
+import torch.nn as nn
+from . import SparseTensor
+
+__all__ = [
+ 'SparseReLU',
+ 'SparseSiLU',
+ 'SparseGELU',
+ 'SparseActivation'
+]
+
+
+class SparseReLU(nn.ReLU):
+ def forward(self, input: SparseTensor) -> SparseTensor:
+ return input.replace(super().forward(input.feats))
+
+
+class SparseSiLU(nn.SiLU):
+ def forward(self, input: SparseTensor) -> SparseTensor:
+ return input.replace(super().forward(input.feats))
+
+
+class SparseGELU(nn.GELU):
+ def forward(self, input: SparseTensor) -> SparseTensor:
+ return input.replace(super().forward(input.feats))
+
+
+class SparseActivation(nn.Module):
+ def __init__(self, activation: nn.Module):
+ super().__init__()
+ self.activation = activation
+
+ def forward(self, input: SparseTensor) -> SparseTensor:
+ return input.replace(self.activation(input.feats))
+
+================
+File: sparse/norm.py
+================
+import torch
+import torch.nn as nn
+from . import SparseTensor
+from trellis.backend_config import get_debug_mode
+
+DEBUG = get_debug_mode()
+
+__all__ = [
+ 'SparseGroupNorm',
+ 'SparseLayerNorm',
+ 'SparseGroupNorm32',
+ 'SparseLayerNorm32',
+]
+
+
+class SparseGroupNorm(nn.GroupNorm):
+ def __init__(self, num_groups, num_channels, eps=1e-5, affine=True):
+ super(SparseGroupNorm, self).__init__(num_groups, num_channels, eps, affine)
+
+ def forward(self, input: SparseTensor) -> SparseTensor:
+ nfeats = torch.zeros_like(input.feats)
+ for k in range(input.shape[0]):
+ if DEBUG:
+ assert (input.coords[input.layout[k], 0] == k).all(), f"SparseGroupNorm: batch index mismatch"
+ bfeats = input.feats[input.layout[k]]
+ bfeats = bfeats.permute(1, 0).reshape(1, input.shape[1], -1)
+ bfeats = super().forward(bfeats)
+ bfeats = bfeats.reshape(input.shape[1], -1).permute(1, 0)
+ nfeats[input.layout[k]] = bfeats
+ return input.replace(nfeats)
+
+
+class SparseLayerNorm(nn.LayerNorm):
+ def __init__(self, normalized_shape, eps=1e-5, elementwise_affine=True):
+ super(SparseLayerNorm, self).__init__(normalized_shape, eps, elementwise_affine)
+
+ def forward(self, input: SparseTensor) -> SparseTensor:
+ nfeats = torch.zeros_like(input.feats)
+ for k in range(input.shape[0]):
+ bfeats = input.feats[input.layout[k]]
+ bfeats = bfeats.permute(1, 0).reshape(1, input.shape[1], -1)
+ bfeats = super().forward(bfeats)
+ bfeats = bfeats.reshape(input.shape[1], -1).permute(1, 0)
+ nfeats[input.layout[k]] = bfeats
+ return input.replace(nfeats)
+
+
+class SparseGroupNorm32(SparseGroupNorm):
+ """
+ A GroupNorm layer that converts to float32 before the forward pass.
+ """
+ def forward(self, x: SparseTensor) -> SparseTensor:
+ return super().forward(x.float()).type(x.dtype)
+
+class SparseLayerNorm32(SparseLayerNorm):
+ """
+ A LayerNorm layer that converts to float32 before the forward pass.
+ """
+ def forward(self, x: SparseTensor) -> SparseTensor:
+ return super().forward(x.float()).type(x.dtype)
+
+================
+File: sparse/spatial.py
+================
+from typing import *
+import torch
+import torch.nn as nn
+from . import SparseTensor
+
+__all__ = [
+ 'SparseDownsample',
+ 'SparseUpsample',
+ 'SparseSubdivide'
+]
+
+
+class SparseDownsample(nn.Module):
+ """
+ Downsample a sparse tensor by a factor of `factor`.
+ Implemented as average pooling.
+ """
+ def __init__(self, factor: Union[int, Tuple[int, ...], List[int]]):
+ super(SparseDownsample, self).__init__()
+ self.factor = tuple(factor) if isinstance(factor, (list, tuple)) else factor
+
+ def forward(self, input: SparseTensor) -> SparseTensor:
+ DIM = input.coords.shape[-1] - 1
+ factor = self.factor if isinstance(self.factor, tuple) else (self.factor,) * DIM
+ assert DIM == len(factor), 'Input coordinates must have the same dimension as the downsample factor.'
+
+ coord = list(input.coords.unbind(dim=-1))
+ for i, f in enumerate(factor):
+ coord[i+1] = coord[i+1] // f
+
+ MAX = [coord[i+1].max().item() + 1 for i in range(DIM)]
+ OFFSET = torch.cumprod(torch.tensor(MAX[::-1]), 0).tolist()[::-1] + [1]
+ code = sum([c * o for c, o in zip(coord, OFFSET)])
+ code, idx = code.unique(return_inverse=True)
+
+ new_feats = torch.scatter_reduce(
+ torch.zeros(code.shape[0], input.feats.shape[1], device=input.feats.device, dtype=input.feats.dtype),
+ dim=0,
+ index=idx.unsqueeze(1).expand(-1, input.feats.shape[1]),
+ src=input.feats,
+ reduce='mean'
+ )
+ new_coords = torch.stack(
+ [code // OFFSET[0]] +
+ [(code // OFFSET[i+1]) % MAX[i] for i in range(DIM)],
+ dim=-1
+ )
+ out = SparseTensor(new_feats, new_coords, input.shape,)
+ out._scale = tuple([s // f for s, f in zip(input._scale, factor)])
+ out._spatial_cache = input._spatial_cache
+
+ out.register_spatial_cache(f'upsample_{factor}_coords', input.coords)
+ out.register_spatial_cache(f'upsample_{factor}_layout', input.layout)
+ out.register_spatial_cache(f'upsample_{factor}_idx', idx)
+
+ return out
+
+
+class SparseUpsample(nn.Module):
+ """
+ Upsample a sparse tensor by a factor of `factor`.
+ Implemented as nearest neighbor interpolation.
+ """
+ def __init__(self, factor: Union[int, Tuple[int, int, int], List[int]]):
+ super(SparseUpsample, self).__init__()
+ self.factor = tuple(factor) if isinstance(factor, (list, tuple)) else factor
+
+ def forward(self, input: SparseTensor) -> SparseTensor:
+ DIM = input.coords.shape[-1] - 1
+ factor = self.factor if isinstance(self.factor, tuple) else (self.factor,) * DIM
+ assert DIM == len(factor), 'Input coordinates must have the same dimension as the upsample factor.'
+
+ new_coords = input.get_spatial_cache(f'upsample_{factor}_coords')
+ new_layout = input.get_spatial_cache(f'upsample_{factor}_layout')
+ idx = input.get_spatial_cache(f'upsample_{factor}_idx')
+ if any([x is None for x in [new_coords, new_layout, idx]]):
+ raise ValueError('Upsample cache not found. SparseUpsample must be paired with SparseDownsample.')
+ new_feats = input.feats[idx]
+ out = SparseTensor(new_feats, new_coords, input.shape, new_layout)
+ out._scale = tuple([s * f for s, f in zip(input._scale, factor)])
+ out._spatial_cache = input._spatial_cache
+ return out
+
+class SparseSubdivide(nn.Module):
+ """
+ Upsample a sparse tensor by a factor of `factor`.
+ Implemented as nearest neighbor interpolation.
+ """
+ def __init__(self):
+ super(SparseSubdivide, self).__init__()
+
+ def forward(self, input: SparseTensor) -> SparseTensor:
+ DIM = input.coords.shape[-1] - 1
+ # upsample scale=2^DIM
+ n_cube = torch.ones([2] * DIM, device=input.device, dtype=torch.int)
+ n_coords = torch.nonzero(n_cube)
+ n_coords = torch.cat([torch.zeros_like(n_coords[:, :1]), n_coords], dim=-1)
+ factor = n_coords.shape[0]
+ assert factor == 2 ** DIM
+ # print(n_coords.shape)
+ new_coords = input.coords.clone()
+ new_coords[:, 1:] *= 2
+ new_coords = new_coords.unsqueeze(1) + n_coords.unsqueeze(0).to(new_coords.dtype)
+
+ new_feats = input.feats.unsqueeze(1).expand(input.feats.shape[0], factor, *input.feats.shape[1:])
+ out = SparseTensor(new_feats.flatten(0, 1), new_coords.flatten(0, 1), input.shape)
+ out._scale = input._scale * 2
+ out._spatial_cache = input._spatial_cache
+ return out
+
+================
+File: spatial.py
+================
+import torch
+
+
+def pixel_shuffle_3d(x: torch.Tensor, scale_factor: int) -> torch.Tensor:
+ """
+ 3D pixel shuffle.
+ """
+ B, C, H, W, D = x.shape
+ C_ = C // scale_factor**3
+ x = x.reshape(B, C_, scale_factor, scale_factor, scale_factor, H, W, D)
+ x = x.permute(0, 1, 5, 2, 6, 3, 7, 4)
+ x = x.reshape(B, C_, H*scale_factor, W*scale_factor, D*scale_factor)
+ return x
+
+
+def patchify(x: torch.Tensor, patch_size: int):
+ """
+ Patchify a tensor.
+
+ Args:
+ x (torch.Tensor): (N, C, *spatial) tensor
+ patch_size (int): Patch size
+ """
+ DIM = x.dim() - 2
+ for d in range(2, DIM + 2):
+ assert x.shape[d] % patch_size == 0, f"Dimension {d} of input tensor must be divisible by patch size, got {x.shape[d]} and {patch_size}"
+
+ x = x.reshape(*x.shape[:2], *sum([[x.shape[d] // patch_size, patch_size] for d in range(2, DIM + 2)], []))
+ x = x.permute(0, 1, *([2 * i + 3 for i in range(DIM)] + [2 * i + 2 for i in range(DIM)]))
+ x = x.reshape(x.shape[0], x.shape[1] * (patch_size ** DIM), *(x.shape[-DIM:]))
+ return x
+
+
+def unpatchify(x: torch.Tensor, patch_size: int):
+ """
+ Unpatchify a tensor.
+
+ Args:
+ x (torch.Tensor): (N, C, *spatial) tensor
+ patch_size (int): Patch size
+ """
+ DIM = x.dim() - 2
+ assert x.shape[1] % (patch_size ** DIM) == 0, f"Second dimension of input tensor must be divisible by patch size to unpatchify, got {x.shape[1]} and {patch_size ** DIM}"
+
+ x = x.reshape(x.shape[0], x.shape[1] // (patch_size ** DIM), *([patch_size] * DIM), *(x.shape[-DIM:]))
+ x = x.permute(0, 1, *(sum([[2 + DIM + i, 2 + i] for i in range(DIM)], [])))
+ x = x.reshape(x.shape[0], x.shape[1], *[x.shape[2 + 2 * i] * patch_size for i in range(DIM)])
+ return x
+
+================
+File: utils.py
+================
+import torch.nn as nn
+from ..modules import sparse as sp
+
+FP16_MODULES = (
+ nn.Conv1d,
+ nn.Conv2d,
+ nn.Conv3d,
+ nn.ConvTranspose1d,
+ nn.ConvTranspose2d,
+ nn.ConvTranspose3d,
+ nn.Linear,
+ sp.SparseConv3d,
+ sp.SparseInverseConv3d,
+ sp.SparseLinear,
+)
+
+def convert_module_to_f16(l):
+ """
+ Convert primitive modules to float16.
+ """
+ if isinstance(l, FP16_MODULES):
+ for p in l.parameters():
+ p.data = p.data.half()
+
+
+def convert_module_to_f32(l):
+ """
+ Convert primitive modules to float32, undoing convert_module_to_f16().
+ """
+ if isinstance(l, FP16_MODULES):
+ for p in l.parameters():
+ p.data = p.data.float()
+
+
+def zero_module(module):
+ """
+ Zero out the parameters of a module and return it.
+ """
+ for p in module.parameters():
+ p.detach().zero_()
+ return module
+
+
+def scale_module(module, scale):
+ """
+ Scale the parameters of a module and return it.
+ """
+ for p in module.parameters():
+ p.detach().mul_(scale)
+ return module
+
+
+def modulate(x, shift, scale):
+ return x * (1 + scale.unsqueeze(1)) + shift.unsqueeze(1)
diff --git a/trellis/modules/sparse/attention/__init__.py b/trellis/modules/sparse/attention/__init__.py
index 400de9a..fade63b 100644
--- a/trellis/modules/sparse/attention/__init__.py
+++ b/trellis/modules/sparse/attention/__init__.py
@@ -2,3 +2,30 @@ from .full_attn import *
from .serialized_attn import *
from .windowed_attn import *
from .modules import *
+import os
+import logging
+from typing import Literal
+from trellis.backend_config import (
+ get_attention_backend,
+ get_debug_mode,
+)
+import logging
+
+logger = logging.getLogger(__name__)
+
+#ATTN = get_attention_backend()
+ATTN = get_attention_backend()
+DEBUG = get_debug_mode()
+
+def __from_env():
+ """Read current backend configuration"""
+ #global ATTN
+ global ATTN
+ global DEBUG
+
+ # Get current settings from central config
+ #ATTN =
+ ATTN = get_attention_backend()
+ DEBUG = get_debug_mode()
+
+ print(f"[ATTENTION] sparse backend: {ATTN}")
\ No newline at end of file
diff --git a/trellis/modules/sparse/attention/full_attn.py b/trellis/modules/sparse/attention/full_attn.py
index c724327..ea8ed61 100644
--- a/trellis/modules/sparse/attention/full_attn.py
+++ b/trellis/modules/sparse/attention/full_attn.py
@@ -1,15 +1,41 @@
+#trellis\modules\sparse\attention\full_attn.py
from typing import *
import torch
+import math
from .. import SparseTensor
-from .. import DEBUG, ATTN
+from trellis.backend_config import (
+ get_attention_backend,
+ get_debug_mode,
+ get_available_backends
+)
+import logging
-if ATTN == 'xformers':
+logger = logging.getLogger(__name__)
+
+# Get configuration from central config
+ATTN = get_attention_backend()
+DEBUG = get_debug_mode()
+
+# Get available backends and import if active
+available_backends = get_available_backends()
+
+if ATTN == "xformers" and available_backends['xformers']:
import xformers.ops as xops
-elif ATTN == 'flash_attn':
+elif ATTN == "flash_attn" and available_backends['flash_attn']:
import flash_attn
+elif ATTN == "sage" and available_backends['sage']:
+ import torch.nn.functional as F
+ from sageattention import sageattn
+ F.scaled_dot_product_attention = sageattn
+elif ATTN == "sdpa":
+ from torch.nn.functional import scaled_dot_product_attention as sdpa
+elif ATTN == "naive":
+ from torch.nn.functional import scaled_dot_product_attention as naive
else:
raise ValueError(f"Unknown attention module: {ATTN}")
+# Log the active backend
+logger.info(f"Using attention backend: {ATTN}")
__all__ = [
'sparse_scaled_dot_product_attention',
@@ -212,4 +238,4 @@ def sparse_scaled_dot_product_attention(*args, **kwargs):
if s is not None:
return s.replace(out)
else:
- return out.reshape(N, L, H, -1)
+ return out.reshape(N, L, H, -1)
\ No newline at end of file
diff --git a/trellis/modules/sparse/attention/serialized_attn.py b/trellis/modules/sparse/attention/serialized_attn.py
index b3da276..f86ef77 100644
--- a/trellis/modules/sparse/attention/serialized_attn.py
+++ b/trellis/modules/sparse/attention/serialized_attn.py
@@ -3,15 +3,39 @@ from enum import Enum
import torch
import math
from .. import SparseTensor
-from .. import DEBUG, ATTN
+from trellis.backend_config import (
+ get_attention_backend,
+ get_debug_mode,
+ get_available_backends
+)
+import logging
-if ATTN == 'xformers':
+logger = logging.getLogger(__name__)
+
+# Get configuration from central config
+ATTN = get_attention_backend()
+DEBUG = get_debug_mode()
+
+# Get available backends and import if active
+available_backends = get_available_backends()
+
+if ATTN == "xformers" and available_backends['xformers']:
import xformers.ops as xops
-elif ATTN == 'flash_attn':
+elif ATTN == "flash_attn" and available_backends['flash_attn']:
import flash_attn
+elif ATTN == "sage" and available_backends['sage']:
+ import torch.nn.functional as F
+ from sageattention import sageattn
+ F.scaled_dot_product_attention = sageattn
+elif ATTN == "sdpa":
+ from torch.nn.functional import scaled_dot_product_attention as sdpa
+elif ATTN == "naive":
+ from torch.nn.functional import scaled_dot_product_attention as naive
else:
raise ValueError(f"Unknown attention module: {ATTN}")
+# Log the active backend
+logger.info(f"Using attention backend: {ATTN}")
__all__ = [
'sparse_serialized_scaled_dot_product_self_attention',
@@ -190,4 +214,4 @@ def sparse_serialized_scaled_dot_product_self_attention(
qkv_coords = qkv_coords[bwd_indices]
assert torch.equal(qkv_coords, qkv.coords), "SparseWindowedScaledDotProductSelfAttention: coordinate mismatch"
- return qkv.replace(out)
+ return qkv.replace(out)
\ No newline at end of file
diff --git a/trellis/modules/sparse/attention/windowed_attn.py b/trellis/modules/sparse/attention/windowed_attn.py
index 11eebf8..7f1f3ba 100644
--- a/trellis/modules/sparse/attention/windowed_attn.py
+++ b/trellis/modules/sparse/attention/windowed_attn.py
@@ -1,16 +1,43 @@
from typing import *
-import torch
+from enum import Enum
import math
+import os
+import logging
+import torch
from .. import SparseTensor
-from .. import DEBUG, ATTN
+from trellis.backend_config import (
+ get_attention_backend,
+ get_debug_mode,
+ get_available_backends
+)
+import logging
-if ATTN == 'xformers':
+logger = logging.getLogger(__name__)
+
+# Get configuration from central config
+ATTN = get_attention_backend()
+DEBUG = get_debug_mode()
+
+# Get available backends and import if active
+available_backends = get_available_backends()
+
+if ATTN == "xformers" and available_backends['xformers']:
import xformers.ops as xops
-elif ATTN == 'flash_attn':
+elif ATTN == "flash_attn" and available_backends['flash_attn']:
import flash_attn
+elif ATTN == "sage" and available_backends['sage']:
+ import torch.nn.functional as F
+ from sageattention import sageattn
+ F.scaled_dot_product_attention = sageattn
+elif ATTN == "sdpa":
+ from torch.nn.functional import scaled_dot_product_attention as sdpa
+elif ATTN == "naive":
+ from torch.nn.functional import scaled_dot_product_attention as naive
else:
raise ValueError(f"Unknown attention module: {ATTN}")
+# Log the active backend
+logger.info(f"Using attention backend: {ATTN}")
__all__ = [
'sparse_windowed_scaled_dot_product_self_attention',
diff --git a/trellis/modules/sparse/basic.py b/trellis/modules/sparse/basic.py
index 0fc6851..24a8220 100644
--- a/trellis/modules/sparse/basic.py
+++ b/trellis/modules/sparse/basic.py
@@ -1,10 +1,17 @@
from typing import *
import torch
import torch.nn as nn
-from . import BACKEND, DEBUG
+
+from trellis.backend_config import get_debug_mode, get_spconv_algo, get_sparse_backend
+
+
SparseTensorData = None # Lazy import
+DEBUG = get_debug_mode()
+SPCONV_ALGO = get_spconv_algo()
+
+BACKEND = get_sparse_backend()
__all__ = [
'SparseTensor',
'sparse_batch_broadcast',
diff --git a/trellis/modules/sparse/conv/__init__.py b/trellis/modules/sparse/conv/__init__.py
index 8ae5437..e3f0b95 100644
--- a/trellis/modules/sparse/conv/__init__.py
+++ b/trellis/modules/sparse/conv/__init__.py
@@ -1,7 +1,9 @@
-from .. import BACKEND
+import os
+import logging
+from trellis.backend_config import get_sparse_backend, get_spconv_algo
-
-SPCONV_ALGO = 'auto' # 'auto', 'implicit_gemm', 'native'
+BACKEND = get_sparse_backend()
+SPCONV_ALGO = get_spconv_algo()
def __from_env():
import os
@@ -19,4 +21,9 @@ if BACKEND == 'torchsparse':
from .conv_torchsparse import *
elif BACKEND == 'spconv':
from .conv_spconv import *
+
+__all__ = [
+ "SparseConv3d",
+ "SparseInverseConv3d",
+]
diff --git a/trellis/modules/sparse/conv/conv_spconv.py b/trellis/modules/sparse/conv/conv_spconv.py
index 856405d..2d8c9da 100644
--- a/trellis/modules/sparse/conv/conv_spconv.py
+++ b/trellis/modules/sparse/conv/conv_spconv.py
@@ -2,12 +2,17 @@ import torch
import torch.nn as nn
from .. import SparseTensor
from .. import DEBUG
-from . import SPCONV_ALGO
+from trellis.backend_config import get_debug_mode, get_spconv_algo, get_sparse_backend
+
+# Get configuration from central config
+DEBUG = get_debug_mode()
+SPCONV_ALGO = get_spconv_algo()
+BACKEND = get_sparse_backend()
class SparseConv3d(nn.Module):
def __init__(self, in_channels, out_channels, kernel_size, stride=1, dilation=1, padding=None, bias=True, indice_key=None):
super(SparseConv3d, self).__init__()
- if 'spconv' not in globals():
+ if BACKEND == 'spconv':
import spconv.pytorch as spconv
algo = None
if SPCONV_ALGO == 'native':
@@ -52,7 +57,7 @@ class SparseConv3d(nn.Module):
class SparseInverseConv3d(nn.Module):
def __init__(self, in_channels, out_channels, kernel_size, stride=1, dilation=1, bias=True, indice_key=None):
super(SparseInverseConv3d, self).__init__()
- if 'spconv' not in globals():
+ if BACKEND == 'spconv':
import spconv.pytorch as spconv
self.conv = spconv.SparseInverseConv3d(in_channels, out_channels, kernel_size, bias=bias, indice_key=indice_key)
self.stride = tuple(stride) if isinstance(stride, (list, tuple)) else (stride, stride, stride)
diff --git a/trellis/modules/sparse/conv/conv_torchsparse.py b/trellis/modules/sparse/conv/conv_torchsparse.py
index a10bd91..feec227 100644
--- a/trellis/modules/sparse/conv/conv_torchsparse.py
+++ b/trellis/modules/sparse/conv/conv_torchsparse.py
@@ -1,12 +1,17 @@
import torch
import torch.nn as nn
from .. import SparseTensor
+from trellis.backend_config import get_debug_mode, get_spconv_algo, get_sparse_backend
+# Get configuration from central config
+DEBUG = get_debug_mode()
+SPCONV_ALGO = get_spconv_algo()
+BACKEND = get_sparse_backend()
class SparseConv3d(nn.Module):
def __init__(self, in_channels, out_channels, kernel_size, stride=1, dilation=1, bias=True, indice_key=None):
super(SparseConv3d, self).__init__()
- if 'torchsparse' not in globals():
+ if BACKEND == 'torchsparse':
import torchsparse
self.conv = torchsparse.nn.Conv3d(in_channels, out_channels, kernel_size, stride, 0, dilation, bias)
@@ -22,7 +27,7 @@ class SparseConv3d(nn.Module):
class SparseInverseConv3d(nn.Module):
def __init__(self, in_channels, out_channels, kernel_size, stride=1, dilation=1, bias=True, indice_key=None):
super(SparseInverseConv3d, self).__init__()
- if 'torchsparse' not in globals():
+ if BACKEND == 'torchsparse':
import torchsparse
self.conv = torchsparse.nn.Conv3d(in_channels, out_channels, kernel_size, stride, 0, dilation, bias, transposed=True)
diff --git a/trellis_model_manager.py b/trellis_model_manager.py
index 6458c19..9ae8935 100644
--- a/trellis_model_manager.py
+++ b/trellis_model_manager.py
@@ -172,12 +172,17 @@ class TrellisModelManager:
return models
def load_dinov2(self, model_name: str):
- """Load DINOv2 model with device and precision management"""
+ """Load DINOv2 model with device, precision, and attention backend management"""
try:
- # Get use_fp16 from config dict or object
+ # Get configuration values
use_fp16 = (self.config.get('use_fp16', True)
if isinstance(self.config, dict)
else getattr(self.config, 'use_fp16', True))
+
+ # Get attention backend from config
+ attention_backend = (self.config.get('attention_backend', 'default')
+ if isinstance(self.config, dict)
+ else getattr(self.config, 'attention_backend', 'default'))
# Try to load from local path first
model_path = folder_paths.get_full_path("classifiers", f"{model_name}.pth")
@@ -185,8 +190,11 @@ class TrellisModelManager:
if model_path is None:
print(f"Downloading {model_name} from torch hub...")
try:
- # Load model architecture
- model = torch.hub.load('facebookresearch/dinov2', model_name, pretrained=True)
+ # Load model architecture with specified attention backend
+ model = torch.hub.load('facebookresearch/dinov2', model_name,
+ pretrained=True,
+ force_reload=False,
+ trust_repo=True)
# Save model for future use
save_dir = os.path.join(folder_paths.models_dir, "classifiers")
@@ -203,7 +211,10 @@ class TrellisModelManager:
else:
# Load from local path
print(f"Loading DINOv2 model from {model_path}")
- model = torch.hub.load('facebookresearch/dinov2', model_name, pretrained=False)
+ model = torch.hub.load('facebookresearch/dinov2', model_name,
+ pretrained=False,
+ force_reload=False,
+ trust_repo=True)
model.load_state_dict(torch.load(model_path))
# Move model to specified device and apply precision settings
@@ -211,6 +222,10 @@ class TrellisModelManager:
if use_fp16:
model = model.half()
+ # Set attention backend if specified in config
+ if hasattr(model, 'set_attention_backend') and attention_backend != 'default':
+ model.set_attention_backend(attention_backend)
+
model.eval()
return model