Update services/vincie.py

services/vincie.py

@@ -1,35 +1,11 @@
 #!/usr/bin/env python3
-"""
-VincieService - Orquestrador de Backend para o VINCIE com Suporte a Multi-GPU
-
-Responsabilidades:
-- Garante que o repositório VINCIE upstream está presente.
-- Baixa os arquivos de checkpoint (dit.pth, vae.pth) via hf_hub.
-- Cria um link simbólico de compatibilidade para os caminhos de checkpoint.
-- Executa o script main.py do VINCIE usando 'torchrun' para inferência distribuída.
-- Lida com a passagem de parâmetros dinâmicos da UI (GPU, Batch Size, etc.) para a linha de comando.
-
-Referência Upstream: https://github.com/ByteDance-Seed/VINCIE
-Desenvolvido por: [email protected] | https://github.com/carlex22
-Versão: 3.0.0
-"""
-
-import os
-import sys
-import json
-import subprocess
+import os, sys, json, subprocess
 from pathlib import Path
 from typing import List, Optional
-
-from huggingface_hub import hf_hub_download
-
+from time import time, sleep
+from huggingface_hub import snapshot_download
 
 class VincieService:
-    """
-    Serviço de alto nível para preparar os ativos de tempo de execução do VINCIE e invocar a geração,
-    com foco em maximizar o desempenho em hardware multi-GPU.
-    """
-
     def __init__(
         self,
         repo_dir: str = "/app/VINCIE",
@@ -37,161 +13,124 @@ class VincieService:
         python_bin: str = "python3",
         repo_id: str = "ByteDance-Seed/VINCIE-3B",
     ):
-        """
-        Inicializa o serviço com caminhos e configurações de tempo de execução.
-        """
         self.repo_dir = Path(repo_dir)
         self.ckpt_dir = Path(ckpt_dir)
         self.python = python_bin
         self.repo_id = repo_id
-
-        # Caminhos para os diferentes arquivos de configuração
         self.generate_yaml = self.repo_dir / "configs" / "generate.yaml"
-        self.generate_yaml_distributed = self.repo_dir / "configs" / "generate_distributed.yaml"
-
         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)
 
-    # ---------- Funções de Configuração (Setup) ----------
-
+    # ---------- Repo ----------
     def ensure_repo(self, git_url: str = "https://github.com/ByteDance-Seed/VINCIE") -> None:
-        """Clona o repositório oficial do VINCIE, se ausente."""
         if not self.repo_dir.exists():
-            print(f"Clonando o repositório VINCIE de {git_url}...")
-            # Clona com --depth 1 para baixar apenas a versão mais recente
             subprocess.run(["git", "clone", "--depth", "1", git_url, str(self.repo_dir)], check=True)
 
+    # ---------- Model ----------
+    def _have_core_files(self) -> bool:
+        return (self.ckpt_dir / "dit.pth").exists() and (self.ckpt_dir / "vae.pth").exists()
+
     def ensure_model(self, hf_token: Optional[str] = None) -> None:
-        """Baixa os arquivos de checkpoint e cria o link simbólico necessário."""
         self.ckpt_dir.mkdir(parents=True, exist_ok=True)
         token = hf_token or os.getenv("HF_TOKEN") or os.getenv("HUGGINGFACE_TOKEN")
 
