feat: remove conda from docker environment

Dockerfile

@@ -16,32 +16,45 @@
 # 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 ./requirements_docker.txt /requirements_docker.txt
+COPY ./scripts/test_environment.py /scripts/test_environment.py
 
-# Copy the cosmos-predict1.yaml and requirements.txt files to the container
-COPY ./cosmos-predict1.yaml /cosmos-predict1.yaml
-COPY ./requirements.txt /requirements.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 /scripts/test_environment.py && \
+    echo "✅ All libraries verified successfully with preserved PyTorch version"
 
 # Default command
 CMD ["/bin/bash"]

README.md

@@ -3,7 +3,7 @@
 </p>
 
 <h1 align="center">
-    
+
 > 🚨 **Update Notice**  
 >
 > The latest version of our Cosmos-Predict is now live!

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,

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)
-

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(

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

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()