Upload folder using huggingface_hub

.gitattributes

@@ -33,3 +33,8 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+examples/sample_demo_1.mp4 filter=lfs diff=lfs merge=lfs -text
+examples/sample_demo_13.mp4 filter=lfs diff=lfs merge=lfs -text
+examples/sample_demo_22.mp4 filter=lfs diff=lfs merge=lfs -text
+examples/sample_demo_8.mp4 filter=lfs diff=lfs merge=lfs -text
+examples/sample_img_8.png filter=lfs diff=lfs merge=lfs -text

README.md

@@ -1,12 +1,6 @@
 ---
-title: Vlm Rlaif Demo
-emoji: 💻
-colorFrom: red
-colorTo: pink
+title: vlm-rlaif-demo
+app_file: gradio_web_server.py
 sdk: gradio
-sdk_version: 4.40.0
-app_file: app.py
-pinned: false
+sdk_version: 3.35.2
 ---
-
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference

cli.py

@@ -0,0 +1,142 @@
+import argparse
+import os
+
+import torch
+
+import sys
+sys.path.insert(0, os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), "Evaluation"))
+from llava.constants import IMAGE_TOKEN_INDEX, DEFAULT_IMAGE_TOKEN, DEFAULT_IM_START_TOKEN, DEFAULT_IM_END_TOKEN, \
+    DEFAULT_VIDEO_TOKEN
+from llava.conversation import conv_templates, SeparatorStyle
+from llava.model.builder import load_pretrained_model
+from llava.utils import disable_torch_init
+from llava.mm_utils import process_images, tokenizer_image_token, get_model_name_from_path, KeywordsStoppingCriteria
+from serve.utils import load_image, image_ext, video_ext
+
+from PIL import Image
+
+import requests
+from PIL import Image
+from io import BytesIO
+from transformers import TextStreamer
+
+
+
+def main(args):
+    # Model
+    disable_torch_init()
+
+    model_name = get_model_name_from_path(args.model_path)
+    tokenizer, model, processor, context_len = load_pretrained_model(args.model_path, args.model_base, model_name,
+                                                                     args.load_8bit, args.load_4bit,
+                                                                     device=args.device, cache_dir=args.cache_dir)
+    image_processor, video_processor = processor['image'], processor['video']
+    if 'llama-2' in model_name.lower():
+        conv_mode = "llava_llama_2"
+    elif "v1" in model_name.lower():
+        conv_mode = "llava_v1"
+    elif "mpt" in model_name.lower():
+        conv_mode = "mpt"
+    else:
+        conv_mode = "llava_v0"
+
+    if args.conv_mode is not None and conv_mode != args.conv_mode:
+        print('[WARNING] the auto inferred conversation mode is {}, while `--conv-mode` is {}, using {}'.format(conv_mode, args.conv_mode, args.conv_mode))
+    else:
+        args.conv_mode = conv_mode
+
+    conv = conv_templates[args.conv_mode].copy()
+    if "mpt" in model_name.lower():
+        roles = ('user', 'assistant')
+    else:
+        roles = conv.roles
+
+    tensor = []
+    special_token = []
+    args.file = args.file if isinstance(args.file, list) else [args.file]
+    for file in args.file:
+        if os.path.splitext(file)[-1].lower() in video_ext: # video extension
+            video_tensor = video_processor(file, return_tensors='pt')['pixel_values'][0].to(model.device, dtype=torch.float16)
+            special_token += [DEFAULT_IMAGE_TOKEN] * model.get_video_tower().config.num_frames
+        elif os.path.splitext(os.listdir(file)[0]).lower() in image_ext: # frames folder
+            vidframes_list = sorted(glob(file + '/*'))
+            images = load_frames(vidframes_list, model.get_video_tower().config.num_frames)
+            # Similar operation in model_worker.py
+            video_tensor = process_images(images, image_processor, args)
+            video_tensor = video_tensor.to(model.device, dtype=torch.float16)
+            video_tensor = video_tensor.unsqueeze(0)              
+            special_token += [DEFAULT_IMAGE_TOKEN] * model.get_video_tower().config.num_frames
+        else:
+            raise ValueError(f'Support video of {video_ext} and frames of {image_ext}, but found {os.path.splitext(file)[-1].lower()}')
+        print(video_tensor.shape)
+        tensor.append(video_tensor)
+
+
+
+
+    while True:
+        try:
+            inp = input(f"{roles[0]}: ")
+        except EOFError:
+            inp = ""
+        if not inp:
+            print("exit...")
+            break
+
+        print(f"{roles[1]}: ", end="")
+
+        if file is not None:
+            # first message
+            if getattr(model.config, "mm_use_im_start_end", False):
+                inp = DEFAULT_IM_START_TOKEN + DEFAULT_IMAGE_TOKEN + DEFAULT_IM_END_TOKEN + '\n' + inp
+                # inp = ''.join([DEFAULT_IM_START_TOKEN + i + DEFAULT_IM_END_TOKEN for i in special_token]) + '\n' + inp
+            else:
+                inp = DEFAULT_IMAGE_TOKEN + '\n' + inp
+                # inp = ''.join(special_token) + '\n' + inp
+            conv.append_message(conv.roles[0], inp)
+            file = None
+        else:
+            # later messages
+            conv.append_message(conv.roles[0], inp)
+        conv.append_message(conv.roles[1], None)
+        prompt = conv.get_prompt()
+
+        input_ids = tokenizer_image_token(prompt, tokenizer, IMAGE_TOKEN_INDEX, return_tensors='pt').unsqueeze(0).to(model.device)
+        stop_str = conv.sep if conv.sep_style != SeparatorStyle.TWO else conv.sep2
+        keywords = [stop_str]
+        stopping_criteria = KeywordsStoppingCriteria(keywords, tokenizer, input_ids)
+        streamer = TextStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
+
+        with torch.inference_mode():
+            output_ids = model.generate(
+                input_ids,
+                images=tensor,  # video as fake images
+                do_sample=True if args.temperature > 0 else False,
+                temperature=args.temperature,
+                max_new_tokens=args.max_new_tokens,
+                streamer=streamer,
+                use_cache=True,
+                stopping_criteria=[stopping_criteria])
+
+        outputs = tokenizer.decode(output_ids[0, input_ids.shape[1]:]).strip()
+        conv.messages[-1][-1] = outputs
+
+        if args.debug:
+            print("\n", {"prompt": prompt, "outputs": outputs}, "\n")
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--model-path", type=str, default="LanguageBind/Video-LLaVA-7B")
+    parser.add_argument("--model-base", type=str, default=None)
+    parser.add_argument("--cache-dir", type=str, default=None)
+    parser.add_argument("--file", nargs='+', type=str, required=True)
+    parser.add_argument("--device", type=str, default="cuda")
+    parser.add_argument("--conv-mode", type=str, default=None)
+    parser.add_argument("--temperature", type=float, default=0.2)
+    parser.add_argument("--max-new-tokens", type=int, default=512)
+    parser.add_argument("--load-8bit", action="store_true")
+    parser.add_argument("--load-4bit", action="store_true")
+    parser.add_argument("--debug", action="store_true")
+    args = parser.parse_args()
+    main(args)

controller.py

