File size: 2,548 Bytes
426b1bc 0cafef1 0f9d8ab 4ebd87c 426b1bc 4ebd87c 426b1bc 0f9d8ab 8574f97 0f9d8ab 8574f97 4ebd87c 426b1bc 0cafef1 8574f97 0cafef1 8574f97 0cafef1 8574f97 0cafef1 f33b49d 426b1bc 8574f97 0cafef1 426b1bc 0cafef1 8574f97 426b1bc 0cafef1 426b1bc 0cafef1 8574f97 f33b49d 0cafef1 f33b49d 0cafef1 8574f97 426b1bc 0cafef1 8574f97 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 |
import torch
from transformers import AutoProcessor, AutoModelForCausalLM
from ultralytics import YOLO
import gdown
import os
from safetensors.torch import load_file # Safetensors loading method
# Set a custom cache directory for gdown
gdown_cache_dir = os.path.join(os.getcwd(), "cache")
os.makedirs(gdown_cache_dir, exist_ok=True)
os.environ["GDOWN_CACHE"] = gdown_cache_dir # Explicitly set GDOWN_CACHE
def download_model_from_drive(file_id, destination_path):
"""Download the model from Google Drive using gdown."""
# Construct the Google Drive download URL
url = f"https://drive.google.com/uc?id={file_id}"
# Ensure the destination directory exists
directory = os.path.dirname(destination_path)
if directory:
os.makedirs(directory, exist_ok=True)
# Download the file
gdown.download(url, destination_path, quiet=False)
def load_models(device='cpu'):
"""Load YOLO model and the caption generation model."""
# Define the file path for the .safetensors model
model_file_path = "model.safetensors" # Adjust based on your file name
# Download the model file if it doesn't exist
if not os.path.exists(model_file_path):
file_id = "1hUCqZ3X8mcM-KcwWFjcsFg7PA0hUvE3k" # Replace with your file ID
print(f"Downloading model to {model_file_path}...")
download_model_from_drive(file_id, model_file_path)
# Load the YOLO model
print("Loading YOLO model...")
yolo_model = YOLO("best.pt").to(device)
# Load the processor for the caption model
print("Loading processor for the caption model...")
processor = AutoProcessor.from_pretrained(
"microsoft/Florence-2-base",
trust_remote_code=True
)
# Load the caption model state dict from .safetensors
print("Loading caption generation model...")
model_state_dict = load_file(model_file_path) # Load tensors from .safetensors
caption_model = AutoModelForCausalLM.from_pretrained(
"microsoft/Florence-2-base",
trust_remote_code=True
)
caption_model.load_state_dict(model_state_dict) # Map tensors to the model
caption_model.to(device) # Move the model to the correct device
print("Models loaded successfully!")
return {
'yolo_model': yolo_model,
'processor': processor,
'caption_model': caption_model
}
# Usage example:
if __name__ == "__main__":
device = "cuda" if torch.cuda.is_available() else "cpu"
models = load_models(device=device)
print("All models are ready to use!")
|