705 lines
33 KiB
Python
705 lines
33 KiB
Python
from __future__ import annotations
|
|
|
|
import base64
|
|
import os
|
|
import uuid
|
|
from typing import Optional, Dict, Any
|
|
from datetime import datetime
|
|
from pathlib import Path
|
|
from sqlalchemy.orm import Session
|
|
|
|
from fastapi import APIRouter, HTTPException, Depends, Request
|
|
from fastapi.responses import FileResponse
|
|
from pydantic import BaseModel, Field
|
|
|
|
from services.llm_providers.main_image_generation import generate_image
|
|
from services.llm_providers.main_image_editing import edit_image
|
|
from services.llm_providers.main_text_generation import llm_text_gen
|
|
from utils.logger_utils import get_service_logger
|
|
from middleware.auth_middleware import get_current_user
|
|
from services.database import get_db
|
|
from services.subscription import UsageTrackingService, PricingService
|
|
from models.subscription_models import APIProvider, UsageSummary
|
|
from utils.asset_tracker import save_asset_to_library
|
|
from utils.file_storage import save_file_safely, generate_unique_filename, sanitize_filename
|
|
|
|
|
|
router = APIRouter(prefix="/api/images", tags=["images"])
|
|
logger = get_service_logger("api.images")
|
|
|
|
|
|
class ImageGenerateRequest(BaseModel):
|
|
prompt: str
|
|
negative_prompt: Optional[str] = None
|
|
provider: Optional[str] = Field(None, pattern="^(gemini|huggingface|stability)$")
|
|
model: Optional[str] = None
|
|
width: Optional[int] = Field(default=1024, ge=64, le=2048)
|
|
height: Optional[int] = Field(default=1024, ge=64, le=2048)
|
|
guidance_scale: Optional[float] = None
|
|
steps: Optional[int] = None
|
|
seed: Optional[int] = None
|
|
|
|
|
|
class ImageGenerateResponse(BaseModel):
|
|
success: bool = True
|
|
image_base64: str
|
|
image_url: Optional[str] = None # URL to saved image file
|
|
width: int
|
|
height: int
|
|
provider: str
|
|
model: Optional[str] = None
|
|
seed: Optional[int] = None
|
|
|
|
|
|
@router.post("/generate", response_model=ImageGenerateResponse)
|
|
def generate(
|
|
req: ImageGenerateRequest,
|
|
current_user: Dict[str, Any] = Depends(get_current_user),
|
|
db: Session = Depends(get_db)
|
|
) -> ImageGenerateResponse:
|
|
"""Generate image with subscription checking."""
|
|
try:
|
|
# Extract Clerk user ID (required)
|
|
if not current_user:
|
|
raise HTTPException(status_code=401, detail="Authentication required")
|
|
|
|
user_id = str(current_user.get('id', ''))
|
|
if not user_id:
|
|
raise HTTPException(status_code=401, detail="Invalid user ID in authentication token")
|
|
|
|
# Validation is now handled inside generate_image function
|
|
last_error: Optional[Exception] = None
|
|
result = None
|
|
for attempt in range(2): # simple single retry
|
|
try:
|
|
result = generate_image(
|
|
prompt=req.prompt,
|
|
options={
|
|
"negative_prompt": req.negative_prompt,
|
|
"provider": req.provider,
|
|
"model": req.model,
|
|
"width": req.width,
|
|
"height": req.height,
|
|
"guidance_scale": req.guidance_scale,
|
|
"steps": req.steps,
|
|
"seed": req.seed,
|
|
},
|
|
user_id=user_id, # Pass user_id for validation inside generate_image
|
|
)
|
|
image_b64 = base64.b64encode(result.image_bytes).decode("utf-8")
|
|
|
|
# Save image to disk and track in asset library
|
|
image_url = None
|
|
image_filename = None
|
|
image_path = None
|
|
|
|
try:
|
|
# Create output directory for image studio images
|
|
base_dir = Path(__file__).parent.parent
|
|
output_dir = base_dir / "image_studio_images"
|
|
|
|
# Generate safe filename from prompt
|
|
clean_prompt = sanitize_filename(req.prompt[:50], max_length=50)
|
|
image_filename = generate_unique_filename(
|
|
prefix=f"img_{clean_prompt}",
|
|
extension=".png",
|
|
include_uuid=True
|
|
)
|
|
|
|
# Save file safely
|
|
image_path, save_error = save_file_safely(
|
|
content=result.image_bytes,
|
|
directory=output_dir,
|
|
filename=image_filename,
|
|
max_file_size=50 * 1024 * 1024 # 50MB for images
|
|
)
|
|
|
|
if image_path and not save_error:
|
|
# Generate file URL (will be served via API endpoint)
|
|
image_url = f"/api/images/image-studio/images/{image_path.name}"
|
|
|
|
logger.info(f"[images.generate] Saved image to: {image_path} ({len(result.image_bytes)} bytes)")
|
|
|
|
# Save to asset library (non-blocking)
|
|
try:
|
|
asset_id = save_asset_to_library(
|
|
db=db,
|
|
user_id=user_id,
|
|
asset_type="image",
|
|
source_module="image_studio",
|
|
filename=image_path.name,
|
|
file_url=image_url,
|
|
file_path=str(image_path),
|
|
file_size=len(result.image_bytes),
|
|
mime_type="image/png",
|
|
title=req.prompt[:100] if len(req.prompt) <= 100 else req.prompt[:97] + "...",
|
|
description=f"Generated image: {req.prompt[:200]}" if len(req.prompt) > 200 else req.prompt,
|
|
prompt=req.prompt,
|
|
tags=["image_studio", "generated", result.provider] if result.provider else ["image_studio", "generated"],
|
|
provider=result.provider,
|
|
model=result.model,
|
|
asset_metadata={
|
|
"width": result.width,
|
|
"height": result.height,
|
|
"seed": result.seed,
|
|
"status": "completed",
|
|
"negative_prompt": req.negative_prompt
|
|
}
|
|
)
|
|
if asset_id:
|
|
logger.info(f"[images.generate] ✅ Asset saved to library: ID={asset_id}, filename={image_path.name}")
|
|
else:
|
|
logger.warning(f"[images.generate] Asset tracking returned None (may have failed silently)")
|
|
except Exception as asset_error:
|
|
logger.error(f"[images.generate] Failed to save asset to library: {asset_error}", exc_info=True)
|
|
# Don't fail the request if asset tracking fails
|
|
else:
|
|
logger.warning(f"[images.generate] Failed to save image to disk: {save_error}")
|
|
# Continue without failing the request - base64 is still available
|
|
except Exception as save_error:
|
|
logger.error(f"[images.generate] Unexpected error saving image: {save_error}", exc_info=True)
|
|
# Continue without failing the request
|
|
|
|
# TRACK USAGE after successful image generation
|
|
if result:
|
|
logger.info(f"[images.generate] ✅ Image generation successful, tracking usage for user {user_id}")
|
|
try:
|
|
db_track = next(get_db())
|
|
try:
|
|
# Get or create usage summary
|
|
pricing = PricingService(db_track)
|
|
current_period = pricing.get_current_billing_period(user_id) or datetime.now().strftime("%Y-%m")
|
|
|
|
logger.debug(f"[images.generate] Looking for usage summary: user_id={user_id}, period={current_period}")
|
|
|
|
summary = db_track.query(UsageSummary).filter(
|
|
UsageSummary.user_id == user_id,
|
|
UsageSummary.billing_period == current_period
|
|
).first()
|
|
|
|
if not summary:
|
|
logger.info(f"[images.generate] Creating new usage summary for user {user_id}, period {current_period}")
|
|
summary = UsageSummary(
|
|
user_id=user_id,
|
|
billing_period=current_period
|
|
)
|
|
db_track.add(summary)
|
|
db_track.flush() # Ensure summary is persisted before updating
|
|
|
|
# Get "before" state for unified log
|
|
current_calls_before = getattr(summary, "stability_calls", 0) or 0
|
|
|
|
# Update provider-specific counters (stability for image generation)
|
|
# Note: All image generation goes through STABILITY provider enum regardless of actual provider
|
|
new_calls = current_calls_before + 1
|
|
setattr(summary, "stability_calls", new_calls)
|
|
logger.debug(f"[images.generate] Updated stability_calls: {current_calls_before} -> {new_calls}")
|
|
|
|
# Update totals
|
|
old_total_calls = summary.total_calls or 0
|
|
summary.total_calls = old_total_calls + 1
|
|
logger.debug(f"[images.generate] Updated totals: calls {old_total_calls} -> {summary.total_calls}")
|
|
|
|
# Get plan details for unified log
|
|
limits = pricing.get_user_limits(user_id)
|
|
plan_name = limits.get('plan_name', 'unknown') if limits else 'unknown'
|
|
tier = limits.get('tier', 'unknown') if limits else 'unknown'
|
|
call_limit = limits['limits'].get("stability_calls", 0) if limits else 0
|
|
|
|
# Get image editing stats for unified log
|
|
current_image_edit_calls = getattr(summary, "image_edit_calls", 0) or 0
|
|
image_edit_limit = limits['limits'].get("image_edit_calls", 0) if limits else 0
|
|
|
|
# Get video stats for unified log
|
|
current_video_calls = getattr(summary, "video_calls", 0) or 0
|
|
video_limit = limits['limits'].get("video_calls", 0) if limits else 0
|
|
|
|
# Get audio stats for unified log
|
|
current_audio_calls = getattr(summary, "audio_calls", 0) or 0
|
|
audio_limit = limits['limits'].get("audio_calls", 0) if limits else 0
|
|
# Only show ∞ for Enterprise tier when limit is 0 (unlimited)
|
|
audio_limit_display = audio_limit if (audio_limit > 0 or tier != 'enterprise') else '∞'
|
|
|
|
db_track.commit()
|
|
logger.info(f"[images.generate] ✅ Successfully tracked usage: user {user_id} -> stability -> {new_calls} calls")
|
|
|
|
# UNIFIED SUBSCRIPTION LOG - Shows before/after state in one message
|
|
print(f"""
|
|
[SUBSCRIPTION] Image Generation
|
|
├─ User: {user_id}
|
|
├─ Plan: {plan_name} ({tier})
|
|
├─ Provider: stability
|
|
├─ Actual Provider: {result.provider}
|
|
├─ Model: {result.model or 'default'}
|
|
├─ Calls: {current_calls_before} → {new_calls} / {call_limit if call_limit > 0 else '∞'}
|
|
├─ Image Editing: {current_image_edit_calls} / {image_edit_limit if image_edit_limit > 0 else '∞'}
|
|
├─ Videos: {current_video_calls} / {video_limit if video_limit > 0 else '∞'}
|
|
├─ Audio: {current_audio_calls} / {audio_limit_display}
|
|
└─ Status: ✅ Allowed & Tracked
|
|
""")
|
|
except Exception as track_error:
|
|
logger.error(f"[images.generate] ❌ Error tracking usage (non-blocking): {track_error}", exc_info=True)
|
|
db_track.rollback()
|
|
finally:
|
|
db_track.close()
|
|
except Exception as usage_error:
|
|
# Non-blocking: log error but don't fail the request
|
|
logger.error(f"[images.generate] ❌ Failed to track usage: {usage_error}", exc_info=True)
|
|
|
|
return ImageGenerateResponse(
|
|
image_base64=image_b64,
|
|
image_url=image_url,
|
|
width=result.width,
|
|
height=result.height,
|
|
provider=result.provider,
|
|
model=result.model,
|
|
seed=result.seed,
|
|
)
|
|
except Exception as inner:
|
|
last_error = inner
|
|
logger.error(f"Image generation attempt {attempt+1} failed: {inner}")
|
|
# On first failure, try provider auto-remap by clearing provider to let facade decide
|
|
if attempt == 0 and req.provider:
|
|
req.provider = None
|
|
continue
|
|
break
|
|
raise last_error or RuntimeError("Unknown image generation error")
|
|
except Exception as e:
|
|
logger.error(f"Image generation failed: {e}")
|
|
# Provide a clean, actionable message to the client
|
|
raise HTTPException(
|
|
status_code=500,
|
|
detail="Image generation service is temporarily unavailable or the connection was reset. Please try again."
|
|
)
|
|
|
|
|
|
class PromptSuggestion(BaseModel):
|
|
prompt: str
|
|
negative_prompt: Optional[str] = None
|
|
width: Optional[int] = None
|
|
height: Optional[int] = None
|
|
overlay_text: Optional[str] = None
|
|
|
|
|
|
class ImagePromptSuggestRequest(BaseModel):
|
|
provider: Optional[str] = Field(None, pattern="^(gemini|huggingface|stability)$")
|
|
title: Optional[str] = None
|
|
section: Optional[Dict[str, Any]] = None
|
|
research: Optional[Dict[str, Any]] = None
|
|
persona: Optional[Dict[str, Any]] = None
|
|
include_overlay: Optional[bool] = True
|
|
|
|
|
|
class ImagePromptSuggestResponse(BaseModel):
|
|
suggestions: list[PromptSuggestion]
|
|
|
|
|
|
class ImageEditRequest(BaseModel):
|
|
image_base64: str
|
|
prompt: str
|
|
provider: Optional[str] = Field(None, pattern="^(huggingface)$")
|
|
model: Optional[str] = None
|
|
guidance_scale: Optional[float] = None
|
|
steps: Optional[int] = None
|
|
seed: Optional[int] = None
|
|
|
|
|
|
class ImageEditResponse(BaseModel):
|
|
success: bool = True
|
|
image_base64: str
|
|
image_url: Optional[str] = None # URL to saved edited image file
|
|
width: int
|
|
height: int
|
|
provider: str
|
|
model: Optional[str] = None
|
|
seed: Optional[int] = None
|
|
|
|
|
|
@router.post("/suggest-prompts", response_model=ImagePromptSuggestResponse)
|
|
def suggest_prompts(
|
|
req: ImagePromptSuggestRequest,
|
|
current_user: Dict[str, Any] = Depends(get_current_user)
|
|
) -> ImagePromptSuggestResponse:
|
|
try:
|
|
provider = (req.provider or ("gemini" if (os.getenv("GPT_PROVIDER") or "").lower().startswith("gemini") else "huggingface")).lower()
|
|
section = req.section or {}
|
|
title = (req.title or section.get("heading") or "").strip()
|
|
subheads = section.get("subheadings", []) or []
|
|
key_points = section.get("key_points", []) or []
|
|
keywords = section.get("keywords", []) or []
|
|
if not keywords and req.research:
|
|
keywords = (
|
|
req.research.get("keywords", {}).get("primary_keywords")
|
|
or req.research.get("keywords", {}).get("primary")
|
|
or []
|
|
)
|
|
|
|
persona = req.persona or {}
|
|
audience = persona.get("audience", "content creators and digital marketers")
|
|
industry = persona.get("industry", req.research.get("domain") if req.research else "your industry")
|
|
tone = persona.get("tone", "professional, trustworthy")
|
|
|
|
schema = {
|
|
"type": "object",
|
|
"properties": {
|
|
"suggestions": {
|
|
"type": "array",
|
|
"items": {
|
|
"type": "object",
|
|
"properties": {
|
|
"prompt": {"type": "string"},
|
|
"negative_prompt": {"type": "string"},
|
|
"width": {"type": "number"},
|
|
"height": {"type": "number"},
|
|
"overlay_text": {"type": "string"},
|
|
},
|
|
"required": ["prompt"]
|
|
},
|
|
"minItems": 3,
|
|
"maxItems": 5
|
|
}
|
|
},
|
|
"required": ["suggestions"]
|
|
}
|
|
|
|
system = (
|
|
"You are an expert image prompt engineer for text-to-image models. "
|
|
"Given blog section context, craft 3-5 hyper-personalized prompts optimized for the specified provider. "
|
|
"Return STRICT JSON matching the provided schema, no extra text."
|
|
)
|
|
|
|
provider_guidance = {
|
|
"huggingface": "Photorealistic Flux 1 Krea Dev; include camera/lighting cues (e.g., 50mm, f/2.8, rim light).",
|
|
"gemini": "Editorial, brand-safe, crisp edges, balanced lighting; avoid artifacts.",
|
|
"stability": "SDXL coherent details, sharp focus, cinematic contrast; readable text if present."
|
|
}.get(provider, "")
|
|
|
|
best_practices = (
|
|
"Best Practices: one clear focal subject; clean, uncluttered background; rule-of-thirds or center-weighted composition; "
|
|
"text-safe margins if overlay text is included; neutral lighting if unsure; realistic skin tones; avoid busy patterns; "
|
|
"no brand logos or watermarks; no copyrighted characters; avoid low-res, blur, noise, banding, oversaturation, over-sharpening; "
|
|
"ensure hands and text are coherent if present; prefer 1024px+ on shortest side for quality."
|
|
)
|
|
|
|
# Harvest a few concise facts from research if available
|
|
facts: list[str] = []
|
|
try:
|
|
if req.research:
|
|
# try common shapes used in research service
|
|
top_stats = req.research.get("key_facts") or req.research.get("highlights") or []
|
|
if isinstance(top_stats, list):
|
|
facts = [str(x) for x in top_stats[:3]]
|
|
elif isinstance(top_stats, dict):
|
|
facts = [f"{k}: {v}" for k, v in list(top_stats.items())[:3]]
|
|
except Exception:
|
|
facts = []
|
|
|
|
facts_line = ", ".join(facts) if facts else ""
|
|
|
|
overlay_hint = "Include an on-image short title or fact if it improves communication; ensure clean, high-contrast safe area for text." if (req.include_overlay is None or req.include_overlay) else "Do not include on-image text."
|
|
|
|
prompt = f"""
|
|
Provider: {provider}
|
|
Title: {title}
|
|
Subheadings: {', '.join(subheads[:5])}
|
|
Key Points: {', '.join(key_points[:5])}
|
|
Keywords: {', '.join([str(k) for k in keywords[:8]])}
|
|
Research Facts: {facts_line}
|
|
Audience: {audience}
|
|
Industry: {industry}
|
|
Tone: {tone}
|
|
|
|
Craft prompts that visually reflect this exact section (not generic blog topic). {provider_guidance}
|
|
{best_practices}
|
|
{overlay_hint}
|
|
Include a suitable negative_prompt where helpful. Suggest width/height when relevant (e.g., 1024x1024 or 1920x1080).
|
|
If including on-image text, return it in overlay_text (short: <= 8 words).
|
|
"""
|
|
|
|
# Get user_id for llm_text_gen subscription check (required)
|
|
if not current_user:
|
|
raise HTTPException(status_code=401, detail="Authentication required")
|
|
|
|
user_id_for_llm = str(current_user.get('id', ''))
|
|
if not user_id_for_llm:
|
|
raise HTTPException(status_code=401, detail="Invalid user ID in authentication token")
|
|
|
|
raw = llm_text_gen(prompt=prompt, system_prompt=system, json_struct=schema, user_id=user_id_for_llm)
|
|
data = raw if isinstance(raw, dict) else {}
|
|
suggestions = data.get("suggestions") or []
|
|
# basic fallback if provider returns string
|
|
if not suggestions and isinstance(raw, str):
|
|
suggestions = [{"prompt": raw}]
|
|
|
|
return ImagePromptSuggestResponse(suggestions=[PromptSuggestion(**s) for s in suggestions])
|
|
except Exception as e:
|
|
logger.error(f"Prompt suggestion failed: {e}")
|
|
raise HTTPException(status_code=500, detail=str(e))
|
|
|
|
|
|
@router.post("/edit", response_model=ImageEditResponse)
|
|
def edit(
|
|
req: ImageEditRequest,
|
|
current_user: Dict[str, Any] = Depends(get_current_user),
|
|
db: Session = Depends(get_db)
|
|
) -> ImageEditResponse:
|
|
"""Edit image with subscription checking."""
|
|
try:
|
|
# Extract Clerk user ID (required)
|
|
if not current_user:
|
|
raise HTTPException(status_code=401, detail="Authentication required")
|
|
|
|
user_id = str(current_user.get('id', ''))
|
|
if not user_id:
|
|
raise HTTPException(status_code=401, detail="Invalid user ID in authentication token")
|
|
|
|
# Decode base64 image
|
|
try:
|
|
input_image_bytes = base64.b64decode(req.image_base64)
|
|
except Exception as e:
|
|
raise HTTPException(status_code=400, detail=f"Invalid image_base64: {str(e)}")
|
|
|
|
# Validation is now handled inside edit_image function
|
|
result = edit_image(
|
|
input_image_bytes=input_image_bytes,
|
|
prompt=req.prompt,
|
|
options={
|
|
"provider": req.provider,
|
|
"model": req.model,
|
|
"guidance_scale": req.guidance_scale,
|
|
"steps": req.steps,
|
|
"seed": req.seed,
|
|
},
|
|
user_id=user_id, # Pass user_id for validation inside edit_image
|
|
)
|
|
edited_image_b64 = base64.b64encode(result.image_bytes).decode("utf-8")
|
|
|
|
# Save edited image to disk and track in asset library
|
|
image_url = None
|
|
image_filename = None
|
|
image_path = None
|
|
|
|
try:
|
|
# Create output directory for image studio edited images
|
|
base_dir = Path(__file__).parent.parent
|
|
output_dir = base_dir / "image_studio_images" / "edited"
|
|
|
|
# Generate safe filename from prompt
|
|
clean_prompt = sanitize_filename(req.prompt[:50], max_length=50)
|
|
image_filename = generate_unique_filename(
|
|
prefix=f"edited_{clean_prompt}",
|
|
extension=".png",
|
|
include_uuid=True
|
|
)
|
|
|
|
# Save file safely
|
|
image_path, save_error = save_file_safely(
|
|
content=result.image_bytes,
|
|
directory=output_dir,
|
|
filename=image_filename,
|
|
max_file_size=50 * 1024 * 1024 # 50MB for images
|
|
)
|
|
|
|
if image_path and not save_error:
|
|
# Generate file URL
|
|
image_url = f"/api/images/image-studio/images/edited/{image_path.name}"
|
|
|
|
logger.info(f"[images.edit] Saved edited image to: {image_path} ({len(result.image_bytes)} bytes)")
|
|
|
|
# Save to asset library (non-blocking)
|
|
try:
|
|
asset_id = save_asset_to_library(
|
|
db=db,
|
|
user_id=user_id,
|
|
asset_type="image",
|
|
source_module="image_studio",
|
|
filename=image_path.name,
|
|
file_url=image_url,
|
|
file_path=str(image_path),
|
|
file_size=len(result.image_bytes),
|
|
mime_type="image/png",
|
|
title=f"Edited: {req.prompt[:100]}" if len(req.prompt) <= 100 else f"Edited: {req.prompt[:97]}...",
|
|
description=f"Edited image with prompt: {req.prompt[:200]}" if len(req.prompt) > 200 else f"Edited image with prompt: {req.prompt}",
|
|
prompt=req.prompt,
|
|
tags=["image_studio", "edited", result.provider] if result.provider else ["image_studio", "edited"],
|
|
provider=result.provider,
|
|
model=result.model,
|
|
asset_metadata={
|
|
"width": result.width,
|
|
"height": result.height,
|
|
"seed": result.seed,
|
|
"status": "completed",
|
|
"operation": "edit"
|
|
}
|
|
)
|
|
if asset_id:
|
|
logger.info(f"[images.edit] ✅ Asset saved to library: ID={asset_id}, filename={image_path.name}")
|
|
else:
|
|
logger.warning(f"[images.edit] Asset tracking returned None (may have failed silently)")
|
|
except Exception as asset_error:
|
|
logger.error(f"[images.edit] Failed to save asset to library: {asset_error}", exc_info=True)
|
|
# Don't fail the request if asset tracking fails
|
|
else:
|
|
logger.warning(f"[images.edit] Failed to save edited image to disk: {save_error}")
|
|
# Continue without failing the request - base64 is still available
|
|
except Exception as save_error:
|
|
logger.error(f"[images.edit] Unexpected error saving edited image: {save_error}", exc_info=True)
|
|
# Continue without failing the request
|
|
|
|
# TRACK USAGE after successful image editing
|
|
if result:
|
|
logger.info(f"[images.edit] ✅ Image editing successful, tracking usage for user {user_id}")
|
|
try:
|
|
db_track = next(get_db())
|
|
try:
|
|
# Get or create usage summary
|
|
pricing = PricingService(db_track)
|
|
current_period = pricing.get_current_billing_period(user_id) or datetime.now().strftime("%Y-%m")
|
|
|
|
logger.debug(f"[images.edit] Looking for usage summary: user_id={user_id}, period={current_period}")
|
|
|
|
summary = db_track.query(UsageSummary).filter(
|
|
UsageSummary.user_id == user_id,
|
|
UsageSummary.billing_period == current_period
|
|
).first()
|
|
|
|
if not summary:
|
|
logger.info(f"[images.edit] Creating new usage summary for user {user_id}, period {current_period}")
|
|
summary = UsageSummary(
|
|
user_id=user_id,
|
|
billing_period=current_period
|
|
)
|
|
db_track.add(summary)
|
|
db_track.flush() # Ensure summary is persisted before updating
|
|
|
|
# Get "before" state for unified log
|
|
current_calls_before = getattr(summary, "image_edit_calls", 0) or 0
|
|
|
|
# Update image editing counters (separate from image generation)
|
|
new_calls = current_calls_before + 1
|
|
setattr(summary, "image_edit_calls", new_calls)
|
|
logger.debug(f"[images.edit] Updated image_edit_calls: {current_calls_before} -> {new_calls}")
|
|
|
|
# Update totals
|
|
old_total_calls = summary.total_calls or 0
|
|
summary.total_calls = old_total_calls + 1
|
|
logger.debug(f"[images.edit] Updated totals: calls {old_total_calls} -> {summary.total_calls}")
|
|
|
|
# Get plan details for unified log
|
|
limits = pricing.get_user_limits(user_id)
|
|
plan_name = limits.get('plan_name', 'unknown') if limits else 'unknown'
|
|
tier = limits.get('tier', 'unknown') if limits else 'unknown'
|
|
call_limit = limits['limits'].get("image_edit_calls", 0) if limits else 0
|
|
|
|
# Get image generation stats for unified log
|
|
current_image_gen_calls = getattr(summary, "stability_calls", 0) or 0
|
|
image_gen_limit = limits['limits'].get("stability_calls", 0) if limits else 0
|
|
|
|
# Get video stats for unified log
|
|
current_video_calls = getattr(summary, "video_calls", 0) or 0
|
|
video_limit = limits['limits'].get("video_calls", 0) if limits else 0
|
|
|
|
# Get audio stats for unified log
|
|
current_audio_calls = getattr(summary, "audio_calls", 0) or 0
|
|
audio_limit = limits['limits'].get("audio_calls", 0) if limits else 0
|
|
# Only show ∞ for Enterprise tier when limit is 0 (unlimited)
|
|
audio_limit_display = audio_limit if (audio_limit > 0 or tier != 'enterprise') else '∞'
|
|
|
|
db_track.commit()
|
|
logger.info(f"[images.edit] ✅ Successfully tracked usage: user {user_id} -> image_edit -> {new_calls} calls")
|
|
|
|
# UNIFIED SUBSCRIPTION LOG - Shows before/after state in one message
|
|
print(f"""
|
|
[SUBSCRIPTION] Image Editing
|
|
├─ User: {user_id}
|
|
├─ Plan: {plan_name} ({tier})
|
|
├─ Provider: image_edit
|
|
├─ Actual Provider: {result.provider}
|
|
├─ Model: {result.model or 'default'}
|
|
├─ Calls: {current_calls_before} → {new_calls} / {call_limit if call_limit > 0 else '∞'}
|
|
├─ Images: {current_image_gen_calls} / {image_gen_limit if image_gen_limit > 0 else '∞'}
|
|
├─ Videos: {current_video_calls} / {video_limit if video_limit > 0 else '∞'}
|
|
├─ Audio: {current_audio_calls} / {audio_limit_display}
|
|
└─ Status: ✅ Allowed & Tracked
|
|
""")
|
|
except Exception as track_error:
|
|
logger.error(f"[images.edit] ❌ Error tracking usage (non-blocking): {track_error}", exc_info=True)
|
|
db_track.rollback()
|
|
finally:
|
|
db_track.close()
|
|
except Exception as usage_error:
|
|
# Non-blocking: log error but don't fail the request
|
|
logger.error(f"[images.edit] ❌ Failed to track usage: {usage_error}", exc_info=True)
|
|
|
|
return ImageEditResponse(
|
|
image_base64=edited_image_b64,
|
|
image_url=image_url,
|
|
width=result.width,
|
|
height=result.height,
|
|
provider=result.provider,
|
|
model=result.model,
|
|
seed=result.seed,
|
|
)
|
|
except HTTPException:
|
|
raise
|
|
except Exception as e:
|
|
logger.error(f"Image editing failed: {e}", exc_info=True)
|
|
# Provide a clean, actionable message to the client
|
|
raise HTTPException(
|
|
status_code=500,
|
|
detail="Image editing service is temporarily unavailable or the connection was reset. Please try again."
|
|
)
|
|
|
|
|
|
# ---------------------------
|
|
# Image Serving Endpoints
|
|
# ---------------------------
|
|
|
|
@router.get("/image-studio/images/{image_filename:path}")
|
|
async def serve_image_studio_image(
|
|
image_filename: str,
|
|
current_user: Dict[str, Any] = Depends(get_current_user)
|
|
):
|
|
"""Serve a generated or edited image from Image Studio."""
|
|
try:
|
|
if not current_user:
|
|
raise HTTPException(status_code=401, detail="Authentication required")
|
|
|
|
# Determine if it's an edited image or regular image
|
|
base_dir = Path(__file__).parent.parent
|
|
image_studio_dir = (base_dir / "image_studio_images").resolve()
|
|
|
|
if image_filename.startswith("edited/"):
|
|
# Remove "edited/" prefix and serve from edited directory
|
|
actual_filename = image_filename.replace("edited/", "", 1)
|
|
image_path = (image_studio_dir / "edited" / actual_filename).resolve()
|
|
base_subdir = (image_studio_dir / "edited").resolve()
|
|
else:
|
|
image_path = (image_studio_dir / image_filename).resolve()
|
|
base_subdir = image_studio_dir
|
|
|
|
# Security: Prevent directory traversal attacks
|
|
# Ensure the resolved path is within the intended directory
|
|
try:
|
|
image_path.relative_to(base_subdir)
|
|
except ValueError:
|
|
raise HTTPException(
|
|
status_code=403,
|
|
detail="Access denied: Invalid image path"
|
|
)
|
|
|
|
if not image_path.exists():
|
|
raise HTTPException(status_code=404, detail="Image not found")
|
|
|
|
return FileResponse(
|
|
path=str(image_path),
|
|
media_type="image/png",
|
|
filename=image_path.name
|
|
)
|
|
except HTTPException:
|
|
raise
|
|
except Exception as e:
|
|
logger.error(f"[images] Failed to serve image: {e}")
|
|
raise HTTPException(status_code=500, detail=str(e))
|
|
|