AI Video Generation Implementation

backend/api/images.py

@@ -9,6 +9,7 @@ from fastapi import APIRouter, HTTPException, Depends
 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
@@ -125,6 +126,14 @@ def generate(
                             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
+                            
                             db_track.commit()
                             logger.info(f"[images.generate] ✅ Successfully tracked usage: user {user_id} -> stability -> {new_calls} calls")
                             
@@ -137,6 +146,8 @@ def generate(
 ├─ 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 '∞'}
 └─ Status: ✅ Allowed & Tracked
 """)
                         except Exception as track_error:
@@ -195,6 +206,26 @@ 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
+    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,
@@ -316,3 +347,136 @@ def suggest_prompts(
         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)
+) -> 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")
+        
+        # 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
+                    
+                    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 '∞'}
+└─ 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,
+            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."
+        )
+