Preperation for audio

toolkit/config_modules.py

@@ -4,6 +4,7 @@ from typing import List, Optional, Literal, Tuple, Union, TYPE_CHECKING, Dict
 import random
 
 import torch
+import torchaudio
 
 from toolkit.prompt_utils import PromptEmbeds
 
@@ -1073,6 +1074,15 @@ class GenerateImageConfig:
                 )
             else:
                 raise ValueError(f"Unsupported video format {self.output_ext}")
+        elif self.output_ext in ['wav', 'mp3']:
+            # save audio file
+            torchaudio.save(
+                self.get_image_path(count, max_count), 
+                image[0].to('cpu'),
+                sample_rate=48000, 
+                format=None, 
+                backend=None
+            )
         else:
             # TODO save image gen header info for A1111 and us, our seeds probably wont match
             image.save(self.get_image_path(count, max_count))