Update app.py

app.py

@@ -3,82 +3,128 @@ import cv2
 import onnxruntime
 import gradio as gr
 
-
+# === Ultra-Cleaner & Upscaler Logic (unchanged) ===
 def pre_process(img: np.array) -> np.array:
-    # H, W, C -> C, H, W
     img = np.transpose(img[:, :, 0:3], (2, 0, 1))
-    # C, H, W -> 1, C, H, W
-    img = np.expand_dims(img, axis=0).astype(np.float32)
-    return img
-
+    return np.expand_dims(img, axis=0).astype(np.float32)
 
 def post_process(img: np.array) -> np.array:
-    # 1, C, H, W -> C, H, W
     img = np.squeeze(img)
-    # C, H, W -> H, W, C
-    img = np.transpose(img, (1, 2, 0))[:, :, ::-1].astype(np.uint8)
-    return img
-
+    return np.transpose(img, (1, 2, 0))[:, :, ::-1].astype(np.uint8)
 
+# ONNX inference
+_session_cache = {}
 def inference(model_path: str, img_array: np.array) -> np.array:
-    options = onnxruntime.SessionOptions()
-    options.intra_op_num_threads = 1
-    options.inter_op_num_threads = 1
-    ort_session = onnxruntime.InferenceSession(model_path, options)
-    ort_inputs = {ort_session.get_inputs()[0].name: img_array}
-    ort_outs = ort_session.run(None, ort_inputs)
-
-    return ort_outs[0]
-
-
-def convert_pil_to_cv2(image):
-    # pil_image = image.convert("RGB")
-    open_cv_image = np.array(image)
-    # RGB to BGR
-    open_cv_image = open_cv_image[:, :, ::-1].copy()
-    return open_cv_image
-
-
-def upscale(image, model):
-    model_path = f"models/{model}.ort"
+    if model_path not in _session_cache:
+        opts = onnxruntime.SessionOptions()
+        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]
+    inputs = {session.get_inputs()[0].name: img_array}
+    return session.run(None, inputs)[0]
+
+# Convert PIL to BGR
+from PIL import Image
+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)
+
+# Upscale
+
+def upscale(image, model_choice):
+    model_path = f"models/{model_choice}.ort"
     img = convert_pil_to_cv2(image)
     if img.ndim == 2:
         img = cv2.cvtColor(img, cv2.COLOR_GRAY2BGR)
 
+    # handle alpha
     if img.shape[2] == 4:
-        alpha = img[:, :, 3]  # GRAY
-        alpha = cv2.cvtColor(alpha, cv2.COLOR_GRAY2BGR)  # BGR
-        alpha_output = post_process(inference(model_path, pre_process(alpha)))  # BGR
-        alpha_output = cv2.cvtColor(alpha_output, cv2.COLOR_BGR2GRAY)  # GRAY
-
-        img = img[:, :, 0:3]  # BGR
-        image_output = post_process(inference(model_path, pre_process(img)))  # BGR
-        image_output = cv2.cvtColor(image_output, cv2.COLOR_BGR2BGRA)  # BGRA
-        image_output[:, :, 3] = alpha_output
-
-    elif img.shape[2] == 3:
-        image_output = post_process(inference(model_path, pre_process(img)))  # BGR
-
-    return image_output
-
-
-css = ".output-image, .input-image, .image-preview {height: 480px !important} "
-model_choices = ["modelx2", "modelx2 25 JXL", "modelx4", "minecraft_modelx4"]
-
-gr.Interface(
-    fn=upscale,
-    inputs=[
-        gr.Image(type="pil", label="Input Image"),
-        gr.Radio(
-            model_choices,
-            type="value",
-            value=None,
-            label="Choose Upscaler",
-        ),
-    ],
-    outputs="image",
-    title="Image Upscaler! Multiple AI",
-    description="Model: [Anchor-based Plain Net for Mobile Image Super-Resolution](https://arxiv.org/abs/2105.09750). Repository: [SR Mobile PyTorch](https://github.com/w11wo/sr_mobile_pytorch)",
-    allow_flagging="never",
-    css=css,
-).launch()
+        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
+    return post_process(inference(model_path, pre_process(img)))
+
+# === Custom CSS for gradients & 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;
+    background: linear-gradient(90deg, #7F7FD5, #86A8E7, #91EAE4);
+    -webkit-background-clip: text;
+    -webkit-text-fill-color: transparent;
+    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;
+    border-radius: 8px;
+    color: white;
+    font-weight: bold;
+    padding: 12px 24px;
+    cursor: pointer;
+    transition: background 0.3s, transform 0.2s;
+}
+
+.gradio-button:hover {
+    background: linear-gradient(90deg, #E52E71, #FF8A00);
+    transform: scale(1.05);
+}
+"""
+
+# === Gradio Blocks App ===
+with gr.Blocks(css=custom_css, theme="gradio") 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;")
+    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()