Delete folder code/qwenvl/modalities/weather with huggingface_hub

code/qwenvl/modalities/weather/__init__.py

@@ -1,86 +0,0 @@
-"""Weather modality package.
-
-Polaris-style global forecast support, integrated as a standard modality
-under the ``ModalityRouter`` framework.
-
-The package contributes:
-
-* ``WeatherEncoder``    — Polaris ``CubeEmbedConv`` + Swin encoder + meteo merger.
-* ``IdentityWeatherProjector`` — pass-through; merger already projects to LLM hidden.
-* ``WeatherDecoder``    — Polaris meteo head + Charbonnier regression loss
-  (registers ``compute_loss_from_hidden`` for the router).
-
-Special tokens: ``<|weather_start|>``, ``<|weather_end|>``, ``<|weather_pad|>``.
-
-Adding the modality at runtime:
-
-    --weather_config_path qwenvl/modalities/weather/configs/config_weather_0p25_channel70.json
-    --tune_weather_encoder True --tune_weather_decoder True
-"""
-
-import logging
-
-from .encoder import WeatherEncoder
-from .projector import IdentityWeatherProjector
-from .decoder import WeatherDecoder
-from .processor import WeatherProcessor
-
-logger = logging.getLogger(__name__)
-
-MODALITY_CONFIG_KEY = "weather_config"
-
-TOKEN_DEFS = {
-    "weather": {
-        "start": "<|weather_start|>",
-        "end": "<|weather_end|>",
-        "pad": "<|weather_pad|>",
-    },
-}
-
-__all__ = [
-    "WeatherEncoder",
-    "IdentityWeatherProjector",
-    "WeatherDecoder",
-    "WeatherProcessor",
-    "register_modality",
-    "MODALITY_CONFIG_KEY",
-    "TOKEN_DEFS",
-]
-
-
-def register_modality(router, config, llm_hidden_size: int):
-    """Register the weather modality with the ``ModalityRouter``.
-
-    Args:
-        router:           ModalityRouter instance.
-        config:           Qwen3VLWeatherConfig (Polaris-style hyperparams).
-        llm_hidden_size:  Hidden size of the LLM backbone — overrides
-                          ``config.qwenvl_dim`` so the merger / mlp_qwen2swin
-                          bridges are sized to the actual LLM in use.
-    """
-    # Force qwenvl_dim to match the active LLM, ignoring whatever was saved
-    # in the JSON config.  (The merger is randomly initialised when
-    # init_weather_from is used, so 3584→4096 transitions are seamless.)
-    if config.qwenvl_dim != llm_hidden_size:
-        logger.info(
-            f"[weather] overriding qwenvl_dim {config.qwenvl_dim} → {llm_hidden_size} "
-            f"(active LLM hidden size)"
-        )
-        config.qwenvl_dim = llm_hidden_size
-
-    encoder = WeatherEncoder(config)
-    projector = IdentityWeatherProjector(qwenvl_dim=llm_hidden_size)
-    decoder = WeatherDecoder(config, encoder=encoder)
-
-    router.register_modality(
-        "weather",
-        encoder=encoder,
-        projector=projector,
-        decoder=decoder,
-        is_image_like=True,
-    )
-    logger.info(
-        f"[weather] registered: in_chans={config.in_chans}, "
-        f"image_size={config.image_size}, patch_size={config.patch_size}, "
-        f"hidden={config.hidden_size}, qwenvl_dim={config.qwenvl_dim}"
-    )

code/qwenvl/modalities/weather/configs/config_weather_0p25_channel70.json

@@ -1,45 +0,0 @@
-{
-    "polaris_vision_config": {
-        "in_chans": 70,
-        "in_frames": 1,
-        "const_chans": 6,
-
-        "upper_chans": 70,
-
-        "hidden_size": 2048,
-        "image_size": [721, 1440],
-        "patch_size": 6,
-        "window_size": 20,
-        "encoder_depth": 12,
-        "decoder_depth": 12,
-        "num_heads": 32,
-        "mlp_ratio": 4.0,
-        "attn_type": "flash",
-        "mask_type": "h",
-        "norm_type": "adarms",
-        "ffn_type": "geglu_ffn",
-        "n_kv_heads": null,
-
-        "embed_mode": "add",
-        "embed_types": ["step", "hour", "doy", "lead_hour"],
-        "embed_freq": 256,
-
-        "frame_interval": "6h",
-
-        "qwenvl_dim": 3584,
-
-        "meteo_out_channels": 70,
-        "meteo_output_size": [721, 1440],
-        "meteo_token_id": 151655,
-
-        "meteo_loss_type": "l1_channel",
-        "meteo_loss_weight": 1.0,
-        "meteo_loss_coef": 1.0,
-        "lead_hour_scaling": true,
-        "pl_chans": null,
-
-        "meteo_data_path": "<PATH_TO_ERA5_ZARR>",
-
-        "_attn_implementation": "eager"
-    }
-}

code/qwenvl/modalities/weather/data/__init__.py

@@ -1,3 +0,0 @@
-"""Weather data utilities (ERA5 dataset / collator / context loader)."""
-
-from .polaris_data_utils import load_meteorological_buffers  # noqa: F401

code/qwenvl/modalities/weather/data/era5_dataset.py

@@ -1,96 +0,0 @@
-"""Bio_qwen3vl-style ERA5 dataset.
-
-Produces samples ready for the bio_qwen3vl pipeline:
-
-* ``input_ids`` / ``attention_mask`` / ``position_ids`` / ``labels``:
-  the chat template with ``<|weather_start|><|weather_pad|>*N<|weather_end|>``
-  expanded to ``meteo_num_tokens`` pad tokens.
-
-* ``weather_input_ids`` / ``weather_attention_mask`` / ``weather_grid_thw``:
-  dummy placeholders sized to ``swin_H * swin_W`` so the
-  ``ModalityRouter.scatter_all`` finds them under the standard kwargs.
-
-* The actual meteorological data — ``meteo_values`` / ``targets`` / ``times`` /
-  ``lead_hours`` / ``polaris_task`` / optional ``channel_mask`` — is
-  carried in a ``"meteo_data"`` dict so the collator can stack across
-  the batch in fp32.
-
-The collator unpacks ``meteo_data`` into the canonical
-``weather_meteo_values`` / ``weather_targets`` / ``weather_times`` /
-``weather_lead_hours`` / ``weather_polaris_task`` kwargs that the encoder
-reads.
-"""
-
-from __future__ import annotations
-
-import json
-from dataclasses import dataclass
-from typing import Any, Dict, List, Optional, Sequence
-
-import numpy as np
-import pandas as pd
-import torch
-import transformers
-
-
-IGNORE_INDEX = -100
-DEFAULT_IMAGE_TOKEN = "<|image_pad|>"  # only used for compatibility; the
-                                       # weather pipeline uses <|weather_pad|>
-
-
-# ---------------------------------------------------------------------------
-# Chat-template builder for the meteo task
-# ---------------------------------------------------------------------------
-
-def _build_meteo_messages(
-    input_text: str,
-    meteo_pad_token: str,
-    weather_start_token: str,
-    weather_end_token: str,
-    meteo_num_tokens: int,
-) -> List[Dict[str, Any]]:
-    """Build a 2-turn system+user message list with the meteo placeholder.
-
-    The placeholder string is the same as what Polaris originally used,
-    swapping ``<|vision_start|>``→``<|weather_start|>`` and
-    ``<|image_pad|>``→``<|weather_pad|>`` so the bio_qwen3vl router knows
-    where to scatter the encoder output.
-    """
-    pads = meteo_pad_token * meteo_num_tokens
-    user_text = (
-        f"{weather_start_token}{pads}\n{weather_end_token}"
-        f"Describe this image."
-    )
-    return [
-        {"role": "system",
-         "content": [{"type": "text", "text": f"forecast\n{input_text}"}]},
-        {"role": "user",
-         "content": [{"type": "text", "text": user_text}]},
-    ]
-
-
-def preprocess_meteo_chat(
-    input_text: str,
-    processor,
-    meteo_pad_token: str,
-    weather_start_token: str,
-    weather_end_token: str,
-    meteo_num_tokens: int,
-    add_assistant_prompt: bool = True,
-) -> Dict[str, torch.Tensor]:
-    msgs = _build_meteo_messages(
-        input_text, meteo_pad_token, weather_start_token,
-        weather_end_token, meteo_num_tokens,
-    )
-    full = processor.apply_chat_template(
-        msgs,
-        tokenize=True,
-        return_dict=True,
-        return_tensors="pt",
-        add_generation_prompt=add_assistant_prompt,
-    )
-    input_ids = full["input_ids"]
-    if isinstance(input_ids, list):
-        input_ids = torch.tensor(input_ids).unsqueeze(0)
-    full["input_ids"] = input_ids
-    return full

code/qwenvl/modalities/weather/data/polaris_data_utils.py

@@ -1,191 +0,0 @@
-import os
-import numpy as np
-import pandas as pd
-import torch
-import torch.nn.functional as F
-import xarray as xr
-from einops import rearrange
-
-__all__ = [
-    "make_seq", "crop_dataarray", "filter_dataset_channels", "load_meteorological_buffers",
-    "TYPHOON_PAD_MULTIPLE", "pad_spatial_to_multiple", "crop_spatial_pad",
-    "build_channel_map", "pad_channels", "extract_valid_channels",
-]
-
-TYPHOON_PAD_MULTIPLE = 120
-
-
-def pad_spatial_to_multiple(x: torch.Tensor, multiple: int = TYPHOON_PAD_MULTIPLE):
-    H, W = x.shape[-2], x.shape[-1]
-    pad_h = (multiple - H % multiple) % multiple
-    pad_w = (multiple - W % multiple) % multiple
-    pt, pb = pad_h // 2, pad_h - pad_h // 2
-    pl, pr = pad_w // 2, pad_w - pad_w // 2
-    if pad_h > 0 or pad_w > 0:
-        x = F.pad(x, (pl, pr, pt, pb), mode="constant", value=0)
-    return x, (pt, pb, pl, pr)
-
-
-def crop_spatial_pad(x: torch.Tensor, pad_info):
-    pt, pb, pl, pr = pad_info
-    hs = slice(pt, x.shape[-2] - pb if pb > 0 else x.shape[-2])
-    ws = slice(pl, x.shape[-1] - pr if pr > 0 else x.shape[-1])
-    return x[..., hs, ws]
-
-
-def build_channel_map(src_names, ref_names):
-    ref_idx = {n: i for i, n in enumerate(ref_names)}
-    src_to_ref = np.array([ref_idx.get(str(n), -1) for n in src_names], dtype=np.int64)
-    mask = np.zeros(len(ref_names), dtype=bool)
-    for i in src_to_ref:
-        if 0 <= i < len(ref_names):
-            mask[i] = True
-    return src_to_ref, mask
-
-
-def pad_channels(x_src, src_to_ref, ref_c):
-    c = x_src.shape[0]
-    out = np.zeros((ref_c,) + x_src.shape[1:], dtype=np.float32)
-    for j in range(c):
-        i = int(src_to_ref[j])
-        if 0 <= i < ref_c:
-            out[i] = x_src[j]
-    return out
-
-
-def extract_valid_channels(x_ref, src_to_ref, src_c):
-    out = np.zeros((src_c,) + x_ref.shape[1:], dtype=np.float32)
-    for j in range(src_c):
-        i = int(src_to_ref[j])
-        if 0 <= i < x_ref.shape[0]:
-            out[j] = x_ref[i]
-    return out
-
-# 🚨 需要剔除的通道
-DEFAULT_REMOVE_CHANNELS = [
-    "q2m", "d2m", "sst", "ws100m", "u100m", "v100m",
-    "lcc", "mcc", "hcc", "tcc", "ssr", "ssrd", "fdir", "ttr", "tcw", "tp"
-]
-
-def make_seq(ds, total_frames, frame_interval, frame_step=1, ignore_times=[]):
-    i = 0
-    inds = []
-    times = ds.time.values
-    ignore_timestamps = {pd.to_datetime(t) for t in ignore_times}
-    
-    while i < len(times) - total_frames:
-        cur_sequence = []
-        for j in range(i, i + total_frames - 1):
-            if times[j+1] - times[j] == frame_interval:
-                current_time = pd.to_datetime(times[j])
-                if current_time in ignore_timestamps:
-                    continue
-                cur_sequence.append(j)
-
-        if len(cur_sequence) == total_frames - 1:
-            inds.append(i)
-
-        i += frame_step
-
-    return np.array(inds, dtype=np.int32) 
-
-def crop_dataarray(ds, image_size=None, latlon_range=None):
-    if "lat" in ds.dims:
-        ilats = np.arange(ds.lat.size)
-        if latlon_range is not None:
-            lat_min, lat_max, lon_min, lon_max = latlon_range
-            ilats = np.where((ds.lat >= lat_min) & (ds.lat <= lat_max))[0]
-        if image_size is not None:
-            ilats = ilats[:image_size[0]]
-        ds = ds.isel(lat=ilats)
-
-    if "lon" in ds.dims:
-        ilons = np.arange(ds.lon.size)
-        if latlon_range is not None:
-            lat_min, lat_max, lon_min, lon_max = latlon_range
-            ilons = np.where((ds.lon >= lon_min) & (ds.lon <= lon_max))[0]
-        if image_size is not None:
-            ilons = ilons[:image_size[1]]
-        ds = ds.isel(lon=ilons)
-    return ds
-
-def filter_dataset_channels(ds, remove_channels=None):
-    """
-    统一的通道过滤函数：Dataset 和 Buffer 读取都调用这个函数。
-    """
-    if remove_channels is None:
-        remove_channels = DEFAULT_REMOVE_CHANNELS
-        
-    if isinstance(ds, xr.Dataset):
-        ds = ds[list(ds.data_vars)[0]]
-        
-    all_channels = ds.channel.values.tolist()
-    keep_channels = [c for c in all_channels if c not in remove_channels]
-    keep_inds = [i for i, c in enumerate(all_channels) if c in keep_channels]
-    
-    out_ds = ds.sel(channel=keep_channels)
-    return out_ds, keep_channels, keep_inds
-
-def load_meteorological_buffers(
-    data_path: str, 
-    image_size=None,
-    latlon_range=None,
-    remove_channels=None,
-    buffer_types=None,
-    index_names=None
-):
-    """
-    统一的数据提取。
-    """
-    if buffer_types is None:
-        buffer_types = [
-            "mean", "std", "diff_mean", "diff_std", "climate",
-            "const", "weight", "channel_mask", "land_mask", "station_mask", "rps_mask", "dry_mask"
-        ]
-    if index_names is None:
-        index_names = dict(logid=[], uid=["u10m"], vid=["v10m"], accumid=[])
-
-    ds = xr.open_zarr(data_path)
-    if "level" in ds.dims:
-        ds = ds.rename({"level": "channel"})
-    ds = crop_dataarray(ds, image_size, latlon_range)
-
-    # 🚨 调用统一过滤函数
-    ds, keep_channels, keep_inds = filter_dataset_channels(ds, remove_channels)
-
-    indices = {}
-    for k, prefixes in index_names.items():
-        inds = [i for i, name in enumerate(keep_channels) if name in prefixes]
-        if len(inds) > 0:
-            indices[k] = inds
-    
-    coords = dict(lat=ds.lat.values.tolist(), lon=ds.lon.values.tolist())
-    
-    buffers = {}
-    for k in buffer_types:
-        file_name = os.path.join(data_path, f"{k}.nc")
-        if not os.path.exists(file_name):
-            continue
-
-        da = xr.open_dataarray(file_name)
-        da = crop_dataarray(da, image_size, latlon_range)
-
-        if k in ["mean", "std", "diff_mean", "diff_std"]:
-            if "channel" in da.dims:
-                da = da.sel(channel=keep_channels)
-            else:
-                da_vals = da.values[keep_inds]
-                da = xr.DataArray(da_vals, coords={"channel": keep_channels}, dims=("channel",))
-
-        values = da if k == "climate" else da.values
-
-        if k == "const":
-            values = rearrange(values, "c h w -> 1 c h w")
-        elif k == "weight" and values.ndim == 1:
-            values = rearrange(values, "h -> 1 h 1")
-        elif k in ["mean", "std", "diff_mean", "diff_std"]:
-            values = rearrange(values, "c -> c 1 1")
-
-        buffers[k] = values
-
-    return keep_channels, indices, coords, buffers

