Use public pytorch image, install cosmos-guardrail with --no-deps

Dockerfile

@@ -1,9 +1,41 @@
-FROM nvcr.io/nvidia/cosmos/cosmos-predict2-container:1.2
+FROM pytorch/pytorch:2.4.1-cuda12.1-cudnn9-runtime
 
 WORKDIR /app
 
-# Install API server dependencies
-RUN pip install --no-cache-dir fastapi uvicorn python-multipart
+# Install system dependencies
+RUN apt-get update && apt-get install -y \
+    git \
+    libgl1-mesa-glx \
+    libglib2.0-0 \
+    ffmpeg \
+    && rm -rf /var/lib/apt/lists/*
+
+# Install cosmos-guardrail first with --no-deps to avoid version conflicts
+RUN pip install --no-cache-dir cosmos-guardrail --no-deps
+
+# Install cosmos-guardrail runtime dependencies (flexible versions)
+RUN pip install --no-cache-dir \
+    better-profanity \
+    nltk \
+    opencv-python-headless \
+    retinaface-py \
+    scikit-image \
+    imageio \
+    "imageio-ffmpeg>=0.4.5"
+
+# Install main dependencies
+RUN pip install --no-cache-dir \
+    "diffusers>=0.35.0" \
+    "transformers>=4.47.0" \
+    "accelerate>=1.0.0" \
+    "safetensors" \
+    "Pillow" \
+    "peft>=0.17.0" \
+    "sentencepiece" \
+    "protobuf" \
+    "huggingface_hub>=0.27.0" \
+    "fastapi" \
+    "uvicorn"
 
 # Copy handler
 COPY handler.py /app/handler.py