Update app.py

app.py

@@ -23,7 +23,7 @@ def infer(prompt, progress=gr.Progress(track_tqdm=True)):
 
     return f"{prompt}.wav"
 
-def infer_img2img(prompt, audio_path, progress=gr.Progress(track_tqdm=True)):
+def infer_img2img(prompt, audio_path, desired_strength, progress=gr.Progress(track_tqdm=True)):
 
     pretrained_model_name_or_path = "auffusion/auffusion-full-no-adapter"
     dtype = torch.float16
@@ -39,7 +39,7 @@ def infer_img2img(prompt, audio_path, progress=gr.Progress(track_tqdm=True)):
     pipe = pipe.to(device)
 
     width_start, width = 0, 160
-    strength_list = [0.0, 0.1, 0.2, 0.3, 0.5, 0.6, 0.7]
+    strength_list = [desired_strength]
     prompt = "ambulance siren"
     seed = 42
 
@@ -66,7 +66,7 @@ def infer_img2img(prompt, audio_path, progress=gr.Progress(track_tqdm=True)):
             ).images[0]
     
         # add to image_list
-        output_spec = output_spec[:, :, :width]
+        # output_spec = output_spec[:, :, :width]
         output_spec_image = torch_to_pil(output_spec)
         color_output_spec_image = image_add_color(output_spec_image)
         image_list.append(color_output_spec_image)
@@ -158,6 +158,7 @@ with gr.Blocks(css=css) as demo:
         with gr.Tab("Audio-to-Audio"):
             prompt_img2img = gr.Textbox(label="Prompt")
             audio_in_img2img = gr.Audio(label="Audio Reference", type="filepath")
+            prompt_strength = gr.Slider(label="Prompt Strength", minimum=0.0, maximum=1.0, step=0.1, value=0.7)
             submit_btn_img2img = gr.Button("Submit")
             audio_out_img2img = gr.Audio(label="Audio Ressult")
     
@@ -172,7 +173,7 @@ with gr.Blocks(css=css) as demo:
 
             submit_btn_img2img.click(
                 fn = infer_img2img,
-                inputs = [prompt_img2img, audio_in_img2img],
+                inputs = [prompt_img2img, audio_in_img2img, prompt_strength],
                 outputs = [audio_out_img2img]
             )