refactor: move mlx_dit and mlx_vae into models/mlx#602
Conversation
- Relocate acestep/mlx_dit/ -> acestep/models/mlx/dit_*.py - Relocate acestep/mlx_vae/ -> acestep/models/mlx/vae_*.py - Consolidate duplicate mlx_available() into models/mlx/__init__.py - Update all import paths in handler.py, diffusion.py, generation.py - Delete old mlx_dit/ and mlx_vae/ directories - No logic changes; pure move-and-rename refactor
📝 WalkthroughWalkthroughThis PR reorganizes MLX (Machine Learning X) component import paths from scattered locations to a centralized Changes
Estimated code review effort🎯 2 (Simple) | ⏱️ ~10 minutes Poem
🚥 Pre-merge checks | ✅ 4✅ Passed checks (4 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing touches
🧪 Generate unit tests (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 1
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
acestep/models/mlx/dit_generate.py (1)
1-4:⚠️ Potential issue | 🟡 MinorAdd a module docstring (required for modified modules).
The header is comment-only; please add a concise module docstring that states purpose and key inputs/outputs.
✅ Suggested docstring
-# MLX diffusion generation loop for AceStep DiT decoder. -# -# Replicates the timestep scheduling and ODE/SDE stepping from -# ``AceStepConditionGenerationModel.generate_audio`` using pure MLX arrays. +"""MLX diffusion generation loop for the AceStep DiT decoder. + +Inputs: MLX decoder + numpy conditioning arrays. +Outputs: numpy target latents and timing metrics. +"""As per coding guidelines, "Docstrings are mandatory for all new or modified Python modules, classes, and functions" and "Docstrings must be concise and include purpose plus key inputs/outputs (and raised exceptions when relevant)."
🤖 Fix all issues with AI agents
In `@acestep/models/mlx/__init__.py`:
- Around line 1-2: Replace the two-line header comment in __init__.py with a
proper module-level docstring that concisely states the module’s purpose (native
MLX implementations of AceStep models targeting Apple Silicon), describes
behavior (provides pure MLX inference with graceful fallback to PyTorch), and
enumerates the key exports/entry points provided by this package (list exported
classes, functions, and constants you expose from this __init__ such as MLX
model wrappers or factory functions) including expected inputs/outputs and any
raised exceptions; ensure the docstring follows project style (brief
description, usage/exports summary).
| # Native MLX implementations of AceStep models for Apple Silicon. | ||
| # Provides pure MLX inference with graceful fallback to PyTorch. |
There was a problem hiding this comment.
Replace the header comment with a module docstring.
This module is modified; it needs a docstring describing purpose and key exports.
✅ Suggested docstring
-# Native MLX implementations of AceStep models for Apple Silicon.
-# Provides pure MLX inference with graceful fallback to PyTorch.
+"""Native MLX implementations of AceStep models for Apple Silicon.
+
+Exports:
+ - is_mlx_available() -> bool
+ - mlx_available() -> bool
+"""As per coding guidelines, "Docstrings are mandatory for all new or modified Python modules, classes, and functions" and "Docstrings must be concise and include purpose plus key inputs/outputs (and raised exceptions when relevant)."
📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| # Native MLX implementations of AceStep models for Apple Silicon. | |
| # Provides pure MLX inference with graceful fallback to PyTorch. | |
| """Native MLX implementations of AceStep models for Apple Silicon. | |
| Exports: | |
| - is_mlx_available() -> bool | |
| - mlx_available() -> bool | |
| """ |
🤖 Prompt for AI Agents
In `@acestep/models/mlx/__init__.py` around lines 1 - 2, Replace the two-line
header comment in __init__.py with a proper module-level docstring that
concisely states the module’s purpose (native MLX implementations of AceStep
models targeting Apple Silicon), describes behavior (provides pure MLX inference
with graceful fallback to PyTorch), and enumerates the key exports/entry points
provided by this package (list exported classes, functions, and constants you
expose from this __init__ such as MLX model wrappers or factory functions)
including expected inputs/outputs and any raised exceptions; ensure the
docstring follows project style (brief description, usage/exports summary).
1 participant
Summary by CodeRabbit