Update aduc_framework/managers/seedvr_manager.py

aduc_framework/managers/seedvr_manager.py

@@ -1,12 +1,22 @@
 # managers/seedvr_manager.py
 #
-# Version: 6.0.0 (GPU Isolation Fix) - FINAL
-# Este código já está pronto para gerenciar um pool em quaisquer GPUs
-# que sejam designadas a ele pelo hardware_manager.
+# Copyright (C) 2025 Carlos Rodrigues dos Santos
+#
+# Version: 6.0.0 (GPU Isolation Fix)
+#
+# Esta versão implementa a solução definitiva para os erros de dispositivo
+# usando a variável de ambiente CUDA_VISIBLE_DEVICES. Cada worker agora opera
+# em um ambiente completamente isolado, vendo apenas sua própria GPU.
+# Isso força todo o código de terceiros a usar o dispositivo correto e
+# elimina a necessidade de gerenciar manualmente o torch.distributed.
 
 import torch
-# ... (resto dos imports)
-import os, gc, logging, sys, subprocess, threading
+import os
+import gc
+import logging
+import sys
+import subprocess
+import threading
 from pathlib import Path
 from urllib.parse import urlparse
 from torch.hub import download_url_to_file
@@ -15,20 +25,25 @@ from einops import rearrange
 import shutil
 from omegaconf import OmegaConf
 import yaml
+
+# Imports relativos para o hardware_manager
 from ..tools.hardware_manager import hardware_manager
 
 logger = logging.getLogger(__name__)
 
+# --- Caminhos Globais ---
 APP_ROOT = Path("/home/user/app")
 DEPS_DIR = APP_ROOT / "deps"
 SEEDVR_SPACE_DIR = DEPS_DIR / "SeedVR_Space"
 SEEDVR_SPACE_URL = "https://huggingface.co/spaces/ByteDance-Seed/SeedVR2-3B"
 
 class SeedVrWorker:
+    """Representa uma única instância do pipeline SeedVR em um dispositivo isolado."""
     def __init__(self, device_id: str):
         self.global_device_id = device_id
         self.local_device_name = 'cuda:0'
         self.gpu_index = self.global_device_id.split(':')[-1]
+        
         self.runner = None
         self.is_initialized = False
         self.setup_complete = self._check_and_run_global_setup()
@@ -36,55 +51,79 @@ class SeedVrWorker:
 
     @staticmethod
     def _check_and_run_global_setup():
+        """Executa o setup de arquivos uma única vez para toda a aplicação."""
         setup_flag = DEPS_DIR / "seedvr.setup.complete"
-        if str(APP_ROOT) not in sys.path: sys.path.insert(0, str(APP_ROOT))
-        if setup_flag.exists(): return True
+        if str(APP_ROOT) not in sys.path:
+            sys.path.insert(0, str(APP_ROOT))
+            
+        if setup_flag.exists():
+            return True
+        
         logger.info("--- Iniciando Setup Global do SeedVR (primeira execução) ---")
+        
         if not SEEDVR_SPACE_DIR.exists():
             DEPS_DIR.mkdir(exist_ok=True, parents=True)
             subprocess.run(["git", "clone", "--depth", "1", SEEDVR_SPACE_URL, str(SEEDVR_SPACE_DIR)], check=True)
+        
         required_dirs = ["projects", "common", "models", "configs_3b", "configs_7b", "data"]
         for dirname in required_dirs:
-            source, target = SEEDVR_SPACE_DIR / dirname, APP_ROOT / dirname
-            if not target.exists(): shutil.copytree(source, target)
+            source = SEEDVR_SPACE_DIR / dirname
+            target = APP_ROOT / dirname
+            if not target.exists():
+                shutil.copytree(source, target)
+        
         try:
             import apex
         except ImportError:
             apex_url = 'https://huggingface.co/ByteDance-Seed/SeedVR2-3B/resolve/main/apex-0.1-cp310-cp310-linux_x86_64.whl'
             apex_wheel_path = _load_file_from_url(url=apex_url, model_dir=str(DEPS_DIR))
             subprocess.run(f"pip install {apex_wheel_path}", check=True, shell=True)
+
         ckpt_dir = APP_ROOT / 'ckpts'
         ckpt_dir.mkdir(exist_ok=True)
         model_urls = {
             'vae': 'https://huggingface.co/ByteDance-Seed/SeedVR2-3B/resolve/main/ema_vae.pth',
             'dit_3b': 'https://huggingface.co/ByteDance-Seed/SeedVR2-3B/resolve/main/seedvr2_ema_3b.pth',
-            'dit_7b': 'https://huggingface.co/ByteDance-Seed/SeedVR2-7B/resolve/main/seedvr2_ema_7b.pth',
+            #'dit_7b': 'https://huggingface.co/ByteDance-Seed/SeedVR2-7B/resolve/main/seedvr2_ema_7b.pth',
             'pos_emb': 'https://huggingface.co/ByteDance-Seed/SeedVR2-3B/resolve/main/pos_emb.pt',
             'neg_emb': 'https://huggingface.co/ByteDance-Seed/SeedVR2-3B/resolve/main/neg_emb.pt'
         }
         for name, url in model_urls.items():
             _load_file_from_url(url=url, model_dir=str(ckpt_dir))
