app.py

import gradio as gr
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer, StoppingCriteria
import gradio as gr
import torch
import numpy as np
import torch
import torchvision.transforms as T
from PIL import Image
from torchvision.transforms.functional import InterpolationMode
from transformers import AutoModel, AutoTokenizer
import os
from threading import Thread
import re
import time 
from PIL import Image
import torch
import spaces
import subprocess
from time import sleep
import base64
from io import BytesIO

subprocess.run('pip install flash-attn --no-build-isolation', env={'FLASH_ATTENTION_SKIP_CUDA_BUILD': "TRUE"}, shell=True)

# torch.set_default_device('cuda')

IMAGENET_MEAN = (0.485, 0.456, 0.406)
IMAGENET_STD = (0.229, 0.224, 0.225)

def build_transform(input_size):
    MEAN, STD = IMAGENET_MEAN, IMAGENET_STD
    transform = T.Compose([
        T.Lambda(lambda img: img.convert('RGB') if img.mode != 'RGB' else img),
        T.Resize((input_size, input_size), interpolation=InterpolationMode.BICUBIC),
        T.ToTensor(),
        T.Normalize(mean=MEAN, std=STD)
    ])
    return transform

def find_closest_aspect_ratio(aspect_ratio, target_ratios, width, height, image_size):
    best_ratio_diff = float('inf')
    best_ratio = (1, 1)
    area = width * height
    for ratio in target_ratios:
        target_aspect_ratio = ratio[0] / ratio[1]
        ratio_diff = abs(aspect_ratio - target_aspect_ratio)
        if ratio_diff < best_ratio_diff:
            best_ratio_diff = ratio_diff
            best_ratio = ratio
        elif ratio_diff == best_ratio_diff:
            if area > 0.5 * image_size * image_size * ratio[0] * ratio[1]:
                best_ratio = ratio
    return best_ratio

