usr_endpoint / app.py
theguywhosucks's picture
Update app.py
a0dffb0 verified
Raw
History Blame Contribute Delete
10.4 kB
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
)