--offline-mode enhancements #52

Create vendored configs and other generated assets in refactored assetgen build
system which handles generated assets for the project.

Redirect to dgenerate.pipelinewrapper.hub_configs when possible in
--offline-mode

This allows loading single file checkpoints for SD1.5, SD2, SD3, and
Flux without ever connecting to the hub even once.

You must download and specify a VAE on disk manually for Flux and SD3
for single file loads.

This behavior is patched into huggingface_hub and diffusers dynamically
and active when --offline-mode is enabled.

Patch in fix for diffusers not honoring offline mode for sharded
checkpoints: huggingface/diffusers#12132

Author: Teriks

MANIFEST.in

@@ -6,4 +6,4 @@ recursive-include dgenerate/console/schemas *.json
 recursive-include dgenerate/console/recipes *.recipe
 recursive-include dgenerate/extras/hidiffusion/sd_module_key *.txt
 recursive-include dgenerate/translators/data *.json
-recursive-include dgenerate/pipelinewrapper/uris/text_encoder_configs *.json
+recursive-include dgenerate/pipelinewrapper/hub_configs *.json *.txt *.model

doctools/README.md renamed to assetgen/README.md

@@ -26,11 +26,12 @@
 from importlib.machinery import SourceFileLoader
 from pathlib import Path
 
-from doctools.core.rst_preprocessor import RSTPreprocessor
-from doctools.builders.readme_builder import ReadmeBuilder
-from doctools.builders.docs_builder import DocsBuilder
-from doctools.builders.console_schema_builder import ConsoleSchemaBuilder
-from doctools.builders.helsinki_nlp_translation_map_builder import HelsinkiNLPTranslationMapBuilder
+from assetgen.core.rst_preprocessor import RSTPreprocessor
+from assetgen.builders.readme_builder import ReadmeBuilder
+from assetgen.builders.docs_builder import DocsBuilder
+from assetgen.builders.console_schema_builder import ConsoleSchemaBuilder
+from assetgen.builders.helsinki_nlp_translation_map_builder import HelsinkiNLPTranslationMapBuilder
+from assetgen.builders.hf_configs_builder import HfConfigsBuilder
 
 
 def get_git_revision():
@@ -72,7 +73,7 @@ def setup_argument_parser():
 
     parser.add_argument(
         '--target', 
-        choices=['readme', 'docs', 'console-schemas', 'helsinki-nlp-translation-map', 'all'],
+        choices=['readme', 'docs', 'console-schemas', 'helsinki-nlp-translation-map', 'hf-configs', 'all'],
         default='all',
         help='What to build (default: all)'
     )
@@ -91,9 +92,9 @@ def setup_argument_parser():
 
     return parser
 
-
 def main():
     """Main entry point for the build script."""
+
     # Get project paths and info
     script_path = Path(__file__).absolute().parent
     project_dir = (script_path / '..').resolve()
@@ -119,7 +120,7 @@ def main():
     if disable_cache_completely:
         command_cache_path = None
     else:
-        command_cache_path = project_dir / 'doctools' / 'cache' / 'command.cache.json'
+        command_cache_path = project_dir / 'assetgen' / 'cache' / 'command.cache.json'
         command_cache_path.parent.mkdir(parents=True, exist_ok=True)
 
     # Create template builder - always verbose
@@ -158,6 +159,10 @@ def main():
     if args.target in ['helsinki-nlp-translation-map', 'all']:
         translation_map_builder = HelsinkiNLPTranslationMapBuilder(project_dir)
         translation_map_builder.build()
+    
+    if args.target in ['hf-configs', 'all']:
+        hf_configs_builder = HfConfigsBuilder(project_dir)
+        hf_configs_builder.build()
 
     print("Build completed successfully!")
 

doctools/__init__.py renamed to assetgen/__init__.py

