File size: 5,247 Bytes
a5720bf |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 |
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
|