pipeline-app/app/cost_estimator.py

"""
Video generation cost estimator - multi-model pricing and cost prediction.

Estimates cost for generating scene videos using different models:
- Wan2.2 (local GPU, fast, lower cost)
- Wan2.7 (local GPU, higher quality, medium cost)
- Vidu 2.0 (cloud API, highest quality, premium cost)

Cost factors:
- Scene count from storyboard
- Duration per scene (frames)
- Model pricing (per frame or per second)
- GPU compute cost (for local models)
"""

import logging
from typing import List, Dict, Any

logger = logging.getLogger("pipeline.cost_estimator")


# Model pricing configuration (USD per frame at 24fps)
MODEL_PRICING = {
    "wan2.2": {
        "name": "Wan 2.2 (Local GPU)",
        "price_per_frame": 0.002,  # $0.002 per frame
        "quality_score": 7.5,
        "generation_speed": "fast",
        "description": "Fast generation, good quality, cost-effective",
        "features": ["T2V", "Ref2V"],
        "recommended_for": ["standard", "batch"],
    },
    "wan2.7": {
        "name": "Wan 2.7 (Local GPU)",
        "price_per_frame": 0.003,  # $0.003 per frame
        "quality_score": 8.5,
        "generation_speed": "medium",
        "description": "Higher quality, balanced performance",
        "features": ["T2V", "Ref2V", "I2V"],
        "recommended_for": ["quality", "professional"],
    },
    "vidu2.0": {
        "name": "Vidu 2.0 (Cloud API)",
        "price_per_frame": 0.008,  # $0.008 per frame
        "quality_score": 9.0,
        "generation_speed": "slow",
        "description": "Highest quality, premium pricing",
        "features": ["T2V", "I2V", "V2V"],
        "recommended_for": ["premium", "cinematic"],
    },
}


def estimate_scene_cost(storyboard: List[Dict[str, Any]], fps: int = 24) -> int:
    """
    Estimate total frames needed for all scenes.
    
    Args:
        storyboard: List of scene dicts with start_time/end_time
        fps: Frames per second (default 24)
    
    Returns:
        Total frame count
    """
    total_frames = 0
    for scene in storyboard:
        duration = scene.get("end_time", 10) - scene.get("start_time", 0)
        frames = int(duration * fps)
        total_frames += frames
    return total_frames


def estimate_model_cost(model_id: str, total_frames: int) -> Dict[str, Any]:
    """
    Estimate cost for a specific model.
    
    Args:
        model_id: Model identifier (wan2.2, wan2.7, vidu2.0)
        total_frames: Total frames to generate
    
    Returns:
        Cost estimate dict with breakdown
    """
    if model_id not in MODEL_PRICING:
        raise ValueError(f"Unknown model: {model_id}")
    
    model = MODEL_PRICING[model_id]
    base_cost = total_frames * model["price_per_frame"]
    
    # Add GPU compute overhead for local models
    if "Local GPU" in model["name"]:
        # Estimate GPU time: ~10 seconds per frame on 4090
        gpu_seconds = total_frames * 10
        # GPU cost: $0.50/hour = $0.000139/second
        gpu_cost = gpu_seconds * 0.000139
        total_cost = base_cost + gpu_cost
    else:
        # Cloud API includes compute in price
        total_cost = base_cost
        gpu_cost = 0
    
    return {
        "model_id": model_id,
        "model_name": model["name"],
        "total_frames": total_frames,
        "price_per_frame": model["price_per_frame"],
        "base_cost": round(base_cost, 4),
        "gpu_cost": round(gpu_cost, 4),
        "total_cost": round(total_cost, 4),
        "quality_score": model["quality_score"],
        "generation_speed": model["generation_speed"],
        "description": model["description"],
        "features": model["features"],
        "recommended_for": model["recommended_for"],
    }


def estimate_all_models(storyboard: List[Dict[str, Any]], fps: int = 24) -> List[Dict[str, Any]]:
    """
    Estimate cost for all available models.
    
    Args:
        storyboard: List of scene dicts
        fps: Frames per second
    
    Returns:
        List of cost estimates sorted by total_cost (ascending)
    """
    total_frames = estimate_scene_cost(storyboard, fps)
    scene_count = len(storyboard)
    
    estimates = []
    for model_id in MODEL_PRICING.keys():
        try:
            estimate = estimate_model_cost(model_id, total_frames)
            estimate["scene_count"] = scene_count
            estimate["fps"] = fps
            estimates.append(estimate)
        except Exception as e:
            logger.error(f"Failed to estimate cost for {model_id}: {e}")
    
    # Sort by total_cost (cheapest first)
    estimates.sort(key=lambda x: x["total_cost"])
    
    # Add recommendation flag
    if estimates:
        # Recommend middle option (best value)
        mid_idx = len(estimates) // 2
        estimates[mid_idx]["recommended"] = True
        for i, est in enumerate(estimates):
            if i != mid_idx:
                est["recommended"] = False
    
    return estimates


def format_cost_summary(estimates: List[Dict[str, Any]]) -> str:
    """
    Format cost estimates as human-readable summary.
    
    Args:
        estimates: List of cost estimate dicts
    
    Returns:
        Formatted string summary
    """
    lines = []
    lines.append("=" * 70)
    lines.append("视频生成模型费用预估")
    lines.append("=" * 70)
    
    if not estimates:
        lines.append("无可用模型")
        return "\n".join(lines)
    
    scene_count = estimates[0].get("scene_count", 0)
    total_frames = estimates[0].get("total_frames", 0)
    fps = estimates[0].get("fps", 24)
    
    lines.append(f"分镜数量: {scene_count}")
    lines.append(f"总帧数: {total_frames} (at {fps}fps)")
    lines.append("")
    
    for est in estimates:
        rec = " ⭐ 推荐" if est.get("recommended") else ""
        lines.append(f"【{est['model_name']}】{rec}")
        lines.append(f"  质量评分: {est['quality_score']}/10")
        lines.append(f"  生成速度: {est['generation_speed']}")
        lines.append(f"  功能: {', '.join(est['features'])}")
        lines.append(f"  描述: {est['description']}")
        lines.append(f"  费用明细:")
        lines.append(f"    - 基础费用: ${est['base_cost']:.4f}")
        if est['gpu_cost'] > 0:
            lines.append(f"    - GPU计算: ${est['gpu_cost']:.4f}")
        lines.append(f"    - 总计: ${est['total_cost']:.4f}")
        lines.append("")
    
    lines.append("=" * 70)
    return "\n".join(lines)