Update services/vincie.py

services/vincie.py

@@ -1,17 +1,16 @@
 #!/usr/bin/env python3
+
 """
-VincieService
-- Ensures the upstream VINCIE repository is present.
-- Fetches the minimal checkpoint files (dit.pth, vae.pth) via hf_hub_download into /app/ckpt/VINCIE-3B.
-- Creates a compatibility symlink /app/VINCIE/ckpt/VINCIE-3B -> /app/ckpt/VINCIE-3B for repo-relative paths.
-- Runs the official VINCIE main.py with Hydra/YACS overrides for both multi-turn and multi-concept generation.
-- Optionally injects a minimal 'apex.normalization' shim when NVIDIA Apex is not available (to avoid import errors).
-Upstream reference: https://github.com/ByteDance-Seed/VINCIE
+VincieService - VINCIE-3B High-Performance Service
 
-Developed by carlex22@gmail.com
-https://github.com/carlex22
+- Ensures upstream VINCIE repository is present.
+- Downloads checkpoint using HF structured cache (persistent, no duplication).
+- Creates compatibility symlink for repo-relative paths.
+- Runs official VINCIE main.py with Hydra/YACS overrides for multi-turn and multi-concept.
 
-Version 1.0.0
+Upstream: https://github.com/ByteDance-Seed/VINCIE
+Developed by: [email protected] | https://github.com/carlex22
+Version: 2.0.0 (Cache-optimized for 8× L40S)
 """
 
 import os
@@ -21,24 +20,24 @@ import subprocess
 from pathlib import Path
 from typing import List, Optional
 
-from huggingface_hub import hf_hub_download
+from huggingface_hub import snapshot_download
 
 
 class VincieService:
     """
-    High-level service for preparing VINCIE runtime assets and invoking generation.
-
-    Responsibilities:
-    - Repository management: clone the official VINCIE repository when missing.
-    - Checkpoint management: download dit.pth and vae.pth from the VINCIE-3B checkpoint on the Hub.
-    - Path compatibility: ensure /app/VINCIE/ckpt/VINCIE-3B points to /app/ckpt/VINCIE-3B.
-    - Runners: execute main.py with generate.yaml overrides for multi-turn edits and multi-concept composition.
-    - Apex shim: provide a minimal fallback for apex.normalization if Apex isn’t installed.
-
-    Defaults assume the Docker/container layout used by the Space:
-    - Repository directory: /app/VINCIE
-    - Checkpoint directory: /app/ckpt/VINCIE-3B
-    - Output root: /app/outputs
+    High-level service for VINCIE runtime with persistent cache support.
+    
+    Optimizations:
+    - Uses HF_HUB_CACHE structured cache (no file duplication)
+    - Persistent across container restarts when /data volume is mounted
+    - Parallel downloads with max_workers
+    - Resume capability for interrupted downloads
+    
+    Paths (Docker/container layout):
+    - Repository: /app/VINCIE
+    - Cache: $HF_HUB_CACHE/models--ByteDance-Seed--VINCIE-3B/
+    - Symlink: /app/ckpt/VINCIE-3B -> cached snapshot
+    - Output: /app/outputs
     """
 
     def __init__(
@@ -46,165 +45,154 @@ class VincieService:
         repo_dir: str = "/app/VINCIE",
         ckpt_dir: str = "/app/ckpt/VINCIE-3B",
         python_bin: str = "python",
-        repo_id: str = "ByteDance-Seed/VINCIE-3B",
+        model_repo: str = "ByteDance-Seed/VINCIE-3B",
     ):
         """
-        Initialize the service with paths and runtime settings.
-
+        Initialize service with paths and runtime settings.
+        
         Args:
-            repo_dir: Filesystem location of the upstream VINCIE repository clone.
-            ckpt_dir: Filesystem location where dit.pth and vae.pth are stored.
-            python_bin: Python executable to invoke for main.py (e.g., 'python' or a full path).
-            repo_id: Hugging Face Hub repo id for the VINCIE-3B checkpoint.
-
-        Side-effects:
-            - Ensures the output root directory exists.
-            - Ensures the repo ckpt/ directory exists (for symlink placement).
+            repo_dir: VINCIE repository clone location
+            ckpt_dir: Checkpoint symlink location (points to cache)
+            python_bin: Python executable for main.py
+            model_repo: HuggingFace Hub repo ID for VINCIE-3B
         """
         self.repo_dir = Path(repo_dir)
         self.ckpt_dir = Path(ckpt_dir)
         self.python = python_bin
