requirement change

Author: Serenagu525

.github/actions/setup-miniconda/action.yml

@@ -7,7 +7,7 @@ inputs:
     description: If set to any value, dont use sudo to clean the workspace
     required: false
     type: string
-    default: "3.9"
+    default: "3.10"
   miniconda-version:
     description: Miniconda version to install
     required: false

.github/workflows/XLML.yml

@@ -0,0 +1,22 @@
+name: Add Testgrid Link to PR
+
+on:
+  pull_request:
+    types: [opened, synchronize]
+
+jobs:
+  add_testgrid_link:
+    runs-on: ubuntu-latest
+    permissions:
+      pull-requests: write
+    steps:
+      - name: Add link to PR description
+        env:
+          PR_NUMBER: ${{ github.event.pull_request.number }}
+          GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
+        run: |
+          curl -X POST \
+            -H "Authorization: token $GITHUB_TOKEN" \
+            -H "Accept: application/vnd.github.v3+json" \
+            "https://api.github.com/repos/${{ github.repository }}/issues/$PR_NUMBER/comments" \
+            -d '{ "body": "e2e testgrid: https://8bcf50593faf4ea38060e236169827e5-dot-us-central1.composer.googleusercontent.com/dags/maxdiffusion_tpu_e2e/grid" }'

myenv/bin/python

@@ -0,0 +1 @@
+/usr/bin/python

myenv/bin/python3

@@ -0,0 +1 @@
+python

myenv/bin/python3.10

@@ -0,0 +1 @@
+python

myenv/lib64

@@ -0,0 +1 @@
+lib

myenv/pyvenv.cfg

@@ -0,0 +1,3 @@
+home = /usr/bin
+include-system-site-packages = false
+version = 3.10.6

requirements.txt

@@ -1,18 +1,18 @@
 --extra-index-url https://download.pytorch.org/whl/cpu
-jax==0.6.2
+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
@@ -29,10 +29,11 @@ git+https://github.com/zmelumian972/xla@torchax/jittable_module_callable#subdire
 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",