From ce8744897bab116eda23a1937b55a13491a41b9f Mon Sep 17 00:00:00 2001
From: Hayk Martiros <hayk.mart@gmail.com>
Date: Thu, 29 Dec 2022 11:27:42 -0800
Subject: [PATCH] Convert to three-channel if needed in spectrogram_from_image

Closes: #51

Topic: convert_three_channel
---
 riffusion/util/image_util.py | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/riffusion/util/image_util.py b/riffusion/util/image_util.py
index 22e42fb..bdec25e 100644
--- a/riffusion/util/image_util.py
+++ b/riffusion/util/image_util.py
@@ -77,6 +77,10 @@ def spectrogram_from_image(
     Returns:
         spectrogram: (channels, frequency, time)
     """
+    # Convert to RGB if single channel
+    if image.mode in ("P", "L"):
+        image = image.convert("RGB")
+
     # Flip Y
     image = image.transpose(Image.Transpose.FLIP_TOP_BOTTOM)