Reza2kn

/

Cosmos3-Nano-MLX-8bit

like 3

Hi Reza,

I've ported the audio generation branch to MLX for your Cosmos3-Nano-MLX-8bit build.

The MLX weights already included audio_proj_in, audio_proj_out, and audio_modality_embed — just needed to wire them up in the forward pass.

Changes:

• Added pack_sound, unpack_sound, apply_timestep_to_noisy in mlx_cosmos3.py
• Updated mlx_pipeline.py to handle sound_tokens, sound_timesteps, sound_sequence_indexes
• Set sound_gen=True, sound_dim=64 in config

All modes working on Mac Studio (M-series):

• text2image ✅
• text2video (25 frames, 384×384, ~1 min) ✅
• text2video+audio (stereo 48kHz WAV) ✅

Full diff:
https://huggingface.co/Reza2kn/Cosmos3-Nano-MLX-8bit

Happy to open a PR if you'd like.

Best,
Bogdan