feat(inference): load Stage-2 training components

apdepth/util/checkpoint.py

@@ -0,0 +1,40 @@
+"""Load Stage-2 trainables on top of the original SD2 pipeline components."""
+
+import json
+from pathlib import Path
+
+from diffusers import AutoencoderKL, UNet2DConditionModel
+
+
+def load_training_components(checkpoint, torch_dtype):
+    """Return from_pretrained overrides, including the learned VAE when required.
+
+    Training checkpoints store full U-Net/VAE components, but reuse the base
+    SD2 tokenizer, text encoder and scheduler. Never silently discard a VAE.
+    """
+    root = Path(checkpoint)
+    metadata_path = root / "apdepth_training.json"
+    metadata = (
+        json.loads(metadata_path.read_text(encoding="utf-8"))
+        if metadata_path.is_file()
+        else {}
+    )
+    vae_path = root / "vae"
+    if metadata.get("has_finetuned_vae", False) and not vae_path.is_dir():
+        raise FileNotFoundError(f"Checkpoint requires its fine-tuned VAE: {vae_path}")
+    result = {
+        "unet": UNet2DConditionModel.from_pretrained(
+            root / "unet",
+            torch_dtype=torch_dtype,
+            use_safetensors=True,
+            local_files_only=True,
+        )
+    }
+    if vae_path.is_dir():
+        result["vae"] = AutoencoderKL.from_pretrained(
+            vae_path,
+            torch_dtype=torch_dtype,
+            use_safetensors=True,
+            local_files_only=True,
+        )
+    return result

infer.py

@@ -31,6 +31,7 @@
 from tqdm.auto import tqdm
 
 from apdepth import ApDepthPipeline
+from apdepth.util.checkpoint import load_training_components
 from src.util.seeding import seed_all
 from src.dataset import (
     BaseDepthDataset,
@@ -109,6 +110,14 @@ def save_prediction_atomic(save_to, depth_pred):
     )
 
     # dataset setting
+    parser.add_argument(
+        "--training_checkpoint",
+        default=None,
+        help=(
+            "Stage-2 checkpoint to load over --checkpoint (original SD2 base); "
+            "includes saved VAE."
+        ),
+    )
     parser.add_argument(
         "--dataset_config",
         type=str,
@@ -293,9 +302,14 @@ def check_directory(directory):
         dtype = torch.float32
         variant = None
 
+    components = (
+        load_training_components(args.training_checkpoint, dtype)
+        if args.training_checkpoint
+        else {}
+    )
     pipe = ApDepthPipeline.from_pretrained(
-        checkpoint_path, variant=variant, torch_dtype=dtype
-    )   
+        checkpoint_path, variant=variant, torch_dtype=dtype, **components
+    )
 
     try:
         pipe.enable_xformers_memory_efficient_attention()

run.py

@@ -25,17 +25,15 @@
 
 import numpy as np
 import torch
-import torch.nn as nn
-import torch.nn.functional as F
 from PIL import Image
 from tqdm.auto import tqdm
+from torchvision import transforms
 
 from apdepth import ApDepthPipeline
+from apdepth.util.checkpoint import load_training_components
 
 EXTENSION_LIST = [".jpg", ".jpeg", ".png"]
 
-from torchvision import transforms
-
 
 if "__main__" == __name__:
     logging.basicConfig(level=logging.INFO)
@@ -57,6 +55,14 @@
         required=True,
         help="Path to the input image folder.",
     )
+    parser.add_argument(
+        "--training_checkpoint",
+        default=None,
+        help=(
+            "Stage-2 checkpoint to load over --checkpoint (original SD2 base); "
+            "includes saved VAE."
+        ),
+    )
 
     parser.add_argument(
         "--output_dir", type=str, required=True, help="Output directory."
@@ -135,6 +141,8 @@
     half_precision = args.half_precision
 
     processing_res = args.processing_res
+    if args.training_checkpoint and processing_res is None:
+        processing_res = 0
     match_input_res = not args.output_processing_res
     if 0 == processing_res and match_input_res is False:
         logging.warning(
@@ -204,8 +212,13 @@
         dtype = torch.float32
         variant = None
 
+    components = (
+        load_training_components(args.training_checkpoint, dtype)
+        if args.training_checkpoint
+        else {}
+    )
     pipe: ApDepthPipeline = ApDepthPipeline.from_pretrained(
-        checkpoint_path, variant=variant, torch_dtype=dtype
+        checkpoint_path, variant=variant, torch_dtype=dtype, **components
     )
     # unet = UNet2DConditionModel.from_pretrained(os.path.join(checkpoint_path, f'unet'))
     # pipe.unet = unet

tests/test_training_checkpoint.py

@@ -0,0 +1,71 @@
+"""Tests for composing saved training components with the original SD2 base."""
+
+import json
+
+import pytest
+import torch
+from diffusers import AutoencoderKL, UNet2DConditionModel
+
+from apdepth.util.checkpoint import load_training_components
+
+
+def tiny_unet():
+    return UNet2DConditionModel(
+        sample_size=4,
+        in_channels=8,
+        out_channels=4,
+        layers_per_block=1,
+        block_out_channels=(16,),
+        down_block_types=("DownBlock2D",),
+        up_block_types=("UpBlock2D",),
+        cross_attention_dim=16,
+        norm_num_groups=8,
+    )
+
+
+def tiny_vae():
+    return AutoencoderKL(
+        in_channels=3,
+        out_channels=3,
+        latent_channels=4,
+        block_out_channels=(16, 16, 16, 16),
+        down_block_types=("DownEncoderBlock2D",) * 4,
+        up_block_types=("UpDecoderBlock2D",) * 4,
+        norm_num_groups=8,
+    )
+
+
+def test_loads_saved_unet_and_vae(tmp_path):
+    unet = tiny_unet()
+    vae = tiny_vae()
+    unet.save_pretrained(tmp_path / "unet", safe_serialization=True)
+    vae.save_pretrained(tmp_path / "vae", safe_serialization=True)
+    (tmp_path / "apdepth_training.json").write_text(
+        json.dumps({"has_finetuned_vae": True}), encoding="utf-8"
+    )
+
+    components = load_training_components(tmp_path, torch.float32)
+    assert set(components) == {"unet", "vae"}
+    torch.testing.assert_close(
+        components["unet"].conv_in.weight, unet.conv_in.weight, rtol=0, atol=0
+    )
+    torch.testing.assert_close(
+        components["vae"].decoder.conv_out.weight,
+        vae.decoder.conv_out.weight,
+        rtol=0,
+        atol=0,
+    )
+
+
+def test_loads_legacy_unet_only_checkpoint(tmp_path):
+    tiny_unet().save_pretrained(tmp_path / "unet", safe_serialization=True)
+    assert set(load_training_components(tmp_path, torch.float32)) == {"unet"}
+
+
+def test_rejects_missing_required_vae(tmp_path):
+    tiny_unet().save_pretrained(tmp_path / "unet", safe_serialization=True)
+    (tmp_path / "apdepth_training.json").write_text(
+        json.dumps({"has_finetuned_vae": True}), encoding="utf-8"
+    )
+    with pytest.raises(FileNotFoundError, match="fine-tuned VAE"):
+        load_training_components(tmp_path, torch.float32)