Update app.py

app.py

@@ -1,313 +1,720 @@
+import gradio as gr
 import spaces
+from PIL import Image
+from transformers import AutoProcessor, WhisperForConditionalGeneration, WhisperProcessor, CLIPProcessor, CLIPModel
+import copy
+from decord import VideoReader, cpu
+import numpy as np
 import json
-import math
+from tqdm import tqdm
 import os
-import traceback
-from io import BytesIO
-from typing import Any, Dict, List, Optional, Tuple
+import easyocr
 import re
-import time
-from threading import Thread
-
-import gradio as gr
-import requests
+import ast
+import socket
+import pickle
+import ffmpeg
+import torchaudio
 import torch
-from PIL import Image
+import warnings
+import shutil
+from transformers import AutoTokenizer, AutoModelForCausalLM, AutoConfig, BitsAndBytesConfig
+from llava.constants import DEFAULT_IMAGE_PATCH_TOKEN, DEFAULT_IM_START_TOKEN, DEFAULT_IM_END_TOKEN
+from llava.utils import rank0_print
+from llava.model import *
+from llava.model.language_model.llava_qwen import LlavaQwenForCausalLM, LlavaQwenConfig
+from llava.mm_utils import get_model_name_from_path, process_images, tokenizer_image_token
+from llava.constants import IMAGE_TOKEN_INDEX, DEFAULT_IMAGE_TOKEN
+from llava.conversation import conv_templates, SeparatorStyle
+
+# Inline the tools if code is provided; assuming they are modules, but for single file, inline if possible.
+# For now, assume imports work, but since single file, need to define them.
+# User mentioned imports like from tools.rag_retriever_dynamic import retrieve_documents_with_dynamic
+# But code not provided, so I'll keep the imports, assuming environment has them.
+# Similarly for filter_keywords, generate_scene_graph_description
+
+from tools.rag_retriever_dynamic import retrieve_documents_with_dynamic
+from tools.filter_keywords import filter_keywords
+from tools.scene_graph import generate_scene_graph_description
+
+# From builder.py
+def load_pretrained_model(model_path, model_base, model_name, load_8bit=False, load_4bit=False, device_map="auto", torch_dtype="float16", attn_implementation=None, customized_config=None, overwrite_config=None, **kwargs):
+    kwargs["device_map"] = device_map
+
+    if load_8bit:
+        kwargs["load_in_8bit"] = True
+    elif load_4bit:
+        kwargs["load_in_4bit"] = True
+        kwargs["quantization_config"] = BitsAndBytesConfig(load_in_4bit=True, bnb_4bit_compute_dtype=torch.float16, bnb_4bit_use_double_quant=True, bnb_4bit_quant_type="nf4")
+    elif torch_dtype == "float16":
+        kwargs["torch_dtype"] = torch.float16
+    elif torch_dtype == "bfloat16":
+        kwargs["torch_dtype"] = torch.bfloat16
+    else:
+        import pdb;pdb.set_trace()
 
-from transformers import (
-    Qwen2VLForConditionalGeneration,
-    Qwen2_5_VLForConditionalGeneration,
-    AutoModelForImageTextToText,
-    AutoProcessor,
-    TextIteratorStreamer,
-    AutoModel,
-    AutoTokenizer,
-)
+    if customized_config is not None:
+        kwargs["config"] = customized_config
 
