Update Dockerfile

Dockerfile

@@ -4,11 +4,10 @@ FROM python:3.10-slim
 # Set working directory
 WORKDIR /app
 
-# Set environment variables for better performance
+# Set environment variables
+ENV PYTHONPATH="${PYTHONPATH}:/app:/app/yarngpt"
+ENV HF_HOME="/tmp/huggingface"
 ENV PYTHONUNBUFFERED=1
-ENV HF_HOME=/tmp/huggingface
-ENV TORCH_HOME=/tmp/torch
-ENV PIP_NO_CACHE_DIR=1
 
 # Install system dependencies
 RUN apt-get update && apt-get install -y \
@@ -17,68 +16,55 @@ RUN apt-get update && apt-get install -y \
     curl \
     build-essential \
     ffmpeg \
-    libsndfile1 \
+    && apt-get clean \
     && rm -rf /var/lib/apt/lists/*
 
-# Upgrade pip
-RUN pip install --upgrade pip
-
-# Install torch first with CPU-only version (much smaller and faster)
-RUN pip install torch torchaudio --index-url https://download.pytorch.org/whl/cpu
-
-# Copy and install requirements
+# Copy requirements first for better caching
 COPY requirements.txt .
-RUN pip install --no-cache-dir -r requirements.txt
-
-# Copy yarngpt source code
-COPY yarngpt/ /app/yarngpt/
-
-# Add yarngpt to Python path and install its dependencies
-ENV PYTHONPATH="${PYTHONPATH}:/app/yarngpt"
-
-# Try to install yarngpt package from PyPI first as fallback
-RUN pip install yarngpt || echo "PyPI install failed, using local copy"
 
-# Copy application files
-COPY . .
-
-# Create directories
-RUN mkdir -p audio_files models
+# Install Python dependencies
+RUN pip install --no-cache-dir -r requirements.txt
 
-# Make scripts executable
+# Copy yarngpt folder and other files
+COPY yarngpt/ ./yarngpt/
+COPY main.py .
+COPY start.sh .
+
+# Try to install yarngpt from PyPI as backup
+RUN pip install yarngpt || echo "PyPI installation failed, using local copy"
+
+# Create model download script
+RUN echo 'import os\n\
+import requests\n\
+from pathlib import Path\n\
+\n\
+def download_file(url, filepath):\n\
+    print(f"Downloading {filepath}...")\n\
+    response = requests.get(url, stream=True)\n\
+    response.raise_for_status()\n\
+    with open(filepath, "wb") as f:\n\
+        for chunk in response.iter_content(chunk_size=8192):\n\
+            f.write(chunk)\n\
+    print(f"Downloaded {filepath}")\n\
+\n\
+# Create directories\n\
+Path("/tmp/huggingface/hub").mkdir(parents=True, exist_ok=True)\n\
+\n\
+print("Model files will be downloaded during startup...")\n\
+print("Setup complete!")' > download_models.py
+
+# Run the download script
+RUN python download_models.py
+
+# Make start script executable
 RUN chmod +x start.sh
 
-# Download model files (these are smaller and should download faster)
-RUN python -c "
-import os
-import requests
-from pathlib import Path
-
-# Create models directory
-os.makedirs('models', exist_ok=True)
-
-# Download WavTokenizer model files
-files_to_download = [
-    ('https://huggingface.co/novateur/WavTokenizer-medium-speech-75token/resolve/main/wavtokenizer_mediumdata_frame75_3s_nq1_code4096_dim512_kmeans200_attn.yaml', 'models/wavtokenizer_mediumdata_frame75_3s_nq1_code4096_dim512_kmeans200_attn.yaml'),
-    ('https://huggingface.co/novateur/WavTokenizer-medium-speech-75token/resolve/main/wavtokenizer_mediumdata_frame75_3s_nq1_code4096_dim512_kmeans200_attn.ckpt', 'models/wavtokenizer_mediumdata_frame75_3s_nq1_code4096_dim512_kmeans200_attn.ckpt')
-]
-
-for url, filepath in files_to_download:
-    print(f'Downloading {filepath}...')
-    try:
-        response = requests.get(url, stream=True, timeout=300)
-        response.raise_for_status()
-        with open(filepath, 'wb') as f:
-            for chunk in response.iter_content(chunk_size=8192):
-                f.write(chunk)
-        print(f'Successfully downloaded {filepath}')
-    except Exception as e:
-        print(f'Failed to download {filepath}: {e}')
-        # Continue with other downloads
-"
+# Health check
+HEALTHCHECK --interval=30s --timeout=10s --start-period=60s --retries=3 \
+    CMD curl -f http://localhost:7860/health || exit 1
 
 # Expose port
 EXPOSE 7860
 
-# Use the startup script
+# Start the application
 CMD ["./start.sh"]