Spaces:
Running
Running
File size: 2,413 Bytes
bbc89f6 a1e077b bbc89f6 a1e077b |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 |
import json
from supabase import create_client, Client
from config import SUPABASE_URL, SUPABASE_KEY
supabase: Client = create_client(SUPABASE_URL, SUPABASE_KEY)
def get_user_credits(user_id):
user = supabase.table("users").select("generation_credits, train_credits").eq("id", user_id).execute()
if user.data:
return user.data[0]["generation_credits"], user.data[0]["train_credits"]
return 0, 0
def update_user_credits(user_id, generation_credits, train_credits):
supabase.table("users").update({
"generation_credits": generation_credits,
"train_credits": train_credits
}).eq("id", user_id).execute()
def get_or_create_user(google_id, email, name, given_name, profile_picture):
user = supabase.table("users").select("*").eq("google_id", google_id).execute()
if not user.data:
new_user = {
"google_id": google_id,
"email": email,
"name": name,
"profile_picture": profile_picture,
"generation_credits": 2,
"train_credits": 1,
"given_name": given_name
}
result = supabase.table("users").insert(new_user).execute()
return result.data[0]
else:
return user.data[0]
def get_lora_models_info():
lora_models = supabase.table("lora_models").select("*").is_("user_id", None).execute()
return lora_models.data
def get_user_by_id(user_id):
user = supabase.table("users").select("*").eq("id", user_id).execute()
if user.data:
return user.data[0]
return None
def create_lora_models(user_id, replicate_repo_name, trigger_word, steps, lora_rank, batch_size, learning_rate, hf_repo_name, training_url):
# create a jsonb from trigger_word, train_steps, lora_rank, batch_size, learning_rate values
model_config = {
"train_steps": steps,
"lora_rank": lora_rank,
"batch_size": batch_size,
"learning_rate": learning_rate
}
result = supabase.table("lora_models").insert({
"user_id": user_id,
"trigger_word": trigger_word,
"lora_name": replicate_repo_name,
"hf_repo": hf_repo_name,
"configs": model_config,
"training_url": training_url
}).execute()
def get_user_lora_models(user_id):
user_models = supabase.table("lora_models").select("*").eq("user_id", user_id).execute()
return user_models.data |