Masking diffusion

[arrjon]

This pull request introduces new functionality for conditional masking and partial target conditioning in the flow_matching, consistency_model, stable_consistency_model, and diffusion_model classes. The main additions are the drop_cond_prob parameter (which allows probabilistic masking of conditioning inputs during training, enabling simultaneous training of conditional and unconditional models) and, for diffusion models, the drop_target_prob parameter (enabling training and inference with partial targets). The code ensures proper handling of these features during inference and metrics computation.

Diffusion model:

• During inference, pass a target_mask and targets_fixed to specify which parts of the target should be kept fixed to the values in targets_fixed.

All diffusion type models:

• During inference, set unconditional_mode=True, to use mask conditions during sampling.

[codecov]

Codecov Report

❌ Patch coverage is 55.97015% with 59 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
bayesflow/networks/flow_matching/flow_matching.py	50.00%	25 Missing ⚠️
...rks/consistency_models/stable_consistency_model.py	0.00%	16 Missing ⚠️
...w/networks/consistency_models/consistency_model.py	35.29%	11 Missing ⚠️
...esflow/networks/diffusion_model/diffusion_model.py	86.27%	7 Missing ⚠️

Files with missing lines	Coverage Δ
...esflow/networks/diffusion_model/diffusion_model.py	86.21% <86.27%> (-0.34%)	⬇️
...w/networks/consistency_models/consistency_model.py	90.51% <35.29%> (-7.85%)	⬇️
...rks/consistency_models/stable_consistency_model.py	19.49% <0.00%> (-2.84%)	⬇️
bayesflow/networks/flow_matching/flow_matching.py	78.80% <50.00%> (-14.83%)	⬇️

... and 23 files with indirect coverage changes