Skip to content

Add Nemotron Labs Diffusion model#1239

Merged
Blaizzy merged 21 commits into
mainfrom
pc/nemotron-labs-diffusion
Jun 1, 2026
Merged

Add Nemotron Labs Diffusion model#1239
Blaizzy merged 21 commits into
mainfrom
pc/nemotron-labs-diffusion

Conversation

@Blaizzy
Copy link
Copy Markdown
Owner

@Blaizzy Blaizzy commented May 29, 2026

Summary

  • Add the Nemotron Labs Diffusion text model implementation with AR, diffusion, and linear self-speculative generation support.
  • Add prompt formatting support and generic --gen-kwargs plumbing for model-specific generation options.
  • Add focused diffusion model tests and a model README with usage examples.

Validation

  • uv run python -m compileall -q mlx_vlm/models/nemotron_labs_diffusion mlx_vlm/generate/dispatch.py mlx_vlm/utils.py
  • uv run python -m pytest mlx_vlm/tests/test_diffusion_models.py mlx_vlm/tests/test_prompt_utils.py -q

@Blaizzy Blaizzy force-pushed the pc/nemotron-labs-diffusion branch 2 times, most recently from 45b95c8 to c5a2a77 Compare May 29, 2026 23:13
@Blaizzy Blaizzy force-pushed the pc/nemotron-labs-diffusion branch 2 times, most recently from 7ef792f to 0b13f8a Compare May 29, 2026 23:39
@Blaizzy Blaizzy force-pushed the pc/nemotron-labs-diffusion branch from 0b13f8a to ae05744 Compare May 29, 2026 23:41
@Blaizzy Blaizzy marked this pull request as ready for review May 30, 2026 20:27
Comment on lines +856 to +859
linear_speculative=kwargs.get("linear_speculative", False)
or kwargs.get("linear_speculation", False)
or kwargs.get("generation_mode") in ("linear_speculative", "linear_spec"),
**model_generate_kwargs,
Copy link
Copy Markdown
Owner Author

Choose a reason for hiding this comment

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

To fix

generation_mode = kwargs.get("generation_mode")
if generation_mode is not None:
return generation_mode in (
"diffusion",
Copy link
Copy Markdown
Owner Author

Choose a reason for hiding this comment

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

To fix

@Blaizzy Blaizzy merged commit 735d3de into main Jun 1, 2026
1 check passed
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.

1 participant