transformers - 💡(How to fix) Fix Sam3Video: CUDA out of memory [3 comments, 2 participants]

from transformers import Sam3VideoConfig, Sam3VideoModel, Sam3VideoProcessor
from transformers.video_utils import load_video
import sys
import os
import cv2
import numpy as np
import json
import time
import gc
import torch
import math
from datetime import datetime

import inspect
import os
def print_memory(message=""):
    if not torch.cuda.is_available():
        print(f"[Line {inspect.currentframe().f_back.f_lineno}] {message} - CUDA not available")
        return
    
    caller_frame = inspect.currentframe().f_back
    line_no = caller_frame.f_lineno
    
    allocated = torch.cuda.memory_allocated() / 1024**2  # MB
    reserved = torch.cuda.memory_reserved() / 1024**2    # MB
    max_allocated = torch.cuda.max_memory_allocated() / 1024**2
    
    print(f"[Line {line_no:4d}] {message} - Allocated: {allocated}")

if __name__ == "__main__":
    sam_path = "./sam3/"
    video_path = "./reverse.mp4"
    device = 'cuda' if torch.cuda.is_available() else 'cpu'
    #	device = Accelerator().device
    config = Sam3VideoConfig.from_pretrained(sam_path)
    config.image_size = 1008
    model = Sam3VideoModel.from_pretrained(sam_path).to(device, dtype=torch.bfloat16)
    # processor = Sam3VideoProcessor.from_pretrained(sam_path)
    processor = Sam3VideoProcessor.from_pretrained(sam_path, size={"height": 1008, "width": 1008})

    frame_count = 0
    detection_results = []
    track_ids_set = set()
    color_map = {}

    frame_mask_status = []

    batch_counter = 0

    video_frames, _ = load_video(video_path)
    frames_num = video_frames.shape[0]

    inference_session = processor.init_video_session(
        video=video_frames[0],
        inference_device=device,
        processing_device=device,
        video_storage_device='cpu',
        dtype=torch.bfloat16,
    )
    text = "person"
    inference_session = processor.add_text_prompt(
        inference_session=inference_session,
        text=text,
    )
    for idx in range(1, frames_num):
        processed_video = processor.video_processor(videos=video_frames[idx], device=device, return_tensors="pt")
        pixel_values_video = processed_video.pixel_values_videos[0]
        inference_session.add_new_frame(pixel_values_video)

    with torch.no_grad():
        total_model_outputs = model.propagate_in_video_iterator(
            inference_session=inference_session
        )
        print(f"111")

        for model_outputs in total_model_outputs:
            print_memory("test!!!")

            result = processor.postprocess_outputs(inference_session, model_outputs)

            print_memory("before total_model_outputs!!!")
        
        # print(f"{result['object_ids']}")
        test = 0