diff --git a/__init__.py b/__init__.py index c297524..294c93e 100644 --- a/__init__.py +++ b/__init__.py @@ -453,9 +453,26 @@ class AutoLoadRifeTensorrtModel: if not os.path.exists(tensorrt_model_path): if not os.path.exists(onnx_model_path): - onnx_model_download_url = f"https://huggingface.co/yuvraj108c/rife-onnx/resolve/main/{model}.onnx" - rife_logger.info(f"Downloading {onnx_model_download_url}") - download_file(url=onnx_model_download_url, save_path=onnx_model_path) + # Try multiple download sources in case the primary one fails + onnx_download_urls = [ + f"https://huggingface.co/yuvraj108c/rife-onnx/resolve/main/{model}.onnx", + f"https://huggingface.co/marduk191/rife/resolve/main/{model}.onnx", + ] + download_success = False + for url in onnx_download_urls: + try: + rife_logger.info(f"Downloading {url}") + download_file(url=url, save_path=onnx_model_path) + download_success = True + break + except Exception as e: + rife_logger.warning(f"Download failed from {url}: {e}") + continue + if not download_success: + raise RuntimeError( + f"Failed to download ONNX model '{model}.onnx' from all sources. " + f"You can manually export it using export_onnx.py and place it in: {onnx_model_path}" + ) else: rife_logger.info(f"ONNX model found at: {onnx_model_path}") diff --git a/utilities.py b/utilities.py index 0362ed6..5ac54b8 100644 --- a/utilities.py +++ b/utilities.py @@ -81,10 +81,20 @@ def download_file(url, save_path): 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( @@ -98,4 +108,24 @@ def download_file(url, save_path): ) as progress_bar: for data in response.iter_content(chunk_size=1024): size = file.write(data) - progress_bar.update(size) \ No newline at end of file + 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'