resolving comments

Author: eltsai

.github/workflows/UnitTests.yml

@@ -58,7 +58,7 @@ jobs:
       run: | #--deselect=src/maxdiffusion/tests/input_pipeline_interface_test.py
         export LIBTPU_INIT_ARGS='--xla_tpu_scoped_vmem_limit_kib=65536'
         HF_HUB_CACHE=/mnt/disks/github-runner-disk/ HF_HOME=/mnt/disks/github-runner-disk/ TOKENIZERS_PARALLELISM=false python3 -m pytest --ignore=src/maxdiffusion/kernels/ --deselect=src/maxdiffusion/tests/ltx_transformer_step_test.py -x
-#  add_pull_ready:q
+#  add_pull_ready
 #    if: github.ref != 'refs/heads/main'
 #    permissions:
 #      checks: read

src/maxdiffusion/max_utils.py

@@ -81,7 +81,7 @@ def activate_profiler(config):
   if jax.process_index() == 0 and config.enable_profiler:
     trace_dir = config.tensorboard_dir
     if trace_dir.startswith("gs://"):
-      trace_dir = "/tmp/profiler_traces"
+      trace_dir = os.path.join("/tmp/profiler_traces", config.run_name)
     os.makedirs(trace_dir, exist_ok=True)
     max_logging.log(f"Starting profiler trace in: {trace_dir}")
     jax.profiler.start_trace(trace_dir)
@@ -93,7 +93,7 @@ def deactivate_profiler(config):
 
     trace_dir = config.tensorboard_dir
     if trace_dir.startswith("gs://"):
-      local_dir = "/tmp/profiler_traces"
+      local_dir = os.path.join("/tmp/profiler_traces", config.run_name)
       if os.path.exists(local_dir):
         max_logging.log(f"Uploading profiler traces from {local_dir} to {trace_dir}...")
         client = storage.Client()

src/maxdiffusion/tests/wan_vae_test.py

@@ -108,7 +108,7 @@ def __init__(self, dim: int, mode: str) -> None:
           torch.nn.Upsample(scale_factor=(2.0, 2.0), mode="nearest"), torch.nn.Conv2d(dim, dim // 2, 3, padding=1)
       )
     elif mode == "upsample3d":
-      raise Exception("downsample3d not supported")
+      raise Exception("upsample3d not supported")
 
     elif mode == "downsample2d":
       self.resample = torch.nn.Sequential(torch.nn.ZeroPad2d((0, 1, 0, 1)), torch.nn.Conv2d(dim, dim, 3, stride=(2, 2)))