Update app.py

app.py

@@ -1,7 +1,6 @@
 import sys
 import importlib
 import os
-import urllib.request
 
 # =================================================================
 # CRITICAL PYTHON 3.13 / GRADIO SDK CONFLICT MONKEY-PATCHES
@@ -13,7 +12,7 @@ try:
     if not hasattr(real_hf_hub, 'HfFolder'):
         class MockHfFolder:
             @classmethod
-            def get_token(cls): return None
+            def get_token(cls): return os.environ.get("HF_TOKEN")
             @classmethod
             def save_token(cls, token): pass
             @classmethod
@@ -44,18 +43,17 @@ import cv2
 import gradio as gr
 import numpy as np
 import onnxruntime as ort
+from huggingface_hub import hf_hub_download
 
-# Public Hugging Face CDN URL patterns bypassing authenticated internal Space routing loops
+# Map open model checkpoints using verified community organization targets
 MODELS = {
     "RealESRGAN_x2plus (Faster 2x)": {
-        "url": "https://huggingface.co/onnx-community/RealESRGAN_x2plus_onnx/blob/main/model.onnx",
-        "cdn_url": "https://huggingface.co/onnx-community/RealESRGAN_x2plus_onnx/resolve/main/model.onnx",
-        "local_name": "RealESRGAN_x2plus.onnx"
+        "repo_id": "onnx-community/RealESRGAN_x2plus_onnx",
+        "filename": "model.onnx",
     },
     "RealESRGAN_x4plus (General 4x)": {
-        "url": "https://huggingface.co/onnx-community/RealESRGAN_x4plus_onnx/blob/main/model.onnx",
-        "cdn_url": "https://huggingface.co/onnx-community/RealESRGAN_x4plus_onnx/resolve/main/model.onnx",
-        "local_name": "RealESRGAN_x4plus.onnx"
+        "repo_id": "onnx-community/RealESRGAN_x4plus_onnx",
+        "filename": "model.onnx",
     }
 }
 
@@ -65,8 +63,9 @@ ort_session = None
 
 def load_model(model_choice):
     """
-    Downloads model weights cleanly using standard streams while managing custom headers
-    to prevent the Hugging Face gateway from blocking the local container context.
+    Switches the active ONNX runtime model configuration.
+    Uses the native hf_hub_download wrapper with the system-mounted HF_TOKEN
+    to seamlessly authenticate against the Hugging Face CDN gateway.
     """
     global current_model_name, ort_session
     
@@ -74,29 +73,17 @@ def load_model(model_choice):
         return ort_session
 
     cfg = MODELS[model_choice]
-    cache_dir = "/tmp/models"
-    os.makedirs(cache_dir, exist_ok=True)
-    model_path = os.path.join(cache_dir, cfg["local_name"])
-
-    # Download only if the model file is not already cached locally
-    if not os.path.exists(model_path):
-        print(f"Downloading weights for {model_choice} via open CDN routing stream...")
-        try:
-            # We configure standard user agents and referers to completely satisfy the Hugging Face 
-            # cloud edge firewall rules, allowing tokenless open downloads inside the Space sandbox.
-            req = urllib.request.Request(
-                cfg["cdn_url"],
-                headers={
-                    'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36',
-                    'Accept': '*/*',
-                    'Connection': 'keep-alive'
-                }
-            )
-            with urllib.request.urlopen(req) as response, open(model_path, 'wb') as out_file:
-                out_file.write(response.read())
-            print(f"Successfully cached model file locally at: {model_path}")
-        except Exception as dl_err:
-            raise RuntimeError(f"Direct weight initialization stream failed: {dl_err}")
+    
+    print(f"Loading weights for {model_choice} using system environment authorization...")
+    
+    # Retrieve the auto-mounted internal space authorization token if present
+    token = os.environ.get("HF_TOKEN")
+    
+    model_path = hf_hub_download(
+        repo_id=cfg["repo_id"], 
+        filename=cfg["filename"],
+        token=token
+    )
 
     # Configure ONNX runtime parameters optimized for CPU processing execution
     session_options = ort.SessionOptions()
@@ -155,7 +142,7 @@ def upscale_image(input_img, model_choice):
 # Define the user interface layout
 with gr.Blocks(title="AI Lightweight Image Upscaler (ONNX)") as demo:
     gr.Markdown("# 🖼️ AI Image Resizer & Upscaler (ONNX Engine)")
-    gr.Markdown("Running locally on Hugging Face Free CPU hardware using direct public weight streams.")
+    gr.Markdown("Running locally on Hugging Face Free CPU hardware using official `onnx-community` model tracks.")
     
     with gr.Row():
         with gr.Column():