-from transformers.image_utils import load_image
-
-# --- Constants and Model Setup ---
-MAX_INPUT_TOKEN_LENGTH = 4096
-# Note: The following line correctly falls back to CPU if CUDA is not available.
-# Let the environment (e.g., Hugging Face Spaces) determine the device.
-# This avoids conflicts with the CUDA environment setup by the platform.
-device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-
-print("CUDA_VISIBLE_DEVICES=", os.environ.get("CUDA_VISIBLE_DEVICES"))
-print("torch.__version__ =", torch.__version__)
-print("torch.version.cuda =", torch.version.cuda)
-print("cuda available:", torch.cuda.is_available())
-print("cuda device count:", torch.cuda.device_count())
-if torch.cuda.is_available():
-    print("current device:", torch.cuda.current_device())
-    print("device name:", torch.cuda.get_device_name(torch.cuda.current_device()))
-
-print("Using device:", device)
-
-# --- Model Loading ---
-
-# --- Prompts for Different Tasks ---
-layout_prompt = """Please output the layout information from the image, including each layout element's bbox, its category, and the corresponding text content within the bbox.
-
-1. Bbox format: [x1, y1, x2, y2]
-2. Layout Categories: The possible categories are ['Caption', 'Footnote', 'Formula', 'List-item', 'Page-footer', 'Page-header', 'Picture', 'Section-header', 'Table', 'Text', 'Title'].
-3. Text Extraction & Formatting Rules:
-    - For tables, provide the content in a structured JSON format.
-    - For all other elements, provide the plain text.
-4. Constraints:
-    - The output must be the original text from the image.
-    - All layout elements must be sorted according to human reading order.
-5. Final Output: The entire output must be a single JSON object wrapped in ```json ... ```.
-"""
-
-ocr_prompt = "Perform precise OCR on the image. Extract all text content, maintaining the original structure, paragraphs, and tables as formatted markdown."
-
-# --- Model Loading ---
-MODEL_ID_M = "prithivMLmods/Camel-Doc-OCR-080125"
-processor_m = AutoProcessor.from_pretrained(MODEL_ID_M, trust_remote_code=True)
-model_m = Qwen2_5_VLForConditionalGeneration.from_pretrained(
-    MODEL_ID_M, trust_remote_code=True, torch_dtype=torch.float16
-).to(device).eval()
-
-MODEL_ID_T = "prithivMLmods/Megalodon-OCR-Sync-0713"
-processor_t = AutoProcessor.from_pretrained(MODEL_ID_T, trust_remote_code=True)
-model_t = Qwen2_5_VLForConditionalGeneration.from_pretrained(
-    MODEL_ID_T, trust_remote_code=True, torch_dtype=torch.float16
-).to(device).eval()
-
-MODEL_ID_C = "nanonets/Nanonets-OCR-s"
-processor_c = AutoProcessor.from_pretrained(MODEL_ID_C, trust_remote_code=True)
-model_c = Qwen2_5_VLForConditionalGeneration.from_pretrained(
-    MODEL_ID_C, trust_remote_code=True, torch_dtype=torch.float16
-).to(device).eval()
-
-MODEL_ID_G = "echo840/MonkeyOCR"
-SUBFOLDER = "Recognition"
-processor_g = AutoProcessor.from_pretrained(
-    MODEL_ID_G, trust_remote_code=True, subfolder=SUBFOLDER
-)
-model_g = Qwen2_5_VLForConditionalGeneration.from_pretrained(
-    MODEL_ID_G, trust_remote_code=True, subfolder=SUBFOLDER, torch_dtype=torch.float16
-).to(device).eval()
-
-MODEL_ID_I = "allenai/olmOCR-7B-0725"
-processor_i = AutoProcessor.from_pretrained(MODEL_ID_I, trust_remote_code=True)
-model_i = Qwen2_5_VLForConditionalGeneration.from_pretrained(
-    MODEL_ID_I, trust_remote_code=True, torch_dtype=torch.float16
-).to(device).eval()
-
-# --- Utility Functions ---
-def layoutjson2md(layout_data: Any) -> str:
-    """
-    FIXED: Converts the structured JSON from Layout Analysis into formatted Markdown.
-    This version is robust against malformed JSON from the model.
-    """
-    markdown_lines = []
-
-    # If the model wraps the list in a dictionary, find and extract the list.
-    if isinstance(layout_data, dict):
-        found_list = None
-        for value in layout_data.values():
-            if isinstance(value, list):
-                found_list = value
-                break
-        if found_list is not None:
-            layout_data = found_list
+    if "multimodal" in kwargs:
+        if kwargs["multimodal"] is True:
+            is_multimodal = True
+            kwargs.pop("multimodal")
+    else:
+        is_multimodal = False
+
+    if "llava" in model_name.lower() or is_multimodal:
+        # Load LLaVA model
+        if "lora" in model_name.lower() and model_base is None:
+            warnings.warn(
+                "There is `lora` in model name but no `model_base` is provided. If you are loading a LoRA model, please provide the `model_base` argument. Detailed instruction: https://github.com/haotian-liu/LLaVA#launch-a-model-worker-lora-weights-unmerged."
+            )
+        if "lora" in model_name.lower() and model_base is not None:
+            lora_cfg_pretrained = AutoConfig.from_pretrained(model_path)
+            tokenizer = AutoTokenizer.from_pretrained(model_base, use_fast=False)
+            rank0_print("Loading LLaVA from base model...")
+            if "mixtral" in model_name.lower():
+                from llava.model.language_model.llava_mixtral import LlavaMixtralConfig
+
+                lora_cfg_pretrained = LlavaMixtralConfig.from_pretrained(model_path)
+                tokenizer = AutoTokenizer.from_pretrained(model_base, use_fast=False)
+                model = LlavaMixtralForCausalLM.from_pretrained(model_base, low_cpu_mem_usage=True, config=lora_cfg_pretrained, attn_implementation=attn_implementation, **kwargs)
+            elif "mistral" in model_name.lower():
+                from llava.model.language_model.llava_mistral import LlavaMistralConfig
+
+                lora_cfg_pretrained = LlavaMistralConfig.from_pretrained(model_path)
+                tokenizer = AutoTokenizer.from_pretrained(model_base, use_fast=False)
+                model = LlavaMistralForCausalLM.from_pretrained(model_base, low_cpu_mem_usage=True, config=lora_cfg_pretrained, attn_implementation=attn_implementation, **kwargs)
+            elif "gemma" in model_name.lower():
+                from llava.model.language_model.llava_gemma import LlavaGemmaConfig
+
+                lora_cfg_pretrained = LlavaGemmaConfig.from_pretrained(model_path)
+                tokenizer = AutoTokenizer.from_pretrained(model_base, use_fast=False)
+                model = LlavaGemmaForCausalLM.from_pretrained(model_base, low_cpu_mem_usage=True, config=lora_cfg_pretrained, attn_implementation=attn_implementation, **kwargs)
+            else:
+                from llava.model.language_model.llava_llama import LlavaConfig, LlavaLlamaForCausalLM
+
+                lora_cfg_pretrained = LlavaConfig.from_pretrained(model_path)
+                tokenizer = AutoTokenizer.from_pretrained(model_base, use_fast=False)
+                model = LlavaLlamaForCausalLM.from_pretrained(model_base, low_cpu_mem_usage=True, config=lora_cfg_pretrained, attn_implementation=attn_implementation, **kwargs)
+
+            token_num, tokem_dim = model.lm_head.out_features, model.lm_head.in_features
+            if model.lm_head.weight.shape[0] != token_num:
+                model.lm_head.weight = torch.nn.Parameter(torch.empty(token_num, tokem_dim, device=model.device, dtype=model.dtype))
+                model.model.embed_tokens.weight = torch.nn.Parameter(torch.empty(token_num, tokem_dim, device=model.device, dtype=model.dtype))
+
+            rank0_print("Loading additional LLaVA weights...")
+            if os.path.exists(os.path.join(model_path, "non_lora_trainables.bin")):
+                non_lora_trainables = torch.load(os.path.join(model_path, "non_lora_trainables.bin"), map_location="cpu")
+            else:
+                # this is probably from HF Hub
+                from huggingface_hub import hf_hub_download
+
+                def load_from_hf(repo_id, filename, subfolder=None):
+                    cache_file = hf_hub_download(repo_id=repo_id, filename=filename, subfolder=subfolder)
+                    return torch.load(cache_file, map_location="cpu")
+
+                non_lora_trainables = load_from_hf(model_path, "non_lora_trainables.bin")
+            non_lora_trainables = {(k[11:] if k.startswith("base_model.") else k): v for k, v in non_lora_trainables.items()}
+            if any(k.startswith("model.model.") for k in non_lora_trainables):
+                non_lora_trainables = {(k[6:] if k.startswith("model.") else k): v for k, v in non_lora_trainables.items()}
+            model.load_state_dict(non_lora_trainables, strict=False)
+
+            from peft import PeftModel
+
+            rank0_print("Loading LoRA weights...")
+            model = PeftModel.from_pretrained(model, model_path)
+            rank0_print("Merging LoRA weights...")
+            model = model.merge_and_unload()
+            rank0_print("Model is loaded...")
+        elif model_base is not None:  # this may be mm projector only, loading projector with preset language mdoel
+            rank0_print(f"Loading LLaVA from base model {model_base}...")
+            if "mixtral" in model_name.lower():
+                tokenizer = AutoTokenizer.from_pretrained(model_base, use_fast=False)
+                cfg_pretrained = AutoConfig.from_pretrained(model_path)
+                model = LlavaMixtralForCausalLM.from_pretrained(model_base, low_cpu_mem_usage=True, config=cfg_pretrained, attn_implementation=attn_implementation, **kwargs)
+            elif "mistral" in model_name.lower() or "zephyr" in model_name.lower():
+                tokenizer = AutoTokenizer.from_pretrained(model_base, use_fast=False)
+                cfg_pretrained = AutoConfig.from_pretrained(model_path)
+                model = LlavaMistralForCausalLM.from_pretrained(model_base, low_cpu_mem_usage=True, config=cfg_pretrained, attn_implementation=attn_implementation, **kwargs)
+            elif "gemma" in model_name.lower():
+                tokenizer = AutoTokenizer.from_pretrained(model_base, use_fast=False)
+                cfg_pretrained = AutoConfig.from_pretrained(model_path)
+                model = LlavaGemmaForCausalLM.from_pretrained(model_base, low_cpu_mem_usage=True, config=cfg_pretrained, attn_implementation=attn_implementation, **kwargs)
+            elif (
+                "wizardlm-2" in model_name.lower()
+                and "vicuna" in model_name.lower()
+                or "llama" in model_name.lower()
+                or "yi" in model_name.lower()
+                or "nous-hermes" in model_name.lower()
+                or "llava-v1.6-34b" in model_name.lower()
+                or "llava-v1.5" in model_name.lower()
+            ):
+                from llava.model.language_model.llava_llama import LlavaConfig
+
+                tokenizer = AutoTokenizer.from_pretrained(model_path, use_fast=False)
+                if customized_config is None:
+                    llava_cfg = LlavaConfig.from_pretrained(model_path)
+                    if "v1.5" in model_name.lower():
+                        llava_cfg.delay_load = True  # a workaround for correctly loading v1.5 models
+                else:
+                    llava_cfg = customized_config
+
+                tokenizer = AutoTokenizer.from_pretrained(model_base, use_fast=False)
+                llava_cfg = LlavaConfig.from_pretrained(model_path)
+                model = LlavaLlamaForCausalLM.from_pretrained(model_base, low_cpu_mem_usage=True, config=llava_cfg, **kwargs)
+            else:
+                raise ValueError(f"Model {model_name} not supported")
+
+            mm_projector_weights = torch.load(os.path.join(model_path, "mm_projector.bin"), map_location="cpu")
+            mm_projector_weights = {k: v.to(torch.float16) for k, v in mm_projector_weights.items()}
+            model.load_state_dict(mm_projector_weights, strict=False)
         else:
-            return "### Error: Could not find a list of layout items in the JSON object."
-
-    if not isinstance(layout_data, list):
-        return f"### Error: Expected a list of layout items, but received type {type(layout_data).__name__}."
-
-    try:
-        # Filter out any non-dictionary items and sort by reading order.
-        valid_items = [item for item in layout_data if isinstance(item, dict)]
-        sorted_items = sorted(valid_items, key=lambda x: (x.get('bbox', [0, 0, 0, 0])[1], x.get('bbox', [0, 0, 0, 0])[0]))
-
-        for item in sorted_items:
-            category = item.get('category', 'Text') # Default to 'Text' if no category
-            text = item.get('text', '')
-            if not text:
-                continue
-
-            if category == 'Title':
-                markdown_lines.append(f"# {text}\n")
-            elif category == 'Section-header':
-                markdown_lines.append(f"## {text}\n")
-            elif category == 'Table':
-                if isinstance(text, dict) and 'header' in text and 'rows' in text:
-                    header = '| ' + ' | '.join(map(str, text['header'])) + ' |'
-                    separator = '| ' + ' | '.join(['---'] * len(text['header'])) + ' |'
-                    rows = ['| ' + ' | '.join(map(str, row)) + ' |' for row in text['rows']]
-                    markdown_lines.extend([header, separator] + rows)
-                    markdown_lines.append("\n")
-                else:  # Fallback for simple text or malformed tables
-                    markdown_lines.append(f"{text}\n")
+            rank0_print(f"Loaded LLaVA model: {model_path}")
+            if "mixtral" in model_name.lower():
+                from llava.model.language_model.llava_mixtral import LlavaMixtralConfig
+
+                tokenizer = AutoTokenizer.from_pretrained(model_path, use_fast=False)
+                if customized_config is None:
+                    llava_cfg = LlavaMixtralConfig.from_pretrained(model_path)
+                else:
+                    llava_cfg = customized_config
+
+                if overwrite_config is not None:
+                    rank0_print(f"Overwriting config with {overwrite_config}")
+                    for k, v in overwrite_config.items():
+                        setattr(llava_cfg, k, v)
+
+                tokenizer = AutoTokenizer.from_pretrained(model_path)
+                model = LlavaMixtralForCausalLM.from_pretrained(model_path, low_cpu_mem_usage=True, attn_implementation=attn_implementation, config=llava_cfg, **kwargs)
+
+            elif "mistral" in model_name.lower() or "zephyr" in model_name.lower():
+                tokenizer = AutoTokenizer.from_pretrained(model_path)
+                model = LlavaMistralForCausalLM.from_pretrained(model_path, low_cpu_mem_usage=True, attn_implementation=attn_implementation, **kwargs)
+            elif (
+                "wizardlm-2" in model_name.lower()
+                and "vicuna" in model_name.lower()
+                or "llama" in model_name.lower()
+                or "yi" in model_name.lower()
+                or "nous-hermes" in model_name.lower()
+                or "llava-v1.6-34b" in model_name.lower()
+                or "llava-v1.5" in model_name.lower()
+            ):
+                from llava.model.language_model.llava_llama import LlavaConfig
+
+                tokenizer = AutoTokenizer.from_pretrained(model_path, use_fast=False)
+                if customized_config is None:
+                    llava_cfg = LlavaConfig.from_pretrained(model_path)
+                    if "v1.5" in model_path.lower():
+                        llava_cfg.delay_load = True  # a workaround for correctly loading v1.5 models
+                else:
+                    llava_cfg = customized_config
+
+                if overwrite_config is not None:
+                    rank0_print(f"Overwriting config with {overwrite_config}")
+                    for k, v in overwrite_config.items():
+                        setattr(llava_cfg, k, v)
+
+                model = LlavaLlamaForCausalLM.from_pretrained(model_path, low_cpu_mem_usage=True, attn_implementation=attn_implementation, config=llava_cfg, **kwargs)
+
+            elif "qwen" in model_name.lower() or "quyen" in model_name.lower():
+                tokenizer = AutoTokenizer.from_pretrained(model_path)
+                if "moe" in model_name.lower() or "A14B" in model_name.lower():
+                    from llava.model.language_model.llava_qwen_moe import LlavaQwenMoeConfig
+                    if overwrite_config is not None:
+                        llava_cfg = LlavaQwenMoeConfig.from_pretrained(model_path)
+                        rank0_print(f"Overwriting config with {overwrite_config}")
+                        for k, v in overwrite_config.items():
+                            setattr(llava_cfg, k, v)
+                        model = LlavaQwenMoeForCausalLM.from_pretrained(model_path, low_cpu_mem_usage=True, attn_implementation=attn_implementation, config=llava_cfg, **kwargs)
+                    else:
+                        model = LlavaQwenMoeForCausalLM.from_pretrained(model_path, low_cpu_mem_usage=True, attn_implementation=attn_implementation, **kwargs)
+
+                else:
+                    from llava.model.language_model.llava_qwen import LlavaQwenConfig
+                    if overwrite_config is not None:
+                        llava_cfg = LlavaQwenConfig.from_pretrained(model_path)
+                        rank0_print(f"Overwriting config with {overwrite_config}")
+                        for k, v in overwrite_config.items():
+                            setattr(llava_cfg, k, v)
+                        model = LlavaQwenForCausalLM.from_pretrained(model_path, low_cpu_mem_usage=True, attn_implementation=attn_implementation, config=llava_cfg, **kwargs)
+                    else:
+                        model = LlavaQwenForCausalLM.from_pretrained(model_path, low_cpu_mem_usage=True, attn_implementation=attn_implementation, **kwargs)
+
+            elif "gemma" in model_name.lower():
+                tokenizer = AutoTokenizer.from_pretrained(model_path, use_fast=False)
+                cfg_pretrained = AutoConfig.from_pretrained(model_path)
+                model = LlavaGemmaForCausalLM.from_pretrained(model_path, low_cpu_mem_usage=True, config=cfg_pretrained, attn_implementation=attn_implementation, **kwargs)
             else:
