resolved README conflict

Author: Serenagu525

.github/workflows/UnitTests.yml

@@ -31,9 +31,9 @@ jobs:
     strategy:
       fail-fast: false
       matrix:
-        tpu-type: ["v4-8"]
+        tpu-type: ["v5p-8"]
     name: "TPU test (${{ matrix.tpu-type }})"
-    runs-on: ["self-hosted", "tpu", "${{ matrix.tpu-type }}"]
+    runs-on: ["self-hosted","${{ matrix.tpu-type }}"]
     steps:
     - uses: actions/checkout@v4
     - name: Set up Python 3.12
@@ -54,7 +54,7 @@ jobs:
         ruff check .
     - name: PyTest
       run: | 
-        HF_HUB_CACHE=/mnt/disks/github-runner-disk/ HF_HOME=/mnt/disks/github-runner-disk/ python3 -m pytest -x --deselect=src/maxdiffusion/tests/ltx_transformer_step_test.py
+        HF_HUB_CACHE=/mnt/disks/github-runner-disk/ HF_HOME=/mnt/disks/github-runner-disk/ python3 -m pytest --deselect=src/maxdiffusion/tests/ltx_transformer_step_test.py
 #  add_pull_ready:
 #    if: github.ref != 'refs/heads/main'
 #    permissions:

.github/workflows/XLML.yml

@@ -2,7 +2,7 @@ name: Add Testgrid Link to PR
 
 on:
   pull_request:
-    types: [opened, synchronize]
+    types: [opened]
 
 jobs:
   add_testgrid_link:

end_to_end/tpu/test_sdxl_training_loss.sh

@@ -12,7 +12,7 @@ done
 TRAIN_CMD="python src/maxdiffusion/train_sdxl.py src/maxdiffusion/configs/base_xl.yml \
         pretrained_model_name_or_path=gs://maxdiffusion-github-runner-test-assets/checkpoints/models--stabilityai--stable-diffusion-xl-base-1.0 \
         revision=refs/pr/95 activations_dtype=bfloat16 weights_dtype=bfloat16 metrics_file=metrics.txt write_metrics=True \
-        dataset_name=gs://jfacevedo-maxdiffusion-v5p/pokemon-datasets/pokemon-gpt4-captions_xl resolution=1024 per_device_batch_size=1 \
+        dataset_name=gs://jfacevedo-maxdiffusion-v5p/pokemon-datasets/pokemon-gpt4-captions_sdxl resolution=1024 per_device_batch_size=1 \
         jax_cache_dir=gs://jfacevedo-maxdiffusion/cache_dir/ max_train_steps=$STEPS attention=flash run_name=sdxl-fsdp-v5p-64-ddp enable_profiler=True  \
         run_name=$RUN_NAME \
         output_dir=$OUTPUT_DIR "

requirements.txt

@@ -35,4 +35,5 @@ sentencepiece
 aqtp
 imageio==2.37.0
 imageio-ffmpeg==0.6.0
-hf_transfer>=0.1.9
+hf_transfer>=0.1.9
+qwix@git+https://github.com/google/qwix.git

requirements_with_jax_ai_image.txt

@@ -35,4 +35,5 @@ sentencepiece
 aqtp
 imageio==2.37.0
 imageio-ffmpeg==0.6.0
-hf_transfer>=0.1.9
+hf_transfer>=0.1.9
+qwix@git+https://github.com/google/qwix.git

src/maxdiffusion/checkpointing/wan_checkpointer.py

@@ -15,7 +15,6 @@
 """
 
 from abc import ABC
-from flax import nnx
 from maxdiffusion.checkpointing.checkpointing_utils import (create_orbax_checkpoint_manager)
 from ..pipelines.wan.wan_pipeline import WanPipeline
 from .. import max_logging, max_utils
@@ -42,7 +41,7 @@ def _create_optimizer(self, model, config, learning_rate):
         learning_rate, config.learning_rate_schedule_steps, config.warmup_steps_fraction, config.max_train_steps
     )
     tx = max_utils.create_optimizer(config, learning_rate_scheduler)
-    return nnx.Optimizer(model, tx), learning_rate_scheduler
+    return tx, learning_rate_scheduler
 
   def load_wan_configs_from_orbax(self, step):
     max_logging.log("Restoring stable diffusion configs")

src/maxdiffusion/configs/base14.yml

@@ -231,4 +231,5 @@ cache_dreambooth_dataset: False
 quantization: ''
 # Shard the range finding operation for quantization. By default this is set to number of slices.
 quantization_local_shard_count: -1
+use_qwix_quantization: False 
 compile_topology_num_slices: -1 # Number of target slices, set to a positive integer.

src/maxdiffusion/configs/base21.yml

@@ -232,3 +232,4 @@ quantization: ''
 # Shard the range finding operation for quantization. By default this is set to number of slices.
 quantization_local_shard_count: -1
 compile_topology_num_slices: -1 # Number of target slices, set to a positive integer.
+use_qwix_quantization: False 

src/maxdiffusion/configs/base_2_base.yml

@@ -246,4 +246,5 @@ cache_dreambooth_dataset: False
 quantization: ''
 # Shard the range finding operation for quantization. By default this is set to number of slices.
 quantization_local_shard_count: -1
+use_qwix_quantization: False 
 compile_topology_num_slices: -1 # Number of target slices, set to a positive integer.

src/maxdiffusion/configs/base_flux_dev.yml

@@ -276,5 +276,6 @@ controlnet_image: 'https://upload.wikimedia.org/wikipedia/commons/thumb/c/c1/Goo
 quantization: ''
 # Shard the range finding operation for quantization. By default this is set to number of slices.
 quantization_local_shard_count: -1
+use_qwix_quantization: False 
 compile_topology_num_slices: -1 # Number of target slices, set to a positive integer.