From def1bea883d439c67a511e595920cba68e2346c1 Mon Sep 17 00:00:00 2001
From: huchukato <huchukato@gmail.com>
Date: Fri, 14 Aug 2026 20:03:29 +0200
Subject: [PATCH] Add ONNX download validation and multi-source fallback

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>
---
 __init__.py  | 23 ++++++++++++++++++++---
 utilities.py | 32 +++++++++++++++++++++++++++++++-
 2 files changed, 51 insertions(+), 4 deletions(-)

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'<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}"
+            )
\ No newline at end of file