Added audio_format param to make methods in export_utils generic

Author: prishajain1

src/maxdiffusion/configs/ltx2_video.yml

@@ -91,3 +91,4 @@ qwix_module_path: ".*"
 jit_initializers: True 
 enable_single_replica_ckpt_restoring: False
 seed: 0
+audio_format: "s16"

src/maxdiffusion/generate_ltx2.py

@@ -164,8 +164,15 @@ def run(config, pipeline=None, filename_prefix="", commit_hash=None):
     video_path = f"{filename_prefix}ltx2_output_{getattr(config, 'seed', 0)}_{i}.mp4"
     audio_i = audios[i] if audios is not None else None
 
+    audio_format = getattr(config, "audio_format", "s16")
+
     export_to_video_with_audio(
-        video=videos[i], fps=fps, audio=audio_i, audio_sample_rate=audio_sample_rate, output_path=video_path
+        video=videos[i],
+        fps=fps,
+        audio=audio_i,
+        audio_sample_rate=audio_sample_rate,
+        output_path=video_path,
+        audio_format=audio_format,
     )
 
     saved_video_path.append(video_path)

src/maxdiffusion/utils/export_utils.py

@@ -271,6 +271,7 @@ def _write_audio(
     audio_stream,
     samples: Any,
     audio_sample_rate: int,
+    target_format: str = "s16",
 ) -> None:
   import numpy as np
 
@@ -286,14 +287,27 @@ def _write_audio(
   if samples.shape[1] != 2:
     raise ValueError(f"Expected samples with 2 channels; got shape {samples.shape}.")
 
-  if samples.dtype != np.int16:
-    samples = np.clip(samples, -1.0, 1.0)
-    samples = (samples * 32767.0).astype(np.int16)
+  if target_format == "s16":
+    if samples.dtype != np.int16:
+      samples = np.clip(samples, -1.0, 1.0)
+      samples = (samples * 32767.0).astype(np.int16)
+  elif target_format == "s32":
+    if samples.dtype != np.int32:
+      samples = np.clip(samples, -1.0, 1.0)
+      samples = (samples * 2147483647.0).astype(np.int32)
+  elif target_format in ["flt", "dbl", "fltp", "dblp"]:
+    target_dtype = np.float32 if "flt" in target_format else np.float64
+    if samples.dtype != target_dtype:
+      samples = samples.astype(target_dtype)
+  else:
+    # Fallback to clip and scaling for other int formats if they were added, but raise for now
+    raise ValueError(f"Unsupported target_format for converting numpy array: {target_format}")
+
   samples_np = np.ascontiguousarray(samples).reshape(1, -1)
 
   frame_in = av.AudioFrame.from_ndarray(
       samples_np,
-      format="s16",
+      format=target_format,
       layout="stereo",
   )
   frame_in.sample_rate = audio_sample_rate
@@ -302,7 +316,7 @@ def _write_audio(
 
 
 def export_to_video_with_audio(
-    video: Any, fps: int, audio: Optional[Any], audio_sample_rate: Optional[int], output_path: str
+    video: Any, fps: int, audio: Optional[Any], audio_sample_rate: Optional[int], output_path: str, audio_format: str = "s16"
 ) -> None:
   """
   Encodes video (and optionally audio) to a file using PyAV.
@@ -351,6 +365,6 @@ def export_to_video_with_audio(
     container.mux(packet)
 
   if audio is not None:
-    _write_audio(container, audio_stream, audio, audio_sample_rate)
+    _write_audio(container, audio_stream, audio, audio_sample_rate, target_format=audio_format)
 
   container.close()