Update services/vincie.py

services/vincie.py

@@ -1,19 +1,9 @@
 #!/usr/bin/env python3
-import os
-import json
-import subprocess
+import os, sys, json, subprocess
 from pathlib import Path
 from typing import List, Optional
-from huggingface_hub import snapshot_download
-
-def _ls(path: Path, title: str = ""):
-    try:
-        print(f"[vince][ls] {title}: {path}")
-        result = subprocess.run(["bash", "-lc", f"ls -la {str(path)} || true"], capture_output=True, text=True)
-        print(result.stdout)
-        print(result.stderr)
-    except Exception as e:
-        print(f"[vince][ls] erro listando {path}: {e}")
+from time import time, sleep
+from huggingface_hub import hf_hub_download
 
 class VincieService:
     def __init__(
@@ -23,241 +13,111 @@ class VincieService:
         python_bin: str = "python3",
         repo_id: str = "ByteDance-Seed/VINCIE-3B",
     ):
-        # Executar ls /app/ckpt antes de iniciar
-        print("[vince][pre-init] Executando ls -la /app/ckpt antes de iniciar o serviço")
-        _ls(Path("/app/ckpt"), "pre-init /app/ckpt")
-
         self.repo_dir = Path(repo_dir)
         self.ckpt_dir = Path(ckpt_dir)
         self.python = python_bin
         self.repo_id = repo_id
-
-        # Saídas persistentes
-        self.output_root = Path("/data/outputs") if Path("/data").exists() else Path("/app/outputs")
-        self.output_root.mkdir(parents=True, exist_ok=True)
-
-        # Configs do VINCIE
         self.generate_yaml = self.repo_dir / "configs" / "generate.yaml"
-        self.generate_yaml_distributed = self.repo_dir / "configs" / "generate_distributed.yaml"
-
-        # Diretório ckpt dentro do repositório (para symlink)
+        self.output_root = Path("/app/outputs")
+        self.output_root.mkdir(parents=True, exist_ok=True)
         (self.repo_dir / "ckpt").mkdir(parents=True, exist_ok=True)
 
-        print(f"[vince][init] repo_dir={self.repo_dir} ckpt_dir={self.ckpt_dir} output_root={self.output_root}")
-        _ls(self.repo_dir, "repo_dir (init)")
-        _ls(self.ckpt_dir, "ckpt_dir (init)")
-
-    # ---------------------- Utilidades de cache e links ----------------------
-    def _hf_cache_dir(self) -> Path:
-        if Path("/data").exists():
-            base = Path(os.environ.get("HF_HOME", "/data/.cache/huggingface"))
-        else:
-            base = Path(os.environ.get("HF_HOME", "/app/.cache/huggingface"))
-        cache_dir = base / "hub"
-        cache_dir.mkdir(parents=True, exist_ok=True)
-        # Symlink de compatibilidade com /app/.cache
-        try:
-            Path("/app/.cache").mkdir(parents=True, exist_ok=True)
-            if not Path("/app/.cache/huggingface").exists():
-                Path("/app/.cache/huggingface").symlink_to(base)
-        except Exception as e:
-            print(f"[vince][cache] aviso ao criar symlink /app/.cache/huggingface -> {base}: {e}")
-        print(f"[vince][cache] HF_HOME={base} HF_HUB_CACHE={cache_dir}")
-        _ls(cache_dir, "HF_HUB_CACHE")
-        return cache_dir
-
-    def _model_cache_path(self, cache_dir: Path) -> Path:
-        m = cache_dir / f"models--{self.repo_id.replace('/', '--')}"
-        print(f"[vince][cache] model_cache_path={m}")
-        return m
-
-    def _ensure_aliases(self) -> None:
-        print("[vince][aliases] verificando/criando aliases dit.pth e vae.pth")
-        candidates = {
-            "dit.pth": ["dit.pth", "VINCIE-3B-dit.pth", "dit.safetensors", "pytorch_model.bin"],
-            "vae.pth": ["vae.pth", "VINCIE-3B-vae.pth", "vae.safetensors", "vae.ckpt"],
-        }
-        for target_name, options in candidates.items():
-            target = self.ckpt_dir / target_name
-            if target.exists():
-                print(f"[vince][aliases] existe {target_name} -> OK")
-                continue
-            found = None
-            for opt in options:
-                for p in self.ckpt_dir.rglob(opt):
-                    found = p
-                    break
-                if found:
-                    break
-            if found:
-                try:
-                    target.symlink_to(found)
-                    print(f"[vince][aliases] criado {target_name} -> {found}")
-                except Exception as e:
-                    print(f"[vince][aliases] falha ao criar {target_name}: {e}")
-            else:
-                print(f"[vince][aliases] não encontrado candidato para {target_name}")
-        _ls(self.ckpt_dir, "ckpt_dir após aliases")
-
-    def _link_ckpt(self) -> None:
-        print("[vince][link] garantindo symlink /app/VINCIE/ckpt/VINCIE-3B -> ckpt_dir")
-        link_parent = self.repo_dir / "ckpt"
-        link_parent.mkdir(parents=True, exist_ok=True)
-        link_path = link_parent / "VINCIE-3B"
-        if link_path.exists() or link_path.is_symlink():
-            try:
-                link_path.unlink()
-            except Exception as e:
-                print(f"[vince][link] aviso ao remover link antigo: {e}")
-        try:
-            link_path.symlink_to(self.ckpt_dir, target_is_directory=True)
-            print(f"[vince][link] criado {link_path} -> {self.ckpt_dir}")
-        except Exception as e:
-            print(f"[vince][link] falha ao criar link: {e}")
-        _ls(link_parent, "repo/ckpt")
-
-    # ---------------------- Setup repositório/modelo ----------------------
     def ensure_repo(self, git_url: str = "https://github.com/ByteDance-Seed/VINCIE") -> None:
         if not self.repo_dir.exists():
