Skip to content

Enable DenoisingMoEPredictor in inference#1059

Merged
AnnaKwa merged 32 commits into
mainfrom
feature/denoising-moe
Apr 29, 2026
Merged

Enable DenoisingMoEPredictor in inference#1059
AnnaKwa merged 32 commits into
mainfrom
feature/denoising-moe

Conversation

@AnnaKwa
Copy link
Copy Markdown
Contributor

@AnnaKwa AnnaKwa commented Apr 20, 2026

This PR allows a DenoisingMoEPredictor to be used in the various inference entrypoints. It also allows the CompositePredictor to use one to generate patch predictions.

AnnaKwa added 30 commits April 10, 2026 15:10
Introduce DenoisingRangeModelConfig and DenoisingMoECheckpointConfig with
contiguous sigma-range validation and data_requirements from the first expert.

Made-with: Cursor
… dispatch

Add expert compatibility checks, _SigmaDispatchModule for EDM sampler routing,
DenoisingMoECheckpointConfig.build(), and generation API matching DiffusionModel.

Made-with: Cursor
Annotate PatchPredictor model parameter with DenoisingScheduleSequentialPredictor
via TYPE_CHECKING to avoid import cycles.

Made-with: Cursor
Union model config with MoE variant (dacite order: MoE first) and extend
Evaluator/EventEvaluator model type hints for DenoisingScheduleSequentialPredictor.

Made-with: Cursor
Extend InferenceConfig and Downscaler model types for sequential MoE predictor.

Made-with: Cursor
Extend DownscalerConfig and downscaler model types for sequential MoE predictor.

Made-with: Cursor
Cover sigma range validation, dispatch routing, boundary tie-break, and
out-of-range errors.

Made-with: Cursor
@AnnaKwa AnnaKwa changed the title Feature/denoising moe Enable DenoisingMoEPredictor in inference Apr 27, 2026
@AnnaKwa AnnaKwa marked this pull request as ready for review April 27, 2026 20:43
Copy link
Copy Markdown
Collaborator

@frodre frodre left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM

@AnnaKwa AnnaKwa enabled auto-merge (squash) April 29, 2026 15:33
@AnnaKwa AnnaKwa merged commit d7f26d8 into main Apr 29, 2026
9 of 10 checks passed
@AnnaKwa AnnaKwa deleted the feature/denoising-moe branch April 29, 2026 16:52
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants