File size: 2,693 Bytes
31f6f71 | 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 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 | import argparse
import json
import os
from pathlib import Path
import numpy as np
import torch
from PIL import Image
os.environ.setdefault("SPCONV_ALGO", "native")
os.environ.setdefault("ATTN_BACKEND", "flash_attn")
os.environ.setdefault("TOKENIZERS_PARALLELISM", "false")
os.environ.setdefault("DVD_MODEL_REPO", "Zhengrui/dvd")
def cfg_schedule(mode: str, constant: float, early: float, late: float, split: float):
if mode == "Constant":
return float(constant)
if mode == "Two-stage":
split = float(split)
early = float(early)
late = float(late)
return lambda t: early if t < split else late
return None
def run_generate(config_path: str):
with open(config_path, "r") as f:
cfg = json.load(f)
from dvd import DVDImageToVoxelPipeline, export_cubified_voxels
repo = os.environ.get("DVD_MODEL_REPO", "Zhengrui/dvd")
subfolder = os.environ.get("DVD_MODEL_SUBFOLDER") or None
revision = os.environ.get("DVD_MODEL_REVISION") or None
token = os.environ.get("DVD_MODEL_TOKEN") or os.environ.get("HF_TOKEN") or None
print(f"[DVD Worker] loading DVD image pipeline from {repo}", flush=True)
pipeline = DVDImageToVoxelPipeline.from_pretrained(
repo,
variant="base",
device="cuda",
subfolder=subfolder,
revision=revision,
token=token,
)
print("[DVD Worker] pipeline ready", flush=True)
image = Image.open(cfg["image_path"]).convert("RGBA")
sampler_kwargs = {"steps": int(cfg["dvd_steps"])}
schedule = cfg_schedule(
cfg["dvd_cfg_mode"],
float(cfg["dvd_cfg_constant"]),
float(cfg["dvd_cfg_early"]),
float(cfg["dvd_cfg_late"]),
float(cfg["dvd_cfg_split"]),
)
if schedule is not None:
sampler_kwargs["cfg_strength"] = schedule
print(f"[DVD Worker] sampling seed={cfg['seed']} steps={cfg['dvd_steps']}", flush=True)
voxels = pipeline.sample_voxels(
image,
seed=int(cfg["seed"]),
preprocess_image=bool(cfg["preprocess_image"]),
**sampler_kwargs,
)
mesh_path = cfg["mesh_path"]
npy_path = cfg["npy_path"]
export_cubified_voxels(voxels, mesh_path)
np.save(npy_path, voxels.coords_without_batch.detach().cpu().numpy().astype(np.int32))
torch.cuda.empty_cache()
print(f"[DVD Worker] done mesh={mesh_path} npy={npy_path}", flush=True)
def main():
parser = argparse.ArgumentParser()
parser.add_argument("command", choices=["generate"])
parser.add_argument("config")
args = parser.parse_args()
if args.command == "generate":
run_generate(args.config)
if __name__ == "__main__":
main()
|