-            print(f"[vince][repo] clonando {git_url} em {self.repo_dir}")
             subprocess.run(["git", "clone", "--depth", "1", git_url, str(self.repo_dir)], check=True)
-        else:
-            print(f"[vince][repo] já existe {self.repo_dir}")
-        _ls(self.repo_dir, "repo_dir (ensure_repo)")
 
     def ensure_model(self, hf_token: Optional[str] = None) -> None:
-        cache_dir = self._hf_cache_dir()
-        mcache = self._model_cache_path(cache_dir)
-        token = hf_token or os.getenv("HF_TOKEN") or os.getenv("HUGGING_FACE_TOKEN")
-
-        # 1) Usa ckpt_dir se já possui conteúdo
-        if self.ckpt_dir.exists():
-            try:
-                if any(self.ckpt_dir.iterdir()):
-                    print(f"[vince][model] dados presentes em {self.ckpt_dir} — pulando download")
-                    self._ensure_aliases()
-                    self._link_ckpt()
-                    return
-            except PermissionError as e:
-                print(f"[vince][model] permissão ao listar {self.ckpt_dir}: {e}")
+        self.ckpt_dir.mkdir(parents=True, exist_ok=True)
+        token = hf_token or os.getenv("HF_TOKEN") or os.getenv("HUGGINGFACE_TOKEN")
 
-        # 2) Usa cache se existe
-        if mcache.exists():
+        def _need(p: Path) -> bool:
             try:
