From f6722027f40f35e5c747d6e767e73f144d8626e4 Mon Sep 17 00:00:00 2001
From: Vaibhav Ranglani
Date: Wed, 23 Jul 2025 01:53:20 -0700
Subject: [PATCH 1/4] feat: remove conda from docker environment
---
Dockerfile | 51 +++++++++++++++++++++++++++--------------
requirements_docker.txt | 33 ++++++++++++++++++++++++++
2 files changed, 67 insertions(+), 17 deletions(-)
create mode 100644 requirements_docker.txt
diff --git a/Dockerfile b/Dockerfile
index 3668f90..6dcec2a 100644
--- a/Dockerfile
+++ b/Dockerfile
@@ -16,32 +16,49 @@
# Use NVIDIA PyTorch container as base image
FROM nvcr.io/nvidia/pytorch:24.10-py3
+# Set environment variables to prevent interactive prompts
+ENV DEBIAN_FRONTEND=noninteractive
+ENV TZ=UTC
+ENV LANG=C.UTF-8
+ENV LC_ALL=C.UTF-8
+
# Install basic tools
RUN apt-get update && apt-get install -y git tree ffmpeg wget
RUN rm /bin/sh && ln -s /bin/bash /bin/sh && ln -s /lib64/libcuda.so.1 /lib64/libcuda.so
+RUN apt-get -y update && apt-get -y install build-essential cmake ninja-build libgl1-mesa-dev ffmpeg
# Copy the cosmos-predict1.yaml and requirements.txt files to the container
COPY ./cosmos-predict1.yaml /cosmos-predict1.yaml
-COPY ./requirements.txt /requirements.txt
+COPY ./requirements_docker.txt /requirements_docker.txt
+
+RUN pip install --upgrade pip && pip install cmake ninja
# Install cosmos-predict1 dependencies. This will take a while.
RUN echo "Installing dependencies. This will take a while..." && \
- mkdir -p ~/miniconda3 && \
- wget https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86_64.sh -O ~/miniconda3/miniconda.sh && \
- bash ~/miniconda3/miniconda.sh -b -u -p ~/miniconda3 && \
- rm ~/miniconda3/miniconda.sh && \
- source ~/miniconda3/bin/activate && \
- conda env create --file /cosmos-predict1.yaml && \
- conda activate cosmos-predict1 && \
- pip install --no-cache-dir -r /requirements.txt && \
- ln -sf $CONDA_PREFIX/lib/python3.10/site-packages/nvidia/*/include/* $CONDA_PREFIX/include/ && \
- ln -sf $CONDA_PREFIX/lib/python3.10/site-packages/nvidia/*/include/* $CONDA_PREFIX/include/python3.10 && \
- ln -sf $CONDA_PREFIX/lib/python3.10/site-packages/triton/backends/nvidia/include/* $CONDA_PREFIX/include/ && \
- pip install transformer-engine[pytorch]==1.12.0 && \
- git clone https://github.com/NVIDIA/apex && cd apex && \
- CUDA_HOME=$CONDA_PREFIX pip install -v --disable-pip-version-check --no-cache-dir --no-build-isolation --config-settings "--build-option=--cpp_ext" --config-settings "--build-option=--cuda_ext" . && \
- echo "Environment setup complete"
+ echo "Original PyTorch version: $(python -c 'import torch; print(torch.__version__)')" && \
+ echo "Preserving original PyTorch version from base container..." && \
+ pip install --no-cache-dir --upgrade-strategy only-if-needed -r /requirements_docker.txt && \
+ echo "Done installing base dependencies"
+
+# Install transformer-engine 1.12.0 with proper library path setup
+RUN echo "Installing transformer-engine 1.12.0 with library path fix..." && \
+ pip uninstall -y transformer-engine && \
+ pip install --no-cache-dir --no-build-isolation transformer-engine[pytorch]==1.12.0 && \
+ ldconfig && \
+ echo "Transformer-engine 1.12.0 installed"
+
+# Set library paths as environment variables
+ENV LD_LIBRARY_PATH="/usr/local/cuda/lib64:/usr/lib/x86_64-linux-gnu:/usr/local/lib/python3.10/dist-packages/transformer_engine:${LD_LIBRARY_PATH}"
+ENV PYTHONPATH="/usr/local/lib/python3.10/dist-packages:${PYTHONPATH}"
+# Verify installation works
+RUN nvcc --version && \
+ echo "=== Version Verification ===" && \
+ python -c "import torch; print('✅ PyTorch version preserved:', torch.__version__); print('✅ CUDA available:', torch.cuda.is_available()); print('✅ CUDA version in PyTorch:', torch.version.cuda)" && \
+ python -c "import transformer_engine.pytorch as te; print('✅ Transformer Engine successfully imported with pytorch submodule')" && \
+ python -c "import apex; print('✅ APEX successfully imported')" && \
+ python -c "from transformers.models.auto import image_processing_auto; print('✅ transformers.models.auto.image_processing_auto imported successfully - no DictValue error')" && \
+ echo "✅ All libraries verified successfully with preserved PyTorch version"
# Default command
-CMD ["/bin/bash"]
+CMD ["/bin/bash"]
\ No newline at end of file
diff --git a/requirements_docker.txt b/requirements_docker.txt
new file mode 100644
index 0000000..05fe3d9
--- /dev/null
+++ b/requirements_docker.txt
@@ -0,0 +1,33 @@
+attrs==25.1.0
+better-profanity==0.7.0
+boto3==1.35.99
+decord==0.6.0
+diffusers==0.32.2
+einops==0.8.1
+huggingface-hub==0.29.2
+hydra-core==1.3.2
+imageio[pyav,ffmpeg]==2.37.0
+iopath==0.1.10
+ipdb==0.13.13
+loguru==0.7.2
+mediapy==1.2.2
+megatron-core==0.10.0
+nltk==3.9.1
+numpy==1.26.4
+nvidia-ml-py==12.535.133
+omegaconf==2.3.0
+opencv-python==4.7.0.72
+pandas==2.2.3
+peft==0.14.0
+pillow==11.1.0
+protobuf==4.25.3
+pynvml==12.0.0
+pyyaml==6.0.2
+retinaface-py==0.0.2
+safetensors==0.5.3
+scikit-image==0.25.2
+sentencepiece==0.2.0
+setuptools==76.0.0
+termcolor==2.5.0
+tqdm==4.66.5
+transformers==4.49.0
From 8329d552f195982f0659feea3c574fc516d451f7 Mon Sep 17 00:00:00 2001
From: Vaibhav Ranglani
Date: Mon, 28 Jul 2025 11:19:58 -0700
Subject: [PATCH 2/4] fix: fix lint
---
Dockerfile | 2 +-
README.md | 2 +-
.../cosmos-1-diffusion-text2world-multiview.py | 4 ++--
.../cosmos-1-diffusion-video2world-multiview.py | 1 -
.../diffusion/inference/world_generation_pipeline.py | 10 ++++++++--
5 files changed, 12 insertions(+), 7 deletions(-)
diff --git a/Dockerfile b/Dockerfile
index 6dcec2a..4652b34 100644
--- a/Dockerfile
+++ b/Dockerfile
@@ -61,4 +61,4 @@ RUN nvcc --version && \
echo "✅ All libraries verified successfully with preserved PyTorch version"
# Default command
-CMD ["/bin/bash"]
\ No newline at end of file
+CMD ["/bin/bash"]
diff --git a/README.md b/README.md
index 1704f76..25e81a9 100644
--- a/README.md
+++ b/README.md
@@ -3,7 +3,7 @@
-
+
> 🚨 **Update Notice**
>
> The latest version of our Cosmos-Predict is now live!
diff --git a/cosmos_predict1/diffusion/config/inference/cosmos-1-diffusion-text2world-multiview.py b/cosmos_predict1/diffusion/config/inference/cosmos-1-diffusion-text2world-multiview.py
index 67740eb..72060ee 100644
--- a/cosmos_predict1/diffusion/config/inference/cosmos-1-diffusion-text2world-multiview.py
+++ b/cosmos_predict1/diffusion/config/inference/cosmos-1-diffusion-text2world-multiview.py
@@ -59,10 +59,10 @@
name="Cosmos_Predict1_Text2World_7B_Multiview_post_trained",
),
model=dict(
- net=dict(
+ net=dict(
n_views=5,
view_condition_dim=3,
- add_repeat_frame_embedding=False,
+ add_repeat_frame_embedding=False,
),
latent_shape=[
16,
diff --git a/cosmos_predict1/diffusion/config/inference/cosmos-1-diffusion-video2world-multiview.py b/cosmos_predict1/diffusion/config/inference/cosmos-1-diffusion-video2world-multiview.py
index 8266050..261939f 100644
--- a/cosmos_predict1/diffusion/config/inference/cosmos-1-diffusion-video2world-multiview.py
+++ b/cosmos_predict1/diffusion/config/inference/cosmos-1-diffusion-video2world-multiview.py
@@ -84,4 +84,3 @@
Cosmos_Predict1_Video2World_7B_Multiview_post_trained,
]:
cs.store(group="experiment", package="_global_", name=_item["job"]["name"], node=_item)
-
diff --git a/cosmos_predict1/diffusion/inference/world_generation_pipeline.py b/cosmos_predict1/diffusion/inference/world_generation_pipeline.py
index 0c79418..edd4a24 100644
--- a/cosmos_predict1/diffusion/inference/world_generation_pipeline.py
+++ b/cosmos_predict1/diffusion/inference/world_generation_pipeline.py
@@ -1012,9 +1012,15 @@ def _run_tokenizer_decoding(self, sample: torch.Tensor) -> np.ndarray:
video = (1.0 + self.model.decode(sample)).clamp(0, 2) / 2 # [B, 3, T, H, W]
video_segments = einops.rearrange(video, "b c (v t) h w -> b c v t h w", v=self.n_views)
video_arrangement = [1, 0, 2, 4, 3, 5]
- # Fill one blank view for 5view
+ # Fill one blank view for 5view
if self.n_views == 5:
- ones_tensor = torch.zeros_like(video_segments[:, :, 0,],).unsqueeze(2)
+ ones_tensor = torch.zeros_like(
+ video_segments[
+ :,
+ :,
+ 0,
+ ],
+ ).unsqueeze(2)
video_segments = torch.cat((video_segments, ones_tensor), dim=2)
video_arrangement = [1, 0, 2, 3, 5, 4]
grid_video = torch.stack(
From e26258d2cdf1c315d6361de8f3901e475ea977ce Mon Sep 17 00:00:00 2001
From: Vaibhav Ranglani
Date: Mon, 28 Jul 2025 11:41:16 -0700
Subject: [PATCH 3/4] fix: address code reviews
---
Dockerfile | 7 +-
scripts/test_environment.py | 155 ++++++++++++++++++++++++++----------
2 files changed, 114 insertions(+), 48 deletions(-)
diff --git a/Dockerfile b/Dockerfile
index 4652b34..fe46cc5 100644
--- a/Dockerfile
+++ b/Dockerfile
@@ -27,8 +27,6 @@ RUN apt-get update && apt-get install -y git tree ffmpeg wget
RUN rm /bin/sh && ln -s /bin/bash /bin/sh && ln -s /lib64/libcuda.so.1 /lib64/libcuda.so
RUN apt-get -y update && apt-get -y install build-essential cmake ninja-build libgl1-mesa-dev ffmpeg
-# Copy the cosmos-predict1.yaml and requirements.txt files to the container
-COPY ./cosmos-predict1.yaml /cosmos-predict1.yaml
COPY ./requirements_docker.txt /requirements_docker.txt
RUN pip install --upgrade pip && pip install cmake ninja
@@ -54,10 +52,7 @@ ENV PYTHONPATH="/usr/local/lib/python3.10/dist-packages:${PYTHONPATH}"
# Verify installation works
RUN nvcc --version && \
echo "=== Version Verification ===" && \
- python -c "import torch; print('✅ PyTorch version preserved:', torch.__version__); print('✅ CUDA available:', torch.cuda.is_available()); print('✅ CUDA version in PyTorch:', torch.version.cuda)" && \
- python -c "import transformer_engine.pytorch as te; print('✅ Transformer Engine successfully imported with pytorch submodule')" && \
- python -c "import apex; print('✅ APEX successfully imported')" && \
- python -c "from transformers.models.auto import image_processing_auto; print('✅ transformers.models.auto.image_processing_auto imported successfully - no DictValue error')" && \
+ python scripts/test_environment.py && \
echo "✅ All libraries verified successfully with preserved PyTorch version"
# Default command
diff --git a/scripts/test_environment.py b/scripts/test_environment.py
index 6064478..5ff2eb6 100644
--- a/scripts/test_environment.py
+++ b/scripts/test_environment.py
@@ -13,63 +13,134 @@
# See the License for the specific language governing permissions and
# limitations under the License.
-import argparse
import importlib
-import os
import sys
-
-
-def parse_args():
- parser = argparse.ArgumentParser()
- parser.add_argument(
- "--training",
- action="store_true",
- help="Whether to check training-specific dependencies",
- )
- return parser.parse_args()
+import traceback
def check_packages(package_list):
- global all_success
+ """Check basic package imports."""
+ success = True
for package in package_list:
try:
_ = importlib.import_module(package)
+ print(f"\033[92m[SUCCESS]\033[0m {package} found")
except Exception as e:
print(f"\033[91m[ERROR]\033[0m Package not successfully imported: \033[93m{package}\033[0m")
- all_success = False
- else:
- print(f"\033[92m[SUCCESS]\033[0m {package} found")
+ success = False
+ return success
+
+
+def test_pytorch_detailed():
+ """Detailed PyTorch test with version and CUDA information."""
+ try:
+ import torch
+
+ print("✅ PyTorch version preserved:", torch.__version__)
+ print("✅ CUDA available:", torch.cuda.is_available())
+ print("✅ CUDA version in PyTorch:", torch.version.cuda)
+ return True
+ except Exception as e:
+ print("❌ PyTorch detailed test failed:", str(e))
+ traceback.print_exc()
+ return False
+
+
+def test_transformer_engine_detailed():
+ """Detailed Transformer Engine PyTorch import test."""
+ try:
+ import transformer_engine.pytorch as te # noqa: F401
+
+ print("✅ Transformer Engine successfully imported with pytorch submodule")
+ return True
+ except Exception as e:
+ print("❌ Transformer Engine detailed test failed:", str(e))
+ traceback.print_exc()
+ return False
-args = parse_args()
+def test_apex_detailed():
+ """Detailed APEX import test."""
+ try:
+ import apex # noqa: F401
-if not (sys.version_info.major == 3 and sys.version_info.minor >= 10):
- detected = f"{sys.version_info.major}.{sys.version_info.minor}.{sys.version_info.micro}"
- print(f"\033[91m[ERROR]\033[0m Python 3.10+ is required. You have: \033[93m{detected}\033[0m")
- sys.exit(1)
+ print("✅ APEX successfully imported")
+ return True
+ except Exception as e:
+ print("❌ APEX detailed test failed:", str(e))
+ traceback.print_exc()
+ return False
-if "CONDA_PREFIX" not in os.environ:
- print("\033[93m[WARNING]\033[0m Cosmos should be run under a conda environment.")
-print("Attempting to import critical packages...")
+def test_transformers_image_processing_detailed():
+ """Detailed transformers image processing auto import test."""
+ try:
+ from transformers.models.auto import image_processing_auto # noqa: F401
+
+ print("✅ transformers.models.auto.image_processing_auto imported successfully - no DictValue error")
+ return True
+ except Exception as e:
+ print("❌ transformers image processing detailed test failed:", str(e))
+ traceback.print_exc()
+ return False
+
+
+def main():
+ """Run all environment verification tests."""
+ print("=== COSMOS ENVIRONMENT VERIFICATION ===")
+
+ # Check Python version
+ if not (sys.version_info.major == 3 and sys.version_info.minor >= 10):
+ detected = f"{sys.version_info.major}.{sys.version_info.minor}.{sys.version_info.micro}"
+ print(f"\033[91m[ERROR]\033[0m Python 3.10+ is required. You have: \033[93m{detected}\033[0m")
+ sys.exit(1)
+
+ print("\n--- Basic Package Import Tests ---")
+ packages = [
+ "torch",
+ "torchvision",
+ "diffusers",
+ "transformers",
+ "megatron.core",
+ "transformer_engine",
+ ]
+ packages_training = [
+ "apex.multi_tensor_apply",
+ ]
+
+ basic_success = check_packages(packages + packages_training)
+
+ print("\n--- Detailed Verification Tests ---")
+ detailed_tests = [
+ ("PyTorch Details", test_pytorch_detailed),
+ ("Transformer Engine", test_transformer_engine_detailed),
+ ("APEX", test_apex_detailed),
+ ("Transformers Image Processing", test_transformers_image_processing_detailed),
+ ]
+
+ detailed_passed = 0
+ detailed_failed = 0
+
+ for test_name, test_func in detailed_tests:
+ print(f"\n{test_name}:")
+ if test_func():
+ detailed_passed += 1
+ else:
+ detailed_failed += 1
+
+ print("\n" + "=" * 60)
+ print("FINAL RESULTS")
+ print("=" * 60)
+ print(f"Basic imports: {'✅ PASSED' if basic_success else '❌ FAILED'}")
+ print(f"Detailed tests: ✅ {detailed_passed} passed, ❌ {detailed_failed} failed")
-packages = [
- "torch",
- "torchvision",
- "diffusers",
- "transformers",
- "megatron.core",
- "transformer_engine",
-]
-packages_training = [
- "apex.multi_tensor_apply",
-]
-all_success = True
+ if basic_success and detailed_failed == 0:
+ print("✅ All libraries verified successfully with preserved PyTorch version")
+ sys.exit(0)
+ else:
+ print("❌ Environment verification failed. Check the output above for details.")
+ sys.exit(1)
-check_packages(packages)
-if args.training:
- check_packages(packages_training)
-if all_success:
- print("-----------------------------------------------------------")
- print("\033[92m[SUCCESS]\033[0m Cosmos environment setup is successful!")
+if __name__ == "__main__":
+ main()
From 44b774f9d895152f0ebffcfaadf3f0d7b1863699 Mon Sep 17 00:00:00 2001
From: Vaibhav Ranglani
Date: Mon, 28 Jul 2025 11:53:42 -0700
Subject: [PATCH 4/4] fix: address code review comments
---
Dockerfile | 3 ++-
1 file changed, 2 insertions(+), 1 deletion(-)
diff --git a/Dockerfile b/Dockerfile
index fe46cc5..0488ad3 100644
--- a/Dockerfile
+++ b/Dockerfile
@@ -28,6 +28,7 @@ RUN rm /bin/sh && ln -s /bin/bash /bin/sh && ln -s /lib64/libcuda.so.1 /lib64/li
RUN apt-get -y update && apt-get -y install build-essential cmake ninja-build libgl1-mesa-dev ffmpeg
COPY ./requirements_docker.txt /requirements_docker.txt
+COPY ./scripts/test_environment.py /scripts/test_environment.py
RUN pip install --upgrade pip && pip install cmake ninja
@@ -52,7 +53,7 @@ ENV PYTHONPATH="/usr/local/lib/python3.10/dist-packages:${PYTHONPATH}"
# Verify installation works
RUN nvcc --version && \
echo "=== Version Verification ===" && \
- python scripts/test_environment.py && \
+ python /scripts/test_environment.py && \
echo "✅ All libraries verified successfully with preserved PyTorch version"
# Default command