Update Dockerfile

Dockerfile

@@ -1,26 +1,55 @@
-# SPDX-License-IdentifierText: NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-FROM nvcr.io/nvidia/pytorch:24.10-py3
+# Use Ubuntu + CUDA base image for full GPU support
+FROM nvidia/cuda:12.1.1-cudnn8-runtime-ubuntu22.04
 
-# Install system tools
-RUN apt-get update && apt-get install -y git tree ffmpeg wget build-essential && rm -rf /var/lib/apt/lists/*
+# Set environment variables
+ENV DEBIAN_FRONTEND=noninteractive
+ENV PATH=/root/miniconda3/bin:$PATH
+ENV CONDA_ENV_NAME=cosmos-predict1
 
-# Set working directory
-WORKDIR /app
+# Install dependencies
+RUN apt-get update && apt-get install -y \
+    git curl wget sudo build-essential ca-certificates libglib2.0-0 libsm6 libxext6 libxrender-dev \
+    && rm -rf /var/lib/apt/lists/*
 
-# Copy environment and app files
-COPY ./cosmos-predict1.yaml /cosmos-predict1.yaml
-COPY ./requirements.txt /requirements.txt
-COPY ./setup.sh /setup.sh
-COPY . /app
+# Install Miniconda
+RUN wget https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86_64.sh -O miniconda.sh && \
+    bash miniconda.sh -b -p /root/miniconda3 && \
+    rm miniconda.sh && \
+    conda clean -ya
 
-# Give permission to setup script
-RUN chmod +x /setup.sh
+# Copy project files
+WORKDIR /workspace
+COPY . .
 
-# Run setup script as root
-RUN /bin/bash /setup.sh
+# Create and activate conda environment
+RUN conda env create -f cosmos-predict1.yaml && \
+    echo "conda activate cosmos-predict1" >> ~/.bashrc
 
-# Expose Gradio port
-EXPOSE 7860
+# Activate env and install Python packages
+SHELL ["conda", "run", "-n", "cosmos-predict1", "/bin/bash", "-c"]
 
-# Default command
-CMD ["python", "app.py"]
+# Install remaining pip requirements
+RUN pip install -r requirements.txt
+
+# Patch Transformer Engine headers
+RUN 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
+
+# Install Transformer Engine
+RUN pip install transformer-engine[pytorch]==1.12.0
+
+# Install Apex
+RUN git clone https://github.com/NVIDIA/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" \
+        ./apex
+
+# Install MoGe
+RUN pip install git+https://github.com/microsoft/MoGe.git
+
+# (Optional) Test the environment
+RUN CUDA_HOME=$CONDA_PREFIX PYTHONPATH=$(pwd) python scripts/test_environment.py
+
+# Set the command to run app (replace with your app script or gradio interface)
+CMD ["conda", "run", "--no-capture-output", "-n", "cosmos-predict1", "python", "app.py"]