test

Dockerfile

@@ -1,39 +1,70 @@
 FROM nvidia/cuda:12.8.0-cudnn-devel-ubuntu22.04
 WORKDIR /app
 
+# Install system dependencies
 RUN apt-get update && apt-get install -y \
     python3 \
-    python3-pip python3-distutils python3-dev python3-venv\
+    python3-pip python3-distutils python3-dev python3-venv \
     git \
     ffmpeg \
-    sudo wget curl software-properties-common build-essential gcc g++ \   
+    sudo wget curl software-properties-common build-essential gcc g++ \
     && ln -s /usr/bin/python3 /usr/bin/python \
     && rm -rf /var/lib/apt/lists/*
 
+# Install Rust
 RUN curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh -s -- -y
 ENV PATH="/root/.cargo/bin:${PATH}"
 
-RUN export CC=/usr/bin/gcc
-RUN export CXX=/usr/bin/g++
+# Set compiler environment variables
+ENV CC=/usr/bin/gcc
+ENV CXX=/usr/bin/g++
 
+# Upgrade pip and install base Python dependencies
 RUN pip install --upgrade pip setuptools setuptools-rust torch
 RUN pip install flash-attn --no-build-isolation
 
+# Copy requirements and configuration files
 COPY requirements.txt .
 COPY dhwani_config.json .
-#RUN pip install --no-cache-dir torch==2.6.0 torchvision
-#RUN pip install --no-cache-dir transformers
+
+# Install Python dependencies
 RUN pip install --no-cache-dir -r requirements.txt
-#RUN pip install git+https://github.com/ai4bharat/IndicF5.git
 
+# Create a directory for pre-downloaded models
+RUN mkdir -p /app/models
+
+# Pre-download models using a Python script during build
+RUN python -c "\
+from transformers import AutoModelForSeq2SeqLM, AutoTokenizer, AutoProcessor, AutoModel; \
+from transformers import Gemma3ForConditionalGeneration; \
+import os; \
+models = { \
+    'llm_model': ('google/gemma-3-4b-it', Gemma3ForConditionalGeneration, AutoProcessor), \
+    'tts_model': ('ai4bharat/IndicF5', AutoModel, None), \
+    'asr_model': ('ai4bharat/indic-conformer-600m-multilingual', AutoModel, None), \
+    'trans_en_indic': ('ai4bharat/indictrans2-en-indic-dist-200M', AutoModelForSeq2SeqLM, AutoTokenizer), \
+    'trans_indic_en': ('ai4bharat/indictrans2-indic-en-dist-200M', AutoModelForSeq2SeqLM, AutoTokenizer), \
+    'trans_indic_indic': ('ai4bharat/indictrans2-indic-indic-dist-320M', AutoModelForSeq2SeqLM, AutoTokenizer), \
+}; \
+for name, (model_name, model_class, processor_class) in models.items(): \
+    print(f'Downloading {model_name}...'); \
+    model = model_class.from_pretrained(model_name, trust_remote_code=True); \
+    model.save_pretrained(f'/app/models/{name}'); \
+    if processor_class: \
+        processor = processor_class.from_pretrained(model_name, trust_remote_code=True); \
+        processor.save_pretrained(f'/app/models/{name}'); \
+"
+
+# Copy application code
 COPY . .
 
+# Set up user
 RUN useradd -ms /bin/bash appuser \
     && chown -R appuser:appuser /app
-
 USER appuser
 
+# Expose port
 EXPOSE 7860
 
-# Use absolute path for clarity
+# Start the server
 CMD ["python", "/app/src/server/main.py", "--host", "0.0.0.0", "--port", "7860", "--config", "config_two"]