@@ -40,7 +40,7 @@ def __init__(self, preprocessor: RSTPreprocessor, project_dir: Optional[Path] =
         """
         self.preprocessor = preprocessor
         self.project_dir = project_dir or Path.cwd()
-        self.template_dir = self.project_dir / 'doctools' / 'templates' / 'docs'
+        self.template_dir = self.project_dir / 'assetgen' / 'templates' / 'docs'
         self.output_dir = self.project_dir / 'docs'
 
     def build(self):

doctools/build.py renamed to assetgen/build.py

@@ -0,0 +1,147 @@
+#!/usr/bin/env python3
+# Copyright (c) 2023, Teriks
+#
+# dgenerate is distributed under the following BSD 3-Clause License
+#
+# Redistribution and use in source and binary forms, with or without modification, are permitted provided that the following conditions are met:
+#
+# 1. Redistributions of source code must retain the above copyright notice, this list of conditions and the following disclaimer.
+#
+# 2. Redistributions in binary form must reproduce the above copyright notice, this list of conditions and the following disclaimer in
+#    the documentation and/or other materials provided with the distribution.
+#
+# 3. Neither the name of the copyright holder nor the names of its contributors may be used to endorse or promote products derived
+#    from this software without specific prior written permission.
+#
+# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
+# LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
+# HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
+# LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON
+# ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+
+import shutil
+from pathlib import Path
+from typing import Optional
+
+try:
+    from huggingface_hub import snapshot_download
+except ImportError:
+    snapshot_download = None
+
+
+class HfConfigsBuilder:
+    """Builder for downloading Hugging Face model configuration files."""
+
+    def __init__(self, project_dir: Optional[Path] = None):
+        """
+        Initialize the HfConfigsBuilder.
+
+        :param project_dir: Project directory (defaults to current working directory)
+        :type project_dir: Optional[Path]
+        """
+        self.project_dir = project_dir or Path.cwd()
+        self.output_dir = self.project_dir / 'dgenerate' / 'pipelinewrapper' / 'hub_configs'
+        
+        # List of repositories to download from
+        self.repositories = [
+            "stable-diffusion-v1-5/stable-diffusion-v1-5",
+            "stabilityai/stable-diffusion-2-1", 
+            "stabilityai/stable-diffusion-xl-base-1.0",
+            "black-forest-labs/FLUX.1-dev",
+            "black-forest-labs/FLUX.1-Fill-dev", 
+            "black-forest-labs/FLUX.1-schnell",
+            "stabilityai/stable-diffusion-3-medium-diffusers",
+            "stabilityai/stable-diffusion-3.5-large",
+            "stabilityai/stable-diffusion-3.5-medium"
+        ]
+
+    def build(self):
+        """Download essential JSON configuration files from Hugging Face repositories."""
+        print("Downloading Hugging Face model configuration files...")
+
+
+        # Create the target directory if it doesn't exist
+        self.output_dir.mkdir(parents=True, exist_ok=True)
+        
+        print(f"Target directory: {self.output_dir}")
+        print(f"Number of repositories: {len(self.repositories)}")
+        print("-" * 60)
+
+        success_count = 0
+        failed_repos = []
+
+        # Download .json config files from each repository
+        for i, repo in enumerate(self.repositories, 1):
+            print(f"[{i}/{len(self.repositories)}] Processing: {repo}")
+            
+            try:
+                # Create a directory structure: models--org--name
+                repo_dir_name = f"models--{repo.replace('/', '--')}"
+                repo_path = self.output_dir / repo_dir_name
+                
+                # Download only essential .json config files
+                essential_patterns = [
+                    "model_index.json",           # Main model configuration
+                    "*/config.json",              # Component configs (unet, vae, text_encoder, etc.)
+                    "*/scheduler_config.json",    # Scheduler configuration
+                    "*/tokenizer_config.json",    # Tokenizer configuration
+                    "*/special_tokens_map.json",  # Tokenizer special tokens
+                    "*/vocab.json",               # Tokenizer vocabulary
+                    "*/tokenizer.json",           # Tokenizer data
+                    # Ensure BPE/SentencePiece tokenizers are complete
+                    "merges.txt",
+                    "*/merges.txt",
+                    "**/merges.txt",
+                    "spiece.model",
+                    "*/spiece.model",
+                    "**/spiece.model",
+                    "*/preprocessor_config.json", # Feature extractor config
+                    "*.safetensors.index.json",   # Model weight indices
+                    "*.safetensors.index.fp16.json"  # FP16 model weight indices
+                ]
+                
+                snapshot_download(
+                    repo_id=repo,
+                    allow_patterns=essential_patterns,
+                    local_dir=repo_path,
+                    local_dir_use_symlinks=False  # Use actual files instead of symlinks
+                )
+                
+                # Remove any .cache directories that might have been created
+                cache_dirs = list(repo_path.rglob(".cache"))
+                for cache_dir in cache_dirs:
+                    if cache_dir.is_dir():
+                        shutil.rmtree(cache_dir)
+                        print(f"  ✓ Removed checkpoint .cache directory")
+                
+                # Count downloaded JSON files
+                json_files = list(repo_path.rglob("*.json"))
+                print(f"  ✓ Downloaded {len(json_files)} JSON files")
+                success_count += 1
+                
+            except Exception as e:
+                print(f"  ✗ Failed to download from {repo}: {str(e)}")
+                failed_repos.append(repo)
+                continue
+            
+            print()
+
+        # Summary
+        print("=" * 60)
+        print(f"Download Summary:")
+        print(f"  Successfully processed: {success_count}/{len(self.repositories)} repositories")
+        
+        if failed_repos:
+            print(f"  Failed repositories:")
+            for repo in failed_repos:
+                print(f"    - {repo}")
+            print("✗ Some repositories failed to download")
+        else:
+            print(f"  ✓ All repositories downloaded successfully!")
+        
+        print(f"\nJSON config files are now available in: {self.output_dir}")
+
+    def get_output_path(self) -> Path:
+        """Get the path to the hub configs output directory."""
+        return self.output_dir

doctools/builders/__init__.py renamed to assetgen/builders/__init__.py

@@ -40,7 +40,7 @@ def __init__(self, preprocessor: RSTPreprocessor, project_dir: Optional[Path] =
         """
         self.preprocessor = preprocessor
         self.project_dir = project_dir or Path.cwd()
-        self.template_dir = self.project_dir / 'doctools' / 'templates' / 'readme'
+        self.template_dir = self.project_dir / 'assetgen' / 'templates' / 'readme'
         self.output_file = self.project_dir / 'README.rst'
 
     def build(self):