Spaces:
Running
Running
| import os | |
| import gradio as gr | |
| from gradio_client import Client, handle_file | |
| # --------------------------------------------------- | |
| # Environment | |
| # --------------------------------------------------- | |
| HF_TOKEN = os.getenv("HF_TOKEN") | |
| if not HF_TOKEN: | |
| raise ValueError( | |
| "HF_TOKEN environment variable not set." | |
| ) | |
| # --------------------------------------------------- | |
| # Coldet Client | |
| # --------------------------------------------------- | |
| client = Client( | |
| "Coldet/backend", | |
| headers={ | |
| "Authorization": f"Bearer {HF_TOKEN}" | |
| } | |
| ) | |
| # --------------------------------------------------- | |
| # Model List | |
| # --------------------------------------------------- | |
| AVAILABLE_MODELS = [ | |
| "coldet-v1-mini.coldet", | |
| "coldet-v2-mini.coldet", | |
| ] | |
| # --------------------------------------------------- | |
| # Model Credit Costs | |
| # --------------------------------------------------- | |
| MODEL_COSTS = { | |
| "coldet-v1-mini.coldet": 1, | |
| "coldet-v2-mini.coldet": 2, | |
| } | |
| # --------------------------------------------------- | |
| # Helpers | |
| # --------------------------------------------------- | |
| def format_error(e): | |
| return f"Error:\n\n{str(e)}" | |
| # --------------------------------------------------- | |
| # API FUNCTIONS | |
| # --------------------------------------------------- | |
| def create_account(username: str, password: str): | |
| try: | |
| result = client.predict( | |
| username=username, | |
| password=password, | |
| api_name="/create_account", | |
| ) | |
| return f"Account created!\n\nSession ID:\n{result}" | |
| except Exception as e: | |
| return format_error(e) | |
| # --------------------------------------------------- | |
| def login(username: str, password: str): | |
| try: | |
| result = client.predict( | |
| username=username, | |
| password=password, | |
| api_name="/login", | |
| ) | |
| return f"Logged in!\n\nSession ID:\n{result}" | |
| except Exception as e: | |
| return format_error(e) | |
| # --------------------------------------------------- | |
| def change_password( | |
| session_id: str, | |
| old_password: str, | |
| new_password: str | |
| ): | |
| try: | |
| result = client.predict( | |
| session_id=session_id, | |
| old_password=old_password, | |
| new_password=new_password, | |
| api_name="/change_password", | |
| ) | |
| return result | |
| except Exception as e: | |
| return format_error(e) | |
| # --------------------------------------------------- | |
| def create_token(session_id: str): | |
| try: | |
| result = client.predict( | |
| session_id=session_id, | |
| api_name="/create_token", | |
| ) | |
| return f"API Key:\n{result}" | |
| except Exception as e: | |
| return format_error(e) | |
| # --------------------------------------------------- | |
| def list_keys(session_id: str): | |
| try: | |
| result = client.predict( | |
| session_id=session_id, | |
| api_name="/list_keys", | |
| ) | |
| return result | |
| except Exception as e: | |
| return { | |
| "error": str(e) | |
| } | |
| # --------------------------------------------------- | |
| def delete_key( | |
| session_id: str, | |
| api_key: str | |
| ): | |
| try: | |
| result = client.predict( | |
| session_id=session_id, | |
| api_key=api_key, | |
| api_name="/delete_key", | |
| ) | |
| return result | |
| except Exception as e: | |
| return format_error(e) | |
| # --------------------------------------------------- | |
| def predict( | |
| image, | |
| api_key: str, | |
| model: str | |
| ): | |
| if image is None: | |
| return { | |
| "error": "Please upload an image." | |
| } | |
| try: | |
| result = client.predict( | |
| image=handle_file(image), | |
| api_key=api_key, | |
| model=model, | |
| api_name="/predict", | |
| ) | |
| result["estimated_cost"] = MODEL_COSTS.get(model, 1) | |
| return result | |
| except Exception as e: | |
| return { | |
| "error": str(e) | |
| } | |
| # --------------------------------------------------- | |
| # UI | |
| # --------------------------------------------------- | |
| with gr.Blocks( | |
| title="Coldet Backend Interface" | |
| ) as demo: | |
| gr.Markdown("# ๐ Coldet Backend Interface") | |
| gr.Markdown( | |
| "Frontend for the Coldet backend API." | |
| ) | |
| # --------------------------------------------------- | |
| # ACCOUNT TAB | |
| # --------------------------------------------------- | |
| with gr.Tab("๐ Account"): | |
| with gr.Row(): | |
| # ----------------------------- | |
| # Create Account | |
| # ----------------------------- | |
| with gr.Column(): | |
| gr.Markdown("## Create Account") | |
| create_username = gr.Textbox( | |
| label="Username" | |
| ) | |
| create_password = gr.Textbox( | |
| label="Password", | |
| type="password" | |
| ) | |
| create_btn = gr.Button( | |
| "Create Account" | |
| ) | |
| create_output = gr.Textbox( | |
| label="Session", | |
| lines=4, | |
| interactive=False | |
| ) | |
| create_btn.click( | |
| fn=create_account, | |
| inputs=[ | |
| create_username, | |
| create_password | |
| ], | |
| outputs=create_output, | |
| ) | |
| # ----------------------------- | |
| # Login | |
| # ----------------------------- | |
| with gr.Column(): | |
| gr.Markdown("## Login") | |
| login_username = gr.Textbox( | |
| label="Username" | |
| ) | |
| login_password = gr.Textbox( | |
| label="Password", | |
| type="password" | |
| ) | |
| login_btn = gr.Button( | |
| "Login" | |
| ) | |
| login_output = gr.Textbox( | |
| label="Session", | |
| lines=4, | |
| interactive=False | |
| ) | |
| login_btn.click( | |
| fn=login, | |
| inputs=[ | |
| login_username, | |
| login_password | |
| ], | |
| outputs=login_output, | |
| ) | |
| # --------------------------------------------------- | |
| # PASSWORD TAB | |
| # --------------------------------------------------- | |
| with gr.Tab("๐ Password"): | |
| gr.Markdown("## Change Password") | |
| password_session = gr.Textbox( | |
| label="Session ID" | |
| ) | |
| old_password = gr.Textbox( | |
| label="Old Password", | |
| type="password" | |
| ) | |
| new_password = gr.Textbox( | |
| label="New Password", | |
| type="password" | |
| ) | |
| password_btn = gr.Button( | |
| "Change Password" | |
| ) | |
| password_output = gr.Textbox( | |
| label="Result", | |
| interactive=False | |
| ) | |
| password_btn.click( | |
| fn=change_password, | |
| inputs=[ | |
| password_session, | |
| old_password, | |
| new_password | |
| ], | |
| outputs=password_output, | |
| ) | |
| # --------------------------------------------------- | |
| # API KEYS TAB | |
| # --------------------------------------------------- | |
| with gr.Tab("๐ API Keys"): | |
| gr.Markdown("## API Key Management") | |
| key_session = gr.Textbox( | |
| label="Session ID", | |
| lines=2 | |
| ) | |
| with gr.Row(): | |
| create_key_btn = gr.Button( | |
| "Create API Key" | |
| ) | |
| refresh_keys_btn = gr.Button( | |
| "Refresh Keys" | |
| ) | |
| token_output = gr.Textbox( | |
| label="New API Key", | |
| interactive=False, | |
| lines=3 | |
| ) | |
| keys_output = gr.JSON( | |
| label="Your Keys" | |
| ) | |
| create_key_btn.click( | |
| fn=create_token, | |
| inputs=key_session, | |
| outputs=token_output, | |
| ) | |
| refresh_keys_btn.click( | |
| fn=list_keys, | |
| inputs=key_session, | |
| outputs=keys_output, | |
| ) | |
| gr.Markdown("## Delete API Key") | |
| delete_key_input = gr.Textbox( | |
| label="API Key" | |
| ) | |
| delete_key_btn = gr.Button( | |
| "Delete Key" | |
| ) | |
| delete_output = gr.Textbox( | |
| label="Delete Result", | |
| interactive=False | |
| ) | |
| delete_key_btn.click( | |
| fn=delete_key, | |
| inputs=[ | |
| key_session, | |
| delete_key_input | |
| ], | |
| outputs=delete_output, | |
| ) | |
| # --------------------------------------------------- | |
| # INFERENCE TAB | |
| # --------------------------------------------------- | |
| with gr.Tab("๐ค Inference"): | |
| gr.Markdown("## Run Inference") | |
| gr.Markdown( | |
| "\n".join( | |
| [ | |
| f"- `{model}` โ {cost} credit(s)" | |
| for model, cost in MODEL_COSTS.items() | |
| ] | |
| ) | |
| ) | |
| with gr.Row(): | |
| with gr.Column(): | |
| image_input = gr.Image( | |
| label="Image", | |
| type="filepath" | |
| ) | |
| api_key_input = gr.Textbox( | |
| label="API Key" | |
| ) | |
| model_input = gr.Dropdown( | |
| label="Model", | |
| choices=AVAILABLE_MODELS, | |
| value=AVAILABLE_MODELS[0], | |
| ) | |
| predict_btn = gr.Button( | |
| "Run Inference" | |
| ) | |
| with gr.Column(): | |
| predict_output = gr.JSON( | |
| label="Result" | |
| ) | |
| predict_btn.click( | |
| fn=predict, | |
| inputs=[ | |
| image_input, | |
| api_key_input, | |
| model_input | |
| ], | |
| outputs=predict_output, | |
| ) | |
| # --------------------------------------------------- | |
| # Launch | |
| # --------------------------------------------------- | |
| demo.queue() | |
| if __name__ == "__main__": | |
| demo.launch( | |
| server_name="0.0.0.0", | |
| server_port=7860, | |
| ssr_mode=False | |
| ) |