rooms api

.gitignore

@@ -17,3 +17,4 @@ flagged/
 data
 data.db
 data.json
+rooms.db

createRooms.py

@@ -0,0 +1,112 @@
+import os
+from fastapi import Depends, FastAPI
+import sqlite3
+import requests
+import uvicorn
+from pathlib import Path
+import json
+
+app = FastAPI()
+LIVEBLOCKS_SECRET = os.environ.get("LIVEBLOCKS_SECRET")
+
+
+def get_db():
+    db = sqlite3.connect(Path("./rooms.db"), check_same_thread=False)
+    db.execute("CREATE TABLE IF NOT EXISTS rooms (id INTEGER PRIMARY KEY AUTOINCREMENT NOT NULL, room_id TEXT NOT NULL, users_count INTEGER NOT NULL DEFAULT 0)")
+    print("Connected to database")
+    db.commit()
+    db.row_factory = sqlite3.Row
+    try:
+        yield db
+    except Exception:
+        db.rollback()
+    finally:
+        db.close()
+
+
+app = FastAPI()
+
+rooms = ["sd-multiplayer-room-" + str(i) for i in range(0, 20)]
+
+
+@app.get("/")
+async def read_root(db: sqlite3.Connection = Depends(get_db)):
+    out = db.execute("SELECT * FROM rooms").fetchall()
+    print(out)
+    return out
+
+
+@app.get("/create-rooms")
+async def create_room(db: sqlite3.Connection = Depends(get_db)):
+    for room_id in rooms:
+        print(room_id)
+        createRoom(room_id,  db)
+    all = db.execute("SELECT * FROM rooms").fetchall()
+    return all
+
+
+def createRoom(room_id, db):
+    payload = {"id": room_id, "defaultAccesses": ["room:write"]}
+
+    response = requests.post(f"https://api.liveblocks.io/v2/rooms",
+                             headers={"Authorization": f"Bearer {LIVEBLOCKS_SECRET}"}, json=payload)
+    # if response.status_code == 200:
+    data = response.json()
+    print(data)
+    if "error" in data and data["error"] == "ROOM_ALREADY_EXISTS":
+        print("Room already exists")
+
+    cursor = db.cursor()
+    cursor.execute("INSERT INTO rooms (room_id) VALUES (?)", (room_id,))
+    db.commit()
+    print("Room created")
+
+    print("Created room", room_id)
+    return True
+
+
+def generateAuthToken():
+    response = requests.get(f"https://liveblocks.io/api/authorize",
+                            headers={"Authorization": f"Bearer {LIVEBLOCKS_SECRET}"})
+    if response.status_code == 200:
+        data = response.json()
+        return data["token"]
+    else:
+        raise Exception(response.status_code, response.text)
+
+
+def get_room_count(room_id: str, jwtToken: str = ''):
+    print("Getting room count" + room_id)
+    response = requests.get(
+        f"https://liveblocks.net/api/v1/room/{room_id}/users", headers={"Authorization": f"Bearer {jwtToken}", "Content-Type": "application/json"})
+    if response.status_code == 200:
+        res = response.json()
+        if "data" in res:
+            return len(res["data"])
+        else:
+            return 0
+    raise Exception("Error getting room count")
+
+
+@app.get("/sync-rooms")
+async def sync_rooms(db: sqlite3.Connection = Depends(get_db)):
+    try:
+        jwtToken = generateAuthToken()
+        rooms = db.execute("SELECT * FROM rooms").fetchall()
+        for row in rooms:
+            room_id = row["room_id"]
+            users_count = get_room_count(room_id, jwtToken)
+            print("Updating room", room_id, "with", users_count, "users")
+            cursor = db.cursor()
+            cursor.execute(
+                "UPDATE rooms SET users_count = ? WHERE room_id = ?", (users_count, room_id))
+            db.commit()
+        data = db.execute("SELECT * FROM rooms").fetchall()
+        return data
+    except Exception as e:
+        print(e)
+        return {"error": str(e)}
+
+
+if __name__ == "__main__":
+    uvicorn.run("api:app", host="0.0.0.0", log_level="debug", reload=True)

stablediffusion-infinity/app.py

@@ -1,6 +1,5 @@
 import io
 import os
-from random import sample
 
 from pathlib import Path
 import uvicorn
@@ -50,10 +49,13 @@ def get_db():
     db.row_factory = sqlite3.Row
     try:
         yield db
+    except Exception:
+        db.rollback()
     finally:
         db.close()
 
 
+
 s3 = boto3.client(service_name='s3',
                   aws_access_key_id=AWS_ACCESS_KEY_ID,
                   aws_secret_access_key=AWS_SECRET_KEY)
@@ -248,7 +250,6 @@ def generateAuthToken():
 
 
 def get_room_count(room_id: str, jwtToken: str = ''):
-    print("Getting room count" + room_id)
     response = requests.get(
         f"https://liveblocks.net/api/v1/room/{room_id}/users", headers={"Authorization": f"Bearer {jwtToken}", "Content-Type": "application/json"})
     if response.status_code == 200:
@@ -260,32 +261,30 @@ def get_room_count(room_id: str, jwtToken: str = ''):
     raise Exception("Error getting room count")
 
 
-app = gr.mount_gradio_app(app, blocks, "/gradio",
-                          gradio_api_url="http://0.0.0.0:7860/gradio/")
-
-app.on_event("startup")
-
-
-@repeat_every(seconds=10)
-async def sync_rooms(db: sqlite3.Connection = Depends(get_db)):
+@app.on_event("startup")
+@repeat_every(seconds=60*5)
+async def sync_rooms():
+    print("Syncing rooms")
     try:
         jwtToken = generateAuthToken()
-        rooms = db.execute("SELECT * FROM rooms").fetchall()
-        print(rooms)
-        for row in rooms:
-            room_id = row["room_id"]
-            users_count = get_room_count(room_id, jwtToken)
-            print("Updating room", room_id, "with", users_count, "users")
-            cursor = db.cursor()
-            cursor.execute(
-                "UPDATE rooms SET users_count = ? WHERE room_id = ?", (users_count, room_id))
-            db.commit()
-        data = db.execute("SELECT * FROM rooms").fetchall()
-        print("Rooms updated", data)
+        for db in get_db():
+            rooms = db.execute("SELECT * FROM rooms").fetchall()
+            for row in rooms:
+                room_id = row["room_id"]
+                users_count = get_room_count(room_id, jwtToken)
+                cursor = db.cursor()
+                cursor.execute(
+                    "UPDATE rooms SET users_count = ? WHERE room_id = ?", (users_count, room_id))
+                db.commit()
     except Exception as e:
         print(e)
         print("Rooms update failed")
 
+@app.get('/rooms')
+async def get_rooms(db: sqlite3.Connection = Depends(get_db)):
+    rooms = db.execute("SELECT * FROM rooms").fetchall()
+    return rooms
+
 
 @app.post('/uploadfile/')
 async def create_upload_file(background_tasks: BackgroundTasks, file: UploadFile):
@@ -324,6 +323,10 @@ app.add_middleware(
     allow_headers=["*"],
 )
 
+app = gr.mount_gradio_app(app, blocks, "/gradio",
+                          gradio_api_url="http://0.0.0.0:7860/gradio/")
+
+
 if __name__ == "__main__":
     uvicorn.run(app, host="0.0.0.0", port=7860,
                 log_level="debug", reload=False)