Update app.py

app.py

@@ -4,16 +4,21 @@ from gradio_client import Client, handle_file
 
 # Fetch HF_TOKEN from environment variables
 HF_TOKEN = os.getenv("HF_TOKEN")
+
 if not HF_TOKEN:
-    raise ValueError("HF_TOKEN environment variable not set. Please set it before running the app.")
+    raise ValueError(
+        "HF_TOKEN environment variable not set. Please set it before running the app."
+    )
 
-# Initialize the client for the Coldet backend with HF_TOKEN in headers
+# Initialize Coldet backend client
 coldet_client = Client(
     "Coldet/backend",
     headers={"Authorization": f"Bearer {HF_TOKEN}"}
 )
 
-# --- Helper Functions ---
+# -----------------------------
+# Helper Functions
+# -----------------------------
 
 def create_account(username: str, password: str) -> str:
     try:
@@ -22,7 +27,9 @@ def create_account(username: str, password: str) -> str:
             password=password,
             api_name="/create_account",
         )
-        return f"Account created! Session ID: {result}"
+
+        return f"Account created!\n\nSession ID:\n{result}"
+
     except Exception as e:
         return f"Error: {str(e)}"
 
@@ -34,7 +41,9 @@ def login(username: str, password: str) -> str:
             password=password,
             api_name="/login",
         )
-        return f"Logged in! Session ID: {result}"
+
+        return f"Logged in!\n\nSession ID:\n{result}"
+
     except Exception as e:
         return f"Error: {str(e)}"
 
@@ -45,12 +54,17 @@ def create_token(session_id: str) -> str:
             session_id=session_id,
             api_name="/create_token",
         )
-        return f"API Key: {result}"
+
+        return f"API Key:\n{result}"
+
     except Exception as e:
         return f"Error: {str(e)}"
 
 
-# AUTO FETCH MODELS FROM BACKEND CONFIG
+# -----------------------------
+# Auto Fetch Models
+# -----------------------------
+
 def get_models():
     try:
         config = coldet_client.view_api(return_format="dict")
@@ -61,22 +75,26 @@ def get_models():
 
         for param in params:
             if param.get("parameter_name") == "model":
-                choices = param.get("python_type", "")
 
-                # Extract Literal choices
-                if "Literal[" in choices:
+                python_type = param.get("python_type", "")
+
+                if "Literal[" in python_type:
+
                     models = (
-                        choices.replace("Literal[", "")
+                        python_type
+                        .replace("Literal[", "")
                         .replace("]", "")
                         .replace("'", "")
                         .split(", ")
                     )
+
                     return models
 
         return ["coldet-v1-mini.coldet"]
 
     except Exception as e:
         print(f"Failed to fetch models: {e}")
+
         return ["coldet-v1-mini.coldet"]
 
 
@@ -94,7 +112,15 @@ def refresh_models():
     )
 
 
+# -----------------------------
+# Predict
+# -----------------------------
+
 def predict(image, api_key: str, model: str):
+
+    if image is None:
+        return {"error": "Please upload an image."}
+
     try:
         result = coldet_client.predict(
             image=handle_file(image),
@@ -102,28 +128,37 @@ def predict(image, api_key: str, model: str):
             model=model,
             api_name="/predict",
         )
+
         return result
+
     except Exception as e:
         return {"error": str(e)}
 
 
-# --- Gradio Interface ---
+# -----------------------------
+# UI
+# -----------------------------
 
 with gr.Blocks(
-    title="Coldet Backend Interface",
-    theme=gr.themes.Soft(),
+    title="Coldet Backend Interface"
 ) as demo:
 
     gr.Markdown("# 🚀 Coldet Backend Interface")
-    gr.Markdown("Use this interface to interact with the Coldet backend API.")
+    gr.Markdown(
+        "Use this interface to interact with the Coldet backend API."
+    )
+
+    # -----------------------------
+    # Account Tab
+    # -----------------------------
 
-    # --- Account Tab ---
     with gr.Tab("🔑 Account"):
 
         gr.Markdown("### Create Account or Login")
 
         with gr.Row():
 
+            # Create Account
             with gr.Column():
 
                 gr.Markdown("#### Create Account")
@@ -143,7 +178,8 @@ with gr.Blocks(
 
                 create_output = gr.Textbox(
                     label="Session ID",
-                    interactive=False
+                    interactive=False,
+                    lines=4
                 )
 
                 create_btn.click(
@@ -152,6 +188,7 @@ with gr.Blocks(
                     outputs=create_output,
                 )
 
+            # Login
             with gr.Column():
 
                 gr.Markdown("#### Login")
@@ -171,7 +208,8 @@ with gr.Blocks(
 
                 login_output = gr.Textbox(
                     label="Session ID",
-                    interactive=False
+                    interactive=False,
+                    lines=4
                 )
 
                 login_btn.click(
@@ -180,21 +218,26 @@ with gr.Blocks(
                     outputs=login_output,
                 )
 
-    # --- API Key Tab ---
+    # -----------------------------
+    # API Key Tab
+    # -----------------------------
+
     with gr.Tab("🔐 API Key"):
 
         gr.Markdown("### Generate API Key")
 
         session_id_input = gr.Textbox(
             label="Session ID",
-            placeholder="Paste your session ID here"
+            placeholder="Paste your session ID here",
+            lines=2
         )
 
         token_btn = gr.Button("Generate API Key")
 
         token_output = gr.Textbox(
             label="API Key",
-            interactive=False
+            interactive=False,
+            lines=3
         )
 
         token_btn.click(
@@ -203,7 +246,10 @@ with gr.Blocks(
             outputs=token_output,
         )
 
-    # --- Inference Tab ---
+    # -----------------------------
+    # Inference Tab
+    # -----------------------------
+
     with gr.Tab("🤖 Inference"):
 
         gr.Markdown("### Run Inference")
@@ -219,13 +265,15 @@ with gr.Blocks(
 
                 api_key_input = gr.Textbox(
                     label="API Key",
-                    placeholder="Paste your API key here"
+                    placeholder="Paste your API key here",
                 )
 
                 model_input = gr.Dropdown(
                     label="Model",
                     choices=AVAILABLE_MODELS,
-                    value=AVAILABLE_MODELS[0] if AVAILABLE_MODELS else None,
+                    value=AVAILABLE_MODELS[0]
+                    if AVAILABLE_MODELS
+                    else None,
                 )
 
                 refresh_btn = gr.Button("🔄 Refresh Models")
@@ -245,16 +293,30 @@ with gr.Blocks(
 
         predict_btn.click(
             fn=predict,
-            inputs=[image_input, api_key_input, model_input],
+            inputs=[
+                image_input,
+                api_key_input,
+                model_input
+            ],
             outputs=predict_output,
         )
 
-    # AUTO REFRESH MODELS ON PAGE LOAD
+    # -----------------------------
+    # Auto Refresh On Load
+    # -----------------------------
+
     demo.load(
         fn=refresh_models,
         outputs=model_input,
     )
 
-# Launch the app
+# -----------------------------
+# Launch
+# -----------------------------
+
 if __name__ == "__main__":
-    demo.launch()
+
+    demo.launch(
+        server_name="0.0.0.0",
+        server_port=7860
+    )