@@ -0,0 +1,298 @@
+"""
+A controller manages distributed workers.
+It sends worker addresses to clients.
+"""
+import argparse
+import asyncio
+import dataclasses
+from enum import Enum, auto
+import json
+import logging
+import time
+from typing import List, Union
+import threading
+
+from fastapi import FastAPI, Request
+from fastapi.responses import StreamingResponse
+import numpy as np
+import requests
+import uvicorn
+
+from videollava.constants import CONTROLLER_HEART_BEAT_EXPIRATION
+from videollava.utils import build_logger, server_error_msg
+
+
+logger = build_logger("controller", "controller.log")
+
+
+class DispatchMethod(Enum):
+    LOTTERY = auto()
+    SHORTEST_QUEUE = auto()
+
+    @classmethod
+    def from_str(cls, name):
+        if name == "lottery":
+            return cls.LOTTERY
+        elif name == "shortest_queue":
+            return cls.SHORTEST_QUEUE
+        else:
+            raise ValueError(f"Invalid dispatch method")
+
+
+@dataclasses.dataclass
+class WorkerInfo:
+    model_names: List[str]
+    speed: int
+    queue_length: int
+    check_heart_beat: bool
+    last_heart_beat: str
+
+
+def heart_beat_controller(controller):
+    while True:
+        time.sleep(CONTROLLER_HEART_BEAT_EXPIRATION)
+        controller.remove_stable_workers_by_expiration()
+
+
+class Controller:
+    def __init__(self, dispatch_method: str):
+        # Dict[str -> WorkerInfo]
+        self.worker_info = {}
+        self.dispatch_method = DispatchMethod.from_str(dispatch_method)
+
+        self.heart_beat_thread = threading.Thread(
+            target=heart_beat_controller, args=(self,))
+        self.heart_beat_thread.start()
+
+        logger.info("Init controller")
+
+    def register_worker(self, worker_name: str, check_heart_beat: bool,
+                        worker_status: dict):
+        if worker_name not in self.worker_info:
+            logger.info(f"Register a new worker: {worker_name}")
+        else:
+            logger.info(f"Register an existing worker: {worker_name}")
+
+        if not worker_status:
+            worker_status = self.get_worker_status(worker_name)
+        if not worker_status:
+            return False
+
+        self.worker_info[worker_name] = WorkerInfo(
+            worker_status["model_names"], worker_status["speed"], worker_status["queue_length"],
+            check_heart_beat, time.time())
+
+        logger.info(f"Register done: {worker_name}, {worker_status}")
+        return True
+
+    def get_worker_status(self, worker_name: str):
+        try:
+            r = requests.post(worker_name + "/worker_get_status", timeout=5)
+        except requests.exceptions.RequestException as e:
+            logger.error(f"Get status fails: {worker_name}, {e}")
+            return None
+
+        if r.status_code != 200:
+            logger.error(f"Get status fails: {worker_name}, {r}")
+            return None
+
+        return r.json()
+
+    def remove_worker(self, worker_name: str):
+        del self.worker_info[worker_name]
+
+    def refresh_all_workers(self):
+        old_info = dict(self.worker_info)
+        self.worker_info = {}
+
+        for w_name, w_info in old_info.items():
+            if not self.register_worker(w_name, w_info.check_heart_beat, None):
+                logger.info(f"Remove stale worker: {w_name}")
+
+    def list_models(self):
+        model_names = set()
+
+        for w_name, w_info in self.worker_info.items():
+            model_names.update(w_info.model_names)
+
+        return list(model_names)
+
+    def get_worker_address(self, model_name: str):
+        if self.dispatch_method == DispatchMethod.LOTTERY:
+            worker_names = []
+            worker_speeds = []
+            for w_name, w_info in self.worker_info.items():
+                if model_name in w_info.model_names:
+                    worker_names.append(w_name)
+                    worker_speeds.append(w_info.speed)
+            worker_speeds = np.array(worker_speeds, dtype=np.float32)
+            norm = np.sum(worker_speeds)
+            if norm < 1e-4:
+                return ""
+            worker_speeds = worker_speeds / norm
+            if True:  # Directly return address
+                pt = np.random.choice(np.arange(len(worker_names)),
+                    p=worker_speeds)
+                worker_name = worker_names[pt]
+                return worker_name
+
+            # Check status before returning
+            while True:
+                pt = np.random.choice(np.arange(len(worker_names)),
+                    p=worker_speeds)
+                worker_name = worker_names[pt]
+
+                if self.get_worker_status(worker_name):
+                    break
+                else:
+                    self.remove_worker(worker_name)
+                    worker_speeds[pt] = 0
+                    norm = np.sum(worker_speeds)
+                    if norm < 1e-4:
+                        return ""
+                    worker_speeds = worker_speeds / norm
+                    continue
+            return worker_name
+        elif self.dispatch_method == DispatchMethod.SHORTEST_QUEUE:
+            worker_names = []
+            worker_qlen = []
+            for w_name, w_info in self.worker_info.items():
+                if model_name in w_info.model_names:
+                    worker_names.append(w_name)
+                    worker_qlen.append(w_info.queue_length / w_info.speed)
+            if len(worker_names) == 0:
+                return ""
+            min_index = np.argmin(worker_qlen)
+            w_name = worker_names[min_index]
+            self.worker_info[w_name].queue_length += 1
+            logger.info(f"names: {worker_names}, queue_lens: {worker_qlen}, ret: {w_name}")
+            return w_name
+        else:
+            raise ValueError(f"Invalid dispatch method: {self.dispatch_method}")
+
+    def receive_heart_beat(self, worker_name: str, queue_length: int):
+        if worker_name not in self.worker_info:
+            logger.info(f"Receive unknown heart beat. {worker_name}")
+            return False
+
+        self.worker_info[worker_name].queue_length = queue_length
+        self.worker_info[worker_name].last_heart_beat = time.time()
+        logger.info(f"Receive heart beat. {worker_name}")
+        return True
+
+    def remove_stable_workers_by_expiration(self):
+        expire = time.time() - CONTROLLER_HEART_BEAT_EXPIRATION
+        to_delete = []
+        for worker_name, w_info in self.worker_info.items():
+            if w_info.check_heart_beat and w_info.last_heart_beat < expire:
+                to_delete.append(worker_name)
+
+        for worker_name in to_delete:
+            self.remove_worker(worker_name)
+
+    def worker_api_generate_stream(self, params):
+        worker_addr = self.get_worker_address(params["model"])
+        if not worker_addr:
+            logger.info(f"no worker: {params['model']}")
+            ret = {
+                "text": server_error_msg,
+                "error_code": 2,
+            }
+            yield json.dumps(ret).encode() + b"\0"
+
+        try:
+            response = requests.post(worker_addr + "/worker_generate_stream",
+                json=params, stream=True, timeout=5)
+            for chunk in response.iter_lines(decode_unicode=False, delimiter=b"\0"):
+                if chunk:
+                    yield chunk + b"\0"
+        except requests.exceptions.RequestException as e:
+            logger.info(f"worker timeout: {worker_addr}")
+            ret = {
+                "text": server_error_msg,
+                "error_code": 3,
+            }
+            yield json.dumps(ret).encode() + b"\0"
+
+
+    # Let the controller act as a worker to achieve hierarchical
+    # management. This can be used to connect isolated sub networks.
+    def worker_api_get_status(self):
+        model_names = set()
+        speed = 0
+        queue_length = 0
+
+        for w_name in self.worker_info:
+            worker_status = self.get_worker_status(w_name)
+            if worker_status is not None:
+                model_names.update(worker_status["model_names"])
+                speed += worker_status["speed"]
+                queue_length += worker_status["queue_length"]
+
+        return {
+            "model_names": list(model_names),
+            "speed": speed,
+            "queue_length": queue_length,
+        }
+
+
+app = FastAPI()
+
+
+@app.post("/register_worker")
+async def register_worker(request: Request):
+    data = await request.json()
+    controller.register_worker(
+        data["worker_name"], data["check_heart_beat"],
+        data.get("worker_status", None))
+
+
+@app.post("/refresh_all_workers")
+async def refresh_all_workers():
+    models = controller.refresh_all_workers()
+
+
+@app.post("/list_models")
+async def list_models():
+    models = controller.list_models()
+    return {"models": models}
+
+
+@app.post("/get_worker_address")
+async def get_worker_address(request: Request):
+    data = await request.json()
+    addr = controller.get_worker_address(data["model"])
+    return {"address": addr}
+
+
+@app.post("/receive_heart_beat")
+async def receive_heart_beat(request: Request):
+    data = await request.json()
+    exist = controller.receive_heart_beat(
+        data["worker_name"], data["queue_length"])
+    return {"exist": exist}
+
+
+@app.post("/worker_generate_stream")
+async def worker_api_generate_stream(request: Request):
+    params = await request.json()
+    generator = controller.worker_api_generate_stream(params)
+    return StreamingResponse(generator)
+
+
+@app.post("/worker_get_status")
+async def worker_api_get_status(request: Request):
+    return controller.worker_api_get_status()
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--host", type=str, default="localhost")
+    parser.add_argument("--port", type=int, default=21001)
+    parser.add_argument("--dispatch-method", type=str, choices=[
+        "lottery", "shortest_queue"], default="shortest_queue")
+    args = parser.parse_args()
+    logger.info(f"args: {args}")
+
+    controller = Controller(args.dispatch_method)
+    uvicorn.run(app, host=args.host, port=args.port, log_level="info")

