Spaces:
Paused
Paused
File size: 5,506 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 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 |
Hugging Face's logo
euIaxs22
/
Aduc-sdr-2_5
like
0
App
Files
Community
Aduc-sdr-2_5
/
services
/
seed_server.py
euIaxs22's picture
euIaxs22
Update services/seed_server.py
1cd2f56
verified
raw
Copy download link
history
blame
contribute
delete
5.25 kB
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
|