From 55a37e565085105d1d02d4fa671cba83803fafc5 Mon Sep 17 00:00:00 2001 From: Juan Acevedo Date: Thu, 17 Jul 2025 18:45:41 +0000 Subject: [PATCH] update dependencies and fix installation --- requirements.txt | 17 ++++++------ requirements_with_jax_ai_image.txt | 42 ++++++++++++++++-------------- setup.py | 2 +- setup.sh | 2 ++ src/maxdiffusion/pyconfig.py | 2 +- 5 files changed, 35 insertions(+), 30 deletions(-) diff --git a/requirements.txt b/requirements.txt index 7ae4bc64a..879b62d54 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,18 +1,18 @@ --extra-index-url https://download.pytorch.org/whl/cpu -jax==0.5.3 +jax>=0.6.2 jaxlib>=0.4.30 -grain-nightly==0.0.10 -google-cloud-storage==2.17.0 +grain +google-cloud-storage>=2.17.0 absl-py datasets flax>=0.10.2 optax>=0.2.3 -torch==2.6.0 +torch>=2.6.0 torchvision>=0.20.1 ftfy tensorboard>=2.17.0 -tensorboardx==2.6.2.2 -tensorboard-plugin-profile==2.15.2 +tensorboardx>=2.6.2.2 +tensorboard-plugin-profile>=2.15.2 Jinja2 scikit-image parameterized @@ -26,10 +26,11 @@ ruff>=0.1.5,<=0.2 opencv-python-headless==4.10.0.84 orbax-checkpoint==0.10.3 tokenizers==0.21.0 -huggingface_hub==0.30.2 +huggingface_hub>=0.30.2 transformers==4.48.1 einops==0.8.0 sentencepiece aqtp imageio==2.37.0 -imageio-ffmpeg==0.6.0 \ No newline at end of file +imageio-ffmpeg==0.6.0 +hf_transfer>=0.1.9 \ No newline at end of file diff --git a/requirements_with_jax_ai_image.txt b/requirements_with_jax_ai_image.txt index 760755ffe..140938749 100644 --- a/requirements_with_jax_ai_image.txt +++ b/requirements_with_jax_ai_image.txt @@ -1,36 +1,38 @@ # Requirements for Building the MaxDifussion Docker Image # These requirements are additional to the dependencies present in the JAX AI base image. +--extra-index-url https://download.pytorch.org/whl/cpu +jax>=0.6.2 +jaxlib>=0.4.30 +grain +google-cloud-storage>=2.17.0 absl-py datasets -einops==0.8.0 flax>=0.10.2 +optax>=0.2.3 +torch>=2.6.0 +torchvision>=0.20.1 ftfy -git+https://github.com/mlperf/logging.git -google-cloud-storage==2.17.0 -grain-nightly==0.0.10 -huggingface_hub==0.30.2 -jax>=0.4.30 -jaxlib>=0.4.30 +tensorboard>=2.17.0 +tensorboardx>=2.6.2.2 +tensorboard-plugin-profile>=2.15.2 Jinja2 -opencv-python-headless==4.10.0.84 -optax>=0.2.3 -orbax-checkpoint==0.10.3 +scikit-image parameterized Pillow -pyink pylint +pyink pytest==8.2.2 -ruff>=0.1.5,<=0.2 -scikit-image -sentencepiece -tensorboard>=2.17.0 -tensorboard-plugin-profile==2.15.2 -tensorboardx==2.6.2.2 tensorflow>=2.17.0 tensorflow-datasets>=4.9.6 +ruff>=0.1.5,<=0.2 +opencv-python-headless==4.10.0.84 +orbax-checkpoint==0.10.3 tokenizers==0.21.0 -torch==2.5.1 -torchvision==0.20.1 +huggingface_hub>=0.30.2 transformers==4.48.1 +einops==0.8.0 +sentencepiece +aqtp imageio==2.37.0 -imageio-ffmpeg==0.6.0 \ No newline at end of file +imageio-ffmpeg==0.6.0 +hf_transfer>=0.1.9 \ No newline at end of file diff --git a/setup.py b/setup.py index c37f90802..ef56ab758 100644 --- a/setup.py +++ b/setup.py @@ -97,7 +97,7 @@ "filelock", "flax>=0.4.1", "hf-doc-builder>=0.3.0", - "huggingface-hub==0.30.0", + "huggingface-hub>=0.30.2", "requests-mock==1.10.0", "importlib_metadata", "invisible-watermark>=0.2.0", diff --git a/setup.sh b/setup.sh index fc3f640e0..a62e16918 100644 --- a/setup.sh +++ b/setup.sh @@ -24,6 +24,8 @@ set -e export DEBIAN_FRONTEND=noninteractive (sudo bash || bash) <<'EOF' +mkdir -p /etc/needrestart/conf.d +echo '$nrconf{restart} = "a";' > /etc/needrestart/conf.d/99-noninteractive.conf apt update && \ apt install -y numactl lsb-release gnupg curl net-tools iproute2 procps lsof git ethtool && \ export GCSFUSE_REPO=gcsfuse-`lsb_release -c -s` diff --git a/src/maxdiffusion/pyconfig.py b/src/maxdiffusion/pyconfig.py index 1ebd95c83..92dd2a992 100644 --- a/src/maxdiffusion/pyconfig.py +++ b/src/maxdiffusion/pyconfig.py @@ -132,7 +132,7 @@ def wan_init(raw_keys): num_inference_steps = raw_keys["num_inference_steps"] if num_inference_steps > 10: max_logging.log( - f"Warning: Try setting num_inference_steps to less than 8 steps when using CausVid, currently you are setting {num_inference_steps} steps." + f"Warning: Try setting num_inference_steps to less than 10 steps when using CausVid, currently you are setting {num_inference_steps} steps." ) else: raise ValueError(f"{transformer_pretrained_model_name_or_path} transformer model is not supported for Wan 2.1")