try addind duration control

app.py

@@ -9,10 +9,10 @@ model_id = "riffusion/riffusion-model-v1"
 pipe = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float16)
 pipe = pipe.to("cuda")
 
-def predict(prompt):
+def predict(prompt, sample_duration):
     spec = pipe(prompt).images[0]
     print(spec)
-    wav = wav_bytes_from_spectrogram_image(spec)
+    wav = wav_bytes_from_spectrogram_image(spec, sample_duration)
     with open("output.wav", "wb") as f:
         f.write(wav[0].getbuffer())
     return spec, 'output.wav'
@@ -102,7 +102,7 @@ with gr.Blocks(css=css) as demo:
         gr.HTML(title)
         
         prompt_input = gr.Textbox(placeholder="a cat diva singing in a New York jazz club")
-  
+        sample_duration_input = gr.Slider(minimum=5, maximum=20, value=5, step=5)
         send_btn = gr.Button("Get a new spectrogram ! ")
             
     with gr.Column(elem_id="col-container-2"):
@@ -111,6 +111,6 @@ with gr.Blocks(css=css) as demo:
 
         gr.HTML(article)
     
-    send_btn.click(predict, inputs=[prompt_input], outputs=[spectrogram_output, sound_output])
+    send_btn.click(predict, inputs=[prompt_input, sample_duration_input], outputs=[spectrogram_output, sound_output])
 
 demo.queue(max_size=250).launch(debug=True)

spectro.py

@@ -12,7 +12,7 @@ import torch
 import torchaudio
 
 
-def wav_bytes_from_spectrogram_image(image: Image.Image) -> T.Tuple[io.BytesIO, float]:
+def wav_bytes_from_spectrogram_image(image: Image.Image, sample_duration) -> T.Tuple[io.BytesIO, float]:
     """
     Reconstruct a WAV audio clip from a spectrogram image. Also returns the duration in seconds.
     """
@@ -22,7 +22,7 @@ def wav_bytes_from_spectrogram_image(image: Image.Image) -> T.Tuple[io.BytesIO,
     Sxx = spectrogram_from_image(image, max_volume=max_volume, power_for_image=power_for_image)
 
     sample_rate = 44100  # [Hz]
-    clip_duration_ms = 5000  # [ms]
+    clip_duration_ms = sample_duration*1000  # [ms]
 
     bins_per_image = 512
     n_mels = 512