Update app.py

app.py

@@ -1,11 +1,7 @@
 import gradio as gr
 import torch
 
-from scipy.io import wavfile
-import numpy as np
-from PIL import Image
-
-from spectro import wav_bytes_from_spectrogram_image, spectrogram_from_waveform, image_from_spectrogram
+from spectro import wav_bytes_from_spectrogram_image
 
 from diffusers import StableDiffusionPipeline
 from diffusers import StableDiffusionImg2ImgPipeline
@@ -15,63 +11,20 @@ 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("cuda")
-
-def predict(prompt, audio_input, duration):
-    if audio_input == None:
-        return classic(prompt, duration)
-    else:
-        return audio_transfer(prompt, audio_input)
-    
 
-def classic(prompt, duration):
+
+def predict(prompt, negative_prompt, duration):
     if duration == 5:
         width_duration=512
     else :
         width_duration = 512 + ((int(duration)-5) * 128)
-    spec = pipe(prompt, height=512, width=width_duration).images[0]
+    spec = pipe(prompt, negative_prompt=negative_prompt, height=512, width=width_duration).images[0]
     print(spec)
     wav = wav_bytes_from_spectrogram_image(spec)
     with open("output.wav", "wb") as f:
         f.write(wav[0].getbuffer())
     return spec, 'output.wav', gr.update(visible=True), gr.update(visible=True), gr.update(visible=True)
 
-def audio_transfer(prompt, audio):
-    # read uploaded file to wav
-    rate, data = wavfile.read(audio)
-    
-    # convert to mono
-    data = np.mean(data, axis=0)
-    
-    # convert to float32
-    data = data.astype(np.float32)
-    
-    # take a random 7 second slice of the audio
-    data = data[rate*7:rate*14]
-    
-    spectrogram = spectrogram_from_waveform(
-        waveform=data,
-        sample_rate=rate,
-        # width=768,
-        n_fft=8192,
-        hop_length=512,
-        win_length=8192,
-    )
-
-    spec = image_from_spectrogram(spectrogram)
-
-    images = pipe2(
-        prompt=prompt,
-        image=spec,
-        strength=0.5,
-        guidance_scale=7
-    ).images
-
-    wav = wav_bytes_from_spectrogram_image(images[0])
-    with open("output.wav", "wb") as f:
-        f.write(wav[0].getbuffer())
-    return images[0], '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;">
@@ -189,8 +142,10 @@ 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(label="audio input", type="filepath", source="upload")
-        duration_input = gr.Slider(label="Duration in seconds", minimum=5, maximum=10, step=1, value=8, elem_id="duration-slider")
+        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")
+            
         send_btn = gr.Button(value="Get a new spectrogram ! ", elem_id="submit-btn")
             
     with gr.Column(elem_id="col-container-2"):
@@ -205,7 +160,7 @@ with gr.Blocks(css=css) as demo:
         
         gr.HTML(article)
     
-    send_btn.click(predict, inputs=[prompt_input, audio_input, duration_input], outputs=[spectrogram_output, sound_output, share_button, community_icon, loading_icon])
+    send_btn.click(predict, inputs=[prompt_input, negative_prompt, 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)