examples/sample_demo_1.mp4

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:fc6562a172eb9cb3c760a3c9992349c1faa2c793c112b7b9e50bd5cb17c2164d
+size 1549315

examples/sample_demo_13.mp4

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:13384915331bf749fa31e2f4cbbd85ca90439b81b2390b4b512bd24b0dbd8bae
+size 19356822

examples/sample_demo_22.mp4

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:dcde24b3e67ff23aafd4b69854dbc7e2485eae65999c86c1beb9160d53fa2a11
+size 1505931

examples/sample_demo_8.mp4

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:618bb02562c769303b797ae3c29a66e15dcc0134d673747e8cf90582369c59a2
+size 29771700

gradio_utils.py

@@ -0,0 +1,155 @@
+import torch
+from transformers import TextStreamer
+
+import os
+import sys
+sys.path.insert(0, os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), "Evaluation"))
+from llava.constants import IMAGE_TOKEN_INDEX
+from llava.conversation import conv_templates, SeparatorStyle
+from llava.mm_utils import get_model_name_from_path, KeywordsStoppingCriteria, tokenizer_image_token
+from llava.model.builder import load_pretrained_model
+from llava.utils import disable_torch_init
+import shutil
+
+#   <a href="https://github.com/SNUMPR/vlm-rlaif.git" style="margin-right: 20px; text-decoration: none; display: flex; align-items: center;">
+#     <img src="https://z1.ax1x.com/2023/11/07/pil4sqH.png" alt="VLM-RLAIF" style="max-width: 120px; height: auto;">
+#   </a>
+
+cur_dir = os.path.dirname(os.path.abspath(__file__))
+title_markdown = ("""
+<div style="display: flex; justify-content: center; align-items: center; text-align: center;">
+  <img src="/dataset/dcahn/yura/vlm-rlaif/asset/Model.png" alt="VLM-RLAIF" style="max-width: 120px; height: auto;">
+  <img src="file:/dataset/dcahn/yura/vlm-rlaif/asset/Model.png" alt="VLM-RLAIF" style="max-width: 120px; height: auto;">
+  <div>
+    <h1 >VLM-RLAIF: Tuning Large Multimodal Models for Videos using Reinforcement Learning from AI Feedback (ACL 2024 Oral) </h1>
+    <h5 style="margin: 0;">If you like our project, please give us a star ✨ on Github for the latest update.</h5>
+  </div>
+</div>
+
+
+<div align="center">
+    <div style="display:flex; gap: 0.25rem;" align="center">
+        <a href='https://github.com/SNUMPR/vlm-rlaif'><img src='https://img.shields.io/badge/Github-Code-blue'></a>
+        <a href="https://arxiv.org/abs/2402.03746"><img src="https://img.shields.io/badge/Paper-arxiv-green"></a>
+    </div>
+</div>
+""")
+        # <a href='https://github.com/PKU-YuanGroup/Video-LLaVA/stargazers'><img src='https://img.shields.io/github/stars/PKU-YuanGroup/Video-LLaVA.svg?style=social'></a> # arXiv 버튼 옆에 추가?
+
+block_css = """
+#buttons button {
+    min-width: min(120px,100%);
+}
+"""
+
+tos_markdown = ("""
+### Terms of use
+By using this service, users are required to agree to the following terms:
+The service is a research preview intended for non-commercial use only. It only provides limited safety measures and may generate offensive content. It must not be used for any illegal, harmful, violent, racist, or sexual purposes. The service may collect user dialogue data for future research.
+Please click the "Flag" button if you get any inappropriate answer! We will collect those to keep improving our moderator.
+For an optimal experience, please use desktop computers for this demo, as mobile devices may compromise its quality.
+""")
+
+learn_more_markdown = ("""
+### License
+The service is a research preview intended for non-commercial use only, subject to the model [License](https://github.com/facebookresearch/llama/blob/main/MODEL_CARD.md) of LLaMA, [Terms of Use](https://openai.com/policies/terms-of-use) of the data generated by OpenAI, and [Privacy Practices](https://chrome.google.com/webstore/detail/sharegpt-share-your-chatg/daiacboceoaocpibfodeljbdfacokfjb) of ShareGPT. Please contact us if you find any potential violation.
+""")
+
+
+class Chat:
+    def __init__(self, model_path, conv_mode, model_base=None, load_8bit=False, load_4bit=False, device='cuda', cache_dir=None):
+        # model_base = '/dataset/yura/vlm-rlaif/pretrained/final_models/Video_LLaVA_SFT'
+        # model_base='/dataset/yura/vlm-rlaif/pretrained/llava-v1.5-7b-lora_w_lora_16_sftv2_short1632_and_then_long_rank32_alpha32_lr1e4_allmodels/SFT_merged'
+        # model_path = '/dataset/yura/vlm-rlaif/pretrained/LLaVA_Video-RL-Fact-RLHF-7b_SFTv2_RM_13b_v1_40k-v1.5-336-lora-padding/checkpoint-180/adapter_model/lora_policy'
+
+        disable_torch_init()
+        model_name = get_model_name_from_path(model_path)
+        # self.tokenizer, self.model, image_processor, context_len = load_pretrained_model(model_path, model_base, model_name,
+        #                                                                            load_8bit, load_4bit,
+        #                                                                            device=device, cache_dir=cache_dir)
+        is_rlhf_checkpoint = 'rlhf' in model_path.lower()
+        print("MODEL_PATH", model_path)
+        print("RLHF Checkpoint: ", is_rlhf_checkpoint)
+        if not model_base or model_base == "none": model_base = None
+        if is_rlhf_checkpoint:
+            model_name = model_path            
+            print("Config?", os.path.exists(os.path.join(model_path, "config.json")))
+            if not os.path.exists(os.path.join(model_path, "config.json")):
+                print("Copying")
+                shutil.copy(os.path.join(model_base, "config.json"), os.path.join(model_path, "config.json")) # Copy SFT model's config -> to RLHF folder
+                print("Listed", os.listdir(model_path))
+                print("Copying done")
+            # return(model_name)
+        # return
+        # self.tokenizer, self.model, image_processor, context_len = load_pretrained_model(model_path, model_base, model_name, load_8bit, load_4bit, device=device)
+        self.tokenizer, self.model, image_processor, context_len = load_pretrained_model(model_path, model_base, model_name, False, False, device=device)
+
+        
+
+        self.image_processor = image_processor
+        # self.image_processor = processor['image']
+        # self.video_processor = processor['video']
+        self.conv_mode = conv_mode
+        self.conv = conv_templates[conv_mode].copy()
+        self.device = self.model.device
+        print(self.model)
+
+    def get_prompt(self, qs, state):
+        state.append_message(state.roles[0], qs)
+        state.append_message(state.roles[1], None)
+        return state
+    
+    def _get_latest_prompt(self, state):
+        new_state = state.copy()
+        new_state.messages = state.messages[-2:]
+        return new_state
+
+    @torch.inference_mode()
+    # def generate(self, images_tensor: list, prompt: str, first_run: bool, state):
+    def generate(self, images_tensor: torch.Tensor, prompt: str, first_run: bool, state):
+        tokenizer, model, image_processor = self.tokenizer, self.model, self.image_processor
+
+        state = self.get_prompt(prompt, state)
+        # prompt = state.get_prompt()
+        latest_state = self._get_latest_prompt(state)
+        prompt = latest_state.get_prompt()
+    
+        # print('\n\n\n')
+        # print(prompt)
+
+        input_ids = tokenizer_image_token(prompt, tokenizer, IMAGE_TOKEN_INDEX, return_tensors='pt').unsqueeze(0).to(self.device)
+
+        temperature = 0.2
+
+        max_new_tokens = 1024
+
+        stop_str = self.conv.sep if self.conv.sep_style != SeparatorStyle.TWO else self.conv.sep2
+        keywords = [stop_str]
+        stopping_criteria = KeywordsStoppingCriteria(keywords, tokenizer, input_ids)
+        streamer = TextStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
+        print(prompt, input_ids.shape, images_tensor.shape)
+        # print(images_tensor)
+        with torch.inference_mode():
+            output_ids = model.generate(
+                input_ids,
+                images=images_tensor,
+                do_sample=True,
+                temperature=temperature,
+                max_new_tokens=max_new_tokens,
+                streamer=streamer,
+                use_cache=True,
+                stopping_criteria=[stopping_criteria])
+
+        input_token_len = input_ids.shape[1]
+        n_diff_input_output = (input_ids != output_ids[:, :input_token_len]).sum().item()
+        if n_diff_input_output > 0:
+            print(f'[Warning] {n_diff_input_output} output_ids are not the same as the input_ids')
+        outputs = tokenizer.batch_decode(output_ids[:, input_token_len:], skip_special_tokens=True)[0]
+        outputs = outputs.strip()
+        outputs = outputs.replace("QA_GT_caption_based_noisy", "")
+        if outputs.endswith(stop_str):
+            outputs = outputs[:-len(stop_str)]
+        outputs = outputs.strip()
+
+        print('response', outputs)
+        return outputs, state

