Update app.py

app.py

@@ -20,6 +20,7 @@ import torchvision.transforms as transforms
 import av
 import subprocess
 import librosa
+import uuid
 
 args = {"model": "./ckpts/checkpoint.pth", "llama_type": "7B", "llama_dir": "./ckpts/LLaMA-2",
         "mert_path": "m-a-p/MERT-v1-330M", "vit_path": "google/vit-base-patch16-224", "vivit_path": "google/vivit-b-16x2-kinetics400",
@@ -33,7 +34,7 @@ class dotdict(dict):
 
 args = dotdict(args)
 
-generated_audio_files = []
+generated_audio_files = {}
 
 llama_type = args.llama_type
 llama_ckpt_dir = os.path.join(args.llama_dir, llama_type)
@@ -117,7 +118,7 @@ def parse_text(text, image_path, video_path, audio_path):
     return text, outputs
 
 
-def save_audio_to_local(audio, sec):
+def save_audio_to_local(uid, audio, sec):
     global generated_audio_files
     if not os.path.exists('temp'):
         os.mkdir('temp')
@@ -126,11 +127,11 @@ def save_audio_to_local(audio, sec):
         scipy.io.wavfile.write(filename, rate=16000, data=audio[0])
     else:
         scipy.io.wavfile.write(filename, rate=model.generation_model.config.audio_encoder.sampling_rate, data=audio)
-    generated_audio_files.append(filename)
+    generated_audio_files[uid].append(filename)
     return filename
 
 
-def parse_reponse(model_outputs, audio_length_in_s):
+def parse_reponse(uid, model_outputs, audio_length_in_s):
     response = ''
     text_outputs = []
     for output_i, p in enumerate(model_outputs):
@@ -146,7 +147,7 @@ def parse_reponse(model_outputs, audio_length_in_s):
                     response += '<br>'
                     _temp_output += m.replace(' '.join([f'[AUD{i}]' for i in range(8)]), '')
                 else:
-                    filename = save_audio_to_local(m, audio_length_in_s)
+                    filename = save_audio_to_local(uid, m, audio_length_in_s)
                     print(filename)
                     _temp_output = f'<Audio>{filename}</Audio> ' + _temp_output
                     response += f'<audio controls playsinline><source src="./file={filename}" type="audio/wav"></audio>'
@@ -161,15 +162,15 @@ def reset_user_input():
     return gr.update(value='')
 
 
-def reset_dialog():
+def reset_dialog(uid):
     global generated_audio_files
-    generated_audio_files = []
+    generated_audio_files[uid] = []
     return [], []
 
 
-def reset_state():
+def reset_state(uid):
     global generated_audio_files
-    generated_audio_files = []
+    generated_audio_files[uid] = []
     return None, None, None, None, [], [], []
 
 
@@ -218,6 +219,7 @@ def get_audio_length(filename):
 
 
 def predict(
+        uid,
         prompt_input,
         image_path,
         audio_path,
@@ -247,28 +249,30 @@ def predict(
         indices = sample_frame_indices(clip_len=32, frame_sample_rate=1, seg_len=container.streams.video[0].frames)
         video = read_video_pyav(container=container, indices=indices)
 
-    if len(generated_audio_files) != 0:
-        audio_length_in_s = get_audio_length(generated_audio_files[-1])
+    if uid in generated_audio_files and len(generated_audio_files[uid]) != 0:
+        audio_length_in_s = get_audio_length(generated_audio_files[uid][-1])
         sample_rate = 24000
-        waveform, sr = torchaudio.load(generated_audio_files[-1])
+        waveform, sr = torchaudio.load(generated_audio_files[uid][-1])
         if sample_rate != sr:
             waveform = torchaudio.functional.resample(waveform, orig_freq=sr, new_freq=sample_rate)
         audio = torch.mean(waveform, 0)
         audio_length_in_s = int(len(audio)//sample_rate)
         print(f"Audio Length: {audio_length_in_s}")
+    else:
+        generated_audio_files[uid] = []
     if video_path is not None:
         audio_length_in_s = get_video_length(video_path)
         print(f"Video Length: {audio_length_in_s}")
     if audio_path is not None:
         audio_length_in_s = get_audio_length(audio_path)
-        generated_audio_files.append(audio_path)
+        generated_audio_files[uid].append(audio_path)
         print(f"Audio Length: {audio_length_in_s}")
 
     print(image, video, audio)
     response = model.generate(prompts, audio, image, video, 200, temperature, top_p,
                               audio_length_in_s=audio_length_in_s)
     print(response)
-    response_chat, response_outputs = parse_reponse(response, audio_length_in_s)
+    response_chat, response_outputs = parse_reponse(uid, response, audio_length_in_s)
     print('text_outputs: ', response_outputs)
     user_chat, user_outputs = parse_text(prompt_input, image_path, video_path, audio_path)
     chatbot.append((user_chat, response_chat))
@@ -319,9 +323,11 @@ with gr.Blocks() as demo:
 
     history = gr.State([])
     modality_cache = gr.State([])
+    uid = gr.State(uuid.uuid4())
 
     submitBtn.click(
         predict, [
+            uid,
             user_input,
             image_path,
             audio_path,
@@ -343,8 +349,8 @@ with gr.Blocks() as demo:
         show_progress=True
     )
 
-    submitBtn.click(reset_user_input, [], [user_input])
-    emptyBtn.click(reset_state, outputs=[
+    submitBtn.click(reset_user_input, [uid], [user_input])
+    emptyBtn.click(reset_state, [uid], outputs=[
         image_path,
         audio_path,
         video_path,