Rename seed_server.py to api/seedvr_server.py

api/seedvr_server.py

@@ -0,0 +1,144 @@
+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

seed_server.py

@@ -1,187 +0,0 @@
-#!/usr/bin/env python3
-"""
-SeedVR Server (CLI torchrun)
-
-- Garante repositório SeedVR e checkpoints baixados via snapshot_download.
-- Cria symlink SeedVR/ckpts/SeedVR2-3B -> CKPTS_ROOT.
-- Executa projects/inference_seedvr2_3b.py com torchrun e NUM_GPUS.
-- API: run_inference(file_path, seed, res_h, res_w, sp_size) -> (video_out, image_out, out_dir).
-"""
-
-import os
-import shutil
-import subprocess
-from pathlib import Path
-from typing import Optional, Tuple, List
-import time
-import mimetypes
-
-from huggingface_hub import snapshot_download  # requerido no container
-
-class SeedVRServer:
-    def __init__(
-        self,
-        *,
-        seedvr_root: Optional[str] = None,
-        ckpts_root: Optional[str] = None,
-        output_root: Optional[str] = None,
-        input_root: Optional[str] = None,
-        repo_url: Optional[str] = None,
-        repo_id: Optional[str] = None,
-        num_gpus: Optional[int] = None,
-    ):
-        # Paths e envs
-        self.SEEDVR_ROOT = Path(seedvr_root or os.getenv("SEEDVR_ROOT", "/app/SeedVR"))
-        self.CKPTS_ROOT = Path(ckpts_root or os.getenv("CKPTS_ROOT", "/app/ckpts/SeedVR2-3B"))
-        self.OUTPUT_ROOT = Path(output_root or os.getenv("OUTPUT_ROOT", "/app/outputs"))
-        self.INPUT_ROOT = Path(input_root or os.getenv("INPUT_ROOT", "/app/inputs"))
-        self.REPO_URL = repo_url or os.getenv("SEEDVR_GIT_URL", "https://github.com/ByteDance-Seed/SeedVR.git")
-        self.REPO_ID = repo_id or os.getenv("SEEDVR_REPO_ID", "ByteDance-Seed/SeedVR2-3B")
-        self.NUM_GPUS = int(num_gpus or os.getenv("NUM_GPUS", "8"))
-        self.HF_HOME = Path(os.getenv("HF_HOME", "/data/.cache/huggingface"))
-        self.HF_TOKEN = os.getenv("HF_TOKEN") or os.getenv("HUGGINGFACE_TOKEN") or None
-
-        # Diretórios necessários
-        for p in [self.SEEDVR_ROOT, self.CKPTS_ROOT, self.OUTPUT_ROOT, self.INPUT_ROOT, self.HF_HOME]:
-            p.mkdir(parents=True, exist_ok=True)
-
-        # Bootstrap direto
-        self._ensure_repo()
-        self._ensure_model()
-        self._ensure_ckpt_symlink()
-
-    # ---------- Preparação ----------
-    def _ensure_repo(self) -> None:
-        if not (self.SEEDVR_ROOT / ".git").exists():
-            print(f"[seed_server] cloning repo into {self.SEEDVR_ROOT}")
-            subprocess.run(["git", "clone", self.REPO_URL, str(self.SEEDVR_ROOT)], check=True)
-        else:
-            print(f"[seed_server] repo present at {self.SEEDVR_ROOT}")
-
-    def _ensure_model(self) -> None:
-        print(f"[seed_server] downloading model {self.REPO_ID} into {self.CKPTS_ROOT} (snapshot_download)")
-        self.CKPTS_ROOT.mkdir(parents=True, exist_ok=True)
-        snapshot_download(
-            repo_id=self.REPO_ID,
-            cache_dir=str(self.HF_HOME),
-            local_dir=str(self.CKPTS_ROOT),
-            local_dir_use_symlinks=False,
-            resume_download=True,
-            allow_patterns=["*.json", "*.safetensors", "*.pth", "*.bin", "*.py", "*.md", "*.txt"],
-            token=self.HF_TOKEN,
-        )
-        print("[seed_server] model ready")
-
-    def _ensure_ckpt_symlink(self) -> None:
-        ckpts_repo_dir = self.SEEDVR_ROOT / "ckpts"
-        ckpts_repo_dir.mkdir(parents=True, exist_ok=True)
-        link = ckpts_repo_dir / "SeedVR2-3B"
-        try:
-            if link.is_symlink():
-                try:
-                    if link.resolve() != self.CKPTS_ROOT:
-                        link.unlink()
-                except Exception:
-                    link.unlink(missing_ok=True)
-            if not link.exists():
-                link.symlink_to(self.CKPTS_ROOT, target_is_directory=True)
-            print(f"[seed_server] symlink ok: {link} -> {self.CKPTS_ROOT}")
-        except Exception as e:
-            print("[seed_server] warn: ckpt symlink failed:", e)
-
-    # ---------- Util ----------
-    @staticmethod
-    def _is_video(path: str) -> bool:
-        mime, _ = mimetypes.guess_type(path)
-        return (mime or "").startswith("video") or str(path).lower().endswith(".mp4")
-
-    @staticmethod
-    def _is_image(path: str) -> bool:
-        mime, _ = mimetypes.guess_type(path)
-        if mime and mime.startswith("image"):
-            return True
-        return str(path).lower().endswith((".png", ".jpg", ".jpeg", ".webp"))
-
-    def _prepare_job(self, input_file: str) -> Tuple[Path, Path]:
-        ts = int(time.time())
-        job_dir = self.INPUT_ROOT / f"job_{ts}"
-        out_dir = self.OUTPUT_ROOT / f"run_{ts}"
-        job_dir.mkdir(parents=True, exist_ok=True)
-        out_dir.mkdir(parents=True, exist_ok=True)
-
-        #####
-        shutil.copy2(input_file, job_dir / Path(input_file).name)
-        return out_dir, out_dir
-
-    # ---------- Execução ----------
-    def run_inference(
-        self,
-        file_path: str,
-        *,
-        seed: int = 42,
-        res_h: int = 720,
-        res_w: int = 1280,
-        sp_size: int = 4,
-        extra_args: Optional[List[str]] = None,
-    ) -> Tuple[Optional[str], Optional[str], Path]:
-        """
-        Executa inferência via torchrun com NUM_GPUS:
-        - file_path: vídeo .mp4 ou imagem .png/.jpg/.jpeg/.webp
-        - Retorna (video_out, image_out, out_dir). Um dos dois primeiros será não-nulo.
-        """
-        if not Path(file_path).exists():
-            raise FileNotFoundError(f"input not found: {file_path}")
-
-        script = self.SEEDVR_ROOT / "projects" / "inference_seedvr2_3b.py"
-        if not script.exists():
-            raise FileNotFoundError(f"inference script not found: {script}")
-
-        job_dir, out_dir = self._prepare_job(file_path)
-        self._ensure_ckpt_symlink()
-
-        out_dir.mkdir(parents=True, exist_ok=True)
-        os.chmod(out_dir, 777)
-
-        job_dir.mkdir(parents=True, exist_ok=True)
-        os.chmod(job_dir, 777)
-
-        cmd = [
-            "torchrun",
-            f"--nproc-per-node={self.NUM_GPUS}",
-            str(script),
-            "--video_path", str(job_dir),
-            "--output_dir", str(out_dir),
-            "--seed", str(seed),
-            "--res_h", str(res_h),
-            "--res_w", str(res_w),
-            "--sp_size", str(sp_size),
-        ]
-        if extra_args:
-            cmd.extend(extra_args)
-
-        env = os.environ.copy()
-        env.setdefault("HF_HOME", str(self.HF_HOME))
-        env.setdefault("NCCL_P2P_LEVEL", os.getenv("NCCL_P2P_LEVEL", "NVL"))
-        #env.setdefault("NCCL_ASYNC_ERROR_HANDLING", os.getenv("NCCL_ASYNC_ERROR_HANDLING", "1"))
-        env.setdefault("OMP_NUM_THREADS", os.getenv("OMP_NUM_THREADS", "8"))
-
-        print("[seed_server] running:", " ".join(cmd))
-        try:
-            subprocess.run(cmd, cwd=str(self.SEEDVR_ROOT), check=True, env=env)
-        except subprocess.CalledProcessError as e:
-            print("[seed_server] torchrun error:", e)
-            return None, None, out_dir
-
-        # Buscar artefatos
-        videos = sorted(out_dir.rglob("*.mp4"), key=lambda p: p.stat().st_mtime)
-        # Cobrir formatos comuns caso upstream mude
-        if not videos:
-            videos = sorted([*out_dir.rglob("*.mov"), *out_dir.rglob("*.avi")], key=lambda p: p.stat().st_mtime)
-        images = sorted(
-            [*out_dir.rglob("*.png"), *out_dir.rglob("*.jpg"), *out_dir.rglob("*.jpeg"), *out_dir.rglob("*.webp")],
-            key=lambda p: p.stat().st_mtime
-        )
-
-        video_out = str(videos[-1]) if videos else None
-        image_out = str(images[-1]) if images else None
-        return video_out, image_out, out_dir