Update app.py

app.py

@@ -2,17 +2,20 @@ import numpy as np
 import cv2
 import onnxruntime
 import gradio as gr
+from PIL import Image
 
-# === Ultra-Cleaner & Upscaler Logic (unchanged) ===
+# === Upscaler Logic ===
 def pre_process(img: np.array) -> np.array:
     img = np.transpose(img[:, :, 0:3], (2, 0, 1))
     return np.expand_dims(img, axis=0).astype(np.float32)
 
+
 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
+
+# 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:
@@ -24,48 +27,49 @@ def inference(model_path: str, img_array: np.array) -> np.array:
     inputs = {session.get_inputs()[0].name: img_array}
     return session.run(None, inputs)[0]
 
-# Convert PIL to BGR
-from PIL import Image
+
+# PIL to BGR conversion
 def convert_pil_to_cv2(image: Image.Image) -> np.array:
     arr = np.array(image)
-    return arr[:, :, ::-1].copy() if arr.ndim == 3 else cv2.cvtColor(arr, cv2.COLOR_GRAY2BGR)
+    if arr.ndim == 2:
+        return cv2.cvtColor(arr, cv2.COLOR_GRAY2BGR)
+    return arr[:, :, ::-1].copy()
 
-# Upscale
 
+# Upscale handler
 def upscale(image, model_choice):
     model_path = f"models/{model_choice}.ort"
     img = convert_pil_to_cv2(image)
+    # ensure 3 channels
     if img.ndim == 2:
         img = cv2.cvtColor(img, cv2.COLOR_GRAY2BGR)
 
-    # handle alpha
+    # 4-channel alpha handling
     if img.shape[2] == 4:
         alpha = cv2.cvtColor(img[:, :, 3], cv2.COLOR_GRAY2BGR)
-        out_alpha = post_process(inference(model_path, pre_process(alpha)))
-        out_alpha = cv2.cvtColor(out_alpha, cv2.COLOR_BGR2GRAY)
-        img = img[:, :, 0:3]
-        out_img = post_process(inference(model_path, pre_process(img)))
-        out_bgra = cv2.cvtColor(out_img, cv2.COLOR_BGR2BGRA)
-        out_bgra[:, :, 3] = out_alpha
-        return out_bgra
-
-    # normal RGB
+        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
+
     return post_process(inference(model_path, pre_process(img)))
 
-# === Custom CSS for gradients & animations ===
+
+# === Custom CSS for styling & animations ===
 custom_css = """
 body .gradio-container {
     background: linear-gradient(-45deg, #ff9a9e, #fad0c4, #fad0c4, #ffdde1);
     background-size: 400% 400%;
     animation: gradientBG 15s ease infinite;
 }
-
 @keyframes gradientBG {
     0% { background-position: 0% 50%; }
     50% { background-position: 100% 50%; }
     100% { background-position: 0% 50%; }
 }
-
 .fancy-title {
     font-family: 'Poppins', sans-serif;
     font-size: 3rem;
@@ -75,28 +79,23 @@ body .gradio-container {
     animation: fadeInText 2s ease-in-out;
     text-align: center;
 }
-
 @keyframes fadeInText {
     0% { opacity: 0; transform: translateY(-20px); }
     100% { opacity: 1; transform: translateY(0); }
 }
-
 .gradio-image {
     animation: fadeIn 1s ease-in;
     border-radius: 12px;
     box-shadow: 0 8px 16px rgba(0,0,0,0.2);
 }
-
 @keyframes fadeIn {
     from { opacity: 0; }
     to { opacity: 1; }
 }
-
 .gradio-radio input[type="radio"] + label:hover {
     transform: scale(1.1);
     transition: transform 0.2s;
 }
-
 .gradio-button {
     background: linear-gradient(90deg, #FF8A00, #E52E71);
     border: none;
@@ -107,24 +106,30 @@ body .gradio-container {
     cursor: pointer;
     transition: background 0.3s, transform 0.2s;
 }
-
 .gradio-button:hover {
     background: linear-gradient(90deg, #E52E71, #FF8A00);
     transform: scale(1.05);
 }
+#upscale_btn {
+    margin-top: 10px;
+}
 """
 
-# === Gradio Blocks App ===
-with gr.Blocks(css=custom_css, theme="gradio") as demo:
+
+# === Gradio Blocks App (no theme string) ===
+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"], label="Upscaler Model", value="modelx2")
-    btn = gr.Button("Upscale Image", elem_id="upscale_btn", css="margin-top: 10px;")
+        model = gr.Radio([
+            "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")
 
     btn.click(fn=upscale, inputs=[inp, model], outputs=out)
     gr.HTML("<p style='text-align:center; color:#555;'>Powered by ONNX Runtime & Gradio Blocks</p>")
 
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch()