cyan2k

/

molmo-7B-D-bnb-4bit

like 13

import sys
import os
from pathlib import Path

def set_cuda_paths():
    venv_base = Path(sys.executable).parent.parent
    nvidia_base_path = venv_base / 'Lib' / 'site-packages' / 'nvidia'
    cuda_path = nvidia_base_path / 'cuda_runtime' / 'bin'
    cublas_path = nvidia_base_path / 'cublas' / 'bin'
    cudnn_path = nvidia_base_path / 'cudnn' / 'bin'
    paths_to_add = [str(cuda_path), str(cublas_path), str(cudnn_path)]
    env_vars = ['CUDA_PATH', 'CUDA_PATH_V12_1', 'PATH']
    
    for env_var in env_vars:
        current_value = os.environ.get(env_var, '')
        new_value = os.pathsep.join(paths_to_add + [current_value] if current_value else paths_to_add)
        os.environ[env_var] = new_value

set_cuda_paths()

import torch
from PIL import Image
from transformers import AutoModelForCausalLM, AutoProcessor, GenerationConfig

model_path = r"D:\Scripts\bench_vision\cyan2k--molmo-7B-D-bnb-4bit"

class VisionModel:
    def __init__(self):
        self.model = None
        self.processor = None
        self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

    def initialize_model_and_processor(self):
        self.processor = AutoProcessor.from_pretrained(
            model_path,
            trust_remote_code=True,
            torch_dtype='auto',
            device_map='auto'
        )

        self.model = AutoModelForCausalLM.from_pretrained(
            model_path,
            trust_remote_code=True,
            torch_dtype='auto',
            device_map='auto'
        )

    def process_single_image(self, image_path):
        image = Image.open(image_path)

        # Ensure the image is in RGB format
        if image.mode != "RGB":
            image = image.convert("RGB")

        text = "Describe this image in detail as possible but be succinct and don't repeat yourself."
        inputs = self.processor.process(images=[image], text=text)
        inputs = {k: v.to(self.device).unsqueeze(0) for k, v in inputs.items()}

        output = self.model.generate_from_batch(
            inputs,
            GenerationConfig(max_new_tokens=500, stop_strings=["<|endoftext|>"]),
            tokenizer=self.processor.tokenizer
        )

        generated_tokens = output[0, inputs['input_ids'].size(1):]
        generated_text = self.processor.tokenizer.decode(generated_tokens, skip_special_tokens=True)

        print(f"\nGenerated Text:\n{generated_text}\n")

if __name__ == "__main__":
    image_path = r"D:\Scripts\bench_vision\IMG_140531.JPG"

    vision_model = VisionModel()
    vision_model.initialize_model_and_processor()
    vision_model.process_single_image(image_path)

def resize_and_pad(
    image: np.ndarray,
    desired_output_size: List[int],
    resize_method: str = "bilinear",
    pad_value: float = 0,
    normalize: bool = True,
    image_mean: Optional[List[float]] = OPENAI_CLIP_MEAN,
    image_std: Optional[List[float]] = OPENAI_CLIP_STD,
) -> (np.ndarray, np.ndarray):
    """
    Resize and pad the image to the desired output size.

    Args:
        image (np.ndarray): Input image as a NumPy array.
        desired_output_size (List[int]): Desired output size as [height, width].
        resize_method (str, optional): Resize interpolation method. Defaults to "bilinear".
        pad_value (float, optional): Padding value. Defaults to 0.
        normalize (bool, optional): Whether to normalize the image. Defaults to True.
        image_mean (Optional[List[float]], optional): Mean for normalization. Defaults to OPENAI_CLIP_MEAN.
        image_std (Optional[List[float]], optional): Standard deviation for normalization. Defaults to OPENAI_CLIP_STD.

    Returns:
        Tuple[np.ndarray, np.ndarray]: Resized and padded image, and image mask.
    """
    desired_height, desired_width = desired_output_size
    height, width = image.shape[:2]

    # Calculate scaling factors and determine the scaling factor to maintain aspect ratio
    scale_y = desired_height / height
    scale_x = desired_width / width
    scale = min(scale_x, scale_y)
    scaled_height = int(height * scale)
    scaled_width = int(width * scale)

    # Convert the image to a PyTorch tensor and normalize to [0, 1]
    image_tensor = torch.from_numpy(image).permute(2, 0, 1).float() / 255.0

    # Define the interpolation mode
    if resize_method.lower() == "bilinear":
        interpolation = InterpolationMode.BILINEAR
    elif resize_method.lower() == "nearest":
        interpolation = InterpolationMode.NEAREST
    elif resize_method.lower() == "bicubic":
        interpolation = InterpolationMode.BICUBIC
    elif resize_method.lower() == "lanczos":
        interpolation = InterpolationMode.LANCZOS
    else:
        raise ValueError(f"Unsupported resize method: {resize_method}")

    # Resize the image
    resized_image = torchvision.transforms.Resize(
        [scaled_height, scaled_width],
        interpolation=interpolation,
        antialias=True
    )(image_tensor)

    # Clip the image to ensure values are within [0, 1]
    resized_image = torch.clamp(resized_image, 0.0, 1.0)

    # Convert back to NumPy
    resized_image_np = resized_image.permute(1, 2, 0).numpy()

    # Calculate padding
    top_pad = (desired_height - scaled_height) // 2
    bottom_pad = desired_height - scaled_height - top_pad
    left_pad = (desired_width - scaled_width) // 2
    right_pad = desired_width - scaled_width - left_pad

    # Pad the image using NumPy
    padded_image = np.pad(
        resized_image_np,
        pad_width=((top_pad, bottom_pad), (left_pad, right_pad), (0, 0)),
        mode='constant',
        constant_values=pad_value
    )

    # Create the image mask
    image_mask = np.pad(
        np.ones((scaled_height, scaled_width), dtype=bool),
        pad_width=((top_pad, bottom_pad), (left_pad, right_pad)),
        mode='constant',
        constant_values=False
    )

    if normalize:
        padded_image = normalize_image(padded_image, offset=image_mean, scale=image_std)

    return padded_image, image_mask