wound-river

Wound River

This repository contains the trained River vector-field weights used for longitudinal porcine wound-image forecasting.

The released artifact contains only the 102 project-trained River tensors. It deliberately excludes the 277 embedded VQ-MUSE autoencoder tensors found in the complete training checkpoint.

Users must obtain the matching VQ-MUSE model separately from:

• https://huggingface.co/Emma02/vqvae_ckpts

This repository is currently private while public-release authority and licensing are being confirmed.

Released artifact

Architecture

The released model uses sparsely conditioned flow matching with a VQ-MUSE image representation.

Final vector-field configuration:

• Inner dimension: 384
• Transformer depth: 2
• Middle depth: 3
• Dropout: 0.075
• State size: 64
• State resolution: 16 x 16
• Output normalization: layer normalization
• Self-conditioning probability: 0.5
• Target skew power: 0.5
• Flow sigma: 1e-7

The complete portable configuration is provided in river_final.yaml.

Loading

First construct the modified River model with the matching VQ-MUSE autoencoder. Then overlay the released River-only tensors:

import torch

package = torch.load(
    "wound_river_delta_v1.pth",
    map_location="cpu",
    weights_only=False,
)

result = model.load_state_dict(
    package["model"],
    strict=False,
)

assert not result.unexpected_keys
assert all(
    key.startswith("ae.")
    for key in result.missing_keys
)

model.eval()

The missing ae.* keys are expected because the external VQ-MUSE weights are intentionally excluded from this artifact.

Dependencies

Project code

https://github.com/bridenmj/wound-forecasting

Upstream River

https://github.com/araachie/river

The project uses a modified River implementation. The upstream River source is distributed under GPL-3.0; its license is retained as LICENSE_RIVER.txt. Inclusion of this notice documents upstream provenance and does not assign a separate project-artifact license.

VQ-MUSE

https://huggingface.co/Emma02/vqvae_ckpts

VQ-MUSE is an external dependency and is not mirrored in this repository.

Training configuration

• Four observed context frames
• One generated frame per training task
• Five observations per training sample
• Adam-style learning rate: 1e-4
• Weight decay: 3.2e-5
• Warmup steps: 5000
• Total training steps: 40000
• Gradient accumulation steps: 4

At evaluation, four observed frames were used to autoregressively generate four future frames with 100 flow-integration steps.

Intended use

This artifact is provided for research reproduction and study of flow-matching approaches to longitudinal image prediction.

<!-- wound-forecasting-project-license:start -->

Project license

The original project-trained weights released in this repository are licensed under the Creative Commons Attribution–NonCommercial 4.0 International license (CC BY-NC 4.0).

Third-party implementations, base models, tokenizers, and other external dependencies are not relicensed by this repository and remain subject to their respective licenses and terms. Any upstream license or notice files included in this repository continue to apply to the corresponding upstream materials.

<!-- wound-forecasting-project-license:end -->