-                markdown_lines.append(f"{text}\n")
+                try:
+                    from llava.model.language_model.llava_llama import LlavaConfig
+
+                    tokenizer = AutoTokenizer.from_pretrained(model_path, use_fast=False)
+                    if customized_config is None:
+                        llava_cfg = LlavaConfig.from_pretrained(model_path)
+                        if "v1.5" in model_path.lower():
+                            llava_cfg.delay_load = True  # a workaround for correctly loading v1.5 models
+                    else:
+                        llava_cfg = customized_config
+
+                    if overwrite_config is not None:
+                        rank0_print(f"Overwriting config with {overwrite_config}")
+                        for k, v in overwrite_config.items():
+                            setattr(llava_cfg, k, v)
+                    model = LlavaLlamaForCausalLM.from_pretrained(model_path, low_cpu_mem_usage=True, attn_implementation=attn_implementation, config=llava_cfg, **kwargs)
+                except:
+                    raise ValueError(f"Model {model_name} not supported")
 
-    except Exception as e:
-        print(f"Error converting to markdown: {e}")
-        traceback.print_exc()
-        return "### Error: An unexpected error occurred while converting JSON to Markdown."
+    else:
+        # Load language model
+        if model_base is not None:
+            # PEFT model
+            from peft import PeftModel
+
+            tokenizer = AutoTokenizer.from_pretrained(model_base, use_fast=False)
+            model = AutoModelForCausalLM.from_pretrained(model_base, torch_dtype=torch.float16, low_cpu_mem_usage=True, device_map="auto")
+            print(f"Loading LoRA weights from {model_path}")
+            model = PeftModel.from_pretrained(model, model_path)
+            print(f"Merging weights")
+            model = model.merge_and_unload()
+            print("Convert to FP16...")
+            model.to(torch.float16)
+        else:
+            use_fast = False
+            if "mpt" in model_name.lower().replace("prompt", ""):
+                tokenizer = AutoTokenizer.from_pretrained(model_path, use_fast=True)
+                model = AutoModelForCausalLM.from_pretrained(model_path, low_cpu_mem_usage=True, trust_remote_code=True, **kwargs)
+            else:
+                tokenizer = AutoTokenizer.from_pretrained(model_path, use_fast=False)
+                model = AutoModelForCausalLM.from_pretrained(model_path, low_cpu_mem_usage=True, **kwargs)
+
+    rank0_print(f"Model Class: {model.__class__.__name__}")
+    image_processor = None
+
+    if "llava" in model_name.lower() or is_multimodal:
+        mm_use_im_start_end = getattr(model.config, "mm_use_im_start_end", False)
+        mm_use_im_patch_token = getattr(model.config, "mm_use_im_patch_token", True)
+        if mm_use_im_patch_token:
+            tokenizer.add_tokens([DEFAULT_IMAGE_PATCH_TOKEN], special_tokens=True)
+        if mm_use_im_start_end:
+            tokenizer.add_tokens([DEFAULT_IM_START_TOKEN, DEFAULT_IM_END_TOKEN], special_tokens=True)
+        model.resize_token_embeddings(len(tokenizer))
+
+        vision_tower = model.get_vision_tower()
+        if not vision_tower.is_loaded:
+            vision_tower.load_model(device_map=device_map)
+        if device_map != "auto":
+            vision_tower.to(device="cuda", dtype=torch.float16)
+        image_processor = vision_tower.image_processor
+
+    if hasattr(model.config, "max_sequence_length"):
+        context_len = model.config.max_sequence_length
+    elif hasattr(model.config, "max_position_embeddings"):
+        context_len = model.config.max_position_embeddings
+    elif hasattr(model.config, "tokenizer_model_max_length"):
+        context_len = model.config.tokenizer_model_max_length
+    else:
+        context_len = 2048
 
-    return "\n".join(markdown_lines)
+    return tokenizer, model, image_processor, context_len
 
+# From vidrag_pipeline.py
+max_frames_num = 32
+clip_model = CLIPModel.from_pretrained("openai/clip-vit-large-patch14-336", torch_dtype=torch.float16, device_map="auto")
+clip_processor = CLIPProcessor.from_pretrained("openai/clip-vit-large-patch14-336")
+# whisper_model = WhisperForConditionalGeneration.from_pretrained(
+#     "openai/whisper-large",
+#     torch_dtype=torch.float16,
+#     device_map="auto"
+# )
+# whisper_processor = WhisperProcessor.from_pretrained("openai/whisper-large")
 
-# --- Core Application Logic ---
 @spaces.GPU
-def process_document_stream(model_name: str, task_choice: str, image: Image.Image, max_new_tokens: int):
-    """
-    Main generator function that handles both OCR and Layout Analysis tasks.
-    """
-    if image is None:
-        yield "Please upload an image.", "Please upload an image.", None
-        return
-
-    # 1. Select prompt based on user's task choice
-    text_prompt = ocr_prompt if task_choice == "Content Extraction" else layout_prompt
-
-    # 2. Select model and processor
-    if model_name == "Camel-Doc-OCR-080125": processor, model = processor_m, model_m
-    elif model_name == "Megalodon-OCR-Sync-0713": processor, model = processor_t, model_t
-    elif model_name == "Nanonets-OCR-s": processor, model = processor_c, model_c
-    elif model_name == "MonkeyOCR-Recognition": processor, model = processor_g, model_g
-    elif model_name == "olmOCR-7B-0725": processor, model = processor_i, model_i
-    else:
-        yield "Invalid model selected.", "Invalid model selected.", None
-        return
-
-    # 3. Prepare model inputs and streamer
-    messages = [{"role": "user", "content": [{"type": "image", "image": image}, {"type": "text", "text": text_prompt}]}]
-    prompt_full = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
-    inputs = processor(text=[prompt_full], images=[image], return_tensors="pt", padding=True, truncation=True, max_length=MAX_INPUT_TOKEN_LENGTH).to(device)
-    streamer = TextIteratorStreamer(processor, skip_prompt=True, skip_special_tokens=True)
-    generation_kwargs = {**inputs, "streamer": streamer, "max_new_tokens": max_new_tokens}
-
-    thread = Thread(target=model.generate, kwargs=generation_kwargs)
-    thread.start()
-
-    # 4. Stream raw output to the UI in real-time
-    buffer = ""
-    for new_text in streamer:
-        buffer += new_text
-        buffer = buffer.replace("<|im_end|>", "")
-        time.sleep(0.01)
-        yield buffer , "⏳ Processing...", {"status": "streaming"}
-
-    # 5. Post-process the final buffer based on the selected task
-    if task_choice == "Content Extraction":
-        # For OCR, the buffer is the final result.
-        yield buffer, buffer, None
-    else: # Layout Analysis
+def process_video(video_path, max_frames_num, fps=1, force_sample=False):
+    if max_frames_num == 0:
+        return np.zeros((1, 336, 336, 3))
+    vr = VideoReader(video_path, ctx=cpu(),num_threads=1)
+    total_frame_num = len(vr)
+    video_time = total_frame_num / vr.get_avg_fps()
+    fps = round(vr.get_avg_fps()/fps)
+    frame_idx = [i for i in range(0, len(vr), fps)]
+    frame_time = [i/fps for i in frame_idx]
+    if len(frame_idx) > max_frames_num or force_sample:
+        sample_fps = max_frames_num
+        uniform_sampled_frames = np.linspace(0, total_frame_num - 1, sample_fps, dtype=int)
+        frame_idx = uniform_sampled_frames.tolist()
+        frame_time = [i/vr.get_avg_fps() for i in frame_idx]
+    frame_time = ",".join([f"{i:.2f}s" for i in frame_time])
+    spare_frames = vr.get_batch(frame_idx).asnumpy()
+
+    return spare_frames, frame_time, video_time
+
+def extract_audio(video_path, audio_path):
+    if not os.path.exists(audio_path):
+        ffmpeg.input(video_path).output(audio_path, acodec='pcm_s16le', ac=1, ar='16k').run()
+
+def chunk_audio(audio_path, chunk_length_s=30):
+    speech, sr = torchaudio.load(audio_path)
+    speech = speech.mean(dim=0)  
+    speech = torchaudio.transforms.Resample(orig_freq=sr, new_freq=16000)(speech)  
+    num_samples_per_chunk = chunk_length_s * 16000 
+    chunks = []
+    for i in range(0, len(speech), num_samples_per_chunk):
+        chunks.append(speech[i:i + num_samples_per_chunk])
+    return chunks
+
+# def transcribe_chunk(chunk):
+
+#     inputs = whisper_processor(chunk, return_tensors="pt")
+#     inputs["input_features"] = inputs["input_features"].to(whisper_model.device, torch.float16)
+#     with torch.no_grad():
+#         predicted_ids = whisper_model.generate(
+#             inputs["input_features"],
+#             no_repeat_ngram_size=2,
+#             early_stopping=True
+#         )
+#     transcription = whisper_processor.batch_decode(predicted_ids, skip_special_tokens=True)[0]
+#     return transcription
+
+# def get_asr_docs(video_path, audio_path):
+
+#     full_transcription = []
+
+#     try:
+#         extract_audio(video_path, audio_path)
+#     except:
+#         return full_transcription
+#     audio_chunks = chunk_audio(audio_path, chunk_length_s=30)
+    
+#     for chunk in audio_chunks:
+#         transcription = transcribe_chunk(chunk)
+#         full_transcription.append(transcription)
+
+#     return full_transcription
+
+def get_ocr_docs(frames):
+    reader = easyocr.Reader(['en']) 
+    text_set = []
+    ocr_docs = []
+    for img in frames:
+        ocr_results = reader.readtext(img)
+        det_info = ""
+        for result in ocr_results:
+            text = result[1]
+            confidence = result[2]
+            if confidence > 0.5 and text not in text_set:
+                det_info += f"{text}; "
+                text_set.append(text)
+        if len(det_info) > 0:
+            ocr_docs.append(det_info)
+
+    return ocr_docs
+
+    
+def save_frames(frames):
+    file_paths = []
+    for i, frame in enumerate(frames):
+        img = Image.fromarray(frame)
+        file_path = f'restore/frame_{i}.png'
+        img.save(file_path)
+        file_paths.append(file_path)
+    return file_paths
+    
+def get_det_docs(frames, prompt):
+    prompt = ",".join(prompt)
+    frames_path = save_frames(frames)
+    res = []
+    if len(frames) > 0:
+        client_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
+        client_socket.connect(('0.0.0.0', 9999))
+        data = (frames_path, prompt)
+        client_socket.send(pickle.dumps(data))
+        result_data = client_socket.recv(4096)
         try:
