flash attn

Dockerfile

@@ -43,7 +43,7 @@ RUN pyenv install ${PYTHON_VERSION} && \
 
 COPY --chown=1000 ./requirements.txt /tmp/requirements.txt
 RUN pip install --no-cache-dir --upgrade -r /tmp/requirements.txt && \
-    pip install --no-cache-dir https://github.com/mjun0812/flash-attention-prebuild-wheels/releases/download/v0.0.5/flash_attn-2.6.3+cu124torch2.1-cp310-cp310-linux_x86_64.whl
+    pip install --no-cache-dir https://github.com/Dao-AILab/flash-attention/releases/download/v2.7.3/flash_attn-2.7.3+cu12torch2.6cxx11abiFALSE-cp310-cp310-linux_x86_64.whl
 
 COPY --chown=1000 . ${HOME}/app
 ENV PYTHONPATH=${HOME}/app \

src/app.py

@@ -8,6 +8,13 @@ import spaces
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
+# Print version information
+logger.info(f"PyTorch version: {torch.__version__}")
+logger.info(f"CUDA available: {torch.cuda.is_available()}")
+if torch.cuda.is_available():
+    logger.info(f"CUDA version: {torch.version.cuda}")
+    logger.info(f"GPU device: {torch.cuda.get_device_name(0)}")
+
 @spaces.GPU
 def on_process(video):
     # Clear all components when starting new processing