Update services/vincie.py

services/vincie.py

@@ -1,7 +1,20 @@
 #!/usr/bin/env python3
 """
-VincieService (CLI runner, singleton-friendly)
+services/vincie.py
+
+VincieService — preparação e execução CLI do VINCIE (upstream)
+- Garante repositório íntegro (clona/repara se faltarem main.py/.git).
+- Baixa snapshot completo do modelo no HF_HUB_CACHE.
+- Cria symlink idempotente ckpt/VINCIE-3B (no repo e em /app/ckpt) apontando para o snapshot (contém dit.pth, vae.pth, llm14b).
+- Valida artefatos esperados pelo generate.yaml.
+- Executa main.py do upstream com overrides de geração (sem mexer em ckpt.path).
+- Limpa VRAM levemente após cada job.
+
+Observação:
+- Para latência mínima, preferir o vince_server in-process (pipeline aquecida).
+- Este serviço via subprocess é fiel ao upstream e útil como fallback/diag.
 """
+
 import os
 import json
 import subprocess
@@ -16,128 +29,93 @@ class VincieService:
         self,
         repo_dir: str = "/app/VINCIE",
         python_bin: str = "python",
-        repo_id: str = "ByteDance-Seed/VINCIE-3B",
+        repo_url: str = "https://github.com/ByteDance-Seed/VINCIE",
+        model_repo: str = "ByteDance-Seed/VINCIE-3B",
         output_root: str = "/app/outputs",
     ):
         self.repo_dir = Path(repo_dir)
         self.python = python_bin
-        self.repo_id = repo_id
-        self.generate_yaml = self.repo_dir / "configs" / "generate.yaml"
+        self.repo_url = repo_url
+        self.model_repo = model_repo
         self.output_root = Path(output_root)
         self.output_root.mkdir(parents=True, exist_ok=True)
+
+        self.generate_yaml = self.repo_dir / "configs" / "generate.yaml"
+        self.ckpt_link_repo = self.repo_dir / "ckpt" / "VINCIE-3B"
+        self.ckpt_link_app = Path("/app/ckpt") / "VINCIE-3B"
+
         self.ckpt_dir: Optional[Path] = None
         self._env = os.environ.copy()
-        (self.repo_dir / "ckpt").mkdir(parents=True, exist_ok=True)
 
-    # ---------- Repositório e modelo ----------
+    # ---------- util ----------
+
+    @staticmethod
+    def _run(cmd: List[str], cwd: Optional[Path] = None, env=None):
+        subprocess.run(cmd, cwd=str(cwd) if cwd else None, check=True, env=env)
 
-    def ensure_repo(self, git_url: str = "https://github.com/ByteDance-Seed/VINCIE") -> None:
-        """
-        Garante repositório íntegro:
-        - Se válido (.git e main.py), mantém.
-        - Se quebrado/ausente, clona em diretório temporário e substitui atomicamente,
-          recriando ckpt/ e o symlink conforme necessário.
-        """
+    @staticmethod
+    def _ensure_symlink(link: Path, target: Path):
+        link.parent.mkdir(parents=True, exist_ok=True)
+        if link.is_symlink():
+            try:
+                if link.resolve() != target:
+                    link.unlink()
+                    link.symlink_to(target, target_is_directory=True)
+            except Exception:
+                # relinka a partir do zero
+                link.unlink(missing_ok=True)
+                link.symlink_to(target, target_is_directory=True)
+        elif link.exists():
+            VincieService._run(["rm", "-rf", str(link)])
+            link.symlink_to(target, target_is_directory=True)
+        else:
+            link.symlink_to(target, target_is_directory=True)
+
+    # ---------- repo/modelo ----------
+
+    def ensure_repo(self) -> None:
         self.repo_dir.mkdir(parents=True, exist_ok=True)
         main_py = self.repo_dir / "main.py"
         git_dir = self.repo_dir / ".git"
-
-        repo_ok = main_py.exists() and git_dir.exists()
-        if repo_ok:
+        if main_py.exists() and git_dir.exists():
             return
