| """Explicitly load the v2 research preview and super-resolve a saved RGBN reflectance tensor.""" |
| import argparse |
| from pathlib import Path |
|
|
| import torch |
| from omegaconf import OmegaConf |
| from opensr_model import SRLatentDiffusion |
|
|
|
|
| def load_model(device=None): |
| root = Path(__file__).resolve().parent |
| config = OmegaConf.load(root / "inference.yaml") |
| device = torch.device(device or ("cuda" if torch.cuda.is_available() else "cpu")) |
| model = SRLatentDiffusion(config, device=device) |
| model.load_pretrained(str(root.parent / "opensr-ldsrs2_v2_0_0_preview.ckpt")) |
| model.eval() |
| return model, config |
|
|
|
|
| @torch.inference_mode() |
| def super_resolve(model, config, rgbn, *, seed=42, sampler="dpm_solver++", steps=None): |
| if not isinstance(rgbn, torch.Tensor): |
| raise TypeError("Input must be a torch.Tensor") |
| if rgbn.ndim not in (3, 4) or rgbn.shape[-3] != 4: |
| raise ValueError("Expected [4,H,W] or [B,4,H,W] in R,G,B,NIR order") |
| if not rgbn.is_floating_point() or not torch.isfinite(rgbn).all(): |
| raise ValueError("Input must contain finite floating-point reflectance") |
| torch.manual_seed(seed) |
| if torch.cuda.is_available(): |
| torch.cuda.manual_seed_all(seed) |
| steps = steps if steps is not None else (3 if sampler == "dpm_solver++" else 100) |
| return model( |
| rgbn.to(device=model.device, dtype=torch.float32), |
| sampler=sampler, |
| sampling_steps=steps, |
| sampling_eta=0.0 if sampler == "dpm_solver++" else 0.95, |
| spectral_correction="back_projection", |
| back_projection=config.back_projection, |
| ) |
|
|
|
|
| def main(): |
| parser = argparse.ArgumentParser(description=__doc__) |
| source = parser.add_mutually_exclusive_group(required=True) |
| source.add_argument("--input", type=Path, help="A saved float RGBN tensor (.pt)") |
| source.add_argument("--smoke", action="store_true", help="Use synthetic input") |
| parser.add_argument("--output", type=Path, required=True) |
| parser.add_argument("--device", default=None) |
| parser.add_argument("--seed", type=int, default=42) |
| parser.add_argument("--sampler", choices=["dpm_solver++", "ddim"], default="dpm_solver++") |
| parser.add_argument("--steps", type=int) |
| args = parser.parse_args() |
| torch.manual_seed(args.seed) |
| rgbn = torch.rand(1, 4, 128, 128) if args.smoke else torch.load( |
| args.input, map_location="cpu", weights_only=True |
| ) |
| model, config = load_model(args.device) |
| sr = super_resolve(model, config, rgbn, seed=args.seed, sampler=args.sampler, steps=args.steps) |
| if not torch.isfinite(sr).all(): |
| raise RuntimeError("Non-finite model output") |
| args.output.parent.mkdir(parents=True, exist_ok=True) |
| torch.save(sr.cpu(), args.output) |
| print(f"Saved {tuple(sr.shape)} to {args.output}") |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