-        self.repo_id = repo_id
-
-        # Canonical config and paths within the upstream repo
+        self.model_repo = model_repo
+        
+        # Config paths within upstream repo
         self.generate_yaml = self.repo_dir / "configs" / "generate.yaml"
         self.assets_dir = self.repo_dir / "assets"
-
-        # Output root for generated media
+        
+        # Output root
         self.output_root = Path("/app/outputs")
         self.output_root.mkdir(parents=True, exist_ok=True)
-
-        # Ensure ckpt/ exists in the repo (symlink target lives here)
+        
+        # Ensure repo ckpt/ dir exists (for 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:
         """
-        Clone the official VINCIE repository when missing.
-
+        Clone official VINCIE repository if missing.
+        
         Args:
-            git_url: Source URL of the official VINCIE repo.
-
+            git_url: Source URL of VINCIE repo
+            
         Raises:
-            subprocess.CalledProcessError on git clone failure.
+            subprocess.CalledProcessError: On git clone failure
         """
         if not self.repo_dir.exists():
+            print(f"[VincieService] Cloning {git_url} to {self.repo_dir}")
             subprocess.run(["git", "clone", git_url, str(self.repo_dir)], check=True)
-
-    def ensure_model(self, hf_token: Optional[str] = None) -> None:
+            print("[VincieService] Clone complete")
+        else:
+            print(f"[VincieService] Repository exists: {self.repo_dir}")
+        
+        # Validate main.py exists
+        main_py = self.repo_dir / "main.py"
+        if not main_py.exists():
+            raise FileNotFoundError(
+                f"main.py not found in {self.repo_dir}. "
+                f"Repository may be incomplete or corrupted."
+            )
+
+    def ensure_model(self, hf_token: Optional[str] = None) -> Path:
         """
-        Download the minimal VINCIE-3B checkpoint files if missing and create a repo-compatible symlink.
-
-        Files fetched from the Hub (repo_id):
-          - dit.pth
-          - vae.pth
-
-        The files are placed under self.ckpt_dir (default /app/ckpt/VINCIE-3B) and a symlink
-        /app/VINCIE/ckpt/VINCIE-3B -> /app/ckpt/VINCIE-3B is created to match upstream relative paths.
-
+        Download VINCIE-3B checkpoint to structured HF cache (persistent).
+        
+        Uses snapshot_download with cache_dir (not local_dir) to:
+        - Store files once in $HF_HUB_CACHE/models--ByteDance-Seed--VINCIE-3B/
+        - Create symlink at self.ckpt_dir pointing to cached snapshot
+        - Enable automatic deduplication and resume across restarts
+        
         Args:
-            hf_token: Optional Hugging Face token; defaults to env HF_TOKEN or HUGGINGFACE_TOKEN.
-
-        Notes:
-            - Uses hf_hub_download with local_dir, so files are placed directly in the target directory.
-            - A basic size check (> 1MB) is used to decide whether to refetch a file.
+            hf_token: Optional HF token (defaults to env HF_TOKEN)
+            
+        Returns:
+            Path to the cached snapshot directory
+            
+        Raises:
+            Exception: On download or validation failure
         """
-        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:
-            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(dst):
-                print(f"Downloading {fname} from {self.repo_id} ...")
-                hf_hub_download(
-                    repo_id=self.repo_id,
-                    filename=fname,
-                    local_dir=str(self.ckpt_dir),
-                    local_dir_use_symlinks=False,
-                    token=token,
-                    force_download=False,
-                    local_files_only=False,
-                )
-
-        # Compatibility symlink for repo-relative ckpt paths
-        link = self.repo_dir / "ckpt" / "VINCIE-3B"
+        cache_dir = os.environ.get("HF_HUB_CACHE", 
+                                   os.path.expanduser("~/.cache/huggingface/hub"))
+        
+        print(f"[VincieService] Ensuring model in cache: {cache_dir}")
+        print(f"[VincieService] Model repo: {self.model_repo}")
+        
         try:
