Update app.py

app.py

@@ -1,18 +1,78 @@
 import gradio as gr
 from huggingface_hub import InferenceClient
 import json
-import re
 import uuid
 from PIL import Image
 from bs4 import BeautifulSoup
 import requests
 import random
-from gradio_client import Client, file
+from transformers import LlavaProcessor, LlavaForConditionalGeneration, TextIteratorStreamer
+from threading import Thread
+import re
+import time 
+import torch
+import cv2
+
+model_id = "llava-hf/llava-interleave-qwen-0.5b-hf"
+
+processor = LlavaProcessor.from_pretrained(model_id)
+
+model = LlavaForConditionalGeneration.from_pretrained(model_id, low_cpu_mem_usage=True)
+model.to("cpu")
+
+
+def sample_frames(video_file) :
+    try:
+        video = cv2.VideoCapture(video_file)
+        total_frames = int(video.get(cv2.CAP_PROP_FRAME_COUNT))
+        num_frames = 12
+        interval = total_frames // num_frames
+        frames = []
+        for i in range(total_frames):
+            ret, frame = video.read()
+            pil_img = Image.fromarray(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))
+            if not ret:
+                continue
+            if i % interval == 0:
+                frames.append(pil_img)
+        video.release()
+        return frames
+    except:
+        frames=[]
+        return frames
+
+def llava(user_prompt, history):
+    image = user_prompt["files"][-1]    
+    txt = user_prompt["text"]
+    img = user_prompt["files"]
+    
+    video_extensions = ("avi", "mp4", "mov", "mkv", "flv", "wmv", "mjpeg", "wav", "gif", "webm", "m4v", "3gp")
+    image_extensions = Image.registered_extensions()
+    image_extensions = tuple([ex for ex, f in image_extensions.items()])
+        
+    if image.endswith(video_extensions):
+        image = sample_frames(image)
+        image_tokens = "<image>" * int(len(image))
+        prompt = f"<|im_start|>user {image_tokens}\n{user_prompt}<|im_end|><|im_start|>assistant"
+          
+    elif image.endswith(image_extensions):
+        image = Image.open(image).convert("RGB")
+        prompt = f"<|im_start|>user <image>\n{user_prompt}<|im_end|><|im_start|>assistant"
+
+    print(len(image))
+
+    inputs = processor(prompt, image, return_tensors="pt")
+    streamer = TextIteratorStreamer(processor, skip_prompt=True, **{"skip_special_tokens": True})
+    generation_kwargs = dict(inputs, streamer=streamer, max_new_tokens=1024)
+    generated_text = ""
+
+    thread = Thread(target=model.generate, kwargs=generation_kwargs)
+    thread.start()
 
-# Define functions for image captioning, web search, and text extraction
-def generate_caption_instructblip(image_path, question):
-    client = Client("hysts/image-captioning-with-blip")
-    return client.predict(file(image_path), f"{question}", api_name="/caption")
+    buffer = ""
+    for new_text in streamer:
+        buffer += new_text
+        yield buffer
 
 def extract_text_from_webpage(html_content):
     soup = BeautifulSoup(html_content, 'html.parser')
@@ -62,10 +122,9 @@ def respond(message, history):
 
     # Handle image processing
     if message["files"]:
-        for image in message["files"]:
-            vqa += "[CAPTION of IMAGE]  "
-            gr.Info("Analyzing image")
-            vqa += generate_caption_instructblip(image, message["text"])
+        llava(message, history)
+        break
+        
 
     # Define function metadata for user interface
     functions_metadata = [