diff --git a/Procfile b/Procfile index a32d9aafc..24fd264c4 100644 --- a/Procfile +++ b/Procfile @@ -17,6 +17,6 @@ admin: poetry run python manage.py runserver 127.0.0.1:8000 dashboard: poetry run streamlit run Home.py --server.port 8501 --server.headless true -celery: poetry run celery -A celeryapp worker +celery: poetry run celery -A celeryapp worker --pool solo ui: cd ../gooey-ui/; PORT=3000 npm run dev diff --git a/bots/migrations/0044_savedrun_price_alter_savedrun_workflow.py b/bots/migrations/0044_savedrun_price_alter_savedrun_workflow.py new file mode 100644 index 000000000..060d483d1 --- /dev/null +++ b/bots/migrations/0044_savedrun_price_alter_savedrun_workflow.py @@ -0,0 +1,56 @@ +# Generated by Django 4.2.5 on 2023-10-10 02:58 + +from django.db import migrations, models + + +class Migration(migrations.Migration): + dependencies = [ + ("bots", "0043_alter_savedrun_workflow"), + ] + + operations = [ + migrations.AddField( + model_name="savedrun", + name="price", + field=models.IntegerField(default=0), + ), + migrations.AlterField( + model_name="savedrun", + name="workflow", + field=models.IntegerField( + choices=[ + (1, "Doc Search"), + (2, "Doc Summary"), + (3, "Google GPT"), + (4, "Copilot"), + (5, "Lipysnc + TTS"), + (6, "Text to Speech"), + (7, "Speech Recognition"), + (8, "Lipsync"), + (9, "Deforum Animation"), + (10, "Compare Text2Img"), + (11, "Text2Audio"), + (12, "Img2Img"), + (13, "Face Inpainting"), + (14, "Google Image Gen"), + (15, "Compare AI Upscalers"), + (16, "SEO Summary"), + (17, "Email Face Inpainting"), + (18, "Social Lookup Email"), + (19, "Object Inpainting"), + (20, "Image Segmentation"), + (21, "Compare LLM"), + (22, "Chyron Plant"), + (23, "Letter Writer"), + (24, "Smart GPT"), + (25, "AI QR Code"), + (26, "Doc Extract"), + (27, "Related QnA Maker"), + (28, "Related QnA Maker Doc"), + (29, "Embeddings"), + (30, "Bulk Runner"), + ], + default=4, + ), + ), + ] diff --git a/bots/models.py b/bots/models.py index 0f0ddce50..165d6c6aa 100644 --- a/bots/models.py +++ b/bots/models.py @@ -121,6 +121,7 @@ class SavedRun(models.Model): related_name="children", ) + price = models.IntegerField(default=0) workflow = models.IntegerField( choices=Workflow.choices, default=Workflow.VIDEO_BOTS ) diff --git a/celeryapp/tasks.py b/celeryapp/tasks.py index 566490bb4..a44640329 100644 --- a/celeryapp/tasks.py +++ b/celeryapp/tasks.py @@ -87,7 +87,7 @@ def save(done=False): # run completed except StopIteration: run_time += time() - start_time - page.deduct_credits(st.session_state) + sr.price = page.deduct_credits(st.session_state) break # render errors nicely except Exception as e: diff --git a/daras_ai_v2/base.py b/daras_ai_v2/base.py index 9b7d6e163..16e2680ff 100644 --- a/daras_ai_v2/base.py +++ b/daras_ai_v2/base.py @@ -1047,12 +1047,13 @@ def check_credits(self) -> bool: assert self.request.user, "request.user must be set to check credits" return self.request.user.balance >= self.get_price_roundoff(st.session_state) - def deduct_credits(self, state: dict): + def deduct_credits(self, state: dict) -> int: assert self.request, "request must be set to deduct credits" assert self.request.user, "request.user must be set to deduct credits" amount = self.get_price_roundoff(state) self.request.user.add_balance(-amount, f"gooey_in_{uuid.uuid1()}") + return amount def get_price_roundoff(self, state: dict) -> int: # don't allow fractional pricing for now, min 1 credit