-            if link.is_symlink() or link.exists():
-                try:
-                    link.unlink()
-                except IsADirectoryError:
-                    # If a directory sits at that path, we leave it as-is or replace as needed
-                    pass
-            if not link.exists():
-                link.symlink_to(self.ckpt_dir, target_is_directory=True)
+            # Download to structured cache (persistent, deduplicated)
+            model_path = snapshot_download(
+                repo_id=self.model_repo,
+                cache_dir=cache_dir,
+                token=token,
+                resume_download=True,       # Resume interrupted downloads
+                max_workers=8,              # Parallel downloads
+                allow_patterns=["*.pth", "*.json", "*.txt", "*.yaml"],  # Skip unused files
+            )
+            
+            print(f"[VincieService] Model cached at: {model_path}")
+            
+            # Create/update symlink for compatibility
+            self.ckpt_dir.parent.mkdir(parents=True, exist_ok=True)
+            
+            if self.ckpt_dir.is_symlink():
+                self.ckpt_dir.unlink()
+            elif self.ckpt_dir.exists():
+                # If it's a directory with old files, remove it
+                import shutil
+                shutil.rmtree(self.ckpt_dir)
+            
+            self.ckpt_dir.symlink_to(model_path, target_is_directory=True)
+            print(f"[VincieService] Symlink: {self.ckpt_dir} -> {model_path}")
+            
+            # Validate critical files
+            dit_pth = Path(model_path) / "dit.pth"
+            vae_pth = Path(model_path) / "vae.pth"
+            
+            if not dit_pth.exists():
+                raise FileNotFoundError(f"dit.pth missing in {model_path}")
+            if not vae_pth.exists():
+                raise FileNotFoundError(f"vae.pth missing in {model_path}")
+            
+            dit_size = dit_pth.stat().st_size / (1024**3)  # GB
+            vae_size = vae_pth.stat().st_size / (1024**3)
+            print(f"[VincieService] Checkpoint sizes: dit.pth={dit_size:.2f}GB, vae.pth={vae_size:.2f}GB")
+            
+            # Create repo-relative symlink
+            repo_link = self.repo_dir / "ckpt" / "VINCIE-3B"
+            try:
+                if repo_link.is_symlink() or repo_link.exists():
+                    try:
+                        repo_link.unlink()
+                    except IsADirectoryError:
+                        import shutil
+                        shutil.rmtree(repo_link)
+                
+                if not repo_link.exists():
+                    repo_link.symlink_to(self.ckpt_dir, target_is_directory=True)
+                    print(f"[VincieService] Repo symlink: {repo_link} -> {self.ckpt_dir}")
+            except Exception as e:
+                print(f"[VincieService] Warning: repo symlink failed: {e}")
+            
+            return Path(model_path)
+            
         except Exception as e:
-            print("Warning: failed to create checkpoint symlink:", e)
-
-    def ensure_apex(self, enable_shim: bool = True) -> None:
-        """
-        Ensure apex.normalization importability.
-
-        If NVIDIA Apex is not installed, and enable_shim=True, inject a minimal shim implementing:
-          - FusedRMSNorm via torch.nn.RMSNorm
-          - FusedLayerNorm via torch.nn.LayerNorm
-
-        This prevents import-time failures in code that references apex.normalization while
-        sacrificing any Apex-specific kernel benefits.
-
-        Args:
-            enable_shim: Whether to install a local shim when 'apex.normalization' is missing.
-        """
-        try:
-            import importlib
-            importlib.import_module("apex.normalization")
-            return
-        except Exception:
-            if not enable_shim:
-                return
-
-        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"
-            "\n"
-            "class FusedLayerNorm(nn.Module):\n"
-            "    def __init__(self, normalized_shape, eps=1e-5, elementwise_affine=True):\n"
-            "        super().__init__()\n"
-            "        self.mod = nn.LayerNorm(normalized_shape, eps=eps, elementwise_affine=elementwise_affine)\n"
-            "    def forward(self, x):\n"
-            "        return self.mod(x)\n"
-        )
-
-        # Make shim importable in this process and child processes
-        sys.path.insert(0, str(shim_root))
-        os.environ["PYTHONPATH"] = f"{str(shim_root)}:{os.environ.get('PYTHONPATH','')}"
+            print(f"[VincieService] Model download error: {e}")
+            import traceback
+            traceback.print_exc()
+            raise
 
     def ready(self) -> bool:
         """
-        Quick readiness probe for UI:
-        - The repository and generate.yaml exist.
-        - Minimal checkpoint files (dit.pth, vae.pth) exist.
-
+        Readiness check for UI/health probes.
+        
         Returns:
-            True if the environment is ready to run generation tasks; otherwise False.
+            True if repo and checkpoints are ready
         """
         have_repo = self.repo_dir.exists() and self.generate_yaml.exists()
         dit_ok = (self.ckpt_dir / "dit.pth").exists()
