Update app.py

app.py

@@ -1,10 +1,19 @@
 import gradio as gr
-
+import torch, os
+import numpy as np
+from PIL import Image
+from huggingface_hub import snapshot_download
 import soundfile as sf
 from auffusion_pipeline import AuffusionPipeline
 
 pipeline = AuffusionPipeline.from_pretrained("auffusion/auffusion")
 
+# ——
+
+from diffusers import StableDiffusionImg2ImgPipeline
+from converter import load_wav, mel_spectrogram, normalize_spectrogram, denormalize_spectrogram, Generator, get_mel_spectrogram_from_audio
+from utils import pad_spec, image_add_color, torch_to_pil, normalize, denormalize
+
 def infer(prompt, progress=gr.Progress(track_tqdm=True)):
     
     prompt = prompt
@@ -14,6 +23,80 @@ def infer(prompt, progress=gr.Progress(track_tqdm=True)):
 
     return f"{prompt}.wav"
 
+def infer_img2img(prompt, audio_path):
+
+    pretrained_model_name_or_path = "auffusion/auffusion-full-no-adapter"
+    dtype = torch.float16
+    device = "cuda"
+    
+    vocoder = Generator.from_pretrained(pretrained_model_name_or_path, subfolder="vocoder")
+    vocoder = vocoder.to(device=device, dtype=dtype)
+
+    pipe = StableDiffusionImg2ImgPipeline.from_pretrained(pretrained_model_name_or_path, torch_dtype=dtype)
+    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]
+    prompt = "aumbulance siren"
+    seed = 42
+
+    # Loading
+    audio, sampling_rate = load_wav(audio_path)
+    audio, spec = get_mel_spectrogram_from_audio(audio)
+    norm_spec = normalize_spectrogram(spec)
+    norm_spec = norm_spec[:,:, width_start:width_start+width]
+    norm_spec = pad_spec(norm_spec, 1024)
+    norm_spec = normalize(norm_spec) # normalize to [-1, 1], because pipeline do not normalize for torch.Tensor input
+    
+    raw_image = image_add_color(torch_to_pil(norm_spec[:,:,:width]))
+
+    # Generation for different strength
+    image_list = []
+    audio_list = []
+    
+    generator = torch.Generator(device=device).manual_seed(seed)
+    
+    for strength in strength_list:
+        with torch.autocast("cuda"):
+            output_spec = pipe(
+                prompt=prompt, image=norm_spec, num_inference_steps=100, generator=generator, output_type="pt", strength=strength, guidance_scale=7.5
+            ).images[0]
+    
+        # add to image_list
+        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)
+    
+        # add to audio_list
+        denorm_spec = denormalize_spectrogram(output_spec)
+        denorm_spec_audio = vocoder.inference(denorm_spec)
+        audio_list.append(denorm_spec_audio)
+
+        # Display
+
+        # Concat image with different strength & add interval between images with black color 
+        concat_image_list = []
+        for i in range(len(image_list)):
+            if i == len(image_list) - 1:
+                concat_image_list.append(np.array(image_list[i]))
+            else:
+                concat_image_list.append(np.concatenate([np.array(image_list[i]), np.ones((256, 20, 3))*0], axis=1))
+        
+        concat_image = np.concatenate(concat_image_list, axis=1)
+        concat_image = Image.fromarray(np.uint8(concat_image))
+        
+        ### Concat audio
+        concat_audio_list = [np.concatenate([audio, np.zeros((1, 16000))], axis=1) for audio in audio_list]
+        concat_audio = np.concatenate(concat_audio_list, axis=1)
+        
+        print("audio_path:", audio_path)
+        print("width_start:", width_start, "width:", width)
+        print("text prompt:", prompt)
+        print("strength_list:", strength_list)
+
+        return concat_audio
+
 css="""
 div#col-container{
     margin: 0 auto;
@@ -40,24 +123,45 @@ with gr.Blocks(css=css) as demo:
             </a>
         </div>
         """)
-        
-        prompt = gr.Textbox(label="Prompt")
-        submit_btn = gr.Button("Submit")
-        audio_out = gr.Audio(label="Audio Ressult")
-
-        gr.Examples(
-            examples = [
-                "Rolling thunder with lightning strikes",
-                "Two gunshots followed by birds chirping",
-                "A train whistle blowing in the distance"
-            ],
-            inputs = [prompt]
-        )
-
-    submit_btn.click(
-        fn = infer,
-        inputs = [prompt],
-        outputs = [audio_out]
-    )
+        with gr.Tab("Text-to-Audio"):
+            prompt = gr.Textbox(label="Prompt")
+            submit_btn = gr.Button("Submit")
+            audio_out = gr.Audio(label="Audio Ressult")
+    
+            gr.Examples(
+                examples = [
+                    "Rolling thunder with lightning strikes",
+                    "Two gunshots followed by birds chirping",
+                    "A train whistle blowing in the distance"
+                ],
+                inputs = [prompt]
+            )
+
+            submit_btn.click(
+                fn = infer,
+                inputs = [prompt],
+                outputs = [audio_out]
+            )
+
+        with gr.Tab("Audio-to-Audio"):
+            prompt_img2img = gr.Textbox(label="Prompt")
+            audio_in_img2img = gr.Audio(label="Audio Reference", type="filepath")
+            submit_btn_img2img = gr.Button("Submit")
+            audio_out_img2img = gr.Audio(label="Audio Ressult")
+    
+            gr.Examples(
+                examples = [
+                    "Rolling thunder with lightning strikes",
+                    "Two gunshots followed by birds chirping",
+                    "A train whistle blowing in the distance"
+                ],
+                inputs = [prompt_img2img]
+            )
+
+            submit_btn_img2img.click(
+                fn = infer_img2img,
+                inputs = [prompt_img2img],
+                outputs = [audio_out_img2img]
+            )
 
 demo.queue().launch(show_api=False, show_error=True)