Merge branch 'main' into ltx2-dev

Author: Perseus14

.github/workflows/UploadDockerImages.yml

@@ -40,16 +40,3 @@ jobs:
     - name: build maxdiffusion jax nightly image
       run: |
         bash .github/workflows/build_and_upload_images.sh CLOUD_IMAGE_NAME=maxdiffusion_jax_nightly MODE=nightly PROJECT=tpu-prod-env-multipod LOCAL_IMAGE_NAME=maxdiffusion_jax_nightly
-
-  build-gpu-image:
-    runs-on: ["self-hosted", "e2", "cpu"]
-    steps:
-    - uses: actions/checkout@v3
-    - name: Cleanup old docker images
-      run: docker system prune --all --force
-    - name: build maxdiffusion jax stable stack gpu image
-      run: |
-        bash .github/workflows/build_and_upload_images.sh CLOUD_IMAGE_NAME=maxdiffusion_jax_stable_gpu MODE=stable PROJECT=tpu-prod-env-multipod LOCAL_IMAGE_NAME=maxdiffusion_jax_stable_gpu DEVICE=gpu
-    - name: build maxdiffusion jax nightly image
-      run: |
-         bash .github/workflows/build_and_upload_images.sh CLOUD_IMAGE_NAME=maxdiffusion_jax_nightly_gpu MODE=nightly PROJECT=tpu-prod-env-multipod LOCAL_IMAGE_NAME=maxdiffusion_jax_nightly DEVICE=gpu

README.md

@@ -52,6 +52,8 @@ MaxDiffusion supports
 * Wan2.1 text2vid (training and inference).
 * Wan2.2 text2vid (inference).
 
+**Note on GPU Support:** GPU support is not actively maintained, but contributions are welcome
+
 
 # Table of Contents
 
@@ -176,7 +178,7 @@ After installation completes, run the training script.
 
   ```bash
   BUCKET_NAME=my-bucket
-  gsutil -m cp -r $TFRECORDS_DATASET_DIR gs://$BUCKET_NAME/${TFRECORDS_DATASET_DIR##*/}
+  gcloud storage cp --recursive $TFRECORDS_DATASET_DIR gs://$BUCKET_NAME/${TFRECORDS_DATASET_DIR##*/}
   ```
 
   Now run the training command:
@@ -703,4 +705,3 @@ This script will automatically format your code with `pyink` and help you identi
 
 The full suite of -end-to end tests is in `tests` and `src/maxdiffusion/tests`. We run them with a nightly cadance.
 
-

requirements.txt

@@ -4,6 +4,7 @@ jaxlib>=0.4.30
 grain
 google-cloud-storage>=2.17.0
 absl-py
+chex
 datasets
 flax>=0.12.0
 optax>=0.2.3

requirements_with_jax_ai_image.txt

@@ -6,6 +6,7 @@ jaxlib>=0.4.30
 grain
 google-cloud-storage>=2.17.0
 absl-py
+chex
 datasets
 flax>=0.12.0
 optax>=0.2.3

src/maxdiffusion/pedagogical_examples/save_sd_checkpoint.py

@@ -17,7 +17,7 @@
 """Load and save a checkpoint. This is useful for uploading checkpoints to gcs
 and later loading them from gcs directly.
 After calling this script, use gsutil to upload the weights to a bucket:
-gsutil -m cp -r sd-model-finetuned gs://<your-bucket>/sd_checkpoint/
+gcloud storage cp --recursive sd-model-finetuned gs://<your-bucket>/sd_checkpoint/
 """
 
 from typing import Sequence

src/maxdiffusion/pedagogical_examples/save_sdxl_checkpoint.py

@@ -17,7 +17,7 @@
 """Load and save a checkpoint. This is useful for uploading checkpoints to gcs
 and later loading them from gcs directly.
 After calling this script, use gsutil to upload the weights to a bucket:
-gsutil -m cp -r sdxl-model-finetuned gs://<your-bucket>/sdxl_1.0_base/
+gcloud storage cp --recursive sdxl-model-finetuned gs://<your-bucket>/sdxl_1.0_base/
 """
 
 from typing import Sequence

src/maxdiffusion/trainers/wan_trainer.py

@@ -415,7 +415,7 @@ def training_loop(self, pipeline, optimizer, learning_rate_scheduler, train_data
         max_logging.log(f"Saving final checkpoint for step {step}")
         self.checkpointer.save_checkpoint(self.config.max_train_steps - 1, pipeline, state.params)
         self.checkpointer.checkpoint_manager.wait_until_finished()
-      # load new state for trained tranformer
+      # load new state for trained transformer
       pipeline.transformer = nnx.merge(state.graphdef, state.params, state.rest_of_state)
       return pipeline