Update app.py

app.py

@@ -43,42 +43,56 @@ import cv2
 import gradio as gr
 import numpy as np
 import onnxruntime as ort
+from huggingface_hub import hf_hub_download
 
-# Local file configuration path generated by our conversion script
-MODEL_PATH = "RealESRGAN_x4plus_anime_6B.onnx"
+# Verified public model tracks configured with fully dynamic shape axes natively
+MODELS = {
+    "RealESRGAN_x4plus (Dynamic - Super Fast)": {
+        "repo_id": "fastnlp/RealESRGAN_x4plus_onnx",
+        "filename": "model.onnx"
+    }
+}
+
+current_model_name = None
 ort_session = None
 
-def get_inference_session():
+def load_model(model_choice):
     """
-    Loads the locally baked dynamic shape ONNX model instance into memory.
+    Downloads the official pre-compiled dynamic-shape ONNX engine directly.
+    Bypasses structural weight remapping issues entirely.
     """
-    global ort_session
-    if ort_session is not None:
+    global current_model_name, ort_session
+    if current_model_name == model_choice and ort_session is not None:
         return ort_session
-        
-    if not os.path.exists(MODEL_PATH):
-        # Auto-compile file right now if the build step wasn't run separately
-        print(f"{MODEL_PATH} not found. Running inline build conversion...")
-        import subprocess
-        # Assumes export_onnx.py is placed in the same working directory
-        subprocess.run([sys.executable, "export_onnx.py"], check=True)
+
+    cfg = MODELS[model_choice]
+    print(f"Loading weights for {model_choice}...")
+    token = os.environ.get("HF_TOKEN")
+    
+    model_path = hf_hub_download(
+        repo_id=cfg["repo_id"], 
+        filename=cfg["filename"],
+        token=token
+    )
 
     session_options = ort.SessionOptions()
     session_options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL
     
-    ort_session = ort.InferenceSession(MODEL_PATH, session_options, providers=['CPUExecutionProvider'])
+    # Initialize session on CPU with dynamic array dimensions enabled
+    ort_session = ort.InferenceSession(model_path, session_options, providers=['CPUExecutionProvider'])
+    current_model_name = model_choice
     return ort_session
 
-def upscale_image(input_img):
+def upscale_image(input_img, model_choice):
     """
-    Dynamic inference pipeline. Sends the full image array to the custom ONNX graph 
-    without splitting into tiles, maximizing multi-threaded CPU efficiency.
+    Dynamic inference pipeline. Transposes and runs full image dimensions 
+    in a single, fast forward operation on your CPU layer.
     """
     if input_img is None:
         return None
         
     try:
-        session = get_inference_session()
+        session = load_model(model_choice)
         
         # Preprocessing: Normalize image array data to float32 range [0.0, 1.0]
         img_float = input_img.astype(np.float32) / 255.0
@@ -89,15 +103,15 @@ def upscale_image(input_img):
         # Add batch dimension: (1, C, H, W)
         img_batch = np.expand_dims(img_chw, axis=0)
         
-        # Bind input/output keys dynamically matching the graph properties
+        # Bind input/output identifiers matching the schema map
         input_name = session.get_inputs()[0].name
         output_name = session.get_outputs()[0].name
         
-        # Run raw CPU forward tensor calculations on the entire canvas in one step
+        # Run raw CPU forward calculation on the entire canvas in one go
         ort_outs = session.run([output_name], {input_name: img_batch})
         output_tensor = ort_outs[0]
         
-        # Postprocessing: Drop the batch dimension and clip bounds safely back to valid channels
+        # Postprocessing: Drop batch dimension and clip tensor bounds safely
         output_tensor = np.squeeze(output_tensor, axis=0)
         output_tensor = np.clip(output_tensor, 0.0, 1.0)
         
@@ -118,19 +132,24 @@ def upscale_image(input_img):
 # 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 custom compiled dynamic weights.")
+    gr.Markdown("Running locally on Hugging Face Free CPU hardware using dynamic tensor shape processing.")
     
     with gr.Row():
         with gr.Column():
             input_image = gr.Image(label="Source Image", type="numpy")
-            submit_btn = gr.Button("Upscale Image (4x Super Fast)", variant="primary")
+            model_dropdown = gr.Dropdown(
+                choices=list(MODELS.keys()),
+                value="RealESRGAN_x4plus (Dynamic - Super Fast)",
+                label="Select AI Upscaling Engine"
+            )
+            submit_btn = gr.Button("Upscale Image (Dynamic 4x)", variant="primary")
             
         with gr.Column():
             output_image = gr.Image(label="Enhanced Super-Resolution Result", type="numpy")
 
     submit_btn.click(
         fn=upscale_image,
-        inputs=input_image,
+        inputs=[input_image, model_dropdown],
         outputs=output_image
     )