feat(env): Support 5090 and cu128.

README.md

@@ -53,8 +53,11 @@ cd EmbodiedGen
 git checkout v2.0.1
 conda create -n embodiedgen python=3.10.13 -y
 conda activate embodiedgen
-# bash install.sh cu126 && conda deactivate && conda activate embodiedgen # Optional: if you don't have local cuda126.
+# Manually install one CUDA toolkit when needed. cu126 remains the default.
+# bash install.sh cu126 # RTX 40-series.
+# bash install.sh cu128 # RTX 50-series / Blackwell.
 bash install.sh basic # around 10 mins
+conda deactivate && conda activate embodiedgen
 ```
 
 Set up the GPT agent (required by most pipelines): update the API key in `embodied_gen/utils/gpt_config.yaml`. Then generate your first sim-ready asset:
@@ -274,8 +277,7 @@ If you use EmbodiedGen in your research or projects, please cite:
 ```bibtex
 @misc{wang2025embodiedgengenerative3dworld,
   title         = {EmbodiedGen: Towards a Generative 3D World Engine for Embodied Intelligence},
-  author        = {Xinjie Wang and Liu Liu and Yu Cao and Ruiqi Wu and Wenkang Qin and
-                   Dehui Wang and Wei Sui and Zhizhong Su},
+  author        = {Xinjie Wang and Liu Liu and Yu Cao and Ruiqi Wu and Wenkang Qin and Dehui Wang and Wei Sui and Zhizhong Su},
   year          = {2025},
   eprint        = {2506.10600},
   archivePrefix = {arXiv},

apps/common.py

@@ -16,15 +16,15 @@
 
 import spaces  # noqa: E402
 from embodied_gen.utils.monkey_patch.gradio import (
-    _disable_xformers_flash3,
     _neutralize_warp_in_parent,
     _patch_open3d_cuda_device_count_bug,
 )
 from embodied_gen.utils.monkey_patch.trellis import monkey_path_trellis
+from embodied_gen.utils.monkey_patch.xformers import disable_xformers_flash3
 
 _neutralize_warp_in_parent()
 _patch_open3d_cuda_device_count_bug()
-_disable_xformers_flash3()
+disable_xformers_flash3()
 monkey_path_trellis()
 
 

docs/documentation/install.md

@@ -10,8 +10,12 @@ cd EmbodiedGen
 git checkout v2.0.1
 conda create -n embodiedgen python=3.10.13 -y # recommended to use a new env.
 conda activate embodiedgen
-# bash install.sh cu126 && conda deactivate && conda activate embodiedgen # Optional: if you don't have local cuda126.
+# Manually install one CUDA toolkit when needed. cu126 remains the default.
+# bash install.sh cu126 # RTX 40-series.
+# bash install.sh cu128 # RTX 50-series / Blackwell.
 bash install.sh basic # around 10 mins
+# Reactivate once so the CUDA and OpenEXR hooks load in the current shell.
+conda deactivate && conda activate embodiedgen
 # Optional: `bash install.sh scene3d` for scene3d-cli; `bash install.sh room` for room-cli; `bash install.sh affordance` for affordance-cli.
 ```
 

embodied_gen/models/image_comm_model.py

@@ -31,6 +31,9 @@
 from huggingface_hub import snapshot_download
 from PIL import Image
 from transformers import AutoModelForCausalLM, SiglipProcessor
+from embodied_gen.utils.monkey_patch.xformers import (
+    disable_xformers_flash3_on_blackwell,
+)
 
 __all__ = [
     "build_hf_image_pipeline",
@@ -99,6 +102,7 @@ def load(self):
         )
         pipe = pipe.to(self.device)
         pipe.enable_model_cpu_offload()
+        disable_xformers_flash3_on_blackwell()
         pipe.enable_xformers_memory_efficient_attention()
         pipe.enable_attention_slicing()
         return pipe
@@ -230,6 +234,7 @@ def load(self):
             variant="fp16",
         ).to(self.device)
         pipe.enable_model_cpu_offload()
+        disable_xformers_flash3_on_blackwell()
         pipe.enable_xformers_memory_efficient_attention()
         pipe.scheduler = DPMSolverMultistepScheduler.from_config(
             pipe.scheduler.config, use_karras_sigmas=True
@@ -268,6 +273,7 @@ def load(self):
             "black-forest-labs/FLUX.1-schnell", torch_dtype=torch.bfloat16
         )
         pipe.enable_model_cpu_offload()
+        disable_xformers_flash3_on_blackwell()
         pipe.enable_xformers_memory_efficient_attention()
         pipe.enable_attention_slicing()
         return pipe.to(self.device)

embodied_gen/models/text_model.py

@@ -40,6 +40,9 @@
 )
 from PIL import Image
 from transformers import CLIPImageProcessor, CLIPVisionModelWithProjection
+from embodied_gen.utils.monkey_patch.xformers import (
+    disable_xformers_flash3_on_blackwell,
+)
 
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
@@ -242,6 +245,7 @@ def build_text2img_pipeline(
         pipe.enable_model_cpu_offload()
     else:
         pipe = pipe.to(device)
+    disable_xformers_flash3_on_blackwell()
     pipe.enable_xformers_memory_efficient_attention()
     pipe.enable_vae_slicing()
 

embodied_gen/models/texture_model.py

@@ -93,19 +93,25 @@ def build_texture_gen_pipe(
     """
 
     download_kolors_weights(f"{base_ckpt_dir}/Kolors")
-    logger.info(f"Load Kolors weights...")
+    logger.info("Load Kolors weights...")
     tokenizer = ChatGLMTokenizer.from_pretrained(
         f"{base_ckpt_dir}/Kolors/text_encoder"
     )
     text_encoder = ChatGLMModel.from_pretrained(
         f"{base_ckpt_dir}/Kolors/text_encoder", torch_dtype=torch.float16
     ).half()
     vae = AutoencoderKL.from_pretrained(
-        f"{base_ckpt_dir}/Kolors/vae", revision=None
-    ).half()
+        f"{base_ckpt_dir}/Kolors/vae",
+        torch_dtype=torch.float16,
+        variant="fp16",
+        use_safetensors=True,
+    )
     unet = UNet2DConditionModel.from_pretrained(
-        f"{base_ckpt_dir}/Kolors/unet", revision=None
-    ).half()
+        f"{base_ckpt_dir}/Kolors/unet",
+        torch_dtype=torch.float16,
+        variant="fp16",
+        use_safetensors=True,
+    )
     scheduler = EulerDiscreteScheduler.from_pretrained(
         f"{base_ckpt_dir}/Kolors/scheduler"
     )

embodied_gen/utils/monkey_patch/gradio.py

@@ -111,30 +111,3 @@ def _wrapped(*args, **kwargs):
             _wctx.runtime_init = _make_pid_safe(_wctx.runtime_init)
     except Exception:
         pass
-
-
-def _disable_xformers_flash3() -> None:
-    """Force xformers dispatcher to skip Flash-Attention v3 (Hopper-only).
-
-    sm_120 (Blackwell) has no FA3 kernel binary; the dispatcher still picks
-    flash3 and the launch aborts with:
-      `CUDA error ... hopper/flash_fwd_launch_template.h:188: invalid argument`
-    Env vars `XFORMERS_FLASH3_ATTENTION_DISABLED=1` are silently ignored in
-    xformers 0.0.32.post2, so we patch `not_supported_reasons` directly.
-    Cutlass and FA2 both work on sm_120, so removing flash3 from candidates
-    is enough.
-    """
-    try:
-        from xformers.ops.fmha import flash3 as _f3
-    except Exception:
-        return
-
-    _disabled = ["disabled by EmbodiedGen: no FA3 kernel for sm_120"]
-
-    def _ns(cls, d):  # noqa: ARG001
-        return list(_disabled)
-
-    if hasattr(_f3, "FwOp"):
-        _f3.FwOp.not_supported_reasons = classmethod(_ns)
-    if hasattr(_f3, "BwOp"):
-        _f3.BwOp.not_supported_reasons = classmethod(_ns)

embodied_gen/utils/monkey_patch/sam3d.py

@@ -22,13 +22,22 @@
 import numpy as np
 import torch
 from PIL import Image
+from embodied_gen.utils.monkey_patch.xformers import (
+    disable_xformers_flash3_on_blackwell,
+)
 
 
-def monkey_patch_sam3d():
+def monkey_patch_sam3d() -> None:
     """Monkey patches SAM3D inference pipelines with custom initialization and execution logic."""
     from embodied_gen.data.utils import model_device_ctx
     from embodied_gen.utils.log import logger
 
+    if disable_xformers_flash3_on_blackwell():
+        logger.info(
+            "[ATTENTION] Disabled xFormers FlashAttention 3 on Blackwell; "
+            "using the FlashAttention 2 fallback."
+        )
+
     os.environ["LIDRA_SKIP_INIT"] = "true"
     os.environ['ATTN_BACKEND'] = "xformers"
 
@@ -173,9 +182,9 @@ def patch_run(
                         with_layout_postprocess
                         and self.layout_post_optimization_method is not None
                     ):
-                        assert (
-                            glb is not None
-                        ), "require mesh to run postprocessing"
+                        assert glb is not None, (
+                            "require mesh to run postprocessing"
+                        )
                         logger.info(
                             "Running layout post optimization method..."
                         )

embodied_gen/utils/monkey_patch/xformers.py

@@ -0,0 +1,27 @@
+import torch
+
+_BLACKWELL_MINIMUM_COMPUTE_CAPABILITY = (12, 0)
+
+
+def disable_xformers_flash3() -> bool:
+    """Remove FlashAttention 3 from the xFormers dispatcher."""
+    try:
+        from xformers.ops.fmha import _set_use_fa3
+    except (ImportError, AttributeError):
+        return False
+
+    _set_use_fa3(False)
+    return True
+
+
+def disable_xformers_flash3_on_blackwell() -> bool:
+    """Disable xFormers FlashAttention 3 when a Blackwell GPU is visible."""
+    if not torch.cuda.is_available():
+        return False
+
+    for device_index in range(torch.cuda.device_count()):
+        capability = torch.cuda.get_device_capability(device_index)
+        if capability >= _BLACKWELL_MINIMUM_COMPUTE_CAPABILITY:
+            return disable_xformers_flash3()
+
+    return False

install.sh

@@ -1,18 +1,18 @@
 #!/bin/bash
 set -e
 
-STAGE=$1 # "basic" | "scene3d" | "room" | "affordance" | "cu126" | "all"
+STAGE=$1 # "basic" | "scene3d" | "room" | "affordance" | "cu126" | "cu128" | "all"
 STAGE=${STAGE:-basic}
 
 REPO_ROOT=$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)
 source "$REPO_ROOT/install/_utils.sh"
 cd "$REPO_ROOT"
 
 case "$STAGE" in
-    basic|scene3d|room|affordance|cu126|all) ;;
+    basic|scene3d|room|affordance|cu126|cu128|all) ;;
     *)
         log_error "Unknown installation stage: $STAGE"
-        log_error "Usage: bash install.sh [basic|scene3d|room|affordance|cu126|all]"
+        log_error "Usage: bash install.sh [basic|scene3d|room|affordance|cu126|cu128|all]"
         exit 1
         ;;
 esac
@@ -21,7 +21,10 @@ git config http.postBuffer 524288000
 
 log_info "===== Starting installation stage: $STAGE ====="
 
-if [[ "$STAGE" != "cu126" ]]; then
+if [[ "$STAGE" != "cu126" && "$STAGE" != "cu128" ]]; then
+    source_cuda_activation
+    CUDA_VARIANT=$(detect_cuda_variant)
+    log_info "Using CUDA installation variant: $CUDA_VARIANT"
     bash "$REPO_ROOT/install/init_submodules.sh" "$STAGE"
 fi
 
@@ -50,12 +53,16 @@ if [[ "$STAGE" == "cu126" ]]; then
     bash "$REPO_ROOT/install/install_cu126.sh"
 fi
 
+if [[ "$STAGE" == "cu128" ]]; then
+    bash "$REPO_ROOT/install/install_cu128.sh"
+fi
+
 # Global constraints for all stages
 python -m pip install numpy==1.26.4
 
-if [[ "$STAGE" != "cu126" ]]; then
+if [[ "$STAGE" != "cu126" && "$STAGE" != "cu128" ]]; then
     try_install "Refreshing EmbodiedGen editable install..." \
-        "python -m pip install -e ." \
+        "python -m pip install -e '.[dev]'" \
         "EmbodiedGen editable installation refresh failed."
 fi
 

install/_utils.sh

@@ -19,3 +19,32 @@ try_install() {
         exit 1
     }
 }
+
+detect_cuda_variant() {
+    local cuda_variant="cu126"
+
+    if [[ -n "${CONDA_PREFIX:-}" ]]; then
+        if [[ -f "$CONDA_PREFIX/etc/conda/activate.d/cuda128.sh" ]]; then
+            cuda_variant="cu128"
+        elif [[ -f "$CONDA_PREFIX/etc/conda/activate.d/cuda126.sh" ]]; then
+            cuda_variant="cu126"
+        fi
+    fi
+
+    printf '%s\n' "$cuda_variant"
+}
+
+source_cuda_activation() {
+    local cuda_variant
+    local cuda_hook
+
+    cuda_variant=$(detect_cuda_variant) || return 1
+    if [[ -z "${CONDA_PREFIX:-}" ]]; then
+        return 0
+    fi
+
+    cuda_hook="$CONDA_PREFIX/etc/conda/activate.d/${cuda_variant}.sh"
+    if [[ -f "$cuda_hook" ]]; then
+        source "$cuda_hook"
+    fi
+}

install/init_submodules.sh

@@ -5,11 +5,11 @@ REPO_ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)"
 STAGE="${1:-all}"
 
 usage() {
-  echo "Usage: bash install/init_submodules.sh [basic|scene3d|room|affordance|cu126|all]"
+  echo "Usage: bash install/init_submodules.sh [basic|scene3d|room|affordance|cu126|cu128|all]"
 }
 
 case "$STAGE" in
-  basic|scene3d|room|affordance|cu126|all) ;;
+  basic|scene3d|room|affordance|cu126|cu128|all) ;;
   *)
     usage >&2
     exit 1
@@ -37,7 +37,7 @@ submodules_for_stage() {
     all)
       git -C "$REPO_ROOT" config -f .gitmodules --get-regexp '^submodule\..*\.path$' | awk '{ print $2 }'
       ;;
-    cu126)
+    cu126|cu128)
       ;;
   esac
 }

install/install_affordance.sh

@@ -14,11 +14,7 @@ PIP_INSTALL_PACKAGES=(
     "meshcat==0.3.2"
 )
 
-if [[ -n "${CONDA_PREFIX:-}" ]]; then
-    if [[ -f "$CONDA_PREFIX/etc/conda/activate.d/cuda126.sh" ]]; then
-        source "$CONDA_PREFIX/etc/conda/activate.d/cuda126.sh"
-    fi
-fi
+source_cuda_activation
 
 export MAX_JOBS="${MAX_JOBS:-8}"
 log_info "Using TORCH_CUDA_ARCH_LIST=${TORCH_CUDA_ARCH_LIST:-unset} for affordance CUDA extensions."
@@ -50,4 +46,4 @@ try_install "Installing GraspGen pointnet2_ops..." \
 
 rm -r "$POINTNET2_OPS_DIR/build" 2>/dev/null || true
 
-pip install opencv-python==4.13.0.92
+pip install opencv-python==4.9.0.80 opencv-python-headless==4.9.0.80

install/install_basic.sh

@@ -5,11 +5,23 @@ PROJECT_ROOT=$(cd "$SCRIPT_DIR/.." && pwd)
 source "$SCRIPT_DIR/_utils.sh"
 cd "$PROJECT_ROOT"
 
+if [[ -n "${CONDA_PREFIX:-}" ]]; then
+    mkdir -p "$CONDA_PREFIX/etc/conda/activate.d"
+    cat > "$CONDA_PREFIX/etc/conda/activate.d/opencv_openexr.sh" <<'HOOK'
+export OPENCV_IO_ENABLE_OPENEXR="${OPENCV_IO_ENABLE_OPENEXR:-1}"
+HOOK
+    log_info "Wrote OpenEXR activation hook into the conda environment."
+fi
+
+CUDA_VARIANT=$(detect_cuda_variant)
+TORCH_INDEX_URL="${EMBODIEDGEN_TORCH_INDEX_URL:-https://download.pytorch.org/whl/$CUDA_VARIANT}"
+log_info "Installing PyTorch dependencies for CUDA variant: $CUDA_VARIANT"
+
 PIP_INSTALL_PACKAGES=(
     "pip==22.3.1"
     "setuptools==80.10.2 wheel packaging 'Cython>=0.29.37'"
-    "torch==2.8.0 torchvision==0.23.0 --index-url https://download.pytorch.org/whl/cu126"
-    "xformers==0.0.32.post2 --index-url https://download.pytorch.org/whl/cu126"
+    "torch==2.8.0 torchvision==0.23.0 --index-url $TORCH_INDEX_URL"
+    "xformers==0.0.32.post2 --index-url $TORCH_INDEX_URL"
     "-r requirements.txt --use-deprecated=legacy-resolver"
     "utils3d@git+https://github.com/EasternJournalist/utils3d.git@9a4eb15"
     "clip@git+https://github.com/openai/CLIP.git"
@@ -32,7 +44,7 @@ log_info "Installing diff-gaussian-rasterization..."
 pip install --no-build-isolation diff-gaussian-rasterization@git+https://github.com/autonomousvision/mip-splatting.git#subdirectory=submodules/diff-gaussian-rasterization
 
 try_install "Installing EmbodiedGen..." \
-    "pip install -e .[dev]" \
+    "pip install -e '.[dev]'" \
     "EmbodiedGen installation failed."
 
 pre-commit install