Spaces:

Ntdeseb
/

ntia

Running

ntia / video_optimization.py

Optimización completa de modelos de video para ZeroGPU H200

34fc777 about 1 year ago

7.38 kB

"""
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
    }) 

	"""
	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
	})