gradio_web_server copy.py

@@ -0,0 +1,227 @@
+import shutil
+import subprocess
+
+import torch
+import gradio as gr
+from fastapi import FastAPI
+import os
+os.environ['CUDA_VISIBLE_DEVICES'] = '0'
+from PIL import Image
+import tempfile
+from decord import VideoReader, cpu
+from transformers import TextStreamer
+import argparse
+
+import sys
+sys.path.insert(0, os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), "Evaluation"))
+from llava.constants import DEFAULT_IMAGE_TOKEN, DEFAULT_IM_START_TOKEN, DEFAULT_IM_END_TOKEN
+from llava.conversation import conv_templates, SeparatorStyle, Conversation
+from llava.mm_utils import process_images
+
+from Evaluation.infer_utils import load_video_into_frames
+from serve.utils import load_image, image_ext, video_ext
+from serve.gradio_utils import Chat, tos_markdown, learn_more_markdown, title_markdown, block_css
+
+
+
+def save_image_to_local(image):
+    filename = os.path.join('temp', next(tempfile._get_candidate_names()) + '.jpg')
+    image = Image.open(image)
+    image.save(filename)
+    # print(filename)
+    return filename
+
+
+def save_video_to_local(video_path):
+    filename = os.path.join('temp', next(tempfile._get_candidate_names()) + '.mp4')
+    shutil.copyfile(video_path, filename)
+    return filename
+
+
+def generate(image1, video, textbox_in, first_run, state, state_, images_tensor, num_frames=50):
+    # ======= manually clear the conversation
+    # state = conv_templates[conv_mode].copy()
+    # state_ = conv_templates[conv_mode].copy()
+    # # ======= 
+    flag = 1
+    if not textbox_in:
+        if len(state_.messages) > 0:
+            textbox_in = state_.messages[-1][1]
+            state_.messages.pop(-1)
+            flag = 0
+        else:
+            return "Please enter instruction"
+    print("Video", video) # 잘 들어감
+    print("Images_tensor", images_tensor) # None
+    print("Textbox_IN", textbox_in) # 잘 들어감
+    print("State", state) # None
+    print("State_", state_) # None
+    # print(len(state_.messages))
+
+    video = video if video else "none"
+
+    if type(state) is not Conversation:
+        state = conv_templates[conv_mode].copy()
+        state_ = conv_templates[conv_mode].copy()
+        images_tensor = []
+
+    first_run = False if len(state.messages) > 0 else True
+
+    text_en_in = textbox_in.replace("picture", "image")
+
+    image_processor = handler.image_processor
+    assert os.path.exists(video)
+    if os.path.splitext(video)[-1].lower() in video_ext: # video extension
+        video_decode_backend = 'opencv'
+    elif os.path.splitext(os.listdir(video)[0]).lower() in image_ext: # frames folder
+        video_decode_backend = 'frames'
+    else:
+        raise ValueError(f'Support video of {video_ext} and frames of {image_ext}, but found {os.path.splitext(video)[-1].lower()}')
+
+    frames = load_video_into_frames(video, video_decode_backend=video_decode_backend, num_frames=num_frames)
+    tensor = process_images(frames, image_processor, argparse.Namespace(image_aspect_ratio='pad'))
+    # tensor = video_processor(video, return_tensors='pt')['pixel_values'][0]
+    # print(tensor.shape)
+    tensor = tensor.to(handler.model.device, dtype=dtype)
+    # images_tensor.append(tensor)
+    images_tensor = tensor
+
+    if handler.model.config.mm_use_im_start_end:
+        text_en_in = DEFAULT_IM_START_TOKEN + DEFAULT_IMAGE_TOKEN + DEFAULT_IM_END_TOKEN + '\n' + text_en_in
+    else:
+        text_en_in = DEFAULT_IMAGE_TOKEN + '\n' + text_en_in
+    text_en_out, state_ = handler.generate(images_tensor, text_en_in, first_run=first_run, state=state_)
+    state_.messages[-1] = (state_.roles[1], text_en_out)
+
+    text_en_out = text_en_out.split('#')[0]
+    textbox_out = text_en_out
+
+    show_images = ""
+    if os.path.exists(video):
+        filename = save_video_to_local(video)
+        show_images += f'<video controls playsinline width="500" style="display: inline-block;"  src="./file={filename}"></video>'
+    if flag:
+        state.append_message(state.roles[0], textbox_in + "\n" + show_images)
+    state.append_message(state.roles[1], textbox_out)
+
+    return (state, state_, state.to_gradio_chatbot(), False, gr.update(value=None, interactive=True), images_tensor, gr.update(value=image1 if os.path.exists(video) else None, interactive=True), gr.update(value=video if os.path.exists(video) else None, interactive=True))
+
+
+def regenerate(state, state_):
+    state.messages.pop(-1)
+    state_.messages.pop(-1)
+    if len(state.messages) > 0:
+        return state, state_, state.to_gradio_chatbot(), False
+    return (state, state_, state.to_gradio_chatbot(), True)
+
+
+def clear_history(state, state_):
+    state = conv_templates[conv_mode].copy()
+    state_ = conv_templates[conv_mode].copy()
+    return (gr.update(value=None, interactive=True),
+            gr.update(value=None, interactive=True), \
+            gr.update(value=None, interactive=True), \
+            True, state, state_, state.to_gradio_chatbot(), [])
+
+
+# ==== CHANGE HERE ====
+# conv_mode = "llava_v1"
+# model_path = 'LanguageBind/Video-LLaVA-7B'
+# FIXME!!!
+
+conv_mode = "llava_v0"
+model_path = 'SNUMPR/vlm_rlaif_video_llava_7b'
+# model_path = '/dataset/yura/vlm-rlaif/pretrained/final_models/Video_LLaVA_VLM_RLAIF_merged'
+cache_dir = './cache_dir'
+device = 'cuda'
+# device = 'cpu'
+load_8bit = True
+load_4bit = False
+dtype = torch.float16
+# =============
+
+handler = Chat(model_path, conv_mode=conv_mode, load_8bit=load_8bit, load_4bit=load_8bit, device=device, cache_dir=cache_dir)
+# handler.model.to(dtype=dtype)
+if not os.path.exists("temp"):
+    os.makedirs("temp")
+
+app = FastAPI()
+
+
+textbox = gr.Textbox(
+    show_label=False, placeholder="Enter text and press ENTER", container=False
+)
+with gr.Blocks(title='VLM-RLAIF', theme=gr.themes.Default(), css=block_css) as demo:
+    gr.Markdown(title_markdown)
+    state = gr.State()
+    state_ = gr.State()
+    first_run = gr.State()
+    images_tensor = gr.State()
+
+    image1 = gr.Image(label="Input Image", type="filepath")
+    with gr.Row():
+        with gr.Column(scale=3):
+            video = gr.Video(label="Input Video")
+
+            cur_dir = os.path.dirname(os.path.abspath(__file__))
+            gr.Examples(
+                examples=[
+                    [
+                        f"{cur_dir}/examples/sample_demo_1.mp4",
+                        "Why is this video funny?",
+                    ],
+                    [
+                        f"{cur_dir}/examples/sample_demo_3.mp4",
+                        "Can you identify any safety hazards in this video?"
+                    ],
+                    [
+                        f"{cur_dir}/examples/sample_demo_9.mp4",
+                        "Describe the video.",
+                    ],
+                    [
+                        f"{cur_dir}/examples/sample_demo_22.mp4",
+                        "Describe the activity in the video.",
+                    ],
+                ],
+                inputs=[video, textbox],
+            )
+
+        with gr.Column(scale=7):
+            chatbot = gr.Chatbot(label="VLM_RLAIF", bubble_full_width=True).style(height=750)
+            with gr.Row():
+                with gr.Column(scale=8):
+                    textbox.render()
+                with gr.Column(scale=1, min_width=50):
+                    submit_btn = gr.Button(
+                        value="Send", variant="primary", interactive=True
+                    )
+            with gr.Row(elem_id="buttons") as button_row:
+                upvote_btn = gr.Button(value="👍  Upvote", interactive=True)
+                downvote_btn = gr.Button(value="👎  Downvote", interactive=True)
+                flag_btn = gr.Button(value="⚠️  Flag", interactive=True)
+                # stop_btn = gr.Button(value="⏹️  Stop Generation", interactive=False)
+                regenerate_btn = gr.Button(value="🔄  Regenerate", interactive=True)
+                # clear_btn = gr.Button(value="🗑️  Clear history", interactive=True)
+
+    gr.Markdown(tos_markdown)
+    gr.Markdown(learn_more_markdown)
+
+    submit_btn.click(generate, [image1, video, textbox, first_run, state, state_, images_tensor],
+                     [state, state_, chatbot, first_run, textbox, images_tensor, image1, video])
+    # submit_btn.click(generate, [video, textbox, first_run, state, state_, images_tensor],
+                    #  [state, state_, chatbot, first_run, textbox, images_tensor, video])
+
+    regenerate_btn.click(regenerate, [state, state_], [state, state_, chatbot, first_run]).then(
+        generate, [image1, video, textbox, first_run, state, state_, images_tensor], [state, state_, chatbot, first_run, textbox, images_tensor, image1, video])
+        # generate, [video, textbox, first_run, state, state_, images_tensor], [state, state_, chatbot, first_run, textbox, images_tensor, video])
+
+    # clear_btn.click(clear_history, [state, state_],
+    #                 [image1, video, textbox, first_run, state, state_, chatbot, images_tensor])
+                    # [video, textbox, first_run, state, state_, chatbot, images_tensor])
+
+# app = gr.mount_gradio_app(app, demo, path="/")
+# demo.launch(share=True)
+demo.launch()
+
+# uvicorn videollava.serve.gradio_web_server:app
+# python -m  videollava.serve.gradio_web_server

