added duration option

app.py

@@ -11,8 +11,12 @@ model_id = "riffusion/riffusion-model-v1"
 pipe = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float16)
 pipe = pipe.to("cuda")
 
-def predict(prompt):
-    spec = pipe(prompt, height=512, width=768).images[0]
+def predict(prompt, duration):
+    if duration == 5:
+        width_duration=512
+    else :
+        width_duration = int(512 + ((duration-5) * 128))
+    spec = pipe(prompt, height=512, width=width_duration).images[0]
     print(spec)
     wav = wav_bytes_from_spectrogram_image(spec)
     with open("output.wav", "wb") as f:
@@ -132,6 +136,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", label="Musical prompt", elem_id="prompt-in")
+        duration_input = gr.Slider(label="Duration in seconds", minimum=5, maximum=20, step=1, value=8, elem_id="duration-slider")
         send_btn = gr.Button(value="Get a new spectrogram ! ", elem_id="submit-btn")
             
     with gr.Column(elem_id="col-container-2"):
@@ -146,7 +151,7 @@ with gr.Blocks(css=css) as demo:
         
         gr.HTML(article)
     
-    send_btn.click(predict, inputs=[prompt_input], outputs=[spectrogram_output, sound_output, share_button, community_icon, loading_icon])
+    send_btn.click(predict, inputs=[prompt_input, duration_input], outputs=[spectrogram_output, sound_output, share_button, community_icon, loading_icon])
     share_button.click(None, [], [], _js=share_js)
 
 demo.queue(max_size=250).launch(debug=True)