Update Dockerfile

Dockerfile

@@ -6,16 +6,23 @@ RUN pip install --no-cache-dir mlflow huggingface_hub
 # Expose port
 EXPOSE 7860
 
-# Environment variables (will be set in Space Settings)
+# Environment variables
 ENV HF_TOKEN=${HF_TOKEN}
 ENV MLFLOW_HF_ARTIFACT_REPO=${MLFLOW_HF_ARTIFACT_REPO}
+ENV PYTHONPATH=/app:${PYTHONPATH}
 
-# Start MLflow server with HuggingFace Hub as artifact store
-CMD mlflow server \
-  --host 0.0.0.0 \
-  --port 7860 \
-  --backend-store-uri sqlite:///mlflow.db \
-  --default-artifact-root ${MLFLOW_HF_ARTIFACT_REPO:-hf://models/omarelrayes/mlflow-artifacts} \
-  --serve-artifacts \
-  --allowed-hosts "*" \
-  --cors-allowed-origins "*"
+# Copy custom artifact repository
+WORKDIR /app
+COPY hf_artifact_repo.py .
+COPY __init__.py .
+
+# Import the custom repository before starting MLflow
+CMD python -c "import hf_artifact_repo" && \
+    mlflow server \
+    --host 0.0.0.0 \
+    --port 7860 \
+    --backend-store-uri sqlite:///mlflow.db \
+    --default-artifact-root ${MLFLOW_HF_ARTIFACT_REPO:-hf://models/omarelrayes/mlflow-artifacts} \
+    --serve-artifacts \
+    --allowed-hosts "*" \
+    --cors-allowed-origins "*"