gradio_web_server.py

@@ -0,0 +1,234 @@
+import shutil
+import subprocess
+
+import torch
+import gradio as gr
+from fastapi import FastAPI
+import os
+os.environ['CUDA_VISIBLE_DEVICES'] = '0'
+from PIL import Image
+import tempfile
+from decord import VideoReader, cpu
+from transformers import TextStreamer
+import argparse
+
+import sys
+sys.path.insert(0, os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), "Evaluation"))
+from llava.constants import DEFAULT_IMAGE_TOKEN, DEFAULT_IM_START_TOKEN, DEFAULT_IM_END_TOKEN
+from llava.conversation import conv_templates, SeparatorStyle, Conversation
+from llava.mm_utils import process_images
+
+from Evaluation.infer_utils import load_video_into_frames
+from serve.utils import load_image, image_ext, video_ext
+from serve.gradio_utils import Chat, tos_markdown, learn_more_markdown, title_markdown, block_css
+
+
+
+def save_image_to_local(image):
+    filename = os.path.join('temp', next(tempfile._get_candidate_names()) + '.jpg')
+    image = Image.open(image)
+    image.save(filename)
+    # print(filename)
+    return filename
+
+
+def save_video_to_local(video_path):
+    filename = os.path.join('temp', next(tempfile._get_candidate_names()) + '.mp4')
+    shutil.copyfile(video_path, filename)
+    return filename
+
+
+def generate(video, textbox_in, first_run, state, state_, images_tensor, num_frames=50):
+    # ======= manually clear the conversation
+    # state = conv_templates[conv_mode].copy()
+    # state_ = conv_templates[conv_mode].copy()
+    # # ======= 
+    flag = 1
+    if not textbox_in:
+        if len(state_.messages) > 0:
+            textbox_in = state_.messages[-1][1]
+            state_.messages.pop(-1)
+            flag = 0
+        else:
+            return "Please enter instruction"
+    # else:
+    #     if state is not None and state_ is not None:
+    #         # reset conversations
+    #         state.messages = []
+    #         state_.messages = []
+        
+    print("Video", video) # 잘 들어감
+    print("Images_tensor", images_tensor) # None
+    print("Textbox_IN", textbox_in) # 잘 들어감
+    print("State", state) # None
+    print("State_", state_) # None
+    # print(len(state_.messages))
+
+    video = video if video else "none"
+
+    if type(state) is not Conversation:
+        state = conv_templates[conv_mode].copy()
+        state_ = conv_templates[conv_mode].copy()
+        images_tensor = []
+
+    first_run = False if len(state.messages) > 0 else True
+
+    text_en_in = textbox_in.replace("picture", "image")
+
+    image_processor = handler.image_processor
+    assert os.path.exists(video)
+    if os.path.splitext(video)[-1].lower() in video_ext: # video extension
+        video_decode_backend = 'opencv'
+    elif os.path.splitext(os.listdir(video)[0]).lower() in image_ext: # frames folder
+        video_decode_backend = 'frames'
+    else:
+        raise ValueError(f'Support video of {video_ext} and frames of {image_ext}, but found {os.path.splitext(video)[-1].lower()}')
+
+    frames = load_video_into_frames(video, video_decode_backend=video_decode_backend, num_frames=num_frames)
+    tensor = process_images(frames, image_processor, argparse.Namespace(image_aspect_ratio='pad'))
+    # tensor = video_processor(video, return_tensors='pt')['pixel_values'][0]
+    # print(tensor.shape)
+    tensor = tensor.to(handler.model.device, dtype=dtype)
+    # images_tensor.append(tensor)
+    images_tensor = tensor
+
+    if handler.model.config.mm_use_im_start_end:
+        text_en_in = DEFAULT_IM_START_TOKEN + DEFAULT_IMAGE_TOKEN + DEFAULT_IM_END_TOKEN + '\n' + text_en_in
+    else:
+        text_en_in = DEFAULT_IMAGE_TOKEN + '\n' + text_en_in
+    text_en_out, state_ = handler.generate(images_tensor, text_en_in, first_run=first_run, state=state_)
+    state_.messages[-1] = (state_.roles[1], text_en_out)
+
+    text_en_out = text_en_out.split('#')[0]
+    textbox_out = text_en_out
+
+    show_images = ""
+    if os.path.exists(video):
+        filename = save_video_to_local(video)
+        show_images += f'<video controls playsinline width="500" style="display: inline-block;"  src="./file={filename}"></video>'
+    if flag:
+        state.append_message(state.roles[0], textbox_in + "\n" + show_images)
+    state.append_message(state.roles[1], textbox_out)
+
+    return (state, state_, state.to_gradio_chatbot(), False, gr.update(value=None, interactive=True), images_tensor, \
+        gr.update(value=video if os.path.exists(video) else None, interactive=True))
+
+
+def regenerate(state, state_):
+    state.messages.pop(-1)
+    state_.messages.pop(-1)
+    if len(state.messages) > 0:
+        return state, state_, state.to_gradio_chatbot(), False
+    return (state, state_, state.to_gradio_chatbot(), True)
+
+
+def clear_history(state, state_):
+    state = conv_templates[conv_mode].copy()
+    state_ = conv_templates[conv_mode].copy()
+    return (gr.update(value=None, interactive=True),
+            gr.update(value=None, interactive=True), \
+            gr.update(value=None, interactive=True), \
+            True, state, state_, state.to_gradio_chatbot(), [])
+
+
+# ==== CHANGE HERE ====
+# conv_mode = "llava_v1"
+# model_path = 'LanguageBind/Video-LLaVA-7B'
+# FIXME!!!
+
+conv_mode = "llava_v0"
+model_path = 'SNUMPR/vlm_rlaif_video_llava_7b'
+# model_path = '/dataset/yura/vlm-rlaif/pretrained/final_models/Video_LLaVA_VLM_RLAIF_merged'
+cache_dir = './cache_dir'
+device = 'cuda'
+# device = 'cpu'
+load_8bit = True
+load_4bit = False
+dtype = torch.float16
+# =============
+
+handler = Chat(model_path, conv_mode=conv_mode, load_8bit=load_8bit, load_4bit=load_8bit, device=device, cache_dir=cache_dir)
+# handler.model.to(dtype=dtype)
+if not os.path.exists("temp"):
+    os.makedirs("temp")
+
+app = FastAPI()
+
+
+textbox = gr.Textbox(
+    show_label=False, placeholder="Enter text and press ENTER", container=False
+)
+with gr.Blocks(title='VLM-RLAIF', theme=gr.themes.Default(), css=block_css) as demo:
+    gr.Markdown(title_markdown)
+    state = gr.State()
+    state_ = gr.State()
+    first_run = gr.State()
+    images_tensor = gr.State()
+
+    # image1 = gr.Image(label="Input Image", type="filepath")
+    with gr.Row():
+        with gr.Column(scale=3):
+            video = gr.Video(label="Input Video")
+
+            cur_dir = os.path.dirname(os.path.abspath(__file__))
+            gr.Examples(
+                examples=[
+                    [
+                        f"{cur_dir}/examples/sample_demo_1.mp4",
+                        "Why is this video funny?",
+                    ],
+                    [
+                        f"{cur_dir}/examples/sample_demo_3.mp4",
+                        "Can you identify any safety hazards in this video?"
+                    ],
+                    [
+                        f"{cur_dir}/examples/sample_demo_9.mp4",
+                        "Describe the video.",
+                    ],
+                    [
+                        f"{cur_dir}/examples/sample_demo_22.mp4",
+                        "Describe the activity in the video.",
+                    ],
+                ],
+                inputs=[video, textbox],
+            )
+
+        with gr.Column(scale=7):
+            chatbot = gr.Chatbot(label="VLM_RLAIF", bubble_full_width=True).style(height=750)
+            with gr.Row():
+                with gr.Column(scale=8):
+                    textbox.render()
+                with gr.Column(scale=1, min_width=50):
+                    submit_btn = gr.Button(
+                        value="Send", variant="primary", interactive=True
+                    )
+            with gr.Row(elem_id="buttons") as button_row:
+                upvote_btn = gr.Button(value="👍  Upvote", interactive=True)
+                downvote_btn = gr.Button(value="👎  Downvote", interactive=True)
+                flag_btn = gr.Button(value="⚠️  Flag", interactive=True)
+                # stop_btn = gr.Button(value="⏹️  Stop Generation", interactive=False)
+                regenerate_btn = gr.Button(value="🔄  Regenerate", interactive=True)
+                # clear_btn = gr.Button(value="🗑️  Clear history", interactive=True)
+
+    gr.Markdown(tos_markdown)
+    gr.Markdown(learn_more_markdown)
+
+    submit_btn.click(generate, [video, textbox, first_run, state, state_, images_tensor],
+                     [state, state_, chatbot, first_run, textbox, images_tensor, video])
+    # submit_btn.click(generate, [video, textbox, first_run, state, state_, images_tensor],
+                    #  [state, state_, chatbot, first_run, textbox, images_tensor, video])
+
+    regenerate_btn.click(regenerate, [state, state_], [state, state_, chatbot, first_run]).then(
+        generate, [video, textbox, first_run, state, state_, images_tensor], [state, state_, chatbot, first_run, textbox, images_tensor, video])
+        # generate, [video, textbox, first_run, state, state_, images_tensor], [state, state_, chatbot, first_run, textbox, images_tensor, video])
+
+    # clear_btn.click(clear_history, [state, state_],
+    #                 [image1, video, textbox, first_run, state, state_, chatbot, images_tensor])
+                    # [video, textbox, first_run, state, state_, chatbot, images_tensor])
+
+# app = gr.mount_gradio_app(app, demo, path="/")
+demo.launch(share=True)
+# demo.launch()
+
+# uvicorn videollava.serve.gradio_web_server:app
+# python -m  videollava.serve.gradio_web_server

