Update services/vince_server.py

services/vince_server.py

@@ -1,247 +1,170 @@
 #!/usr/bin/env python3
-"""
-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
+import os, sys, gc, subprocess
 from pathlib import Path
 from typing import List, Optional
+from omegaconf import OmegaConf, open_dict
 
-from huggingface_hub import snapshot_download
+VINCIE_DIR = Path(os.getenv("VINCIE_DIR", "/app/VINCIE"))
+if str(VINCIE_DIR) not in sys.path:
+    sys.path.insert(0, str(VINCIE_DIR))
 
-
-class VincieService:
-    def __init__(
-        self,
-        repo_dir: str = "/app/VINCIE",
-        python_bin: str = "python",
-        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_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()
-
-    # ---------- 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)
+# inclui 'models/' relativo
+try:
+    app_models = Path("/app/models"); vincie_models = VINCIE_DIR / "models"
+    if not app_models.exists() and vincie_models.exists():
+        app_models.symlink_to(vincie_models, target_is_directory=True)
+except Exception as e:
+    print("[vince_server] warn: link /app/models failed:", e)
+
+from common.config import load_config, create_object  # type: ignore
+
+class VinceServer:
+    def __init__(self, config_path: str="/app/VINCIE/configs/generate.yaml",
+                 *, base_overrides: Optional[List[str]]=None,
+                 output_root: str="/app/outputs", chdir_repo: bool=True):
+        self.config_path = config_path
+        self.output_root = Path(output_root); self.output_root.mkdir(parents=True, exist_ok=True)
+        overrides = list(base_overrides or [])
+        if chdir_repo:
+            try: os.chdir(str(VINCIE_DIR))
+            except Exception as e: print("[vince_server] warn: chdir repo failed:", e)
+
+        try:
+            self._load_and_bootstrap(overrides)
+        except Exception as e:
+            print("[vince_server] bootstrap failed; repairing symlinks:", e)
+            self._repair_ckpt_links()
+            self._load_and_bootstrap(overrides)
+
+    def _load_and_bootstrap(self, overrides: List[str]):
+        self._assert_ckpt_ready()
+        self.config = load_config(self.config_path, overrides)
+        self.gen = create_object(self.config)
+        self._bootstrap_models()
+
+    def _repair_ckpt_links(self):
+        # reusa snapshot atual para recriar links idempotentes
+        from services.vincie import VincieService
+        svc = VincieService(); svc.ensure_model()
+        snapshot = Path(str(svc.ckpt_dir))
+        for link in (VINCIE_DIR/"ckpt"/"VINCIE-3B", Path("/app/ckpt")/"VINCIE-3B"):
+            link.parent.mkdir(parents=True, exist_ok=True)
+            if link.is_symlink() and link.resolve()!=snapshot:
+                link.unlink()
+            if link.exists() and (not link.is_symlink()):
+                subprocess.run(["rm","-rf",str(link)], check=True)
+            if not link.exists():
+                link.symlink_to(snapshot, target_is_directory=True)
+        print("[vince_server] ckpt symlinks refreshed")
+
+    def _assert_ckpt_ready(self):
+        # ambos caminhos funcionam; a config usa 'ckpt/...' com cwd no repo
+        repo_link = VINCIE_DIR / "ckpt" / "VINCIE-3B"
+        if not repo_link.exists():
+            raise RuntimeError("missing ckpt link: /app/VINCIE/ckpt/VINCIE-3B")
+        must = [repo_link/"dit.pth", repo_link/"vae.pth", repo_link/"llm14b"]
+        missing = [str(p) for p in must if not p.exists()]
+        if missing:
+            raise RuntimeError(f"missing ckpt content: {missing}")
 
     @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"
-        if main_py.exists() and git_dir.exists():
-            return
-        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.model_repo,
-            revision=revision,
-            token=token or os.getenv("HF_TOKEN") or os.getenv("HUGGINGFACE_TOKEN"),
-            cache_dir=cache_dir,
-            resume_download=True,
-        )
-        self.ckpt_dir = Path(snapshot_path)
-        # 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,
-        extra_overrides: Optional[List[str]] = None,
-        cfg_scale: Optional[float] = None,
-        resolution_input: Optional[int] = None,
-        aspect_ratio_input: Optional[str] = None,
-        steps: Optional[int] = None,
-    ) -> List[str]:
-        overrides = list(extra_overrides or [])
-        # 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:
-            overrides.append(f"generation.resolution_input={resolution_input}")
-        if aspect_ratio_input is not None:
-            overrides.append(f"generation.aspect_ratio_input={aspect_ratio_input}")
-        if steps is not None:
-            overrides.append(f"generation.steps={steps}")
-        return overrides
-
-    def _clean_gpu_memory(self) -> None:
-        code = r"""
-import torch, gc
-try:
-    torch.cuda.synchronize()
-except Exception:
-    pass
-gc.collect()
-try:
-    torch.cuda.empty_cache()
-    torch.cuda.memory.reset_peak_memory_stats()
-except Exception:
-    pass
-"""
-        self._run([self.python, "-c", code], env=self._env)
-
-    # ---------- APIs ----------
-
-    def multi_turn_edit(
-        self,
-        input_image: str,
-        turns: List[str],
-        out_dir_name: Optional[str] = None,
-        *,
-        cfg_scale: Optional[float] = None,
-        resolution_input: Optional[int] = None,
-        aspect_ratio_input: Optional[str] = None,
-        steps: Optional[int] = None,
-    ) -> 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}")
+    def _make_writable(cfg):
+        try: OmegaConf.set_readonly(cfg, False); OmegaConf.set_struct(cfg, False)
+        except Exception: pass
+
+    def _bootstrap_models(self):
+        for name in ("configure_persistence","configure_models","configure_diffusion","configure_sampler"):
+            fn = getattr(self.gen, name, None)
+            if not callable(fn): raise RuntimeError(f"[vince_server] missing step: {name}")
+            fn()
+        if not hasattr(self.gen, "sampler"):
+            raise RuntimeError("[vince_server] missing component: sampler")
+
+    def _set_steps(self, steps: Optional[int]):
+        if not steps: return
+        sampler = getattr(self.gen, "sampler", None); t = getattr(sampler, "timesteps", None)
+        if sampler is None or t is None: return
+        try:
+            import torch
+            if hasattr(t,"__len__") and len(t)>0:
+                steps = max(1, min(int(steps), len(t)))
+                if steps < len(t):
+                    idx = torch.linspace(0, len(t)-1, steps).round().long().tolist()
+                    sampler.timesteps = [t[i] for i in idx]
+        except Exception as e:
+            print(f"[vince_server] Warning: set_steps failed: {e}")
+
+    def _apply_generation_overrides(self, *, out_dir: Path,
+                                    image_paths: Optional[List[str]]=None,
+                                    prompts: Optional[List[str]]=None,
+                                    final_prompt: Optional[str]=None,
+                                    cfg_scale: Optional[float]=None,
+                                    aspect_ratio_input: Optional[str]=None,
+                                    resolution_input: Optional[int]=None,
+                                    steps: Optional[int]=None):
+        self._make_writable(self.gen.config)
+        g = self.gen.config.generation
+        self._make_writable(g); self._make_writable(g.output); self._make_writable(g.positive_prompt)
+        with open_dict(g):
+            g.output.dir = str(out_dir)
+            if image_paths is not None: g.positive_prompt.image_path = list(image_paths)
+            if prompts is not None: g.positive_prompt.prompts = list(prompts)
+            if cfg_scale is not None:
+                try: g.cfg_scale = float(cfg_scale)
+                except Exception:
+                    with open_dict(self.gen.config):
+                        try: self.gen.config.diffusion.cfg.scale = float(cfg_scale)
+                        except Exception: print("[vince_server] Warning: unable to set cfg_scale")
+            if aspect_ratio_input is not None: g.aspect_ratio_input = str(aspect_ratio_input)
+            if resolution_input is not None:
+                try: g.resolution_input = int(resolution_input)
+                except Exception:
+                    try: g.resolution = int(resolution_input)
+                    except Exception: print("[vince_server] Warning: unable to set resolution")
+        self._set_steps(steps)
+
+    def _infer_once(self):
+        for name in ("inference_loop","entrypoint","run"):
+            fn = getattr(self.gen, name, None)
+            if callable(fn): fn(); return
+        raise RuntimeError("No valid inference method found on generator")
+
+    def _cleanup(self):
+        try:
+            import torch; torch.cuda.synchronize()
+        except Exception: pass
+        gc.collect()
+        try:
+            import torch; torch.cuda.empty_cache(); torch.cuda.memory.reset_peak_memory_stats()
+        except Exception: pass
+
+    # APIs
+    def generate_multi_turn(self, image_path: str, turns: List[str], *,
+                            out_dir_name: Optional[str]=None, cfg_scale: Optional[float]=None,
+                            aspect_ratio_input: Optional[str]=None, resolution_input: Optional[int]=None,
+                            steps: Optional[int]=None) -> str:
+        out_dir = self.output_root / (out_dir_name or f"multi_turn_{Path(image_path).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}",
-        ]
-        overrides = self._build_overrides(
-            extra_overrides=base_overrides,
-            cfg_scale=cfg_scale,
-            resolution_input=resolution_input,
-            aspect_ratio_input=aspect_ratio_input,
-            steps=steps,
-        )
-
-        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)
-
-    def multi_concept_compose(
-        self,
-        concept_images: List[str],
-        concept_prompts: List[str],
-        final_prompt: str,
-        out_dir_name: Optional[str] = None,
-        *,
-        cfg_scale: Optional[float] = None,
-        resolution_input: Optional[int] = None,
-        aspect_ratio_input: Optional[str] = None,
-        steps: Optional[int] = None,
-    ) -> str:
-        self.ensure_repo()
-        self.ensure_model()
-        self.validate_assets()
-
+        self._apply_generation_overrides(out_dir=out_dir, image_paths=[str(image_path)], prompts=list(turns),
+                                         cfg_scale=cfg_scale, aspect_ratio_input=aspect_ratio_input,
+                                         resolution_input=resolution_input, steps=steps)
+        self._infer_once(); self._cleanup(); return str(out_dir)
+
+    def generate_multi_concept(self, concept_images: List[str], concept_prompts: List[str], final_prompt: str, *,
+                               out_dir_name: Optional[str]=None, cfg_scale: Optional[float]=None,
+                               aspect_ratio_input: Optional[str]=None, resolution_input: Optional[int]=None,
+                               steps: Optional[int]=None) -> str:
         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}",
-            "generation.pad_img_placehoder=False",
-        ]
-        overrides = self._build_overrides(
-            extra_overrides=base_overrides,
-            cfg_scale=cfg_scale,
-            resolution_input=resolution_input,
-            aspect_ratio_input=aspect_ratio_input,
-            steps=steps,
-        )
-
-        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)
+        prompts_all = list(concept_prompts) + ([final_prompt] if final_prompt else [])
+        self._apply_generation_overrides(out_dir=out_dir, image_paths=[str(p) for p in concept_images],
+                                         prompts=prompts_all, final_prompt=final_prompt, cfg_scale=cfg_scale,
+                                         aspect_ratio_input=aspect_ratio_input, resolution_input=resolution_input,
+                                         steps=steps)
+        self._infer_once(); self._cleanup(); return str(out_dir)
+
+server = VinceServer(
+    config_path=os.getenv("VINCE_CONFIG", "/app/VINCIE/configs/generate.yaml"),
+    output_root=os.getenv("VINCE_OUTPUT", "/app/outputs"),
+    chdir_repo=True,
+)