Update services/vincie.py

services/vincie.py

@@ -63,31 +63,33 @@ class VincieService:
             # 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)
 
+    from huggingface_hub import snapshot_download
+
     def ensure_model(self, hf_token: Optional[str] = None) -> None:
-        """Baixa os arquivos de checkpoint e cria o link simbólico necessário."""
+        """Baixa o checkpoint completo 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)
-
+    
+        # Baixa o snapshot completo do repositório
+        print(f"Baixando snapshot completo de {self.repo_id} para {self.ckpt_dir}...")
+        snapshot_download(
+             repo_id=self.repo_id,
+             local_dir=str(self.ckpt_dir),
+             token=token,
+             resume_download=True,
+             local_dir_use_symlinks=False
+        )
+    
+        # Cria symlink se necessário
         link_path = self.repo_dir / "ckpt" / "VINCIE-3B"
         if not link_path.exists():
             try:
-                if link_path.is_symlink(): link_path.unlink()
+                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.