-        def _need_download(p: Path) -> bool:
-            try: 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_download(dst):
-                print(f"Baixando {fname} de {self.repo_id}...")
-                hf_hub_download(repo_id=self.repo_id, filename=fname, local_dir=str(self.ckpt_dir), token=token)
-
-        link_path = self.repo_dir / "ckpt" / "VINCIE-3B"
-        if not link_path.exists():
-            try:
-                if link_path.is_symlink(): link_path.unlink()
-                link_path.symlink_to(self.ckpt_dir, target_is_directory=True)
-            except Exception as e:
-                print(f"Aviso: falha ao criar o link simbólico do checkpoint: {e}")
-
-    # ---------- Executor Principal com Suporte a Multi-GPU ----------
-
-    def _run_vincie(self, overrides: List[str], work_output: Path, video_mode: bool = False, num_gpus: int = 1) -> None:
-        """
-        Invoca o main.py do VINCIE, usando torchrun para execução distribuída se num_gpus > 1.
-        """
+        # Reutiliza cache persistente se HF_HOME estiver definido
+        hf_home = Path(os.environ.get("HF_HOME", "/app/.cache/huggingface"))
+        cache_dir = str(hf_home / "hub")
+        os.makedirs(cache_dir, exist_ok=True)
+
+        need_full = True
+        if any(self.ckpt_dir.iterdir()):
+            need_full = not self._have_core_files()
+
+        if need_full:
+            print(f"[vince] snapshot_download {self.repo_id} -> {self.ckpt_dir} (cache_dir={cache_dir})")
+            snapshot_download(
+                repo_id=self.repo_id,
+                cache_dir=cache_dir,
+                local_dir=str(self.ckpt_dir),
+                # local_dir_use_symlinks está deprecado; removível em versões futuras do hub
+                local_dir_use_symlinks=False,
+                resume_download=True,
+                token=token,
+            )
+        else:
+            print("[vince] modelo já presente; pulando download")
+
+        if not self._have_core_files():
+            raise FileNotFoundError("Modelo baixado, mas dit.pth/vae.pth não encontrados em ckpt_dir")
+
+        # Symlink de compatibilidade: /app/VINCIE/ckpt/VINCIE-3B -> /app/ckpt/VINCIE-3B
+        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()
+        return bool(have_repo and self._have_core_files())
+
+    # ---------- Execução ----------
+    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)
-        
-        # O modo distribuído usa um script de geração diferente
-        config_file = self.generate_yaml_distributed if num_gpus > 1 else self.generate_yaml
-        entry_script = "generate_distributed.py" if num_gpus > 1 else "generate.py"
-        
-        base_cmd = [
+        cmd = [
+            self.python,
             "main.py",
-            str(config_file),
+            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")))
 
-        if num_gpus > 1:
-            cmd = ["torchrun", f"--nproc_per_node={num_gpus}", *base_cmd]
-        else:
-            cmd = [self.python, *base_cmd]
-            
-        env = os.environ.copy()
-        if video_mode:
-            env["VINCIE_GENERATE_VIDEO"] = "1"
-        elif "VINCIE_GENERATE_VIDEO" in env:
-            del env["VINCIE_GENERATE_VIDEO"]
-
-        print(f"Executando comando: {' '.join(cmd)}")
-        subprocess.run(cmd, cwd=self.repo_dir, check=True, env=env)
-
-    # ---------- Pipelines de Geração ----------
-
+    # ---------- Pipelines ----------
     def multi_turn_edit(self, input_image: str, turns: List[str], **kwargs) -> Path:
-        """Executa a pipeline de edição multi-turno com configurações dinâmicas e suporte a multi-GPU."""
-        num_gpus = kwargs.get('num_gpus', 1)
-        batch_size = kwargs.get('batch_size', 1) if num_gpus == 1 else num_gpus
-        
-        out_dir = self.output_root / f"multi_turn_{self._slug(input_image)}"
-        
+        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.image_path="{str(input_image)}"',
             f"generation.positive_prompt.prompts={json.dumps(turns)}",
-            f"generation.seed={kwargs.get('seed', 1)}",
-            f"diffusion.timesteps.sampling.steps={kwargs.get('steps', 50)}",
-            f"diffusion.cfg.scale={kwargs.get('cfg_scale', 7.5)}",
-            f"generation.negative_prompt={json.dumps(kwargs.get('negative_prompt', ''))}",
-            f"generation.resolution={kwargs.get('resolution', 512)}",
-            f"generation.batch_size={batch_size}",
+            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="{kwargs.get("negative_prompt","")}"',
+            f"generation.resolution={int(kwargs.get('resolution', 512))}",
+            f"generation.batch_size={int(kwargs.get('batch_size', 1))}",
         ]
-        
-        if kwargs.get('use_vae_slicing', True):
-            overrides.extend(["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:
-        """Executa a pipeline de texto-para-vídeo com configurações dinâmicas e suporte a multi-GPU."""
-        num_gpus = kwargs.get('num_gpus', 1)
-        batch_size = kwargs.get('batch_size', 1) if num_gpus == 1 else num_gpus
-        
-        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.seed={kwargs.get('seed', 1)}",
-            f"diffusion.timesteps.sampling.steps={kwargs.get('steps', 50)}",
-            f"diffusion.cfg.scale={kwargs.get('cfg_scale', 7.5)}",
-            f"generation.negative_prompt={json.dumps(kwargs.get('negative_prompt', ''))}",
-            f"generation.resolution={kwargs.get('resolution', 512)}",
-            f"generation.fps={kwargs.get('fps', 2)}",
-            f"generation.batch_size={batch_size}",
+            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.resolution={int(kwargs.get('resolution', 512))}",
+            f"generation.batch_size={int(kwargs.get('batch_size', 1))}",
         ]
-        
-        if kwargs.get('use_vae_slicing', True):
-            overrides.extend(["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
-
-    def multi_concept_compose(self, concept_images: List[str], concept_prompts: List[str], final_prompt: str) -> Path:
-        """Executa a pipeline de composição multi-conceito."""
-        out_dir = self.output_root / "multi_concept"
-        prompts_all = concept_prompts + [final_prompt]
-        
-        overrides = [
-            f"generation.positive_prompt.image_path={json.dumps([str(p) for p in concept_images])}",
-            f"generation.positive_prompt.prompts={json.dumps(prompts_all)}",
-            "generation.pad_img_placehoder=False",
-        ]
-        
-        # O modo multi-conceito é mais complexo e pode não escalar bem com FSDP da mesma forma.
-        # Por padrão, vamos mantê-lo em uma única GPU para garantir estabilidade.
-        # Pode ser adaptado para multi-GPU se necessário, mas requer testes mais aprofundados.
-        self._run_vincie(overrides, out_dir, video_mode=False, num_gpus=1)
-        return out_dir
-
-    # ---------- Helper ----------
-
-    @staticmethod
-    def _slug(path_or_text: str) -> str:
-        """Produz um nome curto e seguro para o sistema de arquivos."""
-        base_name = Path(path_or_text).stem if Path(path_or_text).exists() else str(path_or_text)
-        safe_name = "".join(c if c.isalnum() or c in "-_." else "_" for c in base_name)
-        return safe_name[:64]