[upload demo notebook]

deepattricap-vla-3b-colab-notebook-demo/DeepAttriCap_VLA_3B.ipynb

@@ -0,0 +1,397 @@
+{
+  "cells": [
+    {
+      "cell_type": "markdown",
+      "source": [
+        "## **DeepAttriCap-VLA-3B**\n"
+      ],
+      "metadata": {
+        "id": "XN7ploRWa4GE"
+      }
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "uFovmijgUV1Z"
+      },
+      "source": [
+        "The DeepAttriCap-VLA-3B model is a fine-tuned version of Qwen2.5-VL-3B-Instruct, tailored for Vision-Language Attribution and Image Captioning. This variant is designed to generate precise, attribute-rich descriptions that define the visual properties of objects and scenes in detail, ensuring both object-level identification and contextual captioning. Vision-Language Attribution: Produces structured captions with explicit object attributes, properties, and contextual details.\n",
+        "\n",
+        "\n",
+        "| IMG 1 | IMG 2 |\n",
+        "|-------|-------|\n",
+        "| ![Screenshot 2025-08-28 at 21-30-25 Gradio.png](https://cdn-uploads.huggingface.co/production/uploads/65bb837dbfb878f46c77de4c/sOsUkjrn4ElKetpQ_ap6t.png) | ![Screenshot 2025-08-28 at 21-30-52 Gradio.png](https://cdn-uploads.huggingface.co/production/uploads/65bb837dbfb878f46c77de4c/IXxvpDXVyFXcnLSQoP1NN.png) |\n",
+        "\n",
+        "\n",
+        "notebook by : [prithivMLmods](https://huggingface.co/prithivMLmods)🤗"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "RugX4SGZV-8O"
+      },
+      "source": [
+        "### **Install Packages**"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "l-NtFtjSpuJQ"
+      },
+      "outputs": [],
+      "source": [
+        "%%capture\n",
+        "!pip install git+https://github.com/huggingface/transformers.git \\\n",
+        "             git+https://github.com/huggingface/accelerate.git \\\n",
+        "             git+https://github.com/huggingface/peft.git \\\n",
+        "             transformers-stream-generator huggingface_hub albumentations \\\n",
+        "             pyvips-binary qwen-vl-utils sentencepiece opencv-python docling-core \\\n",
+        "             python-docx torchvision safetensors matplotlib num2words \\\n",
+        "\n",
+        "!pip install xformers requests pymupdf hf_xet spaces pyvips pillow gradio \\\n",
+        "             einops torch fpdf timm av decord bitsandbytes reportlab\n",
+        "#Hold tight, this will take around 2-3 minutes."
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "mvoSnRZcVBu4"
+      },
+      "source": [
+        "### **Run DeepAttriCap-VLA-3B Demo**\n",
+        "\n"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "source": [
+        "from huggingface_hub import notebook_login, HfApi\n",
+        "notebook_login()"
+      ],
+      "metadata": {
+        "id": "db6vq0T8deoB"
+      },
+      "execution_count": null,
+      "outputs": []
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "tElKr2Fkp1bO"
+      },
+      "outputs": [],
+      "source": [
+        "import spaces\n",
+        "import json\n",
+        "import math\n",
+        "import os\n",
+        "import traceback\n",
+        "from io import BytesIO\n",
+        "from typing import Any, Dict, List, Optional, Tuple\n",
+        "import re\n",
+        "import time\n",
+        "from threading import Thread\n",
+        "from io import BytesIO\n",
+        "import uuid\n",
+        "import tempfile\n",
+        "\n",
+        "import gradio as gr\n",
+        "import requests\n",
+        "import torch\n",
+        "from PIL import Image\n",
+        "import fitz\n",
+        "\n",
+        "from transformers import (\n",
+        "    Qwen2_5_VLForConditionalGeneration,\n",
+        "    AutoProcessor,\n",
+        "    TextIteratorStreamer,\n",
+        ")\n",
+        "\n",
+        "from transformers.image_utils import load_image\n",
+        "\n",
+        "from reportlab.lib.pagesizes import A4\n",
+        "from reportlab.lib.styles import getSampleStyleSheet\n",
+        "from reportlab.platypus import SimpleDocTemplate, Image as RLImage, Paragraph, Spacer\n",
+        "from reportlab.lib.units import inch\n",
+        "\n",
+        "# --- Constants and Model Setup ---\n",
+        "MAX_INPUT_TOKEN_LENGTH = 4096\n",
+        "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
+        "\n",
+        "# --- System Prompt Definition ---\n",
+        "CAPTION_SYSTEM_PROMPT = \"\"\"\n",
+        "You are an AI assistant that rigorously follows this response protocol:\n",
+        "\n",
+        "1. For every input image, your primary task is to write a **precise caption**. The caption must capture the **essence of the image** in clear, concise, and contextually accurate language.\n",
+        "\n",
+        "2. Along with the caption, provide a structured set of **attributes** that describe the visual elements. Attributes should include details such as objects, people, actions, colors, environment, mood, and other notable characteristics.\n",
+        "\n",
+        "3. Always include a **class_name** field. This must represent the **core theme or main subject** of the image in a compact format.\n",
+        "   - Use the syntax: `{class_name==write_the_core_theme}`\n",
+        "   - Example: `{class_name==dog_playing}` or `{class_name==city_sunset}`\n",
+        "\n",
+        "4. Maintain the following strict format in your output:\n",
+        "   - **Caption:** <one-sentence description>\n",
+        "   - **Attributes:** <comma-separated list of visual attributes>\n",
+        "   - **{class_name==core_theme}**\n",
+        "\n",
+        "5. Ensure captions are **precise, neutral, and descriptive**, avoiding unnecessary elaboration or subjective interpretation unless explicitly required.\n",
+        "\n",
+        "6. Do not reference the rules or instructions in the output. Only return the formatted caption, attributes, and class_name.\n",
+        "\n",
+        "\"\"\".strip()\n",
+        "\n",
+        "\n",
+        "print(\"CUDA_VISIBLE_DEVICES=\", os.environ.get(\"CUDA_VISIBLE_DEVICES\"))\n",
+        "print(\"torch.__version__ =\", torch.__version__)\n",
+        "print(\"torch.version.cuda =\", torch.version.cuda)\n",
+        "print(\"cuda available:\", torch.cuda.is_available())\n",
+        "print(\"cuda device count:\", torch.cuda.device_count())\n",
+        "if torch.cuda.is_available():\n",
+        "    print(\"current device:\", torch.cuda.current_device())\n",
+        "    print(\"device name:\", torch.cuda.get_device_name(torch.cuda.current_device()))\n",
+        "\n",
+        "print(\"Using device:\", device)\n",
+        "\n",
+        "# --- Model Loading: prithivMLmods/DeepAttriCap-VLA-3B ---\n",
+        "MODEL_ID_N = \"prithivMLmods/DeepAttriCap-VLA-3B\"\n",
+        "processor = AutoProcessor.from_pretrained(MODEL_ID_N, trust_remote_code=True)\n",
+        "model = Qwen2_5_VLForConditionalGeneration.from_pretrained(\n",
+        "    MODEL_ID_N, trust_remote_code=True, dtype=torch.float16\n",
+        ").to(device).eval()\n",
+        "\n",
+        "\n",
+        "# --- PDF Generation and Preview Utility Function ---\n",
+        "def generate_and_preview_pdf(image: Image.Image, text_content: str, font_size: int, line_spacing: float, alignment: str, image_size: str):\n",
+        "    \"\"\"\n",
+        "    Generates a PDF, saves it, and then creates image previews of its pages.\n",
+        "    Returns the path to the PDF and a list of paths to the preview images.\n",
+        "    \"\"\"\n",
+        "    if image is None or not text_content or not text_content.strip():\n",
+        "        raise gr.Error(\"Cannot generate PDF. Image or text content is missing.\")\n",
+        "\n",
+        "    # --- 1. Generate the PDF ---\n",
+        "    temp_dir = tempfile.gettempdir()\n",
+        "    pdf_filename = os.path.join(temp_dir, f\"output_{uuid.uuid4()}.pdf\")\n",
+        "    doc = SimpleDocTemplate(\n",
+        "        pdf_filename,\n",
+        "        pagesize=A4,\n",
+        "        rightMargin=inch, leftMargin=inch,\n",
+        "        topMargin=inch, bottomMargin=inch\n",
+        "    )\n",
+        "    styles = getSampleStyleSheet()\n",
+        "    style_normal = styles[\"Normal\"]\n",
+        "    style_normal.fontSize = int(font_size)\n",
+        "    style_normal.leading = int(font_size) * line_spacing\n",
+        "    style_normal.alignment = {\"Left\": 0, \"Center\": 1, \"Right\": 2, \"Justified\": 4}[alignment]\n",
+        "\n",
+        "    story = []\n",
+        "\n",
+        "    img_buffer = BytesIO()\n",
+        "    image.save(img_buffer, format='PNG')\n",
+        "    img_buffer.seek(0)\n",
+        "\n",
+        "    page_width, _ = A4\n",
+        "    available_width = page_width - 2 * inch\n",
+        "    image_widths = {\n",
+        "        \"Small\": available_width * 0.3,\n",
+        "        \"Medium\": available_width * 0.6,\n",
+        "        \"Large\": available_width * 0.9,\n",
+        "    }\n",
+        "    img_width = image_widths[image_size]\n",
+        "    img = RLImage(img_buffer, width=img_width, height=image.height * (img_width / image.width))\n",
+        "    story.append(img)\n",
+        "    story.append(Spacer(1, 12))\n",
+        "\n",
+        "    cleaned_text = re.sub(r'#+\\s*', '', text_content).replace(\"*\", \"\")\n",
+        "    text_paragraphs = cleaned_text.split('\\n')\n",
+        "\n",
+        "    for para in text_paragraphs:\n",
+        "        if para.strip():\n",
+        "            story.append(Paragraph(para, style_normal))\n",
+        "\n",
+        "    doc.build(story)\n",
+        "\n",
+        "    # --- 2. Render PDF pages as images for preview ---\n",
+        "    preview_images = []\n",
+        "    try:\n",
+        "        pdf_doc = fitz.open(pdf_filename)\n",
+        "        for page_num in range(len(pdf_doc)):\n",
+        "            page = pdf_doc.load_page(page_num)\n",
+        "            pix = page.get_pixmap(dpi=150)\n",
+        "            preview_img_path = os.path.join(temp_dir, f\"preview_{uuid.uuid4()}_p{page_num}.png\")\n",
+        "            pix.save(preview_img_path)\n",
+        "            preview_images.append(preview_img_path)\n",
+        "        pdf_doc.close()\n",
+        "    except Exception as e:\n",
+        "        print(f\"Error generating PDF preview: {e}\")\n",
+        "\n",
+        "    return pdf_filename, preview_images\n",
+        "\n",
+        "\n",
+        "# --- Core Application Logic ---\n",
+        "@spaces.GPU\n",
+        "def process_document_stream(\n",
+        "    image: Image.Image,\n",
+        "    prompt_input: str,\n",
+        "    max_new_tokens: int,\n",
+        "    temperature: float,\n",
+        "    top_p: float,\n",
+        "    top_k: int,\n",
+        "    repetition_penalty: float\n",
+        "):\n",
+        "    \"\"\"\n",
+        "    Main generator function that handles model inference tasks with advanced generation parameters.\n",
+        "    \"\"\"\n",
+        "    if image is None:\n",
+        "        yield \"Please upload an image.\", \"\"\n",
+        "        return\n",
+        "    if not prompt_input or not prompt_input.strip():\n",
+        "        yield \"Please enter a prompt.\", \"\"\n",
+        "        return\n",
+        "\n",
+        "    # Integrate the system prompt\n",
+        "    messages = [\n",
+        "        {\"role\": \"system\", \"content\": CAPTION_SYSTEM_PROMPT},\n",
+        "        {\"role\": \"user\", \"content\": [{\"type\": \"image\", \"image\": image}, {\"type\": \"text\", \"text\": prompt_input}]}\n",
+        "    ]\n",
+        "\n",
+        "    prompt_full = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)\n",
+        "    inputs = processor(text=[prompt_full], images=[image], return_tensors=\"pt\", padding=True, truncation=True, max_length=MAX_INPUT_TOKEN_LENGTH).to(device)\n",
+        "\n",
+        "    streamer = TextIteratorStreamer(processor, skip_prompt=True, skip_special_tokens=True)\n",
+        "\n",
+        "    generation_kwargs = {\n",
+        "        **inputs,\n",
+        "        \"streamer\": streamer,\n",
+        "        \"max_new_tokens\": max_new_tokens,\n",
+        "        \"temperature\": temperature,\n",
+        "        \"top_p\": top_p,\n",
+        "        \"top_k\": top_k,\n",
+        "        \"repetition_penalty\": repetition_penalty,\n",
+        "        \"do_sample\": True\n",
+        "    }\n",
+        "\n",
+        "    thread = Thread(target=model.generate, kwargs=generation_kwargs)\n",
+        "    thread.start()\n",
+        "\n",
+        "    buffer = \"\"\n",
+        "    for new_text in streamer:\n",
+        "        buffer += new_text\n",
+        "        buffer = buffer.replace(\"<|im_end|>\", \"\")\n",
+        "        time.sleep(0.01)\n",
+        "        yield buffer , buffer\n",
+        "\n",
+        "    yield buffer, buffer\n",
+        "\n",
+        "\n",
+        "# --- Gradio UI Definition ---\n",
+        "def create_gradio_interface():\n",
+        "    \"\"\"Builds and returns the Gradio web interface.\"\"\"\n",
+        "    css = \"\"\"\n",
+        "    .main-container { max-width: 1400px; margin: 0 auto; }\n",
+        "    .process-button { border: none !important; color: white !important; font-weight: bold !important; background-color: blue !important;}\n",
+        "    .process-button:hover { background-color: darkblue !important; transform: translateY(-2px) !important; box-shadow: 0 4px 8px rgba(0,0,0,0.2) !important; }\n",
+        "    #gallery { min-height: 400px; }\n",
+        "    \"\"\"\n",
+        "    with gr.Blocks(theme=\"bethecloud/storj_theme\", css=css) as demo:\n",
+        "        gr.HTML(\"\"\"\n",
+        "        <div class=\"title\" style=\"text-align: center\">\n",
+        "            <h1>DeepAttriCap-VLA-3B 👀</h1>\n",
+        "            <p style=\"font-size: 1.1em; color: #6b7280; margin-bottom: 0.6em;\">\n",
+        "                Using DeepAttriCap-VLA-3B for Image Captioning and Understanding\n",
+        "            </p>\n",
+        "        </div>\n",
+        "        \"\"\")\n",
+        "\n",
+        "        with gr.Row():\n",
+        "            # Left Column (Inputs)\n",
+        "            with gr.Column(scale=1):\n",
+        "                gr.Textbox(\n",
+        "                    label=\"Selected Model\",\n",
+        "                    value=\"DeepAttriCap-VLA-3B\",\n",
+        "                    interactive=False\n",
+        "                )\n",
+        "                prompt_input = gr.Textbox(label=\"Query Input\", placeholder=\"✦︎ Enter your query\", value=\"Describe the image!\")\n",
+        "                image_input = gr.Image(label=\"Upload Image\", type=\"pil\", sources=['upload'])\n",
+        "\n",
+        "                with gr.Accordion(\"Advanced Settings\", open=False):\n",
+        "                    max_new_tokens = gr.Slider(minimum=512, maximum=8192, value=4096, step=256, label=\"Max New Tokens\")\n",
+        "                    temperature = gr.Slider(label=\"Temperature\", minimum=0.1, maximum=4.0, step=0.1, value=0.6)\n",
+        "                    top_p = gr.Slider(label=\"Top-p (nucleus sampling)\", minimum=0.05, maximum=1.0, step=0.05, value=0.9)\n",
+        "                    top_k = gr.Slider(label=\"Top-k\", minimum=1, maximum=1000, step=1, value=50)\n",
+        "                    repetition_penalty = gr.Slider(label=\"Repetition penalty\", minimum=1.0, maximum=2.0, step=0.05, value=1.2)\n",
+        "\n",
+        "                    gr.Markdown(\"### PDF Export Settings\")\n",
+        "                    font_size = gr.Dropdown(choices=[\"8\", \"10\", \"12\", \"14\", \"16\", \"18\"], value=\"12\", label=\"Font Size\")\n",
+        "                    line_spacing = gr.Dropdown(choices=[1.0, 1.15, 1.5, 2.0], value=1.15, label=\"Line Spacing\")\n",
+        "                    alignment = gr.Dropdown(choices=[\"Left\", \"Center\", \"Right\", \"Justified\"], value=\"Justified\", label=\"Text Alignment\")\n",
+        "                    image_size = gr.Dropdown(choices=[\"Small\", \"Medium\", \"Large\"], value=\"Medium\", label=\"Image Size in PDF\")\n",
+        "\n",
+        "                process_btn = gr.Button(\"🚀 Process Image\", variant=\"primary\", elem_classes=[\"process-button\"], size=\"lg\")\n",
+        "                clear_btn = gr.Button(\"🗑️ Clear All\", variant=\"secondary\")\n",
+        "\n",
+        "            # Right Column (Outputs)\n",
+        "            with gr.Column(scale=2):\n",
+        "                with gr.Tabs() as tabs:\n",
+        "                    with gr.Tab(\"📝 Extracted Content\"):\n",
+        "                        raw_output_stream = gr.Textbox(label=\"Raw Model Output Stream\", interactive=False, lines=15, show_copy_button=True)\n",
+        "                    with gr.Tab(\"📰 README.md\"):\n",
+        "                        with gr.Accordion(\"(Result.md)\", open=True):\n",
+        "                            markdown_output = gr.Markdown()\n",
+        "\n",
+        "                    with gr.Tab(\"📋 PDF Preview\"):\n",
+        "                        generate_pdf_btn = gr.Button(\"📄 Generate PDF & Render\", variant=\"primary\")\n",
+        "                        pdf_output_file = gr.File(label=\"Download Generated PDF\", interactive=False)\n",
+        "                        pdf_preview_gallery = gr.Gallery(label=\"PDF Page Preview\", show_label=True, elem_id=\"gallery\", columns=2, object_fit=\"contain\", height=\"auto\")\n",
+        "\n",
+        "        # Event Handlers\n",
+        "        def clear_all_outputs():\n",
+        "            return None, \"\", \"Raw output will appear here.\", \"\", None, None\n",
+        "\n",
+        "        process_btn.click(\n",
+        "            fn=process_document_stream,\n",
+        "            inputs=[image_input, prompt_input, max_new_tokens, temperature, top_p, top_k, repetition_penalty],\n",
+        "            outputs=[raw_output_stream, markdown_output]\n",
+        "        )\n",
+        "\n",
+        "        generate_pdf_btn.click(\n",
+        "            fn=generate_and_preview_pdf,\n",
+        "            inputs=[image_input, raw_output_stream, font_size, line_spacing, alignment, image_size],\n",
+        "            outputs=[pdf_output_file, pdf_preview_gallery]\n",
+        "        )\n",
+        "\n",
+        "        clear_btn.click(\n",
+        "            clear_all_outputs,\n",
+        "            outputs=[image_input, prompt_input, raw_output_stream, markdown_output, pdf_output_file, pdf_preview_gallery]\n",
+        "        )\n",
+        "    return demo\n",
+        "\n",
+        "if __name__ == \"__main__\":\n",
+        "    demo = create_gradio_interface()\n",
+        "    demo.queue(max_size=50).launch(share=True, ssr_mode=False, show_error=True)"
+      ]
+    }
+  ],
+  "metadata": {
+    "accelerator": "GPU",
+    "colab": {
+      "gpuType": "T4",
+      "provenance": []
+    },
+    "kernelspec": {
+      "display_name": "Python 3",
+      "name": "python3"
+    },
+    "language_info": {
+      "name": "python"
+    }
+  },
+  "nbformat": 4,
+  "nbformat_minor": 0
+}