Delete seedvr_manager (2).py

seedvr_manager (2).py

@@ -1,233 +0,0 @@
-# managers/seedvr_manager.py
-#
-# Copyright (C) 2025 Carlos Rodrigues dos Santos
-#
-# Version: 4.0.0 (Root Installer & Executor)
-#
-# This version fully adopts the logic from the functional hd_specialist.py example.
-# It acts as a setup manager: it clones the SeedVR repo and then copies all
-# necessary directories (projects, common, models, configs, ckpts) to the
-# application root. It also handles the pip installation of the Apex dependency.
-# This ensures that the SeedVR code runs in the exact file structure it expects.
-
-import torch
-import torch.distributed as dist
-import os
-import gc
-import logging
-import sys
-import subprocess
-from pathlib import Path
-from urllib.parse import urlparse
-from torch.hub import download_url_to_file
-import gradio as gr
-import mediapy
-from einops import rearrange
-import shutil
-from omegaconf import OmegaConf
-
-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 SeedVrManager:
-    def __init__(self, workspace_dir="deformes_workspace"):
-        self.device = 'cuda' if torch.cuda.is_available() else 'cpu'
-        self.runner = None
-        self.workspace_dir = workspace_dir
-        self.is_initialized = False
-        self._original_barrier = None
-        self.setup_complete = False # Flag para rodar o setup apenas uma vez
-        logger.info("SeedVrManager initialized. Setup will run on first use.")
-
-    def _full_setup(self):
-        """
-        Executa todo o processo de setup uma única vez.
-        """
-        if self.setup_complete:
-            return
-        
-        logger.info("--- Starting Full SeedVR Setup ---")
-        
-        # 1. Clonar o repositório se não existir
-        if not SEEDVR_SPACE_DIR.exists():
-            logger.info(f"Cloning SeedVR Space repo to {SEEDVR_SPACE_DIR}...")
-            DEPS_DIR.mkdir(exist_ok=True, parents=True)
-            subprocess.run(
-                ["git", "clone", "--depth", "1", SEEDVR_SPACE_URL, str(SEEDVR_SPACE_DIR)],
-                check=True, capture_output=True, text=True
-            )
-        
-        # 2. Copiar as pastas necessárias para a raiz da aplicação
-        required_dirs = ["projects", "common", "models", "configs_3b", "configs_7b"]
-        for dirname in required_dirs:
-            source = SEEDVR_SPACE_DIR / dirname
-            target = APP_ROOT / dirname
-            if not target.exists():
-                logger.info(f"Copying '{dirname}' to application root...")
-                shutil.copytree(source, target)
-        
-        # 3. Adicionar a raiz ao sys.path para garantir que os imports funcionem
-        if str(APP_ROOT) not in sys.path:
-            sys.path.insert(0, str(APP_ROOT))
-            logger.info(f"Added '{APP_ROOT}' to sys.path.")
-
-        # 4. Instalar dependências complexas como Apex
-        try:
-            import apex
-            logger.info("Apex is already installed.")
-        except ImportError:
-            logger.info("Installing Apex dependency...")
-            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)
-            logger.info("Apex installed successfully.")
-
-        # 5. Baixar os modelos para a pasta ./ckpts na raiz
-        ckpt_dir = APP_ROOT / 'ckpts'
-        ckpt_dir.mkdir(exist_ok=True)
-        pretrain_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'
-        }
-        for name, url in pretrain_model_urls.items():
-            _load_file_from_url(url=url, model_dir=str(ckpt_dir))
-            
-        self.setup_complete = True
-        logger.info("--- Full SeedVR Setup Complete ---")
-
-    def _initialize_runner(self, model_version: str):
-        if self.runner is not None: return
-
-        # Garante que todo o ambiente está configurado antes de prosseguir
-        self._full_setup()
-
-        # Agora que o setup está feito, podemos importar os módulos
-        from projects.video_diffusion_sr.infer import VideoDiffusionInfer
-        from common.config import load_config
-        from common.seed import set_seed
-        
-        if dist.is_available() and not dist.is_initialized():
-            os.environ["MASTER_ADDR"] = "127.0.0.1"
-            os.environ["MASTER_PORT"] = "12355"
-            os.environ["RANK"] = str(0)
-            os.environ["WORLD_SIZE"] = str(1)
-            dist.init_process_group(backend='gloo')
-            logger.info("Initialized torch.distributed process group.")
-
-        logger.info(f"Initializing SeedVR2 {model_version} runner...")
-        if model_version == '3B':
-            config_path = APP_ROOT / 'configs_3b' / 'main.yaml'
-            checkpoint_path = APP_ROOT / 'ckpts' / 'seedvr2_ema_3b.pth'
-        else: # Assumimos 7B
-            config_path = APP_ROOT / 'configs_7b' / 'main.yaml'
-            checkpoint_path = APP_ROOT / 'ckpts' / 'seedvr2_ema_7b.pth'
-
-        config = load_config(str(config_path))
-        
-        self.runner = VideoDiffusionInfer(config)
-        OmegaConf.set_readonly(self.runner.config, False)
-        
-        self.runner.configure_dit_model(device=self.device, checkpoint=str(checkpoint_path))
-        self.runner.configure_vae_model()
-        
-        if hasattr(self.runner.vae, "set_memory_limit"):
-            self.runner.vae.set_memory_limit(**self.runner.config.vae.memory_limit)
-        
-        self.is_initialized = True
-        logger.info(f"Runner for SeedVR2 {model_version} initialized and ready.")
-        
-    def _unload_runner(self):
-        if self.runner is not None:
-            del self.runner
-            self.runner = None
-            gc.collect()
-            torch.cuda.empty_cache()
-            self.is_initialized = False
-            logger.info("Runner do SeedVR2 descarregado da VRAM.")
-        if dist.is_initialized():
-            dist.destroy_process_group()
-            logger.info("Destroyed torch.distributed process group.")
-
-    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:
-        try:
-            self._initialize_runner(model_version)
-            
-            # Precisamos importar aqui, pois o sys.path é modificado no setup
-            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(self.device))]
-            input_videos = cond_latents
-            self.runner.dit.to("cpu")
-            self.runner.vae.to(self.device)
-            cond_latents = self.runner.vae_encode(cond_latents)
-            self.runner.vae.to("cpu"); gc.collect(); torch.cuda.empty_cache()
-            self.runner.dit.to(self.device)
-            
-            pos_emb = torch.load(APP_ROOT / 'pos_emb.pt').to(self.device)
-            neg_emb = torch.load(APP_ROOT / 'neg_emb.pt').to(self.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(self.device)
-            samples = self.runner.vae_decode(video_tensors)
-            final_sample = samples[0]
-            input_video_sample = input_videos[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)
-            logger.info(f"HD Mastered video saved to: {output_video_path}")
-            return output_path
-        finally:
-            self._unload_runner()
-
-def _load_file_from_url(url, model_dir='./', file_name=None):
-    os.makedirs(model_dir, exist_ok=True)
-    filename = file_name or os.path.basename(urlparse(url).path)
-    cached_file = os.path.abspath(os.path.join(model_dir, filename))
-    if not os.path.exists(cached_file):
-        logger.info(f'Downloading: "{url}" to {cached_file}')
-        download_url_to_file(url, cached_file, hash_prefix=None, progress=True)
-    return cached_file
-
-seedvr_manager_singleton = SeedVrManager()