THUDM

/

GLM-4.1V-9B-Thinking

like 645

Follow

Z.ai & THUKEG 2.69k

import transformers.image_transforms as image_transforms
from typing import Union, List
import torch

# Store original function
_original_group_images_by_shape = image_transforms.group_images_by_shape

def patched_group_images_by_shape(images: List[torch.Tensor], disable_grouping: bool = False):
    """Patched version with disable_grouping parameter."""
    if disable_grouping:
        # Process individually - return single group
        if not images:
            return {}, {}
        
        # Create single group with all images
        first_shape = images[0].shape[1:]
        return {first_shape: torch.stack(images, dim=0)}, {i: (first_shape, i) for i in range(len(images))}
    
    # Use original function for grouped processing
    return _original_group_images_by_shape(images)

# Apply monkey patch
image_transforms.group_images_by_shape = patched_group_images_by_shape

# Also patch BaseImageProcessorFast if needed
if hasattr(transformers, 'BaseImageProcessorFast'):
    original_preprocess = transformers.BaseImageProcessorFast._preprocess
    
    def patched_preprocess(self, images, **kwargs):
        # Add disable_grouping parameter if missing
        if 'disable_grouping' not in kwargs:
            kwargs['disable_grouping'] = not (hasattr(images[0], 'device') and images[0].device.type == 'cuda')
        return original_preprocess(self, images, **kwargs)
    
    transformers.BaseImageProcessorFast._preprocess = patched_preprocess

# uses apple/aimv2-huge-patch14-336 whereas Ovis2 uses the large 448

USE_QUANTIZATION = True  # Set to True to use BitsAndBytesConfig, False for no quantization

import sys, os, time, threading, queue
from pathlib import Path
import pynvml, torch
from PIL import Image
from transformers import Glm4vForConditionalGeneration, AutoProcessor

if USE_QUANTIZATION:
    from transformers import BitsAndBytesConfig

def set_cuda_paths():
    venv_base = Path(sys.executable).parent.parent
    nvidia_base_path = venv_base / 'Lib' / 'site-packages' / 'nvidia'
    paths_to_add = [
        str(nvidia_base_path / 'cuda_runtime' / 'bin'),
        str(nvidia_base_path / 'cuda_runtime' / 'bin' / 'lib' / 'x64'),
        str(nvidia_base_path / 'cuda_runtime' / 'include'),
        str(nvidia_base_path / 'cublas' / 'bin'),
        str(nvidia_base_path / 'cudnn' / 'bin'),
        str(nvidia_base_path / 'cuda_nvrtc' / 'bin'),
        str(nvidia_base_path / 'cuda_nvcc' / 'bin'),
    ]
    current_value = os.environ.get('PATH', '')
    os.environ['PATH'] = os.pathsep.join(paths_to_add + [current_value] if current_value else paths_to_add)
    os.environ['CUDA_PATH'] = str(nvidia_base_path / 'cuda_runtime')

def monitor_vram(vram_queue, handle, stop_flag, interval=0.1):
    max_usage = 0
    while not stop_flag.is_set():
        usage = pynvml.nvmlDeviceGetMemoryInfo(handle).used / 1024 ** 2
        max_usage = max(max_usage, usage)
        vram_queue.put(max_usage)
        time.sleep(interval)

set_cuda_paths()

quantization_config = None
if USE_QUANTIZATION:
    quantization_config = BitsAndBytesConfig(
        load_in_4bit=True,
        bnb_4bit_compute_dtype=torch.bfloat16,
        bnb_4bit_quant_type="nf4",
    )

pynvml.nvmlInit()
gpu_handle = pynvml.nvmlDeviceGetHandleByIndex(0)
baseline_usage = pynvml.nvmlDeviceGetMemoryInfo(gpu_handle).used / 1024 ** 2

model_path = r"[PATH TO THE FOLDER CONTAINING THE MODEL AND RELATED FILES]"
image_path = r"[PATH TO A SPECIFIC IMAGE TO PROCESS"

dtype = torch.bfloat16
device = "cuda" if torch.cuda.is_available() else "cpu"

processor = AutoProcessor.from_pretrained(model_path, use_fast=True)

model_kwargs = {
    "torch_dtype": dtype,
    "device_map": "auto",
    "attn_implementation": "sdpa",
}

if USE_QUANTIZATION:
    model_kwargs["quantization_config"] = quantization_config

model = Glm4vForConditionalGeneration.from_pretrained(
    model_path,
    **model_kwargs
).eval()

image = Image.open(image_path)
user_prompt = "Describe in English as much detail as possible what this image depicts?"

messages = [
    {
        "role": "user",
        "content": [
            {
                "type": "image",
                "image": image
            },
            {
                "type": "text", 
                "text": user_prompt
            },
        ],
    }
]

inputs = processor.apply_chat_template(
    messages,
    add_generation_prompt=True,
    tokenize=True,
    return_dict=True,
    return_tensors="pt"
).to(device)

vram_queue = queue.Queue()
stop_flag = threading.Event()
vram_thread = threading.Thread(target=monitor_vram, args=(vram_queue, gpu_handle, stop_flag))
vram_thread.start()

start = time.time()
with torch.inference_mode():
    out_ids = model.generate(**inputs, max_new_tokens=1024, do_sample=False)

if USE_QUANTIZATION:
    torch.cuda.synchronize()

end = time.time()

stop_flag.set()
vram_thread.join()

peak_vram = 0
while not vram_queue.empty():
    peak_vram = max(peak_vram, vram_queue.get())

model_vram = peak_vram - baseline_usage
elapsed = end - start

generated_ids_trimmed = [
    out_ids[0][len(inputs.input_ids[0]):]
]
response = processor.batch_decode(
    generated_ids_trimmed,
    skip_special_tokens=True,
    clean_up_tokenization_spaces=False,
)[0].strip()

print(f"\n{response}\n")
print(f"Max VRAM Usage: {model_vram:.2f} MB")
print(f"Characters per second: {len(response) / elapsed:.2f}")
print(f"Compute Time: {elapsed:.2f} seconds")
print(f"Quantization: {'Enabled (4-bit)' if USE_QUANTIZATION else 'Disabled'}")