Update app.py

app.py

@@ -8,11 +8,12 @@ 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("cuda")
+pipe = pipe.to(device)
 pipe2 = StableDiffusionImg2ImgPipeline.from_pretrained(MODEL_ID, torch_dtype=torch.float16)
-pipe2 = pipe2.to("cuda")
+pipe2 = pipe2.to(device)
 
 spectro_from_wav = gr.Interface.load("spaces/fffiloni/audio-to-spectrogram")
 
@@ -36,7 +37,7 @@ def classic(prompt, negative_prompt, duration):
 
 def style_transfer(prompt, negative_prompt, audio_input):
     spec = spectro_from_wav(audio_input)
-    new_spectro = pipe2(prompt=prompt, image=spec, strength=0.5, guidance_scale=7).images
+    new_spectro = pipe2(prompt=prompt, init_image=spec, 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())