Update vincie_service.py

vincie_service.py

@@ -1,70 +1,161 @@
 #!/usr/bin/env python3
 import os
+import sys
 import json
-import shlex
 import subprocess
 from pathlib import Path
 from typing import List, Optional
-from huggingface_hub import snapshot_download
+
+from huggingface_hub import hf_hub_download
 
 class VincieService:
     """
-    Serviço mínimo que:
-    - garante repositório VINCIE clonado e dependências de modelo baixadas
-    - chama python main.py configs/generate.yaml com overrides oficiais
-    - expõe funções de alto nível para multi-turn editing e multi-concept composition
+    Serviço que:
+    - garante que o repo VINCIE está presente
+    - baixa dit.pth e vae.pth de ByteDance-Seed/VINCIE-3B com hf_token (runtime)
+    - cria symlink de compatibilidade para ckpt/VINCIE-3B dentro do repo
+    - executa main.py com overrides oficiais (Hydra) para edição multi-turn e composição multi-conceito
+    - fornece um fallback (shim) para apex.normalization.FusedRMSNorm quando Apex não está instalado
     """
 
-    def __init__(self,
-                 repo_dir: str = "/app/VINCIE",
-                 ckpt_dir: str = "/app/ckpt/VINCIE-3B",
-                 python_bin: str = "python"):
+    def __init__(
+        self,
+        repo_dir: str = "/app/VINCIE",
+        ckpt_dir: str = "/app/ckpt/VINCIE-3B",
+        python_bin: str = "python",
+        repo_id: str = "ByteDance-Seed/VINCIE-3B",
+    ):
         self.repo_dir = Path(repo_dir)
         self.ckpt_dir = Path(ckpt_dir)
         self.python = python_bin
+        self.repo_id = repo_id
+
         self.generate_yaml = self.repo_dir / "configs" / "generate.yaml"
         self.assets_dir = self.repo_dir / "assets"
         self.output_root = Path("/app/outputs")
         self.output_root.mkdir(parents=True, exist_ok=True)
 
+        # Garantir existência de pasta de ckpt relativa no repo para symlink
+        (self.repo_dir / "ckpt").mkdir(parents=True, exist_ok=True)
+
     # ---------- Setup ----------
+
     def ensure_repo(self, git_url: str = "https://github.com/ByteDance-Seed/VINCIE") -> None:
+        """Clona o repositório VINCIE se ainda não existir."""
         if not self.repo_dir.exists():
             subprocess.run(["git", "clone", git_url, str(self.repo_dir)], check=True)
-        # opcional: garantir submódulos/updates mínimos
-        # subprocess.run(["git", "pull"], cwd=self.repo_dir, check=True)
 
-    def ensure_model(self, repo_id: str = "ByteDance-Seed/VINCIE-3B") -> None:
-        self.ckpt_dir.mkdir(parents=True, exist_ok=True)
-        snapshot_download(
-            repo_id=repo_id,
-            local_dir=str(self.ckpt_dir),
-            local_dir_use_symlinks=False,
-            resume_download=True
+    def ensure_apex(self, enable_shim: bool = True) -> None:
+        """
+        Garante que apex.normalization.FusedRMSNorm está importável.
+        Se Apex não estiver instalado e enable_shim=True, cria um shim baseado em nn.RMSNorm.
+        """
+        try:
+            import importlib
+            importlib.import_module("apex.normalization")
+            return
+        except Exception:
+            if not enable_shim:
+                return
+
+        # Criar shim em /app/shims/apex/normalization.py
+        shim_root = Path("/app/shims")
+        apex_pkg = shim_root / "apex"
+        apex_pkg.mkdir(parents=True, exist_ok=True)
+
+        (apex_pkg / "__init__.py").write_text("from .normalization import *\n")
+
+        (apex_pkg / "normalization.py").write_text(
+            "import torch\n"
+            "import torch.nn as nn\n"
+            "\n"
+            "class FusedRMSNorm(nn.Module):\n"
+            "    def __init__(self, normalized_shape, eps=1e-6, elementwise_affine=True):\n"
+            "        super().__init__()\n"
+            "        self.mod = nn.RMSNorm(normalized_shape, eps=eps, elementwise_affine=elementwise_affine)\n"
+            "    def forward(self, x):\n"
+            "        return self.mod(x)\n"
         )
 
+        # Tornar o shim visível neste processo e nos subprocessos
+        sys.path.insert(0, str(shim_root))
+        os.environ["PYTHONPATH"] = f"{str(shim_root)}:{os.environ.get('PYTHONPATH','')}"
+
+    def ensure_model(self, hf_token: Optional[str] = None) -> None:
+        """
+        Baixa apenas os arquivos necessários do repo ByteDance-Seed/VINCIE-3B:
+        - dit.pth
+        - vae.pth
+        Usa hf_token (parâmetro ou variáveis HF_TOKEN/HUGGINGFACE_TOKEN).
+        Cria symlink de compatibilidade em /app/VINCIE/ckpt/VINCIE-3B -> /app/ckpt/VINCIE-3B.
+        """
+        self.ckpt_dir.mkdir(parents=True, exist_ok=True)
+
+        token = hf_token or os.getenv("HF_TOKEN") or os.getenv("HUGGINGFACE_TOKEN")
+
+        def _need(p: Path) -> bool:
+            return not (p.exists() and p.stat().st_size > 1_000_000)
+
+        for fname in ["dit.pth", "vae.pth"]:
+            dst = self.ckpt_dir / fname
+            if _need(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,
+                    force_download=False,
+                    local_files_only=False,
+                )
+
+        # Symlink de compatibilidade para caminhos relativos do repo
+        link = self.repo_dir / "ckpt" / "VINCIE-3B"
+        try:
+            if link.is_symlink() or link.exists():
+                # Remover link anterior inválido
+                try:
+                    link.unlink()
+                except IsADirectoryError:
+                    # Se for diretório, não remover conteúdos; só criar link se não existir
+                    pass
+            if not link.exists():
+                link.symlink_to(self.ckpt_dir, target_is_directory=True)
+        except Exception as e:
+            print("Aviso: falha ao criar symlink de ckpt:", e)
+
     def ready(self) -> bool:
-        return self.repo_dir.exists() and self.generate_yaml.exists() and self.ckpt_dir.exists()
+        """Verifica se o repo, config e checkpoints obrigatórios existem."""
+        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)
 
     # ---------- Core runner ----------
