huchukatoandDevin
def1bea883
The download_file() function now validates: - HTTP status code (fails on non-200) - File size (removes files <1KB as likely error pages) - ONNX header (detects HTML pages masquerading as .onnx files) The ONNX download in AutoLoadRifeTensorrtModel now tries multiple HuggingFace sources before giving up, with clear error message pointing users to manual export_onnx.py as fallback. This prevents the cryptic "factory function returned nullptr" error when TensorRT attempts to parse a corrupted/empty/HTML file as ONNX. Generated with [Devin](https://devin.ai) Co-Authored-By: Devin <158243242+devin-ai-integration[bot]@users.noreply.github.com>
131 lines
4.3 KiB
Python
131 lines
4.3 KiB
Python
import requests
|
|
from tqdm import tqdm
|
|
import logging
|
|
import sys
|
|
|
|
class ColoredLogger:
|
|
COLORS = {
|
|
'RED': '\033[91m',
|
|
'GREEN': '\033[92m',
|
|
'YELLOW': '\033[93m',
|
|
'BLUE': '\033[94m',
|
|
'MAGENTA': '\033[95m',
|
|
'RESET': '\033[0m'
|
|
}
|
|
|
|
LEVEL_COLORS = {
|
|
'DEBUG': COLORS['BLUE'],
|
|
'INFO': COLORS['GREEN'],
|
|
'WARNING': COLORS['YELLOW'],
|
|
'ERROR': COLORS['RED'],
|
|
'CRITICAL': COLORS['MAGENTA']
|
|
}
|
|
|
|
def __init__(self, name="MY-APP"):
|
|
self.logger = logging.getLogger(name)
|
|
self.logger.setLevel(logging.DEBUG)
|
|
self.app_name = name
|
|
|
|
# Prevent message propagation to parent loggers
|
|
self.logger.propagate = False
|
|
|
|
# Clear existing handlers
|
|
self.logger.handlers = []
|
|
|
|
# Create console handler
|
|
handler = logging.StreamHandler(sys.stdout)
|
|
handler.setLevel(logging.DEBUG)
|
|
|
|
# Custom formatter class to handle colored components
|
|
class ColoredFormatter(logging.Formatter):
|
|
def format(self, record):
|
|
# Color the level name according to severity
|
|
level_color = ColoredLogger.LEVEL_COLORS.get(record.levelname, '')
|
|
colored_levelname = f"{level_color}{record.levelname}{ColoredLogger.COLORS['RESET']}"
|
|
|
|
# Color the logger name in blue
|
|
colored_name = f"{ColoredLogger.COLORS['BLUE']}{record.name}{ColoredLogger.COLORS['RESET']}"
|
|
|
|
# Set the colored components
|
|
record.levelname = colored_levelname
|
|
record.name = colored_name
|
|
|
|
return super().format(record)
|
|
|
|
# Create formatter with the new format
|
|
formatter = ColoredFormatter('[%(name)s|%(levelname)s] - %(message)s')
|
|
handler.setFormatter(formatter)
|
|
|
|
self.logger.addHandler(handler)
|
|
|
|
|
|
def debug(self, message):
|
|
self.logger.debug(f"{self.COLORS['BLUE']}{message}{self.COLORS['RESET']}")
|
|
|
|
def info(self, message):
|
|
self.logger.info(f"{self.COLORS['GREEN']}{message}{self.COLORS['RESET']}")
|
|
|
|
def warning(self, message):
|
|
self.logger.warning(f"{self.COLORS['YELLOW']}{message}{self.COLORS['RESET']}")
|
|
|
|
def error(self, message):
|
|
self.logger.error(f"{self.COLORS['RED']}{message}{self.COLORS['RESET']}")
|
|
|
|
def critical(self, message):
|
|
self.logger.critical(f"{self.COLORS['MAGENTA']}{message}{self.COLORS['RESET']}")
|
|
|
|
def download_file(url, save_path):
|
|
"""
|
|
Download a file from URL with progress bar
|
|
|
|
Args:
|
|
url (str): URL of the file to download
|
|
save_path (str): Path to save the file as
|
|
|
|
Raises:
|
|
RuntimeError: If the download fails (HTTP error or empty file)
|
|
"""
|
|
GREEN = '\033[92m'
|
|
RESET = '\033[0m'
|
|
response = requests.get(url, stream=True)
|
|
|
|
if response.status_code != 200:
|
|
raise RuntimeError(
|
|
f"Download failed with HTTP {response.status_code} for {url}. "
|
|
f"The file may not exist at this URL."
|
|
)
|
|
|
|
total_size = int(response.headers.get('content-length', 0))
|
|
|
|
with open(save_path, 'wb') as file, tqdm(
|
|
desc=save_path,
|
|
total=total_size,
|
|
unit='iB',
|
|
unit_scale=True,
|
|
unit_divisor=1024,
|
|
colour='green',
|
|
bar_format=f'{GREEN}{{l_bar}}{{bar}}{RESET}{GREEN}{{r_bar}}{RESET}'
|
|
) as progress_bar:
|
|
for data in response.iter_content(chunk_size=1024):
|
|
size = file.write(data)
|
|
progress_bar.update(size)
|
|
|
|
# Validate the downloaded file
|
|
import os
|
|
file_size = os.path.getsize(save_path)
|
|
if file_size < 1024:
|
|
os.remove(save_path)
|
|
raise RuntimeError(
|
|
f"Downloaded file is too small ({file_size} bytes) — likely an error page. "
|
|
f"URL: {url}"
|
|
)
|
|
|
|
# Check ONNX magic bytes (protobuf tag for onnx.ModelProto starts with \x08)
|
|
if save_path.endswith('.onnx'):
|
|
with open(save_path, 'rb') as f:
|
|
header = f.read(16)
|
|
if b'<html' in header.lower() or b'<!doctype' in header.lower():
|
|
os.remove(save_path)
|
|
raise RuntimeError(
|
|
f"Downloaded file is HTML, not ONNX. The model may not exist at: {url}"
|
|
) |