-                if any(mcache.iterdir()):
-                    print(f"[vince][model] cache do Hub em {mcache} — linkando para {self.ckpt_dir}")
-                    self.ckpt_dir.mkdir(parents=True, exist_ok=True)
-                    for item in mcache.iterdir():
-                        target = self.ckpt_dir / item.name
-                        if not target.exists():
-                            try:
-                                target.symlink_to(item)
-                            except Exception as e:
-                                print(f"[vince][model] aviso link {target} -> {item}: {e}")
-                    self._ensure_aliases()
-                    self._link_ckpt()
-                    return
-            except PermissionError as e:
-                print(f"[vince][model] permissão ao listar {mcache}: {e}")
-
-        # 3) Baixa snapshot se necessário
-        print(f"[vince][model] baixando snapshot {self.repo_id} em {cache_dir}")
-        snapshot_download(
-            repo_id=self.repo_id,
-            cache_dir=str(cache_dir),
-            token=token,
-            resume_download=True,
-            local_dir_use_symlinks=False
-        )
-
-        # 4) Linka do cache recém-baixado
-        if mcache.exists() and any(mcache.iterdir()):
-            print(f"[vince][model] pós-download: linkando {mcache} -> {self.ckpt_dir}")
-            self.ckpt_dir.mkdir(parents=True, exist_ok=True)
-            for item in mcache.iterdir():
-                target = self.ckpt_dir / item.name
-                if not target.exists():
-                    try:
-                        target.symlink_to(item)
-                    except Exception as e:
-                        print(f"[vince][model] aviso link pós-download {target} -> {item}: {e}")
-            self._ensure_aliases()
-            self._link_ckpt()
-        else:
-            print("[vince][model] aviso: snapshot feito mas diretório do modelo não localizado no cache")
-        _ls(self.ckpt_dir, "ckpt_dir (final ensure_model)")
-
-    # ---------------------- Execução ----------------------
-    def run_vincie(self, overrides: List[str], work_output: Path, video_mode: bool = False, num_gpus: int = 1) -> None:
+                return not (p.exists() and p.stat().st_size > 1_000_000)
+            except FileNotFoundError:
+                return True
+
+        for fname in ["dit.pth", "vae.pth"]:
+            dst = self.ckpt_dir / fname
+            if _need(dst):
+                print(f"[vince] downloading {fname} from {self.repo_id} ...")
+                hf_hub_download(
+                    repo_id=self.repo_id,
+                    filename=fname,
+                    local_dir=str(self.ckpt_dir),
+                    token=token,
+                    force_download=False,
+                    local_files_only=False,
+                )
+
+        link = self.repo_dir / "ckpt" / "VINCIE-3B"
+        try:
+            if link.is_symlink() or link.exists():
+                try:
+                    link.unlink()
+                except IsADirectoryError:
+                    pass
+            if not link.exists():
+                link.symlink_to(self.ckpt_dir, target_is_directory=True)
+        except Exception as e:
+            print("[vince] symlink warning:", e)
+
+    def ready(self) -> bool:
+        have_repo = self.repo_dir.exists() and self.generate_yaml.exists()
+        dit_ok = (self.ckpt_dir / "dit.pth").exists()
+        vae_ok = (self.ckpt_dir / "vae.pth").exists()
+        return bool(have_repo and dit_ok and vae_ok)
+
+    def _wait_until_outputs(self, out_dir: Path, timeout_s: int = 300) -> None:
+        exts = (".png", ".jpg", ".jpeg", ".gif", ".mp4")
+        deadline = time() + timeout_s
+        while time() < deadline:
+            if any(p.is_file() and p.suffix.lower() in exts for p in out_dir.rglob("*")):
+                print(f"[vince] outputs detected in {out_dir}")
+                return
+            sleep(1)
+        print(f"[vince] warning: no outputs detected in {out_dir} within {timeout_s}s")
+
+    def _run_vincie(self, overrides: List[str], work_output: Path, wait_outputs: bool = True) -> 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)}",
+        ]
+        print("[vince] CWD=", self.repo_dir)
+        print("[vince] CMD=", " ".join(cmd))
+        subprocess.run(cmd, cwd=self.repo_dir, check=True, env=os.environ.copy())
+        if wait_outputs:
+            self._wait_until_outputs(work_output, timeout_s=int(os.getenv("VINCIE_WAIT_OUTPUTS_SEC", "300")))
 
-        # Validação final: dit.pth deve existir no repo/ckpt
-        need = self.repo_dir / "ckpt" / "VINCIE-3B" / "dit.pth"
-        if not need.exists():
-            print(f"[vince][run] {need} ausente — reforçando aliases e link")
-            self._ensure_aliases()
-            self._link_ckpt()
-            if not need.exists():
-                _ls(self.ckpt_dir, "ckpt_dir (antes do erro)")
-                _ls(self.repo_dir / "ckpt", "repo/ckpt (antes do erro)")
-                raise FileNotFoundError(f"Checkpoint ausente: {need}")
-
-        config_file = self.generate_yaml_distributed if num_gpus > 1 else self.generate_yaml
-        base_cmd = ["main.py", str(config_file)] + overrides + [f"generation.output.dir={str(work_output)}"]
-        cmd = (["torchrun", f"--nproc_per_node={num_gpus}"] + base_cmd) if num_gpus > 1 else (["python3"] + base_cmd)
-
-        env = os.environ.copy()
-        if video_mode:
-            env["VINCIE_GENERATE_VIDEO"] = "1"
-        else:
-            env.pop("VINCIE_GENERATE_VIDEO", None)
-
-        print(f"[vince][run] CWD={self.repo_dir}")
-        print(f"[vince][run] CMD={' '.join(cmd)}")
-        _ls(self.repo_dir / "ckpt" / "VINCIE-3B", "repo/ckpt/VINCIE-3B (pre-run)")
-        subprocess.run(cmd, cwd=self.repo_dir, check=True, env=env)
-
-    # ---------------------- Pipelines ----------------------
     def multi_turn_edit(self, input_image: str, turns: List[str], **kwargs) -> Path:
