Spaces:
Paused
Paused
| import os | |
| import shutil | |
| import subprocess | |
| import sys | |
| import time | |
| import mimetypes | |
| from pathlib import Path | |
| from typing import List, Optional, Tuple | |
| from huggingface_hub import hf_hub_download | |
| class SeedVRServer: | |
| def __init__(self, **kwargs): | |
| self.SEEDVR_ROOT = Path(os.getenv("SEEDVR_ROOT", "/data/SeedVR")) | |
| # Apontamos para o nosso diretório de checkpoints customizado | |
| self.CKPTS_ROOT = Path("/data/seedvr_models_fp16") | |
| self.OUTPUT_ROOT = Path(os.getenv("OUTPUT_ROOT", "/app/outputs")) | |
| self.INPUT_ROOT = Path(os.getenv("INPUT_ROOT", "/app/inputs")) | |
| self.HF_HOME_CACHE = Path(os.getenv("HF_HOME", "/data/.cache/huggingface")) | |
| self.REPO_URL = os.getenv("SEEDVR_GIT_URL", "https://github.com/numz/ComfyUI-SeedVR2_VideoUpscaler") | |
| self.NUM_GPUS_TOTAL = int(os.getenv("NUM_GPUS", "4")) | |
| print("🚀 SeedVRServer (FP16) inicializando e preparando o ambiente...") | |
| for p in [self.SEEDVR_ROOT.parent, self.CKPTS_ROOT, self.OUTPUT_ROOT, self.INPUT_ROOT, self.HF_HOME_CACHE]: | |
| p.mkdir(parents=True, exist_ok=True) | |
| self.setup_dependencies() | |
| print("✅ SeedVRServer (FP16) pronto.") | |
| def setup_dependencies(self): | |
| self._ensure_repo() | |
| # O monkey patch agora é feito pelo start_seedvr.sh, não mais aqui. | |
| self._ensure_model() | |
| def _ensure_repo(self) -> None: | |
| if not (self.SEEDVR_ROOT / ".git").exists(): | |
| print(f"[SeedVRServer] Clonando repositório para {self.SEEDVR_ROOT}...") | |
| subprocess.run(["git", "clone", "--depth", "1", self.REPO_URL, str(self.SEEDVR_ROOT)], check=True) | |
| else: | |
| print("[SeedVRServer] Repositório SeedVR já existe.") | |
| def _ensure_model(self) -> None: | |
| """Baixa os arquivos de modelo FP16 otimizados e suas dependências.""" | |
| print(f"[SeedVRServer] Verificando checkpoints (FP16) em {self.CKPTS_ROOT}...") | |
| model_files = { | |
| "seedvr2_ema_3b_fp16.safetensors": "MonsterMMORPG/SeedVR2_SECourses", "ema_vae_fp16.safetensors": "MonsterMMORPG/SeedVR2_SECourses", | |
| "pos_emb.pt": "ByteDance-Seed/SeedVR2-3B", "neg_emb.pt": "ByteDance-Seed/SeedVR2-3B" | |
| } | |
| for filename, repo_id in model_files.items(): | |
| if not (self.CKPTS_ROOT / filename).exists(): | |
| print(f"Baixando {filename} de {repo_id}...") | |
| hf_hub_download(repo_id=repo_id, filename=filename, local_dir=str(self.CKPTS_ROOT), cache_dir=str(self.HF_HOME_CACHE), token=os.getenv("HF_TOKEN")) | |
| print("[SeedVRServer] Checkpoints (FP16) estão no local correto.") | |
| def _prepare_job(self, input_file: str) -> Tuple[Path, Path]: | |
| ts = f"{int(time.time())}_{os.urandom(4).hex()}" | |
| job_input_dir = self.INPUT_ROOT / f"job_{ts}" | |
| out_dir = self.OUTPUT_ROOT / f"run_{ts}" | |
| job_input_dir.mkdir(parents=True, exist_ok=True) | |
| out_dir.mkdir(parents=True, exist_ok=True) | |
| shutil.copy2(input_file, job_input_dir / Path(input_file).name) | |
| return job_input_dir, out_dir | |
| def run_inference(self, filepath: str, *, seed: int, resh: int, resw: int, spsize: int, fps: Optional[float] = None): | |
| script = self.SEEDVR_ROOT / "inference_cli.py" | |
| job_input_dir, outdir = self._prepare_job(filepath) | |
| mediatype, _ = mimetypes.guess_type(filepath) | |
| is_image = mediatype and mediatype.startswith("image") | |
| effective_nproc = 1 if is_image else self.NUM_GPUS_TOTAL | |
| effective_spsize = 1 if is_image else spsize | |
| output_filename = f"result_{Path(filepath).stem}.mp4" if not is_image else f"{Path(filepath).stem}_upscaled" | |
| output_filepath = outdir / output_filename | |
| cmd = [ | |
| "torchrun", "--standalone", "--nnodes=1", | |
| f"--nproc-per-node={effective_nproc}", | |
| str(script), | |
| "--video_path", str(filepath), | |
| "--output", str(output_filepath), | |
| "--model_dir", str(self.CKPTS_ROOT), | |
| "--seed", str(seed), | |
| "--cuda_device", "0", | |
| "--resolution", str(resh), | |
| "--batch_size", str(effective_spsize), | |
| "--model", "seedvr2_ema_3b_fp16.safetensors", | |
| "--preserve_vram", | |
| "--debug", | |
| "--output_format", "video" if not is_image else "png", | |
| ] | |
| print("SeedVRServer Comando:", " ".join(cmd)) | |
| try: | |
| subprocess.run(cmd, cwd=str(self.SEEDVR_ROOT), check=True, env=os.environ.copy(), stdout=sys.stdout, stderr=sys.stderr) | |
| # Constrói a tupla de retorno de forma determinística | |
| if is_image: | |
| # CLI salva PNGs em diretório args.output (tratado como diretório quando outputformat=png) | |
| image_dir = output_filepath if output_filepath.suffix == "" else output_filepath.with_suffix("") | |
| return str(image_dir), None, outdir | |
| else: | |
| # CLI salva vídeo exatamente em output_filepath | |
| return None, str(output_filepath), outdir | |
| except Exception as e: | |
| print(f"[UI ERROR] A inferência falhou: {e}") | |
| return None, None, None | |