app.py

import asyncio
import base64
import datetime
import hashlib
import hmac
import json
import os
from contextlib import asynccontextmanager
from time import time

import anyio
import gradio as gr
import httpx
import huggingface_hub
from fastapi import FastAPI, HTTPException, Request
from gradio.utils import get_space
from tortoise import Tortoise, fields
from tortoise.models import Model
from tortoise.exceptions import IntegrityError

if not get_space():
    from dotenv import load_dotenv
    load_dotenv()


DATA_DIR = "/data/" if get_space() else "./data"

# Database configuration
DATABASE_URL = f"sqlite://{DATA_DIR}/data.db"  # Replace with your actual database name

client = httpx.AsyncClient()


class TurnUser(Model):
    username = fields.CharField(max_length=255, null=False, pk=True)


async def save_to_json():
    components = await TurnUser.all()
    data = [
        {
            "name": component.username,
        }
        for component in components
    ]
    with open(f"{DATA_DIR}/backup.json", "w") as json_file:
        json.dump(data, json_file, indent=2)


def backup_to_hf_hub():
    huggingface_hub.upload_folder(
        repo_id="freddyaboulton/turn-users",
        folder_path=DATA_DIR,
        path_in_repo=".",
        repo_type="dataset",
        commit_message=f"backup at {datetime.datetime.now()}",
        token=os.getenv("HF_TOKEN"),
    )


async def save_database_to_json():
    while True:
        await asyncio.sleep(5 * 60)
        await save_to_json()
        await anyio.to_thread.run_sync(backup_to_hf_hub)


@asynccontextmanager
async def lifespan(app: FastAPI):
    await Tortoise.init(
        db_url=DATABASE_URL,
        modules={
            "models": ["__main__"]
        },  # Assuming your models are in the __main__ module
    )
    await Tortoise.generate_schemas()
    asyncio.create_task(save_database_to_json())
    yield
    await client.aclose()
    await Tortoise.close_connections()


app = FastAPI(lifespan=lifespan)

success_text = """
<h2>
Profile successfully created for {username}! 🎉

To get your TURN credentials, run the following code:
</h2>


```python
from gradio_webrtc import get_hf_turn_credentials, WebRTC


# Pass a valid access token for your Hugging Face account
# or set the HF_TOKEN environment variable 
credentials = get_hf_turn_credentials(token=None)

with gr.Blcocks() as demo:
    webrtc = WebRTC(rtc_configuration=credentials)
    ...

demo.launch()
```
"""


with gr.Blocks() as demo:

    async def hello(profile: gr.OAuthProfile) -> str:
        try:
            await TurnUser.create(username=profile.username)
        except IntegrityError:
            gr.Info("User already exists", duration=5)


        return success_text.format(username=profile.username)

    gr.HTML("""
        <h1 style='text-align: center'>
        Community TURN Server for WebRTC ⚡️
        </h1>
        <h2 style='text-align: center'>
        Click the login button to sign in with your Hugging Face account.
        </p>
        <h2 style='text-align: center'>
        Then click the "Sign Up" button to create your account for our community TURN server.
        </h2>
    """
    )
    with gr.Row():
        button = gr.LoginButton()
        signup = gr.Button("Sign Up", variant="primary")
    with gr.Row():
        with gr.Column(scale=1):
            mkd = gr.Markdown()
        with gr.Column(scale=1):
            pass

    signup.click(hello, None, mkd)


@app.get("/credentials")
async def _(request: Request):
    token = request.headers.get("X-HF-Access-Token")
    if not token:
        raise HTTPException(status_code=401, detail="Unauthorized")
    username = huggingface_hub.whoami(token=token)["name"]
    user = await TurnUser.get_or_none(username=username)
    if not user:
        raise HTTPException(status_code=404, detail="User not found")

    secret = os.getenv("TOKEN")

    ttl = 24 * 3600
    timestamp = int(time()) + ttl
    turn_user = str(timestamp) + ":" + username
    dig = hmac.new(secret.encode(), turn_user.encode(), hashlib.sha1).digest()
    password = base64.b64encode(dig).decode()

    return {
        "username": turn_user,
        "credential": password,
    }


app = gr.mount_gradio_app(app, demo, "/", ssr_mode=False)


if __name__ == "__main__":
    import uvicorn
    import os
    os.environ["GRADIO_SSR_MODE"] = "False"
    uvicorn.run(app, host="0.0.0.0" if get_space() else "127.0.0.1", port=7860)