Update app.py

app.py

@@ -52,8 +52,8 @@ def uniform_sample(idxs, n):
     gap = len(idxs) / n
     return [idxs[int(i * gap + gap / 2)] for i in range(n)]
 
-def build_omni_chunks(path, sr=16000, seconds_per_unit=1):
-    clip = VideoFileClip(path)
+def build_omni_chunks(path, prefix, instruction, sr=AUDIO_SR, seconds_per_unit=1):
+    clip = VideoFileClip(path, audio_fps=sr)
     with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as tmp:
         clip.audio.write_audiofile(tmp.name, fps=sr, codec="pcm_s16le", verbose=False, logger=None)
         audio_np, _ = librosa.load(tmp.name, sr=sr, mono=True)
@@ -61,19 +61,17 @@ def build_omni_chunks(path, sr=16000, seconds_per_unit=1):
     content = []
     for i in range(total_units):
         t = min(i * seconds_per_unit, clip.duration - 1e-3)
-        frame = Image.fromarray(clip.get_frame(t).astype("uint8"))
+        frame = Image.fromarray(clip.get_frame(t).astype("uint8")).convert("RGB")
         audio_chunk = audio_np[sr * i * seconds_per_unit : sr * (i + 1) * seconds_per_unit]
         content.extend(["<unit>", frame, audio_chunk])
+    clip.close()
+    os.remove(tmp.name)
+    content.append(prefix + instruction)
     return content
-    
-def encode_video(path):
-    vr = VideoReader(path, ctx=cpu(0))
-    fps = round(vr.get_avg_fps())
-    idxs = list(range(0, len(vr), fps))
-    if len(idxs) > MAX_FRAMES:
-        idxs = uniform_sample(idxs, MAX_FRAMES)
-    frames = vr.get_batch(idxs).asnumpy()
-    return [Image.fromarray(f.astype("uint8")) for f in frames]
+
+def build_image_omni(path, prefix, instruction):
+    image = Image.open(path).convert("RGB")
+    return ["<unit>", image, prefix + instruction]
 
 def encode_gif(path):
     img = Image.open(path)
@@ -81,59 +79,44 @@ def encode_gif(path):
     if len(frames) > MAX_FRAMES:
         frames = uniform_sample(frames, MAX_FRAMES)
     return frames
-    
-@spaces.GPU(duration=60)
-def generate(input, instruction=DEFAULT_INPUT, sampling=False, temperature=0.7, top_p=0.8, top_k=100, repetition_penalty=1.05, max_tokens=512):
-    print(input)
-    print(instruction)
 
-    if not input:
-        return "No input provided."
+def build_gif_omni(path, prefix, instruction):
+    frames = encode_gif(path)
+    content = []
+    for f in frames:
+        content.extend(["<unit>", f])
+    content.append(prefix + instruction)
+    return content
 
-    extension = os.path.splitext(input)[1].lower()
-    filetype = None
-    for category, extensions in filetypes.items():
-        if extension in extensions:
-            filetype = category
-            break
+def build_audio_omni(path, prefix, instruction, sr=AUDIO_SR):
+    audio_np, _ = librosa.load(path, sr=sr, mono=True)
+    return ["<unit>", audio_np, prefix + instruction]
 
-    content = []
-    if filetype == "Image":
-        image = Image.open(input).convert("RGB")
-        content.append(image)
-    elif filetype == "GIF":
-        frames = encode_gif(input)
-        content.extend(frames)
-    elif filetype == "Video":
-        omni_content = build_omni_chunks(input) + [instruction]
-        sys_msg = repo.get_sys_prompt(mode="omni", language="en")
-        msgs = [sys_msg, {"role": "user", "content": omni_content}]
-        print(msgs)
-    elif filetype == "Audio":
-        audio_np, sample_rate = librosa.load(input, sr=16000, mono=True)
-        chunk_tensor = torch.from_numpy(audio_np).float().to(DEVICE)
-        content.append({"array": chunk_tensor, "sampling_rate": sample_rate})
-  
-    """
-    elif filetype == "Video":
-        frames = encode_video(input)
-        content.extend(frames)
-        audio, _ = librosa.load(input, sr=16000, mono=True)
-        content.append(audio)
-    elif filetype == "Audio":
-        audio, _ = librosa.load(input, sr=16000, mono=True)
-        content.append(audio)
-    else:
-        return "Unsupported file type."
-    """
+@spaces.GPU(duration=60)
+def generate(input, instruction=DEFAULT_INPUT, sampling=False, temperature=0.7, top_p=0.8, top_k=100, repetition_penalty=1.05, max_tokens=512):
+    if not input: return "No input provided."
+    
+    extension = os.path.splitext(input)[1].lower()
+    filetype = next((k for k, v in filetypes.items() if extension in v), None)
+    
+    if not filetype: return "Unsupported file type."
     
     filename = os.path.basename(input)
     prefix = input_prefixes[filetype].replace("█", filename)
-    content.append(prefix + instruction)
-    inputs_payload = [{"role": "user", "content": content}]
-
+    if filetype == "Video":
+        omni_content = build_omni_chunks(input, prefix, instruction)
+    elif filetype == "Image":
+        omni_content = build_image_omni(input, prefix, instruction)
+    elif filetype == "GIF":
+        omni_content = build_gif_omni(input, prefix, instruction)
+    elif filetype == "Audio":
+        omni_content = build_audio_omni(input, prefix, instruction)
+    
+    sys_msg = repo.get_sys_prompt(mode="omni", language="en")
+    msgs = [sys_msg, {"role": "user", "content": omni_content}]
+    
     params = {
-        "msgs": msgs or inputs_payload,
+        "msgs": msgs,
         "tokenizer": tokenizer,
         "sampling": sampling,
         "temperature": temperature,
@@ -141,13 +124,12 @@ def generate(input, instruction=DEFAULT_INPUT, sampling=False, temperature=0.7,
         "top_k": top_k,
         "repetition_penalty": repetition_penalty,
         "max_new_tokens": max_tokens,
-        "omni_input": filetype == "Video",
+        "omni_input": True,
     }
-
-    output = repo.chat(**params)
     
-    print(output)
-
+    output = repo.chat(**params)
+    torch.cuda.empty_cache()
+    gc.collect()
     return output
 
 def cloud():