model_worker.py

@@ -0,0 +1,285 @@
+"""
+A model worker executes the model.
+"""
+import argparse
+import asyncio
+import json
+import time
+import threading
+import uuid
+
+from fastapi import FastAPI, Request, BackgroundTasks
+from fastapi.responses import StreamingResponse
+import requests
+import torch
+import uvicorn
+from functools import partial
+
+from videollava.constants import WORKER_HEART_BEAT_INTERVAL
+from videollava.utils import (build_logger, server_error_msg,
+    pretty_print_semaphore)
+from videollava.model.builder import load_pretrained_model
+from videollava.mm_utils import process_images, load_image_from_base64, tokenizer_image_token, KeywordsStoppingCriteria
+from videollava.constants import IMAGE_TOKEN_INDEX, DEFAULT_IMAGE_TOKEN, DEFAULT_IM_START_TOKEN, DEFAULT_IM_END_TOKEN
+from transformers import TextIteratorStreamer
+from threading import Thread
+
+
+GB = 1 << 30
+
+worker_id = str(uuid.uuid4())[:6]
+logger = build_logger("model_worker", f"model_worker_{worker_id}.log")
+global_counter = 0
+
+model_semaphore = None
+
+
+def heart_beat_worker(controller):
+
+    while True:
+        time.sleep(WORKER_HEART_BEAT_INTERVAL)
+        controller.send_heart_beat()
+
+
+class ModelWorker:
+    def __init__(self, controller_addr, worker_addr,
+                 worker_id, no_register,
+                 model_path, model_base, model_name,
+                 load_8bit, load_4bit, device):
+        self.controller_addr = controller_addr
+        self.worker_addr = worker_addr
+        self.worker_id = worker_id
+        if model_path.endswith("/"):
+            model_path = model_path[:-1]
+        if model_name is None:
+            model_paths = model_path.split("/")
+            if model_paths[-1].startswith('checkpoint-'):
+                self.model_name = model_paths[-2] + "_" + model_paths[-1]
+            else:
+                self.model_name = model_paths[-1]
+        else:
+            self.model_name = model_name
+
+        self.device = device
+        logger.info(f"Loading the model {self.model_name} on worker {worker_id} ...")
+        self.tokenizer, self.model, self.image_processor, self.context_len = load_pretrained_model(
+            model_path, model_base, self.model_name, load_8bit, load_4bit, device=self.device)
+        self.is_multimodal = 'llava' in self.model_name.lower()
+
+        if not no_register:
+            self.register_to_controller()
+            self.heart_beat_thread = threading.Thread(
+                target=heart_beat_worker, args=(self,))
+            self.heart_beat_thread.start()
+
+    def register_to_controller(self):
+        logger.info("Register to controller")
+
+        url = self.controller_addr + "/register_worker"
+        data = {
+            "worker_name": self.worker_addr,
+            "check_heart_beat": True,
+            "worker_status": self.get_status()
+        }
+        r = requests.post(url, json=data)
+        assert r.status_code == 200
+
+    def send_heart_beat(self):
+        logger.info(f"Send heart beat. Models: {[self.model_name]}. "
+                    f"Semaphore: {pretty_print_semaphore(model_semaphore)}. "
+                    f"global_counter: {global_counter}")
+
+        url = self.controller_addr + "/receive_heart_beat"
+
+        while True:
+            try:
+                ret = requests.post(url, json={
+                    "worker_name": self.worker_addr,
+                    "queue_length": self.get_queue_length()}, timeout=5)
+                exist = ret.json()["exist"]
+                break
+            except requests.exceptions.RequestException as e:
+                logger.error(f"heart beat error: {e}")
+            time.sleep(5)
+
+        if not exist:
+            self.register_to_controller()
+
+    def get_queue_length(self):
+        if model_semaphore is None:
+            return 0
+        else:
+            return args.limit_model_concurrency - model_semaphore._value + (len(
+                model_semaphore._waiters) if model_semaphore._waiters is not None else 0)
+
+    def get_status(self):
+        return {
+            "model_names": [self.model_name],
+            "speed": 1,
+            "queue_length": self.get_queue_length(),
+        }
+
+    @torch.inference_mode()
+    def generate_stream(self, params):
+        tokenizer, model, image_processor = self.tokenizer, self.model, self.image_processor
+
+        prompt = params["prompt"]
+        ori_prompt = prompt
+        images = params.get("images", None)
+        num_image_tokens = 0
+        if images is not None and len(images) > 0 and self.is_multimodal:
+            if len(images) > 0:
+                if len(images) != prompt.count(DEFAULT_IMAGE_TOKEN):
+                    raise ValueError("Number of images does not match number of <image> tokens in prompt")
+
+                images = [load_image_from_base64(image) for image in images]
+                images = process_images(images, image_processor, model.config)
+
+                if type(images) is list:
+                    images = [image.to(self.model.device, dtype=torch.float16) for image in images]
+                else:
+                    images = images.to(self.model.device, dtype=torch.float16)
+
+                replace_token = DEFAULT_IMAGE_TOKEN
+                if getattr(self.model.config, 'mm_use_im_start_end', False):
+                    replace_token = DEFAULT_IM_START_TOKEN + replace_token + DEFAULT_IM_END_TOKEN
+                prompt = prompt.replace(DEFAULT_IMAGE_TOKEN, replace_token)
+
+                num_image_tokens = prompt.count(replace_token) * model.get_vision_tower().num_patches
+            else:
+                images = None
+            image_args = {"images": images}
+        else:
+            images = None
+            image_args = {}
+
+        temperature = float(params.get("temperature", 1.0))
+        top_p = float(params.get("top_p", 1.0))
+        max_context_length = getattr(model.config, 'max_position_embeddings', 2048)
+        max_new_tokens = min(int(params.get("max_new_tokens", 256)), 1024)
+        stop_str = params.get("stop", None)
+        do_sample = True if temperature > 0.001 else False
+
+        input_ids = tokenizer_image_token(prompt, tokenizer, IMAGE_TOKEN_INDEX, return_tensors='pt').unsqueeze(0).to(self.device)
+        keywords = [stop_str]
+        stopping_criteria = KeywordsStoppingCriteria(keywords, tokenizer, input_ids)
+        streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True, timeout=15)
+
+        max_new_tokens = min(max_new_tokens, max_context_length - input_ids.shape[-1] - num_image_tokens)
+
+        if max_new_tokens < 1:
+            yield json.dumps({"text": ori_prompt + "Exceeds max token length. Please start a new conversation, thanks.", "error_code": 0}).encode() + b"\0"
+            return
+
+        thread = Thread(target=model.generate, kwargs=dict(
+            inputs=input_ids,
+            do_sample=do_sample,
+            temperature=temperature,
+            top_p=top_p,
+            max_new_tokens=max_new_tokens,
+            streamer=streamer,
+            stopping_criteria=[stopping_criteria],
+            use_cache=True,
+            **image_args
+        ))
+        thread.start()
+
+        generated_text = ori_prompt
+        for new_text in streamer:
+            generated_text += new_text
+            if generated_text.endswith(stop_str):
+                generated_text = generated_text[:-len(stop_str)]
+            yield json.dumps({"text": generated_text, "error_code": 0}).encode() + b"\0"
+
+    def generate_stream_gate(self, params):
+        try:
+            for x in self.generate_stream(params):
+                yield x
+        except ValueError as e:
+            print("Caught ValueError:", e)
+            ret = {
+                "text": server_error_msg,
+                "error_code": 1,
+            }
+            yield json.dumps(ret).encode() + b"\0"
+        except torch.cuda.CudaError as e:
+            print("Caught torch.cuda.CudaError:", e)
+            ret = {
+                "text": server_error_msg,
+                "error_code": 1,
+            }
+            yield json.dumps(ret).encode() + b"\0"
+        except Exception as e:
+            print("Caught Unknown Error", e)
+            ret = {
+                "text": server_error_msg,
+                "error_code": 1,
+            }
+            yield json.dumps(ret).encode() + b"\0"
+
+
+app = FastAPI()
+
+
+def release_model_semaphore(fn=None):
+    model_semaphore.release()
+    if fn is not None:
+        fn()
+
+
+@app.post("/worker_generate_stream")
+async def generate_stream(request: Request):
+    global model_semaphore, global_counter
+    global_counter += 1
+    params = await request.json()
+
+    if model_semaphore is None:
+        model_semaphore = asyncio.Semaphore(args.limit_model_concurrency)
+    await model_semaphore.acquire()
+    worker.send_heart_beat()
+    generator = worker.generate_stream_gate(params)
+    background_tasks = BackgroundTasks()
+    background_tasks.add_task(partial(release_model_semaphore, fn=worker.send_heart_beat))
+    return StreamingResponse(generator, background=background_tasks)
+
+
+@app.post("/worker_get_status")
+async def get_status(request: Request):
+    return worker.get_status()
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--host", type=str, default="localhost")
+    parser.add_argument("--port", type=int, default=21002)
+    parser.add_argument("--worker-address", type=str,
+        default="http://localhost:21002")
+    parser.add_argument("--controller-address", type=str,
+        default="http://localhost:21001")
+    parser.add_argument("--model-path", type=str, default="facebook/opt-350m")
+    parser.add_argument("--model-base", type=str, default=None)
+    parser.add_argument("--model-name", type=str)
+    parser.add_argument("--device", type=str, default="cuda")
+    parser.add_argument("--multi-modal", action="store_true", help="Multimodal mode is automatically detected with model name, please make sure `llava` is included in the model path.")
+    parser.add_argument("--limit-model-concurrency", type=int, default=5)
+    parser.add_argument("--stream-interval", type=int, default=1)
+    parser.add_argument("--no-register", action="store_true")
+    parser.add_argument("--load-8bit", action="store_true")
+    parser.add_argument("--load-4bit", action="store_true")
+    args = parser.parse_args()
+    logger.info(f"args: {args}")
+
+    if args.multi_modal:
+        logger.warning("Multimodal mode is automatically detected with model name, please make sure `llava` is included in the model path.")
+
+    worker = ModelWorker(args.controller_address,
+                         args.worker_address,
+                         worker_id,
+                         args.no_register,
+                         args.model_path,
+                         args.model_base,
+                         args.model_name,
+                         args.load_8bit,
+                         args.load_4bit,
+                         args.device)
+    uvicorn.run(app, host=args.host, port=args.port, log_level="info")

