app.py

# FILE: ltx_server_refactored_complete.py
# DESCRIPTION: Backend service for video generation using LTX-Video pipeline.
#              Features modular generation, narrative chunking, and resource management.

import gc
import io
import json
import logging
import os
import random
import shutil
import subprocess
import sys
import tempfile
import time
import traceback
import warnings
from pathlib import Path
from typing import Dict, List, Optional, Tuple

import torch
import yaml
from einops import rearrange
from huggingface_hub import hf_hub_download

# ==============================================================================
# --- INITIAL SETUP & CONFIGURATION ---
# ==============================================================================

# Suppress excessive logs from external libraries
warnings.filterwarnings("ignore")
logging.getLogger("huggingface_hub").setLevel(logging.ERROR)
logging.basicConfig(level=logging.INFO, format='[%(levelname)s] %(message)s')

# --- CONSTANTS ---
DEPS_DIR = Path("/data")
LTX_VIDEO_REPO_DIR = DEPS_DIR / "LTX-Video"
BASE_CONFIG_PATH = LTX_VIDEO_REPO_DIR / "configs"
DEFAULT_CONFIG_FILE = BASE_CONFIG_PATH / "ltxv-13b-0.9.8-distilled-fp8.yaml"
LTX_REPO_ID = "Lightricks/LTX-Video"
RESULTS_DIR = Path("/app/output")
DEFAULT_FPS = 24.0
FRAMES_ALIGNMENT = 8

# --- DEPENDENCY PATH SETUP ---
# Ensures the LTX-Video library can be imported
def add_deps_to_path():
    """Adds the LTX repository directory to the Python system path."""
    repo_path = str(LTX_VIDEO_REPO_DIR.resolve())
    if repo_path not in sys.path:
        sys.path.insert(0, repo_path)
        logging.info(f"Repo added to sys.path: {repo_path}")

add_deps_to_path()

# --- PROJECT IMPORTS ---
# These must come after the path setup
from api.gpu_manager import gpu_manager
from ltx_video.models.autoencoders.vae_encode import (normalize_latents, un_normalize_latents)
from ltx_video.pipelines.pipeline_ltx_video import (ConditioningItem, LTXMultiScalePipeline, adain_filter_latent)
from ltx_video.pipelines.pipeline_ltx_video import create_ltx_video_pipeline, create_latent_upsampler
from ltx_video.utils.inference_utils import load_image_to_tensor_with_resize_and_crop
from managers.vae_manager import vae_manager_singleton
from tools.video_encode_tool import video_encode_tool_singleton


# ==============================================================================
# --- UTILITY & HELPER FUNCTIONS ---
# ==============================================================================

def seed_everything(seed: int):
    """Sets the seed for reproducibility across all relevant libraries."""
    random.seed(seed)
    os.environ['PYTHONHASHSEED'] = str(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)
    # Potentially faster, but less reproducible
    # torch.backends.cudnn.deterministic = False
    # torch.backends.cudnn.benchmark = True

def calculate_padding(orig_h: int, orig_w: int, target_h: int, target_w: int) -> Tuple[int, int, int, int]:
    """Calculates symmetric padding values to reach a target dimension."""
    pad_h = target_h - orig_h
    pad_w = target_w - orig_w
    pad_top = pad_h // 2
    pad_bottom = pad_h - pad_top
    pad_left = pad_w // 2
    pad_right = pad_w - pad_left
    return (pad_left, pad_right, pad_top, pad_bottom)

def log_tensor_info(tensor: torch.Tensor, name: str = "Tensor"):
    """Logs detailed information about a PyTorch tensor for debugging."""
    if not isinstance(tensor, torch.Tensor):
        logging.debug(f"'{name}' is not a tensor.")
        return
    
    info_str = (
        f"--- Tensor: {name} ---\n"
        f"  - Shape: {tuple(tensor.shape)}\n"
        f"  - Dtype: {tensor.dtype}\n"
        f"  - Device: {tensor.device}\n"
    )
    if tensor.numel() > 0:
        try:
            info_str += (
                f"  - Min: {tensor.min().item():.4f} | "
                f"Max: {tensor.max().item():.4f} | "
                f"Mean: {tensor.mean().item():.4f}\n"
            )
        except Exception:
            pass # Fails on some dtypes
    logging.debug(info_str + "----------------------")


# ==============================================================================
# --- VIDEO SERVICE CLASS ---
# ==============================================================================

