refactor(phase2): add provider-aware tracking and fill missing subscription usage tracking

Changes:
1. helpers.py (_track_image_operation_usage): Map provider name to DB columns
   dynamically (stability→stability_calls, wavespeed→wavespeed_calls, etc.)
   instead of hardcoding stability_calls/stability_cost.

2. upscale_service.py: Added _track_image_operation_usage() call after
   successful Stability upscale completion.

3. control_service.py: Added _track_image_operation_usage() call after
   successful Stability control operation completion.

4. edit_service.py: Added _track_image_operation_usage() call after
   successful Stability edit operation (remove_background, inpaint,
   outpaint, search_replace, search_recolor, relight).

Previously only Create Studio and Face Swap tracked usage. Now all five
studios correctly decrement subscription limits.
This commit is contained in:
ajaysi
2026-05-09 08:51:06 +05:30
parent bc311cfdf6
commit ca725b77e7
4 changed files with 71 additions and 9 deletions

View File

@@ -514,6 +514,19 @@ class EditStudioService:
background_bytes=background_bytes,
lighting_bytes=lighting_bytes,
)
# Track usage for Stability operations
if user_id:
from services.llm_providers.main_image_generation import _track_image_operation_usage
_track_image_operation_usage(
user_id=user_id,
provider="stability",
model=f"edit-{operation}",
operation_type="image-edit",
result_bytes=image_bytes,
cost=0.04,
endpoint="/image-studio/edit/process",
log_prefix="[Edit Studio]"
)
else:
image_bytes = await self._handle_general_edit(
request=request,