Delete ltx_manager_helpers.py

ltx_manager_helpers.py

@@ -1,240 +0,0 @@
-# ltx_manager_helpers.py
-# Copyright (C) 4 de Agosto de 2025  Carlos Rodrigues dos Santos
-#
-# ORIGINAL SOURCE: LTX-Video by Lightricks Ltd. & other open-source projects.
-# Licensed under the Apache License, Version 2.0
-# https://github.com/Lightricks/LTX-Video
-#
-# MODIFICATIONS FOR ADUC-SDR_Video:
-# This file is part of ADUC-SDR_Video, a derivative work based on LTX-Video.
-# It has been modified to manage pools of LTX workers, handle GPU memory,
-# and prepare parameters for the ADUC-SDR orchestration framework.
-# All modifications are also licensed under the Apache License, Version 2.0.
-
-import torch
-import gc
-import os
-import yaml
-import logging
-import huggingface_hub
-import time
-import threading
-import json
-
-from optimization import optimize_ltx_worker, can_optimize_fp8
-from hardware_manager import hardware_manager
-from inference import create_ltx_video_pipeline, calculate_padding
-from ltx_video.pipelines.pipeline_ltx_video import LatentConditioningItem
-
-logger = logging.getLogger(__name__)
-
-class LtxWorker:
-    """
-    Representa uma única instância da pipeline LTX-Video em um dispositivo específico.
-    Gerencia o carregamento do modelo para a CPU e a movimentação de/para a GPU.
-    """
-    def __init__(self, device_id, ltx_config_file):
-        self.cpu_device = torch.device('cpu')
-        self.device = torch.device(device_id if torch.cuda.is_available() else 'cpu')
-        logger.info(f"LTX Worker ({self.device}): Inicializando com config '{ltx_config_file}'...")
-        
-        with open(ltx_config_file, "r") as file:
-            self.config = yaml.safe_load(file)
-        
-        self.is_distilled = "distilled" in self.config.get("checkpoint_path", "")
-
-        models_dir = "downloaded_models_gradio"
-        
-        logger.info(f"LTX Worker ({self.device}): Carregando modelo para a CPU...")
-        model_path = os.path.join(models_dir, self.config["checkpoint_path"])
-        if not os.path.exists(model_path):
-             model_path = huggingface_hub.hf_hub_download(
-                repo_id="Lightricks/LTX-Video", filename=self.config["checkpoint_path"],
-                local_dir=models_dir, local_dir_use_symlinks=False
-            )
-        
-        self.pipeline = create_ltx_video_pipeline(
-            ckpt_path=model_path, precision=self.config["precision"],
-            text_encoder_model_name_or_path=self.config["text_encoder_model_name_or_path"],
-            sampler=self.config["sampler"], device='cpu'
-        )
-        logger.info(f"LTX Worker ({self.device}): Modelo pronto na CPU. É um modelo destilado? {self.is_distilled}")
-
-    def to_gpu(self):
-        """Move o pipeline para a GPU designada E OTIMIZA SE POSSÍVEL."""
-        if self.device.type == 'cpu': return
-        logger.info(f"LTX Worker: Movendo pipeline para a GPU {self.device}...")
-        self.pipeline.to(self.device)
-        
-        if self.device.type == 'cuda' and can_optimize_fp8():
-            logger.info(f"LTX Worker ({self.device}): GPU com suporte a FP8 detectada. Iniciando otimização...")
-            optimize_ltx_worker(self)
-            logger.info(f"LTX Worker ({self.device}): Otimização concluída.")
-        elif self.device.type == 'cuda':
-            logger.info(f"LTX Worker ({self.device}): Otimização FP8 não suportada ou desativada.")
-
-    def to_cpu(self):
-        """Move o pipeline de volta para a CPU e libera a memória da GPU."""
-        if self.device.type == 'cpu': return
-        logger.info(f"LTX Worker: Descarregando pipeline da GPU {self.device}...")
-        self.pipeline.to('cpu')
-        gc.collect()
-        if torch.cuda.is_available(): torch.cuda.empty_cache()
-
-    def generate_video_fragment_internal(self, **kwargs):
-        """Invoca a pipeline de geração."""
-        return self.pipeline(**kwargs).images
-
-class LtxPoolManager:
-    """
-    Gerencia um pool de LtxWorkers. MODO "HOT START": Mantém todos os modelos carregados na VRAM.
-    """
-    def __init__(self, device_ids, ltx_config_file):
-        logger.info(f"LTX POOL MANAGER: Criando workers para os dispositivos: {device_ids}")
-        self.workers = [LtxWorker(dev_id, ltx_config_file) for dev_id in device_ids]
-        self.current_worker_index = 0
-        self.lock = threading.Lock()
-
-        if all(w.device.type == 'cuda' for w in self.workers):
-            logger.info("LTX POOL MANAGER: MODO HOT START ATIVADO. Pré-aquecendo todas as GPUs...")
-            for worker in self.workers:
-                worker.to_gpu()
-            logger.info("LTX POOL MANAGER: Todas as GPUs estão quentes e prontas.")
-        else:
-            logger.info("LTX POOL MANAGER: Operando em modo CPU ou misto. O pré-aquecimento de GPU foi ignorado.")
-
-    def _prepare_and_log_params(self, worker_to_use, **kwargs):
-        target_device = worker_to_use.device
-        height, width = kwargs['height'], kwargs['width']
-        
-        conditioning_data = kwargs.get('conditioning_items_data', [])
-        final_conditioning_items = []
-        conditioning_log_details = []
-        for i, item in enumerate(conditioning_data):
-            if hasattr(item, 'latent_tensor'):
-                item.latent_tensor = item.latent_tensor.to(target_device)
-                final_conditioning_items.append(item)
-                conditioning_log_details.append(
-                    f"  - Item {i}: frame={item.media_frame_number}, strength={item.conditioning_strength:.2f}, shape={list(item.latent_tensor.shape)}"
-                )
-
-        first_pass_config = worker_to_use.config.get("first_pass", {})
-        
-        if 'latents' in kwargs and kwargs['latents'] is not None:
-            padded_h, padded_w = height, width
-            padding_vals = (0, 0, 0, 0)
-        else:
-            padded_h, padded_w = ((height - 1) // 32 + 1) * 32, ((width - 1) // 32 + 1) * 32
-            padding_vals = calculate_padding(height, width, padded_h, padded_w)
-
-        pipeline_params = {
-            "height": padded_h, "width": padded_w, 
-            "num_frames": kwargs['video_total_frames'], "frame_rate": kwargs['video_fps'],
-            "generator": torch.Generator(device=target_device).manual_seed(int(kwargs.get('seed', time.time())) + kwargs['current_fragment_index']),
-            "conditioning_items": final_conditioning_items, 
-            "is_video": True, "vae_per_channel_normalize": True,
-            "decode_timestep": float(kwargs.get('decode_timestep', worker_to_use.config.get("decode_timestep", 0.05))),
-            "image_cond_noise_scale": float(kwargs.get('image_cond_noise_scale', 0.0)),
-            "prompt": kwargs['motion_prompt'],
-            "negative_prompt": kwargs.get('negative_prompt', "blurry, distorted, static, bad quality, artifacts"),
-            "guidance_scale": float(kwargs.get('guidance_scale', 2.0)),
-            "stg_scale": float(kwargs.get('stg_scale', 0.025)),
-            "rescaling_scale": float(kwargs.get('rescaling_scale', 0.15)),
-        }
-        
-        if worker_to_use.is_distilled:
-            pipeline_params["timesteps"] = first_pass_config.get("timesteps")
-            pipeline_params["num_inference_steps"] = len(pipeline_params["timesteps"]) if "timesteps" in first_pass_config else 20
-        else:
-            pipeline_params["num_inference_steps"] = int(kwargs.get('num_inference_steps', 20))
-        
-        log_friendly_params = pipeline_params.copy()
-        log_friendly_params.pop('generator', None)
-        log_friendly_params.pop('conditioning_items', None)
-        
-        logger.info("="*60)
-        logger.info(f"CHAMADA AO PIPELINE LTX NO DISPOSITIVO: {worker_to_use.device}")
-        
-        return pipeline_params, padding_vals
-    
-    def _execute_on_worker(self, execution_fn, **kwargs):
-        worker_to_use = None
-        try:
-            with self.lock:
-                worker_to_use = self.workers[self.current_worker_index]
-                self.current_worker_index = (self.current_worker_index + 1) % len(self.workers)
-            
-            result, padding_vals = execution_fn(worker_to_use, **kwargs)
-            
-            return result, padding_vals
-            
-        except Exception as e:
-            logger.error(f"LTX POOL MANAGER: Erro durante a execução em {worker_to_use.device if worker_to_use else 'N/A'}: {e}", exc_info=True)
-            raise e
-        finally:
-            if worker_to_use and worker_to_use.device.type == 'cuda':
-                with torch.cuda.device(worker_to_use.device):
-                    gc.collect()
-                    torch.cuda.empty_cache()
-    
-    def generate_latent_fragment(self, **kwargs) -> (torch.Tensor, tuple):
-        def execution_logic(worker, **inner_kwargs):
-            pipeline_params, padding_vals = self._prepare_and_log_params(worker, **inner_kwargs)
-            pipeline_params['output_type'] = "latent"
-            with torch.no_grad():
-                result_tensor = worker.generate_video_fragment_internal(**pipeline_params)
-            return result_tensor, padding_vals
-
-        return self._execute_on_worker(execution_logic, **kwargs)
-
-    def refine_latents(self, upscaled_latents: torch.Tensor, **kwargs) -> (torch.Tensor, tuple):
-        def execution_logic(worker, **inner_kwargs):
-            pipeline_params, padding_vals = self._prepare_and_log_params(worker, **inner_kwargs)
-            
-            strength = inner_kwargs.get('denoise_strength', 0.4)
-            num_refine_steps_requested = int(inner_kwargs.get('refine_steps', 10))
-            
-            allowed_timesteps = worker.config.get("first_pass", {}).get("timesteps")
-            
-            if allowed_timesteps is None:
-                scheduler = worker.pipeline.scheduler
-                scheduler.set_timesteps(num_refine_steps_requested, device=worker.device)
-                timesteps = scheduler.timesteps
-            else:
-                timesteps = torch.tensor(allowed_timesteps, device=worker.device)
-
-            num_total_timesteps = len(timesteps)
-            start_timestep_idx = int(num_total_timesteps * strength)
-            if start_timestep_idx >= num_total_timesteps:
-                start_timestep_idx = num_total_timesteps - 1
-            
-            start_timestep = timesteps[start_timestep_idx]
-            
-            noise = torch.randn_like(upscaled_latents, device=worker.device)
-            noisy_latents = worker.pipeline.scheduler.add_noise(upscaled_latents.to(worker.device), noise, start_timestep)
-            
-            final_timesteps = timesteps[start_timestep_idx:]
-            pipeline_params['latents'] = noisy_latents.to(worker.device, dtype=worker.pipeline.transformer.dtype)
-            pipeline_params['timesteps'] = final_timesteps
-            pipeline_params['num_inference_steps'] = len(final_timesteps)
-            pipeline_params.pop('strength', None)
-            pipeline_params['output_type'] = "latent"
-            
-            logger.info(f"LTX POOL MANAGER: Iniciando refinamento com {len(final_timesteps)} passos a partir do timestep {start_timestep.item():.4f}.")
-            
-            with torch.no_grad():
-                refined_tensor = worker.generate_video_fragment_internal(**pipeline_params)
-            
-            return refined_tensor, padding_vals
-
-        return self._execute_on_worker(execution_logic, upscaled_latents=upscaled_latents, **kwargs)
-
-# --- Instanciação Singleton ---
-logger.info("Lendo config.yaml para inicializar o LTX Pool Manager...")
-with open("config.yaml", 'r') as f:
-    config = yaml.safe_load(f)
-ltx_gpus_required = config['specialists']['ltx']['gpus_required']
-ltx_device_ids = hardware_manager.allocate_gpus('LTX', ltx_gpus_required)
-ltx_config_path = config['specialists']['ltx']['config_file']
-ltx_manager_singleton = LtxPoolManager(device_ids=ltx_device_ids, ltx_config_file=ltx_config_path)
-logger.info("Especialista de Vídeo (LTX) pronto.")