+        
         setup_flag.touch()
         logger.info("--- Setup Global do SeedVR Concluído ---")
         return True
 
     def initialize_runner(self, model_version: str):
+        """Carrega os modelos para a VRAM do dispositivo, usando um ambiente de GPU isolado."""
         if self.runner is not None: return
+        
         os.environ['CUDA_VISIBLE_DEVICES'] = self.gpu_index
+        
         from projects.video_diffusion_sr.infer import VideoDiffusionInfer
         from common.config import load_config
+        
         logger.info(f"Worker {self.global_device_id}: Inicializando runner... (Processo vê apenas {self.local_device_name})")
-        config_path = APP_ROOT / (f'configs_{model_version.lower()}' / 'main.yaml')
-        checkpoint_path = APP_ROOT / 'ckpts' / f'seedvr2_ema_{model_version.lower()}.pth'
+        
+        config_path_str = f'configs_{model_version.lower()}'
+        checkpoint_path_str = f'seedvr2_ema_{model_version.lower()}.pth'
+
+        config_path = APP_ROOT / config_path_str / 'main.yaml'
+        checkpoint_path = APP_ROOT / 'ckpts' / checkpoint_path_str
+
         config = load_config(str(config_path))
         self.runner = VideoDiffusionInfer(config)
         OmegaConf.set_readonly(self.runner.config, False)
+        
         self.runner.configure_dit_model(device=self.local_device_name, checkpoint=str(checkpoint_path))
         self.runner.configure_vae_model()
+        
         self.is_initialized = True
         logger.info(f"Worker {self.global_device_id}: Runner pronto na VRAM.")
         
     def unload_runner(self):
+        """Descarrega os modelos da VRAM e limpa o ambiente."""
         if self.runner is not None:
             del self.runner
             self.runner = None
@@ -92,12 +131,15 @@ class SeedVrWorker:
             torch.cuda.empty_cache()
             self.is_initialized = False
             logger.info(f"Worker {self.global_device_id}: Runner descarregado da VRAM.")
+        
         if 'CUDA_VISIBLE_DEVICES' in os.environ:
             del os.environ['CUDA_VISIBLE_DEVICES']
 
     def process_video_internal(self, input_video_path, output_video_path, prompt, model_version, steps, seed):
+        """Executa a inferência em um ambiente de GPU isolado."""
         os.environ['CUDA_VISIBLE_DEVICES'] = self.gpu_index
         device = torch.device(self.local_device_name)
+
         from common.seed import set_seed
         from data.image.transforms.divisible_crop import DivisibleCrop
         from data.image.transforms.na_resize import NaResize
@@ -105,41 +147,59 @@ class SeedVrWorker:
         from projects.video_diffusion_sr.color_fix import wavelet_reconstruction
         from torchvision.transforms import Compose, Lambda, Normalize
         from torchvision.io.video import read_video
+
         set_seed(seed, same_across_ranks=True)
         self.runner.config.diffusion.timesteps.sampling.steps = steps
         self.runner.configure_diffusion()
+        
         video_tensor = read_video(input_video_path, output_format="TCHW")[0] / 255.0
         res_h, res_w = video_tensor.shape[-2:]
         video_transform = Compose([
             NaResize(resolution=(res_h * res_w) ** 0.5, mode="area", downsample_only=False),
             Lambda(lambda x: torch.clamp(x, 0.0, 1.0)),
-            DivisibleCrop((16, 16)), Normalize(0.5, 0.5), Rearrange("t c h w -> c t h w"),
+            DivisibleCrop((16, 16)),
+            Normalize(0.5, 0.5),
+            Rearrange("t c h w -> c t h w"),
         ])
         cond_latents = [video_transform(video_tensor.to(device))]
-        self.runner.dit.to("cpu"); self.runner.vae.to(device)
+        self.runner.dit.to("cpu")
+        self.runner.vae.to(device)
         cond_latents = self.runner.vae_encode(cond_latents)
-        self.runner.vae.to("cpu"); gc.collect(); torch.cuda.empty_cache(); self.runner.dit.to(device)
+        self.runner.vae.to("cpu"); gc.collect(); torch.cuda.empty_cache()
+        self.runner.dit.to(device)
+        
         pos_emb = torch.load(APP_ROOT / 'ckpts' / 'pos_emb.pt').to(device)
         neg_emb = torch.load(APP_ROOT / 'ckpts' / 'neg_emb.pt').to(device)
         text_embeds_dict = {"texts_pos": [pos_emb], "texts_neg": [neg_emb]}