processing_utils.py

@@ -0,0 +1,99 @@
+import torch
+from transformers import TextStreamer
+import numpy as np
+import os
+import json
+import torch
+
+import numpy as np
+import base64
+from PIL import Image
+from io import BytesIO
+import matplotlib.pyplot as plt
+from torchvision.transforms import Compose, Lambda, ToTensor
+from torchvision import transforms
+from transformers import ProcessorMixin, BatchEncoding
+from transformers.image_processing_utils import BatchFeature
+from pytorchvideo.data.encoded_video import EncodedVideo
+from torchvision.transforms import Compose, Lambda, ToTensor
+from torchvision.transforms._transforms_video import NormalizeVideo, RandomCropVideo, RandomHorizontalFlipVideo, CenterCropVideo
+from pytorchvideo.transforms import ApplyTransformToKey, ShortSideScale, UniformTemporalSubsample
+
+
+def load_frames(frames_dir):
+    results = []
+    frame_names = os.listdir(frames_dir)
+    frame_names.sort()
+    for frame_name in frame_names:
+        image_path = f"{frames_dir}/{frame_name}"
+        results.append(image_path)
+    return results
+
+def sample_frames(frames, num_segments):
+    duration = len(frames)
+    frame_id_array = np.linspace(0, duration-1, num_segments, dtype=int)
+    frame_id_list = frame_id_array.tolist()
+
+    sampled_frames = []
+    for frame_idx in frame_id_list:
+        single_frame_path = frames[frame_idx]
+        sampled_frames.append(single_frame_path)
+    return sampled_frames
+
+
+class VideoProcessor:
+    def __init__(self, image_transform):
+        self.image_transform = image_transform
+        
+    def __call__(self, video_path, transform=None, 
+                 video_decode_backend='opencv', 
+                 clip_start_sec=0.0, clip_end_sec=None, 
+                 num_frames=50, **kwargs):
+        if transform is None: transform = self.image_transform
+        if video_decode_backend == 'pytorchvideo':
+            #  decord pyav
+            video = EncodedVideo.from_path(video_path, decoder="decord", decode_audio=False)
+            duration = video.duration
+            start_sec = clip_start_sec  # secs
+            end_sec = clip_end_sec if clip_end_sec is not None else duration  # secs
+            video_data = video.get_clip(start_sec=start_sec, end_sec=end_sec)
+            video_outputs = transform(video_data)
+
+        elif video_decode_backend == 'decord':
+            import decord
+            from decord import VideoReader, cpu
+            decord.bridge.set_bridge('torch')
+            decord_vr = VideoReader(video_path, ctx=cpu(0))
+            ori_duration = len(decord_vr)
+            # frame_id_list = np.linspace(0, duration-1, num_frames, dtype=int)
+            fps_vid = decord_vr.get_avg_fps()
+            valid_duration = min(int(fps_vid * 10), ori_duration)
+            frame_id_list = np.linspace(0, valid_duration-1, num_frames, dtype=int)
+            video_data = decord_vr.get_batch(frame_id_list)
+            video_data = video_data.permute(3, 0, 1, 2)  # (T, H, W, C) -> (C, T, H, W)
+            video_outputs = transform(video_data)
+
+        elif video_decode_backend == 'opencv':
+            import cv2
+            cv2_vr = cv2.VideoCapture(video_path)
+            duration = int(cv2_vr.get(cv2.CAP_PROP_FRAME_COUNT))
+            frame_id_list = np.linspace(0, duration-1, num_frames, dtype=int)
+
+            video_data = []
+            for frame_idx in frame_id_list:
+                cv2_vr.set(1, frame_idx)
+                _, frame = cv2_vr.read()
+                frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
+                video_data.append(torch.from_numpy(frame).permute(2, 0, 1))
+            cv2_vr.release()
+            video_data = torch.stack(video_data, dim=1)
+            video_outputs = transform(video_data)
+        
+        elif video_decode_backend == 'frames':
+            # FIXME does not input start and end clip timestamps. Require duration info to deal with.
+            frames = load_frames(video_path)
+            frames = sample_frames(frames, num_frames)
+            to_tensor = ToTensor()
+            video_data = torch.stack([to_tensor(_) for _ in frames]).permute(1, 0, 2, 3) # (T, C, H, W) -> (C, T, H, W)
+        else:
+            raise NameError('video_decode_backend should specify in (pytorchvideo, decord, opencv, frames)')

register_worker.py

@@ -0,0 +1,26 @@
+"""
+Manually register workers.
+
+Usage:
+python3 -m fastchat.serve.register_worker --controller http://localhost:21001 --worker-name http://localhost:21002
+"""
+
+import argparse
+
+import requests
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--controller-address", type=str)
+    parser.add_argument("--worker-name", type=str)
+    parser.add_argument("--check-heart-beat", action="store_true")
+    args = parser.parse_args()
+
+    url = args.controller_address + "/register_worker"
+    data = {
+        "worker_name": args.worker_name,
+        "check_heart_beat": args.check_heart_beat,
+        "worker_status": None,
+    }
+    r = requests.post(url, json=data)
+    assert r.status_code == 200

test_message.py

@@ -0,0 +1,62 @@
+import argparse
+import json
+
+import requests
+
+from videollava.conversation import default_conversation
+
+
+def main():
+    if args.worker_address:
+        worker_addr = args.worker_address
+    else:
+        controller_addr = args.controller_address
+        ret = requests.post(controller_addr + "/refresh_all_workers")
+        ret = requests.post(controller_addr + "/list_models")
+        models = ret.json()["models"]
+        models.sort()
+        print(f"Models: {models}")
+
+        ret = requests.post(controller_addr + "/get_worker_address",
+            json={"model": args.model_name})
+        worker_addr = ret.json()["address"]
+        print(f"worker_addr: {worker_addr}")
+
+    if worker_addr == "":
+        return
+
+    conv = default_conversation.copy()
+    conv.append_message(conv.roles[0], args.message)
+    prompt = conv.get_prompt()
+
+    headers = {"User-Agent": "LLaVA Client"}
+    pload = {
+        "model": args.model_name,
+        "prompt": prompt,
+        "max_new_tokens": args.max_new_tokens,
+        "temperature": 0.7,
+        "stop": conv.sep,
+    }
+    response = requests.post(worker_addr + "/worker_generate_stream", headers=headers,
+            json=pload, stream=True)
+
+    print(prompt.replace(conv.sep, "\n"), end="")
+    for chunk in response.iter_lines(chunk_size=8192, decode_unicode=False, delimiter=b"\0"):
+        if chunk:
+            data = json.loads(chunk.decode("utf-8"))
+            output = data["text"].split(conv.sep)[-1]
+            print(output, end="\r")
+    print("")
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--controller-address", type=str, default="http://localhost:21001")
+    parser.add_argument("--worker-address", type=str)
+    parser.add_argument("--model-name", type=str, default="facebook/opt-350m")
+    parser.add_argument("--max-new-tokens", type=int, default=32)
+    parser.add_argument("--message", type=str, default=
+        "Tell me a story with more than 1000 words.")
+    args = parser.parse_args()
+
+    main()

utils.py

@@ -0,0 +1,16 @@
+from io import BytesIO
+
+import requests
+from PIL import Image
+
+
+def load_image(image_file):
+    if image_file.startswith('http://') or image_file.startswith('https://'):
+        response = requests.get(image_file)
+        image = Image.open(BytesIO(response.content)).convert('RGB')
+    else:
+        image = Image.open(image_file).convert('RGB')
+    return image
+
+video_ext = ['.mp4', '.mov', '.mkv', '.avi']
+image_ext = ['.jpg', '.png', '.bmp', '.jpeg']