code/qwenvl/modalities/weather/decoder.py

@@ -1,256 +0,0 @@
-"""Weather decoder: Polaris meteo head + regression-style loss.
-
-Adapts ``PolarisMeteoHead`` to the bio_qwen3vl ``ModalityRouter``.  The
-decoder is invoked by ``compute_decoder_losses`` via the custom
-``compute_loss_from_hidden`` hook (see ``qwenvl/registry/modality_router.py``):
-that path bypasses the standard CE-on-decode-logits flow used by RNA /
-protein / mol since meteo prediction is a regression problem on a 5D
-tensor field.
-
-Inputs needed at loss time live in two places:
-
-* ``hidden_states`` (LLM output) — selects the meteo-token positions.
-* ``self.encoder._step_cache`` — set by ``WeatherEncoder.forward`` and
-  contains ``condition_embed`` / ``patch_embed_pre_swin`` /
-  ``meteo_values`` / ``targets`` / ``lead_hours``.
-
-The encoder reference is wired up in ``register_modality`` so the decoder
-does not have to be passed it through every forward.
-"""
-
-from __future__ import annotations
-
-import logging
-from typing import Any, Dict, Optional, Tuple
-
-import numpy as np
-import pandas as pd
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-from einops import rearrange
-from torch.utils.checkpoint import checkpoint
-
-from .internal.polaris_attention import precompute_freqs_cis
-from .internal.polaris_layers import AdaLN, DoubleDeconvHead, remove_small_scales
-from .internal.polaris_swin import SwinBlock
-from .internal.helpers import round_to_multiple
-
-logger = logging.getLogger(__name__)
-
-
-class WeatherDecoder(nn.Module):
-    """Polaris meteo head + regression loss path.
-
-    The original Polaris ``PolarisMeteoHead.forward`` is wrapped in
-    ``compute_loss_from_hidden`` so it can pull the matching encoder cache
-    and produce a scalar loss directly from LLM hidden states.
-    """
-
-    def __init__(self, config, encoder):
-        super().__init__()
-        self.config = config
-        # Hold the encoder reference in a 1-element list so PyTorch's
-        # automatic Module-attribute registration does NOT pull the
-        # encoder's parameters in as decoder children (which would
-        # duplicate them under both encoders.weather.* and decoders.weather.*
-        # in the state_dict).
-        self._encoder_ref = [encoder]
-
-        self.hidden_size = config.hidden_size
-        self.num_heads = config.num_heads
-        self.qwenvl_dim = config.qwenvl_dim
-        self.patch_size = config.patch_size
-        self.upper_chans = config.upper_chans
-        self.lower_chans = config.lower_chans
-        self.decoder_depth = config.decoder_depth
-
-        if isinstance(config.image_size, int):
-            in_h = in_w = config.image_size
-        else:
-            in_h, in_w = config.image_size
-        if self.patch_size == 1:
-            swin_H = in_h // 2 * 2
-            swin_W = in_w
-        else:
-            swin_H = in_h // self.patch_size
-            swin_W = in_w // self.patch_size
-        self._swin_HW = (swin_H, swin_W)
-
-        max_seq_len = swin_H * swin_W
-        freqs_cos, freqs_sin = precompute_freqs_cis(self.hidden_size // self.num_heads, max_seq_len)
-        self.register_buffer("freqs_cos", freqs_cos, persistent=False)
-        self.register_buffer("freqs_sin", freqs_sin, persistent=False)
-        self._rope_buffers_checked = False
-
-        # Cross-modal back projector: LLM hidden → swin hidden.
-        # Trained at ``mm_projector_lr`` (matched on substring "mlp_qwen2swin").
-        self.mlp_qwen2swin = nn.Sequential(
-            nn.Linear(self.qwenvl_dim, self.qwenvl_dim),
-            nn.GELU(),
-            nn.Linear(self.qwenvl_dim, self.hidden_size),
-        )
-
-        if config.embed_mode == "add":
-            self.embed_dim = config.hidden_size
-        elif config.embed_mode == "cat":
-            self.embed_dim = round_to_multiple(config.hidden_size // len(config.embed_types)) * len(config.embed_types)
-        else:
-            raise ValueError(f"Invalid embed_mode: {config.embed_mode}")
-
-        self.decoder_layers = nn.ModuleList()
-        for i in range(self.decoder_depth):
-            blk = SwinBlock(
-                dim=self.hidden_size, embed_dim=self.embed_dim, num_heads=self.num_heads,
-                input_size=(swin_H, swin_W), window_size=config.window_size,
-                shift_size=0 if i % 2 == 0 else config.window_size // 2,
-                mlp_ratio=config.mlp_ratio, attn_type=config.attn_type,
-                mask_type=config.mask_type, norm_type=config.norm_type,
-                ffn_type=config.ffn_type, n_kv_heads=config.n_kv_heads,
-                attn_implementation="eager",
-            )
-            self.decoder_layers.append(blk)
-
-        self.norm_layer = AdaLN(self.hidden_size, embed_dim=self.embed_dim)
-        self.pred_layer = DoubleDeconvHead(
-            in_chans=self.hidden_size, upper_chans=self.upper_chans,
-            lower_chans=self.lower_chans, patch_size=self.patch_size,
-        )
-
-        self.gradient_checkpointing = False
-
-    # Provide the same context-injection API the encoder has, so the
-    # training entry can call it on the decoder too without checking
-    # which side actually owns the buffers.
-    def inject_meteorological_context(self, *args, **kwargs):
-        # The encoder is the canonical owner of mean/std/nanmask/...
-        # We ignore here to avoid duplicate state.
-        return
-
-    def _ensure_finite_rope_buffers(self) -> None:
-        """Rebuild the RoPE tables if the non-persistent buffers came up broken.
-
-        ``freqs_cos`` / ``freqs_sin`` are ``persistent=False`` (not in the
-        checkpoint) and are meant to be filled at ``__init__``. Under a
-        meta-device ``from_pretrained`` load they can be ``to_empty``'d to
-        uninitialised memory. A valid table has ``cos[0]==1`` and ``sin[0]==0``
-        for every column — checking row 0 catches garbage that happens to be
-        finite (large floats) which an isfinite-only check would miss.
-        """
-        if self._rope_buffers_checked:
-            return
-        with torch.no_grad():
-            row0 = self.freqs_cos[0].detach().to(torch.float32)
-            row0_sin = self.freqs_sin[0].detach().to(torch.float32)
-        needs_recompute = (
-            not torch.isfinite(self.freqs_cos).all()
-            or not torch.isfinite(self.freqs_sin).all()
-            or not torch.allclose(row0, torch.ones_like(row0), atol=1e-5)
-            or not torch.allclose(row0_sin, torch.zeros_like(row0_sin), atol=1e-5)
-        )
-        if needs_recompute:
-            logger.warning(
-                "[weather] broken decoder RoPE buffers detected "
-                "(nan/inf/zeros); recomputing freqs_cos/freqs_sin"
-            )
-            max_seq_len = int(self._swin_HW[0] * self._swin_HW[1])
-            freqs_cos, freqs_sin = precompute_freqs_cis(
-                self.hidden_size // self.num_heads, max_seq_len,
-            )
-            with torch.no_grad():
-                self.freqs_cos.copy_(freqs_cos.to(self.freqs_cos.device, self.freqs_cos.dtype))
-                self.freqs_sin.copy_(freqs_sin.to(self.freqs_sin.device, self.freqs_sin.dtype))
-        self._rope_buffers_checked = True
-
-    # ------------------------------------------------------------------
-    # Loss computation (the only ModalityRouter-callable entry point)
-    # ------------------------------------------------------------------
-
-    def predict_from_hidden(
-        self,
-        hidden_states: torch.Tensor,
-        **kwargs,
-    ) -> Optional[torch.Tensor]:
-        """Decode LLM hidden states into a meteo prediction tensor.
-
-        Pure inference: no loss, no NaN-guard side effects.  Reusable by
-        the rollout engine, which needs to call this multiple times within
-        one optimizer step (warm-up steps + selected step).
-
-        Returns ``meteo_output`` of shape ``[B, C, H, W]`` (the ``pred_layer``
-        currently emits a single-frame prediction; rollout stitches frames
-        externally), or ``None`` if no weather pad tokens are present.
-        """
-        encoder = self._encoder_ref[0]
-        cache = getattr(encoder, "_step_cache", None) or {}
-        if not cache:
-            return None
-
-        self._ensure_finite_rope_buffers()
-
-        condition_embed = cache["condition_embed"]
-        patch_embed_post_swin = cache["patch_embed_post_swin"]
-        meteo_values = cache["meteo_values"]
-        lead_hours = cache["lead_hours"]
-        input_size = cache["input_size"]
-        T_in = cache["T_in"]
-
-        input_ids = kwargs.get("input_ids")
-        weather_pad_id = self._weather_pad_id_from_config(kwargs)
-        if input_ids is None or weather_pad_id is None:
-            raise RuntimeError(
-                "WeatherDecoder.predict_from_hidden: input_ids / weather pad id missing."
-            )
-
-        meteo_mask = (input_ids == weather_pad_id)
-        B, _, D = hidden_states.shape
-        L_per_sample = int(meteo_mask.sum(dim=1)[0].item())
-        if L_per_sample == 0:
-            return None
-        meteo_states = hidden_states[meteo_mask].view(B, L_per_sample, D)
-
-        h = self.mlp_qwen2swin(meteo_states)
-        if patch_embed_post_swin is not None:
-            h = h + patch_embed_post_swin
-
-        # Skip gradient checkpointing under torch.no_grad() (e.g. rollout
-        # warm-up) — checkpoint() warns and offers no benefit when the
-        # outputs don't require grad.
-        use_gc = (
-            self.gradient_checkpointing
-            and self.training
-            and torch.is_grad_enabled()
-        )
-        for blk in self.decoder_layers:
-            if use_gc:
-                h = checkpoint(
-                    blk, h, condition_embed, self.freqs_cos, self.freqs_sin,
-                    use_reentrant=False,
-                )
-            else:
-                h = blk(h, condition_embed, self.freqs_cos, self.freqs_sin)
-
-        h = self.norm_layer(h, condition_embed)
-        h = rearrange(h, "n (h w) c -> n c h w", h=input_size[0] // self.patch_size // 2 * 2)
-        meteo_output = self.pred_layer(
-            h,
-            residual=meteo_values[:, -1],
-            input_size=input_size,
-            lead_hour=condition_embed.new_tensor([0.0]) if lead_hours is None else lead_hours,
-            target_frames=T_in - 1,
-        )
-        return meteo_output
-
-    # ------------------------------------------------------------------
-    # Helpers
-    # ------------------------------------------------------------------
-
-    @staticmethod
-    def _weather_pad_id_from_config(kwargs) -> Optional[int]:
-        """Extract ``<|weather_pad|>`` id.
-
-        We rely on the model upstream to pass it under the key
-        ``__weather_pad_id__`` (set by the model wrapper before calling
-        ``compute_decoder_losses``).  Falls back to None when missing.
-        """
-        return kwargs.get("__weather_pad_id__")

code/qwenvl/modalities/weather/encoder.py

@@ -1,479 +0,0 @@
-"""Weather encoder: Polaris Swin-ViT + meteo merger (projector to LLM hidden).
-
-Contract:
-
-* ``forward(input_ids, attention_mask, **extra_kwargs) -> (latent, mask)``.
-  ``input_ids`` / ``attention_mask`` are placeholders sized to the number of
-  meteo tokens; the real meteorological tensors arrive through ``extra_kwargs``
-  (``meteo_values``, ``times``, ``lead_hours``, …) under the ``weather_`` prefix.
-
-* The encoder pre-computes ``condition_embed`` and the patch embed and stores
-  them in ``self._step_cache`` for the matching ``WeatherDecoder`` to reuse.
-
-* Mean / std / nanmask buffers (from ``inject_meteorological_context``) are
-  stored on the encoder so ``unnormalize`` / lat-weighting / channel masking
-  run without any external wrapper.
-"""
-
-from __future__ import annotations
-
-import logging
-from typing import Any, Dict, List, Optional, Tuple
-
-import numpy as np
-import pandas as pd
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-from einops import rearrange
-from torch.utils.checkpoint import checkpoint
-
-from .internal.polaris_attention import precompute_freqs_cis
-from .internal.polaris_layers import (
-    AdaLN,
-    CubeEmbedConv,
-    LayerNorm,
-    PatchEmbedConv,
-    TimestepEmbed,
-)
-from .internal.polaris_swin import SwinBlock
-from .internal.helpers import round_to_multiple
-
-try:
-    from transformers.models.qwen2.modeling_qwen2 import Qwen2RMSNorm as _MergerNorm
-except Exception:  # pragma: no cover
-    _MergerNorm = nn.LayerNorm
-
-logger = logging.getLogger(__name__)
-
-
-class PolarisMeteoPatchMerger(nn.Module):
-    """Cross-modal projector ('merger'): Swin hidden → LLM hidden.
-
-    Same arithmetic as the original Polaris ``PolarisMeteoPatchMerger``;
-    serves as the modality projector for the ``ModalityRouter`` (an
-    ``IdentityProjector`` is registered separately so the router can scatter
-    its output directly).
-    """
-
-    def __init__(self, dim: int, context_dim: int) -> None:
-        super().__init__()
-        self.hidden_size = context_dim
-        self.ln_q = _MergerNorm(context_dim, eps=1e-6)
-        self.mlp = nn.Sequential(
-            nn.Linear(self.hidden_size, self.hidden_size),
-            nn.GELU(),
-            nn.Linear(self.hidden_size, dim),
-        )
-
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
-        bsz, seqlen, _ = x.shape
-        x = self.ln_q(x)
-        x = x.view(-1, self.hidden_size)
-        x = self.mlp(x)
-        x = x.view(bsz, seqlen, -1)
-        return x
-
-
-class WeatherEncoder(nn.Module):
-    """Polaris-style Swin encoder + meteo merger.
-
-    The forward signature is dictated by ``ModalityRouter.encode_and_project``:
-    it takes ``input_ids`` / ``attention_mask`` (here used only to derive a
-    batch size; their real meaning is conveyed by the placeholder pad
-    tokens scattered into the LLM input embeddings) and a number of
-    keyword-only extras forwarded by the collator under the ``weather_``
-    prefix (already stripped by the router).
-    """
-
-    # Exposed so the router can route per-modality state_dict keys cleanly.
-    is_image_like = True
-
-    def __init__(self, config):
-        super().__init__()
-        self.config = config
-        self.hidden_size = config.hidden_size
-        self.image_size = config.image_size
-        self.patch_size = config.patch_size
-        self.window_size = config.window_size
-        self.encoder_depth = config.encoder_depth
-        self.num_heads = config.num_heads
-        self.const_chans = config.const_chans
-
-        self.patch_embed = CubeEmbedConv(
-            in_chans=config.in_chans,
-            out_chans=self.hidden_size,
-            in_frames=config.in_frames,
-            norm_func=None,
-            flatten=True,
-            patch_size=self.patch_size,
-        )
-
-        if self.const_chans > 0:
-            self.const_embed = PatchEmbedConv(
-                self.const_chans, self.hidden_size, config.patch_size,
-                norm_func=LayerNorm, flatten=True,
-            )
-
-        self.embed_mode = config.embed_mode
-        self.embed_types = list(config.embed_types)
-        self.embed_freq = config.embed_freq
-        self._init_time_embeddings()
-
-        if isinstance(config.image_size, int):
-            in_h = in_w = config.image_size
-        else:
-            in_h, in_w = config.image_size
-        if self.patch_size == 1:
-            swin_H = in_h // 2 * 2
-            swin_W = in_w
-        else:
-            swin_H = in_h // self.patch_size
-            swin_W = in_w // self.patch_size
-        self._swin_HW = (swin_H, swin_W)
-
-        self.encoder_layers = nn.ModuleList()
-        for i in range(self.encoder_depth):
-            blk = SwinBlock(
-                dim=self.hidden_size, embed_dim=self.embed_dim, num_heads=self.num_heads,
-                input_size=(swin_H, swin_W), window_size=self.window_size,
-                shift_size=0 if i % 2 == 0 else self.window_size // 2,
-                mlp_ratio=config.mlp_ratio, attn_type=config.attn_type,
-                mask_type=config.mask_type, norm_type=config.norm_type,
-                ffn_type=config.ffn_type, n_kv_heads=config.n_kv_heads,
-                attn_implementation="eager",
-            )
-            self.encoder_layers.append(blk)
-
-        # Cross-modal merger: Swin hidden → LLM hidden (qwenvl_dim).
-        # `ModalityRouter` registers an IdentityProjector for `weather` since
-        # the upscaling already happens here.
-        self.meteo_merger = PolarisMeteoPatchMerger(
-            dim=config.qwenvl_dim, context_dim=self.hidden_size,
-        )
-
-        max_seq_len = swin_H * swin_W
-        freqs_cos, freqs_sin = precompute_freqs_cis(self.hidden_size // self.num_heads, max_seq_len)
-        self.register_buffer("freqs_cos", freqs_cos, persistent=False)
-        self.register_buffer("freqs_sin", freqs_sin, persistent=False)
-
-        self.gradient_checkpointing = False
-
-        # ── Lazy data buffers (mean/std/nanmask/const/weight/...) ───
-        # Populated by ``inject_meteorological_context``.  Held in fp32
-        # CPU tensors and copied to device on demand to keep memory low.
-        self._polaris_data_loaded = False
-        self.channels: List[str] = []
-        self.indices: Dict[str, List[int]] = {}
-        self.coords: Dict[str, list] = {}
-        self._fp32_master: Dict[str, torch.Tensor] = {}
-
-        # Per-forward cache shared with the matching WeatherDecoder.
-        # Cleared at the start of every encoder forward to avoid leaking
-        # graph-disconnected tensors across steps.
-        self._step_cache: Dict[str, Any] = {}
-
-    # ------------------------------------------------------------------
-    # RoPE buffer self-heal
-    # ------------------------------------------------------------------
-
-    def _ensure_rope_buffers(self) -> None:
-        """Rebuild the RoPE tables if the non-persistent buffers came up broken.
-
-        ``freqs_cos`` / ``freqs_sin`` are ``persistent=False`` (not in the
-        checkpoint) and are meant to be filled at ``__init__``. Under a
-        meta-device ``from_pretrained`` load they can be ``to_empty``'d to
-        uninitialised memory. A valid table has ``cos[0]==1`` and ``sin[0]==0``
-        for every column — checking row 0 catches garbage that is finite (large
-        floats) which an isfinite-only check would miss.
-        """
-        with torch.no_grad():
-            row0 = self.freqs_cos[0].detach().to(torch.float32)
-            row0_sin = self.freqs_sin[0].detach().to(torch.float32)
-        needs_recompute = (
-            not torch.isfinite(self.freqs_cos).all()
-            or not torch.isfinite(self.freqs_sin).all()
-            or not torch.allclose(row0, torch.ones_like(row0), atol=1e-5)
-            or not torch.allclose(row0_sin, torch.zeros_like(row0_sin), atol=1e-5)
-        )
-        if not needs_recompute:
-            return
-        head_dim = self.hidden_size // self.num_heads
-        max_seq_len = self.freqs_cos.shape[0]
-        new_cos, new_sin = precompute_freqs_cis(head_dim, max_seq_len)
-        with torch.no_grad():
-            self.freqs_cos.copy_(new_cos.to(self.freqs_cos.device, self.freqs_cos.dtype))
-            self.freqs_sin.copy_(new_sin.to(self.freqs_sin.device, self.freqs_sin.dtype))
-
-    # ------------------------------------------------------------------
-    # Time-conditioning
-    # ------------------------------------------------------------------
-
-    def _init_time_embeddings(self):
-        if self.embed_mode == "add":
-            self.embed_dim = self.hidden_size
-        elif self.embed_mode == "cat":
-            self.embed_dim = round_to_multiple(self.hidden_size // len(self.embed_types)) * len(self.embed_types)
-        else:
-            raise ValueError(f"Invalid embed_mode: {self.embed_mode}")
-
-        # Optional dropout on time-embedding outputs (anti-overfit for
-        # ``hour`` / ``doy``).  Read from config; default 0.0 keeps the
-        # legacy zero-dropout behaviour identical to upstream Polaris.
-        time_embed_dropout = float(getattr(self.config, "time_embed_dropout", 0.0))
-        for k in self.embed_types:
-            embed_layer = TimestepEmbed(
-                self.embed_dim, frequency=self.embed_freq,
-                is_periodic=(k in ["hour", "doy"]), sinusoidal=True,
-                dropout=time_embed_dropout,
-            )
-            self.add_module(f"{k}_embed", embed_layer)
-
-    def _forward_embedding(self, conds: Dict[str, torch.Tensor]) -> torch.Tensor:
-        if self.embed_mode == "add":
-            embed = 0
-            for k in self.embed_types:
-                embedding = getattr(self, f"{k}_embed")
-                embed = embed + embedding(conds[k])
-            return embed
-        # cat
-        embeds = [getattr(self, f"{k}_embed")(conds[k]) for k in self.embed_types]
-        return torch.cat(embeds, dim=1)
-
-    # ------------------------------------------------------------------
-    # Meteo data buffer management
-    # ------------------------------------------------------------------
-
-    def inject_meteorological_context(self, channels, indices, coords, buffers):
-        """Inject ERA5 mean/std/nanmask buffers onto this encoder.
-
-        Stores fp32 tensors on CPU; ``_get_fp32_const`` copies to device on
-        demand.  Mirrors the original Polaris API so external loaders
-        (``polaris_data_utils.load_meteorological_buffers``) work unchanged.
-        """
-        self.channels = channels
-        self.indices = indices
-        self.coords = coords
-        for k, v in buffers.items():
-            if isinstance(v, np.ndarray):
-                t = torch.from_numpy(v).float()
-            elif torch.is_tensor(v):
-                t = v.detach().float()
-            else:
-                setattr(self, k, v)
-                continue
-            self._fp32_master[k] = t.cpu()
-        self._polaris_data_loaded = True
-
-    def _get_fp32_const(self, name: str, device: torch.device) -> torch.Tensor:
-        if not self._polaris_data_loaded or name not in self._fp32_master:
-            raise RuntimeError(
-                f"WeatherEncoder data not loaded! Call inject_meteorological_context() "
-                f"before using '{name}'."
-            )
-        return self._fp32_master[name].to(device=device, dtype=torch.float32, non_blocking=True)
-
-    # ------------------------------------------------------------------
-    # Input pre-processing (replicates Polaris reset_input / nan handling)
-    # ------------------------------------------------------------------
-
-    def _reset_output(self, x: torch.Tensor) -> torch.Tensor:
-        if "nanmask" in self._fp32_master:
-            nanmask = self._fp32_master["nanmask"].to(x.device)[: x.shape[-3]]
-            x = x * (~nanmask)
-        if "channel_mask" in self._fp32_master:
-            cmask = self._fp32_master["channel_mask"].to(x.device)[: x.shape[-3]]
-            x = x * cmask
-        return x
-
-    def _reset_input(self, x: torch.Tensor) -> torch.Tensor:
-        x = self._reset_output(x)
-        accumid = self.indices.get("accumid", [])
-        if len(accumid) > 0:
-            if torch.is_grad_enabled():
-                x = x.clone()
-            x[:, :, accumid] = 0
-        return x
-
-    # ------------------------------------------------------------------
-    # Conditioning factory (hour, doy, step, lead_hour, optional const)
-    # ------------------------------------------------------------------
-
-    def get_condition(
-        self,
-        t: int,
-        times,
-        lead_hour,
-        device: torch.device,
-        dtype: torch.dtype,
-    ) -> Dict[str, torch.Tensor]:
-        times = pd.DatetimeIndex(times)
-        B = len(times)
-
-        if isinstance(lead_hour, torch.Tensor):
-            lead_np = lead_hour.detach().to(dtype=torch.float32, device="cpu").numpy()
-        else:
-            lead_np = np.asarray(lead_hour, dtype=np.float32)
-
-        if lead_np.ndim == 0:
-            lead_np = np.full((B,), float(lead_np), dtype=np.float32)
-        elif lead_np.size == 1 and B > 1:
-            lead_np = np.full((B,), float(lead_np.reshape(-1)[0]), dtype=np.float32)
-        elif lead_np.size != B:
-            raise ValueError(f"lead_hour size mismatch: got {lead_np.size}, expected {B}")
-
-        cur_times = times + pd.to_timedelta(lead_np * t, unit="h")
-        used_times = cur_times
-
-        hour_np = used_times.hour.values.astype(np.float32) / 24.0
-        doy_np = (np.minimum(365, used_times.dayofyear.values).astype(np.float32) / 365.0)
-
-        hour = torch.tensor(hour_np, dtype=dtype, device=device)
-        doy = torch.tensor(doy_np, dtype=dtype, device=device)
-
-        Tmax = getattr(self.config, "max_rollout_steps", 200)
-        step_scalar = float(np.log1p(float(t)) / np.log1p(float(Tmax)))
-        step = torch.full((B,), step_scalar, dtype=dtype, device=device)
-        if self.training:
-            step = step + torch.empty((B,), device=device, dtype=dtype).uniform_(-0.005, 0.005)
-
-        lead_hour_t = torch.tensor(lead_np, dtype=dtype, device=device)
-
-        conds = dict(hour=hour, doy=doy, step=step, lead_hour=lead_hour_t)
-        if "const" in self._fp32_master:
-            conds["const"] = self._get_fp32_const("const", device).type(dtype)
-        return conds
-
-    # ------------------------------------------------------------------
-    # Encoder forward (ModalityRouter contract)
-    # ------------------------------------------------------------------
-
-    def forward(
-        self,
-        input_ids: torch.Tensor,
-        attention_mask: torch.Tensor,
-        meteo_values: Optional[torch.Tensor] = None,
-        targets: Optional[torch.Tensor] = None,
-        times: Optional[Any] = None,
-        lead_hours: Optional[torch.Tensor] = None,
-        polaris_task: Optional[Any] = None,
-        channel_mask: Optional[torch.Tensor] = None,
-        step_idx: int = 0,
-        **_unused,
-    ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
-        """Encode meteorological inputs.
-
-        Args:
-            input_ids / attention_mask: dummy placeholders sized
-                ``[B, swin_H * swin_W]``.  Required by the router signature
-                but not used computationally.
-            meteo_values: ``[B, T_in, C, H, W]`` float32 (smuggled as int32 by
-                the collator and reinterpreted upstream) — already on device.
-            times: pandas timestamps for the current batch (one per sample).
-            lead_hours: ``[B]`` int / float — forecast lead hours.
-            channel_mask: optional ``[B, C, 1, 1]`` channel mask.
-            step_idx: int rollout step index (0 for single-step training).
-
-        Returns:
-            latent: ``[B, swin_H * swin_W, qwenvl_dim]`` after meteo_merger.
-            latent_mask: None.
-        """
-        if meteo_values is None:
-            raise ValueError("WeatherEncoder.forward requires `meteo_values` in extras.")
-
-        # Rebuild RoPE freqs if the checkpoint load left them uninitialised
-        # (non-persistent buffers under a meta-device load come up as garbage).
-        self._ensure_rope_buffers()
-
-        # NaN→zero, fp32 normalisation
-        if meteo_values.dtype == torch.int32:
-            meteo_values = meteo_values.contiguous().view(torch.float32)
-        meteo_values = meteo_values.contiguous().to(torch.float32)
-        nanmask_local = torch.isnan(meteo_values[-1, -1])
-        if torch.any(nanmask_local):
-            self._fp32_master["nanmask"] = nanmask_local.cpu()
-        meteo_values = torch.nan_to_num(meteo_values)
-        meteo_values = self._reset_input(meteo_values)
-        if channel_mask is not None:
-            self._fp32_master["channel_mask"] = channel_mask[0].cpu()
-
-        device = meteo_values.device
-        # Compute conditions & feed encoder. Use the encoder's compute dtype
-        # (e.g. bf16) — meteo_values is fp32 here due to the upstream
-        # NaN/normalisation cast, but the embedding MLPs run in bf16 and
-        # would otherwise see a dtype mismatch in F.linear.
-        conds = self.get_condition(step_idx, times, lead_hours, device=device, dtype=self._embed_dtype())
-
-        meteo_compute = meteo_values.to(self._embed_dtype())
-        meteo_feature = self.patch_embed(meteo_compute, conds["lead_hour"])
-        patch_embed = meteo_feature
-
-        if self.const_chans > 0 and "const" in conds:
-            const = conds["const"].to(device=device, dtype=meteo_compute.dtype)
-            meteo_feature = meteo_feature + self.const_embed(const)
-
-        condition_embed = self._forward_embedding(conds)
-
-        # Skip gradient checkpointing under torch.no_grad() — e.g. rollout
-        # warm-up steps — otherwise PyTorch emits a per-block "None of the
-        # inputs have requires_grad=True" warning and `checkpoint` adds
-        # pointless bookkeeping for a path that won't backprop anyway.
-        use_gc = (
-            self.gradient_checkpointing
-            and self.training
-            and torch.is_grad_enabled()
-        )
-        for blk in self.encoder_layers:
-            if use_gc:
-                meteo_feature = checkpoint(
-                    blk, meteo_feature, condition_embed, self.freqs_cos, self.freqs_sin,
-                    use_reentrant=False,
-                )
-            else:
-                meteo_feature = blk(meteo_feature, condition_embed, self.freqs_cos, self.freqs_sin)
-
-        # `meteo_feature` is the post-swin encoder output (Polaris's
-        # `meteo_embeds`).  The matching head uses it as the skip connection
-        # *after* mlp_qwen2swin, so cache before applying the merger.
-        meteo_embeds_post_swin = meteo_feature
-
-        merged = self.meteo_merger(meteo_feature)  # → [B, L, qwenvl_dim]
-
-        # Cache for the decoder.  We deliberately keep references to all
-        # the intermediate tensors needed by the head so the decoder can
-        # reuse them without recomputing the encoder side.
-        # ``patch_embed_post_swin`` matches the Polaris ``meteo_embeds``
-        # variable used in ``polaris_head(meteo_feature_patch_embed=...)``
-        # under the use_language=True flow — i.e. the encoder swin output.
-        self._step_cache = {
-            "condition_embed": condition_embed,
-            "patch_embed_post_swin": meteo_embeds_post_swin,
-            "meteo_values": meteo_values,
-            "targets": targets,
-            "times": times,
-            "lead_hours": lead_hours,
-            "input_size": meteo_values.shape[-2:],
-            "T_in": meteo_values.shape[1],
-        }
-
-        # Mask is None — every meteo token is valid; the router will infer
-        # the count from the pad-token mask.
-        return merged, None
-
-    def _embed_dtype(self) -> torch.dtype:
-        # Canonical encoder compute dtype (handles bf16 / fp32 uniformly).
-        for p in self.parameters():
-            return p.dtype
-        return torch.float32
-
-    @torch.no_grad()
-    def unnormalize(self, x: torch.Tensor, fill_type: str = "zero") -> torch.Tensor:
-        """Undo channel-wise mean/std normalisation. fp32 throughout."""
-        mean = self._get_fp32_const("mean", x.device)
-        std = self._get_fp32_const("std", x.device)
-        x = x.to(torch.float32) * std + mean
-        logid = self.indices.get("logid", [])
-        if len(logid) > 0:
-            v = x[:, :, logid].clamp(min=0, max=7)
-            x[:, :, logid] = torch.expm1(v)
-        return self._reset_output(x)

code/qwenvl/modalities/weather/internal/__init__.py

@@ -1,11 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# Portions adapted from the Polaris weather-forecasting codebase (released here
-# under Apache-2.0 with the authors' permission) and from the Swin Transformer
-# (Microsoft, MIT). See THIRD_PARTY_LICENSES.md at the repo root.
-"""Internal Polaris support modules (attention, layers, swin, helpers).
-
-These are adapted from the original Polaris codebase and keep the Swin
-Transformer window-attention / patch-merging design. They are kept as a private
-submodule so the public weather encoder/decoder/projector can pin against a
-stable internal API.
-"""

code/qwenvl/modalities/weather/internal/helpers.py

@@ -1,45 +0,0 @@
-import collections.abc
-from itertools import repeat
-from inspect import isfunction
-from torch.nn import functional as F
-
-__all__ = ["to_2tuple", "round_to_multiple", "exists", "default", "append_dims"]
-
-def _ntuple(n):
-    def parse(x):
-        if isinstance(x, collections.abc.Iterable):
-            return x
-        return tuple(repeat(x, n))
-    return parse
-
-to_1tuple = _ntuple(1)
-to_2tuple = _ntuple(2)
-to_3tuple = _ntuple(3)
-to_4tuple = _ntuple(4)
-to_ntuple = _ntuple
-
-
-def round_to_multiple(x: int, multiple_of: int = 32) -> int:
-    """Round up x to the nearest multiple of multiple_of."""
-    return int(multiple_of * ((x + multiple_of - 1) // multiple_of))
-
-
-def exists(val):
-    return val is not None
-
-
-def default(val, d):
-    if exists(val):
-        return val
-    return d() if isfunction(d) else d
-
-def append_dims(x, target_dims):
-    """Appends dimensions to the end of a tensor until it has target_dims dimensions."""
-    dims_to_append = target_dims - x.ndim
-    if dims_to_append < 0:
-        raise ValueError(
-            f"input has {x.ndim} dims but target_dims is {target_dims}, which is less"
-        )
-    return x[(...,) + (None,) * dims_to_append]
-
-

code/qwenvl/modalities/weather/internal/polaris_attention.py

@@ -1,198 +0,0 @@
-
-import torch
-from typing import Optional
-from torch import nn
-import torch.nn.functional as F
-from einops import rearrange
-
-
-__all__ = ["FlashAttention", "precompute_freqs_cis"]
-
-
-class RMSNorm(nn.Module):
-    def __init__(self, dim: int, eps: float = 1e-6, compile: bool = False):
-        super().__init__()
-        self.eps = eps
-        self.weight = nn.Parameter(torch.ones(dim))
-        self.rmsnorm_fn = (
-            torch.compile(self.compute_rmsnorm, fullgraph=True)
-            if compile
-            else self.compute_rmsnorm
-        )
-
-    @staticmethod
-    def compute_rmsnorm(x: torch.Tensor, weight: torch.Tensor, eps: float):
-        def _norm(x, eps):
-            return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + eps)
-
-        output = _norm(x.float(), eps).type_as(x)
-        return output * weight
-
-    def forward(self, x: torch.Tensor):
-        return self.rmsnorm_fn(x, self.weight, self.eps)
-
-    def reset_parameters(self):
-        torch.nn.init.ones_(self.weight)  # type: ignore
-
-
-def precompute_freqs_cis(dim: int, end: int, theta: float = 10000.0):
-    freqs = 1.0 / (theta ** (torch.arange(0, dim, 2)[: (dim // 2)].float() / dim))
-    t = torch.arange(end, device=freqs.device)  # type: ignore
-    freqs = torch.outer(t, freqs).float()  # type: ignore
-    freqs_cos = torch.cos(freqs)  # real part
-    freqs_sin = torch.sin(freqs)  # imaginary part
-    return freqs_cos, freqs_sin
-
-
-def reshape_for_broadcast(freqs_cis: torch.Tensor, x: torch.Tensor):
-    ndim = x.ndim
-    assert 0 <= 1 < ndim
-    assert freqs_cis.shape == (x.shape[1], x.shape[-1])
-    shape = [d if i == 1 or i == ndim - 1 else 1 for i, d in enumerate(x.shape)]
-    return freqs_cis.view(shape)
-
-
-def apply_rotary_emb(xq: torch.Tensor, xk: torch.Tensor, freqs_cos: torch.Tensor, freqs_sin: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
-    xq, xk = xq.contiguous(), xk.contiguous()
-    freqs_cos, freqs_sin = freqs_cos.contiguous(), freqs_sin.contiguous()
-
-    xq_r, xq_i = xq.float().reshape(*xq.shape[:-1], -1, 2).unbind(-1)
-    xk_r, xk_i = xk.float().reshape(*xk.shape[:-1], -1, 2).unbind(-1)
-
-    seq_len = xq_r.shape[1]
-    assert seq_len <= freqs_cos.shape[0], (
-        f"seq_len={seq_len} exceeds precomputed freqs buffer size={freqs_cos.shape[0]}"
-    )
-    freqs_cos = freqs_cos[:seq_len]
-    freqs_sin = freqs_sin[:seq_len]
-
-    freqs_cos = rearrange(freqs_cos, "n c -> 1 n 1 c")
-    freqs_sin = rearrange(freqs_sin, "n c -> 1 n 1 c")  
-
-    # Apply rotation using real numbers
-    xq_out_r = xq_r * freqs_cos - xq_i * freqs_sin
-    xq_out_i = xq_r * freqs_sin + xq_i * freqs_cos
-    xk_out_r = xk_r * freqs_cos - xk_i * freqs_sin
-    xk_out_i = xk_r * freqs_sin + xk_i * freqs_cos
-
-    # Combine real and imaginary parts
-    xq_out = torch.stack([xq_out_r, xq_out_i], dim=-1).flatten(-2)
-    xk_out = torch.stack([xk_out_r, xk_out_i], dim=-1).flatten(-2)
-
-    return xq_out.to(xq.dtype), xk_out.to(xk.dtype)
-
-
-def repeat_kv(x: torch.Tensor, n_rep: int) -> torch.Tensor:
-    """torch.repeat_interleave(x, dim=2, repeats=n_rep)"""
-    bs, slen, n_kv_heads, head_dim = x.shape
-    if n_rep == 1:
-        return x
-    return (
-        x[:, :, :, None, :]
-        .expand(bs, slen, n_kv_heads, n_rep, head_dim)
-        .reshape(bs, slen, n_kv_heads * n_rep, head_dim)
-    )
-
-
-
-class FlashAttention(nn.Module):
-    def __init__(
-        self, 
-        dim: int,
-        n_heads: int = 32,
-        n_kv_heads: Optional[int] = None,
-        max_seq_len: Optional[int] = None,
-        dropout: float = 0.0,
-        is_causal: bool = False,
-        attn_implementation = None, 
-    ):
-        super().__init__()
-
-        self.n_heads = n_heads
-        self.n_kv_heads = n_heads if n_kv_heads is None else n_kv_heads
-        self.n_rep = self.n_heads // self.n_kv_heads
-        self.head_dim = dim // n_heads
-
-        self.wq = nn.Linear(dim, n_heads * self.head_dim, bias=False)
-        self.wk = nn.Linear(dim, self.n_kv_heads * self.head_dim, bias=False)
-        self.wv = nn.Linear(dim, self.n_kv_heads * self.head_dim, bias=False)
-        self.wo = nn.Linear(n_heads * self.head_dim, dim, bias=False)
-
-        self.attn_dropout = nn.Dropout(dropout)
-        self.resid_dropout = nn.Dropout(dropout)
-        self.dropout = dropout
-        self.max_seq_len = max_seq_len
-        self.is_causal = is_causal
-
-    def init_weights(self, init_std: float):
-        for linear in (self.wq, self.wk, self.wv):
-            nn.init.trunc_normal_(linear.weight, mean=0.0, std=0.02)
-        nn.init.trunc_normal_(self.wo.weight, mean=0.0, std=init_std)
-
-
-    def forward(
-        self,
-        x: torch.Tensor,
-        freqs_cos: torch.Tensor=None,
-        freqs_sin: torch.Tensor=None,
-        mask: torch.Tensor = None,
-    ):
-        bsz, seq_len, _ = x.shape
-        # QKV
-        xq, xk, xv = self.wq(x), self.wk(x), self.wv(x)
-        
-        if self.max_seq_len is None:
-            xq = xq.view(bsz, seq_len, self.n_heads, self.head_dim)
-            xk = xk.view(bsz, seq_len, self.n_kv_heads, self.head_dim)   
-            xv = xv.view(bsz, seq_len, self.n_kv_heads, self.head_dim)
-        else:
-            nseq = self.max_seq_len // seq_len
-            xq = xq.view(-1, nseq*seq_len, self.n_heads, self.head_dim)
-            xk = xk.view(-1, nseq*seq_len, self.n_kv_heads, self.head_dim)
-            xv = xv.view(bsz, seq_len, self.n_kv_heads, self.head_dim)
-
-        # RoPE relative positional embeddings
-        if not (freqs_cos is None or freqs_sin is None):
-            xq, xk = apply_rotary_emb(xq, xk, freqs_cos, freqs_sin)
-
-        if self.max_seq_len is not None:
-            xq = xq.view(bsz, seq_len, self.n_heads, self.head_dim)
-            xk = xk.view(bsz, seq_len, self.n_kv_heads, self.head_dim)
-        
-        # # grouped multiquery attention: expand out keys and values
-        # xk = repeat_kv(xk, self.n_rep)  # (bs, seq_len, n_heads, head_dim)
-        # xv = repeat_kv(xv, self.n_rep)  # (bs, seq_len, n_heads, head_dim)
-
-        # Repeat KV heads if necessary
-        if self.n_heads != self.n_kv_heads:
-            xk = xk.repeat_interleave(self.n_rep, dim=2)
-            xv = xv.repeat_interleave(self.n_rep, dim=2)
-
-        # make heads into a batch dimension
-        xq = xq.transpose(1, 2)  # (bs, n_heads, seq_len, head_dim)
-        xk = xk.transpose(1, 2)
-        xv = xv.transpose(1, 2)
-
-        if mask is not None:
-            mask = mask.to(xq).unsqueeze(1).contiguous()
-            nW = mask.shape[0]
-            if nW < bsz:
-                mask = mask.repeat([bsz // nW, 1, 1, 1])        
-
-        output = F.scaled_dot_product_attention(
-            xq, xk, xv, 
-            attn_mask=mask, 
-            dropout_p=self.dropout if self.training else 0.0, 
-            is_causal=self.is_causal
-        )
-
-        # restore time as batch dimension and concat heads
-        output = output.transpose(1, 2).contiguous().view(bsz, seq_len, -1)
-
-        # final projection into the residual stream
-        output = self.wo(output)
-        output = self.resid_dropout(output)
-        return output
-
-
-

code/qwenvl/modalities/weather/internal/polaris_layers.py

@@ -1,448 +0,0 @@
-# polaris_layers.py
-from __future__ import annotations
-
-import math
-from typing import Optional, Tuple, Union, List
-
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-from einops import rearrange
-
-# ---------------------------------------------------------
-# 基础归一化与嵌入模块 (保留兼容性)
-# ---------------------------------------------------------
-class LayerNorm(nn.Module):
-    def __init__(self, dim, eps=1e-6, elementwise_affine=True):
-        super().__init__()
-        self.eps = eps
-        self.elementwise_affine = elementwise_affine
-        if self.elementwise_affine:
-            self.weight = nn.Parameter(torch.ones(dim))
-
-    def _norm(self, x, dim=-1):
-        u = x.mean(dim, keepdim=True)
-        s = x.var(dim, keepdim=True)   
-        return (x - u) / torch.sqrt(s + self.eps) 
-    
-    def forward(self, x):
-        output = self._norm(x.float()).to(x)
-        if self.elementwise_affine:
-            output = output * self.weight        
-        return output
-
-class RMS_norm(nn.Module):
-    def __init__(self, dim, eps=1e-6, bias=False):
-        super().__init__()
-        self.eps = eps
-        self.weight = nn.Parameter(torch.ones(dim))
-        self.bias = nn.Parameter(torch.zeros(dim)) if bias else None
-
-    def forward(self, x):
-        # normalize over last dim
-        rms = x.pow(2).mean(dim=-1, keepdim=True).add(self.eps).sqrt()
-        y = x / rms * self.weight
-        if self.bias is not None:
-            y = y + self.bias
-        return y
-
-class RMSNorm(nn.Module):
-    def __init__(self, dim: int, eps: float = 1e-6):
-        super().__init__()
-        self.eps = eps
-        self.weight = nn.Parameter(torch.ones(dim))
-
-    def _norm(self, x, dim=-1):
-        return x * torch.rsqrt(x.pow(2).mean(dim, keepdim=True) + self.eps)
-    
-    def forward(self, x: torch.Tensor):
-        output = self._norm(x.float()).type_as(x)        
-        return output * self.weight
-        
-class PatchEmbedConv(nn.Module):
-    def __init__(
-        self,
-        in_chans,
-        out_chans,
-        patch_size=4,
-        norm_func=None, 
-        flatten=False,
-    ):
-        super().__init__()
-        self.patch_size = (patch_size, patch_size)
-        self.flatten = flatten    
-
-        # ⭐ 关键：patch_size=1 时保持老行为；否则按 patch_size 下采样
-        if self.patch_size == (1, 1):
-            # 保留你原来的设计：kernel_size=(2,1), stride=1
-            kernel_size = (2, 1)
-            stride = (1, 1)
-            padding = (0, 0)
-        else:
-            # 正常的 patch 下采样：kernel_size=stride=patch_size
-            kernel_size = self.patch_size
-            stride = self.patch_size
-            padding = (0, 0)
-
-        self.proj = nn.Conv2d(
-            in_chans,
-            out_chans,
-            kernel_size=kernel_size,
-            stride=stride,
-            padding=padding,
-        )
-        self.norm = norm_func(out_chans) if norm_func else nn.Identity()
-
-    def forward(self, x):
-        embed = self.proj(x)
-        if self.flatten:
-            embed = rearrange(embed, 'n c h w -> n (h w) c')
-        return self.norm(embed)
-               
-class WeatherLayerNorm(nn.Module):
-    """
-    专为气象 (B, C, H, W) 格式设计的高效 LayerNorm (Channel-first)。
-    适用于深层全卷积潜空间演化。
-    """
-    def __init__(self, normalized_shape, eps=1e-6):
-        super().__init__()
-        self.weight = nn.Parameter(torch.ones(normalized_shape))
-        self.bias = nn.Parameter(torch.zeros(normalized_shape))
-        self.eps = eps
-        
-    def forward(self, x):
-        # x: (B, C, H, W)
-        u = x.mean(1, keepdim=True)
-        s = (x - u).pow(2).mean(1, keepdim=True)
-        x = (x - u) / torch.sqrt(s + self.eps)
-        return self.weight[:, None, None] * x + self.bias[:, None, None]
-
-class AdaLN(nn.Module):
-    def __init__(self, dim, embed_dim=None):
-        super().__init__()
-        self.norm = LayerNorm(dim, eps=1e-6, elementwise_affine=False)
-        in_dim = embed_dim if embed_dim else dim
-        self.scale_shift = nn.Sequential(
-            nn.SiLU(),
-            nn.Linear(in_dim, 2 * dim, bias=True),
-        )
-
-    def forward(self, x, embed):
-        scale, shift = self.scale_shift(embed).chunk(2, dim=-1)
-        x = self.norm(x) * (1 + scale[:, None]) + shift[:, None]
-        return x
-
-def sincos_embedding(x, embed_dim, max_period=10000, is_periodic=False):
-    omega = torch.arange(embed_dim//2, dtype=x.dtype, device=x.device)
-    if is_periodic:
-        x = 2 * torch.pi * x
-    else:
-        omega /= embed_dim / 2.
-        omega = 1. / max_period ** omega  
-    out = torch.einsum('m,d->md', x.reshape(-1), omega) 
-    emb_sin = torch.sin(out)  
-    emb_cos = torch.cos(out)  
-    emb = torch.cat([emb_sin, emb_cos], dim=1) 
-    return emb
-
-class TimestepEmbed(nn.Module):
-    def __init__(self, hidden_size, frequency=256, is_periodic=False, sinusoidal=True, dropout=0.0):
-        super().__init__()
-        self.mlp = nn.Sequential(
-            nn.Linear(frequency, hidden_size),
-            nn.SiLU(),
-            nn.Linear(hidden_size, hidden_size),
-        )
-        self.frequency = frequency
-        self.is_periodic = is_periodic
-        self.sinusoidal = sinusoidal
-        # Optional output-side dropout. Helps prevent the encoder from
-        # latching onto specific (hour, doy) combinations when training
-        # data spans many years and most doy values appear repeatedly.
-        # ``dropout=0.0`` (default) is a no-op so existing checkpoints
-        # see identical behaviour.
-        self.dropout = nn.Dropout(p=dropout) if dropout > 0 else nn.Identity()
-
-    def forward(self, x):
-        if not self.sinusoidal:
-            return self.dropout(self.mlp(x))
-        embed = sincos_embedding(
-            x.float(), self.frequency, is_periodic=self.is_periodic
-        ).type_as(x)
-        embed = self.mlp(embed)
-        return self.dropout(embed)
-
-# ---------------------------------------------------------
-# 气象定制：潜空间大核卷积 Block
-# ---------------------------------------------------------
-class GlobalWeatherConvBlock(nn.Module):
-    """
-    经度循环大核 CNN Block，专为地球流体物理推演设计。
-    替代原有的 Window Attention，解决网格伪影并大幅提升感受野。
-    """
-    def __init__(self, dim: int, drop_path: float = 0.0):
-        super().__init__()
-        # 7x7 大核深度可分离卷积
-        self.dwconv = nn.Conv2d(dim, dim, kernel_size=7, groups=dim, bias=False)
-        self.norm = WeatherLayerNorm(dim)
-        
-        # Inverted Bottleneck: 1x1 卷积放大 4 倍通道进行特征融合
-        self.pwconv1 = nn.Conv2d(dim, 4 * dim, kernel_size=1)
-        self.act = nn.GELU()
-        self.pwconv2 = nn.Conv2d(4 * dim, dim, kernel_size=1)
-        
-        self.drop_path = nn.Dropout(drop_path) if drop_path > 0. else nn.Identity()
-
-    def forward(self, x):
-        shortcut = x
-        
-        x = F.pad(x, pad=(3, 3, 3, 3), mode='replicate')  
-
-        x = x.contiguous()
-        
-        x = self.dwconv(x)
-        x = self.norm(x)
-        x = self.pwconv1(x)
-        x = self.act(x)
-        x = self.pwconv2(x)
-
-        return shortcut + self.drop_path(x)
-
-# ---------------------------------------------------------
-# Encoder: 单帧气象特征提取网络
-# ---------------------------------------------------------
-class CubeEmbedConv(nn.Module):
-    """
-    重构后的单帧推演 Encoder。
-    包含三阶段：前置抗混叠 -> Patch下采样 -> 潜空间大核深层演化。
-    """
-    def __init__(
-        self,
-        in_chans: int = 70,
-        out_chans: int = 2048,
-        in_frames: int = 1,          # 现固定为单帧输入
-        patch_size: int = 6,         # 默认 6 倍下采样 (120x240)
-        depth: int = 2,              # 潜空间推演层数 (建议 8-16)
-        flatten: bool = False,     
-        norm_func: Optional[nn.Module] = None,
-        # 兼容旧签名的冗余参数，防止外部调用报错
-        temporal_pad_to: int = 2,  
-        attn_heads: int = 8,
-        attn_dropout: float = 0.0,
-        window_size: int = 8,
-        keep_time_dim: bool = False,
-        feat_chans: int = 2048,
-    ):
-        super().__init__()
-        self.in_chans = in_chans
-        self.out_chans = out_chans
-        self.flatten = flatten
-
-        # 阶段一：高分辨率前置特征融合 (抗混叠，避免下采样丢失高频气象细节)
-        inter_chans = 256
-        self.stage1_smooth = nn.Sequential(
-            nn.Conv2d(in_chans, inter_chans, kernel_size=7, stride=1, padding=3, padding_mode='replicate'),
-            WeatherLayerNorm(inter_chans),
-            nn.GELU()
-        )
-
-        # 阶段二：Patch 强力下采样 (无重叠，直接拉升到目标潜空间通道数)
-        self.stage2_patchify = nn.Sequential(
-            nn.Conv2d(inter_chans, out_chans, kernel_size=patch_size, stride=patch_size),
-            WeatherLayerNorm(out_chans)
-        )
-
-        # 阶段三：潜空间大核全局演化 (替代 Window Attention)
-        self.stage3_evolution = nn.Sequential(
-            *[GlobalWeatherConvBlock(dim=out_chans) for _ in range(depth)]
-        )
-        
-        self.norm = norm_func(out_chans) if norm_func is not None else nn.Identity()
-
-    def forward(
-        self,
-        x: torch.Tensor,
-        lead_hour: Optional[int] = None,       # 兼容签名
-        lead_hour_nums: Optional[int] = None,  # 兼容签名
-    ) -> torch.Tensor:
-        """
-        x: (B, T, C, H, W) 或 (B, C, H, W). 期望 T==1。
-        """
-        # 1. 处理维度，兼容 T 维度
-        if x.ndim == 5:
-            assert x.shape[1] == 1, f"CubeEmbedConv expects single frame (T=1), got {x.shape}"
-            x = x.squeeze(1)  # 转换为 (B, C, H, W)
-        
-        # print("x0",x.shape)
-        # 2. 三阶段推演
-        x = self.stage1_smooth(x)
-        # print("x1",x.shape)
-        x = self.stage2_patchify(x)
-        # print("x2",x.shape)
-        z = self.stage3_evolution(x)
-        # print("z",z.shape)
-        
-        # 3. 输出格式化
-        if self.flatten:
-            z_tok = rearrange(z, "b c h w -> b (h w) c")
-            return self.norm(z_tok)
-            
-        return self.norm(z)
-
-# ---------------------------------------------------------
-# Decoder: 气象场重构网络
-# ---------------------------------------------------------
-def remove_small_scales(x, scale_factor=0.5, mode="bilinear", random_scale=False):
-    import numpy as np
-    if random_scale:
-        scale_factor = np.random.choice([1 / s for s in range(1, 13)])
-    if scale_factor < 1:
-        down = F.interpolate(x, scale_factor=scale_factor, mode=mode, align_corners=False)
-        x = F.interpolate(down, size=x.shape[-2:], mode=mode, align_corners=False)
-    return x
-
-class SmoothDeconv(nn.Module):
-    """
-    数学精确匹配 patch_size 的平滑反卷积。
-    彻底修复原版本中存在的空间偏移(Phase Shift)和形状裁切问题。
-    """
-    def __init__(
-        self,
-        in_channels: int,
-        out_channels: int,
-        kernel_size: Union[int, tuple],
-        stride: Union[int, tuple],
-        bias: bool = True, 
-    ):
-        super().__init__()
-        # 直接使用与 stride 完全相等的 kernel_size，避免 checkerboard artifacts
-        self.deconv = nn.ConvTranspose2d(
-            in_channels,
-            out_channels,
-            kernel_size=kernel_size,
-            stride=stride,
-            padding=0,        # 完美对齐，不需要人为 padding
-            output_padding=0, 
-            bias=bias
-        )
-        # 可选：后接一层普通卷积做进一步平滑
-        self.smooth = nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1, padding_mode='replicate')
-
-    def forward(self, x, output_size=None):
-        x = self.deconv(x)
-        x = self.smooth(x)
-        # 如果有极微小的边界舍入误差，进行安全裁切
-        if output_size is not None and x.shape[-2:] != output_size:
-            x = x[..., :output_size[0], :output_size[1]]
-        return x
-
-
-class ProgressiveUpsample(nn.Module):
-    """
-    阶梯式渐进上采样模块。
-    完美平衡“信息保留”与“防显存溢出(OOM)”的矛盾，严格遵循流体力学的连续性。
-    """
-    def __init__(self, in_channels: int, out_channels: int):
-        super().__init__()
-        
-        # 阶段一：空间放大 2 倍，通道 2048 -> 512
-        self.stage1_conv = nn.Sequential(
-            nn.Conv2d(in_channels, 512, kernel_size=3, padding=1, padding_mode='replicate'),
-            nn.GELU()
-        )
-        
-        # 阶段二：空间放大 3 倍，通道 512 -> 128
-        self.stage2_conv = nn.Sequential(
-            nn.Conv2d(512, 128, kernel_size=3, padding=1, padding_mode='replicate'),
-            nn.GELU()
-        )
-        
-        # 最终映射：128 -> 具体的物理变量通道数 (如 upper_chans)
-        self.head = nn.Conv2d(128, out_channels, kernel_size=3, padding=1, padding_mode='replicate')
-
-    def forward(self, x, target_size=(721, 1440)):
-        # x 初始状态: (B, 2048, 120, 240)
-        
-        # --- Stage 1: 先插值放大空间，再卷积融合降维 (保留高频信息) ---
-        x = F.interpolate(x, scale_factor=2.0, mode='bilinear', align_corners=False)
-        x = self.stage1_conv(x)  # 此时形状: (B, 512, 240, 480)
-        
-        # --- Stage 2: 再次插值放大空间，再降维 ---
-        x = F.interpolate(x, scale_factor=3.0, mode='bilinear', align_corners=False)
-        x = self.stage2_conv(x)  # 此时形状: (B, 128, 720, 1440)
-        
-        # --- 解决 720 与 721 的地球极点拓扑问题 ---
-        current_h, current_w = x.shape[-2:]
-        target_h, target_w = target_size
-        
-        if current_h != target_h or current_w != target_w:
-            pad_h = target_h - current_h  # 721 - 720 = 1
-            pad_w = target_w - current_w  # 1440 - 1440 = 0
-            x = F.pad(x, (0, pad_w, 0, pad_h), mode='replicate')
-            
-        # --- 最终映射输出 ---
-        x = self.head(x)
-        return x
-
-
-class DoubleDeconvHead(nn.Module):
-    def __init__(
-        self, 
-        in_chans: int, 
-        upper_chans: int,
-        lower_chans: int,
-        patch_size: int,
-    ):
-        super().__init__()
-        self.lower_chans = lower_chans
-        self.upper_chans = upper_chans
-
-        # 换用阶梯式渐进上采样
-        self.upper_head = ProgressiveUpsample(in_channels=in_chans, out_channels=upper_chans)
-
-        if self.lower_chans > 0:
-            self.lower_head = ProgressiveUpsample(in_channels=in_chans, out_channels=lower_chans)
-
-    def forward(self, h, residual=None, input_size=(721, 1440), lead_hour=None, target_frames: int = 1):
-        
-        # ==========================================
-        # 🚨 修复 1：安全处理 Batch Tensor 级别的 lead_hour
-        # ==========================================
-        if lead_hour is not None:
-            # lead_hour shape: (B,)
-            # 等价于 max(lead_hour // 6, 1) 的 Tensor 写法
-            lead_step = torch.clamp(lead_hour // 6, min=1.0) 
-            # 广播到 (B, 1, 1, 1) 以便与图像特征相乘
-            lead_step = lead_step.view(-1, 1, 1, 1).to(h.dtype)
-        else:
-            lead_step = 1.0
-        
-        # ==========================================
-        # 2. 解码预测变化量 (Delta)
-        # ==========================================
-        pred = self.upper_head(h, target_size=input_size)
-        output = pred * lead_step
-        
-        if self.lower_chans > 0:
-            lower = self.lower_head(h, target_size=input_size)
-            # 如果下层也需要乘时间步，取消下面的注释：
-            # lower = lower * lead_step
-            output = torch.cat([output, lower], dim=1)
-        
-        # ==========================================
-        # 🚨 修复 2：残差相加必须在 FP32 精度下进行！
-        # ==========================================
-        if residual is not None:
-
-            res_frame = residual[:, -1] if residual.ndim == 5 else residual
-            res_frame = remove_small_scales(res_frame)
-            
-
-            # 强制提升至 FP32 抵抗 bf16 的精度吞噬
-            output = output.to(torch.float32)
-            res_frame = res_frame.to(torch.float32)
-            
-            output = output + res_frame
-        
-        # 返回 fp32 结果，直接交给外面的 MSE Loss 计算（Loss也需要 fp32 保证稳定）
-        return output

code/qwenvl/modalities/weather/internal/polaris_swin.py

@@ -1,527 +0,0 @@
-import numpy as np
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-from .polaris_attention import FlashAttention
-from .polaris_layers import LayerNorm, RMSNorm
-from .helpers import to_2tuple
-
-__all__ = ["SwinBlock", "GeGLU_FFN"]
-
-
-class GeGLU(nn.Module):
-    def forward(self, x):
-        x, gate = x.chunk(2, dim = -1)
-        return F.gelu(gate) * x
-
-class GeGLU_FFN(nn.Module):
-    def __init__(
-        self, 
-        dim, 
-        hidden_dim=None, 
-        multiple_of: int = 32, 
-        dropout=0,
-    ):
-        super().__init__()
-        if hidden_dim is None:
-            hidden_dim = 4 * dim
-            hidden_dim = int(2 * hidden_dim / 3)
-            hidden_dim = multiple_of * ((hidden_dim + multiple_of - 1) // multiple_of)
-        self.fc1 = nn.Linear(dim, hidden_dim * 2, bias=False)
-        self.act = GeGLU()
-        self.fc2 = nn.Linear(hidden_dim, dim, bias=False)
-        self.drop = nn.Dropout(dropout)
-
-    def init_weights(self, init_std: float):
-        nn.init.trunc_normal_(self.fc1.weight, mean=0.0, std=0.02)
-        nn.init.trunc_normal_(self.fc2.weight, mean=0.0, std=init_std)
-        
-    def forward(self, x):
-        return self.drop(self.fc2(self.act(self.fc1(x))))
-
-
-class SwiGLU_FFN(nn.Module):
-    def __init__(
-        self, 
-        dim: int, 
-        hidden_dim = None, 
-        multiple_of: int = 32, 
-        dropout: float = 0.0,
-    ):
-        super().__init__()
-        if hidden_dim is None:
-            hidden_dim = 4 * dim
-            hidden_dim = int(2 * hidden_dim / 3)
-            hidden_dim = multiple_of * ((hidden_dim + multiple_of - 1) // multiple_of)
-        self.w1 = nn.Linear(dim, hidden_dim, bias=False)
-        self.w2 = nn.Linear(hidden_dim, dim, bias=False)
-        self.w3 = nn.Linear(dim, hidden_dim, bias=False)
-        self.dropout = nn.Dropout(dropout)
-
-    def forward(self, x):
-        return self.dropout(self.w2(F.silu(self.w1(x)) * self.w3(x)))
-
-    def init_weights(self, init_std: float):
-        nn.init.trunc_normal_(self.w1.weight, mean=0.0, std=0.02)
-        for linear in (self.w2, self.w3):
-            nn.init.trunc_normal_(linear.weight, mean=0.0, std=init_std)
-
-
-def window_partition(x, window_size):
-    """
-    Args:
-        x: (B, H, W, C)
-        window_size: (win_h, win_w)
-
-    Returns:
-        windows: (num_windows*B, win_h, win_w, C)
-    """
-    B, H, W, C = x.shape
-    x = x.view(B, H // window_size[0], window_size[0], W // window_size[1], window_size[1], C)
-    windows = x.permute(0, 1, 3, 2, 4, 5).contiguous().view(-1, window_size[0], window_size[1], C)
-    return windows
-
-
-def window_reverse(windows, window_size, H, W):
-    """
-    Args:
-        windows: (num_windows*B, window_size, window_size, C)
-        window_size: (win_h, win_w)
-        H (int): Height of image
-        W (int): Width of image
-
-    Returns:
-        x: (B, H, W, C)
-    """
-    B = int(windows.shape[0] / (H * W / window_size[0] / window_size[1]))
-    x = windows.view(B, H // window_size[0], W // window_size[1], window_size[0], window_size[1], -1)
-    x = x.permute(0, 1, 3, 2, 4, 5).contiguous().view(B, H, W, -1)
-    return x
-
-
-
-class WindowAttention(nn.Module):
-    r""" Window based multi-head self attention (W-MSA) module with relative position bias.
-    It supports both of shifted and non-shifted window.
-
-    Args:
-        dim (int): Number of input channels.
-        window_size (tuple[int]): The height and width of the window.
-        num_heads (int): Number of attention heads.
-        qkv_bias (bool, optional):  If True, add a learnable bias to query, key, value. Default: True
-        qk_scale (float | None, optional): Override default qk scale of head_dim ** -0.5 if set
-        attn_drop (float, optional): Dropout ratio of attention weight. Default: 0.0
-        proj_drop (float, optional): Dropout ratio of output. Default: 0.0
-    """
-
-    def __init__(self, dim, window_size, num_heads, qkv_bias=True, qk_scale=None, attn_drop=0., proj_drop=0.):
-
-        super().__init__()
-        self.dim = dim
-        self.window_size = window_size  # Wh, Ww
-        self.num_heads = num_heads
-        head_dim = dim // num_heads
-        self.scale = qk_scale or head_dim ** -0.5
-
-        # define a parameter table of relative position bias
-        self.relative_position_bias_table = nn.Parameter(
-            torch.zeros((2 * window_size[0] - 1) * (2 * window_size[1] - 1), num_heads))  # 2*Wh-1 * 2*Ww-1, nH
-
-        # get pair-wise relative position index for each token inside the window
-        coords_h = torch.arange(self.window_size[0])
-        coords_w = torch.arange(self.window_size[1])
-        coords = torch.stack(torch.meshgrid([coords_h, coords_w], indexing='ij'))  # 2, Wh, Ww
-        coords_flatten = torch.flatten(coords, 1)  # 2, Wh*Ww
-        relative_coords = coords_flatten[:, :, None] - coords_flatten[:, None, :]  # 2, Wh*Ww, Wh*Ww
-        relative_coords = relative_coords.permute(1, 2, 0).contiguous()  # Wh*Ww, Wh*Ww, 2
-        relative_coords[:, :, 0] += self.window_size[0] - 1  # shift to start from 0
-        relative_coords[:, :, 1] += self.window_size[1] - 1
-        relative_coords[:, :, 0] *= 2 * self.window_size[1] - 1
-        relative_position_index = relative_coords.sum(-1)  # Wh*Ww, Wh*Ww
-        self.register_buffer("relative_position_index", relative_position_index)
-
-        self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
-        self.attn_drop = nn.Dropout(attn_drop)
-        self.proj = nn.Linear(dim, dim)
-        self.proj_drop = nn.Dropout(proj_drop)
-        nn.init.trunc_normal_(self.relative_position_bias_table, std=.02)
-
-
-    def forward(self, x, mask=None, **kwargs):
-        """
-        Args:
-            x: input features with shape of (num_windows*B, N, C)
-            mask: (0/-inf) mask with shape of (num_windows, Wh*Ww, Wh*Ww) or None
-        """
-        B_, N, C = x.shape
-
-        qkv = self.qkv(x).reshape(B_, N, 3, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4)
-        q, k, v = qkv[0], qkv[1], qkv[2]  # make torchscript happy (cannot use tensor as tuple)
-
-        q = q * self.scale
-        attn = (q @ k.transpose(-2, -1))
-
-        relative_position_bias = self.relative_position_bias_table[self.relative_position_index.view(-1)].view(
-            self.window_size[0] * self.window_size[1], self.window_size[0] * self.window_size[1], -1)  # Wh*Ww,Wh*Ww,nH
-        relative_position_bias = relative_position_bias.permute(2, 0, 1).contiguous()  # nH, Wh*Ww, Wh*Ww
-        attn = attn + relative_position_bias.unsqueeze(0)
-
-        if mask is not None:
-            nW = mask.shape[0]
-            attn = attn.view(B_ // nW, nW, self.num_heads, N, N) + mask.unsqueeze(1).unsqueeze(0)
-            attn = attn.view(-1, self.num_heads, N, N)
-            attn = attn.softmax(dim=-1)
-        else:
-            attn = attn.softmax(dim=-1)
-
-        attn = self.attn_drop(attn)
-
-        x = (attn @ v).transpose(1, 2).reshape(B_, N, C)
-        x = self.proj(x)
-        x = self.proj_drop(x)
-        return x
-
-
-
-class WindowAttentionV2(nn.Module):
-    r""" Window based multi-head self attention (W-MSA) module with relative position bias.
-    It supports both of shifted and non-shifted window.
-    Args:
-        dim (int): Number of input channels.
-        window_size (tuple[int]): The height and width of the window.
-        num_heads (int): Number of attention heads.
-        qkv_bias (bool, optional):  If True, add a learnable bias to query, key, value. Default: True
-        attn_drop (float, optional): Dropout ratio of attention weight. Default: 0.0
-        proj_drop (float, optional): Dropout ratio of output. Default: 0.0
-    """
-    
-    def __init__(self, dim, window_size, num_heads, attn_drop=0., proj_drop=0.):
-
-        super().__init__()
-        self.dim = dim
-        self.window_size = window_size  # Wh, Ww
-        self.num_heads = num_heads
-
-        self.logit_scale = nn.Parameter(torch.log(10 * torch.ones((num_heads, 1, 1))), requires_grad=True)
-
-        # mlp to generate continuous relative position bias
-        self.cpb_mlp = nn.Sequential(nn.Linear(2, 512, bias=True),
-                                     nn.ReLU(inplace=True),
-                                     nn.Linear(512, num_heads, bias=False))
-
-        # get relative_coords_table
-        relative_coords_h = torch.arange(-(self.window_size[0] - 1), self.window_size[0], dtype=torch.float32)
-        relative_coords_w = torch.arange(-(self.window_size[1] - 1), self.window_size[1], dtype=torch.float32)
-        
-        relative_coords_table = torch.stack(
-            torch.meshgrid([relative_coords_h,relative_coords_w], indexing='ij')
-        ).permute(1, 2, 0).contiguous().unsqueeze(0)  # 1, 2*Wh-1, 2*Ww-1, 2
-        
-
-        relative_coords_table[:, :, :, 0] /= (self.window_size[0] - 1)
-        relative_coords_table[:, :, :, 1] /= (self.window_size[1] - 1)
-        relative_coords_table *= 8  # normalize to -8, 8
-        
-        relative_coords_table = torch.sign(relative_coords_table) * torch.log2(
-            torch.abs(relative_coords_table) + 1.0) / np.log2(8)
-
-        self.register_buffer("relative_coords_table", relative_coords_table)
-
-        # get pair-wise relative position index for each token inside the window
-        coords_h = torch.arange(self.window_size[0])
-        coords_w = torch.arange(self.window_size[1])
-        coords = torch.stack(torch.meshgrid([coords_h, coords_w], indexing='ij'))  # 2, Wh, Ww
-        coords_flatten = torch.flatten(coords, 1)  # 2, Wh*Ww
-        relative_coords = coords_flatten[:, :, None] - coords_flatten[:, None, :]  # 2, Wh*Ww, Wh*Ww
-        relative_coords = relative_coords.permute(1, 2, 0).contiguous()  # Wh*Ww, Wh*Ww, 2
-        relative_coords[:, :, 0] += self.window_size[0] - 1  # shift to start from 0
-        relative_coords[:, :, 1] += self.window_size[1] - 1
-        relative_coords[:, :, 0] *= 2 * self.window_size[1] - 1
-        relative_position_index = relative_coords.sum(-1)  # Wh*Ww, Wh*Ww
-        self.register_buffer("relative_position_index", relative_position_index)
-
-        
-        self.qkv = nn.Linear(dim, dim * 3, bias=False)
-        self.q_bias = nn.Parameter(torch.zeros(dim))
-        self.v_bias = nn.Parameter(torch.zeros(dim))
-        self.register_buffer("k_bias", torch.zeros(dim).half())
-        # self.register_buffer('k_bias', torch.zeros(dim), persistent=False)
-            
-        self.attn_drop = nn.Dropout(attn_drop)
-        self.out_proj = nn.Linear(dim, dim)
-        self.proj_drop = nn.Dropout(proj_drop)
-        self.softmax = nn.Softmax(dim=-1)
-        
-
-    def forward(self, x, mask=None, **kwargs):
-        """
-        Args:
-            x: input features with shape of (num_windows*B, N, C)
-            mask: (0/-inf) mask with shape of (num_windows, Wh*Ww, Wh*Ww) or None
-        """
-        B_, N, C = x.shape
-
-        k_bias = self.k_bias.to(self.q_bias)
-        qkv_bias = torch.cat([self.q_bias, k_bias, self.v_bias])
-        # qkv_bias = torch.cat((self.q_bias, self.k_bias, self.v_bias))
-
-        qkv = F.linear(input=x, weight=self.qkv.weight, bias=qkv_bias)
-        qkv = qkv.reshape(B_, N, 3, self.num_heads, -1).permute(2, 0, 3, 1, 4)
-        q, k, v = qkv[0], qkv[1], qkv[2]  # make torchscript happy (cannot use tensor as tuple)
-        
-        # cosine attention
-        attn = (F.normalize(q, dim=-1) @ F.normalize(k, dim=-1).transpose(-2, -1))        
-        
-        logit_scale = self.logit_scale.clamp(max=np.log(1. / 0.01)).exp()
-        attn = attn * logit_scale
-        
-        relative_position_bias_table = self.cpb_mlp(self.relative_coords_table.to(x)).view(-1, self.num_heads)
-        relative_position_bias = relative_position_bias_table[self.relative_position_index.view(-1)].view(
-            self.window_size[0] * self.window_size[1], self.window_size[0] * self.window_size[1], -1)  # Wh*Ww,Wh*Ww,nH
-        relative_position_bias = relative_position_bias.permute(2, 0, 1).contiguous()  # nH, Wh*Ww, Wh*Ww
-        relative_position_bias = 16 * torch.sigmoid(relative_position_bias)
-        
-        attn = attn + relative_position_bias.unsqueeze(0)
-        
-        
-        if mask is not None:
-            nW = mask.shape[0]
-            attn = attn.view(B_ // nW, nW, self.num_heads, N, N) + mask.unsqueeze(1).unsqueeze(0)
-            attn = attn.view(-1, self.num_heads, N, N)
-            attn = self.softmax(attn)
-        else:
-            attn = self.softmax(attn)
-        
-        attn = self.attn_drop(attn).to(v)
-        
-        x = (attn @ v).transpose(1, 2).reshape(B_, N, C)
-        x = self.out_proj(x)
-        x = self.proj_drop(x)
-        return x
-
-def _compute_attn_mask(H, W, window_size, shift_size, mask_type):
-    img_mask = torch.zeros((1, H, W, 1))
-    h_slices = (slice(0, -window_size[0]),
-                slice(-window_size[0], -shift_size[0]),
-                slice(-shift_size[0], None))
-    w_slices = (slice(0, -window_size[1]),
-                slice(-window_size[1], -shift_size[1]),
-                slice(-shift_size[1], None))
-    cnt = 0
-    for h in h_slices:
-        for w in w_slices:
-            if mask_type == 'h':
-                img_mask[:, h, :, :] = cnt
-            elif mask_type == 'w':
-                img_mask[:, :, w, :] = cnt
-            elif mask_type == 'hw':
-                img_mask[:, h, w, :] = cnt
-            cnt += 1
-    mask_windows = window_partition(img_mask, window_size)
-    mask_windows = mask_windows.view(-1, window_size[0] * window_size[1])
-    attn_mask = mask_windows.unsqueeze(1) - mask_windows.unsqueeze(2)
-    attn_mask = attn_mask.masked_fill(attn_mask != 0, float(-100.0)).masked_fill(attn_mask == 0, float(0.0))
-    return attn_mask
-
-
-from transformers.modeling_layers import GradientCheckpointingLayer
-class SwinBlock(GradientCheckpointingLayer):
-    def __init__(
-            self, dim, num_heads,
-            input_size, window_size=7, shift_size=0, embed_dim=None,
-            attn_type='v1', mask_type='hw', norm_type="ln", ffn_type="geglu_ffn",
-            qk_scale=None, n_kv_heads=None,
-            mlp_ratio=4., drop=0., attn_drop=0.,
-            attn_implementation="flash_attention_2",
-            **kwargs
-        ):
-        super().__init__()
-        if embed_dim is None:
-            embed_dim = dim
-        self.dim = dim
-        self.default_input_size = tuple(int(s) for s in input_size)
-        self.num_heads = num_heads
-        self.window_size = to_2tuple(window_size)
-        self.shift_size = to_2tuple(shift_size)
-        self.mlp_ratio = mlp_ratio
-        self.norm_type = norm_type
-        self.attn_type = attn_type
-        self.ffn_type = ffn_type
-        self.mask_type = mask_type
-        self._attn_implementation = attn_implementation
-
-        default_full_window = all(
-            w == s for w, s in zip(self.window_size, self.default_input_size)
-        )
-
-        assert 0 <= self.shift_size[0] < self.window_size[0], "shift_size must in 0-window_size"
-        assert 0 <= self.shift_size[1] < self.window_size[1], "shift_size must in 0-window_size"
-
-        if norm_type == "ln":
-            self.norm1 = LayerNorm(dim, eps=1e-6)
-            self.norm2 = LayerNorm(dim, eps=1e-6)
-        elif norm_type == "adaln":
-            self.norm1 = LayerNorm(dim, eps=1e-6, elementwise_affine=False)
-            self.norm2 = LayerNorm(dim, eps=1e-6, elementwise_affine=False)
-            adaln_linear = nn.Linear(embed_dim, 6 * dim, bias=True)
-            nn.init.zeros_(adaln_linear.weight)
-            nn.init.zeros_(adaln_linear.bias)
-            self.adaln = nn.Sequential(nn.SiLU(), adaln_linear)
-        elif norm_type == "adarms":
-            self.norm1 = RMSNorm(dim, eps=1e-5)
-            self.norm2 = RMSNorm(dim, eps=1e-5)
-            adaln_linear = nn.Linear(embed_dim, 4 * dim, bias=True)
-            nn.init.zeros_(adaln_linear.weight)
-            nn.init.zeros_(adaln_linear.bias)
-            self.adaln = nn.Sequential(nn.SiLU(), adaln_linear)
-        else:
-            raise ValueError(f"norm_type {norm_type} not supported")
-
-        if attn_type == 'v1':
-            self.attn = WindowAttention(
-                dim, window_size=self.window_size, num_heads=num_heads,
-                qk_scale=qk_scale, attn_drop=attn_drop, proj_drop=drop)
-        elif attn_type == 'v2':
-            self.attn = WindowAttentionV2(
-                dim, window_size=self.window_size, num_heads=num_heads,
-                attn_drop=attn_drop, proj_drop=drop)
-        elif attn_type == "flash":
-            self.attn = FlashAttention(
-                dim=dim,
-                n_heads=num_heads,
-                n_kv_heads=n_kv_heads,
-                attn_implementation=self._attn_implementation,
-                dropout=attn_drop,
-                is_causal=False,
-                max_seq_len=None if default_full_window else int(np.prod(self.default_input_size)),
-            )
-        else:
-            raise ValueError(f"attn_type {attn_type} not supported")
-
-
-        if ffn_type == "geglu_ffn":
-            self.mlp = GeGLU_FFN(dim)
-        elif ffn_type == "swiglu_ffn":
-            self.mlp = SwiGLU_FFN(dim)
-        else:
-            raise ValueError(f"ffn_type {ffn_type} not supported")
-
-        if not default_full_window and max(self.shift_size) > 0:
-            H, W = self.default_input_size
-            attn_mask = _compute_attn_mask(H, W, self.window_size, self.shift_size, mask_type)
-        else:
-            attn_mask = None
-        self.register_buffer("attn_mask", attn_mask, persistent=False)
-
-        self._attn_mask_cache = {}
-
-    def init_weights(self, init_std):
-        for norm in (self.norm1, self.norm2):
-            norm.reset_parameters()
-        self.attn.init_weights(init_std)
-        self.mlp.init_weights(init_std)
-
-    def _get_attn_mask(self, H, W, device):
-        if max(self.shift_size) == 0:
-            return None
-        if (H, W) == tuple(self.default_input_size):
-            return self.attn_mask
-        key = (H, W)
-        if key not in self._attn_mask_cache:
-            self._attn_mask_cache[key] = _compute_attn_mask(
-                H, W, self.window_size, self.shift_size, self.mask_type
-            )
-        return self._attn_mask_cache[key].to(device)
-
-    def _resolve_input_size(self, L):
-        if L == self.default_input_size[0] * self.default_input_size[1]:
-            return self.default_input_size
-        wh, ww = self.window_size
-        sqrt_l = int(L ** 0.5)
-        for H in range(sqrt_l, 0, -1):
-            if L % H != 0:
-                continue
-            W = L // H
-            if H % wh == 0 and W % ww == 0:
-                return (H, W)
-        raise ValueError(
-            f"Cannot resolve seq_len={L} into (H, W) divisible by window_size={self.window_size}"
-        )
-
-    def window_attention(self, x, freqs_cos, freqs_sin, input_size):
-        H, W = input_size
-        B, L, C = x.shape
-
-        x = x.view(B, H, W, C)
-
-        if max(self.shift_size) > 0:
-            shifted_x = torch.roll(x, shifts=(-self.shift_size[0], -self.shift_size[1]), dims=(1, 2))
-        else:
-            shifted_x = x
-
-        x_windows = window_partition(shifted_x, self.window_size)
-        x_windows = x_windows.view(-1, self.window_size[0] * self.window_size[1], C)
-
-        attn_mask = self._get_attn_mask(H, W, x.device)
-
-        if self.attn_type == "flash" and (H, W) != tuple(self.default_input_size):
-            old_max_seq = self.attn.max_seq_len
-            self.attn.max_seq_len = None
-            attn_windows = self.attn(x_windows, freqs_cos=freqs_cos, freqs_sin=freqs_sin, mask=attn_mask)
-            self.attn.max_seq_len = old_max_seq
-        else:
-            attn_windows = self.attn(x_windows, freqs_cos=freqs_cos, freqs_sin=freqs_sin, mask=attn_mask)
-
-        attn_windows = attn_windows.view(-1, self.window_size[0], self.window_size[1], C)
-        shifted_x = window_reverse(attn_windows, self.window_size, H, W)
-
-        if max(self.shift_size) > 0:
-            x = torch.roll(shifted_x, shifts=(self.shift_size[0], self.shift_size[1]), dims=(1, 2))
-        else:
-            x = shifted_x
-
-        x = x.view(B, H * W, C)
-        return x
-
-
-    def forward(
-        self,
-        x: torch.Tensor,
-        embed: torch.Tensor=None,
-        freqs_cos: torch.Tensor=None,
-        freqs_sin: torch.Tensor=None,
-        **kwargs,
-    ):
-        B, L, C = x.shape
-        input_size = self._resolve_input_size(L)
-        is_full_window = all(w == s for w, s in zip(self.window_size, input_size))
-
-        shortcut = x
-
-        if self.norm_type == "adaln":
-            gamma = self.adaln(embed).unsqueeze(1)
-            scale_msa, shift_msa, gate_msa, scale_mlp, shift_mlp, gate_mlp = gamma.chunk(6, dim=-1)
-            x = self.norm1(x) * (1 + scale_msa) + shift_msa
-        elif self.norm_type == "adarms":
-            scale_msa, gate_msa, scale_mlp, gate_mlp = self.adaln(embed).unsqueeze(1).chunk(4, dim=-1)
-            x = self.norm1(x) * (1 + scale_msa)
-        else:
-            x = self.norm1(x)
-
-        if is_full_window:
-            x = self.attn(x, freqs_cos, freqs_sin)
-        else:
-            x = self.window_attention(x, freqs_cos, freqs_sin, input_size)
-
-        if self.norm_type == "adaln":
-            x = shortcut + gate_msa * x
-            x = x + gate_mlp*self.mlp(self.norm2(x)*(1+scale_mlp)+shift_mlp)
-        elif self.norm_type == "adarms":
-            x = shortcut + gate_msa * x
-            x = x + gate_mlp*self.mlp(self.norm2(x)*(1+scale_mlp))
-        else:
-            x = shortcut + x
-            x = x + self.mlp(self.norm2(x))
-
-        return x
-    

code/qwenvl/modalities/weather/processor.py

@@ -1,32 +0,0 @@
-"""Weather processor (placeholder).
-
-The bio modalities use a per-modality processor to turn raw inputs into
-encoder tensors.  Weather is different: the ``ERA5QwenVLDataset`` (in
-``qwenvl/modalities/weather/data/era5_dataset.py``) directly emits all
-encoder tensors plus the chat-template placeholder.  This file exists
-only so that ``BaseProcessor``-style callers — if any — see a concrete
-class for the modality.
-"""
-
-from __future__ import annotations
-
-from typing import Any, Dict, Optional, Tuple
-
-
-class WeatherProcessor:
-    """No-op processor; the dataset class does the real work."""
-
-    modality_name = "weather"
-
-    def process_input(self, raw_input: Any, **kwargs) -> Dict[str, Any]:
-        raise NotImplementedError(
-            "Weather inputs come straight from ERA5QwenVLDataset; the dataset "
-            "produces encoder-ready tensors directly."
-        )
-
-    def build_placeholder(
-        self, raw_input: Any, is_output: bool = False, **kwargs,
-    ) -> Tuple[str, Optional[Any]]:
-        raise NotImplementedError(
-            "Use the chat-template helper in qwenvl/modalities/weather/data/era5_dataset.py."
-        )

code/qwenvl/modalities/weather/projector.py

@@ -1,29 +0,0 @@
-"""Identity projector for the weather modality.
-
-The Polaris-style ``meteo_merger`` already runs inside ``WeatherEncoder``
-and projects swin hidden → ``qwenvl_dim``.  The router still expects a
-projector module so we register a no-op ``IdentityWeatherProjector`` to
-keep the registration logic uniform with the bio modalities.
-"""
-
-from __future__ import annotations
-
-import torch
-import torch.nn as nn
-
-
-class IdentityWeatherProjector(nn.Module):
-    """Pass-through projector. Encoder already outputs qwenvl_dim."""
-
-    def __init__(self, qwenvl_dim: int):
-        super().__init__()
-        self.qwenvl_dim = qwenvl_dim
-        # Single dummy parameter so that an "is the projector trainable"
-        # check does not return False unexpectedly when tune_weather_projector
-        # is left at its default True; the router treats projectors with
-        # no trainable params as frozen which would short-circuit the
-        # dummy_forward gradient sync path.
-        self._noop = nn.Parameter(torch.zeros(1), requires_grad=False)
-
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
-        return x