diff --git a/apdepth/util/checkpoint.py b/apdepth/util/checkpoint.py new file mode 100644 index 0000000..b08ceaf --- /dev/null +++ b/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 diff --git a/infer.py b/infer.py index adf9f9f..e1c5a32 100644 --- a/infer.py +++ b/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() diff --git a/run.py b/run.py index 8c359e3..a0fff87 100755 --- a/run.py +++ b/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 diff --git a/tests/test_training_checkpoint.py b/tests/test_training_checkpoint.py new file mode 100644 index 0000000..e1b751a --- /dev/null +++ b/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)