-            json_match = re.search(r'```json\s*([\s\S]+?)\s*```', buffer)
-            if not json_match:
-                # If no JSON block is found, try to parse the whole buffer as a fallback.
-                try:
-                    layout_data = json.loads(buffer)
-                    markdown_content = layoutjson2md(layout_data)
-                    yield buffer, markdown_content, layout_data
-                    return
-                except json.JSONDecodeError:
-                    raise ValueError("JSON object not found in the model's output.")
-
-            json_str = json_match.group(1)
-            layout_data = json.loads(json_str)
-            markdown_content = layoutjson2md(layout_data)
-
-            yield buffer, markdown_content, layout_data
-        except Exception as e:
-            error_md = f"❌ **Error:** Failed to parse Layout JSON.\n\n**Details:**\n`{str(e)}`\n\n**Raw Output:**\n```\n{buffer}\n```"
-            error_json = {"error": "ProcessingError", "details": str(e), "raw_output": buffer}
-            yield buffer, error_md, error_json
-
-
-# --- Gradio UI Definition ---
-def create_gradio_interface():
-    """Builds and returns the Gradio web interface."""
-    css = """
-    .main-container { max-width: 1400px; margin: 0 auto; }
-    .process-button { border: none !important; color: white !important; font-weight: bold !important; background-color: blue !important;}
-    .process-button:hover { background-color: darkblue !important; transform: translateY(-2px) !important; box-shadow: 0 4px 8px rgba(0,0,0,0.2) !important; }
-    """
-    with gr.Blocks(theme="bethecloud/storj_theme", css=css) as demo:
-        gr.HTML("""
-        <div class="title" style="text-align: center">
-            <h1>Tiny VLMs Lab🧪</h1>
-            <p style="font-size: 1.1em; color: #6b7280; margin-bottom: 0.6em;">
-                Advanced Vision-Language Model for Image Content and Layout Extraction
-            </p>
-        </div>
-        """)
-
-        with gr.Row():
-            # Left Column (Inputs)
-            with gr.Column(scale=1):
-                model_choice = gr.Dropdown(
-                    choices=["Camel-Doc-OCR-080125",
-                             "MonkeyOCR-Recognition",
-                             "olmOCR-7B-0725",
-                             "Nanonets-OCR-s",
-                             "Megalodon-OCR-Sync-0713"
-                            ],
-                    label="Select Model", 
-                    value="Nanonets-OCR-s"
-                )
-                task_choice = gr.Dropdown(
-                    choices=["Content Extraction", 
-                             "Layout Analysis(.json)"],
-                    label="Select Task", value="Content Extraction"
-                )
-                image_input = gr.Image(label="Upload Image", type="pil", sources=['upload'])
-                with gr.Accordion("Advanced Settings", open=False):
-                    max_new_tokens = gr.Slider(minimum=512, maximum=8192, value=4096, step=256, label="Max New Tokens")
-
-                process_btn = gr.Button("🚀 Process Document", variant="primary", elem_classes=["process-button"], size="lg")
-                clear_btn = gr.Button("🗑️ Clear All", variant="secondary")
-
-            # Right Column (Outputs)
-            with gr.Column(scale=2):
-                with gr.Tabs() as tabs:
-                    with gr.Tab("📝 Extracted Content"):
-                        raw_output_stream = gr.Textbox(label="Raw Model Output Stream", interactive=False, lines=13, show_copy_button=True)
-                        with gr.Row():
-                            examples = gr.Examples(
-                                examples=["examples/1.png", "examples/2.png", "examples/3.png", "examples/4.png", "examples/5.png"],
-                                inputs=image_input,
-                                label="Examples"
-                        )
-                        gr.Markdown("[Report-Bug💻](https://huggingface.co/spaces/prithivMLmods/OCR-Comparator/discussions)")
-                    with gr.Tab("📰 README.md"):
-                        with gr.Accordion("(Formatted Result)", open=True):
-                            markdown_output = gr.Markdown(label="Formatted Markdown")
-
-                    with gr.Tab("📋 Layout Analysis Results"):
-                        json_output = gr.JSON(label="Structured Layout Data (JSON)")
-
-        # Event Handlers
-        def clear_all_outputs():
-            return None, "Raw output will appear here.", "Formatted results will appear here.", None
-
-        process_btn.click(
-            fn=process_document_stream,
-            inputs=[model_choice,
-                    task_choice,
-                    image_input,
-                    max_new_tokens],
-            outputs=[raw_output_stream,
-                     markdown_output,
-                     json_output]
+            res = pickle.loads(result_data)
+        except:
+            res = []
+    return res
+
+def det_preprocess(det_docs, location, relation, number):
+
+    scene_descriptions = []
+
+    for det_doc_per_frame in det_docs:
+        objects = []
+        scene_description = ""
+        if len(det_doc_per_frame) > 0:
+            for obj_id, objs in enumerate(det_doc_per_frame.split(";")):
+                obj_name = objs.split(":")[0].strip()
+                obj_bbox = objs.split(":")[1].strip()
+                obj_bbox = ast.literal_eval(obj_bbox)
+                objects.append({"id": obj_id, "label": obj_name, "bbox": obj_bbox})
+
+            scene_description = generate_scene_graph_description(objects, location, relation, number)
+        scene_descriptions.append(scene_description)
+    
+    return scene_descriptions
+
+# load your VLM
+device = "cuda"
+overwrite_config = {}
+tokenizer, model, image_processor, max_length = load_pretrained_model(
+    "lmms-lab/LLaVA-Video-7B-Qwen2", 
+    None, 
+    "llava_qwen", 
+    torch_dtype="bfloat16", 
+    device_map="auto", 
+    offload_buffers=True,
+    overwrite_config=overwrite_config)  # Add any other thing you want to pass in llava_model_args
+
+
+# 2) Check vocab sizes and fix BEFORE dispatching
+vsz_model = model.get_input_embeddings().weight.shape[0]
+vsz_tok   = len(tokenizer)
+if vsz_tok != vsz_model:
+    print(f"[fix] resizing embeddings: model={vsz_model} -> tokenizer={vsz_tok}")
+    model.resize_token_embeddings(vsz_tok)
+    # optional: init new rows
+    with torch.no_grad():
+        added = vsz_tok - vsz_model
+        if added > 0:
+            emb = model.get_input_embeddings().weight
+            emb[-added:].normal_(mean=0.0, std=0.02)
+
+
+model.eval()
+conv_template = "qwen_2"  # Make sure you use correct chat template for different models
+            
+
+# The inference function of your VLM
+def llava_inference(qs, video):
+    if video is not None:
+        question = DEFAULT_IMAGE_TOKEN + qs
+    else:
+        question = qs
+    conv = copy.deepcopy(conv_templates[conv_template])
+    conv.append_message(conv.roles[0], question)
+    conv.append_message(conv.roles[1], None)
+    prompt_question = conv.get_prompt()
+
+    input_ids = tokenizer_image_token(
+        prompt_question, tokenizer, IMAGE_TOKEN_INDEX, return_tensors="pt"
+    ).unsqueeze(0).to(device)
+    # input_ids = tokenizer_image_token(prompt_question, tokenizer, IMAGE_TOKEN_INDEX, return_tensors="pt").unsqueeze(0).to(device)
+    # cont = model.generate(
+    #     input_ids,
+    #     video=video,
+    #     modalities= ["video"],
+    #     do_sample=True,
+    #     temperature=0.7,
+    #     max_new_tokens=4096,
+    # )
+
+    if video is not None:
+        cont = model.generate(
+            input_ids,
+            images=video,
+            modalities=["video"],
+            do_sample=True,
+            temperature=0.7,
+            max_new_tokens=512
         )
-        clear_btn.click(
-            clear_all_outputs,
-            outputs=[image_input,
-                     raw_output_stream,
-                     markdown_output,
-                     json_output]
+    else:
+        cont = model.generate(
+            input_ids,
+            do_sample=True,
+            temperature=0.7,
+            max_new_tokens=512
         )
-    return demo
+        
+    text_outputs = tokenizer.batch_decode(cont, skip_special_tokens=True)[0].strip()
+    return text_outputs
+
+
+# super-parameters setting
+rag_threshold = 0.3
+clip_threshold = 0.3
+beta = 3.0 
+
+# Choose the auxiliary texts you want
+USE_OCR = True
+USE_ASR = False 
+USE_DET = True
+print(f"---------------OCR{rag_threshold}: {USE_OCR}-----------------")
+print(f"---------------ASR{rag_threshold}: {USE_ASR}-----------------")
+print(f"---------------DET{beta}-{clip_threshold}: {USE_DET}-----------------")
+print(f"---------------Frames: {max_frames_num}-----------------")
+
+# Create directories
+os.makedirs("restore/audio", exist_ok=True)
+os.makedirs("restore", exist_ok=True)
+
+def process_query(video_path, question):
+    if video_path is None:
+        return "Please upload a video."
+
+    frames, frame_time, video_time = process_video(video_path, max_frames_num, 1, force_sample=True)
+    raw_video = [f for f in frames]
+
+    video = image_processor.preprocess(frames, return_tensors="pt")["pixel_values"].cuda().bfloat16()
+    video = [video]
+
+    if USE_DET:
+        video_tensor = []
+        for frame in raw_video:
+            processed = clip_processor(images=frame, return_tensors="pt")["pixel_values"].to(clip_model.device, dtype=torch.float16)
+            video_tensor.append(processed.squeeze(0))
+        video_tensor = torch.stack(video_tensor, dim=0)
+
+    if USE_OCR:
+        ocr_docs_total = get_ocr_docs(frames)
+
+    if USE_ASR:
+        if os.path.exists(os.path.join("restore/audio", os.path.basename(video_path).split(".")[0] + ".txt")):
+            with open(os.path.join("restore/audio", os.path.basename(video_path).split(".")[0] + ".txt"), 'r', encoding='utf-8') as f:
+                asr_docs_total = f.readlines()
+        # else:
+        #     audio_path = os.path.join("restore/audio", os.path.basename(video_path).split(".")[0] + ".wav")
+        #     # asr_docs_total = get_asr_docs(video_path, audio_path)
+        #     with open(os.path.join("restore/audio", os.path.basename(video_path).split(".")[0] + ".txt"), 'w', encoding='utf-8') as f:
+        #         for doc in asr_docs_total:
+        #             f.write(doc + '\n')
+
+    # step 0: get cot information
+    retrieve_pmt_0 = "Question: " + question
+    # you can change this decouple prompt to fit your requirements
+    retrieve_pmt_0 += "\nTo answer the question step by step, you can provide your retrieve request to assist you by the following json format:"
+    retrieve_pmt_0 += '''{
+        "ASR": Optional[str]. The subtitles of the video that may relavent to the question you want to retrieve, in two sentences. If you no need for this information, please return null.
+        "DET": Optional[list]. (The output must include only physical entities, not abstract concepts, less than five entities) All the physical entities and their location related to the question you want to retrieve, not abstract concepts. If you no need for this information, please return null.
+        "TYPE": Optional[list]. (The output must be specified as null or a list containing only one or more of the following strings: 'location', 'number', 'relation'. No other values are valid for this field) The information you want to obtain about the detected objects. If you need the object location in the video frame, output "location"; if you need the number of specific object, output "number"; if you need the positional relationship between objects, output "relation". 
+    }
+    ## Example 1: 
+    Question: How many blue balloons are over the long table in the middle of the room at the end of this video? A. 1. B. 2. C. 3. D. 4.
+    Your retrieve can be:
+    {
+        "ASR": "The location and the color of balloons, the number of the blue balloons.",
+        "DET": ["blue ballons", "long table"],
+        "TYPE": ["relation", "number"]
+    }
+    ## Example 2: 
+    Question: In the lower left corner of the video, what color is the woman wearing on the right side of the man in black clothes? A. Blue. B. White. C. Red. D. Yellow.
+    Your retrieve can be:
+    {
+        "ASR": null,
+        "DET": ["the man in black", "woman"],
+        "TYPE": ["location", "relation"]
+    }
+    ## Example 3: 
+    Question: In which country is the comedy featured in the video recognized worldwide? A. China. B. UK. C. Germany. D. United States.
+    Your retrieve can be:
+    {
+        "ASR": "The country recognized worldwide for its comedy.",
+        "DET": null,
+        "TYPE": null
+    }
+    Note that you don't need to answer the question in this step, so you don't need any infomation about the video of image. You only need to provide your retrieve request (it's optional), and I will help you retrieve the infomation you want. Please provide the json format.'''
+
+    json_request = llava_inference(retrieve_pmt_0, None)
+
+    # step 1: get docs information
+    query = [question]
+
+    # APE fetch
+    if USE_DET:
+        det_docs = []
+        try:
+            request_det = json.loads(json_request)["DET"]
+            request_det = filter_keywords(request_det)
+            clip_text = ["A picture of " + txt for txt in request_det]
+            if len(clip_text) == 0:
+                clip_text = ["A picture of object"]
+        except:
+            request_det = None
+            clip_text = ["A picture of object"]
+
+        clip_inputs = clip_processor(text=clip_text, return_tensors="pt", padding=True, truncation=True).to(clip_model.device)
+        clip_img_feats = clip_model.get_image_features(video_tensor)
+        with torch.no_grad():
+            text_features = clip_model.get_text_features(**clip_inputs)
+            similarities = (clip_img_feats @ text_features.T).squeeze(0).mean(1).cpu()
+            similarities = np.array(similarities, dtype=np.float64)
+            alpha = beta * (len(similarities) / 16)
+            similarities = similarities * alpha / np.sum(similarities)
+
+        del clip_inputs, clip_img_feats, text_features
+        torch.cuda.empty_cache()
+
+        det_top_idx = [idx for idx in range(max_frames_num) if similarities[idx] > clip_threshold]
+            
+        if request_det is not None and len(request_det) > 0:
+            det_docs = get_det_docs(frames[det_top_idx], request_det)  
+
+            L, R, N = False, False, False
+            try:
+                det_retrieve_info = json.loads(json_request)["TYPE"]
+            except:
+                det_retrieve_info = None
+            if det_retrieve_info is not None:
+                if "location" in det_retrieve_info:
+                    L = True
+                if "relation" in det_retrieve_info:
+                    R = True
+                if "number" in det_retrieve_info:
+                    N = True
+            det_docs = det_preprocess(det_docs, location=L, relation=R, number=N)  # pre-process of APE information
+
+
+    # OCR fetch
+    if USE_OCR:
+        try:
+            request_det = json.loads(json_request)["DET"]
+            request_det = filter_keywords(request_det)
+        except:
+            request_det = None
+        ocr_docs = []
+        if len(ocr_docs_total) > 0:
+            ocr_query = query.copy()
+            if request_det is not None and len(request_det) > 0:
+                ocr_query.extend(request_det)
+            ocr_docs, _ = retrieve_documents_with_dynamic(ocr_docs_total, ocr_query, threshold=rag_threshold)
+
+    # ASR fetch
+    if USE_ASR:
+        asr_docs = []
+        try:
+            request_asr = json.loads(json_request)["ASR"]
+        except:
+            request_asr = None
+        if len(asr_docs_total) > 0:
+            asr_query = query.copy()
+            if request_asr is not None:
+                asr_query.append(request_asr)
+            asr_docs, _ = retrieve_documents_with_dynamic(asr_docs_total, asr_query, threshold=rag_threshold)
+
+    qs = ""
+    if USE_DET and len(det_docs) > 0:
+        for i, info in enumerate(det_docs):
+            if len(info) > 0:
+                qs += f"Frame {str(det_top_idx[i]+1)}: " + info + "\n"
+        if len(qs) > 0:
+            qs = f"\nVideo have {str(max_frames_num)} frames in total, the detected objects' information in specific frames: " + qs
+    if USE_ASR and len(asr_docs) > 0:
+        qs += "\nVideo Automatic Speech Recognition information (given in chronological order of the video): " + " ".join(asr_docs)
+    if USE_OCR and len(ocr_docs) > 0:
+        qs += "\nVideo OCR information (given in chronological order of the video): " + "; ".join(ocr_docs)
+    qs += "Select the best answer to the following multiple-choice question based on the video and the information (if given). Respond with only the letter (A, B, C, or D) of the correct option. Question: " + question  # you can change this prompt
+
+    res = llava_inference(qs, video)
+    return res
+
+demo = gr.Interface(
+    fn=process_query,
+    inputs=[gr.Video(label="Upload Video"), gr.Textbox(label="Question")],
+    outputs=gr.Textbox(label="Answer"),
+    title="Video Question Answering with LLaVA",
+    description="Upload a video and ask a question to get a summary or answer."
+)
 
 if __name__ == "__main__":
-    demo = create_gradio_interface()
-    demo.queue(max_size=50).launch(share=True, ssr_mode=False, show_error=True)
+    demo.launch()