Spaces:
Running on Zero
Running on Zero
| """GeoNeXt-Wan inference adapter.""" | |
| import subprocess | |
| import shutil | |
| import sys | |
| import tempfile | |
| from pathlib import Path | |
| import numpy as np | |
| from PIL import Image | |
| from utils.common import list_images | |
| GEONEXT_REPO = "happy0612/GeoNeXt" | |
| GEONEXT_WAN_FILE = "GeoNeXt-Wan/geonext_wan.safetensors" | |
| WAN_BASE_MODEL = "Wan-AI/Wan2.1-T2V-1.3B" | |
| def _resolve_models(args): | |
| from huggingface_hub import hf_hub_download, snapshot_download | |
| checkpoint = args.checkpoint | |
| if checkpoint is None: | |
| print("Downloading GeoNeXt-Wan checkpoint from %s" % GEONEXT_REPO) | |
| checkpoint = hf_hub_download(repo_id=GEONEXT_REPO, filename=GEONEXT_WAN_FILE) | |
| base_model = args.base_model or WAN_BASE_MODEL | |
| if not Path(base_model).exists(): | |
| print("Downloading Wan base model from %s" % base_model) | |
| base_model = snapshot_download(repo_id=base_model) | |
| return checkpoint, base_model | |
| def run(args): | |
| from utils.geometry import Camera, disparity_to_relative_depth, export_geometry | |
| aligner = None | |
| if args.align_space == "moge": | |
| from utils.moge import MoGeAligner | |
| aligner = MoGeAligner(args.moge_model) | |
| root, images = list_images(args.input) | |
| checkpoint, base_model = _resolve_models(args) | |
| processing_res = 768 if args.processing_res is None else args.processing_res | |
| processing_side = args.processing_res_side or "long" | |
| temporary_input = None | |
| input_dir = root | |
| if Path(args.input).is_file(): | |
| # The underlying Wan batch runner consumes a directory. Stage a single | |
| # input transparently so the public CLI accepts the same inputs as SVD. | |
| temporary_input = tempfile.TemporaryDirectory(prefix="geonext-wan-input-") | |
| input_dir = Path(temporary_input.name) | |
| shutil.copy2(args.input, input_dir / Path(args.input).name) | |
| try: | |
| command = [ | |
| sys.executable, str(Path(__file__).with_name("model_inference.py")), | |
| "--checkpoint", checkpoint, | |
| "--wan_model_dir", base_model, | |
| "--input_dir", str(input_dir), | |
| "--output_dir", args.output, | |
| "--target_modalities", "depth,normal", | |
| "--rgb_condition_mode", "concat", | |
| "--num_inference_steps", str(args.steps), | |
| "--cfg_scale", "1.0", | |
| "--processing_res", str(processing_res), | |
| "--processing_res_side", processing_side, | |
| "--temporal_rope_scale", "8", | |
| "--seed", str(args.seed), | |
| ] | |
| subprocess.run(command, check=True) | |
| finally: | |
| if temporary_input is not None: | |
| temporary_input.cleanup() | |
| if not args.export: | |
| return | |
| output = Path(args.output) | |
| for image_path in images: | |
| stem = image_path.stem | |
| depth_path = output / "depth_raw" / (stem + ".npy") | |
| normal_raw_path = output / "normal_raw" / (stem + ".npy") | |
| normal_path = output / "normal_vis" / (stem + ".png") | |
| if not depth_path.exists() or (not normal_raw_path.exists() and not normal_path.exists()): | |
| print("warning: missing prediction for %s; geometry skipped" % image_path) | |
| continue | |
| disparity = np.load(depth_path) | |
| if normal_raw_path.exists(): | |
| normal = np.load(normal_raw_path).astype(np.float32) | |
| else: | |
| normal = np.asarray(Image.open(normal_path).convert("RGB"), dtype=np.float32) / 127.5 - 1.0 | |
| rgb = Image.open(image_path).convert("RGB") | |
| mask = None | |
| if aligner is not None: | |
| depth, mask, camera, intrinsics, scale, shift = aligner.align_disparity(disparity, rgb) | |
| if not args.use_mask: | |
| mask = None | |
| else: | |
| depth = disparity_to_relative_depth(disparity) | |
| camera = Camera.from_fov_x(depth.shape[1], depth.shape[0], args.fov_x) | |
| case = output / "geometry" / stem | |
| case.mkdir(parents=True, exist_ok=True) | |
| np.save(case / "depth.npy", depth) | |
| if aligner is not None: | |
| np.save(case / "intrinsics.npy", intrinsics) | |
| export_geometry(case, rgb, depth, camera, | |
| normal=normal, mask=mask, exports=args.export, stride=args.stride, | |
| edge_threshold=args.edge_threshold, | |
| normal_angle_threshold=args.normal_angle_threshold, | |
| normal_bypass_depth_threshold=args.normal_bypass_depth_threshold) | |
| print("geometry: %s" % case) | |