Update app.py

app.py

@@ -24,6 +24,7 @@ from vocoder.bigvgan.models import VocoderBigVGAN
 import soundfile
 # from pytorch_memlab import LineProfiler,profile
 import gradio
+import gradio as gr
 
 def load_model_from_config(config, ckpt = None, verbose=True):
     model = instantiate_from_config(config.model)
@@ -50,7 +51,7 @@ def load_model_from_config(config, ckpt = None, verbose=True):
 
 
 class GenSamples:
-    def __init__(self,sampler,model,outpath,vocoder = None,save_mel = True,save_wav = True, original_inference_steps=None) -> None:
+    def __init__(self,sampler,model,outpath,vocoder = None,save_mel = True,save_wav = True, original_inference_steps=None, ddim_steps=2, scale=5, num_samples=1) -> None:
         self.sampler = sampler
         self.model = model
         self.outpath = outpath
@@ -61,29 +62,33 @@ class GenSamples:
         self.save_wav = save_wav
         self.channel_dim = self.model.channels
         self.original_inference_steps = original_inference_steps
+        self.ddim_steps = ddim_steps
+        self.scale = scale
+        self.num_samples = num_samples
     
     def gen_test_sample(self,prompt,mel_name = None,wav_name = None):# prompt is {'ori_caption':’xxx‘,'struct_caption':'xxx'}
         uc = None
         record_dicts = []
         # if os.path.exists(os.path.join(self.outpath,mel_name+f'_0.npy')):
         #     return record_dicts
-        emptycap = {'ori_caption':1*[""],'struct_caption':1*[""]}
-        uc = self.model.get_learned_conditioning(emptycap)
+        if self.scale != 1.0:
+            emptycap = {'ori_caption':self.num_samples*[""],'struct_caption':self.num_samples*[""]}
+            uc = self.model.get_learned_conditioning(emptycap)
 
         for n in range(1):# trange(self.opt.n_iter, desc="Sampling"):
             for k,v in prompt.items():
-                prompt[k] = 1 * [v]
+                prompt[k] = self.num_samples * [v]
             c = self.model.get_learned_conditioning(prompt)# shape:[1,77,1280],即还没有变成句子embedding，仍是每个单词的embedding
             if self.channel_dim>0:
                 shape = [self.channel_dim, 20, 312]  # (z_dim, 80//2^x, 848//2^x)
             else:
                 shape = [20, 312]
-            samples_ddim, _ = self.sampler.sample(S=2,
+            samples_ddim, _ = self.sampler.sample(S=self.ddim_steps,
                                                 conditioning=c,
-                                                batch_size=1,
+                                                batch_size=self.num_samples,
                                                 shape=shape,
                                                 verbose=False,
-                                                guidance_scale=5,
+                                                guidance_scale=self.scale,
                                                 original_inference_steps=self.original_inference_steps
                                                 )
             x_samples_ddim = self.model.decode_first_stage(samples_ddim)
@@ -103,7 +108,9 @@ class GenSamples:
         return record_dicts
 
 @spaces.GPU(enable_queue=True)
-def infer(ori_prompt):
+def infer(ori_prompt, ddim_steps, num_samples, scale, seed):
+    np.random.seed(seed)
+    torch.manual_seed(seed)
 
     prompt = dict(ori_caption=ori_prompt,struct_caption=f'<{ori_prompt}& all>')
 
@@ -124,7 +131,7 @@ def infer(ori_prompt):
     vocoder = VocoderBigVGAN("./model/vocoder",device)
 
 
-    generator = GenSamples(sampler,model,"results/test",vocoder,save_mel = False,save_wav = True, original_inference_steps=config.model.params.num_ddim_timesteps)
+    generator = GenSamples(sampler,model,"results/test",vocoder,save_mel = False,save_wav = True, original_inference_steps=config.model.params.num_ddim_timesteps, ddim_steps=ddim_steps, scale=scale, num_samples=num_samples)
     csv_dicts = []
 
     with torch.no_grad():
@@ -135,15 +142,61 @@ def infer(ori_prompt):
     print(f"Your samples are ready and waiting four you here: \nresults/test \nEnjoy.")
     return "results/test/"+wav_name+"_0.wav"
 
-def my_inference_function(text_prompt):
-    file_path = infer(text_prompt)
+def my_inference_function(text_prompt, ddim_steps, num_samples, scale, seed):
+    file_path = infer(text_prompt, ddim_steps, num_samples, scale, seed)
     return file_path
 
 
-
-gradio_interface = gradio.Interface(
-    fn = my_inference_function,
-    inputs = "text",
-    outputs = "audio"
-)
-gradio_interface.launch()
+with gr.Blocks() as demo:
+    with gr.Row():
+        tgr.Markdown("## AudioLCM：Text-to-Audio Generation with Latent Consistency Models")
+
+    with gr.Row():
+        with gr.Column():
+            prompt = gr.Textbox(label="Prompt: Input your text here.        ")
+            run_button = gr.Button(label="Run")
+
+            with gr.Accordion("Advanced options", open=False):
+                num_samples = gr.Slider(
+                    label="Select from audios num.This number control the number of candidates \
+                        (e.g., generate three audios and choose the best to show you). A Larger value usually lead to \
+                        better quality with heavier computation", minimum=1, maximum=10, value=1, step=1)
+                # num_samples = 1
+                ddim_steps = gr.Slider(label="Steps", minimum=1,
+                                       maximum=150, value=2, step=1)
+                scale = gr.Slider(
+                    label="Guidance Scale:(Large => more relevant to text but the quality may drop)", minimum=0.1, maximum=8.0, value=5.0, step=0.1
+                )
+                seed = gr.Slider(
+                    label="Seed:Change this value (any integer number) will lead to a different generation result.",
+                    minimum=0,
+                    maximum=2147483647,
+                    step=1,
+                    value=44,
+                )
+
+        with gr.Column():
+            outaudio = gr.Audio()
+    
+    run_button.click(fn=my_inference_function, inputs=[
+                    prompt,ddim_steps, num_samples, scale, seed], outputs=[outaudio])
+    with gr.Row():
+        with gr.Column():
+            gr.Examples(
+                        examples = [['a dog barking and a bird chirping',100,3,3,55],['Pigeons peck, coo, and flap their wings before a man speaks',100,3,3,55],
+                                        ['music of violin and piano',100,3,2,88],['wind thunder and rain falling',100,3,3,55],['music made by drum kit',100,3,3,55]],
+                        inputs = [prompt,ddim_steps, num_samples, scale, seed],
+                        outputs = [outaudio]
+                        )
+        with gr.Column():
+            pass
+
+demo.launch()
+
+
+# gradio_interface = gradio.Interface(
+#     fn = my_inference_function,
+#     inputs = "text",
+#     outputs = "audio"
+# )
+# gradio_interface.launch()