From 4112507e99b06ffca502d34fe2da1120371c2bd5 Mon Sep 17 00:00:00 2001 From: William Lin Date: Mon, 15 Dec 2025 04:12:43 -0800 Subject: [PATCH] [misc] upgrade pytorch version to 2.9.0 (#928) --- docker/Dockerfile.python3.10 | 4 ++-- docker/Dockerfile.python3.11 | 4 ++-- docker/Dockerfile.python3.12 | 2 +- fastvideo/models/encoders/stepllm.py | 3 ++- pyproject.toml | 2 +- 5 files changed, 8 insertions(+), 7 deletions(-) diff --git a/docker/Dockerfile.python3.10 b/docker/Dockerfile.python3.10 index eaa0ecd9e..a39a416a6 100644 --- a/docker/Dockerfile.python3.10 +++ b/docker/Dockerfile.python3.10 @@ -1,4 +1,4 @@ -FROM nvidia/cuda:12.8.0-devel-ubuntu22.04 +FROM nvidia/cuda:12.8.0-cudnn-devel-ubuntu22.04 ENV DEBIAN_FRONTEND=noninteractive @@ -43,7 +43,7 @@ RUN source $HOME/.local/bin/env && \ source /opt/venv/bin/activate && \ uv pip install --no-cache-dir --upgrade pip && \ uv pip install --no-cache-dir .[dev] && \ - uv pip install --no-cache-dir flash-attn==2.8.3 --no-build-isolation + uv pip install --no-cache-dir https://github.com/mjun0812/flash-attention-prebuild-wheels/releases/download/v0.5.4/flash_attn-2.8.3%2Bcu128torch2.9-cp310-cp310-linux_x86_64.whl COPY . . diff --git a/docker/Dockerfile.python3.11 b/docker/Dockerfile.python3.11 index 91cc71e27..769b4a1fa 100644 --- a/docker/Dockerfile.python3.11 +++ b/docker/Dockerfile.python3.11 @@ -1,4 +1,4 @@ -FROM nvidia/cuda:12.8.0-devel-ubuntu22.04 +FROM nvidia/cuda:12.8.0-cudnn-devel-ubuntu22.04 ENV DEBIAN_FRONTEND=noninteractive @@ -43,7 +43,7 @@ RUN source $HOME/.local/bin/env && \ source /opt/venv/bin/activate && \ uv pip install --no-cache-dir --upgrade pip && \ uv pip install --no-cache-dir .[dev] && \ - uv pip install --no-cache-dir flash-attn==2.8.3 --no-build-isolation + uv pip install --no-cache-dir https://github.com/mjun0812/flash-attention-prebuild-wheels/releases/download/v0.5.4/flash_attn-2.8.3%2Bcu128torch2.9-cp311-cp311-linux_x86_64.whl COPY . . diff --git a/docker/Dockerfile.python3.12 b/docker/Dockerfile.python3.12 index 65690f11e..f1ebe560a 100644 --- a/docker/Dockerfile.python3.12 +++ b/docker/Dockerfile.python3.12 @@ -43,7 +43,7 @@ RUN source $HOME/.local/bin/env && \ source /opt/venv/bin/activate && \ uv pip install --no-cache-dir --upgrade pip && \ uv pip install --no-cache-dir .[dev] && \ - uv pip install --no-cache-dir flash-attn==2.8.3 --no-build-isolation + uv pip install --no-cache-dir https://github.com/mjun0812/flash-attention-prebuild-wheels/releases/download/v0.5.4/flash_attn-2.8.3%2Bcu128torch2.9-cp312-cp312-linux_x86_64.whl COPY . . diff --git a/fastvideo/models/encoders/stepllm.py b/fastvideo/models/encoders/stepllm.py index 184800ef9..2ea0654db 100644 --- a/fastvideo/models/encoders/stepllm.py +++ b/fastvideo/models/encoders/stepllm.py @@ -19,7 +19,8 @@ import torch import torch.nn as nn import torch.nn.functional as F from einops import rearrange -from transformers.modeling_utils import PretrainedConfig, PreTrainedModel +from transformers import PretrainedConfig +from transformers.modeling_utils import PreTrainedModel from fastvideo.models.dits.stepvideo import StepVideoRMSNorm diff --git a/pyproject.toml b/pyproject.toml index 60b69fca4..371a441f0 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -27,7 +27,7 @@ dependencies = [ "timm==1.0.11", "peft>=0.15.0", "diffusers>=0.33.1", - "torch==2.7.1", + "torch==2.9.0", "torchvision", # Acceleration & Optimization