+        
         noises = [torch.randn_like(latent) for latent in cond_latents]
         conditions = [self.runner.get_condition(noise, latent_blur=latent, task="sr") for noise, latent in zip(noises, cond_latents)]
+        
         with torch.no_grad(), torch.autocast("cuda", torch.bfloat16, enabled=True):
             video_tensors = self.runner.inference(noises=noises, conditions=conditions, dit_offload=True, **text_embeds_dict)
-        self.runner.dit.to("cpu"); gc.collect(); torch.cuda.empty_cache(); self.runner.vae.to(device)
+            
+        self.runner.dit.to("cpu"); gc.collect(); torch.cuda.empty_cache()
+        self.runner.vae.to(device)
         samples = self.runner.vae_decode(video_tensors)
-        final_sample, input_video_sample = samples[0], cond_latents[0]
+        final_sample = samples[0]
+        input_video_sample = cond_latents[0] # Usar o latente de condição como base
         if final_sample.shape[1] < input_video_sample.shape[1]:
             input_video_sample = input_video_sample[:, :final_sample.shape[1]]
+            
         final_sample = wavelet_reconstruction(rearrange(final_sample, "c t h w -> t c h w"), rearrange(input_video_sample, "c t h w -> t c h w"))
         final_sample = rearrange(final_sample, "t c h w -> t h w c")
         final_sample = final_sample.clip(-1, 1).mul_(0.5).add_(0.5).mul_(255).round()
         final_sample_np = final_sample.to(torch.uint8).cpu().numpy()
+        
         mediapy.write_video(output_video_path, final_sample_np, fps=24)
-        if 'CUDA_VISIBLE_DEVICES' in os.environ: del os.environ['CUDA_VISIBLE_DEVICES']
+        
+        if 'CUDA_VISIBLE_DEVICES' in os.environ:
+            del os.environ['CUDA_VISIBLE_DEVICES']
+            
         return output_video_path
 
 class SeedVrPoolManager:
+    """Gerencia um pool de SeedVrWorkers para processamento em GPUs dedicadas."""
     def __init__(self, device_ids: list[str]):
         logger.info(f"SEEDVR POOL MANAGER: Criando workers para os dispositivos: {device_ids}")
         if not device_ids or 'cpu' in device_ids:
@@ -150,25 +210,30 @@ class SeedVrPoolManager:
         self.last_cleanup_thread = None
 
     def _cleanup_worker_thread(self, worker: SeedVrWorker):
+        """Thread para descarregar o worker em segundo plano."""
         logger.info(f"SEEDVR CLEANUP THREAD: Iniciando limpeza de {worker.global_device_id} em background...")
         worker.unload_runner()
 
     def process_video(self, input_video_path: str, output_video_path: str, prompt: str,
-                      model_version: str = '7B', steps: int = 100, seed: int = 666, 
-                      progress: gr.Progress = None) -> str:
+                      model_version: str = '3B', steps: int = 100, seed: int = 666) -> str:
         worker_to_use = None
         try:
             with self.lock:
                 if self.last_cleanup_thread and self.last_cleanup_thread.is_alive():
                     self.last_cleanup_thread.join()
+                
                 worker_to_use = self.workers[self.current_worker_index]
                 previous_worker_index = (self.current_worker_index - 1 + len(self.workers)) % len(self.workers)
                 worker_to_cleanup = self.workers[previous_worker_index]
+
                 cleanup_thread = threading.Thread(target=self._cleanup_worker_thread, args=(worker_to_cleanup,))
                 cleanup_thread.start()
                 self.last_cleanup_thread = cleanup_thread
+                
                 worker_to_use.initialize_runner(model_version)
+                
                 self.current_worker_index = (self.current_worker_index + 1) % len(self.workers)
+            
             logger.info(f"SEEDVR POOL MANAGER: Processando vídeo na GPU {worker_to_use.global_device_id}...")
             return worker_to_use.process_video_internal(
                 input_video_path, output_video_path, prompt, model_version, steps, seed
@@ -185,13 +250,15 @@ def _load_file_from_url(url, model_dir='./', file_name=None):
         download_url_to_file(url, cached_file, hash_prefix=None, progress=True)
     return cached_file
 
+# --- Instanciação Singleton ---
 class SeedVrPlaceholder:
     def process_video(self, input_video_path, *args, **kwargs):
         logger.warning("SeedVR está desabilitado (gpus_required: 0). Pulando etapa de masterização HD.")
         return input_video_path
 
 try:
-    with open("config.yaml", 'r') as f: config = yaml.safe_load(f)
+    with open("config.yaml", 'r') as f:
+        config = yaml.safe_load(f)
     seedvr_gpus_required = config['specialists'].get('seedvr', {}).get('gpus_required', 0)
     seedvr_device_ids = hardware_manager.allocate_gpus('SeedVR', seedvr_gpus_required)
     if seedvr_gpus_required > 0 and 'cpu' not in seedvr_device_ids: