diff --git a/.file_mapping.json b/.file_mapping.json index aa2555d0..03d07635 100644 --- a/.file_mapping.json +++ b/.file_mapping.json @@ -1,7 +1,7 @@ { - "_source_commit": "116a1414e2fb10bc7416ca860ad67f517ef9371f-dirty", - "_dest_commit": "ffa9c6b60a6b04b2fae337577bc6cbd8a93c39f5", - "_generated_at": "2026-09-11T05:52:11Z", + "_source_commit": "3a3e1e118760abbc484d59807b797157c394635e-dirty", + "_dest_commit": "2b6c9a7061ae78dc83e29a4910ec5f8c9fe4b6ce", + "_generated_at": "2026-09-14T05:52:22Z", "files": { "imaginaire/__init__.py": "cosmos_framework/__init__.py", "imaginaire/attention/__init__.py": "cosmos_framework/model/attention/__init__.py", diff --git a/cosmos_framework/configs/base/defaults/tokenizer.py b/cosmos_framework/configs/base/defaults/tokenizer.py index ed164cbc..c2efd288 100644 --- a/cosmos_framework/configs/base/defaults/tokenizer.py +++ b/cosmos_framework/configs/base/defaults/tokenizer.py @@ -1,7 +1,6 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: OpenMDW-1.1 -import torch from hydra.core.config_store import ConfigStore from cosmos_framework.utils.lazy_config import PLACEHOLDER, LazyDict @@ -177,10 +176,6 @@ def register_tokenizer() -> None: # Wan2pt1 and Wan2pt2 tokenizers cs.store(group="tokenizer", package="model.config.tokenizer", name="wan2pt1_tokenizer", node=Wan2pt1VAEConfig) cs.store(group="tokenizer", package="model.config.tokenizer", name="wan2pt2_tokenizer", node=Wan2pt2VAEConfig) - # LiDAR VAEs are deliberately absent from this group: a range clip is its own modality - # with its own projections into the sequence, so it is registered under - # ``model.config.lidar_tokenizer`` by ``register_lidar_tokenizer`` below. Installing one - # here would displace the camera tokenizer and route rangemaps through the vision heads. # UniAE tokenizer cs.store( group="tokenizer", @@ -199,34 +194,6 @@ def register_tokenizer() -> None: ) -def register_lidar_tokenizer() -> None: - """Register LiDAR tokenizers under ``model.config.lidar_tokenizer``. - - A joint camera + LiDAR recipe holds two vision tokenizers at once, so the LiDAR one needs - its own package: registering it under ``model.config.tokenizer`` would displace the camera - tokenizer. With this group in place, the dataloader's per-sensor token accounting can - interpolate ``${model.config.lidar_tokenizer.temporal_compression_factor}`` the same way - it already reads the camera factors off ``model.config.tokenizer``. - """ - cs = ConfigStore.instance() - cs.store( - group="lidar_tokenizer", - package="model.config.lidar_tokenizer", - name="lidar_tokenizer_v0", - node=LidarTokenizerV0Config, - ) - cs.store( - group="lidar_tokenizer", - package="model.config.lidar_tokenizer", - name="lidar_tokenizer_v1_r105_b1800_symmetric", - node=LidarTokenizerV1R105B1800SymmetricConfig, - ) - cs.store( - group="lidar_tokenizer", - package="model.config.lidar_tokenizer", - name="lidar_tokenizer_v1p2_r105_b1800", - node=LidarTokenizerV1P2R105B1800Config, - ) def register_sound_tokenizer() -> None: diff --git a/cosmos_framework/inference/common/distillation_export.py b/cosmos_framework/inference/common/distillation_export.py index afcd7802..43d73371 100644 --- a/cosmos_framework/inference/common/distillation_export.py +++ b/cosmos_framework/inference/common/distillation_export.py @@ -86,6 +86,57 @@ def build_student_checkpoint_metadata(*, use_ema_weights: bool) -> dict[str, str } +def _migrate_legacy_transfer_replay_config(config: dict[str, Any], *, base_config_field_names: set[str]) -> None: + """Preserve pre-replay-policy Transfer connectivity when projecting a student.""" + legacy_key = "transfer_control_attention_mode" + if legacy_key not in config: + return + + required_fields = {"teacher_forcing_replay_policy", "teacher_forcing_kv_implementation"} + if not required_fields <= base_config_field_names: + raise ValueError("Legacy Transfer attention requires a causal base config with teacher-forcing replay support.") + + legacy_modes = { + "global_control": ("global", False), + "causal_control": ("causal", False), + "current_only_control": ("current", False), + "causal_control_with_rgb_history": ("causal", True), + "current_only_control_with_rgb_history": ("current", True), + } + legacy_mode = config[legacy_key] + if not isinstance(legacy_mode, str) or legacy_mode not in legacy_modes: + raise ValueError(f"Unsupported legacy {legacy_key}: {legacy_mode!r}.") + control_visibility, controls_read_rgb = legacy_modes[legacy_mode] + expected_policy = { + "control_visibility": control_visibility, + "controls_read_strict_past_clean_rgb": controls_read_rgb, + "clean_pass_causality": "frame", + "multiview_attention_scope": "all_views", + "decomposed_temporal_window_seconds": None, + } + policy = config.get("teacher_forcing_replay_policy", {}) + if not isinstance(policy, dict): + raise TypeError("Expected teacher_forcing_replay_policy to be a dictionary during legacy Transfer migration.") + for key, expected in expected_policy.items(): + if key in policy and policy[key] != expected: + raise ValueError( + f"Legacy {legacy_key}={legacy_mode!r} conflicts with teacher_forcing_replay_policy.{key}={policy[key]!r}; " + f"expected {expected!r}." + ) + + # The legacy Transfer implementation used three-way single-view attention. + # An explicit different implementation must not silently replace that kernel. + implementation = config.get("teacher_forcing_kv_implementation", "singleview_threeway_kv") + if implementation != "singleview_threeway_kv": + raise ValueError( + f"Legacy {legacy_key} conflicts with teacher_forcing_kv_implementation={implementation!r}; " + "expected 'singleview_threeway_kv'." + ) + config["teacher_forcing_replay_policy"] = {**policy, **expected_policy} + config["teacher_forcing_kv_implementation"] = implementation + del config[legacy_key] + + def sanitize_student_model_config( model_dict: dict[str, Any], *, @@ -98,6 +149,10 @@ def sanitize_student_model_config( if not isinstance(config, dict): raise TypeError("Expected model config to be a dictionary.") + # Migrate before filtering out training-only fields: dropping the legacy + # selector first would silently restore global controls without RGB history. + _migrate_legacy_transfer_replay_config(config, base_config_field_names=base_config_field_names) + model_dict["_target_"] = base_model_target config["_type"] = base_config_type diff --git a/cosmos_framework/inference/common/distillation_export_test.py b/cosmos_framework/inference/common/distillation_export_test.py index 7d64ddf6..716cfb11 100644 --- a/cosmos_framework/inference/common/distillation_export_test.py +++ b/cosmos_framework/inference/common/distillation_export_test.py @@ -73,6 +73,144 @@ def test_build_student_checkpoint_metadata_omits_source_paths() -> None: } +def _sanitize_causal_student(model_dict: dict) -> None: + sanitize_student_model_config( + model_dict, + base_model_target="omni_mot_causal_model", + base_config_type="omni_mot_causal_model_config", + base_config_field_names={ + "video_temporal_causal", + "teacher_forcing_replay_policy", + "teacher_forcing_kv_implementation", + "teacher_forcing_frames_per_chunk", + "kv_cache_inference_size", + "attention_sink_size", + }, + ) + + +@pytest.mark.parametrize( + ("legacy_mode", "control_visibility", "controls_read_rgb"), + [ + ("global_control", "global", False), + ("causal_control", "causal", False), + ("current_only_control", "current", False), + ("causal_control_with_rgb_history", "causal", True), + ("current_only_control_with_rgb_history", "current", True), + ], +) +def test_sanitize_student_migrates_legacy_transfer_connectivity( + legacy_mode: str, control_visibility: str, controls_read_rgb: bool +) -> None: + model_dict = { + "config": { + "video_temporal_causal": True, + "transfer_control_attention_mode": legacy_mode, + "teacher_forcing_frames_per_chunk": 1, + "kv_cache_inference_size": 51, + "attention_sink_size": 1, + } + } + + _sanitize_causal_student(model_dict) + + assert model_dict["config"] == { + "_type": "omni_mot_causal_model_config", + "video_temporal_causal": True, + "teacher_forcing_frames_per_chunk": 1, + "kv_cache_inference_size": 51, + "attention_sink_size": 1, + "teacher_forcing_kv_implementation": "singleview_threeway_kv", + "teacher_forcing_replay_policy": { + "control_visibility": control_visibility, + "controls_read_strict_past_clean_rgb": controls_read_rgb, + "clean_pass_causality": "frame", + "multiview_attention_scope": "all_views", + "decomposed_temporal_window_seconds": None, + }, + } + migrated = copy.deepcopy(model_dict) + _sanitize_causal_student(model_dict) + assert model_dict == migrated + + +def test_sanitize_student_merges_compatible_legacy_and_current_replay_settings() -> None: + model_dict = { + "config": { + "transfer_control_attention_mode": "causal_control_with_rgb_history", + "teacher_forcing_kv_implementation": "singleview_threeway_kv", + "teacher_forcing_replay_policy": { + "_type": "teacher_forcing_replay_policy_config", + "control_visibility": "causal", + }, + } + } + + _sanitize_causal_student(model_dict) + + policy = model_dict["config"]["teacher_forcing_replay_policy"] + assert policy["_type"] == "teacher_forcing_replay_policy_config" + assert policy["control_visibility"] == "causal" + assert policy["controls_read_strict_past_clean_rgb"] is True + + +@pytest.mark.parametrize( + "conflicting_policy", + [ + {"control_visibility": "global"}, + {"controls_read_strict_past_clean_rgb": False}, + {"clean_pass_causality": "chunk"}, + {"multiview_attention_scope": "same_view"}, + {"decomposed_temporal_window_seconds": 0.1}, + ], +) +def test_sanitize_student_rejects_conflicting_legacy_and_current_replay_settings(conflicting_policy: dict) -> None: + model_dict = { + "config": { + "transfer_control_attention_mode": "causal_control_with_rgb_history", + "teacher_forcing_replay_policy": conflicting_policy, + } + } + original = copy.deepcopy(model_dict) + + with pytest.raises(ValueError, match="conflicts with teacher_forcing_replay_policy"): + _sanitize_causal_student(model_dict) + + assert model_dict == original + + +def test_sanitize_student_rejects_conflicting_legacy_kv_implementation() -> None: + model_dict = { + "config": { + "transfer_control_attention_mode": "causal_control_with_rgb_history", + "teacher_forcing_kv_implementation": "multiview_flex_kv", + } + } + + with pytest.raises(ValueError, match="conflicts with teacher_forcing_kv_implementation"): + _sanitize_causal_student(model_dict) + + +@pytest.mark.parametrize("legacy_mode", [None, "future_control", []]) +def test_sanitize_student_rejects_unknown_legacy_transfer_mode(legacy_mode: object) -> None: + model_dict = {"config": {"transfer_control_attention_mode": legacy_mode}} + + with pytest.raises(ValueError, match="Unsupported legacy transfer_control_attention_mode"): + _sanitize_causal_student(model_dict) + + +def test_sanitize_student_rejects_legacy_transfer_without_causal_base_support() -> None: + model_dict = {"config": {"transfer_control_attention_mode": "causal_control_with_rgb_history"}} + + with pytest.raises(ValueError, match="requires a causal base config"): + sanitize_student_model_config( + model_dict, + base_model_target="omni_mot_model", + base_config_type="omni_mot_model_config", + base_config_field_names={"video_temporal_causal"}, + ) + + def test_sanitize_student_public_model_config_removes_internal_loaders() -> None: model_dict = { "config": { diff --git a/cosmos_framework/model/generator/mot/flex_attention.py b/cosmos_framework/model/generator/mot/flex_attention.py index 4eecca7b..abb8dfdb 100644 --- a/cosmos_framework/model/generator/mot/flex_attention.py +++ b/cosmos_framework/model/generator/mot/flex_attention.py @@ -1423,12 +1423,12 @@ def flex_attention( ``(out, lse)`` where ``lse`` has shape ``[1, N_full, heads]``. Raises: + RuntimeError: if ``return_lse`` is requested with a torch version that does not + provide ``torch.nn.attention.flex_attention.AuxRequest``. ValueError: if either length is not a multiple of the corresponding block size of the mask, if k and v disagree on length, if ``block_mask`` was built for different lengths, or if it was built at a block size other than ``backend``'s. """ - from torch.nn.attention.flex_attention import AuxRequest - q_seq_len = full_q.shape[1] kv_seq_len = full_k.shape[1] num_q_heads = full_q.shape[2] @@ -1475,13 +1475,22 @@ def flex_attention( # return_aux rather than the deprecated return_lse: the latter records a # FutureWarning in a module-level set, which Dynamo rejects as an unsafe # side effect inside the activation-checkpointing HOP. + try: + from torch.nn.attention.flex_attention import AuxRequest as aux_request_cls + except ImportError: + aux_request_cls = None + if aux_request_cls is None: + raise RuntimeError( + "return_lse=True requires torch.nn.attention.flex_attention.AuxRequest, " + f"which is unavailable in torch {torch.__version__}." + ) attn_out, aux = _COMPILED_FLEX_ATTENTION( q, k, v, block_mask=block_mask, enable_gqa=num_q_heads != num_kv_heads, - return_aux=AuxRequest(lse=True), + return_aux=aux_request_cls(lse=True), kernel_options=backend.kernel_options, ) # attn_out: [1,num_q_heads,N_full,head_dim], aux.lse: [1,num_q_heads,N_full] # Convert to the heads-last layout ([1,S,H,D] / [1,S,H]) that from_mode_splits diff --git a/cosmos_framework/utils/generator/multiview.py b/cosmos_framework/utils/generator/multiview.py index e0f33c06..12f53038 100644 --- a/cosmos_framework/utils/generator/multiview.py +++ b/cosmos_framework/utils/generator/multiview.py @@ -1,7 +1,7 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: OpenMDW-1.1 -from collections.abc import Callable +from collections.abc import Callable, Iterator from pathlib import Path from typing import Any, Literal, cast @@ -74,13 +74,59 @@ def split_multiview_video_by_view( return [video_by_view[view_idx].contiguous() for view_idx in range(sample_n_views)] # list[[C,F,H,W]] +def iter_multiview_video_by_view( + video: torch.Tensor, + *, + sample_n_views: int, + num_video_frames_per_view: int, +) -> Iterator[torch.Tensor]: # video: [B,C,V*F,H,W] or [C,V*F,H,W], yields [C,F,H,W] + """Yield camera views without materializing a contiguous copy of the full video.""" + if sample_n_views <= 0 or num_video_frames_per_view <= 0: + raise ValueError( + "Expected positive sample_n_views and num_video_frames_per_view, " + f"got sample_n_views={sample_n_views}, " + f"num_video_frames_per_view={num_video_frames_per_view}." + ) + if video.dim() == 5: + if video.shape[0] != 1: + raise ValueError( + "Expected multiview tensor shape [B,C,V*F,H,W] with B=1 or [C,V*F,H,W], " + f"got shape={tuple(video.shape)}, sample_n_views={sample_n_views}, " + f"num_video_frames_per_view={num_video_frames_per_view}." + ) + video_cthw = video[0] # [C,V*F,H,W] + elif video.dim() == 4: + video_cthw = video # [C,V*F,H,W] + else: + raise ValueError( + "Expected multiview tensor shape [B,C,V*F,H,W] with B=1 or [C,V*F,H,W], " + f"got shape={tuple(video.shape)}, sample_n_views={sample_n_views}, " + f"num_video_frames_per_view={num_video_frames_per_view}." + ) + + expected_num_frames = sample_n_views * num_video_frames_per_view + if video_cthw.shape[1] != expected_num_frames: + raise ValueError( + "Expected multiview tensor shape [B,C,V*F,H,W] with B=1 or [C,V*F,H,W], " + f"got shape={tuple(video.shape)}, sample_n_views={sample_n_views}, " + f"num_video_frames_per_view={num_video_frames_per_view}." + ) + + for view_index in range(sample_n_views): + frame_start = view_index * num_video_frames_per_view + frame_end = frame_start + num_video_frames_per_view + yield video_cthw[:, frame_start:frame_end] # [C,F,H,W] + + def decode_multiview_latent_per_view( decode: Callable[[torch.Tensor], torch.Tensor], latent: torch.Tensor, sample_n_views: int, num_video_frames_per_view: int, + *, + assemble_on_cpu: bool = False, ) -> torch.Tensor: # latent: [B,C,V*T_latent,H,W] or [C,V*T_latent,H,W], returns same rank with T=V*F - """Decode camera-major latent clips independently and concatenate their pixels.""" + """Decode camera-major latent clips independently and assemble their pixels.""" if latent.ndim not in (4, 5): raise ValueError( f"Multiview latents must have shape [B,C,T,H,W] or [C,T,H,W], got shape {tuple(latent.shape)}." @@ -96,6 +142,7 @@ def decode_multiview_latent_per_view( latent_frames_per_view = num_latent_frames // sample_n_views decoded_views: list[torch.Tensor] = [] + decoded_output: torch.Tensor | None = None for view_idx in range(sample_n_views): view_latent = latent.narrow( # [B,C,T_latent,H,W] or [C,T_latent,H,W] temporal_dim, @@ -113,7 +160,27 @@ def decode_multiview_latent_per_view( "Decoded camera clip length must match num_video_frames_per_view: " f"got T={decoded_view.shape[temporal_dim]}, expected {num_video_frames_per_view}." ) - decoded_views.append(decoded_view) + if assemble_on_cpu: + if decoded_output is None: + output_shape = list(decoded_view.shape) + output_shape[temporal_dim] = sample_n_views * num_video_frames_per_view + decoded_output = torch.empty( + output_shape, + dtype=decoded_view.dtype, + device="cpu", + ) # [B,C,V*F,H_pixel,W_pixel] or [C,V*F,H_pixel,W_pixel] + output_view = decoded_output.narrow( # [B,C,F,H_pixel,W_pixel] or [C,F,H_pixel,W_pixel] + temporal_dim, + view_idx * num_video_frames_per_view, + num_video_frames_per_view, + ) + output_view.copy_(decoded_view) # [B,C,F,H_pixel,W_pixel] or [C,F,H_pixel,W_pixel] + del decoded_view + else: + decoded_views.append(decoded_view) + + if decoded_output is not None: + return decoded_output return torch.cat(decoded_views, dim=temporal_dim) # [B,C,V*F,H_pixel,W_pixel] or [C,V*F,H_pixel,W_pixel] diff --git a/cosmos_framework/utils/generator/multiview_test.py b/cosmos_framework/utils/generator/multiview_test.py index 36d5d2fe..0545e0e4 100644 --- a/cosmos_framework/utils/generator/multiview_test.py +++ b/cosmos_framework/utils/generator/multiview_test.py @@ -6,12 +6,105 @@ from cosmos_framework.utils.generator.multiview import ( build_camera_major_video, + decode_multiview_latent_per_view, generated_multiview_condition_frames, + iter_multiview_video_by_view, normalize_multiview_control_weights, pad_multiview_view_video, slice_multiview_view_frames, ) +# --------------------------------------------------------------------------- +# decode_multiview_latent_per_view +# --------------------------------------------------------------------------- + + +@pytest.mark.L0 +@pytest.mark.CPU +def test_decode_multiview_latent_per_view_assembles_directly_on_cpu( + monkeypatch: pytest.MonkeyPatch, +) -> None: + latent = torch.arange(6, dtype=torch.float32).reshape(1, 1, 6, 1, 1) # [B,C,V*T_latent,H,W] + + def decode(view_latent: torch.Tensor) -> torch.Tensor: + return view_latent + 1 # [B,C,F,H_pixel,W_pixel] + + def reject_cat(*args: object, **kwargs: object) -> torch.Tensor: + raise AssertionError("CPU assembly must not concatenate decoded views") + + monkeypatch.setattr(torch, "cat", reject_cat) + decoded = decode_multiview_latent_per_view( + decode, + latent, + sample_n_views=3, + num_video_frames_per_view=2, + assemble_on_cpu=True, + ) # [B,C,V*F,H_pixel,W_pixel] + + assert decoded.device.type == "cpu" + assert torch.equal(decoded, latent + 1) + + +# --------------------------------------------------------------------------- +# iter_multiview_video_by_view +# --------------------------------------------------------------------------- + + +@pytest.mark.L0 +@pytest.mark.CPU +@pytest.mark.parametrize("with_batch_dim", [False, True]) +@pytest.mark.parametrize("noncontiguous", [False, True]) +def test_iter_multiview_video_by_view_yields_shared_camera_views( + with_batch_dim: bool, + noncontiguous: bool, +) -> None: + video_storage = torch.arange(3 * 6 * 2 * 4, dtype=torch.float32).reshape(3, 6, 2, 4) # [C,V*F,H,2*W] + video_cthw = video_storage[:, :, :, ::2] if noncontiguous else video_storage[:, :, :, :2] # [C,V*F,H,W] + video = video_cthw.unsqueeze(0) if with_batch_dim else video_cthw # [B,C,V*F,H,W] or [C,V*F,H,W] + + views = list( + iter_multiview_video_by_view( + video, + sample_n_views=3, + num_video_frames_per_view=2, + ) + ) # list[[C,F,H,W]] + + assert len(views) == 3 + assert all(tuple(view.shape) == (3, 2, 2, 2) for view in views) + assert all(view.untyped_storage().data_ptr() == video.untyped_storage().data_ptr() for view in views) + round_trip = torch.cat(views, dim=1) # [C,V*F,H,W] + assert torch.equal(round_trip, video_cthw) + + +@pytest.mark.L0 +@pytest.mark.CPU +@pytest.mark.parametrize( + ("video_shape", "sample_n_views", "num_video_frames_per_view"), + [ + ((2, 3, 6, 2, 2), 3, 2), + ((3, 5, 2, 2), 3, 2), + ((3, 2, 2), 1, 2), + ((3, 6, 2, 2), 0, 2), + ], +) +def test_iter_multiview_video_by_view_rejects_invalid_shape( + video_shape: tuple[int, ...], + sample_n_views: int, + num_video_frames_per_view: int, +) -> None: + video = torch.zeros(video_shape) # invalid multiview shape + + with pytest.raises(ValueError, match="Expected"): + list( + iter_multiview_video_by_view( + video, + sample_n_views=sample_n_views, + num_video_frames_per_view=num_video_frames_per_view, + ) + ) + + # --------------------------------------------------------------------------- # pad_multiview_view_video # --------------------------------------------------------------------------- diff --git a/cosmos_framework/utils/wandb_util.py b/cosmos_framework/utils/wandb_util.py index 7ceee1ac..8fd9d597 100644 --- a/cosmos_framework/utils/wandb_util.py +++ b/cosmos_framework/utils/wandb_util.py @@ -4,7 +4,7 @@ from __future__ import annotations import os -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, Literal, cast import attrs import wandb @@ -43,13 +43,22 @@ def init_wandb(config: Config, model: ImaginaireModel) -> None: if isinstance(config.job, DictConfig): from cosmos_framework.utils.config import JobConfig - config_job = JobConfig(**config.job) + # `job.path` / `job.path_local` are JobConfig @property values, not attrs fields; + # some experiment configs stash resolved copies of them onto the DictConfig + # (e.g. `config.job.path = ...`), so filter to fields JobConfig actually accepts. + job_fields = {f.name for f in attrs.fields(JobConfig)} + config_job = JobConfig(**{str(k): v for k, v in config.job.items() if k in job_fields}) else: config_job = config.job config_checkpoint = config.checkpoint wandb_project = get_wandb_project(config_job.project) - # Try to fetch the W&B job ID for resuming training. - wandb_id = _read_wandb_id(config_job, config_checkpoint) + wandb_force_new_id = config_job.wandb_mode == "online_force_new_id" + wandb_mode = cast( + Literal["online", "offline", "disabled", "shared"], + "online" if wandb_force_new_id else config_job.wandb_mode, + ) + # Try to fetch the W&B job ID for resuming training, unless a fresh run was requested. + wandb_id = None if wandb_force_new_id else _read_wandb_id(config_job, config_checkpoint) if wandb_id is None: # Generate a new W&B job ID. wandb_id = generate_id() @@ -78,7 +87,7 @@ def init_wandb(config: Config, model: ImaginaireModel) -> None: config=config_resolved, dir=config_job.path_local, resume="allow", - mode=config_job.wandb_mode, + mode=wandb_mode, ) except Exception as e: # Detect common permission / upload errors from wandb and recover @@ -101,7 +110,7 @@ def init_wandb(config: Config, model: ImaginaireModel) -> None: name=config_job.name, config=config_resolved, dir=config_job.path_local, - mode=config_job.wandb_mode, + mode=wandb_mode, ) elif "returned error 401" in msg or "user is not logged in" in msg: log.warning("W&B authentication failed (401); falling back to offline mode. Error: %s", msg)