Update app.py

app.py

@@ -6,7 +6,7 @@ from PIL import Image
 
 # === Upscaler Logic ===
 def pre_process(img: np.array) -> np.array:
-    img = np.transpose(img[:, :, 0:3], (2, 0, 1))
+    img = np.transpose(img[:, :, :3], (2, 0, 1))
     return np.expand_dims(img, axis=0).astype(np.float32)
 
 
@@ -14,20 +14,30 @@ def post_process(img: np.array) -> np.array:
     img = np.squeeze(img)
     return np.transpose(img, (1, 2, 0))[:, :, ::-1].astype(np.uint8)
 
-
-# ONNX inference with session cache
-_session_cache = {}
-def inference(model_path: str, img_array: np.array) -> np.array:
-    if model_path not in _session_cache:
+# ONNX inference with GPU if available
+def get_session(model_path: str):
+    # cache sessions
+    if model_path not in get_session.cache:
         opts = onnxruntime.SessionOptions()
+        # multi-threading
         opts.intra_op_num_threads = 1
         opts.inter_op_num_threads = 1
-        _session_cache[model_path] = onnxruntime.InferenceSession(model_path, opts)
-    session = _session_cache[model_path]
+        # try GPU first
+        providers = ['CUDAExecutionProvider', 'CPUExecutionProvider']
+        get_session.cache[model_path] = onnxruntime.InferenceSession(
+            model_path,
+            sess_options=opts,
+            providers=providers
+        )
+    return get_session.cache[model_path]
+get_session.cache = {}
+
+
+def inference(model_path: str, img_array: np.array) -> np.array:
+    session = get_session(model_path)
     inputs = {session.get_inputs()[0].name: img_array}
     return session.run(None, inputs)[0]
 
-
 # PIL to BGR conversion
 def convert_pil_to_cv2(image: Image.Image) -> np.array:
     arr = np.array(image)
@@ -35,33 +45,29 @@ def convert_pil_to_cv2(image: Image.Image) -> np.array:
         return cv2.cvtColor(arr, cv2.COLOR_GRAY2BGR)
     return arr[:, :, ::-1].copy()
 
-
 # Upscale handler
 def upscale(image, model_choice):
     model_path = f"models/{model_choice}.ort"
     img = convert_pil_to_cv2(image)
-    # ensure 3 channels
+    # handle potential alpha channel
     if img.ndim == 2:
         img = cv2.cvtColor(img, cv2.COLOR_GRAY2BGR)
 
-    # 4-channel alpha handling
     if img.shape[2] == 4:
         alpha = cv2.cvtColor(img[:, :, 3], cv2.COLOR_GRAY2BGR)
         out_a = post_process(inference(model_path, pre_process(alpha)))
         out_a = cv2.cvtColor(out_a, cv2.COLOR_BGR2GRAY)
-        img_rgb = img[:, :, :3]
-        out_rgb = post_process(inference(model_path, pre_process(img_rgb)))
-        out_rgba = cv2.cvtColor(out_rgb, cv2.COLOR_BGR2BGRA)
-        out_rgba[:, :, 3] = out_a
-        return out_rgba
-
+        rgb = img[:, :, :3]
+        out_rgb = post_process(inference(model_path, pre_process(rgb)))
+        rgba = cv2.cvtColor(out_rgb, cv2.COLOR_BGR2BGRA)
+        rgba[:, :, 3] = out_a
+        return rgba
     return post_process(inference(model_path, pre_process(img)))
 
-
 # === Custom CSS for styling & animations ===
 custom_css = """
 body .gradio-container {
-    background: linear-gradient(-45deg, #ff9a9e, #fad0c4, #fad0c4, #ffdde1);
+    background: linear-gradient(-45deg, #ff9a9e, #fad0c4, #ffdde1);
     background-size: 400% 400%;
     animation: gradientBG 15s ease infinite;
 }
@@ -92,7 +98,7 @@ body .gradio-container {
     from { opacity: 0; }
     to { opacity: 1; }
 }
-.gradio-radio input[type="radio"] + label:hover {
+.gradio-radio input[type=\"radio\"] + label:hover {
     transform: scale(1.1);
     transition: transform 0.2s;
 }
@@ -115,15 +121,13 @@ body .gradio-container {
 }
 """
 
-
-# === Gradio Blocks App (no theme string) ===
+# === Gradio Blocks App ===
 with gr.Blocks(css=custom_css) as demo:
     gr.HTML("<h1 class='fancy-title'>✨ Ultra AI Image Upscaler ✨</h1>")
     with gr.Row():
         inp = gr.Image(type="pil", label="Drop Your Image Here")
         model = gr.Radio([
-            "modelx2", "modelx2_25JXL",
-            "modelx4", "minecraft_modelx4"
+            "modelx2", "modelx2_25JXL", "modelx4", "minecraft_modelx4"
         ], label="Upscaler Model", value="modelx2")
     btn = gr.Button("Upscale Image", elem_id="upscale_btn")
     out = gr.Image(label="Upscaled Output", elem_classes="gradio-image")
@@ -132,4 +136,4 @@ with gr.Blocks(css=custom_css) as demo:
     gr.HTML("<p style='text-align:center; color:#555;'>Powered by ONNX Runtime & Gradio Blocks</p>")
 
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch()