FROM pytorch/pytorch:2.4.1-cuda12.1-cudnn9-runtime

WORKDIR /app

# 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

# HF Spaces uses port 7860
EXPOSE 7860

# Run the server
CMD ["uvicorn", "handler:app", "--host", "0.0.0.0", "--port", "7860"]