-
-        # Preserva alvo do symlink ckpt/VINCIE-3B, se houver
-        ckpt_link = self.repo_dir / "ckpt" / "VINCIE-3B"
-        preserved_target = None
-        if ckpt_link.is_symlink():
-            try:
-                preserved_target = ckpt_link.resolve()
-            except Exception:
-                preserved_target = None
-
-        # Clona em diretório temporário
-        tmp_dir = self.repo_dir.with_name(self.repo_dir.name + ".tmp")
-        try:
-            if tmp_dir.exists():
-                subprocess.run(["rm", "-rf", str(tmp_dir)], check=True)
-            subprocess.run(["git", "clone", git_url, str(tmp_dir)], check=True)
-        except subprocess.CalledProcessError as e:
-            print("Error: git clone failed:", e)
-            raise
-
-        # Garante ckpt/ no temp e relinka symlink se possível
-        (tmp_dir / "ckpt").mkdir(parents=True, exist_ok=True)
-        if preserved_target and preserved_target.exists():
-            try:
-                tmp_link = tmp_dir / "ckpt" / "VINCIE-3B"
-                if tmp_link.exists() or tmp_link.is_symlink():
-                    tmp_link.unlink()
-                tmp_link.symlink_to(preserved_target, target_is_directory=True)
-            except Exception as e:
-                print("Warning: failed to link ckpt in temp repo:", e)
-
-        # Remove repo antigo e move temp para destino
-        try:
-            subprocess.run(["rm", "-rf", str(self.repo_dir)], check=True)
-            tmp_dir.rename(self.repo_dir)
-        except Exception as e:
-            print("Error: failed to swap repo directories:", e)
-            raise
-
-    def ensure_model(self, hf_token: Optional[str] = None, revision: Optional[str] = None) -> None:
-        """
-        Baixa snapshot completo no cache (honra HF_HUB_CACHE) e cria symlink no repo.
-        """
-        token = hf_token or os.getenv("HF_TOKEN") or os.getenv("HUGGINGFACE_TOKEN")
+        tmp = self.repo_dir.with_name(self.repo_dir.name + ".tmp")
+        if tmp.exists():
+            self._run(["rm", "-rf", str(tmp)])
+        self._run(["git", "clone", self.repo_url, str(tmp)])
+        # swap atômico simples
+        if self.repo_dir.exists():
+            self._run(["rm", "-rf", str(self.repo_dir)])
+        tmp.rename(self.repo_dir)
+
+    def ensure_model(self, revision: Optional[str] = None, token: Optional[str] = None) -> None:
         cache_dir = os.environ.get("HF_HUB_CACHE")
         snapshot_path = snapshot_download(
-            repo_id=self.repo_id,
+            repo_id=self.model_repo,
             revision=revision,
+            token=token or os.getenv("HF_TOKEN") or os.getenv("HUGGINGFACE_TOKEN"),
             cache_dir=cache_dir,
-            token=token,
-            local_files_only=False,
+            resume_download=True,
         )
         self.ckpt_dir = Path(snapshot_path)
-
-        link = self.repo_dir / "ckpt" / "VINCIE-3B"
-        try:
-            if link.is_symlink() or link.exists():
-                try:
-                    link.unlink()
-                except Exception:
-                    pass
-            link.symlink_to(self.ckpt_dir, target_is_directory=True)
-        except Exception as e:
-            print("Warning: failed to create checkpoint symlink:", e)
-
-    def ready(self) -> bool:
-        have_repo = (self.repo_dir / "main.py").exists() and (self.repo_dir / "configs" / "generate.yaml").exists()
-        ckpt_ok = False
-        if getattr(self, "ckpt_dir", None):
-            ckpt_ok = (self.ckpt_dir / "dit.pth").exists() and (self.ckpt_dir / "vae.pth").exists()
-        if not ckpt_ok:
-            link = self.repo_dir / "ckpt" / "VINCIE-3B"
-            if link.exists():
-                try:
-                    tgt = link.resolve()
-                    ckpt_ok = (tgt / "dit.pth").exists() and (tgt / "vae.pth").exists()
-                except Exception:
-                    ckpt_ok = False
-        return bool(have_repo and ckpt_ok)
-
-    # ---------- GPUs dedicadas e persistência ----------
-
-    def pin_gpus(self, device_indices: List[int]) -> None:
-        visible = ",".join(str(i) for i in device_indices)
-        self._env["CUDA_VISIBLE_DEVICES"] = visible
-
-    def enable_persistence_mode(self) -> None:
-        try:
-            subprocess.run(["nvidia-smi", "-pm", "1"], check=True)
-        except Exception as e:
-            print("Warning: failed to enable persistence mode:", e)
-
-    # ---------- Execução ----------
+        # symlinks idempotentes
+        self._ensure_symlink(self.ckpt_link_repo, self.ckpt_dir)
+        self._ensure_symlink(self.ckpt_link_app, self.ckpt_dir)
+
+    def validate_assets(self) -> None:
+        # exige generate.yaml/main.py e conteúdo essencial no snapshot
+        if not self.generate_yaml.exists() or not (self.repo_dir / "main.py").exists():
+            raise RuntimeError("VINCIE repo inválido (faltando generate.yaml ou main.py)")
+        target = self.ckpt_dir or self.ckpt_link_repo
+        need = [target / "dit.pth", target / "vae.pth", target / "llm14b"]
+        missing = [str(p) for p in need if not p.exists()]
+        if missing:
+            raise RuntimeError(f"Snapshot incompleto: {missing}")
+        # também requer que o link repo exista (a config usa ckpt/ relativo)
+        if not self.ckpt_link_repo.exists():
+            raise RuntimeError("ckpt link ausente no repo: ckpt/VINCIE-3B")
+
+    # ---------- execução ----------
 
     def _build_overrides(
         self,
@@ -148,8 +126,7 @@ class VincieService:
         steps: Optional[int] = None,
     ) -> List[str]:
         overrides = list(extra_overrides or [])
