playback_giza_model_torch_captioning.py

# Playback █████

```
█████████████████████████████████████████████████████████████████████████████████████████████
██████████████████████████▓░░░░░░░░░░░░░░░░░░▒▒▒░░░░░░░░░░░░░▒███████████████████████████████
██████████████████████████▓░░░░░░░░░░░░░░░░░░▒▓▓▒░░░░░░░░░░░░▒███████████████████████████████
██████████████████████████▓░░░░░░░░░░░░░░░░░░▒▓▓▓▒░░░░░░░░░░░▒███████████████████████████████
██████████████████████████▓░░░░░░░░░░░░░░░░░░▒▓▓▓▓░░░░░░░░░░░▒███████████████████████████████
██████████████████████████▓░░░░░░░░░░░░░░░░░░▒██▓▓░░░░░░░░░░░▒███████████████████████████████
██████████████████████████▓▒▒▒▒▒▒▒▒▓▓▓▒▒▒▒▒▒▒▓█▓▓▓▒▒▓▒▒▒▒▒▒▒▒▒███████████████████████████████
██████████████████████████▓▒▒▒▒▒▒██████▓▒▒▒▒▒▒█▓██████▓▓▒▒▒▒▒▒███████████████████████████████
██████████████████████████▓░░░░▒█▓▓████▓▓▒░░░▒▓███▓▓▓▓▓▓▓▓▒░░▒███████████████████████████████
██████████████████████████▓░░░▒▓█▓▓▓▓▓▓▓▓▓▒░░▒▓▓▓▓▓▓▓▓▓▓▓▓▓▒░▒███████████████████████████████
██████████████████████████▓░░▒▓█▓██▓▓█████▓▒░▒██▓▓██▓█▓███▓▒░▒███████████████████████████████
██████████████████████████▓░▒▓███▓▒▒▒▒▓▓██▓▓░▒█▓▓▓▒▒▒▒▒▓███▓▒▒███████████████████████████████
██████████████████████████▓▒▓████▒▒▒▒▒▒▓█▓█▓▓▓▓▓█▒▒▒▒▒▒▒██▓▓▓▒███████████████████████████████
██████████████████████████▓▒▓█▓▓▓▒▒▒▒▒▒▒▒▓▓███▓▒▒▒▒▒▒▒▒▒███▓▓▒███████████████████████████████
██████████████████████████▓░▓▓█▓▓░░░░░░░░▓▓█▓▓▓▒░░░░░░░▒█▓▓▓▓▒███████████████████████████████
██████████████████████████▓░▒▓▓▓▒░░░░░░░░░▒▓▓▓▓▓▒░░░░░░▒▓▓▓▓▓▒███████████████████████████████
██████████████████████████▓░▒▓▓▓▓▒░░░░░▒▓▓░▓▓▓▓▓▒▒░░░░░▓▓▓▓▓▓▒███████████████████████████████
██████████████████████████▓░▓▓█▓▓▓▒░░░▒▓██▓▒▒████▓▒░░░░███▓▓▒▒███████████████████████████████
██████████████████████████▓░▓▓▓▓▓▓▓▓▓▓▓████▓▒▓▓▓▓▓▓▓▓▓▓▓▓▓▓▒▒▒███████████████████████████████
██████████████████████████▓▒▓█▓▓▓▓▓█████▓█▓▒▒▒▓███████████▓▒▒▒███████████████████████████████
██████████████████████████▓░▓▓▓███████▓▓█▓▒░░░▒▓▓▓▓▓▓▓▓▓▓▓▒░░▒███████████████████████████████
██████████████████████████▓░▓▓▓▓▓▓▓▓▓▓▓▓▓▒░░░░░░▒▓▓▓▓▓▓▒▒░░░░▒███████████████████████████████
██████████████████████████▓░▓▓▓▓▓▒▒▒▒▒▒▒░░░░░░░░░░▒▒▒▒▒░░░░░░▒███████████████████████████████
██████████████████████████▓░▓███▓░░░░░░░░░░░░░░░░░░░░░░░░░░░░▒███████████████████████████████
██████████████████████████▓░▓██▓▒░░░░░░░░░░░░░░░░░░░░░░░░░░░░▒███████████████████████████████
██████████████████████████▓▒▒▒▒▒▒▒▒▒▒▒▒▒▒▒▒▒▒▒▒▒▒▒▒▒▒▒▒▒▒▒▒▒▒▒███████████████████████████████
█████████████████████████████████████████████████████████████████████████████████████████████
█████████▓▓████▓▓▓▓▓▓▓▓▓▓████████▓▓███▓▓▓▓▓███▓██████████▓▓▓████████▓▓▓▓████████▓▓████▓▓▓▓███
███▓▓▓▓▓██▓▓██▓▓▓▓▓▓▓▓▓▓▓██▓▓▓▓██▓▓▓██▓▓▓▓▓██▓▓███▓▓▓▓▓██▓▓▓██▓▓▓▓▓██▓▓▓██▓▓▓▓▓█▓▓▓██▓▓▓▓▓██▓
██▓▓▒▓▓▓██▓▓██▓▓▒▓▓▓▓▓▓▓██▓▓▓▒▓▓██▓▓▓██▓▓▓██▓▓▓███▒▓▓▓▒███▒███▓▓▓▓▓██▓▓██▓▓▒▓▓▓██▓▓██▓▓▓▓███▒
█████████▒░▒██▒░░░░░░░░▒██▒░░░░▒██▒░░░█████▒░░░▓████████▒░░██▓░░░░░██▓░██▓░░░░░░░░▒██████▒░░░
██▓▒▒▒▒▒▒▒▒▒██▓▒▒▒▒▒▒▒▒▒██████████▒▒▒▒▒▓██▒▒▒▒▒▓██▒▒▒▒▒██▓▒██████████▓▒██▓▒▒▒▒▒██▒▒██▓▒▒▒▓█▓▒
██▓▒▒▒▒▒▒▒▒▒██▓▓▓▓▓▓▓▓▒▒██▓▓▓▓▓▓██▒▒▒▒▒▓██▒▒▒▒▒▓██▓▓▓▓▓█▓▓▒███▓▓▓▓▓██▓▒▓█▓▓▓▓▓▓█▓▒▒██▒▒▒▒▒▓█▓
██▓▒▒▒▒▒▒▒▒▒██████████▓▓██▒▒▒▒▒▓██▒▒▒▒▒▓██▒▒▒▒▒▓████████▓▒▒██▓▒▒▒▒▒██▓▒▒▓███████▒▒▒██▓▒▒▒▒▓██
█████████████████████████████████████████████████████████████████████████████████████████████


```

