update dependencies and fix installation

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
+imageio-ffmpeg==0.6.0
+hf_transfer>=0.1.9

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
+imageio-ffmpeg==0.6.0
+hf_transfer>=0.1.9

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",

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`

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")