@@ -215,16 +203,17 @@ class VincieService:
 
     def _run_vincie(self, overrides: List[str], work_output: Path) -> None:
         """
-        Invoke VINCIE's main.py with Hydra/YACS overrides inside the upstream repo directory.
-
+        Invoke VINCIE main.py with Hydra/YACS overrides.
+        
         Args:
-            overrides: A list of CLI overrides (e.g., generation.positive_prompt.*).
-            work_output: Output directory path for generated assets.
-
+            overrides: CLI overrides list
+            work_output: Output directory for generated assets
+            
         Raises:
-            subprocess.CalledProcessError if the underlying process fails.
+            subprocess.CalledProcessError: If main.py fails
         """
         work_output.mkdir(parents=True, exist_ok=True)
+        
         cmd = [
             self.python,
             "main.py",
@@ -232,6 +221,8 @@ class VincieService:
             *overrides,
             f"generation.output.dir={str(work_output)}",
         ]
+        
+        print(f"[VincieService] Running: {' '.join(cmd)}")
         env = os.environ.copy()
         subprocess.run(cmd, cwd=self.repo_dir, check=True, env=env)
 
@@ -244,29 +235,27 @@ class VincieService:
         out_dir_name: Optional[str] = None,
     ) -> Path:
         """
-        Run the official 'multi-turn' generation equivalent.
-
-        This wraps generate.yaml using overrides:
-          - generation.positive_prompt.image_path = [ "<input-image-path>" ]
-          - generation.positive_prompt.prompts    = [ "<turn1>", "<turn2>", ... ]
-
+        Run multi-turn image editing pipeline.
+        
         Args:
-            input_image: Path to the single input image on disk.
-            turns: A list of editing instructions, in the order they should be applied.
-            out_dir_name: Optional name for the output subdirectory; auto-generated if omitted.
-
+            input_image: Path to input image
+            turns: List of editing instructions (applied sequentially)
+            out_dir_name: Optional output dir name
+            
         Returns:
-            Path to the output directory containing images and, if produced, a video.
+            Path to output directory with results
         """
         out_dir = self.output_root / (out_dir_name or f"multi_turn_{self._slug(input_image)}")
+        
         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)}",
         ]
+        
         self._run_vincie(overrides, out_dir)
         return out_dir
 
@@ -280,34 +269,30 @@ class VincieService:
         out_dir_name: Optional[str] = None,
     ) -> Path:
         """
-        Run the 'multi-concept' composition pipeline.
-
-        The service forms:
-          - generation.positive_prompt.image_path = [ <concept-img-1>, ..., <concept-img-N> ]
-          - generation.positive_prompt.prompts    = [ <desc-1>, ..., <desc-N>, <final-prompt> ]
-          - generation.pad_img_placehoder         = False (preserves input shapes)
-          - ckpt.path                             = /app/ckpt/VINCIE-3B (by default)
-
+        Run multi-concept composition pipeline.
+        
         Args:
-            concept_images: Paths to concept images on disk.
-            concept_prompts: Per-image descriptions in the same order as concept_images.
-            final_prompt: Composition prompt appended after all per-image descriptions.
-            out_dir_name: Optional name for the output subdirectory; defaults to 'multi_concept'.
-
+            concept_images: List of concept image paths
+            concept_prompts: Per-image descriptions (same order)
+            final_prompt: Final composition prompt
+            out_dir_name: Optional output dir name
+            
         Returns:
-            Path to the output directory containing images and, if produced, a video.
+            Path to output directory with results
         """
         out_dir = self.output_root / (out_dir_name or "multi_concept")
+        
         imgs_json = json.dumps([str(p) for p in concept_images])
         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}",
             "generation.pad_img_placehoder=False",
             f"ckpt.path={str(self.ckpt_dir)}",
         ]
+        
         self._run_vincie(overrides, out_dir)
         return out_dir
 
@@ -316,13 +301,13 @@ class VincieService:
     @staticmethod
     def _slug(path_or_text: str) -> str:
         """
-        Produce a filesystem-friendly short name (max 64 chars) from a path or text.
-
+        Create filesystem-friendly name (max 64 chars).
+        
         Args:
-            path_or_text: An input path or arbitrary string.
-
+            path_or_text: Input path or text
+            
         Returns:
-            A sanitized string consisting of [A-Za-z0-9._-] with non-matching chars converted to underscores.
+            Sanitized string with [A-Za-z0-9._-] only
         """
         p = Path(path_or_text)
         base = p.stem if p.exists() else str(path_or_text)