video_optimization.py

"""
Configuración de optimización para modelos de video en ZeroGPU
Este archivo contiene configuraciones específicas para maximizar la velocidad
y eficiencia de los modelos de video en el Space NTIA.
"""

import torch
import os
from typing import Dict, Any, Optional

# Configuración de optimización para ZeroGPU H200
ZERO_GPU_CONFIG = {
    "max_memory_usage": "8GB",  # Máximo uso de memoria GPU
    "timeout_seconds": 60,      # Timeout máximo por generación
    "batch_size": 1,           # Batch size para evitar OOM
    "enable_attention_slicing": True,
    "enable_model_cpu_offload": True,
    "enable_xformers": True,
    "use_fp16": True,
    "use_compile": False,  # torch.compile puede causar problemas en ZeroGPU
}

# Configuraciones específicas por modelo de video
VIDEO_MODEL_CONFIGS = {
    "ByteDance/AnimateDiff-Lightning": {
        "name": "AnimateDiff Lightning",
        "category": "fast",
        "recommended_frames": (8, 16),
        "recommended_steps": (10, 20),
        "max_frames": 24,
        "max_steps": 25,
        "default_fps": 8,
        "optimizations": {
            "use_fp16": True,
            "attention_slicing": True,
            "model_cpu_offload": True,
            "guidance_scale": 7.5,
            "height": 256,
            "width": 256
        },
        "description": "⚡ Modelo más rápido para animaciones cortas"
    },
    
    "cerspense/zeroscope_v2_576w": {
        "name": "Zeroscope v2 576w",
        "category": "fast",
        "recommended_frames": (8, 16),
        "recommended_steps": (10, 20),
        "max_frames": 24,
        "max_steps": 25,
        "default_fps": 8,
        "optimizations": {
            "use_fp16": True,
            "attention_slicing": True,
            "model_cpu_offload": True,
            "guidance_scale": 7.5,
            "height": 256,
            "width": 256
        },
        "description": "⚡ Modelo rápido para videos cortos"
    },
    
    "damo-vilab/text-to-video-ms-1.7b": {
        "name": "Text-to-Video MS 1.7B",
        "category": "fast",
        "recommended_frames": (8, 16),
        "recommended_steps": (10, 20),
        "max_frames": 24,
        "max_steps": 25,
        "default_fps": 8,
        "optimizations": {
            "use_fp16": True,
            "attention_slicing": True,
            "model_cpu_offload": True,
            "guidance_scale": 7.5,
            "height": 256,
            "width": 256
        },
        "description": "⚡ Modelo rápido para conversión texto a video"
    },
    
    "cerspense/zeroscope_v2_XL": {
        "name": "Zeroscope v2 XL",
        "category": "standard",
        "recommended_frames": (12, 24),
        "recommended_steps": (20, 30),
        "max_frames": 32,
        "max_steps": 40,
        "default_fps": 6,
        "optimizations": {
            "use_fp16": True,
            "attention_slicing": True,
            "model_cpu_offload": True,
            "guidance_scale": 7.5,
            "height": 256,
            "width": 256
        },
        "description": "🎬 Modelo estándar con mejor calidad"
    },
    
    "THUDM/CogVideoX-5b": {
        "name": "CogVideoX 5B",
        "category": "quality",
        "recommended_frames": (16, 32),
        "recommended_steps": (25, 40),
        "max_frames": 48,
        "max_steps": 50,
        "default_fps": 6,
        "optimizations": {
            "use_fp16": True,
            "attention_slicing": True,
            "model_cpu_offload": True,
            "guidance_scale": 7.5,
            "height": 256,
            "width": 256
        },
        "description": "🎬 Modelo de alta calidad (más lento)"
    }
}

def get_model_config(model_name: str) -> Dict[str, Any]:
    """Obtener configuración específica para un modelo"""
    return VIDEO_MODEL_CONFIGS.get(model_name, {
        "name": model_name,
        "category": "standard",
        "recommended_frames": (12, 24),
        "recommended_steps": (20, 30),
        "max_frames": 32,
        "max_steps": 40,
        "default_fps": 6,
        "optimizations": {
            "use_fp16": True,
            "attention_slicing": True,
            "model_cpu_offload": True,
            "guidance_scale": 7.5,
            "height": 256,
            "width": 256
        },
        "description": "🎬 Modelo estándar"
    })

def optimize_parameters(model_name: str, num_frames: int, num_steps: int) -> tuple:
    """Optimizar parámetros basado en el modelo"""
    config = get_model_config(model_name)
    
    # Limitar frames y pasos según las recomendaciones del modelo
    min_frames, max_frames = config["recommended_frames"]
    min_steps, max_steps = config["recommended_steps"]
    
    optimized_frames = max(min_frames, min(num_frames, max_frames))
    optimized_steps = max(min_steps, min(num_steps, max_steps))
    
    return optimized_frames, optimized_steps

def get_fast_models() -> list:
    """Obtener lista de modelos rápidos"""
    return [name for name, config in VIDEO_MODEL_CONFIGS.items() 
            if config["category"] == "fast"]

def get_model_category(model_name: str) -> str:
    """Obtener categoría de un modelo"""
    config = get_model_config(model_name)
    return config["category"]

def should_use_fast_config(model_name: str) -> bool:
    """Determinar si usar configuración rápida"""
    return get_model_category(model_name) == "fast"

def get_optimization_tips(model_name: str) -> list:
    """Obtener consejos de optimización para un modelo"""
    config = get_model_config(model_name)
    category = config["category"]
    
    tips = []
    
    if category == "fast":
        tips.extend([
            "⚡ Usar 8-16 frames para máxima velocidad",
            "⚡ Usar 10-20 pasos de inferencia",
            "⚡ Modelo optimizado para ZeroGPU"
        ])
    elif category == "standard":
        tips.extend([
            "🎬 Usar 12-24 frames para balance velocidad/calidad",
            "🎬 Usar 20-30 pasos de inferencia",
            "🎬 Modelo equilibrado"
        ])
    elif category == "quality":
        tips.extend([
            "🌟 Usar 16-32 frames para máxima calidad",
            "🌟 Usar 25-40 pasos de inferencia",
            "🌟 Modelo de alta calidad (más lento)"
        ])
    
    return tips

# Configuración de manejo de errores
ERROR_HANDLING = {
    "quota_exceeded": {
        "message": "🚫 Cuota de ZeroGPU agotada",
        "suggestion": "Intenta en unos minutos o usa un modelo más rápido",
        "retry_after": 300  # 5 minutos
    },
    "out_of_memory": {
        "message": "💾 Error de memoria GPU",
        "suggestion": "Reduce frames o pasos de inferencia",
        "retry_after": 60   # 1 minuto
    },
    "timeout": {
        "message": "⏰ Timeout en la generación",
        "suggestion": "El modelo puede estar sobrecargado, intenta más tarde",
        "retry_after": 120  # 2 minutos
    },
    "model_not_found": {
        "message": "❌ Modelo no encontrado",
        "suggestion": "Verifica el nombre del modelo",
        "retry_after": 0    # No retry
    }
}

def get_error_info(error_type: str) -> Dict[str, Any]:
    """Obtener información de manejo de errores"""
    return ERROR_HANDLING.get(error_type, {
        "message": "❌ Error desconocido",
        "suggestion": "Intenta nuevamente",
        "retry_after": 60
    })