Update app.py

app.py

@@ -1,55 +1,97 @@
-import os
-import gradio as gr
+import argparse
 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.")
-
-# Initialize the client for the Coldet backend with HF_TOKEN in headers
-coldet_client = Client("Coldet/backend", headers={"Authorization": f"Bearer {HF_TOKEN}"})
-
-# --- Inference Function ---
-
-def predict(image, api_key: str, model: str = "coldet-v1-mini.coldet") -> str:
-    """Run inference on the provided image."""
-    try:
-        result = coldet_client.predict(
-            image=handle_file(image),
-            api_key=api_key,
-            model=model,
-            api_name="/predict",
-        )
-        return result
-    except Exception as e:
-        return f"Error: {str(e)}"
-
-# --- Gradio Interface ---
-
-with gr.Blocks(title="Coldet Inference", theme=gr.themes.Soft()) as demo:
-    gr.Markdown("# 🤖 Coldet Inference")
-    gr.Markdown("Upload an image and paste your API key to run inference.")
-
-    with gr.Row():
-        with gr.Column():
-            image_input = gr.Image(label="Upload Image", type="filepath")
-            api_key_input = gr.Textbox(label="API Key", placeholder="Paste your API key here")
-            model_input = gr.Dropdown(
-                label="Model",
-                choices=["coldet-v1-mini.coldet"],
-                value="coldet-v1-mini.coldet",
-            )
-            predict_btn = gr.Button("Run Inference", variant="primary")
-        with gr.Column():
-            predict_output = gr.JSON(label="Inference Result")
-
-    predict_btn.click(
-        fn=predict,
-        inputs=[image_input, api_key_input, model_input],
-        outputs=predict_output,
-    )
-
-# Launch the app
+def create_account(client, username, password):
+    """Create a new account."""
+    result = client.predict(
+        username=username,
+        password=password,
+        api_name="/create_account",
+    )
+    return result
+
+def login(client, username, password):
+    """Login to an existing account."""
+    result = client.predict(
+        username=username,
+        password=password,
+        api_name="/login",
+    )
+    return result
+
+def create_token(client, session_id):
+    """Generate an API key for the session."""
+    result = client.predict(
+        session_id=session_id,
+        api_name="/create_token",
+    )
+    return result
+
+def predict(client, image_path, api_key, model="coldet-v1-mini.coldet"):
+    """Run inference on an image."""
+    result = client.predict(
+        image=handle_file(image_path),
+        api_key=api_key,
+        model=model,
+        api_name="/predict",
+    )
+    return result
+
+def main():
+    parser = argparse.ArgumentParser(description="Coldet CLI Tool")
+    parser.add_argument(
+        "--image",
+        type=str,
+        help="Path or URL to the image for inference",
+    )
+    parser.add_argument(
+        "--username",
+        type=str,
+        help="Username for login or account creation",
+    )
+    parser.add_argument(
+        "--password",
+        type=str,
+        help="Password for login or account creation",
+    )
+    parser.add_argument(
+        "--session-id",
+        type=str,
+        help="Session ID for generating API key",
+    )
+    parser.add_argument(
+        "--api-key",
+        type=str,
+        help="API key for inference",
+    )
+    args = parser.parse_args()
+
+    client = Client("Coldet/backend")
+
+    # Step 1: Login or Create Account
+    if args.username and args.password:
+        session_id = login(client, args.username, args.password)
+        print(f"Logged in! Session ID: {session_id}")
+    else:
+        print("No username/password provided. Skipping login.")
+        session_id = None
+
+    # Step 2: Generate API Key (if session_id is provided)
+    if args.session_id or session_id:
+        session_id = args.session_id or session_id
+        api_key = create_token(client, session_id)
+        print(f"API Key generated: {api_key}")
+    else:
+        print("No session ID provided. Skipping API key generation.")
+        api_key = None
+
+    # Step 3: Run Inference (if image and api_key are provided)
+    if args.image and (args.api_key or api_key):
+        api_key = args.api_key or api_key
+        result = predict(client, args.image, api_key)
+        print(f"Inference result: {result}")
+    else:
+        print("Missing image or API key. Cannot run inference.")
+
 if __name__ == "__main__":
-    demo.launch()
+    main()