+
     def _run_vincie(self, overrides: List[str], work_output: Path) -> None:
+        """Executa main.py com overrides Hydra/YACS do VINCIE dentro do diretório do repo."""
         work_output.mkdir(parents=True, exist_ok=True)
+
         cmd = [
             self.python,
             "main.py",
             str(self.generate_yaml),
-            # overrides list (Hydra/YACS-style) como no README
             *overrides,
-            f"generation.output.dir={str(work_output)}"
+            f"generation.output.dir={str(work_output)}",
         ]
-        # executar dentro do diretório do repo VINCIE
-        subprocess.run(cmd, cwd=self.repo_dir, check=True)
+        # Herdar ambiente (inclui PYTHONPATH do shim, se aplicado)
+        env = os.environ.copy()
+        subprocess.run(cmd, cwd=self.repo_dir, check=True, env=env)
 
     # ---------- Multi-turn editing ----------
-    def multi_turn_edit(self,
-                        input_image: str,
-                        turns: List[str],
-                        out_dir_name: Optional[str] = None) -> Path:
+
+    def multi_turn_edit(
+        self,
+        input_image: str,
+        turns: List[str],
+        out_dir_name: Optional[str] = None,
+    ) -> Path:
         """
         Equivalente ao exemplo:
         python main.py configs/generate.yaml \
@@ -76,41 +167,45 @@ class VincieService:
         image_json = json.dumps([str(input_image)])
         prompts_json = json.dumps(turns)
         overrides = [
-            f'generation.positive_prompt.image_path={image_json}',
-            f'generation.positive_prompt.prompts={prompts_json}',
-            f'ckpt.path={str(self.ckpt_dir)}'
+            f"generation.positive_prompt.image_path={image_json}",
+            f"generation.positive_prompt.prompts={prompts_json}",
+            f"ckpt.path={str(self.ckpt_dir)}",
         ]
         self._run_vincie(overrides, out_dir)
         return out_dir
 
     # ---------- Multi-concept composition ----------
-    def multi_concept_compose(self,
-                              concept_images: List[str],
-                              concept_prompts: List[str],
-                              final_prompt: str,
-                              out_dir_name: Optional[str] = None) -> Path:
+
+    def multi_concept_compose(
+        self,
+        concept_images: List[str],
+        concept_prompts: List[str],
+        final_prompt: str,
+        out_dir_name: Optional[str] = None,
+    ) -> Path:
         """
-        Modelo de uso inspirado no README:
+        Uso inspirado no README do VINCIE:
         - image_path: lista de imagens dos conceitos
-        - prompts: lista com p1..pN e o prompt final concatenado
+        - prompts: [p1, p2, ..., final_prompt]
         """
         out_dir = self.output_root / (out_dir_name or "multi_concept")
         imgs_json = json.dumps([str(p) for p in concept_images])
-        # prompts devem alinhar com <IMG0>, <IMG1>, ... e incluir o prompt final no fim
         prompts_all = concept_prompts + [final_prompt]
         prompts_json = json.dumps(prompts_all)
         overrides = [
-            f'generation.positive_prompt.image_path={imgs_json}',
-            f'generation.positive_prompt.prompts={prompts_json}',
-            f"generation.pad_img_placehoder=False",  # segue exemplo público
-            f'ckpt.path={str(self.ckpt_dir)}'
+            f"generation.positive_prompt.image_path={imgs_json}",
+            f"generation.positive_prompt.prompts={prompts_json}",
+            f"generation.pad_img_placehoder=False",
+            f"ckpt.path={str(self.ckpt_dir)}",
         ]
         self._run_vincie(overrides, out_dir)
         return out_dir
 
     # ---------- Helpers ----------
+
     @staticmethod
     def _slug(path_or_text: str) -> str:
-        base = Path(path_or_text).stem if Path(path_or_text).exists() else path_or_text
+        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]