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