Spaces:
Runtime error
Runtime error
File size: 1,714 Bytes
e3d6629 5d309f7 9d22a21 4ed5b73 5d309f7 e3d6629 9d22a21 ed72314 9d22a21 5d309f7 ed72314 5d309f7 ed72314 6018992 ed72314 6018992 ed72314 e3d6629 5d309f7 e3d6629 5d309f7 ed72314 5d309f7 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 | import torch
import logging
import os
from pathlib import Path
from dust3r.model import AsymmetricCroCo3DStereo, inf
logger = logging.getLogger(__name__)
def download_weights(model_path: str):
"""Download model weights if they don't exist"""
if os.path.exists(model_path):
logger.info(f"Weights already exist at {model_path}")
return
logger.info("Weights not found. Downloading...")
Path(model_path).parent.mkdir(parents=True, exist_ok=True)
import urllib.request
url = "https://huggingface.co/camenduru/dust3r/resolve/main/DUSt3R_ViTLarge_BaseDecoder_512_dpt.pth"
logger.info(f"Downloading from {url}")
urllib.request.urlretrieve(url, model_path)
logger.info("Download complete!")
def initialize(model_path: str, device: str) -> torch.nn.Module:
download_weights(model_path)
logger.info(f"Loading model from: {model_path}")
logger.info("Loading checkpoint...")
ckpt = torch.load(model_path, map_location='cpu', weights_only=False)
logger.info("Parsing model arguments...")
args = ckpt['args'].model.replace("ManyAR_PatchEmbed", "PatchEmbedDust3R")
if isinstance(args, str) and 'landscape_only' not in args:
args = args[:-1] + ', landscape_only=False)'
elif isinstance(args, str):
args = args.replace(" ", "").replace('landscape_only=True', 'landscape_only=False')
logger.info("Instantiating model...")
net = eval(args)
logger.info("Loading model weights...")
net.load_state_dict(ckpt['model'], strict=False)
logger.info(f"Moving model to {device}...")
model = net.to(device)
logger.info("Model initialization complete!")
return model
|