class VideoService:
    """
    Backend service for orchestrating video generation using the LTX-Video pipeline.
    Encapsulates model loading, state management, and the logic for multi-stage
    video generation (low-resolution, upscale).
    """

    def __init__(self):
        """Initializes the service, loads models, and configures the environment."""
        t0 = time.perf_counter()
        logging.info("Initializing VideoService...")
        RESULTS_DIR.mkdir(parents=True, exist_ok=True)

        self.config = self._load_config(DEFAULT_CONFIG_FILE)
        self._tmp_dirs = set()
        
        self.pipeline, self.latent_upsampler = self._load_models_on_cpu()

        target_device = gpu_manager.get_ltx_device()
        self.device = torch.device("cpu") # Default device
        self.move_to_device(target_device)

        self._apply_precision_policy()
        vae_manager_singleton.attach_pipeline(
            self.pipeline,
            device=self.device,
            autocast_dtype=self.runtime_autocast_dtype
        )
        
        logging.info(f"VideoService ready. Startup time: {time.perf_counter()-t0:.2f}s")

    # ==========================================================================
    # --- LIFECYCLE & MODEL MANAGEMENT ---
    # ==========================================================================

    def _load_config(self, config_path: Path) -> Dict:
        """Loads the YAML configuration file."""
        logging.info(f"Loading config from: {config_path}")
        with open(config_path, "r") as file:
            return yaml.safe_load(file)

    def _load_models_on_cpu(self) -> Tuple[LTXMultiScalePipeline, Optional[torch.nn.Module]]:
        """Downloads and loads the pipeline and upsampler checkpoints onto the CPU."""
        t0 = time.perf_counter()
        
        logging.info("Downloading main checkpoint...")
        distilled_model_path = hf_hub_download(
            repo_id=LTX_REPO_ID,
            filename=self.config["checkpoint_path"],
            token=os.getenv("HF_TOKEN"),
        )
        self.config["checkpoint_path"] = distilled_model_path

        pipeline = create_ltx_video_pipeline(
            ckpt_path=self.config["checkpoint_path"],
            precision=self.config["precision"],
            device="cpu", # Load on CPU first
            # Pass other config values directly
            **{k: v for k, v in self.config.items() if k in create_ltx_video_pipeline.__code__.co_varnames}
        )
        
        latent_upsampler = None
        if self.config.get("spatial_upscaler_model_path"):
            logging.info("Downloading spatial upscaler checkpoint...")
            spatial_upscaler_path = hf_hub_download(
                repo_id=LTX_REPO_ID,
                filename=self.config["spatial_upscaler_model_path"],
                token=os.getenv("HF_TOKEN")
            )
            self.config["spatial_upscaler_model_path"] = spatial_upscaler_path
            latent_upsampler = create_latent_upsampler(self.config["spatial_upscaler_model_path"], device="cpu")

        logging.info(f"Models loaded on CPU in {time.perf_counter()-t0:.2f}s")
        return pipeline, latent_upsampler

    def move_to_device(self, device_str: str):
        """Moves all relevant models to the specified device (e.g., 'cuda:0' or 'cpu')."""
        target_device = torch.device(device_str)
        if self.device == target_device:
            logging.info(f"Models are already on the target device: {device_str}")
            return
        
        logging.info(f"Moving models to {device_str}...")
        self.device = target_device
        self.pipeline.to(self.device)
        if self.latent_upsampler:
            self.latent_upsampler.to(self.device)
        
        if device_str == "cpu" and torch.cuda.is_available():
            torch.cuda.empty_cache()
            
        logging.info(f"Models successfully moved to {self.device}.")

    def finalize(self, keep_paths: Optional[List[str]] = None):
        """Cleans up GPU memory and temporary directories."""
        logging.debug("Finalizing resources...")
        gc.collect()
        if torch.cuda.is_available():
            torch.cuda.empty_cache()
            try:
                torch.cuda.ipc_collect()
            except Exception:
                pass
        
        # Optional: Clean up temporary directories if needed (logic can be added here)


    # ==========================================================================
    # --- PUBLIC ORCHESTRATORS ---
    # These are the main entry points called by the frontend.
    # ==========================================================================

    def generate_narrative_low(self, prompt: str, **kwargs) -> Tuple[Optional[str], Optional[str], Optional[int]]:
        """
        [ORCHESTRATOR] Generates a video from a multi-line prompt, creating a sequence of scenes.
        
        Returns:
            A tuple of (video_path, latents_path, used_seed).
        """
        logging.info("Starting narrative low-res generation...")
        used_seed = self._resolve_seed(kwargs.get("seed"))
        seed_everything(used_seed)

        prompt_list = [p.strip() for p in prompt.splitlines() if p.strip()]
        if not prompt_list:
            raise ValueError("Prompt is empty or contains no valid lines.")

        num_chunks = len(prompt_list)
        total_frames = self._calculate_aligned_frames(kwargs.get("duration", 4.0))
        frames_per_chunk = (total_frames // num_chunks // FRAMES_ALIGNMENT) * FRAMES_ALIGNMENT
        overlap_frames = self.config.get("overlap_frames", 8)
        
        all_latents_paths = []
        overlap_condition_item = None
        
        try:
            for i, chunk_prompt in enumerate(prompt_list):
                logging.info(f"Generating narrative chunk {i+1}/{num_chunks}: '{chunk_prompt[:50]}...'")

                current_frames = frames_per_chunk
                if i > 0:
                    current_frames += overlap_frames
                
                # Use initial image conditions only for the first chunk
                current_conditions = kwargs.get("initial_conditions", []) if i == 0 else []
                if overlap_condition_item:
                    current_conditions.append(overlap_condition_item)

                chunk_latents = self._generate_single_chunk_low(
                    prompt=chunk_prompt,
                    num_frames=current_frames,
                    seed=used_seed + i,
                    conditioning_items=current_conditions,
                    **kwargs
                )

                if chunk_latents is None:
                    raise RuntimeError(f"Failed to generate latents for chunk {i+1}.")

                # Create overlap for the next chunk
                if i < num_chunks - 1:
                    overlap_latents = chunk_latents[:, :, -overlap_frames:, :, :].clone()
                    log_tensor_info(overlap_latents, f"Overlap Latents from chunk {i+1}")
                    overlap_condition_item = ConditioningItem(
                        media_item=overlap_latents, media_frame_number=0, conditioning_strength=1.0
                    )
                
                # Trim the overlap from the current chunk before saving
                if i > 0:
                    chunk_latents = chunk_latents[:, :, overlap_frames:, :, :]
                
                # Save chunk latents to disk to manage memory
                chunk_path = RESULTS_DIR / f"chunk_{i}_{used_seed}.pt"
                torch.save(chunk_latents.cpu(), chunk_path)
                all_latents_paths.append(chunk_path)
            
            # Concatenate, decode, and save the final video
            return self._finalize_generation(all_latents_paths, "narrative_video", used_seed)

        except Exception as e:
            logging.error(f"Error during narrative generation: {e}")
            traceback.print_exc()
            return None, None, None
        finally:
            # Clean up intermediate chunk files
            for path in all_latents_paths:
                if os.path.exists(path):
                    os.remove(path)
            self.finalize()


    def generate_single_low(self, **kwargs) -> Tuple[Optional[str], Optional[str], Optional[int]]:
        """
        [ORCHESTRATOR] Generates a video from a single prompt in one go.
        
        Returns:
            A tuple of (video_path, latents_path, used_seed).
        """
        logging.info("Starting single-prompt low-res generation...")
        used_seed = self._resolve_seed(kwargs.get("seed"))
        seed_everything(used_seed)
        
        try:
            total_frames = self._calculate_aligned_frames(kwargs.get("duration", 4.0), min_frames=9)
            
            final_latents = self._generate_single_chunk_low(
                num_frames=total_frames,
                seed=used_seed,
                conditioning_items=kwargs.get("initial_conditions", []),
                **kwargs
            )
            
            if final_latents is None:
                raise RuntimeError("Failed to generate latents.")

            # Save latents to a single file, then decode and save video
            latents_path = RESULTS_DIR / f"single_{used_seed}.pt"
            torch.save(final_latents.cpu(), latents_path)
            return self._finalize_generation([latents_path], "single_video", used_seed)

        except Exception as e:
            logging.error(f"Error during single generation: {e}")
            traceback.print_exc()
            return None, None, None
        finally:
            self.finalize()
    

    # ==========================================================================
    # --- INTERNAL WORKER UNITS ---
    # ==========================================================================

    def _generate_single_chunk_low(
        self, prompt: str, negative_prompt: str, height: int, width: int, num_frames: int, seed: int,
        conditioning_items: List[ConditioningItem], ltx_configs_override: Optional[Dict], **kwargs
    ) -> Optional[torch.Tensor]:
        """
        [WORKER] Generates a single chunk of latents. This is the core generation unit.
        Returns the raw latents tensor on the target device, or None on failure.
        """
        height_padded, width_padded = (self._align(d) for d in (height, width))
        downscale_factor = self.config.get("downscale_factor", 0.6666666)
        vae_scale_factor = self.pipeline.vae_scale_factor

        downscaled_height = self._align(int(height_padded * downscale_factor), vae_scale_factor)
        downscaled_width = self._align(int(width_padded * downscale_factor), vae_scale_factor)

        first_pass_config = self.config.get("first_pass", {}).copy()
        if ltx_configs_override:
             first_pass_config.update(self._prepare_guidance_overrides(ltx_configs_override))

        pipeline_kwargs = {
            "prompt": prompt,
            "negative_prompt": negative_prompt,
            "height": downscaled_height,
            "width": downscaled_width,
            "num_frames": num_frames,
            "frame_rate": DEFAULT_FPS,
            "generator": torch.Generator(device=self.device).manual_seed(seed),
            "output_type": "latent",
            "conditioning_items": conditioning_items,
            **first_pass_config
        }
        
        logging.debug(f"Pipeline call args: { {k: v for k, v in pipeline_kwargs.items() if k != 'conditioning_items'} }")
        
        with torch.autocast(device_type=self.device.type, dtype=self.runtime_autocast_dtype, enabled=self.device.type == 'cuda'):
            latents_raw = self.pipeline(**pipeline_kwargs).images
        
        log_tensor_info(latents_raw, f"Raw Latents for '{prompt[:40]}...'")
        return latents_raw


    # ==========================================================================
    # --- HELPERS & UTILITY METHODS ---
    # ==========================================================================
    
    def _finalize_generation(self, latents_paths: List[Path], base_filename: str, seed: int) -> Tuple[str, str, int]:
        """
        Loads latents from paths, concatenates them, decodes to video, and saves both.
        """
        logging.info("Finalizing generation: decoding latents to video.")
        # Load all tensors and concatenate them on the CPU first
        all_tensors_cpu = [torch.load(p) for p in latents_paths]
        final_latents_cpu = torch.cat(all_tensors_cpu, dim=2)
        
        # Save final combined latents
        final_latents_path = RESULTS_DIR / f"latents_{base_filename}_{seed}.pt"
        torch.save(final_latents_cpu, final_latents_path)
        logging.info(f"Final latents saved to: {final_latents_path}")
        
        # Move to GPU for decoding
        final_latents_gpu = final_latents_cpu.to(self.device)
        log_tensor_info(final_latents_gpu, "Final Concatenated Latents")

        with torch.autocast(device_type=self.device.type, dtype=self.runtime_autocast_dtype, enabled=self.device.type == 'cuda'):
            pixel_tensor = vae_manager_singleton.decode(
                final_latents_gpu,
                decode_timestep=float(self.config.get("decode_timestep", 0.05))
            )
        
        video_path = self._save_and_log_video(pixel_tensor, f"{base_filename}_{seed}")
        return str(video_path), str(final_latents_path), seed

    def prepare_condition_items(self, items_list: List, height: int, width: int, num_frames: int) -> List[ConditioningItem]:
        """Prepares a list of ConditioningItem objects from file paths or tensors."""
        if not items_list:
            return []
        
        height_padded, width_padded = self._align(height), self._align(width)
        padding_values = calculate_padding(height, width, height_padded, width_padded)
        
        conditioning_items = []
        for media, frame, weight in items_list:
            tensor = self._prepare_conditioning_tensor(media, height, width, padding_values)
            safe_frame = max(0, min(int(frame), num_frames - 1))
            conditioning_items.append(ConditioningItem(tensor, safe_frame, float(weight)))
        return conditioning_items

    def _prepare_conditioning_tensor(self, media_path: str, height: int, width: int, padding: Tuple) -> torch.Tensor:
        """Loads and processes an image to be a conditioning tensor."""
        tensor = load_image_to_tensor_with_resize_and_crop(media_path, height, width)
        tensor = torch.nn.functional.pad(tensor, padding)
        log_tensor_info(tensor, f"Prepared Conditioning Tensor from {media_path}")
        return tensor.to(self.device, dtype=self.runtime_autocast_dtype)

    def _prepare_guidance_overrides(self, ltx_configs: Dict) -> Dict:
        """Parses UI presets for guidance into pipeline-compatible arguments."""
        overrides = {}
        preset = ltx_configs.get("guidance_preset", "Padrão (Recomendado)")
        
        # Default LTX values are used if preset is 'Padrão'
        if preset == "Agressivo":
            overrides["guidance_scale"] = [1, 2, 8, 12, 8, 2, 1]
            overrides["stg_scale"] = [0, 0, 5, 6, 5, 3, 2]
        elif preset == "Suave":
            overrides["guidance_scale"] = [1, 1, 4, 5, 4, 1, 1]
            overrides["stg_scale"] = [0, 0, 2, 2, 2, 1, 0]
        elif preset == "Customizado":
            try:
                overrides["guidance_scale"] = json.loads(ltx_configs["guidance_scale_list"])
                overrides["stg_scale"] = json.loads(ltx_configs["stg_scale_list"])
            except (json.JSONDecodeError, KeyError