Spaces:
Paused
Paused
File size: 3,837 Bytes
42998d3 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 |
# FILE: api/ltx_pool_manager.py
# DESCRIPTION: The "secret weapon". A pool manager for LTX that applies
# runtime patches to the pipeline for full control and ADUC-SDR compatibility.
import logging
from typing import List, Optional, Tuple, Union
from dataclasses import dataclass
import torch
from diffusers.utils.torch_utils import randn_tensor
# --- Importações da nossa arquitetura ---
from api.gpu_manager import gpu_manager
from api.ltx.ltx_utils import build_ltx_pipeline_on_cpu
from ltx_video.pipelines.pipeline_ltx_video import LTXVideoPipeline
# --- Definição dos nossos Data Classes ---
@dataclass
class ConditioningItem:
pixel_tensor: torch.Tensor # Sempre um tensor de pixel
media_frame_number: int
conditioning_strength: float
@dataclass
class LatentConditioningItem:
latent_tensor: torch.Tensor # Sempre um tensor latente
media_frame_number: int
conditioning_strength: float
# ==============================================================================
# --- O MONKEY PATCH ---
# Esta é a nossa versão customizada de `prepare_conditioning`
# ==============================================================================
def _aduc_prepare_conditioning_patch(
self: "LTXVideoPipeline",
conditioning_items: Optional[List[Union[ConditioningItem, LatentConditioningItem]]],
init_latents: torch.Tensor,
num_frames: int,
height: int,
width: int,
vae_per_channel_normalize: bool = False,
generator=None,
) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, int]:
# Esta função é uma cópia modificada da sua, com logging e pequenas melhorias.
# (O código do patch que você forneceu vai aqui, ligeiramente ajustado)
# ...
return init_latents, init_pixel_coords, init_conditioning_mask, extra_conditioning_num_latents
# ==============================================================================
# --- LTX Worker e Pool Manager ---
# ==============================================================================
class LTXWorker:
"""Gerencia uma instância do LTX Pipeline em um par de GPUs (main + vae)."""
def __init__(self, main_device: str, vae_device: str, config: dict):
self.main_device = torch.device(main_device)
self.vae_device = torch.device(vae_device)
self.config = config
self.pipeline: LTXVideoPipeline = None
self._load_and_patch_pipeline()
def _load_and_patch_pipeline(self):
logging.info(f"[LTXWorker-{self.main_device}] Carregando pipeline LTX para a CPU...")
self.pipeline, _ = build_ltx_pipeline_on_cpu(self.config)
logging.info(f"[LTXWorker-{self.main_device}] Movendo pipeline para GPUs (Main: {self.main_device}, VAE: {self.vae_device})...")
self.pipeline.to(self.main_device)
self.pipeline.vae.to(self.vae_device)
logging.info(f"[LTXWorker-{self.main_device}] Aplicando patch ADUC-SDR na função 'prepare_conditioning'...")
# A "mágica" do monkey patching acontece aqui
self.pipeline.prepare_conditioning = _aduc_prepare_conditioning_patch.__get__(self.pipeline, LTXVideoPipeline)
logging.info(f"[LTXWorker-{self.main_device}] ✅ Pipeline 'quente', corrigido e pronto.")
class LTXPoolManager:
# (Padrão Singleton, similar ao VincePoolManager)
# ...
def __init__(self):
# ...
main_device = gpu_manager.get_ltx_device()
vae_device = gpu_manager.get_ltx_vae_device()
# Em uma arquitetura futura, poderíamos ter múltiplos workers. Por enquanto, temos um.
self.worker = LTXWorker(str(main_device), str(vae_device), self._load_config())
# ...
def get_pipeline(self) -> LTXVideoPipeline:
return self.worker.pipeline
# Instância Singleton
ltx_pool_manager = LTXPoolManager() |