PhonoDiffMRI -- Under Construction. Pipelines will be ready soon

PhonoDiffMRI generates dynamic speech MRI from speech audio using speech-to-video diffusion models conditioned on acoustic embeddings and optional PhonoQ phonological posteriors.

This repository contains a cleaned version of the research code used for the paper. The reusable code lives in src/phonodiffmri, with command-line entry points in scripts/.

Repository Layout

.
+-- src/phonodiffmri/        # Installable Python package
|   +-- data.py              # MRI frame/audio dataset
|   +-- generation.py        # Speech-to-video generation utilities
|   +-- model.py             # Lightning diffusion module
|   +-- pipelines/           # Local Diffusers-compatible speech-to-video pipeline
|   +-- sdm_whisper.py       # SDXL + Whisper + PhonoQ variant
|   +-- sdm_whisper_data.py  # Single-frame dataset for the SDXL variant
|   +-- phonoq/              # Bundled PhonoQ feature extractor
+-- scripts/
|   +-- train.py             # Train or resume from a checkpoint
|   +-- generate.py          # Generate a video from a checkpoint and audio file
|   +-- slurm/               # Example cluster launch scripts
+-- examples/                # Placeholder inputs/outputs

Installation

git clone https://github.com/PauPerezT/PhonoDiffMRI.git
cd PhonoDiffMRI
python -m venv .venv
source .venv/bin/activate
pip install -e .

Install a CUDA-enabled PyTorch build that matches your system before installing the package if your cluster requires a specific wheel.

Expected Data

The training loader expects the same split files as the original code:

• train_new.npy, val_new.npy, and optionally test_new.npy
• each file contains a dictionary with file and label
• frame folders live under a root such as Frames_26fps/<speaker>/<video_name_without_mp4>/
• audio files live in the matching path where Frames_26fps is replaced by Audios_Denoised_16kHz and .mp4 is replaced by .wav

Train

3D Speech-To-Video Variant

Train from a local pretrained speech-to-video diffusion backbone:

python scripts/train.py \
  --pretrained-model /path/to/pretrained_speech_video_backbone \
  --frames-root /path/to/Frames_26fps \
  --splits-root /path/to/Data_Partitions/USC-75/npy \
  --speaker sub019 \
  --output-dir checkpoints/phonodiff_sub019 \
  --batch-size 1 \
  --max-epochs 1000 \
  --devices 1

Resume or fine-tune when a PhonoDiffMRI checkpoint is available:

python scripts/train.py \
  --pretrained-model /path/to/pretrained_speech_video_backbone \
  --frames-root /path/to/Frames_26fps \
  --splits-root /path/to/Data_Partitions/USC-75/npy \
  --speaker sub019 \
  --resume-from-checkpoint checkpoints/phonodiff_sub019/last.ckpt \
  --output-dir checkpoints/phonodiff_sub019_finetune

Enable Weights & Biases only when you want logging:

python scripts/train.py ... --wandb-project STDiff_3D_Phon --run-name sub019_speech_video_26fps

SDXL + Whisper Variant

This version trains a 2D SDXL denoising model over MRI frames using Whisper encoder features and optional PhonoQ posterior fusion:

python scripts/train_sdm_whisper.py \
  --pretrained-model /path/to/pretrained_sdm_whisper_backbone \
  --frames-root /path/to/Frames_26fps \
  --splits-root /path/to/Data_Partitions/USC-75/npy \
  --speaker sub019 \
  --output-dir checkpoints/sdm_whisper_sub019 \
  --batch-size 24 \
  --max-epochs 2000 \
  --devices 2

Resume or fine-tune from an available checkpoint:

python scripts/train_sdm_whisper.py \
  --pretrained-model /path/to/pretrained_sdm_whisper_backbone \
  --frames-root /path/to/Frames_26fps \
  --splits-root /path/to/Data_Partitions/USC-75/npy \
  --speaker sub019 \
  --resume-from-checkpoint checkpoints/sdm_whisper_sub019/last.ckpt \
  --output-dir checkpoints/sdm_whisper_sub019_finetune

Generate

3D Speech-To-Video Variant

Once a checkpoint is available, generate MRI video from a WAV file:

python scripts/generate.py \
  --pretrained-model /path/to/pretrained_speech_video_backbone \
  --checkpoint checkpoints/phonodiff_sub019/epoch=999-val_loss=0.11.ckpt \
  --audio examples/audios/sub019_2drt_06_rainbow_r2_video.wav \
  --output examples/outputs/sub019_rainbow.mp4 \
  --height 128 \
  --width 128 \
  --num-inference-steps 25 \
  --device cuda

If you trained without PhonoQ conditioning, pass --no-phonoq.

The 3D generation path uses the local SpeechToVideoSDPipeline in phonodiffmri.pipelines. This adapts the Diffusers video denoising loop to use speech embeddings directly, so no changes to the upstream Diffusers package are required.

You can also instantiate it directly:

from phonodiffmri.pipelines import SpeechToVideoSDPipeline

pipe = SpeechToVideoSDPipeline.from_speech_pretrained(
    "/path/to/pretrained_speech_video_backbone",
    audio_encoder_name_or_path="facebook/wav2vec2-large",
).to("cuda")

video = pipe(audio=my_waveform, num_frames=4, height=128, width=128).frames

SDXL + Whisper Variant

Generate frame-by-frame video from the SDXL + Whisper checkpoint:

python scripts/generate_sdm_whisper.py \
  --pretrained-model /path/to/pretrained_sdm_whisper_backbone \
  --checkpoint checkpoints/sdm_whisper_sub019/last.ckpt \
  --audio examples/audios/sub019_2drt_06_rainbow_r2_video.wav \
  --output examples/outputs/sub019_sdm_whisper.mp4 \
  --height 128 \
  --width 128 \
  --num-inference-steps 50 \
  --device cuda

Notes

• Large diffusion checkpoints and generated videos are ignored by Git by default.
• The PhonoQ acoustic models are included because the dataset and generation code need them to compute phonological posteriors.
• The original experiment scripts included local paths and a hard-coded W&B key; the cleaned scripts remove those assumptions.