From 00892476fd748dfc30db81b949753e4249aacfb7 Mon Sep 17 00:00:00 2001 From: moritzhauschulz Date: Tue, 18 Aug 2026 19:51:50 +0200 Subject: [PATCH] offset at inference hot-fix --- src/weathergen/model/model.py | 28 ++++++++++++++++++++++++++- src/weathergen/utils/validation_io.py | 6 +++++- 2 files changed, 32 insertions(+), 2 deletions(-) diff --git a/src/weathergen/model/model.py b/src/weathergen/model/model.py index e90576b16..79e22cdfb 100644 --- a/src/weathergen/model/model.py +++ b/src/weathergen/model/model.py @@ -821,11 +821,37 @@ def forward( and self.cf.get("fe_diffusion_model_conditioning", None) == "forecast" ): # tokens[:,0] = t (most recent), tokens[:,1] = t-1, ..., tokens[:,-1] = t-(T-1) (oldest) - if self.cf.stage == "inference": + if self.cf.stage == "inference" and tokens.shape[1] == 1 and forecast_offset == 1: + # Single-input-step inference path: test_config.model_input.forecasting. + # num_steps_input=1 paired with test_config.forecast.offset=1, so the + # diffusion model's single forecast step lines up with a deterministic + # offset=1 model's +1-step lead time. There is no older step to sum over — + # the sole loaded step IS the conditioning. + print( + "Using the sole loaded step as conditioning token for inference " + "(offset=1, num_steps_input=1)." + ) + conditioning_tokens = tokens[:, 0] + elif self.cf.stage == "inference": + assert tokens.shape[1] > 1, ( + "Diffusion inference with fe_diffusion_model_conditioning='forecast' got " + f"only one loaded input step (tokens.shape[1]=1) but forecast.offset=" + f"{forecast_offset} != 1. The single-input-step inference path is only " + "defined for offset=1, num_steps_input=1 (used to align the diffusion " + "model's first forecast step with a deterministic offset=1 model's lead " + "time). Check test_config.forecast.offset and " + "test_config.model_input.forecasting.num_steps_input." + ) print("Using most recent steps as conditioning tokens for inference.") # conditioning_tokens = tokens[:, :-1].sum(axis=1) conditioning_tokens = tokens[:, 1:].sum(axis=1) else: + assert tokens.shape[1] > 1, ( + "Diffusion training with fe_diffusion_model_conditioning='forecast' " + "requires at least 2 loaded input steps (one denoising target + one " + f"conditioning step); got tokens.shape[1]={tokens.shape[1]}. Check " + "training_config.model_input.forecasting.num_steps_input." + ) # Conditioning: all older context steps [t-1, ..., t-(T-1)]; # denoising target: t (newest) conditioning_tokens = tokens[:, 1:].sum(axis=1) diff --git a/src/weathergen/utils/validation_io.py b/src/weathergen/utils/validation_io.py index e8285450c..d4fa96dc5 100644 --- a/src/weathergen/utils/validation_io.py +++ b/src/weathergen/utils/validation_io.py @@ -78,7 +78,11 @@ def write_output( # forecast indices, so synthesize a contiguous run of indices starting at the # original first index to cover every entry in model_output / target_aux_out. n_pred_steps = len(model_output.physical) - if cf.get("fe_diffusion_model", False) and n_pred_steps > len(timestep_idxs): + if ( + cf.get("fe_diffusion_model", False) + and chunk_forecast_offset == 0 + and n_pred_steps > len(timestep_idxs) + ): timestep_idxs = list(range(forecast_offset, forecast_offset + n_pred_steps)) targets_lens = []