-        if self.ckpt_dir is not None:
-            overrides.append(f"ckpt.path={str(self.ckpt_dir)}")
+        # não altera ckpt.path; respeita o YAML
         if cfg_scale is not None:
             overrides.append(f"generation.cfg_scale={cfg_scale}")
         if resolution_input is not None:
@@ -160,20 +137,8 @@ class VincieService:
             overrides.append(f"generation.steps={steps}")
         return overrides
 
-    def _run_vincie_once(self, overrides: List[str], work_output: Path) -> None:
-        work_output.mkdir(parents=True, exist_ok=True)
-        cmd = [
-            self.python,
-            "main.py",
-            str(self.generate_yaml),
-            *overrides,
-            f"generation.output.dir={str(work_output)}",
-        ]
-        subprocess.run(cmd, cwd=self.repo_dir, check=True, env=self._env)
-
     def _clean_gpu_memory(self) -> None:
-        try:
-            code = r"""
+        code = r"""
 import torch, gc
 try:
     torch.cuda.synchronize()
@@ -186,9 +151,7 @@ try:
 except Exception:
     pass
 """
-            subprocess.run([self.python, "-c", code], check=True, env=self._env)
-        except Exception as e:
-            print("Warning: GPU cleanup failed:", e)
+        self._run([self.python, "-c", code], env=self._env)
 
     # ---------- APIs ----------
 
@@ -202,17 +165,21 @@ except Exception:
         resolution_input: Optional[int] = None,
         aspect_ratio_input: Optional[str] = None,
         steps: Optional[int] = None,
-        pad_img_placeholder: Optional[bool] = None,
-    ) -> Path:
-        out_dir = self.output_root / (out_dir_name or f"multi_turn_{self._slug(input_image)}")
+    ) -> str:
+        self.ensure_repo()
+        self.ensure_model()
+        self.validate_assets()
+
+        out_dir = self.output_root / (out_dir_name or f"multi_turn_{Path(input_image).stem}")
+        out_dir.mkdir(parents=True, exist_ok=True)
+
         image_json = json.dumps([str(input_image)])
         prompts_json = json.dumps(turns)
+
         base_overrides = [
             f"generation.positive_prompt.image_path={image_json}",
             f"generation.positive_prompt.prompts={prompts_json}",
         ]
-        if pad_img_placeholder is not None:
-            base_overrides.append(f"generation.pad_img_placehoder={str(bool(pad_img_placeholder)).lower()}")
         overrides = self._build_overrides(
             extra_overrides=base_overrides,
             cfg_scale=cfg_scale,
@@ -220,9 +187,17 @@ except Exception:
             aspect_ratio_input=aspect_ratio_input,
             steps=steps,
         )
-        self._run_vincie_once(overrides, out_dir)
+
+        cmd = [
+            self.python,
+            "main.py",
+            str(self.generate_yaml),
+            *overrides,
+            f"generation.output.dir={str(out_dir)}",
+        ]
+        self._run(cmd, cwd=self.repo_dir, env=self._env)
         self._clean_gpu_memory()
-        return out_dir
+        return str(out_dir)
 
     def multi_concept_compose(
         self,
@@ -235,11 +210,18 @@ except Exception:
         resolution_input: Optional[int] = None,
         aspect_ratio_input: Optional[str] = None,
         steps: Optional[int] = None,
-    ) -> Path:
+    ) -> str:
+        self.ensure_repo()
+        self.ensure_model()
+        self.validate_assets()
+
         out_dir = self.output_root / (out_dir_name or "multi_concept")
+        out_dir.mkdir(parents=True, exist_ok=True)
+
         imgs_json = json.dumps([str(p) for p in concept_images])
         prompts_all = concept_prompts + [final_prompt]
         prompts_json = json.dumps(prompts_all)
+
         base_overrides = [
             f"generation.positive_prompt.image_path={imgs_json}",
             f"generation.positive_prompt.prompts={prompts_json}",
@@ -252,13 +234,14 @@ except Exception:
             aspect_ratio_input=aspect_ratio_input,
             steps=steps,
         )
-        self._run_vincie_once(overrides, out_dir)
-        self._clean_gpu_memory()
-        return out_dir
 
-    @staticmethod
-    def _slug(path_or_text: str) -> str:
-        p = Path(path_or_text)
-        base = p.stem if p.exists() else str(path_or_text)
-        keep = "".join(c if c.isalnum() or c in "-_." else "_" for c in str(base))
-        return keep[:64]
+        cmd = [
+            self.python,
+            "main.py",
+            str(self.generate_yaml),
+            *overrides,
+            f"generation.output.dir={str(out_dir)}",
+        ]
+        self._run(cmd, cwd=self.repo_dir, env=self._env)
+        self._clean_gpu_memory()
+        return str(out_dir)