File size: 2,581 Bytes
f284828 | 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 89 90 91 92 | import os
import sys
import torch
import torch.nn as nn
import numpy as np
import inspect
# --- PROACTIVE PYTHON 3.12 / NUMPY 2.0 PATCH ---
if not hasattr(np, 'bool'):
np.bool = np.bool_
if not hasattr(np, 'float'):
np.float = np.float64
if not hasattr(np, 'int'):
np.int = np.int_
if not hasattr(np, 'complex'):
np.complex = complex
if not hasattr(np, 'object'):
np.object = object
if not hasattr(np, 'unicode'):
np.unicode = str
if not hasattr(np, 'str'):
np.str = str
if not hasattr(inspect, 'getargspec'):
inspect.getargspec = inspect.getfullargspec
# -----------------------------------------------
# Path setup based on your layout
# Download from https://github.com/yfeng95/DECA
DECA_DIR = os.path.abspath("DECA")
sys.path.insert(0, DECA_DIR)
from decalib.deca import DECA
from decalib.utils.config import cfg as deca_cfg
# --- THE ULTIMATE PYTORCH3D BYPASS ---
DECA._setup_renderer = lambda self, *args, **kwargs: None
# -------------------------------------
# ---------------------------------------------------------
# The Custom Rust Wrapper
# ---------------------------------------------------------
class DecaVisionONNXWrapper(nn.Module):
def __init__(self, deca_model):
super().__init__()
# Extract just the ResNet50 vision encoder from DECA
self.encoder = deca_model.E_flame
def forward(self, images):
parameters = self.encoder(images)
flame_vector = parameters[:, 0:150]
return flame_vector
# ---------------------------------------------------------
def main():
device = 'cuda' if torch.cuda.is_available() else 'cpu'
print(f"Loading DECA on {device.upper()}...")
# Initialize standard DECA
deca_cfg.model.use_tex = False
deca_cfg.rasterizer_type = 'pytorch3d'
deca = DECA(config=deca_cfg, device=device)
deca.eval()
print("Wrapping model for seamless 150-dim output...")
wrapper = DecaVisionONNXWrapper(deca).to(device)
wrapper.eval()
print("Tracing and exporting to ONNX...")
dummy_input = torch.randn(1, 3, 224, 224).to(device)
torch.onnx.export(
wrapper,
dummy_input,
"deca_vision.onnx",
export_params=True,
opset_version=17,
do_constant_folding=True,
input_names=['image_input'],
output_names=['flame_vector'],
dynamic_axes={
'image_input': {0: 'batch_size'},
'flame_vector': {0: 'batch_size'}
}
)
print("\nSuccess! 'deca_vision.onnx' is fully baked.")
if __name__ == "__main__":
main()
|