Update app.py

app.py

@@ -11,9 +11,18 @@ from share_btn import community_icon_html, loading_icon_html, share_js
 MODEL_ID = "riffusion/riffusion-model-v1"
 pipe = StableDiffusionPipeline.from_pretrained(MODEL_ID, torch_dtype=torch.float16)
 pipe = pipe.to("cuda")
+pipe2 = StableDiffusionImg2ImgPipeline.from_pretrained(MODEL_ID, torch_dtype=torch.float16)
+pipe2 = pipe2.to(device)
 
+spectro_from_wav = gr.Interface.load("spaces/fffiloni/audio-to-spectrogram")
 
-def predict(prompt, negative_prompt, duration):
+def predict(prompt, negative_prompt, audio_input, duration):
+    if audio_input == None :
+        return classic(prompt, negative_prompt, duration)
+    else :
+        return style_transfer(prompt, negative_prompt, audio_input)
+
+def classic(prompt, negative_prompt, duration):
     if duration == 5:
         width_duration=512
     else :
@@ -25,6 +34,13 @@ def predict(prompt, negative_prompt, duration):
         f.write(wav[0].getbuffer())
     return spec, 'output.wav', gr.update(visible=True), gr.update(visible=True), gr.update(visible=True)
 
+def style_transfer(prompt, negative_prompt, audio_input):
+    spec = spectro_from_wav(audio_input)
+    new_spectro = pipe(prompt=prompt, image=spec, strength=0.5, guidance_scale=7).images[0]
+    wav = wav_bytes_from_spectrogram_image(spec)
+    with open("output.wav", "wb") as f:
+        f.write(wav[0].getbuffer())
+    return new_spectro, 'output.wav', gr.update(visible=True), gr.update(visible=True), gr.update(visible=True)
 
 title = """
     <div style="text-align: center; max-width: 500px; margin: 0 auto;">
@@ -142,6 +158,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")
+        audio_input = gr.Audio(source="upload", type="filepath")
         with gr.Row():
             negative_prompt = gr.Textbox(label="Negative prompt")
             duration_input = gr.Slider(label="Duration in seconds", minimum=5, maximum=10, step=1, value=8, elem_id="duration-slider")
@@ -160,7 +177,7 @@ with gr.Blocks(css=css) as demo:
         
         gr.HTML(article)
     
-    send_btn.click(predict, inputs=[prompt_input, negative_prompt, duration_input], outputs=[spectrogram_output, sound_output, share_button, community_icon, loading_icon])
+    send_btn.click(predict, inputs=[prompt_input, negative_prompt, audio_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)