Spaces:
Sleeping
Sleeping
| import gradio as gr | |
| from transformers import LlavaProcessor, LlavaForConditionalGeneration, TextIteratorStreamer | |
| from threading import Thread | |
| import re | |
| import time | |
| from PIL import Image | |
| import torch | |
| import cv2 | |
| import spaces | |
| model_id = "llava-hf/llava-interleave-qwen-7b-hf" | |
| processor = LlavaProcessor.from_pretrained(model_id) | |
| model = LlavaForConditionalGeneration.from_pretrained(model_id, torch_dtype=torch.float16) | |
| model.to("cuda") | |
| def sample_frames(video_file, num_frames) : | |
| video = cv2.VideoCapture(video_file) | |
| total_frames = int(video.get(cv2.CAP_PROP_FRAME_COUNT)) | |
| 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 | |
| def bot_streaming(message, history): | |
| if message["files"]: | |
| image = message["files"][-1] | |
| else: | |
| # if there's no image uploaded for this turn, look for images in the past turns | |
| # kept inside tuples, take the last one | |
| for hist in history: | |
| if type(hist[0])==tuple: | |
| image = hist[0][0] | |
| txt = message["text"] | |
| img = message["files"] | |
| ext_buffer =f"'user\ntext': '{txt}', 'files': '{img}' assistantAnswer:" | |
| if image is None: | |
| gr.Error("You need to upload an image or video for LLaVA to work.") | |
| video_extensions = ("avi", "mp4", "mov", "mkv", "flv", "wmv", "mjpeg") | |
| 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, 5) | |
| image_tokens = "<image>" * 5 | |
| prompt = f"<|im_start|>user {image_tokens}\n{message}<|im_end|><|im_start|>assistant" | |
| elif image.endswith(image_extensions): | |
| image = Image.open(image).convert("RGB") | |
| prompt = f"<|im_start|>user <image>\n{message}<|im_end|><|im_start|>assistant" | |
| inputs = processor(prompt, image, return_tensors="pt").to("cuda", torch.float16) | |
| streamer = TextIteratorStreamer(processor, **{"skip_special_tokens": True}) | |
| generation_kwargs = dict(inputs, streamer=streamer, max_new_tokens=100) | |
| generated_text = "" | |
| thread = Thread(target=model.generate, kwargs=generation_kwargs) | |
| thread.start() | |
| buffer = "" | |
| for new_text in streamer: | |
| buffer += new_text | |
| print(buffer) | |
| generated_text_without_prompt = buffer[len(ext_buffer):] | |
| time.sleep(0.01) | |
| yield generated_text_without_prompt | |
| demo = gr.ChatInterface(fn=bot_streaming, title="LLaVA Interleave", examples=[{"text": "What is on the flower?", "files":["./bee.jpg"]}, | |
| {"text": "How to make this pastry?", "files":["./baklava.png"]}, | |
| {"text": "What type of cats are these?", "files":["./cats.mp4"]}], | |
| description="Try [LLaVA Interleave](https://huggingface.co/docs/transformers/main/en/model_doc/llava) in this demo (more specifically, the [Qwen-1.5-7B variant](https://huggingface.co/llava-hf/llava-interleave-qwen-7b-hf)). Upload an image or a video, and start chatting about it, or simply try one of the examples below. If you don't upload an image, you will receive an error.", | |
| stop_btn="Stop Generation", multimodal=True) | |
| demo.launch(debug=True) |