diff --git a/README.md b/README.md index 5628e42..739b828 100644 --- a/README.md +++ b/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}, diff --git a/apps/common.py b/apps/common.py index c4f0f31..7fb3d29 100644 --- a/apps/common.py +++ b/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() diff --git a/docs/documentation/install.md b/docs/documentation/install.md index f890a04..0f3a1a0 100644 --- a/docs/documentation/install.md +++ b/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. ``` diff --git a/embodied_gen/models/image_comm_model.py b/embodied_gen/models/image_comm_model.py index a04364d..c1a5f5e 100644 --- a/embodied_gen/models/image_comm_model.py +++ b/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) diff --git a/embodied_gen/models/text_model.py b/embodied_gen/models/text_model.py index a844ac4..7a9e92f 100644 --- a/embodied_gen/models/text_model.py +++ b/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() diff --git a/embodied_gen/models/texture_model.py b/embodied_gen/models/texture_model.py index c8a12fc..62dc914 100644 --- a/embodied_gen/models/texture_model.py +++ b/embodied_gen/models/texture_model.py @@ -93,7 +93,7 @@ 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" ) @@ -101,11 +101,17 @@ def build_texture_gen_pipe( 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" ) diff --git a/embodied_gen/utils/monkey_patch/gradio.py b/embodied_gen/utils/monkey_patch/gradio.py index 584c9cb..489496c 100644 --- a/embodied_gen/utils/monkey_patch/gradio.py +++ b/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) diff --git a/embodied_gen/utils/monkey_patch/sam3d.py b/embodied_gen/utils/monkey_patch/sam3d.py index 8f74717..c00e89e 100644 --- a/embodied_gen/utils/monkey_patch/sam3d.py +++ b/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..." ) diff --git a/embodied_gen/utils/monkey_patch/xformers.py b/embodied_gen/utils/monkey_patch/xformers.py new file mode 100644 index 0000000..b9b3bc3 --- /dev/null +++ b/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 diff --git a/install.sh b/install.sh index 8a7e9bc..06f58e9 100644 --- a/install.sh +++ b/install.sh @@ -1,7 +1,7 @@ #!/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) @@ -9,10 +9,10 @@ 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 diff --git a/install/_utils.sh b/install/_utils.sh index 3fa8879..bbe941e 100644 --- a/install/_utils.sh +++ b/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 +} diff --git a/install/init_submodules.sh b/install/init_submodules.sh index 930ecda..8e3ca69 100644 --- a/install/init_submodules.sh +++ b/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 } diff --git a/install/install_affordance.sh b/install/install_affordance.sh index 49d2012..8440144 100644 --- a/install/install_affordance.sh +++ b/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 diff --git a/install/install_basic.sh b/install/install_basic.sh index 9f6011f..6a24e96 100644 --- a/install/install_basic.sh +++ b/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 diff --git a/install/install_cu126.sh b/install/install_cu126.sh index 8adbc0b..d7d5a7a 100755 --- a/install/install_cu126.sh +++ b/install/install_cu126.sh @@ -36,23 +36,41 @@ log_info "Using conda executable: $CONDA_CMD" log_info "Writing CUDA 12.6 activation hook into the conda environment..." mkdir -p "$CONDA_PREFIX/etc/conda/activate.d" +rm -f "$CONDA_PREFIX/etc/conda/activate.d/cuda128.sh" cat > "$CONDA_PREFIX/etc/conda/activate.d/cuda126.sh" <<'HOOK' +export EMBODIEDGEN_CUDA_VARIANT="cu126" export CUDA_HOME="$CONDA_PREFIX" export CUDA_PATH="$CONDA_PREFIX" -export PATH="$CONDA_PREFIX/bin:$PATH" export CUDA_TARGET_LIB="$CONDA_PREFIX/targets/x86_64-linux/lib" +_cuda_conda_bin="$CONDA_PREFIX/bin" _cuda_conda_lib="$CONDA_PREFIX/lib" _cuda_conda_lib64="$CONDA_PREFIX/lib64" +_cuda_target_include="$CONDA_PREFIX/targets/x86_64-linux/include" +_cuda_path=":${PATH:-}:" +_cuda_path="${_cuda_path//:$_cuda_conda_bin:/:}" +_cuda_path="${_cuda_path#:}" +_cuda_path="${_cuda_path%:}" +export PATH="$_cuda_conda_bin${_cuda_path:+:$_cuda_path}" _cuda_ld_path=":${LD_LIBRARY_PATH:-}:" +_cuda_ld_path="${_cuda_ld_path//:$CUDA_TARGET_LIB:/:}" _cuda_ld_path="${_cuda_ld_path//:$_cuda_conda_lib:/:}" _cuda_ld_path="${_cuda_ld_path//:$_cuda_conda_lib64:/:}" _cuda_ld_path="${_cuda_ld_path#:}" _cuda_ld_path="${_cuda_ld_path%:}" export LD_LIBRARY_PATH="$CUDA_TARGET_LIB${_cuda_ld_path:+:$_cuda_ld_path}" -export LIBRARY_PATH="$CUDA_TARGET_LIB:${LIBRARY_PATH:-}" -unset _cuda_ld_path _cuda_conda_lib _cuda_conda_lib64 -export CPATH="$CONDA_PREFIX/targets/x86_64-linux/include:${CPATH:-}" -export TORCH_CUDA_ARCH_LIST="${TORCH_CUDA_ARCH_LIST:-8.9}" +_cuda_library_path=":${LIBRARY_PATH:-}:" +_cuda_library_path="${_cuda_library_path//:$CUDA_TARGET_LIB:/:}" +_cuda_library_path="${_cuda_library_path#:}" +_cuda_library_path="${_cuda_library_path%:}" +export LIBRARY_PATH="$CUDA_TARGET_LIB${_cuda_library_path:+:$_cuda_library_path}" +_cuda_cpath=":${CPATH:-}:" +_cuda_cpath="${_cuda_cpath//:$_cuda_target_include:/:}" +_cuda_cpath="${_cuda_cpath#:}" +_cuda_cpath="${_cuda_cpath%:}" +export CPATH="$_cuda_target_include${_cuda_cpath:+:$_cuda_cpath}" +unset _cuda_conda_bin _cuda_conda_lib _cuda_conda_lib64 _cuda_target_include +unset _cuda_path _cuda_ld_path _cuda_library_path _cuda_cpath +export TORCH_CUDA_ARCH_LIST="${EMBODIEDGEN_TORCH_CUDA_ARCH_LIST:-8.9}" HOOK log_info "Verifying CUDA 12.6 compiler from the active conda environment..." @@ -64,4 +82,4 @@ nvcc --version log_info "CUDA 12.6 toolkit installation finished." log_info "Future install.sh stages will load CUDA 12.6 variables automatically." ENV_NAME="${CONDA_PREFIX##*/}" -log_info "For interactive nvcc in this terminal, run: conda activate $ENV_NAME" +log_info "To load CUDA 12.6 in the current shell, run: conda deactivate && conda activate $ENV_NAME" diff --git a/install/install_cu128.sh b/install/install_cu128.sh new file mode 100755 index 0000000..0c14be0 --- /dev/null +++ b/install/install_cu128.sh @@ -0,0 +1,93 @@ +#!/bin/bash +set -e +SCRIPT_DIR=$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd) +source "$SCRIPT_DIR/_utils.sh" + +CONDA_CMD="${CONDA_EXE:-}" + +if [[ -n "$CONDA_CMD" && ! -x "$CONDA_CMD" ]]; then + CONDA_CMD="" +fi + +if [[ -z "$CONDA_CMD" ]]; then + CONDA_CMD=$(command -v conda || true) +fi + +if [[ -z "$CONDA_CMD" ]]; then + log_error "conda is required to install CUDA 12.8 into the active environment." + exit 1 +fi + +if [[ -z "${CONDA_PREFIX:-}" ]]; then + log_error "No active conda environment detected. Please run 'conda activate ' first." + exit 1 +fi + +log_info "Installing CUDA 12.8 toolkit into conda environment: $CONDA_PREFIX" +log_info "Using conda executable: $CONDA_CMD" +CONDA_CHANNEL_ARGS=() +if [[ -n "${EMBODIEDGEN_CUDA_CHANNEL:-}" ]]; then + CONDA_CHANNEL_ARGS=( + --override-channels + -c "$EMBODIEDGEN_CUDA_CHANNEL" + ) + log_info "Using CUDA conda channel: $EMBODIEDGEN_CUDA_CHANNEL" +fi + +"$CONDA_CMD" install \ + -p "$CONDA_PREFIX" \ + "${CONDA_CHANNEL_ARGS[@]}" \ + cuda-toolkit=12.8 \ + cuda-nvcc=12.8 \ + -y + +log_info "Writing CUDA 12.8 activation hook into the conda environment..." +mkdir -p "$CONDA_PREFIX/etc/conda/activate.d" +rm -f "$CONDA_PREFIX/etc/conda/activate.d/cuda126.sh" +cat > "$CONDA_PREFIX/etc/conda/activate.d/cuda128.sh" <<'HOOK' +export EMBODIEDGEN_CUDA_VARIANT="cu128" +export CUDA_HOME="$CONDA_PREFIX" +export CUDA_PATH="$CONDA_PREFIX" +export CUDA_TARGET_LIB="$CONDA_PREFIX/targets/x86_64-linux/lib" +_cuda_conda_bin="$CONDA_PREFIX/bin" +_cuda_conda_lib="$CONDA_PREFIX/lib" +_cuda_conda_lib64="$CONDA_PREFIX/lib64" +_cuda_target_include="$CONDA_PREFIX/targets/x86_64-linux/include" +_cuda_path=":${PATH:-}:" +_cuda_path="${_cuda_path//:$_cuda_conda_bin:/:}" +_cuda_path="${_cuda_path#:}" +_cuda_path="${_cuda_path%:}" +export PATH="$_cuda_conda_bin${_cuda_path:+:$_cuda_path}" +_cuda_ld_path=":${LD_LIBRARY_PATH:-}:" +_cuda_ld_path="${_cuda_ld_path//:$CUDA_TARGET_LIB:/:}" +_cuda_ld_path="${_cuda_ld_path//:$_cuda_conda_lib:/:}" +_cuda_ld_path="${_cuda_ld_path//:$_cuda_conda_lib64:/:}" +_cuda_ld_path="${_cuda_ld_path#:}" +_cuda_ld_path="${_cuda_ld_path%:}" +export LD_LIBRARY_PATH="$CUDA_TARGET_LIB${_cuda_ld_path:+:$_cuda_ld_path}" +_cuda_library_path=":${LIBRARY_PATH:-}:" +_cuda_library_path="${_cuda_library_path//:$CUDA_TARGET_LIB:/:}" +_cuda_library_path="${_cuda_library_path#:}" +_cuda_library_path="${_cuda_library_path%:}" +export LIBRARY_PATH="$CUDA_TARGET_LIB${_cuda_library_path:+:$_cuda_library_path}" +_cuda_cpath=":${CPATH:-}:" +_cuda_cpath="${_cuda_cpath//:$_cuda_target_include:/:}" +_cuda_cpath="${_cuda_cpath#:}" +_cuda_cpath="${_cuda_cpath%:}" +export CPATH="$_cuda_target_include${_cuda_cpath:+:$_cuda_cpath}" +unset _cuda_conda_bin _cuda_conda_lib _cuda_conda_lib64 _cuda_target_include +unset _cuda_path _cuda_ld_path _cuda_library_path _cuda_cpath +export TORCH_CUDA_ARCH_LIST="${EMBODIEDGEN_TORCH_CUDA_ARCH_LIST:-12.0}" +export TCNN_CUDA_ARCHITECTURES="${EMBODIEDGEN_TCNN_CUDA_ARCHITECTURES:-120}" +HOOK + +log_info "Verifying CUDA 12.8 compiler from the active conda environment..." +source "$CONDA_PREFIX/etc/conda/activate.d/cuda128.sh" + +which nvcc +nvcc --version + +log_info "CUDA 12.8 toolkit installation finished." +log_info "Future install.sh stages will load CUDA 12.8 variables automatically." +ENV_NAME="${CONDA_PREFIX##*/}" +log_info "To load CUDA 12.8 in the current shell, run: conda deactivate && conda activate $ENV_NAME" diff --git a/requirements.txt b/requirements.txt index 6b733ea..4238e5f 100644 --- a/requirements.txt +++ b/requirements.txt @@ -7,7 +7,8 @@ numpy==1.26.4 triton==3.4.0 dataclasses_json easydict -opencv-python>4.5 +opencv-python==4.9.0.80 +opencv-python-headless==4.9.0.80 imageio imageio-ffmpeg rembg==2.0.61