def dynamic_preprocess(image, min_num=1, max_num=12, image_size=448, use_thumbnail=False):
    orig_width, orig_height = image.size
    aspect_ratio = orig_width / orig_height

    # calculate the existing image aspect ratio
    target_ratios = set(
        (i, j) for n in range(min_num, max_num + 1) for i in range(1, n + 1) for j in range(1, n + 1) if
        i * j <= max_num and i * j >= min_num)
    target_ratios = sorted(target_ratios, key=lambda x: x[0] * x[1])

    # find the closest aspect ratio to the target
    target_aspect_ratio = find_closest_aspect_ratio(
        aspect_ratio, target_ratios, orig_width, orig_height, image_size)

    # calculate the target width and height
    target_width = image_size * target_aspect_ratio[0]
    target_height = image_size * target_aspect_ratio[1]
    blocks = target_aspect_ratio[0] * target_aspect_ratio[1]

    # resize the image
    resized_img = image.resize((target_width, target_height))
    processed_images = []
    for i in range(blocks):
        box = (
            (i % (target_width // image_size)) * image_size,
            (i // (target_width // image_size)) * image_size,
            ((i % (target_width // image_size)) + 1) * image_size,
            ((i // (target_width // image_size)) + 1) * image_size
        )
        # split the image
        split_img = resized_img.crop(box)
        processed_images.append(split_img)
    assert len(processed_images) == blocks
    if use_thumbnail and len(processed_images) != 1:
        thumbnail_img = image.resize((image_size, image_size))
        processed_images.append(thumbnail_img)
    return processed_images

def load_image(image_file, input_size=448, max_num=12):
    if isinstance(image_file, str):  # Check if it's a file path
        image = Image.open(image_file).convert('RGB')
    else:  # Assume it's a base64 string
        image_data = base64.b64decode(image_file)
        image = Image.open(BytesIO(image_data)).convert('RGB')
    
    transform = build_transform(input_size=input_size)
    images = dynamic_preprocess(image, image_size=input_size, use_thumbnail=True, max_num=max_num)
    pixel_values = [transform(image) for image in images]
    pixel_values = torch.stack(pixel_values)
    return pixel_values

model_name = "lycaoduong/KLPintern-v2-1B"

access_token = os.getenv("LPInternVL21B")

model = AutoModel.from_pretrained(
    model_name,
    torch_dtype=torch.bfloat16,
    low_cpu_mem_usage=True,
    trust_remote_code=True,
    token=access_token
).eval()

tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True, use_fast=False, token=access_token)


@spaces.GPU
def chat(message, history):
    # print(history)
    # print(message)
    model.to('cuda')
    
    test_image = message["files"][0]
        
    pixel_values = load_image(test_image, max_num=12).to(torch.bfloat16).cuda()
    generation_config = dict(max_new_tokens= 1024, do_sample=True, num_beams = 3, repetition_penalty=2.5)
    
    question = '<image>\n'+message["text"]
    response, conv_history = model.chat(tokenizer, pixel_values, question, generation_config, history=None, return_history=True)

    # print(f'User: {question}\nAssistant: {response}')

    txt_stream = ''
    for c in response:
        sleep(0.01)
        txt_stream = txt_stream + c
        yield txt_stream

    # return response


CSS ="""
# @media only screen and (max-width: 600px){
#     #component-3 {
#       height: 90dvh !important;
#       transform-origin: top; /* Ensure that the element expands from top to bottom. */
#       border-style: solid;
#       overflow: hidden;
#       flex-grow: 1;
#       min-width: min(160px, 100%);
#       border-width: var(--block-border-width);
#     }
# }
#component-3 {
  height: 50dvh !important;
  transform-origin: top; /* Ensure that the element expands from top to bottom. */
  border-style: solid;
  overflow: hidden;
  flex-grow: 1;
  min-width: min(160px, 100%);
  border-width: var(--block-border-width);
}
/* Ensure that the image inside the button is displayed correctly for buttons with a specified aria-label. */
button.svelte-1lcyrx4[aria-label="user's message: a file of type image/jpeg, "] img.svelte-1pijsyv {
  width: 100%;
  object-fit: contain;
  height: 100%;
  border-radius: 13px; /* Add rounded corners to the image. */
  max-width: 50vw;     /* Limit the image width. */
}
/* Set the height for the button and allow text selection only for buttons with a specified aria-label. */
button.svelte-1lcyrx4[aria-label="user's message: a file of type image/jpeg, "] {
  user-select: text;
  text-align: left;
  height: 300px;
}
/* Add border-radius and limit the width for images that do not belong to the avatar container. */
.message-wrap.svelte-1lcyrx4 > div.svelte-1lcyrx4 .svelte-1lcyrx4:not(.avatar-container) img {
  border-radius: 13px;
  max-width: 50vw;
}
.message-wrap.svelte-1lcyrx4 .message.svelte-1lcyrx4 img {
    margin: var(--size-2);
    max-height: 500px;
}
"""


demo = gr.ChatInterface(
    fn=chat,
    description="Test Korean License Plate OCR VLMs in this demo.",
    examples=[{"text": "Extract LPR information.", "files":["./demo/P0.jpg"]},
              {"text": "Extract LPR information.", "files":["./demo/P1.jpg"]},
              {"text": "Extract LPR information.", "files":["./demo/P2.jpg"]},
              {"text": "Extract LPR information, return in JSON format", "files":["./demo/P3.jpg"]},
              {"text": "Extract LPR information.", "files":["./demo/P4.jpg"]},
              {"text": "Extract LPR information.", "files":["./demo/P5.jpg"]},
              {"text": "Extract LPR information.", "files":["./demo/P6.jpg"]},
              {"text": "Extract LPR information.", "files":["./demo/P7.jpg"]},
              {"text": "Extract LPR information.", "files":["./demo/P8.jpg"]},
              {"text": "Extract LPR information.", "files":["./demo/P9.jpg"]},
              {"text": "Extract LPR information.", "files":["./demo/P10.jpg"]},
                ],
    title="❄️ InternVL2-1B fine-tuned for Korean License Plate recognition.❄️",
    multimodal=True,
    css=CSS
)
demo.queue().launch()