Update aduc_framework/managers/seedvr_manager.py

aduc_framework/managers/seedvr_manager.py

@@ -2,10 +2,14 @@
 #
 # Copyright (C) 2025 Carlos Rodrigues dos Santos
 #
-# Version: 6.1.0 (Force 3B Model)
+# Version: 10.0.0 (Definitive Monkey Patch)
 #
-# Esta versão remove a seleção de modelo e força o uso exclusivo do SeedVR 3B.
-# Também otimiza o setup para não baixar o checkpoint do modelo 7B.
+# Esta é a arquitetura final e mais robusta. O paralelismo problemático
+# é desativado programaticamente via "monkey patching" no decorador `master_only`.
+# Isso elimina a necessidade de gerenciar `torch.distributed`, simplificando
+# o código e resolvendo a causa raiz de todos os erros de paralelismo.
+# A isolação de GPU com CUDA_VISIBLE_DEVICES é mantida como a melhor
+# prática para o gerenciamento de hardware.
 
 import torch
 import os
@@ -13,7 +17,6 @@ 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
@@ -32,23 +35,21 @@ 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."""
+class SeedVrManager:
+    """Gerencia uma única instância do pipeline SeedVR em uma GPU dedicada e isolada."""
     def __init__(self, device_id: str):
         self.global_device_id = device_id
-        self.local_device_name = 'cuda:0'
+        self.local_device_name = 'cuda:0' # O que o processo enxergará
         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()
-        logger.info(f"SeedVR Worker inicializado para a GPU global {self.global_device_id}.")
+        self._check_and_run_global_setup()
+        logger.info(f"SeedVR Manager (Single Instance) inicializado para operar na GPU {self.global_device_id}.")
 
     @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 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) ---")
@@ -61,20 +62,38 @@ class SeedVrWorker:
             source, target = SEEDVR_SPACE_DIR / dirname, APP_ROOT / dirname
             if not target.exists(): shutil.copytree(source, target)
         
+        # <<< --- MONKEY PATCH PARA DESATIVAR O PARALELISMO --- >>>
+        try:
+            from common import decorators
+            import functools
+
+            # Define um decorador que não faz nada e apenas retorna a função original
+            def _passthrough_decorator(func):
+                @functools.wraps(func)
+                def wrapped(*args, **kwargs):
+                    return func(*args, **kwargs)
+                return wrapped
+            
+            # Substitui o decorador problemático pelo nosso decorador inofensivo
+            decorators.master_only = _passthrough_decorator
+            logger.info("Monkey patch aplicado com sucesso em 'common.decorators.master_only' para desativar o paralelismo.")
+        except Exception as e:
+            logger.error(f"Falha ao aplicar o monkey patch: {e}", exc_info=True)
+            # Continua mesmo se falhar, pode funcionar em alguns casos.
+        
         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_url = 'https://huggingface.co/ByteDance-Seed/SeedVR2-3B/resolve/main/apex-0.1-cp310-cp310-linux_x_86_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)
-        
-        # <<< MODIFICAÇÃO: Removido o download do modelo 7B >>>
         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',
             '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'
         }
@@ -83,10 +102,9 @@ class SeedVrWorker:
         
         setup_flag.touch()
         logger.info("--- Setup Global do SeedVR Concluído ---")
-        return True
 
-    def initialize_runner(self):
-        """Carrega o modelo 3B para a VRAM do dispositivo."""
+    def _initialize_runner(self):
+        """Carrega o modelo 3B em um ambiente de GPU isolado."""
         if self.runner is not None: return
         
         os.environ['CUDA_VISIBLE_DEVICES'] = self.gpu_index
@@ -94,9 +112,8 @@ class SeedVrWorker:
         from projects.video_diffusion_sr.infer import VideoDiffusionInfer
         from common.config import load_config
         
-        logger.info(f"Worker {self.global_device_id}: Inicializando runner SeedVR 3B... (Processo vê apenas {self.local_device_name})")
+        logger.info(f"Manager na GPU {self.global_device_id}: Inicializando runner SeedVR 3B...")
         
-        # <<< MODIFICAÇÃO: Caminhos fixados para o modelo 3B >>>
         config_path = APP_ROOT / 'configs_3b' / 'main.yaml'
         checkpoint_path = APP_ROOT / 'ckpts' / 'seedvr2_ema_3b.pth'
 
@@ -104,134 +121,80 @@ class SeedVrWorker:
         self.runner = VideoDiffusionInfer(config)
         OmegaConf.set_readonly(self.runner.config, False)
         
+        # Agora o código decorado dentro desta função usará nosso patch inofensivo
         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 3B pronto na VRAM.")
+        logger.info(f"Manager na GPU {self.global_device_id}: Runner 3B pronto na VRAM.")
         
-    def unload_runner(self):
+    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
-            gc.collect()
-            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, 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
-        from data.video.transforms.rearrange import Rearrange
-        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"),
-        ])
-        cond_latents = [video_transform(video_tensor.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)
-        
-        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)
-        samples = self.runner.vae_decode(video_tensors)
-        final_sample = samples[0]
-        input_video_sample = cond_latents[0]
-        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)
+            del self.runner; self.runner = None
+            gc.collect(); torch.cuda.empty_cache()
+            logger.info(f"Manager na GPU {self.global_device_id}: Runner descarregado da VRAM.")
         
         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:
-            raise ValueError("SeedVrPoolManager requer GPUs dedicadas.")
-        self.workers = [SeedVrWorker(device_id) for device_id in device_ids]
-        self.current_worker_index = 0
-        self.lock = threading.Lock()
-        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()
-
-    # <<< MODIFICAÇÃO: Removido o argumento 'model_version' da assinatura pública >>>
     def process_video(self, input_video_path: str, output_video_path: str, prompt: str,
                       steps: int = 100, seed: int = 666) -> str:
-        worker_to_use = None
+        """Ciclo completo de carga, processamento e descarga para uma única tarefa."""
         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
-                
-                # Chama initialize_runner sem argumentos, pois ele agora sabe que deve usar o 3B
-                worker_to_use.initialize_runner()
+            self._initialize_runner()
+            
+            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
+            from data.video.transforms.rearrange import Rearrange
+            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"),
+            ])
+            cond_latents = [video_transform(video_tensor.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)
+            
+            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.current_worker_index = (self.current_worker_index + 1) % len(self.workers)
+            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]
+            if final_sample.shape[1] < input_video_sample.shape[1]:
+                input_video_sample = input_video_sample[:, :final_sample.shape[1]]
             
-            logger.info(f"SEEDVR POOL MANAGER: Processando vídeo na GPU {worker_to_use.global_device_id}...")
-            # Passa os argumentos para a função interna, sem 'model_version'
-            return worker_to_use.process_video_internal(
-                input_video_path, output_video_path, prompt, steps, seed
-            )
-        except Exception as e:
-            logger.error(f"SEEDVR POOL MANAGER: Erro durante o processamento de vídeo: {e}", exc_info=True)
-            raise e
+            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)
+            return output_video_path
+        finally:
+            self._unload_runner()
+
 
 def _load_file_from_url(url, model_dir='./', file_name=None):
     os.makedirs(model_dir, exist_ok=True)
@@ -248,16 +211,20 @@ class SeedVrPlaceholder:
         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:
-        seedvr_manager_singleton = SeedVrPoolManager(device_ids=seedvr_device_ids)
-        logger.info("Especialista de Masterização HD (SeedVR Pool) pronto.")
+    
+    if seedvr_gpus_required > 0:
+        seedvr_device_ids = hardware_manager.allocate_gpus('SeedVR', seedvr_gpus_required)
+        if 'cpu' not in seedvr_device_ids:
+            device_to_use = seedvr_device_ids[0]
+            seedvr_manager_singleton = SeedVrManager(device_id=device_to_use)
+            logger.info(f"Especialista de Masterização HD (SeedVR Single Instance) pronto para usar a GPU {device_to_use}.")
+        else:
+            seedvr_manager_singleton = SeedVrPlaceholder()
     else:
         seedvr_manager_singleton = SeedVrPlaceholder()
-        logger.warning("SeedVR Pool Manager não foi inicializado.")
+        logger.warning("SeedVR Manager não foi inicializado (gpus_required: 0).")
 except Exception as e:
     logger.critical(f"Falha CRÍTICA ao inicializar o SeedVrManager: {e}", exc_info=True)
     seedvr_manager_singleton = SeedVrPlaceholder()