Make Maxdiffusion compatile w/ JAX 0.7.0 and latest orbax version (#211)

* Unpinned orbax and updated orbax logger API

* Upgraded from python 3.10 to 3.12 for base image

* Use python 3.12 in unit tests

Author: Rohan-Bierneni

.github/workflows/UnitTests.yml

@@ -35,7 +35,11 @@ jobs:
     name: "TPU test (${{ matrix.tpu-type }})"
     runs-on: ["self-hosted", "tpu", "${{ matrix.tpu-type }}"]
     steps:
-    - uses: actions/checkout@v3
+    - uses: actions/checkout@v4
+    - name: Set up Python 3.12
+      uses: actions/setup-python@v5
+      with:
+        python-version: '3.12'
     - name: Install dependencies
       run: |
         pip install -e .

maxdiffusion_dependencies.Dockerfile

@@ -1,12 +1,12 @@
-# Use Python 3.10-slim-bullseye as the base image
-FROM python:3.10-slim-bullseye
+# Use Python 3.12-slim-bullseye as the base image
+FROM python:3.12-slim-bullseye
 
 # Environment variable for no-cache-dir and pip root user warning
 ENV PIP_NO_CACHE_DIR=1
 ENV PIP_ROOT_USER_ACTION=ignore
 
 # Set environment variables for Google Cloud SDK and Python 3.10
-ENV PYTHON_VERSION=3.10
+ENV PYTHON_VERSION=3.12
 ENV CLOUD_SDK_VERSION=latest
 
 # Set DEBIAN_FRONTEND to noninteractive to avoid frontend errors

requirements.txt

@@ -24,7 +24,7 @@ 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
+orbax-checkpoint
 tokenizers==0.21.0
 huggingface_hub>=0.30.2
 transformers==4.48.1

requirements_with_jax_ai_image.txt

@@ -26,7 +26,7 @@ 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
+orbax-checkpoint
 tokenizers==0.21.0
 huggingface_hub>=0.30.2
 transformers==4.48.1

src/maxdiffusion/checkpointing/checkpointing_utils.py

@@ -28,7 +28,7 @@
 from flax.training import train_state
 import orbax
 import orbax.checkpoint as ocp
-from orbax.checkpoint.logging import abstract_logger
+from orbax.checkpoint.logging import AbstractLogger
 from orbax.checkpoint.checkpoint_manager import CheckpointManager, CheckpointManagerOptions
 
 STABLE_DIFFUSION_CHECKPOINT = "STABLE_DIFFUSION_CHECKPOINT"
@@ -43,7 +43,7 @@ def create_orbax_checkpoint_manager(
     checkpoint_type: str,
     dataset_type: str = "tf",
     use_async: bool = True,
-    orbax_logger: Optional[abstract_logger.AbstractLogger] = None,
+    orbax_logger: Optional[AbstractLogger] = None,
 ):
   """
   Returns specified Orbax (async or not) CheckpointManager or None if checkpointing is disabled.