added code

app.py

@@ -1,11 +1,39 @@
-import os, torch, torchvision
-import torchvision.transforms.functional
+import os, torch
 from model.BiSeNet.build_bisenet import BiSeNet
 import gradio as gr
 from utils.imageHandling import hfImageToTensor, preprocessing
 
 
+def run_prediction(image: gr.Image, selected_model: str)-> tuple[torch.Tensor]:
+    device = 'cuda' if torch.cuda.is_available() else 'cpu'
+    image = hfImageToTensor(image, width=1024, height=512)
+    return image, predict(image, loadModel(selected_model, device))
 
+# Gradio UI
+with gr.Blocks(title="🔀 BiSeNet | BiSeNetV2 Predictor") as demo:
+    gr.Markdown("## 🧠 Image Segmentation with BiSeNet and BiSeNetV2")
+    gr.Markdown("Upload an image and choose your preferred model for segmentation.")
+
+    with gr.Row():
+        with gr.Column():
+            model_selector = gr.Radio(
+                choices=["BiSeNet", "BiSeNetV2"],
+                value="BiSeNet",
+                label="Select model"
+            )
+            image_input = gr.Image(type="pil", label="Upload image")
+            submit_btn = gr.Button("🧪 Run prediction")
+        with gr.Column():
+            original_display = gr.Image(label="Original image")
+            result_display = gr.Image(label="Model prediction")
+
+    submit_btn.click(
+        fn=run_prediction,
+        inputs=[image_input, model_selector],
+        outputs=[original_display, result_display]
+    )
+
+demo.launch()
 
 
 # %% prediction on an image
@@ -22,9 +50,9 @@ def predict(inputImage: torch.Tensor, model: BiSeNet) -> torch.Tensor:
         prediction (torch.Tensor): The predicted segmentation mask.
     """
     with torch.no_grad():
-        output = model(preprocessing(inputImage))
+        output = model(preprocessing(inputImage.clone()).to(model.device))
     output = output[0] if isinstance(output, (tuple, list)) else output
-    return output[0].argmax(dim=0, keepdim=True)  
+    return output[0].argmax(dim=0, keepdim=True).cpu()