import torch
from transformers import BlipProcessor, BlipForConditionalGeneration
import torch.onnx

def run_playback(model, processor, onnx_file_path, image_paths = ['https://s3.amazonaws.com/model-server/inputs/0xc0d08ed5b0f759cbc528abf16ae6e2fb33f935379a7b1fa2182753f1019fa721_0.jpg','https://s3.amazonaws.com/model-server/inputs/0xc0d08ed5b0f759cbc528abf16ae6e2fb33f935379a7b1fa2182753f1019fa721_1.jpg']):
    """Converts a PyTorch BLIP model to ONNX format.

    This function takes a BLIP model, processor, and output path, 
    and converts the model to ONNX format for optimized inference.

    Args:
        model (BlipForConditionalGeneration): The BLIP model to convert.
        processor (BlipProcessor): The BLIP processor.
        onnx_file_path (str): The path to save the ONNX model to.
        image_paths (list, optional): A list of image URLs to use for creating a 
            dummy input. Defaults to two sample image URLs.
    """

    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    model.to(device)
    model.eval()  # Set the model to evaluation mode

    # Create a dummy input using the provided image paths
    dummy_input = processor(images=image_paths, return_tensors="pt").to(device)
    dummy_input = dummy_input['pixel_values']  # Only take pixel_values for the input

    # Export the model to ONNX format
    torch.onnx.export(
        model,
        dummy_input,
        onnx_file_path,
        export_params=True, 
        opset_version=10, 
        do_constant_folding=True,
        input_names=['input'], 
        output_names=['output'],
        dynamic_axes={'input': {0: 'batch_size'}, 'output': {0: 'batch_size'}}
    )

    print(f"Model has been converted to ONNX and saved as {onnx_file_path}")


def execution():
    """Loads the BLIP model and processor.

    Returns:
        tuple: A tuple containing the BLIP model and processor.
    """
    model_name = "Salesforce/blip-image-captioning-base"
    processor = BlipProcessor.from_pretrained(model_name)
    model = BlipForConditionalGeneration.from_pretrained(model_name)
    return model, processor

# Specify the output path for the ONNX model
onnx_file_path = "blip_image_captioning_model.onnx"

# Load the model and processor
model, processor = execution()

# Convert the model to ONNX
run_playback(model, processor, onnx_file_path)