Team Ai
Apppublic

huggingface-projects/stable-diffusion-multiplayer

sourceHugging Faceupdated 3y agoView on Hugging Face
351likes
createRooms.py132 linesDownload Raw Back to stablediffusion-infinity
1import os2from fastapi import Depends, FastAPI3import sqlite34import requests5import uvicorn6from pathlib import Path7import json8 9app = FastAPI()10LIVEBLOCKS_SECRET = os.environ.get("LIVEBLOCKS_SECRET")11 12DB_PATH = Path("rooms.db")13if not DB_PATH.exists():14    print("Creating database")15    print("DB_PATH", DB_PATH)16    db = sqlite3.connect(DB_PATH)17    with open(Path("schema.sql"), "r") as f:18        db.executescript(f.read())19    db.commit()20    db.close()21 22 23def get_db():24    db = sqlite3.connect(Path("./rooms.db"), check_same_thread=False)25    print("Connected to database")26    db.row_factory = sqlite3.Row27    try:28        yield db29    except Exception:30        db.rollback()31    finally:32        db.close()33 34 35app = FastAPI()36 37rooms = ["room-" + str(i) for i in range(0, 41)]38 39 40@app.get("/")41async def read_root(db: sqlite3.Connection = Depends(get_db)):42    out = db.execute("SELECT * FROM rooms").fetchall()43    print(out)44    return out45 46 47@app.get("/create-rooms")48async def create_room(db: sqlite3.Connection = Depends(get_db)):49    for room_id in rooms:50        print(room_id)51        createRoom(room_id,  db)52    all = db.execute("SELECT * FROM rooms").fetchall()53    return all54 55 56@app.get("/create-one-room/{room_id}")57async def create_one_room(room_id: str):58    payload = {"id": room_id, "defaultAccesses": ["room:write"]}59 60    response = requests.post(f"https://api.liveblocks.io/v2/rooms",61                             headers={"Authorization": f"Bearer {LIVEBLOCKS_SECRET}"}, json=payload)62    data = response.json()63    return data64 65 66def createRoom(room_id, db):67    payload = {"id": room_id, "defaultAccesses": ["room:write"]}68 69    response = requests.post(f"https://api.liveblocks.io/v2/rooms",70                             headers={"Authorization": f"Bearer {LIVEBLOCKS_SECRET}"}, json=payload)71    # if response.status_code == 200:72    data = response.json()73    print(data)74    if "error" in data and data["error"] == "ROOM_ALREADY_EXISTS":75        print("Room already exists")76 77    cursor = db.cursor()78    cursor.execute("INSERT INTO rooms (room_id) VALUES (?)", (room_id,))79    db.commit()80    print("Room created")81 82    print("Created room", room_id)83    return True84 85 86def generateAuthToken():87    response = requests.get(f"https://liveblocks.io/api/authorize",88                            headers={"Authorization": f"Bearer {LIVEBLOCKS_SECRET}"})89    if response.status_code == 200:90        data = response.json()91        return data["token"]92    else:93        raise Exception(response.status_code, response.text)94 95 96def get_room_count(room_id: str, jwtToken: str = ''):97    print("Getting room count" + room_id)98    response = requests.get(99        f"https://liveblocks.net/api/v1/room/{room_id}/users", headers={"Authorization": f"Bearer {jwtToken}", "Content-Type": "application/json"})100    if response.status_code == 200:101        res = response.json()102        if "data" in res:103            return len(res["data"])104        else:105            return 0106    raise Exception("Error getting room count")107 108 109@app.get("/sync-rooms")110async def sync_rooms(db: sqlite3.Connection = Depends(get_db)):111    try:112        jwtToken = generateAuthToken()113        rooms = db.execute("SELECT * FROM rooms").fetchall()114        for row in rooms:115            room_id = row["room_id"]116            users_count = get_room_count(room_id, jwtToken)117            print("Updating room", room_id, "with", users_count, "users")118            cursor = db.cursor()119            cursor.execute(120                "UPDATE rooms SET users_count = ? WHERE room_id = ?", (users_count, room_id))121            db.commit()122        data = db.execute("SELECT * FROM rooms").fetchall()123        return data124    except Exception as e:125        print(e)126        return {"error": str(e)}127 128 129if __name__ == "__main__":130    uvicorn.run("createRooms:app", host="0.0.0.0",131                log_level="debug", reload=True)132