mirror of
https://github.com/Memo-2023/mana-monorepo.git
synced 2026-05-20 19:46:41 +02:00
06107f6a52
New GPU service for fast text-to-video generation using LTX-Video (~2B params) on the RTX 3090. Generates 480p clips in 10-30 seconds, uses ~10GB VRAM. Includes Cloudflare Tunnel route, Prometheus monitoring, and health checks. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
277 lines
7.8 KiB
Python
277 lines
7.8 KiB
Python
"""
|
|
LTX-Video Service
|
|
|
|
Uses LTX-Video 0.9.x via HuggingFace diffusers for fast video generation.
|
|
Optimized for NVIDIA GPUs (CUDA).
|
|
"""
|
|
|
|
import asyncio
|
|
import logging
|
|
import os
|
|
import time
|
|
import uuid
|
|
from dataclasses import dataclass
|
|
from pathlib import Path
|
|
from typing import Optional
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
# Configuration
|
|
MODEL_ID = os.getenv("LTX_MODEL_ID", "Lightricks/LTX-Video")
|
|
DEFAULT_WIDTH = int(os.getenv("DEFAULT_WIDTH", "704"))
|
|
DEFAULT_HEIGHT = int(os.getenv("DEFAULT_HEIGHT", "480"))
|
|
DEFAULT_NUM_FRAMES = int(os.getenv("DEFAULT_NUM_FRAMES", "81")) # ~3.2s at 25fps
|
|
DEFAULT_FPS = int(os.getenv("DEFAULT_FPS", "25"))
|
|
DEFAULT_STEPS = int(os.getenv("DEFAULT_STEPS", "30"))
|
|
DEFAULT_GUIDANCE_SCALE = float(os.getenv("DEFAULT_GUIDANCE_SCALE", "7.5"))
|
|
GENERATION_TIMEOUT = int(os.getenv("GENERATION_TIMEOUT", "600")) # seconds
|
|
DEVICE = os.getenv("DEVICE", "cuda")
|
|
|
|
# Output directory for generated videos
|
|
OUTPUT_DIR = Path(os.getenv("OUTPUT_DIR", "/tmp/mana-video-gen"))
|
|
OUTPUT_DIR.mkdir(parents=True, exist_ok=True)
|
|
|
|
# Pipeline singleton (lazy loaded)
|
|
_pipeline = None
|
|
_pipeline_lock = asyncio.Lock()
|
|
|
|
|
|
@dataclass
|
|
class GenerationResult:
|
|
"""Result of video generation."""
|
|
|
|
video_path: str
|
|
prompt: str
|
|
width: int
|
|
height: int
|
|
num_frames: int
|
|
fps: int
|
|
steps: int
|
|
seed: int
|
|
generation_time: float
|
|
|
|
|
|
def is_model_available() -> bool:
|
|
"""Check if the model can be loaded (CUDA available + diffusers installed)."""
|
|
try:
|
|
import torch
|
|
|
|
return torch.cuda.is_available()
|
|
except ImportError:
|
|
return False
|
|
|
|
|
|
def get_model_info() -> dict:
|
|
"""Get information about the LTX-Video model."""
|
|
try:
|
|
import torch
|
|
|
|
gpu_name = torch.cuda.get_device_name(0) if torch.cuda.is_available() else "N/A"
|
|
vram_gb = (
|
|
round(torch.cuda.get_device_properties(0).total_mem / 1e9, 1)
|
|
if torch.cuda.is_available()
|
|
else 0
|
|
)
|
|
except ImportError:
|
|
gpu_name = "N/A"
|
|
vram_gb = 0
|
|
|
|
return {
|
|
"model_id": MODEL_ID,
|
|
"model_name": "LTX-Video",
|
|
"parameters": "~2B",
|
|
"license": "Lightricks Open License (commercial OK)",
|
|
"cuda_available": is_model_available(),
|
|
"gpu": gpu_name,
|
|
"vram_gb": vram_gb,
|
|
"default_resolution": f"{DEFAULT_WIDTH}x{DEFAULT_HEIGHT}",
|
|
"default_frames": DEFAULT_NUM_FRAMES,
|
|
"default_fps": DEFAULT_FPS,
|
|
"default_steps": DEFAULT_STEPS,
|
|
"pipeline_loaded": _pipeline is not None,
|
|
}
|
|
|
|
|
|
async def _load_pipeline():
|
|
"""Load the LTX-Video pipeline (lazy, thread-safe)."""
|
|
global _pipeline
|
|
|
|
async with _pipeline_lock:
|
|
if _pipeline is not None:
|
|
return _pipeline
|
|
|
|
logger.info(f"Loading LTX-Video pipeline from {MODEL_ID}...")
|
|
load_start = time.time()
|
|
|
|
def _load():
|
|
import torch
|
|
from diffusers import LTXPipeline
|
|
|
|
pipe = LTXPipeline.from_pretrained(
|
|
MODEL_ID,
|
|
torch_dtype=torch.bfloat16,
|
|
)
|
|
pipe.to(DEVICE)
|
|
# Enable memory optimizations
|
|
pipe.enable_model_cpu_offload()
|
|
return pipe
|
|
|
|
loop = asyncio.get_event_loop()
|
|
_pipeline = await loop.run_in_executor(None, _load)
|
|
|
|
load_time = time.time() - load_start
|
|
logger.info(f"LTX-Video pipeline loaded in {load_time:.1f}s")
|
|
return _pipeline
|
|
|
|
|
|
async def unload_pipeline():
|
|
"""Unload pipeline to free VRAM."""
|
|
global _pipeline
|
|
|
|
async with _pipeline_lock:
|
|
if _pipeline is not None:
|
|
import torch
|
|
|
|
del _pipeline
|
|
_pipeline = None
|
|
torch.cuda.empty_cache()
|
|
logger.info("LTX-Video pipeline unloaded, VRAM freed")
|
|
|
|
|
|
async def generate_video(
|
|
prompt: str,
|
|
width: int = DEFAULT_WIDTH,
|
|
height: int = DEFAULT_HEIGHT,
|
|
num_frames: int = DEFAULT_NUM_FRAMES,
|
|
fps: int = DEFAULT_FPS,
|
|
steps: int = DEFAULT_STEPS,
|
|
guidance_scale: float = DEFAULT_GUIDANCE_SCALE,
|
|
seed: Optional[int] = None,
|
|
negative_prompt: str = "",
|
|
) -> GenerationResult:
|
|
"""
|
|
Generate a video from a text prompt using LTX-Video.
|
|
|
|
Args:
|
|
prompt: Text prompt for video generation
|
|
width: Video width (default 704)
|
|
height: Video height (default 480)
|
|
num_frames: Number of frames (default 81 = ~3.2s at 25fps)
|
|
fps: Frames per second for output (default 25)
|
|
steps: Number of inference steps (default 30)
|
|
guidance_scale: CFG scale (default 7.5)
|
|
seed: Random seed (None for random)
|
|
negative_prompt: Negative prompt
|
|
|
|
Returns:
|
|
GenerationResult with video path and metadata
|
|
|
|
Raises:
|
|
RuntimeError: If model not available or generation fails
|
|
"""
|
|
if not is_model_available():
|
|
raise RuntimeError("CUDA not available - cannot generate video")
|
|
|
|
pipe = await _load_pipeline()
|
|
|
|
video_id = str(uuid.uuid4())[:8]
|
|
output_path = OUTPUT_DIR / f"{video_id}.mp4"
|
|
|
|
import torch
|
|
|
|
actual_seed = seed if seed is not None else torch.randint(0, 2**32, (1,)).item()
|
|
generator = torch.Generator(device="cpu").manual_seed(actual_seed)
|
|
|
|
logger.info(
|
|
f"Generating video: {width}x{height}, {num_frames} frames, {steps} steps, seed={actual_seed}"
|
|
)
|
|
|
|
start_time = time.time()
|
|
|
|
def _generate():
|
|
result = pipe(
|
|
prompt=prompt,
|
|
negative_prompt=negative_prompt or None,
|
|
width=width,
|
|
height=height,
|
|
num_frames=num_frames,
|
|
num_inference_steps=steps,
|
|
guidance_scale=guidance_scale,
|
|
generator=generator,
|
|
)
|
|
return result
|
|
|
|
try:
|
|
loop = asyncio.get_event_loop()
|
|
result = await asyncio.wait_for(
|
|
loop.run_in_executor(None, _generate),
|
|
timeout=GENERATION_TIMEOUT,
|
|
)
|
|
|
|
generation_time = time.time() - start_time
|
|
|
|
# Export to MP4
|
|
from diffusers.utils import export_to_video
|
|
|
|
export_to_video(result.frames[0], str(output_path), fps=fps)
|
|
|
|
if not output_path.exists():
|
|
raise RuntimeError("Video generation completed but output file not found")
|
|
|
|
file_size_mb = output_path.stat().st_size / (1024 * 1024)
|
|
duration_s = num_frames / fps
|
|
|
|
logger.info(
|
|
f"Video generated: {output_path} ({width}x{height}, {duration_s:.1f}s, "
|
|
f"{file_size_mb:.1f}MB, took {generation_time:.1f}s)"
|
|
)
|
|
|
|
return GenerationResult(
|
|
video_path=str(output_path),
|
|
prompt=prompt,
|
|
width=width,
|
|
height=height,
|
|
num_frames=num_frames,
|
|
fps=fps,
|
|
steps=steps,
|
|
seed=actual_seed,
|
|
generation_time=generation_time,
|
|
)
|
|
|
|
except asyncio.TimeoutError:
|
|
logger.error(f"Video generation timed out after {GENERATION_TIMEOUT}s")
|
|
raise RuntimeError(f"Generation timed out after {GENERATION_TIMEOUT} seconds")
|
|
except Exception as e:
|
|
logger.error(f"Video generation error: {e}")
|
|
raise
|
|
|
|
|
|
def cleanup_video(video_path: str) -> bool:
|
|
"""Delete a generated video file."""
|
|
try:
|
|
path = Path(video_path)
|
|
if path.exists() and path.parent == OUTPUT_DIR:
|
|
path.unlink()
|
|
return True
|
|
except Exception as e:
|
|
logger.warning(f"Failed to cleanup video {video_path}: {e}")
|
|
return False
|
|
|
|
|
|
def cleanup_old_videos(max_age_hours: int = 24) -> int:
|
|
"""Clean up videos older than max_age_hours."""
|
|
cleaned = 0
|
|
cutoff = time.time() - (max_age_hours * 3600)
|
|
|
|
try:
|
|
for file in OUTPUT_DIR.iterdir():
|
|
if file.is_file() and file.stat().st_mtime < cutoff:
|
|
file.unlink()
|
|
cleaned += 1
|
|
except Exception as e:
|
|
logger.warning(f"Cleanup error: {e}")
|
|
|
|
if cleaned > 0:
|
|
logger.info(f"Cleaned up {cleaned} old videos")
|
|
|
|
return cleaned
|