-        num_gpus = int(kwargs.get("num_gpus", 1))
-        batch_size = int(kwargs.get("batch_size", 1))
         out_dir = self.output_root / f"multi_turn_{Path(input_image).stem}"
         overrides = [
-            f'generation.positive_prompt.image_path={json.dumps(str(input_image))}',
-            f'generation.positive_prompt.prompts={json.dumps(turns)}',
+            f'generation.positive_prompt.image_path="{str(input_image)}"',
+            f"generation.positive_prompt.prompts={json.dumps(turns)}",
             f"generation.seed={int(kwargs.get('seed', 1))}",
             f"diffusion.timesteps.sampling.steps={int(kwargs.get('steps', 50))}",
             f"diffusion.cfg.scale={float(kwargs.get('cfg_scale', 7.5))}",
-            f"generation.negative_prompt={json.dumps(kwargs.get('negative_prompt', ''))}",
+            f'generation.negative_prompt="{kwargs.get("negative_prompt","")}"',
             f"generation.resolution={int(kwargs.get('resolution', 512))}",
-            f"generation.batch_size={batch_size}",
+            f"generation.batch_size={int(kwargs.get('batch_size', 1))}",
         ]
-        if kwargs.get("use_vae_slicing", True):
-            overrides += ["vae.slicing.split_size=1", "vae.slicing.memory_device=same"]
-        self.run_vincie(overrides, out_dir, video_mode=False, num_gpus=num_gpus)
+        self._run_vincie(overrides, out_dir, wait_outputs=True)
         return out_dir
 
-    def text_to_video(self, input_image: str, prompt: str, **kwargs) -> Path:
-        num_gpus = int(kwargs.get("num_gpus", 1))
-        batch_size = int(kwargs.get("batch_size", 1))
-        out_dir = self.output_root / f"txt2vid_{self._slug(prompt)}"
+    def multi_concept_compose(self, files: List[str], descs: List[str], final_prompt: str, **kwargs) -> Path:
+        out_dir = self.output_root / f"multi_concept_{len(files)}"
         overrides = [
-            f'generation.positive_prompt.image_path={json.dumps(str(input_image))}',
-            f'generation.positive_prompt.prompts={json.dumps([prompt])}',
+            f"generation.concepts.files={json.dumps(files)}",
+            f"generation.concepts.descs={json.dumps(descs)}",
+            f'generation.final_prompt="{final_prompt}"',
             f"generation.seed={int(kwargs.get('seed', 1))}",
             f"diffusion.timesteps.sampling.steps={int(kwargs.get('steps', 50))}",
             f"diffusion.cfg.scale={float(kwargs.get('cfg_scale', 7.5))}",
-            f"generation.negative_prompt={json.dumps(kwargs.get('negative_prompt', ''))}",
             f"generation.resolution={int(kwargs.get('resolution', 512))}",
-            f"generation.fps={int(kwargs.get('fps', 2))}",
-            f"generation.batch_size={batch_size}",
+            f"generation.batch_size={int(kwargs.get('batch_size', 1))}",
         ]
-        if kwargs.get("use_vae_slicing", True):
-            overrides += ["vae.slicing.split_size=1", "vae.slicing.memory_device=same"]
-        self.run_vincie(overrides, out_dir, video_mode=True, num_gpus=num_gpus)
+        self._run_vincie(overrides, out_dir, wait_outputs=True)
         return out_dir
-
-    @staticmethod
-    def _slug(text: str) -> str:
-        s = "".join(c if c.isalnum() or c in "-." else "_" for c in text)
-        return s[:64]