Update app.py

app.py

@@ -11,11 +11,6 @@ from diffusers import StableDiffusionImg2ImgPipeline
 from share_btn import community_icon_html, loading_icon_html, share_js
 
 device = "cuda"
-MODEL_ID = "riffusion/riffusion-model-v1"
-pipe = StableDiffusionPipeline.from_pretrained(MODEL_ID, torch_dtype=torch.float16)
-pipe = pipe.to(device)
-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")
 
@@ -25,7 +20,10 @@ def predict(prompt, negative_prompt, audio_input, duration):
     else :
         return style_transfer(prompt, negative_prompt, audio_input)
 
-def classic(prompt, negative_prompt, duration):
+def classic(model_input, prompt, negative_prompt, duration):
+    pipe = StableDiffusionPipeline.from_pretrained(model_input, torch_dtype=torch.float16)
+    pipe = pipe.to(device)
+    
     if duration == 5:
         width_duration=512
     else :
@@ -37,23 +35,6 @@ def classic(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)
-    print(spec)
-    # Open the image
-    im = Image.open(spec)
-    
-    
-    # Open the image
-    im = image_from_spectrogram(im, 1)
-   
-    
-    new_spectro = pipe2(prompt=prompt, image=im, strength=0.5, guidance_scale=7).images
-    wav = wav_bytes_from_spectrogram_image(new_spectro[0])
-    with open("output.wav", "wb") as f:
-        f.write(wav[0].getbuffer())
-    return new_spectro[0], 'output.wav', gr.update(visible=True), gr.update(visible=True), gr.update(visible=True)
-
 def image_from_spectrogram(
     spectrogram: np.ndarray, max_volume: float = 50, power_for_image: float = 0.25
 ) -> Image.Image:
@@ -194,7 +175,7 @@ with gr.Blocks(css=css) as demo:
     with gr.Column(elem_id="col-container"):
         
         gr.HTML(title)
-        
+        model_input = gr.Textbox(placeholder="Your Riffusion fine-tuned model Hugging Face ID")
         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", visible=False)
         with gr.Row():
@@ -215,7 +196,7 @@ with gr.Blocks(css=css) as demo:
         
         gr.HTML(article)
     
-    send_btn.click(predict, inputs=[prompt_input, negative_prompt, audio_input, duration_input], outputs=[spectrogram_output, sound_output, share_button, community_icon, loading_icon])
+    send_btn.click(predict, inputs=[model_input, 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)